[SPARK-57996][SQL] CoalesceShufflePartitions should coalesce partitioning-aware UnionExec children as a single group

[ulysses-you]

What changes were proposed in this pull request?

CoalesceShufflePartitions.collectCoalesceGroups treats every UnionExec as a node whose children do not need compatible partitioning, so each union child becomes an independent coalesce group and may be coalesced to a different number of partitions.

UnionExec.outputPartitioning can report a real partitioning (HashPartitioning / KeyedPartitioning / SinglePartition) when its children are compatibly partitioned. When it does, downstream operators rely on that partitioning, so the children must stay co-partitioned. This PR makes childrenNeedCompatiblePartitioning return true for a UnionExec that is not a plain union (i.e. reports a real partitioning), so its shuffle stages are coalesced together as a single group. UnionExec.isPlainUnion is promoted to private[spark] and reused for this check.

Why are the changes needed?

When a partitioning-aware union's children are coalesced independently, they end up with different partition counts and are no longer co-partitioned. This does not produce a wrong result: AdaptiveSparkPlanExec runs ValidateRequirements after each AQEShuffleReadRule and reverts the entire CoalesceShufflePartitions application when the co-partitioning check fails. The net effect is that shuffle partition coalescing is silently lost whenever a union reports a real partitioning.

Does this PR introduce any user-facing change?

No. Results are unchanged; the executed plan may now coalesce shuffle partitions that were previously left un-coalesced.

How was this patch tested?

New test in AdaptiveQueryExecSuite. Without the fix the coalescing is reverted (no AQEShuffleReadExec); with the fix both union children are coalesced to the same number of partitions.

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

Generated-by: Claude Opus 4.8

🤖 Generated with Claude Code

[ulysses-you]

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

[dongjoon-hyun]

coalesces -> coalesce.

[dongjoon-hyun]

private[spark] -> private[sql] because we prefer to use the narrowest one like the following in this file?

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

Line 1415 in 5a63d3a

private[execution] val executionContext = ExecutionContext.fromExecutorService(

[dongjoon-hyun]

+1, LGTM with two minor nit comments.

[viirya]

The fix looks correct to me. I traced the full mechanism to convince myself it holds end to end: with both children's shuffle stages in one coalesce group, ShufflePartitionsUtil.coalescePartitions hands them identical partition boundaries, so each child reports the same CoalescedHashPartitioning after the read is inserted. UnionExec.comparePartitioning remaps the second child's attributes before the equality check, so the union keeps reporting the real partitioning post-coalescing, the downstream aggregate's distribution requirement stays satisfied, and ValidateRequirements no longer reverts. The fallback paths are also safe: when the group can't be formed (a non-stage leaf under the union, an unsupported shuffle origin, or differing partition counts), the union subtree is simply left un-coalesced and co-partitioning is preserved. Since spark.sql.unionOutputPartitioning defaults to true, the pre-fix rule-level revert -- one partitioning-aware union with unevenly sized children silently discarding the entire plan's coalescing -- hits default configurations, so this is a nice find.

One observation worth recording (non-blocking, no change requested): the new predicate is local -- it checks whether the union reports a real partitioning, not whether anything downstream actually relies on it. When the union output is re-shuffled anyway, the old independent coalescing was legal and would have stuck, while the single-group requirement is stricter and can give up coalescing for the whole union subtree in mixed cases (e.g. one child with a bucketed-scan leaf or an unsupported shuffle inside). That situation is narrow, and reliance can't be determined at this point in the rule, so the conservative choice here is right -- just worth a note in case someone revisits this for finer-grained grouping later.

Also +1 to @dongjoon-hyun's two pending nits (test name grammar, private[sql]).

LGTM otherwise.

[dongjoon-hyun]

Merge Summary:

• merged into master 217fb4e
• merged into branch-4.x 3e227d6

Posted by merge_spark_pr.py

[dongjoon-hyun]

Thank you, @ulysses-you and all.