From 61bb7903d5efcafea8d22a2bd2a1d6021418a437 Mon Sep 17 00:00:00 2001 From: Hemanth Boyina Date: Wed, 2 Sep 2026 15:28:21 +0530 Subject: [PATCH] [SPARK-59141][SQL] Columnar UnionExec ignores outputPartitioning and returns wrong results ### What changes were proposed in this pull request? `UnionExec.doExecuteColumnar` now gets the same partitioning split that `doExecute` already has: a union whose `outputPartitioning` is index-co-locatable builds a `SQLPartitioningAwareUnionRDD` that interleaves same-index partitions across children, and only an `UnknownPartitioning` or `KeyedPartitioning` union concatenates. `SQLPartitioningAwareUnionRDD` is generic over its element type, so it works for `RDD[ColumnarBatch]` with no new RDD type. ### Why are the changes needed? Wrong results on the default configuration. `doExecuteColumnar` concatenated its children unconditionally while `outputPartitioning` kept reporting the partitioning the children agreed on, so a parent that dropped its exchange on the strength of that report read a concatenation instead of an interleaving. A union of two columnar-only bucketed scans reports `HashPartitioning(k, 4)`, which lets a downstream `GROUP BY k` skip its shuffle; the concatenated columnar union then puts each key in two partitions, so the shuffle-free aggregate emits each group twice. The row path has been partitioning-aware since SPARK-52921 (`spark.sql.unionOutputPartitioning`, on by default), so this reaches back to 4.1.0. ### Does this PR introduce _any_ user-facing change? Yes, it fixes wrong results. Any union of columnar children that reports an index-co-locatable partitioning was affected. ### How was this patch tested? A new case in `DataFrameSetOperationsSuite` builds a union of two bucketed (columnar-only) scans feeding a shuffle-free aggregate, asserts the union is columnar-only, reports a `HashPartitioning`, and has no exchange, then compares against the same query with `spark.sql.unionOutputPartitioning` off. It fails with ten rows where five are expected without the fix and passes with it. `DataFrameSetOperationsSuite`, `UnionCodegenSuite`, `BucketedReadWithoutHiveSupportSuite` and `KeyGroupedPartitioningSuite` pass. --- .../execution/basicPhysicalOperators.scala | 12 ++++- .../sql/DataFrameSetOperationsSuite.scala | 49 +++++++++++++++++++ 2 files changed, 60 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 f0186af264d90..4c1006b92985d 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 @@ -1255,7 +1255,17 @@ 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())) + // Mirror `doExecute`: a union whose `outputPartitioning` is index-co-locatable must interleave + // same-index partitions so it actually delivers the partitioning it reports. Only an + // `UnknownPartitioning` or `KeyedPartitioning` union concatenates its children. + outputPartitioning match { + case _: UnknownPartitioning | _: KeyedPartitioning => + sparkContext.union(children.map(_.executeColumnar())) + case _ => + val nonEmptyRdds = children.map(_.executeColumnar()).filter(!_.partitions.isEmpty) + new SQLPartitioningAwareUnionRDD( + sparkContext, nonEmptyRdds, outputPartitioning.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..d6225ea9195c3 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 @@ -2025,6 +2025,55 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP } } + test("SPARK-59141: columnar union interleaves partitions to honor outputPartitioning") { + withSQLConf( + // Keep the bucketed scans co-located instead of collapsing them to a plain scan. + SQLConf.AUTO_BUCKETED_SCAN_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTable("t1", "t2") { + // Bucketed parquet scans are batch-readable, so `FileSourceScanExec` is columnar-only + // (`supportsColumnar` true, `supportsRowBased` false). A union of two of them is therefore + // columnar-only too, which forces the columnar execution path (`doExecuteColumnar`). + 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 sqlText = + "SELECT k, count(*) c FROM (SELECT k FROM t1 UNION ALL SELECT k FROM t2) GROUP BY k" + val union = sql(sqlText) + val plan = union.queryExecution.executedPlan + + val unionExec = plan.collect { case u: UnionExec => u } + assert(unionExec.size == 1) + // The union must run columnar (row path unavailable), so it uses `doExecuteColumnar`. + assert(unionExec.head.supportsColumnar && !unionExec.head.supportsRowBased, + s"expected a columnar-only union but got\n$plan") + // Both children are bucketed by `k` into 4 buckets, so the union reports + // HashPartitioning(k, 4), which is index-co-locatable. + assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning], + s"expected a HashPartitioning pass-through but got " + + s"${unionExec.head.outputPartitioning}\n$plan") + // The aggregate reuses that partitioning, so there is no exchange to re-shuffle the rows; + // correctness then depends entirely on the union actually co-locating same-key partitions. + assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty, + s"group-by should reuse the union partitioning (no shuffle) but got\n$plan") + + // Oracle: the same query with union output partitioning disabled inserts a correct shuffle + // before the aggregate. + val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + sql(sqlText).collect() + } + // Without the fix, `doExecuteColumnar` concatenates the children instead of interleaving + // same-index partitions, so each key is counted once per table and the shuffle-free + // aggregate returns ten rows (each key twice) instead of five. + checkAnswer(union, correctResult) + checkAnswer(union, + Row(0, 8) :: Row(1, 8) :: Row(2, 8) :: Row(3, 8) :: Row(4, 8) :: Nil) + } + } + } + test("SPARK-51262: exceptAll after dropDuplicates with subset should not throw") { // Data where dropDuplicates(subset) produces deterministic results - to avoid test flakiness. val df1 = spark.createDataFrame(Seq(