You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by pt...@apache.org on 2026/10/06 12:42:45 UTC
(spark) branch branch-4.2 updated: [SPARK-59887][SPARK-59901][4.2][SQL] Fix wrong results when a storage-partitioned join pairs a transform of a join-key expression with one of a column
This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new d0af88ad0b1c [SPARK-59887][SPARK-59901][4.2][SQL] Fix wrong results when a storage-partitioned join pairs a transform of a join-key expression with one of a column
d0af88ad0b1c is described below
commit d0af88ad0b1c4eadb7619dabff66f4bdfb404a94
Author: Peter Toth <pe...@gmail.com>
AuthorDate: Tue Oct 6 14:42:28 2026 +0200
[SPARK-59887][SPARK-59901][4.2][SQL] Fix wrong results when a storage-partitioned join pairs a transform of a join-key expression with one of a column
### 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 ma [...]
- `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:
```sql
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 [...]
- **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)`. [...]
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)
Closes #59238 from peter-toth/SPARK-59887-spj-expression-key-identity-4.2.
Authored-by: Peter Toth <pe...@gmail.com>
Signed-off-by: Peter Toth <pe...@gmail.com>
---
.../catalyst/expressions/TransformExpression.scala | 10 -
.../sql/catalyst/plans/physical/partitioning.scala | 149 +++++--
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 108 ++++-
.../datasources/v2/GroupPartitionsExec.scala | 14 +-
.../execution/exchange/EnsureRequirements.scala | 28 +-
.../connector/KeyGroupedPartitioningSuite.scala | 450 ++++++++++++++++++++-
.../exchange/EnsureRequirementsSuite.scala | 66 +++
7 files changed, 766 insertions(+), 59 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala
index 191de9bf5062..bfa94914f4ae 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala
@@ -123,16 +123,6 @@ case class TransformExpression(
}
}
- /**
- * Re-targets this partition transform expression at `attr`. A partition transform expression
- * has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), so this replaces that
- * attribute and keeps the rest of the expression: any field path above the leaf comes from this
- * expression, not from the re-targeted key. Re-targeting is therefore only faithful when the
- * source and the target key expressions have the same path shape.
- */
- def withReference(attr: Attribute): TransformExpression =
- transform { case _: AttributeReference => attr }.asInstanceOf[TransformExpression]
-
// Return a Reducer for a reducible function on another reducible function
private def reducer(
thisFunction: ReducibleFunction[_, _],
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index 05824238f83f..76e930395ea2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -23,6 +23,7 @@ import scala.collection.mutable
import org.apache.spark.{SparkException, SparkUnsupportedOperationException}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
import org.apache.spark.sql.connector.catalog.functions.Reducer
import org.apache.spark.sql.internal.SQLConf
@@ -625,6 +626,16 @@ case class KeyedPartitioning(
if (mayContainUnknownPartitionKeys && joinKeyPositions.length < expressions.length) {
return result
}
+ // No position covers a cluster key, so there is nothing to project onto. Projecting onto
+ // none would coalesce every partition into one, which is never the answer, and it is the
+ // same refusal `EnsureRequirements.resolveKeyedPartitioning` makes for its own caller.
+ // `keyPositions` maps an expression's single *leaf* onto the clustering, so this is
+ // reachable for a clustering that names the partition expression itself, where
+ // `satisfies0` says yes and this finds nothing. That covers an expression with several
+ // leaves too, e.g. `b + 1`. Return the unprojected spec, as above.
+ if (joinKeyPositions.isEmpty && expressions.nonEmpty) {
+ return result
+ }
// If allowing join keys to be subset of clustering keys, we should create a new
// `KeyedPartitioning` here that is grouped on the join keys instead, and use that as
// the returned shuffle spec. The construction carries the unknown-keys marker across:
@@ -946,12 +957,22 @@ trait ShuffleSpec {
* co-partitioned with the `other`, and therefore no shuffle is required when joining the two
* sides.
*
- * Note that Spark assumes this to be reflexive, symmetric and transitive.
+ * Note that Spark assumes this to be reflexive, symmetric and transitive. Some specs are not
+ * reflexive:
+ * - a [[RangeShuffleSpec]], whose boundaries are sampled;
+ * - a [[KeyedShuffleSpec]] with a partition expression that `keyPositions` maps to no cluster
+ * key, e.g. `years(ts)` clustered on `years(ts)`, `bucket(4, b + c)` or `b + 1`;
+ * - a [[KeyedShuffleSpec]] with a partition expression that is neither a column nor a transform,
+ * e.g. `-b` or `s.a`.
+ * None of them answers [[canCreatePartitioning]], which keeps the planner from picking one as
+ * the layout to shuffle the other children onto.
*/
def isCompatibleWith(other: ShuffleSpec): Boolean
/**
- * Whether this shuffle spec can be used to create partitionings for the other children.
+ * Whether this shuffle spec can be used to create partitionings for the other children. A
+ * [[ShuffleSpecCollection]] can when any of its members can, and the planner builds the
+ * partitioning from a member that can.
*/
def canCreatePartitioning: Boolean
@@ -1099,10 +1120,10 @@ case class CoalescedHashShuffleSpec(
* keys, and this is this side's own expression marked with the pairing
* that reduced it (`TransformExpression.reducedWith`). Stored as-is from
* the expression pair that produced the reducer; the only structural
- * consumer, `GroupPartitionsExec`, re-targets it at each reported
- * `KeyedPartitioning`'s own key attribute in `outputPartitioning` and
+ * consumer, `GroupPartitionsExec`, rebuilds it over each reported
+ * `KeyedPartitioning`'s own argument in `outputPartitioning` and
* normalizes it positionally in `doCanonicalize`, so the attribute it
- * carries here is not load-bearing.
+ * carries here does not matter.
*/
case class KeyReducer(reducer: Reducer[_, _], reducedExpression: TransformExpression)
@@ -1115,11 +1136,13 @@ case class KeyReducer(reducer: Reducer[_, _], reducedExpression: TransformExpres
* It is a case class so that structurally identical reducers compare by value, and plans holding
* them stay canonicalization-equal.
*
- * @param transform the partition transform, re-targeted at the identity side's key attribute
+ * @param transform the partition transform, rebuilt over the identity side's key attribute
*/
case class IdentityReducer(transform: TransformExpression) extends Reducer[Any, Any] {
- // `transform` has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), which is
- // bound to ordinal 0 of the single-value row `reduce` evaluates it against.
+ // `transform` is over a single bare column, which is bound to ordinal 0 of the single-value row
+ // `reduce` evaluates it against. `KeyedShuffleSpec.isExpressionCompatible` admits only bare
+ // column arguments here, and `keyPositions` maps an expression to a position only when it has
+ // one leaf.
@transient private lazy val bound: Expression =
BindReferences.bindReference(transform, AttributeSeq(transform.references.toSeq))
@@ -1148,8 +1171,18 @@ case class KeyedShuffleSpec(
* keys. For instance, if cluster keys are [a, b, b] and partition expressions are
* [bucket(4, a), years(b)], the result will be [(0), (1, 2)].
*
- * Note that we only allow each partition expression to contain a single partition key.
- * Therefore the mapping here is very similar to that from `HashShuffleSpec`.
+ * Only an expression with a single leaf maps to a position. A scan reports no other expression
+ * (`KeyedPartitioning.supportsExpressions`), but two planner paths build one over the other
+ * side's join key, e.g. `bucket(4, b + c)` from `b + c`. A one-side shuffle does, and so does a
+ * broadcast hash join, which reports the streamed side's layout over the build side's join key.
+ * No single cluster key stands for such an expression, so it maps to nothing, and
+ * `canCreatePartitioning` and `areKeysCompatible` turn the spec away. Under
+ * `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled`,
+ * `KeyedPartitioning.createShuffleSpec` projects the position away from an unmarked layout
+ * where another position maps.
+ *
+ * This says which cluster key an expression is a function of, not which function it is.
+ * `bucket(4, -b)` maps to `b` just as `bucket(4, b)` does, see `argumentsAreColumns`.
*/
lazy val keyPositions: Seq[mutable.BitSet] = {
val distKeyToPos = mutable.Map.empty[Expression, mutable.BitSet]
@@ -1158,8 +1191,11 @@ case class KeyedShuffleSpec(
}
partitioning.expressions.map { e =>
val leaves = e.collectLeaves()
- assert(leaves.size == 1, s"Expected exactly one child from $e, but found ${leaves.size}")
- distKeyToPos.getOrElse(leaves.head.canonicalized, mutable.BitSet.empty)
+ if (leaves.size == 1) {
+ distKeyToPos.getOrElse(leaves.head.canonicalized, mutable.BitSet.empty)
+ } else {
+ mutable.BitSet.empty
+ }
}
}
@@ -1174,9 +1210,10 @@ case class KeyedShuffleSpec(
// 3.2 for each pair of partition expressions at the same index, the corresponding
// partition keys must share overlapping positions in their respective clustering keys.
// 3.3 each pair of partition expressions at the same index must describe one key space: two
- // bare references, the same transform function, or the two sides of one reduce. A pair
- // the join would first reduce onto one key space does not count, see
- // `areKeysCompatible`'s `allowReduce`.
+ // bare references, the same transform function, or two transforms reduced through the
+ // same pairing. A transform of an expression counts only against one over the same
+ // argument shape. A pair the join would first reduce onto one key space does not count,
+ // see `areKeysCompatible`'s `allowReduce`.
// 4. the partition values from both sides are following the same order.
case otherSpec @ KeyedShuffleSpec(otherPartitioning, otherDistribution, _) =>
distribution.clustering.length == otherDistribution.clustering.length &&
@@ -1266,7 +1303,16 @@ case class KeyedShuffleSpec(
left: Expression,
right: Expression,
allowReduce: Boolean): Boolean = {
- if (TransformExpression.hasReducedKeys(left) || TransformExpression.hasReducedKeys(right)) {
+ val overExpression = Seq(left, right).exists {
+ case t: TransformExpression => !argumentsAreColumns(t)
+ case _ => false
+ }
+ if (overExpression && !argumentShape(left).exists(argumentShape(right).contains)) {
+ // A transform of an expression pairs only with a transform over the same argument shape.
+ // See `argumentsAreColumns`.
+ false
+ } else if (
+ TransformExpression.hasReducedKeys(left) || TransformExpression.hasReducedKeys(right)) {
// Reduced keys are in a key space that neither transform names, so comparing the transforms
// says nothing about whether the two sides are laid out the same way. The pair that was
// reduced together is laid out the same way, since its two sides came out of one reduce onto
@@ -1291,6 +1337,47 @@ case class KeyedShuffleSpec(
}
}
+ /**
+ * Whether every argument of `transform` is a bare column. `isExpressionCompatible` asks it of
+ * both sides before it compares anything, `canCreatePartitioning` asks it of a template, and
+ * `rebuiltOver` asserts it.
+ *
+ * `keyPositions` pairs two sides up by the column each expression is over. That pairing
+ * describes one key space when the argument is the column itself. `bucket(4, -b)` is a function
+ * of `b`, but not the same function of it as `bucket(4, x)` is of `x`. Nor is `bucket(4, s.a)`
+ * the same function of `s` as `bucket(4, t.b)` is of `t`. So such a transform pairs only with a
+ * transform over the same `argumentShape`. Over the same shape, the two compare as two
+ * transforms of columns do, e.g. `bucket(4, -b)` with `bucket(4, -c)` on a join `b = c`. A
+ * reduce of `bucket(8, -b)` onto `bucket(4, -c)` is fine too, since it maps the partition keys
+ * and evaluates no argument. An identity side is not reduced onto such a transform, since that
+ * would evaluate `-id` on keys the query never computes it for, which can overflow.
+ */
+ private def argumentsAreColumns(transform: TransformExpression): Boolean =
+ transform.children.forall(_.isInstanceOf[Attribute])
+
+ /**
+ * `transform`, whose arguments are bare columns, with each replaced by `key`, which may be any
+ * expression. Any other argument would lose what surrounds its column, e.g. the `-` of
+ * `bucket(4, -b)`.
+ */
+ private def rebuiltOver(transform: TransformExpression, key: Expression)
+ : TransformExpression = {
+ assert(argumentsAreColumns(transform), s"Cannot rebuild $transform over $key")
+ transform.copy(children = transform.children.map(_ => key))
+ }
+
+ /**
+ * The arguments of a transform with their one column replaced by a placeholder of its type, or
+ * `None` when they reference other than one column or are not deterministic, or `expression` is
+ * not a transform. Two transforms with the same function and shape are the same function of
+ * their columns, so a join that pairs the columns up pairs the keys up too.
+ */
+ private def argumentShape(expression: Expression): Option[Seq[Expression]] = expression match {
+ case t: TransformExpression if t.references.size == 1 && t.deterministic =>
+ Some(t.children.map(QueryPlan.normalizeExpressions(_, t.references.toSeq)))
+ case _ => None
+ }
+
/**
* Whether a join may reduce one or both sides' partition keys onto a common key space, which is
* what lets two different transforms be compatible in the first place.
@@ -1340,8 +1427,8 @@ case class KeyedShuffleSpec(
if (otherReducer.isEmpty) {
// Only this side reduces. The reducer contract is r(f1(x)) = f2(x) where "=" matches
// both value and data type, so the reduced keys equal the target transform applied to
- // this side's child. Report the target transform (re-targeted at the reporting
- // partitioning's key attribute by `GroupPartitionsExec`) instead of the un-reduced
+ // this side's child. Report the target transform (rebuilt over the reporting
+ // partitioning's own argument by `GroupPartitionsExec`) instead of the un-reduced
// `e1`, whose type can differ from the reduced keys. A connector that violates the
// contract with a reducer of a different result type fails the reduced-types check in
// `EnsureRequirements`, since the other side's keys are typed by the target transform.
@@ -1368,15 +1455,16 @@ case class KeyedShuffleSpec(
(thisResult, otherResult)
// Identity transform on this side, arbitrary transform on the other side: create a reducer
- // that applies the other's transform to the raw identity values. Each partition expression
- // is guaranteed to have exactly one leaf child (asserted in keyPositions), which
- // `IdentityReducer` binds to ordinal 0.
+ // that applies the other's transform to the raw identity values. `areKeysCompatible` admitted
+ // this pair, so `t`'s argument is one bare column. `isExpressionCompatible` refuses any other
+ // argument here, and `keyPositions` maps an expression with more leaves to no position.
+ // `IdentityReducer` binds `a` to ordinal 0.
case (a: AttributeReference, t: TransformExpression) =>
- (Some(KeyReducer(IdentityReducer(t.withReference(a)), t)), None)
+ (Some(KeyReducer(IdentityReducer(rebuiltOver(t, a)), t)), None)
// Symmetric: identity transform on the other side.
case (t: TransformExpression, a: AttributeReference) =>
- (None, Some(KeyReducer(IdentityReducer(t.withReference(a)), t)))
+ (None, Some(KeyReducer(IdentityReducer(rebuiltOver(t, a)), t)))
case (_, _) => (None, None)
}
@@ -1396,8 +1484,14 @@ case class KeyedShuffleSpec(
// cannot rewrite it. This also keeps the unprojected spec returned by `createShuffleSpec`
// for a marked narrowing projection from being chosen as the best spec.
keyPositions.forall(_.nonEmpty) &&
- partitioning.expressions.forall { e =>
- e.isInstanceOf[AttributeReference] || e.isInstanceOf[TransformExpression]
+ partitioning.expressions.forall {
+ case _: AttributeReference => true
+ // `createPartitioning` replaces a transform's argument with the other child's cluster key,
+ // which drops whatever surrounds the key, e.g. the `-` of `bucket(4, -b)` or the `.a` of
+ // `bucket(4, s.a)`. Such a spec's own child would then be laid out differently from the
+ // children shuffled onto it.
+ case t: TransformExpression => argumentsAreColumns(t)
+ case _ => false
} &&
// Shuffling another child onto these keys evaluates the partition expressions per row to
// decide where each row goes, and reduced keys are not what those expressions compute.
@@ -1409,7 +1503,8 @@ case class KeyedShuffleSpec(
val newExpressions = partitioning.expressions.zip(keyPositions).map {
case (te: TransformExpression, positionSet) =>
- te.copy(children = te.children.map(_ => clustering(positionSet.head)))
+ // `canCreatePartitioning` admits only a template over bare columns, see there.
+ rebuiltOver(te, clustering(positionSet.head))
case (_, positionSet) => clustering(positionSet.head)
}
// The child re-shuffled onto this layout may hold keys outside the declared set, so every
@@ -1474,7 +1569,7 @@ case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec {
}
override def canCreatePartitioning: Boolean =
- specs.forall(_.canCreatePartitioning)
+ specs.exists(_.canCreatePartitioning)
override def createPartitioning(clustering: Seq[Expression]): Partitioning = {
// as we only consider # of partitions as the cost now, it doesn't matter which one we choose
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
index 17621e67d1c4..ca03d6b1cb37 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst
import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException}
import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression}
+import org.apache.spark.sql.catalyst.expressions.{Abs, Attribute, AttributeReference, DirectShufflePartitionID, Expression, GetStructField, TransformExpression}
import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.connector.catalog.functions.{FlipLowBitFunction, Reducer, ReducibleFunction, ScalarFunction}
@@ -646,6 +646,112 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
}
}
+ test("SPARK-59887: a spec over a transform of an expression pairs only with the same shape") {
+ // A join pairs its two sides up by the column each partition expression is over. Over `a = b`
+ // it would pair `bucket(4, -a)` with `bucket(4, b)`, and over `s = t` it would pair
+ // `bucket(4, s.x)` with `bucket(4, t.y)`. A row with `a = b` sits in a different bucket on each
+ // side. `bucket(4, -a)` and `bucket(4, -b)` are the same function of `a` and of `b`, so they
+ // pair up. The expression is `-a`, since `keyPositions` maps only an expression with a single
+ // leaf, and `a + 1` has two.
+ val fn = new FakeBucket
+ val a = $"a".long
+ val b = $"b".long
+ val structType = new StructType().add("x", LongType).add("y", LongType)
+ val s = $"s".struct(structType)
+ val t = $"t".struct(structType)
+ // A spec clustered on the column its one partition expression references.
+ def keyed(expression: Expression): KeyedShuffleSpec =
+ KeyedShuffleSpec(
+ KeyedPartitioning(Seq(expression), Seq(InternalRow(0L), InternalRow(1L))),
+ ClusteredDistribution(expression.references.toSeq))
+ def bucket(numBuckets: Int, argument: Expression): TransformExpression =
+ TransformExpression(fn, Seq(argument), Some(numBuckets))
+ def spec(numBuckets: Int, argument: Expression): KeyedShuffleSpec =
+ keyed(bucket(numBuckets, argument))
+ // Whether the two sides of one reduce, `bucket(8, left)` with `bucket(4, right)`, pair up.
+ def pairedByReduce(left: Expression, right: Expression): Boolean =
+ keyed(bucket(8, left).reducedTogetherWith(bucket(4, right)))
+ .isCompatibleWith(keyed(bucket(4, right).reducedTogetherWith(bucket(8, left))))
+ val column = spec(4, b)
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ assert(spec(4, a).isCompatibleWith(column), "two columns pair up")
+ assert(spec(4, a).canCreatePartitioning)
+ assert(spec(8, a).areKeysCompatible(column, allowReduce = true),
+ "the fixture reduces two columns")
+ assert(pairedByReduce(a, b), "two columns reduced together share one key space")
+ assert(keyed(b).areKeysCompatible(spec(4, a), allowReduce = true),
+ "an identity side reduces onto a column's transform")
+
+ Seq(
+ ("an expression", -a, -b, Abs(b)),
+ ("a struct field", GetStructField(s, 0), GetStructField(t, 0), GetStructField(t, 1))
+ ).foreach { case (shape, argument, sameShape, otherShape) =>
+ val over = spec(4, argument)
+ assert(over.isCompatibleWith(spec(4, sameShape)), s"over $shape, the same shape")
+ assert(over.isCompatibleWith(over), s"over $shape, with itself")
+ assert(!over.isCompatibleWith(spec(4, otherShape)), s"over $shape, another shape")
+ assert(!over.isCompatibleWith(column), s"over $shape, against a column")
+ assert(!column.isCompatibleWith(over), s"over $shape, the other way round")
+ assert(!over.canCreatePartitioning, s"over $shape, no layout to shuffle onto")
+
+ // Over the same shape, such a pair reduces as two columns do, whichever side reduces. It
+ // does not reduce onto another shape or a column, and an identity side does not reduce
+ // onto it.
+ assert(spec(8, argument).areKeysCompatible(spec(4, sameShape), allowReduce = true),
+ s"over $shape, reducing the same shape")
+ assert(spec(4, sameShape).areKeysCompatible(spec(8, argument), allowReduce = true),
+ s"over $shape, reducing the same shape the other way round")
+ assert(!spec(8, argument).areKeysCompatible(spec(4, otherShape), allowReduce = true),
+ s"over $shape, reducing another shape")
+ assert(!spec(8, argument).areKeysCompatible(column, allowReduce = true), s"over $shape")
+ assert(!column.areKeysCompatible(spec(8, argument), allowReduce = true), s"over $shape")
+ assert(!over.areKeysCompatible(spec(8, b), allowReduce = true), s"over $shape")
+ assert(!keyed(b).areKeysCompatible(over, allowReduce = true),
+ s"over $shape, reducing an identity side")
+ // Keys an earlier join reduced are in the space of that reduce. Two transforms reduced
+ // through the same pairing pair over the same shape, as two columns do. `bucket(12)`
+ // reduced with `bucket(8)` holds buckets of 4, with `bucket(18)` buckets of 6, and
+ // unreduced buckets of 12.
+ assert(pairedByReduce(argument, sameShape), s"over $shape, reduced together")
+ assert(!pairedByReduce(argument, otherShape),
+ s"over $shape, another shape reduced together")
+ val reducedWith8 = keyed(bucket(12, argument).reducedTogetherWith(bucket(8, a)))
+ assert(!reducedWith8.isCompatibleWith(
+ keyed(bucket(12, sameShape).reducedTogetherWith(bucket(18, b)))),
+ s"over $shape, reduced through different pairings")
+ assert(!reducedWith8.isCompatibleWith(spec(12, sameShape)),
+ s"over $shape, reduced against unreduced")
+ assert(!spec(12, sameShape).isCompatibleWith(reducedWith8),
+ s"over $shape, unreduced against reduced")
+
+ // One such member does not keep a sibling from serving as the layout.
+ assert(ShuffleSpecCollection(Seq(over, spec(4, a))).canCreatePartitioning)
+ }
+ }
+ }
+
+ test("SPARK-59901: an expression over two columns maps to no cluster key") {
+ // A one-side shuffle builds its partition expression over the other side's join key, which can
+ // reference two columns. No single cluster key stands for it.
+ val a = $"a".long
+ val b = $"b".long
+ Seq(a + b, TransformExpression(new FakeBucket, Seq(a + b), Some(4))).foreach { e =>
+ val twoColumns = KeyedShuffleSpec(
+ KeyedPartitioning(Seq(e), Seq(InternalRow(0L), InternalRow(1L))),
+ ClusteredDistribution(Seq(a, b)))
+ assert(twoColumns.keyPositions.forall(_.isEmpty), s"$e")
+ assert(!twoColumns.isCompatibleWith(twoColumns), s"$e")
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(!twoColumns.canCreatePartitioning, s"$e")
+ }
+ }
+ }
+
test("createShuffleSpec: a marked narrowing projection yields an unusable spec") {
val a = $"a".int
val b = $"b".int
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
index 81a64dc15ba1..79973b599ede 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
@@ -76,7 +76,7 @@ case class GroupPartitionsExec(
// There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they
// can only differ in `expressions`. `partitionKeys` must match, so they are calculated
// only once via `groupedPartitions`. When reducers are applied, the stored reduced
- // expressions are re-targeted at each `KeyedPartitioning`'s own key attribute and reported
+ // expressions are rebuilt over each `KeyedPartitioning`'s own argument and reported
// instead of the original ones. Their data types match the reduced partition keys for the
// identity-vs-transform and single-side-transform reducers; for the both-sides-reduce
// shape no single transform describes the keys (see `KeyedShuffleSpec.reducersBothWays`).
@@ -112,9 +112,15 @@ case class GroupPartitionsExec(
projectedExpressions.zip(exprs).map {
case (expr, Some(KeyReducer(_, reduced))) =>
// `reduced` was stored from the single spec that `createKeyedShuffleSpec`
- // picked (`collectFirst`); re-target it at this `KeyedPartitioning`'s own key
- // attribute so that every `KeyedPartitioning` in a collection keeps its own.
- reduced.withReference(expr.references.head)
+ // picked (`collectFirst`), which need not be this `KeyedPartitioning`. Only
+ // its function says what the reduced keys are, so this member's own argument
+ // takes the place of `reduced`'s. A side shuffled onto this layout reports
+ // `bucket(8, -b)` next to `bucket(8, id)`, or `-b` next to an identity `id`.
+ val argument = expr match {
+ case t: TransformExpression => t.children
+ case e => Seq(e)
+ }
+ reduced.copy(children = argument)
case (expr, None) => expr
}
case None => projectedExpressions
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index 63dd2e639aa6..d85b0493aebb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -206,7 +206,11 @@ case class EnsureRequirements(
val shouldConsiderMinParallelism = specs.forall(p =>
!p._2.canCreatePartitioning || children(p._1).isInstanceOf[ShuffleExchangeLike]
)
- // Choose all the specs that can be used to shuffle other children
+ // Choose all the specs that can be used to shuffle other children. For a collection only the
+ // members that can serve as the layout count, here and below. Any other member would build a
+ // partitioning its own child is not laid out on.
+ def servingMembers(spec: ShuffleSpec): Seq[ShuffleSpec] =
+ flattenSpec(spec).filter(_.canCreatePartitioning)
val candidateSpecs = specs.filter { case (index, spec) =>
spec.canCreatePartitioning &&
(!shouldConsiderMinParallelism ||
@@ -228,8 +232,9 @@ case class EnsureRequirements(
} else {
candidateSpecs
}
- // Pick the spec with the best parallelism
- Some(finalCandidateSpecs.values.maxBy(_.numPartitions))
+ // Pick the spec with the best parallelism. A collection counts by its first serving member,
+ // which is its first member whenever that one serves.
+ Some(finalCandidateSpecs.values.maxBy(servingMembers(_).head.numPartitions))
}
// Check if the following conditions are satisfied:
@@ -261,16 +266,17 @@ case class EnsureRequirements(
// a compatible child and the partitioning built for a re-shuffled child both have to come
// from one member, otherwise the sides end up grouped on different keys, or on a key set the
// child does not even have. Pick that member once, preferring the finest when several
- // qualify. Only the branch that shuffles a child reads these, hence `lazy`.
+ // qualify. Only the branch that shuffles a child reads these, hence `lazy`. A child counts as
+ // matched only through a member that can serve as the layout.
lazy val matchedIndexes = bestSpecOpt.toSeq.flatMap { best =>
- childrenIndexes.filter(i => best.isCompatibleWith(specs(i)))
+ childrenIndexes.filter(i => servingMembers(best).exists(_.isCompatibleWith(specs(i))))
}
lazy val bestMemberOpt = bestSpecOpt.flatMap { best =>
val matchedMembers = matchedIndexes.map(i => flattenSpec(specs(i)))
// No member serving every matched child means there is no layout to align them on, so they
// all take the ordinary shuffle. That needs three or more clustered children, since with
// two the member that reported the match serves both, and no operator has three today.
- flattenSpec(best)
+ servingMembers(best)
.filter(m => matchedMembers.forall(_.exists(m.isCompatibleWith)))
.maxByOption(_.numPartitions)
}
@@ -928,9 +934,13 @@ case class EnsureRequirements(
val clustering = distribution.clustering
val satisfies = if (SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION)) {
- attributes.length == clustering.length && attributes.zip(clustering).forall {
- case (l, r) => l.semanticEquals(r)
- }
+ // Only an expression with a single leaf covers that leaf. One with several, e.g. `b + c`,
+ // maps to no position (`KeyedShuffleSpec.keyPositions`). The spec turns it away or
+ // projects it away, so a layout with one does not cover the cluster keys.
+ partitioning.expressions.forall(_.collectLeaves().size == 1) &&
+ attributes.length == clustering.length && attributes.zip(clustering).forall {
+ case (l, r) => l.semanticEquals(r)
+ }
} else {
partitioning.satisfies(distribution)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index 47c66b298e6c..6904476f754a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.execution.{
UnionExec}
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec}
import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike, ValidateRequirements}
-import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.functions.{col, max}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.SQLConf._
@@ -337,6 +337,33 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
private def createBucketedIdTables(bucketCounts: Int*): Unit =
bucketCounts.foreach(n => createBucketedIdTable(s"bucket$n", n))
+ /** Creates `ident`, identity-partitioned on `id`, holding `ids`. */
+ private def createIdentityIdTable(ids: Seq[Long] = 0L until 8L): Unit = {
+ createTable("ident", Array(Column.create("id", LongType)), Array(identity("id")))
+ sql("INSERT INTO testcat.ns.ident VALUES " + ids.map(i => s"($i)").mkString(", "))
+ }
+
+ /** Creates `name`, unpartitioned, with a column `b` holding `bs`. */
+ private def createPlainTable(bs: Seq[Long] = 0L until 8L, name: String = "plain"): Unit = {
+ createTable(name, Array(Column.create("b", LongType)), Array.empty)
+ sql(s"INSERT INTO testcat.ns.$name VALUES " + bs.map(b => s"($b)").mkString(", "))
+ }
+
+ /** Creates `plain2`, unpartitioned, with columns `b` and `c`, where `b + c` runs from 1 to 3. */
+ private def createPlain2Table(): Unit = {
+ createTable("plain2", Array(Column.create("b", LongType), Column.create("c", LongType)),
+ Array.empty)
+ sql("INSERT INTO testcat.ns.plain2 VALUES (0, 1), (0, 2), (1, 1), (1, 2)")
+ }
+
+ /** Creates `ps` with a struct column `s` of 8 rows, `(a = i, b = 7 - i)`. */
+ private def createStructTable(): Unit = {
+ createTable("ps", Array(Column.create("s", new StructType().add("a", LongType)
+ .add("b", LongType))), Array.empty)
+ sql("INSERT INTO testcat.ns.ps VALUES " +
+ (0 until 8).map(i => s"(named_struct('a', ${i}L, 'b', ${7 - i}L))").mkString(", "))
+ }
+
/** Joins `bucket12`, `bucket8` and `bucket<third>` on `id`, in that order. */
private def threeWayBucketJoinDF(third: Int): DataFrame =
sql("SELECT b12.id FROM testcat.ns.bucket12 b12 " +
@@ -949,6 +976,415 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
}
}
+ // Each key is a function of `b` other than `b` itself. A check on the leaves of the key would
+ // reject `b + 1` for its literal, but not `-b`, so both are tested.
+ Seq("p.b + 1" -> ((b: Long) => b + 1), "-p.b" -> ((b: Long) => -b)).foreach {
+ case (key, keyOf) =>
+ // `plain` holds `b` in 0..7, and `keys`, bucketed on `id`, holds the key of each `b`.
+ def createKeyTables(): Unit = {
+ createPlainTable()
+ createTable("keys", Array(Column.create("id", LongType)), Array(bucket(4, "id")))
+ sql("INSERT INTO testcat.ns.keys VALUES " +
+ (0 until 8).map(b => s"(${keyOf(b)})").mkString(", "))
+ }
+
+ test("SPARK-59887: a side shuffled onto a transform of an expression is not paired as it " +
+ s"is ($key)") {
+ createKeyTables()
+ Seq(4, 8).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 8))
+
+ // The first join shuffles `plain` onto the bucket of its key. That is a function of `b`,
+ // but not the same function of it as the third table's `bucket(n, id)` is of `id`. The
+ // second join clusters `plain` on `b`. Pairing the two there, as they stand or after
+ // reducing `bucket(8, id)` onto 4 buckets, puts a row with `b = id` on a different
+ // partition on each side.
+ Seq(4 -> "false", 8 -> "true").foreach { case (third, allowCompatibleTransforms) =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> allowCompatibleTransforms) {
+ val df = sql(
+ s"""SELECT k.id, p.b, t.id FROM testcat.ns.keys k
+ |JOIN testcat.ns.plain p ON k.id = $key
+ |JOIN testcat.ns.bucket$third t ON p.b = t.id""".stripMargin)
+ checkAnswer(df, (0 until 8).map(b => Row(keyOf(b), b.toLong, b.toLong)))
+ // One shuffle per join: `plain` for the first, the first join's output for the
+ // second, and none for the third table. The second join does not shuffle onto the
+ // bucket of the key, since that would move both of its sides.
+ assert(collectShuffles(stripAQEPlan(df.queryExecution.executedPlan)).size == 2)
+ }
+ }
+ }
+
+ test(s"SPARK-59887: a side is not shuffled onto a transform of an expression ($key)") {
+ createKeyTables()
+ createTable("xy", Array(Column.create("x", LongType), Column.create("y", LongType)),
+ Array.empty)
+ sql("INSERT INTO testcat.ns.xy VALUES " +
+ (0 until 8).map(b => s"(${keyOf(b)}, $b)").mkString(", "))
+
+ // The first join shuffles `plain` onto the bucket of its key, so its output offers that
+ // layout next to `bucket(4, id)`. The second join must not shuffle `xy` onto the first
+ // one. That builds `bucket(4, y)` for `xy`, not the bucket of the same function of `y`. A
+ // row with `b = y` then lands in a different bucket on each side. The join order decides
+ // which of the two layouts comes first, so both orders run.
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ Seq(
+ s"testcat.ns.plain p JOIN testcat.ns.keys k ON $key = k.id",
+ s"testcat.ns.keys k JOIN testcat.ns.plain p ON k.id = $key").foreach { firstJoin =>
+ val df = sql(
+ s"""SELECT k.id, p.b, q.x FROM $firstJoin
+ |JOIN testcat.ns.xy q ON k.id = q.x AND p.b = q.y""".stripMargin)
+ checkAnswer(df, (0 until 8).map(b => Row(keyOf(b), b.toLong, keyOf(b))))
+ // `xy` is shuffled onto the usable layout, `bucket(4, x)`, and never onto a bucket of
+ // `y`. So `plain` and `xy` are the only sides that move.
+ val shuffles = collectShuffles(stripAQEPlan(df.queryExecution.executedPlan))
+ assert(shuffles.size == 2)
+ assert(!keyedPartitioningsOf(shuffles).exists(_.references.exists(_.name == "y")))
+ }
+ }
+ }
+
+ test("SPARK-59887: a broadcast join's layout over a join key is not paired as it is " +
+ s"($key)") {
+ createKeyTables()
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+
+ // An inner broadcast join reports the streamed side's layout over the build side's join
+ // key too, so the first join also offers the bucket of `plain`'s key. That needs no conf
+ // beyond the defaults. The second join clusters `plain` on `b`, so pairing that layout with
+ // `bucket(4, id)` as it stands puts a row with `b = id` on a different partition on each
+ // side.
+ val df = sql(
+ s"""SELECT /*+ BROADCAST(p) */ k.id, p.b, t.id FROM testcat.ns.keys k
+ |JOIN testcat.ns.plain p ON k.id = $key
+ |JOIN testcat.ns.bucket4 t ON p.b = t.id""".stripMargin)
+ checkAnswer(df, (0 until 8).map(b => Row(keyOf(b), b.toLong, b.toLong)))
+ assert(collect(df.queryExecution.executedPlan) { case j: BroadcastHashJoinExec => j }
+ .nonEmpty, "the first join is a broadcast join")
+ }
+ }
+
+ test("SPARK-59887: a layout over a transform of a struct field is not paired nor shuffled onto") {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createStructTable()
+ // `l` shuffles `ps` onto `bucket(4, s.a)`, and `r` shuffles it onto `bucket(4, s.b)`.
+ val l = "SELECT x.id, ps.s FROM testcat.ns.bucket4 x JOIN testcat.ns.ps ps ON x.id = ps.s.a"
+ val r = "SELECT y.id, qs.s FROM testcat.ns.bucket4 y JOIN testcat.ns.ps qs ON y.id = qs.s.b"
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ // Both layouts are functions of `s`, but not the same one, so the join on `s` must not pair
+ // them as they stand. A row lands in bucket `a % 4` on one side and `b % 4` on the other.
+ checkAnswer(
+ sql(s"SELECT l.id, r.id FROM ($l) l JOIN ($r) r ON l.s = r.s"),
+ (0 until 8).map(i => Row(i.toLong, (7 - i).toLong)))
+ // Nor may a join shuffle a plain `ps` onto `l`'s layout. That builds `bucket(4, s)` for it,
+ // which drops the `.a` and hands the whole struct to a bucket function over a long.
+ checkAnswer(
+ sql(s"SELECT l.id FROM ($l) l JOIN testcat.ns.ps r ON l.s = r.s"),
+ (0 until 8).map(i => Row(i.toLong)))
+ }
+ }
+
+ test("SPARK-59887: a reduce keeps what each member's keys are computed from") {
+ Seq(4, 2).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 8))
+ createIdentityIdTable()
+ createPlainTable()
+ createStructTable()
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ // The first join shuffles `plain` onto `bucket(4, b + 1)`, or onto `b + 1` for the
+ // identity-partitioned `ident`. The second join reduces the first one's output onto
+ // `bucket(2)`, `plain`'s member included. That member is still over `b + 1`. So the third
+ // join, on the bare `b`, must neither pair it with `bucket(2, id)` nor shuffle a plain
+ // side onto it.
+ for {
+ first <- Seq("bucket4", "ident")
+ (third, thirdKey) <- Seq("bucket2" -> "id", "plain" -> "b")
+ } {
+ val df = sql(
+ s"""SELECT f.id, p.b, t2.id, t3.$thirdKey FROM testcat.ns.$first f
+ |JOIN testcat.ns.plain p ON f.id = p.b + 1
+ |JOIN testcat.ns.bucket2 t2 ON f.id = t2.id
+ |JOIN testcat.ns.$third t3 ON p.b = t3.$thirdKey""".stripMargin)
+ checkAnswer(df, (0 until 7).map(b => Row(b + 1L, b.toLong, b + 1L, b.toLong)))
+ }
+
+ // The same for a struct field. `ps`'s member is still over `s.a` after the reduce, so the
+ // third join must not shuffle a plain `ps` onto it, which would build `bucket(2, s)`.
+ val df = sql(
+ """SELECT x.id FROM testcat.ns.bucket4 x
+ |JOIN testcat.ns.ps ps ON x.id = ps.s.a
+ |JOIN testcat.ns.bucket2 t2 ON x.id = t2.id
+ |JOIN testcat.ns.ps r ON ps.s = r.s""".stripMargin)
+ checkAnswer(df, (0 until 8).map(i => Row(i.toLong)))
+ }
+ }
+
+ // `ident` also holds an id that the key's function overflows on under ANSI.
+ Seq(
+ ("p.b + 1", 0L until 8L, Long.MaxValue, (b: Long) => b + 1),
+ ("-p.b", -7L to 0L, Long.MinValue, (b: Long) => -b)
+ ).foreach { case (key, bs, overflowingId, keyOf) =>
+ test(s"SPARK-59887: an identity side is not reduced onto a transform of an expression ($key)") {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createPlainTable(bs)
+ createIdentityIdTable(bs :+ overflowingId)
+
+ // The first join shuffles `plain` onto the bucket of its key. Reducing `ident` onto it would
+ // evaluate the key's function on `ident`'s partition keys while planning, which overflows for
+ // that id. The query itself never computes it, since no `b` matches that id.
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""SELECT k.id, p.b, i.id FROM testcat.ns.bucket4 k
+ |JOIN testcat.ns.plain p ON k.id = $key
+ |JOIN testcat.ns.ident i ON p.b = i.id""".stripMargin)
+ checkAnswer(df,
+ bs.filter(b => (0L until 8L).contains(keyOf(b))).map(b => Row(keyOf(b), b, b)))
+ }
+ }
+ }
+
+ test("SPARK-59887: a reduce keeps a member over two columns") {
+ createIdentityIdTable()
+ createBucketedIdTable("bucket2", 2, numIds = 8)
+ createPlain2Table()
+
+ // The broadcast join reports `ident`'s layout over `p.b + p.c` too, and the second join reduces
+ // it onto `bucket(2)`. That member has to stay over `p.b + p.c`. Reported over `p.b`, it would
+ // serve the `GROUP BY p.b` as it stands, while the rows of one `b` sit in two buckets.
+ withSQLConf(SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ checkAnswer(
+ sql("""SELECT /*+ BROADCAST(p) */ p.b, max(p.c) FROM testcat.ns.ident i
+ |JOIN testcat.ns.plain2 p ON i.id = p.b + p.c
+ |JOIN testcat.ns.bucket2 t2 ON i.id = t2.id
+ |GROUP BY p.b""".stripMargin),
+ Seq(Row(0L, 2L), Row(1L, 2L)))
+ }
+ }
+
+ test("SPARK-59887: two layouts over the same shape of a join key pair as two columns do") {
+ Seq("orders", "returns").foreach(createBucketedIdTable(_, 16, numIds = 16))
+ createTable("customers", Array(Column.create("id", IntegerType)), Array.empty)
+ sql("INSERT INTO testcat.ns.customers VALUES " + (0 until 16).map(i => s"($i)").mkString(", "))
+
+ // Each broadcast join reports its fact table's layout over the dimension's key too, widened to
+ // the fact key's type, `bucket(16, cast(c.id as bigint))`. That needs no conf beyond the
+ // defaults. The join on `c1.id = c2.id` pairs the two as they stand, since they are the same
+ // function of their columns.
+ val dimensions = sql(
+ """SELECT a.id FROM
+ | (SELECT /*+ BROADCAST(c1) */ c1.id FROM testcat.ns.orders o
+ | JOIN testcat.ns.customers c1 ON o.id = c1.id) a
+ |JOIN
+ | (SELECT /*+ BROADCAST(c2) */ c2.id FROM testcat.ns.returns r
+ | JOIN testcat.ns.customers c2 ON r.id = c2.id) b
+ |ON a.id = b.id""".stripMargin)
+ checkAnswer(dimensions, (0 until 16).map(Row(_)))
+ assert(collectAllShuffles(dimensions.queryExecution.executedPlan).isEmpty)
+
+ // The same for two one-side shuffles onto `bucket(n, -p.b)` and `bucket(4, -q.b)`, joined on
+ // `p.b = q.b`. With `bucket(4)` on both sides they pair as they stand. With `bucket(8)` on one
+ // side, `allowCompatibleTransforms` reduces it onto `bucket(4)`. Either way only the two
+ // one-side shuffles are left. The key is `-p.b`, since a key with a literal such as `p.b + 1`
+ // maps to no cluster key on this branch.
+ Seq(4, 8).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 8))
+ createPlainTable(-7L to 0L)
+ // A second unpartitioned table, so that the two one-side shuffles are not one reused exchange.
+ createPlainTable(-7L to 0L, name = "plainq")
+ Seq(4, 8).foreach { n =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val legs = sql(
+ s"""SELECT l.b FROM
+ | (SELECT p.b FROM testcat.ns.bucket$n t1
+ | JOIN testcat.ns.plain p ON t1.id = -p.b) l
+ |JOIN
+ | (SELECT q.b FROM testcat.ns.bucket4 t2
+ | JOIN testcat.ns.plainq q ON t2.id = -q.b) r
+ |ON l.b = r.b""".stripMargin)
+ checkAnswer(legs, (-7 to 0).map(b => Row(b.toLong)))
+ assert(collectAllShuffles(legs.queryExecution.executedPlan).size == 2, s"bucket$n")
+ }
+ }
+ }
+
+ test("SPARK-59887: the same shape reduced together pairs only through the same pairing") {
+ Seq(12, 8, 18).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 24))
+ createPlainTable(-23L to 0L)
+ // A second unpartitioned table, so that the two one-side shuffles are not one reused exchange.
+ createPlainTable(-23L to 0L, name = "plainq")
+
+ // Each leg shuffles its plain table onto `bucket(12, -b)`. The first leg then reduces
+ // `bucket(12)` with `bucket(8)`, so its `-b` member holds buckets of 4. The second leg reduces
+ // with `bucket(8)` too, with `bucket(18)`, or not at all. Only with `bucket(8)` does it hold
+ // buckets of 4 as well, and only then does the final join pair the two as they stand.
+ // Otherwise it holds buckets of 6 or of 12, and the final join shuffles both sides.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ Seq(Some(8) -> 2, Some(18) -> 4, None -> 4).foreach { case (other, shuffles) =>
+ val reduce = other.fold("")(n => s"JOIN testcat.ns.bucket$n u ON t.id = u.id")
+ val df = sql(
+ s"""SELECT l.b FROM
+ | (SELECT p.b FROM testcat.ns.bucket12 t JOIN testcat.ns.plain p ON t.id = -p.b
+ | JOIN testcat.ns.bucket8 u ON t.id = u.id) l
+ |JOIN
+ | (SELECT q.b FROM testcat.ns.bucket12 t JOIN testcat.ns.plainq q ON t.id = -q.b
+ | $reduce) r
+ |ON l.b = r.b""".stripMargin)
+ checkAnswer(df, (-23 to 0).map(b => Row(b.toLong)))
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size == shuffles, s"$other")
+ }
+ }
+ }
+
+ test("SPARK-59887: a member that cannot serve does not rank its collection") {
+ createTable("t1", Array(Column.create("a", LongType), Column.create("c", LongType)),
+ Array(identity("a"), identity("c")))
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 1), (1, 2), (1, 3), (2, 4), (2, 5)")
+ createTable("p", Array(Column.create("y", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.p VALUES (0), (1), (2), (3), (4)")
+ createTable("q", Array(Column.create("m", LongType), Column.create("n", LongType)),
+ Array(identity("m")))
+ sql("INSERT INTO testcat.ns.q VALUES (1, 0), (2, 3), (3, 9)")
+
+ // The broadcast join reports `[a, c]` and `[a, y + 1]`. The first projects to `[a]`, 2
+ // partitions, and can serve as the layout. The second keeps 5 partitions and cannot. So `q`,
+ // with 3 partitions and no exchange, ranks first, and the other side is shuffled onto it.
+ // On this branch `[a, y + 1]` never satisfies the distribution, so this passes without the
+ // change too. "a collection ranks by its first member that can serve as the layout" in
+ // `EnsureRequirementsSuite` covers the ranking here.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(
+ """SELECT /*+ BROADCAST(p) */ * FROM testcat.ns.t1
+ |JOIN testcat.ns.p ON t1.c = p.y + 1
+ |JOIN testcat.ns.q ON t1.a = q.m AND p.y = q.n""".stripMargin)
+ checkAnswer(df, Seq(Row(1L, 1L, 0L, 1L, 0L), Row(2L, 4L, 3L, 2L, 3L)))
+ assert(collectAllShuffles(df.queryExecution.executedPlan)
+ .map(_.outputPartitioning.numPartitions) == Seq(3))
+ }
+ }
+
+ test("SPARK-59901: a join key over two columns is not paired through one of them") {
+ createIdentityIdTable()
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createPlain2Table()
+ val expected = Seq(Row(1L, 0L, 1L), Row(2L, 0L, 2L), Row(2L, 1L, 1L), Row(3L, 1L, 2L))
+
+ // A one-side shuffle builds its partition expression over the other side's join key. That is
+ // `b + c` here, or `bucket(4, b + c)` over a bucketed side. No single cluster key stands for
+ // either, so a spec over one pairs with nothing and is no layout to shuffle onto. On master,
+ // AQE validates the first plan and fails there, but not on this branch. The second join of
+ // the second query asks what to shuffle onto.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ checkAnswer(
+ sql("""SELECT i.id, p.b, p.c FROM testcat.ns.ident i
+ |JOIN testcat.ns.plain2 p ON i.id = p.b + p.c""".stripMargin),
+ expected)
+ checkAnswer(
+ sql("""SELECT b4.id, q.b, q.c FROM testcat.ns.bucket4 b4
+ |JOIN testcat.ns.plain2 p ON b4.id = p.b + p.c
+ |JOIN testcat.ns.plain2 q ON p.b = q.b AND p.c = q.c""".stripMargin),
+ expected)
+ }
+ }
+
+ test("SPARK-59901: a join key over two columns does not cover its columns") {
+ Seq("ident2a", "ident2b").foreach { name =>
+ createTable(name, Array(Column.create("x", LongType), Column.create("d", LongType)),
+ Array(identity("x"), identity("d")))
+ sql(s"INSERT INTO testcat.ns.$name VALUES (1, 1), (2, 1), (3, 2), (4, 2)")
+ }
+ Seq("pa", "pb").foreach { name =>
+ createTable(name, Array(Column.create("b", LongType), Column.create("c", LongType),
+ Column.create("d", LongType)), Array.empty)
+ sql(s"INSERT INTO testcat.ns.$name VALUES (0, 1, 1), (1, 1, 1), (1, 2, 2), (2, 2, 2)")
+ }
+
+ // Each broadcast join reports its identity side's layout over `[b + c, d]` too. A projection
+ // to the join keys drops `b + c`, which maps to no position, so it does not cover `b` and `c`
+ // for `spark.sql.requireAllClusterKeysForCoPartition`. Otherwise the join on `(b, c, d)` would
+ // be paired on `d` alone.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(
+ """SELECT l.b, l.c, l.d FROM
+ | (SELECT /*+ BROADCAST(p) */ p.b, p.c, p.d FROM testcat.ns.ident2a i
+ | JOIN testcat.ns.pa p ON i.x = p.b + p.c AND i.d = p.d) l
+ |JOIN
+ | (SELECT /*+ BROADCAST(q) */ q.b, q.c, q.d FROM testcat.ns.ident2b j
+ | JOIN testcat.ns.pb q ON j.x = q.b + q.c AND j.d = q.d) r
+ |ON l.b = r.b AND l.c = r.c AND l.d = r.d""".stripMargin)
+ checkAnswer(df, Seq(Row(0L, 1L, 1L), Row(1L, 1L, 1L), Row(1L, 2L, 2L), Row(2L, 2L, 2L)))
+ assert(collectAllGroupPartitions(df.queryExecution.executedPlan).isEmpty)
+
+ // With the identity side's columns kept, a sibling layout such as `[x, d]` satisfies the
+ // join's distribution, so no leg is shuffled first. Then only the coverage check keeps
+ // `[b + c, d]` from pairing on `d` alone.
+ val kept = sql(
+ """SELECT l.b, l.c, l.d, l.x, l.id, r.x, r.id FROM
+ | (SELECT /*+ BROADCAST(p) */ p.b, p.c, p.d, i.x, i.d AS id FROM testcat.ns.ident2a i
+ | JOIN testcat.ns.pa p ON i.x = p.b + p.c AND i.d = p.d) l
+ |JOIN
+ | (SELECT /*+ BROADCAST(q) */ q.b, q.c, q.d, j.x, j.d AS id FROM testcat.ns.ident2b j
+ | JOIN testcat.ns.pb q ON j.x = q.b + q.c AND j.d = q.d) r
+ |ON l.b = r.b AND l.c = r.c AND l.d = r.d""".stripMargin)
+ checkAnswer(kept, Seq(Row(0L, 1L, 1L, 1L, 1L, 1L, 1L), Row(1L, 1L, 1L, 2L, 1L, 2L, 1L),
+ Row(1L, 2L, 2L, 3L, 2L, 3L, 2L), Row(2L, 2L, 2L, 4L, 2L, 4L, 2L)))
+ assert(collectAllGroupPartitions(kept.queryExecution.executedPlan).isEmpty)
+ }
+ }
+
+ test("SPARK-59901: a join keyed on a join-key expression does not run in one partition") {
+ createTable("i1", Array(Column.create("id", LongType)), Array(identity("id")))
+ sql("INSERT INTO testcat.ns.i1 VALUES " + (-2 to 4).map(i => s"($i)").mkString(", "))
+ createTable("pa", Array(Column.create("b", LongType), Column.create("c", LongType)),
+ Array.empty)
+ sql("INSERT INTO testcat.ns.pa VALUES (0, 1), (1, 1), (1, 2), (2, 2)")
+ createTable("rk", Array(Column.create("k", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.rk VALUES " + (-2 to 4).map(i => s"($i)").mkString(", "))
+
+ // The broadcast join reports `i1`'s layout over its join key, e.g. `[b + c]`. The second join
+ // is keyed on that expression itself, so under `requireAllClusterKeysForDistribution` the
+ // layout satisfies it, but no position maps to a cluster key. Projecting onto the positions
+ // that map would run the join in one partition. `-b` is the same with one leaf.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ Seq(
+ ("-p.b", "-l.b", Seq(Row(0L, 1L, 0L), Row(1L, 1L, -1L), Row(1L, 2L, -1L),
+ Row(2L, 2L, -2L))),
+ ("p.b + 1", "l.b + 1", Seq(Row(0L, 1L, 1L), Row(1L, 1L, 2L), Row(1L, 2L, 2L),
+ Row(2L, 2L, 3L))),
+ ("p.b + p.c", "l.b + l.c", Seq(Row(0L, 1L, 1L), Row(1L, 1L, 2L), Row(1L, 2L, 3L),
+ Row(2L, 2L, 4L)))).foreach { case (innerKey, outerKey, expected) =>
+ val df = sql(
+ s"""SELECT l.b, l.c, r.k FROM
+ | (SELECT /*+ BROADCAST(p) */ p.b, p.c FROM testcat.ns.i1 i
+ | JOIN testcat.ns.pa p ON i.id = $innerKey) l
+ |JOIN testcat.ns.rk r ON $outerKey = r.k""".stripMargin)
+ checkAnswer(df, expected)
+ val plan = df.queryExecution.executedPlan
+ assert(collectAllGroupPartitions(plan).isEmpty, innerKey)
+ assert(collectShuffles(plan).size == 2, innerKey)
+ }
+ }
+ }
+
test("SPARK-59121: two reduced partitionings are not compatible by their transforms") {
// 24 ids, so that each leg's two sides really report different key sets and the join reduces
// them. With 12 ids `bucket(18, id)` is the identity and the leg would be co-partitioned as it
@@ -4545,21 +4981,19 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
case class Query(left: String, right: String, on: String, order: String, select: String,
expected: Seq[Row], sortingNodes: Int)
val ids = (94 to 100).map(i => Row(i.toLong))
- val singleColumnKeys = Seq(
+ val queries = Seq(
Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b", "p.b",
(6 to 0 by -1).map(b => Row(b.toLong)), 0),
Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b DESC", "p.b",
(0 to 6).map(b => Row(b.toLong)), 1),
Query("ident", "plain", "i.id = p.s.a", "p.s.a", "i.id", ids, 0),
Query("ident2", "plain2", "i.x = 100 - p.b AND i.y = p.c", "100 - p.b, p.c DESC", "p.b, p.c",
- Seq(Row(6L, 2L), Row(6L, 1L), Row(5L, 1L)), 1))
- // AQE's plan validation fails on `b + c` until SPARK-59901 is fixed, so it runs with AQE off.
- val twoColumnKey =
- Query("ident", "plain", "i.id = p.b + p.c", "p.b + p.c", "i.id", ids, 0)
+ Seq(Row(6L, 2L), Row(6L, 1L), Row(5L, 1L)), 1),
+ Query("ident", "plain", "i.id = p.b + p.c", "p.b + p.c", "i.id", ids, 0))
for {
- (query, aqe) <-
- singleColumnKeys.flatMap(q => Seq(q -> true, q -> false)) :+ (twoColumnKey -> false)
+ query <- queries
+ aqe <- Seq(true, false)
broadcast <- Seq(false, true)
} {
withSQLConf(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
index d941c8406061..004a946d2c66 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
@@ -1439,6 +1439,72 @@ class EnsureRequirementsSuite extends SharedSparkSession {
requiredChildOrdering = Seq(Seq.empty)
)
+ test("SPARK-59887: a collection ranks by its first member that can serve as the layout") {
+ val a = AttributeReference("a", IntegerType)()
+ val c = AttributeReference("c", IntegerType)()
+ val y = AttributeReference("y", IntegerType)()
+ val m = AttributeReference("m", IntegerType)()
+ val n = AttributeReference("n", IntegerType)()
+ val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(1, 3), InternalRow(2, 4),
+ InternalRow(2, 5))
+ // Clustered on `(a, y)`, `[a, -y]` keeps its 5 partitions and cannot serve as the layout,
+ // since `-y` is no column. `[a, c]` projects to `[a]`, 2 partitions, and can. The first member
+ // is the one that cannot, so ranking the collection by its first member would rank it at 5.
+ val keyed = new DummySparkPlanWithBatchScanChild(
+ outputPartitioning = PartitioningCollection(Seq(
+ KeyedPartitioning(Seq(a, UnaryMinus(y)), keys),
+ KeyedPartitioning(Seq(a, c), keys))))
+ val other = new DummySparkPlanWithBatchScanChild(
+ outputPartitioning = KeyedPartitioning(Seq(m), Seq(InternalRow(1), InternalRow(2),
+ InternalRow(3))))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val smj = SortMergeJoinExec(Seq(a, y), Seq(m, n), Inner, None, keyed, other)
+ val planned = EnsureRequirements.apply(smj).asInstanceOf[SortMergeJoinExec]
+
+ // The collection ranks at 2, its serving member, so the other side wins with 3, and the
+ // collection's side is shuffled onto it.
+ val shuffles = planned.left.collect { case s: ShuffleExchangeExec => s }
+ assert(shuffles.map(_.outputPartitioning.numPartitions) === Seq(3))
+ assert(planned.right.collect { case s: ShuffleExchangeExec => s }.isEmpty)
+ }
+ }
+
+ test("SPARK-59887: a child matched only through a member that cannot serve is shuffled") {
+ val a = AttributeReference("a", IntegerType)()
+ val c = AttributeReference("c", IntegerType)()
+ val y = AttributeReference("y", IntegerType)()
+ val m = AttributeReference("m", IntegerType)()
+ val n = AttributeReference("n", IntegerType)()
+ val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(1, 3), InternalRow(2, 0),
+ InternalRow(2, 1))
+ // Clustered on `(a, y)`, `[a, c]` projects to `[a]`, 2 partitions, and serves as the layout.
+ // `[a, bucket(4, -y)]` cannot, since its argument is no column. The other side pairs with that
+ // second member only. So it is shuffled onto `[a]`, and the keyed side keeps its layout.
+ // Without this change both sides stay as they are, paired through `[a, bucket(4, -y)]`. That
+ // plan is correct, and giving it up is a cost the description lists.
+ val keyed = DummySparkPlan(
+ outputPartitioning = PartitioningCollection(Seq(
+ KeyedPartitioning(Seq(a, c), keys),
+ KeyedPartitioning(Seq(a, bucket(4, UnaryMinus(y))), keys))))
+ val other = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(m, bucket(4, UnaryMinus(n))), keys))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val smj = SortMergeJoinExec(Seq(a, y), Seq(m, n), Inner, None, keyed, other)
+ val planned = EnsureRequirements.apply(smj).asInstanceOf[SortMergeJoinExec]
+
+ assert(planned.left.collect { case s: ShuffleExchangeExec => s }.isEmpty)
+ val shuffles = planned.right.collect { case s: ShuffleExchangeExec => s }
+ assert(shuffles.map(_.outputPartitioning.numPartitions) === Seq(2))
+ }
+ }
+
test("SPARK-59080: pushed-down positions index into the child's own partition expressions") {
val nL = AttributeReference("nL", IntegerType)()
val iL = AttributeReference("iL", IntegerType)()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org