Return empty list from TaskSource after finish

trinodb · May 21, 2022 · cb18890 · cb18890
1 parent 3b9e47d
commit cb18890
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/StageTaskSourceFactory.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/StageTaskSourceFactory.java
@@ -214,6 +214,9 @@ public SingleDistributionTaskSource(ListMultimap<PlanNodeId, ExchangeSourceHandl
         @Override
         public ListenableFuture<List<TaskDescriptor>> getMoreTasks()
         {
+            if (finished) {
+                return immediateFuture(ImmutableList.of());
+            }
             List<TaskDescriptor> result = ImmutableList.of(new TaskDescriptor(
                     0,
                     ImmutableListMultimap.of(),
@@ -291,6 +294,9 @@ public ArbitraryDistributionTaskSource(
         @Override
         public ListenableFuture<List<TaskDescriptor>> getMoreTasks()
         {
+            if (finished) {
+                return immediateFuture(ImmutableList.of());
+            }
             NodeRequirements nodeRequirements = new NodeRequirements(Optional.empty(), ImmutableSet.of(), taskMemory);
 
             ImmutableList.Builder<TaskDescriptor> result = ImmutableList.builder();