[SPARK-59285][SQL] Hold a KeyedPartitioning's shared partition layout in one value by peter-toth · Pull Request #58552 · apache/spark

peter-toth · 2026-09-06T18:29:17Z

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:

GroupPartitionsExec.outputPartitioning reports the keys EnsureRequirements merged, and picks its member with collectFirst, which need not be the member the planner merged from.
PartitioningPreservingUnaryExecNode.projectKeyedPartitionings projects kps.head once and stamps every alias alternative onto it with copy(expressions = ...).
KeyedShuffleSpec.createPartitioning puts the other child's expressions over these keys.
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)

…g erased ### What changes were proposed in this pull request? `InternalRowComparableWrapper`'s factory builds every partition key row at `comparableTypes`, the given types with struct field names and every nullability erased and nothing else touched. It is `asNullable` for the nullability half plus a positional `StructType` rename for the naming half. Three places that report a key list's types alongside it erase them the same way, so a partitioning's key types are always the types its keys are compared at: `KeyedPartitioning.keyDataTypes`, in its no-key fallback, and `KeyedPartitioning.projectKeys` and `reduceKeys`, the latter over types that come from a connector's `Reducer`. The erasure is idempotent and keeps the list it was given, so a caller that hands its own types in and the result back to another factory gets one instance, and `equals` keeps its reference fast path. ### Why are the changes needed? `InternalRowComparableWrapper.equals` compares its `dataTypes` before its values, so two key rows of one value never matched when the columns they came from were named differently. A storage-partitioned join is about the opposite: a key value belongs where its value says, not where its column name says. Two consequences, both measured on master. **A join between two keyed sides whose struct key fields are named differently throws.** `identity` carries the column's own struct type into the key type, and the analyzer accepts an equi-join across `struct<a:int>` and `struct<b:int>` -- `BinaryComparison.sameType` is `DataType.equalsStructurally(_, _, ignoreNullability = true)`, so no `Cast` is inserted. Both sides are keyed and the values match, but the key rows never matched, so the co-partitioned fast path in `KeyedShuffleSpec.isCompatibleWith` was out, and a pair of attributes has no reducer, so the reduced-types check threw `STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES` for a join that reduced nothing: s1(id struct<a:int>, v string) partitioned by identity(id), keys named_struct('a',1), ('a',2) s2(k struct<b:int>, w string) partitioned by identity(k), keys named_struct('b',1), ('b',2) SELECT s1.v, s2.w FROM s1 JOIN s2 ON s1.id = s2.k -- threw, now runs with no shuffle **A union over a key that two children hold under two namings silently drops rows.** `KeyedPartitioning.concat` puts the children's key rows in one list and asks whether any repeats. Rows of two namings never matched, so the merged partitioning reported unique keys when they were not. Nothing regroups it then, so `KeyedShuffleSpec.canCreatePartitioning` accepts it and the other side is shuffled straight onto those keys. `KeyGroupedPartitioner`'s map holds one partition per key, so the union partition holding the earlier copy of the repeated key receives no rows at all: items(id struct<a:int>, name string) partitioned by identity(id), keys ('a',1), ('a',2) purchases(item_id struct<b:int>, ...) unpartitioned t3(c struct<b:int>) partitioned by identity(c), key ('b',1) s4(k4 struct<b:int>, w string) unpartitioned, rows (('b',1),'x'), (('b',2),'y') SELECT u.k, s.w FROM ( SELECT p.item_id AS k FROM purchases p LEFT JOIN items i ON p.item_id = i.id UNION ALL SELECT c AS k FROM t3 ) u JOIN s4 s ON u.k = s.k4 returns 2 rows on master and 3 with this change. An inner join loses a row it should return. `ShuffleExchangeExec` already knew about this and worked around it, re-wrapping the partitioner's map keys through the same factory as its per-row lookup keys so the naming could not decide (SPARK-59054). Doing the erasure where rows are built removes that workaround's reason. This PR leaves the re-wrap in place under a different one: the stored keys were built at `keyDataTypes`, and re-wrapping is what makes that list and the lookups' list agree. A mismatch there is silent, since `KeyGroupedPartitioner.getPartition` answers a miss with the key's hash. ### Why this shape The factory is a chokepoint: six sites in production build partition-key wrappers, none of them wants un-erased types, and the only readers of a wrapper's `dataTypes` are its own `equals` and `keyDataTypes`. No site can opt out. The class already erased half of this: `structTypeCache` names every top-level field `"f"`, and `RowOrdering.createNaturalAscendingOrdering` forces `nullable = true`, so `hashCode` was naming-blind while `equals` was not. What is erased is the naming and nothing more. A collation, a decimal precision, a `char` length and a UDT all decide where a value belongs, so they still tell two rows apart. The erasure is exactly `DataType.equalsStructurally(_, _, ignoreNullability = true)` expressed as a canonical value rather than a predicate, and the new suite pins it against that primitive. A value is what is needed rather than a predicate, because the result is a `NonFateSharingCache` key, the `dataTypes` field two wrappers compare, and what `keyDataTypes` reports. Nothing that compares or hashes a row reads a field name: `GenerateOrdering.genComparisons` rebuilds each field's `SortOrder` positionally, and `Murmur3HashFunction` hashes through the field types. So the erasure cannot move a row or change a sort order, and the `KeyedPartitioning.toGrouped` / `GroupPartitionsExec.groupAndSortByKeys` sort contract is unaffected. It can only make more keys compare equal, so `isGrouped` moves toward "not unique" and a regroup is added, never skipped. ### Does this PR introduce _any_ user-facing change? Yes, three things. - The join above ran into an error and now returns its rows without a shuffle. - The union above dropped a row and now returns it. - `STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES` prints a struct key's field names positionally, ``STRUCT<`0`: INT>`` rather than `STRUCT<a: INT>`. Deliberate: the message fires only on a real structural mismatch now, and the connector's own names would point a reader at a difference that is not the cause. ### How was this patch tested? `InternalRowComparableWrapperSuite`, new: - "comparableTypes erases the naming and nothing else", over 22 type pairs including collation, decimal precision, `char` length, two UDTs, field metadata, struct nullability and nested arrays and maps, each checked against `DataType.equalsStructurally(ignoreNullability = true)`. - "erasing is idempotent, and keeps the list it was given". - "two rows of one value are equal however their columns were named", which also asserts they hash alike and collapse in a set. `KeyGroupedPartitioningSuite`, three end-to-end tests, all failing on master: - "two keyed sides whose struct field names differ join without a shuffle" -- the first query above. Fails with `STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES`. - "a union of a key that repeats across children under two namings joins right" -- the second query above. Fails on the answer, 2 rows where 3 are correct, and then on `isGrouped`. - "a shuffled side keeps its own struct field names over shared empty keys" -- a join whose two members carry the two sides' own expressions over one shared, pruned-to-nothing key list. Fails because the members answer for two key spaces. `ShuffleSpecSuite`, "reduceKeys reports the types its keys are compared at", for a `Reducer` whose result type names a struct field. Each production hunk was ablated in turn and each has a test that fails without it. `KeyGroupedPartitioningSuite`, `KeyGroupedPartitioningRuntimeFilterSuite`, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `PlannerSuite`, `ExchangeSuite`, `ExplainSuite`, `DataFrameSetOperationsSuite`, `DataSourceV2Suite`, `DataSourceV2CatalystRuntimeFilterSuite`, `ProjectedOrderingAndPartitioningSuite`, `DistributionSuite`, `ShuffleSpecSuite`, `TransformExpressionSuite` and `InternalRowComparableWrapperSuite`, 533 tests. Scalastyle and scalafmt clean. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5)

… in one value ### 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)

peter-toth · 2026-09-06T18:29:32Z

Draft on purpose. This is stacked on #58501, so its diff includes that PR's commit, and it should not be reviewed or merged until two others land:

[SPARK-59187][SQL] Compare partition key rows at types with the naming erased #58501 (SPARK-59187), the bugfix this sits on. It erases the naming where partition key rows are built, which is what lets keyDataTypes be one answer here rather than two.
[SPARK-59050][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins #58339 (SPARK-59050), which adds mayContainUnknownPartitionKeys to KeyedPartitioning. That is another field the layout should hold, so doing this first would mean redoing it.

I will rebase onto master and un-draft once both are in. Only the second commit here is this ticket's.

peter-toth added 2 commits September 6, 2026 18:25

Navigation Menu

Uh oh!

[SPARK-59285][SQL] Hold a KeyedPartitioning's shared partition layout in one value - #58552

[SPARK-59285][SQL] Hold a KeyedPartitioning's shared partition layout in one value#58552
peter-toth wants to merge 2 commits into
apache:masterfrom
peter-toth:SPARK-59285-keylayout

peter-toth commented Sep 6, 2026

Uh oh!

peter-toth commented Sep 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

peter-toth commented Sep 6, 2026

What changes were proposed in this pull request?

Why are the changes needed?

Why this shape

Does this PR introduce any user-facing change?

How was this patch tested?

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

Uh oh!

peter-toth commented Sep 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant