Package org.elasticsearch.xpack.esql.optimizer.rules.logical

Class PushAggregateThroughUnionAll

java.lang.Object
org.elasticsearch.xpack.esql.rule.Rule<Aggregate,LogicalPlan>
org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.OptimizerRule<Aggregate>
org.elasticsearch.xpack.esql.optimizer.rules.logical.PushAggregateThroughUnionAll
public class PushAggregateThroughUnionAll extends OptimizerRules.OptimizerRule<Aggregate>
Decomposes an Aggregate whose child is a direct-leaf UnionAll (heterogeneous FROM) into per-branch partial aggregates combined by a final merge aggregate.

Two kinds of aggregate are decomposed, via two combine strategies:

  • Algebraic (COUNT, SUM, MIN, MAX): each branch emits the aggregate's partial final value and the coordinator merges them with another plain aggregate (COUNT→SUM, SUM→SUM, MIN→MIN, MAX→MAX).
  • Intermediate-state (COUNT_DISTINCT, PERCENTILE, STDDEV): these are not algebraically decomposable, so each branch emits its mergeable intermediate state (HLL sketch, t-digest, Welford state) as a PARTIAL_AGG column via ToPartial, and the coordinator merges those with FromPartial.

MEDIAN and AVG are not handled here: they are surrogates rewritten before this rule runs (MEDIAN→PERCENTILE, AVG→SUM/COUNT), so they decompose transitively through the cases above.

Filters: a per-aggregate filter (COUNT_DISTINCT(x) WHERE ...) pushes in both paths. For the algebraic path the filtered aggregate is emitted as-is per branch. For the intermediate-state path the filter is carried on the ToPartial node (the inner aggregate stays unfiltered); the physical layer then wraps the inner aggregator in a filtered supplier so the branch filters before folding into its intermediate state. See ToPartial#supplierWithInnerFilter and the planner's aggregate factory.

Transformation sketch (two-branch case). c/s take the algebraic path; d (a heavy aggregate) takes the intermediate-state path, so its partial column is typed PARTIAL_AGG and is produced by ToPartial per branch and merged by FromPartial on the coordinator: 


 -- Before:
 Aggregate[c = COUNT(*), s = SUM(salary), d = COUNT_DISTINCT(salary), dep] BY [dep]
   UnionAll[[dep{r1}, salary{r2}]]
     EsRelation[[dep{f1}, salary{f2}]]
     ExternalRelation[[dep{f3}, salary{f4}]]

 -- After:
 Aggregate[c = SUM($$partial$$c), s = SUM($$partial$$s), d = FromPartial($$partial$$d, COUNT_DISTINCT),
           dep = Alias($$g_dep)] BY [$$g_dep]
   UnionAll[[$$partial$$c, $$partial$$s, $$partial$$d{PARTIAL_AGG}, $$g_dep]]
     Aggregate[$$partial$$c = COUNT(*), $$partial$$s = SUM(salary{f2}),
               $$partial$$d = ToPartial(COUNT_DISTINCT(salary{f2})), Alias(dep{f1})] BY [Alias(dep{f1})]
       EsRelation[[dep{f1}, salary{f2}]]
     Aggregate[$$partial$$c = COUNT(*), $$partial$$s = SUM(salary{f4}),
               $$partial$$d = ToPartial(COUNT_DISTINCT(salary{f4})), Alias(dep{f3})] BY [Alias(dep{f3})]
       ExternalRelation[[dep{f3}, salary{f4}]]

All partial aggregate aliases and shared grouping aliases use pre-allocated NameIds that are consistent across branches, so the outer UnionAll output and the outer Aggregate's grouping/aggregate references resolve correctly. The outer Aggregate preserves the output IDs of the original Aggregate so that plan nodes above it remain valid.