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

[ulysses-you]

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 rebuilt over 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 node is decided for one child -- its grouping, the partitioning it reports and the child
partitioning it was decided over -- so it comes back through GroupPartitionsExec.apply rather than
through copy, which would carry that decision over. An aligned node is turned away. So is a node
whose child no longer reports the partitioning it was decided for: the positions name keys in that
space, and the node has already given up its claim for it.

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, when it is the sort
that aggregate was planned against, the ordering its required child ordering names. 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, and it needs the regrouping not to report that ordering: a sort lands between the pair
only then, which the defaults do not leave.

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 where the ordering is derived from the partition keys; and a
comparison against the plan bypassPartialAggregation builds for the same query, whose grouping
projects what the fold's does, the sort the partial aggregate holds staying under the regrouping
where no sort is crossed, and the pair folded over collect_set. The positions moving into a child
that reports the same keys in another order are covered in GroupPartitionsExecSuite, where that
composition can be built directly.

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 sorts the fold takes down to one. Each test asserts the shape the
defaults leave, where the ordering the source derives and the key ordering the regrouping reports
satisfy both aggregates and the pair folds with no sort at all; the shapes that need a sort between
the pair, which needs the regrouping not to report its ordering, are the ones that set
preserveKeyOrderingOnCoalesce.enabled=false, or both orderings off for the sort the fold takes
with the partial aggregate.

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. Dropping the lower sort unconditionally fails
the new sort test, and dropping the child-partitioning turn-away fails the new position test,
re-parenting onto a key space the node was not decided for.

Ran sql/Test/compile, KeyGroupedPartitioningSuite and GroupPartitionsExecSuite (225), and
AdaptivePartialAggregationSuite, CombineAdjacentAggregationSuite, DataFrameAggregateSuite,
EnsureRequirementsSuite, PushDownLocalSortSuite, RemoveRedundantSortsSuite,
RemoveRedundantWindowGroupLimitsSuite, ReplaceHashWithSortAggSuite and SQLMetricsSuite (393),
618 in total, no failures.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (deepseek-flash v4.1)

[peter-toth]

Thanks for the PR, @ulysses-you!

The rule now walks a chain of GroupPartitionsExec and local sorts to reach the partial aggregate, and withKeyPositionsFor moves the grouping's projected positions from the aggregate's key space into the child's. I traced that arithmetic and it is right: positions.map(newChildKp.expressions) reproduces the expression list the node reported before, and the projected key rows, dataTypes, isGrouped and isCollapsed come out the same, so outputPartitioning is unchanged and the combined aggregate's ClusteredDistribution stays satisfied. I also checked the direction your calibration would expect me to: the fold cannot regress, since it replaces two hash/sort passes with one over the same stream and the bail covers the shape where the kept sort would grow from the partial aggregate's output to the whole scan.

I built the branch and ran 21 probes of my own outside the suite's shapes - count-distinct, an aggregate over a storage-partitioned join grouping by either side's key, an alias on the grouping key, a grouping key that is not a plain attribute, partially-clustered distribution, a filter and a projection below, a reversed grouping order, both coalesce-ordering configs, WITH ROLLUP, a limit above, and AQE. In every one the rule-on answer matched the rule-off answer and ValidateRequirements.validate held. Nothing blocking from me.

Non-blocking

• 1. The bail's derived-ordering path is untested: the description singles this path out as the reason the bail is load-bearing "for a grouping that derives its ordering, not only for one whose source declares it", but the only bail test uses a source that declares its ordering. I measured that the derived path is a distinct reachable shape and that nothing covers it. One extra withSQLConf closes it. inline: KeyGroupedPartitioningSuite.scala:6820

Minor

• 2. withKeyPositionsFor cites a producer whose expression list is discarded: KeyedPartitioning.project builds the projected expressions, then AliasAwareOutputExpression.projectKeyedPartitionings replaces them with its per-position cross-product. What actually makes the lookup work is that a partial aggregate's resultExpressions are groupingAttributes ++ bufferAttributes, so hasAlias is false there and the alternatives are the child's own expressions. inline: GroupPartitionsExec.scala:562

[peter-toth]

Finding 1. This pins the declared-ordering arm of the bail. The description says the bail matters for more than that:

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.

I measured that on 98632fd and it holds, on the fixture of the sort-aggregate test at :6717 rather than this one: the same table, the same three rows, the same query, only the conf differs.

V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED	aggregates	sorts
false (that test's arm)	1, the fold	1
true	2, the bail	1

So the conf alone flips that query from the fold to the bail, and nothing holds that arm. It also means the mutation check the description reports, that dropping the bail fails this test, only covers the declared-ordering source; the derived shape would stay silent under the same mutation.

An extra arm on the :6717 test is enough, since it needs neither a second table nor a declared ordering. This is what I ran, and it passes as written:

      // The ordering derived from the partition keys satisfies the partial aggregate the way a
      // declared one does, so the bail holds for it too.
      withSQLConf(
          SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true",
          SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") {
        val df = sql(query)
        checkAnswer(df, expected)
        val plan = df.queryExecution.executedPlan
        val aggs = collect(plan) { case agg: BaseAggregateExec => agg }
        assert(aggs.size == 2,
          s"the derived ordering leaves the partial aggregate no sort to give up either:\n$plan")
      }

[ulysses-you]

Added the arm in aafc7cb, on the sort-aggregate test: with partitionKeyOrdering.enabled the
scan's derived ordering satisfies the partial aggregate, so it holds no sort of its own and the bail
keeps the pair. It asserts what the declared-ordering test does: 2 aggregates, one local sort reading
the grouping, the partial aggregate on the scan. Dropping the bail now fails that arm as well, so the
derived shape no longer goes unheld under that mutation.

[peter-toth]

Finding 2. The parenthetical names a producer whose expression list does not survive.

KeyedPartitioning.project does build positions.map(expressions) (partitioning.scala:786), but AliasAwareOutputExpression.projectKeyedPartitionings throws that away: it calls project only for the layout and then replaces the expressions with its per-position cross-product, projected.copy(expressions = projectedExprs) (AliasAwareOutputExpression.scala:149-158). The alternatives come from projectExpression, which under hasAlias yields the node's output attributes, not the child's expressions.

What actually makes the lookup work here is one step further out: a partial aggregate's resultExpressions are groupingAttributes ++ aggBufferAttributes (AggUtils.planAggregateWithoutDistinct), all plain attributes, so hasAlias is false on it and the no-alias branch keeps kp.expressions(i) as the child's own. That is also what makes indexOf safe rather than merely lucky: the no-alias branch filters per position by references.subsetOf(outputSet) and projectablePositions is ascending, so a found index is the position the aggregate projected it from and cannot be a coincidental hit at another position.

Worth saying, because the reachability of the whole method rests on it, and project reads as if it were the guarantee.

[ulysses-you]

Reworked in aafc7cb. The comment now credits
AliasAwareOutputExpression.projectKeyedPartitionings for keeping the projected expressions the
child's own, and names the condition it rests on: a partial aggregate has no aliases, its
resultExpressions being groupingAttributes ++ bufferAttributes, so each position's alternative
stays the child's expression rather than the one project built and the projection replaced.

[ulysses-you]

Thanks for the review @peter-toth . Both points are addressed in aafc7cb.

[ulysses-you]

thank you @peter-toth , also cc @cloud if you have other comments

[uros-b]

Thank you @ulysses-you and @peter-toth!

[dongjoon-hyun]

+1, LGTM.

I traced the key-position translation, the attribute binding of the sort that crosses the aggregate, the enableSortedMerge carry-over and the marked-layout (mayContainUnknownPartitionKeys) path, and did not find a correctness hole. A few non-blocking comments, any of which can be a follow-up:

1. detachAggregate runs before the compatibility checks. group.withKeyPositionsFor(child) goes through GroupPartitionsExec.apply -> computeGrouping, which walks every split of the child and builds an InternalRowComparableWrapper per split. When combinedAggregate / isPartialAgg then rejects the pair (a PartialMerge + Final object-hash pair, a logicalLink mismatch, ...), that work is thrown away. In the HashAggregateExec arm, evaluating combinedAggregate first would avoid it.

2. withKeyPositionsFor reads the live child.outputPartitioning rather than the stored childPartitioning that joinKeyPositions was computed against. Where the two diverge, outputPartitioning deliberately drops the keyed claim, but this path translates the positions against the live one and hands back a node as if the claim still held. No rule is known to perform such a rewrite today, so this is only about staying symmetric with checkChildStillMatches.

3. The sort handling keys off the presence of a local sort above the aggregate, not off its ordering. Today the only inserter of a local sort between the pair is EnsureRequirements for SortAggregateExec.requiredChildOrdering, and isPartialAgg pins the two groupings equal, so the crossed sort always covers the one being dropped -- PushDownLocalSort cannot get in between either, since it treats neither the aggregates nor GroupPartitionsExec as order-preserving. That invariant lives in the comment rather than in the code; a SortOrder.orderingSatisfies(sort.sortOrder, aggregate.requiredChildOrdering.head) guard would keep it honest if a rule starts inserting sorts there.

4. Two shapes the fold reaches that the suite does not pin: the composing branch of joinKeyPositions.fold(plannedInNewChild)(_.map(plannedInNewChild)) (the fold path always starts from None, the partial aggregate having already narrowed), and the hasUpperSort = false branch where the sort below the partial aggregate stays under the grouping.

[ulysses-you]

Thanks @dongjoon-hyun. All four are in 4174278.

1. detachAggregate now takes the aggregate the pair folds into and asks isCompatibleAggregates at the leaf, before the chain is rebuilt, so the grouping is re-decided over the new child only for a pair that folds. That check lives there and nowhere else: isPartialAgg is now the mode half alone and combinedAggregate no longer repeats the compatibility part.

2. withKeyPositionsFor turns away a child that no longer reports the partitioning this node was decided for, and reads the expressions from the stored childPartitioning. The positions name keys in that space, so moving them against another one re-targets the grouping. Same judgement as outputPartitioning and checkChildStillMatches.

3. The sort below the partial aggregate goes only when it satisfies that aggregate's requiredChildOrdering. The check is vacuous for the hash and object-hash arms, whose required ordering is empty; there the crossing rests on the pair being compatible, which pins the two groupings.

4. Both shapes are pinned.

   • The composing branch needs a node that already projects, and no planner shape reaches the fold with one: a partial aggregate's reported keys are its grouping attributes, which are the cluster keys it is asked to satisfy, so every position survives and EnsureRequirements hands out positions only for a coalescing node. Measured over the suite's seven tests: all 11 crossings arrive with joinKeyPositions = None. It is covered directly in GroupPartitionsExecSuite now, on a node projecting position 1 of a child reporting the same keys in the other order.
   • The sort that stays is covered in KeyGroupedPartitioningSuite: a hash pair with a local sort under the partial aggregate, which the fold leaves under the regrouping.

Both new tests discriminate by mutation: dropping the lower sort unconditionally fails the sort test, and dropping the child-partitioning turn-away fails the position test, re-parenting with JoinKeyPositions: [0] over a child whose key at that position is another one.

Ran KeyGroupedPartitioningSuite and GroupPartitionsExecSuite (224) and the nine related suites (392), 616 in total, no failures.

[ulysses-you]

Rebased onto master, which now carries SPARK-59396: partitionKeyOrdering.enabled and preserveKeyOrderingOnCoalesce.enabled default to true, so the ordering a scan derives from its partition keys is the grouping these tests are about, and a regrouping reports the key ordering it keeps. The tests now assert the default shapes: the pair folds with no sort at all, the regrouping reading the scan, and the narrowed case keeps the partial aggregate's sort under the regrouping.

Two shapes need a sort between the pair, which needs the regrouping not to report its ordering, so the bail test and the bail arm set preserveKeyOrderingOnCoalesce.enabled=false for it, and the arm that covers the sort going with the partial aggregate sets both orderings off, the configuration that shape needs.

ReplaceHashWithSortAgg plans the combined aggregate sort-based once the regrouping reports an ordering that satisfies the grouping, so the collect_set arm asserts the complete mode and the answer rather than the aggregate kind.

Ran KeyGroupedPartitioningSuite and GroupPartitionsExecSuite (225) and the nine related suites (393), 618 in total, no failures.

[ulysses-you]

Merge Summary:

• merged into master f0be43c
• merged into branch-4.x 8948295

Posted by merge_spark_pr.py

[ulysses-you]

thank you all!