James Xu created SPARK-57928:
--------------------------------
Summary: [SQL] Collapse redundant Partial+Final aggregate pair
when shuffle is skipped
Key: SPARK-57928
URL: https://issues.apache.org/jira/browse/SPARK-57928
Project: Spark
Issue Type: Improvement
Components: Optimizer
Affects Versions: 4.3.0
Reporter: James Xu
---
Problem:
Spark plans aggregations in two phases: a Partial aggregate that pre-aggregates
within each partition, followed by a shuffle and a Final aggregate that merges
across partitions. EnsureRequirements inserts the shuffle only when the child's
output partitioning does not already satisfy
ClusteredDistribution(groupingKeys). When the shuffle is skipped — for example,
because the aggregate sits directly on top of a SortMergeJoin keyed on exactly
the grouping columns — the Final aggregate receives data that is already fully
aggregated within each partition. The Final pass provides no cross-partition
merging, yet it still materialises intermediate buffers and processes every row
a second time.
A representative pattern is an aggregation over a join where the join keys
match the GROUP BY columns:
SELECT t1.user_id, t1.region, SUM(t1.amount) AS total
FROM events t1
JOIN users t2 ON t1.user_id = t2.user_id AND t1.region = t2.region
GROUP BY t1.user_id, t1.region
When the join is a SortMergeJoin on (user_id, region), the output is already
hash-partitioned on those keys. EnsureRequirements skips the shuffle before the
Final aggregate, but leaves both Partial and Final nodes in the plan. On large
inputs with high row counts per partition key, the redundant aggregate pair
causes severe disk spill and widespread fallback to sort-based aggregation.
---
Root Cause:
AggUtils.planAggregateWithoutDistinct unconditionally emits a Partial+Final
pair, expecting EnsureRequirements to insert an Exchange between them. When
EnsureRequirements instead finds that the child's partitioning already
satisfies ClusteredDistribution(groupingKeys) and skips the shuffle, no rule
subsequently detects that the Final node is now doing redundant work. The two
nodes remain in the physical plan and both execute in full.
---
Solution:
Add a new physical rule, RemoveRedundantAggregates, that runs after
EnsureRequirements in both the standard preparations sequence and the AQE
queryStagePreparationRules. The rule pattern-matches a Final-mode aggregate
sitting directly over a Partial-mode aggregate of the same concrete type with
no Exchange between them, and replaces both with a single Complete-mode node.
The absence of an Exchange node between the two aggregates is structural proof
that no shuffle was inserted — if one had been, it would appear as the direct
child of the Final aggregate and the match would not fire.
Complete mode drives the same updateExpressions path as Partial, evaluating
aggregate functions directly on raw input rows and emitting final results in
one pass. This is semantically equivalent to Partial→Final on the same
partition without an intermediate shuffle, but eliminates the buffer
materialisation between the two phases. The replacement node preserves the same
output attributes as the original Final aggregate, so all downstream operators
bind correctly.
The rule applies to HashAggregateExec, ObjectHashAggregateExec, and
SortAggregateExec. It must run after EnsureRequirements; running it before
would collapse pairs that still need a shuffle.
---
Expected Impact:
For queries that aggregate over a SortMergeJoin (or any operator providing
ClusteredDistribution) on the grouping keys, with large input sizes and low
group cardinality relative to input volume:
- Disk spill from the redundant Final aggregate pass is eliminated entirely,
since Complete mode processes raw input in a single pass with no intermediate
buffer materialisation.
- Tasks that previously fell back to sort-based aggregation due to spill
pressure are no longer subject to that pressure from the second pass.
- Overall stage wall-clock time is reduced proportionally to the elimination of
the redundant pass and its associated spill I/O.
The optimisation applies to any query where a SortMergeJoin, bucketed scan, or
other operator providing ClusteredDistribution on the grouping keys sits below
a two-phase aggregate.
--
This message was sent by Atlassian Jira
(v8.20.10#820010)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
