peter-toth opened a new pull request, #59323:
URL: https://github.com/apache/spark/pull/59323
### What changes were proposed in this pull request?
`KeyedShuffleSpec.canCreatePartitioning` now honors
`spark.sql.requireAllClusterKeysForCoPartition`, as
`HashShuffleSpec.canCreatePartitioning` does. With the conf on (the default), a
keyed spec can be the layout of a co-partitioned join only if its partition
keys cover every join key.
It uses the storage-partitioned join's coverage rule. Every join key must be
the reference of some partition expression, in any order. The rule moves from a
local `def` in `EnsureRequirements.createKeyedShuffleSpecs` to
`KeyedPartitioning.referencesAllClusterKeys`, and both places call it.
The `spark.sql.sources.v2.bucketing.shuffle.enabled` docs and the migration
guide say so too.
### Why are the changes needed?
`spark.sql.requireAllClusterKeysForCoPartition` keeps a join from
co-partitioning on only some of its join keys, to avoid skew and low
parallelism. With `spark.sql.sources.v2.bucketing.shuffle.enabled` on, a join
could still do that. Take `k1(a, b)` and `k2(a, b)`, identity-partitioned on
`a` with the same `a` values, and an unpartitioned `plain(a, b)`:
```sql
SELECT * FROM k1 x JOIN k2 y ON x.a = y.a AND x.b = y.b;
SELECT * FROM k1 x JOIN plain y ON x.a = y.a AND x.b = y.b;
```
- The first query took no shuffle. The storage-partitioned join declined,
since `a` does not cover the join key `b`. Then the shuffle path paired both
sides as they stood, so the join still ran on `k1`'s partitions.
- The second query shuffled `plain` onto `k1`'s layout, so the join ran on
`a` only, with `k1`'s partition count.
The shuffle path picks its layout with `canCreatePartitioning`, and only the
keyed spec did not check the conf. This came with SPARK-41471 in 4.0.0.
SPARK-58558 made the storage-partitioned join check the coverage, but did not
change the shuffle path.
### Does this PR introduce _any_ user-facing change?
Yes, with `spark.sql.sources.v2.bucketing.shuffle.enabled` on. A join no
longer runs on partition keys that cover only some of the join keys:
- If a side covers every join key, the join runs on that side's layout.
- Otherwise both sides are shuffled on all the join keys. Both queries above
now take two shuffles.
A cogroup changes the same way. The results are the same. Setting
`spark.sql.requireAllClusterKeysForCoPartition` to `false` allows the previous
plans. The migration guide has an entry for 4.4.
### How was this patch tested?
- New test in `ShuffleSpecSuite`: `canCreatePartitioning` for the clustering
`[a, b]`, of a keyed spec over `[a]`, `[a, b]`, `[b, a]`, `[bucket(4, a), b]`
and `[bucket(4, a)]` with the conf on, and over `[a]` with it off.
- New test in `KeyGroupedPartitioningSuite`, "SPARK-59971: a side is laid
out on another only if its keys cover every join key". It runs both queries
above with the conf on, and the first one with it off. It fails without the
fix. The first query then takes no shuffle, and the second takes one.
- Tests that need a layout covering only some of the join keys now turn
`spark.sql.requireAllClusterKeysForCoPartition` off:
- SPARK-52246, both SPARK-54439 tests and three SPARK-59887 tests, which
assert a one-side shuffle onto such keys.
- The two SPARK-59080 tests and the SPARK-59256 test, which pick between
collection members, one of which covers only some of the join keys. With the
conf on, that member cannot serve, so a wrong pick would no longer fail them.
- The SPARK-59887 test "a member that cannot serve does not rank its
collection" now buckets `q`. Otherwise, with the conf off, the
storage-partitioned join pairs `q` with `t1`'s `[a]`. It still fails when the
ranking counts a member that cannot serve.
- Ran `ShuffleSpecSuite`, `EnsureRequirementsSuite` and
`KeyGroupedPartitioningSuite`, 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]
