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 <[email protected]>
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 <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../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: [email protected]
For additional commands, e-mail: [email protected]
