[SPARK-58594][SQL] Union output partitioning should merge KeyedPartitioning in PartitioningCollection children - #57798
Draft
peter-toth wants to merge 2 commits into
Conversation
…ioning in PartitioningCollection children ### What changes were proposed in this pull request? `UnionExec.outputPartitioning` merges its children's `KeyedPartitioning`s only when every child reports a *single* one. This PR lets that merge accept a `PartitioningCollection` per child. Each child's partitioning is flattened to its `KeyedPartitioning` members, and the merge fires when every child offers at least one. For each of the first child's candidates that has a match in every other child, one merged `KeyedPartitioning` is produced exactly as before. A single survivor is returned bare (today's shape); several are returned as a `PartitioningCollection` of alternatives; none falls through to the co-located pass-through case. `doExecute` is extended accordingly: a `PartitioningCollection` of `KeyedPartitioning`s describes the concatenated layout, so it takes the same `sparkContext.union` arm as a single `KeyedPartitioning` rather than the co-located `SQLPartitioningAwareUnionRDD` one. ### Why are the changes needed? An inner `ShuffledJoin` reports `PartitioningCollection(left, right)`, so a storage-partitioned join inside a `UNION ALL` leg makes that leg report a collection of `KeyedPartitioning`s. Such a child misses the all-children-are-a-single-`KeyedPartitioning` test and is filtered out of the co-located case, so the union falls back to `UnknownPartitioning` and the storage-partitioned join is lost. On the join-union-join shape in the new tests this costs two shuffles. A projection that aliases a partition key to more than one output name (`SELECT id, id AS k`) reports a collection for the same reason. ### Does this PR introduce _any_ user-facing change? No, other than better plans: affected queries lose two shuffles. `spark.sql.unionOutputPartitioning` still gates the whole behaviour. ### How was this patch tested? New tests in `KeyGroupedPartitioningSuite` (union of two storage-partitioned joins; the aliasing shape; a merge that needs `GroupPartitionsExec` because the legs share a key; and an AQE variant) and in `DataFrameSetOperationsSuite` (the merged partitioning's shape). All five fail without the `UnionExec` change. Also ran `DataFrameSetOperationsSuite`, `KeyGroupedPartitioningSuite`, `UnionCodegenSuite`, `EnsureRequirementsSuite`, `PlannerSuite` and `AdaptiveQueryExecSuite`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5)
…d alternatives All members of the merged `PartitioningCollection` share one `partitionKeys` reference, so they describe a single physical layout under different names. `isNarrowed` is a fact about that layout -- whether partitions holding distinct keys in a finer-grained partitioning were collapsed into these -- so it cannot legitimately differ between alternatives, the way `isGrouped` cannot. Computing it per alternative from just the matched `KeyedPartitioning`s let members disagree whenever a child's own candidates disagreed, and a member reading `false` would silently drop `GroupPartitionsExec`'s skew protection. Hoist it beside the merged keys and `isGrouped`, taken across every candidate of every child, which is both uniform by construction and the fail-safe direction. No behaviour change for any reachable input: `projectKeyedPartitionings` already assigns one `isNarrowed` to all the candidates it builds, so the OR equals what the per-alternative computation produced.
peter-toth
marked this pull request as draft
August 5, 2026 19:22
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
UnionExec.outputPartitioningmerges its children'sKeyedPartitionings only when every child reports a single one. This PR lets that merge accept aPartitioningCollectionper child.Each child's partitioning is flattened to its
KeyedPartitioningmembers, and the merge fires when every child offers at least one. For each of the first child's candidates that has a match in every other child, one mergedKeyedPartitioningis produced exactly as before. A single survivor is returned bare (today's shape); several are returned as aPartitioningCollectionof alternatives; none falls through to the co-located pass-through case.doExecuteis extended accordingly: aPartitioningCollectionofKeyedPartitionings describes the concatenated layout, so it takes the samesparkContext.unionarm as a singleKeyedPartitioningrather than the co-locatedSQLPartitioningAwareUnionRDDone.Why are the changes needed?
An inner
ShuffledJoinreportsPartitioningCollection(left, right), so a storage-partitioned join inside aUNION ALLleg makes that leg report a collection ofKeyedPartitionings. Such a child misses the all-children-are-a-single-KeyedPartitioningtest and is filtered out of the co-located case, so the union falls back toUnknownPartitioningand the storage-partitioned join is lost. On the join-union-join shape in the new tests this costs two shuffles.A projection that aliases a partition key to more than one output name (
SELECT id, id AS k) reports a collection for the same reason.Does this PR introduce any user-facing change?
No, other than better plans: affected queries lose two shuffles.
spark.sql.unionOutputPartitioningstill gates the whole behaviour.How was this patch tested?
New tests in
KeyGroupedPartitioningSuite(union of two storage-partitioned joins; the aliasing shape; a merge that needsGroupPartitionsExecbecause the legs share a key; and an AQE variant) and inDataFrameSetOperationsSuite(the merged partitioning's shape). All five fail without theUnionExecchange. Also ranDataFrameSetOperationsSuite,KeyGroupedPartitioningSuite,UnionCodegenSuite,EnsureRequirementsSuite,PlannerSuiteandAdaptiveQueryExecSuite.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)