ulysses-you opened a new pull request, #59535:
URL: https://github.com/apache/spark/pull/59535
### What changes were proposed in this pull request?
A `UnionExec` whose children all report a `KeyedPartitioning` merges their
keys into one layout
(`KeyedPartitioning.concat`), so an operator above it can group the union's
partitions
(`GroupPartitionsExec`) instead of shuffling them. That merge is given up
whenever one child's
layout carries `mayContainUnknownPartitionKeys`: such a child routes a row
whose key it does not
declare into the partition of some declared key, by
`nonNegativeMod(key.hashCode, numPartitions)`,
and another child may declare exactly that key, so the merged claim could
promise co-location the
marked child cannot honor.
This PR keeps the merge where that cannot happen.
`UnionExec.canMergeMarkedKeys` asks two things
of the legs:
- a marked leg declares every key the other legs declare. Otherwise its rows
of such a key land in
a partition the merge names for another key while the other leg's land in
its own;
- marked legs declare the same keys in the same order, so a key no leg
declares reaches the same
declared key from all of them (the fallback reads the key's hash and the
partition count, which a
shared declaration fixes).
Every key's rows then share one group, which is what a clustering asks of
the merged layout. What
the layout cannot claim is that a partition holds only the key it names, so
`KeyedPartitioning.concat` now carries the marker onto the merged layout,
read from the children
rather than inherited from the first one. Pairing keeps reading the marker
there and admits exactly
what it admitted before (`KeyedShuffleSpec.areKeysCompatible`: identical
declared keys against
another marked side, a subset of them against an accurate one).
The `mayContainUnknownPartitionKeys` carry rule (1) in `KeyLayout`'s doc is
updated accordingly.
### Why are the changes needed?
A union leg that comes out of a one-side-shuffled join is marked although it
declares every key
the other legs hold, and today the whole union gives up its keyed
partitioning, so the operator
above it pays an exchange it does not need. For instance, with `a`
partitioned by `identity(id)`
over keys `{1, 2, 3}`, `s` partitioned the same way over `{1, 2}`, and an
unpartitioned `t`:
```sql
SELECT id, count(*) FROM (
SELECT t.id AS id FROM a RIGHT OUTER JOIN t ON a.id = t.id
UNION ALL
SELECT id FROM s
) GROUP BY id
```
The right-outer join exposes a marked layout declaring `{1, 2, 3}`, which
covers `s`'s keys. Before
this PR the union reports `UnknownPartitioning` and the aggregate adds its
own shuffle; after it,
the keys merge and the aggregate groups the union's partitions, leaving the
join's one-side shuffle
as the only exchange in the plan.
### Does this PR introduce _any_ user-facing change?
No, other than the plan change above: a union of keyed legs, one of which
may hold keys outside its
declared set, can now keep a keyed partitioning and save the exchange a
consumer would otherwise
add. Results are unchanged.
### How was this patch tested?
New tests in `KeyGroupedPartitioningSuite`:
- `SPARK-60277: SPJ: a union merges a marked leg that declares what the
other legs declare` -- the
query above, asserting the merged marked layout, the grouping, and a
single exchange, with the
marked leg holding out-of-set rows so the counts verify their co-location.
Its negative control
points the second leg at exactly the key the marked leg holds out-of-set,
where the merge must
still be refused and the aggregate pays its own exchange;
- `SPARK-60277: SPJ: two marked legs route the keys they do not declare
alike` -- two marked legs
declaring the same keys and holding different out-of-set keys, one of
which both hold, so the
counts fail if the legs were free to route it differently.
`SPARK-59057: a shuffle built from a template that collapsed nothing needs
no opt-in` asserted the
old refusal; its marked leg declares every key the other leg does, so it now
asserts the merge and
the exchange it saves, keeping what it was written to exercise.
Ran `KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`,
`UnionCodegenSuite`,
`ValidateRequirementsSuite`, `GroupPartitionsExecSuite` and
`DataFrameSetOperationsSuite`: 477
tests, all passing, plus scalastyle for `catalyst`, `sql` and `sql/Test`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5
🤖 Generated with [Claude Code](https://claude.com/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]
