Enable Estimated Stats on Query Plan for Instrumentation Monitoring

presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java

@@ -4302,7 +4302,7 @@ public static Consumer<Plan> assertRemoteExchangesCount(int expectedRemoteExchan
                     .size();
             if (actualRemoteExchangesCount != expectedRemoteExchangesCount) {
                 Metadata metadata = queryRunner.getCoordinator().getMetadata();
-                String formattedPlan = textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata.getFunctionAndTypeManager(), StatsAndCosts.empty(), session, 0);
+                String formattedPlan = textLogicalPlan(plan.getRoot(), plan.getTypes(), StatsAndCosts.empty(), metadata.getFunctionAndTypeManager(), session, 0);
                 throw new AssertionError(format(
                         "Expected [\n%s\n] remote exchanges but found [\n%s\n] remote exchanges. Actual plan is [\n\n%s\n]",
                         expectedRemoteExchangesCount,

presto-main/src/main/java/com/facebook/presto/cost/PlanNodeStatsEstimate.java

@@ -24,8 +24,10 @@
 import com.facebook.presto.spi.statistics.PlanStatistics;
 import com.facebook.presto.spi.statistics.PlanStatisticsWithSourceInfo;
 import com.facebook.presto.spi.statistics.SourceInfo;
+import com.facebook.presto.sql.Serialization;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import com.google.common.collect.ImmutableMap;
 import org.pcollections.HashTreePMap;
 import org.pcollections.PMap;
@@ -194,6 +196,7 @@ public VariableStatsEstimate getVariableStatistics(VariableReferenceExpression v
         return variableStatistics.getOrDefault(variable, VariableStatsEstimate.unknown());
     }
 
+    @JsonSerialize(keyUsing = Serialization.VariableReferenceExpressionSerializer.class)
     @JsonProperty
     public Map<VariableReferenceExpression, VariableStatsEstimate> getVariableStatistics()
     {

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

@@ -75,6 +75,7 @@
 
 import static com.facebook.presto.execution.QueryState.QUEUED;
 import static com.facebook.presto.execution.StageInfo.getAllStages;
+import static com.facebook.presto.sql.planner.planPrinter.PlanPrinter.graphvizDistributedPlan;
 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.collect.ImmutableList.toImmutableList;
@@ -144,6 +145,7 @@ public void queryCreatedEvent(BasicQueryInfo queryInfo)
                                 Optional.empty(),
                                 Optional.empty(),
                                 Optional.empty(),
+                                Optional.empty(),
                                 ImmutableList.of(),
                                 queryInfo.getSession().getTraceToken())));
     }
@@ -167,6 +169,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         Optional.empty(),
                         Optional.empty(),
                         Optional.empty(),
+                        Optional.empty(),
                         ImmutableList.of(),
                         queryInfo.getSession().getTraceToken()),
                 new QueryStatistics(
@@ -270,6 +273,7 @@ private QueryMetadata createQueryMetadata(QueryInfo queryInfo)
                 queryInfo.getSelf(),
                 createTextQueryPlan(queryInfo),
                 createJsonQueryPlan(queryInfo),
+                createGraphvizQueryPlan(queryInfo),
                 queryInfo.getOutputStage().flatMap(stage -> stageInfoCodec.toJsonWithLengthLimit(stage, maxJsonLimit)),
                 queryInfo.getRuntimeOptimizedStages().orElse(ImmutableList.of()).stream()
                         .map(stageId -> String.valueOf(stageId.getId()))
@@ -403,7 +407,8 @@ private Optional<String> createJsonQueryPlan(QueryInfo queryInfo)
         try {
             if (queryInfo.getOutputStage().isPresent()) {
                 return Optional.of(jsonDistributedPlan(
-                        queryInfo.getOutputStage().get()));
+                        queryInfo.getOutputStage().get(),
+                        functionAndTypeManager));
             }
         }
         catch (Exception e) {
@@ -413,6 +418,23 @@ private Optional<String> createJsonQueryPlan(QueryInfo queryInfo)
         return Optional.empty();
     }
 
+    private Optional<String> createGraphvizQueryPlan(QueryInfo queryInfo)
+    {
+        try {
+            if (queryInfo.getOutputStage().isPresent()) {
+                return Optional.of(graphvizDistributedPlan(
+                        queryInfo.getOutputStage().get(),
+                        functionAndTypeManager,
+                        queryInfo.getSession().toSession(sessionPropertyManager)));
+            }
+        }
+        catch (Exception e) {
+            // Don't fail to create event if the graphviz plan can not be created
+            log.warn(e, "Error creating graphviz plan for query %s: %s", queryInfo.getQueryId(), e);
+        }
+        return Optional.empty();
+    }
+
     private static QueryIOMetadata getQueryIOMetadata(QueryInfo queryInfo)
     {
         ImmutableList.Builder<QueryInputMetadata> inputs = ImmutableList.builder();

presto-main/src/main/java/com/facebook/presto/execution/scheduler/LegacySqlQueryScheduler.java

@@ -17,6 +17,7 @@
 import com.facebook.airlift.log.Logger;
 import com.facebook.airlift.stats.TimeStat;
 import com.facebook.presto.Session;
+import com.facebook.presto.cost.StatsAndCosts;
 import com.facebook.presto.execution.BasicStageExecutionStats;
 import com.facebook.presto.execution.LocationFactory;
 import com.facebook.presto.execution.PartialResultQueryManager;
@@ -609,6 +610,7 @@ private Optional<PlanFragment> performRuntimeOptimizations(StreamingSubPlan subP
             newRoot = optimizer.optimize(newRoot, session, variableAllocator.getTypes(), variableAllocator, idAllocator, warningCollector);
         }
         if (newRoot != fragment.getRoot()) {
+            StatsAndCosts estimatedStatsAndCosts = fragment.getStatsAndCosts();
             return Optional.of(
                     // The partitioningScheme should stay the same
                     // even if the root's outputVariable layout is changed.
@@ -621,8 +623,8 @@ private Optional<PlanFragment> performRuntimeOptimizations(StreamingSubPlan subP
                             fragment.getPartitioningScheme(),
                             fragment.getStageExecutionDescriptor(),
                             fragment.isOutputTableWriterFragment(),
-                            fragment.getStatsAndCosts(),
-                            Optional.of(jsonFragmentPlan(newRoot, fragment.getVariables(), functionAndTypeManager, session))));
+                            estimatedStatsAndCosts,
+                            Optional.of(jsonFragmentPlan(newRoot, fragment.getVariables(), estimatedStatsAndCosts, functionAndTypeManager, session))));
         }
         return Optional.empty();
     }

presto-main/src/main/java/com/facebook/presto/execution/scheduler/SqlQueryScheduler.java

@@ -17,6 +17,7 @@
 import com.facebook.airlift.log.Logger;
 import com.facebook.airlift.stats.TimeStat;
 import com.facebook.presto.Session;
+import com.facebook.presto.cost.StatsAndCosts;
 import com.facebook.presto.execution.BasicStageExecutionStats;
 import com.facebook.presto.execution.ExecutionFailureInfo;
 import com.facebook.presto.execution.LocationFactory;
@@ -497,6 +498,7 @@ private Optional<PlanFragment> performRuntimeOptimizations(StreamingSubPlan subP
             newRoot = optimizer.optimize(newRoot, session, variableAllocator.getTypes(), variableAllocator, idAllocator, warningCollector);
         }
         if (newRoot != fragment.getRoot()) {
+            StatsAndCosts estimatedStatsAndCosts = fragment.getStatsAndCosts();
             return Optional.of(
                     // The partitioningScheme should stay the same
                     // even if the root's outputVariable layout is changed.
@@ -509,8 +511,8 @@ private Optional<PlanFragment> performRuntimeOptimizations(StreamingSubPlan subP
                             fragment.getPartitioningScheme(),
                             fragment.getStageExecutionDescriptor(),
                             fragment.isOutputTableWriterFragment(),
-                            fragment.getStatsAndCosts(),
-                            Optional.of(jsonFragmentPlan(newRoot, fragment.getVariables(), functionAndTypeManager, session))));
+                            estimatedStatsAndCosts,
+                            Optional.of(jsonFragmentPlan(newRoot, fragment.getVariables(), estimatedStatsAndCosts, functionAndTypeManager, session))));
         }
         return Optional.empty();
     }

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

@@ -127,7 +127,7 @@ public String getPlan(Session session, Statement statement, Type planType, List<
         switch (planType) {
             case LOGICAL:
                 Plan plan = getLogicalPlan(session, statement, parameters, warningCollector);
-                return PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata.getFunctionAndTypeManager(), plan.getStatsAndCosts(), session, 0, verbose);
+                return PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), plan.getStatsAndCosts(), metadata.getFunctionAndTypeManager(), session, 0, verbose);
             case DISTRIBUTED:
                 SubPlan subPlan = getDistributedPlan(session, statement, parameters, warningCollector);
                 return PlanPrinter.textDistributedPlan(subPlan, metadata.getFunctionAndTypeManager(), session, verbose);
@@ -153,10 +153,10 @@ public String getGraphvizPlan(Session session, Statement statement, Type planTyp
         switch (planType) {
             case LOGICAL:
                 Plan plan = getLogicalPlan(session, statement, parameters, warningCollector);
-                return graphvizLogicalPlan(plan.getRoot(), plan.getTypes(), session, metadata.getFunctionAndTypeManager());
+                return graphvizLogicalPlan(plan.getRoot(), plan.getTypes(), plan.getStatsAndCosts(), metadata.getFunctionAndTypeManager(), session);
             case DISTRIBUTED:
                 SubPlan subPlan = getDistributedPlan(session, statement, parameters, warningCollector);
-                return graphvizDistributedPlan(subPlan, session, metadata.getFunctionAndTypeManager());
+                return graphvizDistributedPlan(subPlan, metadata.getFunctionAndTypeManager(), session);
         }
         throw new IllegalArgumentException("Unhandled plan type: " + planType);
     }
@@ -179,7 +179,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);
+                return jsonDistributedPlan(subPlan, metadata.getFunctionAndTypeManager());
             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/planner/BasePlanFragmenter.java

@@ -193,7 +193,7 @@ private SubPlan buildFragment(PlanNode root, FragmentProperties properties, Plan
                 StageExecutionDescriptor.ungroupedExecution(),
                 outputTableWriterFragment,
                 statsAndCosts.getForSubplan(root),
-                Optional.of(jsonFragmentPlan(root, fragmentVariableTypes, metadata.getFunctionAndTypeManager(), session)));
+                Optional.of(jsonFragmentPlan(root, fragmentVariableTypes, statsAndCosts.getForSubplan(root), metadata.getFunctionAndTypeManager(), session)));
 
         return new SubPlan(fragment, properties.getChildren());
     }

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

@@ -14,14 +14,23 @@
 package com.facebook.presto.sql.planner.planPrinter;
 
 import com.facebook.airlift.json.JsonCodec;
+import com.facebook.airlift.json.JsonCodecFactory;
+import com.facebook.airlift.json.JsonObjectMapperProvider;
+import com.facebook.presto.cost.PlanNodeStatsEstimate;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.spi.SourceLocation;
+import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.facebook.presto.sql.Serialization;
 import com.facebook.presto.sql.planner.plan.PlanFragmentId;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonRawValue;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
 
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
@@ -30,21 +39,33 @@
 public class JsonRenderer
         implements Renderer<String>
 {
-    private static final JsonCodec<JsonRenderedNode> CODEC = JsonCodec.jsonCodec(JsonRenderedNode.class);
-    private static final JsonCodec<Map<PlanFragmentId, JsonPlanFragment>> PLAN_MAP_CODEC = JsonCodec.mapJsonCodec(PlanFragmentId.class, JsonPlanFragment.class);
+    private static JsonCodec<JsonRenderedNode> codec = JsonCodec.jsonCodec(JsonRenderedNode.class);
+    private static JsonCodec<Map<PlanFragmentId, JsonPlanFragment>> planMapCodec = JsonCodec.mapJsonCodec(PlanFragmentId.class, JsonPlanFragment.class);
+
+    public JsonRenderer(FunctionAndTypeManager functionAndTypeManager)
+    {
+        JsonObjectMapperProvider provider = new JsonObjectMapperProvider();
+        provider.setJsonSerializers(ImmutableMap.of(VariableReferenceExpression.class, new Serialization.VariableReferenceExpressionSerializer()));
+        provider.setKeyDeserializers(ImmutableMap.of(VariableReferenceExpression.class, new Serialization.VariableReferenceExpressionDeserializer(functionAndTypeManager)));
+
+        JsonCodecFactory codecFactory = new JsonCodecFactory(provider, true);
+        this.codec = codecFactory.jsonCodec(JsonRenderedNode.class);
+        this.planMapCodec = codecFactory.mapJsonCodec(PlanFragmentId.class, JsonPlanFragment.class);
+    }
 
     @Override
     public String render(PlanRepresentation plan)
     {
-        return CODEC.toJson(renderJson(plan, plan.getRoot()));
+        return codec.toJson(renderJson(plan, plan.getRoot()));
     }
 
     public String render(Map<PlanFragmentId, JsonPlanFragment> fragmentJsonMap)
     {
-        return PLAN_MAP_CODEC.toJson(fragmentJsonMap);
+        return planMapCodec.toJson(fragmentJsonMap);
     }
 
-    private JsonRenderedNode renderJson(PlanRepresentation plan, NodeRepresentation node)
+    @VisibleForTesting
+    public JsonRenderedNode renderJson(PlanRepresentation plan, NodeRepresentation node)
     {
         List<JsonRenderedNode> children = node.getChildren().stream()
                 .map(plan::getNode)
@@ -62,7 +83,8 @@ private JsonRenderedNode renderJson(PlanRepresentation plan, NodeRepresentation
                 children,
                 node.getRemoteSources().stream()
                         .map(PlanFragmentId::toString)
-                        .collect(toImmutableList()));
+                        .collect(toImmutableList()),
+                node.getEstimatedStats());
     }
 
     public static class JsonRenderedNode
@@ -74,9 +96,10 @@ public static class JsonRenderedNode
         private final String details;
         private final List<JsonRenderedNode> children;
         private final List<String> remoteSources;
+        private final List<PlanNodeStatsEstimate> estimates;
 
         @JsonCreator
-        public JsonRenderedNode(Optional<SourceLocation> sourceLocation, String id, String name, String identifier, String details, List<JsonRenderedNode> children, List<String> remoteSources)
+        public JsonRenderedNode(Optional<SourceLocation> sourceLocation, String id, String name, String identifier, String details, List<JsonRenderedNode> children, List<String> remoteSources, List<PlanNodeStatsEstimate> estimates)
         {
             this.sourceLocation = sourceLocation;
             this.id = requireNonNull(id, "id is null");
@@ -85,6 +108,7 @@ public JsonRenderedNode(Optional<SourceLocation> sourceLocation, String id, Stri
             this.details = requireNonNull(details, "details is null");
             this.children = requireNonNull(children, "children is null");
             this.remoteSources = requireNonNull(remoteSources, "id is null");
+            this.estimates = requireNonNull(estimates, "estimate is null");
         }
 
         @JsonProperty
@@ -128,6 +152,38 @@ public List<String> getRemoteSources()
         {
             return remoteSources;
         }
+
+        @JsonProperty
+        public List<PlanNodeStatsEstimate> getEstimates()
+        {
+            return estimates;
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            if (this == obj) {
+                return true;
+            }
+            if (obj == null || getClass() != obj.getClass()) {
+                return false;
+            }
+            JsonRenderedNode other = (JsonRenderedNode) obj;
+            return Objects.equals(this.name, other.name) &&
+                    Objects.equals(this.id, other.id) &&
+                    Objects.equals(this.sourceLocation, other.sourceLocation) &&
+                    Objects.equals(this.identifier, other.identifier) &&
+                    Objects.equals(this.details, other.details) &&
+                    Objects.equals(this.children, other.children) &&
+                    Objects.equals(this.estimates, other.estimates) &&
+                    Objects.equals(this.remoteSources, other.remoteSources);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(name, id, sourceLocation, identifier, details, children, estimates, remoteSources);
+        }
     }
 
     public static class JsonPlanFragment