Show stats in EXPLAIN/EXPLAIN ANALYZE with intermediate aggregation

Author: raunaqmorarka

core/trino-main/src/main/java/io/trino/cost/AggregationStatsRule.java

@@ -49,13 +49,13 @@ public Pattern<AggregationNode> getPattern()
     @Override
     protected Optional<PlanNodeStatsEstimate> doCalculate(AggregationNode node, Context context)
     {
-        if (node.getGroupingSetCount() != 1 || node.getStep() == INTERMEDIATE) {
+        if (node.getGroupingSetCount() != 1) {
             return Optional.empty();
         }
 
         PlanNodeStatsEstimate estimate;
 
-        if (node.getStep() == PARTIAL) {
+        if (node.getStep() == PARTIAL || node.getStep() == INTERMEDIATE) {
             estimate = partialGroupBy(context.statsProvider().getStats(node.getSource()),
                     node.getGroupingKeys(),
                     node.getAggregations());
@@ -101,7 +101,8 @@ public static double getRowsCount(PlanNodeStatsEstimate sourceStats, Collection<
 
     private static PlanNodeStatsEstimate partialGroupBy(PlanNodeStatsEstimate sourceStats, Collection<Symbol> groupBySymbols, Map<Symbol, Aggregation> aggregations)
     {
-        // Pessimistic assumption of no reduction from PARTIAL aggregation, forwarding of the source statistics. This makes the CBO estimates in the EXPLAIN plan output easier to understand,
+        // Pessimistic assumption of no reduction from PARTIAL and INTERMEDIATE aggregation, forwarding of the source statistics.
+        // This makes the CBO estimates in the EXPLAIN plan output easier to understand,
         // even though partial aggregations are added after the CBO rules have been run.
         PlanNodeStatsEstimate.Builder result = PlanNodeStatsEstimate.builder();
         result.setOutputRowCount(sourceStats.getOutputRowCount());

core/trino-main/src/test/java/io/trino/cost/TestAggregationStatsRule.java

@@ -185,4 +185,46 @@ public void testAggregationWithMoreGroupingSets()
                         .build())
                 .check(PlanNodeStatsAssertion::outputRowsCountUnknown);
     }
+
+    @Test
+    void testAggregationStep()
+    {
+        testAggregationStep(AggregationNode.Step.PARTIAL);
+        testAggregationStep(AggregationNode.Step.INTERMEDIATE);
+        testAggregationStep(AggregationNode.Step.FINAL);
+        testAggregationStep(AggregationNode.Step.SINGLE);
+    }
+
+    void testAggregationStep(AggregationNode.Step step)
+    {
+        tester().assertStatsFor(pb -> pb
+                        .aggregation(ab -> ab
+                                .step(step)
+                                .addAggregation(pb.symbol("sum", BIGINT), aggregation("sum", ImmutableList.of(new Reference(BIGINT, "x"))), ImmutableList.of(BIGINT))
+                                .singleGroupingSet(pb.symbol("y", BIGINT), pb.symbol("z", BIGINT))
+                                .source(pb.values(pb.symbol("y", BIGINT), pb.symbol("z", BIGINT)))))
+                .withSourceStats(PlanNodeStatsEstimate.builder()
+                        .setOutputRowCount(100)
+                        .addSymbolStatistics(new Symbol(BIGINT, "y"), SymbolStatsEstimate.builder()
+                                .setLowValue(1)
+                                .setHighValue(10)
+                                .setDistinctValuesCount(5)
+                                .setNullsFraction(0.3)
+                                .build())
+                        .addSymbolStatistics(new Symbol(BIGINT, "z"), SymbolStatsEstimate.builder()
+                                .setLowValue(0)
+                                .setHighValue(3)
+                                .setDistinctValuesCount(3)
+                                .setNullsFraction(0)
+                                .build())
+                        .build())
+                .check(check -> {
+                    if (step.isOutputPartial()) {
+                        check.outputRowsCount(100);
+                    }
+                    else {
+                        check.outputRowsCount(18);
+                    }
+                });
+    }
 }