Support format arguments in EXPLAIN ANALYZE

presto-docs/src/main/sphinx/sql/explain-analyze.rst

@@ -7,7 +7,7 @@ Synopsis
 
 .. code-block:: none
 
-    EXPLAIN ANALYZE [VERBOSE] statement
+    EXPLAIN ANALYZE [VERBOSE] [(format <TEXT|JSON>)] statement
 
 Description
 -----------
@@ -17,6 +17,8 @@ along with the cost of each operation.
 
 The ``VERBOSE`` option will give more detailed information and low-level statistics;
 understanding these may require knowledge of Presto internals and implementation details.
+The format of the output can be set by the user with the ``format`` option. The default
+output format is ``TEXT``.
 
 .. note::
 

presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java

@@ -472,7 +472,8 @@ private Optional<String> createJsonQueryPlan(QueryInfo queryInfo)
             if (queryInfo.getOutputStage().isPresent()) {
                 return Optional.of(jsonDistributedPlan(
                         queryInfo.getOutputStage().get(),
-                        functionAndTypeManager));
+                        functionAndTypeManager,
+                        queryInfo.getSession().toSession(sessionPropertyManager)));
             }
         }
         catch (Exception e) {

presto-main/src/main/java/com/facebook/presto/operator/ExplainAnalyzeOperator.java

@@ -20,13 +20,17 @@
 import com.facebook.presto.execution.StageId;
 import com.facebook.presto.execution.StageInfo;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
+import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.plan.PlanNodeId;
+import com.facebook.presto.sql.tree.ExplainFormat;
 import com.google.common.collect.ImmutableList;
 
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
+import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
+import static com.facebook.presto.sql.planner.planPrinter.PlanPrinter.jsonDistributedPlan;
 import static com.facebook.presto.sql.planner.planPrinter.PlanPrinter.textDistributedPlan;
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
@@ -42,28 +46,31 @@ public static class ExplainAnalyzeOperatorFactory
         private final QueryPerformanceFetcher queryPerformanceFetcher;
         private final FunctionAndTypeManager functionAndTypeManager;
         private final boolean verbose;
+        private final ExplainFormat.Type format;
         private boolean closed;
 
         public ExplainAnalyzeOperatorFactory(
                 int operatorId,
                 PlanNodeId planNodeId,
                 QueryPerformanceFetcher queryPerformanceFetcher,
                 FunctionAndTypeManager functionAndTypeManager,
-                boolean verbose)
+                boolean verbose,
+                ExplainFormat.Type format)
         {
             this.operatorId = operatorId;
             this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
             this.queryPerformanceFetcher = requireNonNull(queryPerformanceFetcher, "queryPerformanceFetcher is null");
             this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionManager is null");
             this.verbose = verbose;
+            this.format = requireNonNull(format, "format is null");
         }
 
         @Override
         public Operator createOperator(DriverContext driverContext)
         {
             checkState(!closed, "Factory is already closed");
             OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, ExplainAnalyzeOperator.class.getSimpleName());
-            return new ExplainAnalyzeOperator(operatorContext, queryPerformanceFetcher, functionAndTypeManager, verbose);
+            return new ExplainAnalyzeOperator(operatorContext, queryPerformanceFetcher, functionAndTypeManager, verbose, format);
         }
 
         @Override
@@ -75,27 +82,30 @@ public void noMoreOperators()
         @Override
         public OperatorFactory duplicate()
         {
-            return new ExplainAnalyzeOperatorFactory(operatorId, planNodeId, queryPerformanceFetcher, functionAndTypeManager, verbose);
+            return new ExplainAnalyzeOperatorFactory(operatorId, planNodeId, queryPerformanceFetcher, functionAndTypeManager, verbose, format);
         }
     }
 
     private final OperatorContext operatorContext;
     private final QueryPerformanceFetcher queryPerformanceFetcher;
     private final FunctionAndTypeManager functionAndTypeManager;
     private final boolean verbose;
+    private final ExplainFormat.Type format;
     private boolean finishing;
     private boolean outputConsumed;
 
     public ExplainAnalyzeOperator(
             OperatorContext operatorContext,
             QueryPerformanceFetcher queryPerformanceFetcher,
             FunctionAndTypeManager functionAndTypeManager,
-            boolean verbose)
+            boolean verbose,
+            ExplainFormat.Type format)
     {
         this.operatorContext = requireNonNull(operatorContext, "operatorContext is null");
         this.queryPerformanceFetcher = requireNonNull(queryPerformanceFetcher, "queryPerformanceFetcher is null");
         this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionManager is null");
         this.verbose = verbose;
+        this.format = requireNonNull(format, "format is null");
     }
 
     @Override
@@ -144,8 +154,18 @@ public Page getOutput()
         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);
+                break;
+            case JSON:
+                plan = jsonDistributedPlan(queryInfo.getOutputStage().get().getSubStages().get(0), functionAndTypeManager, operatorContext.getSession());
+                break;
+            default:
+                throw new PrestoException(GENERIC_INTERNAL_ERROR, "Explain format not supported: " + format);
+        }
 
-        String plan = textDistributedPlan(queryInfo.getOutputStage().get().getSubStages().get(0), functionAndTypeManager, operatorContext.getSession(), verbose);
         BlockBuilder builder = VARCHAR.createBlockBuilder(null, 1);
         VARCHAR.writeString(builder, plan);
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/queryplan/JsonPrestoQueryPlanFunctions.java

@@ -38,6 +38,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -197,7 +198,7 @@ private static JsonRenderedNode scrubJsonPlan(JsonRenderedNode node)
         List<JsonRenderedNode> newChildren = node.getChildren().stream().map(x -> scrubJsonPlan(x)).collect(Collectors.toList());
         String newDetails = scrubDetails(node.getDetails());
 
-        return new JsonRenderedNode(node.getSourceLocation(), "PLANID", newName, newIdentifier, newDetails, newChildren, node.getRemoteSources(), ImmutableList.of());
+        return new JsonRenderedNode(node.getSourceLocation(), "PLANID", newName, newIdentifier, newDetails, newChildren, node.getRemoteSources(), ImmutableList.of(), Optional.empty());
     }
 
     private static String scrubName(String name)

presto-main/src/main/java/com/facebook/presto/sql/analyzer/QueryExplainer.java

@@ -186,7 +186,7 @@ public String getJsonPlan(Session session, Statement statement, Type planType, L
                 return jsonLogicalPlan(plan.getRoot(), plan.getTypes(), metadata.getFunctionAndTypeManager(), plan.getStatsAndCosts(), session);
             case DISTRIBUTED:
                 SubPlan subPlan = getDistributedPlan(session, statement, parameters, warningCollector);
-                return jsonDistributedPlan(subPlan, metadata.getFunctionAndTypeManager());
+                return jsonDistributedPlan(subPlan, metadata.getFunctionAndTypeManager(), session);
             default:
                 throw new PrestoException(NOT_SUPPORTED, format("Unsupported explain plan type %s for JSON format", planType));
         }

presto-main/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java

@@ -94,6 +94,7 @@
 import com.facebook.presto.sql.tree.Except;
 import com.facebook.presto.sql.tree.Execute;
 import com.facebook.presto.sql.tree.Explain;
+import com.facebook.presto.sql.tree.ExplainFormat;
 import com.facebook.presto.sql.tree.ExplainType;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.ExpressionRewriter;
@@ -273,6 +274,8 @@
 import static com.facebook.presto.sql.planner.ExpressionDeterminismEvaluator.isDeterministic;
 import static com.facebook.presto.sql.planner.ExpressionInterpreter.evaluateConstantExpression;
 import static com.facebook.presto.sql.planner.ExpressionInterpreter.expressionOptimizer;
+import static com.facebook.presto.sql.tree.ExplainFormat.Type.JSON;
+import static com.facebook.presto.sql.tree.ExplainFormat.Type.TEXT;
 import static com.facebook.presto.sql.tree.ExplainType.Type.DISTRIBUTED;
 import static com.facebook.presto.sql.tree.FrameBound.Type.CURRENT_ROW;
 import static com.facebook.presto.sql.tree.FrameBound.Type.FOLLOWING;
@@ -1135,9 +1138,20 @@ protected Scope visitExplain(Explain node, Optional<Scope> scope)
                 throws SemanticException
         {
             checkState(node.isAnalyze(), "Non analyze explain should be rewritten to Query");
-            if (node.getOptions().stream().anyMatch(option -> !option.equals(new ExplainType(DISTRIBUTED)))) {
-                throw new SemanticException(NOT_SUPPORTED, node, "EXPLAIN ANALYZE only supports TYPE DISTRIBUTED option");
-            }
+            List<ExplainFormat.Type> formats = node.getOptions().stream()
+                    .filter(option -> option instanceof ExplainFormat)
+                    .map(ExplainFormat.class::cast)
+                    .map(ExplainFormat::getType)
+                    .collect(Collectors.toList());
+            checkState(formats.size() <= 1, "only a single format option is supported in EXPLAIN ANALYZE");
+            formats.stream().findFirst().ifPresent(format -> checkState(format.equals(TEXT) || format.equals(JSON),
+                    "only TEXT and JSON formats are supported in EXPLAIN ANALYZE"));
+            checkState(node.getOptions().stream()
+                    .filter(option -> option instanceof ExplainType)
+                    .findFirst()
+                    .map(ExplainType.class::cast)
+                    .map(ExplainType::getType)
+                    .orElse(DISTRIBUTED).equals(DISTRIBUTED), "only DISTRIBUTED type is supported in EXPLAIN ANALYZE");
             process(node.getStatement(), scope);
             analysis.setUpdateType(null);
             return createAndAssignScope(node, scope, Field.newUnqualified(node.getLocation(), "Query Plan", VARCHAR));

presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java

@@ -998,7 +998,8 @@ public PhysicalOperation visitExplainAnalyze(ExplainAnalyzeNode node, LocalExecu
                     node.getId(),
                     analyzeContext.getQueryPerformanceFetcher(),
                     metadata.getFunctionAndTypeManager(),
-                    node.isVerbose());
+                    node.isVerbose(),
+                    node.getFormat());
             return new PhysicalOperation(operatorFactory, makeLayout(node), context, source);
         }
 

presto-main/src/main/java/com/facebook/presto/sql/planner/LogicalPlanner.java

@@ -59,6 +59,7 @@
 import com.facebook.presto.sql.tree.CreateTableAsSelect;
 import com.facebook.presto.sql.tree.Delete;
 import com.facebook.presto.sql.tree.Explain;
+import com.facebook.presto.sql.tree.ExplainFormat;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.Identifier;
 import com.facebook.presto.sql.tree.Insert;
@@ -101,6 +102,7 @@
 import static com.facebook.presto.sql.planner.plan.TableWriterNode.UpdateTarget;
 import static com.facebook.presto.sql.planner.plan.TableWriterNode.WriterTarget;
 import static com.facebook.presto.sql.relational.Expressions.constant;
+import static com.facebook.presto.sql.tree.ExplainFormat.Type.TEXT;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -196,7 +198,12 @@ private RelationPlan createExplainAnalyzePlan(Analysis analysis, Explain stateme
         PlanNode root = underlyingPlan.getRoot();
         Scope scope = analysis.getScope(statement);
         VariableReferenceExpression outputVariable = newVariable(variableAllocator, scope.getRelationType().getFieldByIndex(0));
-        root = new ExplainAnalyzeNode(getSourceLocation(statement), idAllocator.getNextId(), root, outputVariable, statement.isVerbose());
+        ExplainFormat.Type type = statement.getOptions()
+                .stream().filter(option -> option instanceof ExplainFormat)
+                .findFirst().map(ExplainFormat.class::cast)
+                .map(ExplainFormat::getType)
+                .orElse(TEXT);
+        root = new ExplainAnalyzeNode(getSourceLocation(statement), idAllocator.getNextId(), root, outputVariable, statement.isVerbose(), type);
         return new RelationPlan(root, scope, ImmutableList.of(outputVariable));
     }
 

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java

@@ -215,7 +215,7 @@ public PlanNode visitGroupId(GroupIdNode node, RewriteContext<Void> context)
         public PlanNode visitExplainAnalyze(ExplainAnalyzeNode node, RewriteContext<Void> context)
         {
             PlanNode source = context.rewrite(node.getSource());
-            return new ExplainAnalyzeNode(node.getSourceLocation(), node.getId(), source, canonicalize(node.getOutputVariable()), node.isVerbose());
+            return new ExplainAnalyzeNode(node.getSourceLocation(), node.getId(), source, canonicalize(node.getOutputVariable()), node.isVerbose(), node.getFormat());
         }
 
         @Override

presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ExplainAnalyzeNode.java

@@ -17,6 +17,7 @@
 import com.facebook.presto.spi.plan.PlanNode;
 import com.facebook.presto.spi.plan.PlanNodeId;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.facebook.presto.sql.tree.ExplainFormat.Type;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableList;
@@ -36,16 +37,18 @@ public class ExplainAnalyzeNode
     private final PlanNode source;
     private final VariableReferenceExpression outputVariable;
     private final boolean verbose;
+    private final Type format;
 
     @JsonCreator
     public ExplainAnalyzeNode(
             Optional<SourceLocation> sourceLocation,
             @JsonProperty("id") PlanNodeId id,
             @JsonProperty("source") PlanNode source,
             @JsonProperty("outputVariable") VariableReferenceExpression outputVariable,
-            @JsonProperty("verbose") boolean verbose)
+            @JsonProperty("verbose") boolean verbose,
+            @JsonProperty("format") Type format)
     {
-        this(sourceLocation, id, Optional.empty(), source, outputVariable, verbose);
+        this(sourceLocation, id, Optional.empty(), source, outputVariable, verbose, format);
     }
 
     public ExplainAnalyzeNode(
@@ -54,12 +57,14 @@ public ExplainAnalyzeNode(
             Optional<PlanNode> statsEquivalentPlanNode,
             PlanNode source,
             VariableReferenceExpression outputVariable,
-            boolean verbose)
+            boolean verbose,
+            Type format)
     {
         super(sourceLocation, id, statsEquivalentPlanNode);
         this.source = requireNonNull(source, "source is null");
         this.outputVariable = requireNonNull(outputVariable, "outputVariable is null");
         this.verbose = verbose;
+        this.format = requireNonNull(format, "options is null");
     }
 
     @JsonProperty("outputVariable")
@@ -80,6 +85,12 @@ public boolean isVerbose()
         return verbose;
     }
 
+    @JsonProperty("format")
+    public Type getFormat()
+    {
+        return format;
+    }
+
     @Override
     public List<VariableReferenceExpression> getOutputVariables()
     {
@@ -101,12 +112,12 @@ public <R, C> R accept(InternalPlanVisitor<R, C> visitor, C context)
     @Override
     public PlanNode replaceChildren(List<PlanNode> newChildren)
     {
-        return new ExplainAnalyzeNode(getSourceLocation(), getId(), getStatsEquivalentPlanNode(), Iterables.getOnlyElement(newChildren), outputVariable, isVerbose());
+        return new ExplainAnalyzeNode(getSourceLocation(), getId(), getStatsEquivalentPlanNode(), Iterables.getOnlyElement(newChildren), outputVariable, isVerbose(), format);
     }
 
     @Override
     public PlanNode assignStatsEquivalentPlanNode(Optional<PlanNode> statsEquivalentPlanNode)
     {
-        return new ExplainAnalyzeNode(getSourceLocation(), getId(), statsEquivalentPlanNode, source, outputVariable, isVerbose());
+        return new ExplainAnalyzeNode(getSourceLocation(), getId(), statsEquivalentPlanNode, source, outputVariable, isVerbose(), format);
     }
 }

presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/HashCollisionPlanNodeStats.java

@@ -36,12 +36,17 @@ public HashCollisionPlanNodeStats(
             PlanNodeId planNodeId,
             Duration planNodeScheduledTime,
             Duration planNodeCpuTime,
+            Duration planNodeBlockedWallTime,
+            Duration planNodeAddInputWallTime,
+            Duration planNodeGetOutputWallTime,
+            Duration planNodeFinishWallTime,
             long planNodeInputPositions,
             DataSize planNodeInputDataSize,
             long planNodeRawInputPositions,
             DataSize planNodeRawInputDataSize,
             long planNodeOutputPositions,
             DataSize planNodeOutputDataSize,
+            DataSize planNodePeakMemorySize,
             Map<String, OperatorInputStats> operatorInputStats,
             long planNodeNullJoinBuildKeyCount,
             long planNodeJoinBuildKeyCount,
@@ -50,9 +55,9 @@ public HashCollisionPlanNodeStats(
             Optional<DynamicFilterStats> dynamicFilterStats,
             Map<String, OperatorHashCollisionsStats> operatorHashCollisionsStats)
     {
-        super(planNodeId, planNodeScheduledTime, planNodeCpuTime, planNodeInputPositions, planNodeInputDataSize, planNodeRawInputPositions, planNodeRawInputDataSize,
-                planNodeOutputPositions, planNodeOutputDataSize, operatorInputStats, planNodeNullJoinBuildKeyCount, planNodeJoinBuildKeyCount, planNodeNullJoinProbeKeyCount,
-                planNodeJoinProbeKeyCount, dynamicFilterStats);
+        super(planNodeId, planNodeScheduledTime, planNodeCpuTime, planNodeBlockedWallTime, planNodeAddInputWallTime, planNodeGetOutputWallTime, planNodeFinishWallTime,
+                planNodeInputPositions, planNodeInputDataSize, planNodeRawInputPositions, planNodeRawInputDataSize, planNodeOutputPositions, planNodeOutputDataSize,
+                planNodePeakMemorySize, operatorInputStats, planNodeNullJoinBuildKeyCount, planNodeJoinBuildKeyCount, planNodeNullJoinProbeKeyCount, planNodeJoinProbeKeyCount, dynamicFilterStats);
         this.operatorHashCollisionsStats = requireNonNull(operatorHashCollisionsStats, "operatorHashCollisionsStats is null");
     }
 
@@ -101,12 +106,17 @@ public PlanNodeStats mergeWith(PlanNodeStats other)
                 merged.getPlanNodeId(),
                 merged.getPlanNodeScheduledTime(),
                 merged.getPlanNodeCpuTime(),
+                merged.getPlanNodeBlockedWallTime(),
+                merged.getPlanNodeAddInputWallTime(),
+                merged.getPlanNodeGetOutputWallTime(),
+                merged.getPlanNodeFinishWallTime(),
                 merged.getPlanNodeInputPositions(),
                 merged.getPlanNodeInputDataSize(),
                 merged.getPlanNodeRawInputPositions(),
                 merged.getPlanNodeRawInputDataSize(),
                 merged.getPlanNodeOutputPositions(),
                 merged.getPlanNodeOutputDataSize(),
+                merged.getPlanNodePeakMemorySize(),
                 merged.operatorInputStats,
                 merged.getPlanNodeNullJoinBuildKeyCount(),
                 merged.getPlanNodeJoinBuildKeyCount(),