Enable push partial aggregation though join

core/trino-main/src/main/java/io/trino/sql/planner/OptimizerConfig.java

@@ -75,7 +75,7 @@ public class OptimizerConfig
     private boolean preferPartialAggregation = true;
     private boolean pushAggregationThroughOuterJoin = true;
     private boolean enableIntermediateAggregations;
-    private boolean pushPartialAggregationThroughJoin;
+    private boolean pushPartialAggregationThroughJoin = true;
     private boolean preAggregateCaseAggregationsEnabled = true;
     private boolean enableForcedExchangeBelowGroupId = true;
     private boolean optimizeTopNRanking = true;

core/trino-main/src/main/java/io/trino/sql/planner/PlanOptimizers.java

@@ -1000,11 +1000,13 @@ public PlanOptimizers(
                 ruleStats,
                 statsCalculator,
                 costCalculator,
-                ImmutableSet.of(
-                        new PushPartialAggregationThroughJoin(),
-                        new PushPartialAggregationThroughExchange(plannerContext),
-                        new PruneJoinColumns(),
-                        new PruneJoinChildrenColumns())));
+                ImmutableSet.<Rule<?>>builder()
+                        .addAll(new PushPartialAggregationThroughJoin().rules())
+                        .add(new PushPartialAggregationThroughExchange(plannerContext),
+                                new PruneJoinColumns(),
+                                new PruneJoinChildrenColumns(),
+                                new RemoveRedundantIdentityProjections())
+                        .build()));
         builder.add(new IterativeOptimizer(
                 plannerContext,
                 ruleStats,

core/trino-main/src/main/java/io/trino/sql/planner/iterative/rule/PushPartialAggregationThroughExchange.java

@@ -157,6 +157,9 @@ private PlanNode pushPartial(AggregationNode aggregation, ExchangeNode exchange,
 
             SymbolMapper symbolMapper = mappingsBuilder.build();
             AggregationNode mappedPartial = symbolMapper.map(aggregation, source, context.getIdAllocator().getNextId());
+            mappedPartial = AggregationNode.builderFrom(mappedPartial)
+                    .setIsInputReducingAggregation(true)
+                    .build();
 
             Assignments.Builder assignments = Assignments.builder();
 

core/trino-main/src/main/java/io/trino/sql/planner/plan/AggregationNode.java

@@ -53,6 +53,11 @@ public class AggregationNode
     private final Optional<Symbol> hashSymbol;
     private final Optional<Symbol> groupIdSymbol;
     private final List<Symbol> outputs;
+    /**
+     * Indicates whether it is beneficial (e.g. reduces remote exchange input) to retain this aggregation
+     * as an auxiliary step when making a decision to push down partial aggregation more aggressively.
+     */
+    private final Optional<Boolean> isInputReducingAggregation;
 
     public static AggregationNode singleAggregation(
             PlanNodeId id,
@@ -63,6 +68,19 @@ public static AggregationNode singleAggregation(
         return new AggregationNode(id, source, aggregations, groupingSets, ImmutableList.of(), SINGLE, Optional.empty(), Optional.empty());
     }
 
+    public AggregationNode(
+            PlanNodeId id,
+            PlanNode source,
+            Map<Symbol, Aggregation> aggregations,
+            GroupingSetDescriptor groupingSets,
+            List<Symbol> preGroupedSymbols,
+            Step step,
+            Optional<Symbol> hashSymbol,
+            Optional<Symbol> groupIdSymbol)
+    {
+        this(id, source, aggregations, groupingSets, preGroupedSymbols, step, hashSymbol, groupIdSymbol, Optional.empty());
+    }
+
     @JsonCreator
     public AggregationNode(
             @JsonProperty("id") PlanNodeId id,
@@ -72,7 +90,8 @@ public AggregationNode(
             @JsonProperty("preGroupedSymbols") List<Symbol> preGroupedSymbols,
             @JsonProperty("step") Step step,
             @JsonProperty("hashSymbol") Optional<Symbol> hashSymbol,
-            @JsonProperty("groupIdSymbol") Optional<Symbol> groupIdSymbol)
+            @JsonProperty("groupIdSymbol") Optional<Symbol> groupIdSymbol,
+            @JsonProperty("isInputReducingAggregation") Optional<Boolean> isInputReducingAggregation)
     {
         super(id);
 
@@ -104,6 +123,7 @@ public AggregationNode(
         outputs.addAll(aggregations.keySet());
 
         this.outputs = outputs.build();
+        this.isInputReducingAggregation = requireNonNull(isInputReducingAggregation, "exchangeInputAggregation is null");
     }
 
     public List<Symbol> getGroupingKeys()
@@ -207,6 +227,12 @@ public Optional<Symbol> getGroupIdSymbol()
         return groupIdSymbol;
     }
 
+    @JsonProperty("isInputReducingAggregation")
+    public boolean isInputReducingAggregation()
+    {
+        return isInputReducingAggregation.orElse(false);
+    }
+
     public boolean hasOrderings()
     {
         return aggregations.values().stream()
@@ -513,6 +539,7 @@ public static class Builder
         private Step step;
         private Optional<Symbol> hashSymbol;
         private Optional<Symbol> groupIdSymbol;
+        private Optional<Boolean> isInputReducingAggregation;
 
         public Builder(AggregationNode node)
         {
@@ -525,6 +552,7 @@ public Builder(AggregationNode node)
             this.step = node.getStep();
             this.hashSymbol = node.getHashSymbol();
             this.groupIdSymbol = node.getGroupIdSymbol();
+            this.isInputReducingAggregation = node.isInputReducingAggregation;
         }
 
         public Builder setId(PlanNodeId id)
@@ -575,6 +603,12 @@ public Builder setGroupIdSymbol(Optional<Symbol> groupIdSymbol)
             return this;
         }
 
+        public Builder setIsInputReducingAggregation(boolean isInputReducingAggregation)
+        {
+            this.isInputReducingAggregation = Optional.of(isInputReducingAggregation);
+            return this;
+        }
+
         public AggregationNode build()
         {
             return new AggregationNode(
@@ -585,7 +619,8 @@ public AggregationNode build()
                     preGroupedSymbols,
                     step,
                     hashSymbol,
-                    groupIdSymbol);
+                    groupIdSymbol,
+                    isInputReducingAggregation);
         }
     }
 }

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

@@ -67,7 +67,7 @@ public void testDefaults()
                 .setEnableForcedExchangeBelowGroupId(true)
                 .setEnableIntermediateAggregations(false)
                 .setPushAggregationThroughOuterJoin(true)
-                .setPushPartialAggregationThroughJoin(false)
+                .setPushPartialAggregationThroughJoin(true)
                 .setPreAggregateCaseAggregationsEnabled(true)
                 .setDistinctAggregationsStrategy(null)
                 .setPreferPartialAggregation(true)
@@ -126,7 +126,7 @@ public void testExplicitPropertyMappings()
                 .put("optimizer.push-table-write-through-union", "false")
                 .put("optimizer.dictionary-aggregation", "true")
                 .put("optimizer.push-aggregation-through-outer-join", "false")
-                .put("optimizer.push-partial-aggregation-through-join", "true")
+                .put("optimizer.push-partial-aggregation-through-join", "false")
                 .put("optimizer.pre-aggregate-case-aggregations.enabled", "false")
                 .put("optimizer.enable-intermediate-aggregations", "true")
                 .put("optimizer.force-single-node-output", "true")
@@ -181,7 +181,7 @@ public void testExplicitPropertyMappings()
                 .setPushTableWriteThroughUnion(false)
                 .setDictionaryAggregation(true)
                 .setPushAggregationThroughOuterJoin(false)
-                .setPushPartialAggregationThroughJoin(true)
+                .setPushPartialAggregationThroughJoin(false)
                 .setPreAggregateCaseAggregationsEnabled(false)
                 .setEnableIntermediateAggregations(true)
                 .setDistinctAggregationsStrategy(MARK_DISTINCT)