[SPARK-57881][SQL] UnionExec outputPartitioning supports KeyedPartitioning

[pan3793]

What changes were proposed in this pull request?

Extend spark.sql.unionOutputPartitioning to KeyedPartitioning, so a UNION ALL of V2 storage-partitioned tables can do a shuffle-free SPJ.

Changes in UnionExec:

• comparePartitioning: add a KeyedPartitioning case matching on partition expressions (keys not compared — children carry different key sets that are merged below).

• outputPartitioning: when all children report compatible KeyedPartitionings, emit a merged one whose partitionKeys is the concatenation of the children's keys (one per physical partition), recomputing isGrouped and propagating isNarrowed (sticky: set if any child is narrowed).

• doExecute: a KeyedPartitioning union uses the plain concatenating UnionRDD, not the index-based SQLPartitioningAwareUnionRDD (which is only correct for HashPartitioning/SinglePartition and would mix keys).

The merged descriptor is consumed by the existing SPJ logic in EnsureRequirements: duplicate keys (overlapping child sets, isGrouped=false) are coalesced by GroupPartitionsExec; join-leg key-set mismatches are reconciled via the push-part-values path.

Why are the changes needed?

A UNION ALL of V2 bucketed tables reports UnknownPartitioning today, so any join on the partition key shuffles — defeating SPJ for sharded/bucketed sources.

SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data
FROM (SELECT id, data FROM t1 UNION ALL SELECT id, data FROM t2) u
JOIN t3 ON u.id = t3.id;

Before: both sides shuffle. After: the union reports a merged KeyedPartitioning over id; SPJ runs shuffle-free (with GroupPartitionsExec coalescing overlapping keys).

Does this PR introduce any user-facing change?

Results are unchanged. The executed plan may drop a previously inserted shuffle.

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

Generated-by: GLM 5.2

[pan3793]

cc @peter-toth, could you please take a look?

[peter-toth]

The 3-arg constructor leaves isNarrowed = false, so the merge drops the flag whenever a child KeyedPartitioning is narrowed (e.g. a union leg that is a ProjectExec/aggregate narrowing a composite partition key -- ProjectExec is a PartitioningPreservingUnaryExecNode, basicPhysicalOperators.scala:49). comparePartitioning only matches on expressions, so a narrowed child and a non-narrowed child with the same expressions are still considered compatible and merged here.

AliasAwareOutputExpression.scala:135-139 calls this flag out as needing to be "sticky ... a subsequent [node] that passes all positions through would otherwise recompute isNarrowed=false, silently dropping the protection" -- this merge is exactly that dropping. The consequence is in KeyedPartitioning.groupedSatisfies (partitioning.scala:588-593): a narrowed, non-grouped KP is meant to satisfy ClusteredDistribution only under v2BucketingAllowKeysSubsetOfPartitionKeys, precisely because GroupPartitionsExec will then merge partitions that held distinct keys in the finer partitioning (skew risk). Reporting isNarrowed=false lets that SPJ proceed without the opt-in.

Propagate the flag from the children:

Suggested change

	KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped)
	val isNarrowed = partitionings.exists {
	case k: KeyedPartitioning => k.isNarrowed
	case _ => false
	}
	KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed)

[pan3793]

@peter-toth, thank you for pointing out this correctness issue, applied this fix.

[peter-toth]

All four tests partition by identity("id"). A variant over a bucket transform (Array(bucket(4, "id"))) would cover the PR's motivating bucketed-source case end-to-end.

[pan3793]

added a new test SPARK-57881: storage-partitioned join over union: bucket transform partitioning

[yadavay-amzn]

LGTM, one test-coverage suggestion (non-blocking)

Traced the merge end to end and the core is sound - the partition-index-to-key alignment holds: sparkContext.union concatenates children in order, mergedKeys concatenates partitionKeys in the same child order, and each child's DataSourceRDD keeps one partition per key (pruned splits stay as None placeholders in filteredPartitions, so a child's partition count always equals its partitionKeys.length). So the descriptor and the RDD layout stay aligned. Building on peter-toth's points, not repeating them.

One gap worth closing: none of the added tests exercise a union child whose partitions are all pruned (an empty / None-padded leg), and because every current case is order-insensitive, no test would fail if that index-to-key alignment were ever broken. A case where one union leg's keys are entirely runtime-pruned (or a leg contributes zero live partitions) while the merged KeyedPartitioning still carries its keys would directly guard the concatenation invariant the shuffle-free path depends on.

What I verified:

• mergedKeys (basicPhysicalOperators.scala:957) and sparkContext.union (:1200) both concatenate in children order, so mergedKeys(i) describes union RDD partition i.
• filteredPartitions: Seq[Option[InputPartition]] (BatchScanExec.scala:79) pads pruned splits with None rather than dropping them, so per-child partition count == partitionKeys.length even after runtime filtering - the alignment holds under pruning.
• The KeyedPartitioning doExecute branch correctly avoids SQLPartitioningAwareUnionRDD (index-interleaving would mix keys); plain concat matches the merged descriptor.

[pan3793]

@yadavay-amzn, thank you for the review, I added a new test case SPARK-57881: storage-partitioned join over union: a union leg is entirely runtime-pruned to cover the case you mentioned.

[peter-toth]

LGTM, pending CI.

[yadavay-amzn]

LGTM, thanks @pan3793 !

[pan3793]

Merge Summary:

• merged into master f8ee9ef
• merged into branch-4.x 46c1826

Posted by merge_spark_pr.py

[pan3793]

thank you @peter-toth @yadavay-amzn, merged