(spark) branch branch-4.2 updated: [SPARK-59141][4.2][SQL] Interleave partitions in columnar UnionExec

[ya...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

LuciferYang pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.2 by this push:
     new 3c705cbacc03 [SPARK-59141][4.2][SQL] Interleave partitions in columnar UnionExec
3c705cbacc03 is described below

commit 3c705cbacc03ce86be01c74b4f2e0d5cd836e1c8
Author: YangJie <ya...@baidu.com>
AuthorDate: Fri Sep 4 18:55:58 2026 +0800

    [SPARK-59141][4.2][SQL] Interleave partitions in columnar UnionExec
    
    ### What changes were proposed in this pull request?
    
    This backports SPARK-59141 (#58445, `d5fbdc173bd` on master) to branch-4.2.
    
    `UnionExec.doExecuteColumnar` 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, and only an `UnknownPartitioning` union concatenates. Both paths go through one helper rather than two copies of the dispatch, which is what let them drift:
    
    ```scala
    private def unionRDDs[T: ClassTag](executeChild: SparkPlan => RDD[T]): RDD[T] = {
      val partitioning = outputPartitioning
      val rdds = children.map(executeChild)
      if (partitioning.isInstanceOf[UnknownPartitioning]) {
        sparkContext.union(rdds)
      } else {
        new SQLPartitioningAwareUnionRDD(
          sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions)
      }
    }
    ```
    
    Two differences from the master commit, both forced by this branch. The arm is an `if` on `UnknownPartitioning` rather than a `match`, because `KeyedPartitioning` does not participate in `UnionExec.outputPartitioning` here: 4.2's `KeyedPartitioning` is not a `HashPartitioningLike`, so `comparePartitioning` rejects it and such a union reports `UnknownPartitioning`, taking the same concatenating arm master routes it to explicitly. And the helper takes the execution action rather than th [...]
    
    ### Why are the changes needed?
    
    Wrong results on default configuration. `doExecuteColumnar` concatenated its children unconditionally, while `outputPartitioning` went on reporting whatever the children agreed on, so a parent that dropped its exchange on the strength of that report read a concatenation instead of an interleaving.
    
    ```scala
    spark.conf.set("spark.sql.sources.bucketing.autoBucketedScanEnabled", "false")
    spark.conf.set("spark.sql.adaptive.enabled", "false")
    
    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")
    
    sql("SELECT k, count(*) c FROM (SELECT k FROM t1 UNION ALL SELECT k FROM t2) GROUP BY k").show()
    ```
    
    Five groups of eight come back as ten rows of four, each group reported twice. There is no exchange under the aggregate, because the union reported `hashpartitioning(k, 4)` from its two bucketed children. `FileSourceScanExec` overrides `supportsColumnar` but not `supportsRowBased`, whose default is its negation, so a batch-readable bucketed scan is columnar-only and so is a union of two of them. `ColumnarToRowExec` therefore goes above the union rather than below it, the union runs `d [...]
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes. The query above returns five rows instead of ten. Any union of columnar-only children that reports an index-co-locatable partitioning was affected the same way.
    
    A co-located columnar union now has as many partitions as it reports rather than the sum of its children's, so the UNION ALL of two 4-bucket tables above runs 4 tasks instead of 8, and each task reads one bucket from each table. The row path has behaved this way since SPARK-52921, but a batch-readable bucketed scan is columnar-only, so bucketed file scans never reached it. Setting `spark.sql.unionOutputPartitioning` to false restores the old layout.
    
    ### How was this patch tested?
    
    The new case from the master commit, ported unchanged into `DataFrameSetOperationsSuite` next to the existing SPARK-52921 union partitioning cases. It builds the shape above, asserts that the union is columnar-only, that it reports a `HashPartitioning`, and that the aggregate has no exchange, then compares against pinned expected rows. Reverting `doExecuteColumnar` to this branch's unconditional `sparkContext.union` makes it fail.
    
    `DataFrameSetOperationsSuite`, `UnionCodegenSuite`, `KeyGroupedPartitioningSuite` and `BucketedReadWithoutHiveSupportSuite` run 231 cases, all passing, and `sql/scalastyle` and `sql/Test/scalastyle` are clean.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Opus 5
    
    Closes #58511 from LuciferYang/SPARK-59141-4.2.
    
    Authored-by: YangJie <ya...@baidu.com>
    Signed-off-by: yangjie01 <ya...@baidu.com>
---
 .../sql/execution/basicPhysicalOperators.scala     | 23 +++++++++-----
 .../spark/sql/DataFrameSetOperationsSuite.scala    | 36 +++++++++++++++++++++-
 2 files changed, 50 insertions(+), 9 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 8c96f1ff9579..307e2f23e9b7 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, Partition, SparkContext, SparkException, TaskContext}
 import org.apache.spark.internal.LogKeys
@@ -1159,25 +1160,31 @@ 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] = {
-    if (isPlainUnion) {
-      sparkContext.union(children.map(_.execute()))
+  // 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)
+    if (partitioning.isInstanceOf[UnknownPartitioning]) {
+      sparkContext.union(rdds)
     } else {
       // This union has a known 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)
+      new SQLPartitioningAwareUnionRDD(
+        sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions)
     }
   }
 
+  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] = {
-    sparkContext.union(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 d838ba4c234f..c0cce9291dd9 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
@@ -22,7 +22,7 @@ import java.util.Locale
 
 import org.apache.spark.sql.catalyst.optimizer.RemoveNoopUnion
 import org.apache.spark.sql.catalyst.plans.logical.Union
-import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning
+import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, UnknownPartitioning}
 import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
@@ -1617,6 +1617,40 @@ 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 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
+        // 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")
+
+        // 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)))
+      }
+    }
+  }
+
   test("SPARK-53550: union partitioning should compare canonicalized attributes") {
     withSQLConf(
       SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org