Use pre-order stages ordering in EXPLAN output

Author: losipiuk

core/trino-main/src/main/java/io/trino/operator/ExplainAnalyzeOperator.java

@@ -28,7 +28,6 @@
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static io.trino.sql.planner.planprinter.PlanPrinter.textDistributedPlan;
 import static java.util.Objects.requireNonNull;
@@ -153,15 +152,14 @@ public Page getOutput()
 
         QueryInfo queryInfo = queryPerformanceFetcher.getQueryInfo(operatorContext.getDriverContext().getTaskId().getQueryId());
         checkState(queryInfo.getStages().isPresent(), "Stages informations is missing");
-        checkState(queryInfo.getStages().get().getOutputStage().getSubStages().size() == 1, "Expected one sub stage of explain node");
+        StagesInfo stagesInfo = queryInfo.getStages().get();
+        checkState(stagesInfo.getOutputStage().getSubStages().size() == 1, "Expected one sub stage of explain node");
 
-        if (!hasFinalStageInfo(queryInfo.getStages().get())) {
+        if (!hasFinalStageInfo(stagesInfo)) {
             return null;
         }
 
-        List<StageInfo> stagesWithoutOutputStage = queryInfo.getStages().orElseThrow().getStages().stream()
-                .filter(stage -> !stage.getStageId().equals(queryInfo.getStages().orElseThrow().getOutputStageId()))
-                .collect(toImmutableList());
+        List<StageInfo> stagesWithoutOutputStage = stagesInfo.getSubStagesDeepPreOrder(stagesInfo.getOutputStageId(), false);
 
         String plan = textDistributedPlan(
                 stagesWithoutOutputStage,

core/trino-main/src/main/java/io/trino/sql/planner/planprinter/PlanPrinter.java

@@ -428,7 +428,7 @@ public static String textDistributedPlan(
             Anonymizer anonymizer,
             NodeVersion version)
     {
-        return textDistributedPlan(stages.getStages(), queryStats, valuePrinter, verbose, anonymizer, version);
+        return textDistributedPlan(stages.getSubStagesDeepPreOrder(stages.getOutputStageId(), true), queryStats, valuePrinter, verbose, anonymizer, version);
     }
 
     public static String textDistributedPlan(

testing/trino-benchmark-queries/src/main/resources/sql/trino/tpch/q17.sql

@@ -1,18 +1,18 @@
-SELECT 
-  sum(l.extendedprice)/7.0 as avg_yearly 
-FROM 
-  "${database}"."${schema}"."${prefix}lineitem" l,
-  "${database}"."${schema}"."${prefix}part" p
-WHERE 
-  p.partkey = l.partkey 
-  AND p.brand = 'Brand#23' 
+SELECT
+  sum(l.extendedprice)/7.0 as avg_yearly
+FROM
+  "lineitem" l,
+  "part" p
+WHERE
+  p.partkey = l.partkey
+  AND p.brand = 'Brand#23'
   AND p.container = 'MED BOX'
   AND l.quantity < (
-    SELECT 
-      0.2*avg(l.quantity) 
-    FROM 
-      "${database}"."${schema}"."${prefix}lineitem" l
-    WHERE 
+    SELECT
+      0.2*avg(l.quantity)
+    FROM
+      "lineitem" l
+    WHERE
     l.partkey = p.partkey
   )
 ;