From f5f0f20b6df5eac675c8fe4ac46de3796ad87d16 Mon Sep 17 00:00:00 2001 From: Hazmi 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 --- .../hive/TestHiveDistributedQueries.java | 21 +++++++++ .../operator/ExplainAnalyzeOperator.java | 12 +++-- .../sql/planner/planPrinter/JsonRenderer.java | 12 +++-- .../sql/planner/planPrinter/PlanPrinter.java | 23 ++++++++-- .../tests/AbstractTestDistributedQueries.java | 44 ++++++++++++++++--- 5 files changed, 95 insertions(+), 17 deletions(-) diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java index 028149a9d62a8..8f365d2ec12cc 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java @@ -17,6 +17,8 @@ import com.facebook.presto.hive.TestHiveEventListenerPlugin.TestingHiveEventListener; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.eventlistener.EventListener; +import com.facebook.presto.spi.plan.PlanFragmentId; +import com.facebook.presto.sql.planner.planPrinter.JsonRenderer; import com.facebook.presto.testing.MaterializedResult; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; @@ -24,6 +26,8 @@ import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; @@ -114,6 +118,23 @@ public void testTrackMaterializedCTEs() checkCTEInfo(explain, "tbl", 2, false, true); checkCTEInfo(explain, "tbl2", 1, false, true); + + JsonRenderer renderer = new JsonRenderer(getQueryRunner().getMetadata().getFunctionAndTypeManager()); + + List> fragmentsList = renderer.deserialize((String) computeActual(materializedSession,"EXPLAIN ANALYZE (format JSON) SELECT * FROM orders").getOnlyValue()); + fragmentsList.forEach(fragments -> { + fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan())); + }); + + fragmentsList = renderer.deserialize((String) computeActual(materializedSession, "EXPLAIN ANALYZE (format JSON) SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders").getOnlyValue()); + fragmentsList.forEach(fragments -> { + fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan())); + }); + + fragmentsList = renderer.deserialize((String) computeActual(materializedSession, "EXPLAIN ANALYZE (format JSON) SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders WHERE orderkey < 0").getOnlyValue()); + fragmentsList.forEach(fragments -> { + fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan())); + }); } // Hive specific tests should normally go in TestHiveIntegrationSmokeTest 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..37b106f292a12 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 @@ -149,7 +149,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 +156,17 @@ 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()); + plan = jsonDistributedPlan(queryInfo.getOutputStage().get().getSubStages(), functionAndTypeManager, + operatorContext.getSession()); 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 index aeca329baae30..570d464577f56 100644 --- 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 @@ -39,9 +39,10 @@ public class JsonRenderer implements Renderer { + private final JsonCodec>> planListCodec; private final JsonCodec> planMapCodec; private final JsonCodec codec; - private final JsonCodec> deserializationCodec; + private final JsonCodec>> deserializationCodec; public JsonRenderer(FunctionAndTypeManager functionAndTypeManager) { @@ -52,10 +53,11 @@ 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.planListCodec = codecFactory.listJsonCodec(codecFactory.mapJsonCodec(PlanFragmentId.class, JsonPlanFragment.class)); + this.deserializationCodec = codecFactory.listJsonCodec(codecFactory.mapJsonCodec(PlanFragmentId.class, JsonPlan.class)); } - public Map deserialize(String serialized) + public List> deserialize(String serialized) { return deserializationCodec.fromJson(serialized); } @@ -70,6 +72,10 @@ public String render(Map fragmentJsonMap) { return planMapCodec.toJson(fragmentJsonMap); } + public String render(List> fragmentJsonMap) + { + return planListCodec.toJson(fragmentJsonMap); + } @VisibleForTesting public JsonRenderedNode renderJson(PlanRepresentation plan, NodeRepresentation node) diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java index 77dff87b9b330..e8280138b4b0c 100644 --- a/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java +++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java @@ -326,6 +326,21 @@ public static String jsonLogicalPlan( return new PlanPrinter(plan, types, stageExecutionStrategy, estimatedStatsAndCosts, stats, functionAndTypeManager, session).toJson(); } + public static String jsonDistributedPlan(List subStages, FunctionAndTypeManager functionAndTypeManager, Session session) + { + ImmutableList.Builder> stageListBuilder = ImmutableList.builder(); + for (StageInfo outputStageInfo : subStages) { + List allStages = getAllStages(Optional.of(outputStageInfo)); + Map aggregatedStats = aggregateStageStats(allStages); + List allFragments = getAllStages(Optional.of(outputStageInfo)).stream() + .map(StageInfo::getPlan) + .map(Optional::get) + .collect(toImmutableList()); + stageListBuilder.add(fragmentMapBuilder(allFragments, Optional.of(aggregatedStats), functionAndTypeManager, session)); + } + return new JsonRenderer(functionAndTypeManager).render(stageListBuilder.build()); + } + public static String jsonDistributedPlan(StageInfo outputStageInfo, FunctionAndTypeManager functionAndTypeManager, Session session) { List allStages = getAllStages(Optional.of(outputStageInfo)); @@ -334,12 +349,12 @@ public static String jsonDistributedPlan(StageInfo outputStageInfo, FunctionAndT .map(StageInfo::getPlan) .map(Optional::get) .collect(toImmutableList()); - return formatJsonFragmentList(allFragments, Optional.of(aggregatedStats), functionAndTypeManager, session); + return new JsonRenderer(functionAndTypeManager).render(fragmentMapBuilder(allFragments, Optional.of(aggregatedStats), functionAndTypeManager, session)); } public static String jsonDistributedPlan(SubPlan plan, FunctionAndTypeManager functionAndTypeManager, Session session) { - return formatJsonFragmentList(plan.getAllFragments(), Optional.empty(), functionAndTypeManager, session); + return new JsonRenderer(functionAndTypeManager).render(fragmentMapBuilder(plan.getAllFragments(), Optional.empty(), functionAndTypeManager, session)); } private String formatSourceLocation(Optional sourceLocation1, Optional sourceLocation2) @@ -360,7 +375,7 @@ private String formatSourceLocation(Optional sourceLocation) return ""; } - private static String formatJsonFragmentList(List fragments, Optional> executionStats, FunctionAndTypeManager functionAndTypeManager, Session session) + private static Map fragmentMapBuilder(List fragments, Optional> executionStats, FunctionAndTypeManager functionAndTypeManager, Session session) { ImmutableSortedMap.Builder fragmentJsonMap = ImmutableSortedMap.naturalOrder(); for (PlanFragment fragment : fragments) { @@ -370,7 +385,7 @@ private static String formatJsonFragmentList(List fragments, Optio JsonPlanFragment jsonPlanFragment = new JsonPlanFragment(printer.toJson()); fragmentJsonMap.put(fragmentId, jsonPlanFragment); } - return new JsonRenderer(functionAndTypeManager).render(fragmentJsonMap.build()); + return fragmentJsonMap.build(); } private static String formatFragment( 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 index 76197a25dae8a..c420361966c1c 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java @@ -370,7 +370,7 @@ public void testExplainAnalyzeVerbose() assertExplainAnalyze("EXPLAIN ANALYZE VERBOSE SELECT rank() OVER (PARTITION BY orderkey ORDER BY clerk DESC) FROM orders WHERE orderkey < 0"); } - private static void assertJsonNodesHaveStats(JsonRenderer.JsonRenderedNode node) + public static void assertJsonNodesHaveStats(JsonRenderer.JsonRenderedNode node) { assertTrue(node.getStats().isPresent()); node.getChildren().forEach(AbstractTestDistributedQueries::assertJsonNodesHaveStats); @@ -380,12 +380,42 @@ private static void assertJsonNodesHaveStats(JsonRenderer.JsonRenderedNode node) public void testExplainAnalyzeFormatJson() { JsonRenderer renderer = new JsonRenderer(getQueryRunner().getMetadata().getFunctionAndTypeManager()); - Map 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> fragmentsList = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT * FROM orders").getOnlyValue()); + fragmentsList.forEach(fragments -> { + 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()); + fragmentsList.forEach(fragments -> { + 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()); + fragmentsList.forEach(fragments -> { + fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan())); + }); + } + + @Test + public void testCTEEnabledExplainAnalyzeFormatJson() + { + JsonRenderer renderer = new JsonRenderer(getQueryRunner().getMetadata().getFunctionAndTypeManager()); + + List> fragmentsList = renderer.deserialize((String) computeActual("EXPLAIN ANALYZE (format JSON) SELECT * FROM orders").getOnlyValue()); + fragmentsList.forEach(fragments -> { + 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()); + fragmentsList.forEach(fragments -> { + 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()); + fragmentsList.forEach(fragments -> { + fragments.values().forEach(planFragment -> assertJsonNodesHaveStats(planFragment.getPlan())); + }); } @Test(expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "EXPLAIN ANALYZE doesn't support statement type: DropTable")