diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/ArbitraryDistributionSplitAssigner.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/ArbitraryDistributionSplitAssigner.java
index 4dfed047bdf3..83469e8472cd 100644
--- a/core/trino-main/src/main/java/io/trino/execution/scheduler/ArbitraryDistributionSplitAssigner.java
+++ b/core/trino-main/src/main/java/io/trino/execution/scheduler/ArbitraryDistributionSplitAssigner.java
@@ -143,6 +143,13 @@ private AssignmentResult assignReplicatedSplits(PlanNodeId planNodeId, List<Spli
                             replicatedSplits.get(replicatedSourceId),
                             true));
                 }
+                for (PlanNodeId partitionedSourceId : partitionedSources) {
+                    assignment.updatePartition(new PartitionUpdate(
+                            0,
+                            partitionedSourceId,
+                            ImmutableList.of(),
+                            true));
+                }
                 assignment.sealPartition(0);
             }
             else {
@@ -232,6 +239,14 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
                             replicatedSplits.get(replicatedSourceId),
                             true));
                 }
+                for (PlanNodeId partitionedSourceId : partitionedSources) {
+                    assignment.updatePartition(new PartitionUpdate(
+                            0,
+                            partitionedSourceId,
+                            ImmutableList.of(),
+                            true));
+                }
+
                 assignment.sealPartition(0);
             }
             else {
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenTaskSource.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenTaskSource.java
index d6ffa62e943f..f782a0e06dbb 100644
--- a/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenTaskSource.java
+++ b/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenTaskSource.java
@@ -273,7 +273,9 @@ private int getSplitPartition(Split split)
 
     private void fail(Throwable failure)
     {
-        callback.failed(failure);
+        synchronized (assignerLock) {
+            callback.failed(failure);
+        }
         close();
     }
 
diff --git a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestArbitraryDistributionSplitAssigner.java b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestArbitraryDistributionSplitAssigner.java
index dc03725e1cad..9575c7edd5f9 100644
--- a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestArbitraryDistributionSplitAssigner.java
+++ b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestArbitraryDistributionSplitAssigner.java
@@ -47,6 +47,7 @@
 import static java.util.Collections.shuffle;
 import static java.util.Objects.requireNonNull;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
@@ -73,6 +74,7 @@ public void testEmpty()
         SplitAssigner splitAssigner = createSplitAssigner(ImmutableSet.of(PARTITIONED_1), ImmutableSet.of(), 100, false);
         TestingTaskSourceCallback callback = new TestingTaskSourceCallback();
         splitAssigner.assign(PARTITIONED_1, ImmutableListMultimap.of(), true).update(callback);
+        assertTrue(callback.isNoMoreSplits(0, PARTITIONED_1));
         splitAssigner.finish().update(callback);
         List<TaskDescriptor> taskDescriptors = callback.getTaskDescriptors();
         assertThat(taskDescriptors).hasSize(1);
@@ -82,6 +84,7 @@ public void testEmpty()
         splitAssigner = createSplitAssigner(ImmutableSet.of(), ImmutableSet.of(REPLICATED_1), 100, false);
         callback = new TestingTaskSourceCallback();
         splitAssigner.assign(REPLICATED_1, ImmutableListMultimap.of(), true).update(callback);
+        assertTrue(callback.isNoMoreSplits(0, REPLICATED_1));
         splitAssigner.finish().update(callback);
         taskDescriptors = callback.getTaskDescriptors();
         assertThat(taskDescriptors).hasSize(1);
@@ -91,7 +94,11 @@ public void testEmpty()
         splitAssigner = createSplitAssigner(ImmutableSet.of(PARTITIONED_1), ImmutableSet.of(REPLICATED_1), 100, true);
         callback = new TestingTaskSourceCallback();
         splitAssigner.assign(REPLICATED_1, ImmutableListMultimap.of(), true).update(callback);
+        assertFalse(callback.isNoMoreSplits(0, PARTITIONED_1));
+        assertFalse(callback.isNoMoreSplits(0, REPLICATED_1));
         splitAssigner.assign(PARTITIONED_1, ImmutableListMultimap.of(), true).update(callback);
+        assertTrue(callback.isNoMoreSplits(0, PARTITIONED_1));
+        assertTrue(callback.isNoMoreSplits(0, REPLICATED_1));
         splitAssigner.finish().update(callback);
         taskDescriptors = callback.getTaskDescriptors();
         assertThat(taskDescriptors).hasSize(1);
@@ -100,7 +107,11 @@ public void testEmpty()
         splitAssigner = createSplitAssigner(ImmutableSet.of(PARTITIONED_1), ImmutableSet.of(REPLICATED_1), 100, true);
         callback = new TestingTaskSourceCallback();
         splitAssigner.assign(PARTITIONED_1, ImmutableListMultimap.of(), true).update(callback);
+        assertFalse(callback.isNoMoreSplits(0, PARTITIONED_1));
+        assertFalse(callback.isNoMoreSplits(0, REPLICATED_1));
         splitAssigner.assign(REPLICATED_1, ImmutableListMultimap.of(), true).update(callback);
+        assertTrue(callback.isNoMoreSplits(0, PARTITIONED_1));
+        assertTrue(callback.isNoMoreSplits(0, REPLICATED_1));
         splitAssigner.finish().update(callback);
         taskDescriptors = callback.getTaskDescriptors();
         assertThat(taskDescriptors).hasSize(1);
@@ -111,7 +122,16 @@ public void testEmpty()
         splitAssigner.assign(REPLICATED_1, ImmutableListMultimap.of(), true).update(callback);
         splitAssigner.assign(PARTITIONED_1, ImmutableListMultimap.of(), true).update(callback);
         splitAssigner.assign(PARTITIONED_2, ImmutableListMultimap.of(), true).update(callback);
+        assertFalse(callback.isNoMoreSplits(0, PARTITIONED_1));
+        assertFalse(callback.isNoMoreSplits(0, REPLICATED_1));
+        assertFalse(callback.isNoMoreSplits(0, PARTITIONED_2));
+        assertFalse(callback.isNoMoreSplits(0, REPLICATED_2));
         splitAssigner.assign(REPLICATED_2, ImmutableListMultimap.of(), true).update(callback);
+        assertTrue(callback.isNoMoreSplits(0, PARTITIONED_1));
+        assertTrue(callback.isNoMoreSplits(0, REPLICATED_1));
+        assertTrue(callback.isNoMoreSplits(0, PARTITIONED_2));
+        assertTrue(callback.isNoMoreSplits(0, REPLICATED_2));
+
         splitAssigner.finish().update(callback);
         taskDescriptors = callback.getTaskDescriptors();
         assertThat(taskDescriptors).hasSize(1);