[SPARK-59187][SQL] Carry the partition key data types on KeyedPartitioning

[peter-toth]

What changes were proposed in this pull request?

KeyedPartitioning gains a keyDataTypes field beside expressions, and keyDataTypes stops sampling the first partition key row.

The types are already computed wherever the keys are built, and were dropped there. KeyedPartitioning.apply computed them only to build the key wrapper factory. project called projectKeys(positions)._2 and threw away the projected types. GroupPartitionsExec.grouping computed the reduced types, PartitionGrouping did not carry them, and outputPartitioning re-derived them by sampling. All three now pass what they know.

Every copy keeps the types, which is right at each site: KeyedShuffleSpec.createPartitioning and withNewChildrenInternal change the expressions and leave the key rows alone, and canonicalisation and toGrouped do not change what a key holds. Two copies do change the rows. project passes the projected types, and concat, the one place that mixes rows from several partitionings, requires the children to agree on them.

The field is checked against the thing it describes rather than argued about: the constructor requires one type per expression, and requires the first key row to have been built with them. Every key wrapper carries its own types, so this costs one comparison. PartitioningCollection's invariant deliberately does not gain a clause. Its members share one key list, and a wrapper compares its types before its values, so structurally equal keys force equal types, leaving only empty-key members where the types describe nothing.

Two consequences in EnsureRequirements. The reduced-types comparison drops SPARK-59176's empty-side guard, which existed only because a keyless side reported the expressions' types. And the two KeyedPartitioning(clustering, _, _, _) patterns become type patterns reading .expressions, so a future field does not touch them.

Why are the changes needed?

keyDataTypes answered by reading the first key row and falling back to expressionDataTypes when there was none. Three things followed.

A partitioning with no key reported types no key of it would have held. That is SPARK-59176, whose fix keeps such a side out of one comparison rather than giving it an answer. With the field the side answers, and the guard is gone.

Nothing checked that the keys after the first agreed with it. There is now nothing to disagree with, because no key is consulted for types at all.

The divergence at KeyedShuffleSpec.createPartitioning was accidental. copy(expressions = ...) silently left the types describing the old expressions. It is now an explicit property of the copy, and the scaladoc says so on expressionDataTypes, which is the member whose use needs the warning.

Does this PR introduce any user-facing change?

No. A stringArgs override keeps the field out of explain, so plan output is unchanged.

KeyedPartitioning is a catalyst class, so this is not public API. It is still a binary- and source-incompatible change to a case class: the constructor, apply, copy, unapply and the companion's tupled/curried all change shape. MiMa does not flag it, since MimaExcludes blanket-excludes org.apache.spark.sql.catalyst.*. The two destructuring patterns in the repository are converted to type patterns.

How was this patch tested?

A new GroupPartitionsExecSuite test asserts that a reduce's result type reaches the reported partitioning, both with keys and with none, while the expression the node reports keeps its own different type. It uses a both-sides reduce, which is the shape whose reported expression cannot carry the key type: EnsureRequirements refuses the same divergence from a one-side reduce.

The existing SPARK-59176 tests pin the guard removal: both still pass with the comparison back to its plain form, including the one that asserts a contract-breaking reducer still raises STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES.

The type propagation itself is pinned by twelve existing KeyGroupedPartitioningSuite tests, measured by making GroupPartitionsExec report the child's types instead of the reduced ones.

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

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.

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

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

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

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

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

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

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