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 d34fd273d550..6a767e7ff354 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
@@ -128,6 +128,7 @@ private AssignmentResult assignReplicatedSplits(PlanNodeId planNodeId, List<Spli
             assignment.updatePartition(new PartitionUpdate(
                     partitionAssignment.getPartitionId(),
                     planNodeId,
+                    false,
                     splits,
                     noMoreSplits));
         }
@@ -153,6 +154,7 @@ private AssignmentResult assignReplicatedSplits(PlanNodeId planNodeId, List<Spli
                     assignment.updatePartition(new PartitionUpdate(
                             0,
                             replicatedSourceId,
+                            false,
                             replicatedSplits.get(replicatedSourceId),
                             true));
                 }
@@ -160,6 +162,7 @@ private AssignmentResult assignReplicatedSplits(PlanNodeId planNodeId, List<Spli
                     assignment.updatePartition(new PartitionUpdate(
                             0,
                             partitionedSourceId,
+                            false,
                             ImmutableList.of(),
                             true));
                 }
@@ -173,6 +176,7 @@ private AssignmentResult assignReplicatedSplits(PlanNodeId planNodeId, List<Spli
                             assignment.updatePartition(new PartitionUpdate(
                                     partitionAssignment.getPartitionId(),
                                     partitionedSourceNodeId,
+                                    false,
                                     ImmutableList.of(),
                                     true));
                         }
@@ -203,6 +207,7 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
                     assignment.updatePartition(new PartitionUpdate(
                             partitionAssignment.getPartitionId(),
                             partitionedSourceNodeId,
+                            false,
                             ImmutableList.of(),
                             true));
                 }
@@ -232,6 +237,7 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
                     assignment.updatePartition(new PartitionUpdate(
                             partitionAssignment.getPartitionId(),
                             replicatedSourceId,
+                            false,
                             replicatedSplits.get(replicatedSourceId),
                             completedSources.contains(replicatedSourceId)));
                 }
@@ -239,6 +245,7 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
             assignment.updatePartition(new PartitionUpdate(
                     partitionAssignment.getPartitionId(),
                     planNodeId,
+                    true,
                     ImmutableList.of(split),
                     false));
             partitionAssignment.assignSplit(splitSizeInBytes);
@@ -257,6 +264,7 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
                     assignment.updatePartition(new PartitionUpdate(
                             0,
                             replicatedSourceId,
+                            false,
                             replicatedSplits.get(replicatedSourceId),
                             true));
                 }
@@ -264,6 +272,7 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
                     assignment.updatePartition(new PartitionUpdate(
                             0,
                             partitionedSourceId,
+                            false,
                             ImmutableList.of(),
                             true));
                 }
@@ -277,6 +286,7 @@ private AssignmentResult assignPartitionedSplits(PlanNodeId planNodeId, List<Spl
                         assignment.updatePartition(new PartitionUpdate(
                                 partitionAssignment.getPartitionId(),
                                 partitionedSourceNodeId,
+                                false,
                                 ImmutableList.of(),
                                 true));
                     }
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java
index b34687c6acd0..9ed72bca9b43 100644
--- a/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java
+++ b/core/trino-main/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java
@@ -1131,15 +1131,16 @@ public void onSplitAssignment(SplitAssignmentEvent event)
             StageExecution stageExecution = getStageExecution(event.getStageId());
             AssignmentResult assignment = event.getAssignmentResult();
             for (Partition partition : assignment.partitionsAdded()) {
-                Optional<PrioritizedScheduledTask> scheduledTask = stageExecution.addPartition(partition.partitionId(), partition.nodeRequirements());
-                scheduledTask.ifPresent(schedulingQueue::addOrUpdate);
+                stageExecution.addPartition(partition.partitionId(), partition.nodeRequirements());
             }
             for (PartitionUpdate partitionUpdate : assignment.partitionUpdates()) {
-                stageExecution.updatePartition(
+                Optional<PrioritizedScheduledTask> scheduledTask = stageExecution.updatePartition(
                         partitionUpdate.partitionId(),
                         partitionUpdate.planNodeId(),
+                        partitionUpdate.readyForScheduling(),
                         partitionUpdate.splits(),
                         partitionUpdate.noMoreSplits());
+                scheduledTask.ifPresent(schedulingQueue::addOrUpdate);
             }
             assignment.sealedPartitions().forEach(partitionId -> {
                 Optional<PrioritizedScheduledTask> scheduledTask = stageExecution.sealPartition(partitionId);
@@ -1300,10 +1301,10 @@ public boolean isExchangeClosed()
             return exchangeClosed;
         }
 
-        public Optional<PrioritizedScheduledTask> addPartition(int partitionId, NodeRequirements nodeRequirements)
+        public void addPartition(int partitionId, NodeRequirements nodeRequirements)
         {
             if (getState().isDone()) {
-                return Optional.empty();
+                return;
             }
 
             ExchangeSinkHandle exchangeSinkHandle = exchange.addSink(partitionId);
@@ -1323,18 +1324,28 @@ public Optional<PrioritizedScheduledTask> addPartition(int partitionId, NodeRequ
             checkState(partitions.putIfAbsent(partitionId, partition) == null, "partition with id %s already exist in stage %s", partitionId, stage.getStageId());
             getSourceOutputSelectors().forEach((partition::updateExchangeSourceOutputSelector));
             remainingPartitions.add(partitionId);
-
-            return Optional.of(PrioritizedScheduledTask.createSpeculative(stage.getStageId(), partitionId, schedulingPriority));
         }
 
-        public void updatePartition(int partitionId, PlanNodeId planNodeId, List<Split> splits, boolean noMoreSplits)
+        public Optional<PrioritizedScheduledTask> updatePartition(
+                int partitionId,
+                PlanNodeId planNodeId,
+                boolean readyForScheduling,
+                List<Split> splits,
+                boolean noMoreSplits)
         {
             if (getState().isDone()) {
-                return;
+                return Optional.empty();
             }
 
             StagePartition partition = getStagePartition(partitionId);
             partition.addSplits(planNodeId, splits, noMoreSplits);
+            if (readyForScheduling && !partition.isTaskScheduled()) {
+                partition.setTaskScheduled(true);
+                return Optional.of(PrioritizedScheduledTask.createSpeculative(stage.getStageId(), partitionId, schedulingPriority));
+            }
+            else {
+                return Optional.empty();
+            }
         }
 
         public Optional<PrioritizedScheduledTask> sealPartition(int partitionId)
@@ -1760,6 +1771,7 @@ private static class StagePartition
         private final Set<TaskId> runningTasks = new HashSet<>();
         private final Set<PlanNodeId> finalSelectors = new HashSet<>();
         private final Set<PlanNodeId> noMoreSplits = new HashSet<>();
+        private boolean taskScheduled;
         private boolean finished;
 
         public StagePartition(
@@ -1951,6 +1963,17 @@ public boolean isRunning()
             return !runningTasks.isEmpty();
         }
 
+        public boolean isTaskScheduled()
+        {
+            return taskScheduled;
+        }
+
+        public void setTaskScheduled(boolean taskScheduled)
+        {
+            checkArgument(taskScheduled, "taskScheduled must be true");
+            this.taskScheduled = taskScheduled;
+        }
+
         public boolean isFinished()
         {
             return finished;
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/HashDistributionSplitAssigner.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/HashDistributionSplitAssigner.java
index f7775abe553d..81235c8a46a3 100644
--- a/core/trino-main/src/main/java/io/trino/execution/scheduler/HashDistributionSplitAssigner.java
+++ b/core/trino-main/src/main/java/io/trino/execution/scheduler/HashDistributionSplitAssigner.java
@@ -147,7 +147,7 @@ public AssignmentResult assign(PlanNodeId planNodeId, ListMultimap<Integer, Spli
         if (replicatedSources.contains(planNodeId)) {
             replicatedSplits.putAll(planNodeId, splits.values());
             for (Integer partitionId : createdTaskPartitions) {
-                assignment.updatePartition(new PartitionUpdate(partitionId, planNodeId, ImmutableList.copyOf(splits.values()), noMoreSplits));
+                assignment.updatePartition(new PartitionUpdate(partitionId, planNodeId, false, ImmutableList.copyOf(splits.values()), noMoreSplits));
             }
         }
         else {
@@ -164,7 +164,7 @@ public AssignmentResult assign(PlanNodeId planNodeId, ListMultimap<Integer, Spli
                 }
 
                 for (SubPartition subPartition : subPartitions) {
-                    assignment.updatePartition(new PartitionUpdate(subPartition.getId(), planNodeId, ImmutableList.of(split), false));
+                    assignment.updatePartition(new PartitionUpdate(subPartition.getId(), planNodeId, true, ImmutableList.of(split), false));
                 }
             });
         }
@@ -172,7 +172,7 @@ public AssignmentResult assign(PlanNodeId planNodeId, ListMultimap<Integer, Spli
         if (noMoreSplits) {
             completedSources.add(planNodeId);
             for (Integer taskPartition : createdTaskPartitions) {
-                assignment.updatePartition(new PartitionUpdate(taskPartition, planNodeId, ImmutableList.of(), true));
+                assignment.updatePartition(new PartitionUpdate(taskPartition, planNodeId, false, ImmutableList.of(), true));
             }
 
             if (completedSources.containsAll(allSources)) {
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/SingleDistributionSplitAssigner.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/SingleDistributionSplitAssigner.java
index 4cbf38e130bf..a07e5b455db3 100644
--- a/core/trino-main/src/main/java/io/trino/execution/scheduler/SingleDistributionSplitAssigner.java
+++ b/core/trino-main/src/main/java/io/trino/execution/scheduler/SingleDistributionSplitAssigner.java
@@ -56,6 +56,7 @@ public AssignmentResult assign(PlanNodeId planNodeId, ListMultimap<Integer, Spli
             assignment.updatePartition(new PartitionUpdate(
                     0,
                     planNodeId,
+                    true,
                     ImmutableList.copyOf(splits.values()),
                     false));
         }
@@ -63,6 +64,7 @@ public AssignmentResult assign(PlanNodeId planNodeId, ListMultimap<Integer, Spli
             assignment.updatePartition(new PartitionUpdate(
                     0,
                     planNodeId,
+                    false,
                     ImmutableList.of(),
                     true));
             completedSources.add(planNodeId);
diff --git a/core/trino-main/src/main/java/io/trino/execution/scheduler/SplitAssigner.java b/core/trino-main/src/main/java/io/trino/execution/scheduler/SplitAssigner.java
index 908d8a42b2ae..81614f59ee49 100644
--- a/core/trino-main/src/main/java/io/trino/execution/scheduler/SplitAssigner.java
+++ b/core/trino-main/src/main/java/io/trino/execution/scheduler/SplitAssigner.java
@@ -24,6 +24,7 @@
 
 import java.util.List;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -44,11 +45,17 @@ record Partition(int partitionId, NodeRequirements nodeRequirements)
         }
     }
 
-    record PartitionUpdate(int partitionId, PlanNodeId planNodeId, List<Split> splits, boolean noMoreSplits)
+    record PartitionUpdate(
+            int partitionId,
+            PlanNodeId planNodeId,
+            boolean readyForScheduling,
+            List<Split> splits,
+            boolean noMoreSplits)
     {
         public PartitionUpdate
         {
             requireNonNull(planNodeId, "planNodeId is null");
+            checkArgument(!(readyForScheduling && splits.isEmpty()), "partition update with empty splits marked as ready for scheduling");
             splits = ImmutableList.copyOf(requireNonNull(splits, "splits is null"));
         }
     }
diff --git a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestEventDrivenTaskSource.java b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestEventDrivenTaskSource.java
index abb16729ead9..a5320299276d 100644
--- a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestEventDrivenTaskSource.java
+++ b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestEventDrivenTaskSource.java
@@ -665,15 +665,15 @@ public AssignmentResult assign(PlanNodeId planNodeId, ListMultimap<Integer, Spli
                 if (partitions.add(partition)) {
                     result.addPartition(new Partition(partition, new NodeRequirements(Optional.empty(), ImmutableSet.of())));
                     for (PlanNodeId finishedSource : finishedSources) {
-                        result.updatePartition(new PartitionUpdate(partition, finishedSource, ImmutableList.of(), true));
+                        result.updatePartition(new PartitionUpdate(partition, finishedSource, false, ImmutableList.of(), true));
                     }
                 }
-                result.updatePartition(new PartitionUpdate(partition, planNodeId, splits, noMoreSplits));
+                result.updatePartition(new PartitionUpdate(partition, planNodeId, true, splits, noMoreSplits));
             });
             if (noMoreSplits) {
                 finishedSources.add(planNodeId);
                 for (Integer partition : partitions) {
-                    result.updatePartition(new PartitionUpdate(partition, planNodeId, ImmutableList.of(), true));
+                    result.updatePartition(new PartitionUpdate(partition, planNodeId, false, ImmutableList.of(), true));
                 }
             }
             if (finishedSources.containsAll(allSources)) {