Add enhanced cte scheduling mode

presto-docs/src/main/sphinx/admin/cte-materialization.rst

@@ -118,7 +118,6 @@ This setting specifies the Hash function type for CTE materialization.
 
 Use the ``hive.bucket_function_type_for_cte_materialization`` session property to set on a per-query basis.
 
-
 ``query.max-written-intermediate-bytes``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -129,6 +128,17 @@ This setting defines a cap on the amount of data that can be written during CTE
 
 Use the ``query_max_written_intermediate_bytes`` session property to set on a per-query basis.
 
+``enhanced-cte-scheduling-enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+    * **Type:** ``boolean``
+    * **Default value:** ``true``
+
+Flag to enable or disable the enhanced-cte-blocking during CTE Materialization. Enhanced CTE blocking restricts only the table scan stages of the CTE TableScan, rather than blocking entire plan sections, including the main query, until the query completes.
+This approach can improve latency in scenarios where parts of the query can execute concurrently with CTE materialization writes.
+
+Use the ``enhanced_cte_scheduling_enabled`` session property to set on a per-query basis.
+
 
 How to Participate in Development
 ---------------------------------

presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java

@@ -1171,6 +1171,15 @@ public void testChainedCteProjectionAndFilterPushDown()
                 generateMaterializedCTEInformation("cte5", 1, false, true)));
     }
 
+    @Test
+    public void testCTEMaterializationWithEnhancedScheduling()
+    {
+        QueryRunner queryRunner = getQueryRunner();
+        String sql = "WITH  temp as (SELECT orderkey FROM ORDERS) " +
+                "SELECT * FROM temp t1 JOIN (SELECT custkey FROM customer) c ON t1.orderkey=c.custkey";
+        verifyResults(queryRunner, sql, ImmutableList.of(generateMaterializedCTEInformation("temp", 1, false, true)));
+    }
+
     @Test
     public void testWrittenIntemediateByteLimit()
             throws Exception

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -234,6 +234,7 @@ public final class SystemSessionProperties
     public static final String QUERY_RETRY_MAX_EXECUTION_TIME = "query_retry_max_execution_time";
     public static final String PARTIAL_RESULTS_ENABLED = "partial_results_enabled";
     public static final String PARTIAL_RESULTS_COMPLETION_RATIO_THRESHOLD = "partial_results_completion_ratio_threshold";
+    public static final String ENHANCED_CTE_SCHEDULING_ENABLED = "enhanced-cte-scheduling-enabled";
     public static final String PARTIAL_RESULTS_MAX_EXECUTION_TIME_MULTIPLIER = "partial_results_max_execution_time_multiplier";
     public static final String OFFSET_CLAUSE_ENABLED = "offset_clause_enabled";
     public static final String VERBOSE_EXCEEDED_MEMORY_LIMIT_ERRORS_ENABLED = "verbose_exceeded_memory_limit_errors_enabled";
@@ -1278,6 +1279,11 @@ public SystemSessionProperties(
                         "Minimum query completion ratio threshold for partial results",
                         featuresConfig.getPartialResultsCompletionRatioThreshold(),
                         false),
+                booleanProperty(
+                        ENHANCED_CTE_SCHEDULING_ENABLED,
+                        "Applicable for CTE Materialization. If enabled, only tablescans of the pending tablewriters are blocked and other stages can continue.",
+                        featuresConfig.getEnhancedCTESchedulingEnabled(),
+                        true),
                 booleanProperty(
                         OFFSET_CLAUSE_ENABLED,
                         "Enable support for OFFSET clause",
@@ -2663,6 +2669,11 @@ public static double getPartialResultsCompletionRatioThreshold(Session session)
         return session.getSystemProperty(PARTIAL_RESULTS_COMPLETION_RATIO_THRESHOLD, Double.class);
     }
 
+    public static boolean isEnhancedCTESchedulingEnabled(Session session)
+    {
+        return isCteMaterializationApplicable(session) & session.getSystemProperty(ENHANCED_CTE_SCHEDULING_ENABLED, Boolean.class);
+    }
+
     public static double getPartialResultsMaxExecutionTimeMultiplier(Session session)
     {
         return session.getSystemProperty(PARTIAL_RESULTS_MAX_EXECUTION_TIME_MULTIPLIER, Double.class);

presto-main/src/main/java/com/facebook/presto/execution/SqlStageExecution.java

@@ -26,10 +26,15 @@
 import com.facebook.presto.metadata.Split;
 import com.facebook.presto.server.remotetask.HttpRemoteTask;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.plan.CteMaterializationInfo;
 import com.facebook.presto.spi.plan.PlanFragmentId;
+import com.facebook.presto.spi.plan.PlanNode;
 import com.facebook.presto.spi.plan.PlanNodeId;
+import com.facebook.presto.spi.plan.TableFinishNode;
+import com.facebook.presto.spi.plan.TableScanNode;
 import com.facebook.presto.split.RemoteSplit;
 import com.facebook.presto.sql.planner.PlanFragment;
+import com.facebook.presto.sql.planner.optimizations.PlanNodeSearcher;
 import com.facebook.presto.sql.planner.plan.RemoteSourceNode;
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableMap;
@@ -60,8 +65,10 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 import static com.facebook.presto.SystemSessionProperties.getMaxFailedTaskPercentage;
+import static com.facebook.presto.SystemSessionProperties.isEnhancedCTESchedulingEnabled;
 import static com.facebook.presto.failureDetector.FailureDetector.State.GONE;
 import static com.facebook.presto.operator.ExchangeOperator.REMOTE_CONNECTOR_ID;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
@@ -557,7 +564,6 @@ private synchronized RemoteTask scheduleTask(InternalNode node, TaskId taskId, M
             // stage finished while we were scheduling this task
             task.abort();
         }
-
         return task;
     }
 
@@ -594,6 +600,59 @@ private static Split createRemoteSplitFor(TaskId taskId, URI remoteSourceTaskLoc
         return new Split(REMOTE_CONNECTOR_ID, new RemoteTransactionHandle(), new RemoteSplit(new Location(splitLocation), remoteSourceTaskId));
     }
 
+    private static String getCteIdFromSource(PlanNode source)
+    {
+        // Traverse the plan node tree to find a TableWriterNode with TemporaryTableInfo
+        return PlanNodeSearcher.searchFrom(source)
+                .where(planNode -> planNode instanceof TableFinishNode)
+                .findFirst()
+                .flatMap(planNode -> ((TableFinishNode) planNode).getCteMaterializationInfo())
+                .map(CteMaterializationInfo::getCteId)
+                .orElseThrow(() -> new IllegalStateException("TemporaryTableInfo has no CTE ID"));
+    }
+
+    public boolean isCTETableFinishStage()
+    {
+        return PlanNodeSearcher.searchFrom(planFragment.getRoot())
+                .where(planNode -> planNode instanceof TableFinishNode &&
+                        ((TableFinishNode) planNode).getCteMaterializationInfo().isPresent())
+                .findSingle()
+                .isPresent();
+    }
+
+    public String getCTEWriterId()
+    {
+        // Validate that this is a CTE TableFinish stage and return the associated CTE ID
+        if (!isCTETableFinishStage()) {
+            throw new IllegalStateException("This stage is not a CTE writer stage");
+        }
+        return getCteIdFromSource(planFragment.getRoot());
+    }
+
+    public boolean requiresMaterializedCTE()
+    {
+        if (!isEnhancedCTESchedulingEnabled(session)) {
+            return false;
+        }
+        // Search for TableScanNodes and check if they reference TemporaryTableInfo
+        return PlanNodeSearcher.searchFrom(planFragment.getRoot())
+                .where(planNode -> planNode instanceof TableScanNode)
+                .findAll().stream()
+                .anyMatch(planNode -> ((TableScanNode) planNode).getCteMaterializationInfo().isPresent());
+    }
+
+    public List<String> getRequiredCTEList()
+    {
+        // Collect all CTE IDs referenced by TableScanNodes with TemporaryTableInfo
+        return PlanNodeSearcher.searchFrom(planFragment.getRoot())
+                .where(planNode -> planNode instanceof TableScanNode)
+                .findAll().stream()
+                .map(planNode -> ((TableScanNode) planNode).getCteMaterializationInfo()
+                        .orElseThrow(() -> new IllegalStateException("TableScanNode has no TemporaryTableInfo")))
+                .map(CteMaterializationInfo::getCteId)
+                .collect(Collectors.toList());
+    }
+
     private void updateTaskStatus(TaskId taskId, TaskStatus taskStatus)
     {
         StageExecutionState stageExecutionState = getState();

presto-main/src/main/java/com/facebook/presto/execution/scheduler/CTEMaterializationTracker.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.execution.scheduler;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+/*
+ * Tracks the completion status of table-finish nodes that write temporary tables for CTE materialization.
+ * CTEMaterializationTracker manages a map of materialized CTEs and their associated materialization futures.
+ * When a stage includes a CTE table finish, it marks the corresponding CTE as materialized and completes
+ * the associated future. This signals the scheduler that some dependency has been resolved, prompting it
+ * to resume scheduling. The scheduler can call the clearAllFutures() method once it starts scheduling again
+ */
+public class CTEMaterializationTracker
+{
+    private final Map<String, SettableFuture<Void>> materializationFutures = new HashMap<>();
+
+    private final Map<String, Boolean> materializedCtes = new HashMap<>();
+
+    public synchronized ListenableFuture<Void> getFutureForCTE(String cteName)
+    {
+        return materializationFutures.compute(cteName, (key, existingFuture) -> {
+            if (existingFuture == null) {
+                return SettableFuture.create();
+            }
+            checkArgument(!existingFuture.isCancelled(), "CTE future was found in a cancelled state");
+            return existingFuture;
+        });
+    }
+
+    public synchronized void clearAllFutures()
+    {
+        materializationFutures.clear();
+    }
+
+    public synchronized void markCTEAsMaterialized(String cteName)
+    {
+        if (materializedCtes.putIfAbsent(cteName, true) == null) {
+            SettableFuture<Void> future = materializationFutures.get(cteName);
+            if (future != null) {
+                checkArgument(!future.isCancelled(), "CTE future was found in a cancelled state");
+                future.set(null); // Notify all listeners
+            }
+        }
+    }
+
+    public synchronized boolean hasBeenMaterialized(String cteName)
+    {
+        return materializedCtes.containsKey(cteName);
+    }
+}

presto-main/src/main/java/com/facebook/presto/execution/scheduler/FixedSourcePartitionedScheduler.java

@@ -74,6 +74,8 @@ public class FixedSourcePartitionedScheduler
 
     private final Queue<Integer> tasksToRecover = new ConcurrentLinkedQueue<>();
 
+    private final CTEMaterializationTracker cteMaterializationTracker;
+
     @GuardedBy("this")
     private boolean closed;
 
@@ -87,13 +89,15 @@ public FixedSourcePartitionedScheduler(
             int splitBatchSize,
             OptionalInt concurrentLifespansPerTask,
             NodeSelector nodeSelector,
-            List<ConnectorPartitionHandle> partitionHandles)
+            List<ConnectorPartitionHandle> partitionHandles,
+            CTEMaterializationTracker cteMaterializationTracker)
     {
         requireNonNull(stage, "stage is null");
         requireNonNull(splitSources, "splitSources is null");
         requireNonNull(bucketNodeMap, "bucketNodeMap is null");
         checkArgument(!requireNonNull(nodes, "nodes is null").isEmpty(), "nodes is empty");
         requireNonNull(partitionHandles, "partitionHandles is null");
+        this.cteMaterializationTracker = cteMaterializationTracker;
 
         this.stage = stage;
         this.nodes = ImmutableList.copyOf(nodes);
@@ -179,6 +183,29 @@ public ScheduleResult schedule()
     {
         // schedule a task on every node in the distribution
         List<RemoteTask> newTasks = ImmutableList.of();
+
+        // CTE Materialization Check
+        if (stage.requiresMaterializedCTE()) {
+            List<ListenableFuture<?>> blocked = new ArrayList<>();
+            List<String> requiredCTEIds = stage.getRequiredCTEList();
+            for (String cteId : requiredCTEIds) {
+                if (!cteMaterializationTracker.hasBeenMaterialized(cteId)) {
+                    // Add CTE materialization future to the blocked list
+                    ListenableFuture<Void> materializationFuture = cteMaterializationTracker.getFutureForCTE(cteId);
+                    blocked.add(materializationFuture);
+                }
+            }
+            // If any CTE is not materialized, return a blocked ScheduleResult
+            if (!blocked.isEmpty()) {
+                return ScheduleResult.blocked(
+                        false,
+                        newTasks,
+                        whenAnyComplete(blocked),
+                        BlockedReason.WAITING_FOR_CTE_MATERIALIZATION,
+                        0);
+            }
+        }
+        // schedule a task on every node in the distribution
         if (!scheduledTasks) {
             newTasks = Streams.mapWithIndex(
                     nodes.stream(),
@@ -191,9 +218,8 @@ public ScheduleResult schedule()
             // notify listeners that we have scheduled all tasks so they can set no more buffers or exchange splits
             stage.transitionToFinishedTaskScheduling();
         }
-
-        boolean allBlocked = true;
         List<ListenableFuture<?>> blocked = new ArrayList<>();
+        boolean allBlocked = true;
         BlockedReason blockedReason = BlockedReason.NO_ACTIVE_DRIVER_GROUP;
 
         if (groupedLifespanScheduler.isPresent()) {

presto-main/src/main/java/com/facebook/presto/execution/scheduler/ScheduleResult.java

@@ -57,13 +57,19 @@ public enum BlockedReason
          * grouped execution where there are multiple lifespans per task).
          */
         MIXED_SPLIT_QUEUES_FULL_AND_WAITING_FOR_SOURCE,
+
+        /**
+         * Waiting for the completion of CTE materialization by the table writer.
+         */
+        WAITING_FOR_CTE_MATERIALIZATION,
         /**/;
 
         public BlockedReason combineWith(BlockedReason other)
         {
             switch (this) {
                 case WRITER_SCALING:
                     throw new IllegalArgumentException("cannot be combined");
+                case WAITING_FOR_CTE_MATERIALIZATION:
                 case NO_ACTIVE_DRIVER_GROUP:
                     return other;
                 case SPLIT_QUEUES_FULL:

presto-main/src/main/java/com/facebook/presto/execution/scheduler/SectionExecutionFactory.java

@@ -168,7 +168,8 @@ public SectionExecution createSectionExecutions(
             boolean summarizeTaskInfo,
             RemoteTaskFactory remoteTaskFactory,
             SplitSourceFactory splitSourceFactory,
-            int attemptId)
+            int attemptId,
+            CTEMaterializationTracker cteMaterializationTracker)
     {
         // Only fetch a distribution once per section to ensure all stages see the same machine assignments
         Map<PartitioningHandle, NodePartitionMap> partitioningCache = new HashMap<>();
@@ -184,7 +185,8 @@ public SectionExecution createSectionExecutions(
                 summarizeTaskInfo,
                 remoteTaskFactory,
                 splitSourceFactory,
-                attemptId);
+                attemptId,
+                cteMaterializationTracker);
         StageExecutionAndScheduler rootStage = getLast(sectionStages);
         rootStage.getStageExecution().setOutputBuffers(outputBuffers);
         return new SectionExecution(rootStage, sectionStages);
@@ -203,7 +205,8 @@ private List<StageExecutionAndScheduler> createStreamingLinkedStageExecutions(
             boolean summarizeTaskInfo,
             RemoteTaskFactory remoteTaskFactory,
             SplitSourceFactory splitSourceFactory,
-            int attemptId)
+            int attemptId,
+            CTEMaterializationTracker cteMaterializationTracker)
     {
         ImmutableList.Builder<StageExecutionAndScheduler> stageExecutionAndSchedulers = ImmutableList.builder();
 
@@ -238,7 +241,8 @@ private List<StageExecutionAndScheduler> createStreamingLinkedStageExecutions(
                     summarizeTaskInfo,
                     remoteTaskFactory,
                     splitSourceFactory,
-                    attemptId);
+                    attemptId,
+                    cteMaterializationTracker);
             stageExecutionAndSchedulers.addAll(subTree);
             childStagesBuilder.add(getLast(subTree).getStageExecution());
         }
@@ -260,7 +264,8 @@ private List<StageExecutionAndScheduler> createStreamingLinkedStageExecutions(
                 stageExecution,
                 partitioningHandle,
                 tableWriteInfo,
-                childStageExecutions);
+                childStageExecutions,
+                cteMaterializationTracker);
         stageExecutionAndSchedulers.add(new StageExecutionAndScheduler(
                 stageExecution,
                 stageLinkage,
@@ -279,7 +284,8 @@ private StageScheduler createStageScheduler(
             SqlStageExecution stageExecution,
             PartitioningHandle partitioningHandle,
             TableWriteInfo tableWriteInfo,
-            Set<SqlStageExecution> childStageExecutions)
+            Set<SqlStageExecution> childStageExecutions,
+            CTEMaterializationTracker cteMaterializationTracker)
     {
         Map<PlanNodeId, SplitSource> splitSources = splitSourceFactory.createSplitSources(plan.getFragment(), session, tableWriteInfo);
         int maxTasksPerStage = getMaxTasksPerStage(session);
@@ -383,7 +389,8 @@ else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
                         splitBatchSize,
                         getConcurrentLifespansPerNode(session),
                         nodeScheduler.createNodeSelector(session, connectorId, nodePredicate),
-                        connectorPartitionHandles);
+                        connectorPartitionHandles,
+                        cteMaterializationTracker);
                 if (plan.getFragment().getStageExecutionDescriptor().isRecoverableGroupedExecution()) {
                     stageExecution.registerStageTaskRecoveryCallback(taskId -> {
                         checkArgument(taskId.getStageExecutionId().getStageId().equals(stageId), "The task did not execute this stage");