Skip to content

[SPARK-59564][SQL] Combine adjacent aggregation across a GroupPartitionsExec - #58884

Open
ulysses-you wants to merge 1 commit into
apache:masterfrom
ulysses-you:SPARK-59564
Open

ulysses-you wants to merge 1 commit into
apache:masterfrom
ulysses-you:SPARK-59564

Conversation

@ulysses-you

@ulysses-you ulysses-you commented Sep 17, 2026

Copy link
Copy Markdown
Contributor

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 v4.1)

…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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant