fix: [Spark 4.1] preserve union output partitioning in CometUnionExec#4207
Open
andygrove wants to merge 5 commits intoapache:mainfrom
Open
fix: [Spark 4.1] preserve union output partitioning in CometUnionExec#4207andygrove wants to merge 5 commits intoapache:mainfrom
andygrove wants to merge 5 commits intoapache:mainfrom
Conversation
…apache#4122) On Spark 4.1, SPARK-52921 (UNION_OUTPUT_PARTITIONING) lets UnionExec report a non-trivial hash partitioning when all children share the same partitioning, and downstream plans skip otherwise-required shuffles. CometUnionExec was (a) concatenating partitions via `sparkContext.union`, which breaks that partitioning contract, and (b) reading `outputPartitioning` from the frozen `originalPlan` snapshot, so post-AQE coalescing was invisible. The result was silent data-loss for EXCEPT ALL / INTERSECT ALL where both sides are GROUP BY queries. Override `outputPartitioning` to recompute from the live children, and route `doExecuteColumnar` through SQLPartitioningAwareUnionRDD on 4.1+ via a new `ShimCometUnionExec` shim. Pre-4.1 shims preserve the existing `sparkContext.union` behavior.
…ldens apache#4122 removed the inputs/intersect-all.sql and inputs/except-all.sql hunks from dev/diffs/4.1.1.diff but left two paired whitespace-trimming hunks on analyzer-results/intersect-all.sql.out and results/intersect-all.sql.out. The goldens came out trimmed while the upstream .sql still had trailing spaces, so SQLQueryTestSuite echoed the untrimmed SQL and failed to match the trimmed golden. Restore both .out files to upstream by regenerating the diff.
parthchandra
reviewed
May 4, 2026
| // is stale relative to the RDDs (e.g. children were coalesced by AQE but the reported | ||
| // partitioning was not). Fall back to plain concat in that case. | ||
| if (nonEmpty.isEmpty || nonEmpty.exists(_.partitions.length != numPartitions)) { | ||
| sc.union(rdds) |
Contributor
There was a problem hiding this comment.
if _.partitions.length != numPartitions fires, then we should probably log a warning message.
| val df = sql("""SELECT v FROM tab3 GROUP BY v | ||
| |EXCEPT ALL | ||
| |SELECT k FROM tab4 GROUP BY k""".stripMargin) | ||
| checkAnswer(df, Seq(Row(3))) |
Contributor
There was a problem hiding this comment.
use checkSparkAnswerAndOperator ?
Member
Author
There was a problem hiding this comment.
The query does have some operators that cannot be converted. I updated these tests to check for CometUnionExec though
| val df = sql("""SELECT v FROM tab1 GROUP BY v | ||
| |INTERSECT ALL | ||
| |SELECT k FROM tab2 GROUP BY k""".stripMargin) | ||
| checkAnswer(df, Seq(Row(2), Row(3), Row(null))) |
Contributor
There was a problem hiding this comment.
use checkSparkAnswerAndOperator here as well?
Log a warning when CometUnionExec falls back to plain SparkContext.union because child partition counts diverge from the declared output partitioning, so the unexpected state is observable. Strengthen CometSetOpWithGroupBySuite by comparing results to vanilla Spark via checkSparkAnswer and asserting CometUnionExec is present in the executed plan, instead of asserting hardcoded row literals.
Member
Author
|
Could you take another look @parthchandra |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #4122.
Rationale for this change
On Spark 4.1, SPARK-52921 added
UNION_OUTPUT_PARTITIONING: when all children of aUnionExecshare the same hash/single partitioning, the union itself reports that same partitioning. Downstream operators (e.g. a final hash aggregate) then skip an otherwise-required shuffle, and Spark's row-basedUnionExec.doExecutekeeps the partitioning invariant by routing throughSQLPartitioningAwareUnionRDD(each output partition unions partition i from every child).CometUnionExecsilently broke both halves of that contract:doExecuteColumnarusedsparkContext.union(...), which concatenates partitions — partition i of the output only holds partition i of a single child.outputPartitioningdelegated to the frozenoriginalPlansnapshot captured atCometExecRuletime, so AQE's post-stage coalescing was invisible.The result:
EXCEPT ALL/INTERSECT ALLwhose sides are themselvesGROUP BYaggregates lost rows silently (e.g.EXCEPT ALLreturning{2, 3}instead of{3}). Two Spark 4.1.1SQLQueryTestSuitefiles (except-all.sql,intersect-all.sql) were disabled for Comet because of this.What changes are included in this PR?
CometUnionExec.outputPartitioningto recompute from the livechildrenrather thanoriginalPlan.doExecuteColumnarthrough a newShimCometUnionExec.unionRDDshelper that usesSQLPartitioningAwareUnionRDDon Spark 4.1+ when a known partitioning is declared (with a partition-count sanity check and a safe fallback to plain concat), and retainssparkContext.unionbehavior on pre-4.1 Spark whereUnionExec.outputPartitioningis alwaysUnknownPartitioning.CometSetOpWithGroupBySuitecovering the two queries from the Spark SQL tests.spark.comet.enabled = falseguards at the top ofexcept-all.sqlandintersect-all.sqlindev/diffs/4.1.1.diff.How are these changes tested?
CometSetOpWithGroupBySuitepasses on Spark 3.5 and Spark 4.1.1 profiles.CometExecSuite(246 tests) passes on Spark 3.5.