[SPARK-59187][SQL] Compare partition key rows at types with the naming erased

[peter-toth]

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.

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:

t1(k1 struct<a:int>)           partitioned by identity(k1), keys ('a',1), ('a',2)
t2(k2 struct<b:int>)           partitioned by identity(k2), 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 k1 AS k FROM t1 UNION ALL SELECT k2 AS k FROM t2
) 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.

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.

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.

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.

Generated-by: Claude Code (Opus 5)

[peter-toth]

cc @dongjoon-hyun, @ulysses-you

[dongjoon-hyun]

Thank you for the detailed write-up. The direction (carrying the types instead of sampling a key row, and dropping the SPARK-59176 guard) looks right to me, and I could not find a regression path in the EnsureRequirements comparison itself.

However, I found two regressions that share one root cause: the new keyDataTypes invariant is not enforced when the key list is empty, and two readers pick members of a PartitioningCollection by different rules. Details inline, plus a few doc nits.

[dongjoon-hyun]

This looks like a regression. grouping takes keyDataTypes from the member found by collectFirst (line ~177) and this line stamps it on every member, while EnsureRequirements.createKeyedShuffleSpec builds expectedPartitionKeys from the first member that satisfies the distribution. Members of a PartitioningCollection with empty key lists can carry different keyDataTypes (nothing checks or normalizes them, see my comment on checkKeyedPartitioningInvariant), so the two can disagree and the new constructor require fires.

Reachable shape, with pushPartValues, partitionFilter and allowCompatibleTransforms on:

1. leg1 = t1 JOIN t2, both identity-partitioned on a: string with disjoint keys -> partition filter intersects to nothing, two members KP(a, [], [String]).
2. leg2 = t3 JOIN t4, both bucket(4, b: string), disjoint -> two members KP(bucket(4,b), [], [Int]).
3. leg1 JOIN leg2 ON a = b: KeyedShuffleSpec.isCompatibleWith is true (Nil == Nil, numPartitions 0 == 0, attribute vs transform is compatible via canReduceKeys), so the push-down branch and its type check are skipped and fromPartitionings produces a collection mixing [String] and [Int].
4. ... FULL OUTER JOIN t5(bucket(4, c)) ON b = c: EnsureRequirements picks the bucket(4,b) member, merged keys are Int rows. grouping picks member a, so reducedDataTypes = [String], and this line builds KeyedPartitioning(bucket(4,b), <Int keys>, [String], ...) -> IllegalArgumentException from the require at partitioning.scala:603.

Before this PR the types were read from the key rows, so the query ran. The inner-join variant passes the require (merged keys are empty) but carries [String] on the bucket member, and a later push-down join then throws STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES.

[peter-toth]

Confirmed, and reproduced. Thank you.

Chasing it end to end turned up a fourth site that puts a partitioning's expressions over keys it did not build: PartitioningPreservingUnaryExecNode.projectKeyedPartitionings, which projects kps.head once and then stamps every alias alternative onto it with copy(expressions = ...). A bucket(4, id) member ended up reporting the identity(id) member's LongType. So an independent keyDataTypes field has to be decided at four places, and two of them mix members.

That is what killed the design. This PR is reshaped: there is no keyDataTypes field, it stays derived, and the constructor require is gone with it. GroupPartitionsExec stamps no types.

Carrying the types on the partitioning moved to SPARK-59285. There they live on a shared KeyLayout that the collection's members hold by reference, so no member can report another member's types, and the four re-targeting sites have nothing to decide.

[dongjoon-hyun]

This require is reachable from UnionExec.comparePartitioning, the only caller, which compares children by semanticEquals on expressions alone and otherwise falls back to super.outputPartitioning. Semantically equal expressions do not imply equal keyDataTypes; this PR itself documents KeyedShuffleSpec.createPartitioning keeping the keyed side's struct field names. So a query that used to plan now fails with IllegalArgumentException.

Concrete case with shuffle-one-side on: purchases p LEFT JOIN items i ON p.item_id = i.id, items identity-partitioned on id: struct<a:int>, purchases unkeyed. The purchases side is shuffled via createPartitioning, and being LeftOuter its KeyedPartitioning(item_id, <items keys>, [struct<a:int>]) becomes the join output. SELECT item_id ... UNION ALL SELECT c FROM t3 with t3 identity-partitioned on c: struct<b:int>. BinaryComparison.sameType ignores struct field names so no Cast is inserted (the existing test SPARK-59054: shuffle one side: struct partition keys with different field names plans exactly this join). UnionExec remaps both expressions to the union output attribute, semanticEquals holds, and concat receives [struct<a:int>] vs [struct<b:int>].

Before this PR the mixed-type concat ran (with a latent isGrouped miscount, since wrappers of different types never compare equal). I think the check belongs in UnionExec.comparePartitioning next to the expression comparison, so that a mismatch takes the existing fallback instead of throwing here.

[peter-toth]

Confirmed. The require is gone, since there is no keyDataTypes field to check.

Your parenthesis turned out to be the important part. I measured the latent isGrouped miscount, and it is not latent: it silently drops rows. The union reports unique keys, so nothing regroups it, 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

2 rows on master where 3 are correct, with v2.bucketing.shuffle.enabled and union.output.partitioning.enabled on and AQE off. An inner join loses a row, with no error.

That is now half of what this PR fixes, and it is the test SPARK-59187: a union of a key that repeats across children under two namings joins right. Thank you for spotting it.

[dongjoon-hyun]

I don't think this premise holds. An empty-key member's keyDataTypes is read: by GroupPartitionsExec.grouping (via collectFirst), by the reduced-types comparison in EnsureRequirements, and by PushDownUtils. Those readers pick a member by different rules (collectFirst vs. the first member that satisfies), so if members disagree the answer depends on which one is consulted. That is what produces the GroupPartitionsExec failure I described above.

A one-line require(rep.keyDataTypes == first.keyDataTypes, ...) next to the existing isCollapsed check (same O(members) cost) would make this structural, or fromPartitionings could normalize the field the way it interns partitionKeys.

[peter-toth]

You are right, and the PR no longer claims it. There is no field, so keyDataTypes is derived per member again.

Part of what you describe is closed by the erasure: two members over one key space that differ only in how their struct fields are named now give one answer, whichever reader consults whichever member.

Part of it is not. Two members can describe genuinely different key spaces, [String] and [Int] in your example, and nothing here refuses that. Your require, or normalizing in fromPartitionings, is in SPARK-59285: fromPartitionings interns one shared layout and refuses a member that describes another space, which also covers isGrouped that the field-by-field check leaves out.

I kept it out of this PR because I measured the mixed collection on master and it is not a live bug. The plan passes ValidateRequirements and the query returns the right answer, so that part is hardening.

[dongjoon-hyun]

nit: the justification does not hold. InternalRowComparableWrapper has no toString, so the keys print as InternalRowComparableWrapper@<hash> and no type is legible beside them. The hide is also asymmetric: TreeNode.jsonFields / asCode use productIterator, so the field shows up in toJSON but not in explain.

The case this PR exists for (a marked days(...) expression of DateType over LongType keys, possibly with no key at all) is exactly where explain would show a misleading expression type with no way to see the real one, and two partitionings that differ only in keyDataTypes print identically in require/assert messages. No golden file or test asserts this string, so I would drop the override. If it is kept for output stability, a one-line comment saying so would be clearer than arguing the information is worthless.

[peter-toth]

Dropped, along with the field it was hiding. stringArgs is master's again.

Your point stands for SPARK-59285, and it lands the other way there. The types are erased, so printing them would put a struct field named 0 into every explain. The override is kept there with a comment saying that, rather than arguing the information is worthless.

[dongjoon-hyun]

nit: this reads as if KeyedShuffleSpec.createPartitioning were a case where the expressions have stopped describing the keys, but it only does partitioning.copy(expressions = newExpressions) and sets no marker, so expressionsDescribeKeys stays true there and the last paragraph's "expressionsDescribeKeys is what keeps them sound" only covers the reduce case. Something like: "May differ from keyDataTypes in two cases: (a) a both-sides reduce marks the expressions (expressionsDescribeKeys); (b) KeyedShuffleSpec.createPartitioning re-targets the expressions at the other child's attributes, so struct field names can differ while the expressions still describe the keys."

[peter-toth]

Taken, and your (b) changed shape rather than going away.

KeyedShuffleSpec.createPartitioning re-targeting the expressions no longer makes the two answers differ, because the erasure brings both to one answer. So the doc now lists the naming, since expressionDataTypes is not erased at all, and the both-sides reduce. It also says plainly that a one-side reduce keeps the two equal up to the naming, because the expression the partitioning reports is the target transform and EnsureRequirements refuses a reducer whose result type disagrees with it.

Your underlying point is what I got wrong: expressionsDescribeKeys only covers the reduce case. The ShuffleExchangeExec paragraph now says that, and the comment at the partitioner build site gives the reason the re-wrap earns its keep, rather than the struct-field-name one it used to give.

[dongjoon-hyun]

Optional, for consideration: the (types, keys) pair now appears as two constructor args here, two tuple-returning helpers (projectKeys, reduceKeys), two PartitionGrouping fields and the fold seeds in EnsureRequirements, and the pairing is guaranteed only by a head-key require that is vacuous when the key list is empty. A small value object (say TypedKeys(dataTypes, keys)) returned by projectKeys/reduceKeys and held here and in PartitionGrouping would make the pairing structural and is the deeper fix for both issues above. Fine as a follow-up if you prefer to keep this PR small.

[peter-toth]

Taken, and it is the deeper fix. I owe you two things on it.

First, I tried TypedKeys and it did not settle the pairing. A value object per partitioning still lets two members of a PartitioningCollection hold two different pairs. What settles it is KeyLayout(partitionKeys, dataTypes, isGrouped, isCollapsed) shared by reference, so the collection's invariant is one eq plus the arity, createPartitioning has nothing to decide, and PartitionGrouping becomes the layout it will report. That is SPARK-59285, and the attempt is kept as history.

Second, projectKeys and reduceKeys still return tuples in this PR, and the fold seeds in EnsureRequirements still pair them by hand. My own cleanup pass flagged the same thing independently, so I am not going to argue it is fine. It is deliberately left for that ticket to keep this one backportable.

[dongjoon-hyun]

nit: this paragraph, the constructor comment and the concat scaladoc each inventory the copy sites (project, concat, toGrouped, fromPartitionings) and argue the design. The lists will silently rot at the next copy(partitionKeys = ...) (GroupPartitionsExec already builds one directly). I'd keep the contract only, e.g. "The types the partitionKeys rows were built with, one per expression; kept even when there is no key row.", and on concat: "Children must agree on keyDataTypes; the constructor only checks the first key."

[peter-toth]

Taken. Those inventories went with the field.

keyDataTypes' doc no longer lists copy sites, and concat has no keyDataTypes clause to document. What is left there is the contract plus the one thing a reader cannot derive: with no key row the expressions answer, and after a both-sides reduce that answer is a type no key of the partitioning holds.

[peter-toth]

Moved to draft for now, I'm gonna fix this more holistically...

[peter-toth]

@dongjoon-hyun thank you for the review, and sorry for the churn. The goal of this PR changed, so your two regressions no longer apply to the code, and your latent concat finding turned out to be the more serious of the two bugs it fixes.

What it is now: a bugfix only. Partition key rows compare at types with the naming erased, which closes two defects measured on master.

• A join between two keyed sides whose struct key fields are named differently throws STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES, for a join that reduced nothing.
• A union over a key that two children hold under two namings silently drops rows. This is the isGrouped miscount from your concat comment, and it is not latent. An inner join over such a union returns 2 rows where 3 are correct.

The JIRA is retyped to a Bug and rewritten, and it affects 4.2.0 onwards.

What moved out: carrying the key data types on the partitioning, which was the original goal and the old title. Your review is why. An independent field has to be decided at four sites that put a partitioning's expressions over keys they did not build, and two of them mix members, so the field needs to be shared rather than per-member. That is SPARK-59285, as a KeyLayout the collection's members hold by reference, and it will be a separate PR stacked on this one. The SPARK-59176 guard in EnsureRequirements therefore stays here for now.

I answered each of your comments in its thread, and said where each one landed in SPARK-59285 rather than here.

[ulysses-you]

I'm a bit concerned about we may miss call comparableTypes in future since the related code exists in many key place. Can we narrow the comparableTypes in the InternalRowComparableWrapper file ? It seems for now, the two key code path depends on it:

1. InternalRowComparableWrapper.hashCode/equals
2. InternalRowComparableWrapper.getInternalRowComparableWrapperFactory

We can normalize the data types inside these method.

[peter-toth]

Good point, and I took it. Fixed in 487da27.

Not quite where you suggested, though. hashCode and equals do not call comparableTypes, they read the dataTypes field, and that field is already erased: the primary constructor is private and both entry points erase, so no wrapper can exist over a raw list. Normalising inside equals would also put the erasure on a per-comparison path, and the merge's HashSet and the distinct in KeyedPartitioning.apply compare once per key.

What I did instead answers the same worry a level up. The factory is now a Factory class that answers for the schema it settled on, so a caller reporting a type list beside the rows it built reads factory.dataTypes rather than erasing on its own:

• KeyedPartitioning.projectKeys and reduceKeys do that now, which also drops a second erasure each was paying;
• comparableTypes is private[catalyst], which is the tight scope, since its only callers are partitioning.scala and its own suite.

One caller is left, keyDataTypes' fallback for a partitioning with no key row. There is no factory there to read the types off, and building one just to ask would be two cache lookups for no row. SPARK-59285 carries the types on the partitioning, and that caller goes with it.

[ulysses-you]

the change lgtm, two nits:

1. the description does not match the PR template
2. the description says "What is erased is the naming and nothing more," but field metadata is erased too