From 7926ff680c2472bf03a44be0f5b37d1fe0650160 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 1 Sep 2026 16:44:22 +0800 Subject: [PATCH 1/3] [SPARK-59141][SQL] Interleave partitions in columnar UnionExec `UnionExec.doExecuteColumnar` concatenated its children unconditionally while `outputPartitioning` reported whatever the children agreed on, so a parent that dropped its exchange on the strength of that report read a concatenation. Two bucketed parquet tables unioned under an aggregate return each group twice on default configuration. Give `doExecuteColumnar` the same split `doExecute` has; `SQLPartitioningAwareUnionRDD` is generic, so no new RDD is needed. --- .../execution/basicPhysicalOperators.scala | 11 +++++- .../sql/DataFrameSetOperationsSuite.scala | 36 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) 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 bea86501e6f3a..35f41d2e0c23e 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 @@ -1259,7 +1259,16 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup override def supportsRowBased: Boolean = children.forall(_.supportsRowBased) protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { - sparkContext.union(children.map(_.executeColumnar())) + // Same split as `doExecute`: a union that reports an index-co-locatable partitioning has to + // interleave same-index partitions, or a parent that skipped an exchange on that report reads + // a concatenation instead. + outputPartitioning match { + case _: UnknownPartitioning | _: KeyedPartitioning => + sparkContext.union(children.map(_.executeColumnar())) + case partitioning => + val nonEmptyRdds = children.map(_.executeColumnar()).filter(!_.partitions.isEmpty) + new SQLPartitioningAwareUnionRDD(sparkContext, nonEmptyRdds, partitioning.numPartitions) + } } override protected def withNewChildrenInternal(newChildren: IndexedSeq[SparkPlan]): UnionExec = diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index 32b6da3a30b15..d66dcafe4641e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -1772,6 +1772,42 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP } } + test("SPARK-59141: columnar union interleaves the partitions it reports as co-located") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false") { + withTable("t1", "t2") { + spark.range(0, 20, 1, 1).selectExpr("id % 5 AS k") + .write.bucketBy(4, "k").saveAsTable("t1") + spark.range(20, 40, 1, 1).selectExpr("id % 5 AS k") + .write.bucketBy(4, "k").saveAsTable("t2") + + val query = + "SELECT k, count(*) c FROM (SELECT k FROM t1 UNION ALL SELECT k FROM t2) GROUP BY k" + val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + sql(query).collect() + } + + val df = sql(query) + val plan = df.queryExecution.executedPlan + // `FileSourceScanExec` overrides `supportsColumnar` but not `supportsRowBased`, whose + // default is its negation, so a bucketed batch-readable scan is columnar-only and so is + // this union. It therefore runs `doExecuteColumnar` under a `ColumnarToRowExec` while + // reporting the bucketed `HashPartitioning` that lets the aggregate drop its exchange. + val unionExec = plan.collect { case u: UnionExec => u } + assert(unionExec.size == 1) + assert(unionExec.head.supportsColumnar && !unionExec.head.supportsRowBased, + "this shape must take the columnar path, or the test exercises nothing") + assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning], + "this shape must report a co-locatable partitioning, or the test exercises nothing") + assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the aggregate's exchange must have been dropped, or the test exercises nothing") + + checkAnswer(df, correctResult.toImmutableArraySeq) + } + } + } + test("SPARK-57881: union partitioning - keyed partitioning") { withSQLConf("spark.sql.catalog.testcat" -> classOf[InMemoryCatalog].getName) { sql("CREATE TABLE testcat.ns.t1 (id bigint, data string) PARTITIONED BY (id)") From 215ef11d49a83fd147bc1fd3b1b2a2f6551ba586 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Wed, 2 Sep 2026 23:50:20 +0800 Subject: [PATCH 2/3] Share the partitioning split between the row and columnar paths --- .../execution/basicPhysicalOperators.scala | 53 ++++++++----------- 1 file changed, 23 insertions(+), 30 deletions(-) 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 35f41d2e0c23e..552beed1e36dc 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 @@ -23,6 +23,7 @@ import java.util.concurrent.TimeUnit._ import scala.collection.mutable import scala.concurrent.ExecutionContext import scala.concurrent.duration.Duration +import scala.reflect.ClassTag import org.apache.spark.{InterruptibleIterator, SparkException, TaskContext} import org.apache.spark.internal.LogKeys @@ -1234,42 +1235,34 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // decides which output columns to materialize. override def usedInputs: AttributeSet = AttributeSet.empty - protected override def doExecute(): RDD[InternalRow] = { - outputPartitioning match { - case _: UnknownPartitioning | _: KeyedPartitioning => - // An `UnknownPartitioning` union simply concatenates its children. A - // `KeyedPartitioning` union does the same: its merged partition keys describe the - // concatenated layout (one key per physical partition), and a downstream - // `GroupPartitionsExec` regroups partitions that share a key. This differs from an - // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware - // union RDD interleaves same-index partitions across children. - sparkContext.union(children.map(_.execute())) - case _ => - // This union has a known, index-co-locatable partitioning, i.e., its children have the - // same partitioning in semantics so this union can choose not to change the partitioning - // by using a custom partitioning aware union RDD. - val nonEmptyRdds = children.map(_.execute()).filter(!_.partitions.isEmpty) - new SQLPartitioningAwareUnionRDD( - sparkContext, nonEmptyRdds, outputPartitioning.numPartitions) - } + // Shared by both execution paths so they cannot report one partitioning and build another. The + // split lived only in `doExecute` before, which is how `doExecuteColumnar` came to concatenate + // while this node advertised its children's `HashPartitioning`. + private def unionRDDs[T: ClassTag](rdds: Seq[RDD[T]]): RDD[T] = outputPartitioning match { + case _: UnknownPartitioning | _: KeyedPartitioning => + // An `UnknownPartitioning` union simply concatenates its children. A + // `KeyedPartitioning` union does the same: its merged partition keys describe the + // concatenated layout (one key per physical partition), and a downstream + // `GroupPartitionsExec` regroups partitions that share a key. This differs from an + // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware + // union RDD interleaves same-index partitions across children. + sparkContext.union(rdds) + case partitioning => + // This union has a known, index-co-locatable partitioning, i.e., its children have the + // same partitioning in semantics so this union can choose not to change the partitioning + // by using a custom partitioning aware union RDD. + new SQLPartitioningAwareUnionRDD( + sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions) } + protected override def doExecute(): RDD[InternalRow] = unionRDDs(children.map(_.execute())) + override def supportsColumnar: Boolean = children.forall(_.supportsColumnar) override def supportsRowBased: Boolean = children.forall(_.supportsRowBased) - protected override def doExecuteColumnar(): RDD[ColumnarBatch] = { - // Same split as `doExecute`: a union that reports an index-co-locatable partitioning has to - // interleave same-index partitions, or a parent that skipped an exchange on that report reads - // a concatenation instead. - outputPartitioning match { - case _: UnknownPartitioning | _: KeyedPartitioning => - sparkContext.union(children.map(_.executeColumnar())) - case partitioning => - val nonEmptyRdds = children.map(_.executeColumnar()).filter(!_.partitions.isEmpty) - new SQLPartitioningAwareUnionRDD(sparkContext, nonEmptyRdds, partitioning.numPartitions) - } - } + protected override def doExecuteColumnar(): RDD[ColumnarBatch] = + unionRDDs(children.map(_.executeColumnar())) override protected def withNewChildrenInternal(newChildren: IndexedSeq[SparkPlan]): UnionExec = copy(children = newChildren) From 8eae58a878b43754989de75e67258104a2a2dc15 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 3 Sep 2026 01:33:16 +0800 Subject: [PATCH 3/3] Read outputPartitioning before the children execute, and pin the test's expected rows --- .../execution/basicPhysicalOperators.scala | 52 ++++++++++--------- .../sql/DataFrameSetOperationsSuite.scala | 14 +++-- 2 files changed, 34 insertions(+), 32 deletions(-) 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 552beed1e36dc..83fd1863e1508 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 @@ -1012,7 +1012,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // Intersect across all children, anchored on the first child's set. Every surviving member // is shared by all children, so their `numPartitions` agree; a `PartitioningCollection` // built from a subset of one child's members therefore keeps its uniform-numPartitions - // invariant, and the co-located `doExecute` arm's invariant holds. + // invariant, and the co-located arm in `unionRDDs` keeps its invariant. val head = candidateSets.head val intersection = head.filter { c => candidateSets.tail.forall(_.exists(comparePartitioning(c, _))) @@ -1025,8 +1025,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } // True when the codegen path applies: `outputPartitioning` is `UnknownPartitioning`, - // and `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `doExecute`. - // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `doExecute`, but + // and `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `unionRDDs`. + // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `unionRDDs`, but // 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. @@ -1235,34 +1235,38 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // decides which output columns to materialize. override def usedInputs: AttributeSet = AttributeSet.empty - // Shared by both execution paths so they cannot report one partitioning and build another. The - // split lived only in `doExecute` before, which is how `doExecuteColumnar` came to concatenate - // while this node advertised its children's `HashPartitioning`. - private def unionRDDs[T: ClassTag](rdds: Seq[RDD[T]]): RDD[T] = outputPartitioning match { - case _: UnknownPartitioning | _: KeyedPartitioning => - // An `UnknownPartitioning` union simply concatenates its children. A - // `KeyedPartitioning` union does the same: its merged partition keys describe the - // concatenated layout (one key per physical partition), and a downstream - // `GroupPartitionsExec` regroups partitions that share a key. This differs from an - // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware - // union RDD interleaves same-index partitions across children. - sparkContext.union(rdds) - case partitioning => - // This union has a known, index-co-locatable partitioning, i.e., its children have the - // same partitioning in semantics so this union can choose not to change the partitioning - // by using a custom partitioning aware union RDD. - new SQLPartitioningAwareUnionRDD( - sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions) + // Shared by `doExecute` and `doExecuteColumnar` so the two cannot report one partitioning and + // build another. `outputPartitioning` is read once, before the children execute: a child's + // partitioning can sharpen once it has run, as `InMemoryTableScanExec` does over a + // not-yet-materialized AQE cached plan. + private def unionRDDs[T: ClassTag](executeChild: SparkPlan => RDD[T]): RDD[T] = { + val partitioning = outputPartitioning + val rdds = children.map(executeChild) + partitioning match { + case _: UnknownPartitioning | _: KeyedPartitioning => + // An `UnknownPartitioning` union simply concatenates its children. A + // `KeyedPartitioning` union does the same: its merged partition keys describe the + // concatenated layout (one key per physical partition), and a downstream + // `GroupPartitionsExec` regroups partitions that share a key. This differs from an + // index-co-locatable partitioning (e.g. `HashPartitioning`), where a partitioning-aware + // union RDD interleaves same-index partitions across children. + sparkContext.union(rdds) + case _ => + // This union has a known, index-co-locatable partitioning, i.e., its children have the + // same partitioning in semantics so this union can choose not to change the partitioning + // by using a custom partitioning aware union RDD. + new SQLPartitioningAwareUnionRDD( + sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions) + } } - protected override def doExecute(): RDD[InternalRow] = unionRDDs(children.map(_.execute())) + protected override def doExecute(): RDD[InternalRow] = unionRDDs(_.execute()) override def supportsColumnar: Boolean = children.forall(_.supportsColumnar) override def supportsRowBased: Boolean = children.forall(_.supportsRowBased) - protected override def doExecuteColumnar(): RDD[ColumnarBatch] = - unionRDDs(children.map(_.executeColumnar())) + protected override def doExecuteColumnar(): RDD[ColumnarBatch] = unionRDDs(_.executeColumnar()) override protected def withNewChildrenInternal(newChildren: IndexedSeq[SparkPlan]): UnionExec = copy(children = newChildren) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index d66dcafe4641e..4cce6129dc5ac 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -1782,13 +1782,8 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP spark.range(20, 40, 1, 1).selectExpr("id % 5 AS k") .write.bucketBy(4, "k").saveAsTable("t2") - val query = - "SELECT k, count(*) c FROM (SELECT k FROM t1 UNION ALL SELECT k FROM t2) GROUP BY k" - val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { - sql(query).collect() - } - - val df = sql(query) + val df = sql( + "SELECT k, count(*) c FROM (SELECT k FROM t1 UNION ALL SELECT k FROM t2) GROUP BY k") val plan = df.queryExecution.executedPlan // `FileSourceScanExec` overrides `supportsColumnar` but not `supportsRowBased`, whose // default is its negation, so a bucketed batch-readable scan is columnar-only and so is @@ -1803,7 +1798,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty, "the aggregate's exchange must have been dropped, or the test exercises nothing") - checkAnswer(df, correctResult.toImmutableArraySeq) + // Interleaving puts a key's bucket from each table in one partition, so its eight rows + // land together; concatenating would split them and the exchange-free aggregate would + // report the key twice with four. + checkAnswer(df, (0L until 5L).map(k => Row(k, 8L))) } } }