peter-toth opened a new pull request, #58814: URL: https://github.com/apache/spark/pull/58814
### What changes were proposed in this pull request? `EnsureRequirements.checkKeyGroupCompatible` asks its two rebuilt children whether they still declare the same aligned partition key sequence before it commits to a storage-partitioned join, and declines when they do not. The check is pairwise rather than a per-side `satisfies`, through `KeyLayout.describesSameKeys`, which compares the key rows and the key types. ### Why are the changes needed? `GroupPartitionsExec` deliberately gives up a marked layout's keyed claim, reporting `UnknownPartitioning`, when the regrouping is not the identity: a marked side realigned onto differently-ordered merged keys, or a grouping that applies a reducer or a non-identity projection. That is the right answer to a claim the regrouping invalidates. The gap is when it is answered. The give-up happens inside the node, which `checkKeyGroupCompatible` builds through `applyGroupPartitions` **after** it has decided the pairing is compatible, and nothing re-checks the rebuilt children. So a committed join can carry a child that no longer satisfies its required distribution. `ValidateRequirements` then rejects the whole stage, and every `AQEShuffleReadRule` and `OptimizeSkewedJoin` drops its result on a stage that does not validate, so partition coalescing, local read and skew join are all off for it. A per-side recheck is not viable, and this is why the fix is pairwise. Partially clustered distribution deliberately leaves both rebuilt children ungrouped yet value-aligned, both onto one `mergedPartitionKeys` by `alignToExpectedKeys`. A `KeyedPartitioning` only satisfies a `ClusteredDistribution` once grouped, so a `satisfies` gate would reject that whole family. What the two sides owe each other is the key sequence `alignToExpectedKeys` constructs, each key repeated as many times as the merge expects, whichever side replicates, and that is what is asked here. Nothing that was accepted before is refused on the `compatibleAsIs` path, where the children are the ones the pairing read. `KeyedShuffleSpec.isCompatibleWith` already ends in `describesSameKeys`, and all keyed members of a `PartitioningCollection` share one `KeyLayout`, so whichever member the spec matched on declares what the representative does. ### Does this PR introduce _any_ user-facing change? **Yes**, one plan change. 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". ### How was this patch tested? The `SPARK-59050` regrouping test already reaches this state, so it is the test: its expectation moves from three shuffles to four, and it now asserts `ValidateRequirements.validate` on the executed plan ahead of the shuffle count, because that names the reason the fourth shuffle exists. It is the only test in `KeyGroupedPartitioningSuite`'s 183 that the change moves. **A differential sweep measures the rest.** 307200 generated storage-partitioned join plans, 40 partitioning shapes a side crossed with six join types over 32 configuration cells, run against master 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. | | master | this change | |---|---|---| | plans that skip both shuffles yet disagree on the keys | 2024 | **0** | | plans `ValidateRequirements` rejects | 6724 | 4700 | | plans the planner cannot build at all | 29952 | 29952 | **Zero cases are added to any of the three**, and the 2024 co-partitioning violations that go are exactly the same 2024 cases as the validation failures that go. That equality is the whole of this change in one number: every plan that skipped both shuffles on keys the two sides did not share was also a plan the validator rejected. Green: `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `ValidateRequirementsSuite`, `GroupPartitionsExecSuite`, `PlannerSuite`, `AdaptiveQueryExecSuite`, `JoinSuite` and the seven plan-stability suites, 1679 tests. `dev/lint-scala` is clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
