peter-toth opened a new pull request, #58311:
URL: https://github.com/apache/spark/pull/58311
### What changes were proposed in this pull request?
`ShuffleExchangeExec.getPartitioner` built its `KeyGroupedPartitioner` from
`KeyedPartitioning.toGrouped`, i.e. from the *sorted* partition keys. It now
builds it from `partitionKeys` in the order the partitioning declares, and
asserts the keys are unique.
`KeyedShuffleSpec.canCreatePartitioning` additionally refuses an ungrouped
partitioning, so a spec whose keys are not unique is never chosen as the
template for shuffling the other child.
### Why are the changes needed?
`partitionKeys` is a physical layout indexed by partition id: partition `i`
holds key `partitionKeys(i)`. With
`spark.sql.sources.v2.bucketing.shuffle.enabled`, `EnsureRequirements` shuffles
only the non-keyed side of a storage-partitioned join and reuses the keyed
side's `KeyedPartitioning` as the target
(`KeyedShuffleSpec.createPartitioning`). Re-deriving the order in
`getPartitioner` therefore contradicts the side the shuffle is supposed to
co-partition with, and the join reads rows from partitions holding different
keys on each side.
Those keys are not always sorted. Two producers report them unsorted:
- `UnionExec` concatenates its children's keys in child order, so children
partitioned by `identity(id)` holding `[3, 4]` and `[1, 2]` merge to `[3, 4, 1,
2]`.
- a narrowing projection through `PartitioningPreservingUnaryExecNode`
projects the keys onto a subset of key positions, which preserves neither
sortedness nor uniqueness.
Where uniqueness is lost the result no longer satisfies a
`ClusteredDistribution`, so `EnsureRequirements` interposes a
`GroupPartitionsExec` that re-groups and sorts. Where only sortedness is lost
the partitioning still satisfies, nothing is interposed, and the unsorted keys
are the real layout of that side.
The result is silently wrong. Nothing catches it: `createPartitioning`
passes the keyed side's `partitionKeys` *reference* through unchanged, so
`PartitioningCollection.fromPartitionings` interns it on `eq` and never runs
the `require` added by SPARK-56877, and `ValidateRequirements` compares the
same declared keys on both sides.
To reproduce, two v2 tables partitioned by `identity(id)` that report one
input partition per key, one holding ids `3, 4` and the other `1, 2`:
```sql
CREATE TABLE testcat.ns.nt1 (id BIGINT, data STRING) PARTITIONED BY (id);
CREATE TABLE testcat.ns.nt2 (id BIGINT, data STRING) PARTITIONED BY (id);
INSERT INTO testcat.ns.nt1 VALUES (3, 'c'), (4, 'd');
INSERT INTO testcat.ns.nt2 VALUES (1, 'a'), (2, 'b');
CREATE TABLE t1 (id BIGINT, x STRING) USING parquet;
INSERT INTO t1 VALUES (1, 'x'), (2, 'x'), (3, 'x'), (4, 'x');
SET spark.sql.sources.v2.bucketing.shuffle.enabled=true;
SELECT u.id, u.data, t1.x
FROM (SELECT * FROM testcat.ns.nt1 UNION ALL SELECT * FROM testcat.ns.nt2) u
JOIN t1 ON u.id = t1.id;
```
keyed side (not re-grouped) partition 0 -> id 3 | 1 -> 4 | 2 -> 1 | 3
-> 2
shuffled side, before partition 0 -> id 1 | 1 -> 2 | 2 -> 3 | 3
-> 4
shuffled side, after partition 0 -> id 3 | 1 -> 4 | 2 -> 1 | 3
-> 2
Before the change the query returns no rows; the correct answer is four.
The affected branches are `master`, `branch-4.x` and `branch-4.3`.
`branch-4.2` sorts in `getPartitioner` too but has no producer of unsorted
keys, and `branch-4.1` and `branch-4.0` do not sort at all -- their
`uniquePartitionValues` only deduplicates.
SPARK-52246 fixed the same symptom for a different trigger, where the join
keys are a subset of the partition keys. This one needs no subset keys.
It is also distinct from SPARK-58988, which fixed the projected branch of
`createShuffleSpec`. There a `GroupPartitionsExec` *is* interposed on the keyed
side and sorts it, so the shuffled side's declared keys have to be sorted to
match, and the mismatch surfaced as a loud `require` failure rather than wrong
rows. Both are needed: reverting either one leaves its own case broken.
### Does this PR introduce _any_ user-facing change?
Yes, it fixes wrong results, but only for queries that opt into
`spark.sql.sources.v2.bucketing.shuffle.enabled`, which is off by default.
### How was this patch tested?
Added two tests to `KeyGroupedPartitioningSuite`, one per producer of
unsorted keys: the union case above, and a narrowing projection that drops a
partition column. Both fail without the change, with `Correct Answer - 4`
against `Spark Answer - 0` and `Correct Answer - 2` against `Spark Answer - 0`
respectively. The plan shape is identical either way (one shuffle, no
`GroupPartitionsExec`), which the tests assert, so `checkAnswer` is what
discriminates. The union test also re-checks the answer with adaptive execution
on, since the plan-shape assertions require it off.
Added a `ShuffleSpecSuite` test for the `canCreatePartitioning` gate,
covering a grouped spec with unsorted keys (accepted) and an ungrouped one
(refused); it fails if the new clause is removed.
Ran `ShuffleSpecSuite`, `KeyGroupedPartitioningSuite`,
`EnsureRequirementsSuite`, `GroupPartitionsExecSuite`,
`ProjectedOrderingAndPartitioningSuite`, `DataSourceV2Suite`,
`WriteDistributionAndOrderingSuite`, `ValidateRequirementsSuite`,
`PlannerSuite` and `AdaptiveQueryExecSuite`: 595 tests, 0 failures. Also
verified with a temporary assertion that no covered scenario reaches
`getPartitioner` with duplicate partition keys, and that a duplicating
narrowing projection is handled by a `GroupPartitionsExec` under both
`allowKeysSubsetOfPartitionKeys` and `requireAllClusterKeysForDistribution`.
### Was this patch authored or co-authored using generative AI tooling?
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]