[SPARK-59289][SQL] Decide a co-partitioned child's GroupPartitionsExec once, where the join's merged keys are known

[peter-toth]

What changes were proposed in this pull request?

GroupPartitionsExec is how storage-partitioned join lines two scans up: it coalesces, projects, reorders and pads a scan's partitions. Today the planner decides that node twice and the node then re-derives it many times. This PR makes the planner decide once and the node carry the answer.

Where the node comes from. ensureDistributionAndOrdering resolved every child on its own first, so it placed a GroupPartitionsExec over a co-partitioned child before knowing the parent was a join. checkKeyGroupCompatible then rewrote or replaced it, descending past a local sort, merging new parameters in and dropping any grouping stacked on top. The two sites also worked in different index spaces, the placeholder's positions against the raw partition keys and the spec's against the node's already projected report, which is what joinKeyPositions.orElse and the reducer fast path existed to paper over.

The two paths are now apart. The per-child loop resolves only the children that answer for themselves, through a lifted resolveChild. A new coPartitionChildren owns the rest end to end: it peels every grouping this rule inserted off both children, plans the pairing from the sources' own reports, and builds each node once.

How the children are carried. The co-partitioned set used to be a Seq[Int] of child indexes beside an isCoPartitioned flag, threaded alongside Seq[Distribution]. It is one value now: a Seq[Option[ClusteredDistribution]] aligned with children, holding what each child has to satisfy and nothing for a child that answers for itself. Readers zip it instead of indexing, and it carries the type the split point already knew, so three asserts, four asInstanceOf[ClusteredDistribution] and a dead branch in candidatesFor go with it. ensureDistributionAndOrdering is 53 lines and holds no var.

Who decides what. coPartitionChildren recognises only the three shapes that need nothing arranged between the children, each at the same altitude: every side a single small partition, a storage-partitioned join, or two compatible pass-through specs. The shuffle is shuffleToCoPartition, and the layout it drags everyone onto is pickCoPartitionTarget, returning a CoPartitionTarget that names the winning member and, per child, the member pairing with it. ShuffleSpecCollection answers isCompatibleWith existentially and so never names the member the two sides agreed on, which is why that pairing used to be rediscovered at three levels. It is decided once and carried.

What the node holds. grouping (the index groups it emits and what they say about the layout) and outputPartitioning are constructor fields now, derived once by GroupPartitionsExec.apply, the way ShuffleExchangeExec holds the partitioning it produces. They were per-instance lazy vals, so every copy and withNewChildren threw the memo away.

What satisfies answers. KeyedPartitioning.satisfies returned true for a partitioning that only satisfies after a projecting node, for the storage-partitioned join's benefit alone. It is strict now, and the loose question moved to keysCanSatisfy, which keysMaySatisfy and mayGroupToSatisfy compose with the strict one. EnsureRequirements.resolveKeyedPartitioning picks the positions to project onto through a new KeyedPartitioning.positionsCoveringClusterKeys, which keysSatisfy also reads, so the predicate that decides whether a projection is needed and the one that picks the positions cannot drift apart. Whether a projection is permitted at all is one mayProjectToClusterKeys, where it used to be spelled twice in opposite polarity.

A small vocabulary on ClusteredDistribution. Five Partitioning.satisfies0 implementations spelled out the same if (requireAllClusterKeys) areAllClusterKeysMatched(exprs) else exprs.forall(x => clustering.exists(_.semanticEquals(x))), so the flag was read in five places and the membership test was written out nine times across the file. The pair is now matchesClusterKeys, beside the areAllClusterKeysMatched it wraps, and the membership test it is built from is isClusterKey. allClusterKeysAmong replaces the mirror fold, which asks whether a set of attributes covers the clustering, and AQEUtils reaches for it too. areAllClusterKeysMatched on both distributions, and the StatefulOpClusteredDistribution arm that spelled the same thing out at a call site, are one Seq.corresponds each. The five callers are HashPartitioningLike, NullAwareHashPartitioning, CoalescedNullAwareHashPartitioning, RangePartitioning and ShufflePartitionIdPassThrough, and each rewrite is one call over the same expressions in the same order. The prose in these files called the same thing an "operation key" in about forty places, which was a second name for what the API already calls a cluster key, so that is gone too.

What a second run sees. Peeling and re-planning would re-decide a pairing an earlier run already settled, on an input that by then holds the keyed shuffle that run inserted. So checkKeyGroupCompatible asks first whether the two children satisfy their distributions and line up with each other as they arrive, and hands them back untouched when they do. That question decides nothing new on a plan this rule has not seen, since a source that needs a node does not satisfy its distribution as it stands.

EnsureRequirements
  ensureDistributionAndOrdering        the two paths, and the ordering step
    resolveChild                       a child that answers for itself
    coPartitionChildren                children that only answer together
      checkKeyGroupCompatible          try the join, build both nodes once
        alreadyCoPartitioned           unless a previous run settled it
      resolveEachChild                 else each on its own
      shuffleToCoPartition             and shuffle whoever is left over
        pickCoPartitionTarget          onto one layout, picked once

GroupPartitionsExec(child, grouping, outputPartitioning, <recipe>, enableSortedMerge)
  GroupPartitionsExec.apply            the only way to build one
    computeGrouping                    pure, takes the child's KeyedPartitioning
    computeOutputPartitioning          pure, takes the same

Retires rewriteGroupPartitions, applyGroupPartitions and innermostGroupPartition's sort-rebuilding half, along with the orElse and the two index spaces they served.

Why are the changes needed?

A join could commit to a pairing its own children then refuse. GroupPartitionsExec gives up its keyed claim when it turns out to regroup a layout that pins undeclared rows to hash(key) % numPartitions, and only the node knows the permutation it performs, so that answer arrived after checkKeyGroupCompatible had skipped both shuffles. The result is a plan ValidateRequirements rejects, and every AQEShuffleReadRule and OptimizeSkewedJoin drops its result on such a stage, so partition coalescing, local read and skew join are all off for it. Now the site that builds the two nodes asks them before returning. This closes SPARK-59272.

The same derivation ran many times. Measured on KeyGroupedPartitioningSuite's "join with two partition keys and matching & sorted partitions": the base derives the grouping 16 times over 48 node instances, this derives it 8 times over 64. A lazy val belongs to one instance, and the columnar and codegen rules rebuild the tree after this rule, so the derivations followed the instances. As a field they follow the planner's decisions instead.

The caller was compensating for the predicate. resolveKeyedPartitioning asked in two steps whether a member needed a node at all, the cheap full-coverage test and then the projection that decides whether a narrowing merges anything. keysSatisfy owns both now and the caller is admitted.find(_.satisfiesAsIs). ValidateRequirements also becomes a real guard: on a KeyedPartitioning([a, b]) with two partitions sharing a = 1 under a ClusteredDistribution([a]), validate was true before and is false now, and the plan spreads rows sharing the cluster key either way.

The rule was not idempotent. Applying it to its own output changed the plan, which matters because AQE re-plans every query stage and two AQE rules hand the tree back to it. Measured over KeyGroupedPartitioningSuite by re-running the rule on each of its own results, 2713 applications in all: the base differs on 2, and it differs by adding a shuffle. Without the check described above this differs on 16, all of them a grouping node added over a shuffle. With it, on none.

Does this PR introduce any user-facing change?

Yes, three plan changes. No public API changes. What does move is GroupPartitionsExec's case class signature and the predicates on KeyedPartitioning and ClusteredDistribution, in sql.execution and sql.catalyst, and MimaExcludes treats both packages as internals in its defaultExcludes section.

A join declines rather than leaving an unvalidatable plan. On the SPARK-59050: SPJ: regrouping a marked layout must not keep the unknown-keyed claim query the second join now takes a keyed one-side shuffle, three shuffles to four, and ValidateRequirements.validate on the executed plan goes from false to true. The alternative was not "no shuffle", it was "no shuffle and no AQE".

satisfies accepts a partitioning whose expression is a cluster key, such as a years(ts) under a clustering naming years(ts). The old reference-level test refused it unless requireAllClusterKeys was set, while the requireAllClusterKeys arm accepted the same shape, so this removes an inconsistency. It is a widening on the default configuration.

A re-planned stage keeps the plan it had. The rule is idempotent now, where the base can add a shuffle on the second pass, so an AQE stage's plan no longer drifts from the one first planned.

How was this patch tested?

Eleven tests, ten of them new and one a rewrite of the existing peel test. Each fails once the decision it pins is put back, so none of them is vacuous.

test	the decision put back
DistributionSuite: satisfies is strict about a projection that merges partitions	the loose keysSatisfy
DistributionSuite: satisfies accepts an expression that is itself a cluster key	the reference-only coverage test
ShuffleSpecSuite: a collection whose members all need grouping still yields them	the collection filter on strict satisfies
EnsureRequirementsSuite: SPARK-58996 only a local sort is looked through	descending through a global sort
a local sort with no grouping under it is left alone	peeling a sort with nothing under it
a grouped side is paired on its own key order	toGrouped for a grouped source too
pushing join key positions into a node re-derives its grouping	copy instead of the rebuild
single-partition children still honour a required partition count	leaving both children alone
a partitioning is never projected onto no position	projecting onto the empty position set
a second pass leaves a pairing this rule already made alone	re-deciding a settled pairing
pushed positions compose with the ones a node already carries	replacing a node's positions instead of composing

One existing expectation moved, Seq(true, true, true) to Seq(true, true, true, true) in the SPARK-59050 regrouping test, and it now asserts ValidateRequirements.validate ahead of the shuffle count because that names the reason the shuffle exists. Three test names lost the word "operation" with the rest of the prose.

A differential sweep is the strongest evidence here. 307200 generated storage-partitioned join plans, 40 partitioning shapes a side crossed with six join types over 32 configuration cells, run against the base and against this change. Each plan is checked twice, that ValidateRequirements passes and that two sides whose shuffles were both skipped really declare the same key sequence. The cases are keyed and set-differenced, not just counted.

	base	this change
plans that skip both shuffles yet disagree on the keys	2024	0
plans ValidateRequirements rejects	6724	5296
plans the planner cannot build at all	29952	29356

Nothing is worse on any of the three, and the sets line up exactly. The 2024 co-partitioning violations are the same 2024 cases as the removed validation failures, which is SPARK-59272 in one number. The 596 fewer planner failures are the same 596 cases that now appear among the rejected plans, so a shape the base could not plan at all is now planned, badly. Zero cases are added to any of the three.

The remaining 29356 failures are shared with the base and are the generator's, not the planner's. They come from Partitioning.createShuffleSpec, which throws when a co-partitioned child reports UnknownPartitioning, and that needs a marked layout that is not grouped. Nothing builds one. The marker is only ever put on a layout KeyedShuffleSpec.createPartitioning has just laid out one partition per key, and both paths that would regroup a marked layout refuse it.

Green: DistributionSuite, ShuffleSpecSuite, EnsureRequirementsSuite, ValidateRequirementsSuite, PlannerSuite, GroupPartitionsExecSuite, ProjectedOrderingAndPartitioningSuite, KeyGroupedPartitioningSuite, DataSourceV2CatalystRuntimeFilterSuite, AdaptiveQueryExecSuite, ExchangeSuite, JoinSuite, CoalesceShufflePartitionsSuite and DataSourceV2Suite, 1624 tests. dev/lint-scala is clean.

All seven plan-stability suites are green as well, 322 tests, so no golden TPCDS or TPCH plan changed. That is the net for a stray plan change out of the EnsureRequirements restructuring, since the rule runs on every query and those two workloads have no storage-partitioned join.

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

Generated-by: Claude Code

[peter-toth]

#58531 (SPARK-59256) and #58552 (SPARK-59285) have both merged, so this is rebased onto master and is one commit again. The whole diff is for review, not just part of it.

Still draft while I keep iterating on it.

[peter-toth]

This is now ready for review.

cc @dongjoon-hyun, @ulysses-you, @cloud-fan, @szehon-ho

[dongjoon-hyun]

Thank you for the thorough write-up, @peter-toth. I went through the whole diff. The direction (decide once in the planner, let the node carry the answer) looks right, and I could confirm that pickCoPartitionTarget and keysCanSatisfy preserve the previous decisions except for the widening you describe. I have one concern that I think needs to be addressed before merging, plus a few smaller ones. The first item is based on reading the code paths; I did not reproduce it by running a query.

1. withNewChildInternal keeps a stale grouping / outputPartitioning when AQE replaces the child

GroupPartitionsExec.withNewChildInternal is copy(child = newChild), so both derived fields survive a child rewrite. The factory doc says this holds "because nothing in the tree hands this node a child that reports a different partitioning", but two AQE rules do exactly that:

• OptimizeShuffleWithLocalRead.createProbeSideLocalRead runs transformDown over the whole stage plan and replaces the probe-side ShuffleQueryStageExec of any BroadcastHashJoinExec with an AQEShuffleReadExec. For a one-mapper-per-task local read its outputPartitioning is the pre-shuffle child's partitioning, otherwise UnknownPartitioning.
• CoalesceShufflePartitions.collectCoalesceGroups coalesces the shuffle stages under a node whose children need no compatible partitioning (e.g. a BHJ) independently, and for a KeyedPartitioning shuffle the coalesced AQEShuffleReadExec reports UnknownPartitioning.

Scenario: a first AQE round plans an SPJ with a one-side keyed shuffle of X onto keyed table Y. On re-optimization Y turns out to be small and the join becomes a BHJ, with X's materialized keyed shuffle stage as its probe side; the BHJ's output partitioning is that keyed layout. A later operator then puts a GroupPartitionsExec over the BHJ output: the push branch of a second SPJ join (an identity grouping keeps the keyed claim even over a marked layout), or a single-child operator on a subset of the keys under allowJoinKeysSubsetOfPartitionKeys. When optimizeQueryStage then applies one of the two rules above:

• Before this PR, copy minted a fresh instance, the lazy vals were re-derived from the new child, the node reported UnknownPartitioning (case o => o), the parent's requirement was no longer satisfied, and ValidateRequirements.validate(applied, ...) reverted the rule.
• With this PR, the node keeps reporting the old keyed partitioning, so validation passes, and at execution grouping.partitions indexes into a child RDD whose partition count and layout have changed (number of mappers, or the coalesced count). That is an exception at best and silently wrong join output at worst.

The old safety net depended on re-derivation, which is what this PR removes. Suggestions:

• Make withNewChildInternal re-derive when newChild.outputPartitioning != child.outputPartitioning (the common wrapper insertions keep the partitioning, so the perf win stays). That needs the derivation inputs back, i.e. keep expectedPartitionKeys (transient) rather than only expectedKeyCount, or recover them from grouping.partitions. The re-derivation also has to tolerate a child with no keyed partitioning without throwing, as the old outputPartitioning did, so that ValidateRequirements can reject the plan instead.
• Add an AQE test where a GroupPartitionsExec sits over a keyed shuffle stage and a local read / coalesce is attempted. KeyGroupedPartitioningSuite has essentially one test with AQE enabled today.

2. satisfies can now project the partition keys on every call

keysSatisfy's second branch calls numPartitionsProjectedOn whenever allowJoinKeysSubsetOfPartitionKeys is on and isFunctionOfClusterKeys is false. That allocates a row per partition and hashes it with the uncached hashCode, and satisfies is asked by EnsureRequirements, ValidateRequirements and every AQE rule application. Previously that projection happened once per position set through the memo in resolveKeyedPartitioning, which this PR also removes, so the same set can now be projected up to three times there (satisfies, eligible, maxBy). Since KeyedPartitioning is immutable, a small transient memo on the position set inside the class would restore the old cost.

3. The assert in withJoinKeyPositions

withJoinKeyPositions asserts expectedKeyCount.isEmpty && reducers.isEmpty && !distributePartitions. On a re-run, an aligned node from an earlier pass can pass through resolveChild unchanged and reach shuffleToCoPartition with a projected paired spec, and the planner then dies on the assert. The pushed positions index the layout g reports, so wrapping (GroupPartitionsExec(g, Some(positions))) is a sound fallback for that shape instead of asserting.

4. Nit: outputPartitioning is the only non-@transient derived field

It used to be @transient lazy val, and every other derived parameter of the node is @transient. This matches ShuffleExchangeExec, so it may be intentional, but a short note would help.

For what it is worth, things I checked and found fine: all keyed members of a PartitioningCollection share one KeyLayout, so the relaxed filter in PartitioningCollection.createShuffleSpec does not widen what ValidateRequirements accepts; QueryPlan.doCanonicalize normalizes the stored KeyedPartitioning through mapExpressions, so exchange reuse keeps working; and the DSv2 scan keeps pruned partitions as None, so runtime filters do not shift the indices the grouping holds.

[ulysses-you]

Thanks Peter. Reviewed at db78554 (base 2870b66).

What I verified independently: DistributionSuite + ShuffleSpecSuite (48 passed),
EnsureRequirementsSuite + GroupPartitionsExecSuite + ValidateRequirementsSuite (111 passed),
KeyGroupedPartitioningSuite (183 passed). I traced the two paths, the position derivations,
the satisfies family, the marked-layout handling and the re-run behaviour against the base,
and I re-derived the two claims the sweep rests on (no operator reports a keyed member beside
a non-keyed one, and a marked layout that is not grouped is unreachable). Both hold in the
code as far as I can see. I agree with the direction.

The extra shuffle on the SPARK-59050 query is the right side of that trade: the alternative
was not "no shuffle", it was "no shuffle and no AQE", and a valid plan is the precondition
for coalescing, local read and skew join.

Two items I would like settled before merge, both below, plus a test ask and three nits.
Neither of the two is a repro I could build, so treat them as questions about an invariant
rather than as bugs I can point at.

One practical ask: the 307200-plan sweep is described in the body but the harness is not in
the PR. If anyone follows up on the marker or on the idempotency machinery, re-running that
sweep is the acceptance test, and today it cannot be reproduced from the branch. Could you
attach the script or describe how to regenerate it?

Also: the body says this closes SPARK-59272 while the title is SPARK-59289. Which ticket
should the release note and any backport hang off?

Scope note: the "operation key" -> "cluster key" rename and the matchesClusterKeys /
isClusterKey / allClusterKeysAmong consolidation are behavior neutral and account for a large
share of partitioning.scala's diff. If this needs a backport, splitting them out would
shorten both the review and the cherry-pick.

[ulysses-you]

grouping and outputPartitioning are carried over a child rewrite (:302-303), and the factory doc (:404-415) states the invariant that nothing in the tree hands this node a child reporting a different partitioning. I verified it for the codegen and columnar wrappers (ColumnarToRowExec, RowToColumnarExec, InputAdapter, WholeStageCodegenExec all report child.outputPartitioning), but not for AQE, and I could not settle it by reading.

A probe at this head shows what the carried fields do once the invariant is broken:

• withNewChildren(Seq(childReportingOnePartition)) still reports numPartitions 2 and groups the two old input partitions, silently, while child.outputPartitioning reports 1.
• A Java serialization round trip leaves grouping == null (NPE on groupedPartitions) and an outputPartitioning whose KeyLayout.partitionKeys is null (KeyLayout:493 is @transient) while numPartitions still reports the old count, so toString and equals change across the wire.

The base held all three of these as @transient ... lazy val (base :77, :226, :309), so a rewritten or deserialized node recomputed them from its new child and could not be stale.

The rewriter I could not rule out is an AQEShuffleReadExec landing between this node and a keyed shuffle child: CoalesceShufflePartitions.isSupported accepts a keyed ENSURE_REQUIREMENTS shuffle (CoalesceShufflePartitions:38-44), and AQEShuffleReadExec reports UnknownPartitioning for a keyed shuffle (AQEShuffleReadExec:100-105). With the stale keyed claim the join above still satisfies its distribution, so ValidateRequirements.validate (AdaptiveSparkPlanExec:205-213) would accept the coalesced plan, where the base rejects it and reverts the rewrite.

Since I could not build a plan with a GroupPartitionsExec directly over a keyed ShuffleExchangeExec, I am not calling this blocking. Could we either re-derive in withNewChildInternal when newChild.outputPartitioning ne child.outputPartitioning, or name in the doc the rules that guarantee the invariant? Either way a future child-rewriting rule fails loudly instead of reading a layout its child no longer has.

[ulysses-you]

This filter also feeds ValidateRequirements (ValidateRequirements:57-68), which is the gate that makes AQE drop a rewrite that breaks co-partitioning (AdaptiveSparkPlanExec:209), and it now admits members whose spec describes a layout the child does not have.

The soundness argument in the comment holds for two grouped sides. For an ungrouped side a duplicated key row means the group is split across the partitions carrying that key, and two such sides are co-partitioned only when their splits line up (partial clustering makes them line up by construction; an ordinary ungrouped source does not).

I could not find a producer of a collection holding an ungrouped keyed member next to a satisfying member: scans report a bare KeyedPartitioning (DataSourceV2ScanExecBase:120-134), and checkKeyedPartitioningInvariant forces all keyed members of a collection to share one layout. One producer worth ruling out explicitly, since it is the natural guess: unions do not report this mixture. UnionExec's all-keyed arm returns a bare KeyedPartitioning.concat (basicPhysicalOperators.scala:994, :1006) and its co-located arm keeps only HashPartitioningLike and SinglePartition members (basicPhysicalOperators.scala:1014, :1016-1020), so a keyed child beside a hash child falls back to UnknownPartitioning.

So this is not a finding. But given the filter is now load-bearing for a safety gate, please either narrow it to what the base admitted, or pin the shape as unreachable with a test. If narrowing: satisfies || (keyed && isGrouped && keysMaySatisfy) restores the base admission set while keeping the new ShuffleSpecSuite test, which needs a fallback because strict filtering would leave the collection empty.

[ulysses-you]

keepArrivedPairing only fires with partially clustered distribution off, and the PR pins one shape. Since idempotency is now an explicit decision rather than a property of the reuse-at-depth mechanism, could we get an assertion set over a handful of SPJ queries, EnsureRequirements.apply(planned) == planned, including a partially clustered one so the re-plan path runs? The 2713 applications / 0 differs number is a measurement, not a test.

[ulysses-you]

This log fires inside pushed, but the pairing can still be discarded by the describesSameKeys gate at :915, which cannot be asked earlier because it needs the built nodes. So the log can announce a pushdown that never lands. The bestPair.isEmpty early return at :689 avoids exactly this for the other discarding case, so this reads as a miss. Moving the log after the gate would keep the messages truthful, and this line is how a reader explains the extra shuffle in the SPARK-59050 query.

[ulysses-you]

positions.map(_) is positions.map(oldPositions) with the Seq used as a Function1 (SeqOps is a PartialFunction). Typer output: ((x$1: Seq[Int]) => positions.mapInt). That is the composition the comment intends, but it reads as positions.map(identity) and would stop compiling if either side became an Array or a Set. positions.map(g.joinKeyPositions.get) says the same thing explicitly.

[ulysses-you]

The doc lists expectedPartitionKeys among the parameters that have to come back through the factory, but that name only exists on apply; the field it feeds is expectedKeyCount (:78). One word so a reader can find the field.

[peter-toth]

@ulysses-you you are right, the two are better kept apart. Sorry for the muddle: on #58339 I asked to take SPARK-59272 inside this refactor, and your original split was the better call.

SPARK-59272 is now its own PR: #58814. It is the pairwise check your ticket describes, 28 lines in checkKeyGroupCompatible, and the SPARK-59050 regrouping test is the only one of KeyGroupedPartitioningSuite's 183 that it moves. On the 307200-plan sweep it accounts for the whole 2024 -> 0 line on its own: the co-partitioning violations it removes are exactly the validation failures it removes, and it adds no case anywhere.

I will rebase this PR onto that one, and onto master once it lands. So the release note for the pairing gate hangs off SPARK-59272, and this PR stays SPARK-59289, the refactor.

On a backport, since you asked which ticket it would hang off: the gate applies to branch-4.x as it stands, because the marker, alignToExpectedKeys and describesSameKeys are all already there. The older branches would need more, describesSameKeys is not on them.