diff --git a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java index dd5792f52..d01af7425 100644 --- a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java +++ b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java @@ -47,7 +47,6 @@ import org.apache.wayang.core.platform.PartialExecution; import org.apache.wayang.core.platform.Platform; import org.apache.wayang.core.platform.lineage.ExecutionLineageNode; -import org.apache.wayang.core.util.WayangCollections; import org.apache.wayang.jdbc.channels.SqlQueryChannel; import org.apache.wayang.jdbc.compiler.FunctionCompiler; import org.apache.wayang.jdbc.operators.JdbcExecutionOperator; @@ -338,16 +337,20 @@ private Collection createExecutionLineageNodes( * unless it is not comprising a {@link JdbcExecutionOperator} and/or not in the * given {@link ExecutionStage}. * - * @param task whose follow-up {@link ExecutionTask} is requested; should have - * a single follower + * @param task whose follow-up {@link ExecutionTask} is requested * @param stage in which the follow-up {@link ExecutionTask} should be * @return the said follow-up {@link ExecutionTask} or {@code null} if none */ private static ExecutionTask findJdbcExecutionOperatorTaskInStage(final ExecutionTask task, final ExecutionStage stage) { - assert task.getNumOuputChannels() == 1; + if (task.getNumOuputChannels() != 1) { + return null; + } final Channel outputChannel = task.getOutputChannel(0); - final ExecutionTask consumer = WayangCollections.getSingle(outputChannel.getConsumers()); + if (outputChannel == null || outputChannel.getConsumers().size() != 1) { + return null; + } + final ExecutionTask consumer = outputChannel.getConsumers().iterator().next(); return consumer.getStage() == stage && consumer.getOperator() instanceof JdbcExecutionOperator ? consumer : null; diff --git a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java index 8f7b3d8a2..ca73e7663 100644 --- a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java +++ b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java @@ -40,6 +40,7 @@ import java.sql.SQLException; import java.util.Collections; +import java.util.Set; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.Mockito.mock; @@ -244,4 +245,62 @@ void testExecuteWithProjectionAndFilters() throws SQLException { sqlQueryChannelInstance.getSqlQuery() ); } + + @Test + void testExecuteWithNoTerminalConsumers() { + assertQueryWithTerminalConsumers(0); + } + + @Test + void testExecuteWithMultipleTerminalConsumers() { + assertQueryWithTerminalConsumers(2); + } + + private void assertQueryWithTerminalConsumers(int consumerCount) { + Configuration configuration = new Configuration(); + Job job = mock(Job.class); + when(job.getConfiguration()).thenReturn(configuration); + CrossPlatformExecutor executionState = new CrossPlatformExecutor(job, new NoInstrumentationStrategy()); + when(job.getCrossPlatformExecutor()).thenReturn(executionState); + SqlQueryChannel.Descriptor descriptor = HsqldbPlatform.getInstance().getSqlQueryChannelDescriptor(); + ExecutionStage stage = mock(ExecutionStage.class); + + JdbcTableSource source = new HsqldbTableSource("customer"); + ExecutionTask sourceTask = new ExecutionTask(source); + sourceTask.setStage(stage); + sourceTask.setOutputChannel(0, new SqlQueryChannel(descriptor, source.getOutput(0))); + + JdbcFilterOperator filter = new HsqldbFilterOperator(new PredicateDescriptor( + record -> true, Record.class).withSqlImplementation("age >= 18")); + ExecutionTask filterTask = new ExecutionTask(filter); + filterTask.setStage(stage); + sourceTask.getOutputChannel(0).addConsumer(filterTask, 0); + // A reusable channel lets the stage boundary have multiple downstream consumers. + SqlQueryChannel output = new SqlQueryChannel(descriptor, filter.getOutput(0)) { + @Override + public boolean isReusable() { + return true; + } + }; + filterTask.setOutputChannel(0, output); + when(stage.getStartTasks()).thenReturn(Collections.singleton(sourceTask)); + when(stage.getTerminalTasks()).thenReturn(Collections.singleton(filterTask)); + when(stage.getAllTasks()).thenReturn(Set.of(sourceTask, filterTask)); + + ExecutionStage nextStage = mock(ExecutionStage.class); + for (int i = 0; i < consumerCount; i++) { + ExecutionTask consumer = new ExecutionTask(new SqlToStreamOperator(HsqldbPlatform.getInstance())); + consumer.setStage(nextStage); + output.addConsumer(consumer, 0); + } + + JdbcExecutor executor = new JdbcExecutor(HsqldbPlatform.getInstance(), job); + try { + executor.execute(stage, new DefaultOptimizationContext(job), executionState); + SqlQueryChannel.Instance result = (SqlQueryChannel.Instance) executionState.getChannelInstance(output); + assertEquals("SELECT * FROM customer WHERE age >= 18", result.getSqlQuery()); + } finally { + executor.dispose(); + } + } }