[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 re-derived its answers from the children's outputPartitioning on every call, and that answer moves between planning and execution, so the fusion gate, metrics and unionRDDs could disagree with each other.

What moves is now decided once per node and kept in a TreeNodeTag: whether the union is a plain concatenation (the old outputPartitioning body becomes a private rawPartitioning) and the two codegen confs. A new StampUnionDecisions rule asks for that answer right after EnsureRequirements in the standard and the AQE preparation pipeline, and again after each phase that can add a UnionExec of its own (the injected columnar and query-stage rules). spark.sql.unionOutputPartitioning is recorded one rule earlier, by SnapshotUnionOutputPartitioningConf, so EnsureRequirements and the stamp behind it cannot sample it separately. Every barrier in one preparation takes the same UnionConfSnapshot, read once, so a union an injected rule created or rebuilt after the first pass is stamped from the values the exchanges above it were planned against and not from the conf as it is by then. Reads before the barrier answer from what they see and write nothing, so observing an unprepared plan decides nothing for the prepared one.

A tag rather than a field, because withNewChildren ends in copyTagsFrom: the copy CollapseCodegenStages puts inside the codegen shell inherits the stamped values instead of taking its own. The gate's remaining terms read the children, so they stay per-instance and are re-derived when a rule installs different children.

Why are the changes needed?

Two failures on default configuration, both from those reads disagreeing.

A crash, when a fused union loses the metric its generated code needs:

java.util.NoSuchElementException: key not found: numOutputRows
  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 reports UnknownPartitioning while its inner AdaptiveSparkPlanExec has no final plan, so the union looks plain when CollapseCodegenStages gates on it and is fused. The copy insertInputAdapter puts in the shell evaluates the gate only after the cache stages finalise; by then both children report the same HashPartitioning, so metrics comes back empty under code that asks for numOutputRows. Flipping spark.sql.codegen.wholeStage.union.enabled between planning and execution reaches the same crash with no cache involved (given one child that is not CodegenSupport), which is why the codegen confs are stamped too. Registering the metric unconditionally stops both crashes, but leaves the wrong answer below in place and puts a 0-valued row count on every union that falls back to doExecute.

And a wrong answer, because rawPartitioning read spark.sql.unionOutputPartitioning live while 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()   // the union reports HashPartitioning(k, 4), so the aggregate's exchange is elided
spark.conf.set("spark.sql.unionOutputPartitioning", "false")
df.collect()   // the union concatenates instead, so each group lands in two partitions

Five groups of eight come back as ten rows of four. A prepared node can no longer change what it reports or how it executes, which fixes it. Deciding on first read would not have: executedPlan is prepareForExecution(sparkPlan.clone()) and clone ends in copyTagsFrom, so a read on sparkPlan would have decided for the prepared plan, and where nothing reads the node during preparation the first read lands at execution anyway.

Half of that gap stays open: a union stamped non-plain still derives rawPartitioning per call, so a partitioning change on one child that its siblings do not mirror can leave it concatenating at execution after reporting something concrete at planning. AQE reverts such a change when it breaks a parent's requirement, so what remains is a rule injected at one of the extension points, and a check at execution cannot close that, since the node cannot tell whether a parent relied on what it reported.

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

Does this PR introduce any user-facing change?

Yes. The first query above fails on 4.2.0 and now returns rows; the second returned duplicated groups and now returns the correct ones.

The three confs these decisions read (spark.sql.unionOutputPartitioning, spark.sql.codegen.wholeStage.union.enabled, spark.sql.codegen.wholeStage.union.maxChildren) no longer reach a plan that has already been prepared, only plans prepared afterwards. A DataFrame built before a change but first prepared after it observes the new value.

Three more changes come with deciding early, none of them affecting results:

• A union whose children only agree on partitioning after planning 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.
• DisableUnnecessaryBucketedScan runs after the first stamping pass, so a union over two bucketed scans with a projection on each side is stamped from the bucketed partitioning and then loses it: the gate answers partitioning-aware, and the union does not fuse, whereas deriving at the gate would have fused it. Closing that means splitting the two decisions, which is a separate change.
• CoalesceShufflePartitions keeps a non-plain union's children in one coalesce group, so if OptimizeSkewInRebalancePartitions splits only one of them, both lose coalescing where each used to be coalesced independently.

How was this patch tested?

In UnionCodegenSuite, for the two failures above and the decisions they turn on:

• a fused union keeps numOutputRows, and reports UnknownPartitioning rather than its children's HashPartitioning
• a union planned with spark.sql.unionOutputPartitioning on still executes that way after the conf is flipped off
• the same for spark.sql.codegen.wholeStage.union.enabled, over exchange children so the shell really holds a copy
• the same for spark.sql.codegen.wholeStage.union.maxChildren, over a three-child union whose cap is lowered under it
• an already-decided gate answers against the children a later withNewChildren gives it
• a read on the unprepared plan does not decide for the prepared one
• a read before any barrier leaves the gate provisional, so an unstamped union registers numOutputRows whichever way the stamp later lands
• a partitioning-aware union follows its children's coalesced partition count, so the count unionRDDs hands SQLPartitioningAwareUnionRDD is never a stamped one
• a union fused under AQE keeps both codegen decisions through execution, where the shell is built while the query runs

And in the same suite, for the barriers and the snapshot, driving the new rules by hand:

• a later stamping pass fills in a union an extension made, and leaves an earlier decision alone
• the stamp answers from the conf recorded before EnsureRequirements, not from the value live when it runs
• a replacement node whose own tag kept copyTagsFrom from bringing the stamp across is stamped by a late barrier from the preparation's conf, not from the value live when it runs
• a barrier reaching a union no snapshot pass saw records that preparation's codegen confs rather than the values live when it runs
• a snapshot pass ahead of each injected rule fills in a union the rule before it created
• a recorded conf answers the codegen gate before any decision is stamped

In SparkSessionExtensionSuite, through the real pipelines instead of calling the rules:

• a union an injected columnar rule adds is stamped with AQE off, and so is one the columnar hook adds inside AQE post stage creation
• so is one an injected query stage prep rule adds. That case reads the node from an injected stage-optimizer rule, since the barrier in postStageCreationRules otherwise stands behind the one after the prep rules and would stamp it anyway
• a union one injected prep rule adds is recorded before the next prep rule reads it, and the same for two injected post planner strategy rules, and again through the codegen gate rather than the partitioning
• with spark.sql.unionOutputPartitioning flipped after the AdaptiveSparkPlanExec wrapper has read its snapshot and before the stage carrying the union is created, the barrier there still stamps from the snapshot
• the same for the two codegen confs, over a three-child union an injected columnar rule rebuilds: it still fuses, and the copy inside the shell counts its rows

Each fails when the barrier or snapshot pass it covers is removed. The last two also fail when that barrier reads the conf live rather than taking the wrapper's snapshot: the partitioning one reports UnknownPartitioning(0), the codegen one fuses nothing.

One case in AdaptiveQueryExecSuite pins the order of the AQE list: the snapshot pass immediately before EnsureRequirements, the first stamp immediately after it, and the barrier for the injected prep rules last.

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

[LuciferYang]

friendly ping @dongjoon-hyun cloud you take another look when you have time? Thanks ~

[LuciferYang]

friendly ping @dongjoon-hyun

[LuciferYang]

also cc @cloud-fan

[cloud-fan]

Review summary

The TreeNodeTag mechanism is broader than the state being preserved: an AQE plan update can force a child-derived decision before a supported rule changes the union children, and generic tag copying then carries that old eligibility into the final codegen pass. The decisions should instead live in private per-instance UnionExec state, with an immutable snapshot transferred explicitly only for CollapseCodegenStages' InputAdapter rebuild. That preserves the configuration and metrics fixes while making ordinary child rewrites recompute against their actual topology. The two nearby comments also overstate the current lock and codegen guarantees.

Findings

3 total: 0 P0, 1 P1, 0 P2, 2 P3.

Blocking (P1)

• Replace generic tag propagation with an explicit codegen handoff — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1101 — see inline.

Nit (P3)

• Narrow the decisionLock lock-order claim — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1027 — see inline.
• Describe isPlainUnion as one codegen prerequisite — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1037 — see inline.

Existing discussions

• existing discussion — The current full-reason latch closes the discussed configuration route, while its broader copy lifetime creates the distinct supported extension route.
• existing discussion — The current implementation depends on broad tag copying, and a supported post-metrics rewrite can therefore carry an allow decision to children that fail the original gates.

Verification

• The execution-start plan traversal can force UnionExec.metrics before the final AQE query-stage and columnar rewrites that precede CollapseCodegenStages.
• TreeNode child replacement copies tags while UnionExec's codegen reason checks partition-index dependence, multi-RDD topology, columnar support, and child output types.

[cloud-fan]

Blocking (P1): CODEGEN_FAILURE_REASON is derived from the current children, but storing it in a TreeNodeTag lets generic withNewChildren copies inherit it after those children change. Under AQE, SparkPlanInfo can force metrics before a supported query-stage or columnar rule replaces the children, and the final CollapseCodegenStages pass then trusts the cached None. A newly added SparkPartitionID can read the global UnionRDD index instead of the child-local index; a dual-mode columnar replacement can instead hit the row-codegen assertion. Please keep these decisions in private per-instance UnionExec state and transfer an immutable snapshot explicitly only through a dedicated copy path used for the InputAdapter rebuild. Ordinary child rewrites should get fresh state, and UnionCodegenSuite should cover a post-metrics rewrite.

Recommended change: Replace the UnionExec TreeNodeTag decisions with private per-instance state, and add a UnionExec-specific codegen copy path that transfers an immutable decision snapshot only for CollapseCodegenStages' intended InputAdapter rebuild.

Why this works: An early metrics read stores an allow decision in generic TreeNode metadata; withNewChildren can then copy it after an extension changes the children, so the final codegen pass skips gates that the replacement topology fails.

Scope: UnionExec decision state, CollapseCodegenStages' InputAdapter rebuild, and focused AQE extension coverage in UnionCodegenSuite.

Compatibility: Keep the intended per-plan configuration stability and preserve both decisions across the specific InputAdapter rebuild; ordinary child rewrites must recompute from their actual topology.

Risks: Failing to transfer the snapshot on the intended InputAdapter copy can reintroduce the missing numOutputRows crash this PR fixes. Reusing the snapshot on any other child-changing path can leave another stale eligibility or partitioning decision.

Constraints: Do not add decision state to UnionExec's case-class parameters or reintroduce live-conf divergence between planning, metrics, and execution. Preserve partitioning-aware, multi-RDD, partition-index-dependent, and columnar fallbacks.

Success: Generic child rewrites receive fresh decision state, only the intended codegen copy receives the frozen snapshot, the existing cache and conf-flip regressions remain fixed, and a post-metrics query-stage or columnar rewrite neither changes partition-index values nor hits the row-codegen assertion.

[LuciferYang]

Fixed in ef44aa2, though not by the design you prescribed. The hole is real, and I turned it into a test before changing anything: with the reason on a tag, rebuilt.supportCodegen came back true after withNewChildren installed a child the gate rejects.

Instead of moving both decisions into per-instance state with a handoff from CollapseCodegenStages, I split the reason by what makes each term move. The two confs stay latched in a tag, since they move only because the read happens at a different time. Everything that reads the children (nested union, multi-RDD child, partition-index dependence, child count, supportsColumnar, type mismatch) is memoized per instance, so a copy starts cold and answers against its own children.

The copy in the codegen shell still agrees with the gate, and by construction rather than by routing: InputAdapter delegates output and supportsColumnar to its child, the other terms walk the subtree through it, and each of them is fixed for a given set of children (InMemoryTableScanExec.supportsColumnar is an override val, AdaptiveSparkPlanExec's is a constructor val). So this needs no UnionExec-shaped hook inside a generic rule, and the guarantee does not rest on insertInputAdapter remaining the only rebuild path between the gate and execution. The new case is a plain unit test with no injected extension: force metrics, hand the node a nested union through withNewChildren, assert the gate flips.

isPlainUnion stays latched on purpose. Stale-plain is the conservative direction (the node reports UnknownPartitioning), per-plan stability is what fixes the wrong answer this PR is about, and making it per-instance would let CoalesceShufflePartitions' grouping decision and a later rebuilt node disagree again. Glad to reconsider if you read that one differently.

[cloud-fan]

The split addresses my stale codegen-reason concern: the child-derived gates now recompute on the rebuilt node, and the focused test covers that path. I still see a separate issue with when isPlainUnion is first latched, which I will leave as a separate finding.

[LuciferYang]

Done in 3ed87a98166, with a trim in 773c28999eb. StampUnionDecisions runs right after EnsureRequirements in both QueryExecution.preparations and AQE's queryStagePreparationRules, and nothing else writes the decision, so a read before that point answers from the state it sees and decides nothing.

Both halves reproduce as test failures before the change. Reading queryExecution.sparkPlan's union partitioning left the prepared plan reporting UnknownPartitioning(0), the pre-EnsureRequirements answer riding in through clone and copyTagsFrom. And with whole-stage codegen off over a root union, where nothing consults the node during preparation, a conf flip afterwards moved the executed layout from four partitions to eight.

I put the stamp immediately after EnsureRequirements rather than at the end of preparation. Late stamping has a worse failure: EnsureRequirements can elide a parent's exchange on the strength of a concrete pass-through, and a stamp taken afterwards could freeze plain, leaving the union to concatenate under a parent that no longer shuffles. Stamping early can only freeze an answer a later rule would have sharpened, which costs an exchange elision rather than a result, and every rule below the stamp then reads the frozen value, including the EnsureRequirements re-runs inside OptimizeSkewedJoin.

Your constraint about pre-preparation reads sent me back to the codegen confs, which were still pinned on first read. They are stamp-only now too, so the whole decision is one immutable value written in one place. That also left decisionLock with nothing to serialize, so it and its lock-order comment are gone, and the two tags became one Decisions(plainUnion, unionCodegenEnabled, maxChildren), written by a single setTagValue, so a concurrent reader can no longer see half a decision. The cost of the read-no-write change is that a UnionExec built after the stamp point by an injected rule no longer gets the conf pinned by the gate's own read, so for that node a conf flip can still reach the copy in the codegen shell. It seemed the right way round, since such a node has no protection on the partitioning half either.

On the ordering question: reordering the gate does not buy anything once the decision is stamped, because every UnionExec is asked during preparation whether or not the gate would reject it, and isPlainUnion no longer writes. I tried it and reverted it in dc85f998986; all it changed was the reason reported for a union failing several gates.

I cannot confirm the other half of your question, and I think the answer is the opposite of hopeful. The stamp lands after EnsureRequirements, which is still before a cached child's inner AQE plan finalises, so such a union is stamped plain and keeps reporting UnknownPartitioning for the rest of the query. UnionCodegenSuite's first SPARK-59122 case shows it: after collect() the fused union reports UnknownPartitioning while its children have both reached the same concrete layout, which the test now asserts as the premise of that check. So the elimination is not merely unavailable already, it is given up. Deferring past that point means deciding at execution, which is the crash this PR fixes; outputPartitioning's scaladoc records the trade rather than hiding it.

[cloud-fan]

Thanks for spelling out the remaining boundary. I agree the stamp-only approach fixes pre-preparation reads, but a UnionExec introduced by the supported query-stage or columnar extension hooks after this pass still has live decisions and can recreate the shell-copy/metrics mismatch. I have kept this as a P2 finding; please add a stamping barrier after each late UnionExec-producing extension phase (without overwriting decisions already used by earlier planning consumers).

[cloud-fan]

The late barriers address the runtime issue, and the documented DisableUnnecessaryBucketedScan fusion tradeoff can remain a follow-up. The direct StampUnionDecisions test still does not validate that the non-AQE or AQE extension pipelines invoke the post-hook barriers; please add extension-driven coverage that fails if those barriers are removed or misordered.

[LuciferYang]

Three cases added in SparkSessionExtensionSuite, each building a UnionExec through a real hook: an injected columnar rule with AQE off, an injected query stage prep rule, and the columnar hook inside AQE post stage creation. Each one prepares a union over a repartition, then turns spark.sql.unionOutputPartitioning off and reads the node again. A stamped node keeps answering from the decision it was prepared with; one no barrier reached derives a decision at that read and comes back UnknownPartitioning.

I then removed the barriers one at a time to see what each case actually pins:

removed	columnar, AQE off	prep rule	post stage creation
preparations, after the columnar rules	fails	passes	passes
AQE, after the injected prep rules	passes	passes	passes
AQE postStageCreationRules	passes	passes	fails
both AQE listings	passes	fails	fails

Two of the listings are pinned one for one. The AQE pair is pinned jointly, because the barrier in postStageCreationRules stands behind the one after the prep rules: nothing an injected prep rule creates can reach execution undecided even with that listing gone, so a case that fails on it alone does not exist. I would rather tell you that than write one that looks like it does. What the earlier listing buys is upstream of execution: CoalesceShufflePartitions.childrenNeedCompatiblePartitioning reads isPlainUnion to decide whether a union's children coalesce as one group, and it runs before the post-stage barrier, so stamping earlier keeps that reader and the final decision in step. The test name and comment say exactly that much and no more.

An injected stage-optimizer rule needs no case of its own: it runs ahead of the same postStageCreationRules barrier, which the third case pins.

Your two nits are in as well. rawPartitioning's scaladoc now describes it as the child-derived candidate and gives the plain answer's two grounds separately, and the whereas.

[cloud-fan]

The extension cases cover the final stamping paths, but the query-stage-prep case still passes when its early barrier is removed because the later post-stage barrier stamps the node. Please add a check that observes the union from a query-stage optimizer rule, before that later barrier, so the required ordering has its own failure signal.

[LuciferYang]

Done in 65e3f96ee6c. The case now injects a stage-optimizer rule alongside the prep rule; it reads the union with spark.sql.unionOutputPartitioning turned off and puts the conf back, so a concrete answer can only come from a decision stamped before the stage optimizers ran, and the union an injected prep rule adds carries no recorded conf of its own to answer from.

Removing the barrier after the injected prep rules now fails that case with ListBuffer(UnknownPartitioning(0)), and it is the only case that fails, so each of the three listings has its own signal.

[cloud-fan]

Confirmed: the query-stage-preparation test now observes the injected union from a query-stage optimizer before the later post-stage barrier, so removing the early post-prep stamp fails the case. Resolved.

[cloud-fan]

Nit (P3): This lock-order claim is stronger than the implementation. Both derivations call child outputPartitioning, supportsColumnar, or output while holding decisionLock; ordinary children such as BroadcastHashJoinExec, AQEShuffleReadExec, and FileSourceScanExec implement some of those as lazy vals, so their first read synchronizes on the child. Please state the narrower lock that must be avoided here, especially the AQE final-plan lock, rather than claiming that only a descendant union lock can be acquired.

[LuciferYang]

You are right that it was stronger than the code. Narrowed it: what must not be reachable under decisionLock is the AQE final-plan lock, since CoalesceShufflePartitions reads isPlainUnion while holding it. The comment now also says that child lazy vals such as AQEShuffleReadExec.outputPartitioning do take their own instance monitor, and why that is harmless here.

[cloud-fan]

The stamp-only implementation removed decisionLock and its lock-order comment, so this concern no longer applies. Resolved.

[cloud-fan]

Nit (P3): isPlainUnion establishes only the partitioning prerequisite, so saying that “the codegen path applies” is too strong. deriveCodegenFailureReason can still reject a plain union for nested unions, multi-RDD children, partition-index-dependent expressions, the child limit, columnar support, or a type mismatch. Please describe this as satisfying the partitioning gate while remaining subject to the other codegen checks.

[LuciferYang]

Fixed. It now says isPlainUnion satisfies the partitioning gate on the codegen path rather than the whole of it, and that supportCodegenFailureReason still applies its other checks.

[cloud-fan]

Confirmed: the Scaladoc now says isPlainUnion satisfies only the partitioning gate and leaves the remaining checks to supportCodegenFailureReason. Resolved.

[cloud-fan]

Review summary

The core fix is directionally sound, but the decision is still tied to first observation rather than plan preparation. That leaves observer timing able to change the prepared topology or lets a post-preparation SQLConf change choose the execution layout. I also found four inaccurate comments around the new tests and tag-copy contract. These are P2 and P3 findings, so no P0 or P1 blocker was found, but the preparation-boundary issue should be resolved before merge.

Findings

5 total: 0 P0, 0 P1, 1 P2, 4 P3.

Non-blocking (P2)

• Stamp UnionExec decisions during preparation — sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:8124
  PLAIN_UNION_DECISION is still initialized by the first caller, so neither side of the new prepared-plan contract is enforced. Reading outputPartitioning on queryExecution.sparkPlan can stamp the pre-EnsureRequirements answer and have clone/copyTagsFrom carry it into executedPlan; if nothing reads it during preparation, a later SQLConf flip can instead decide the layout at first execution. Please initialize these decisions at an explicit post-EnsureRequirements preparation point in both normal and AQE planning, and keep pre-preparation reads provisional.

  Recommended change: Add a UnionExec decision-initialization physical rule at the defined post-EnsureRequirements point in both standard and AQE preparation, and expose only the minimal package-private UnionExec hook needed for that rule to stamp the plain-layout and codegen-configuration snapshots. Keep copied snapshots for the later InputAdapter rebuild, but do not let arbitrary pre-preparation reads initialize them.

  Why this works: Before the preparation marker is set, outputPartitioning derives a provisional answer without writing PLAIN_UNION_DECISION. The new preparation rule traverses the prepared plan after requirements are established, marks each UnionExec prepared, and initializes its immutable decisions from that topology and SQLConf. Later metrics, codegen, execution, and withNewChildren copies all consume the stamped values.

  Scope: sql/core/src/main/scala/org/apache/spark/sql/execution, sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive, sql/core/src/test/scala/org/apache/spark/sql/execution

  Compatibility: A union stamped plain remains conservative if its children later sharpen, and a union stamped non-plain continues to rely on supported AQE requirement reconciliation for later child divergence.

  Risks: Stamping before a later preparation rule changes relevant child partitioning would preserve the wrong topology. Failing to carry the stamped codegen snapshot through the intended InputAdapter copy would reintroduce the numOutputRows crash.

  Constraints: Run the rule after EnsureRequirements in both non-AQE and AQE pipelines and before consumers that require a stable answer. Preserve the intentionally conservative plain latch after preparation and the per-instance recomputation of children-derived codegen gates. Do not decide partitioning from PlanLater children or reintroduce live SQLConf reads during execution.

  Success: Reading outputPartitioning on queryExecution.sparkPlan cannot change queryExecution.executedPlan. Changing spark.sql.unionOutputPartitioning after executedPlan is prepared cannot change that plan's UnionRDD layout or partition count. The InputAdapter copy retains the same plain-layout and codegen-config decisions as the UnionExec that CollapseCodegenStages gated. Children-dependent codegen checks still recompute for ordinary child rewrites.

Nit (P3)

• Narrow the InputAdapter copy claim — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:746 — see inline.
• Do not describe withNewChildren as every rule path — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:777 — see inline.
• Qualify the tag propagation guarantee — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1357 — see inline.
• Do not equate an unfused union with empty metrics — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:70 — see inline.

Re-review status

Prior AI findings: 3 addressed, 0 still present; additional unresolved findings in this review: 5.

New attribution: 1 newly introduced, 4 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — The author agreed that an unprepared sparkPlan read must not latch the final decision, but the pinned implementation still writes the tag on first read and QueryExecution still copies it into the prepared clone.

PR description suggestions

• Update the claim that the three configurations cannot affect an already prepared plan: with first-read initialization, a prepared but unread root UnionExec can still observe a later UNION_OUTPUT_PARTITIONING value until the implementation stamps it during preparation.

Decision challenges

Avoid latching layout from a codegen-only metrics read

I verified that SparkPlanInfo to metrics to supportCodegenFailureReason calls isPlainUnion before the later child and topology gates, so a cached AQE union can latch plain even when it is ultimately not code-generated and its children later expose a shared partitioning. I could not cheaply confirm whether supported AQE paths have already made exchange elimination unavailable in every such case. Can you confirm that, or defer the isPlainUnion read until after the other codegen rejection gates?

[cloud-fan]

Nit (P3): insertInputAdapter also recursively rewrites descendants of a CodegenSupport child, so the UnionExec can still be copied even when its direct children support codegen. Please narrow this to the property these Exchange children guarantee for this test: a real InputAdapter rebuild and copy occurs.

[cloud-fan]

Confirmed: the comment now says CodegenSupport children can also lead to a copy and uses the Exchange children only to guarantee a real InputAdapter rebuild for this test. Resolved.

[cloud-fan]

Nit (P3): withNewChildren is the path for recursively changed children, but a transform can return an arbitrary replacement node directly. Please describe this test as covering the child-rewrite path, not every rule replacement.

[cloud-fan]

Confirmed: the comment now limits the case to the withNewChildren child-rewrite path and distinguishes arbitrary replacement nodes. Resolved.

[cloud-fan]

Nit (P3): copyTagsFrom copies tags only when the target tag map is empty. Please qualify the guarantee: fresh or tagless withNewChildren copies inherit this latch, while arbitrary transform replacements do not necessarily do so.

[cloud-fan]

Confirmed: the DECISIONS comment now qualifies tag inheritance to fresh targets with no existing tags and notes that already-tagged replacements are not overwritten. Resolved.

[cloud-fan]

Nit (P3): Being unfused does not imply empty UnionExec metrics: CollapseCodegenStages has additional outer gates absent from supportCodegenFailureReason. Please describe the helper selected shape instead of claiming all unfused unions omit numOutputRows.

[cloud-fan]

Confirmed: the test now ties empty metrics to this nested-union rejection instead of making a claim about every unfused union. Resolved.

[cloud-fan]

Thanks for the ping. I reviewed the latest revision and will post my findings separately.

[cloud-fan]

Review summary

I found no P0 or P1 blocker in the pinned revision. The preparation-time snapshot is the right core direction and closes the built-in first-reader failures, but one P2 lifecycle gap remains for UnionExec nodes introduced by supported physical extension rules after the only stamp. The other P2 items are an inaccurate plan-construction/configuration contract and missing same-plan maxChildren coverage; three P3 findings are localized comment fixes.

Findings

6 total: 0 P0, 0 P1, 3 P2, 3 P3.

Non-blocking (P2)

• Stamp unions introduced by late physical rules — sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala:39 — already raised in an existing discussion.
• Describe the config boundary as plan preparation — sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:2911 — see inline.
• Test maxChildren changes after plan preparation — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1084 — see inline.

Nit (P3)

• Fix the rawPartitioning Scaladoc grammar — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:981 — see inline.
• Correct the AQE metrics-read ordering comment — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1135 — see inline.
• Fix the fusedUnions helper sentence — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:70 — see inline.

Re-review status

Prior AI findings: 5 addressed, 0 still present; additional unresolved findings in this review: 6.

New attribution: 5 newly introduced, 0 late catch, 1 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• Suppressed duplicate: Stamp unions introduced by late physical rules — P2 at sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala:39 — existing discussion

PR description suggestions

• In the user-facing change section, replace 'only plans built afterwards' with 'only plans prepared afterwards'; a DataFrame or sparkPlan built earlier but first prepared after the configuration change observes the new value.

[cloud-fan]

Nit (P3): answers on whether is not grammatical here. Please say that isPlainUnion answers whether rawPartitioning is UnknownPartitioning, or that its answer is based on whether this method returns unknown.

[cloud-fan]

Confirmed: the Scaladoc now describes rawPartitioning as the child-derived candidate and no longer contains the grammatical issue. Resolved.

[cloud-fan]

Nit (P3): This event order is reversed: newQueryStage applies optimizeQueryStage and postStageCreationRules (including CollapseCodegenStages) before withFinalPlanUpdate calls onUpdatePlan. The early metrics read comes from SQLExecution building the initial SparkPlanInfo before execution. Please use that trigger to justify keeping the child-derived gates per instance.

[cloud-fan]

Confirmed: the comment now uses the initial SQLExecution SparkPlanInfo metrics read as the trigger and no longer relies on the reversed event order. Resolved.

[cloud-fan]

Non-blocking (P2): The new boundary is physical preparation, not plan construction. A DataFrame/sparkPlan can be built under one value, changed before executedPlan is first requested, and then stamped with the new value. Please update this and the other two captured union-conf descriptions to say that changes affect plans prepared afterward, not plans built afterward.

[cloud-fan]

Confirmed: all three configuration descriptions now say the value is fixed during physical preparation, so a DataFrame first prepared after a change observes the new value. Resolved.

[cloud-fan]

Non-blocking (P2): maxChildren is now part of the prepared-plan snapshot, but no test distinguishes this from the old live read. The existing cap tests set it before planning, and the same-plan flip covers only union.enabled. Please add the exchange-backed equivalent: prepare a two-child union at cap 2, lower the cap before collect(), and assert the shell copy still supports codegen and has numOutputRows.

[cloud-fan]

Confirmed: the new exchange-backed case prepares three children at cap 3, lowers the cap to 2, and verifies codegen plus the counted numOutputRows. Resolved.

[cloud-fan]

Nit (P3): This should read matches one that an InputAdapter left inside the stage; the relative clause is currently missing that.

[cloud-fan]

Confirmed: the comment now reads matches one that an InputAdapter left inside the stage. Resolved.

[cloud-fan]

Review summary

The preparation-time immutable-decision design is directionally sound and fixes the core wrong-result and codegen-metrics failures. I found no blocking P0/P1 issue, but the new late extension-barrier wiring still lacks a discriminating integration test, and two Scaladoc statements should be corrected.

Findings

3 total: 0 P0, 0 P1, 1 P2, 2 P3.

Non-blocking (P2)

• Exercise the late barriers through extension pipelines — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:893 — see inline.

Nit (P3)

• Document the independent config-disabled branch — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:981 — see inline.
• Use whereas for the contrast — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1098 — see inline.

Re-review status

Prior AI findings: 6 addressed, 0 still present; additional unresolved findings in this review: 3.

New attribution: 3 newly introduced, 0 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — The stale child-derived gate and late live-decision defects are fixed, but the latest author turn confirms the new AQE barrier positions are not asserted. The direct rule test therefore leaves a related integration test defect.
• existing discussion — The grammatical problem is fixed, but the requested replacement wording is semantically false when UNION_OUTPUT_PARTITIONING is disabled. This is a related new documentation defect.

PR description suggestions

• Revise the testing section so it does not imply coverage of the real extension-pipeline barrier wiring until integration tests exercise those paths; the current late-stamping case covers direct rule invocation only.

[cloud-fan]

Non-blocking (P2): This calls StampUnionDecisions directly, so it remains green if any of the new post-extension barriers is removed or misordered. Please add integration coverage that creates a UnionExec through real SparkSessionExtensions for the classic columnar, AQE query-stage preparation, and AQE post-stage columnar paths; those are the wiring points whose regression would recreate the decision/metrics mismatch.

Recommended change: Add extension-driven regression coverage that creates a fresh UnionExec from the non-AQE columnar hook and from the AQE query-stage preparation and post-stage optimization/columnar hooks, then observes a configuration change after preparation to prove each node retained the decision taken by its post-hook barrier.

Why this works: Use the existing SparkSessionExtensions test-session pattern and injected Rule or ColumnarRule implementations to introduce a semantically valid UnionExec at each late phase. Make every case discriminate the barrier by changing a captured union configuration afterward and asserting the extension-created node's final layout or codegen/metric behavior; the case must fail if its corresponding StampUnionDecisions entry is deleted or moved before the hook.

Scope: sql/core/src/test/scala/org/apache/spark/sql

Compatibility: The preparation-time snapshot, tag propagation for intended copies, child-derived codegen gates, and all documented result-preserving optimization tradeoffs remain unchanged.

Risks: A test can pass vacuously if an earlier barrier stamps the original plan rather than the UnionExec created by the targeted hook. An injected rule that duplicates or reshapes rows without a fixed oracle can turn the test into a plan-shape assertion rather than a lifecycle regression test.

Constraints: Use supported SparkSessionExtensions hooks and a real preparation or AQE stage lifecycle. Ensure each extension creates a fresh tagless UnionExec after the barrier whose placement is under test. Keep production rule ordering and write-once decision semantics unchanged.

Success: Removing or moving the post-columnar StampUnionDecisions entry in QueryExecution makes a focused non-AQE case fail. Removing or moving the barrier after AQE queryStagePrepRules makes a focused AQE case fail. Removing or moving the barrier after AQE stage-optimizer or columnar processing makes a focused AQE stage case fail. The existing cache, configuration-flip, and write-once regression cases continue to pass.

[LuciferYang]

Three extension-driven cases now sit in SparkSessionExtensionSuite, one per hook, and I posted the barrier-by-barrier measurement in the thread above. Short version: the listing after the columnar rules in preparations and the one in postStageCreationRules each fail a case on their own, while the listing after the injected prep rules has the post-stage one behind it, so no case can fail on that one alone.

[cloud-fan]

The non-AQE and post-stage cases now pin their barriers, but the query-stage-prep case still passes when only its early barrier is removed. Please observe the added UnionExec from an injected query-stage optimizer rule, before postStageCreationRules, so deleting or moving the post-prep stamp fails this case too.

[cloud-fan]

Confirmed: the query-stage-preparation case now injects ObserveUnionPartitioning as a query-stage optimizer. It reads the added union before postStageCreationRules, so removing the post-prep stamp makes the test observe UnknownPartitioning. Resolved.

[cloud-fan]

Nit (P3): These are not equivalent: when UNION_OUTPUT_PARTITIONING is disabled, rawPartitioning may still return a concrete child-derived partitioning while isPlainUnion is true and outputPartitioning reports unknown. Please describe rawPartitioning as the child-derived candidate and call out the independent configuration-disabled branch.

[LuciferYang]

Agreed, the two are not equivalent. The scaladoc now calls rawPartitioning the child-derived candidate and gives the plain answer's two grounds separately, the conf being off and this coming back UnknownPartitioning, with the consequence you point out: under that conf the candidate can still be concrete while the union reports unknown.

[cloud-fan]

Confirmed: rawPartitioning is now documented as the child-derived candidate, with the configuration-disabled branch stated independently. Resolved.

[cloud-fan]

Nit (P3): where is contrastive here; use whereas.

[LuciferYang]

Fixed.

[cloud-fan]

Confirmed: the contrast now uses whereas. Resolved.

[cloud-fan]

Review summary

The preparation-time decision design resolves the previously reported plan-copy, late-extension, and codegen-metric failures, and I found no blocking P0/P1 issue. Two P2 issues remain: the configuration value consumed by EnsureRequirements can still differ from the value frozen by the following stamp during a concurrent session-config update, and the AQE query-stage-preparation test does not distinguish its early barrier from the later post-stage fallback.

Findings

2 total: 0 P0, 0 P1, 2 P2, 0 P3.

Non-blocking (P2)

• Exercise the early AQE stamp before query-stage optimizers — sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala:666 — already raised in an existing discussion.
• Use one config snapshot across requirements and stamping — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1067 — see inline.

Re-review status

Prior AI findings: 2 addressed, 1 still present; additional unresolved findings in this review: 1.

New attribution: 0 newly introduced, 1 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

• Exercise the early AQE stamp before query-stage optimizers — sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala:666

Existing discussions

• Suppressed duplicate: Exercise the early AQE stamp before query-stage optimizers — P2 at sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala:666 — existing discussion
• existing discussion — The production stale-codegen decision problem and late barriers are fixed, but the latest same-thread request for extension-driven barrier coverage has a remaining exact overlap: the query-stage-prep test can still be satisfied by the later post-stage stamp and does not protect the early AQE ordering.

[cloud-fan]

Non-blocking (P2): EnsureRequirements and this later stamp sample UNION_OUTPUT_PARTITIONING separately. Because the shared SQLConf can be updated concurrently, a true-to-false flip between them lets the parent omit an exchange using a concrete union partitioning, then freezes plain concatenating execution here. An aggregate above the union can consequently emit duplicate groups. Please capture the configuration input before EnsureRequirements and use that same value when deriving the post-requirements decision.

Recommended change: Introduce a preparation-only configuration-input snapshot for UnionExec before each initial EnsureRequirements pass, then derive and publish the final plain-layout decision after requirements using that same captured input. Add deterministic concurrency coverage that changes the session setting between those phases.

Why this works: A pre-requirements rule records only the relevant configuration inputs on each unstamped UnionExec without deriving child partitioning. isPlainUnion consults a completed decision first, then the preparation input snapshot, and only uses the live session configuration before preparation. EnsureRequirements therefore plans against the captured setting; its rebuilt tagless copies inherit the input through normal tag propagation; the following StampUnionDecisions combines that same setting with the post-requirements child partitioning into the existing immutable decision. Fresh unions created by later supported hooks continue to be stamped at their post-hook barriers, and write-once completed decisions remain unchanged.

Scope: sql/core/src/main/scala/org/apache/spark/sql/execution, sql/core/src/test/scala/org/apache/spark/sql/execution

Compatibility: Keep the preparation-time configuration lifecycle, the partitioning decision derived after EnsureRequirements, the late-extension barriers, and the accepted exchange-elision, bucket-scan fusion, and skew-coalescing tradeoffs.

Risks: If the pre-requirements snapshot derives rawPartitioning as well as configuration, it can freeze the pre-exchange answer and lose valid partitioning-aware execution. If requirements rebuilds a UnionExec without propagating the input snapshot, the two phases can diverge again. Late extension-created unions must retain their existing post-hook stamping behavior and must not overwrite an earlier completed decision.

Constraints: Reads of an unprepared plan outside physical preparation remain non-mutating and reflect live state. Do not hold the shared SQLConf settings monitor across physical rule execution. Preserve the existing write-once decision, codegen-copy propagation, and documented result-preserving optimization tradeoffs. Capture configuration inputs without capturing child-derived partitioning before EnsureRequirements.

Success: EnsureRequirements and final UnionExec execution use the same UNION_OUTPUT_PARTITIONING value even when another thread changes the session setting between their rule invocations. A parent exchange elided from a concrete union partitioning cannot be followed by plain concatenating execution because of an intervening configuration update. Pre-preparation inspection does not pin a decision for the prepared clone, and plans first prepared after an ordinary non-concurrent configuration change still observe the new value. Late extension barriers and rebuilt codegen-shell copies preserve the intended completed decision without restamping it.

[LuciferYang]

Real, and the window is one rule wide. Closed in 65e3f96ee6c: a new SnapshotUnionOutputPartitioningConf runs before EnsureRequirements in both pipelines and records spark.sql.unionOutputPartitioning on each unstamped UnionExec; isPlainUnion answers from that record until the decision is stamped, and only a node created after that pass reads the conf live. The rule reads the conf once for the whole plan, so every union in one plan answers from the same value.

Only the conf, never a partitioning: EnsureRequirements has not inserted its exchanges yet, so a decision taken there would freeze plain on a union whose children only become co-partitioned after it. That is your first Risk, and it is why the decision itself still waits for the barrier behind EnsureRequirements.

Coverage is UnionCodegenSuite's "the stamp uses the conf the exchanges were planned against". It drives the three rules in order and flips the conf between the second and the third, since the two sit next to each other in the pipeline and no injected rule can run in the window. It asserts both halves: the aggregate's exchange was elided (both remaining exchanges are REPARTITION_BY_NUM), and the union still reports a concrete partitioning after the flip. Reverting the snapshot read to a live one fails it with UnknownPartitioning(0).

One thing I checked before agreeing: whether EnsureRequirements has the same hazard internally, since it reads conf.numShufflePartitions per child at two sites. It does not, because it reconciles the children against each other's actual partitioning afterwards rather than against the conf.

[cloud-fan]

Confirmed: both pipelines now snapshot the union-output-partitioning setting before EnsureRequirements and the later stamp uses that recorded value; the focused flip regression distinguishes a live re-read. Resolved.

[cloud-fan]

Review summary

The preparation-time snapshot and write-once stamping design is coherent in the pinned source, and the latest extension-driven test now observes the query-stage-preparation union before the later post-stage barrier. The previously reported production lifecycle issues are therefore resolved. This review found two new non-blocking coverage gaps: the initial AQE snapshot/EnsureRequirements ordering has no independent regression, and the cached-child flag-parity test reuses an already stamped plan so its disabled half no longer exercises the disabled path.

Findings

2 total: 0 P0, 0 P1, 2 P2, 0 P3.

Non-blocking (P2)

• Add an AQE-specific snapshot-ordering regression — sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:140 — see inline.
• Prepare a fresh cached-child union for each codegen flag — sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:2942 — see inline.

Re-review status

Prior AI findings: 2 addressed, 0 still present; additional unresolved findings in this review: 2.

New attribution: 1 newly introduced, 1 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

[cloud-fan]

Non-blocking (P2): The deterministic snapshot-ordering regression currently drives only QueryExecution.preparations. AQE maintains this separate private sequence, while the new extension tests exercise the later post-preparation and post-stage barriers with a stable initial setting. Removing this snapshot, or moving it after AQE's EnsureRequirements, therefore leaves those tests green and can reopen the race where requirements elides an exchange under one value but the union later stamps a different layout. Could we add an AQE-enabled discriminator that fails specifically when this initial snapshot is absent or late, and cannot be satisfied by the later stamping barriers?

[LuciferYang]

Added in AdaptiveQueryExecSuite: queryStagePreparationRules is now private[adaptive], and the case asserts the snapshot sits immediately before EnsureRequirements, the stamp immediately after it, and the trailing barrier last in the list.

My first version took only the first stamp index, which the trailing barrier satisfies on its own, so deleting the load-bearing one left the test green. Collecting every index closes that.

A behavioural discriminator is not reachable single-threaded. With the two rules adjacent, the only way the answers diverge is another thread writing the conf inside that one-rule window, so what the case pins is the order, as the non-AQE half does.

[cloud-fan]

Confirmed: AdaptiveQueryExecSuite now asserts the AQE snapshot immediately before EnsureRequirements, the first stamp immediately after it, and the trailing extension barrier last. Resolved.

[cloud-fan]

Non-blocking (P2): This preparation-scoped lifecycle makes the cached-child parity case in UnionCodegenSuite vacuous for its disabled half: that test first collects one DataFrame with union codegen enabled, then changes this setting and collects the same DataFrame again. The first collect has already memoized an executed plan with the enabled decision stamped, so the second collect never prepares or executes the disabled UnionExec fallback. Please construct and physically prepare a fresh cached-child union under each flag value, with a plan-path assertion that proves the two preparations observed different decisions before comparing their fixed results.

[LuciferYang]

Fixed, and four more cases in the same suite had the same shape: BHJ child union, the BHJ case feeding a downstream aggregate, SMJ child union, and column pruning. Each compared a plan's output against its own, since the expected value came from the plan the flag-on preparation had already stamped.

All five now build a fresh DataFrame per flag value through assertFlagParity, which takes a builder for exactly this reason, and the helper itself now asserts the flag-off half took no union into codegen, so every parity case gets that.

That assertion needed a discriminator the suite did not have. Asking whether a codegen stage holds a UnionExec somewhere is true even when CollapseCodegenStages left it under an InputAdapter, taking no part in codegen, and fusedUnions asks for the union to root its own stage, which misses one fused into a parent operator's stage: the aggregate case is that shape. The new codegenUnions walks each stage and stops at InputAdapter, and it replaced the former unionInsideWSCG at every call site, several of which were asserting the weaker property while their names claimed fusion.

[cloud-fan]

Confirmed: assertFlagParity now builds a fresh DataFrame under each flag value and proves the disabled plan has no codegen-participating union; the five affected cases use that helper. Resolved.

[dongjoon-hyun]

+1, LGTM for Apache Spark 4.3. For 4.2, I guess we need a narrow fix because spark.sql.codegen.wholeStage.union.enabled is not enabled by default.

[dongjoon-hyun]

cc @HeartSaVioR because this is a correctness issue. I raised SPARK-59122's Priority to Blocker and set Target Version to 4.3.0.

SPARK-59122: a partitioning-aware union keeps its layout when the conf changes between planning and execution
SPARK-59122: the stamp uses the conf the exchanges were planned against

[cloud-fan]

Review summary

The two findings from the prior AI review are addressed, and no prior finding remains. This pass found four independent P2 issues. Most importantly, the preparation snapshot and completed decision are long-lived correctness state, but TreeNodeTag.copyTagsFrom drops every source tag when a replacement already carries unrelated metadata; a supported late extension plus a concurrent configuration change can therefore make parent requirements and UnionExec execution disagree. The other three issues are narrower: the classic rule-order test does not enforce the required post-EnsureRequirements adjacency, an unstamped partitioning query repeats a cardinality-scaled derivation, and late extension barriers traverse whole plans even when no relevant extension is configured.

Findings

4 total: 0 P0, 0 P1, 4 P2, 0 P3.

Non-blocking (P2)

• Pin the first classic stamp immediately after EnsureRequirements — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:956 — see inline.
• Do not carry preparation decisions in long-lived TreeNode tags — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1413 — see inline.
• Reuse the raw partitioning derived for the unstamped branch — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1132 — see inline.
• Condition late stamp barriers on extension presence — sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala:865 — see inline.

Re-review status

Prior AI findings: 2 addressed, 0 still present; additional unresolved findings in this review: 4.

New attribution: 0 newly introduced, 4 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

[cloud-fan]

Non-blocking (P2): This assertion still passes if a partitioning-changing rule is inserted between EnsureRequirements and the first stamp. That is the unsafe window this lifecycle change needs to close: the parent can plan against one answer and the stamp can freeze another. Please assert firstStamp == ensureRequirements + 1 (while retaining the separate later-stamp check) so the focused test fails on that regression.

[LuciferYang]

Done: firstStamp == ensureRequirements + 1, with the later-stamp check left where it was. The comment above the assertion also claimed the AQE lists were out of reach from this suite, which stopped being true when queryStagePreparationRules became package-visible, so it now points at the AdaptiveQueryExecSuite case that asserts the same adjacency for the list AQE builds.

[cloud-fan]

Confirmed: the classic rule-order test now requires the first stamp immediately after EnsureRequirements and keeps the separate later-barrier assertion. Resolved.

[cloud-fan]

Non-blocking (P2): copyTagsFrom is all-or-nothing: if a replacement already carries any tag, neither OUTPUT_PARTITIONING_CONF nor DECISIONS is copied. These values now bridge EnsureRequirements and execution, so a supported late extension can return an equivalent UnionExec with unrelated metadata, lose both values, and be restamped from the live configuration. If the setting changes in that interval, a parent may already have omitted its exchange based on the old concrete partitioning while the replacement executes as a plain concatenation, yielding duplicate aggregate groups. Please keep this preparation state explicitly on UnionExec outside its primary case-class product, forward it through intentional copies, and make late barriers use the same preparation-scoped configuration snapshot rather than silently falling back to the live conf.

Recommended change: Replace OUTPUT_PARTITIONING_CONF and DECISIONS TreeNodeTag values with explicit non-product UnionExec preparation state, and give the preparation rules one immutable configuration snapshot shared by the initial and late barriers. Preserve existing-node state through intentional UnionExec copy paths, prepare genuinely new extension nodes from that shared snapshot, and fail closed if an unprepared union reaches codegen or execution.

Why this works: Represent unprepared, configuration-captured, and prepared states in private UnionExec state outside the primary case-class parameter list, forwarding it deliberately through withNewChildrenInternal and clone-sensitive copies. Capture the union configuration once in a preparation-scoped immutable input shared by the rules before and after EnsureRequirements and by the late extension barriers, so a genuinely new extension union is prepared from the same input rather than the live conf. Add a boundary check so missing preparation context cannot fall back to a new live configuration during codegen or execution.

Scope: Make UnionExec decision transport explicit and fail closed when preparation state is missing.

Compatibility: The preparation-time configuration lifecycle, write-once decisions, InputAdapter copy agreement, and dynamic AQE partition counts remain unchanged for correctly prepared plans.

Risks: A direct UnionExec copy site omitted from explicit state forwarding could fail validation or regress the original mismatch. Including preparation state in structural plan equality or canonicalization could inhibit reuse and destabilize plan output.

Constraints: Unprepared-plan inspection remains non-mutating and may answer from live configuration. The output-partitioning configuration sampled before EnsureRequirements remains the one used for the corresponding completed decision. Only the plain-versus-aware decision is retained; the Partitioning object and AQE partition count remain derived per call. Plan canonicalization, reuse identity, and normal explain output must not acquire preparation-state semantics.

Success: An unrelated TreeNode tag on a replacement cannot erase UnionExec's configuration snapshot or completed decision. Every copy of an existing prepared UnionExec used by requirements, codegen, metrics, or execution observes the same completed decision. A genuinely new extension-created UnionExec is prepared once at its supported barrier. No unprepared UnionExec can silently reach codegen or execution and sample live configuration.

[LuciferYang]

Done for the configuration half. Both barriers in one preparation now take the same UnionConfSnapshot, read once in QueryExecution.preparations and once per AdaptiveSparkPlanExec, and stampDecisions records that value on a node that reached the barrier without one. So the replacement you describe is stamped from what EnsureRequirements read rather than from the conf as it is by then. A new case pins it: an unrelated tag on a replacement, then a late barrier run while the conf says the opposite. It fails with UnknownPartitioning(0) without the change.

I would rather not move the state off tags, though. What loses it is TreeNode.transformUp calling afterRule.copyTagsFrom(this) on a rule's replacement node, and that is the only hook such a node goes through. A field outside the product is not copied there at all, so it would be lost in this case and also in every case where tags do survive today: withNewChildren, mapChildren, and the copy insertInputAdapter puts inside the codegen shell, which is the one that has to agree with the gate. Overriding withNewChildrenInternal covers the copies we make; nothing covers a node a rule constructed.

Failing closed on an unprepared union has the same problem from the other side: outputPartitioning and supportCodegen are read on plans that never went through prepareForExecution, which is what the reading the unprepared plan does not decide the prepared one case covers, and answering those from the live conf is what every union does on master today, so it cannot regress anything. The stamp only has to win where a decision has to hold still.

[cloud-fan]

The shared UnionConfSnapshot fixes the configuration-loss case. I rechecked the remaining tag lifetime and do not see a supported gap in the current pipeline: late extension replacements are followed by a stamp, and the final codegen rebuild copies from the stamped node into a fresh tagless target. I am resolving this without requiring explicit node fields.

[cloud-fan]

Non-blocking (P2): On an enabled, unstamped node, isPlainUnion has already computed rawPartitioning to determine whether it is unknown; when the result is concrete, this arm computes it again. That repeats the per-child AttributeMap construction and candidate intersections during planning. Please retain the first raw result within this invocation and reuse it for the return value, without caching it across calls because AQE partition counts must remain dynamic.

[LuciferYang]

Done. outputPartitioning holds a method-local lazy val, and the decision takes the raw partitioning by name, so a stamped node still never derives it and nothing is cached across calls.

[cloud-fan]

Confirmed: outputPartitioning now reuses one invocation-local raw value, while stamped nodes avoid deriving it and later calls can still observe AQE partition-count changes. Resolved.

[cloud-fan]

Non-blocking (P2): This late barrier, and the two analogous AQE suffix barriers, still traverse the visible plan when no custom rule can introduce a fresh UnionExec. The built-in rules at these positions retain or copy already-stamped nodes, so the default path pays three no-op tree walks (including per-stage AQE walks). Please add each suffix stamp only when its corresponding extension list is non-empty; the post-stage case needs to account for both injected optimizer and columnar rules. The mandatory initial post-EnsureRequirements stamp should remain unconditional.

Recommended change: Append each trailing StampUnionDecisions only when the immediately preceding extension surface can run a custom rule, while preserving the mandatory initial stamp and extension-enabled integration coverage.

Why this works: Build the classic columnar suffix from whether columnar rules are configured; build the AQE prep suffix from injected query-stage-preparation rules; and build the AQE post-stage suffix from injected optimizer or columnar rules. Update rule-order coverage so it asserts the barrier under configured extensions and its absence in default empty-extension plans.

Scope: Make late union-decision barriers conditional on the extension phases they protect.

Compatibility: Prepared unions and extension-created unions retain the same write-once decisions and execution layout as the pinned implementation.

Risks: Omitting a barrier when either an optimizer or columnar extension can create a union would restore live decisions. Tests that inspect default rule lists must not mistake an intentionally absent optional barrier for an ordering regression.

Constraints: Always retain the initial snapshot/EnsureRequirements/stamp sequence. Retain write-once behavior for already stamped unions. A configured extension-created union must be stamped before the next consumer identified by the existing integration tests.

Success: No late stamping traversal runs for an extension phase whose relevant custom rule lists are empty. Every supported configured extension phase that can create UnionExec retains a post-phase stamp. Each extension-enabled barrier keeps an independent regression signal.

[LuciferYang]

I would rather keep these unconditional. The condition is a different list per barrier: the classic suffix guards sessionState.columnarRules, the AQE prep suffix guards adaptiveRulesHolder.queryStagePrepRules, and the post-stage suffix guards the injected optimizer rules and the columnar ones together. Every hook added at one of those positions later has to be added to the matching condition too, and a missed one puts a union back on the live conf silently, for the one user who configured that extension. What it buys is one foreach per phase at planning time, next to ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages, which transform the same tree. The per-stage walks do not compound either, since foreach stops at QueryStageExec and a stage already created is not revisited.

If the walks are worth removing, I would rather prune the traversal than the rule: give UnionExec a nodePatternsInternal of Seq(UNION) and have the barrier ask plan.containsPattern(UNION) first, so a plan with no union costs a bitset test and the invariant stays unconditional. Happy to add that here if you want it.

[cloud-fan]

The safety rationale is persuasive: keeping the barriers unconditional avoids duplicating three extension-list predicates and silently missing a future hook. I do not have evidence that these adjacent plan walks are independently material, so I am resolving this without requiring conditional barriers.

[cloud-fan]

Review summary

The prior review's four concerns are no longer outstanding: the classic adjacency and raw-partitioning reuse requests are implemented, while the tag-lifetime and unconditional-barrier concerns have concrete, evidence-backed dispositions. This pass found four independent non-blocking issues. The most important is a remaining pre-stamp window in supported late extension hooks: forcing a new UnionExec's lazy codegen gate before its barrier can make the source and final codegen copy use contradictory decisions and recreate the missing-metric crash. Separately, AQE lacks behavioral regression signals for two fields in its three-field snapshot. The remaining two findings are local documentation corrections for an overbroad fusion promise and two subjectless lifecycle clauses.

Findings

4 total: 0 P0, 0 P1, 2 P2, 2 P3.

Non-blocking (P2)

• Keep pre-stamp codegen evaluation provisional — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1192 — see inline.
• Cover both codegen snapshot fields under AQE — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:806 — see inline.

Nit (P3)

• Do not promise all-child single-stage fusion — sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:2941 — see inline.
• Give each tag-lifecycle clause an explicit subject — sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:1434 — see inline.

Re-review status

Prior AI findings: 4 addressed, 0 still present; additional unresolved findings in this review: 4.

New attribution: 0 newly introduced, 4 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

[cloud-fan]

Non-blocking (P2): supportCodegenFailureReason memoizes the complete answer even before a preparation decision is stamped. A supported late extension can create a UnionExec and force this through metrics before its post-hook barrier; if the live codegen conf disagrees with the preparation snapshot, the source stays eligible while the withNewChildren copy inherits the stamped fallback decision and returns empty metrics, so doProduce can fail on missing numOutputRows. Please memoize only stable child-derived gates, keep preparation-dependent eligibility provisional until stamping, and ensure a metric materialized before the stamp remains the metric used if the eventual stamped node fuses. A late-extension regression should force the read before the barrier and install the opposite snapshot.

Recommended change: Separate stable child-topology analysis from preparation-dependent eligibility. Keep only child-derived gates memoized per instance; recompute the wrapper that consults stamped decisions until a stamp exists. If metrics are forced before stamping, expose a stable numOutputRows metric conservatively so either eventual stamped outcome is safe, while preserving empty metrics for ordinary already-stamped fallback unions. Add late-extension regressions for both codegen-enable and partitioning decisions.

Why this works: A provisional read can no longer cache live configuration past stampDecisions. The eventual CollapseCodegenStages gate reevaluates against the installed snapshot, and any metrics map materialized before that decision already contains the metric required if the final stamped answer permits fusion.

Scope: Make UnionExec's pre-stamp codegen inspection provisional and add extension-lifecycle coverage for the gate, shell copy, metrics, and final execution decision.

Compatibility: Unprepared plan inspection remains non-mutating, child-derived gates remain per-instance, and supported late hooks still receive their decision at the existing post-hook barrier.

Risks: A dynamic complete-gate computation must not repeatedly redo the expensive child AttributeMap and subtree checks. A metric object exposed before stamping must remain the same object used by generated code after stamping. The repair must not make an unstamped read write a decision inherited by the prepared clone.

Constraints: Preserve preparation-scoped configuration stability and write-once stamped decisions. Preserve per-instance recomputation when withNewChildren installs a different topology. Keep AQE partition counts and raw partitioning dynamic across calls.

Success: A pre-stamp supportCodegen or metrics read cannot determine the post-stamp gate from live configuration. The gate and any InputAdapter-rebuilt copy use the same stamped decision. If the eventual decision fuses the union, numOutputRows is present even when metrics were inspected before stamping. Ordinary unions stamped onto a fallback path continue to omit the unused row-count metric.

[cloud-fan]

Non-blocking (P2): Both configuration-flip regressions explicitly disable AQE, but AQE owns a separate UnionConfSnapshot and replanning lifecycle. If a later AQE round rereads either union-codegen enablement or maxChildren from the live SQLConf, the non-AQE tests and the AQE partitioning/rule-order checks remain green, and the shell copy can again lose numOutputRows. Please add AQE-enabled, exchange-backed cases for both fields that change the live value after the snapshot and assert a real InputAdapter copy retains fusion and the row-count metric.

[cloud-fan]

Nit (P3): This description says the parent and all children are fused into one WholeStageCodegenExec stage whenever the feature is enabled, but UnionExec still has other eligibility gates and unsupported children deliberately remain behind InputAdapter stage boundaries. Please describe the flag as allowing eligible non-partitioning-aware unions to participate in whole-stage codegen, with normal eligibility checks and child stage boundaries still applying.

[cloud-fan]

Nit (P3): The sentences beginning Written before and the semicolon clause beginning travels have no grammatical subject, which makes this tag lifecycle unnecessarily ambiguous. Please rewrite the pre-requirements write, post-requirements read, late-barrier write, and copy propagation as complete sentences that explicitly name the configuration value or tag.

[cloud-fan]

Review summary

Preparation-scoped UnionExec configuration is enforced at phase boundaries but not throughout supported extension-rule composition; make the captured snapshot available before a later rule can consume a fresh union, and behaviorally pin the same lifetime across later AQE stages.

The previous pre-stamp codegen-gate issue, fusion-description issue, and tag-lifecycle wording issue are addressed in the pinned revision. AQE codegen snapshot coverage is still not behaviorally closed, and the same missing lifetime signal also applies to output partitioning. I found no blocking P0/P1 issue; the remaining review consists of two P2 lifecycle/coverage findings and three localized P3 documentation fixes.

Findings

5 total: 0 P0, 0 P1, 2 P2, 3 P3.

Non-blocking (P2)

• Propagate the preparation snapshot between extension rules — sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:180 — see inline.
• Exercise late AQE unions under a changed snapshot — sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:125 — see inline.

Nit (P3)

• Add the missing comma in the conf-flip comment — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:803 — see inline.
• Document the preparation-wide config snapshot — sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:2942 — see inline.
• Fix the fusedUnions helper sentence — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:66 — see inline.

Re-review status

Prior AI findings: 4 addressed, 0 still present; additional unresolved findings in this review: 5.

New attribution: 0 newly introduced, 5 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — This records the issue's blocker priority and target version and summarizes the desired snapshot behavior. It does not ask a separately answerable question; the current source implements a preparation-scoped UnionConfSnapshot and adjacent initial stamp, while this review reports a distinct same-extension-phase gap.
• existing discussion — The codegen-field portion remains active and is incorporated into the broader Pass B test-integrity finding: AQE has a separate three-field snapshot, but the configuration-flip regressions disable AQE and no later-round test would fail on a live re-read. The canonical finding also covers the scanner-identified output-partitioning lifetime gap.

Verification

• The phase-end barriers do not interpose between consecutive injected prep, optimizer, or columnar rules.
• The current AQE tests do not behaviorally distinguish original snapshot reuse for a fresh later-stage union from a live reread of all three fields.

[cloud-fan]

Non-blocking (P2): This stamp runs only after the complete injected prep-rule list. If one rule creates a fresh co-partitioned UnionExec and the next applies requirements planning, that consumer still reads live UNION_OUTPUT_PARTITIONING; the suffix can then stamp the older preparation value. The parent may omit its exchange from the concrete provisional answer, while execution later concatenates under the stamped plain decision, so the same aggregate key can be emitted from multiple partitions. The preparation snapshot needs to be available before the next supported rule can consume the new union.

Recommended change: Carry the full preparation UnionConfSnapshot as provisional UnionExec state and refresh that provisional state after each injected physical or columnar rule, before the next rule runs; retain the existing phase-end completed-decision barriers and write-once semantics.

Why this works: Store or propagate the immutable three-field UnionConfSnapshot separately from the completed child-derived decision. After each supported external rule invocation, traverse only to fill missing provisional snapshot state on newly created unions. Later rules then read the preparation value for partitioning and codegen gates, while the existing suffix stamp combines that same snapshot with the topology present at its defined barrier.

Scope: Make one preparation snapshot visible throughout ordered extension-rule composition without freezing child-derived partitioning before the established completed-decision points.

Compatibility: Preserve live, non-mutating answers on genuinely unprepared plans; write-once completed decisions; dynamic AQE partition counts; and the documented optimization tradeoffs for already prepared unions.

Risks: Missing either the forward or reverse ColumnarRule fold would leave a same-list gap. Propagating a completed decision instead of only configuration could freeze topology too early and lose valid exchange elimination. A traversal inserted after rather than between external rules would reproduce the current defect.

Constraints: Do not derive or retain a Partitioning object before EnsureRequirements; AQE partition counts and child-derived raw partitioning must remain dynamic. Unprepared plan inspection remains non-mutating and may reflect live configuration. Already completed decisions remain write-once and survive intended tagless copies. Do not add preparation state to case-class product identity, canonicalization, or explain output.

Success: A fresh UnionExec created by one injected rule exposes the preparation snapshot to every later rule in that ordered extension phase. Requirements planning and eventual union execution cannot observe opposing UNION_OUTPUT_PARTITIONING values for the same fresh node. Pre-stamp codegen inspection and eventual stamped codegen decisions use the same codegen-enabled and maxChildren values. Child-derived partitioning is still decided only at the established post-requirements or post-extension barrier.

[cloud-fan]

Non-blocking (P2): The AQE configuration-flip case exercises a union already stamped while initialPlan is built, and the extension cases keep the initial value stable until their late barriers have run. Replacing a later barrier's stored unionConf with a fresh SQLConf read would therefore leave those assertions green. Please add exchange-backed AQE cases that flip output partitioning, codegen enablement, and maxChildren after this snapshot, then create a fresh union at a later stage/extension barrier and assert its layout, fusion, shell copy, and numOutputRows.

[LuciferYang]

Added two cases in SparkSessionExtensionSuite, both flipping a conf after the AdaptiveSparkPlanExec wrapper has read its snapshot and before the stage that carries the injected union is created.

• the barrier in AQE post stage creation stamps from the adaptive snapshot: an injected columnar rule adds the union, spark.sql.unionOutputPartitioning goes off once executedPlan is built, and the union still reports its children's HashPartitioning.
• a late AQE barrier records the codegen confs the wrapper was built with: an injected columnar rule rebuilds a three-child union, with spark.sql.codegen.wholeStage.union.enabled off in one iteration and maxChildren at 2 in the other. It asserts the union still fuses, its children inside the shell are InputAdapters, and the shell copy's numOutputRows reads 60.

Replacing that barrier's unionConf with UnionConfSnapshot(conf) turns both red: UnknownPartitioning(0) for the first, no fused union for the second.

The query-stage-preparation barrier has no case of its own. It cannot show the difference: queryStagePreparationRules is a val, so the barrier in it holds the construction snapshot whichever round applies it, and the union an injected prep rule adds is stamped before a flip can land. With that barrier changed to read SQLConf.get per application, a case of the same shape stays green. postStageCreationRules is the list rebuilt per stage, so it is the one whose value could come from a later read.

[cloud-fan]

Nit (P3): This should read Planned with the conf on, the union is fused; the comma separates the planning condition from the subject.

[cloud-fan]

Nit (P3): The value is not read when each union's decision is fixed: it is captured once before the physical preparation/adaptive rule sequence and reused by later barriers. Please describe that preparation-wide capture point here and on wholeStage.union.enabled, so a mid-preparation change is not documented as affecting a later-created union.

[cloud-fan]

Nit (P3): the node the callers here reach for its tags and metrics is ungrammatical; this can say the node whose tags and metrics the callers here inspect.

[cloud-fan]

Review summary

PR tags: bug-fix · correctness-fix · performance

The pinned revision addresses the prior within-list snapshot gap by recording the captured settings ahead of each injected query-stage preparation and post-planner rule. It also resolves the earlier punctuation and configuration-wording requests. The prior late-AQE test request remains open: current tests do not create a fresh union at a late barrier after the session settings diverge from the adaptive snapshot. I found one additional SQLConf description that needs a preparation-time caveat. Two advisory questions remain about fresh unions read before the late stamp in AQE optimizer validation and sequential columnar rules. No blocking P0/P1 defect is established.

Findings

2 total: 0 P0, 0 P1, 1 P2, 1 P3.

Non-blocking (P2)

• Exercise a fresh late AQE union after the settings diverge — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:1160 — remaining in an existing discussion.

Nit (P3)

• Qualify the partitioning pass-through description — sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:8185 — see inline.

Re-review status

Prior AI findings: 4 addressed, 1 still present; additional unresolved findings in this review: 1.

New attribution: 1 newly introduced, 0 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

• Exercise a fresh late AQE union after the settings diverge — sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:1160

Existing discussions

• Remaining: Exercise a fresh late AQE union after the settings diverge — P2 at sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:1160 — existing discussion

Decision challenges

Can AQE validation inspect a fresh union before its snapshot is recorded?

An injected query-stage optimizer can return a new UnionExec, and AQEShuffleReadRule can call ValidateRequirements on that result before postStageCreationRules stamps it. The unstamped union reads current SQLConf, while the later stamp uses the earlier adaptive snapshot. If the setting changes between those reads, can a supported optimizer produce a parent that relies on the union's provisional concrete partitioning? If so, validation could accept a layout the later plain union does not execute; please make the captured setting available before validation or clarify the supported boundary.

Can sequential columnar rules act on a fresh union's provisional layout?

ApplyColumnarRulesAndInsertTransitions runs injected columnar rules in sequence, but the new stamp follows the entire wrapper. A union created by one rule can therefore be read by the next from live SQLConf before the captured snapshot is written to it. Are exchange-changing decisions based on that read supported in this hook? If they are, a later plain stamp could invalidate such a decision; please record the snapshot between rule calls or explain why this shape is outside the supported hook contract.

[cloud-fan]

Nit (P3): This still describes pass-through solely in terms of the children's current partitioning. A union stamped plain during preparation continues to report UnknownPartitioning if its children later acquire a common concrete layout. Could you qualify the description with that preparation-time decision so it matches what a prepared union can report?

[LuciferYang]

Added to the description: "One decided to concatenate keeps reporting the default partitioning if its children come to share one afterwards."

[LuciferYang]

Pushed the two AQE cases the test thread asked for, and the preparation-time caveat on spark.sql.unionOutputPartitioning.

On validation reading a fresh union: which of the two reads is the permissive one decides what is at stake. With the conf on at validation and off in the snapshot, nothing above the union depends on its partitioning, because preparation saw it plain and EnsureRequirements put an exchange wherever a parent required a distribution; a rewrite that ValidateRequirements keeps cannot remove that exchange. With the conf off at validation and on in the snapshot, validation sees UnknownPartitioning, reads the parent's requirement as unmet, and optimizeQueryStage falls back to the plan before the injected rule. That costs a rewrite, not a layout the union will not execute.

On sequential columnar rules: there is no seam to list a pass into. The injected rules share one ApplyColumnarRulesAndInsertTransitions, which is a single rule from the outside, so recording between two of them means changing that rule. The hook is also not where exchanges are planned: under AQE it runs after the stage is created, without AQE after EnsureRequirements, so nothing downstream would reconcile an exchange added or dropped there. What a read between two columnar rules can still move is what the codegen gate answers, and the barrier for that lands before CollapseCodegenStages, so the gate and the copy in the shell agree either way. That is the boundary written on SnapshotUnionPreparationConf.

[LuciferYang]

Rewrote the testing section of the description. It said ten cases in UnionCodegenSuite and three in SparkSessionExtensionSuite where the net diff has fifteen and eight, and it left out the AdaptiveQueryExecSuite case that pins the AQE rule order. It now says what each case pins, grouped by suite, with no counts to go stale.

Catching up on the eight threads from the last two rounds that I never answered in place. All of them are in the pinned revision:

• supportCodegenFailureReason memoized before a stamp: it is a def again, only the child-topology terms are memoized, and metrics registers numOutputRows while a node is unstamped.
• AQE-enabled codegen-conf coverage: a union fused under AQE keeps both codegen decisions through execution covers the stamped node, and the extension case added today covers a fresh union at the late AQE barrier.
• the wholeStage.union.enabled description: it now describes an eligible union on its non-partitioning-aware path taking part, with the other gates and the InputAdapter boundary stated.
• subjectless sentences in the tag lifecycle: rewritten as complete sentences that name the rule doing the write and the value written.
• the window between two injected prep rules: SnapshotUnionPreparationConf.before lists a snapshot pass ahead of each injected prep rule, and of each post planner strategy rule.
• comma before the subject: fixed.
• where the value is captured: both codegen conf descriptions say it is read once per physical preparation.
• ungrammatical helper comment: now reads the node whose tags and metrics the callers here inspect.

@dongjoon-hyun on 4.2: spark.sql.codegen.wholeStage.union.enabled is createWithDefault(false) on branch-4.2, so the crash half is not reachable there by default, while spark.sql.unionOutputPartitioning has defaulted to true since 4.1, so the wrong-result half is.

That makes me doubt a trimmed variant is worth it. Dropping the two codegen fields from the snapshot would not shrink the blast radius: the three non-result behavior changes listed in the description all follow from taking the partitioning decision during preparation, which any fix for the wrong result has to do. What it would drop is the half with no footprint on 4.2, and the half that makes the codegen flag safe for anyone who turns it on there. So I lean towards taking this as it stands to 4.2, as its own PR against the branch. If you see a narrower shape that still closes the wrong result, I would rather hear it before I open that one.

[LuciferYang]

If no new concerns arise within 24 hours, I will merge this one first.