[SPARK-59887][SPARK-59901][SQL][4.2] Fix wrong results when a storage-partitioned join pairs a transform of a join-key expression with one of a column

[peter-toth]

Differences from #59165

This is the branch-4.2 backport of #59165. The description below is the original one. The cherry-pick does not apply, since the storage-partitioned join code on this branch predates master's rework of it. So the fix is ported onto this branch's own code:

• KeyedShuffleSpec.keyPositions maps a partition expression by its leaves here, not by its references. So a key with a literal, e.g. b + 1, already maps to no cluster key and is not paired wrongly on this branch. A single-leaf key such as -b or s.a is. The assertion in keyPositions fails here too for an expression with two leaves, b + 1 included (SPARK-59901). It is replaced by "only an expression with a single leaf maps to a position".
• This branch picks the layout to shuffle onto in EnsureRequirements through candidateSpecs, bestSpecOpt and bestMemberOpt, not pickCoPartitionTarget. With ShuffleSpecCollection.canCreatePartitioning an exists, a collection whose first member cannot serve becomes a candidate here. So three places look at its serving members only:
  • the ranking counts a collection by its first serving member, which is its first member whenever that one serves, as before;
  • a child counts as matched only through a serving member;
  • the layout is picked from the serving members.
• Two new tests in EnsureRequirementsSuite cover the ranking and the matching, since master's tests do not reach them on this branch: "SPARK-59887: a collection ranks by its first member that can serve as the layout" and "SPARK-59887: a child matched only through a member that cannot serve is shuffled". Each fails with its change removed.
• Matching only through members that can serve costs a shuffle in one more shape on this branch. Two legs that line up only through members of the same shape now take a shuffle, e.g. [a, bucket(4, -y)] against [m, bucket(4, -n)], joined on l.a = r.m AND l.y = r.n. With allowJoinKeysSubsetOfPartitionKeys, the shuffle conf on and requireAllClusterKeysForCoPartition off, branch-4.2 took no shuffle, and this change shuffles one side onto 2 partitions. The rows are right either way.
  • Master pairs such legs in checkKeyGroupCompatible, which compares every member of both sides. It takes no shuffle there, before and after this change.
  • This branch's createKeyedShuffleSpec offers only the first member that satisfies the distribution, here [a, bucket(4, c)]. So only the shuffle path paired the legs, through a member that cannot serve as the layout. This change keeps the shuffle path from that, which is also where master's third cost above comes from.
  • A widening cast instead of the negation behaves the same. So "Two sides that both carry the cast still pair" holds here only where the storage-partitioned join path pairs them.
• This branch's createKeyedShuffleSpec checks spark.sql.requireAllClusterKeysForCoPartition by matching the leaves of the layout to the join keys one by one, where master checks coverage. So the coverage change becomes: a layout with an expression with several leaves does not meet the requirement. With the rest of this change alone, the rows are right, but a join on (b, c, d) over [b + c, d] is paired on d alone.
• That needs a sibling layout that satisfies the join's distribution, so that the child is not shuffled first. The coverage test, "SPARK-59901: a join key over two columns does not cover its columns", has none on this branch, since each leg's projection drops the identity side's columns. So it also runs with those columns kept.
• KeyedPartitioning.createShuffleSpec gets master's refusal to project onto no position. It came with SPARK-59289, which is not on this branch, and master's SPARK-59901 fix relies on it.
  • A join keyed on the expression itself, e.g. l.b + l.c = r.k, can satisfy the distribution under spark.sql.requireAllClusterKeysForDistribution, while no position maps to a cluster key.
  • With allowJoinKeysSubsetOfPartitionKeys, the projection would then run the join in one partition. That needs spark.sql.sources.v2.bucketing.shuffle.enabled, or another side over the same expression with requireAllClusterKeysForCoPartition off. A -b key already did that on this branch.
  • A b + 1 or b + c key failed there with the AssertionError before this change, even with the shuffle conf off. Now all three get the shuffle, which can mean 2 shuffles where branch-4.2 had none.
  • The new test "SPARK-59901: a join keyed on a join-key expression does not run in one partition" covers -b, b + 1 and b + c.
• TransformExpressionSuite does not exist on this branch, so its hunk is dropped. So are the doc changes for code this branch does not have.
• Three tests are tailored, since master's version cannot pass here. On this branch a key with a literal maps to no cluster key, so it is shuffled rather than paired:
  • the first ShuffleSpecSuite test uses -a, -b and abs(b) for master's a + 1, b + 1 and b + 2;
  • "two layouts over the same shape of a join key pair as two columns do" and "the same shape reduced together pairs only through the same pairing" use -p.b over negative values for master's p.b + 1.
• Measured on branch-4.2 without the change, with the shuffle conf on: the -p.b key returns 4 of 8 rows through a one-side shuffle and through the bucket(8) reducer path. Through a broadcast join it returns 4 of 8 rows with every conf at its default. The struct field pairing returns 0 of 8 rows, and shuffling onto it fails with a ClassCastException.
• Two parts of the original description need a -p.b key on this branch, since a member with a literal does not satisfy the distribution here:
  • the plan that gets better: ident i JOIN plain p ON i.id = -p.b JOIN xy q ON i.id = q.x AND p.b = q.y takes 2 shuffles instead of 3. With p.b + 1 it already took 2;
  • the third cost: with -p.b branch-4.2 took no shuffle, and this change takes 1. With p.b + 1 one side was already shuffled.
• A single join keyed on p.b + p.c does not fail on this branch, AQE included, so "a single join is enough" does not hold here. So the first query of "SPARK-59901: a join key over two columns is not paired through one of them" passes without the change, and its second query fails with the AssertionError.
• These included tests pass on branch-4.2 without the change:
  • the p.b + 1 variants of the shuffled-onto, xy, broadcast and identity tests;
  • "two layouts over the same shape of a join key pair as two columns do";
  • "a member that cannot serve does not rank its collection", since its p.y + 1 member does not satisfy the distribution here;
  • the SPARK-59905 test, which now runs p.b + p.c with AQE on as on master;
  • the new ranking test, since a collection with a member that cannot serve is no candidate there.
• Each change was removed on its own:
  • without the shape gate, the -p.b shuffled-onto, broadcast and identity tests and the struct field test fail;
  • without the canCreatePartitioning clause, the -p.b shuffled-onto and xy tests, the struct field, reduce and same-pairing tests fail on the assertion in rebuiltOver. The new matching test and the first ShuffleSpecSuite test fail too;
  • without the member filter, the -p.b xy test and both new EnsureRequirementsSuite tests fail;
  • with the ranking by the first member back, the new ranking test fails;
  • with the matching through any member back, the new matching test fails;
  • without the GroupPartitionsExec change, the reduce test and the two-column GROUP BY fail;
  • with the keyPositions assertion back, the ShuffleSpecSuite SPARK-59901 test, "SPARK-59901: a join key over two columns is not paired through one of them" and "SPARK-59901: a join keyed on a join-key expression does not run in one partition" fail. The coverage test passes, since the coverage change keeps [b + c, d] away from keyPositions;
  • with the forall back in ShuffleSpecCollection.canCreatePartitioning, the -p.b xy test fails on the shuffle count, and so do the new matching test and the first ShuffleSpecSuite test;
  • removing the coverage change or the refusal to project onto no position fails only its new test or run.
• Ran ShuffleSpecSuite, DistributionSuite, the KeyGroupedPartitioning* suites, EnsureRequirementsSuite, GroupPartitionsExecSuite, ProjectedOrderingAndPartitioningSuite, PlannerSuite, WriteDistributionAndOrderingSuite, BucketedReadWithoutHiveSupportSuite, the plan stability suites, the *JoinSuite suites and AdaptiveQueryExecSuite on branch-4.2, 1403 tests in all, plus dev/lint-scala.

What changes were proposed in this pull request?

A storage-partitioned join now pairs a partition transform whose argument is not a bare column, i.e. is an expression or a struct field, only with a transform over the same argument. It also stops dropping such an argument where it used to:

• KeyedShuffleSpec.isExpressionCompatible pairs a transform with an argument that is not an Attribute only with a transform over the same argument shape, i.e. the argument with its one column replaced by a placeholder. So bucket(4, b + 1) pairs with bucket(4, c + 1) on b = c, but not with bucket(4, x). Over the same shape the two compare as two transforms of columns do. Under allowCompatibleTransforms, bucket(8, b + 1) reduces onto bucket(4, c + 1), since a reduce maps the partition keys and evaluates no argument. When an earlier join reduced the keys, two such transforms pair only when they were reduced through the same pairing. An identity side is never reduced onto such a transform. TransformExpression keeps its comparisons. It loses withReference, whose last callers now rebuild the transform with copy.
• KeyedShuffleSpec.canCreatePartitioning asks the same question, so such a layout is never the one other children are shuffled onto. createPartitioning replaces a transform's argument with the other child's cluster key, which drops whatever surrounds the key. It and the identity arm of reducersBothWays now rebuild a transform through one helper that asserts the argument is a column.
• ShuffleSpecCollection.canCreatePartitioning is now true when any member can, and EnsureRequirements.pickCoPartitionTarget offers and ranks only the members that can. So one member that cannot serve no longer rules out a usable sibling. That also covers a bare expression such as b + 1, which master already refused.
• GroupPartitionsExec rebuilds a reduced expression over each member's own argument when it reports the members of a collection. It used to re-target the column only.
• KeyedShuffleSpec.keyPositions maps an expression with more than one reference, e.g. bucket(4, b + c), to no position instead of failing an assertion (SPARK-59901).
• EnsureRequirements.createKeyedShuffleSpecs counts only an expression over one column towards spark.sql.requireAllClusterKeysForCoPartition. An expression over two columns maps to no position, so a projection would drop it together with its columns.

Why are the changes needed?

With spark.sql.sources.v2.bucketing.shuffle.enabled on, this query returns 0 rows instead of 7. t1(id) and t3(x) are partitioned by bucket(4, ...), plain(b) is not partitioned, and each holds the values 0 to 7:

SELECT t1.id, p.b, t3.x FROM t1
JOIN plain p ON t1.id = p.b + 1
JOIN t3 ON p.b = t3.x

1. The first join shuffles plain onto t1's partitioning. KeyedShuffleSpec.createPartitioning builds that partitioning from plain's join key, so it is bucket(4, b + 1). This join is correct.
2. The second join clusters plain on the bare b. KeyedShuffleSpec.keyPositions maps a partition expression to a cluster key through its reference, so bucket(4, b + 1) counts as a function of b, the counterpart of x.
3. isSameFunction compares only the function name and the bucket count. So bucket(4, b + 1) is the same as t3's bucket(4, x), and the join pairs the partitions as they stand. A row with b = x sits in bucket (b + 1) % 4 on one side and in x % 4 on the other.

The comparison ignores which column an argument is, and relies on keyPositions to pair the columns up. That is sound only when the argument is the column itself. bucket(4, b + 1) is a function of b, but not the same function of it as bucket(4, x) is of x. A scan never reports a transform of an expression or of a struct field, but two planner paths build one. A one-side shuffle does, under spark.sql.sources.v2.bucketing.shuffle.enabled. An inner broadcast hash join does too, with every conf at its default. BroadcastHashJoinExec.expandOutputPartitioning also reports the streamed side's layout over the build side's join key. The same problem shows up in more shapes, all measured on master:

• A broadcast join. With every conf at its default, SELECT /*+ BROADCAST(p) */ ... FROM t1 JOIN plain p ON t1.id = p.b + 1 JOIN t3 ON p.b = t3.x returns 0 of 7 rows the same way, through the expanded bucket(4, b + 1) layout.
• The reducer path. Under spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled, a bucket(8, x) third table loses the rows the same way.
• A struct field. A join key p.s.a gives bucket(4, s.a), which keyPositions pairs through s. Joining two such sides on s pairs bucket(4, s.a) with bucket(4, s.b) and returns 0 of 8 rows. Shuffling another side onto such a layout builds bucket(4, s), which fails on executors with a ClassCastException.
• Shuffling onto the layout. In plain p JOIN t1 ON p.b + 1 = t1.id JOIN xy q ON t1.id = q.x AND p.b = q.y, the second join shuffles xy onto the first join's bucket(4, b + 1) layout. That builds bucket(4, y) for xy, drops the + 1, and returns 0 of 8 rows.
• After a reduce. When a later join reduces the first join's output from bucket(4) onto bucket(2), GroupPartitionsExec reports the bucket(4, b + 1) member as bucket(2, b). A join on b then pairs it with a bucket(2, x) scan, or shuffles another side onto it, and returns 0 of 7 rows. The same happens to the bare b + 1 member of a side shuffled onto an identity-partitioned table. The struct form fails with the ClassCastException.
• An identity side reduced onto it. Under allowCompatibleTransforms, a later join can reduce an identity-partitioned side onto bucket(4, b + 1). That evaluates bucket(4, id + 1) on the side's partition keys while planning. The keys come out right, but under ANSI a key of Long.MaxValue fails the query with ARITHMETIC_OVERFLOW, although the query never computes id + 1.
• A GROUP BY after a reduce. With allowCompatibleTransforms, SELECT /*+ BROADCAST(p) */ p.b, max(p.c) FROM ident i JOIN plain2 p ON i.id = p.b + p.c JOIN bucket2 t2 ON i.id = t2.id GROUP BY p.b returns 4 groups instead of 2. The reduce reports the p.b + p.c member as bucket(2, p.b), which serves the GROUP BY as it stands.
• A key over two columns. A join key such as p.b + p.c gives bucket(4, b + c), or the bare b + c over an identity-partitioned side. Any spec over it fails planning with an AssertionError in keyPositions (SPARK-59901). AQE validates the plan in OptimizeSkewedJoin, so a single join is enough.

An ORDER BY over such a partition expression had a related problem in another code path. That was SPARK-59905, fixed in #59189.

The one-side shuffle of transform expressions came with SPARK-48012, in 4.0.0. So did the broadcast expansion of a keyed layout, since SPARK-49205 made KeyGroupedPartitioning an expression.

Does this PR introduce any user-facing change?

Yes, it fixes the wrong results, the crash and the planning failures above. Such a query now shuffles the side instead of pairing it or shuffling onto it.

That costs shuffles in three shapes whose plan was already correct:

• A widening cast counts as part of the argument's shape. When t1.id is BIGINT and p.b is INT, the first join reports bucket(4, cast(b as bigint)). A later join p.b = t3.x with an INT bucketed t3 then shuffles both sides. Before, a connector whose bucket function has one canonical name for both types paired the two sides as they stood. This happens with every conf at its default, through a broadcast join. Two sides that both carry the cast still pair.
• Under allowCompatibleTransforms, an identity-partitioned side is no longer reduced onto such a transform, so it is shuffled.
• With spark.sql.sources.v2.bucketing.shuffle.enabled, a join on two keys whose sides pair only through a transform of an expression gets one more shuffle. Take two broadcast joins, t1 JOIN /*+ BROADCAST(p) */ p ON t1.id = p.b + 1 and the same over t2 and q, joined on l.id = r.c AND l.b = r.b. Each layout covers only one join key, so the storage-partitioned join declines under requireAllClusterKeysForCoPartition. Master then keeps both sides, paired on bucket(4, b + 1). This PR cannot shuffle onto that layout, so it shuffles one side onto bucket(4, id). Both plans co-partition on one key only, which requireAllClusterKeysForCoPartition is meant to prevent. That is a separate, pre-existing question, filed as SPARK-59971.

One plan gets better. A collection with a bare expression member, e.g. the b + 1 of a side shuffled onto an identity-partitioned table, used to be refused as a layout as a whole. Now its other members can serve, so ident i JOIN plain p ON i.id = p.b + 1 JOIN xy q ON i.id = q.x AND p.b = q.y takes 2 shuffles instead of 3.

This PR keeps the fix small, since it has to reach the maintenance branches. SPARK-59900 is the follow-up for the rest: shuffling another side onto such a layout, and reducing an identity side onto it.

How was this patch tested?

New tests. All but two of them fail on master:

• ShuffleSpecSuite, "SPARK-59887: a spec over a transform of an expression pairs only with the same shape". For an expression and a struct field argument, it covers the same shape, another shape and a column, both ways and with itself, the reduce in both directions, the two sides of one reduce, two sides reduced through different pairings, an identity side, canCreatePartitioning, and a collection with a usable sibling.
• ShuffleSpecSuite, "SPARK-59901: an expression over two columns maps to no cluster key".
• KeyGroupedPartitioningSuite, "SPARK-59887: a side shuffled onto a transform of an expression is not paired as it is". It runs the query above against a bucket(4) third table, and with allowCompatibleTransforms against a bucket(8) one that the join reduces, and asserts two shuffles.
• KeyGroupedPartitioningSuite, "SPARK-59887: a side is not shuffled onto a transform of an expression". It runs both join orders and asserts that xy is shuffled onto bucket(4, x) and never onto a bucket of y.
• KeyGroupedPartitioningSuite, "SPARK-59887: a broadcast join's layout over a join key is not paired as it is". It runs with every conf at its default and checks that the first join is a broadcast join.
• KeyGroupedPartitioningSuite, "SPARK-59887: a layout over a transform of a struct field is not paired nor shuffled onto". Master returns 0 of 8 rows.
• KeyGroupedPartitioningSuite, "SPARK-59887: a reduce keeps what each member's keys are computed from". It covers a bucketed and an identity-partitioned first table, a bucketed and an unpartitioned third table, and the struct form. Master returns 0 of 7 rows.
• KeyGroupedPartitioningSuite, "SPARK-59887: an identity side is not reduced onto a transform of an expression". It runs with p.b + 1 over a Long.MaxValue key and with -p.b over a Long.MinValue key. Master fails with the ARITHMETIC_OVERFLOW for both.
• KeyGroupedPartitioningSuite, "SPARK-59887: a reduce keeps a member over two columns". It runs the GROUP BY above. Master returns 4 groups instead of 2.
• KeyGroupedPartitioningSuite, "SPARK-59901: a join key over two columns is not paired through one of them". Master fails with the AssertionError.
• KeyGroupedPartitioningSuite, "SPARK-59901: a join key over two columns does not cover its columns". It joins two broadcast joins that report [b + c, d] on (b, c, d) under allowKeysSubsetOfPartitionKeys, and asserts that the join is not paired on d alone. Master fails with the AssertionError.
• KeyGroupedPartitioningSuite, "SPARK-59887: two layouts over the same shape of a join key pair as two columns do". It runs the cast case above with every conf at its default. It also runs two one-side shuffles onto bucket(4, b + 1) and onto bucket(4) or bucket(8) of q.b + 1, where allowCompatibleTransforms reduces the bucket(8) one. It asserts that the join adds no shuffle. It passes on master, which pairs and reduces these as well.
• KeyGroupedPartitioningSuite, "SPARK-59887: the same shape reduced together pairs only through the same pairing". Two legs each shuffle onto bucket(12, b + 1). The first leg reduces with bucket(8), so its keys are buckets of 4. The second leg reduces with bucket(8) too, with bucket(18), or not at all, so its keys are buckets of 4, 6 or 12. It asserts 2 shuffles for the first and 4 for the other two. Master returns 0 of 23 rows for the third.
• KeyGroupedPartitioningSuite, "SPARK-59887: a member that cannot serve does not rank its collection". It pins the ranking: q keeps its layout with 3 partitions, as on master, where the collection is refused as a whole. Without the member filter, the collection would rank first with 5 partitions, and q would be shuffled onto 2.

The SPARK-59905 test from #59189 now also runs its p.b + p.c query with AQE on. Master fails that run with the AssertionError.

The three expression tests each run with the key p.b + 1 and with -p.b. Master returns 0 of 8 rows for p.b + 1 and 4 of 8 for -p.b. The -p.b key is the shape that also fails on 4.2, 4.1 and 4.0. On those branches b + 1 already works, since its literal leaf blocks the pairing.

Each change was removed on its own:

• without the guard in isExpressionCompatible, the one-side shuffle and broadcast pairing tests, the struct field test and the reduce test return wrong rows, and the identity tests fail on the assertion in rebuiltOver;
• without the shape pairing, the same-shape and same-pairing tests fail on the shuffle count;
• with a reduce of such a pair refused, the same-shape test fails on the shuffle count;
• with every pair of reduced keys over an expression refused, the same-pairing test fails on the shuffle count;
• with the reduced marks ignored, or only checked on both sides, the same-pairing test returns 7 of 23 rows;
• without the canCreatePartitioning clause, the one-side shuffle pairing, xy, struct field, reduce and same-pairing tests fail, all on the assertion in rebuiltOver;
• with the forall back in ShuffleSpecCollection.canCreatePartitioning, the xy tests fail on the shuffle count;
• without the member filter in pickCoPartitionTarget, 15 tests fail in KeyGroupedPartitioningSuite and EnsureRequirementsSuite. The filter now also does what the collection-level check before it did, so 8 of them are older tests. Of this PR's tests, 7 fail, 6 of them on the assertion in rebuiltOver. A member over the same shape pairs with itself, so only the filter keeps it from being the layout;
• without the GroupPartitionsExec change, the reduce test returns 0 of 7 rows, and the two-column GROUP BY returns 4 groups;
• without the requireAllClusterKeysForCoPartition change, the coverage test pairs on d alone;
• with the keyPositions assertion back, both SPARK-59901 tests in KeyGroupedPartitioningSuite fail.

The first ShuffleSpecSuite test also fails for each change to isExpressionCompatible above.

The xy tests pass with the guard removed, since the canCreatePartitioning clause and the member filter keep xy off that layout.

Also ran TransformExpressionSuite, ShuffleSpecSuite, DistributionSuite, the KeyGroupedPartitioning* suites, WriteDistributionAndOrderingSuite, PlannerSuite, ProjectedOrderingAndPartitioningSuite, GroupPartitionsExecSuite, EnsureRequirementsSuite, BucketedReadWithoutHiveSupportSuite, the plan stability suites, the *JoinSuite suites and AdaptiveQueryExecSuite, 1955 tests in all, plus dev/lint-scala.

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

Generated-by: Claude Code (Claude Opus 5.5)

[uros-b]

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

[peter-toth]

Merge Summary:

• merged into branch-4.2 d0af88a

Posted by merge_spark_pr.py