Skip to content

Commit

Permalink
Modified ExplainAnalyzeOperator to support multiple substages in output
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@infvg
infvg committed Oct 16, 2024
1 parent 7deadd1 commit f5f0f20
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 17 deletions.
21 changes: 21 additions & 0 deletions presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,17 @@
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;
import com.google.common.collect.ImmutableSet;
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;

Expand Down Expand Up @@ -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<Map<PlanFragmentId, JsonRenderer.JsonPlan>> 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
Expand Down
12 changes: 9 additions & 3 deletions presto-main/src/main/java/com/facebook/presto/operator/ExplainAnalyzeOperator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,18 +149,24 @@ 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());
plan = jsonDistributedPlan(queryInfo.getOutputStage().get().getSubStages(), functionAndTypeManager,
operatorContext.getSession());
break;
default:
throw new PrestoException(GENERIC_INTERNAL_ERROR, "Explain format not supported: " + format);
Expand Down
12 changes: 9 additions & 3 deletions presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/JsonRenderer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,9 +39,10 @@
public class JsonRenderer
implements Renderer<String>
{
private final JsonCodec<List<Map<PlanFragmentId, JsonPlanFragment>>> planListCodec;
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)
{
Expand All @@ -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<PlanFragmentId, JsonPlan> deserialize(String serialized)
public List<Map<PlanFragmentId, JsonPlan>> deserialize(String serialized)
{
return deserializationCodec.fromJson(serialized);
}
Expand All @@ -70,6 +72,10 @@ public String render(Map<PlanFragmentId, JsonPlanFragment> fragmentJsonMap)
{
return planMapCodec.toJson(fragmentJsonMap);
}
public String render(List<Map<PlanFragmentId, JsonPlanFragment>> fragmentJsonMap)
{
return planListCodec.toJson(fragmentJsonMap);
}

@VisibleForTesting
public JsonRenderedNode renderJson(PlanRepresentation plan, NodeRepresentation node)
Expand Down
23 changes: 19 additions & 4 deletions presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,21 @@ public static String jsonLogicalPlan(
return new PlanPrinter(plan, types, stageExecutionStrategy, estimatedStatsAndCosts, stats, functionAndTypeManager, session).toJson();
}

public static String jsonDistributedPlan(List<StageInfo> subStages, FunctionAndTypeManager functionAndTypeManager, Session session)
{
ImmutableList.Builder<Map<PlanFragmentId, JsonPlanFragment>> stageListBuilder = ImmutableList.builder();
for (StageInfo outputStageInfo : subStages) {
List<StageInfo> allStages = getAllStages(Optional.of(outputStageInfo));
Map<PlanNodeId, PlanNodeStats> aggregatedStats = aggregateStageStats(allStages);
List<PlanFragment> 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<StageInfo> allStages = getAllStages(Optional.of(outputStageInfo));
Expand All @@ -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<SourceLocation> sourceLocation1, Optional<SourceLocation> sourceLocation2)
Expand All @@ -360,7 +375,7 @@ private String formatSourceLocation(Optional<SourceLocation> sourceLocation)
return "";
}

private static String formatJsonFragmentList(List<PlanFragment> fragments, Optional<Map<PlanNodeId, PlanNodeStats>> executionStats, FunctionAndTypeManager functionAndTypeManager, Session session)
private static Map<PlanFragmentId, JsonPlanFragment> fragmentMapBuilder(List<PlanFragment> fragments, Optional<Map<PlanNodeId, PlanNodeStats>> executionStats, FunctionAndTypeManager functionAndTypeManager, Session session)
{
ImmutableSortedMap.Builder<PlanFragmentId, JsonPlanFragment> fragmentJsonMap = ImmutableSortedMap.naturalOrder();
for (PlanFragment fragment : fragments) {
Expand All @@ -370,7 +385,7 @@ private static String formatJsonFragmentList(List<PlanFragment> 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(
Expand Down
44 changes: 37 additions & 7 deletions presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -380,12 +380,42 @@ 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());
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<Map<PlanFragmentId, JsonRenderer.JsonPlan>> 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")
Expand Down

0 comments on commit f5f0f20

Please sign in to comment.