Remove StageExecutionDescriptor class

The class is not useful after removal of grouped execution.
Major parts of PlanFragmenter class has been removed in the process, as well as
some tests strictly related to grouped execution

Author: skrzypo987

core/trino-main/src/main/java/io/trino/execution/SqlTaskExecution.java

@@ -35,7 +35,6 @@
 import io.trino.operator.DriverStats;
 import io.trino.operator.PipelineContext;
 import io.trino.operator.PipelineExecutionStrategy;
-import io.trino.operator.StageExecutionDescriptor;
 import io.trino.operator.TaskContext;
 import io.trino.spi.SplitWeight;
 import io.trino.spi.TrinoException;
@@ -223,7 +222,7 @@ private SqlTaskExecution(
                     taskContext,
                     localExecutionPlan.getDriverFactories().stream()
                             .collect(toImmutableMap(DriverFactory::getPipelineId, DriverFactory::getPipelineExecutionStrategy)));
-            this.schedulingLifespanManager = new SchedulingLifespanManager(localExecutionPlan.getPartitionedSourceOrder(), localExecutionPlan.getStageExecutionDescriptor(), this.status);
+            this.schedulingLifespanManager = new SchedulingLifespanManager(localExecutionPlan.getPartitionedSourceOrder(), this.status);
 
             checkArgument(this.driverRunnerFactoriesWithSplitLifeCycle.keySet().equals(partitionedSources),
                     "Fragment is partitioned, but not all partitioned drivers were found");
@@ -779,7 +778,6 @@ private static class SchedulingLifespanManager
         // Note that different drivers in a task may have different pipelineExecutionStrategy.
 
         private final List<PlanNodeId> sourceStartOrder;
-        private final StageExecutionDescriptor stageExecutionDescriptor;
         private final Status status;
 
         private final Map<Lifespan, SchedulingLifespan> lifespans = new HashMap<>();
@@ -790,10 +788,9 @@ private static class SchedulingLifespanManager
 
         private int maxScheduledPlanNodeOrdinal;
 
-        public SchedulingLifespanManager(List<PlanNodeId> sourceStartOrder, StageExecutionDescriptor stageExecutionDescriptor, Status status)
+        public SchedulingLifespanManager(List<PlanNodeId> sourceStartOrder, Status status)
         {
             this.sourceStartOrder = ImmutableList.copyOf(sourceStartOrder);
-            this.stageExecutionDescriptor = stageExecutionDescriptor;
             this.status = requireNonNull(status, "status is null");
         }
 
@@ -884,13 +881,6 @@ public Optional<PlanNodeId> getSchedulingPlanNode()
         {
             checkState(!isDone());
             while (!isDone()) {
-                // Return current plan node if this lifespan is compatible with the plan node.
-                // i.e. One of the following bullet points is true:
-                // * The execution strategy of the plan node is grouped. And lifespan represents a driver group.
-                // * The execution strategy of the plan node is ungrouped. And lifespan is task wide.
-                if (manager.stageExecutionDescriptor.isScanGroupedExecution(manager.sourceStartOrder.get(schedulingPlanNodeOrdinal)) != lifespan.isTaskWide()) {
-                    return Optional.of(manager.sourceStartOrder.get(schedulingPlanNodeOrdinal));
-                }
                 // This lifespan is incompatible with the plan node. As a result, this method should either
                 // return empty to indicate that scheduling for this lifespan is blocked, or skip the current
                 // plan node and went on to the next one. Which one of the two happens is dependent on whether

core/trino-main/src/main/java/io/trino/execution/SqlTaskExecutionFactory.java

@@ -81,7 +81,6 @@ public SqlTaskExecution create(
                         fragment.getRoot(),
                         TypeProvider.copyOf(fragment.getSymbols()),
                         fragment.getPartitioningScheme(),
-                        fragment.getStageExecutionDescriptor(),
                         fragment.getPartitionedSources(),
                         outputBuffer);
             }

core/trino-main/src/main/java/io/trino/execution/scheduler/FaultTolerantStageScheduler.java

@@ -197,8 +197,6 @@ public FaultTolerantStageScheduler(
             int taskRetryAttemptsPerTask,
             int maxTasksWaitingForNodePerStage)
     {
-        checkArgument(!stage.getFragment().getStageExecutionDescriptor().isStageGroupedExecution(), "grouped execution is expected to be disabled");
-
         this.session = requireNonNull(session, "session is null");
         this.stage = requireNonNull(stage, "stage is null");
         this.failureDetector = requireNonNull(failureDetector, "failureDetector is null");

core/trino-main/src/main/java/io/trino/execution/scheduler/FixedSourcePartitionedScheduler.java

@@ -22,12 +22,8 @@
 import io.trino.execution.RemoteTask;
 import io.trino.execution.TableExecuteContextManager;
 import io.trino.execution.scheduler.ScheduleResult.BlockedReason;
-import io.trino.execution.scheduler.group.DynamicLifespanScheduler;
-import io.trino.execution.scheduler.group.FixedLifespanScheduler;
-import io.trino.execution.scheduler.group.LifespanScheduler;
 import io.trino.metadata.InternalNode;
 import io.trino.metadata.Split;
-import io.trino.operator.StageExecutionDescriptor;
 import io.trino.server.DynamicFilterService;
 import io.trino.spi.connector.ConnectorPartitionHandle;
 import io.trino.split.SplitSource;
@@ -43,7 +39,6 @@
 import java.util.function.Supplier;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static io.airlift.concurrent.MoreFutures.whenAnyComplete;
 import static io.trino.execution.scheduler.SourcePartitionedScheduler.newSourcePartitionedSchedulerAsSourceScheduler;
@@ -59,15 +54,13 @@ public class FixedSourcePartitionedScheduler
     private final List<InternalNode> nodes;
     private final List<SourceScheduler> sourceSchedulers;
     private final List<ConnectorPartitionHandle> partitionHandles;
-    private final Optional<LifespanScheduler> groupedLifespanScheduler;
 
     private final PartitionIdAllocator partitionIdAllocator;
     private final Map<InternalNode, RemoteTask> scheduledTasks;
 
     public FixedSourcePartitionedScheduler(
             StageExecution stageExecution,
             Map<PlanNodeId, SplitSource> splitSources,
-            StageExecutionDescriptor stageExecutionDescriptor,
             List<PlanNodeId> schedulingOrder,
             List<InternalNode> nodes,
             BucketNodeMap bucketNodeMap,
@@ -93,20 +86,14 @@ public FixedSourcePartitionedScheduler(
         BucketedSplitPlacementPolicy splitPlacementPolicy = new BucketedSplitPlacementPolicy(nodeSelector, nodes, bucketNodeMap, stageExecution::getAllTasks);
 
         ArrayList<SourceScheduler> sourceSchedulers = new ArrayList<>();
-        checkArgument(
-                partitionHandles.equals(ImmutableList.of(NOT_PARTITIONED)) != stageExecutionDescriptor.isStageGroupedExecution(),
-                "PartitionHandles should be [NOT_PARTITIONED] if and only if all scan nodes use ungrouped execution strategy");
-        int nodeCount = nodes.size();
         int concurrentLifespans = partitionHandles.size();
 
         boolean firstPlanNode = true;
-        Optional<LifespanScheduler> groupedLifespanScheduler = Optional.empty();
 
         partitionIdAllocator = new PartitionIdAllocator();
         scheduledTasks = new HashMap<>();
         for (PlanNodeId planNodeId : schedulingOrder) {
             SplitSource splitSource = splitSources.get(planNodeId);
-            boolean groupedExecutionForScanNode = stageExecutionDescriptor.isScanGroupedExecution(planNodeId);
             // TODO : change anySourceTaskBlocked to accommodate the correct blocked status of source tasks
             //  (ref : https://github.com/trinodb/trino/issues/4713)
             SourceScheduler sourceScheduler = newSourcePartitionedSchedulerAsSourceScheduler(
@@ -115,47 +102,21 @@ public FixedSourcePartitionedScheduler(
                     splitSource,
                     splitPlacementPolicy,
                     Math.max(splitBatchSize / concurrentLifespans, 1),
-                    groupedExecutionForScanNode,
+                    false,
                     dynamicFilterService,
                     tableExecuteContextManager,
                     () -> true,
                     partitionIdAllocator,
                     scheduledTasks);
 
-            if (stageExecutionDescriptor.isStageGroupedExecution() && !groupedExecutionForScanNode) {
-                sourceScheduler = new AsGroupedSourceScheduler(sourceScheduler);
-            }
             sourceSchedulers.add(sourceScheduler);
 
             if (firstPlanNode) {
                 firstPlanNode = false;
-                if (!stageExecutionDescriptor.isStageGroupedExecution()) {
-                    sourceScheduler.startLifespan(Lifespan.taskWide(), NOT_PARTITIONED);
-                    sourceScheduler.noMoreLifespans();
-                }
-                else {
-                    LifespanScheduler lifespanScheduler;
-                    if (bucketNodeMap.isDynamic()) {
-                        // Callee of the constructor guarantees dynamic bucket node map will only be
-                        // used when the stage has no remote source.
-                        //
-                        // When the stage has no remote source, any scan is grouped execution guarantees
-                        // all scan is grouped execution.
-                        lifespanScheduler = new DynamicLifespanScheduler(bucketNodeMap, nodes, partitionHandles);
-                    }
-                    else {
-                        lifespanScheduler = new FixedLifespanScheduler(bucketNodeMap, partitionHandles);
-                    }
-
-                    // Schedule the first few lifespans
-                    lifespanScheduler.scheduleInitial(sourceScheduler);
-                    // Schedule new lifespans for finished ones
-                    stageExecution.addCompletedDriverGroupsChangedListener(lifespanScheduler::onLifespanFinished);
-                    groupedLifespanScheduler = Optional.of(lifespanScheduler);
-                }
+                sourceScheduler.startLifespan(Lifespan.taskWide(), NOT_PARTITIONED);
+                sourceScheduler.noMoreLifespans();
             }
         }
-        this.groupedLifespanScheduler = groupedLifespanScheduler;
         this.sourceSchedulers = sourceSchedulers;
     }
 
@@ -188,15 +149,6 @@ public ScheduleResult schedule()
         List<ListenableFuture<Void>> blocked = new ArrayList<>();
         BlockedReason blockedReason = BlockedReason.NO_ACTIVE_DRIVER_GROUP;
 
-        if (groupedLifespanScheduler.isPresent()) {
-            // Start new driver groups on the first scheduler if necessary,
-            // i.e. when previous ones have finished execution (not finished scheduling).
-            //
-            // Invoke schedule method to get a new SettableFuture every time.
-            // Reusing previously returned SettableFuture could lead to the ListenableFuture retaining too many listeners.
-            blocked.add(groupedLifespanScheduler.get().schedule(sourceSchedulers.get(0)));
-        }
-
         int splitsScheduled = 0;
         Iterator<SourceScheduler> schedulerIterator = sourceSchedulers.iterator();
         List<Lifespan> driverGroupsToStart = ImmutableList.of();
@@ -299,77 +251,4 @@ public InternalNode getNodeForBucket(int bucketId)
             return bucketNodeMap.getAssignedNode(bucketId).get();
         }
     }
-
-    private static class AsGroupedSourceScheduler
-            implements SourceScheduler
-    {
-        private final SourceScheduler sourceScheduler;
-        private boolean started;
-        private boolean completed;
-        private final List<Lifespan> pendingCompleted;
-
-        public AsGroupedSourceScheduler(SourceScheduler sourceScheduler)
-        {
-            this.sourceScheduler = requireNonNull(sourceScheduler, "sourceScheduler is null");
-            pendingCompleted = new ArrayList<>();
-        }
-
-        @Override
-        public void start()
-        {
-            sourceScheduler.start();
-        }
-
-        @Override
-        public ScheduleResult schedule()
-        {
-            return sourceScheduler.schedule();
-        }
-
-        @Override
-        public void close()
-        {
-            sourceScheduler.close();
-        }
-
-        @Override
-        public PlanNodeId getPlanNodeId()
-        {
-            return sourceScheduler.getPlanNodeId();
-        }
-
-        @Override
-        public void startLifespan(Lifespan lifespan, ConnectorPartitionHandle partitionHandle)
-        {
-            pendingCompleted.add(lifespan);
-            if (started) {
-                return;
-            }
-            started = true;
-            sourceScheduler.startLifespan(Lifespan.taskWide(), NOT_PARTITIONED);
-            sourceScheduler.noMoreLifespans();
-        }
-
-        @Override
-        public void noMoreLifespans()
-        {
-            checkState(started);
-        }
-
-        @Override
-        public List<Lifespan> drainCompletedLifespans()
-        {
-            if (!completed) {
-                List<Lifespan> lifespans = sourceScheduler.drainCompletedLifespans();
-                if (lifespans.isEmpty()) {
-                    return ImmutableList.of();
-                }
-                checkState(ImmutableList.of(Lifespan.taskWide()).equals(lifespans));
-                completed = true;
-            }
-            List<Lifespan> result = ImmutableList.copyOf(pendingCompleted);
-            pendingCompleted.clear();
-            return result;
-        }
-    }
 }

core/trino-main/src/main/java/io/trino/execution/scheduler/SqlQueryScheduler.java

@@ -1393,8 +1393,6 @@ public void stateChanged(QueryState newState)
                 NodeSelector nodeSelector = nodeScheduler.createNodeSelector(session, catalogName);
                 SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeSelector, stageExecution::getAllTasks);
 
-                checkArgument(!fragment.getStageExecutionDescriptor().isStageGroupedExecution());
-
                 return newSourcePartitionedSchedulerAsStageScheduler(
                         stageExecution,
                         planNodeId,
@@ -1431,46 +1429,27 @@ else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
                     List<PlanNodeId> schedulingOrder = fragment.getPartitionedSources();
                     Optional<CatalogName> catalogName = partitioningHandle.getConnectorId();
                     checkArgument(catalogName.isPresent(), "No connector ID for partitioning handle: %s", partitioningHandle);
-                    List<ConnectorPartitionHandle> connectorPartitionHandles;
-                    boolean groupedExecutionForStage = fragment.getStageExecutionDescriptor().isStageGroupedExecution();
-                    if (groupedExecutionForStage) {
-                        connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle);
-                        checkState(!ImmutableList.of(NOT_PARTITIONED).equals(connectorPartitionHandles));
-                    }
-                    else {
-                        connectorPartitionHandles = ImmutableList.of(NOT_PARTITIONED);
-                    }
+                    List<ConnectorPartitionHandle> connectorPartitionHandles = ImmutableList.of(NOT_PARTITIONED);
 
                     BucketNodeMap bucketNodeMap;
                     List<InternalNode> stageNodeList;
                     if (fragment.getRemoteSourceNodes().stream().allMatch(node -> node.getExchangeType() == REPLICATE)) {
                         // no remote source
-                        boolean dynamicLifespanSchedule = fragment.getStageExecutionDescriptor().isDynamicLifespanSchedule();
-                        bucketNodeMap = nodePartitioningManager.getBucketNodeMap(session, partitioningHandle, dynamicLifespanSchedule);
-
-                        // verify execution is consistent with planner's decision on dynamic lifespan schedule
-                        verify(bucketNodeMap.isDynamic() == dynamicLifespanSchedule);
+                        bucketNodeMap = nodePartitioningManager.getBucketNodeMap(session, partitioningHandle, false);
 
                         stageNodeList = new ArrayList<>(nodeScheduler.createNodeSelector(session, catalogName).allNodes());
                         Collections.shuffle(stageNodeList);
                     }
                     else {
-                        // cannot use dynamic lifespan schedule
-                        verify(!fragment.getStageExecutionDescriptor().isDynamicLifespanSchedule());
-
                         // remote source requires nodePartitionMap
                         NodePartitionMap nodePartitionMap = partitioningCache.apply(partitioningHandle);
-                        if (groupedExecutionForStage) {
-                            checkState(connectorPartitionHandles.size() == nodePartitionMap.getBucketToPartition().length);
-                        }
                         stageNodeList = nodePartitionMap.getPartitionToNode();
                         bucketNodeMap = nodePartitionMap.asBucketNodeMap();
                     }
 
                     return new FixedSourcePartitionedScheduler(
                             stageExecution,
                             splitSources,
-                            fragment.getStageExecutionDescriptor(),
                             schedulingOrder,
                             stageNodeList,
                             bucketNodeMap,