[SPARK-58317][SQL] Union output partitioning should support PartitioningCollection children

[ulysses-you]

What changes were proposed in this pull request?

UnionExec.outputPartitioning (SPARK-52921) passes a child partitioning through the union when every child is compatibly partitioned, letting a downstream operator (e.g. an aggregate) skip a shuffle. This PR extends that pass-through to children whose outputPartitioning is a PartitioningCollection.

• outputPartitioning now treats each child's partitioning as a set of candidate partitionings (a PartitioningCollection flattens to its members; a single partitioning is a one-element set) and passes through the intersection across all children: empty -> UnknownPartitioning; one member -> that partitioning; many -> a PartitioningCollection. Only index-co-locatable partitionings (HashPartitioningLike / SinglePartition) participate.
• The existing KeyedPartitioning concatenation path is kept as a separate case that fires only when every child is a single KeyedPartitioning. It is a distinct physical strategy (sparkContext.union, numPartitions = sum) and cannot be folded into the co-located intersection (SQLPartitioningAwareUnionRDD, numPartitions = N), because a PartitioningCollection requires uniform numPartitions across its members.
• comparePartitioning is documented as a leaf-only equivalence predicate; collections are flattened before it is called.

Why are the changes needed?

comparePartitioning had no case for PartitioningCollection, so whenever a child reported one, the whole union collapsed to UnknownPartitioning and an extra shuffle was inserted. This is hit by a common query shape - a UNION ALL where one branch is an inner shuffled-hash join (whose outputPartitioning is PartitioningCollection(Hash(k1), Hash(k2))) and the other is a left join (single HashPartitioning), feeding a GROUP BY:

select c1, c2, c3, count(*) from (
  select t1.c1, t1.c2, t1.c3 from t1 join t2 on t1.c1 = t2.c1
  union all
  select t3.c1, t3.c2, t3.c3 from t3 left join t4 on t3.c1 = t4.c1
) group by c1, c2, c3

The union failed to pass through and a redundant ENSURE_REQUIREMENTS shuffle appeared before the aggregate.

Does this PR introduce any user-facing change?

No change in query results. Under the default-on spark.sql.unionOutputPartitioning, affected plans drop a redundant shuffle.

How was this patch tested?

Added tests to DataFrameSetOperationsSuite covering: the mixed inner+left join union (intersection to a single HashPartitioning, shuffle eliminated); all children reporting collections (pass-through as a PartitioningCollection); empty intersection (fall back to UnknownPartitioning); and the collection pass-through under AQE. Existing DataFrameSetOperationsSuite and KeyGroupedPartitioningSuite (97 tests) regressions pass.

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

Generated-by: Claude Code (Opus 4.8)

[ulysses-you]

cc @viirya @dongjoon-hyun @peter-toth thank you

[uros-b]

LGTM, thank you @ulysses-you!

[peter-toth]

Thanks for the PR, @ulysses-you!

This extends the SPARK-52921 union pass-through so a child reporting a PartitioningCollection (e.g. an inner shuffled-hash join's PartitioningCollection(Hash(k1), Hash(k2))) no longer collapses the whole union to UnknownPartitioning: Case B flattens each child to its co-locatable members, intersects across children, and passes through the shared members, while the KeyedPartitioning concatenation stays as a separate Case A.

I traced the correctness-sensitive parts independently and they hold:

• co-location is sound -- comparePartitioning's l == r on HashPartitioningLike matches on expressions and numPartitions, and the intersection requires a match in every child, so all children agree on numPartitions for the co-located SQLPartitioningAwareUnionRDD arm;
• the PartitioningCollection uniform-numPartitions invariant is preserved because the survivors are all drawn from the head child's (already-uniform) member set;
• KeyedPartitioning is correctly kept out of Case B (it isn't a HashPartitioningLike), so its sum-numPartitions concatenation can't leak into the co-located numPartitions = N arm; all-keyed goes through Case A and mixed keyed+hash falls back to UnknownPartitioning.

No findings from me.

[dongjoon-hyun]

This test case seems to pass without this PR, @ulysses-you . Could you double-check this test case?

[ulysses-you]

Good catch, thank you. You are right -- k was pruned, so no PartitioningCollection reached the union and the test passed on master. Fixed in b79f1b2: GROUP BY c1, c2, c3, k keeps the first branch's collection alive, and t4.c1 AS k (the left join's build side) keeps the second branch a single Hash(c1). I verified the revised test now fails on pre-PR code with UnknownPartitioning(0) was not instance of HashPartitioning and passes with the fix.

[dongjoon-hyun]

Since k is unused above the union, ColumnPruning drops it from both branches and each ProjectExec narrows the join's collection to a single HashPartitioning(c1), so no PartitioningCollection reaches the union — this test passes on master without this PR. Referencing k in the group-by keeps the first branch's collection alive; aliasing t4.c1 instead of t3.c1 keeps the second branch a single Hash(c1), so the intersection is single as intended.

Suggested change

	| SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t3.c2, t3.c3, t3.c1 AS k
	| FROM t3 LEFT JOIN t4 ON t3.c1 = t4.c1
	|) GROUP BY c1, c2, c3
	| SELECT /*+ SHUFFLE_HASH(t4) */ t3.c1, t3.c2, t3.c3, t4.c1 AS k
	| FROM t3 LEFT JOIN t4 ON t3.c1 = t4.c1
	|) GROUP BY c1, c2, c3, k

[dongjoon-hyun]

Previously, k is pruned, so no PartitioningCollection reaches the union. Referencing k in the group-by keeps the collection alive under AQE.

Suggested change

	|) GROUP BY c1
	|) GROUP BY c1, k

[ulysses-you]

Applied in b79f1b2 -- SELECT c1, k, ... / GROUP BY c1, k, same reasoning as above so the PartitioningCollection survives ColumnPruning under AQE.

[dongjoon-hyun]

Please try to use the existing one.

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala

Lines 152 to 160 in b65ef3e

	private def flattenPartitioning(partitioning: Partitioning): Seq[Partitioning] = {
	partitioning match {
	case PartitioningCollection(childPartitionings) =>
	childPartitionings.flatMap(flattenPartitioning)
	case rest =>
	rest +: Nil
	}
	}
	}

[ulysses-you]

Done in b79f1b2. Rather than depend on the sql/core helper from catalyst, I lifted a shared PartitioningCollection.flatten into the catalyst companion object and reused it from both UnionExec and PartitioningPreservingUnaryExecNode, removing the duplication.

[dongjoon-hyun]

The main body code looks good to me but I have three comments on the test cases. Could you address them, @ulysses-you ?

[viirya]

I traced the core logic independently and it's correct — the intersection is anchored on the head child's already-uniform candidate set, comparePartitioning matches HashPartitioningLike on both expressions and numPartitions, and a survivor must match in every child, so the co-located union RDD arm and the PartitioningCollection uniform-numPartitions invariant both hold. KeyedPartitioning is correctly kept in Case A and can't leak into Case B. This matches @peter-toth's trace.

That said, I don't think this is ready to merge yet — @dongjoon-hyun's comments are still open and I independently confirmed the most important one:

• The "intersects to single" test (line ~1684) passes on master without this PR, so it doesn't actually exercise the new path. k is unused above the union, so ColumnPruning drops it and each branch's ProjectExec narrows the first join's PartitioningCollection(Hash(c1), Hash(k)) down to a single HashPartitioning(c1) — no PartitioningCollection ever reaches the union, and master's existing all-equal pass-through already yields a HashPartitioning, so the enabled assertion holds on master too. This is a coverage gap (the production logic is correct; the test just isn't hitting Case B's single-intersection path), so it should be fixed before merge. Dongjoon's suggestion is the right fix: reference k in the GROUP BY to keep the first branch's collection alive, and alias t4.c1 AS k in the second branch so it stays a single Hash(c1).
• The AQE test (line ~1827) has the same issue — GROUP BY c1 lets k get pruned, so the collection never reaches the union under AQE either. GROUP BY c1, k fixes it, per Dongjoon's comment.

Non-blocking: the new local flattenPartitioning (basicPhysicalOperators.scala:990) duplicates the identical one in AliasAwareOutputExpression (line 152). That one is private and UnionExec doesn't mix in the trait, so reusing it isn't a one-liner — it'd mean lifting the helper into a shared home (e.g. the PartitioningCollection companion). Worth doing to avoid two copies of the same recursion, but it's a judgment call.

The other two new tests (all-collections pass-through, empty-intersection) do genuinely exercise Case B, and the core change is sound — just the two tests above need to actually cover it.