peter-toth opened a new pull request, #58552: URL: https://github.com/apache/spark/pull/58552
### What changes were proposed in this pull request? `KeyedPartitioning` grows a `KeyLayout`, and its key types move into it. **One value for what a partitioning's members share.** `KeyLayout(partitionKeys, dataTypes, isGrouped, isCollapsed)` holds everything about the partitions a `KeyedPartitioning` describes except the expressions naming them, so `KeyedPartitioning` is `(expressions, layout)`. The members of a `PartitioningCollection` name one layout with their own expressions and share the object by reference, so: - the collection's invariant is one `eq` on the layout in place of a clause per shared field, and it now covers `isGrouped`, which the field-by-field check left out; - `fromPartitionings` merges one canonical layout instead of interning the keys and ORing a flag, and refuses a member that describes another key space, which interning would otherwise retype; - `KeyedShuffleSpec.createPartitioning` has nothing to decide, since a copy that only replaces the expressions keeps the layout; - `GroupPartitionsExec`'s `PartitionGrouping` is the layout it will report plus the child partitions each of its own is built from. **One answer for the key types, including where no key row is left.** `keyDataTypes` reads the layout rather than sampling the first key row and falling back to the partition expressions. A layout is given the types its keys were built at, at the four places one is built. `EnsureRequirements` therefore drops the exception SPARK-59176 added for a side with no key row, since the layout answers for it. ### Why are the changes needed? Two things, one structural and one a defect the structure hides. **The shared part of a `KeyedPartitioning` is currently four fields that every member of a collection has to agree on by hand.** `checkKeyedPartitioningInvariant` compares them clause by clause, and it left `isGrouped` out. Four places put a partitioning's expressions over keys they did not build, and each has to carry the shared fields forward correctly: 1. `GroupPartitionsExec.outputPartitioning` reports the keys `EnsureRequirements` merged, and picks its member with `collectFirst`, which need not be the member the planner merged from. 2. `PartitioningPreservingUnaryExecNode.projectKeyedPartitionings` projects `kps.head` once and stamps every alias alternative onto it with `copy(expressions = ...)`. 3. `KeyedShuffleSpec.createPartitioning` puts the other child's expressions over these keys. 4. `KeyedPartitioning.concat`, for a `UnionExec`. Sharing one layout by reference is what makes all four correct rather than merely lucky, and it turns the invariant into one `eq`. **A side with no key row answers from its partition expressions, and after a both-sides reduce that is a type no key of it holds.** The reduce leaves keys that are `r1(f1(x))` = `r2(f2(x))`, a space neither transform names, so the reported expression is marked and its own type is the un-reduced one. SPARK-59176 worked around it by leaving such a side out of the co-partition type check, which left the check not checking for the shape most likely to need it, and every other reader of `keyDataTypes` still getting the wrong answer. The layout now carries what the reduce produced, so the workaround goes. ### Why this shape The types are on the layout rather than derived, because a partitioning whose partitions were all pruned has no row to read them off and its expressions do not describe a reduced key space. They are on the *layout* rather than on `KeyedPartitioning`, because that is what makes them shared: two members that share a layout share its keys, so the collection's `eq` covers them, and no consumer can pair one member's rows with another member's types. Two alternatives were tried and dropped. An independent `keyDataTypes` field on `KeyedPartitioning` has to be decided at each of the four sites above, and two of them mix members, which is how it produced two reachable regressions in review. A `TypedKeys(dataTypes, keys)` value object does not settle it either, since two members can still hold two different pairs. No plan string changes. `KeyedPartitioning.stringArgs` prints the layout's contents where the value object would print, and deliberately leaves the key types out of that list: they have their naming erased (SPARK-59187), so printing them would put a struct field named `0` into a plan that appears nowhere in the query. ### Does this PR introduce _any_ user-facing change? No. The two queries SPARK-59187's tests cover already run on this PR's base, and this change adds no behaviour of its own beyond making a pruned side report its own key types truthfully. ### How was this patch tested? Four new tests, plus SPARK-59176's two existing ones, which now pass with its exception removed. Ablation: with the exception removed and `keyDataTypes` derived from the rows and expressions again, "SPARK-59176: a leg reduced onto no key at all still joins" fails with the error SPARK-59176 was filed for. - `DistributionSuite`, "fromPartitionings refuses a member that disagrees on isGrouped", for the layout itself. - `GroupPartitionsExecSuite`, "a reduced key space's type reaches the reported partitioning with no key left": a both-sides reduce onto `LongType` under a `DateType` transform, with keys and without. - `KeyGroupedPartitioningSuite`, "two sides whose partitions were all pruned are not one layout": two legs pruned to nothing, one `identity(id)` on `LongType` and one `bucket(4, id)` on `IntegerType`, joined and then joined again through a FULL OUTER that brings real keys in. It asserts that no node reports two key spaces as one layout, that the plan passes `ValidateRequirements`, and the answer. - `KeyGroupedPartitioningSuite`, "two legs whose struct field names differ are still co-partitioned", the shape where two exact type lists differ while the space does not. `KeyGroupedPartitioningSuite`, `KeyGroupedPartitioningRuntimeFilterSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `DataSourceV2CatalystRuntimeFilterSuite`, `DistributionSuite`, `ShuffleSpecSuite`, `TransformExpressionSuite` and `InternalRowComparableWrapperSuite`, 301 tests. Scalastyle and scalafmt clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) -- 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]
