Collect statistics for intermediate data writes separately

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

@@ -161,6 +161,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         0,
                         0,
                         0,
+                        0,
                         ImmutableList.of(),
                         0,
                         true,
@@ -233,6 +234,7 @@ private QueryStatistics createQueryStatistics(QueryInfo queryInfo)
                 queryStats.getOutputPositions(),
                 queryStats.getWrittenOutputLogicalDataSize().toBytes(),
                 queryStats.getWrittenOutputPositions(),
+                queryStats.getWrittenIntermediatePhysicalDataSize().toBytes(),
                 queryStats.getCumulativeUserMemory(),
                 queryStats.getStageGcStatistics(),
                 queryStats.getCompletedDrivers(),

presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java

@@ -439,6 +439,8 @@ private QueryStats getQueryStats(Optional<StageInfo> rootStage)
         long writtenOutputLogicalDataSize = 0;
         long writtenOutputPhysicalDataSize = 0;
 
+        long writtenIntermediatePhysicalDataSize = 0;
+
         ImmutableList.Builder<StageGcStatistics> stageGcStatistics = ImmutableList.builder();
 
         boolean fullyBlocked = rootStage.isPresent();
@@ -469,24 +471,31 @@ private QueryStats getQueryStats(Optional<StageInfo> rootStage)
                 blockedReasons.addAll(stageStats.getBlockedReasons());
             }
 
-            Optional<PlanFragment> plan = stageInfo.getPlan();
-            if (plan.isPresent() && plan.get().getPartitionedSourceNodes().stream().anyMatch(TableScanNode.class::isInstance)) {
-                rawInputDataSize += stageStats.getRawInputDataSize().toBytes();
-                rawInputPositions += stageStats.getRawInputPositions();
+            if (stageInfo.getPlan().isPresent()) {
+                PlanFragment plan = stageInfo.getPlan().get();
+                if (plan.getPartitionedSourceNodes().stream().anyMatch(TableScanNode.class::isInstance)) {
+                    rawInputDataSize += stageStats.getRawInputDataSize().toBytes();
+                    rawInputPositions += stageStats.getRawInputPositions();
 
-                processedInputDataSize += stageStats.getProcessedInputDataSize().toBytes();
-                processedInputPositions += stageStats.getProcessedInputPositions();
-            }
+                    processedInputDataSize += stageStats.getProcessedInputDataSize().toBytes();
+                    processedInputPositions += stageStats.getProcessedInputPositions();
+                }
 
-            writtenOutputPositions += stageInfo.getStageStats().getOperatorSummaries().stream()
-                    .filter(stats -> stats.getOperatorType().equals(TableWriterOperator.class.getSimpleName()))
-                    .mapToLong(OperatorStats::getInputPositions)
-                    .sum();
-            writtenOutputLogicalDataSize += stageInfo.getStageStats().getOperatorSummaries().stream()
-                    .filter(stats -> stats.getOperatorType().equals(TableWriterOperator.class.getSimpleName()))
-                    .mapToLong(stats -> stats.getInputDataSize().toBytes())
-                    .sum();
-            writtenOutputPhysicalDataSize += stageStats.getPhysicalWrittenDataSize().toBytes();
+                if (plan.isMaterializedExchangeSource()) {
+                    writtenOutputPhysicalDataSize += stageStats.getPhysicalWrittenDataSize().toBytes();
+                }
+                else {
+                    writtenOutputPositions += stageInfo.getStageStats().getOperatorSummaries().stream()
+                            .filter(stats -> stats.getOperatorType().equals(TableWriterOperator.class.getSimpleName()))
+                            .mapToLong(OperatorStats::getInputPositions)
+                            .sum();
+                    writtenOutputLogicalDataSize += stageInfo.getStageStats().getOperatorSummaries().stream()
+                            .filter(stats -> stats.getOperatorType().equals(TableWriterOperator.class.getSimpleName()))
+                            .mapToLong(stats -> stats.getInputDataSize().toBytes())
+                            .sum();
+                    writtenOutputPhysicalDataSize += stageStats.getPhysicalWrittenDataSize().toBytes();
+                }
+            }
 
             stageGcStatistics.add(stageStats.getGcInfo());
 
@@ -553,6 +562,8 @@ private QueryStats getQueryStats(Optional<StageInfo> rootStage)
                 succinctBytes(writtenOutputLogicalDataSize),
                 succinctBytes(writtenOutputPhysicalDataSize),
 
+                succinctBytes(writtenIntermediatePhysicalDataSize),
+
                 stageGcStatistics.build(),
 
                 operatorStatsSummary.build());
@@ -1019,6 +1030,7 @@ private static QueryStats pruneQueryStats(QueryStats queryStats)
                 queryStats.getWrittenOutputPositions(),
                 queryStats.getWrittenOutputLogicalDataSize(),
                 queryStats.getWrittenOutputPhysicalDataSize(),
+                queryStats.getWrittenIntermediatePhysicalDataSize(),
                 queryStats.getStageGcStatistics(),
                 ImmutableList.of()); // Remove the operator summaries as OperatorInfo (especially ExchangeClientStatus) can hold onto a large amount of memory
     }

presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java

@@ -91,6 +91,8 @@ public class QueryStats
     private final DataSize writtenOutputLogicalDataSize;
     private final DataSize writtenOutputPhysicalDataSize;
 
+    private final DataSize writtenIntermediatePhysicalDataSize;
+
     private final List<StageGcStatistics> stageGcStatistics;
 
     private final List<OperatorStats> operatorSummaries;
@@ -148,6 +150,8 @@ public QueryStats(
             @JsonProperty("writtenOutputLogicalDataSize") DataSize writtenOutputLogicalDataSize,
             @JsonProperty("writtenOutputPhysicalDataSize") DataSize writtenOutputPhysicalDataSize,
 
+            @JsonProperty("writtenIntermediatePhysicalDataSize") DataSize writtenIntermediatePhysicalDataSize,
+
             @JsonProperty("stageGcStatistics") List<StageGcStatistics> stageGcStatistics,
 
             @JsonProperty("operatorSummaries") List<OperatorStats> operatorSummaries)
@@ -213,6 +217,7 @@ public QueryStats(
         this.writtenOutputPositions = writtenOutputPositions;
         this.writtenOutputLogicalDataSize = requireNonNull(writtenOutputLogicalDataSize, "writtenOutputLogicalDataSize is null");
         this.writtenOutputPhysicalDataSize = requireNonNull(writtenOutputPhysicalDataSize, "writtenOutputPhysicalDataSize is null");
+        this.writtenIntermediatePhysicalDataSize = requireNonNull(writtenIntermediatePhysicalDataSize, "writtenIntermediatePhysicalDataSize is null");
 
         this.stageGcStatistics = ImmutableList.copyOf(requireNonNull(stageGcStatistics, "stageGcStatistics is null"));
 
@@ -264,6 +269,7 @@ public static QueryStats immediateFailureQueryStats()
                 0,
                 new DataSize(0, BYTE),
                 new DataSize(0, BYTE),
+                new DataSize(0, BYTE),
                 ImmutableList.of(),
                 ImmutableList.of());
     }
@@ -515,6 +521,12 @@ public DataSize getWrittenOutputPhysicalDataSize()
         return writtenOutputPhysicalDataSize;
     }
 
+    @JsonProperty
+    public DataSize getWrittenIntermediatePhysicalDataSize()
+    {
+        return writtenIntermediatePhysicalDataSize;
+    }
+
     @JsonProperty
     public List<StageGcStatistics> getStageGcStatistics()
     {

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

@@ -52,6 +52,7 @@ public class PlanFragment
     private final List<RemoteSourceNode> remoteSourceNodes;
     private final PartitioningScheme partitioningScheme;
     private final StageExecutionDescriptor stageExecutionDescriptor;
+    private final boolean materializedExchangeSource;
     private final StatsAndCosts statsAndCosts;
     private final Optional<String> jsonRepresentation;
 
@@ -64,6 +65,7 @@ public PlanFragment(
             @JsonProperty("partitionedSources") List<PlanNodeId> partitionedSources,
             @JsonProperty("partitioningScheme") PartitioningScheme partitioningScheme,
             @JsonProperty("stageExecutionDescriptor") StageExecutionDescriptor stageExecutionDescriptor,
+            @JsonProperty("materializedExchangeSource") boolean materializedExchangeSource,
             @JsonProperty("statsAndCosts") StatsAndCosts statsAndCosts,
             @JsonProperty("jsonRepresentation") Optional<String> jsonRepresentation)
     {
@@ -74,6 +76,7 @@ public PlanFragment(
         this.partitionedSources = ImmutableList.copyOf(requireNonNull(partitionedSources, "partitionedSources is null"));
         this.partitionedSourcesSet = ImmutableSet.copyOf(partitionedSources);
         this.stageExecutionDescriptor = requireNonNull(stageExecutionDescriptor, "stageExecutionDescriptor is null");
+        this.materializedExchangeSource = materializedExchangeSource;
         this.statsAndCosts = requireNonNull(statsAndCosts, "statsAndCosts is null");
         this.jsonRepresentation = requireNonNull(jsonRepresentation, "jsonRepresentation is null");
 
@@ -141,6 +144,12 @@ public StageExecutionDescriptor getStageExecutionDescriptor()
         return stageExecutionDescriptor;
     }
 
+    @JsonProperty
+    public boolean isMaterializedExchangeSource()
+    {
+        return materializedExchangeSource;
+    }
+
     @JsonProperty
     public StatsAndCosts getStatsAndCosts()
     {
@@ -206,17 +215,17 @@ private static void findRemoteSourceNodes(PlanNode node, Builder<RemoteSourceNod
 
     public PlanFragment withBucketToPartition(Optional<int[]> bucketToPartition)
     {
-        return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme.withBucketToPartition(bucketToPartition), stageExecutionDescriptor, statsAndCosts, jsonRepresentation);
+        return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme.withBucketToPartition(bucketToPartition), stageExecutionDescriptor, materializedExchangeSource, statsAndCosts, jsonRepresentation);
     }
 
     public PlanFragment withFixedLifespanScheduleGroupedExecution(List<PlanNodeId> capableTableScanNodes)
     {
-        return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme, StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(capableTableScanNodes), statsAndCosts, jsonRepresentation);
+        return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme, StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(capableTableScanNodes), materializedExchangeSource, statsAndCosts, jsonRepresentation);
     }
 
     public PlanFragment withDynamicLifespanScheduleGroupedExecution(List<PlanNodeId> capableTableScanNodes)
     {
-        return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme, StageExecutionDescriptor.dynamicLifespanScheduleGroupedExecution(capableTableScanNodes), statsAndCosts, jsonRepresentation);
+        return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme, StageExecutionDescriptor.dynamicLifespanScheduleGroupedExecution(capableTableScanNodes), materializedExchangeSource, statsAndCosts, jsonRepresentation);
     }
 
     @Override

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

@@ -149,7 +149,9 @@ public SubPlan createSubPlans(Session session, Plan plan, boolean forceSingleNod
                 idAllocator,
                 new SymbolAllocator(plan.getTypes().allTypes()));
 
-        FragmentProperties properties = new FragmentProperties(new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), plan.getRoot().getOutputSymbols()));
+        FragmentProperties properties = new FragmentProperties(
+                new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), plan.getRoot().getOutputSymbols()),
+                false);
         if (forceSingleNode || isForceSingleNodeOutput(session)) {
             properties = properties.setSingleNodeDistribution();
         }
@@ -239,6 +241,7 @@ private SubPlan reassignPartitioningHandleIfNecessaryHelper(Session session, Sub
                         outputPartitioningScheme.isReplicateNullsAndAny(),
                         outputPartitioningScheme.getBucketToPartition()),
                 fragment.getStageExecutionDescriptor(),
+                fragment.isMaterializedExchangeSource(),
                 fragment.getStatsAndCosts(),
                 fragment.getJsonRepresentation());
 
@@ -315,6 +318,7 @@ private SubPlan buildFragment(PlanNode root, FragmentProperties properties, Plan
                     schedulingOrder,
                     properties.getPartitioningScheme(),
                     StageExecutionDescriptor.ungroupedExecution(),
+                    properties.isMaterializedExchangeSource(),
                     statsAndCosts.getForSubplan(root),
                     Optional.of(jsonFragmentPlan(root, fragmentSymbolTypes, metadata.getFunctionManager(), session)));
 
@@ -416,7 +420,9 @@ else if (exchange.getType() == ExchangeNode.Type.REPARTITION) {
 
             ImmutableList.Builder<SubPlan> builder = ImmutableList.builder();
             for (int sourceIndex = 0; sourceIndex < exchange.getSources().size(); sourceIndex++) {
-                FragmentProperties childProperties = new FragmentProperties(partitioningScheme.translateOutputLayout(exchange.getInputs().get(sourceIndex)));
+                FragmentProperties childProperties = new FragmentProperties(
+                        partitioningScheme.translateOutputLayout(exchange.getInputs().get(sourceIndex)),
+                        context.get().isMaterializedExchangeSource());
                 builder.add(buildSubPlan(exchange.getSources().get(sourceIndex), childProperties, context));
             }
 
@@ -476,9 +482,9 @@ private PlanNode createRemoteMaterializedExchange(ExchangeNode exchange, Rewrite
                     partitioningSymbolAssignments.getConstants(),
                     partitioningMetadata);
 
-            FragmentProperties writeProperties = new FragmentProperties(new PartitioningScheme(
-                    Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()),
-                    write.getOutputSymbols()));
+            FragmentProperties writeProperties = new FragmentProperties(
+                    new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), write.getOutputSymbols()),
+                    true);
             writeProperties.setCoordinatorOnlyDistribution();
 
             List<SubPlan> children = ImmutableList.of(buildSubPlan(write, writeProperties, context));
@@ -675,13 +681,15 @@ private static class FragmentProperties
         private final List<SubPlan> children = new ArrayList<>();
 
         private final PartitioningScheme partitioningScheme;
+        private final boolean materializedExchangeSource;
 
         private Optional<PartitioningHandle> partitioningHandle = Optional.empty();
         private final Set<PlanNodeId> partitionedSources = new HashSet<>();
 
-        public FragmentProperties(PartitioningScheme partitioningScheme)
+        public FragmentProperties(PartitioningScheme partitioningScheme, boolean materializedExchangeSource)
         {
             this.partitioningScheme = partitioningScheme;
+            this.materializedExchangeSource = materializedExchangeSource;
         }
 
         public List<SubPlan> getChildren()
@@ -800,6 +808,11 @@ public PartitioningScheme getPartitioningScheme()
             return partitioningScheme;
         }
 
+        public boolean isMaterializedExchangeSource()
+        {
+            return materializedExchangeSource;
+        }
+
         public PartitioningHandle getPartitioningHandle()
         {
             return partitioningHandle.get();

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

@@ -324,6 +324,7 @@ public static String graphvizLogicalPlan(PlanNode plan, TypeProvider types, Sess
                 ImmutableList.of(plan.getId()),
                 new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), plan.getOutputSymbols()),
                 StageExecutionDescriptor.ungroupedExecution(),
+                false,
                 StatsAndCosts.empty(),
                 Optional.empty());
         return GraphvizPrinter.printLogical(ImmutableList.of(fragment), session);

presto-main/src/test/java/com/facebook/presto/execution/MockQueryExecution.java

@@ -155,6 +155,7 @@ public QueryInfo getQueryInfo()
                         30,
                         new DataSize(31, BYTE),
                         new DataSize(32, BYTE),
+                        new DataSize(33, BYTE),
                         ImmutableList.of(),
                         ImmutableList.of()),
                 Optional.empty(),

presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java

@@ -120,6 +120,7 @@ public MockRemoteTask createTableScanTask(TaskId taskId, InternalNode newNode, L
                 ImmutableList.of(sourceId),
                 new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)),
                 StageExecutionDescriptor.ungroupedExecution(),
+                false,
                 StatsAndCosts.empty(),
                 Optional.empty());
 

presto-main/src/test/java/com/facebook/presto/execution/TaskTestUtils.java

@@ -102,6 +102,7 @@ private TaskTestUtils()
             new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(SYMBOL))
                     .withBucketToPartition(Optional.of(new int[1])),
             StageExecutionDescriptor.ungroupedExecution(),
+            false,
             StatsAndCosts.empty(),
             Optional.empty());
 

presto-main/src/test/java/com/facebook/presto/execution/TestQueryStats.java

@@ -190,6 +190,8 @@ public class TestQueryStats
             new DataSize(31, BYTE),
             new DataSize(32, BYTE),
 
+            new DataSize(33, BYTE),
+
             ImmutableList.of(new StageGcStatistics(
                     101,
                     102,
@@ -260,9 +262,12 @@ public static void assertExpectedQueryStats(QueryStats actual)
         assertEquals(actual.getOutputPositions(), 29);
 
         assertEquals(actual.getWrittenOutputPositions(), 30);
+
         assertEquals(actual.getWrittenOutputLogicalDataSize(), new DataSize(31, BYTE));
         assertEquals(actual.getWrittenOutputPhysicalDataSize(), new DataSize(32, BYTE));
 
+        assertEquals(actual.getWrittenIntermediatePhysicalDataSize(), new DataSize(33, BYTE));
+
         assertEquals(actual.getStageGcStatistics().size(), 1);
         StageGcStatistics gcStatistics = actual.getStageGcStatistics().get(0);
         assertEquals(gcStatistics.getStageId(), 101);

presto-main/src/test/java/com/facebook/presto/execution/TestSqlStageExecution.java

@@ -177,6 +177,7 @@ private static PlanFragment createExchangePlanFragment()
                 ImmutableList.of(planNode.getId()),
                 new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), planNode.getOutputSymbols()),
                 StageExecutionDescriptor.ungroupedExecution(),
+                false,
                 StatsAndCosts.empty(),
                 Optional.empty());
     }

presto-main/src/test/java/com/facebook/presto/execution/TestStageStateMachine.java

@@ -333,6 +333,7 @@ private static PlanFragment createValuesPlan()
                 ImmutableList.of(valuesNodeId),
                 new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)),
                 StageExecutionDescriptor.ungroupedExecution(),
+                false,
                 StatsAndCosts.empty(),
                 Optional.empty());
 

presto-main/src/test/java/com/facebook/presto/execution/scheduler/TestPhasedExecutionSchedule.java

@@ -264,6 +264,7 @@ private static PlanFragment createFragment(PlanNode planNode)
                 ImmutableList.of(planNode.getId()),
                 new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), planNode.getOutputSymbols()),
                 StageExecutionDescriptor.ungroupedExecution(),
+                false,
                 StatsAndCosts.empty(),
                 Optional.empty());
     }

presto-main/src/test/java/com/facebook/presto/execution/scheduler/TestSourcePartitionedScheduler.java

@@ -481,6 +481,7 @@ private static SubPlan createPlan()
                 ImmutableList.of(TABLE_SCAN_NODE_ID),
                 new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)),
                 StageExecutionDescriptor.ungroupedExecution(),
+                false,
                 StatsAndCosts.empty(),
                 Optional.empty());