From 0e64d8cf45d2399c1a02eb3659808ed2d0f74659 Mon Sep 17 00:00:00 2001
From: Hazmi <ialhazmim@gmail.com>
Date: Mon, 14 Oct 2024 21:26:01 +0300
Subject: [PATCH] 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
---
 .../presto/operator/ExplainAnalyzeOperator.java | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 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
index 956d64b58b924..690662ab7880c 100644
--- 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,7 +150,6 @@ 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;
@@ -157,10 +157,21 @@ public Page getOutput()
         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);