Add source fragments lookup to PartitionMemoryEstimatorFactory

trinodb · Oct 6, 2023 · 902d8f6 · 902d8f6
1 parent fc167ae
commit 902d8f6
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 7 deletions.
diff --git a/...in/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java b/...in/src/main/java/io/trino/execution/scheduler/EventDrivenFaultTolerantQueryScheduler.java
@@ -1198,6 +1198,12 @@ private void createStageExecution(SubPlan subPlan, boolean rootFragment, Map<Sta
                 if (eager) {
                     sourceExchanges.values().forEach(sourceExchange -> sourceExchange.setSourceHandlesDeliveryMode(EAGER));
                 }
+
+                Function<PlanFragmentId, PlanFragment> planFragmentLookup = planFragmentId -> {
+                    StageExecution stageExecution = stageExecutions.get(getStageId(planFragmentId));
+                    checkArgument(stageExecution != null, "stage for fragment %s not started yet", planFragmentId);
+                    return stageExecution.getStageInfo().getPlan();
+                };
                 StageExecution execution = new StageExecution(
                         queryStateMachine,
                         taskDescriptorStorage,
@@ -1206,7 +1212,7 @@ private void createStageExecution(SubPlan subPlan, boolean rootFragment, Map<Sta
                         sinkPartitioningScheme,
                         exchange,
                         noMemoryFragment,
-                        noMemoryFragment ? new NoMemoryPartitionMemoryEstimator() : memoryEstimatorFactory.createPartitionMemoryEstimator(session, fragment),
+                        noMemoryFragment ? new NoMemoryPartitionMemoryEstimator() : memoryEstimatorFactory.createPartitionMemoryEstimator(session, fragment, planFragmentLookup),
                         // do not retry coordinator only tasks
                         coordinatorStage ? 1 : maxTaskExecutionAttempts,
                         schedulingPriority,

diff --git a/...src/main/java/io/trino/execution/scheduler/ExponentialGrowthPartitionMemoryEstimator.java b/...src/main/java/io/trino/execution/scheduler/ExponentialGrowthPartitionMemoryEstimator.java
@@ -25,6 +25,7 @@
 import io.trino.spi.ErrorCode;
 import io.trino.spi.memory.MemoryPoolInfo;
 import io.trino.sql.planner.PlanFragment;
+import io.trino.sql.planner.plan.PlanFragmentId;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import org.assertj.core.util.VisibleForTesting;
@@ -34,6 +35,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -117,7 +119,10 @@ void refreshNodePoolMemoryInfos()
         }
 
         @Override
-        public PartitionMemoryEstimator createPartitionMemoryEstimator(Session session, PlanFragment planFragment)
+        public PartitionMemoryEstimator createPartitionMemoryEstimator(
+                Session session,
+                PlanFragment planFragment,
+                Function<PlanFragmentId, PlanFragment> sourceFragmentLookup)
         {
             DataSize defaultInitialMemoryLimit = planFragment.getPartitioning().equals(COORDINATOR_DISTRIBUTION) ?
                     getFaultTolerantExecutionDefaultCoordinatorTaskMemory(session) :

diff --git a/...rino-main/src/main/java/io/trino/execution/scheduler/PartitionMemoryEstimatorFactory.java b/...rino-main/src/main/java/io/trino/execution/scheduler/PartitionMemoryEstimatorFactory.java
@@ -15,9 +15,15 @@
 
 import io.trino.Session;
 import io.trino.sql.planner.PlanFragment;
+import io.trino.sql.planner.plan.PlanFragmentId;
+
+import java.util.function.Function;
 
 @FunctionalInterface
 public interface PartitionMemoryEstimatorFactory
 {
-    PartitionMemoryEstimator createPartitionMemoryEstimator(Session session, PlanFragment planFragment);
+    PartitionMemoryEstimator createPartitionMemoryEstimator(
+            Session session,
+            PlanFragment planFragment,
+            Function<PlanFragmentId, PlanFragment> sourceFragmentLookup);
 }
diff --git a/...test/java/io/trino/execution/scheduler/TestExponentialGrowthPartitionMemoryEstimator.java b/...test/java/io/trino/execution/scheduler/TestExponentialGrowthPartitionMemoryEstimator.java
@@ -36,6 +36,7 @@
 
 import java.net.URI;
 import java.util.Optional;
+import java.util.function.Function;
 
 import static io.airlift.units.DataSize.Unit.GIGABYTE;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
@@ -50,6 +51,10 @@
 
 public class TestExponentialGrowthPartitionMemoryEstimator
 {
+    private static final Function<PlanFragmentId, PlanFragment> THROWING_PLAN_FRAGMENT_LOOKUP = planFragmentId -> {
+        throw new RuntimeException("should not be used");
+    };
+
     @Test
     public void testDefaultInitialEstimation()
     {
@@ -63,10 +68,10 @@ public void testDefaultInitialEstimation()
                 .setSystemProperty(FAULT_TOLERANT_EXECUTION_TASK_MEMORY, "113MB")
                 .build();
 
-        assertThat(estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(COORDINATOR_DISTRIBUTION)).getInitialMemoryRequirements())
+        assertThat(estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(COORDINATOR_DISTRIBUTION), THROWING_PLAN_FRAGMENT_LOOKUP).getInitialMemoryRequirements())
                 .isEqualTo(new MemoryRequirements(DataSize.of(107, MEGABYTE)));
 
-        assertThat(estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(SINGLE_DISTRIBUTION)).getInitialMemoryRequirements())
+        assertThat(estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(SINGLE_DISTRIBUTION), THROWING_PLAN_FRAGMENT_LOOKUP).getInitialMemoryRequirements())
                 .isEqualTo(new MemoryRequirements(DataSize.of(113, MEGABYTE)));
     }
 
@@ -82,7 +87,7 @@ public void testEstimator()
                 .setSystemProperty(FAULT_TOLERANT_EXECUTION_TASK_MEMORY, "107MB")
                 .build();
 
-        PartitionMemoryEstimator estimator = estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(SINGLE_DISTRIBUTION));
+        PartitionMemoryEstimator estimator = estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(SINGLE_DISTRIBUTION), THROWING_PLAN_FRAGMENT_LOOKUP);
 
         assertThat(estimator.getInitialMemoryRequirements())
                 .isEqualTo(new MemoryRequirements(DataSize.of(107, MEGABYTE)));
@@ -207,7 +212,7 @@ private static void testInitialEstimationWithFinishedPartitions(
                 .setSystemProperty(FAULT_TOLERANT_EXECUTION_TASK_MEMORY, defaultInitialTaskMemory.toString())
                 .build();
 
-        PartitionMemoryEstimator estimator = estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(SINGLE_DISTRIBUTION));
+        PartitionMemoryEstimator estimator = estimatorFactory.createPartitionMemoryEstimator(session, getPlanFragment(SINGLE_DISTRIBUTION), THROWING_PLAN_FRAGMENT_LOOKUP);
 
         for (int i = 0; i < recordedPartitionsCount; i++) {
             estimator.registerPartitionFinished(new MemoryRequirements(recordedMemoryUsage), recordedMemoryUsage, true, Optional.empty());