ulysses-you opened a new pull request, #58339:
URL: https://github.com/apache/spark/pull/58339
### What changes were proposed in this pull request?
The one-side shuffle optimization
(`spark.sql.sources.v2.bucketing.shuffle.enabled`) re-shuffles a join side onto
the keyed side's declared `KeyedPartitioning`. `KeyGroupedPartitioner` silently
routes rows whose key is not among the declared partition keys to arbitrary
partitions, so the re-shuffled side's `KeyedPartitioning` may not actually
cover all its rows.
- Add `KeyedPartitioning.mayContainUnknownPartitionKeys`, set on every
partitioning produced by `KeyedShuffleSpec.createPartitioning`, and require
`KeyedShuffleSpec.areKeysCompatible` to only co-partition such a partitioning
with a side whose keys are a subset of the declared keys (in the same order
when both sides are flagged).
- Propagate the marker through `PartitioningPreservingUnaryExecNode`,
`GroupPartitionsExec` and `UnionExec` so it is never silently dropped, and a
downstream storage-partitioned operator never trusts a layout that may not
cover the data.
### Why are the changes needed?
A storage-partitioned join where the preserved side is one-side-shuffled
(e.g. `a RIGHT OUTER JOIN t` with a keyed `a` and a non-keyed `t` holding keys
`a` does not) produces an output whose declared partitioning omits those keys.
A following storage-partitioned join trusts the declaration and silently loses
the matches -- wrong results with no error. Minimal repro (only
`shuffle.enabled` is needed):
- `a` keyed on `id` with keys {1,2}; `t` (v1) with keys {1,2,3}; `u` keyed
on `id` with keys {1,2,3}.
- `SELECT ... FROM (SELECT t.id AS id FROM a RIGHT OUTER JOIN t ON a.id =
t.id) r JOIN u ON r.id = u.id` returns 2 rows instead of 3 (the id=3 match is
lost).
The new bug-repro tests fail before the fix (3 vs 2, 4 vs 3) and pass after.
### Does this PR introduce _any_ user-facing change?
Yes: it fixes a wrong-results bug in
`spark.sql.sources.v2.bucketing.shuffle.enabled`. Storage-partitioned joins
whose preserved side is one-side-shuffled against a smaller keyed side,
followed by another storage-partitioned join on a larger key set, now fall back
to shuffles and return correct results. Previously they could silently lose
matches.
### How was this patch tested?
- `KeyGroupedPartitioningSuite`: identity and bucket one-side-shuffle repros
(assert correct answers, shuffle count and `mayContainUnknownPartitionKeys`
flags), outer-join variants, union marker preservation, subset-keyed partner
compatibility, and a keyed-preserved-side regression (the one-side shuffle is
still used when sound).
- `ShuffleSpecSuite`: unit tests for the `areKeysCompatible` subset/order
rules.
- Confirmed the bug-repro tests fail without the fix (wrong results) and
pass with it.
- Regression suites all pass: `KeyGroupedPartitioningSuite`,
`GroupPartitionsExecSuite`, `ShuffleSpecSuite`, `DistributionSuite`,
`ValidateRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code.
--
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]
