[SPARK-59141][SQL] Interleave partitions in columnar UnionExec - #58445
[SPARK-59141][SQL] Interleave partitions in columnar UnionExec#58445LuciferYang wants to merge 3 commits into
Conversation
`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.
|
Thank you @LuciferYang! The fix is minimal and correct. It mirrors doExecute's partitioning-aware branch exactly, uses SQLPartitioningAwareUnionRDD[T: ClassTag] generically (no new RDD type needed), and avoids the double-read of outputPartitioning that the sibling PR has to patch in the row path. The test anchors the invariants that make the fix meaningful (columnar-only path taken, co-locatable partitioning reported, exchange dropped) and compares against the result from the non-partitioning-aware path. No API, config, version, or MiMa concerns. |
| // 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 { |
There was a problem hiding this comment.
This is a copy of doExecute's dispatch with execute() swapped for executeColumnar(), and the bug this PR fixes is exactly the two copies drifting: SPARK-52921 added the split only to doExecute while doExecuteColumnar already existed, and SPARK-57881 added the KeyedPartitioning arm only there. Since both SQLPartitioningAwareUnionRDD and SparkContext.union are generic over T: ClassTag, how about one helper shared by both paths so they cannot drift again?
private def unionRDDs[T: ClassTag](rdds: Seq[RDD[T]]): RDD[T] = outputPartitioning match {
case _: UnknownPartitioning | _: KeyedPartitioning =>
// (existing comment from doExecute)
sparkContext.union(rdds)
case partitioning =>
// (existing comment from doExecute)
new SQLPartitioningAwareUnionRDD(
sparkContext, rdds.filter(!_.partitions.isEmpty), partitioning.numPartitions)
}
protected override def doExecute(): RDD[InternalRow] = unionRDDs(children.map(_.execute()))
protected override def doExecuteColumnar(): RDD[ColumnarBatch] =
unionRDDs(children.map(_.executeColumnar()))It needs import scala.reflect.ClassTag, and the "Same split as doExecute" comment can go away since there is only one split left.
There was a problem hiding this comment.
Done in 215ef11 and 8eae58a. The helper takes the execution action rather than the already-executed RDDs:
private def unionRDDs[T: ClassTag](executeChild: SparkPlan => RDD[T]): RDD[T] = {
val partitioning = outputPartitioning
val rdds = children.map(executeChild)
partitioning match { ... }
}
protected override def doExecute(): RDD[InternalRow] = unionRDDs(_.execute())
protected override def doExecuteColumnar(): RDD[ColumnarBatch] = unionRDDs(_.executeColumnar())One difference from the snippet worth flagging: a strict Seq[RDD[T]] parameter evaluates children.map(_.execute()) before the match, so outputPartitioning ends up being read after the children have run, which is not what doExecute did before. A child's partitioning can sharpen once it has executed, because InMemoryTableScanExec.cachedPlan unwraps an inner AdaptiveSparkPlanExec only while isFinalPlan, and materializing the cache is what sets that. Reading the partitioning into a val first keeps the arm and its numPartitions on one read, and it also removes the second read the old co-located arm did after executing the children.
DataFrameSetOperationsSuite, UnionCodegenSuite, AdaptiveQueryExecSuite, KeyGroupedPartitioningSuite, CoalesceShufflePartitionsSuite and BucketedReadWithoutHiveSupportSuite run 400 tests, all passing. The new test now pins its expected rows instead of re-running the query with spark.sql.unionOutputPartitioning=false, and I confirmed it fails without the fix.
dongjoon-hyun
left a comment
dongjoon-hyun
left a comment
There was a problem hiding this comment.
+1, LGTM. Thank you, @LuciferYang. The shared unionRDDs helper and the point about reading outputPartitioning before the children execute both look right to me.
One small request before merging: could you refresh the PR description to match the latest commits? The "What changes" section does not mention the shared helper, and "How was this patch tested" still says the new test compares against the same query with spark.sql.unionOutputPartitioning off, while the test now pins its expected rows.
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
LGTM. I verified the premise that makes this a real default-config bug rather than a theoretical one: SparkPlan.supportsRowBased is !supportsColumnar, FileSourceScanExec overrides only supportsColumnar (DataSourceScanExec.scala:735), so a batched bucketed scan is columnar-only, and UnionExec.supportsRowBased = children.forall(...) makes the union columnar-only too -- ColumnarToRowExec lands above it, doExecuteColumnar runs, and the old unconditional sparkContext.union concatenated partitions the union had just advertised as index-co-located. Sharing one split between the two paths is the right fix; the alternative of reporting UnknownPartitioning whenever a union might run columnar would give up SPARK-52921's exchange elimination for the many unions that end up on the row path, since EnsureRequirements reads outputPartitioning long before columnar transitions are decided.
Two things about the helper are better than they may look. Taking the execution action rather than a strict Seq[RDD[T]] is load-bearing: InMemoryTableScanExec.cachedPlan unwraps an inner AdaptiveSparkPlanExec only while isFinalPlan (:104-105), and materializing the cache is what makes it final, so a strict parameter would have moved the outputPartitioning read to after the children run and changed doExecute's existing behavior. Reading it into a val also closes a pre-existing gap in the old co-located arm, which read outputPartitioning twice -- once for the match, once for numPartitions after the children had executed -- so the arm choice and its partition count could in principle come from two different reads.
On the "two copies drift" concern that motivated the shared helper, I swept the other operators defining both doExecute and doExecuteColumnar. The closest analogue is GroupPartitionsExec, whose row path has a SortedMergeCoalescedRDD arm that the columnar path lacks -- but that one is safe, and by the opposite means: it narrows supportsColumnar to exclude exactly the sorted-merge case (:264-265), so the columnar path can never be reached in a state where it would silently drop the k-way merge. Which is a decent argument that UnionExec had to fix the execution side instead, since it cannot narrow supportsColumnar without losing the exchange elimination.
The test's three guard assertions are what make it durable -- pinning that the shape really is columnar-only, really reports a HashPartitioning, and really has no exchange, each with the "or the test exercises nothing" note. Without those it could keep passing while silently testing nothing if FileSourceScanExec's row support or the bucketing behavior ever changes. Pinning the expected rows in the last commit is also an improvement over comparing against unionOutputPartitioning=false, which would agree if both paths broke together.
Only outstanding item is the description refresh @dongjoon-hyun already asked for, which still reads as before: it doesn't mention the shared unionRDDs helper (now the main structural change) and still says the test compares against spark.sql.unionOutputPartitioning off. Worth doing before merge rather than after, since dev/merge_spark_pr.py appends the body to the commit message, so it becomes the permanent record in git history.
The PR description has been updated. |
|
Merge Summary:
Posted by |
### What changes were proposed in this pull request?
`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` or `KeyedPartitioning` union concatenates. `SQLPartitioningAwareUnionRDD` is generic over its element type, so no new RDD is needed.
Both paths go through one helper rather than two copies of the dispatch, which is what let them drift apart:
```scala
private def unionRDDs[T: ClassTag](executeChild: SparkPlan => RDD[T]): RDD[T] = {
val partitioning = outputPartitioning
val rdds = children.map(executeChild)
partitioning match { ... }
}
protected override def doExecute(): RDD[InternalRow] = unionRDDs(_.execute())
protected override def doExecuteColumnar(): RDD[ColumnarBatch] = unionRDDs(_.executeColumnar())
```
The helper takes the execution action rather than a strict `Seq[RDD[T]]`, so `outputPartitioning` is read before the children execute, which is where `doExecute` read it before. A child's partitioning can sharpen once it has run: `InMemoryTableScanExec.cachedPlan` unwraps an inner `AdaptiveSparkPlanExec` only while `isFinalPlan`, and materializing the cache is what sets that. Reading it into a val also drops the second read the old co-located arm did after executing the children, so the arm and its `numPartitions` now come from one read.
### 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. The plan shows why:
```
*(1) HashAggregate(keys=[k#8L], functions=[count(1)])
+- *(1) ColumnarToRow
+- Union
:- FileScan parquet t1[k#8L] Batched: true, Bucketed: true, SelectedBucketsCount: 4 out of 4
+- FileScan parquet t2[k#10L] Batched: true, Bucketed: true, SelectedBucketsCount: 4 out of 4
```
There is no exchange under the aggregate, because the union reported `hashpartitioning(k#8L, 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.
Fixing the other side, by reporting `UnknownPartitioning` whenever the union might run columnar, does not work as well. `EnsureRequirements` reads `outputPartitioning` at `QueryExecution.preparations`, well before `ApplyColumnarRulesAndInsertTransitions` decides anything, so at that point the only available proxy is `supportsColumnar`. That is true for unions that can run either way, such as one over two vectorized cached scans, since `InMemoryTableScanExec` reports both row and columnar support. Using it would give up SPARK-52921's exchange elimination for all of them.
The row path has been partitioning-aware since SPARK-52921, which is where `spark.sql.unionOutputPartitioning` came from, so this reaches back to 4.1.0.
### 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?
A new case in `DataFrameSetOperationsSuite`, alongside 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 the unconditional `sparkContext.union` makes it fail.
`DataFrameSetOperationsSuite`, `UnionCodegenSuite`, `AdaptiveQueryExecSuite`, `KeyGroupedPartitioningSuite`, `CoalesceShufflePartitionsSuite` and `BucketedReadWithoutHiveSupportSuite` run 400 cases, all passing. `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 #58445 from LuciferYang/SPARK-59141.
Authored-by: YangJie <yangjie01@baidu.com>
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
(cherry picked from commit f04a15f)
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
### What changes were proposed in this pull request?
`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` or `KeyedPartitioning` union concatenates. `SQLPartitioningAwareUnionRDD` is generic over its element type, so no new RDD is needed.
Both paths go through one helper rather than two copies of the dispatch, which is what let them drift apart:
```scala
private def unionRDDs[T: ClassTag](executeChild: SparkPlan => RDD[T]): RDD[T] = {
val partitioning = outputPartitioning
val rdds = children.map(executeChild)
partitioning match { ... }
}
protected override def doExecute(): RDD[InternalRow] = unionRDDs(_.execute())
protected override def doExecuteColumnar(): RDD[ColumnarBatch] = unionRDDs(_.executeColumnar())
```
The helper takes the execution action rather than a strict `Seq[RDD[T]]`, so `outputPartitioning` is read before the children execute, which is where `doExecute` read it before. A child's partitioning can sharpen once it has run: `InMemoryTableScanExec.cachedPlan` unwraps an inner `AdaptiveSparkPlanExec` only while `isFinalPlan`, and materializing the cache is what sets that. Reading it into a val also drops the second read the old co-located arm did after executing the children, so the arm and its `numPartitions` now come from one read.
### 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. The plan shows why:
```
*(1) HashAggregate(keys=[k#8L], functions=[count(1)])
+- *(1) ColumnarToRow
+- Union
:- FileScan parquet t1[k#8L] Batched: true, Bucketed: true, SelectedBucketsCount: 4 out of 4
+- FileScan parquet t2[k#10L] Batched: true, Bucketed: true, SelectedBucketsCount: 4 out of 4
```
There is no exchange under the aggregate, because the union reported `hashpartitioning(k#8L, 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.
Fixing the other side, by reporting `UnknownPartitioning` whenever the union might run columnar, does not work as well. `EnsureRequirements` reads `outputPartitioning` at `QueryExecution.preparations`, well before `ApplyColumnarRulesAndInsertTransitions` decides anything, so at that point the only available proxy is `supportsColumnar`. That is true for unions that can run either way, such as one over two vectorized cached scans, since `InMemoryTableScanExec` reports both row and columnar support. Using it would give up SPARK-52921's exchange elimination for all of them.
The row path has been partitioning-aware since SPARK-52921, which is where `spark.sql.unionOutputPartitioning` came from, so this reaches back to 4.1.0.
### 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?
A new case in `DataFrameSetOperationsSuite`, alongside 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 the unconditional `sparkContext.union` makes it fail.
`DataFrameSetOperationsSuite`, `UnionCodegenSuite`, `AdaptiveQueryExecSuite`, `KeyGroupedPartitioningSuite`, `CoalesceShufflePartitionsSuite` and `BucketedReadWithoutHiveSupportSuite` run 400 cases, all passing. `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 #58445 from LuciferYang/SPARK-59141.
Authored-by: YangJie <yangjie01@baidu.com>
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
(cherry picked from commit f04a15f)
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
|
Could you make backporting PRs for branch-4.2 and branch-4.1 in order to make it sure to pass the CIs, please, @LuciferYang ? |
ok |
### 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 the executed RDDs, so `outputPartitioning` is read once, before the children run: `doExecute` used to evaluate the predicate before executing them and then read `outputPartitioning.numPartitions` after, which is the torn read master's version removes. `isPlainUnion` is untouched and still gates whole-stage codegen. ### 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 `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`, `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 <yangjie01@baidu.com> Signed-off-by: yangjie01 <yangjie01@baidu.com>
### 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: ```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 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. ```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 `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 Closes #58512 from LuciferYang/SPARK-59141-4.1. Authored-by: YangJie <yangjie01@baidu.com> Signed-off-by: yangjie01 <yangjie01@baidu.com>
4 participants
What changes were proposed in this pull request?
UnionExec.doExecuteColumnargets the same partitioning split thatdoExecutealready has: a union whoseoutputPartitioningis index-co-locatable builds aSQLPartitioningAwareUnionRDDthat interleaves same-index partitions, and only anUnknownPartitioningorKeyedPartitioningunion concatenates.SQLPartitioningAwareUnionRDDis generic over its element type, so no new RDD is needed.Both paths go through one helper rather than two copies of the dispatch, which is what let them drift apart:
The helper takes the execution action rather than a strict
Seq[RDD[T]], sooutputPartitioningis read before the children execute, which is wheredoExecuteread it before. A child's partitioning can sharpen once it has run:InMemoryTableScanExec.cachedPlanunwraps an innerAdaptiveSparkPlanExeconly whileisFinalPlan, and materializing the cache is what sets that. Reading it into a val also drops the second read the old co-located arm did after executing the children, so the arm and itsnumPartitionsnow come from one read.Why are the changes needed?
Wrong results on default configuration.
doExecuteColumnarconcatenated its children unconditionally, whileoutputPartitioningwent 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.Five groups of eight come back as ten rows of four, each group reported twice. The plan shows why:
There is no exchange under the aggregate, because the union reported
hashpartitioning(k#8L, 4)from its two bucketed children.FileSourceScanExecoverridessupportsColumnarbut notsupportsRowBased, whose default is its negation, so a batch-readable bucketed scan is columnar-only and so is a union of two of them.ColumnarToRowExectherefore goes above the union rather than below it, the union runsdoExecuteColumnar, and the aggregate reads partitions in which a key can appear twice.Fixing the other side, by reporting
UnknownPartitioningwhenever the union might run columnar, does not work as well.EnsureRequirementsreadsoutputPartitioningatQueryExecution.preparations, well beforeApplyColumnarRulesAndInsertTransitionsdecides anything, so at that point the only available proxy issupportsColumnar. That is true for unions that can run either way, such as one over two vectorized cached scans, sinceInMemoryTableScanExecreports both row and columnar support. Using it would give up SPARK-52921's exchange elimination for all of them.The row path has been partitioning-aware since SPARK-52921, which is where
spark.sql.unionOutputPartitioningcame from, so this reaches back to 4.1.0.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.unionOutputPartitioningto false restores the old layout.How was this patch tested?
A new case in
DataFrameSetOperationsSuite, alongside the existing SPARK-52921 union partitioning cases. It builds the shape above, asserts that the union is columnar-only, that it reports aHashPartitioning, and that the aggregate has no exchange, then compares against pinned expected rows. RevertingdoExecuteColumnarto the unconditionalsparkContext.unionmakes it fail.DataFrameSetOperationsSuite,UnionCodegenSuite,AdaptiveQueryExecSuite,KeyGroupedPartitioningSuite,CoalesceShufflePartitionsSuiteandBucketedReadWithoutHiveSupportSuiterun 400 cases, all passing.sql/scalastyleandsql/Test/scalastyleare clean.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5