[SPARK-58594][SQL] Revise KeyedPartitioning handling in UnionExec output partitioning

[peter-toth]

What changes were proposed in this pull request?

UnionExec.outputPartitioning merged its children's KeyedPartitionings only when every child reported a single one and all of them carried semantically equal expressions at equal arity. This replaces that all-or-nothing comparison with a per-position merge, the way SPARK-46367 made PartitioningPreservingUnaryExecNode narrow a projected KeyedPartitioning instead of dropping it.

A new mergeKeyedPartitionings helper does the work. Each child contributes the KeyedPartitionings reachable from its partitioning, so a PartitioningCollection is accepted as well as a single one. A key position of the first child survives when every other child partitions by the same expression, matched on the canonicalized expression so that bucket(4, id) is not confused with years(id) or bucket(8, id). The first child's candidates are then filtered to the admitted namings, narrowed to the surviving positions and deduplicated, and each child's partition keys are projected to its own positions for those expressions before being concatenated in child order. One surviving alternative is returned bare, several as a PartitioningCollection, and none falls through to the co-located pass-through case.

Consequences of merging per position: children that agree on only some of their key positions now merge at a coarser granularity, and children that list the same expressions in a different order merge by permutation, since projectKeys accepts an arbitrary position sequence.

isNarrowed is set when any child contributes fewer distinct key columns than it partitions by, or was already narrowed. Counting distinct positions matters, because two surviving positions can project the same position of another child: that child names one key column with two different expressions across its candidates, which is what the two sides of an equi-join produce. The merged keys stay truthful in that case, since the namings are equal-valued, but that child does lose a key column and the flag has to say so, or groupedSatisfies would accept an ungrouped, narrowed partitioning without spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled and drop the skew protection that config exists to demand. The flag is uniform across the returned alternatives, because it describes the merged layout rather than one naming of it.

The first child's candidate list is filtered rather than rebuilt from the admitted namings. It was already bounded by spark.sql.optimizer.expressionProjectionCandidateLimit when PartitioningPreservingUnaryExecNode produced it, so re-deriving the combinations could resurrect ones that limit had dropped; filtering needs no bound of its own. The trade-off is that when the surviving namings are split across different candidates, no candidate is admitted and the children do not merge.

doExecute is extended too: a PartitioningCollection of KeyedPartitionings describes the concatenated layout, so it takes the same sparkContext.union arm as a single KeyedPartitioning rather than the co-located SQLPartitioningAwareUnionRDD one. A require backstops the invariant that keyed and co-locatable partitionings never share a collection, since a mixed one would have no correct arm. comparePartitioning's KeyedPartitioning branch is removed, as only the pass-through case still uses it.

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 KeyedPartitionings, and a projection that aliases a partition key to several output names does the same. Neither was merged, so the union fell back to UnknownPartitioning and the storage-partitioned join was lost: on the join-union-join shape in the new tests, two shuffles.

Legs that partition by different numbers of columns, or by the same columns in a different order, were not merged either, even though KP([k1, k2]) and KP([k2]) describe a common KP([k2]) once the first leg's keys are projected.

Does this PR introduce any user-facing change?

No, other than better plans. spark.sql.unionOutputPartitioning still gates the behaviour. Note that narrowing only pays off without further configuration when the projected keys stay distinct; when narrowing introduces duplicate keys, groupedSatisfies requires spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled, exactly as for the narrowing projection added by SPARK-46367.

Note also that a union leg only reaches the merge with the partitioning it still has: a partition column pruned from the scan makes V2ScanPartitioningAndOrdering drop the KeyedPartitioning outright, and ColumnPruning pushes through a union, so a leg narrowed because the query selects only the shared column has nothing left to merge. That is the follow-up SPARK-46367 already called out and is unchanged here.

How was this patch tested?

Nine new tests in KeyGroupedPartitioningSuite: a union of two storage-partitioned joins; a leg aliasing a partition key twice; a merge that needs GroupPartitionsExec because the legs share a key; an AQE variant; legs agreeing on fewer key positions with the narrowed keys staying distinct, and with them colliding (parameterized on allowKeysSubsetOfPartitionKeys); legs listing the same key expressions in a different order; three legs with one of them permuted, which pins the per-leg position bookkeeping; and two surviving positions collapsing onto one position of another leg, which pins the isNarrowed accounting (also parameterized on the config). One more in DataFrameSetOperationsSuite covers the merged partitioning's shape. All ten fail on master.

Also ran DataFrameSetOperationsSuite, KeyGroupedPartitioningSuite, UnionCodegenSuite, EnsureRequirementsSuite, PlannerSuite, AdaptiveQueryExecSuite and ProjectedOrderingAndPartitioningSuite.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Opus 5)