[SPARK-59671][SQL] Validate a co-partitioned pair by its children's pairing

[ulysses-you]

What changes were proposed in this pull request?

ValidateRequirements asks every child of a clustered operator to satisfy its distribution on its own. A ClusteredDistribution is the one distribution an operator can owe two children together rather than one by one, so this PR judges such an operator's children by their pairing, and a pair aligned without grouping (which partially clustered distribution builds on purpose) is one the sides agree on while neither is grouped.

• ValidateRequirements asks the per-side check only of children that answer for themselves or report no keyed layout; the rest are judged together, on the specs the planner builds for them (PartitioningCollection.specsForPairing, a new private[sql] helper on the existing object: a keyed member on the planner's own admission keysMaySatisfy, any other member on satisfies, and a pinned count asked as it stands). One member of the first side has to pair with every other side, which is the question EnsureRequirements commits a pair on (committed, over the pair agreeingPairs picked).
• What those specs hold is what createShuffleSpec makes of a member: its own layout, or under spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys the projection onto the keys it covers, the layout the alignment about to be planned emits. That is the question the planner asks when it picks the member a pair is planned on, less the coverage of every operation key it additionally requires there: spark.sql.requireAllClusterKeysForCoPartition is a skew heuristic, and a member covering a subset of the operation keys is a sound pairing.
• What the pairing cannot say is how the two sides hold a key's rows: a spread side and one that repeats the whole group report the same keys as two sides that split the key between them, and no layout distinguishes those, so that rests on the producer, which is the join path: checkKeyGroupCompatible is where such a pair is planned, and every other co-partitioning operator's children are grouped before the rule is done.

The catalyst change is additive: a new helper beside the existing ones, and no change to PartitioningCollection.createShuffleSpec or maySatisfyAfterProjection. Those two still name ValidateRequirements as a caller in their comments, which this change leaves as they are.

Why are the changes needed?

A storage-partitioned join planned by partially clustered distribution aligns its sides without grouping either of them: the side that keeps its splits spreads them, the other replicates its group across them, so both report keys that repeat on purpose. Such a pair fails the per-side check at the join node even though the two sides agree key by key, and AdaptiveSparkPlanExec.optimizeQueryStage validates a stage's whole candidate plan before accepting an AQEShuffleReadRule change, so every shuffle read in that stage stays uncoalesced, unrelated ones included.

Measured on d39cc1784c0 (base) versus this head, both with spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled=true:

Probe	base	head
A pair the rule planned (EnsureRequirements.apply)	ValidateRequirements.validate false	true
The aggregate's shuffle read in the join's stage	not coalesced (AQEShuffleReadExec.hasCoalescedPartition false)	coalesced
A three-table chain whose outer join reads a projection over both key columns	validate false	true, shuffle-free, still coalesced

The planner-side half of this hazard landed in SPARK-59272, which declines a pairing whose sides no longer declare the same aligned key sequence: that closes the pairs a GroupPartitionsExec gives up on, which should not be built at all. A pair aligned without grouping is the other half, and it is built on purpose, so the validator has to read the pairing instead of each child on its own.

Does this PR introduce any user-facing change?

No. No plan changes unless the plan already contains such an alignment, and no new configuration.

How was this patch tested?

New tests, by what each one is there for (the four marked base-failing were re-run against d39cc1784c0 and fail there):

• ValidateRequirementsSuite
  • the exemption and its refusals (base-failing): a keyed pair that repeats its keys position by position passes, while key sets that disagree, a differing key order, a hashed side, a lone clustered child and two sides that each satisfy on their own but do not line up all fail;
  • the rule's own pair (base-failing): a pair planned by EnsureRequirements under partially clustered distribution passes;
  • the collection shape (base-failing): a side reporting several keyed alternatives is judged on whichever of them pairs, not on the first one;
  • guards, which pass on base as well and pin this change's clauses: a subset-keyed pair still passes, a lone clustered child still owes its own grouping, a pinned partition count is still asked of a pair, a pair that lines up still owes its operator an ordering, and a collapsed pair is not admitted on its pairing alone.
• KeyGroupedPartitioningSuite
  • with partially clustered distribution on, the aggregate's shuffle read in the stage holding a two-table join coalesces (base-failing: it is not coalesced); the test also pins that the join side shuffles nothing, that the chain shuffles once, and that the join shares the final stage with that read, which is what makes the coalesce a decision the join can block;
  • a three-table chain whose outer join reads a projection that keeps both key columns stays shuffle-free, passes validation and keeps its coalescing (base-failing at the validation assertion).
• ShuffleSpecSuite: the specs a side offers are the planner's own: a keyed member offers the layout it reports, a member covering part of the operation's keys offers the projection onto the keys it covers under the subset permission, a member whose keys do not cover the clustering offers nothing, and a member that is not keyed offers its own spec.

Ran locally: ValidateRequirementsSuite, KeyGroupedPartitioningSuite, EnsureRequirementsSuite, GroupPartitionsExecSuite, AdaptiveQueryExecSuite, PushDownLocalSortSuite, ShuffleSpecSuite: all green, with scalastyle clean for catalyst and sql, main and test sources.

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

Generated-by: Claude Code (qwen3.8-flash)

[ulysses-you]

cc @dongjoon-hyun @peter-toth @cloud-fan thank you

[peter-toth]

Thanks for the PR, @ulysses-you!

The per-side satisfies check refuses a pair that partially clustered distribution builds on purpose, since neither side is grouped, and because AQE validates a stage's whole candidate plan, one such join keeps every shuffle read in its stage uncoalesced. Judging an operator with two clustered children by their pairing is the right read of what such an operator owes, and for three or more children the new form is stricter than the old specs.tail.forall(_.isCompatibleWith(specs.head)), since it now needs one member of the first side to pair with all the others. I re-measured the new tests against d39cc1784c0: four in ValidateRequirementsSuite fail there, and so do both end-to-end ones, at the assertions the description names (0 did not equal 1 the aggregate's shuffle read must coalesce, and the validate one). 31 catalyst plus 209 sql tests are green on this head. One substantive point below: the new admission drops the isGrouped clause its sibling helper keeps, so the pairing is judged on a layout the plan does not contain, in every configuration rather than only where such a plan is legitimate.

Unrelated heads-up, since it lands in the same file: #58943 tightens KeyedShuffleSpec.isCompatibleWith so that a pair whose keys still need reducing onto one key space is not compatible as it stands. Your pairing check inherits that. I checked it does not reach the shapes here, since a partially clustered pair carries the same transform on both sides and partial clustering rules out reducing anyway.

Non-blocking

• 1. Admission assumes a grouping the plan does not have: keysMaySatisfy answers for an ungrouped member "yes, once something groups it", and the spec is then built from the ungrouped layout, which is the clause SPARK-59289 deliberately kept in maySatisfyAfterProjection three days ago. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1479]
• 2. Three comments name a caller this PR removes: PartitioningCollection.createShuffleSpec, maySatisfyAfterProjection and SinglePartitionShuffleSpec.isCompatibleWith each justify a choice by naming ValidateRequirements as the caller, and none of them is reached from it any more. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1473]
• 3. The ticket's affected versions look too narrow: SPARK-59671 says 5.0.0 only, while the clause the symptom rests on, isGrouped && keysSatisfy in KeyedPartitioning.satisfies0, is on branch-4.3 as well (line 770 there), and GroupPartitionsExec and partially clustered distribution both exist from branch-4.2. Read from the code only, not measured on those branches. Either the affected list wants the older lines, or the description wants a sentence on why they are not affected. Worth saying what the intended branches are in any case: maySatisfyAfterProjection is master and branch-4.x only, so this would not cherry-pick cleanly below 4.4.

Alternatives

• 4. Let the layout carry the distinction instead of trusting the producer: not a request to change this PR, but the shape that would let the validator decide the question finding 1 is about. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:68]

Minor

• 5. The equivalence with the planner is not quite the stated one: the planner also groups an ungrouped member before it builds the spec, which is the difference that matters here. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1467]

[peter-toth]

Finding 1. keysMaySatisfy is if (isGrouped) keysCanSatisfy else mayGroupToSatisfy, so for an ungrouped member it answers "yes, once something groups it", and this then builds the spec from the ungrouped layout. The planner may ask that question, because it is about to insert the node: createKeyedShuffleSpecs.tryCreate builds the spec from partitioning.toGrouped (EnsureRequirements.scala:1135). A finished plan has nobody left to insert one.

That is the clause maySatisfyAfterProjection keeps, and it was a decision rather than an omission. SPARK-59289 made satisfies strict about a projecting node and had to widen this filter back for it; it widened the projection and not the grouping, and wrote down why (partitioning.scala:1442-1450): "A partitioning that is not grouped is not admitted, even though a node would also group it. ... the caller feeds ValidateRequirements as well as the planner, so it does not widen what a finished plan is checked against."

Your first unit test states the cost: two sides reporting [1, 1, 2], neither grouped, no partially-clustered conf set, and validate says yes. That layout pair has the two readings your own comment in validate names. One side spread and the other replicating the whole group, which is sound and is what the rule builds. Or both sides splitting the key between two partitions, which is not: rows of key 1 in left partition 0 never meet rows of key 1 in right partition 1. The second is not reachable from EnsureRequirements today, which is why this is Non-blocking, and it is also why the guarantee now rests entirely on the producer with nothing enforcing it.

The projection half has the same shape: under allowKeysSubsetOfPartitionKeys an ungrouped member contributes project(...).toGrouped, so two sides whose key multiplicities differ can project onto the same grouped key list and pair, while the plan holds neither the projection nor the grouping.

Narrowing the widening to the shape that needs it keeps the rest as strict as it was:

case k: KeyedPartitioning =>
  // An ungrouped side is a plan only partially clustered distribution builds.
  val mayBeUngrouped = SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled
  Option.when((k.isGrouped || mayBeUngrouped) &&
    distribution.requiredNumPartitions.forall(_ == k.numPartitions) &&
    k.keysMaySatisfy(distribution))(k.createShuffleSpec(distribution))

The per-side exemption in validateInternal wants the same condition, and a test for the refusal in the default configuration would pin it, which is what your first test asserts the other way round today. One more reason to spell the condition out rather than admit everything: SPARK-59436 (#58771) will want the same relaxation for a skew-split side, and its own config then belongs in that disjunction.

[peter-toth]

Finding 2. Three comments justify a choice by naming ValidateRequirements as the caller that constrains them, and after this PR none of them is reached from it:

• partitioning.scala:1389-1395, in PartitioningCollection.createShuffleSpec: "The set matters because ValidateRequirements builds a spec from a finished plan through here." It does not any more, since specsForPairing flattens the collection and builds the member specs itself.
• partitioning.scala:1442-1450, in maySatisfyAfterProjection: "the caller feeds ValidateRequirements as well as the planner, so it does not widen what a finished plan is checked against." The finished-plan check now uses a wider admission, so that sentence no longer describes the code (finding 1).
• partitioning.scala:1690-1696, in SinglePartitionShuffleSpec.isCompatibleWith: "The one production caller that can put a collection on the other side is ValidateRequirements' specs.tail.forall(_.isCompatibleWith(specs.head)), and there the stricter answer is the safer one." That line is the one this PR deletes, and specsForPairing returns leaf specs only, so the case ShuffleSpecCollection(specs) => specs.forall(isCompatibleWith) arm is now reachable only from ShuffleSpecSuite. The forall-versus-exists divergence it documents has no production caller left to be safe for.

The description says the first two are left as they are. Each states an invariant rather than a pointer, so all three are worth correcting here.

[peter-toth]

Finding 4. Not a request to change this PR: the narrowing in finding 1 uses a configuration as a proxy for the property this comment names, and the property itself could be carried instead.

What the pairing cannot tell apart is which role an ungrouped side plays, and the node that made it ungrouped does know: GroupPartitionsExec with distributePartitions = true spreads a key's splits, while the other side repeats its whole group. If KeyLayout carried that role, the way it already carries mayContainUnknownPartitionKeys for another per-row property, then isCompatibleWith could require at most one spread side with a repeating partner, and the validator would accept the sound pair and refuse the unsound one on its own. No configuration would be read, and it would cover SPARK-59436's skew-split side for free rather than needing another disjunct.

The counter-arguments are real, which is why this is not a request. KeyLayout is master and branch-4.x only, so a fix that should reach branch-4.2 cannot be built on it. The marker has to be produced by GroupPartitionsExec and kept across copy, and every consumer that compares layouts has to agree what it means. And it is a larger change than the validator question that prompted it. Worth its own ticket if you think the direction is right; happy for it to sit on either of our lists.

[peter-toth]

Finding 5. The doc, and the description, say these specs are "built the way the planner builds them" and that the admission is the planner's "less the coverage of every operation key it also requires there". There is a second difference, and it is the one this change turns on: createKeyedShuffleSpecs.tryCreate builds the spec from the grouped form of an ungrouped member, val grouped = if (partitioning.isGrouped) partitioning else partitioning.toGrouped at EnsureRequirements.scala:1135, because it is about to insert the node that groups it. This helper deliberately does not, which is the right call for a finished plan and worth saying outright instead of implying an equivalence that does not hold.

Same for "the question EnsureRequirements commits a pair on (committed, over the pair agreeingPairs picked)": on the push path committed is sidesDeclareSameKeys, a describesSameKeys comparison of the two children's declared layouts, and only the compatibleAsIs path asks isCompatibleWith.

[dongjoon-hyun]

Two things to add on top of the existing review rather than repeat it.

Finding 1 reproduces, with numbers. On bb183663ba8, the projection half of
his finding is not hypothetical. Taking your own ValidateRequirementsSuite
case "the same keys in a different order are not aligned" and flipping only the
config:

[1,1,2] vs [1,2,1], subsetConf=false -> validate=false
[1,1,2] vs [1,2,1], subsetConf=true  -> validate=true
[1,1,2] (3 partitions) vs [1,2] (2 partitions), subsetConf=true -> validate=true

At the catalyst level specsForPairing reports the divergence directly:

left  child numPartitions = 3    left  spec numPartitions = 2
right child numPartitions = 3    right spec numPartitions = 2
compatible = true

So the config default is the only reason that test still passes, and the second
line means the validator stops noticing that two children have different
partition counts at all.

coPartitioning is not only the join path. It is true for any operator with
more than one child all requiring ClusteredDistribution, which includes
CoGroupExec (objects.scala:637) and FlatMapCoGroupsInBatchExec
(FlatMapCoGroupsInBatchExec.scala:56-59). For a cogroup, "the two sides agree
key by key while neither is grouped" is not a sufficient input: it pairs left
group with right group per partition, so a key split across two partitions on
each side yields two partial cogroups rather than merely a missing shuffle.

The comment in validateInternal argues every non-join co-partitioning
operator's children are grouped before the rule is done, and that holds today -
checkKeyGroupCompatible returns None for a non-join parent
(EnsureRequirements.scala:598-600), so resolveEachChild groups them. But
ValidateRequirements is also the gate for third-party AQEShuffleReadRules
registered through adaptiveRulesHolder
(AdaptiveSparkPlanExec.optimizeQueryStage), which is exactly where an
invariant EnsureRequirements maintains stops being self-evident. Whatever
shape the narrowing in finding 1 takes, scoping the per-side exemption to the
operators the pairing argument actually covers would keep the guard for the
rest.