peter-toth opened a new pull request, #59238:
URL: https://github.com/apache/spark/pull/59238
### Differences from #59165
This is the branch-4.2 backport of #59165. The description below is the
original one. The cherry-pick does not apply, since the storage-partitioned
join code on this branch predates master's rework of it. So the fix is ported
onto this branch's own code:
- `KeyedShuffleSpec.keyPositions` maps a partition expression by its leaves
here, not by its references. So a key with a literal, e.g. `b + 1`, already
maps to no cluster key and is not paired wrongly on this branch. A single-leaf
key such as `-b` or `s.a` is. The assertion in `keyPositions` fails here too
for an expression with two leaves, `b + 1` included (SPARK-59901). It is
replaced by "only an expression with a single leaf maps to a position".
- This branch picks the layout to shuffle onto in `EnsureRequirements`
through `candidateSpecs`, `bestSpecOpt` and `bestMemberOpt`, not
`pickCoPartitionTarget`. With `ShuffleSpecCollection.canCreatePartitioning` an
`exists`, a collection whose first member cannot serve becomes a candidate
here. So three places look at its serving members only:
- the ranking counts a collection by its first serving member, which is
its first member whenever that one serves, as before;
- a child counts as matched only through a serving member;
- the layout is picked from the serving members.
- Two new tests in `EnsureRequirementsSuite` cover the ranking and the
matching, since master's tests do not reach them on this branch: "SPARK-59887:
a collection ranks by its first member that can serve as the layout" and
"SPARK-59887: a child matched only through a member that cannot serve is
shuffled". Each fails with its change removed.
- Matching only through members that can serve costs a shuffle in one more
shape on this branch. Two legs that line up only through members of the same
shape now take a shuffle, e.g. `[a, bucket(4, -y)]` against `[m, bucket(4,
-n)]`, joined on `l.a = r.m AND l.y = r.n`. With
`allowJoinKeysSubsetOfPartitionKeys`, the shuffle conf on and
`requireAllClusterKeysForCoPartition` off, branch-4.2 took no shuffle, and this
change shuffles one side onto 2 partitions. The rows are right either way.
- Master pairs such legs in `checkKeyGroupCompatible`, which compares
every member of both sides. It takes no shuffle there, before and after this
change.
- This branch's `createKeyedShuffleSpec` offers only the first member that
satisfies the distribution, here `[a, bucket(4, c)]`. So only the shuffle path
paired the legs, through a member that cannot serve as the layout. This change
keeps the shuffle path from that, which is also where master's third cost above
comes from.
- A widening cast instead of the negation behaves the same. So "Two sides
that both carry the cast still pair" holds here only where the
storage-partitioned join path pairs them.
- This branch's `createKeyedShuffleSpec` checks
`spark.sql.requireAllClusterKeysForCoPartition` by matching the leaves of the
layout to the join keys one by one, where master checks coverage. So the
coverage change becomes: a layout with an expression with several leaves does
not meet the requirement. With the rest of this change alone, the rows are
right, but a join on `(b, c, d)` over `[b + c, d]` is paired on `d` alone.
- That needs a sibling layout that satisfies the join's distribution, so
that the child is not shuffled first. The coverage test, "SPARK-59901: a join
key over two columns does not cover its columns", has none on this branch,
since each leg's projection drops the identity side's columns. So it also runs
with those columns kept.
- `KeyedPartitioning.createShuffleSpec` gets master's refusal to project
onto no position. It came with SPARK-59289, which is not on this branch, and
master's SPARK-59901 fix relies on it.
- A join keyed on the expression itself, e.g. `l.b + l.c = r.k`, can
satisfy the distribution under
`spark.sql.requireAllClusterKeysForDistribution`, while no position maps to a
cluster key.
- With `allowJoinKeysSubsetOfPartitionKeys`, the projection would then run
the join in one partition. That needs
`spark.sql.sources.v2.bucketing.shuffle.enabled`, or another side over the same
expression with `requireAllClusterKeysForCoPartition` off. A `-b` key already
did that on this branch.
- A `b + 1` or `b + c` key failed there with the `AssertionError` before
this change, even with the shuffle conf off. Now all three get the shuffle,
which can mean 2 shuffles where branch-4.2 had none.
- The new test "SPARK-59901: a join keyed on a join-key expression does
not run in one partition" covers `-b`, `b + 1` and `b + c`.
- `TransformExpressionSuite` does not exist on this branch, so its hunk is
dropped. So are the doc changes for code this branch does not have.
- Three tests are tailored, since master's version cannot pass here. On this
branch a key with a literal maps to no cluster key, so it is shuffled rather
than paired:
- the first `ShuffleSpecSuite` test uses `-a`, `-b` and `abs(b)` for
master's `a + 1`, `b + 1` and `b + 2`;
- "two layouts over the same shape of a join key pair as two columns do"
and "the same shape reduced together pairs only through the same pairing" use
`-p.b` over negative values for master's `p.b + 1`.
- Measured on branch-4.2 without the change, with the shuffle conf on: the
`-p.b` key returns 4 of 8 rows through a one-side shuffle and through the
`bucket(8)` reducer path. Through a broadcast join it returns 4 of 8 rows with
every conf at its default. The struct field pairing returns 0 of 8 rows, and
shuffling onto it fails with a `ClassCastException`.
- Two parts of the original description need a `-p.b` key on this branch,
since a member with a literal does not satisfy the distribution here:
- the plan that gets better: `ident i JOIN plain p ON i.id = -p.b JOIN xy
q ON i.id = q.x AND p.b = q.y` takes 2 shuffles instead of 3. With `p.b + 1` it
already took 2;
- the third cost: with `-p.b` branch-4.2 took no shuffle, and this change
takes 1. With `p.b + 1` one side was already shuffled.
- A single join keyed on `p.b + p.c` does not fail on this branch, AQE
included, so "a single join is enough" does not hold here. So the first query
of "SPARK-59901: a join key over two columns is not paired through one of them"
passes without the change, and its second query fails with the `AssertionError`.
- These included tests pass on branch-4.2 without the change:
- the `p.b + 1` variants of the shuffled-onto, `xy`, broadcast and
identity tests;
- "two layouts over the same shape of a join key pair as two columns do";
- "a member that cannot serve does not rank its collection", since its
`p.y + 1` member does not satisfy the distribution here;
- the SPARK-59905 test, which now runs `p.b + p.c` with AQE on as on
master;
- the new ranking test, since a collection with a member that cannot serve
is no candidate there.
- Each change was removed on its own:
- without the shape gate, the `-p.b` shuffled-onto, broadcast and identity
tests and the struct field test fail;
- without the `canCreatePartitioning` clause, the `-p.b` shuffled-onto and
`xy` tests, the struct field, reduce and same-pairing tests fail on the
assertion in `rebuiltOver`. The new matching test and the first
`ShuffleSpecSuite` test fail too;
- without the member filter, the `-p.b` `xy` test and both new
`EnsureRequirementsSuite` tests fail;
- with the ranking by the first member back, the new ranking test fails;
- with the matching through any member back, the new matching test fails;
- without the `GroupPartitionsExec` change, the reduce test and the
two-column `GROUP BY` fail;
- with the `keyPositions` assertion back, the `ShuffleSpecSuite`
SPARK-59901 test, "SPARK-59901: a join key over two columns is not paired
through one of them" and "SPARK-59901: a join keyed on a join-key expression
does not run in one partition" fail. The coverage test passes, since the
coverage change keeps `[b + c, d]` away from `keyPositions`;
- with the `forall` back in `ShuffleSpecCollection.canCreatePartitioning`,
the `-p.b` `xy` test fails on the shuffle count, and so do the new matching
test and the first `ShuffleSpecSuite` test;
- removing the coverage change or the refusal to project onto no position
fails only its new test or run.
- Ran `ShuffleSpecSuite`, `DistributionSuite`, the `KeyGroupedPartitioning*`
suites, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`,
`ProjectedOrderingAndPartitioningSuite`, `PlannerSuite`,
`WriteDistributionAndOrderingSuite`, `BucketedReadWithoutHiveSupportSuite`, the
plan stability suites, the `*JoinSuite` suites and `AdaptiveQueryExecSuite` on
branch-4.2, 1403 tests in all, plus `dev/lint-scala`.
### What changes were proposed in this pull request?
A storage-partitioned join now pairs a partition transform whose argument is
not a bare column, i.e. is an expression or a struct field, only with a
transform over the same argument. It also stops dropping such an argument where
it used to:
- `KeyedShuffleSpec.isExpressionCompatible` pairs a transform with an
argument that is not an `Attribute` only with a transform over the same
argument shape, i.e. the argument with its one column replaced by a
placeholder. So `bucket(4, b + 1)` pairs with `bucket(4, c + 1)` on `b = c`,
but not with `bucket(4, x)`. Over the same shape the two compare as two
transforms of columns do. Under `allowCompatibleTransforms`, `bucket(8, b + 1)`
reduces onto `bucket(4, c + 1)`, since a reduce maps the partition keys and
evaluates no argument. When an earlier join reduced the keys, two such
transforms pair only when they were reduced through the same pairing. An
identity side is never reduced onto such a transform. `TransformExpression`
keeps its comparisons. It loses `withReference`, whose last callers now rebuild
the transform with `copy`.
- `KeyedShuffleSpec.canCreatePartitioning` asks the same question, so such a
layout is never the one other children are shuffled onto. `createPartitioning`
replaces a transform's argument with the other child's cluster key, which drops
whatever surrounds the key. It and the identity arm of `reducersBothWays` now
rebuild a transform through one helper that asserts the argument is a column.
- `ShuffleSpecCollection.canCreatePartitioning` is now true when any member
can, and `EnsureRequirements.pickCoPartitionTarget` offers and ranks only the
members that can. So one member that cannot serve no longer rules out a usable
sibling. That also covers a bare expression such as `b + 1`, which master
already refused.
- `GroupPartitionsExec` rebuilds a reduced expression over each member's own
argument when it reports the members of a collection. It used to re-target the
column only.
- `KeyedShuffleSpec.keyPositions` maps an expression with more than one
reference, e.g. `bucket(4, b + c)`, to no position instead of failing an
assertion (SPARK-59901).
- `EnsureRequirements.createKeyedShuffleSpecs` counts only an expression
over one column towards `spark.sql.requireAllClusterKeysForCoPartition`. An
expression over two columns maps to no position, so a projection would drop it
together with its columns.
### Why are the changes needed?
With `spark.sql.sources.v2.bucketing.shuffle.enabled` on, this query returns
0 rows instead of 7. `t1(id)` and `t3(x)` are partitioned by `bucket(4, ...)`,
`plain(b)` is not partitioned, and each holds the values 0 to 7:
```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, with every conf at its default.
`BroadcastHashJoinExec.expandOutputPartitioning` also reports the streamed
side's layout over the build side's join key. The same problem shows up in more
shapes, all measured on master:
- **A broadcast join.** With every conf at its default, `SELECT /*+
BROADCAST(p) */ ... FROM t1 JOIN plain p ON t1.id = p.b + 1 JOIN t3 ON p.b =
t3.x` returns 0 of 7 rows the same way, through the expanded `bucket(4, b + 1)`
layout.
- **The reducer path.** Under
`spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled`, a
`bucket(8, x)` third table loses the rows the same way.
- **A struct field.** A join key `p.s.a` gives `bucket(4, s.a)`, which
`keyPositions` pairs through `s`. Joining two such sides on `s` pairs
`bucket(4, s.a)` with `bucket(4, s.b)` and returns 0 of 8 rows. Shuffling
another side onto such a layout builds `bucket(4, s)`, which fails on executors
with a `ClassCastException`.
- **Shuffling onto the layout.** In `plain p JOIN t1 ON p.b + 1 = t1.id JOIN
xy q ON t1.id = q.x AND p.b = q.y`, the second join shuffles `xy` onto the
first join's `bucket(4, b + 1)` layout. That builds `bucket(4, y)` for `xy`,
drops the `+ 1`, and returns 0 of 8 rows.
- **After a reduce.** When a later join reduces the first join's output from
`bucket(4)` onto `bucket(2)`, `GroupPartitionsExec` reports the `bucket(4, b +
1)` member as `bucket(2, b)`. A join on `b` then pairs it with a `bucket(2, x)`
scan, or shuffles another side onto it, and returns 0 of 7 rows. The same
happens to the bare `b + 1` member of a side shuffled onto an
identity-partitioned table. The struct form fails with the `ClassCastException`.
- **An identity side reduced onto it.** Under `allowCompatibleTransforms`, a
later join can reduce an identity-partitioned side onto `bucket(4, b + 1)`.
That evaluates `bucket(4, id + 1)` on the side's partition keys while planning.
The keys come out right, but under ANSI a key of `Long.MaxValue` fails the
query with `ARITHMETIC_OVERFLOW`, although the query never computes `id + 1`.
- **A GROUP BY after a reduce.** With `allowCompatibleTransforms`, `SELECT
/*+ BROADCAST(p) */ p.b, max(p.c) FROM ident i JOIN plain2 p ON i.id = p.b +
p.c JOIN bucket2 t2 ON i.id = t2.id GROUP BY p.b` returns 4 groups instead of
2. The reduce reports the `p.b + p.c` member as `bucket(2, p.b)`, which serves
the `GROUP BY` as it stands.
- **A key over two columns.** A join key such as `p.b + p.c` gives
`bucket(4, b + c)`, or the bare `b + c` over an identity-partitioned side. Any
spec over it fails planning with an `AssertionError` in `keyPositions`
(SPARK-59901). AQE validates the plan in `OptimizeSkewedJoin`, so a single join
is enough.
An `ORDER BY` over such a partition expression had a related problem in
another code path. That was SPARK-59905, fixed in #59189.
The one-side shuffle of transform expressions came with SPARK-48012, in
4.0.0. So did the broadcast expansion of a keyed layout, since SPARK-49205 made
`KeyGroupedPartitioning` an expression.
### Does this PR introduce _any_ user-facing change?
Yes, it fixes the wrong results, the crash and the planning failures above.
Such a query now shuffles the side instead of pairing it or shuffling onto it.
That costs shuffles in three shapes whose plan was already correct:
- A widening cast counts as part of the argument's shape. When `t1.id` is
`BIGINT` and `p.b` is `INT`, the first join reports `bucket(4, cast(b as
bigint))`. A later join `p.b = t3.x` with an `INT` bucketed `t3` then shuffles
both sides. Before, a connector whose bucket function has one canonical name
for both types paired the two sides as they stood. This happens with every conf
at its default, through a broadcast join. Two sides that both carry the cast
still pair.
- Under `allowCompatibleTransforms`, an identity-partitioned side is no
longer reduced onto such a transform, so it is shuffled.
- With `spark.sql.sources.v2.bucketing.shuffle.enabled`, a join on two keys
whose sides pair only through a transform of an expression gets one more
shuffle. Take two broadcast joins, `t1 JOIN /*+ BROADCAST(p) */ p ON t1.id =
p.b + 1` and the same over `t2` and `q`, joined on `l.id = r.c AND l.b = r.b`.
Each layout covers only one join key, so the storage-partitioned join declines
under `requireAllClusterKeysForCoPartition`. Master then keeps both sides,
paired on `bucket(4, b + 1)`. This PR cannot shuffle onto that layout, so it
shuffles one side onto `bucket(4, id)`. Both plans co-partition on one key
only, which `requireAllClusterKeysForCoPartition` is meant to prevent. That is
a separate, pre-existing question, filed as SPARK-59971.
One plan gets better. A collection with a bare expression member, e.g. the
`b + 1` of a side shuffled onto an identity-partitioned table, used to be
refused as a layout as a whole. Now its other members can serve, so `ident i
JOIN plain p ON i.id = p.b + 1 JOIN xy q ON i.id = q.x AND p.b = q.y` takes 2
shuffles instead of 3.
This PR keeps the fix small, since it has to reach the maintenance branches.
SPARK-59900 is the follow-up for the rest: shuffling another side onto such a
layout, and reducing an identity side onto it.
### How was this patch tested?
New tests. All but two of them fail on master:
- `ShuffleSpecSuite`, "SPARK-59887: a spec over a transform of an expression
pairs only with the same shape". For an expression and a struct field argument,
it covers the same shape, another shape and a column, both ways and with
itself, the reduce in both directions, the two sides of one reduce, two sides
reduced through different pairings, an identity side, `canCreatePartitioning`,
and a collection with a usable sibling.
- `ShuffleSpecSuite`, "SPARK-59901: an expression over two columns maps to
no cluster key".
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a side shuffled onto a
transform of an expression is not paired as it is". It runs the query above
against a `bucket(4)` third table, and with `allowCompatibleTransforms` against
a `bucket(8)` one that the join reduces, and asserts two shuffles.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a side is not shuffled onto a
transform of an expression". It runs both join orders and asserts that `xy` is
shuffled onto `bucket(4, x)` and never onto a bucket of `y`.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a broadcast join's layout
over a join key is not paired as it is". It runs with every conf at its default
and checks that the first join is a broadcast join.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a layout over a transform of
a struct field is not paired nor shuffled onto". Master returns 0 of 8 rows.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a reduce keeps what each
member's keys are computed from". It covers a bucketed and an
identity-partitioned first table, a bucketed and an unpartitioned third table,
and the struct form. Master returns 0 of 7 rows.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: an identity side is not
reduced onto a transform of an expression". It runs with `p.b + 1` over a
`Long.MaxValue` key and with `-p.b` over a `Long.MinValue` key. Master fails
with the `ARITHMETIC_OVERFLOW` for both.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a reduce keeps a member over
two columns". It runs the `GROUP BY` above. Master returns 4 groups instead of
2.
- `KeyGroupedPartitioningSuite`, "SPARK-59901: a join key over two columns
is not paired through one of them". Master fails with the `AssertionError`.
- `KeyGroupedPartitioningSuite`, "SPARK-59901: a join key over two columns
does not cover its columns". It joins two broadcast joins that report `[b + c,
d]` on `(b, c, d)` under `allowKeysSubsetOfPartitionKeys`, and asserts that the
join is not paired on `d` alone. Master fails with the `AssertionError`.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: two layouts over the same
shape of a join key pair as two columns do". It runs the cast case above with
every conf at its default. It also runs two one-side shuffles onto `bucket(4, b
+ 1)` and onto `bucket(4)` or `bucket(8)` of `q.b + 1`, where
`allowCompatibleTransforms` reduces the `bucket(8)` one. It asserts that the
join adds no shuffle. It passes on master, which pairs and reduces these as
well.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: the same shape reduced
together pairs only through the same pairing". Two legs each shuffle onto
`bucket(12, b + 1)`. The first leg reduces with `bucket(8)`, so its keys are
buckets of 4. The second leg reduces with `bucket(8)` too, with `bucket(18)`,
or not at all, so its keys are buckets of 4, 6 or 12. It asserts 2 shuffles for
the first and 4 for the other two. Master returns 0 of 23 rows for the third.
- `KeyGroupedPartitioningSuite`, "SPARK-59887: a member that cannot serve
does not rank its collection". It pins the ranking: `q` keeps its layout with 3
partitions, as on master, where the collection is refused as a whole. Without
the member filter, the collection would rank first with 5 partitions, and `q`
would be shuffled onto 2.
The SPARK-59905 test from #59189 now also runs its `p.b + p.c` query with
AQE on. Master fails that run with the `AssertionError`.
The three expression tests each run with the key `p.b + 1` and with `-p.b`.
Master returns 0 of 8 rows for `p.b + 1` and 4 of 8 for `-p.b`. The `-p.b` key
is the shape that also fails on 4.2, 4.1 and 4.0. On those branches `b + 1`
already works, since its literal leaf blocks the pairing.
Each change was removed on its own:
- without the guard in `isExpressionCompatible`, the one-side shuffle and
broadcast pairing tests, the struct field test and the reduce test return wrong
rows, and the identity tests fail on the assertion in `rebuiltOver`;
- without the shape pairing, the same-shape and same-pairing tests fail on
the shuffle count;
- with a reduce of such a pair refused, the same-shape test fails on the
shuffle count;
- with every pair of reduced keys over an expression refused, the
same-pairing test fails on the shuffle count;
- with the reduced marks ignored, or only checked on both sides, the
same-pairing test returns 7 of 23 rows;
- without the `canCreatePartitioning` clause, the one-side shuffle pairing,
`xy`, struct field, reduce and same-pairing tests fail, all on the assertion in
`rebuiltOver`;
- with the `forall` back in `ShuffleSpecCollection.canCreatePartitioning`,
the `xy` tests fail on the shuffle count;
- without the member filter in `pickCoPartitionTarget`, 15 tests fail in
`KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite`. The filter now
also does what the collection-level check before it did, so 8 of them are older
tests. Of this PR's tests, 7 fail, 6 of them on the assertion in `rebuiltOver`.
A member over the same shape pairs with itself, so only the filter keeps it
from being the layout;
- without the `GroupPartitionsExec` change, the reduce test returns 0 of 7
rows, and the two-column `GROUP BY` returns 4 groups;
- without the `requireAllClusterKeysForCoPartition` change, the coverage
test pairs on `d` alone;
- with the `keyPositions` assertion back, both SPARK-59901 tests in
`KeyGroupedPartitioningSuite` fail.
The first `ShuffleSpecSuite` test also fails for each change to
`isExpressionCompatible` above.
The `xy` tests pass with the guard removed, since the
`canCreatePartitioning` clause and the member filter keep `xy` off that layout.
Also ran `TransformExpressionSuite`, `ShuffleSpecSuite`,
`DistributionSuite`, the `KeyGroupedPartitioning*` suites,
`WriteDistributionAndOrderingSuite`, `PlannerSuite`,
`ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite`,
`EnsureRequirementsSuite`, `BucketedReadWithoutHiveSupportSuite`, the plan
stability suites, the `*JoinSuite` suites and `AdaptiveQueryExecSuite`, 1955
tests in all, plus `dev/lint-scala`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
