[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 an operator whose children all owe one is judged on their mutual layout, and a pair aligned without grouping (which partially clustered distribution builds on purpose) is one the sides agree on while neither is grouped.

• The children of such an operator are judged together, on the layouts they report (PartitioningCollection.specsForPairing, through its reportedSpecOf), and one member of the first side has to pair with every other side. That is every multi-child clustered operator, not only a join: one that zips corresponding partitions, a cogroup for instance, reads a layout both children have to hold together. Every other child, an operator with a single clustered child included, keeps the per-side check.
• What a side offers is the layout it reports, restricted to the partition expressions that carry a cluster key: no key is deduped and none is re-sorted, so the count and the order are the member's own. KeyedPartitioning.createShuffleSpec is the layout a node would emit and is deliberately not what the validator reads: it dedups and sorts, and its count is the grouped one, which the plan does not hold.
• The one member a finished plan may report without satisfying the distribution is the ungrouped shape partially clustered distribution spreads, and one producer builds it: EnsureRequirements.checkKeyGroupCompatible, which is entered for a sort-merge or shuffled-hash join. The waiver is confined to that configuration, that producer and the join types a side may be duplicated for, which the producer applies before it spreads one: no side of a FullOuter join may be. The answer is passed down to the pairing so it cannot be read one way in one place and the other way in the next. What is left to the member itself, its partition count and the permission for the collapse it went through, is asked in the spec: a member whose keys were collapsed is admitted ungrouped only where the planner would agree to group it, and an ungrouped member keeps the keys it reports.
• 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.

Against the planner's own admission of a member (EnsureRequirements.createKeyedShuffleSpecs), the coverage of every operation key (spark.sql.requireAllClusterKeysForCoPartition) is not asked here: it is a skew heuristic, and a member whose partitioning keys cover only a subset of the operation's keys is a sound pairing.

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.

SPARK-59688, merged in the meantime, tightened KeyedShuffleSpec.isCompatibleWith to answer for a pair as it stands, which is the question asked here; a pair that took the reduce still answers it through hasSameReducedKeys, so the shapes in this PR are unaffected.

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. Those that assert the new acceptance pass only with the main-code change, since the base asked satisfies of each child on its own and an ungrouped keyed child fails that in every configuration; those that assert a refusal hold on the base as well, which is where the guards are.

• ValidateRequirementsSuite
  • the exemption and its refusals: a keyed pair that repeats its keys position by position passes where partially clustered distribution is on, 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: a pair planned by EnsureRequirements under partially clustered distribution passes;
  • the producer boundary: the same ungrouped pair a shuffled-hash join reads is refused for an as-of join, which is a ShuffledJoin too and builds no spread side;
  • the collection shape: a side reporting several keyed alternatives is judged on whichever of them pairs, not on the first one, and a side offering no member keyed on the join keys does not pair;
  • a multi-child clustered operator that is not a join: a cogroup over the pair a join reads is refused, and the mutual check on the layouts its sides report refuses a differing partition count and a differing key order while accepting two sides that agree;
  • the key a side is judged on: a side partitioned on [a, b] serving an operation on [a] is judged on the key the operation clusters on, as it reports it, so the pair stands and a side whose keys run the other way is refused;
  • the admission's edges: an ungrouped side is refused where nothing builds one; a pair is judged on the layouts its sides report rather than on a projection no node made, for a key order and for a partition count; a collapsed pair is refused while its collapse may not be grouped, partial clustering on or not; a lone clustered child still owes its own grouping; and a pair that lines up still owes its operator an ordering.
• KeyGroupedPartitioningSuite
  • with partially clustered distribution on, the aggregate's shuffle read in the stage holding a two-table join coalesces, which the base blocks by refusing the pair; 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.
• ShuffleSpecSuite: the specs a side offers are the layouts it reports: a keyed member offers its own partitions under the operation's key, keeping its count and the order it reports; an ungrouped one is offered only where something builds one and produces it; a member whose partitioning keys are a superset of the operation's keys offers its own keys rather than a projection onto them; a marked member is not offered through a projection; a member whose keys do not cover the clustering offers nothing; a member that is not keyed offers its own spec; a collapsed member the planner would not group is not offered ungrouped; and a count the operation pinned is asked as it stands.

Ran locally, all green: ShuffleSpecSuite and DistributionSuite (52 tests), and in one run ValidateRequirementsSuite, EnsureRequirementsSuite, GroupPartitionsExecSuite, KeyGroupedPartitioningSuite and ProjectedOrderingAndPartitioningSuite (362 tests), with dev/lint-scala 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.

[ulysses-you]

Thanks @peter-toth! Narrowed in ce59a03. specsForPairing now offers the layout each member reports (KeyedShuffleSpec(k, distribution), nothing projected and nothing grouped), admitted on satisfies, so the member the plan holds is the member judged. An ungrouped one is the single exception, and it is confined to the configuration that builds one through PartitioningCollection.mayUngroupedMember; validateInternal's per-side exemption reads the same predicate, so the two cannot drift. keysSatisfy is asked from outside the class now, so the family doc lists it beside keysMaySatisfy, and that is the question the ungrouped exception asks.

The refusal in the default configuration is pinned in ValidateRequirementsSuite (the first test, after the conf block) and in ShuffleSpecSuite (the member is offered only where something builds one), and the shapes you would expect to still be refused are: a differing key order and a differing partition count, now with both permissions on, so the answer is the layouts' own rather than the admission's.

Classification: regression of this change. The base asked satisfies per side, which an ungrouped keyed child fails in every configuration, so nothing here is a pre-existing hole being made visible.

[ulysses-you]

One correction to my reply above, from this round's producer-boundary review: the spec is not the member's layout verbatim. It is that layout viewed on the key the operation clusters on, so a partition expression carrying no cluster key is left out, with the member's own count and its own key order, no key deduped and none re-sorted. PartitioningCollection.reportedSpecOf states that and why, and the point of the reply stands: createShuffleSpec's grouped layout is not what the validator reads. Head is f95e44a.

[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.

[ulysses-you]

Thanks @peter-toth! All three corrected in ce59a03, each now stating the invariant rather than a caller: PartitioningCollection.createShuffleSpec names the planner's shuffleToCoPartition as the reader of that admission set; maySatisfyAfterProjection says the validator does not read it and points at specsForPairing; SinglePartitionShuffleSpec says no production caller puts a collection on the other side, which I checked by enumerating the isCompatibleWith call sites: pickCoPartitionTarget flattens to leaves before it pairs, and specsForPairing offers leaves.

One more site said the same thing and is corrected with them, outside the list you gave: the SPARK-59289 test comment in ShuffleSpecSuite, which read "the admission set here is what a finished plan is checked against by ValidateRequirements".

[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.

[ulysses-you]

Thanks @peter-toth! Direction agreed, and it wants its own ticket rather than this PR: KeyLayout is master and branch-4.x only, while the repair here should reach further back.

With the narrowing, no configuration stands in for the role any more: what is read says which shape a plan may report, not which part a side plays in the alignment.

[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.

[ulysses-you]

Thanks @peter-toth! Both corrected in ce59a03. The helper's doc states the two things the planner does and it deliberately does not, since they build a layout the plan does not hold: no projection onto the operation keys under allowKeysSubsetOfPartitionKeys (a plan the permission applies to holds the projection already, in the grouping node EnsureRequirements inserted), and no grouping of an ungrouped member or re-sorting of a grouped one, toGrouped being what tryCreate builds from.

satisfiesForPairing now says this is the question the compatibleAsIs path asks, and that the push path commits on the two sides it builds (committed, a describesSameKeys comparison of the layouts the reduce left), which is sound to ask here because a pair that took the reduce answers the strict question through hasSameReducedKeys.

[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.

[cloud-fan]

Review summary

The patch conflates a planning-time candidate layout with the layout present in a finished plan across the validator, its positive fixture, and its contract comments. The validator needs to restore that distinction, keep the ungrouped exception confined to the partially clustered shuffled-join producer mode, and align the tests and comments with the same boundary.

The existing discussions already capture the blocking finished-layout issue, the false planner-equivalence wording, and the three stale caller comments. I found one additional documentation error: the new subset-key comments reverse the configuration's operation-keys-subset-of-source-partition-keys direction. The end-to-end AQE tests are useful counterevidence and should remain; they prove the intended producer-built path while the validator and focused negative coverage are narrowed around it.

Findings

4 total: 0 P0, 1 P1, 0 P2, 3 P3.

Blocking (P1)

• Validate the layout present in the finished plan — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1479 — already raised in an existing discussion.

Nit (P3)

• Correct the claimed planner equivalence — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1461 — already raised in an existing discussion.
• Update the removed validator caller claims — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1473 — already raised in an existing discussion.
• Fix the subset-key direction in the comments — sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/ValidateRequirementsSuite.scala:207 — see inline.

Existing discussions

• Suppressed duplicate: Validate the layout present in the finished plan — P1 at sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1479 — existing discussion
• Suppressed duplicate: Correct the claimed planner equivalence — P3 at sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1461 — existing discussion
• Suppressed duplicate: Update the removed validator caller claims — P3 at sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1473 — existing discussion
• existing discussion — The concrete projection and physical-partition-count examples exactly reproduce the promoted defect, and the non-join observation identifies the same missing producer boundary in the validator's broad coPartitioning exemption.

[cloud-fan]

Nit (P3): This describes the subset in the opposite direction from the configuration contract. allowKeysSubsetOfPartitionKeys permits the operation keys to be a subset of the source partition keys, so an [a, b] partitioning may be projected for an [a] operation. It does not make the planner build an [a] partitioning for an [a, b] operation; with requireAllClusterKeysForCoPartition at its default, that shape is rejected. Please update this fixture comment and the matching helper text to use the operation-subset-of-source direction.

[ulysses-you]

Thanks @cloud-fan! Corrected in ce59a03, in the fixture and in the helper text: the permission applies where the operation's keys are a subset of the source's partitioning keys, so a [a, b] partitioning may be projected for an [a] operation, and not the other way round.

The fixture now uses that direction: the sides report the projection's own layout, which passes, and the same pair one step earlier, reporting the source's keys that no node projected, is refused. The ShuffleSpecSuite arm had the same wording and is corrected with it.

[cloud-fan]

Confirmed. specsForPairing can project and group an ungrouped KeyedPartitioning even though the finished child does not report that layout, and the broad coPartitioning exemption is not limited to the join producer path. I am retaining this as a correctness finding; the repair should validate as-is layouts except for the explicitly safe producer-specific ungrouped case, with coverage for subset-enabled/partial-disabled and non-join operators.

[ulysses-you]

Thanks @peter-toth, @dongjoon-hyun and @cloud-fan! Addressed in ce59a03, which answers both comments above and the inline ones.

The validator judges as-is layouts now. A keyed member contributes the layout it reports -- KeyedShuffleSpec(k, distribution), with no projection onto the operation keys, no grouping and no re-sorting -- and is admitted on satisfies. The explicitly safe ungrouped case is the one exception: it is confined to the partially clustered producer configuration through a single predicate, PartitioningCollection.mayUngroupedMember, which specsForPairing and the per-side exemption both read, so the two cannot drift. The co-partitioning reading is ShuffledJoin-only, so CoGroupExec and FlatMapCoGroupsInBatchExec keep the per-side check.

That closes the measured shapes: a projected key count no longer stands in for the child's own, so three partitions against two is refused; a projection onto distinct sorted keys no longer normalizes a key order; and an ungrouped side is admitted only where something builds one. Coverage for the two configurations asked for is in ValidateRequirementsSuite: subset-enabled/partial-disabled, where the projection's own layout passes and the source's keys, which no node projected, are refused, and a non-join operator, a cogroup over the pair a join reads, refused.

Two consequences, stated rather than left implicit:

• A join's distributions never pin requiredNumPartitions (ShuffledJoin.scala:57), so with the co-partitioning reading limited to the join, specsForPairing's count clause is no longer reachable from the validator. The pinned-count test moved to ShuffleSpecSuite, as a direct test of the helper.
• Where the operation's keys are a subset of the source's partitioning keys and the projection would merge no partition, satisfies admits the member but the pairing refuses it, the expression the operation does not cluster on mapping to no cluster key. The base accepted that shape through a projection it did not hold; here the plan has to hold the grouping node. This is the one place the validator is stricter than the base.

The comment and doc sites raised inline are answered in their threads.

[cloud-fan]

Review summary

ValidateRequirements now conflates a producer-specific ungrouped waiver with generic multi-child compatibility and does not preserve the planner's full admission predicate. The repair should separate those responsibilities so every clustered multi-child consumer retains mutual compatibility validation while only planner-supported join layouts receive the exact ungrouped waiver.

The latest revision fixes the earlier projection/grouping issue and stale caller comments. This re-review found one blocking regression for grouped non-join consumers, two additional correctness gaps in ungrouped admission, and two misleading changed comments.

Findings

5 total: 0 P0, 1 P1, 2 P2, 2 P3.

Blocking (P1)

• Restore compatibility checks for non-join clustered consumers — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:87 — see inline.

Non-blocking (P2)

• Limit the ungrouped waiver to producer-supported joins — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:57 — see inline.
• Preserve the planner's collapsed-layout gate — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1488 — see inline.

Nit (P3)

• Describe the AQE failure as pre-fix behavior — sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:9233 — see inline.
• Correct the all-cluster-keys direction — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1469 — see inline.

Shared repair plans

Shared repair plan 1

Covered findings:

• Limit the ungrouped waiver to producer-supported joins — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:57
• Preserve the planner's collapsed-layout gate — sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1488

Recommended change: Define one explicit finished-plan ungrouped-pair admission contract that is restricted to planner-supported SMJ/SHJ consumers, preserves as-reported layouts, enforces the planner's collapsed/subset configuration gate, and add focused operator/configuration coverage.

Why this works: Separate operator capability from layout eligibility: gate the per-side waiver on the same concrete producer-supported join family, and have the pairing helper reuse an as-held eligibility predicate that includes required partition count, key satisfaction, and the collapsed-layout permission without projecting or grouping the reported member.

Scope: Align finished-plan ungrouped pairing with the concrete producer and its layout/configuration contract.

Compatibility: Finished plans are judged on their reported layouts, and valid planner-produced partially clustered SMJ/SHJ pairs remain accepted.

Risks: An over-narrow predicate could reintroduce the intended AQE coalescing regression for valid partially clustered SMJ/SHJ plans. Calling keysMaySatisfy directly could accidentally project or group a layout that the finished plan does not hold.

Constraints: Validation must continue to judge the exact layout reported by the finished plan. The intended partially clustered spread/replicate pair must remain accepted.

Success: A planner-produced partially clustered SortMergeJoinExec or ShuffledHashJoinExec pair still validates and does not block unrelated AQE shuffle coalescing. An ungrouped SortMergeAsOfJoinExec pair is not admitted solely because partial clustering is enabled. A collapsed ungrouped pair is rejected when subset permission is disabled even when partial clustering is enabled. Pairing never projects, groups, or re-sorts a layout not reported by the finished plan.

Re-review status

Prior AI findings: 4 addressed, 0 still present; additional unresolved findings in this review: 5.

New attribution: 4 newly introduced, 1 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — The revision fixes the original projection/grouping and broad non-join per-side exemption, but the same requested producer boundary remains too broad for SortMergeAsOfJoinExec and the replacement drops generic compatibility checking for grouped non-join children. These are related new defects rather than exact duplicates of the original reproduced cases.
• existing discussion — The local review response correctly retained the as-reported-layout and producer-specific boundary. The revision resolves its original examples, but its ShuffledJoin-wide producer proxy and removal of the non-join compatibility check leave related defects in the same boundary.
• existing discussion — The author's update accurately describes the as-held-layout repair, but its claim that ShuffledJoin-only scoping preserves the rest of the guard is incomplete: SortMergeAsOfJoinExec is also a ShuffledJoin without the producer invariant, and grouped non-join children no longer receive any mutual compatibility check.

PR description suggestions

• Update the description's requireAllClusterKeysForCoPartition paragraph: disabling that coverage check permits partitioning keys that cover only a subset of the operation keys; the partitioning-key-superset case belongs to allowKeysSubsetOfPartitionKeys.

[cloud-fan]

Non-blocking (P2): This waiver is broader than its producer invariant: EnsureRequirements.checkKeyGroupCompatible builds spread/replicate layouts only for SortMergeJoinExec and ShuffledHashJoinExec, while SortMergeAsOfJoinExec also matches ShuffledJoin. With partial clustering enabled, an AQE rule candidate over aligned repeated AS-OF children can therefore pass without either side being known to repeat the whole key group, and the partition-local AS-OF scan can miss or duplicate matches. Please restrict the waiver to the producer-supported join kinds or carry equivalent explicit provenance.

See Shared repair plan 1 in the review body.

[ulysses-you]

Thanks @cloud-fan! Restricted in f95e44a, and not by a second copy of the list: the operator kinds now live in one predicate, ShuffledJoin.partiallyClusteredJoinType, which EnsureRequirements.checkKeyGroupCompatible dispatches on and ValidateRequirements reads through .isDefined, so the two cannot drift apart. A SortMergeAsOfJoinExec is excluded with them, and the pair a shuffled-hash join accepts is pinned as refused for it.

[cloud-fan]

Blocking (P1): For a non-ShuffledJoin, this branch now returns true after per-child satisfaction without the merge-target validator's mutual shuffle-spec compatibility check. CoGroupExec and FlatMapCoGroupsInBatchExec zip corresponding partitions, so individually satisfying children with different counts can fail at execution, while equal-count but misaligned layouts can emit partial or missing groups. Please restore generic compatibility validation for multi-child clustered consumers and keep only the ungrouped waiver join-specific.

Recommended change: Restore a generic mutual compatibility check for every multi-child operator whose requirements are clustered, while retaining the producer-specific ungrouped exception only on the eligible shuffled-join path; add grouped cogroup coverage for incompatible counts and layouts.

Why this works: After individual distribution and ordering validation, compare compatible shuffle specs for all multi-child clustered consumers. Use the specialized as-held pairing only where an eligible join needs the ungrouped waiver; non-join consumers must still prove a mutually compatible grouped layout.

Scope: Reestablish cross-child compatibility as a generic validator responsibility without broadening the ungrouped join waiver.

Compatibility: The join-specific partially clustered waiver remains separate from ordinary grouped compatibility validation.

Risks: A generic check must not apply the join-only ungrouped semantics to operators that emit one result per key group.

Constraints: Single-child operators continue to owe only their own distribution and ordering. Valid grouped multi-child plans remain accepted.

Success: Two individually grouped cogroup children with different partition counts are rejected before execution. Equal-count but mutually misaligned grouped cogroup layouts are rejected. Compatible grouped non-join children and valid shuffled joins continue to validate.

[ulysses-you]

Thanks @cloud-fan! Restored in f95e44a. The pairing is asked of every multi-child clustered operator again, which is what the base did, and only the ungrouped waiver stays producer-scoped: the waiver's condition is the multi-child cluster, the partially clustered configuration and the producer's kinds, computed once and passed into the pairing. The cogroup test now covers the mutual check from three sides, two sides holding the same grouped layout passing, a differing partition count refused, and a differing key order refused.

[cloud-fan]

Nit (P3): This states the regression as present behavior, but the assertion below requires the unrelated AQEShuffleReadExec to have a coalesced partition. Please make this historical (for example, the join used to keep unrelated shuffles uncoalesced) so the explanation agrees with the behavior this test pins.

[ulysses-you]

Thanks @cloud-fan! Reworded in f95e44a: the failure is stated against the base now, "on the base a storage-partitioned join whose sides are aligned but not grouped kept every shuffle in its stage uncoalesced", and the closing sentence says the assertions below pin the read the join's presence must leave alone.

[cloud-fan]

Non-blocking (P2): With partial clustering enabled, this fallback asks keysSatisfy, which does not apply the isCollapsed gate. EnsureRequirements asks keysMaySatisfy; for an ungrouped collapsed layout, mayGroupToSatisfy rejects it while allowKeysSubsetOfPartitionKeys is off. Two matching collapsed sides can therefore validate even though the planner refuses that shape. The added negative leaves partial clustering disabled and exits before this branch, so please preserve the planner's collapsed-layout permission here and cover the two flags together.

See Shared repair plan 1 in the review body.

[ulysses-you]

Thanks @cloud-fan! The permission is back in f95e44a, read once and shared: collapsedLayoutMayBeGrouped is what mayGroupToSatisfy and the ungrouped admission both ask, so a member a finished plan reports ungrouped is admitted only where the planner would agree to group it. The negative covers the two flags together, partially clustered distribution on with the subset permission off, and a unit test pins that the member is not offered in that configuration.

[cloud-fan]

Nit (P3): Disabling requireAllClusterKeysForCoPartition skips allClusterKeysCovered, so it can admit partitioning on [a] for an operation on [a, b]: the partitioning keys cover only a subset of the operation keys. A partition-key superset is the separate projection case controlled by allowKeysSubsetOfPartitionKeys. Please reverse this example/rationale so it describes the configuration used here.

[ulysses-you]

Thanks @cloud-fan! Corrected in f95e44a, in the helper's text and in the PR description: disabling the coverage check admits a member whose partitioning keys cover only a subset of the operation's keys, and the partition keys being a superset is the separate projection case under allowKeysSubsetOfPartitionKeys.

[cloud-fan]

The as-held-layout fixes are in place, but I found two remaining boundary issues: the ungrouped waiver currently includes SortMergeAsOfJoinExec even though EnsureRequirements never builds its spread/replicate pairing, and non-join cogroup operators lost the old mutual compatibility check for already-grouped children. Please separate the producer-specific waiver from the generic multi-child compatibility validation and add focused AS-OF plus grouped-cogroup coverage.

[ulysses-you]

Thanks @peter-toth, @dongjoon-hyun and @cloud-fan! All of the review's findings are in f95e44a.

• The pairing is owed by every multi-child clustered operator again, so a cogroup keeps the mutual check it had on the base, and only the ungrouped waiver is producer-scoped.
• The operator kinds are one predicate now, ShuffledJoin.partiallyClusteredJoinType, read by the producer's dispatch and by the validator, rather than a copy in each file.
• The ungrouped waiver asks the collapse permission the planner applies, so a member whose keys were collapsed is admitted ungrouped only where the planner would agree to group it.
• A side is judged on the key the operation clusters on, as it reports it: the partition expressions carrying no cluster key are left out, and the count and the key order are the member's own. That takes in the subset shape without reading a projection the plan does not hold.
• The producer states the invariant the waiver rests on, where it is decided: one side spreads its splits and the other repeats the whole group, since replicateRightSide is the negation of replicateLeftSide and the branch is reached only when the picked side may replicate. That is what makes pairing the two index by index sound, and the pairing cannot tell that from two sides that split the key between them.

The tests are the suites the change touches: ShuffleSpecSuite and DistributionSuite (52), and ValidateRequirementsSuite, EnsureRequirementsSuite, GroupPartitionsExecSuite, KeyGroupedPartitioningSuite and ProjectedOrderingAndPartitioningSuite (362), with dev/lint-scala clean.

[peter-toth]

Re-checked from scratch through f95e44a — findings 1, 2, 3 and 5 resolved, nothing regressed. 1: the admission is mayServeUngrouped now, which is keysSatisfy plus the collapse permission and assumes no grouping, and reportedSpecOf builds the spec from the layout the plan reports. 2: all three comments name their real callers. 3: the ticket lists 4.4.0 as well. 5: reportedSpecOf's doc states the difference from createShuffleSpec outright rather than implying an equivalence.

I re-measured the two end-to-end tests against the current base, 2d06119539, since the rebase moved past SPARK-59688 and a fails-on-base verdict from last round would have expired: with the four main-code files reverted to that base, both SPARK-59671 tests in KeyGroupedPartitioningSuite still fail, and both pass on this head. So the coalescing claim holds on the base as it now stands.

Finding 4 stays the deferred direction rather than a request. The description states the trust-the-producer decision openly now, which is what I was after.

Non-blocking

• 6. The ungrouped waiver is wider than its producer by the join type (new): partiallyClusteredJoinType names the operator kinds, but the producer additionally requires canDuplicateLeftSide || canDuplicateRightSide before it spreads a side, so a FullOuter join is waived although nothing can build the shape the waiver is for. Measured on this head: ValidateRequirements.validate returns true for a FullOuter SortMergeJoinExec over the same ungrouped pair your as-of-join test refuses, while canDuplicateLeftSide(FullOuter) and canDuplicateRightSide(FullOuter) are both false. This is the half of the producer-boundary thread that the operator-kind fix left open; the helper already returns the join type and mayBeUngrouped discards it, so folding the test in there (not into the helper, which the producer's dispatch reads to decide whether to attempt SPJ at all) closes it:

      ShuffledJoin.partiallyClusteredJoinType(plan).exists { joinType =>
        ShuffledJoin.canDuplicateLeftSide(joinType) ||
          ShuffledJoin.canDuplicateRightSide(joinType)
      }

  I applied exactly that and ran it: the FullOuter probe flips to refused, ValidateRequirementsSuite stays green at 24, and both end-to-end tests still pass. No plan EnsureRequirements builds today reaches the waived shape for FullOuter, so this is the invariant rather than a live wrong result — worth a case in the producer-boundary test either way, since that test pins the operator kind and not the join type.

[cloud-fan]

Review summary

PR tags: bug-fix · performance

The five findings from the previous review are addressed in the current revision. This pass found two additional non-blocking issues: the new shared SMJ/SHJ classifier still admits FullOuter even though neither side may be duplicated, and a new helper comment overstates equivalence with the planner's compatibleAsIs path. The first was newly introduced by the latest producer-boundary revision; the second is a late catch from the preceding revision. No blocking issues remain.

Findings

2 total: 0 P0, 0 P1, 1 P2, 1 P3.

Non-blocking (P2)

• Exclude non-duplicable join types from the ungrouped waiver — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:67 — see inline.

Nit (P3)

• Correct the compatibleAsIs equivalence claim — sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ValidateRequirements.scala:106 — see inline.

Re-review status

Prior AI findings: 5 addressed, 0 still present; additional unresolved findings in this review: 2.

New attribution: 1 newly introduced, 1 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — The current revision resolves the projection and non-join examples, but the feedback's producer-boundary requirement also exposes the remaining FullOuter case: the class matches while no side is duplicable.
• existing discussion — The local review response correctly required a producer-specific ungrouped exception. The current class-only predicate is still broader than the producer's join-type capability for FullOuter.
• existing discussion — The author's update states that the exception is confined to the producer, but the current helper classifies FullOuter SMJ/SHJ although EnsureRequirements cannot produce a spread/repeated FullOuter pair.
• existing discussion — The prior local pushback required the waiver to match the producer-supported joins. AS-OF and non-join cases are fixed, but FullOuter is still in the selected executor classes and cannot produce the required replicate/spread shape.
• existing discussion — The latest author update claims the shared predicate keeps producer and validator aligned. It aligns their executor-class dispatch, but not the producer's subsequent canDuplicateLeftSide/canDuplicateRightSide gate, so FullOuter remains over-admitted.

[cloud-fan]

Non-blocking (P2): partiallyClusteredJoinType(plan).isDefined proves only that this is an SMJ/SHJ, not that this join type can produce the spread/repeat layout the waiver requires. For FullOuter, both duplication capabilities are false, so EnsureRequirements skips partial clustering, yet matching repeated children can pass here and execute split groups partition-by-partition, missing cross-split matches or emitting matched rows as unmatched. Please make the shared predicate exclude join types with no duplicable side and add a focused FullOuter refusal test.

[ulysses-you]

Thanks @cloud-fan! Fixed in 32ddb26. The waiver asks the second gate the producer applies before it spreads a side now (canDuplicateLeftSide || canDuplicateRightSide), folded in at the waiver rather than into partiallyClusteredJoinType: that helper is the producer's entry to key-group checking altogether, so a kind turned away there would lose the storage-partitioned join it can still plan without spreading a side. The producer-boundary test pins the join type as well as the operator kind, the same pair refused for a full outer join and accepted for the inner pair the waiver is for.

One correction on attribution: the over-admission is not new to the latest revision. The waiver was ShuffledJoin-wide before it, so a FullOuter pair was admitted then too, and the operator-kind revision neither introduced nor closed it. It came in with the waiver itself.

[cloud-fan]

Nit (P3): This is not quite the same question as compatibleAsIs. Here, reportedSpecOf may relabel a grouped [a, b] member for an [a] distribution when dropping b merges no partition; the planner represents that candidate as projected, and compatibleAsIs requires both specs to be unprojected. Please describe the relationship without claiming predicate equivalence across this projection case.

[ulysses-you]

Thanks @cloud-fan! Reworded in 32ddb26: the doc says the planner asks this of a pair it takes as it stands, though not by the same predicate, since its compatibleAsIs path reads two unprojected specs while a member here may be relabelled onto the key the operation clusters on.

[cloud-fan]

The earlier AS-OF and non-join issues are addressed, but the shared predicate still returns FullOuter for SMJ/SHJ even though both duplication predicates are false, so validation admits a shape EnsureRequirements cannot produce safely. Please exclude join types with no duplicable side and add a focused full-outer refusal test.

[ulysses-you]

Thanks @peter-toth and @cloud-fan! Both findings of this pass are in 32ddb26.

• The waiver asks the join type's duplication capability now, the second gate the producer applies before it spreads a side, so a FullOuter pair is refused instead of admitted on its kind. The gate is asked at the waiver and not through partiallyClusteredJoinType, since that helper is the producer's entry to key-group checking altogether: a kind turned away there would lose the storage-partitioned join it can still plan without spreading a side.
• satisfiesForPairing's doc describes the relationship to the planner's compatibleAsIs path without claiming the same predicate.
• The producer-boundary test pins the join type alongside the operator kind: an as-of join and a full outer join refuse the same pair, and the inner pair the waiver is for passes.

On attribution, one correction to the review's "newly introduced": the full-outer over-admission came in with the waiver itself, which was ShuffledJoin-wide before the operator-kind revision, so that revision neither introduced nor closed it.

The change is a second commit on top of f95e44a rather than an amend, so its delta reads on its own.

[peter-toth]

Re-checked through 32ddb26. Finding 6 resolved, nothing new.

The waiver now asks canDuplicateLeftSide || canDuplicateRightSide, which is the producer's own first gate at EnsureRequirements.scala:813. I put .isDefined back and the new full-outer assertion is the one that fails, so the narrowing is pinned by its test. ShuffleSpecSuite is green at 32, ValidateRequirementsSuite at 23, and both end-to-end tests pass. Reverting only ValidateRequirements.scala to the base fails 7 of the 11 new unit tests, and the 4 that survive are the pure refusals the description says hold there.

Thanks for working through all of these, @ulysses-you. Nothing left open from my side.