Modified ExplainAnalyzeOperator to support multiple substages in output

The EXPLAIN ANALYZE operator only supports one substage when returning the output. When outputting in TEXT format, modify it to loop through the substages and return the substage ID and its plan. For JSON format, return a list of plans. Resolves: #23798
prestodb · Oct 15, 2024 · dbe88ac · dbe88ac
1 parent 7deadd1
commit dbe88ac
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 12 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/ExplainAnalyzeOperator.java b/presto-main/src/main/java/com/facebook/presto/operator/ExplainAnalyzeOperator.java
@@ -26,6 +26,7 @@
 import com.google.common.collect.ImmutableList;
 
 import java.util.List;
+import java.util.StringJoiner;
 import java.util.concurrent.TimeUnit;
 
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
@@ -149,18 +150,28 @@ public Page getOutput()
 
         QueryInfo queryInfo = queryPerformanceFetcher.getQueryInfo(operatorContext.getDriverContext().getTaskId().getQueryId());
         checkState(queryInfo.getOutputStage().isPresent(), "Output stage is missing");
-        checkState(queryInfo.getOutputStage().get().getSubStages().size() == 1, "Expected one sub stage of explain node");
 
         if (!hasFinalStageInfo(queryInfo.getOutputStage().get())) {
             return null;
         }
         String plan;
         switch (format) {
             case TEXT:
-                plan = textDistributedPlan(queryInfo.getOutputStage().get().getSubStages().get(0), functionAndTypeManager, operatorContext.getSession(), verbose);
+                StringBuilder planStringBuilder = new StringBuilder();
+                for (StageInfo stageInfo : queryInfo.getOutputStage().get().getSubStages()) {
+                    planStringBuilder.append("Stage ID: ").append(stageInfo.getStageId().toString()).append('\n')
+                            .append(textDistributedPlan(stageInfo, functionAndTypeManager, operatorContext.getSession(),
+                                    verbose));
+                }
+                plan = planStringBuilder.toString();
                 break;
             case JSON:
-                plan = jsonDistributedPlan(queryInfo.getOutputStage().get().getSubStages().get(0), functionAndTypeManager, operatorContext.getSession());
+                StringJoiner planStringJoiner = new StringJoiner(",", "[", "]");
+                for (StageInfo stageInfo : queryInfo.getOutputStage().get().getSubStages()) {
+                    planStringJoiner.add(jsonDistributedPlan(stageInfo, functionAndTypeManager,
+                            operatorContext.getSession()));
+                }
+                plan = planStringJoiner.toString();
                 break;
             default:
                 throw new PrestoException(GENERIC_INTERNAL_ERROR, "Explain format not supported: " + format);

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/JsonRenderer.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/JsonRenderer.java
@@ -41,7 +41,7 @@ public class JsonRenderer
 {
     private final JsonCodec<Map<PlanFragmentId, JsonPlanFragment>> planMapCodec;
     private final JsonCodec<JsonRenderedNode> codec;
-    private final JsonCodec<Map<PlanFragmentId, JsonPlan>> deserializationCodec;
+    private final JsonCodec<List<Map<PlanFragmentId, JsonPlan>>> deserializationCodec;
 
     public JsonRenderer(FunctionAndTypeManager functionAndTypeManager)
     {
@@ -52,10 +52,10 @@ public JsonRenderer(FunctionAndTypeManager functionAndTypeManager)
         JsonCodecFactory codecFactory = new JsonCodecFactory(provider, true);
         this.codec = codecFactory.jsonCodec(JsonRenderedNode.class);
         this.planMapCodec = codecFactory.mapJsonCodec(PlanFragmentId.class, JsonPlanFragment.class);
-        this.deserializationCodec = codecFactory.mapJsonCodec(PlanFragmentId.class, JsonPlan.class);
+        this.deserializationCodec = codecFactory.listJsonCodec(codecFactory.mapJsonCodec(PlanFragmentId.class, JsonPlan.class));
     }
 
-    public Map<PlanFragmentId, JsonPlan> deserialize(String serialized)
+    public List<Map<PlanFragmentId, JsonPlan>> deserialize(String serialized)
     {
         return deserializationCodec.fromJson(serialized);
     }

diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java
@@ -380,12 +380,18 @@ private static void assertJsonNodesHaveStats(JsonRenderer.JsonRenderedNode node)
     public void testExplainAnalyzeFormatJson()
     {
         JsonRenderer renderer = new JsonRenderer(getQueryRunner().getMetadata().getFunctionAndTypeManager());
-        Map<PlanFragmentId, JsonRenderer.JsonPlan> fragments = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT * FROM orders").getOnlyValue());
-        fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan()));
-        fragments = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders").getOnlyValue());
-        fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan()));
-        fragments = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders WHERE orderkey < 0").getOnlyValue());
-        fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan()));
+        List<Map<PlanFragmentId, JsonRenderer.JsonPlan>> fragmentsList = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT * FROM orders").getOnlyValue());
+        for (Map<PlanFragmentId, JsonRenderer.JsonPlan> fragments : fragmentsList) {
+            fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan()));
+        }
+        fragmentsList = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders").getOnlyValue());
+        for (Map<PlanFragmentId, JsonRenderer.JsonPlan> fragments : fragmentsList) {
+            fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan()));
+        }
+        fragmentsList = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders WHERE orderkey < 0").getOnlyValue());
+        for (Map<PlanFragmentId, JsonRenderer.JsonPlan> fragments : fragmentsList) {
+            fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan()));
+        }
     }
 
     @Test(expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "EXPLAIN ANALYZE doesn't support statement type: DropTable")