[SPARK-59141][SQL] Interleave partitions in columnar UnionExec

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
@@ -1011,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, _)))
@@ -1024,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.
@@ -1234,33 +1235,38 @@ 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 {
+  // 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(children.map(_.execute()))
+        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.
-        val nonEmptyRdds = children.map(_.execute()).filter(!_.partitions.isEmpty)
         new SQLPartitioningAwareUnionRDD(
-          sparkContext, nonEmptyRdds, outputPartitioning.numPartitions)
+          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)

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala

@@ -1772,6 +1772,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-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)")