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

[LuciferYang]

What changes were proposed in this pull request?

This backports SPARK-59141 (#58445, d5fbdc173bd on master) to branch-4.1.

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:

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 this branch has no KeyedPartitioning; the predicate is the one doExecute already used here. And the helper takes the execution action rather than the executed RDDs, so outputPartitioning is read once, before the children run: doExecute used to test the predicate before executing them and then read outputPartitioning.numPartitions after, which is the torn read master's version removes.

One consequence of that shape is worth naming for anyone diffing against master. On this branch KeyGroupedPartitioning extends HashPartitioningLike, so two children reporting an equal one make the union report it and take the co-located arm, where master routes its successor type KeyedPartitioning to the concatenating arm instead. That is the right reading here: equality on this branch's KeyGroupedPartitioning compares partitionValues positionally, so partition i of each child holds the same key, and index-wise interleaving is what the reported partitioning means. It is also what doExecute has always done for that shape, so the change brings the columnar path in line with the row path rather than inventing a layout.

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.

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 doExecuteColumnar, and the aggregate reads partitions in which a key can appear twice.

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, BucketedReadWithoutHiveSupportSuite, AdaptiveQueryExecSuite, KeyGroupedPartitioningSuite and CoalesceShufflePartitionsSuite run 280 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

[uros-b]

Thank you @LuciferYang and @dongjoon-hyun!

[LuciferYang]

Merge Summary:

• merged into branch-4.1 e8670b9

Posted by merge_spark_pr.py