[SPARK-59045][SQL] Fix SPJ ClassCastException when reducer changes partition key data type

[ulysses-you]

What changes were proposed in this pull request?

KeyedShuffleSpec.reducersBothWays now pairs each Reducer with the reduced partition expression it produces (KeyReducer), and GroupPartitionsExec.outputPartitioning reports the reduced expression instead of the original partition expressions when reducers are applied. The stored expression is re-targeted at each KeyedPartitioning's own key attribute at the use site via the new TransformExpression.withReference, so a chained storage-partitioned join keeps every side's partitioning intact.

GroupPartitionsExec.doCanonicalize additionally normalizes the exprIds inside KeyReducer - plan canonicalization does not reach into the plain case class - and the reducer applied for an identity-vs-transform pair is a named IdentityReducer case class, so structurally identical SPJ subtrees with value-equal reducers still compare equal and exchange/subquery reuse keeps deduplicating them.

Why are the changes needed?

In a storage-partitioned join with compatible transforms whose result types differ (e.g. identity(id) on one side and bucket(N, id) on the other), the reducer maps the partition keys to the other side's value type. GroupPartitionsExec.outputPartitioning used to report the original expressions with the reduced keys, so the two had different data types and computing the key ordering threw:

java.lang.ClassCastException: class java.lang.Integer cannot be cast to class java.lang.Long

This fix covers the reducers where the reduced keys equal a single transform applied to one side (identity-vs-transform and single-side-transform). When both sides of a compatible-transform join reduce their keys, the reduced keys are not expressible as a single transform; that shape is a known gap tracked in SPARK-59121.

Does this PR introduce any user-facing change?

No by default. Under spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled, a storage-partitioned join whose reducer changes the partition key data type previously threw ClassCastException and now succeeds.

How was this patch tested?

Added regression tests in KeyGroupedPartitioningSuite. On the current base (including #58420), the multi-table reduce through the reduceKeys trigger fails with ClassCastException - the second reduce's reducer is derived from the stale reported expression, which no type fix reaches - and the two updated SPARK-59120 tests fail; all pass here. The two SPARK-59120 tests now pin the combined behavior: a second join with a non-reducing side plans with no shuffle onto the reported expression, and the shape gate accepts the type-correct reduced layout so only the unpartitioned side is shuffled onto it. The identity-vs-bucket reducer and the subset-join-key cases failed on the pre-#58420 base and pass on the current one, where #58420 fixed the read side; they pin this PR's reporting. The per-KeyedPartitioning retargeting test passes on base and regresses together with the multi-reduce test if the use-site re-targeting is dropped, and the canonicalization test arrived with KeyReducer in this PR. The first three also assert the storage-partitioned join introduces no shuffle.

Ran KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, ProjectedOrderingAndPartitioningSuite, ShuffleSpecSuite, ValidateRequirementsSuite, and DistributionSuite.

[ulysses-you]

cc @peter-toth@szehon-ho thank you

[peter-toth]

Thanks for the PR, @ulysses-you!

The invariant is right, and the identity-vs-transform case is handled exactly: the reducer is the other side's transform applied to this side's attribute, so reporting it is both type-correct and true, and resultType() = reducerExpr.dataType (partitioning.scala:1395) makes the type match by construction. Your three tests do fail on base, I checked that separately.

Two other reducer shapes are not covered, and one of them still throws the same ClassCastException. I put both into a single follow-up item (3) rather than ask you to grow this PR, since the uncovered one needs a design decision, not a patch. We are happy to take these on if you would rather not carry them.

What I would fix here is 2, a shuffle this patch adds in a chained SPJ, and the title, which currently promises the general fix (1). Everything below was measured in a worktree at 96f3cfcb4c9.

Blocking

• 1.Title and description promise more than the patch delivers: The title reads as the general fix, so SPARK-59045 will be resolved on merge while the two-transform ClassCastException still throws (3). Narrowing both to the identity case, and naming what stays open, keeps the ticket honest. The user-facing-change answer also reads as if this patch could make a query fail, where it turns a failure into a success.
• 2.The reduced expression is stamped onto every KeyedPartitioning below:reduced comes from the single spec that collectFirst picked, but the transform applies it to all of them, so the other side's key attribute is overwritten and a GROUP BY on it shuffles where base did not. Measured, same rows both ways. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:83]

Non-blocking

• 3.Two reducer shapes stay uncovered, one of them still crashes: Where only this side reduces but both sides are transforms, the exact expression exists and is not used; where both sides reduce, no expression can describe the keys and the crash survives. Follow-up material, with a measured crash and a measured wrong result behind it. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1383]
• 4.The second trigger is untested: All three new tests enable allowKeysSubsetOfPartitionKeys, so all three fail on base at createShuffleSpec -> toGrouped. Dropping that config from this test reaches the other site, reduceKeys, and the test still fails on base. [inline: sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:590]

Minor

• 5.Comment states an invariant the transform branch breaks: The parenthetical claim that the reported data type matches the reduced keys holds for the identity branch only. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:73]
• 6.A named type for the reducer pair: The raw tuple now appears in six signatures and forces _1 at the use sites. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:1380]

[peter-toth]

Finding 2.reduced was derived from one KeyedShuffleSpec, but this transform applies it to every KeyedPartitioning in the child's partitioning, so the others lose their own key attribute.

createKeyedShuffleSpec takes the first KeyedPartitioning that satisfies the distribution (collectFirst, EnsureRequirements.scala:819), and the reducers are computed from that one alone. When this node sits on a join - a chained SPJ - the child reports one KeyedPartitioning per side, and all of them get that single reduced expression, which references only the chosen side's attribute. The assert above has the same root: the length is only guaranteed to match for the spec's own KeyedPartitioning.

Three tables bucketed 16, 8 and 4 on id, ids 0..15, joined on a.id and grouped by b.id:

• base: KPs = [bucket(16, id#18L), bucket(8, id#20L)], 0 shuffles.
• this commit: KPs = [bucket(16, id#18L), bucket(16, id#18L)], 1 shuffle, because GROUP BY b.id no longer sees a partitioning on b.id.

Both return the same rows, so this is an avoidable shuffle rather than a wrong answer. Retargeting the reduced expression at each KeyedPartitioning's own key attribute fixes it, and the single-attribute invariant makes that well defined:

projectedExpressions.zip(exprs).map {
case (expr, Some((_, reduced))) =>// `reduced` was derived from the spec's `KeyedPartitioning`; re-target it at this one's key so// that every `KeyedPartitioning` in a collection keeps its own attribute.valattr= expr.references.head
reduced.transform { case_: AttributeReference=> attr }
case (expr, None) => expr
}

[ulysses-you]

Fixed in 2b0ef18. GroupPartitionsExec.outputPartitioning now re-targets the reduced expression at each KeyedPartitioning's own key attribute before applying it, so a chained SPJ keeps every side's partitioning on its own attribute. Added the 16/8/4 + GROUP BY middle-side regression test; it asserts 0 shuffles and the full 16-row result.

[peter-toth]

Finding 3. This branch covers two different shapes, and keeping e1 is wrong in both. Not asking you to fix them here - the second one needs a design decision - but they should not be left implicit either.

There are three reducer shapes, and the reported expression has to be judged per shape:

1. This side is identity, the other is a transform. The reduced expression exists and the branch below builds it. Handled, and correctly.
2. Both sides are transforms, only this side reduces. The reduced expression also exists: r(f1(x)) = f2(x) is the ReducibleFunction contract, so it is the other side's transform retargeted at this side's child. This branch reports e1 instead. SPARK-56046: Reducers with same result types is this shape - days reduces onto years, YearsFunction.reducer returns null - and the reported expression stays days(arrive_time) where years(arrive_time) is what the keys hold.
3. Both sides reduce. The keys are r1(f1(x)) = r2(f2(x)), a space that neither transform describes, so no substitution can be correct here.

Shape 3 still throws the exception this PR fixes. reduceKeys types the keys with reducer.resultType() while e1.dataType is e1.function.resultType(), and the suite already ships a pair where they differ: DaysFunctionWithToYearsReducerWithLongResult (DateType) and YearsFunctionWithToYearsReducerWithLongResult (IntegerType) both reduce to LongType. Take SPARK-56164: Reducers with different result types to original keys and add V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS -> true:

java.lang.ClassCastException: class java.lang.Long cannot be cast to class java.lang.Integer
at org.apache.spark.sql.catalyst.expressions.GenericInternalRow.getInt(rows.scala:170)
at org.apache.spark.sql.catalyst.expressions.GeneratedClass$SpecificOrdering.compare(Unknown Source)
at org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning.toGrouped(partitioning.scala:553)
at org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning.createShuffleSpec(partitioning.scala:629)
at org.apache.spark.sql.execution.exchange.ValidateRequirements$.validateInternal(ValidateRequirements.scala:60)

e2 would not help: years is IntegerType and the keys are LongType.

The other cost of a stale expression is that a second join derives its reducer from a transform the data has already left behind. Three tables bucketed 12, 8 and 6 on id, ids 0..11: join 1 reduces both sides with BucketReducer(4), so the keys become id % 4 while the left side reports bucket(12, id); join 2 compares that stale bucket(12, id) with bucket(6, id), gets BucketReducer(6), and (id % 4) % 6 leaves the keys at id % 4, while the right side has gcd == thisNumBuckets, gets no reducer, and keeps id % 6. The two sides are matched across different key spaces: the query returns 4 of 12 rows, 0 shuffles, no error, and all 12 with allowCompatibleTransforms=false.

valcols=Array(Column.create("id", LongType), Column.create("data", StringType))
createTable("b12", cols, Array(bucket(12, "id")))
createTable("b8", cols, Array(bucket(8, "id")))
createTable("b6", cols, Array(bucket(6, "id")))
valvalues= (0 until 12).map(i =>s"($i, 'v$i')").mkString(", ")
Seq("b12", "b8", "b6").foreach(t => sql(s"INSERT INTO testcat.ns.$t VALUES $values"))
valdf= sql(
"""SELECT /*+ MERGE(a, b, c) */ a.id FROM testcat.ns.b12 a |JOIN testcat.ns.b8 b ON a.id = b.id |JOIN testcat.ns.b6 c ON a.id = c.id""".stripMargin)
withSQLConf(SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key ->"true") {
checkAnswer(df, (0 until 12).map(i =>Row(i.toLong))) // returns only 0, 1, 2, 3
}

That one is pre-existing, base reports the same expression, so it is not something this PR broke. It is not master-only either: the same query returns 8 of 12 rows on branch-4.1, where KeyGroupedPartitionedScan.getOutputKeyGroupedPartitioning reports the original expressions next to the reduced common partition values in the same way.

The shape I would suggest, in case it is useful: extend the substitution to shape 2, which is a small step from what you already do for shape 1; and for shape 3 stop trying to express the keys as a transform - carry the reduced key data types on KeyedPartitioning instead. reduceKeys already computes them one frame up and GroupPartitionsExec.groupedPartitionsTuple already holds them as reducedDataTypes, they are just dropped, and the reported partitioning re-derives its types from expressions.map(_.dataType) (partitioning.scala:545). All eight readers go through expressionDataTypes, so they would pick the carried types up for free. Such a partitioning then has to refuse a further reduction, since its expressions no longer describe its keys - which is exactly what the 12/8/6 case above needs.

Happy to pick this up as a follow-up if you would rather not carry it - tell me which way you prefer.

[ulysses-you]

Thanks for the detailed breakdown. Shape 2 is now handled in 2b0ef18: for a single-side reduce we report the target transform re-targeted at this side's attribute, guarded by e2.reducers(e1).isEmpty, so both-sides-reduce positions keep the original expression bit-for-bit (shape 3 unchanged, still a tracked gap).

[ulysses-you]

And for shape 3 — please go ahead and take it as a follow-up as you offered; we are happy to leave the both-sides-reduce shape (and the 12/8/6 stale-expression case) out of this PR. Happy to review your follow-up, and to file a JIRA for it if that is useful.

[peter-toth]

Finding 4. All three new tests enable this config, so all three fail on base at the same place, createShuffleSpec -> toGrouped. There is a second, independent trigger that none of them reach.

This test does not need the config, the join uses the whole partition key. With it removed the test still fails on base, but through reduceKeys at the second join, where the reducer is bound to the stale un-reduced type:

java.lang.ClassCastException: class java.lang.Integer cannot be cast to class java.lang.Long
at org.apache.spark.sql.catalyst.plans.physical.KeyedShuffleSpec$$anon$1.reduce(partitioning.scala:1394)
at org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning$.reduceKeys(partitioning.scala:725)
at org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning.reduceKeys(partitioning.scala:572)
at org.apache.spark.sql.execution.exchange.EnsureRequirements.checkKeyGroupCompatible(EnsureRequirements.scala:561)

and it passes on this commit. Test 1 and test 3 do need the config, test 3 genuinely joins on a subset.

[ulysses-you]

Done in 2b0ef18. Dropped V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS from the multi-table test (both joins use the whole partition key); it now reaches the reduceKeys trigger and fails on base at KeyedShuffleSpec.reduce. Added a comment noting the two triggers.

[peter-toth]

Finding 5. The parenthetical holds for the identity branch only. For two transforms reducers carries e1, so the reported data type is the un-reduced one and can differ from the keys (finding 3). Worth naming the branch the claim applies to, since the next reader will lean on it.

[ulysses-you]

Fixed in 2b0ef18. The comment now names the branches: the data-type match holds for the identity-vs-transform and single-side-transform reducers, and the both-sides-reduce shape has no single transform expression (pointed at KeyedShuffleSpec.reducers).

[peter-toth]

Finding 6.Option[Seq[Option[(Reducer[_, _], TransformExpression)]]] now appears in six signatures and forces _1 at the use sites, including the explain string. A small named type would read better and would document which element is which:

caseclassKeyReducer(reducer: Reducer[_, _], reducedExpression: TransformExpression)

[ulysses-you]

Done in 2b0ef18. Added case class KeyReducer(reducer, reducedExpression) next to KeyedShuffleSpec and switched all signatures/use sites (including the explain string) to it.

[ulysses-you]

Thanks for the review — all feedback addressed in 2b0ef18. Summary:

• 1 (title/description): narrowed the title to the reducer-changes-key-data-type case and rewrote the description: scope is identity-vs-transform and single-side-transform reducers, and the both-sides-reduce shape is called out as a known gap tracked separately.
• 2: per-KeyedPartitioning re-targeting in outputPartitioning (blocking regression fixed, with a chained-SPJ test).
• 3: shape 2 handled (target transform re-targeted at this side's attribute, guarded by e2.reducers(e1).isEmpty); shape 3 intentionally left open.
• 4: multi-table reduce test now exercises the reduceKeys trigger without the subset-partition-keys config.
• 5/6: comment narrowed to the applicable branches; KeyReducer case class replaces the tuple across all signatures.

Tests: KeyGroupedPartitioningSuite (108), EnsureRequirementsSuite + GroupPartitionsExecSuite (45) all pass.

[peter-toth]

LGTM, approving. All six from my round are addressed, and thanks for taking shape 2 as well - the e2.reducers(e1).isEmpty guard with the re-targeted target transform is exactly the shape I had in mind, and it means only the both-sides-reduce case is left.

I confirmed each one against 2b0ef181bc4: the title now names the reducer rather than the transforms, the user-facing answer reads properly, GroupPartitionsExec.outputPartitioning re-targets the reduced expression per KeyedPartitioning so a chained SPJ keeps every side's own attribute, the multi-table test reaches the reduceKeys trigger without the subset config, the data-type comment names the branches it holds for, and KeyReducer gives the pair a name.

One nit, comment only. In the new reduced expression is retargeted per KeyedPartitioning test, the parenthetical reads "(0 shuffles on base, 1 after the fix)", which is inverted - the test asserts no shuffle, and it is the un-retargeted version that produces one. Worth swapping so the next reader is not misled. Also worth knowing: that test passes on master, since the stamping only exists inside this PR, so it guards this change rather than a pre-existing bug.

I will take the both-sides-reduce shape as the follow-up, together with refusing a further reduction while the reported expressions do not describe the keys, and I will file the ticket myself when I open it - no need for you to.

[szehon-ho]

good catch, thanks!

[dongjoon-hyun]

Overall the single-side-reduce fix looks solid and the test coverage is good. Leaving 7 findings as inline comments: 2 correctness concerns and 5 cleanup/design suggestions.

[dongjoon-hyun]

[correctness] The both-sides-reduce branch keeps the un-reduced e1, so the ClassCastException this PR fixes remains reachable in that shape with no runtime fail-safe. The shape is reachable even with in-tree test functions (BucketFunction's gcd reducer reduces both ways; DaysFunctionWithToYearsReducerWithLongResult / YearsFunctionWithToYearsReducerWithLongResult both reduce to LongType while their transforms report DateType/IntegerType). EnsureRequirements' storagePartitionJoinIncompatibleReducedTypesError check passes here because the LEFT and RIGHT reduced types equal each other, so GroupPartitionsExec reports original-typed expressions over reduced-typed partition keys, and a downstream GROUP BY / second join / shuffle deriving an ordering from expressionDataTypes throws the same CCE.

Since the follow-up is tracked separately, could this branch fail safe until then, e.g. return None (fall back to shuffle) or raise the dedicated error when reducer.resultType() != e1.dataType?

[ulysses-you]

Thanks for the detailed analysis. We explored both fail-safe options and decided to keep reporting the un-reduced expression until the follow-up, for two reasons:

1. return None does not actually fall back to a shuffle. With no reducer on either side, the reduced-types cross-check in EnsureRequirements compares the two sides' original types (DateType vs IntegerType in your example) and throws storagePartitionJoinIncompatibleReducedTypesError at planning time. We verified this by running the SPARK-56164 query with a per-pair None implementation - it failed with exactly that error. So None and raising the dedicated error are indistinguishable here: the query cannot run either way, and the error case is already covered by the existing cross-check without new code.

2. None additionally opens a silent-wrong-results hole the dedicated error does not have: if both transforms' original types are equal while the reducers' result type differs (constructible by a third-party ReducibleFunction pair), both sides return None, the cross-check passes on the equal original types, and the two sides' un-reduced key spaces get merged without any reduction.

Since the reduction itself produces correct key values in this shape (only the reported expression is mis-typed), we kept it and left the gap to the follow-up, which carries the reduced data types on KeyedPartitioning and fixes this shape properly.

[dongjoon-hyun]

[correctness]KeyReducer embeds an exprId-bearing TransformExpression in a non-Expression constructor arg. QueryPlan.doCanonicalize normalizes exprIds only via mapExpressions, which does not recurse into the plain KeyReducer case class (@transient affects serialization only, not equals). Pre-PR, a connector returning value-equal Reducer instances on the ReducibleFunction path (e.g. the test BucketReducer(divisor) case class) allowed two identical SPJ subtrees' GroupPartitionsExec nodes to compare equal after canonicalization; post-PR the side-specific AttributeReference exprIds inside reducedExpression break that equality, so exchange/subquery/stage reuse can silently stop deduplicating those subtrees.

[ulysses-you]

Fixed in 9438115. GroupPartitionsExec.doCanonicalize now normalizes the exprIds inside reducedExpression via QueryPlan.normalizeExpressions against the child's output (the same pattern BatchScanExec uses for keyGroupedPartitioning), so structurally identical SPJ subtrees with value-equal reducers compare equal after canonicalization again. Added a test that builds two GroupPartitionsExecs with differently-numbered exprIds and value-equal BucketReducer instances and asserts their canonical forms are equal.

[peter-toth]

Finding 8 (my numbering) — an addition to this, measured on c01535cc6ed.

The normalization does not fire in the two shapes this PR fixes. Since the third commit stores the raw target transform, reducedExpression is the other join side's expression: t for the identity-vs-transform reducer, e2 for the single-side-transform one. QueryPlan.normalizeExpressions(expr, child.output) rewrites an attribute only when child.output contains its exprId, and this node's child is the reduced side, so the other side's attribute is left untouched. It fires only for the both-sides-reduce shape, where reducedExpression is this side's own e1 — the one shape whose expression is stale anyway.

Measured with two GroupPartitionsExecs over LocalTableScanExec(id#1) and LocalTableScanExec(id#2), a value-equal BucketReducer(2), and reducedExpression built over a third attribute standing in for the other side (oid#11 / oid#12):

PROBE other-side attribute in reducedExpression, canonical equal = false

The same pair with reducedExpression over the node's own attribute — the shape the new test builds — gives true.

There is a second, independent cause. The synthesized identity reducer is new Reducer[Any, Any] {...}, so KeyReducer.equals compares it by reference. Taking the reducer from reducersBothWays for an identity(id) / bucket(2, id) pair, with everything else structurally identical:

PROBE identity-derived reducer, canonical equal = false

That half is pre-existing — the anonymous reducer was already a constructor field of GroupPartitionsExec before this PR — so "with value-equal reducers" in the commit message is accurate. It does mean the identity-vs-transform shape, this PR's headline case, still does not deduplicate.

Fix shape: normalize reducedExpression against its own references instead of the child's output — the scaladoc already says the attribute it carries is not load-bearing, so a positional canonical form is enough — and make the synthesized reducer a named case class over the target transform so it compares by value, normalizing that expression too, or the exprId problem you just fixed reappears inside the reducer. Extending the new test with an identity-derived reducer, and with a reducedExpression over an attribute the child does not output, would pin both.

[ulysses-you]

Fixed in 265008f along both lines you suggested: doCanonicalize now normalizes reducedExpression against its own references (the stored expression references the other join side's key, which this node's child does not output), and the synthesized identity reducer is a named IdentityReducer case class over the re-targeted transform, whose transform is normalized too. Extended the test with both probes - a reduced expression over an attribute the child does not output, and an identity-derived reducer - plus a negative control (structurally different reducers stay unequal) and a mixed Some/None multi-key case.

[dongjoon-hyun]

[design] The reduced key type now has two unvalidated sources of truth: reduceKeys uses reducer.resultType() here, while GroupPartitionsExec.outputPartitioning reports reducedExpression.dataType (the target function's resultType()). Reducer's contract implies they agree, but nothing validates it, so a connector whose Reducer.resultType() disagrees with the target transform's type recreates the exact expressions-vs-keys mismatch this PR fixes, hidden inside the new mechanism. Deriving the reduced types from reducedExpression.dataType, or asserting agreement when constructing KeyReducer, would close the drift by construction.

[ulysses-you]

Agreed on the drift risk. We first tried asserting agreement when constructing KeyReducer, but the in-tree DaysToYearsReducerWithDateResult intentionally violates it (SPARK-56046: Reducers with different result types depends on such a reducer reaching the cross-check error), so a hard assert would break that test. Instead the agreement is now closed by construction / by the existing check in 9438115: the identity-vs-transform reducer reports the transform itself; the single-side branch reports the target transform, whose type the reduced-types cross-check in EnsureRequirements validates against the other side's target-transform-typed keys; and the both-sides-reduce branch reports the original expression (known gap, see the other thread). Comments in reducersBothWays spell this out.

[dongjoon-hyun]

[simplification] This creation-time retargeting is dead work: the only structural consumer of reducedExpression (GroupPartitionsExec.outputPartitioning) unconditionally retargets it again at each KeyedPartitioning's own attribute, and all other reads use only .reducer (reduceKeys, displayName in QueryExecutionErrors). Storing raw e2 and doing the single retarget at the use site is equivalent, removes one tree transform + allocation per expression pair, and stops implying that the attribute chosen here (from whichever spec collectFirst picked) is load-bearing.

[ulysses-you]

Fixed in 9438115. The reducer construction now stores the raw expression, and GroupPartitionsExec.outputPartitioning does the single re-targeting at the use site (via the new TransformExpression.withReference helper).

[dongjoon-hyun]

[efficiency]e2.reducers(e1) is invoked per expression pair solely to test .isEmpty, materializing a reverse Reducer via the catalog function's reducer() call. TransformExpression.reducers is an uncached def, and EnsureRequirements already calls spec.reducers in both directions, so each direction's catalog lookup now runs twice per join-key pair (4 instead of 2) on every SPJ planning pass, including AQE re-planning. A slow or allocating third-party ReducibleFunction.reducer() pays this on a discarded probe; computing both directions once per pair would avoid it.

[ulysses-you]

Fixed in 9438115. The two directions are now computed in one pass, KeyedShuffleSpec.reducersBothWays, so each expression pair's reducers lookup runs once per direction: the reverse lookup a direction used only for the single-side probe is exactly the other direction's reducer. EnsureRequirements now calls it once instead of calling reducers on both specs.

[dongjoon-hyun]

[reuse] This is the third copy of the X.transform { case _: AttributeReference => attr } retargeting idiom (also at KeyedShuffleSpec.reducers in both branches), spread across two modules with no shared helper. This copy must stay exactly consistent with how reducedExpression was built in KeyedShuffleSpec.reducers - a future change to the rebind rule (e.g. per-position rebinding for multi-reference or nested-field transforms) can silently miss one copy, making the reported outputPartitioning stop matching the actual reduced keys. A small helper on TransformExpression (returning TransformExpression, which also removes the asInstanceOf casts) would fit this PR.

[ulysses-you]

Fixed in 9438115. Added TransformExpression.withReference(attr) in catalyst and switched all the copies (the reducer construction in both directions and the use site) to it, which also removes the asInstanceOf casts at the call sites.

[dongjoon-hyun]

[style] Nit: this signature (and object KeyedPartitioning.reduceKeys below) was re-wrapped with the bare return type on its own continuation line, but the one-line forms fit within the 100-char limit after the type shortened (Seq[Option[KeyReducer]] is shorter than Seq[Option[Reducer[_, _]]]), and this wrapping style is not used elsewhere in the file. Keeping the original one-line form avoids the formatting churn.

[ulysses-you]

Fixed

[peter-toth]

Let me review this again today and check if the new keyDataTypes introduced in #58262 can help with this PR.

stale

[peter-toth]

Re-checked through c01535cc6ed — findings 1-6 resolved (title narrowed, per-KeyedPartitioning re-targeting, shape 2 handled, the reduceKeys trigger covered, the comment scoped to its branches, KeyReducer named), nothing regressed. All five SPARK-59045 tests pass here, and KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, ValidateRequirementsSuite, ShuffleSpecSuite and DistributionSuite are green — 193 tests.

As I said above I would, I checked what the keyDataTypes in #58262 does for this PR: on your merge base with only that PR's key-reading half applied, tests 1 and 4 pass — projectKeys and toGrouped stop reading reduced keys at the expressions' declared types — but test 2 still throws. There the reducer for the second join is derived from the stale expression and applied to values it was not built for, which no type fix can reach. So the two changes are complementary rather than overlapping: mine stops the reads from crashing, yours makes the reported expression true, and only yours fixes the chained reduce.

A heads-up on merge order while I am here: this PR conflicts textually with #58262 (reduceKeys, the reduce block in EnsureRequirements, the base types in GroupPartitionsExec) and with #58351 (GroupPartitionsExec.outputPartitioning and the KeyedPartitioning construction in it). Both are mechanical; I found no semantic overlap, and I will rebase whichever of mine lands second.

Blocking

• 7.Description drift (new): the description describes KeyedShuffleSpec.reducers, which c01535cc6ed renamed to reducersBothWays, and it does not mention that commit's canonicalization normalization or the new TransformExpression.withReference. It also says "The tests fail on the base commit and pass after", which holds for three of the five: reduced expression is retargeted per KeyedPartitioning passes on base (measured — it guards this PR's own re-targeting instead, and ablating withReference at the use site fails it together with the multi-reduce test), and canonicalization normalizes the reduced expressions does not compile on base, since it references KeyReducer.

Non-blocking

• 8.doCanonicalize normalizes nothing in the two shapes this PR fixes (new):reducedExpression now holds the other side's expression, whose attribute is not in this node's child.output, so normalizeExpressions leaves it as it is; and the synthesized identity reducer is an anonymous class, so KeyReducer equality fails on the reducer regardless. Both measured — details in my reply on @dongjoon-hyun's thread.

Minor

• 9.withReference keeps the source expression's path above the leaf attribute (new): re-targeting a GetStructField-based transform at a bare attribute reports the source's field ordinal. No producer reaches it today, so this is a scaladoc point. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:106]

[peter-toth]

Finding 9. The premise is right — KeyedPartitioning.supportsExpressions allows one leaf attribute per transform — but it also allows a GetStructField above that attribute, and transform replaces only the attribute, so the path comes from the source expression rather than from the target.

GroupPartitionsExec.outputPartitioning calls this with expr.references.head, which is the bare attribute. Re-targeting the other side's bucket(4, GetStructField(other, 0)) at a struct s whose own key expression is bucket(4, GetStructField(s, 1)) therefore reports field 0 — the source's ordinal read on this side's struct.

I could not reach it from a query, so this is a scaladoc point rather than a live defect: a nested partition reference resolves to an Alias over GetStructField, and supportsExpressions' isReference does not unwrap Alias, so a nested-field partition transform produces no KeyedPartitioning at all today (I measured that for identity("s.a"), which reports UnknownPartitioning). Worth saying in the doc that the caller has to pass the reference of the whole key expression, not just an attribute, so that whoever makes nested keys work does not inherit a silent wrong-field report.

[ulysses-you]

Done in 265008f - the scaladoc now says that any field path above the leaf attribute comes from this expression, and that re-targeting is only faithful when the source and the target key expressions have the same path shape.

[ulysses-you]

Updated the PR description (pushed as 265008f): reducers is now reducersBothWays, it mentions the canonicalization normalization and the named IdentityReducer behind it as well as the new TransformExpression.withReference, and the test claims are reworked - only three of the five tests fail on base (identity-vs-bucket, multi-reduce, subset keys); the per-KeyedPartitioning retargeting test guards this PR's own re-targeting (as you measured, it regresses together with the multi-reduce test when the use-site re-targeting is dropped), and the canonicalization test arrived with KeyReducer in this PR.

[peter-toth]

Re-checked through 265008fc4e8 — findings 7, 8 and 9 resolved, nothing regressed. The description now matches the code, doCanonicalize normalizes reducedExpression and the identity reducer's transform against their own references, IdentityReducer is a named case class, and withReference's scaladoc states the path-shape condition.

I re-measured rather than carrying the earlier verdicts forward, since the base moved:

• on the current base af0d5c3abd1 the identity-vs-bucket, multi-reduce and subset-key tests all fail with ClassCastException: Integer cannot be cast to Long;
• the re-targeting test passes on base and fails once reduced.withReference(expr.references.head) is replaced by reduced;
• the canonicalization test fails once doCanonicalize is removed, so it now covers both halves it was written for;
• 187 tests green across KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, ShuffleSpecSuite and ProjectedOrderingAndPartitioningSuite.

Non-blocking

• 10.Affects versions and backports (late catch):SPARK-59045 lists 5.0.0 as the only affected version, but the same code is on the maintenance branches. KeyedShuffleSpec.reducers has the identity-vs-transform arm and GroupPartitionsExec.outputPartitioning builds KeyedPartitioning(projectedExpressions, groupedPartitions.map(_._1), ...) with no reduced expression on branch-4.2, branch-4.3 and branch-4.x. SPARK-59121, the follow-up for the both-sides shape of the same mechanism, already lists 4.0.0 and 4.2.0. Worth updating the JIRA and saying whether backports are planned. Note v4.3.0-rc1 is already cut, so 4.3.0 would need a respin.

Minor

• 11.Test comment reads backwards (new): "(0 shuffles on base, 1 after the fix)" says the fix adds a shuffle. The 1 belongs to the ablation, not to the fix. inline
• 12.Name the follow-up ticket (new): the both-sides-reduce gap is now SPARK-59121. inline

[peter-toth]

Finding 12. The follow-up has a ticket now, so it can be named here.

SPARK-59121, "Fix wrong results when a storage-partitioned join reduces the partition keys of both sides". It covers this shape plus the silent wrong-results half of it.

Suggested change

	// a ClassCastException. Known gap, tracked in the follow-up for SPARK-59045.
	// a ClassCastException. Known gap, tracked in SPARK-59121.

[ulysses-you]

Done in c161a21 - the gap comment (and the PR description) name SPARK-59121 now.

[peter-toth]

Finding 11. The parenthetical is the wrong way round, it reads as if the fix introduces the shuffle.

Measured on 265008fc4e8: base af0d5c3abd1 gives 0 shuffles and the test passes, this head gives 0 shuffles and passes, and only replacing reduced.withReference(expr.references.head) with reduced in GroupPartitionsExec.outputPartitioning makes it 1 shuffle. So the 1 belongs to the ablation.

Suggested change

	// sees a partitioning on it and the query shuffles (0 shuffles on base, 1 after the fix).
	// sees a partitioning on it and the query shuffles (0 shuffles on base and here, 1 if the
	// use-site re-targeting is dropped).

[ulysses-you]

Fixed in c161a21 with your wording - the 1 belongs to the ablation, not to the fix.

[peter-toth]

#58420 is merged to master, branch-4.x and branch-4.3, so this needs a rebase. Two of its query tests will fail on it, both because this PR makes those shapes work rather than because anything broke:

• SPARK-59120: incompatible reduced key types are reported instead of cast - the error is no longer raised, the query returns the right rows with no shuffle. The error path stays covered by SPARK-56046: Reducers with different result types, so this one can become a "works now" pin.
• SPARK-59120: another child is not shuffled onto reducer-rewritten keys - one shuffle instead of two, onto the now type-correct years(ts) layout.

Measured on 265008fc4e8 without #58420, with the rest of the SPJ suites green there. The full comparison is on the #58420 thread.

With those two updated and the three items from my last round answered, this is good to go from my side.

[ulysses-you]

Rebased onto 7f12c2d and pushed as c161a21; both SPARK-59120 tests are updated as you predicted - the non-reducing identity side now reduces onto the reported years(a.ts) and the query plans with no shuffle, and the shape gate accepts the reduced layout so only the unpartitioned side is shuffled onto it (one shuffle; the second test is renamed accordingly).

I also re-measured the fail-on-base claims against the new base: only the multi-reduce test still fails there (CCE - the second reduce's reducer is derived from the stale expression, which no type fix reaches), while the identity-vs-bucket and subset-key cases pass on it now that #58420 fixed the read side. The description's Tested section reflects that.

For the affected versions (finding 10): SPARK-59045 will be updated to 4.2.0/4.3.0/5.0.0 - the crashes start in 4.2.0 with the KeyedPartitioning/GroupPartitionsExec refactor. The merge covers master (5.0.0) and branch-4.x (4.3.0, which needs a respin since rc1 is cut); no separate branch-4.2 backport for now since the path is opt-in, but happy to cherry-pick if there is demand.

[peter-toth]

Re-checked through c161a21d522 — findings 10, 11 and 12 resolved, nothing new.

I re-measured the two rewritten SPARK-59120 tests against the new base 7f12c2da0b4 rather than reading the assertions, since they are pinning a behaviour change rather than a fix. Both fail there and pass here: the second-join test raises STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES on base, and the shuffle-count test plans 2 shuffles instead of 1. Your read of what changed is exactly right, and both new comments match what the base actually does. Of your own five, the multi-reduce test is the one that still fails on base, as you said; identity-vs-bucket and the subset-join-key case now pass there. I also diffed the round-3 and round-4 PR diffs against each other, and the only differences are the four intended ones, so the rebase carried nothing extra.

230 tests green across ShuffleSpecSuite, DistributionSuite, KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, ValidateRequirementsSuite and ProjectedOrderingAndPartitioningSuite.

One thing left from finding 10, since it is easy to lose before merge: SPARK-59045 still lists 5.0.0 as the only affected version. Your plan for it reads right to me, including leaving branch-4.2 alone — with #58420 backported there, 4.2 falls back cleanly instead of crashing, so the opt-in path is safe without this PR.

Thanks for working through all of these, @ulysses-you — nothing left open from my side.

[szehon-ho]

lgtm, just some ai comment nits

[szehon-ho]

let's remove the comment about 'all producers today'? doesnt seem necessary to understand this function by itself ?

[ulysses-you]

Removed in a74b22f - the scaladoc keeps just the path-shape rule.

[szehon-ho]

optional: quick example to understand this?

[ulysses-you]

Added an example in a74b22f: an identity(id) side joined to a bucket(4, id) side reduces with IdentityReducer(bucket(4, id)), which maps each raw id value to its bucket(4, id) value. The concrete result is up to the connector's bucket algorithm, so the doc does not state one.

[szehon-ho]

(here and other tests), can we not have comments about 'this pr' and just a brief explanation what this is testing

[ulysses-you]

Reworded in a74b22f - the comment now describes what the test pins and what the base did, without self-reference.

[szehon-ho]

same, let's remove 'this pr' and just explain what test does

[ulysses-you]

Reworded in a74b22f, same as the other test.

[szehon-ho]

Two inline review findings from the current head.

[szehon-ho]

The new KeyReducer(reducer, e2) behavior is important but not directly pinned. A chained days(ts) JOIN years(ts) JOIN identity(ts) test would expose regressions: reporting stale days(ts) after the first join causes the identity side to reduce to days while the stored keys are years, potentially returning no rows. Could we extend the existing SPARK-56046 test with this third join and assert both the answer and absence of shuffles?

[ulysses-you]

Done in a74b22f - the SPARK-56046 same-result-types test now also joins a third table partitioned by identity(time), asserting the answer and no shuffles. Verified by ablation: reporting the stale original expression fails it (the identity side then reduces onto the stale expression, and the shape gate refuses the mis-described reduced layout).

[szehon-ho]

After #58420, key ordering reads keyDataTypes, so leaving the original expression here no longer causes the ClassCastException described above; this test passes on the current base. Could this comment instead explain that accurate expressions are required for downstream reduction and layout reuse?

[ulysses-you]

Reworded in a74b22f - it now explains that another join can reduce onto the reported transform, and a reduced layout can serve as another child's shuffle target, only while the expressions describe the keys.

[szehon-ho]

LGTM. I still am not a big fan of too much unrelated code history in the comment, but dont want to block it, feel free to either fix or merge

[szehon-ho]

I think its still too much code history narrated here for a test comment

[ulysses-you]

Trimmed in 6ece99a - both SPARK-59120 test comments now just describe what they pin; the history lives in #58420's description and tests.

[szehon-ho]

Thanks!

[ulysses-you]

Merge Summary:

• merged into master a76ef57
• merged into branch-4.x ca915bf
• merged into branch-4.3 0d253fe

Posted by merge_spark_pr.py

[ulysses-you]

thank you all !

[peter-toth]

Now that this is in, one question about the branch reach.

It went to master, branch-4.3 and branch-4.x, but not to branch-4.2, and the JIRA still lists 5.0.0 as the only affected version.

branch-4.2 looks affected to me. I read the two sites there: KeyedShuffleSpec.reducers still has the identity-vs-transform arm that hands back a bare reducer, and GroupPartitionsExec.outputPartitioning still reports the un-reduced expressions. SPARK-59120 was backported to that branch, so the keyDataTypes half is already in place there.

Would you like a [SPARK-59045][SQL][4.2] backport? I am happy to open one.

I ask because SPARK-59121, the follow-up for the shape where both sides reduce, is a change on top of this one, so it can only reach 4.2 after this does.

[ulysses-you]

@peter-toth created #58451

[peter-toth]

Seems like this causes a test failure on branch-4.3:

This is from https://github.com/apache/spark/actions/runs/33502924519/job/99854844877:

2026-09-01T14:03:48.8434765Z - SPARK-59045: compatible transforms reduce data type with subset join keys *** FAILED *** (326 milliseconds)
2026-09-01T14:03:48.8442493Z List(Exchange hashpartitioning(id#1257917L, 5), ENSURE_REQUIREMENTS, [plan_id=1685013]
2026-09-01T14:03:48.8443966Z +- *(1) Project [id#1257917L, dt#1257918, data#1257919]
2026-09-01T14:03:48.8444868Z +- *(1) Filter isnotnull(id#1257917L)
2026-09-01T14:03:48.8446361Z +- BatchScan testcat.ns.t1[id#1257917L, dt#1257918, data#1257919] class org.apache.spark.sql.connector.catalog.InMemoryBaseTable$InMemoryBatchScan RuntimeFilters: []
2026-09-01T14:03:48.8447458Z , Exchange hashpartitioning(id#1257920L, 5), ENSURE_REQUIREMENTS, [plan_id=1685035]
2026-09-01T14:03:48.8448392Z +- *(2) Project [id#1257920L, dt#1257921, data#1257922]
2026-09-01T14:03:48.8448965Z +- *(2) Filter isnotnull(id#1257920L)
2026-09-01T14:03:48.8450163Z +- BatchScan testcat.ns.t2[id#1257920L, dt#1257921, data#1257922] class org.apache.spark.sql.connector.catalog.InMemoryBaseTable$InMemoryBatchScan RuntimeFilters: []
2026-09-01T14:03:48.8451276Z ) was not empty storage-partitioned join should not shuffle (KeyGroupedPartitioningSuite.scala:659)
2026-09-01T14:03:48.8452107Z org.scalatest.exceptions.TestFailedException:

cc @HeartSaVioR, @dongjoon-hyun

[dongjoon-hyun]

Please feel free to revert it from branch-4.3, @peter-toth .

[peter-toth]

@ulysses-you, do you think you can fix it in a follow-up tailored for branch-4.3?

[ulysses-you]

thank you @peter-toth and @dongjoon-hyun , created followup for 4.3 #58460