This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 3e227d64ff7b [SPARK-57996][SQL] CoalesceShufflePartitions should
coalesce partitioning-aware UnionExec children as a single group
3e227d64ff7b is described below
commit 3e227d64ff7be3add4a4e8f7083ac8eac1d6a9d3
Author: Xiduo You <[email protected]>
AuthorDate: Tue Jul 7 21:47:37 2026 -0700
[SPARK-57996][SQL] CoalesceShufflePartitions should coalesce
partitioning-aware UnionExec children as a single group
### 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. `Unio [...]
### 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](https://claude.com/claude-code)
Closes #57077 from ulysses-you/SPARK-57996.
Authored-by: Xiduo You <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
(cherry picked from commit 217fb4e43ddfb0e93903e7081a8fa3e8bef68e85)
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../adaptive/CoalesceShufflePartitions.scala | 6 ++-
.../sql/execution/basicPhysicalOperators.scala | 2 +-
.../adaptive/AdaptiveQueryExecSuite.scala | 47 ++++++++++++++++++++++
3 files changed, 53 insertions(+), 2 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
index d44667ea2fa3..32021da21171 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
@@ -182,7 +182,11 @@ case class CoalesceShufflePartitions(session:
SparkSession) extends AQEShuffleRe
private def childrenNeedCompatiblePartitioning(p: SparkPlan): Boolean = p
match {
// TODO: match more plan nodes here.
- case _: UnionExec => false
+ // A UnionExec that reports a real outputPartitioning (i.e. it is not a
plain union) has that
+ // partitioning relied on by downstream operators, so its children must
remain co-partitioned
+ // -- coalesce them as a single group. When it reports UnknownPartitioning
(config off, or
+ // children not compatibly partitioned), the children may still be
coalesced independently.
+ case u: UnionExec => !u.isPlainUnion
case _: CartesianProductExec => false
case _: BroadcastHashJoinExec => false
case _: BroadcastNestedLoopJoinExec => false
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
index 4e25bf8bef40..79cca1c4cbd8 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -987,7 +987,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
// codegen is disabled for it (`supportCodegenFailureReason` reports
"partitioning-aware"):
// the per-partition key descriptor is consumed by a downstream
`GroupPartitionsExec`, and
// keeping these unions out of whole-stage codegen matches the
`HashPartitioning` union case.
- private def isPlainUnion: Boolean =
outputPartitioning.isInstanceOf[UnknownPartitioning]
+ private[sql] def isPlainUnion: Boolean =
outputPartitioning.isInstanceOf[UnknownPartitioning]
// Per-child projection from the child's output to the union's output. The
wrapped
// child is always the source `Attribute` (deterministic by construction);
the Alias
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 89d0316a0c59..8cf6fbf921da 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -3599,6 +3599,53 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-57996: CoalesceShufflePartitions should coalesce
partitioning-aware UnionExec " +
+ "children as a single group") {
+ // A UNION ALL of two aggregates on the same key, feeding a downstream
aggregate on that key.
+ // With UNION_OUTPUT_PARTITIONING on, the union reports the shared
HashPartitioning that the
+ // outer aggregate relies on. The two children have very different data
sizes, so coalescing
+ // them independently would produce different partition counts, breaking
co-partitioning and
+ // causing AQE to revert the coalescing. They must be coalesced together
as a single group.
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "1048576",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "2",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_SIZE.key -> "1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 USING parquet AS SELECT id AS c1, uuid() AS c2
FROM range(10)")
+ sql("CREATE TABLE t2 USING parquet AS SELECT id AS c1, uuid() AS c2
FROM range(100)")
+
+ val query =
+ """
+ |SELECT c1, c2 FROM (
+ | SELECT c1, c2 FROM t1 GROUP BY c1, c2
+ | UNION ALL
+ | SELECT c1, c2 FROM t2 GROUP BY c1, c2
+ |) GROUP BY c1, c2
+ |""".stripMargin
+
+ val correctResults = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key
-> "false") {
+ sql(query).collect()
+ }
+
+ withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ val df = sql(query)
+ df.collect()
+ val reads = collect(df.queryExecution.executedPlan) {
+ case r: AQEShuffleReadExec => r
+ }
+ // Both union children are coalesced (the plan is not reverted) and
to the same number
+ // of partitions, so they remain co-partitioned.
+ assert(reads.size === 2)
+ assert(reads.forall(_.hasCoalescedPartition))
+ assert(reads.map(_.partitionSpecs.length).distinct.length === 1)
+ checkAnswer(df, correctResults)
+ }
+ }
+ }
+ }
+
test("SPARK-44065: Optimize BroadcastHashJoin skew") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
