[SPARK-59564][SQL] Combine adjacent aggregation across a GroupPartitionsExec - #58884
Open
ulysses-you wants to merge 1 commit into
Open
ulysses-you wants to merge 1 commit into
ulysses-you wants to merge 1 commit into
Conversation
…onsExec
### What changes were proposed in this pull request?
`CombineAdjacentAggregation` now looks through the `GroupPartitionsExec` and the local sorts
`EnsureRequirements` may have put between the partial and the final aggregate, so the pair is
combined even when the final aggregate's distribution was satisfied without a shuffle.
The grouping is re-parented onto the aggregate's child with the key positions it projects moved into
that child's key space, which is what `GroupPartitionsExec.withKeyPositionsFor` answers. It was
planned against the aggregate, whose partitioning is the child's projected down to what the
aggregate's output keeps, so those positions name a key space the child being handed need not share.
A sort crossed above the aggregate orders the rows the combined aggregate reads, by the grouping the
two aggregates share, so the sort feeding the partial aggregate goes with it. Where the partial
aggregate holds no sort of its own, the pair is left alone instead, being then the only cardinality
reducer before that sort. With no sort crossed at all, the sort below the aggregate stays below it:
the aggregate reads what it read, and the ordering it claims is the one it had.
### Why are the changes needed?
The rule only matched a strictly adjacent pair, so it missed the shape where `EnsureRequirements`
satisfies the final aggregate's `ClusteredDistribution` with a `GroupPartitionsExec` rather than a
shuffle. Folding the pair there removes an aggregation pass, and where a sort fed the partial
aggregate, that sort as well.
### Does this PR introduce _any_ user-facing change?
Yes, the plan changes for an aggregate pair whose child needs a `GroupPartitionsExec`, e.g. for
`GROUP BY id, name` over a v2 table partitioned by `(id, name)` whose keys repeat across splits:
Before:
```
HashAggregate (Final)
+- GroupPartitions
+- HashAggregate (Partial)
+- BatchScan
```
After:
```
HashAggregate (Complete)
+- GroupPartitions
+- BatchScan
```
A grouping on part of the partition keys is covered too: the positions are moved onto the scan's key
space, which is a lookup, the projected expressions being the child's own. Where a sort sits between
the pair it goes with the partial aggregate, so two sorts become one. Where the partial aggregate
holds no sort of its own and a sort above orders the combined aggregate's rows, the pair is left
alone and the plan is unchanged.
That last path is taken by more shapes than a source that declares its ordering: with
`spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` the ordering derived from the partition
keys already satisfies the partial aggregate's, so it holds no sort of its own to give up. The bail is
load-bearing for a grouping that derives its ordering, not only for one whose source declares it.
Results are unchanged.
### How was this patch tested?
New tests in `KeyGroupedPartitioningSuite`: the pair combined across the grouping, including the
object-hash pair and the AQE path; the sort pair with the sort in between; a grouping on part of the
partition keys combined with its positions moved; a sort pair whose grouping covers part of the
partition keys; a grouping the scan reports narrowed itself; the pair kept where the source already
orders the aggregate's input; and a comparison against the plan `bypassPartialAggregation` builds for
the same query, whose grouping projects what the fold's does.
Each asserts the aggregates left and their mode, the grouping's positions where they matter, and
`ValidateRequirements.validate(plan)`, and compares the answer against the rule-off run. The two sort
tests' rule-off arms pin the two sorts the fold takes down to one.
Verified the tests discriminate by mutation: keeping the positions unmoved fails both narrowed-grouping
tests with three rows instead of two, and dropping the sort bail fails the ordered-source test with
one aggregate instead of two.
Ran `sql/Test/compile`, `KeyGroupedPartitioningSuite` (191), and `AdaptivePartialAggregationSuite`,
`CombineAdjacentAggregationSuite`, `DataFrameAggregateSuite`, `EnsureRequirementsSuite`,
`GroupPartitionsExecSuite`, `PushDownLocalSortSuite`, `RemoveRedundantSortsSuite`,
`RemoveRedundantWindowGroupLimitsSuite`, `ReplaceHashWithSortAggSuite` and `SQLMetricsSuite`
(603 in total, no failures).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (deepseek-flash)
Assisted-by: Claude Code (deepseek-flash)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
CombineAdjacentAggregationnow looks through theGroupPartitionsExecand the local sortsEnsureRequirementsmay have put between the partial and the final aggregate, so the pair iscombined even when the final aggregate's distribution was satisfied without a shuffle.
The grouping is re-parented onto the aggregate's child with the key positions it projects moved into
that child's key space, which is what
GroupPartitionsExec.withKeyPositionsForanswers. It wasplanned against the aggregate, whose partitioning is the child's projected down to what the
aggregate's output keeps, so those positions name a key space the child being handed need not share.
A sort crossed above the aggregate orders the rows the combined aggregate reads, by the grouping the
two aggregates share, so the sort feeding the partial aggregate goes with it. Where the partial
aggregate holds no sort of its own, the pair is left alone instead, being then the only cardinality
reducer before that sort. With no sort crossed at all, the sort below the aggregate stays below it:
the aggregate reads what it read, and the ordering it claims is the one it had.
Why are the changes needed?
The rule only matched a strictly adjacent pair, so it missed the shape where
EnsureRequirementssatisfies the final aggregate's
ClusteredDistributionwith aGroupPartitionsExecrather than ashuffle. Folding the pair there removes an aggregation pass, and where a sort fed the partial
aggregate, that sort as well.
Does this PR introduce any user-facing change?
Yes, the plan changes for an aggregate pair whose child needs a
GroupPartitionsExec, e.g. forGROUP BY id, nameover a v2 table partitioned by(id, name)whose keys repeat across splits:Before:
After:
A grouping on part of the partition keys is covered too: the positions are moved onto the scan's key
space, which is a lookup, the projected expressions being the child's own. Where a sort sits between
the pair it goes with the partial aggregate, so two sorts become one. Where the partial aggregate
holds no sort of its own and a sort above orders the combined aggregate's rows, the pair is left
alone and the plan is unchanged.
That last path is taken by more shapes than a source that declares its ordering: with
spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabledthe ordering derived from the partitionkeys already satisfies the partial aggregate's, so it holds no sort of its own to give up. The bail is
load-bearing for a grouping that derives its ordering, not only for one whose source declares it.
Results are unchanged.
How was this patch tested?
New tests in
KeyGroupedPartitioningSuite: the pair combined across the grouping, including theobject-hash pair and the AQE path; the sort pair with the sort in between; a grouping on part of the
partition keys combined with its positions moved; a sort pair whose grouping covers part of the
partition keys; a grouping the scan reports narrowed itself; the pair kept where the source already
orders the aggregate's input; and a comparison against the plan
bypassPartialAggregationbuilds forthe same query, whose grouping projects what the fold's does.
Each asserts the aggregates left and their mode, the grouping's positions where they matter, and
ValidateRequirements.validate(plan), and compares the answer against the rule-off run. The two sorttests' rule-off arms pin the two sorts the fold takes down to one.
Verified the tests discriminate by mutation: keeping the positions unmoved fails both narrowed-grouping
tests with three rows instead of two, and dropping the sort bail fails the ordered-source test with
one aggregate instead of two.
Ran
sql/Test/compile,KeyGroupedPartitioningSuite(191), andAdaptivePartialAggregationSuite,CombineAdjacentAggregationSuite,DataFrameAggregateSuite,EnsureRequirementsSuite,GroupPartitionsExecSuite,PushDownLocalSortSuite,RemoveRedundantSortsSuite,RemoveRedundantWindowGroupLimitsSuite,ReplaceHashWithSortAggSuiteandSQLMetricsSuite(603 in total, no failures).
Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (deepseek-flash v4.1)