[SPARK-59121][SQL] Fix wrong results when a storage-partitioned join reduces the partition keys of both sides

[peter-toth]

What changes were proposed in this pull request?

This builds on #58335, which makes a reduced key-grouped partitioning report a transform that describes its keys for the two reducer shapes where one exists. When both sides of the join reduce there is none. The keys become r1(f1(x)) = r2(f2(x)), a third key space that neither side's transform names. bucket(12, id) joined to bucket(8, id) is the flagship example. BucketFunction.reducer hands both sides BucketReducer(4), the keys become id % 4, and both sides keep reporting the transforms they were built from. That is the gap #58335 names and leaves open, and this PR closes it by saying so rather than by inventing an expression.

TransformExpression gains a fourth field, reducedWith: Option[TransformFunctionId], which names the transform this one's keys were reduced together with. KeyedShuffleSpec.reducersBothWays is the only producer. In its both-sides-reduce branch it now reports e1.reducedTogetherWith(e2) instead of the bare e1. The marker holds no Expression, only a canonical name and a bucket count, so it is safe in a field canonicalization does not descend into.

Putting it on the expression rather than on KeyedPartitioning is what keeps the change small. Every site that derives a partitioning already carries the expressions along. AliasAwareOutputExpression projects them, GroupPartitionsExec.outputPartitioning re-reports them, and TransformExpression.withReference re-targets them, so the marker is inherited, and dropped with the position it belongs to, without a line of new plumbing. GroupPartitionsExec needed no change at all.

Four sites then refuse to reason about such keys, and a fifth refusal falls out of expression equality:

• KeyedShuffleSpec.isExpressionCompatible does not compare marked keys by transform. Two of them are compatible when the same pair was reduced together, which is the pair the join produced. Anything else has to shuffle.
• KeyedShuffleSpec.canCreatePartitioning does not shuffle another child onto marked keys, because that evaluates the reported expressions per row.
• KeyedShuffleSpec.reducersBothWays does not reduce marked keys a second time.
• keysSatisfy does not let marked keys satisfy an OrderedDistribution. An ordering is a claim about the key values, and nothing makes a reducer order-preserving. This one is a local guard. A marked position always carries a transform, and a SQL ORDER BY cannot name one.
• UnionExec.comparePartitioning compares the children's expressions with semanticEquals, so a marked child no longer merges with an unmarked sibling reporting the same transform. No change was needed there.

Clustering is the one thing a marked partitioning still satisfies, and that is sound. The keys remain a function of the same attributes, which is all ClusteredDistribution asks.

Two things promised in the review of SPARK-59120 (#58420) are now delivered, since the marker answers the question they approximated: KeyedPartitioning.expressionsDescribeKeyShape and its use in canCreatePartitioning are replaced by expressionsDescribeKeys, and the three-reader list in the keyDataTypes scaladoc is gone.

Keeping the type check as a second conjunct was considered and dropped. It is not sufficient, since bucket(12) and bucket(8) reducing onto bucket(4) keep their IntegerType and pass it. It is not necessary either. After a one-side reduce the reported expression is the target transform, and EnsureRequirements already refuses a reducer whose resultType() disagrees with the other side's key types, which are that transform's, with STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES. So the keys and the expressions agree at every reachable one-side reduce, and a both-sides reduce is marked. The unit test that pinned the type proxy is replaced by one that pins the marker, and it keeps the proxy's struct case. A partitioning whose key row was built at struct<f> while it declares struct<g>, which is what createPartitioning produces, is still accepted. Re-adding any type comparison to the gate fails that assertion.

Two smaller things came with it. TransformExpression.resolvedFunction refuses a marked expression, so eval throws on one instead of computing the un-reduced transform and misrouting the row. That is a local gate rather than a live check, because every consumer of a reduced partitioning refuses it first and the write path never sees one. It also replaces the copy of the rule that prepends the bucket count as a literal argument in DistributionAndOrderingUtils, and both copies had to be touched here anyway, since the extractor grew a field.

The keyDataTypes scaladoc also records where its no-key fallback stops being truthful. A marked partitioning can end up with no key, for instance when v2BucketingPartitionFilterEnabled intersects two sides that hold disjoint keys, and it then reports the un-reduced transform's type while the other leg of the same pairing reports the reducer's. SPARK-59176 tracks that, with the repro and two ways to fix it. The same query fails on a ClassCastException without this marker, so nothing regresses here.

Why are the changes needed?

Four failures, each one a test here, all measured on this PR's base. All of them need spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled, which is off by default, since that is what admits a reducer at all.

1. Two reduces onto different key spaces look co-partitioned. Four tables holding ids 0 until 24, bucketed by 12 and 8 on one side and by 12 and 18 on the other. The first pair reduces onto id % 4, the second onto id % 6, and both keep reporting bucket(12, id). Joining the two returns 8 of 24 rows with no shuffle. This is also what the pairing in the marker is for. Marking the keys without recording which pair produced them fixes everything else here and still returns those 8 rows.
2. Reducing keys twice. Three tables bucketed by 12, 8 and 6: the first join reduces onto id % 4, then the second derives a bucket(6) reducer from the reported bucket(12, id) and applies it to keys that are already reduced. (id % 4) % 6 leaves the left keys alone while the right side moves to id % 6, so the query loses rows with no shuffle.
3. Reducing keys twice, the other way. Two days/years joins that each reduce both sides onto one year space, then joined to each other with different key sets on the two legs. The outer join derives reducers again and planning throws ClassCastException: Long cannot be cast to Integer.
4. Merging a reduced partitioning in a union. (bucket12 JOIN bucket8) UNION ALL bucket12 reports bucket(12, id) on both sides of the union, so the two are merged although the join side's keys are id % 4. A GROUP BY id above it returns each id twice.

The fifth refusal, canCreatePartitioning, guards a hazard the first one opens up. Once a marked partitioning is no longer compatible with anything but its own pairing, EnsureRequirements takes the one-side-shuffle path instead, and there it would happily shuffle the other child onto the reduced keys, which places rows by evaluating bucket(12, id) against partitions laid out by id % 4. Measured with tables bucketed by 12, 8 and 2 and v2BucketingShuffleEnabled, with the clause removed the query returns 8 of 12 rows.

Does this PR introduce any user-facing change?

Yes, it fixes the wrong results and the crash above. A query that used to reach one of them now shuffles instead, so its plan changes.

It also refuses two plans that happen to be correct today, In both cases nothing can tell the two apart from the outside.

The first is a second join onto an already reduced space. bucket(12) JOIN bucket(8) lands on id % 4. Meeting a bucket(4, id) or bucket(2, id) table, BucketReducer would compose correctly, since those counts divide 4. Meeting a bucket(6, id) table it would not, and that is failure 2 above. The Reducer API cannot say which of the two it is, so both now shuffle: bucket12 JOIN bucket8 JOIN bucket4 goes from 0 shuffles to 2, with the same rows. This only arises where both sides reduced, i.e. where gcd(a, b) < min(a, b). A chain in which one bucket count divides the other takes the one-side path and is unaffected.

The second is two reduces that land on the same space through different pairings, e.g. bucket(12) with bucket(8) and bucket(12) with bucket(20), both of which give id % 4. They are treated as different spaces and the join shuffles.

Both cost a shuffle, not a wrong answer. A follow-up can give them back by letting a ReducibleFunction name the transform it reduces onto, which makes this whole shape disappear rather than be refused.

Queries whose keys were not reduced on both sides are unaffected, and a join of two sides that were reduced together keeps its plan.

How was this patch tested?

Ten new tests. Six query tests in KeyGroupedPartitioningSuite, and one each in TransformExpressionSuite, ShuffleSpecSuite, GroupPartitionsExecSuite and ProjectedOrderingAndPartitioningSuite. One existing test changed, see below.

Four of the six query tests fail on the base commit with wrong rows or a crash, listed as failures 1 to 4 above. A fifth, another side is not shuffled onto reduced keys, returns the right rows there with no shuffle. This PR costs it one shuffle, and it is what pins the canCreatePartitioning clause. Each refusal was ablated, and each ablation fails exactly the tests written for it:

• canCreatePartitioning's clause removed: another side is not shuffled onto reduced keys returns 8 of 12 rows.
• the same-pairing allowance in isExpressionCompatible made unconditional, i.e. the refusal taken too far. two sides reduced onto the same keys still join without a shuffle fails, and so does the existing SPARK-56164: Reducers with different result types to original keys.
• the pairing thrown away, so that any two marked keys count as one space. two reduced partitionings are not compatible by their transforms returns 8 of 24 rows. That is the ablation the test exists for, and it is why the marker records the pair rather than a bit.
• the reducersBothWays guard removed. two sides reduced together are not reduced a second time throws the ClassCastException. No other test in the suite reaches that guard, and instrumenting it to throw on entry showed every one of them misses it. That is why the test is built the way it is, with different key sets on the two legs so that the join computes reducers at all.

SPARK-56164 is the existing test for the both-sides-reduce shape, and it gained one assertion, that the join's own requirements still validate, i.e. that the two sides which were reduced together are still co-partitioned. That is the no-regression half of this change, and it belongs on the test that already describes the shape rather than in a copy of it.

The four unit tests cover what a query test states only indirectly. They are the pairing relation itself, the canCreatePartitioning refusal, that a node which reduces nothing (or reduces another position) inherits the marker from its child, and that a projection drops the marker with the position it belongs to. The last two hold by construction under this representation, which is exactly why they are worth keeping, since an earlier design of this fix carried the marker on KeyedPartitioning and got both wrong.

Green: ShuffleSpecSuite, DistributionSuite, TransformExpressionSuite, KeyGroupedPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, ProjectedOrderingAndPartitioningSuite, ValidateRequirementsSuite, WriteDistributionAndOrderingSuite, PlannerSuite, UnionSuite, DataSourceV2Suite, 488 tests in all. dev/lint-scala is clean.

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

Generated-by: Claude Code

[dongjoon-hyun]

+1, LGTM.

[peter-toth]

@dongjoon-hyun, @ulysses-you, I updated the PR and it is now ready for review.

[ulysses-you]

How the pass was verified

Claim-propagation audit

The marker's claim — "this expression no longer computes the keys; the keys live in the key space identified by the pairing" — was checked across all three sets:

• Producers (exactly one): KeyedShuffleSpec.reducersBothWays marks both sides in the same step (partitioning.scala:1646/1656), so the marker is symmetric by construction. The (None, None) guard (:1619-1625) sits ahead of both the transform-transform and the identity-transform cases, so no reducer is ever derived from or applied to a marked position, including IdentityReducer.
• Carriers: KeyedPartitioning.project, AliasAwareOutputExpression.projectKeyedPartitionings, GroupPartitionsExec.outputPartitioning (both the no-reducer and reduce-other-position shapes), UnionExec.prepareOutputPartitioning's attribute rewrite, TransformExpression.withReference, and GroupPartitionsExec.doCanonicalize's positional normalization all keep the marker, because it is a plain case-class field that no transform/canonicalization descends into; dropping a position drops the marker with it (pinned by the two unit tests). PartitioningCollection members cannot disagree on the marker: marking is per-reduce-step and symmetric, and the collection constructor already requires shared key reference/arity/isCollapsed; the comment at :1013-1017 documents why expressionsDescribeKeys stays out of the invariant list, and satisfies0's exists over members is safe under it.
• Consumers (all enumerated): isExpressionCompatible (gated by pairing), canCreatePartitioning (gated), reducersBothWays (guarded), keysSatisfy's OrderedDistribution arm (gated), UnionExec.comparePartitioning (refuses via semanticEquals, since the field changes expression equality), eval/resolveFunctionCall (refuse, throw rather than misroute), and the executor-side ShuffleExchangeExec key extractor (:475-490), which is only reachable through KeyedShuffleSpec.createPartitioning, which is only reachable through candidateSpecs, which filters on canCreatePartitioning — the chain closes.

The ungated arm is the sound one

keysSatisfy's ClusteredDistribution arm deliberately accepts marked partitionings. The clustering property (same clustering value implies same partition) holds because the reduced keys are still a deterministic function of the same attributes, and clustering asks nothing more. Aggregates/windows/joins on top are all safe under exactly that property; the consumers that need more (ordering, placement, re-reduce) are the gated ones.

Write path unreachable

A KeyedPartitioning becomes a shuffle target only via KeyedShuffleSpec.createPartitioning (the only createPartitioning call site producing one is EnsureRequirements.scala:268); write rebalance/repartition always build HashPartitioning/RoundRobinPartitioning from HasPartitionExpressions, never keyed. So the scaladoc claim "the write path never sees one" holds, and the resolveFunctionCall refusal inDistributionAndOrderingUtils is correctly a local gate (its marked arm is unreachable because write-side transforms come from the catalog, unmarked).

Cross-path axes probed

• codegen/interpreted: doGenCode throws for all TransformExpressions, marked or not — no asymmetry introduced; eval is the only runtimepath and is gated.
• AQE: coalescing a keyed shuffle downgrades outputPartitioning to UnknownPartitioning (AQEShuffleReadExec.outputPartitioning),conservative and pre-existing; local-read conversion of ENSURE_REQUIREMENTS shuffles predates this PR and self-corrects through re-optimization (out of scope).
• Serialization: TransformFunctionId is String + Option[Int], carries no exprIds, canonicalization-invariant — the marker survives the wire as a plain field.
• Ordering: the OrderedDistribution refusal has documented causality (a reduced position always carries a transform; ORDER BY/requiredChildOrdering children are attributes or resolved scalar calls, never TransformExpression, so areAllClusterKeysMatched alreadyrefuses every reachable case) — the comment states this rather than a test, which is the accepted alternative.

[ulysses-you]

+1 to make the condition strict.

The information is already lost at reduce time. After a both-sides reduce the real key space is r1(f1(x)), which no transform names, and the current Reducer API cannot say which space it reduced onto. All the planner has is the reported original expression (e.g. bucket(32, id)).

If the gate were widened to allow deriving another reducer from the reported expression (i.e. reducing already-reduced keys again), it would admit not only the divisible, benign case above but also failure 2 of this PR: bucket(12) JOIN bucket(8) lands on id % 4, then joining bucket(6) derives (id % 4) % 6 — the left keys stay put while the right side moves to id % 6, and the query loses rows silently.

This is not an implementation accident; it is mathematically undecidable from the reduced values: a key of id % 4 = 0 can stand for id = 0, 4, 8, whose id % 6 values are 0, 4, 2 respectively. Mapping from a reduced space to a non-divisible one is simply not well-defined. So unless the exact reduced space is known, "reduce again" is only sound in the divisible case, and the available information cannot tell divisible from non-divisible.

So hasSameReducedKeys admits only the one provably sound case: both sides came out of the same reduce (same pairing), which guarantees the same layout. Everything else has to shuffle. That includes:

1. A further reduce onto a divisor (the 32/12 vs 64/24 case above: id % 4 vs id % 8).
2. Different pairings that happen to land on the same space (12 JOIN 8 and 12 JOIN 20 both give id % 4).
3. Meeting an unreduced table that is already laid out on that space (bucket(12) JOIN bucket(8) reduced onto id % 4, joined to a raw bucket(4, id) table whose keys are already aligned).

The cost in all three is an extra shuffle, never a wrong answer. On the base these shapes ran with 0 shuffles — correct by accident in the divisible cases, losing rows in the non-divisible ones — and the API cannot keep only the good half, so tightening the gate wholesale was the only safe choice available.

[uros-b]

Thank you @peter-toth and @dongjoon-hyun@ulysses-you!

[dongjoon-hyun]

LGTM. The four wrong-result/crash cases are fixed and I could not find a regression. Two small nits inline, both fine to address here or in a follow-up.

[dongjoon-hyun]

nit: the two cases described here can overlap, and the doc does not say what happens then. A both-sides-reduced (marked) partitioning can end up with zero partitionKeys, e.g. an inner join with V2_BUCKETING_PARTITION_FILTER_ENABLED whose two sides hold disjoint keys, so mergeAndDedupPartitionKeys(intersect = true) yields Nil. Then keyDataTypes falls back to expressionDataTypes, i.e. the un-reduced transform's type (DateType for days), while the sibling leg of the same pairing reports the reducer's type (LongType). EnsureRequirements compares the two at the reduced-types check without consulting expressionsDescribeKeys, since reducersBothWays returns (None, None), and throws STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES for a correct (empty) plan.

This is pre-existing on the base commit, so not a blocker. But since this PR rewrites the doc, it would be worth one sentence noting that the no-key fallback is not truthful for a marked expression, or a SPARK- reference if you prefer to track it separately. Carrying the reducer's resultType() in the marker would fix it properly.

[peter-toth]

Thank you, this reproduces. I measured it both ways, on days/years tables where one leg's two sides hold disjoint keys, so the intersect empties it.

On this PR the query fails with STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES. On the base commit the same query fails with ClassCastException: Long cannot be cast to Integer, from applying a reducer to already reduced values. So the shape is broken either way and nothing here regresses.

One correction to the framing. The untruthful fallback is pre-existing, it arrived with SPARK-59120. The failure at the reduced-types check becomes reachable through this PR, because once no reducer is derived for an already reduced pair both sides read keyDataTypes directly. Before that the reducers were computed and both sides reported the reducer's type, so the check passed and the query died later.

Filed as SPARK-59176 with the repro and the two ways to fix it. The scaladoc now points at it, and I have started on it in a separate branch.

[dongjoon-hyun]

nit: I think the fresh-per-call resolveFunctionCall() plus this memoised twin can collapse back into the single lazy val resolvedFunction that was here before, just un-privated and with the reducedWith.isEmpty guard. The sharing hazard the two comments describe does not arise: the only external caller, DistributionAndOrderingUtils.resolveTransformExpression, visits each node once inside expr.transform, and every TransformExpression in a write distribution/ordering is a distinct instance built by V2ExpressionUtils.toCatalyst (distribution and ordering are converted by separate calls). Even if one instance did land in two positions, QueryExecution.cloneWithFreshStatefulExpressions and ExpressionsEvaluator.prepareExpressions already fresh-copy stateful expressions per node and per evaluator. One member and no caveat comments would do the same job.

[peter-toth]

Agreed, and both mechanisms are there. QueryExecution.cloneWithFreshStatefulExpressions maps every node's expressions through freshCopyIfContainsStatefulExpression() before optimization, and ExpressionsEvaluator.prepareExpressions does the same per evaluator. So the hazard the split guarded against cannot arise.

Collapsed back into one lazy val resolvedFunction, un-private, with the reducedWith guard, and both caveat comments are gone.

[peter-toth]

Merge Summary:

• merged into master ed06d80
• merged into branch-4.x 92b20aa

Posted by merge_spark_pr.py

[peter-toth]

Thank you everyone for the review!