Use topological stages ordering in EXPLAIN output

Author: losipiuk

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

@@ -20,6 +20,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -121,6 +122,32 @@ private void collectSubStageIdsPreOrder(StageInfo stageInfo, ImmutableSet.Builde
         });
     }
 
+    @JsonIgnore
+    public List<StageInfo> getSubStagesDeepTopological(StageId root, boolean includeRoot)
+    {
+        ImmutableList.Builder<StageInfo> builder = ImmutableList.builder();
+        getSubStagesDeepTopologicalInner(root, builder, new HashSet<>(), includeRoot);
+
+        return builder.build().reverse();
+    }
+
+    private void getSubStagesDeepTopologicalInner(StageId stageId, ImmutableList.Builder<StageInfo> builder, Set<StageId> visitedFragments, boolean includeRoot)
+    {
+        if (visitedFragments.contains(stageId)) {
+            return;
+        }
+
+        StageInfo stageInfo = stagesById.get(stageId);
+
+        for (StageId childId : stageInfo.getSubStages().reversed()) {
+            getSubStagesDeepTopologicalInner(childId, builder, visitedFragments, true);
+        }
+        if (includeRoot) {
+            builder.add(stageInfo);
+        }
+        visitedFragments.add(stageId);
+    }
+
     public static List<StageInfo> getAllStages(Optional<StagesInfo> stages)
     {
         return stages.map(StagesInfo::getStages).orElse(ImmutableList.of());

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().queryId());
         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.getSubStagesDeepTopological(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.getSubStagesDeepTopological(stages.getOutputStageId(), true), queryStats, valuePrinter, verbose, anonymizer, version);
     }
 
     public static String textDistributedPlan(