[SPARK-59122][SQL] Take UnionExec's plain-union decision once instead of re-deriving it per caller

[LuciferYang]

What changes were proposed in this pull request?

UnionExec reads its children's outputPartitioning to decide three separate things: whether whole-stage codegen fusion applies, whether numOutputRows is registered, and which RDD unionRDDs builds. Those reads can disagree, because the children's answer moves while the plan is prepared and executed. This PR makes the decision once per node, stores it in a TreeNodeTag, and derives outputPartitioning from the stored value.

• the old body of outputPartitioning becomes a private rawPartitioning, still derived from the children;
• isPlainUnion evaluates it on first use and latches the answer in UnionExec.PLAIN_UNION_DECISION, publishing it under a dedicated lock so two concurrent first readers cannot latch different answers;
• outputPartitioning returns super.outputPartitioning for a node that latched plain, and rawPartitioning otherwise;
• supportCodegenFailureReason latches the same way, in UnionExec.CODEGEN_FAILURE_REASON, so the fusion gate and metrics cannot answer differently either;
• the confs those two decisions read (spark.sql.unionOutputPartitioning, spark.sql.codegen.wholeStage.union.enabled, spark.sql.codegen.wholeStage.union.maxChildren) are read where the decision is latched rather than on every call.

TreeNode.withNewChildren ends in copyTagsFrom, which only writes into a tagless node, so the copy that CollapseCodegenStages places inside the codegen shell inherits the decision instead of making its own. A field would not survive that rebuild, which is why the decision lives on a tag rather than in a lazy val.

Why are the changes needed?

Two things go wrong on default configuration, both because those reads disagree.

The first is a crash: a fused union loses the metric its generated code increments.

java.util.NoSuchElementException: key not found: numOutputRows
  at org.apache.spark.sql.execution.SparkPlan.longMetric(SparkPlan.scala:154)
  at org.apache.spark.sql.execution.CodegenSupport.metricTerm(WholeStageCodegenExec.scala:71)
  at org.apache.spark.sql.execution.UnionExec.doProduce(basicPhysicalOperators.scala:1148)
  at org.apache.spark.sql.execution.WholeStageCodegenExec.doCodeGen(WholeStageCodegenExec.scala:676)

spark.range(0, 200, 1, 4).selectExpr("id % 10 AS k", "id AS v")
  .groupBy("k").agg(sum("v").as("s")).createOrReplaceTempView("v")
spark.catalog.cacheTable("v")
spark.sql("SELECT k, abs(s) AS s FROM v UNION ALL SELECT k, s FROM v").collect()

InMemoryTableScanExec.outputPartitioning reports UnknownPartitioning while its inner AdaptiveSparkPlanExec has no final plan. The row-based ProjectExec on one branch keeps the union out of the columnar path, so the union looks plain at the moment CollapseCodegenStages gates on it and is fused. insertInputAdapter then rebuilds the node and puts the copy in the shell, and that copy evaluates the gate for the first time after the cache stages have finalised: both children now report the same HashPartitioning, the gate answers partitioning-aware, metrics comes back empty, and doProduce throws while asking metricTerm for numOutputRows.

The same crash has a second route that needs no cache, because the codegen reason was derived from live confs as well. Plan the query with spark.sql.codegen.wholeStage.union.enabled on so the union is fused, set it to false, then collect: the copy inside the shell evaluates the reason for the first time at execution, answers union-codegen-disabled, and metrics comes back empty for a stage whose code is being generated on the assumption that it is fused. This route needs a child that is not CodegenSupport, an exchange for instance, since withNewChildren returns this when the children compare equal and then no second instance exists to re-derive anything. Latching the whole reason closes it: the gate evaluates it on the original before insertInputAdapter takes the copy.

Registering the metric unconditionally would stop both crashes, but it leaves the disagreement that produces the wrong answer below in place, since a fused union concatenates its children's partitions and a node that went on claiming their HashPartitioning could let a parent skip an exchange it needs. It would also put a 0-valued row count on every union that falls back to doExecute.

The second is a wrong answer with no crash. rawPartitioning used to read spark.sql.unionOutputPartitioning on every call, SparkPlan.conf is the live session conf, and QueryExecution.executedPlan is memoized on first read:

val df = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k"))
  .union(spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k")))
  .groupBy("k").count()

df.explain()   // planned with the conf on: the union reports HashPartitioning(k, 4), so the
               // aggregate's exchange is elided
spark.conf.set("spark.sql.unionOutputPartitioning", "false")
df.collect()   // executed with the conf off: the union concatenates instead of interleaving, so
               // each group lands in two partitions and the aggregate reports it twice

Five groups of eight come back as ten rows of four. Reading the conf where the decision is latched fixes this: once a node is planned, flipping the conf cannot change what it reports or how it executes.

One half of that gap stays open. A union latched non-plain still derives rawPartitioning from its children, and AQE skew splitting through a union leaves the children's partition counts divergent, so such a union can concatenate at execution after reporting something concrete at planning. Adding a check for that turned out to be wrong: nothing in those plans required the reported partitioning, and the node cannot tell at execution time whether anything did, so the check also rejected parents that ask for UnspecifiedDistribution.

This affects the fusion added in SPARK-56482, so branch-4.2 onward carries it.

Does this PR introduce any user-facing change?

Yes. A query of the first shape above fails on 4.2.0 and returns rows after this change; a query of the second shape returns duplicated groups and now returns the correct ones.

In the other direction, a union whose children only agree on partitioning after planning now keeps reporting UnknownPartitioning, so SPARK-52921's exchange elimination no longer applies to that shape. The alternative is a node that concatenates partitions while advertising a partitioning it does not have.

Both decisions are now taken once per node, so the three confs above no longer apply to a plan that has already been prepared. Setting one between explain() and collect() used to change how that plan executed, which is what the second crash route and the wrong answer above both are. A plan built after the change picks up the new value as before.

One AQE-side behavior change comes with the latch. CoalesceShufflePartitions reads the decision to tell whether a union's children have to coalesce compatibly, so a union that latched non-plain keeps its children in one coalesce group; if OptimizeSkewInRebalancePartitions then splits only one of them, the group's partition counts stop matching and both children lose coalescing. Before this change the answer was re-derived, came out plain at that point, and each child was coalesced independently. Results are the same either way.

How was this patch tested?

Three new cases in UnionCodegenSuite. The first covers both halves of the fused-union failure: the union keeps numOutputRows, and it reports UnknownPartitioning rather than its children's HashPartitioning. The second plans a union with spark.sql.unionOutputPartitioning on, flips the conf off, and then collects, so the union still executes by the partitioning it reported. The third does the same with spark.sql.codegen.wholeStage.union.enabled, over a union whose children are exchanges so that the shell really holds a copy. Each fails without the corresponding change: the first and the third with the exception above, the second with ten rows where five are expected.

Locally: UnionCodegenSuite, DataFrameSetOperationsSuite, CoalesceShufflePartitionsSuite, AdaptiveQueryExecSuite and KeyGroupedPartitioningSuite run 395 cases, all passing; sql/scalastyle and sql/Test/scalastyle are clean. The existing case SPARK-56482: partitioning-aware union falls back to non-codegen covers children that agree up front, so it pins that the latch does not disable the partitioning-aware path.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Opus 5

[LuciferYang]

cc @cloud-fan @dongjoon-hyun

[dongjoon-hyun]

Left inline comments. The main ones:

1. doExecuteColumnar still concatenates while a non-plain union advertises HashPartitioning, which gives wrong results with bucketed columnar scans. Pre-existing from SPARK-52921, but the new outputPartitioning doc states the invariant it breaks.
2. The latch changes CoalesceShufflePartitions behavior after a skew split through a union (both children lose coalescing).
3. The numOutputRows crash still reproduces through the codegen confs, since only the isPlainUnion input is latched.

The rest are smaller (lock shape, @transient val, test cleanups).

[dongjoon-hyun]

This invariant only holds on the row path. doExecuteColumnar (L1307) still does sparkContext.union(children.map(_.executeColumnar())) regardless of outputPartitioning, so a non-plain union that stays columnar concatenates while advertising its children's HashPartitioning.

I reproduced wrong results on master with two bucketed Parquet tables (bucketBy(4, "k")):

SELECT k, count(*) FROM (SELECT * FROM t1 UNION ALL SELECT * FROM t2) GROUP BY k

plans as HashAggregate <- ColumnarToRow <- Union(FileScan bucketed, FileScan bucketed) with no exchange and returns 10 rows instead of 5 (spark.sql.unionOutputPartitioning=false gives 5). This predates this PR (SPARK-52921), but since the doc here states the invariant, could we either mirror doExecute in doExecuteColumnar or make supportsColumnar false when !isPlainUnion? A separate JIRA is fine too.

[LuciferYang]

Your reproduction matched mine exactly, down to the ten rows of four. This became SPARK-59141 (#58445) and landed on master, so it arrives here through the merge in 32a2593d7ba, with doExecute and doExecuteColumnar sharing one unionRDDs helper. The bug is reachable on branch-4.2 and branch-4.1 as well, so I have backport PRs open for both (#58511, #58512).

[dongjoon-hyun]

One consumer of isPlainUnion whose behavior this latch changes is CoalesceShufflePartitions.childrenNeedCompatiblePartitioning (L189).

With AQE on, SparkPlanInfo.fromSparkPlan / EnsureRequirements latch the union non-plain while both children are identical rebalance exchanges. If OptimizeSkewInRebalancePartitions later splits only one child, that child reports UnknownPartitioning, but the tag still says non-plain, so both children land in one coalesce group with mixed specs and coalescePartitionsWithSkew bails out for both (Could not apply partition coalescing ...). Before this PR the re-derived answer was plain and each child was coalesced independently.

e.g. df1.hint("rebalance", "k").union(df2.hint("rebalance", "k")).count() with skew only in df1. Results are still correct, but this is a behavior change that isn't mentioned or tested.

[dongjoon-hyun]

This one is still open after the latest push. The CoalesceShufflePartitions consumer at L189 is a behavior change this PR introduces (both children lose coalescing after a skew split through the union), so I'd like to hear your take before moving on: is it acceptable as is, should the consumer read rawPartitioning instead of the latch, or should it be mentioned in the PR description and covered by a test?

[LuciferYang]

Keeping the latch, and I would rather not have the consumer read rawPartitioning.

The reason is narrower than "the latch is better". unionRDDs branches on the latched decision through outputPartitioning, so if childrenNeedCompatiblePartitioning keyed off a freshly derived value, the grouping decision and the arm the union actually takes could come from two different reads. Today they cannot. There is also a functional difference: reading isPlainUnion latches an unlatched union at the earliest AQE consumer, and reading rawPartitioning would not.

I could not construct a wrong result from the fresh read, for what it is worth. Wherever a parent relies on the union's partitioning, rawPartitioning has stayed concrete and the two values agree; where they diverge, nothing relies on it and both arms produce the same rows in a different layout. So this is about keeping the decision locally evident rather than contingent on that analysis.

On the behavior change itself you are right that it is one, and that it was neither documented nor tested. I have added it to the description: after a skew split under a non-plain union both children lose coalescing, where before the re-derived answer was plain and each child was coalesced independently.

One correction to the repro. df1.hint("rebalance", "k").union(df2.hint("rebalance", "k")).count() does not reach it: the ProjectExec that count() puts above the union drops k, so the union latches plain and takes the independent-group path. It needs an aggregate that keeps the key, for example groupBy("k").max("v").

On a test, the cheapest shape I found pins the timing rather than the skew path: cache a child so the union latches while the inner AQE plan is non-final, turn coalescing on, and assert that the two children's read specs may differ. I can add that if you want it, though it does not cover the skew axis, and building a stable one-side-skewed rebalance pair looked more expensive than this behavior change warrants. Your call.

[dongjoon-hyun]

The comment above says agreeing with the withNewChildren copy is the tag's job, but the tag only covers the isPlainUnion input. The copy that insertInputAdapter puts inside the shell still evaluates this lazy val for the first time at execution, so WHOLESTAGE_UNION_CODEGEN_ENABLED / WHOLESTAGE_UNION_MAX_CHILDREN flipped between executedPlan and collect() still gives empty metrics and the same key not found: numOutputRows from doProduce (AQE off, children with an exchange so the copy is fresh).

That's pre-existing, but it suggests the cleaner fix is registering numOutputRows unconditionally like the other CodegenSupport operators, rather than tying metrics to this memo.

[LuciferYang]

Good news: latching the whole reason worked, so I did not need the unconditional-registration fallback. Fixed in 543684188d1.

The reason now lives in a CODEGEN_FAILURE_REASON tag under the same decisionLock, so metrics and doProduce see one answer by construction, and the metric still stays off plans that fall back to doExecute. The trade-off is the one I mentioned: spark.sql.codegen.wholeStage.union.enabled becomes sticky per plan, the same semantics this PR already gives spark.sql.unionOutputPartitioning. Glad to switch to the unconditional registration if you would rather have that.

Your parenthetical about the exchange turned out to be the key part, and I am glad you wrote it. With spark.range(...).union(spark.range(...)) nothing fails, because withNewChildren returns this when the children compare equal, so there is no second instance to re-derive anything; the child has to not be CodegenSupport for insertInputAdapter to produce a real copy. My first repro attempt missed that and came back green. The new test puts a repartition(2) on each side, and reverting the latch makes it fail with the same key not found: numOutputRows.

The PR description now covers both routes.

[dongjoon-hyun]

Since this write happens on first read, reading outputPartitioning on the un-prepared queryExecution.sparkPlan now latches plain (children haven't got their exchanges yet), and executedPlan inherits it through sparkPlan.clone() -> makeCopy -> copyTagsFrom (QueryExecution.scala:395). After EnsureRequirements inserts matching HashPartitioning under both children, the union still reports UnknownPartitioning and the parent's exchange elimination is lost.

QueryTest.checkAnswer(_, planFunction, _) and PlannerSuite inspect sparkPlan this way, and listeners/extensions can too. Before the PR outputPartitioning was side-effect free.

[dongjoon-hyun]

Still open. Could you reply on whether this side effect of reading outputPartitioning on the un-prepared sparkPlan is acceptable, or whether the latch should only be taken once the plan is prepared?

[LuciferYang]

Not acceptable as is, agreed. I do not think "latch only once the plan is prepared" can be expressed from inside the node either, since it has no way to know whether it is being read before or after preparation.

What I would do instead is stamp the decision from a rule appended after EnsureRequirements, in both QueryExecution.preparations and AQE's stage-prep list. That turns "whoever reads first" into a defined point and makes outputPartitioning side-effect free again.

That also answers half of your design comment below: the conf can be a planner-set field, but the partitioning half cannot be decided in SparkStrategies, because the children there are PlanLater placeholders. A decision taken at that point comes out plain for every union, including ones whose children are co-partitioned without any exchange.

Would you rather I do that in this PR, or land the current read-on-first-use and follow up?

[dongjoon-hyun]

This is null after deserialization, so isPlainUnion / outputPartitioning / doExecute NPE on a deserialized UnionExec (e.g. inside a ScalarSubquery / InSubqueryExec plan captured in a task closure), where the previous outputPartitioning still worked. @transient private lazy val keeps it driver-only and survives the round trip, as SPARK-23731 did for FileSourceScanExec.

[dongjoon-hyun]

Still a val after the latest push. Is there a reason not to make it @transient private lazy val? If you'd rather keep it, please say so here.

[LuciferYang]

I would rather keep the val, because lazy val trades this NPE for a lock-ordering problem.

A lazy val's initializer takes the enclosing instance's monitor in Scala 2.13, and that is the same monitor unionedInputRDD's lazy val holds while it builds the children's RDDs. I checked the bytecode on this branch rather than trusting the reference: unionedInputRDD$lzycompute does monitorenter on this and runs children.map(...) and new UnionRDD(...) inside it. supportCodegenFailureReason is another lazy val that calls isPlainUnion from inside its own initializer. So a lazy lock puts its own initialization behind the monitor that the separate lock exists to stay out of, and CoalesceShufflePartitions reads isPlainUnion while holding the AQE lock.

On the NPE I read it as parity with what SparkPlan already does rather than as unreachable. SparkPlan has @transient private val prepareLock = new Object(), taken in prepare() and waitForSubqueries(), which are on the path of every execute*, and @transient val session ... orNull is on the path of conf, sparkContext and metrics. A deserialized plan that anyone uses as a plan has been failing on those long before this, so decisionLock adds no exposure that was not already there. I would not claim more than that: I did not sweep every path that can put a plan in a closure.

If you would still rather not add another one, the shape that avoids the monitor without the NPE is an explicit field forwarded in withNewChildrenInternal, which is your next comment.

[dongjoon-hyun]

rawPartitioning only walks the children (no execute, no lock another thread could hold while waiting on this one), so this can be a single decisionLock.synchronized { getTagValue(...).getOrElse { ...; setTagValue(...); plain } }, like SparkPlan.prepare(). Concurrent first readers then wait behind one derivation instead of each deriving and discarding.

[LuciferYang]

Done in d60079ebca4. It reads much better as one block, and prepare() was the right precedent to point me at.

[dongjoon-hyun]

nit: this paragraph (and "not this memo's job" in the memo comment below) reads as PR history rather than a description of the code. I'd keep the first paragraph plus a one-liner that the non-plain branch may become UnknownPartitioning after AQE skew splitting and is tolerated, and drop the tried-and-reverted sentence.

[LuciferYang]

Good catch, and done in d60079ebca4. Two later passes tightened the same paragraph again after I found it was still claiming more than the code shows: it now states the mechanism and names what reconciles a change, with no tried-and-reverted history left in it.

[dongjoon-hyun]

This and the next test have identical setup and assert two halves of the same decision. One test asserting both metrics.contains("numOutputRows") and outputPartitioning.isInstanceOf[UnknownPartitioning] would do.

[LuciferYang]

Agreed, merged in 9da6b20394e. One test asserts both halves now, and a later pass also pinned spark.sql.unionOutputPartitioning inside it, so it cannot pass vacuously if that default ever flips.

[dongjoon-hyun]

This guard is unreachable: both callers are inside withTempView("v"), which already uncaches via Catalog.dropTempView -> uncacheView at the end of each test.

[LuciferYang]

You are right, and it is gone in 9da6b20394e. The helper's doc comment now just says what withTempView does instead of promising more than dropTempView actually gives.

[dongjoon-hyun]

The answer is fully determined here (4 ids per k from each range), so checkAnswer(planned, (0L until 5L).map(k => Row(k, 8L))) is stronger and avoids a second plan/execute. The parity oracle would still pass if both paths regressed to 10 rows of 4.

[LuciferYang]

Done in 9da6b20394e: it asserts (0L until 5L).map(k => Row(k, 8L)) now and skips the second execution. Your point about both paths regressing together is exactly why, and I made the same change to the SPARK-59141 test for the same reason.

[dongjoon-hyun]

Is this ready back for reviews, @LuciferYang ?

[dongjoon-hyun]

Summary of my review comments so far, for tracking.

Addressed (thanks for the quick turnaround)

• doExecuteColumnar concatenating while a non-plain union advertised HashPartitioning (thread): resolved upstream by SPARK-59141 ([SPARK-59141][SQL] Interleave partitions in columnar UnionExec #58445), now included via the master merge. Both paths share unionRDDs.
• Double-checked locking in isPlainUnion (thread): collapsed to a single synchronized block.
• PR-history wording in the outputPartitioning / memo comments (thread, thread): reworded to describe current behavior. The memo comment now correctly says only the isPlainUnion term is latched; the pre-existing crash through the codegen confs remains, but that is out of scope here.
• Tests (thread, thread, thread): the two fused-union tests are merged, the unreachable isCached guard is gone, and the conf-flip test asserts the literal (k, 8) rows.

Still open, waiting for a reply

1. CoalesceShufflePartitions.childrenNeedCompatiblePartitioning reads the latched isPlainUnion, so after OptimizeSkewInRebalancePartitions splits only one child, both children lose partition coalescing. This is a behavior change introduced by this PR and is neither mentioned in the description nor tested (thread).
2. Reading outputPartitioning on the un-prepared queryExecution.sparkPlan now latches plain, which executedPlan inherits through sparkPlan.clone() -> copyTagsFrom, losing SPARK-52921's exchange elimination for that DataFrame (thread).
3. @transient private val decisionLock is null after deserialization; @transient private lazy val would avoid the NPE (thread).
4. Design: a field forwarded in withNewChildrenInternal (and a planner-set constructor field for UNION_OUTPUT_PARTITIONING) instead of a TreeNodeTag plus a dedicated lock. Fine either way, but the reasoning should be on record (thread).

Item 1 is the one I'd like settled before this merges. The Kafka failure in the latest CI run is unrelated (AvailableNow partition-metadata timeout); sql - slow tests was cancelled and has no result.

[LuciferYang]

Is this ready back for reviews, @LuciferYang ?

Not really… I forgot yesterday that this PR wasn’t fully fixed yet.

[LuciferYang]

Summary of my review comments so far, for tracking.

Addressed (thanks for the quick turnaround)

• doExecuteColumnar concatenating while a non-plain union advertised HashPartitioning (thread): resolved upstream by SPARK-59141 ([SPARK-59141][SQL] Interleave partitions in columnar UnionExec #58445), now included via the master merge. Both paths share unionRDDs.
• Double-checked locking in isPlainUnion (thread): collapsed to a single synchronized block.
• PR-history wording in the outputPartitioning / memo comments (thread, thread): reworded to describe current behavior. The memo comment now correctly says only the isPlainUnion term is latched; the pre-existing crash through the codegen confs remains, but that is out of scope here.
• Tests (thread, thread, thread): the two fused-union tests are merged, the unreachable isCached guard is gone, and the conf-flip test asserts the literal (k, 8) rows.

Still open, waiting for a reply

1. CoalesceShufflePartitions.childrenNeedCompatiblePartitioning reads the latched isPlainUnion, so after OptimizeSkewInRebalancePartitions splits only one child, both children lose partition coalescing. This is a behavior change introduced by this PR and is neither mentioned in the description nor tested (thread).
2. Reading outputPartitioning on the un-prepared queryExecution.sparkPlan now latches plain, which executedPlan inherits through sparkPlan.clone() -> copyTagsFrom, losing SPARK-52921's exchange elimination for that DataFrame (thread).
3. @transient private val decisionLock is null after deserialization; @transient private lazy val would avoid the NPE (thread).
4. Design: a field forwarded in withNewChildrenInternal (and a planner-set constructor field for UNION_OUTPUT_PARTITIONING) instead of a TreeNodeTag plus a dedicated lock. Fine either way, but the reasoning should be on record (thread).

Item 1 is the one I'd like settled before this merges. The Kafka failure in the latest CI run is unrelated (AvailableNow partition-metadata timeout); sql - slow tests was cancelled and has no result.

Thanks for your review. Let me take another look. @dongjoon-hyun