From aad769bb838aca4b92621a79c2a0932253890887 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sun, 30 Aug 2026 23:52:17 +0800 Subject: [PATCH 01/24] [SPARK-59122][SQL] Take UnionExec's plain-union decision once instead of re-deriving it per caller Fused UnionExec threw "key not found: numOutputRows" because the copy insertInputAdapter puts inside the codegen shell re-derived the gate after the cache stages had finalised. Latch the decision in a TreeNodeTag, which withNewChildren copies, and derive outputPartitioning from it. --- .../execution/basicPhysicalOperators.scala | 63 +++++++++++++-- .../sql/execution/UnionCodegenSuite.scala | 77 +++++++++++++++++++ 2 files changed, 132 insertions(+), 8 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index bea86501e6f3a..95da1b277b160 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.optimizer.CollapseProject import org.apache.spark.sql.catalyst.plans.logical.Sample import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} @@ -964,7 +965,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - override def outputPartitioning: Partitioning = { + /** + * The SPARK-52921 pass-through partitioning, derived from the children on every call. Callers + * go through `outputPartitioning`, which reconciles this with the latched + * plain-union decision; see `isPlainUnion`. + */ + private def rawPartitioning: Partitioning = { if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) { return super.outputPartitioning } @@ -1023,13 +1029,46 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // True when the codegen path applies: `outputPartitioning` is `UnknownPartitioning`, - // and `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `doExecute`. - // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `doExecute`, but - // codegen is disabled for it (`supportCodegenFailureReason` reports "partitioning-aware"): - // the per-partition key descriptor is consumed by a downstream `GroupPartitionsExec`, and - // keeping these unions out of whole-stage codegen matches the `HashPartitioning` union case. - private[sql] def isPlainUnion: Boolean = outputPartitioning.isInstanceOf[UnknownPartitioning] + /** + * True when the codegen path applies: this union behaves as a plain concatenation, so + * `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `doExecute`. + * A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `doExecute`, but + * codegen is disabled for it (`supportCodegenFailureReason` reports "partitioning-aware"): + * the per-partition key descriptor is consumed by a downstream `GroupPartitionsExec`, and + * keeping these unions out of whole-stage codegen matches the `HashPartitioning` union case. + * + * DECIDED ONCE and then carried on the node, rather than re-derived per caller. The children's + * partitioning is not stable over a node's lifetime: `InMemoryTableScanExec.outputPartitioning` + * reads `cachedPlan.outputPartitioning`, and an inner `AdaptiveSparkPlanExec` answers + * `UnknownPartitioning` until its final plan exists. A union can therefore look plain while + * `CollapseCodegenStages` decides to fuse it and partitioning-aware by the time the stage runs, + * at which point every consumer of this predicate flips with it. + * + * That flip is observable because `insertInputAdapter` rebuilds the node through + * `withNewChildren` and puts the COPY inside the `WholeStageCodegenExec` it just created, so + * the copy re-derives the predicate later and can answer the opposite of the decision the + * shell was built from: `metrics` comes back empty and `doProduce`'s `metricTerm` throws + * `key not found: numOutputRows`. A `TreeNodeTag` survives that rebuild + * (`TreeNode.withNewChildren` ends in `copyTagsFrom`, which copies into a tagless node), so + * the copy inherits the decision instead of making a new one. + */ + private[sql] def isPlainUnion: Boolean = + getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { + val plain = rawPartitioning.isInstanceOf[UnknownPartitioning] + setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) + plain + } + + /** + * The partitioning this node both reports and executes by, so the RDD shape can never + * contradict the claim. A node latched plain reports `UnknownPartitioning` and concatenates, + * even if its children have since agreed on a concrete partitioning -- a fused union + * concatenates its children's partitions, and claiming their partitioning would let a parent + * skip an exchange it needs. The cost is that SPARK-52921's exchange elimination is lost for a + * union whose children only agree later, which is the conservative direction. + */ + override def outputPartitioning: Partitioning = + if (isPlainUnion) super.outputPartitioning else rawPartitioning // Per-child projection from the child's output to the union's output. The wrapped // child is always the source `Attribute` (deterministic by construction); the Alias @@ -1267,6 +1306,14 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } object UnionExec { + /** + * The once-and-for-all "does this union behave as a plain concatenation" decision, carried on + * the node so that the copy `CollapseCodegenStages.insertInputAdapter` puts inside a + * `WholeStageCodegenExec` cannot re-derive it and answer differently. See + * `UnionExec.isPlainUnion`. + */ + val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") + /** * Codegen operators that return more than one RDD from `inputRDDs()`. * `UnionExec`'s fusion assumes each direct child contributes one RDD. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 6af286b610e2b..c2bc2a37b6b61 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -21,6 +21,8 @@ import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -58,6 +60,34 @@ class UnionCodegenSuite extends SharedSparkSession { case w: WholeStageCodegenExec if w.find(_.isInstanceOf[UnionExec]).isDefined => w }.nonEmpty + /** + * Every node of the final plan, descending through AQE wrappers and query stages. + * `SparkPlan.collect` stops at `AdaptiveSparkPlanExec` and `QueryStageExec`, both of which + * are `LeafExecNode`s, so it cannot see a union that AQE placed inside a stage. + */ + private def allNodes(plan: SparkPlan): Seq[SparkPlan] = plan match { + case a: AdaptiveSparkPlanExec => plan +: allNodes(a.executedPlan) + case q: QueryStageExec => plan +: allNodes(q.plan) + case other => other +: other.children.flatMap(allNodes) + } + + private def fusedUnions(df: DataFrame): Seq[UnionExec] = + allNodes(df.queryExecution.executedPlan).collect { + case w: WholeStageCodegenExec if w.child.isInstanceOf[UnionExec] => + w.child.asInstanceOf[UnionExec] + } + + /** A cached aggregate, so the union's children read an `InMemoryTableScanExec`. */ + private def cacheAggregateView(view: String): Unit = { + spark.catalog.clearCache() + spark.range(0, 200, 1, 4) + .selectExpr("id % 10 AS k", "id AS v") + .groupBy("k") + .agg(sum("v").as("s")) + .createOrReplaceTempView(view) + spark.catalog.cacheTable(view) + } + /** Run query with flag on, then flag off, assert results match. */ protected def assertFlagParity(buildDf: () => DataFrame): Unit = { val onRows = buildDf().collect().toSeq @@ -628,6 +658,53 @@ class UnionCodegenSuite extends SharedSparkSession { } } + test("SPARK-59122: a fused union keeps numOutputRows when a child's partitioning firms up") { + // The children's partitioning is not stable while the plan is being prepared: + // `InMemoryTableScanExec` reads `cachedPlan.outputPartitioning`, and the inner + // `AdaptiveSparkPlanExec` answers `UnknownPartitioning` until its final plan exists. The + // asymmetry matters -- an expression on one branch only keeps a `ProjectExec` from being + // collapsed away, so `comparePartitioning` rejects the pair, the union looks plain and is + // fused. Once the cache stages finalise, both children report the same `HashPartitioning`, + // and re-deriving the decision at that point left `metrics` empty while the generated code + // still incremented it, so `doProduce` threw + // `NoSuchElementException: key not found: numOutputRows`. `SELECT *` or a plain alias is + // collapsed away and does not reproduce this. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("v") { + cacheAggregateView("v") + val df = spark.sql("SELECT k, abs(s) AS s FROM v UNION ALL SELECT k, s FROM v") + // Execute this DataFrame rather than a count over it: the plan being inspected has to be + // the one that ran, and an AQE plan that never ran has no final plan to inspect. + assert(df.collect().length == 20) + val fused = fusedUnions(df) + assert(fused.nonEmpty, + "this shape must actually fuse, or the test is not exercising the defect") + assert(fused.forall(_.metrics.contains("numOutputRows")), + "a fused union must register the metric its generated code increments") + } + } + } + + test("SPARK-59122: a fused union reports UnknownPartitioning, so no parent skips an exchange") { + // The other half of the same decision. A fused union concatenates its children's partitions, + // so if it went on claiming the children's `HashPartitioning` a parent could satisfy a + // clustered distribution from an RDD that does not have it -- a wrong answer rather than a + // crash, which is why the missing metric must not simply be registered unconditionally. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withTempView("v") { + cacheAggregateView("v") + val df = spark.sql("SELECT k, abs(s) AS s FROM v UNION ALL SELECT k, s FROM v") + df.collect() + val fused = fusedUnions(df) + assert(fused.nonEmpty, "this shape must actually fuse") + fused.foreach { u => + assert(u.outputPartitioning.isInstanceOf[UnknownPartitioning], + s"a fused union must not claim a concrete partitioning, got ${u.outputPartitioning}") + } + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From d60079ebca498239f684a850cc29f07e90d9a4df Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 31 Aug 2026 09:06:46 +0800 Subject: [PATCH 02/24] [SPARK-59122][SQL] Address review comments on the plain-union latch - publish the latch under a dedicated `decisionLock`, with the derivation outside it and the first writer winning - read `outputPartitioning` once in `doExecute` so the branch and `numPartitions` come from the same answer - correct four comments whose stated mechanisms did not match the code - reuse `AdaptiveSparkPlanHelper.collect` instead of a hand-written traversal, narrow the test cache helper to this view's own plan, and make `PLAIN_UNION_DECISION` private --- .../execution/basicPhysicalOperators.scala | 106 +++++++++++------- .../sql/execution/UnionCodegenSuite.scala | 41 ++++--- 2 files changed, 83 insertions(+), 64 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 95da1b277b160..b098605589912 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -947,7 +947,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // Compares two leaf partitionings for union pass-through equivalence. Callers pass leaf // partitionings only; a `PartitioningCollection` is flattened to its members by - // `outputPartitioning` before reaching here. + // `rawPartitioning` before reaching here. private def comparePartitioning(left: Partitioning, right: Partitioning): Boolean = { (left, right) match { case (SinglePartition, SinglePartition) => true @@ -955,7 +955,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // For `KeyedPartitioning`, only the partition expressions must match (both sides' // expressions have already been remapped to this union's output attributes by // `prepareOutputPartitioning`). The partition keys are intentionally not compared here: - // children typically carry different key sets, and `outputPartitioning` merges them. + // children typically carry different key sets, and `rawPartitioning` merges them. case (l: KeyedPartitioning, r: KeyedPartitioning) => l.expressions.length == r.expressions.length && l.expressions.zip(r.expressions).forall { case (le, re) => le.semanticEquals(re) } @@ -966,9 +966,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * The SPARK-52921 pass-through partitioning, derived from the children on every call. Callers - * go through `outputPartitioning`, which reconciles this with the latched - * plain-union decision; see `isPlainUnion`. + * The SPARK-52921 pass-through partitioning, derived from the children on every call. + * `isPlainUnion` latches on whether this comes back `UnknownPartitioning`, and + * `outputPartitioning` reports it on the branch where that latch says the union is not a plain + * concatenation. */ private def rawPartitioning: Partitioning = { if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) { @@ -1029,6 +1030,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } + // Guards the `isPlainUnion` latch below. A dedicated object rather than this node's own monitor: + // `unionedInputRDD` is a `lazy val`, so initializing it holds that monitor while it calls + // `inputRDDs()` on the children, which for an `InputAdapter` child is `child.execute()`. + // `@transient`, so null on a deserialized node; every reader of the latch is driver-side, as + // for `AdaptiveSparkPlanExec.lock`. + @transient private val decisionLock = new Object() + /** * True when the codegen path applies: this union behaves as a plain concatenation, so * `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `doExecute`. @@ -1037,35 +1045,49 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * the per-partition key descriptor is consumed by a downstream `GroupPartitionsExec`, and * keeping these unions out of whole-stage codegen matches the `HashPartitioning` union case. * - * DECIDED ONCE and then carried on the node, rather than re-derived per caller. The children's - * partitioning is not stable over a node's lifetime: `InMemoryTableScanExec.outputPartitioning` - * reads `cachedPlan.outputPartitioning`, and an inner `AdaptiveSparkPlanExec` answers - * `UnknownPartitioning` until its final plan exists. A union can therefore look plain while - * `CollapseCodegenStages` decides to fuse it and partitioning-aware by the time the stage runs, - * at which point every consumer of this predicate flips with it. + * DECIDED ONCE and carried on the node rather than re-derived per caller, because the children's + * answer moves: `InMemoryTableScanExec.outputPartitioning` reads through `cachedPlan`, which + * unwraps an inner `AdaptiveSparkPlanExec` to its `executedPlan` only while that plan's + * `isFinalPlan` holds, and `withFinalPlanUpdate` clears `isFinalPlan` on entry. So a union can + * look plain when `CollapseCodegenStages` gates on it and partitioning-aware by the time the + * stage runs. + * + * That matters because `insertWholeStageCodegen` gates on this node but wraps the result of + * `insertInputAdapter`, which rebuilds the node through `withNewChildren` when its subtree + * changed. Re-deriving on the copy can contradict the decision the shell was built on: + * `metrics` on the copy comes back empty, and the `doProduce` that the shell drives throws + * `key not found: numOutputRows` out of `metricTerm`. A memo on the instance would not help, + * because the copy is a new instance; a `TreeNodeTag` does, since `withNewChildren` ends in + * `copyTagsFrom` and the copy carries no tags of its own. * - * That flip is observable because `insertInputAdapter` rebuilds the node through - * `withNewChildren` and puts the COPY inside the `WholeStageCodegenExec` it just created, so - * the copy re-derives the predicate later and can answer the opposite of the decision the - * shell was built from: `metrics` comes back empty and `doProduce`'s `metricTerm` throws - * `key not found: numOutputRows`. A `TreeNodeTag` survives that rebuild - * (`TreeNode.withNewChildren` ends in `copyTagsFrom`, which copies into a tagless node), so - * the copy inherits the decision instead of making a new one. + * Concurrent first readers can derive different answers, since the children's answer moves, so + * the latch is published under `decisionLock` and the first writer wins; the derivation itself + * runs outside the lock. The lock covers this tag only: `TreeNode`'s tag map is unsynchronized, + * and `copyTagsFrom` reads it without taking it. */ - private[sql] def isPlainUnion: Boolean = - getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { + private[sql] def isPlainUnion: Boolean = { + decisionLock.synchronized(getTagValue(UnionExec.PLAIN_UNION_DECISION)).getOrElse { val plain = rawPartitioning.isInstanceOf[UnknownPartitioning] - setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) - plain + decisionLock.synchronized { + getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { + setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) + plain + } + } } + } /** - * The partitioning this node both reports and executes by, so the RDD shape can never - * contradict the claim. A node latched plain reports `UnknownPartitioning` and concatenates, - * even if its children have since agreed on a concrete partitioning -- a fused union - * concatenates its children's partitions, and claiming their partitioning would let a parent - * skip an exchange it needs. The cost is that SPARK-52921's exchange elimination is lost for a - * union whose children only agree later, which is the conservative direction. + * A node latched plain reports `UnknownPartitioning` and concatenates, even if its children + * have since agreed on a concrete partitioning -- a fused union concatenates its children's + * partitions, and claiming their partitioning would let a parent skip an exchange it needs. The + * cost is that SPARK-52921's exchange elimination is lost for a union whose children only agree + * later, which is the conservative direction. + * + * The other branch is still derived per call, so what this node reports is not stable across + * planning and execution: a parent can elide an exchange on a concrete partitioning and + * `doExecute` can then derive `UnknownPartitioning` and concatenate. `doExecute` keeps one + * answer within itself; the cross-phase gap is not closed here. */ override def outputPartitioning: Partitioning = if (isPlainUnion) super.outputPartitioning else rawPartitioning @@ -1086,10 +1108,11 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Memoized: consulted by `supportCodegen` (called multiple times by - // `CollapseCodegenStages`) and by `metrics`. Conf and children are stable - // for a given UnionExec instance; cross-plan staleness is impossible since - // UnionExec is a case class and `withNewChildren` produces a fresh instance. + // Memoized so that `supportCodegen` (called multiple times by + // `CollapseCodegenStages`) and `metrics` see one reason on one instance: + // `conf` is the live session conf, so re-deriving could answer differently. + // Agreeing with the copy `withNewChildren` makes is a separate matter, and + // comes from `isPlainUnion` being latched on a `TreeNodeTag`. @transient private lazy val supportCodegenFailureReason: Option[String] = { if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) { Some("union-codegen-disabled") @@ -1274,7 +1297,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup override def usedInputs: AttributeSet = AttributeSet.empty protected override def doExecute(): RDD[InternalRow] = { - outputPartitioning match { + // One read: the non-plain branch re-derives from the children on every call, so + // `numPartitions` has to come from the same answer that picked the branch. Otherwise + // `SQLPartitioningAwareUnionRDD` can be handed an `UnknownPartitioning(0)`, which gives it + // zero partitions and this union an empty result. + val partitioning = outputPartitioning + partitioning match { case _: UnknownPartitioning | _: KeyedPartitioning => // An `UnknownPartitioning` union simply concatenates its children. A // `KeyedPartitioning` union does the same: its merged partition keys describe the @@ -1288,8 +1316,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // same partitioning in semantics so this union can choose not to change the partitioning // by using a custom partitioning aware union RDD. val nonEmptyRdds = children.map(_.execute()).filter(!_.partitions.isEmpty) - new SQLPartitioningAwareUnionRDD( - sparkContext, nonEmptyRdds, outputPartitioning.numPartitions) + new SQLPartitioningAwareUnionRDD(sparkContext, nonEmptyRdds, partitioning.numPartitions) } } @@ -1306,13 +1333,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } object UnionExec { - /** - * The once-and-for-all "does this union behave as a plain concatenation" decision, carried on - * the node so that the copy `CollapseCodegenStages.insertInputAdapter` puts inside a - * `WholeStageCodegenExec` cannot re-derive it and answer differently. See - * `UnionExec.isPlainUnion`. - */ - val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") + /** The latched "is this a plain concatenation" decision. See `isPlainUnion`. */ + private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") /** * Codegen operators that return more than one RDD from `inputRDDs()`. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index c2bc2a37b6b61..b0fbd144ef595 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -22,7 +22,7 @@ import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -32,7 +32,7 @@ import org.apache.spark.sql.types._ * Tests for `UnionExec` whole-stage codegen fusion: plan-shape assertions, * correctness, type widening, metrics, and fallbacks. */ -class UnionCodegenSuite extends SharedSparkSession { +class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { // Union codegen fusion is off by default; turn it on for this suite. override protected def sparkConf: SparkConf = @@ -61,30 +61,27 @@ class UnionCodegenSuite extends SharedSparkSession { }.nonEmpty /** - * Every node of the final plan, descending through AQE wrappers and query stages. - * `SparkPlan.collect` stops at `AdaptiveSparkPlanExec` and `QueryStageExec`, both of which - * are `LeafExecNode`s, so it cannot see a union that AQE placed inside a stage. + * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query stages, which + * `SparkPlan.collect` does not: both are `LeafExecNode`s, so it cannot see a union that AQE + * placed inside a stage. */ - private def allNodes(plan: SparkPlan): Seq[SparkPlan] = plan match { - case a: AdaptiveSparkPlanExec => plan +: allNodes(a.executedPlan) - case q: QueryStageExec => plan +: allNodes(q.plan) - case other => other +: other.children.flatMap(allNodes) - } - private def fusedUnions(df: DataFrame): Seq[UnionExec] = - allNodes(df.queryExecution.executedPlan).collect { + collect(df.queryExecution.executedPlan) { case w: WholeStageCodegenExec if w.child.isInstanceOf[UnionExec] => w.child.asInstanceOf[UnionExec] } /** A cached aggregate, so the union's children read an `InMemoryTableScanExec`. */ private def cacheAggregateView(view: String): Unit = { - spark.catalog.clearCache() spark.range(0, 200, 1, 4) .selectExpr("id % 10 AS k", "id AS v") .groupBy("k") .agg(sum("v").as("s")) .createOrReplaceTempView(view) + // Both callers need the cache unmaterialized, and `CacheManager` treats caching an + // already-cached plan as a no-op, so drop whatever an earlier test left for this plan. + // `isCached` matches by plan, so it also catches the same plan cached under another name. + if (spark.catalog.isCached(view)) spark.catalog.uncacheTable(view) spark.catalog.cacheTable(view) } @@ -660,15 +657,15 @@ class UnionCodegenSuite extends SharedSparkSession { test("SPARK-59122: a fused union keeps numOutputRows when a child's partitioning firms up") { // The children's partitioning is not stable while the plan is being prepared: - // `InMemoryTableScanExec` reads `cachedPlan.outputPartitioning`, and the inner - // `AdaptiveSparkPlanExec` answers `UnknownPartitioning` until its final plan exists. The - // asymmetry matters -- an expression on one branch only keeps a `ProjectExec` from being - // collapsed away, so `comparePartitioning` rejects the pair, the union looks plain and is - // fused. Once the cache stages finalise, both children report the same `HashPartitioning`, - // and re-deriving the decision at that point left `metrics` empty while the generated code - // still incremented it, so `doProduce` threw - // `NoSuchElementException: key not found: numOutputRows`. `SELECT *` or a plain alias is - // collapsed away and does not reproduce this. + // `InMemoryTableScanExec.cachedPlan` unwraps the inner `AdaptiveSparkPlanExec` only once + // `isFinalPlan` is true, and reports `UnknownPartitioning(0)` until then, so the union looks + // plain and is fused. The projection is what makes that reachable: `supportsColumnar` is + // `children.forall`, so one row-based `ProjectExec` over the columnar scan is enough to make + // it false, and without one `supportCodegenFailureReason` reports `columnar` and nothing + // fuses. `SELECT *` or a plain alias collapses the projection away and does not reproduce + // this. Once the cache stages finalise, both children report the same `HashPartitioning`, and + // re-deriving the decision at that point left `metrics` empty while the generated code still + // incremented it, so `doProduce` threw `key not found: numOutputRows`. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { withTempView("v") { cacheAggregateView("v") From 14cf875725f89339563f58216061022d1a4ae11c Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 31 Aug 2026 17:50:12 +0800 Subject: [PATCH 03/24] [SPARK-59122][SQL] Latch the union output partitioning conf with the decision `rawPartitioning` read `spark.sql.unionOutputPartitioning` on every call and `SparkPlan.conf` is the live session conf, so a plan made with the conf on could execute with it off: the parent aggregate had already lost its exchange because the union reported a concrete `HashPartitioning`, and the union then concatenated, putting one group in two partitions and reporting it twice. Read the conf where the plain-union decision is latched instead, so what the node reports and how it executes cannot be split by a conf change. The other half -- the children answering differently later -- stays as it was. AQE skew splitting through a union does that routinely, and it is safe there only because those parents require no distribution; `doExecute` cannot tell whether anything consumed the reported partitioning, so guarding on the latch alone rejects those plans too. --- .../execution/basicPhysicalOperators.scala | 40 ++++++++++-------- .../sql/execution/UnionCodegenSuite.scala | 41 +++++++++++++++++++ 2 files changed, 64 insertions(+), 17 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index b098605589912..7cf0debc7a83d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -966,16 +966,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * The SPARK-52921 pass-through partitioning, derived from the children on every call. - * `isPlainUnion` latches on whether this comes back `UnknownPartitioning`, and - * `outputPartitioning` reports it on the branch where that latch says the union is not a plain - * concatenation. + * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` latches + * on whether this comes back `UnknownPartitioning`, and `outputPartitioning` reports it on the + * branch where that latch says the union is not a plain concatenation. The + * `UNION_OUTPUT_PARTITIONING` gate is not checked here but in the latch, so that flipping the + * conf after a node has been planned cannot change what it reports or how it executes. */ private def rawPartitioning: Partitioning = { - if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) { - return super.outputPartitioning - } - // Children's partitionings with attributes remapped to this union's output attributes. val partitionings = prepareOutputPartitioning() @@ -1064,10 +1061,16 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * the latch is published under `decisionLock` and the first writer wins; the derivation itself * runs outside the lock. The lock covers this tag only: `TreeNode`'s tag map is unsynchronized, * and `copyTagsFrom` reads it without taking it. + * + * The `UNION_OUTPUT_PARTITIONING` gate is read here rather than in `rawPartitioning`, so it is + * latched with everything else: `conf` is the live session conf, and re-reading it per call let + * a node planned with the conf on execute with it off, concatenating after a parent had already + * skipped an exchange on the strength of a concrete partitioning. */ private[sql] def isPlainUnion: Boolean = { decisionLock.synchronized(getTagValue(UnionExec.PLAIN_UNION_DECISION)).getOrElse { - val plain = rawPartitioning.isInstanceOf[UnknownPartitioning] + val plain = !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || + rawPartitioning.isInstanceOf[UnknownPartitioning] decisionLock.synchronized { getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) @@ -1084,10 +1087,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * cost is that SPARK-52921's exchange elimination is lost for a union whose children only agree * later, which is the conservative direction. * - * The other branch is still derived per call, so what this node reports is not stable across - * planning and execution: a parent can elide an exchange on a concrete partitioning and - * `doExecute` can then derive `UnknownPartitioning` and concatenate. `doExecute` keeps one - * answer within itself; the cross-phase gap is not closed here. + * The other branch is derived from the children, which can answer differently later: AQE skew + * splitting through a union does exactly that, leaving the children's partition counts divergent + * and the intersection empty, so the union concatenates at execution after reporting something + * concrete at planning. That is only safe because nothing required the reported partitioning in + * those plans, and this node cannot tell at execution time whether anything did -- guarding on + * the latch alone rejects `UnspecifiedDistribution` parents too, which is why there is no such + * guard here. */ override def outputPartitioning: Partitioning = if (isPlainUnion) super.outputPartitioning else rawPartitioning @@ -1297,10 +1303,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup override def usedInputs: AttributeSet = AttributeSet.empty protected override def doExecute(): RDD[InternalRow] = { - // One read: the non-plain branch re-derives from the children on every call, so - // `numPartitions` has to come from the same answer that picked the branch. Otherwise - // `SQLPartitioningAwareUnionRDD` can be handed an `UnknownPartitioning(0)`, which gives it - // zero partitions and this union an empty result. + // One read: the non-plain branch re-derives from the children, so `numPartitions` has to come + // from the same answer that picked the branch. Otherwise `SQLPartitioningAwareUnionRDD` can be + // handed an `UnknownPartitioning(0)`, which gives it zero partitions and this union an empty + // result. val partitioning = outputPartitioning partitioning match { case _: UnknownPartitioning | _: KeyedPartitioning => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index b0fbd144ef595..4c88834530dd3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -702,6 +703,46 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a partitioning-aware union keeps its layout when the conf changes between " + + "planning and execution") { + // `spark.sql.unionOutputPartitioning` is read when the plain-union decision is latched, not on + // every `outputPartitioning` call, so a plan is executed by the partitioning it was planned + // against. Reading it per call instead let the parent aggregate lose its exchange during + // planning (the union reported a concrete `HashPartitioning`) and then get a plain + // concatenation at execution, which puts one group in two partitions and reports it twice. + // + // The `collect()` below is deliberately outside the `withSQLConf` block that planned the + // DataFrame: `executedPlan` is memoized on first read, so this is what running a planned query + // under a changed conf looks like. Moving it back inside makes both phases see the same conf + // and the test stops exercising anything. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val left = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k") + val right = spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k") + def build(): DataFrame = + left.repartition(4, col("k")).union(right.repartition(4, col("k"))).groupBy("k").count() + + val expected = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + build().collect().toSeq + } + + val planned = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val df = build() + val plan = df.queryExecution.executedPlan + val unions = plan.collect { case u: UnionExec => u } + assert(unions.size == 1) + assert(!unions.head.isPlainUnion, + "this shape must report a concrete partitioning, or the test exercises nothing") + assert(plan.collect { case s: ShuffleExchangeExec => s }.size == 2, + "only the two repartitions may shuffle; the aggregate's exchange must have been elided") + df + } + + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + checkAnswer(planned, expected) + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From 15b2b858b958a536f79f4913538ee26dba56d9cd Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 31 Aug 2026 19:06:20 +0800 Subject: [PATCH 04/24] [SPARK-59122][SQL] Trim the UnionExec comments Comment-only. Several rounds of review each added a clause, and the blocks ended up restating the same tag-propagation mechanism three times. Keep the pointers a reader needs to navigate -- `InMemoryTableScanExec.outputPartitioning`, `CollapseCodegenStages`, `withNewChildren`/`copyTagsFrom`, `metricTerm` -- plus why `UNION_OUTPUT_PARTITIONING` is read in the latch and why `doExecute` does not guard on the latch alone. Drop the rest. --- .../execution/basicPhysicalOperators.scala | 94 +++++++------------ .../sql/execution/UnionCodegenSuite.scala | 26 ++--- 2 files changed, 44 insertions(+), 76 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 7cf0debc7a83d..afe371e0ff4d3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -966,11 +966,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` latches - * on whether this comes back `UnknownPartitioning`, and `outputPartitioning` reports it on the - * branch where that latch says the union is not a plain concatenation. The - * `UNION_OUTPUT_PARTITIONING` gate is not checked here but in the latch, so that flipping the - * conf after a node has been planned cannot change what it reports or how it executes. + * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` latches on + * whether this comes back `UnknownPartitioning`; `outputPartitioning` reports it when that latch + * says the union is not a plain concatenation. The `UNION_OUTPUT_PARTITIONING` gate lives in the + * latch, not here. */ private def rawPartitioning: Partitioning = { // Children's partitionings with attributes remapped to this union's output attributes. @@ -1027,45 +1026,27 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Guards the `isPlainUnion` latch below. A dedicated object rather than this node's own monitor: - // `unionedInputRDD` is a `lazy val`, so initializing it holds that monitor while it calls - // `inputRDDs()` on the children, which for an `InputAdapter` child is `child.execute()`. - // `@transient`, so null on a deserialized node; every reader of the latch is driver-side, as - // for `AdaptiveSparkPlanExec.lock`. + // Serializes the latch below so concurrent first readers agree on one answer. Not this node's own + // monitor, which `unionedInputRDD`'s `lazy val` holds while it drives `child.execute()`. + // Driver-only, hence `@transient`. @transient private val decisionLock = new Object() /** - * True when the codegen path applies: this union behaves as a plain concatenation, so - * `unionedInputRDD` matches the semantics of `sparkContext.union(...)` in `doExecute`. - * A `KeyedPartitioning` union also uses `sparkContext.union(...)` in `doExecute`, but - * codegen is disabled for it (`supportCodegenFailureReason` reports "partitioning-aware"): - * the per-partition key descriptor is consumed by a downstream `GroupPartitionsExec`, and - * keeping these unions out of whole-stage codegen matches the `HashPartitioning` union case. + * True when this union behaves as a plain concatenation, so `unionedInputRDD` matches + * `sparkContext.union(...)` in `doExecute` and the codegen path applies. A `KeyedPartitioning` + * union also concatenates, but codegen stays off for it: `supportCodegenFailureReason` reports + * "partitioning-aware", because a downstream `GroupPartitionsExec` consumes its key descriptor. * - * DECIDED ONCE and carried on the node rather than re-derived per caller, because the children's - * answer moves: `InMemoryTableScanExec.outputPartitioning` reads through `cachedPlan`, which - * unwraps an inner `AdaptiveSparkPlanExec` to its `executedPlan` only while that plan's - * `isFinalPlan` holds, and `withFinalPlanUpdate` clears `isFinalPlan` on entry. So a union can - * look plain when `CollapseCodegenStages` gates on it and partitioning-aware by the time the - * stage runs. + * Latched, because the answer moves under its consumers. + * `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning` while its inner + * `AdaptiveSparkPlanExec` has no final plan, so a union can look plain when + * `CollapseCodegenStages` gates on it and partitioning-aware by the time the stage runs. The + * shell that gate builds wraps a `withNewChildren` copy, and a copy that re-derives comes back + * with empty `metrics` while `doProduce` asks `metricTerm` for `numOutputRows`. A `TreeNodeTag` + * survives that rebuild where a field would not, since `withNewChildren` ends in `copyTagsFrom`. * - * That matters because `insertWholeStageCodegen` gates on this node but wraps the result of - * `insertInputAdapter`, which rebuilds the node through `withNewChildren` when its subtree - * changed. Re-deriving on the copy can contradict the decision the shell was built on: - * `metrics` on the copy comes back empty, and the `doProduce` that the shell drives throws - * `key not found: numOutputRows` out of `metricTerm`. A memo on the instance would not help, - * because the copy is a new instance; a `TreeNodeTag` does, since `withNewChildren` ends in - * `copyTagsFrom` and the copy carries no tags of its own. - * - * Concurrent first readers can derive different answers, since the children's answer moves, so - * the latch is published under `decisionLock` and the first writer wins; the derivation itself - * runs outside the lock. The lock covers this tag only: `TreeNode`'s tag map is unsynchronized, - * and `copyTagsFrom` reads it without taking it. - * - * The `UNION_OUTPUT_PARTITIONING` gate is read here rather than in `rawPartitioning`, so it is - * latched with everything else: `conf` is the live session conf, and re-reading it per call let - * a node planned with the conf on execute with it off, concatenating after a parent had already - * skipped an exchange on the strength of a concrete partitioning. + * `UNION_OUTPUT_PARTITIONING` is read here rather than in `rawPartitioning` so it is latched too: + * `conf` is live, and re-reading it let a plan made with the conf on execute with it off. */ private[sql] def isPlainUnion: Boolean = { decisionLock.synchronized(getTagValue(UnionExec.PLAIN_UNION_DECISION)).getOrElse { @@ -1081,19 +1062,14 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * A node latched plain reports `UnknownPartitioning` and concatenates, even if its children - * have since agreed on a concrete partitioning -- a fused union concatenates its children's - * partitions, and claiming their partitioning would let a parent skip an exchange it needs. The - * cost is that SPARK-52921's exchange elimination is lost for a union whose children only agree - * later, which is the conservative direction. + * A node latched plain reports `UnknownPartitioning` even once its children agree on a concrete + * one: a fused union concatenates, and claiming their partitioning would let a parent skip an + * exchange it needs. The cost is SPARK-52921's exchange elimination for such a union. * - * The other branch is derived from the children, which can answer differently later: AQE skew - * splitting through a union does exactly that, leaving the children's partition counts divergent - * and the intersection empty, so the union concatenates at execution after reporting something - * concrete at planning. That is only safe because nothing required the reported partitioning in - * those plans, and this node cannot tell at execution time whether anything did -- guarding on - * the latch alone rejects `UnspecifiedDistribution` parents too, which is why there is no such - * guard here. + * The other branch is derived per call and can come back `UnknownPartitioning` later -- AQE skew + * splitting through a union leaves the children's partition counts divergent. Failing there was + * tried and reverted: nothing in those plans required the reported partitioning, and this node + * cannot tell at execution time whether anything did. */ override def outputPartitioning: Partitioning = if (isPlainUnion) super.outputPartitioning else rawPartitioning @@ -1114,11 +1090,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Memoized so that `supportCodegen` (called multiple times by - // `CollapseCodegenStages`) and `metrics` see one reason on one instance: - // `conf` is the live session conf, so re-deriving could answer differently. - // Agreeing with the copy `withNewChildren` makes is a separate matter, and - // comes from `isPlainUnion` being latched on a `TreeNodeTag`. + // Memoized so `supportCodegen` (called repeatedly by `CollapseCodegenStages`) + // and `metrics` see one reason on one instance; `conf` is live, so re-deriving + // could answer differently. Agreeing with the `withNewChildren` copy is the + // `isPlainUnion` tag's job, not this memo's. @transient private lazy val supportCodegenFailureReason: Option[String] = { if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) { Some("union-codegen-disabled") @@ -1303,10 +1278,9 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup override def usedInputs: AttributeSet = AttributeSet.empty protected override def doExecute(): RDD[InternalRow] = { - // One read: the non-plain branch re-derives from the children, so `numPartitions` has to come - // from the same answer that picked the branch. Otherwise `SQLPartitioningAwareUnionRDD` can be - // handed an `UnknownPartitioning(0)`, which gives it zero partitions and this union an empty - // result. + // One read: the non-plain branch is derived per call, so the branch and `numPartitions` have to + // come from the same answer -- an `UnknownPartitioning(0)` here would give + // `SQLPartitioningAwareUnionRDD` zero partitions and this union an empty result. val partitioning = outputPartitioning partitioning match { case _: UnknownPartitioning | _: KeyedPartitioning => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 4c88834530dd3..839e466d66360 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -62,9 +62,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper }.nonEmpty /** - * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query stages, which - * `SparkPlan.collect` does not: both are `LeafExecNode`s, so it cannot see a union that AQE - * placed inside a stage. + * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query stages; + * `SparkPlan.collect` stops at them, since both are `LeafExecNode`s. */ private def fusedUnions(df: DataFrame): Seq[UnionExec] = collect(df.queryExecution.executedPlan) { @@ -79,9 +78,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper .groupBy("k") .agg(sum("v").as("s")) .createOrReplaceTempView(view) - // Both callers need the cache unmaterialized, and `CacheManager` treats caching an - // already-cached plan as a no-op, so drop whatever an earlier test left for this plan. - // `isCached` matches by plan, so it also catches the same plan cached under another name. + // Both callers need the cache unmaterialized, and `CacheManager` no-ops on an already-cached + // plan, so drop whatever an earlier test left for this one. `isCached` matches by plan. if (spark.catalog.isCached(view)) spark.catalog.uncacheTable(view) spark.catalog.cacheTable(view) } @@ -705,16 +703,12 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper test("SPARK-59122: a partitioning-aware union keeps its layout when the conf changes between " + "planning and execution") { - // `spark.sql.unionOutputPartitioning` is read when the plain-union decision is latched, not on - // every `outputPartitioning` call, so a plan is executed by the partitioning it was planned - // against. Reading it per call instead let the parent aggregate lose its exchange during - // planning (the union reported a concrete `HashPartitioning`) and then get a plain - // concatenation at execution, which puts one group in two partitions and reports it twice. - // - // The `collect()` below is deliberately outside the `withSQLConf` block that planned the - // DataFrame: `executedPlan` is memoized on first read, so this is what running a planned query - // under a changed conf looks like. Moving it back inside makes both phases see the same conf - // and the test stops exercising anything. + // `spark.sql.unionOutputPartitioning` is read where the plain-union decision is latched, not on + // every `outputPartitioning` call, so a plan executes by the partitioning it was planned + // against. Reading it per call let the parent aggregate lose its exchange at planning and get a + // plain concatenation at execution, reporting each group twice. The `collect()` stays outside + // the block that planned the DataFrame on purpose: `executedPlan` is memoized on first read, + // and moving it back inside makes both phases see the same conf. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val left = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k") val right = spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k") From 9da6b20394e36ca7f44942c560f91ec935d03ed6 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 3 Sep 2026 13:13:38 +0800 Subject: [PATCH 05/24] Address the remaining review comments on the plain-union latch --- .../execution/basicPhysicalOperators.scala | 31 ++++++------ .../sql/execution/UnionCodegenSuite.scala | 49 +++++++------------ 2 files changed, 32 insertions(+), 48 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index afe371e0ff4d3..e943dcfcb5f0a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1026,9 +1026,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Serializes the latch below so concurrent first readers agree on one answer. Not this node's own - // monitor, which `unionedInputRDD`'s `lazy val` holds while it drives `child.execute()`. - // Driver-only, hence `@transient`. + // Serializes the latch below so concurrent first readers agree on one answer. Held across the + // derivation, which only walks the children, so nothing below can be waiting on it. Not this + // node's own monitor, which `unionedInputRDD`'s `lazy val` holds while it drives + // `child.execute()`. Driver-only, hence `@transient`. @transient private val decisionLock = new Object() /** @@ -1048,16 +1049,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * `UNION_OUTPUT_PARTITIONING` is read here rather than in `rawPartitioning` so it is latched too: * `conf` is live, and re-reading it let a plan made with the conf on execute with it off. */ - private[sql] def isPlainUnion: Boolean = { - decisionLock.synchronized(getTagValue(UnionExec.PLAIN_UNION_DECISION)).getOrElse { + private[sql] def isPlainUnion: Boolean = decisionLock.synchronized { + getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { val plain = !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || rawPartitioning.isInstanceOf[UnknownPartitioning] - decisionLock.synchronized { - getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { - setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) - plain - } - } + setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) + plain } } @@ -1066,10 +1063,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * one: a fused union concatenates, and claiming their partitioning would let a parent skip an * exchange it needs. The cost is SPARK-52921's exchange elimination for such a union. * - * The other branch is derived per call and can come back `UnknownPartitioning` later -- AQE skew - * splitting through a union leaves the children's partition counts divergent. Failing there was - * tried and reverted: nothing in those plans required the reported partitioning, and this node - * cannot tell at execution time whether anything did. + * The other branch is derived per call and can come back `UnknownPartitioning` later, since AQE + * skew splitting through a union leaves the children's partition counts divergent. Such a union + * concatenates, which is tolerated because the plans it arises in ask nothing of its + * partitioning. */ override def outputPartitioning: Partitioning = if (isPlainUnion) super.outputPartitioning else rawPartitioning @@ -1092,8 +1089,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // Memoized so `supportCodegen` (called repeatedly by `CollapseCodegenStages`) // and `metrics` see one reason on one instance; `conf` is live, so re-deriving - // could answer differently. Agreeing with the `withNewChildren` copy is the - // `isPlainUnion` tag's job, not this memo's. + // could answer differently. The `withNewChildren` copy gets its own memo, and + // agrees on the `isPlainUnion` term because that one is latched on a tag. @transient private lazy val supportCodegenFailureReason: Option[String] = { if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) { Some("union-codegen-disabled") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 839e466d66360..3ea1903556585 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -71,16 +71,17 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper w.child.asInstanceOf[UnionExec] } - /** A cached aggregate, so the union's children read an `InMemoryTableScanExec`. */ + /** + * A cached aggregate, so the union's children read an `InMemoryTableScanExec`. The caller needs + * the cache unmaterialized; `withTempView` uncaches this plan on the way out, so nothing is left + * for the next caller to trip over. + */ private def cacheAggregateView(view: String): Unit = { spark.range(0, 200, 1, 4) .selectExpr("id % 10 AS k", "id AS v") .groupBy("k") .agg(sum("v").as("s")) .createOrReplaceTempView(view) - // Both callers need the cache unmaterialized, and `CacheManager` no-ops on an already-cached - // plan, so drop whatever an earlier test left for this one. `isCached` matches by plan. - if (spark.catalog.isCached(view)) spark.catalog.uncacheTable(view) spark.catalog.cacheTable(view) } @@ -654,7 +655,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } - test("SPARK-59122: a fused union keeps numOutputRows when a child's partitioning firms up") { + test("SPARK-59122: a fused union keeps numOutputRows and reports UnknownPartitioning") { // The children's partitioning is not stable while the plan is being prepared: // `InMemoryTableScanExec.cachedPlan` unwraps the inner `AdaptiveSparkPlanExec` only once // `isFinalPlan` is true, and reports `UnknownPartitioning(0)` until then, so the union looks @@ -665,6 +666,11 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // this. Once the cache stages finalise, both children report the same `HashPartitioning`, and // re-deriving the decision at that point left `metrics` empty while the generated code still // incremented it, so `doProduce` threw `key not found: numOutputRows`. + // + // Both halves of the decision are asserted here. Registering the metric unconditionally would + // fix the crash and leave the other half broken: a fused union concatenates its children's + // partitions, so claiming their `HashPartitioning` would let a parent satisfy a clustered + // distribution from an RDD that does not have it. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { withTempView("v") { cacheAggregateView("v") @@ -675,25 +681,9 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val fused = fusedUnions(df) assert(fused.nonEmpty, "this shape must actually fuse, or the test is not exercising the defect") - assert(fused.forall(_.metrics.contains("numOutputRows")), - "a fused union must register the metric its generated code increments") - } - } - } - - test("SPARK-59122: a fused union reports UnknownPartitioning, so no parent skips an exchange") { - // The other half of the same decision. A fused union concatenates its children's partitions, - // so if it went on claiming the children's `HashPartitioning` a parent could satisfy a - // clustered distribution from an RDD that does not have it -- a wrong answer rather than a - // crash, which is why the missing metric must not simply be registered unconditionally. - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { - withTempView("v") { - cacheAggregateView("v") - val df = spark.sql("SELECT k, abs(s) AS s FROM v UNION ALL SELECT k, s FROM v") - df.collect() - val fused = fusedUnions(df) - assert(fused.nonEmpty, "this shape must actually fuse") fused.foreach { u => + assert(u.metrics.contains("numOutputRows"), + "a fused union must register the metric its generated code increments") assert(u.outputPartitioning.isInstanceOf[UnknownPartitioning], s"a fused union must not claim a concrete partitioning, got ${u.outputPartitioning}") } @@ -712,15 +702,10 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val left = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k") val right = spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k") - def build(): DataFrame = - left.repartition(4, col("k")).union(right.repartition(4, col("k"))).groupBy("k").count() - - val expected = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { - build().collect().toSeq - } val planned = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { - val df = build() + val df = left.repartition(4, col("k")) + .union(right.repartition(4, col("k"))).groupBy("k").count() val plan = df.queryExecution.executedPlan val unions = plan.collect { case u: UnionExec => u } assert(unions.size == 1) @@ -732,7 +717,9 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { - checkAnswer(planned, expected) + // Each side contributes four ids per `k`, so the answer is fixed. Comparing against the + // same query run with the conf off would also pass if both paths regressed to ten rows. + checkAnswer(planned, (0L until 5L).map(k => Row(k, 8L))) } } } From 7cdb29cbfbc7aa94ea66f3c611a963231b0b796b Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 5 Sep 2026 13:56:39 +0800 Subject: [PATCH 06/24] [SPARK-59122][SQL] Tighten the latch comments and the two tests --- .../apache/spark/sql/internal/SQLConf.scala | 3 +- .../execution/basicPhysicalOperators.scala | 48 ++++++++++++------- .../sql/execution/UnionCodegenSuite.scala | 25 ++++++---- 3 files changed, 51 insertions(+), 25 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index fececcda81518..d0b8017c1007b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -8100,7 +8100,8 @@ object SQLConf { .internal() .doc("When set to true, the output partitioning of UnionExec will be the same as the " + "input partitioning if its children have same partitioning. Otherwise, it will be a " + - "default partitioning.") + "default partitioning. The value is read when a UnionExec first decides, so a change " + + "applies to plans built afterwards, not to a plan that has already been prepared.") .version("4.1.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 09647d644b9ae..5e94ec9af2ea8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1023,17 +1023,21 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Serializes the latch below so concurrent first readers agree on one answer. Held across the - // derivation, which only walks the children, so nothing below can be waiting on it. Not this - // node's own monitor, which `unionedInputRDD`'s `lazy val` holds while it drives - // `child.execute()`. Driver-only, hence `@transient`. + // Serializes the latch below so concurrent first readers agree on one answer. Private, and + // `isPlainUnion` is its only user, so the only lock taken under it is a nested union's own + // `decisionLock`, always a descendant's. It has to stay that way: no child `outputPartitioning` + // the derivation walks may take a lock, or it would invert `CoalesceShufflePartitions`, which + // reads `isPlainUnion` while holding the AQE lock. `InMemoryTableScanExec` qualifies only + // because it reads `adaptive.executedPlan`, a volatile read, not `finalPhysicalPlan`, which is + // `lock.synchronized`. Driver-only, hence `@transient`. @transient private val decisionLock = new Object() /** - * True when this union behaves as a plain concatenation, so `unionedInputRDD` matches - * `sparkContext.union(...)` in `unionRDDs` and the codegen path applies. A `KeyedPartitioning` - * union also concatenates, but codegen stays off for it: `supportCodegenFailureReason` reports - * "partitioning-aware", because a downstream `GroupPartitionsExec` consumes its key descriptor. + * True when this union behaves as a plain concatenation, so `unionedInputRDD` matches the + * semantics of `sparkContext.union(...)` in `unionRDDs`, and the codegen path applies. A + * `KeyedPartitioning` union also concatenates, but codegen stays off for it: + * `supportCodegenFailureReason` reports "partitioning-aware", because a downstream + * `GroupPartitionsExec` consumes its key descriptor. * * Latched, because the answer moves under its consumers. * `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning` while its inner @@ -1044,7 +1048,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * survives that rebuild where a field would not, since `withNewChildren` ends in `copyTagsFrom`. * * `UNION_OUTPUT_PARTITIONING` is read here rather than in `rawPartitioning` so it is latched too: - * `conf` is live, and re-reading it let a plan made with the conf on execute with it off. + * `conf` is live, and a plan must execute by the partitioning it was planned against. */ private[sql] def isPlainUnion: Boolean = decisionLock.synchronized { getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { @@ -1060,10 +1064,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * one: a fused union concatenates, and claiming their partitioning would let a parent skip an * exchange it needs. The cost is SPARK-52921's exchange elimination for such a union. * - * The other branch is derived per call and can come back `UnknownPartitioning` later, since AQE - * skew splitting through a union leaves the children's partition counts divergent. Such a union - * concatenates, which is tolerated because the plans it arises in ask nothing of its - * partitioning. + * The other branch is derived per call, so `unionRDDs` could take the concatenating arm even + * though `EnsureRequirements` planned the parent against a concrete partitioning: + * `comparePartitioning` compares `HashPartitioningLike` by equality, so a change to one child's + * partitioning that its siblings do not mirror can empty the intersection. This node does not + * re-check it. AQE reconciles it, by validating a partitioning change against the parents' + * requirements and either reverting it or re-running `EnsureRequirements`; an injected rule can + * skip that. */ override def outputPartitioning: Partitioning = if (isPlainUnion) super.outputPartitioning else rawPartitioning @@ -1086,8 +1093,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // Memoized so `supportCodegen` (called repeatedly by `CollapseCodegenStages`) // and `metrics` see one reason on one instance; `conf` is live, so re-deriving - // could answer differently. The `withNewChildren` copy gets its own memo, and - // agrees on the `isPlainUnion` term because that one is latched on a tag. + // could answer differently. The `withNewChildren` copy gets its own memo; it + // re-derives every term but `isPlainUnion`, which it inherits from the tag when + // the original latched first, as the codegen path does (`insertWholeStageCodegen` + // calls `supportCodegen` before `insertInputAdapter` takes the copy). @transient private lazy val supportCodegenFailureReason: Option[String] = { if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) { Some("union-codegen-disabled") @@ -1309,7 +1318,14 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } object UnionExec { - /** The latched "is this a plain concatenation" decision. See `isPlainUnion`. */ + /** + * The latched "is this a plain concatenation" decision. See `isPlainUnion`. + * + * Rebuilds carry it: `withNewChildren` and a transform rule's replacement both go through + * `copyTagsFrom`. A `UnionExec` with no predecessor node at its position starts unlatched: it + * can re-derive the opposite answer and leave `metrics` empty under generated code that + * increments it. + */ private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 3ea1903556585..76e90d5ccce93 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -64,6 +64,10 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper /** * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query stages; * `SparkPlan.collect` stops at them, since both are `LeafExecNode`s. + * + * Stricter than `unionInsideWSCG` on purpose: `w.find` also matches a union that an + * `InputAdapter` left inside the stage unfused, and the callers here assert on `metrics`, which + * an unfused union does not register. */ private def fusedUnions(df: DataFrame): Seq[UnionExec] = collect(df.queryExecution.executedPlan) { @@ -73,8 +77,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper /** * A cached aggregate, so the union's children read an `InMemoryTableScanExec`. The caller needs - * the cache unmaterialized; `withTempView` uncaches this plan on the way out, so nothing is left - * for the next caller to trip over. + * the cache unmaterialized; `withTempView` drops the view on the way out, and `dropTempView` + * uncaches this view's plan. */ private def cacheAggregateView(view: String): Unit = { spark.range(0, 200, 1, 4) @@ -671,7 +675,9 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // fix the crash and leave the other half broken: a fused union concatenates its children's // partitions, so claiming their `HashPartitioning` would let a parent satisfy a clustered // distribution from an RDD that does not have it. - withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { withTempView("v") { cacheAggregateView("v") val df = spark.sql("SELECT k, abs(s) AS s FROM v UNION ALL SELECT k, s FROM v") @@ -696,9 +702,10 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // `spark.sql.unionOutputPartitioning` is read where the plain-union decision is latched, not on // every `outputPartitioning` call, so a plan executes by the partitioning it was planned // against. Reading it per call let the parent aggregate lose its exchange at planning and get a - // plain concatenation at execution, reporting each group twice. The `collect()` stays outside - // the block that planned the DataFrame on purpose: `executedPlan` is memoized on first read, - // and moving it back inside makes both phases see the same conf. + // plain concatenation at execution, reporting each group twice. The `checkAnswer` below stays + // outside the block that planned the DataFrame on purpose: the plan is forced inside that + // block and `executedPlan` is memoized, so the two phases see different confs. Asserting + // inside it, or dropping the second `withSQLConf`, makes the test pass without testing this. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val left = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k") val right = spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k") @@ -709,8 +716,10 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val plan = df.queryExecution.executedPlan val unions = plan.collect { case u: UnionExec => u } assert(unions.size == 1) - assert(!unions.head.isPlainUnion, - "this shape must report a concrete partitioning, or the test exercises nothing") + // Not asserted through `isPlainUnion`: that call latches the decision, which would warm + // a field-based implementation's memo and hide the regression this test is for. The + // exchange count below proves the union reported a concrete partitioning, without + // touching the node. assert(plan.collect { case s: ShuffleExchangeExec => s }.size == 2, "only the two repartitions may shuffle; the aggregate's exchange must have been elided") df From 543684188d1f6e961e03dd7eba3629c5615ca9a5 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 5 Sep 2026 20:00:24 +0800 Subject: [PATCH 07/24] Latch the whole-stage-codegen decision alongside the plain-union one supportCodegenFailureReason read live confs, so the copy insertInputAdapter puts inside the codegen shell could answer differently than the gate did and leave metrics empty under generated code that asks metricTerm for numOutputRows. Store the reason in a TreeNodeTag under decisionLock, the same way isPlainUnion is latched. --- .../execution/basicPhysicalOperators.scala | 42 ++++++++++++------- .../sql/execution/UnionCodegenSuite.scala | 25 +++++++++++ 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 5e94ec9af2ea8..5597b7bab3bfc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1023,13 +1023,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Serializes the latch below so concurrent first readers agree on one answer. Private, and - // `isPlainUnion` is its only user, so the only lock taken under it is a nested union's own - // `decisionLock`, always a descendant's. It has to stay that way: no child `outputPartitioning` - // the derivation walks may take a lock, or it would invert `CoalesceShufflePartitions`, which - // reads `isPlainUnion` while holding the AQE lock. `InMemoryTableScanExec` qualifies only - // because it reads `adaptive.executedPlan`, a volatile read, not `finalPhysicalPlan`, which is - // `lock.synchronized`. Driver-only, hence `@transient`. + // Serializes the two latches below so concurrent first readers agree on one answer. Private to + // this node, so the only lock taken under it is a nested union's own `decisionLock`, always a + // descendant's. It has to stay that way: nothing either derivation walks may take a lock, or it + // would invert `CoalesceShufflePartitions`, which reads `isPlainUnion` while holding the AQE + // lock. That surface is the children's `outputPartitioning`, `supportsColumnar` and `output`; + // `InMemoryTableScanExec` qualifies only because it reads `adaptive.executedPlan`, a volatile + // read, not `finalPhysicalPlan`, which is `lock.synchronized`. Driver-only, hence `@transient`. @transient private val decisionLock = new Object() /** @@ -1091,13 +1091,21 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Memoized so `supportCodegen` (called repeatedly by `CollapseCodegenStages`) - // and `metrics` see one reason on one instance; `conf` is live, so re-deriving - // could answer differently. The `withNewChildren` copy gets its own memo; it - // re-derives every term but `isPlainUnion`, which it inherits from the tag when - // the original latched first, as the codegen path does (`insertWholeStageCodegen` - // calls `supportCodegen` before `insertInputAdapter` takes the copy). - @transient private lazy val supportCodegenFailureReason: Option[String] = { + // Latched for the same reason `isPlainUnion` is: `supportCodegen` and `metrics` must see one + // answer, and `conf` is live. When a child is not `CodegenSupport`, `insertInputAdapter` wraps + // it, so `withNewChildren` returns a real copy whose first evaluation of this would land at + // execution; re-deriving there left `metrics` empty while `doProduce` asked `metricTerm` for + // `numOutputRows`. The first force is not always the gate: under AQE it is a plan-update event + // on the pre-stage-creation tree, so a term added here sees more of the plan than the gate does. + private def supportCodegenFailureReason: Option[String] = decisionLock.synchronized { + getTagValue(UnionExec.CODEGEN_FAILURE_REASON).getOrElse { + val reason = deriveCodegenFailureReason + setTagValue(UnionExec.CODEGEN_FAILURE_REASON, reason) + reason + } + } + + private def deriveCodegenFailureReason: Option[String] = { if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) { Some("union-codegen-disabled") } else if (!isPlainUnion) { @@ -1328,6 +1336,12 @@ object UnionExec { */ private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") + /** + * The latched whole-stage-codegen decision. See `supportCodegenFailureReason`. Carried across + * rebuilds on the same terms as [[PLAIN_UNION_DECISION]]. + */ + private val CODEGEN_FAILURE_REASON = TreeNodeTag[Option[String]]("codegenFailureReason") + /** * Codegen operators that return more than one RDD from `inputRDDs()`. * `UnionExec`'s fusion assumes each direct child contributes one RDD. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 76e90d5ccce93..43c21c9e1734c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -733,6 +733,31 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a fused union keeps numOutputRows when the codegen conf changes between " + + "planning and execution") { + // `supportCodegenFailureReason` used to read `WHOLESTAGE_UNION_CODEGEN_ENABLED` live, and the + // copy that `insertInputAdapter` puts inside the codegen shell evaluated it for the first time + // at execution. Planned with the conf on the union is fused, so the generated code increments + // `numOutputRows`; if the copy re-derives the reason with the conf off, `metrics` comes back + // empty and `doProduce` throws `key not found: numOutputRows`. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val planned = withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") { + // Each child is an exchange, which is not `CodegenSupport`, so `insertInputAdapter` wraps + // it and `withNewChildren` really does produce a copy. With codegen-support children it + // returns `this`, the memo stays the original's warm one, and nothing re-derives. + val df = rangeDF(100).repartition(2).union(rangeDF(100).repartition(2)) + // `fusedUnions` requires the union to be the stage root; `unionInsideWSCG` would also + // match a union that an `InputAdapter` left inside the stage unfused, which is exactly + // the degradation this guard has to catch. + assert(fusedUnions(df).size == 1, "this shape must fuse, or the test exercises nothing") + df + } + withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { + assert(planned.collect().length == 200) + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From ef44aa21596bebdedd689609804d66a4941998d6 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 10 Sep 2026 23:27:17 +0800 Subject: [PATCH 08/24] Latch only the conf half of the codegen decision The reason itself is derived from the children, and a TreeNodeTag outlives them: SparkPlanInfo forces metrics on an AQE plan update, before the rules that run ahead of CollapseCodegenStages, so a rule replacing a child there inherited an allowing answer. Keep the two confs in a tag and memoize the reason per instance instead. Also narrow the decisionLock lock-order comment and the isPlainUnion scaladoc to what the code supports. --- .../execution/basicPhysicalOperators.scala | 76 +++++++++++-------- .../sql/execution/UnionCodegenSuite.scala | 25 ++++++ 2 files changed, 70 insertions(+), 31 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 5597b7bab3bfc..30f2d9a0580c5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1023,29 +1023,30 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Serializes the two latches below so concurrent first readers agree on one answer. Private to - // this node, so the only lock taken under it is a nested union's own `decisionLock`, always a - // descendant's. It has to stay that way: nothing either derivation walks may take a lock, or it - // would invert `CoalesceShufflePartitions`, which reads `isPlainUnion` while holding the AQE - // lock. That surface is the children's `outputPartitioning`, `supportsColumnar` and `output`; - // `InMemoryTableScanExec` qualifies only because it reads `adaptive.executedPlan`, a volatile - // read, not `finalPhysicalPlan`, which is `lock.synchronized`. Driver-only, hence `@transient`. + // Serializes the latches below so concurrent first readers agree on one answer. What must not be + // reachable under it is the AQE final-plan lock: `CoalesceShufflePartitions` reads `isPlainUnion` + // while holding that lock, so the reverse edge deadlocks. `InMemoryTableScanExec` is safe here + // because it reads `adaptive.executedPlan`, a volatile read, not `finalPhysicalPlan`, which is + // `lock.synchronized`. Child `lazy val`s such as `AQEShuffleReadExec.outputPartitioning` do take + // their own instance monitor; that is harmless, as nothing holds one of those and then waits + // here. Driver-only, hence `@transient`. @transient private val decisionLock = new Object() /** * True when this union behaves as a plain concatenation, so `unionedInputRDD` matches the - * semantics of `sparkContext.union(...)` in `unionRDDs`, and the codegen path applies. A - * `KeyedPartitioning` union also concatenates, but codegen stays off for it: - * `supportCodegenFailureReason` reports "partitioning-aware", because a downstream - * `GroupPartitionsExec` consumes its key descriptor. + * semantics of `sparkContext.union(...)` in `unionRDDs`. It satisfies the partitioning gate on + * the codegen path, not the whole of it: `supportCodegenFailureReason` still applies its other + * checks. A `KeyedPartitioning` union also concatenates, but codegen stays off for it, with the + * reason "partitioning-aware", because a downstream `GroupPartitionsExec` consumes its key + * descriptor. * * Latched, because the answer moves under its consumers. * `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning` while its inner * `AdaptiveSparkPlanExec` has no final plan, so a union can look plain when * `CollapseCodegenStages` gates on it and partitioning-aware by the time the stage runs. The - * shell that gate builds wraps a `withNewChildren` copy, and a copy that re-derives comes back - * with empty `metrics` while `doProduce` asks `metricTerm` for `numOutputRows`. A `TreeNodeTag` - * survives that rebuild where a field would not, since `withNewChildren` ends in `copyTagsFrom`. + * shell that gate builds wraps a `withNewChildren` copy, and a copy that re-derived here came + * back with empty `metrics` while `doProduce` asked `metricTerm` for `numOutputRows`. The tag + * carries the answer onto that copy, since `withNewChildren` ends in `copyTagsFrom`. * * `UNION_OUTPUT_PARTITIONING` is read here rather than in `rawPartitioning` so it is latched too: * `conf` is live, and a plan must execute by the partitioning it was planned against. @@ -1091,22 +1092,31 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Latched for the same reason `isPlainUnion` is: `supportCodegen` and `metrics` must see one - // answer, and `conf` is live. When a child is not `CodegenSupport`, `insertInputAdapter` wraps - // it, so `withNewChildren` returns a real copy whose first evaluation of this would land at - // execution; re-deriving there left `metrics` empty while `doProduce` asked `metricTerm` for - // `numOutputRows`. The first force is not always the gate: under AQE it is a plan-update event - // on the pre-stage-creation tree, so a term added here sees more of the plan than the gate does. - private def supportCodegenFailureReason: Option[String] = decisionLock.synchronized { - getTagValue(UnionExec.CODEGEN_FAILURE_REASON).getOrElse { - val reason = deriveCodegenFailureReason - setTagValue(UnionExec.CODEGEN_FAILURE_REASON, reason) - reason + // The conf half of the codegen decision, latched for the reason `isPlainUnion` is: `conf` is + // live, so the gate, `metrics` and the copy `insertInputAdapter` puts inside the codegen shell + // would otherwise be free to read different values. When a child is not `CodegenSupport` that + // copy is real and its first evaluation lands at execution; reading the conf there left + // `metrics` empty while `doProduce` asked `metricTerm` for `numOutputRows`. + private def codegenConfSnapshot: UnionExec.CodegenConfSnapshot = decisionLock.synchronized { + getTagValue(UnionExec.CODEGEN_CONF_SNAPSHOT).getOrElse { + val snapshot = UnionExec.CodegenConfSnapshot( + conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED), + conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) + setTagValue(UnionExec.CODEGEN_CONF_SNAPSHOT, snapshot) + snapshot } } - private def deriveCodegenFailureReason: Option[String] = { - if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) { + // Memoized per instance, not latched on a tag: every term below the confs reads the children, and + // a tag outlives them. `SparkPlanInfo` forces `metrics` on an AQE plan update, before the rules + // that run ahead of `CollapseCodegenStages`, so a rule replacing a child there would inherit an + // allowing answer and fuse a topology that `hasPartitionIndexDependentCodegen` or + // `supportsColumnar` rejects. The copy in the codegen shell still agrees with the gate: + // `InputAdapter` delegates `output` and `supportsColumnar` to its child, the remaining terms walk + // the subtree through it, and each is fixed for a given set of children. + @transient private lazy val supportCodegenFailureReason: Option[String] = { + val confs = codegenConfSnapshot + if (!confs.unionCodegenEnabled) { Some("union-codegen-disabled") } else if (!isPlainUnion) { Some("partitioning-aware") @@ -1116,7 +1126,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup Some("multi-rdd-child") } else if (children.exists(UnionExec.hasPartitionIndexDependentCodegen)) { Some("partition-index-dependent-child") - } else if (children.size > conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) { + } else if (children.size > confs.maxChildren) { Some("max-children-exceeded") } else if (supportsColumnar) { Some("columnar") @@ -1337,10 +1347,14 @@ object UnionExec { private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") /** - * The latched whole-stage-codegen decision. See `supportCodegenFailureReason`. Carried across - * rebuilds on the same terms as [[PLAIN_UNION_DECISION]]. + * The confs `supportCodegenFailureReason` reads, latched by `codegenConfSnapshot`. Only the conf + * values are latched. The terms that read the children are derived per instance instead, so that + * a rule replacing a child cannot inherit an answer taken from the topology it replaced. */ - private val CODEGEN_FAILURE_REASON = TreeNodeTag[Option[String]]("codegenFailureReason") + private case class CodegenConfSnapshot(unionCodegenEnabled: Boolean, maxChildren: Int) + + private val CODEGEN_CONF_SNAPSHOT = + TreeNodeTag[CodegenConfSnapshot]("codegenConfSnapshot") /** * Codegen operators that return more than one RDD from `inputRDDs()`. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 43c21c9e1734c..a0ffa2b32d809 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -758,6 +758,31 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: the codegen gate re-derives when a rule replaces the children") { + // The gate's children-dependent terms must not outlive the children they were taken from. + // `SparkPlanInfo` reads `metrics` on every node when AQE posts a plan update, and that happens + // before the rules running just ahead of `CollapseCodegenStages`; a decision carried from there + // onto the rebuilt node would fuse a topology that the gate rejects. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") { + val df = rangeDF(100).union(rangeDF(100)) + val unions = fusedUnions(df) + assert(unions.size == 1, "this shape must fuse, or the test exercises nothing") + val union = unions.head + // What the plan update does, and what decides the gate for this instance. + assert(union.metrics.contains("numOutputRows")) + assert(union.supportCodegen) + + // A nested union is one of the topologies the gate rejects, and `withNewChildren` is the path + // every rule takes to install a replacement. + val nested = UnionExec(Seq(union.children.head, union.children.head)) + val rebuilt = union.withNewChildren(Seq(nested, union.children.last)).asInstanceOf[UnionExec] + assert(!rebuilt.supportCodegen, "the rebuilt union must answer against its own children") + assert(rebuilt.metrics.isEmpty) + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From 3ed87a98166f2a09ccaf8a27aec87cf406a0dfac Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 12 Sep 2026 13:59:53 +0800 Subject: [PATCH 09/24] Stamp the union decisions during preparation instead of on first read Reading outputPartitioning on queryExecution.sparkPlan used to write the decision, and clone -> makeCopy -> copyTagsFrom carried it into executedPlan, so observing the unprepared plan decided for the prepared one. With nothing reading it during preparation the first read landed at execution instead, where a conf flip decided. A new StampUnionDecisions rule, right after EnsureRequirements in both the standard and the AQE pipeline, asks once; reads before it answer without writing. Also check isPlainUnion last in the codegen gate, so a union rejected on another ground does not fix a decision it has no use for, and narrow four comments to what the code supports. --- .../apache/spark/sql/internal/SQLConf.scala | 8 +- .../spark/sql/execution/QueryExecution.scala | 3 + .../sql/execution/StampUnionDecisions.scala | 42 +++++++ .../adaptive/AdaptiveSparkPlanExec.scala | 4 + .../execution/basicPhysicalOperators.scala | 94 ++++++++++------ .../sql/execution/UnionCodegenSuite.scala | 106 +++++++++++++++--- 6 files changed, 204 insertions(+), 53 deletions(-) create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index e24edc9b33123..7f7f073cc2f4f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2908,7 +2908,9 @@ object SQLConf { .doc("When both this conf and `spark.sql.codegen.wholeStage` are true, " + "UnionExec participates in whole-stage codegen on its " + "non-partitioning-aware path: the parent and all children fuse into " + - "a single WholeStageCodegenExec stage.") + "a single WholeStageCodegenExec stage. The value is read when a UnionExec " + + "first decides, so a change applies to plans built afterwards, not to a " + + "plan that has already been prepared.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .booleanConf @@ -2923,7 +2925,9 @@ object SQLConf { "bytecode size, constant pool growth, JIT compilation time) rather " + "than the JVM per-method bytecode limit. Unions with more children " + "fall back to per-child codegen stages. Only effective when " + - s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true.") + s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true. The value is read when a " + + "UnionExec first decides, so a change applies to plans built afterwards, not " + + "to a plan that has already been prepared.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .intConf diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index 8be200dfc5bb8..d9d122844ee8b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -799,6 +799,9 @@ object QueryExecution { PlanSubqueries(sparkSession), RemoveRedundantProjects, EnsureRequirements(), + // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, and + // the answer to fix is the one the exchanges around it were planned against. + StampUnionDecisions, // This rule must be run after `EnsureRequirements`. InsertSortForLimitAndOffset, // `PushDownLocalSort` pushes a wider local sort down onto a narrower one below it, so a diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala new file mode 100644 index 0000000000000..c2f4539355655 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution + +import org.apache.spark.sql.catalyst.rules.Rule + +/** + * Fixes each [[UnionExec]]'s partitioning decision and codegen conf snapshot at one defined point. + * + * `UnionExec` derives both from state that moves: its children's `outputPartitioning` sharpens as + * AQE finalises the plans behind them, and `conf` is the live session conf. Whoever asked first + * used to decide, which made the answer depend on when it was observed. This rule asks once, + * right after `EnsureRequirements`, so the partitioning a parent's exchange decision was taken + * from is the one `unionRDDs` and the codegen gate use. + * + * It only writes what is not there yet, so re-running it (AQE re-optimizes each round) keeps the + * first answer, and a node rebuilt from a stamped one keeps the tags `copyTagsFrom` gave it. + */ +object StampUnionDecisions extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + plan.foreach { + case u: UnionExec => u.stampDecisions() + case _ => + } + plan + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index 0ef9b02d32d0b..260f77b6aff72 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -135,6 +135,10 @@ case class AdaptiveSparkPlanExec( CoalesceBucketsInJoin, RemoveRedundantProjects, ensureRequirements, + // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, so + // every rule below and the execution itself read the answer the exchanges above it were + // planned against. + StampUnionDecisions, // This rule must be run after `EnsureRequirements`. InsertSortForLimitAndOffset, AdjustShuffleExchangePosition, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 331824601f88d..2237676c8fb12 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1038,41 +1038,58 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Serializes the latches below so concurrent first readers agree on one answer. What must not be - // reachable under it is the AQE final-plan lock: `CoalesceShufflePartitions` reads `isPlainUnion` - // while holding that lock, so the reverse edge deadlocks. `InMemoryTableScanExec` is safe here - // because it reads `adaptive.executedPlan`, a volatile read, not `finalPhysicalPlan`, which is - // `lock.synchronized`. Child `lazy val`s such as `AQEShuffleReadExec.outputPartitioning` do take - // their own instance monitor; that is harmless, as nothing holds one of those and then waits - // here. Driver-only, hence `@transient`. + // Serializes the latches below so concurrent first readers agree on one answer. Nothing called + // under it may take the AQE final-plan lock, which `CoalesceShufflePartitions` holds while it + // reads `isPlainUnion`; `InMemoryTableScanExec` is safe on that count, reading the volatile + // `adaptive.executedPlan` rather than `finalPhysicalPlan`. Monitors are taken under it, since a + // `lazy val` initializes on its own instance, but none belongs to this node or one above it. + // That matters for this node in particular: `metrics` and `supportCodegenFailureReason` hold its + // monitor and then take this lock, so nothing under the lock may force a `lazy val` here. + // `output` and `prepareOutputPartitioning` are `def`s. Driver-only, hence `@transient`. @transient private val decisionLock = new Object() /** * True when this union behaves as a plain concatenation, so `unionedInputRDD` matches the * semantics of `sparkContext.union(...)` in `unionRDDs`. It satisfies the partitioning gate on * the codegen path, not the whole of it: `supportCodegenFailureReason` still applies its other - * checks. A `KeyedPartitioning` union also concatenates, but codegen stays off for it, with the - * reason "partitioning-aware", because a downstream `GroupPartitionsExec` consumes its key - * descriptor. + * checks. When this union merges its children's `KeyedPartitioning` instead, it concatenates all + * the same, but codegen stays off, with the reason "partitioning-aware", because a downstream + * `GroupPartitionsExec` consumes its key descriptor. * * Latched, because the answer moves under its consumers. * `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning` while its inner * `AdaptiveSparkPlanExec` has no final plan, so a union can look plain when * `CollapseCodegenStages` gates on it and partitioning-aware by the time the stage runs. The * shell that gate builds wraps a `withNewChildren` copy, and a copy that re-derived here came - * back with empty `metrics` while `doProduce` asked `metricTerm` for `numOutputRows`. The tag - * carries the answer onto that copy, since `withNewChildren` ends in `copyTagsFrom`. + * back with empty `metrics` while `doProduce` asked `metricTerm` for `numOutputRows`. A fresh + * copy inherits the answer instead, since `withNewChildren` ends in `copyTagsFrom`. * - * `UNION_OUTPUT_PARTITIONING` is read here rather than in `rawPartitioning` so it is latched too: - * `conf` is live, and a plan must execute by the partitioning it was planned against. + * `UNION_OUTPUT_PARTITIONING` is read where the decision is stamped rather than in + * `rawPartitioning`, so it too is fixed once the plan is prepared: `conf` is live, and a plan + * must execute by the partitioning it was planned against. + * + * A read before `StampUnionDecisions` answers from the children as they are then, and does not + * write, so observing an unprepared plan cannot decide anything for the prepared one. */ private[sql] def isPlainUnion: Boolean = decisionLock.synchronized { - getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse { - val plain = !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || - rawPartitioning.isInstanceOf[UnknownPartitioning] - setTagValue(UnionExec.PLAIN_UNION_DECISION, plain) - plain + getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse(derivePlainUnion) + } + + private def derivePlainUnion: Boolean = + !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || + rawPartitioning.isInstanceOf[UnknownPartitioning] + + /** + * Fixes both decisions for the rest of this plan's life. Called by `StampUnionDecisions` right + * after `EnsureRequirements`, so the answer a parent's exchange decision was taken from is the + * one execution uses. Idempotent, and never overwrites: a node that already carries the tags + * keeps them, which is how the copy in the codegen shell stays in step with the gate. + */ + private[execution] def stampDecisions(): Unit = decisionLock.synchronized { + if (getTagValue(UnionExec.PLAIN_UNION_DECISION).isEmpty) { + setTagValue(UnionExec.PLAIN_UNION_DECISION, derivePlainUnion) } + codegenConfSnapshot } /** @@ -1122,19 +1139,22 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Memoized per instance, not latched on a tag: every term below the confs reads the children, and - // a tag outlives them. `SparkPlanInfo` forces `metrics` on an AQE plan update, before the rules - // that run ahead of `CollapseCodegenStages`, so a rule replacing a child there would inherit an - // allowing answer and fuse a topology that `hasPartitionIndexDependentCodegen` or - // `supportsColumnar` rejects. The copy in the codegen shell still agrees with the gate: - // `InputAdapter` delegates `output` and `supportsColumnar` to its child, the remaining terms walk - // the subtree through it, and each is fixed for a given set of children. + // Memoized per instance rather than latched on a tag. Every term below the confs except + // `isPlainUnion` reads the children, and a tag outlives them: `SparkPlanInfo` forces `metrics` on + // an AQE plan update, before the rules that run ahead of `CollapseCodegenStages`, so a rule + // replacing a child there would inherit an allowing answer and fuse a topology that + // `hasPartitionIndexDependentCodegen` or `supportsColumnar` rejects. The copy in the codegen + // shell still agrees with the gate: `InputAdapter` delegates `output` and `supportsColumnar` to + // its child, the other terms walk the subtree through it, and each of those is fixed for a given + // set of children. `isPlainUnion` is not, which is why it is stamped instead. + // + // `isPlainUnion` is checked last of all: a union rejected on any other ground was never going to + // fuse, and asking the question would fix a decision that costs it SPARK-52921's exchange + // elimination for nothing. @transient private lazy val supportCodegenFailureReason: Option[String] = { val confs = codegenConfSnapshot if (!confs.unionCodegenEnabled) { Some("union-codegen-disabled") - } else if (!isPlainUnion) { - Some("partitioning-aware") } else if (children.exists(_.exists(_.isInstanceOf[UnionExec]))) { Some("nested-union") } else if (children.exists(_.exists(UnionExec.isKnownMultiInputRDDCodegen))) { @@ -1148,6 +1168,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } else if (children.exists(c => c.output.zip(output).exists { case (src, tgt) => src.dataType != tgt.dataType })) { Some("type-mismatch") + } else if (!isPlainUnion) { + Some("partitioning-aware") } else { None } @@ -1352,19 +1374,21 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup object UnionExec { /** - * The latched "is this a plain concatenation" decision. See `isPlainUnion`. + * The stamped "is this a plain concatenation" decision. See `isPlainUnion`. * - * Rebuilds carry it: `withNewChildren` and a transform rule's replacement both go through - * `copyTagsFrom`. A `UnionExec` with no predecessor node at its position starts unlatched: it - * can re-derive the opposite answer and leave `metrics` empty under generated code that - * increments it. + * `withNewChildren` copies it onto a rebuilt node, and so does a transform rule's replacement, + * but only where the target carries no tags of its own: `copyTagsFrom` leaves a node that already + * has some untouched. A `UnionExec` reaching execution unstamped therefore answers from the + * children it has then, and can leave `metrics` empty, so `doProduce` fails asking `metricTerm` + * for `numOutputRows`. */ private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") /** * The confs `supportCodegenFailureReason` reads, latched by `codegenConfSnapshot`. Only the conf - * values are latched. The terms that read the children are derived per instance instead, so that - * a rule replacing a child cannot inherit an answer taken from the topology it replaced. + * values are latched here. Every other term of the reason except `isPlainUnion` is derived per + * instance, so a rule replacing a child cannot inherit one taken from the topology it replaced; + * `isPlainUnion` is latched separately, on [[PLAIN_UNION_DECISION]], and does carry over. */ private case class CodegenConfSnapshot(unionCodegenEnabled: Boolean, maxChildren: Int) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index a0ffa2b32d809..bc2875c399ddc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -21,9 +21,9 @@ import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} -import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, UnknownPartitioning} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.{REPARTITION_BY_NUM, ShuffleExchangeExec} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -65,9 +65,9 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query stages; * `SparkPlan.collect` stops at them, since both are `LeafExecNode`s. * - * Stricter than `unionInsideWSCG` on purpose: `w.find` also matches a union that an - * `InputAdapter` left inside the stage unfused, and the callers here assert on `metrics`, which - * an unfused union does not register. + * Stricter than `unionInsideWSCG` on purpose: this matches only a union that is the root of its + * own codegen stage, which is what "fused" means for the callers here, while `w.find` also + * matches one an `InputAdapter` left inside the stage. */ private def fusedUnions(df: DataFrame): Seq[UnionExec] = collect(df.queryExecution.executedPlan) { @@ -667,13 +667,13 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // `children.forall`, so one row-based `ProjectExec` over the columnar scan is enough to make // it false, and without one `supportCodegenFailureReason` reports `columnar` and nothing // fuses. `SELECT *` or a plain alias collapses the projection away and does not reproduce - // this. Once the cache stages finalise, both children report the same `HashPartitioning`, and - // re-deriving the decision at that point left `metrics` empty while the generated code still - // incremented it, so `doProduce` threw `key not found: numOutputRows`. + // this. Once the cache stages finalise, both children report the same concrete layout, and + // re-deriving the decision at that point left `metrics` empty while `doProduce` asked + // `metricTerm` for `numOutputRows`. // // Both halves of the decision are asserted here. Registering the metric unconditionally would // fix the crash and leave the other half broken: a fused union concatenates its children's - // partitions, so claiming their `HashPartitioning` would let a parent satisfy a clustered + // partitions, so claiming their partitioning would let a parent satisfy a clustered // distribution from an RDD that does not have it. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", @@ -688,6 +688,16 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(fused.nonEmpty, "this shape must actually fuse, or the test is not exercising the defect") fused.foreach { u => + // Part of the premise, not the whole of it: the children expose a concrete layout by now, + // so this node is not reporting `UnknownPartitioning` merely for want of anything to + // derive from. `rawPartitioning` also falls back when the children's remapped + // partitionings do not compare equal, and that cannot be asserted here: each side carries + // its own exprIds, and they line up only after the private `prepareOutputPartitioning`. + val childPartitionings = u.children.map(_.outputPartitioning) + assert(childPartitionings.forall(_.isInstanceOf[HashPartitioningLike]), + s"premise: got $childPartitionings") + assert(childPartitionings.map(_.numPartitions).distinct.size == 1, + s"premise: got $childPartitionings") assert(u.metrics.contains("numOutputRows"), "a fused union must register the metric its generated code increments") assert(u.outputPartitioning.isInstanceOf[UnknownPartitioning], @@ -718,10 +728,12 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(unions.size == 1) // Not asserted through `isPlainUnion`: that call latches the decision, which would warm // a field-based implementation's memo and hide the regression this test is for. The - // exchange count below proves the union reported a concrete partitioning, without - // touching the node. - assert(plan.collect { case s: ShuffleExchangeExec => s }.size == 2, - "only the two repartitions may shuffle; the aggregate's exchange must have been elided") + // exchanges below prove the union reported a concrete partitioning, without touching the + // node: only the two repartitions may shuffle, so the aggregate's exchange was elided. + val shuffles = plan.collect { case s: ShuffleExchangeExec => s } + assert(shuffles.size == 2) + assert(shuffles.forall(_.shuffleOrigin == REPARTITION_BY_NUM), + s"expected only the two repartitions, got ${shuffles.map(_.shuffleOrigin)}") df } @@ -743,8 +755,9 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val planned = withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") { // Each child is an exchange, which is not `CodegenSupport`, so `insertInputAdapter` wraps - // it and `withNewChildren` really does produce a copy. With codegen-support children it - // returns `this`, the memo stays the original's warm one, and nothing re-derives. + // it and the union is rebuilt through `withNewChildren`, the copy this test needs. Children + // that do support codegen can still produce one, since `insertInputAdapter` recurses into + // their descendants; exchanges just make it certain. val df = rangeDF(100).repartition(2).union(rangeDF(100).repartition(2)) // `fusedUnions` requires the union to be the stage root; `unionInsideWSCG` would also // match a union that an `InputAdapter` left inside the stage unfused, which is exactly @@ -754,6 +767,16 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { assert(planned.collect().length == 200) + // The row count alone does not discriminate, since the shell was installed at planning and + // keeps emitting; registering `numOutputRows` unconditionally and reading the conf per call + // passes it. This assertion is what fails there, because nothing forces the copy's reason + // before it. It has to sit after the flip, as it does here: taken while the conf was still + // on, it would warm a memoizing implementation with the answer this test needs it not to + // have. + val copy = fusedUnions(planned) + assert(copy.size == 1) + assert(copy.head.supportCodegen, + "the copy in the shell must keep the decision it was planned with") } } } @@ -775,7 +798,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(union.supportCodegen) // A nested union is one of the topologies the gate rejects, and `withNewChildren` is the path - // every rule takes to install a replacement. + // a rule takes when it rewrites children in place. A rule returning an arbitrary replacement + // node is a different path, and one `copyTagsFrom` need not carry the tags along. val nested = UnionExec(Seq(union.children.head, union.children.head)) val rebuilt = union.withNewChildren(Seq(nested, union.children.last)).asInstanceOf[UnionExec] assert(!rebuilt.supportCodegen, "the rebuilt union must answer against its own children") @@ -783,6 +807,56 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: reading the unprepared plan does not decide the prepared one") { + // `QueryExecution.executedPlan` is `prepareForExecution(sparkPlan.clone())`, and `clone` ends + // in `makeCopy`, which calls `copyTagsFrom`. A decision written while answering a read on + // `sparkPlan` would therefore ride into the prepared plan. Here the two answers differ: each + // child is an aggregate whose exchange `EnsureRequirements` has yet to insert, so the union + // passes nothing through before preparation and both children's `HashPartitioning` after it. + // Reads before `StampUnionDecisions` answer without writing, so only preparation decides. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val left = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k").groupBy("k").count() + val right = spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k").groupBy("k").count() + val df = left.union(right) + + val unprepared = df.queryExecution.sparkPlan.collect { case u: UnionExec => u } + assert(unprepared.size == 1) + assert(unprepared.head.outputPartitioning.isInstanceOf[UnknownPartitioning], + "the aggregates have no exchange under them yet, so there is nothing to pass through") + + val prepared = df.queryExecution.executedPlan.collect { case u: UnionExec => u } + assert(prepared.size == 1) + assert(!prepared.head.outputPartitioning.isInstanceOf[UnknownPartitioning], + "the read above must not have decided for the prepared plan, got " + + s"${prepared.head.outputPartitioning}") + checkAnswer(df, (0L until 5L).flatMap(k => Seq(Row(k, 4L), Row(k, 4L)))) + } + } + + test("SPARK-59122: a prepared union keeps its layout when nothing read it during preparation") { + // With whole-stage codegen off, no gate consults the union while the plan is prepared, and a + // root union has no parent to ask for its partitioning either. First-read initialization would + // then decide at execution, under whatever the conf says by then; `StampUnionDecisions` decides + // during preparation instead. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") { + val plan = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + 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"))) + .queryExecution.executedPlan + } + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + // Co-partitioned children pass their four partitions through; a plain concatenation would + // report eight. + assert(plan.execute().getNumPartitions == 4, + "a prepared union must execute by the layout it was prepared with") + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From dc85f998986393cff1f51a92b0de8d066e93e972 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 12 Sep 2026 14:24:10 +0800 Subject: [PATCH 10/24] Put isPlainUnion back as the gate's second term Reordering it to last bought nothing once StampUnionDecisions asks every UnionExec during preparation and isPlainUnion no longer writes on read: the term's position stops deciding whether anything is fixed. All the reorder changed was the reason reported for a union failing several gates. --- .../spark/sql/execution/basicPhysicalOperators.scala | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 2237676c8fb12..e2e4dd25795b5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1147,14 +1147,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // shell still agrees with the gate: `InputAdapter` delegates `output` and `supportsColumnar` to // its child, the other terms walk the subtree through it, and each of those is fixed for a given // set of children. `isPlainUnion` is not, which is why it is stamped instead. - // - // `isPlainUnion` is checked last of all: a union rejected on any other ground was never going to - // fuse, and asking the question would fix a decision that costs it SPARK-52921's exchange - // elimination for nothing. @transient private lazy val supportCodegenFailureReason: Option[String] = { val confs = codegenConfSnapshot if (!confs.unionCodegenEnabled) { Some("union-codegen-disabled") + } else if (!isPlainUnion) { + Some("partitioning-aware") } else if (children.exists(_.exists(_.isInstanceOf[UnionExec]))) { Some("nested-union") } else if (children.exists(_.exists(UnionExec.isKnownMultiInputRDDCodegen))) { @@ -1168,8 +1166,6 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } else if (children.exists(c => c.output.zip(output).exists { case (src, tgt) => src.dataType != tgt.dataType })) { Some("type-mismatch") - } else if (!isPlainUnion) { - Some("partitioning-aware") } else { None } From 773c28999ebcc61bbf084f401c9de875a730037b Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 12 Sep 2026 16:03:49 +0800 Subject: [PATCH 11/24] Hold the union decisions in one tag and drop the lock The two tags were written by the same call at the same instant, so one Decisions(plainUnion, unionCodegenEnabled, maxChildren) says it in one place, and a reader can no longer see half a decision. The conf accessors now derive without writing, like isPlainUnion, which leaves the stamping rule as the only writer of the tag on an existing node and leaves nothing for decisionLock to serialize. Also rename latch to stamp where the wording was left behind, and narrow three comments to what the code supports. --- .../sql/execution/StampUnionDecisions.scala | 13 +- .../execution/basicPhysicalOperators.scala | 120 ++++++++---------- .../sql/execution/UnionCodegenSuite.scala | 10 +- 3 files changed, 68 insertions(+), 75 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index c2f4539355655..f544fcdef083c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -20,16 +20,19 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.catalyst.rules.Rule /** - * Fixes each [[UnionExec]]'s partitioning decision and codegen conf snapshot at one defined point. + * Fixes each [[UnionExec]]'s partitioning decision and the confs its codegen gate reads, at one + * defined point. * * `UnionExec` derives both from state that moves: its children's `outputPartitioning` sharpens as * AQE finalises the plans behind them, and `conf` is the live session conf. Whoever asked first * used to decide, which made the answer depend on when it was observed. This rule asks once, - * right after `EnsureRequirements`, so the partitioning a parent's exchange decision was taken - * from is the one `unionRDDs` and the codegen gate use. + * right after `EnsureRequirements`, so the decision the exchanges around a union were planned + * against is the one `unionRDDs` and the codegen gate use. * - * It only writes what is not there yet, so re-running it (AQE re-optimizes each round) keeps the - * first answer, and a node rebuilt from a stamped one keeps the tags `copyTagsFrom` gave it. + * It only writes what is not there yet, so a second pass over the same nodes keeps the first + * answer, and a node rebuilt from a stamped one keeps the tag `copyTagsFrom` gave it. AQE re-plans + * between rounds, so a union outside a materialized stage is stamped again from what that round + * sees; one already inside a stage is not revisited, since `foreach` stops at `QueryStageExec`. */ object StampUnionDecisions extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index e2e4dd25795b5..57798d752b12a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -978,10 +978,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` latches on - * whether this comes back `UnknownPartitioning`; `outputPartitioning` reports it when that latch - * says the union is not a plain concatenation. The `UNION_OUTPUT_PARTITIONING` gate lives in the - * latch, not here. + * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` answers on + * whether this comes back `UnknownPartitioning`; `outputPartitioning` reports it when that + * decision says the union is not a plain concatenation. The `UNION_OUTPUT_PARTITIONING` gate is + * read with the decision, not here. */ private def rawPartitioning: Partitioning = { // Children's partitionings with attributes remapped to this union's output attributes. @@ -1038,16 +1038,6 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Serializes the latches below so concurrent first readers agree on one answer. Nothing called - // under it may take the AQE final-plan lock, which `CoalesceShufflePartitions` holds while it - // reads `isPlainUnion`; `InMemoryTableScanExec` is safe on that count, reading the volatile - // `adaptive.executedPlan` rather than `finalPhysicalPlan`. Monitors are taken under it, since a - // `lazy val` initializes on its own instance, but none belongs to this node or one above it. - // That matters for this node in particular: `metrics` and `supportCodegenFailureReason` hold its - // monitor and then take this lock, so nothing under the lock may force a `lazy val` here. - // `output` and `prepareOutputPartitioning` are `def`s. Driver-only, hence `@transient`. - @transient private val decisionLock = new Object() - /** * True when this union behaves as a plain concatenation, so `unionedInputRDD` matches the * semantics of `sparkContext.union(...)` in `unionRDDs`. It satisfies the partitioning gate on @@ -1056,7 +1046,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * the same, but codegen stays off, with the reason "partitioning-aware", because a downstream * `GroupPartitionsExec` consumes its key descriptor. * - * Latched, because the answer moves under its consumers. + * Stamped, because the answer moves under its consumers. * `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning` while its inner * `AdaptiveSparkPlanExec` has no final plan, so a union can look plain when * `CollapseCodegenStages` gates on it and partitioning-aware by the time the stage runs. The @@ -1071,29 +1061,31 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * A read before `StampUnionDecisions` answers from the children as they are then, and does not * write, so observing an unprepared plan cannot decide anything for the prepared one. */ - private[sql] def isPlainUnion: Boolean = decisionLock.synchronized { - getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse(derivePlainUnion) - } - - private def derivePlainUnion: Boolean = + private[sql] def isPlainUnion: Boolean = stampedDecisions.map(_.plainUnion).getOrElse { !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || rawPartitioning.isInstanceOf[UnknownPartitioning] + } + + private def stampedDecisions: Option[UnionExec.Decisions] = + getTagValue(UnionExec.DECISIONS) /** - * Fixes both decisions for the rest of this plan's life. Called by `StampUnionDecisions` right - * after `EnsureRequirements`, so the answer a parent's exchange decision was taken from is the - * one execution uses. Idempotent, and never overwrites: a node that already carries the tags - * keeps them, which is how the copy in the codegen shell stays in step with the gate. + * Fixes this node's decisions for the rest of the plan's life. Called by `StampUnionDecisions` + * right after `EnsureRequirements`, so what the exchanges around this union were planned against + * is what execution uses. Nothing else writes this tag on an existing node, and the nodes the + * rule writes are freshly planned and not yet published, so no reader can be looking at one; + * `metrics` and the codegen gate read it later, and a node that already carries it keeps it, + * which is how the copy in the codegen shell stays in step with the gate. */ - private[execution] def stampDecisions(): Unit = decisionLock.synchronized { - if (getTagValue(UnionExec.PLAIN_UNION_DECISION).isEmpty) { - setTagValue(UnionExec.PLAIN_UNION_DECISION, derivePlainUnion) - } - codegenConfSnapshot + private[execution] def stampDecisions(): Unit = if (stampedDecisions.isEmpty) { + setTagValue(UnionExec.DECISIONS, UnionExec.Decisions( + plainUnion = isPlainUnion, + unionCodegenEnabled = conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED), + maxChildren = conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN))) } /** - * A node latched plain reports `UnknownPartitioning` even once its children agree on a concrete + * A node stamped plain reports `UnknownPartitioning` even once its children agree on a concrete * one: a fused union concatenates, and claiming their partitioning would let a parent skip an * exchange it needs. The cost is SPARK-52921's exchange elimination for such a union. * @@ -1124,22 +1116,22 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // The conf half of the codegen decision, latched for the reason `isPlainUnion` is: `conf` is - // live, so the gate, `metrics` and the copy `insertInputAdapter` puts inside the codegen shell - // would otherwise be free to read different values. When a child is not `CodegenSupport` that - // copy is real and its first evaluation lands at execution; reading the conf there left - // `metrics` empty while `doProduce` asked `metricTerm` for `numOutputRows`. - private def codegenConfSnapshot: UnionExec.CodegenConfSnapshot = decisionLock.synchronized { - getTagValue(UnionExec.CODEGEN_CONF_SNAPSHOT).getOrElse { - val snapshot = UnionExec.CodegenConfSnapshot( - conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED), - conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) - setTagValue(UnionExec.CODEGEN_CONF_SNAPSHOT, snapshot) - snapshot - } - } - - // Memoized per instance rather than latched on a tag. Every term below the confs except + // The confs the gate reads, stamped for the reason the plain-union decision is: `conf` is live, + // so the gate, `metrics` and the copy `insertInputAdapter` puts inside the codegen shell would + // otherwise be free to read different values. When a child is not `CodegenSupport` that copy is + // real and its first evaluation lands at execution; reading the conf there left `metrics` empty + // while `doProduce` asked `metricTerm` for `numOutputRows`. A read before the stamp answers from + // the conf as it is then and writes nothing, so observing an unprepared plan cannot pin this + // either. + private def unionCodegenEnabled: Boolean = + stampedDecisions.map(_.unionCodegenEnabled) + .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) + + private def maxCodegenChildren: Int = + stampedDecisions.map(_.maxChildren) + .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) + + // Memoized per instance rather than stamped on the tag. Every term below the confs except // `isPlainUnion` reads the children, and a tag outlives them: `SparkPlanInfo` forces `metrics` on // an AQE plan update, before the rules that run ahead of `CollapseCodegenStages`, so a rule // replacing a child there would inherit an allowing answer and fuse a topology that @@ -1148,8 +1140,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // its child, the other terms walk the subtree through it, and each of those is fixed for a given // set of children. `isPlainUnion` is not, which is why it is stamped instead. @transient private lazy val supportCodegenFailureReason: Option[String] = { - val confs = codegenConfSnapshot - if (!confs.unionCodegenEnabled) { + if (!unionCodegenEnabled) { Some("union-codegen-disabled") } else if (!isPlainUnion) { Some("partitioning-aware") @@ -1159,7 +1150,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup Some("multi-rdd-child") } else if (children.exists(UnionExec.hasPartitionIndexDependentCodegen)) { Some("partition-index-dependent-child") - } else if (children.size > confs.maxChildren) { + } else if (children.size > maxCodegenChildren) { Some("max-children-exceeded") } else if (supportsColumnar) { Some("columnar") @@ -1370,26 +1361,25 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup object UnionExec { /** - * The stamped "is this a plain concatenation" decision. See `isPlainUnion`. - * - * `withNewChildren` copies it onto a rebuilt node, and so does a transform rule's replacement, - * but only where the target carries no tags of its own: `copyTagsFrom` leaves a node that already - * has some untouched. A `UnionExec` reaching execution unstamped therefore answers from the - * children it has then, and can leave `metrics` empty, so `doProduce` fails asking `metricTerm` - * for `numOutputRows`. + * What `StampUnionDecisions` fixes on a `UnionExec`: whether it is a plain concatenation, and the + * two confs the codegen gate reads. Everything else the gate asks is derived per instance, so a + * rule replacing a child cannot inherit an answer taken from the topology it replaced. */ - private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision") + private case class Decisions( + plainUnion: Boolean, + unionCodegenEnabled: Boolean, + maxChildren: Int) /** - * The confs `supportCodegenFailureReason` reads, latched by `codegenConfSnapshot`. Only the conf - * values are latched here. Every other term of the reason except `isPlainUnion` is derived per - * instance, so a rule replacing a child cannot inherit one taken from the topology it replaced; - * `isPlainUnion` is latched separately, on [[PLAIN_UNION_DECISION]], and does carry over. + * The stamped decisions. See `isPlainUnion` and `stampDecisions`. + * + * `withNewChildren` copies the tag onto a rebuilt node, and so does a transform rule's + * replacement, but only where the target carries no tags of its own: `copyTagsFrom` leaves a node + * that already has some untouched. A `UnionExec` reaching execution unstamped therefore answers + * from the state it sees then, and can leave `metrics` empty, so `doProduce` fails asking + * `metricTerm` for `numOutputRows`. */ - private case class CodegenConfSnapshot(unionCodegenEnabled: Boolean, maxChildren: Int) - - private val CODEGEN_CONF_SNAPSHOT = - TreeNodeTag[CodegenConfSnapshot]("codegenConfSnapshot") + private val DECISIONS = TreeNodeTag[Decisions]("unionDecisions") /** * Codegen operators that return more than one RDD from `inputRDDs()`. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index bc2875c399ddc..ba0a34cc45ce1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -709,7 +709,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper test("SPARK-59122: a partitioning-aware union keeps its layout when the conf changes between " + "planning and execution") { - // `spark.sql.unionOutputPartitioning` is read where the plain-union decision is latched, not on + // `spark.sql.unionOutputPartitioning` is read where the plain-union decision is stamped, not on // every `outputPartitioning` call, so a plan executes by the partitioning it was planned // against. Reading it per call let the parent aggregate lose its exchange at planning and get a // plain concatenation at execution, reporting each group twice. The `checkAnswer` below stays @@ -726,10 +726,10 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val plan = df.queryExecution.executedPlan val unions = plan.collect { case u: UnionExec => u } assert(unions.size == 1) - // Not asserted through `isPlainUnion`: that call latches the decision, which would warm - // a field-based implementation's memo and hide the regression this test is for. The - // exchanges below prove the union reported a concrete partitioning, without touching the - // node: only the two repartitions may shuffle, so the aggregate's exchange was elided. + // Asserted through the exchanges rather than through `isPlainUnion`, so that the check + // does not depend on how the decision is stored: only the two repartitions may shuffle, so + // the aggregate's exchange was elided, which it could only be if the union reported a + // concrete partitioning. val shuffles = plan.collect { case s: ShuffleExchangeExec => s } assert(shuffles.size == 2) assert(shuffles.forall(_.shuffleOrigin == REPARTITION_BY_NUM), From 33b79ab0cb8ba8f099752856206b9754f89e17b0 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sun, 13 Sep 2026 15:07:00 +0800 Subject: [PATCH 12/24] Stamp late-created unions too, and answer the rest of the review StampUnionDecisions is now listed after each phase that can add a UnionExec: after the injected columnar rules in both pipelines and after the injected AQE prep rules. The write-once guard is what keeps those passes from moving a decision an earlier consumer planned against, and a test drives the rule directly for both halves. Also: the codegen confs are read at the stamp rather than pinned on first read, so no pre-preparation read can decide anything; the early metrics read is the initial SparkPlanInfo SQLExecution builds, not an AQE plan update; the conf docs describe physical preparation as the boundary; maxChildren gets the same-plan flip test; and outputPartitioning records the fusion a post-stamp partitioning drop costs. --- .../apache/spark/sql/internal/SQLConf.scala | 14 +-- .../spark/sql/execution/QueryExecution.scala | 4 + .../sql/execution/StampUnionDecisions.scala | 27 ++++-- .../adaptive/AdaptiveSparkPlanExec.scala | 9 +- .../execution/basicPhysicalOperators.scala | 31 ++++-- .../sql/execution/UnionCodegenSuite.scala | 96 ++++++++++++++++++- 6 files changed, 151 insertions(+), 30 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 7f7f073cc2f4f..01bed12871c10 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2908,9 +2908,9 @@ object SQLConf { .doc("When both this conf and `spark.sql.codegen.wholeStage` are true, " + "UnionExec participates in whole-stage codegen on its " + "non-partitioning-aware path: the parent and all children fuse into " + - "a single WholeStageCodegenExec stage. The value is read when a UnionExec " + - "first decides, so a change applies to plans built afterwards, not to a " + - "plan that has already been prepared.") + "a single WholeStageCodegenExec stage. The value is read when a UnionExec's " + + "decision is fixed during physical preparation, so a change does not reach a " + + "decision already taken.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .booleanConf @@ -2926,8 +2926,8 @@ object SQLConf { "than the JVM per-method bytecode limit. Unions with more children " + "fall back to per-child codegen stages. Only effective when " + s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true. The value is read when a " + - "UnionExec first decides, so a change applies to plans built afterwards, not " + - "to a plan that has already been prepared.") + "UnionExec's decision is fixed during physical preparation, so a change does not " + + "reach a decision already taken.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .intConf @@ -8125,8 +8125,8 @@ object SQLConf { .internal() .doc("When set to true, the output partitioning of UnionExec will be the same as the " + "input partitioning if its children have same partitioning. Otherwise, it will be a " + - "default partitioning. The value is read when a UnionExec first decides, so a change " + - "applies to plans built afterwards, not to a plan that has already been prepared.") + "default partitioning. The value is read when a UnionExec's decision is fixed during " + + "physical preparation, so a change does not reach a decision already taken.") .version("4.1.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index d9d122844ee8b..79f1f66a5a839 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -827,6 +827,10 @@ object QueryExecution { RemoveRedundantSorts, ApplyColumnarRulesAndInsertTransitions( sparkSession.sessionState.columnarRules, outputsColumnar = false), + // A barrier for a `UnionExec` an injected columnar rule just created, which has no decision + // yet and would otherwise take one wherever it is first asked. A decision already stamped on + // a node is kept. + StampUnionDecisions, CollapseCodegenStages()) ++ (if (subquery) { Nil diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index f544fcdef083c..39e1511c88f49 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -25,14 +25,27 @@ import org.apache.spark.sql.catalyst.rules.Rule * * `UnionExec` derives both from state that moves: its children's `outputPartitioning` sharpens as * AQE finalises the plans behind them, and `conf` is the live session conf. Whoever asked first - * used to decide, which made the answer depend on when it was observed. This rule asks once, - * right after `EnsureRequirements`, so the decision the exchanges around a union were planned - * against is the one `unionRDDs` and the codegen gate use. + * used to decide, which made the answer depend on when it was observed. This rule asks right + * after `EnsureRequirements`, so the decision the exchanges around a union were planned against is + * the one `unionRDDs` and the codegen gate use. * - * It only writes what is not there yet, so a second pass over the same nodes keeps the first - * answer, and a node rebuilt from a stamped one keeps the tag `copyTagsFrom` gave it. AQE re-plans - * between rounds, so a union outside a materialized stage is stamped again from what that round - * sees; one already inside a stage is not revisited, since `foreach` stops at `QueryStageExec`. + * It is listed again after the injected columnar and query-stage rules, the hooks that can add a + * `UnionExec` of their own. One created there has no decision yet, and would otherwise take one + * wherever it is first asked, where the copy in the codegen shell can disagree with the gate. The + * cached-scan branch of stage creation needs no barrier: it rejects a result that is no longer an + * `InMemoryTableScanLike`, which is a leaf. + * + * A tag rather than a constructor field, because a field would land in `argString` and so in + * every `explain` and `PlanStability` golden holding a `Union`, and in `canonicalized`, which + * exchange and cached-plan reuse key on. Writing it in place is safe here, unlike in + * `MarkSingleTaskExecution`, because preparation runs on `sparkPlan.clone()` and every + * shared-subtree boundary is a leaf, so `foreach` cannot reach a node another query owns. + * + * It only writes what is not there yet, so a later pass cannot move a decision already stamped on + * a node, a second pass over the same nodes keeps the first answer, and a node rebuilt from a + * stamped one keeps the tag `copyTagsFrom` gave it. AQE re-plans between rounds, so a union above + * the stages already created is stamped again from what that round sees, where that plan is the one + * adopted; a union inside a stage is not revisited, since `foreach` stops at `QueryStageExec`. */ object StampUnionDecisions extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index 260f77b6aff72..afaa330a502dd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -163,7 +163,10 @@ case class AdaptiveSparkPlanExec( // local sort left dangling right below the extra shuffle that skew join optimization may // insert between two joins. RemoveRedundantSorts - ) ++ context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules + ) ++ context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules :+ + // A barrier for a `UnionExec` an injected prep rule just created. Decisions already stamped + // above are kept. + StampUnionDecisions } // A list of physical optimizer rules to be applied to a new stage before its execution. These @@ -188,6 +191,10 @@ case class AdaptiveSparkPlanExec( private def postStageCreationRules(outputsColumnar: Boolean) = Seq( ApplyColumnarRulesAndInsertTransitions( context.session.sessionState.columnarRules, outputsColumnar), + // A barrier for a `UnionExec` an injected stage-optimizer or columnar rule just created, which + // has no decision yet and would otherwise take one wherever it is first asked. A decision + // already stamped on a node is kept, so this pass cannot move one. + StampUnionDecisions, collapseCodegenStagesRule ) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 57798d752b12a..b8493486f1093 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -978,7 +978,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` answers on + * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` answers * whether this comes back `UnknownPartitioning`; `outputPartitioning` reports it when that * decision says the union is not a plain concatenation. The `UNION_OUTPUT_PARTITIONING` gate is * read with the decision, not here. @@ -1061,7 +1061,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * A read before `StampUnionDecisions` answers from the children as they are then, and does not * write, so observing an unprepared plan cannot decide anything for the prepared one. */ - private[sql] def isPlainUnion: Boolean = stampedDecisions.map(_.plainUnion).getOrElse { + private[execution] def isPlainUnion: Boolean = stampedDecisions.map(_.plainUnion).getOrElse { !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || rawPartitioning.isInstanceOf[UnknownPartitioning] } @@ -1070,11 +1070,11 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup getTagValue(UnionExec.DECISIONS) /** - * Fixes this node's decisions for the rest of the plan's life. Called by `StampUnionDecisions` - * right after `EnsureRequirements`, so what the exchanges around this union were planned against - * is what execution uses. Nothing else writes this tag on an existing node, and the nodes the - * rule writes are freshly planned and not yet published, so no reader can be looking at one; - * `metrics` and the codegen gate read it later, and a node that already carries it keeps it, + * Fixes this node's decisions for the rest of the plan's life. Called by `StampUnionDecisions`, + * first right after `EnsureRequirements`, so what the exchanges around this union were planned + * against is what execution uses. Nothing else writes this tag on an existing node, and the + * nodes the rule writes are freshly planned and not yet published, so no reader can be looking at + * one; `metrics` and the codegen gate read it later, and a node that already carries it keeps it, * which is how the copy in the codegen shell stays in step with the gate. */ private[execution] def stampDecisions(): Unit = if (stampedDecisions.isEmpty) { @@ -1089,6 +1089,17 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * one: a fused union concatenates, and claiming their partitioning would let a parent skip an * exchange it needs. The cost is SPARK-52921's exchange elimination for such a union. * + * Only the decision is stamped, never the `Partitioning` itself. AQE coalescing changes the + * children's `numPartitions` after the stamp, and a stale count is what `unionRDDs` would hand + * `SQLPartitioningAwareUnionRDD`, which builds exactly that many partitions from each child. + * + * The reverse costs fusion. A rule that runs after the stamp and drops a child's partitioning + * leaves the node stamped non-plain, so the codegen gate answers "partitioning-aware" and + * `numOutputRows` goes unregistered, where re-deriving at the gate would have fused it. + * `DisableUnnecessaryBucketedScan` does that to a union over two bucketed scans, once a + * projection on each side makes them row-based; the `columnar` term would reject the bare scans + * anyway. Results are unaffected, since the other branch below re-derives and concatenates. + * * The other branch is derived per call, so `unionRDDs` could take the concatenating arm even * though `EnsureRequirements` planned the parent against a concrete partitioning: * `comparePartitioning` compares `HashPartitioningLike` by equality, so a change to one child's @@ -1132,9 +1143,9 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) // Memoized per instance rather than stamped on the tag. Every term below the confs except - // `isPlainUnion` reads the children, and a tag outlives them: `SparkPlanInfo` forces `metrics` on - // an AQE plan update, before the rules that run ahead of `CollapseCodegenStages`, so a rule - // replacing a child there would inherit an allowing answer and fuse a topology that + // `isPlainUnion` reads the children, and a tag outlives them: `SQLExecution` builds the initial + // `SparkPlanInfo` before execution, forcing `metrics` on every node it visits, so a rule that + // replaces a child after that would inherit an allowing answer and fuse a topology that // `hasPartitionIndexDependentCodegen` or `supportsColumnar` rejects. The copy in the codegen // shell still agrees with the gate: `InputAdapter` delegates `output` and `supportsColumnar` to // its child, the other terms walk the subtree through it, and each of those is fixed for a given diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index ba0a34cc45ce1..0dab65be8c567 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -23,7 +23,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, UnknownPartitioning} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.execution.exchange.{REPARTITION_BY_NUM, ShuffleExchangeExec} +import org.apache.spark.sql.execution.exchange.{EnsureRequirements, REPARTITION_BY_NUM, ShuffleExchangeExec} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -67,7 +67,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper * * Stricter than `unionInsideWSCG` on purpose: this matches only a union that is the root of its * own codegen stage, which is what "fused" means for the callers here, while `w.find` also - * matches one an `InputAdapter` left inside the stage. + * matches one that an `InputAdapter` left inside the stage. */ private def fusedUnions(df: DataFrame): Seq[UnionExec] = collect(df.queryExecution.executedPlan) { @@ -781,11 +781,44 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a fused union keeps numOutputRows when the child cap drops between " + + "planning and execution") { + // `WHOLESTAGE_UNION_MAX_CHILDREN` is on the same snapshot as the enable flag, so the same shape + // has to hold for it: prepared under a cap this union meets, it stays fused even if the cap is + // lowered under it. Reading the cap live would give the shell's copy `max-children-exceeded`, + // empty `metrics`, and `doProduce` failing at `metricTerm`. Three children against a cap of + // two, since the conf refuses anything below two. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val planned = withSQLConf( + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "3") { + // Exchange children again, so the shell really holds a `withNewChildren` copy. + val df = rangeDF(100).repartition(2) + .union(rangeDF(100).repartition(2)) + .union(rangeDF(100).repartition(2)) + val fused = fusedUnions(df) + assert(fused.size == 1 && fused.head.children.size == 3, + s"this shape must fuse as one three-child union, got ${fused.map(_.children.size)}") + df + } + withSQLConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "2") { + assert(planned.collect().length == 300) + val copy = fusedUnions(planned) + assert(copy.size == 1) + assert(copy.head.supportCodegen, + "the copy in the shell must keep the cap it was planned with") + // Not `metrics.contains`, which `collect()` above already proves: an empty `metrics` would + // have thrown at `metricTerm`. The count is what says the fused code ran and counted. + assert(copy.head.metrics("numOutputRows").value == 300) + } + } + } + test("SPARK-59122: the codegen gate re-derives when a rule replaces the children") { // The gate's children-dependent terms must not outlive the children they were taken from. - // `SparkPlanInfo` reads `metrics` on every node when AQE posts a plan update, and that happens - // before the rules running just ahead of `CollapseCodegenStages`; a decision carried from there - // onto the rebuilt node would fuse a topology that the gate rejects. + // `SQLExecution` builds a `SparkPlanInfo` before execution, which reads `metrics` on every + // node; a decision carried from there onto a node whose children a rule then replaced would + // fuse a topology that the gate rejects. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") { @@ -857,6 +890,59 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a later stamping pass fills in a fresh union and keeps stamped ones") { + // `StampUnionDecisions` is listed again after the phases that can add a `UnionExec`, so one an + // injected columnar or query-stage rule created does not answer from whatever the conf says + // wherever it is first asked. A later pass must also not move a decision already taken, which + // is the second half here. The rule is driven directly: injecting an extension needs its own + // session, and what matters is the rule's contract. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + // Pins the property the standard pipeline has to keep: a stamping pass runs after + // `EnsureRequirements`, so the decision is taken from the plan the exchanges were placed in. + // A count would break on a sixth legitimate pass and say nothing about the order. The three + // AQE positions are private to `AdaptiveSparkPlanExec`. + val rules = QueryExecution.preparations(spark, subquery = false) + val firstStamp = rules.indexWhere(_ eq StampUnionDecisions) + val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) + assert(ensureRequirements >= 0 && firstStamp > ensureRequirements, + s"expected a stamping pass after EnsureRequirements, got $ensureRequirements/$firstStamp") + + val stamped = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + 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"))) + val union = df.queryExecution.executedPlan.collect { case u: UnionExec => u } + assert(union.size == 1) + union.head + } + + // A fresh node standing in for one an extension made after the first pass: no decision yet. + val fresh = UnionExec(stamped.children) + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + StampUnionDecisions(fresh) + } + // Read back with the conf the other way round, so the answer can only come from the stamp: + // deriving here would make it non-plain, these children being co-partitioned. + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + assert(fresh.isPlainUnion, "the barrier must have decided the fresh node") + } + + // The other half. The conf a decision was stamped with is the part a second pass could move, + // so the node to watch is one whose gate the conf still answers: plain, and with its reason + // not yet forced. `fusedUnions` returns the copy inside the codegen shell, whose reason no + // preparation rule has asked for, so what it answers below comes from the stamp alone. + val fused = withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") { + val df = rangeDF(100).repartition(2).union(rangeDF(100).repartition(2)) + val union = fusedUnions(df) + assert(union.size == 1, "this shape must fuse, or the test exercises nothing") + union.head + } + withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { + StampUnionDecisions(fused) + assert(fused.supportCodegen, "a second pass must not restamp the conf it was decided with") + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From 2c0a3e54dc57ef9c66a93cdda87b298a25e08b15 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sun, 13 Sep 2026 20:08:41 +0800 Subject: [PATCH 13/24] Pin the late barriers with extension-driven cases Three cases in SparkSessionExtensionSuite drive a union through the real pipelines: an injected columnar rule with AQE off, an injected query stage prep rule, and the columnar hook inside AQE post stage creation. Also the two review nits on rawPartitioning's scaladoc and the "whereas", plus the comment corrections three review rounds turned up. --- .../sql/execution/StampUnionDecisions.scala | 15 ++-- .../execution/basicPhysicalOperators.scala | 18 ++-- .../sql/SparkSessionExtensionSuite.scala | 82 ++++++++++++++++++- .../sql/execution/UnionCodegenSuite.scala | 12 ++- 4 files changed, 108 insertions(+), 19 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index 39e1511c88f49..090decf003561 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -24,10 +24,12 @@ import org.apache.spark.sql.catalyst.rules.Rule * defined point. * * `UnionExec` derives both from state that moves: its children's `outputPartitioning` sharpens as - * AQE finalises the plans behind them, and `conf` is the live session conf. Whoever asked first - * used to decide, which made the answer depend on when it was observed. This rule asks right - * after `EnsureRequirements`, so the decision the exchanges around a union were planned against is - * the one `unionRDDs` and the codegen gate use. + * AQE finalises the plans behind them, and `conf` is the live session conf. Every reader used to + * derive its own answer, so the answer depended on when it was read: the codegen gate could fuse a + * union whose copy in the shell then answered the other way, so `metrics` came back empty and + * `doProduce` failed asking `metricTerm` for `numOutputRows`. This rule asks right after + * `EnsureRequirements`, so the decision the exchanges around a union were planned against is the + * one `unionRDDs` and the codegen gate use. * * It is listed again after the injected columnar and query-stage rules, the hooks that can add a * `UnionExec` of their own. One created there has no decision yet, and would otherwise take one @@ -44,8 +46,9 @@ import org.apache.spark.sql.catalyst.rules.Rule * It only writes what is not there yet, so a later pass cannot move a decision already stamped on * a node, a second pass over the same nodes keeps the first answer, and a node rebuilt from a * stamped one keeps the tag `copyTagsFrom` gave it. AQE re-plans between rounds, so a union above - * the stages already created is stamped again from what that round sees, where that plan is the one - * adopted; a union inside a stage is not revisited, since `foreach` stops at `QueryStageExec`. + * the stages already created is stamped again, from what that round sees and against that round's + * exchanges; a round whose plan loses on cost is discarded whole, stamps included. A union inside a + * stage is not revisited, since `foreach` stops at `QueryStageExec`. */ object StampUnionDecisions extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index b8493486f1093..f6c2526832603 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -978,10 +978,11 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * The SPARK-52921 pass-through partitioning, derived from the children. `isPlainUnion` answers - * whether this comes back `UnknownPartitioning`; `outputPartitioning` reports it when that - * decision says the union is not a plain concatenation. The `UNION_OUTPUT_PARTITIONING` gate is - * read with the decision, not here. + * The SPARK-52921 candidate partitioning derived from the children, which `outputPartitioning` + * reports when the decision says this union is not a plain concatenation. That decision comes out + * plain on either of two grounds: `UNION_OUTPUT_PARTITIONING` being off, which is read where the + * decision is stamped rather than here, or this coming back `UnknownPartitioning`. Under that + * conf the candidate can still be concrete while the union reports unknown. */ private def rawPartitioning: Partitioning = { // Children's partitionings with attributes remapped to this union's output attributes. @@ -1050,9 +1051,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning` while its inner * `AdaptiveSparkPlanExec` has no final plan, so a union can look plain when * `CollapseCodegenStages` gates on it and partitioning-aware by the time the stage runs. The - * shell that gate builds wraps a `withNewChildren` copy, and a copy that re-derived here came - * back with empty `metrics` while `doProduce` asked `metricTerm` for `numOutputRows`. A fresh - * copy inherits the answer instead, since `withNewChildren` ends in `copyTagsFrom`. + * shell that gate builds wraps a `withNewChildren` copy where a child had to be adapted, and a + * copy that re-derived here came back with empty `metrics` while `doProduce` asked `metricTerm` + * for `numOutputRows`. A fresh copy inherits the answer instead, since `withNewChildren` ends in + * `copyTagsFrom`. * * `UNION_OUTPUT_PARTITIONING` is read where the decision is stamped rather than in * `rawPartitioning`, so it too is fixed once the plan is prepared: `conf` is live, and a plan @@ -1095,7 +1097,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * * The reverse costs fusion. A rule that runs after the stamp and drops a child's partitioning * leaves the node stamped non-plain, so the codegen gate answers "partitioning-aware" and - * `numOutputRows` goes unregistered, where re-deriving at the gate would have fused it. + * `numOutputRows` goes unregistered, whereas re-deriving at the gate would have fused it. * `DisableUnnecessaryBucketedScan` does that to a union over two bucketed scans, once a * projection on each side makes them row-based; the `columnar` term would reject the bare scans * anyway. Results are unaffected, since the other branch below re-derives and concatenates. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index 05619efa285f5..f0ec23f0adf08 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, import org.apache.spark.sql.catalyst.parser.{CatalystSqlParser, ParserInterface, SqlStatementSplitResult} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, AggregateHint, ColumnStat, Limit, LocalRelation, LogicalPlan, Project, Range, Sort, SortHint, Statistics, UnresolvedHint} -import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, SinglePartition} +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, SinglePartition, UnknownPartitioning} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.classic.ClassicConversions._ @@ -619,6 +619,67 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { } } + /** + * Prepares a plan whose `UnionExec` was made by an injected rule, then turns + * `UNION_OUTPUT_PARTITIONING` off and reads that node again. A union the `StampUnionDecisions` + * barrier following the hook reached keeps answering from the decision it was prepared with; one + * no barrier reached has no decision to answer from, so this read derives one from the conf as it + * is now and comes back `UnknownPartitioning`. `UnionCodegenSuite` covers the stamping itself by + * calling the rule directly, so it stays green if one of the post-hook listings is dropped; these + * pin the post-hook stamping in each pipeline, the two AQE listings jointly rather than one each. + */ + private def checkInjectedUnionIsStamped( + extensions: Seq[SparkSessionExtensionsProvider], aqeEnabled: Boolean): Unit = { + withSession(extensions) { session => + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, aqeEnabled) + // The union ends up over the projection above a repartition by `k`, so it has a concrete + // partitioning to pass through. A filter would not do: `PushPredicateThroughNonJoin` pushes + // it below the exchange. + val df = session.range(0, 20, 1, 2).selectExpr("id % 5 AS k", "id AS v") + .repartition(4, col("k")).selectExpr("k", "v + 1 AS w") + // Also what makes the final adaptive plan available in the two AQE cases. `w` is `id + 1`, so + // a union that dropped or duplicated a row would show up here. + assert(df.collect().map(_.getLong(1)).sorted.toSeq == (1L to 20L).toSeq, + "the injected union must not drop or duplicate rows") + + val unions = collect(df.queryExecution.executedPlan) { case u: UnionExec => u } + assert(unions.size == 1, s"expected the one union the rule adds, got ${unions.size}") + val union = unions.head + val prepared = union.outputPartitioning + assert(!prepared.isInstanceOf[UnknownPartitioning], + s"this union must be prepared partitioning-aware, got $prepared") + + session.conf.set(SQLConf.UNION_OUTPUT_PARTITIONING.key, false) + assert(union.outputPartitioning == prepared, + "no barrier stamped the union the rule added, so this conf change decided it: " + + s"${union.outputPartitioning}") + } + } + + test("SPARK-59122: the barrier after the injected columnar rules stamps a union they added") { + checkInjectedUnionIsStamped( + create(_.injectColumnar(_ => WrapRootInUnionColumnarRule)), aqeEnabled = false) + } + + test("SPARK-59122: a union an injected query stage prep rule adds is stamped before execution") { + // Attributed to the AQE pipeline rather than to one listing: the barrier in + // `postStageCreationRules` stands behind the one after the prep rules, so this case fails only + // when both are gone. What the earlier one adds is that the answer is fixed before the + // stage-optimizer rules read it, as `CoalesceShufflePartitions` does to decide whether a + // union's children have to be coalesced as one group. + checkInjectedUnionIsStamped( + create(_.injectQueryStagePrepRule(_ => WrapRootInUnion)), aqeEnabled = true) + } + + test("SPARK-59122: the barrier in AQE post stage creation stamps a union added there") { + // With AQE on, the columnar rules reach this plan only through `postStageCreationRules`: the + // root of the plan `QueryExecution.preparations` hands them is `AdaptiveSparkPlanExec`, which + // `WrapRootInUnion` leaves alone. An injected stage-optimizer rule runs ahead of the same + // barrier, so it needs no case of its own. + checkInjectedUnionIsStamped( + create(_.injectColumnar(_ => WrapRootInUnionColumnarRule)), aqeEnabled = true) + } + test("custom aggregate hint") { // The custom hint allows us to replace the aggregate (without grouping keys) with just // Literal. @@ -1387,6 +1448,25 @@ object MyQueryPostPlannerStrategyRule extends Rule[SparkPlan] { } } +/** + * Stands for an extension that introduces a `UnionExec` of its own: replaces a root `ProjectExec` + * with a fresh one-child union, which carries no stamped decision because it is a new instance, and + * leaves the rows alone. Matching only the root keeps it idempotent, since the root is a + * `UnionExec` afterwards; matching a `ProjectExec` keeps it out of the way of + * `postStageCreationRules`, which requires the exchange it is handed to stay an exchange. + */ +object WrapRootInUnion extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan match { + case p: ProjectExec => UnionExec(Seq(p)) + case other => other + } +} + +/** The columnar-rule wrapper for `WrapRootInUnion`. */ +object WrapRootInUnionColumnarRule extends ColumnarRule { + override def postColumnarTransitions: Rule[SparkPlan] = WrapRootInUnion +} + // Example of an Aggregate hint that tells that 'attribute' values are no larger than 'max'. // We will use them to rewrite MAX(attribute) with 'max' constant. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 0dab65be8c567..26d9070d24939 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -836,6 +836,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val nested = UnionExec(Seq(union.children.head, union.children.head)) val rebuilt = union.withNewChildren(Seq(nested, union.children.last)).asInstanceOf[UnionExec] assert(!rebuilt.supportCodegen, "the rebuilt union must answer against its own children") + // Implied by the line above as the code stands, and kept as the pin on that: registering the + // metric unconditionally would leave the line above green, and only this one would fail. assert(rebuilt.metrics.isEmpty) } } @@ -894,13 +896,15 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // `StampUnionDecisions` is listed again after the phases that can add a `UnionExec`, so one an // injected columnar or query-stage rule created does not answer from whatever the conf says // wherever it is first asked. A later pass must also not move a decision already taken, which - // is the second half here. The rule is driven directly: injecting an extension needs its own - // session, and what matters is the rule's contract. + // is the second half here. The rule is driven directly, since what this case is about is its + // contract; that the pipelines still list it after each phase that can add a union is pinned + // from the outside by the extension-driven cases in `SparkSessionExtensionSuite`, which need a + // session of their own. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { // Pins the property the standard pipeline has to keep: a stamping pass runs after // `EnsureRequirements`, so the decision is taken from the plan the exchanges were placed in. - // A count would break on a sixth legitimate pass and say nothing about the order. The three - // AQE positions are private to `AdaptiveSparkPlanExec`. + // A count would break on a sixth legitimate pass and say nothing about the order. The AQE + // lists are private to `AdaptiveSparkPlanExec`, so their first pass has no counterpart here. val rules = QueryExecution.preparations(spark, subquery = false) val firstStamp = rules.indexWhere(_ eq StampUnionDecisions) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) From 65e3f96ee6c134ef2c80d40a25d3a50e2ac912c4 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 14 Sep 2026 04:08:48 +0800 Subject: [PATCH 14/24] Snapshot the partitioning conf before EnsureRequirements The rule and the stamp behind it sampled the conf separately, so a concurrent session-conf change between them could leave a parent's elided exchange over a union that then concatenates. Also gives the barrier after the AQE prep rules its own failure signal, through an injected stage-optimizer rule that observes the union. --- .../apache/spark/sql/internal/SQLConf.scala | 4 +- .../spark/sql/execution/QueryExecution.scala | 4 ++ .../sql/execution/StampUnionDecisions.scala | 29 ++++++++++ .../adaptive/AdaptiveSparkPlanExec.scala | 4 ++ .../execution/basicPhysicalOperators.scala | 43 ++++++++++++--- .../sql/SparkSessionExtensionSuite.scala | 43 ++++++++++++--- .../sql/execution/UnionCodegenSuite.scala | 53 ++++++++++++++++--- 7 files changed, 156 insertions(+), 24 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 01bed12871c10..e477d7bb19397 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -8125,8 +8125,8 @@ object SQLConf { .internal() .doc("When set to true, the output partitioning of UnionExec will be the same as the " + "input partitioning if its children have same partitioning. Otherwise, it will be a " + - "default partitioning. The value is read when a UnionExec's decision is fixed during " + - "physical preparation, so a change does not reach a decision already taken.") + "default partitioning. The value is read during physical preparation, so a change does " + + "not reach a decision already taken.") .version("4.1.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index 79f1f66a5a839..f52c0f6dc7da3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -798,6 +798,10 @@ object QueryExecution { PlanDynamicPruningFilters(sparkSession), PlanSubqueries(sparkSession), RemoveRedundantProjects, + // Must run before `EnsureRequirements`, which asks a `UnionExec` what it reports: it + // records the conf that answer depends on, so the following `StampUnionDecisions` freezes the + // decision under the same value the exchanges were planned against. + SnapshotUnionOutputPartitioningConf, EnsureRequirements(), // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, and // the answer to fix is the one the exchanges around it were planned against. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index 090decf003561..f812fc945a084 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.internal.SQLConf /** * Fixes each [[UnionExec]]'s partitioning decision and the confs its codegen gate reads, at one @@ -59,3 +60,31 @@ object StampUnionDecisions extends Rule[SparkPlan] { plan } } + +/** + * Records on each [[UnionExec]] the `UNION_OUTPUT_PARTITIONING` value it answers from, before + * `EnsureRequirements` asks it what it reports. + * + * Without this, the two phases sample the conf separately: `EnsureRequirements` reads what the + * union reports under the value then, and [[StampUnionDecisions]] freezes the decision under the + * value one rule later. `conf` is the live session conf, so another thread turning it off in that + * window would let a parent drop an exchange over a concrete partitioning and then have the union + * concatenate, which puts one group in two partitions. + * + * Only the conf is recorded, never a partitioning. `EnsureRequirements` has not inserted the + * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only + * become co-partitioned there, which is why the decision itself waits for the barrier behind it. + * + * One read per plan, so every union in it answers from the same value. Writing the tag in place is + * safe for the reason given on [[StampUnionDecisions]]. + */ +object SnapshotUnionOutputPartitioningConf extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + val enabled = plan.conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) + plan.foreach { + case u: UnionExec => u.snapshotOutputPartitioningConf(enabled) + case _ => + } + plan + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index afaa330a502dd..9ba46fae9288d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -134,6 +134,10 @@ case class AdaptiveSparkPlanExec( Seq( CoalesceBucketsInJoin, RemoveRedundantProjects, + // Must run before `ensureRequirements`, which asks a `UnionExec` what it reports: it + // records the conf that answer depends on, so the following `StampUnionDecisions` freezes the + // decision under the same value the exchanges were planned against. + SnapshotUnionOutputPartitioningConf, ensureRequirements, // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, so // every rule below and the execution itself read the answer the exchanges above it were diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index f6c2526832603..e029bb492719a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -980,9 +980,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup /** * The SPARK-52921 candidate partitioning derived from the children, which `outputPartitioning` * reports when the decision says this union is not a plain concatenation. That decision comes out - * plain on either of two grounds: `UNION_OUTPUT_PARTITIONING` being off, which is read where the - * decision is stamped rather than here, or this coming back `UnknownPartitioning`. Under that - * conf the candidate can still be concrete while the union reports unknown. + * plain on either of two grounds: `UNION_OUTPUT_PARTITIONING` being off, which is taken from the + * record `SnapshotUnionOutputPartitioningConf` writes rather than read here, or this coming back + * `UnknownPartitioning`. Under that conf the candidate can still be concrete while the union + * reports unknown. */ private def rawPartitioning: Partitioning = { // Children's partitionings with attributes remapped to this union's output attributes. @@ -1056,21 +1057,39 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * for `numOutputRows`. A fresh copy inherits the answer instead, since `withNewChildren` ends in * `copyTagsFrom`. * - * `UNION_OUTPUT_PARTITIONING` is read where the decision is stamped rather than in - * `rawPartitioning`, so it too is fixed once the plan is prepared: `conf` is live, and a plan - * must execute by the partitioning it was planned against. + * `UNION_OUTPUT_PARTITIONING` is taken from `snapshotOutputPartitioningConf`, recorded before + * `EnsureRequirements`, so the value the exchanges are planned against is the value execution + * uses; a node created after that pass carries no record and reads the live conf. Reading it live + * here would leave one rule between the two: `conf` is live, and another thread setting it in + * that window would let a parent drop an exchange over a concrete partitioning and then have the + * stamp freeze plain concatenation under it. * * A read before `StampUnionDecisions` answers from the children as they are then, and does not * write, so observing an unprepared plan cannot decide anything for the prepared one. */ private[execution] def isPlainUnion: Boolean = stampedDecisions.map(_.plainUnion).getOrElse { - !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) || - rawPartitioning.isInstanceOf[UnknownPartitioning] + !outputPartitioningEnabled || rawPartitioning.isInstanceOf[UnknownPartitioning] } private def stampedDecisions: Option[UnionExec.Decisions] = getTagValue(UnionExec.DECISIONS) + private def outputPartitioningEnabled: Boolean = + getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF) + .getOrElse(conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) + + /** + * Records the conf `isPlainUnion` answers from, read once for the whole plan by + * `SnapshotUnionOutputPartitioningConf` and passed in here, ahead of `EnsureRequirements`, whose + * reads the following stamp has to agree with. Only the conf, never a partitioning: the exchanges + * `EnsureRequirements` adds are not there yet, so a decision taken here would freeze plain on a + * union whose children only become co-partitioned there. + */ + private[execution] def snapshotOutputPartitioningConf(enabled: Boolean): Unit = + if (getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF).isEmpty) { + setTagValue(UnionExec.OUTPUT_PARTITIONING_CONF, enabled) + } + /** * Fixes this node's decisions for the rest of the plan's life. Called by `StampUnionDecisions`, * first right after `EnsureRequirements`, so what the exchanges around this union were planned @@ -1394,6 +1413,14 @@ object UnionExec { */ private val DECISIONS = TreeNodeTag[Decisions]("unionDecisions") + /** + * The `UNION_OUTPUT_PARTITIONING` value `isPlainUnion` answers from until the decision is + * stamped. See `snapshotOutputPartitioningConf`. Written before `EnsureRequirements` and read by + * the stamp after it, so both phases use one value; travels onto rebuilt nodes the same way + * `DECISIONS` does, which is what carries it across the copies `EnsureRequirements` makes. + */ + private val OUTPUT_PARTITIONING_CONF = TreeNodeTag[Boolean]("unionOutputPartitioningConf") + /** * Codegen operators that return more than one RDD from `inputRDDs()`. * `UnionExec`'s fusion assumes each direct child contributes one RDD. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index f0ec23f0adf08..a21ec970f80fe 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -626,7 +626,7 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { * no barrier reached has no decision to answer from, so this read derives one from the conf as it * is now and comes back `UnknownPartitioning`. `UnionCodegenSuite` covers the stamping itself by * calling the rule directly, so it stays green if one of the post-hook listings is dropped; these - * pin the post-hook stamping in each pipeline, the two AQE listings jointly rather than one each. + * pin the post-hook stamping in each pipeline. */ private def checkInjectedUnionIsStamped( extensions: Seq[SparkSessionExtensionsProvider], aqeEnabled: Boolean): Unit = { @@ -662,13 +662,21 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { } test("SPARK-59122: a union an injected query stage prep rule adds is stamped before execution") { - // Attributed to the AQE pipeline rather than to one listing: the barrier in - // `postStageCreationRules` stands behind the one after the prep rules, so this case fails only - // when both are gone. What the earlier one adds is that the answer is fixed before the - // stage-optimizer rules read it, as `CoalesceShufflePartitions` does to decide whether a - // union's children have to be coalesced as one group. + // The barrier in `postStageCreationRules` stands behind the one after the prep rules, so the + // conf flip in `checkInjectedUnionIsStamped` cannot tell them apart. What only the earlier one + // can do is have the answer ready for the stage optimizers, which run in between and read it: + // `CoalesceShufflePartitions` asks whether a union's children have to be coalesced as one + // group. `ObserveUnionPartitioning` reads the node from there, so removing the earlier barrier + // fails this case. + val seen = ListBuffer.empty[Partitioning] checkInjectedUnionIsStamped( - create(_.injectQueryStagePrepRule(_ => WrapRootInUnion)), aqeEnabled = true) + create { extensions => + extensions.injectQueryStagePrepRule(_ => WrapRootInUnion) + extensions.injectQueryStageOptimizerRule(_ => ObserveUnionPartitioning(seen)) + }, aqeEnabled = true) + assert(seen.nonEmpty, "the stage optimizers must have seen the union") + assert(!seen.exists(_.isInstanceOf[UnknownPartitioning]), + s"the barrier after the prep rules must decide before the stage optimizers read: $seen") } test("SPARK-59122: the barrier in AQE post stage creation stamps a union added there") { @@ -1467,6 +1475,27 @@ object WrapRootInUnionColumnarRule extends ColumnarRule { override def postColumnarTransitions: Rule[SparkPlan] = WrapRootInUnion } +/** + * Records what each `UnionExec` reports while the AQE stage optimizers run, which is after the + * barrier at the end of the query stage preparation rules and before the one in + * `postStageCreationRules`. Reads it with `UNION_OUTPUT_PARTITIONING` turned off: the union an + * injected prep rule adds carries no recorded conf of its own, so a concrete answer can only come + * from a decision stamped earlier. Puts the conf back, so nothing downstream sees the flip. + */ +case class ObserveUnionPartitioning(seen: ListBuffer[Partitioning]) extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + plan.foreach { + case u: UnionExec => + val enabled = u.conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) + u.conf.setConf(SQLConf.UNION_OUTPUT_PARTITIONING, false) + try seen += u.outputPartitioning + finally u.conf.setConf(SQLConf.UNION_OUTPUT_PARTITIONING, enabled) + case _ => + } + plan + } +} + // Example of an Aggregate hint that tells that 'attribute' values are no larger than 'max'. // We will use them to rewrite MAX(attribute) with 'max' constant. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 26d9070d24939..70d673f943789 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -709,13 +709,14 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper test("SPARK-59122: a partitioning-aware union keeps its layout when the conf changes between " + "planning and execution") { - // `spark.sql.unionOutputPartitioning` is read where the plain-union decision is stamped, not on - // every `outputPartitioning` call, so a plan executes by the partitioning it was planned - // against. Reading it per call let the parent aggregate lose its exchange at planning and get a - // plain concatenation at execution, reporting each group twice. The `checkAnswer` below stays - // outside the block that planned the DataFrame on purpose: the plan is forced inside that - // block and `executedPlan` is memoized, so the two phases see different confs. Asserting - // inside it, or dropping the second `withSQLConf`, makes the test pass without testing this. + // `spark.sql.unionOutputPartitioning` is read once during preparation, ahead of + // `EnsureRequirements`, not on every `outputPartitioning` call, so a plan executes by the + // partitioning it was planned against. Reading it per call let the parent aggregate lose its + // exchange at planning and get a plain concatenation at execution, reporting each group twice. + // The `checkAnswer` below stays outside the block that planned the DataFrame on purpose: the + // plan is forced inside that block and `executedPlan` is memoized, so the two phases see + // different confs. Asserting inside it, or dropping the second `withSQLConf`, makes the test + // pass without testing this. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val left = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k") val right = spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k") @@ -947,6 +948,44 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: the stamp uses the conf the exchanges were planned against") { + // `EnsureRequirements` asks the union what it reports, and the barrier behind it freezes that + // answer one rule later. `conf` is live, so another thread turning `UNION_OUTPUT_PARTITIONING` + // off in between would leave the parent's elided exchange standing over a union that then + // concatenates. `SnapshotUnionOutputPartitioningConf` records the value ahead of + // `EnsureRequirements` for both to use. Driven rule by rule, because the two sit next to each + // other in the pipeline and no injected rule can run in the window. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + // The pipeline has to keep the two on either side of `EnsureRequirements`; the AQE list is + // private, so only the standard one can be checked from here. + val rules = QueryExecution.preparations(spark, subquery = false) + val snapshot = rules.indexWhere(_ eq SnapshotUnionOutputPartitioningConf) + val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) + assert(snapshot >= 0 && snapshot < ensureRequirements, + s"expected the conf snapshot before EnsureRequirements, got $snapshot/$ensureRequirements") + + 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() + val required = EnsureRequirements()( + SnapshotUnionOutputPartitioningConf(df.queryExecution.sparkPlan.clone())) + assert(required.collect { case s: ShuffleExchangeExec => s.shuffleOrigin } == + Seq(REPARTITION_BY_NUM, REPARTITION_BY_NUM), + "the aggregate's exchange must have been elided, or the window has nothing at stake") + + val union = required.collect { case u: UnionExec => u } + assert(union.size == 1) + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + StampUnionDecisions(required) + assert(!union.head.outputPartitioning.isInstanceOf[UnknownPartitioning], + "the answer must come from the conf snapshot, not from the value read now, got " + + s"${union.head.outputPartitioning}") + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From 708ec9d79227edf7791b0128f035518bb272d0b7 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 15 Sep 2026 12:47:18 +0800 Subject: [PATCH 15/24] Read the codegen confs once per plan, and swap a redundant case The conf snapshot made the "nothing read it during preparation" case pin nothing, so it becomes a case for the other half of that design note: a partitioning-aware union has to follow its children's coalesced partition count, which fails with an AIOOBE if the Partitioning itself is stamped. Also bounds the rule-order assertion above and pins the InputAdapter premise the two conf-flip cases rely on. --- .../sql/execution/StampUnionDecisions.scala | 4 +- .../execution/basicPhysicalOperators.scala | 39 ++++++------ .../sql/execution/UnionCodegenSuite.scala | 60 +++++++++++-------- 3 files changed, 58 insertions(+), 45 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index f812fc945a084..2faa02c4492b3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -53,8 +53,10 @@ import org.apache.spark.sql.internal.SQLConf */ object StampUnionDecisions extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { + val codegenEnabled = plan.conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED) + val maxChildren = plan.conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN) plan.foreach { - case u: UnionExec => u.stampDecisions() + case u: UnionExec => u.stampDecisions(codegenEnabled, maxChildren) case _ => } plan diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index e029bb492719a..2c77000d8e5dc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1093,17 +1093,19 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup /** * Fixes this node's decisions for the rest of the plan's life. Called by `StampUnionDecisions`, * first right after `EnsureRequirements`, so what the exchanges around this union were planned - * against is what execution uses. Nothing else writes this tag on an existing node, and the - * nodes the rule writes are freshly planned and not yet published, so no reader can be looking at - * one; `metrics` and the codegen gate read it later, and a node that already carries it keeps it, - * which is how the copy in the codegen shell stays in step with the gate. + * against is what execution uses; the two confs come from one read per plan there. Nothing else + * writes this tag on an existing node, and the nodes the rule writes are freshly planned and not + * yet published, so no reader can be looking at one; `metrics` and the codegen gate read it + * later, and a node that already carries it keeps it, which is how the copy in the codegen shell + * stays in step with the gate. */ - private[execution] def stampDecisions(): Unit = if (stampedDecisions.isEmpty) { - setTagValue(UnionExec.DECISIONS, UnionExec.Decisions( - plainUnion = isPlainUnion, - unionCodegenEnabled = conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED), - maxChildren = conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN))) - } + private[execution] def stampDecisions(codegenEnabled: Boolean, maxChildren: Int): Unit = + if (stampedDecisions.isEmpty) { + setTagValue(UnionExec.DECISIONS, UnionExec.Decisions( + plainUnion = isPlainUnion, + unionCodegenEnabled = codegenEnabled, + maxChildren = maxChildren)) + } /** * A node stamped plain reports `UnknownPartitioning` even once its children agree on a concrete @@ -1117,17 +1119,14 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * The reverse costs fusion. A rule that runs after the stamp and drops a child's partitioning * leaves the node stamped non-plain, so the codegen gate answers "partitioning-aware" and * `numOutputRows` goes unregistered, whereas re-deriving at the gate would have fused it. - * `DisableUnnecessaryBucketedScan` does that to a union over two bucketed scans, once a - * projection on each side makes them row-based; the `columnar` term would reject the bare scans - * anyway. Results are unaffected, since the other branch below re-derives and concatenates. + * `DisableUnnecessaryBucketedScan` does that to a union over two bucketed scans with a projection + * on each side. Results are unaffected, since the branch below re-derives and concatenates. * - * The other branch is derived per call, so `unionRDDs` could take the concatenating arm even - * though `EnsureRequirements` planned the parent against a concrete partitioning: - * `comparePartitioning` compares `HashPartitioningLike` by equality, so a change to one child's - * partitioning that its siblings do not mirror can empty the intersection. This node does not - * re-check it. AQE reconciles it, by validating a partitioning change against the parents' - * requirements and either reverting it or re-running `EnsureRequirements`; an injected rule can - * skip that. + * That branch is derived per call, so `unionRDDs` can take the concatenating arm even though + * `EnsureRequirements` planned the parent against a concrete partitioning: `comparePartitioning` + * compares `HashPartitioningLike` by equality, so a change to one child's partitioning that its + * siblings do not mirror can empty the intersection. AQE reconciles that, by validating a + * partitioning change against the parents' requirements; an injected rule can skip it. */ override def outputPartitioning: Partitioning = if (isPlainUnion) super.outputPartitioning else rawPartitioning diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 70d673f943789..303f3bdb1dca1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -677,6 +677,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // distribution from an RDD that does not have it. withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { withTempView("v") { cacheAggregateView("v") @@ -770,12 +771,14 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(planned.collect().length == 200) // The row count alone does not discriminate, since the shell was installed at planning and // keeps emitting; registering `numOutputRows` unconditionally and reading the conf per call - // passes it. This assertion is what fails there, because nothing forces the copy's reason - // before it. It has to sit after the flip, as it does here: taken while the conf was still - // on, it would warm a memoizing implementation with the answer this test needs it not to - // have. + // passes it. The `supportCodegen` assertion below is what fails there. val copy = fusedUnions(planned) assert(copy.size == 1) + // The copy this test needs: `insertInputAdapter` wrapped both children, so the shell holds + // a copy rather than the instance the gate answered on. This copy's reason is first forced + // by the `SparkPlanInfo` that `collect()` above builds, with the conf already off, so what + // it answers can only come from the stamp. + assert(copy.head.children.forall(_.isInstanceOf[InputAdapter])) assert(copy.head.supportCodegen, "the copy in the shell must keep the decision it was planned with") } @@ -806,6 +809,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(planned.collect().length == 300) val copy = fusedUnions(planned) assert(copy.size == 1) + assert(copy.head.children.forall(_.isInstanceOf[InputAdapter])) assert(copy.head.supportCodegen, "the copy in the shell must keep the cap it was planned with") // Not `metrics.contains`, which `collect()` above already proves: an empty `metrics` would @@ -871,25 +875,29 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } - test("SPARK-59122: a prepared union keeps its layout when nothing read it during preparation") { - // With whole-stage codegen off, no gate consults the union while the plan is prepared, and a - // root union has no parent to ask for its partitioning either. First-read initialization would - // then decide at execution, under whatever the conf says by then; `StampUnionDecisions` decides - // during preparation instead. + test("SPARK-59122: a partitioning-aware union follows its children's coalesced partition count") { + // Only the decision is stamped, never the `Partitioning`. AQE coalescing changes the children's + // `numPartitions` after the stamp, and `unionRDDs` hands whatever it reports to + // `SQLPartitioningAwareUnionRDD`, which builds exactly that many partitions from each child: a + // count frozen at stamping time asks for partitions the coalesced children no longer have. withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") { - val plan = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { - 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"))) - .queryExecution.executedPlan - } - withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { - // Co-partitioned children pass their four partitions through; a plain concatenation would - // report eight. - assert(plan.execute().getNumPartitions == 4, - "a prepared union must execute by the layout it was prepared with") - } + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "20", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + val left = spark.range(0, 100, 1, 4).selectExpr("id % 10 AS k").groupBy("k").count() + val right = spark.range(100, 200, 1, 4).selectExpr("id % 10 AS k").groupBy("k").count() + val df = left.union(right).groupBy("k").agg(sum("count").as("c")) + checkAnswer(df, (0L until 10L).map(k => Row(k, 20L))) + + val unions = collect(df.queryExecution.executedPlan) { case u: UnionExec => u } + assert(unions.size == 1) + val children = unions.head.children.map(_.outputPartitioning.numPartitions) + assert(children.distinct.size == 1 && children.head < 20, + s"the children must have been coalesced as one group, got $children") + assert(unions.head.outputPartitioning.numPartitions == children.head, + "the union must report what its children report now, got " + + s"${unions.head.outputPartitioning}") } } @@ -909,8 +917,12 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val rules = QueryExecution.preparations(spark, subquery = false) val firstStamp = rules.indexWhere(_ eq StampUnionDecisions) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) - assert(ensureRequirements >= 0 && firstStamp > ensureRequirements, - s"expected a stamping pass after EnsureRequirements, got $ensureRequirements/$firstStamp") + val columnarRules = + rules.indexWhere(_.isInstanceOf[ApplyColumnarRulesAndInsertTransitions]) + assert(ensureRequirements >= 0 && firstStamp > ensureRequirements && + firstStamp < columnarRules, + "expected a stamping pass between EnsureRequirements and the columnar rules, got " + + s"$ensureRequirements/$firstStamp/$columnarRules") val stamped = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { val df = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k")) From 2b0022e3ac547cd01e4abc9c37e6da71eeca5fe3 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Wed, 16 Sep 2026 23:25:44 +0800 Subject: [PATCH 16/24] Prove the two flag states apart in the union codegen tests - assertFlagParity builds a fresh DataFrame per flag value and asserts the off half took no union into codegen; five cases that reused one DataFrame went through it, so they no longer compare a plan against itself - codegenUnions replaces unionInsideWSCG: a union under an InputAdapter is inside a stage without taking part in codegen - AdaptiveQueryExecSuite pins the snapshot, EnsureRequirements and the two stamps in AQE's own preparation list --- .../adaptive/AdaptiveSparkPlanExec.scala | 3 +- .../sql/execution/UnionCodegenSuite.scala | 180 +++++++++++------- .../adaptive/AdaptiveQueryExecSuite.scala | 41 +++- 3 files changed, 150 insertions(+), 74 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index 48d42cc20f044..a34f75e267b87 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -122,7 +122,8 @@ case class AdaptiveSparkPlanExec( // A list of physical plan rules to be applied before creation of query stages. The physical // plan should reach a final status of query stages (i.e., no more addition or removal of // Exchange nodes) after running these rules. - @transient private val queryStagePreparationRules: Seq[Rule[SparkPlan]] = { + // Visible in the package so that a test can assert where the two union barriers below sit. + @transient private[adaptive] val queryStagePreparationRules: Seq[Rule[SparkPlan]] = { // For cases like `df.repartition(a, b).select(c)`, there is no distribution requirement for // the final plan, but we do need to respect the user-specified repartition. Here we ask // `EnsureRequirements` to not optimize out the user-specified repartition-by-col to work diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 303f3bdb1dca1..20fb061dfccf9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, UnknownPartitioning} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.columnar.InMemoryTableScanLike import org.apache.spark.sql.execution.exchange.{EnsureRequirements, REPARTITION_BY_NUM, ShuffleExchangeExec} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -56,18 +57,12 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper case s: WholeStageCodegenExec => s }.size - private def unionInsideWSCG(df: DataFrame): Boolean = - df.queryExecution.executedPlan.collect { - case w: WholeStageCodegenExec if w.find(_.isInstanceOf[UnionExec]).isDefined => w - }.nonEmpty - /** * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query stages; * `SparkPlan.collect` stops at them, since both are `LeafExecNode`s. * - * Stricter than `unionInsideWSCG` on purpose: this matches only a union that is the root of its - * own codegen stage, which is what "fused" means for the callers here, while `w.find` also - * matches one that an `InputAdapter` left inside the stage. + * Stricter than `codegenUnions` on purpose: this matches only a union that is the root of its own + * codegen stage, which is the node the callers here reach for its tags and metrics. */ private def fusedUnions(df: DataFrame): Seq[UnionExec] = collect(df.queryExecution.executedPlan) { @@ -75,6 +70,29 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper w.child.asInstanceOf[UnionExec] } + /** + * The unions that take part in a codegen stage: inside a `WholeStageCodegenExec` with no + * `InputAdapter` between them and the stage root. That is what fusion means. Asking only whether + * a stage holds a `UnionExec` somewhere does not answer it, since `CollapseCodegenStages` leaves + * a non-participating union inside the stage too, under an `InputAdapter`; `fusedUnions` asks for + * the stronger property of rooting the stage. + * + * Ask this of a plan that has run. `CollapseCodegenStages` is a post-stage-creation rule under + * AQE, so a wrapped plan that never executed holds no `WholeStageCodegenExec` at all and the + * answer is empty whatever the confs say. `AdaptiveSparkPlanHelper.collect` is what descends + * through AQE wrappers and query stages; `SparkPlan.collect` stops at them, since both are + * `LeafExecNode`s. + */ + private def codegenUnions(df: DataFrame): Seq[UnionExec] = { + def participating(plan: SparkPlan): Seq[UnionExec] = plan match { + case _: InputAdapter => Nil + case u: UnionExec => u +: u.children.flatMap(participating) + case other => other.children.flatMap(participating) + } + collect(df.queryExecution.executedPlan) { case w: WholeStageCodegenExec => w } + .flatMap(w => participating(w.child)) + } + /** * A cached aggregate, so the union's children read an `InMemoryTableScanExec`. The caller needs * the cache unmaterialized; `withTempView` drops the view on the way out, and `dropTempView` @@ -89,11 +107,25 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper spark.catalog.cacheTable(view) } - /** Run query with flag on, then flag off, assert results match. */ - protected def assertFlagParity(buildDf: () => DataFrame): Unit = { - val onRows = buildDf().collect().toSeq + /** + * Run `buildDf()` with union codegen on, then again with it off, and assert the two agree. + * + * A fresh DataFrame per flag value, and not one built outside: `queryExecution` is memoized and + * the decision is stamped at preparation, so collecting the same DataFrame twice replays the plan + * the first value prepared and compares its output against itself. The off half also has to show + * no union took part in codegen, which is only observable once the plan has executed, hence + * after `checkAnswer`. Whether the on half fused anything is the caller's to assert, since a + * fallback case is unfused either way: both DataFrames are returned, collected, for that. + */ + protected def assertFlagParity(buildDf: () => DataFrame): (DataFrame, DataFrame) = { + val fused = buildDf() + val onRows = fused.collect().toSeq withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - checkAnswer(buildDf(), onRows) + val plain = buildDf() + checkAnswer(plain, onRows) + assert(codegenUnions(plain).isEmpty, + s"expected no union taking part in codegen with the flag off:\n${plain.queryExecution}") + (fused, plain) } } @@ -117,14 +149,14 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper test("SPARK-56482: plain union with filter fuses into one WSCG stage") { val df = rangeDF(100).union(rangeDF(100)).filter(col("id") > 0) assert(wscgCount(df) == 1) - assert(unionInsideWSCG(df)) + assert(codegenUnions(df).nonEmpty) } test("SPARK-56482: flag off restores pre-patch plan shape") { withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { val df = rangeDF(100).union(rangeDF(100)).filter(col("id") > 0) assert(wscgCount(df) >= 2) - assert(!unionInsideWSCG(df)) + assert(codegenUnions(df).isEmpty) } } @@ -133,7 +165,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "2") { val df = rangeDF(10).union(rangeDF(10)).union(rangeDF(10)) - assert(!unionInsideWSCG(df)) + assert(codegenUnions(df).isEmpty) } } @@ -231,7 +263,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val b = rangeDF(3).select(col("id").cast(DecimalType(10, 2)).as("v")) a.union(b) } - assert(unionInsideWSCG(build().filter(col("v") >= 0)), + assert(codegenUnions(build().filter(col("v") >= 0)).nonEmpty, "decimal precision/scale widening should still fuse into one WSCG stage") assertFlagParity(() => build().orderBy("v")) } @@ -254,7 +286,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val a = rangeDF(3).select(col("id").cast(IntegerType).as("v")) val b = rangeDF(3).select(col("id").as("v")) val df = a.union(b).filter(col("v") >= 0) - assert(unionInsideWSCG(df), + assert(codegenUnions(df).nonEmpty, "widened-children Union should fuse with filter into a single WSCG stage") checkAnswer(df, Seq(Row(0L), Row(1L), Row(2L), Row(0L), Row(1L), Row(2L))) } @@ -275,7 +307,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val b = spark.createDataFrame( java.util.Arrays.asList(Row(Row(3)), Row(Row(4))), structOuterNullable) val df = a.union(b) - assert(!unionInsideWSCG(df), + assert(codegenUnions(df).isEmpty, "Nested-nullability mismatch must fall back to non-codegen") val unionExec = df.queryExecution.executedPlan.collectFirst { case u: UnionExec => u @@ -298,7 +330,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val b = spark.createDataFrame( java.util.Arrays.asList(Row(java.util.Arrays.asList(3, 4))), schemaNullable) val df = a.union(b) - assert(!unionInsideWSCG(df), + assert(codegenUnions(df).isEmpty, "Array containsNull mismatch must fall back to non-codegen") val unionExec = df.queryExecution.executedPlan.collectFirst { case u: UnionExec => u @@ -432,16 +464,15 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf( SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - val left = rangeDF(100).select(col("id").as("lk"), col("id").as("lv")) - val right = rangeDF(100).select(col("id").as("rk")) - val bhj = left.join(broadcast(right), col("lk") === col("rk")) - .select("lk", "lv") - val df = bhj.union( - rangeDF(100).select(col("id").as("lk"), col("id").as("lv"))) - val flagOn = df.collect().toSeq - withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - checkAnswer(df, flagOn) + val (fused, _) = assertFlagParity { () => + val left = rangeDF(100).select(col("id").as("lk"), col("id").as("lv")) + val right = rangeDF(100).select(col("id").as("rk")) + val bhj = left.join(broadcast(right), col("lk") === col("rk")) + .select("lk", "lv") + bhj.union(rangeDF(100).select(col("id").as("lk"), col("id").as("lv"))) } + assert(codegenUnions(fused).size == 1, + s"the aliasing this case is about needs the fused path:\n${fused.queryExecution}") } } @@ -449,17 +480,19 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf( SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - val probe = rangeDF(10).select(col("id").as("k")) - val build = rangeDF(20) - .select((col("id") % 5).as("k"), col("id").as("v")) - val bhj = probe.join(broadcast(build), "k") - val df = bhj.union( - rangeDF(0).select(col("id").as("k"), col("id").as("v"))) - val agg = df.groupBy("k").count().orderBy("k") - val flagOn = agg.collect() - withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - checkAnswer(agg, flagOn.toSeq) + val (fused, _) = assertFlagParity { () => + val probe = rangeDF(10).select(col("id").as("k")) + val build = rangeDF(20) + .select((col("id") % 5).as("k"), col("id").as("v")) + val bhj = probe.join(broadcast(build), "k") + val df = bhj.union( + rangeDF(0).select(col("id").as("k"), col("id").as("v"))) + df.groupBy("k").count().orderBy("k") } + // The partial aggregate roots the stage here and the union is fused into it, so this is the + // case `fusedUnions` would miss. + assert(codegenUnions(fused).size == 1, + s"the aliasing this case is about needs the fused path:\n${fused.queryExecution}") } } @@ -473,14 +506,16 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf( SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { - val left = rangeDF(100).select(col("id").as("k")) - val right = rangeDF(100).select(col("id").as("k")) - val smj = left.join(right, "k") - val df = smj.union(rangeDF(100).select(col("id").as("k"))) - val flagOn = df.collect().toSeq - withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - checkAnswer(df, flagOn) + val (fused, _) = assertFlagParity { () => + val left = rangeDF(100).select(col("id").as("k")) + val right = rangeDF(100).select(col("id").as("k")) + val smj = left.join(right, "k") + smj.union(rangeDF(100).select(col("id").as("k"))) } + // Not `fusedUnions`: this union would not root a stage anyway, so that would hold with or + // without the denylist. What the denylist owes is that it takes no part in codegen. + assert(codegenUnions(fused).isEmpty, + s"the denylist has to keep this union out of codegen:\n${fused.queryExecution}") } } @@ -511,11 +546,12 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val cached = rangeDF(100).cache() try { cached.count() - val df = cached.union(rangeDF(100, 200)) - val flagOn = df.collect().toSet - withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - assert(df.collect().toSet == flagOn) - } + val (fused, _) = assertFlagParity(() => cached.union(rangeDF(100, 200))) + assert(codegenUnions(fused).size == 1, s"expected a fused union:\n${fused.queryExecution}") + // The cached child is the point of the case: without it this is a plain range union. + assert(collect(fused.queryExecution.executedPlan) { + case s: InMemoryTableScanLike => s + }.nonEmpty, s"expected a cached child:\n${fused.queryExecution}") } finally { cached.unpersist() } @@ -569,7 +605,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val dfs = (0 until n).map(i => rangeDF(i.toLong, i.toLong + 1L)) val unioned = dfs.reduce((x, y) => x.union(y)) assert(unioned.count() == n.toLong) - assert(!unionInsideWSCG(unioned)) + assert(codegenUnions(unioned).isEmpty) } } @@ -596,7 +632,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val a = rangeDF(10).select(col("id"), rand(42).as("r")) val b = rangeDF(10).select(col("id"), rand(43).as("r")) val df = a.union(b) - assert(!unionInsideWSCG(df), + assert(codegenUnions(df).isEmpty, "Union with Nondeterministic child must not be inside WSCG") // Verify correctness despite fallback assertFlagParity(() => a.union(b).orderBy("id")) @@ -606,19 +642,19 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val a = rangeDF(10).select(col("id"), monotonically_increasing_id().as("mid")) val b = rangeDF(10).select(col("id"), monotonically_increasing_id().as("mid")) val df = a.union(b) - assert(!unionInsideWSCG(df), + assert(codegenUnions(df).isEmpty, "Union with monotonically_increasing_id child must not be inside WSCG") } test("SPARK-56482: column pruning works under union codegen (usedInputs=empty)") { // Union of 2-column children, parent selects only 1 column - val a = rangeDF(10).select(col("id"), (col("id") * 2).as("v")) - val b = rangeDF(10, 20).select(col("id"), (col("id") * 3).as("v")) - val df = a.union(b).select("id").orderBy("id") - val flagOn = df.collect() - withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - checkAnswer(df, flagOn.toSeq) + val (fused, _) = assertFlagParity { () => + val a = rangeDF(10).select(col("id"), (col("id") * 2).as("v")) + val b = rangeDF(10, 20).select(col("id"), (col("id") * 3).as("v")) + a.union(b).select("id").orderBy("id") } + assert(codegenUnions(fused).size == 1, + s"pruning is what this case is about, so the union has to be fused:\n${fused.queryExecution}") } test("SPARK-56482: numOutputRows with empty union children") { @@ -639,16 +675,15 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper test("SPARK-56482: partitioning-aware union falls back to non-codegen") { // After repartition, both children expose a `HashPartitioning` on the same key, // so `UnionExec.outputPartitioning` is non-Unknown and the codegen path is denied. - // AQE is disabled here so the executedPlan exposes the UnionExec directly - // (under AQE the plan is wrapped in `AdaptiveSparkPlanExec`, which does not - // surface its inputPlan via `children`). + // AQE is disabled here for the `collectFirst` below: `SparkPlan.collect` stops at + // `AdaptiveSparkPlanExec`, and before execution there is no final plan to reach anyway. withSQLConf( SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val a = rangeDF(100).repartition(4, col("id")) val b = rangeDF(100, 200).repartition(4, col("id")) val df = a.union(b) - assert(!unionInsideWSCG(df), + assert(codegenUnions(df).isEmpty, "Partitioning-aware union must not fuse into WSCG") val unionExec = df.queryExecution.executedPlan.collectFirst { case u: UnionExec => u @@ -761,9 +796,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // that do support codegen can still produce one, since `insertInputAdapter` recurses into // their descendants; exchanges just make it certain. val df = rangeDF(100).repartition(2).union(rangeDF(100).repartition(2)) - // `fusedUnions` requires the union to be the stage root; `unionInsideWSCG` would also - // match a union that an `InputAdapter` left inside the stage unfused, which is exactly - // the degradation this guard has to catch. + // `fusedUnions` requires the union to be the stage root, which is what this test needs: it + // reaches for that node itself below. assert(fusedUnions(df).size == 1, "this shape must fuse, or the test exercises nothing") df } @@ -970,13 +1004,15 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { - // The pipeline has to keep the two on either side of `EnsureRequirements`; the AQE list is - // private, so only the standard one can be checked from here. + // The three have to stay contiguous, which nothing at the list itself says: an injected rule + // cannot land between them, but an edit to the list can. AQE builds its own list, and + // `AdaptiveQueryExecSuite` asserts the same order there. val rules = QueryExecution.preparations(spark, subquery = false) val snapshot = rules.indexWhere(_ eq SnapshotUnionOutputPartitioningConf) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) - assert(snapshot >= 0 && snapshot < ensureRequirements, - s"expected the conf snapshot before EnsureRequirements, got $snapshot/$ensureRequirements") + assert(snapshot >= 0 && snapshot == ensureRequirements - 1, + s"expected the conf snapshot right before EnsureRequirements at $ensureRequirements, " + + s"got $snapshot") 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"))) @@ -1008,7 +1044,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val a = spark.read.parquet(path).select(col("id"), input_file_name().as("f")) val b = spark.read.parquet(path).select(col("id"), input_file_name().as("f")) val df = a.union(b).filter(col("id") > 0) - assert(unionInsideWSCG(df), + assert(codegenUnions(df).nonEmpty, "Union with input_file_name child should fuse into WSCG") assertFlagParity(() => a.union(b).orderBy("id", "f")) } @@ -1037,7 +1073,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // emission window" requirement. Generating the same fused stage from many // threads reproduces the race. val df = rangeDF(100).union(rangeDF(100)).filter(col("id") > 0) - assert(unionInsideWSCG(df)) + assert(codegenUnions(df).nonEmpty) val wscg = df.queryExecution.executedPlan.collectFirst { case w: WholeStageCodegenExec if w.find(_.isInstanceOf[UnionExec]).isDefined => w }.getOrElse(fail("expected a fused UnionExec stage")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala index 3fad780013fac..66f439d303287 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -46,7 +46,7 @@ import org.apache.spark.sql.execution.columnar.{InMemoryTableScanExec, InMemoryT import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.noop.NoopDataSource import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ENSURE_REQUIREMENTS, Exchange, REPARTITION_BY_COL, REPARTITION_BY_NUM, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ENSURE_REQUIREMENTS, EnsureRequirements, Exchange, REPARTITION_BY_COL, REPARTITION_BY_NUM, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, HashedRelationBroadcastMode, ShuffledHashJoinExec, ShuffledJoin, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.SQLShuffleReadMetricsReporter import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryWrapper} @@ -5324,6 +5324,45 @@ class AdaptiveQueryExecSuite } } + test("SPARK-59122: query stage preparation keeps the union barriers around EnsureRequirements") { + // `EnsureRequirements` asks a `UnionExec` what it reports, and `StampUnionDecisions` freezes + // that answer so every rule below it and the execution read what the exchanges were planned + // against. `SnapshotUnionOutputPartitioningConf` has to run first, or the value the stamp reads + // is whatever `conf` says by then rather than the one `EnsureRequirements` saw. The two sit + // next to `EnsureRequirements` with nothing in between, and nothing at the list itself says the + // three have to stay contiguous: an injected rule cannot land between them, since those are + // appended at the tail, but an edit to the list can. AQE builds its own list, and + // `UnionCodegenSuite` covers the one `QueryExecution` builds when AQE is off. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + // The repartition exchanges are what bring AQE in, with the aggregate's distribution + // requirement as a second reason: `InsertAdaptiveSparkPlan` tests for either and leaves a + // plain union of ranges alone. + 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() + val aqe = df.queryExecution.executedPlan.collectFirst { + case a: AdaptiveSparkPlanExec => a + } + assert(aqe.isDefined, s"expected an AdaptiveSparkPlanExec:\n${df.queryExecution}") + val rules = aqe.get.queryStagePreparationRules + val snapshot = rules.indexWhere(_ eq SnapshotUnionOutputPartitioningConf) + val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) + // Both barriers, not the first one: the list ends with a second `StampUnionDecisions` for a + // union an injected prep rule created, and asking only for the first index would let that one + // stand in for the barrier behind `EnsureRequirements`. + val stamps = rules.zipWithIndex.collect { + case (rule, i) if rule eq StampUnionDecisions => i + } + assert(snapshot >= 0 && snapshot == ensureRequirements - 1, + s"expected the conf snapshot right before EnsureRequirements at $ensureRequirements, " + + s"got $snapshot") + assert(stamps.headOption.contains(ensureRequirements + 1), + s"expected the stamp right behind EnsureRequirements at $ensureRequirements, got $stamps") + assert(stamps.last == rules.length - 1, + s"expected the trailing stamp last of ${rules.length} rules, got $stamps") + } + } + test("SPARK-44065: Optimize BroadcastHashJoin skew") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", From c15becad1fd1acd7642764bc497605da1cf6ac36 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 17 Sep 2026 22:13:39 +0800 Subject: [PATCH 17/24] Share one union conf snapshot across a preparation's barriers Answers review 5235414836: the barriers in one preparation take one UnionConfSnapshot, the classic rule-order case asserts adjacency, and outputPartitioning derives rawPartitioning at most once per call. Also pins the cached-union metric by value and cuts comment clauses that were repeated across the file. --- .../spark/sql/execution/QueryExecution.scala | 9 +- .../sql/execution/StampUnionDecisions.scala | 36 +++++-- .../adaptive/AdaptiveSparkPlanExec.scala | 13 ++- .../execution/basicPhysicalOperators.scala | 72 ++++++++------ .../sql/execution/UnionCodegenSuite.scala | 95 +++++++++++++++---- .../adaptive/AdaptiveQueryExecSuite.scala | 4 +- 6 files changed, 163 insertions(+), 66 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index ce1104ed15117..fe0220971b993 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -818,6 +818,9 @@ object QueryExecution { sparkSession: SparkSession, adaptiveExecutionRule: Option[InsertAdaptiveSparkPlan] = None, subquery: Boolean): Seq[Rule[SparkPlan]] = { + // Read once here so that both union barriers below, and the codegen gate they stamp for, answer + // from the same values however long preparation takes. + val unionConf = UnionConfSnapshot(sparkSession.sessionState.conf) // `AdaptiveSparkPlanExec` is a leaf node. If inserted, all the following rules will be no-op // as the original plan is hidden behind `AdaptiveSparkPlanExec`. adaptiveExecutionRule.toSeq ++ @@ -829,11 +832,11 @@ object QueryExecution { // Must run before `EnsureRequirements`, which asks a `UnionExec` what it reports: it // records the conf that answer depends on, so the following `StampUnionDecisions` freezes the // decision under the same value the exchanges were planned against. - SnapshotUnionOutputPartitioningConf, + new SnapshotUnionOutputPartitioningConf(unionConf), EnsureRequirements(), // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, and // the answer to fix is the one the exchanges around it were planned against. - StampUnionDecisions, + new StampUnionDecisions(unionConf), // This rule must be run after `EnsureRequirements`. InsertSortForLimitAndOffset, // `PushDownLocalSort` pushes a wider local sort down onto a narrower one below it, so a @@ -862,7 +865,7 @@ object QueryExecution { // A barrier for a `UnionExec` an injected columnar rule just created, which has no decision // yet and would otherwise take one wherever it is first asked. A decision already stamped on // a node is kept. - StampUnionDecisions, + new StampUnionDecisions(unionConf), CollapseCodegenStages()) ++ (if (subquery) { Nil diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index 2faa02c4492b3..ae8570ab2d82c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -51,12 +51,10 @@ import org.apache.spark.sql.internal.SQLConf * exchanges; a round whose plan loses on cost is discarded whole, stamps included. A union inside a * stage is not revisited, since `foreach` stops at `QueryStageExec`. */ -object StampUnionDecisions extends Rule[SparkPlan] { +class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { - val codegenEnabled = plan.conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED) - val maxChildren = plan.conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN) plan.foreach { - case u: UnionExec => u.stampDecisions(codegenEnabled, maxChildren) + case u: UnionExec => u.stampDecisions(snapshot) case _ => } plan @@ -77,16 +75,36 @@ object StampUnionDecisions extends Rule[SparkPlan] { * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only * become co-partitioned there, which is why the decision itself waits for the barrier behind it. * - * One read per plan, so every union in it answers from the same value. Writing the tag in place is - * safe for the reason given on [[StampUnionDecisions]]. + * The value comes from the preparation's [[UnionConfSnapshot]], so every union in it answers from + * the same read. Writing the tag in place is safe for the reason given on [[StampUnionDecisions]]. */ -object SnapshotUnionOutputPartitioningConf extends Rule[SparkPlan] { +class SnapshotUnionOutputPartitioningConf(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { - val enabled = plan.conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) plan.foreach { - case u: UnionExec => u.snapshotOutputPartitioningConf(enabled) + case u: UnionExec => u.snapshotOutputPartitioningConf(snapshot.outputPartitioning) case _ => } plan } } + +/** + * The union confs one preparation answers from, read once and shared by every barrier in it. + * + * A barrier that read the live conf instead would let two of them disagree: an injected rule can + * return an equivalent `UnionExec` carrying tags it set itself, and `copyTagsFrom` adds nothing to + * a node that already has one, so such a replacement reaches the late barrier with no record of its + * own and would be stamped from whatever the conf says by then, rather than from what the exchanges + * above it were planned against. + */ +case class UnionConfSnapshot( + outputPartitioning: Boolean, + codegenEnabled: Boolean, + maxChildren: Int) + +object UnionConfSnapshot { + def apply(conf: SQLConf): UnionConfSnapshot = UnionConfSnapshot( + outputPartitioning = conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING), + codegenEnabled = conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED), + maxChildren = conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index a34f75e267b87..2c7b66e83e0c0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -119,6 +119,11 @@ case class AdaptiveSparkPlanExec( conf.costEvaluatorCountLocalSortEnabled) } + // Read once for this execution so that the union barriers in the lists below, which run per + // re-planning round and per stage created, cannot answer from different values. Taken from the + // same session conf this query's `QueryExecution.preparations` reads. + @transient private val unionConf = UnionConfSnapshot(context.session.sessionState.conf) + // A list of physical plan rules to be applied before creation of query stages. The physical // plan should reach a final status of query stages (i.e., no more addition or removal of // Exchange nodes) after running these rules. @@ -138,12 +143,12 @@ case class AdaptiveSparkPlanExec( // Must run before `ensureRequirements`, which asks a `UnionExec` what it reports: it // records the conf that answer depends on, so the following `StampUnionDecisions` freezes the // decision under the same value the exchanges were planned against. - SnapshotUnionOutputPartitioningConf, + new SnapshotUnionOutputPartitioningConf(unionConf), ensureRequirements, // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, so // every rule below and the execution itself read the answer the exchanges above it were // planned against. - StampUnionDecisions, + new StampUnionDecisions(unionConf), // This rule must be run after `EnsureRequirements`. InsertSortForLimitAndOffset, AdjustShuffleExchangePosition, @@ -175,7 +180,7 @@ case class AdaptiveSparkPlanExec( ) ++ context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules :+ // A barrier for a `UnionExec` an injected prep rule just created. Decisions already stamped // above are kept. - StampUnionDecisions + new StampUnionDecisions(unionConf) } // A list of physical optimizer rules to be applied to a new stage before its execution. These @@ -203,7 +208,7 @@ case class AdaptiveSparkPlanExec( // A barrier for a `UnionExec` an injected stage-optimizer or columnar rule just created, which // has no decision yet and would otherwise take one wherever it is first asked. A decision // already stamped on a node is kept, so this pass cannot move one. - StampUnionDecisions, + new StampUnionDecisions(unionConf), collapseCodegenStagesRule ) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 2c77000d8e5dc..1b3d32125dfb9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1059,17 +1059,26 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * * `UNION_OUTPUT_PARTITIONING` is taken from `snapshotOutputPartitioningConf`, recorded before * `EnsureRequirements`, so the value the exchanges are planned against is the value execution - * uses; a node created after that pass carries no record and reads the live conf. Reading it live - * here would leave one rule between the two: `conf` is live, and another thread setting it in - * that window would let a parent drop an exchange over a concrete partitioning and then have the - * stamp freeze plain concatenation under it. + * uses. Reading it live here would leave one rule between the two: `conf` is live, and another + * thread setting it in that window would let a parent drop an exchange over a concrete + * partitioning and then have the stamp freeze plain concatenation under it. A node created after + * that pass carries no record of its own and takes the same preparation's value when it is + * stamped; only a union no barrier reached at all falls back to the live conf. * * A read before `StampUnionDecisions` answers from the children as they are then, and does not * write, so observing an unprepared plan cannot decide anything for the prepared one. */ - private[execution] def isPlainUnion: Boolean = stampedDecisions.map(_.plainUnion).getOrElse { - !outputPartitioningEnabled || rawPartitioning.isInstanceOf[UnknownPartitioning] - } + private[execution] def isPlainUnion: Boolean = isPlainUnion(rawPartitioning) + + /** + * The same decision, for a caller that has already derived the raw partitioning: deriving it + * builds an `AttributeMap` per child and intersects the candidates, and `outputPartitioning` + * needs the value the answer was derived from. By name so that a stamped node never derives it. + */ + private def isPlainUnion(raw: => Partitioning): Boolean = + stampedDecisions.map(_.plainUnion).getOrElse { + !outputPartitioningEnabled || raw.isInstanceOf[UnknownPartitioning] + } private def stampedDecisions: Option[UnionExec.Decisions] = getTagValue(UnionExec.DECISIONS) @@ -1079,11 +1088,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup .getOrElse(conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) /** - * Records the conf `isPlainUnion` answers from, read once for the whole plan by - * `SnapshotUnionOutputPartitioningConf` and passed in here, ahead of `EnsureRequirements`, whose - * reads the following stamp has to agree with. Only the conf, never a partitioning: the exchanges - * `EnsureRequirements` adds are not there yet, so a decision taken here would freeze plain on a - * union whose children only become co-partitioned there. + * Records the conf `isPlainUnion` answers from, read once per preparation into a + * `UnionConfSnapshot` and passed in here. `SnapshotUnionOutputPartitioningConf` does it ahead of + * `EnsureRequirements`, whose reads the following stamp has to agree with, and `stampDecisions` + * does it for a node that pass never saw. Only the conf, never a partitioning: the exchanges + * `EnsureRequirements` adds are not there yet, so a decision taken there would freeze plain on a + * union whose children only become co-partitioned in it. */ private[execution] def snapshotOutputPartitioningConf(enabled: Boolean): Unit = if (getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF).isEmpty) { @@ -1093,18 +1103,23 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup /** * Fixes this node's decisions for the rest of the plan's life. Called by `StampUnionDecisions`, * first right after `EnsureRequirements`, so what the exchanges around this union were planned - * against is what execution uses; the two confs come from one read per plan there. Nothing else + * against is what execution uses; the confs come from one read per preparation. Nothing else * writes this tag on an existing node, and the nodes the rule writes are freshly planned and not * yet published, so no reader can be looking at one; `metrics` and the codegen gate read it * later, and a node that already carries it keeps it, which is how the copy in the codegen shell * stays in step with the gate. */ - private[execution] def stampDecisions(codegenEnabled: Boolean, maxChildren: Int): Unit = + private[execution] def stampDecisions(snapshot: UnionConfSnapshot): Unit = if (stampedDecisions.isEmpty) { + // A node the pass ahead of `EnsureRequirements` never saw takes that pass's value now rather + // than the live conf, so a late barrier cannot decide from a conf changed since. Reachable + // for a union an injected rule added, and for one it rebuilt carrying tags of its own, which + // is enough to stop `copyTagsFrom` from bringing this tag across. + snapshotOutputPartitioningConf(snapshot.outputPartitioning) setTagValue(UnionExec.DECISIONS, UnionExec.Decisions( plainUnion = isPlainUnion, - unionCodegenEnabled = codegenEnabled, - maxChildren = maxChildren)) + unionCodegenEnabled = snapshot.codegenEnabled, + maxChildren = snapshot.maxChildren)) } /** @@ -1128,8 +1143,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * siblings do not mirror can empty the intersection. AQE reconciles that, by validating a * partitioning change against the parents' requirements; an injected rule can skip it. */ - override def outputPartitioning: Partitioning = - if (isPlainUnion) super.outputPartitioning else rawPartitioning + override def outputPartitioning: Partitioning = { + // Derived at most once per call, and only when the decision is not stamped yet: `isPlainUnion` + // needs the same value to answer. Not held across calls -- AQE changes the children's partition + // counts, and the answer has to follow them. + lazy val raw = rawPartitioning + if (isPlainUnion(raw)) super.outputPartitioning else raw + } // Per-child projection from the child's output to the union's output. The wrapped // child is always the source `Attribute` (deterministic by construction); the Alias @@ -1150,10 +1170,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // The confs the gate reads, stamped for the reason the plain-union decision is: `conf` is live, // so the gate, `metrics` and the copy `insertInputAdapter` puts inside the codegen shell would // otherwise be free to read different values. When a child is not `CodegenSupport` that copy is - // real and its first evaluation lands at execution; reading the conf there left `metrics` empty - // while `doProduce` asked `metricTerm` for `numOutputRows`. A read before the stamp answers from - // the conf as it is then and writes nothing, so observing an unprepared plan cannot pin this - // either. + // real and its first evaluation lands at execution, which is where reading the conf produced the + // failure described on `isPlainUnion`. private def unionCodegenEnabled: Boolean = stampedDecisions.map(_.unionCodegenEnabled) .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) @@ -1406,17 +1424,17 @@ object UnionExec { * * `withNewChildren` copies the tag onto a rebuilt node, and so does a transform rule's * replacement, but only where the target carries no tags of its own: `copyTagsFrom` leaves a node - * that already has some untouched. A `UnionExec` reaching execution unstamped therefore answers - * from the state it sees then, and can leave `metrics` empty, so `doProduce` fails asking - * `metricTerm` for `numOutputRows`. + * that already has some untouched. Such a node is stamped when a barrier next reaches it, and one + * that reaches execution unstamped answers from the state it sees then. */ private val DECISIONS = TreeNodeTag[Decisions]("unionDecisions") /** * The `UNION_OUTPUT_PARTITIONING` value `isPlainUnion` answers from until the decision is * stamped. See `snapshotOutputPartitioningConf`. Written before `EnsureRequirements` and read by - * the stamp after it, so both phases use one value; travels onto rebuilt nodes the same way - * `DECISIONS` does, which is what carries it across the copies `EnsureRequirements` makes. + * the stamp after it, so both phases use one value, and written by the stamp itself for a node + * that pass never saw; travels onto rebuilt nodes the same way `DECISIONS` does, which is what + * carries it across the copies `EnsureRequirements` makes. */ private val OUTPUT_PARTITIONING_CONF = TreeNodeTag[Boolean]("unionOutputPartitioningConf") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index 20fb061dfccf9..fb3f23df9a898 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -22,6 +22,7 @@ import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, UnknownPartitioning} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.columnar.InMemoryTableScanLike import org.apache.spark.sql.execution.exchange.{EnsureRequirements, REPARTITION_BY_NUM, ShuffleExchangeExec} @@ -107,6 +108,17 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper spark.catalog.cacheTable(view) } + /** + * Drive a union barrier the way a preparation would: the rules take one `UnionConfSnapshot` for + * the whole pipeline, so a test driving them by hand samples it where the rule it is standing in + * for runs, which is what the surrounding `withSQLConf` decides. + */ + private def stampUnionDecisions(plan: SparkPlan): SparkPlan = + new StampUnionDecisions(UnionConfSnapshot(SQLConf.get))(plan) + + private def snapshotUnionOutputPartitioningConf(plan: SparkPlan): SparkPlan = + new SnapshotUnionOutputPartitioningConf(UnionConfSnapshot(SQLConf.get))(plan) + /** * Run `buildDf()` with union codegen on, then again with it off, and assert the two agree. * @@ -512,8 +524,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val smj = left.join(right, "k") smj.union(rangeDF(100).select(col("id").as("k"))) } - // Not `fusedUnions`: this union would not root a stage anyway, so that would hold with or - // without the denylist. What the denylist owes is that it takes no part in codegen. + // `codegenUnions`, not `fusedUnions`: what the denylist owes is that this union takes no part + // in codegen at any depth, and without it the union would root a stage here. assert(codegenUnions(fused).isEmpty, s"the denylist has to keep this union out of codegen:\n${fused.queryExecution}") } @@ -703,11 +715,13 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // it false, and without one `supportCodegenFailureReason` reports `columnar` and nothing // fuses. `SELECT *` or a plain alias collapses the projection away and does not reproduce // this. Once the cache stages finalise, both children report the same concrete layout, and - // re-deriving the decision at that point left `metrics` empty while `doProduce` asked - // `metricTerm` for `numOutputRows`. + // re-deriving the decision at that point answered "partitioning-aware", so `metrics` came back + // empty while `doProduce` asked `metricTerm` for `numOutputRows`. // - // Both halves of the decision are asserted here. Registering the metric unconditionally would - // fix the crash and leave the other half broken: a fused union concatenates its children's + // Executing the query is what proves that crash is gone: an empty `metrics` throws at + // `metricTerm` while `doProduce` runs. The metric's value then says the fused code ran and + // counted, and the partitioning assertion is the other half, which registering the metric + // unconditionally would have left broken: a fused union concatenates its children's // partitions, so claiming their partitioning would let a parent satisfy a clustered // distribution from an RDD that does not have it. withSQLConf( @@ -734,8 +748,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper s"premise: got $childPartitionings") assert(childPartitionings.map(_.numPartitions).distinct.size == 1, s"premise: got $childPartitionings") - assert(u.metrics.contains("numOutputRows"), - "a fused union must register the metric its generated code increments") + assert(u.metrics("numOutputRows").value == 20, + "a fused union must count the rows its generated code emitted") assert(u.outputPartitioning.isInstanceOf[UnknownPartitioning], s"a fused union must not claim a concrete partitioning, got ${u.outputPartitioning}") } @@ -944,19 +958,20 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // from the outside by the extension-driven cases in `SparkSessionExtensionSuite`, which need a // session of their own. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { - // Pins the property the standard pipeline has to keep: a stamping pass runs after - // `EnsureRequirements`, so the decision is taken from the plan the exchanges were placed in. - // A count would break on a sixth legitimate pass and say nothing about the order. The AQE - // lists are private to `AdaptiveSparkPlanExec`, so their first pass has no counterpart here. + // Pins the property the standard pipeline has to keep: a stamping pass runs immediately after + // `EnsureRequirements`, so the decision is taken from the plan the exchanges were placed in + // and no rule in between can change a partitioning the parent already planned against. A + // count would break on a sixth legitimate pass and say nothing about the order. + // `AdaptiveQueryExecSuite` asserts the same adjacency for the list AQE builds. val rules = QueryExecution.preparations(spark, subquery = false) - val firstStamp = rules.indexWhere(_ eq StampUnionDecisions) + val firstStamp = rules.indexWhere(_.isInstanceOf[StampUnionDecisions]) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) val columnarRules = rules.indexWhere(_.isInstanceOf[ApplyColumnarRulesAndInsertTransitions]) - assert(ensureRequirements >= 0 && firstStamp > ensureRequirements && + assert(ensureRequirements >= 0 && firstStamp == ensureRequirements + 1 && firstStamp < columnarRules, - "expected a stamping pass between EnsureRequirements and the columnar rules, got " + - s"$ensureRequirements/$firstStamp/$columnarRules") + "expected a stamping pass right behind EnsureRequirements and before the columnar " + + s"rules, got $ensureRequirements/$firstStamp/$columnarRules") val stamped = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { val df = spark.range(0, 20, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k")) @@ -969,7 +984,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // A fresh node standing in for one an extension made after the first pass: no decision yet. val fresh = UnionExec(stamped.children) withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { - StampUnionDecisions(fresh) + stampUnionDecisions(fresh) } // Read back with the conf the other way round, so the answer can only come from the stamp: // deriving here would make it non-plain, these children being co-partitioned. @@ -988,7 +1003,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper union.head } withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false") { - StampUnionDecisions(fused) + stampUnionDecisions(fused) assert(fused.supportCodegen, "a second pass must not restamp the conf it was decided with") } } @@ -1008,7 +1023,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // cannot land between them, but an edit to the list can. AQE builds its own list, and // `AdaptiveQueryExecSuite` asserts the same order there. val rules = QueryExecution.preparations(spark, subquery = false) - val snapshot = rules.indexWhere(_ eq SnapshotUnionOutputPartitioningConf) + val snapshot = rules.indexWhere(_.isInstanceOf[SnapshotUnionOutputPartitioningConf]) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) assert(snapshot >= 0 && snapshot == ensureRequirements - 1, s"expected the conf snapshot right before EnsureRequirements at $ensureRequirements, " + @@ -1018,7 +1033,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper .union(spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k"))) .groupBy("k").count() val required = EnsureRequirements()( - SnapshotUnionOutputPartitioningConf(df.queryExecution.sparkPlan.clone())) + snapshotUnionOutputPartitioningConf(df.queryExecution.sparkPlan.clone())) assert(required.collect { case s: ShuffleExchangeExec => s.shuffleOrigin } == Seq(REPARTITION_BY_NUM, REPARTITION_BY_NUM), "the aggregate's exchange must have been elided, or the window has nothing at stake") @@ -1026,7 +1041,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val union = required.collect { case u: UnionExec => u } assert(union.size == 1) withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { - StampUnionDecisions(required) + stampUnionDecisions(required) assert(!union.head.outputPartitioning.isInstanceOf[UnknownPartitioning], "the answer must come from the conf snapshot, not from the value read now, got " + s"${union.head.outputPartitioning}") @@ -1034,6 +1049,44 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a late barrier stamps a replacement from the preparation's conf") { + // An injected rule can return a `UnionExec` of its own in place of a stamped one, and + // `copyTagsFrom` gives nothing to a node already carrying a tag, so such a replacement reaches + // the barrier behind the extension hooks with no record of the value the exchanges above it + // were planned against. Every barrier in one preparation answers from the same + // `UnionConfSnapshot`, so that is still the value the replacement is stamped with. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + 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() + val prepared = stampUnionDecisions(EnsureRequirements()( + snapshotUnionOutputPartitioningConf(df.queryExecution.sparkPlan.clone()))) + val union = prepared.collect { case u: UnionExec => u } + assert(union.size == 1) + // The value a preparation would carry to its late barrier, taken while the conf still says + // what `EnsureRequirements` read above. + val prepConf = UnionConfSnapshot(SQLConf.get) + + // What `transformUp` leaves behind for a rule that replaced the node: `copyTagsFrom` is + // all-or-nothing, so one unrelated tag of its own costs the replacement both of ours. + val replacement = UnionExec(union.head.children) + replacement.setTagValue(TreeNodeTag[Unit]("SPARK-59122-injected"), ()) + replacement.copyTagsFrom(union.head) + + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + assert(replacement.isPlainUnion, + "the replacement must reach the barrier with no decision of its own, or this case " + + "exercises nothing") + new StampUnionDecisions(prepConf)(replacement) + assert(!replacement.outputPartitioning.isInstanceOf[UnknownPartitioning], + "the late barrier must stamp from the preparation's conf, not the value read now, got " + + s"${replacement.outputPartitioning}") + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala index 66f439d303287..4a97a5f962fe7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -5345,13 +5345,13 @@ class AdaptiveQueryExecSuite } assert(aqe.isDefined, s"expected an AdaptiveSparkPlanExec:\n${df.queryExecution}") val rules = aqe.get.queryStagePreparationRules - val snapshot = rules.indexWhere(_ eq SnapshotUnionOutputPartitioningConf) + val snapshot = rules.indexWhere(_.isInstanceOf[SnapshotUnionOutputPartitioningConf]) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) // Both barriers, not the first one: the list ends with a second `StampUnionDecisions` for a // union an injected prep rule created, and asking only for the first index would let that one // stand in for the barrier behind `EnsureRequirements`. val stamps = rules.zipWithIndex.collect { - case (rule, i) if rule eq StampUnionDecisions => i + case (rule, i) if rule.isInstanceOf[StampUnionDecisions] => i } assert(snapshot >= 0 && snapshot == ensureRequirements - 1, s"expected the conf snapshot right before EnsureRequirements at $ensureRequirements, " + From 1955fade5c125aca6e2de40e88af86a669118ea0 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 21 Sep 2026 16:50:22 +0800 Subject: [PATCH 18/24] Keep the codegen gate provisional until a union is stamped --- .../apache/spark/sql/internal/SQLConf.scala | 8 +- .../execution/basicPhysicalOperators.scala | 65 ++++++----- .../sql/execution/UnionCodegenSuite.scala | 101 +++++++++++++++++- 3 files changed, 145 insertions(+), 29 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 781a7594c472a..f9be03cc1e1b9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2935,10 +2935,10 @@ object SQLConf { val WHOLESTAGE_UNION_CODEGEN_ENABLED = buildConf("spark.sql.codegen.wholeStage.union.enabled") .internal() - .doc("When both this conf and `spark.sql.codegen.wholeStage` are true, " + - "UnionExec participates in whole-stage codegen on its " + - "non-partitioning-aware path: the parent and all children fuse into " + - "a single WholeStageCodegenExec stage. The value is read when a UnionExec's " + + .doc("When both this conf and `spark.sql.codegen.wholeStage` are true, an eligible " + + "UnionExec on its non-partitioning-aware path takes part in whole-stage codegen. " + + "The union's other eligibility checks still apply, and a child that does not support " + + "codegen still ends the stage at an InputAdapter. The value is read when a UnionExec's " + "decision is fixed during physical preparation, so a change does not reach a " + "decision already taken.") .version("4.2.0") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 1b3d32125dfb9..e67ede0694a44 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1180,27 +1180,20 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup stampedDecisions.map(_.maxChildren) .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) - // Memoized per instance rather than stamped on the tag. Every term below the confs except - // `isPlainUnion` reads the children, and a tag outlives them: `SQLExecution` builds the initial - // `SparkPlanInfo` before execution, forcing `metrics` on every node it visits, so a rule that - // replaces a child after that would inherit an allowing answer and fuse a topology that - // `hasPartitionIndexDependentCodegen` or `supportsColumnar` rejects. The copy in the codegen - // shell still agrees with the gate: `InputAdapter` delegates `output` and `supportsColumnar` to - // its child, the other terms walk the subtree through it, and each of those is fixed for a given - // set of children. `isPlainUnion` is not, which is why it is stamped instead. - @transient private lazy val supportCodegenFailureReason: Option[String] = { - if (!unionCodegenEnabled) { - Some("union-codegen-disabled") - } else if (!isPlainUnion) { - Some("partitioning-aware") - } else if (children.exists(_.exists(_.isInstanceOf[UnionExec]))) { + // Memoized per instance rather than stamped on the tag. Every term here reads the children, and a + // tag outlives them: `SQLExecution` builds the initial `SparkPlanInfo` before execution, forcing + // `metrics` on every node it visits, so a rule that replaces a child after that would inherit an + // allowing answer and fuse a topology that `hasPartitionIndexDependentCodegen` or + // `supportsColumnar` rejects. The copy in the codegen shell still agrees with the gate: + // `InputAdapter` delegates `output` and `supportsColumnar` to its child, the other terms walk the + // subtree through it, and each of those is fixed for a given set of children. + @transient private lazy val childTopologyFailureReason: Option[String] = { + if (children.exists(_.exists(_.isInstanceOf[UnionExec]))) { Some("nested-union") } else if (children.exists(_.exists(UnionExec.isKnownMultiInputRDDCodegen))) { Some("multi-rdd-child") } else if (children.exists(UnionExec.hasPartitionIndexDependentCodegen)) { Some("partition-index-dependent-child") - } else if (children.size > maxCodegenChildren) { - Some("max-children-exceeded") } else if (supportsColumnar) { Some("columnar") } else if (children.exists(c => @@ -1211,6 +1204,25 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } + // The three preparation-scoped terms are recomputed per call, because memoizing them would let a + // read arriving before `stampDecisions` settle the gate on the live conf. The stamp can install + // the opposite value, and the gate would then keep the memoized answer while the copy + // `insertInputAdapter` builds, a fresh instance carrying the stamped tag, derives the other one. + // Nothing on Spark's own path reads a union that early: every barrier runs inside `preparations`, + // ahead of the `SparkPlanInfo` that forces `metrics`. A late extension hook can, and so can a + // caller inspecting `sparkPlan`. + private def supportCodegenFailureReason: Option[String] = { + if (!unionCodegenEnabled) { + Some("union-codegen-disabled") + } else if (!isPlainUnion) { + Some("partitioning-aware") + } else if (children.size > maxCodegenChildren) { + Some("max-children-exceeded") + } else { + childTopologyFailureReason + } + } + override def supportCodegen: Boolean = { val reason = supportCodegenFailureReason if (reason.isEmpty) true @@ -1223,11 +1235,15 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Registered only when fusion will actually run, so plans that fall back - // to `doExecute` (which never updates the metric) do not surface a - // 0-valued row count in the SQL UI. `doConsume` is the sole incrementer. + // Registered only when fusion will actually run, so plans that fall back to `doExecute` (which + // never updates the metric) do not surface a 0-valued row count in the SQL UI. `doConsume` is the + // sole incrementer. An unstamped node is the exception: its gate is still provisional, the stamp + // can land either way, and a map built without the metric would leave `doProduce` asking + // `metricTerm` for one that is not there. Registering it then costs an unused metric on a union + // an extension inspected and the stamp went on to reject; Spark's own force, the `SparkPlanInfo` + // `SQLExecution` builds, runs after every barrier, so ordinary fallback unions still omit it. override lazy val metrics: Map[String, SQLMetric] = - if (supportCodegenFailureReason.isEmpty) { + if (stampedDecisions.isEmpty || supportCodegenFailureReason.isEmpty) { Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) } else { Map.empty @@ -1431,10 +1447,11 @@ object UnionExec { /** * The `UNION_OUTPUT_PARTITIONING` value `isPlainUnion` answers from until the decision is - * stamped. See `snapshotOutputPartitioningConf`. Written before `EnsureRequirements` and read by - * the stamp after it, so both phases use one value, and written by the stamp itself for a node - * that pass never saw; travels onto rebuilt nodes the same way `DECISIONS` does, which is what - * carries it across the copies `EnsureRequirements` makes. + * stamped. See `snapshotOutputPartitioningConf`. `SnapshotUnionOutputPartitioningConf` writes the + * value before `EnsureRequirements`, and the stamp after it reads what that rule wrote, so both + * phases use one value; `stampDecisions` writes the value itself for a node that pass never saw. + * The tag travels onto rebuilt nodes the same way `DECISIONS` does, which is what carries it + * across the copies `EnsureRequirements` makes. */ private val OUTPUT_PARTITIONING_CONF = TreeNodeTag[Boolean]("unionOutputPartitioningConf") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index fb3f23df9a898..b835a73883c8b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -23,7 +23,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, UnknownPartitioning} import org.apache.spark.sql.catalyst.trees.TreeNodeTag -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper} import org.apache.spark.sql.execution.columnar.InMemoryTableScanLike import org.apache.spark.sql.execution.exchange.{EnsureRequirements, REPARTITION_BY_NUM, ShuffleExchangeExec} import org.apache.spark.sql.functions._ @@ -923,6 +923,36 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a read before the barrier decides neither the gate nor the metric") { + // An extension hook can build a `UnionExec` and force `supportCodegen` or `metrics` on it + // before the barrier behind that hook stamps it. Memoizing the whole gate there settles it on + // the live conf, and the copy `insertInputAdapter` builds, a fresh instance carrying the + // stamped tag, derives the other answer: that disagreement is what `doProduce` used to hit as + // `key not found: numOutputRows`. The preparation-scoped terms are recomputed per call instead, + // and an unstamped node registers the metric, since either stamped outcome is still open. + // Nothing here is executed, so AQE does not enter into it: the children come from `sparkPlan`, + // and the union under test is built by hand. + val planned = rangeDF(100).union(rangeDF(100)) + .queryExecution.sparkPlan.collect { case u: UnionExec => u } + assert(planned.size == 1) + val kids = planned.head.children + + Seq(true, false).foreach { live => + val union = UnionExec(kids) + withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> live.toString) { + assert(union.supportCodegen == live, + "an unstamped read answers from the live conf, or this case starts from nothing") + assert(union.metrics.contains("numOutputRows"), + "an unstamped union must register the metric: the stamp can still fuse it") + } + withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> (!live).toString) { + stampUnionDecisions(union) + assert(union.supportCodegen == !live, + s"the gate must answer from the stamp, not from the read taken with the conf $live") + } + } + } + test("SPARK-59122: a partitioning-aware union follows its children's coalesced partition count") { // Only the decision is stamped, never the `Partitioning`. AQE coalescing changes the children's // `numPartitions` after the stamp, and `unionRDDs` hands whatever it reports to @@ -1087,6 +1117,75 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: the stamp takes the codegen decisions from the preparation's snapshot") { + // The case above pins this for `UNION_OUTPUT_PARTITIONING`. The two codegen fields ride in the + // same `UnionConfSnapshot` and need the same pin: a stamp that read them live would agree with + // the snapshot everywhere else in this suite, because the other cases flip the conf after + // stamping and so pin the gate rather than the stamp. Nothing here is executed, so AQE does not + // enter into it. + val df = rangeDF(100).union(rangeDF(100)).union(rangeDF(100)) + val planned = df.queryExecution.sparkPlan.collect { case u: UnionExec => u } + assert(planned.size == 1 && planned.head.children.size == 3, + s"expected one union of three children, got ${planned.map(_.children.size)}") + val kids = planned.head.children + + // What a preparation would carry to its barriers: fusion allowed, cap high enough for three. + val prepConf = withSQLConf( + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "3") { + UnionConfSnapshot(SQLConf.get) + } + + Seq( + "enablement" -> (SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false"), + "the child cap" -> (SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "2") + ).foreach { case (what, flipped) => + val union = UnionExec(kids) + withSQLConf(flipped) { + new StampUnionDecisions(prepConf)(union) + assert(union.supportCodegen, s"the stamp must take $what from the snapshot") + } + } + } + + test("SPARK-59122: a union fused under AQE keeps both codegen decisions through execution") { + // The AQE counterpart of the two "changes between planning and execution" cases above. + // `CollapseCodegenStages` runs after each stage is created, so the shell and the copy inside it + // are built while the query runs, with the flipped value below already in effect: a gate that + // reads either codegen conf live there would deny fusion, and the copy would come back with + // empty `metrics`. What this case does not pin is where the stamp took its values, because AQE + // stamps the union while it builds `initialPlan`, which is still inside the outer conf; the + // case above is what pins that. That AQE reads its `UnionConfSnapshot` once, at construction, + // rather than per re-planning round is visible only in the code, since reaching a per-stage + // barrier with an unstamped union takes an injected rule. + Seq( + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "2").foreach { flipped => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + // Exchange children are not `CodegenSupport`, so `insertInputAdapter` wraps them and the + // union inside the shell is a `withNewChildren` copy rather than the instance the gate + // answered on. Three of them, so that a cap of two excludes this union; the cap cannot go + // below two. + val df = rangeDF(100).repartition(2) + .union(rangeDF(100).repartition(2)) + .union(rangeDF(100).repartition(2)) + // Constructing the wrapper is what reads AQE's snapshot, and it happens here, before the + // flipped value below. The stages that hold the union are created during `collect()`. + assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) + withSQLConf(flipped) { + assert(df.collect().length == 300) + val fused = fusedUnions(df) + assert(fused.size == 1, + s"the union must still fuse under AQE, got\n${df.queryExecution.executedPlan}") + // The cap iteration needs three children to say anything: at two, a cap of two admits it. + assert(fused.head.children.size == 3) + assert(fused.head.children.forall(_.isInstanceOf[InputAdapter])) + assert(fused.head.metrics("numOutputRows").value == 300) + } + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's From c38613073a230504ca2ae56d8ccc155fe7dadeed Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 22 Sep 2026 13:28:48 +0800 Subject: [PATCH 19/24] Carry one preparation conf snapshot through the injected prep rules --- .../apache/spark/sql/internal/SQLConf.scala | 17 ++-- .../spark/sql/execution/QueryExecution.scala | 2 +- .../sql/execution/StampUnionDecisions.scala | 43 +++++++--- .../adaptive/AdaptiveSparkPlanExec.scala | 5 +- .../execution/basicPhysicalOperators.scala | 59 ++++++++------ .../sql/SparkSessionExtensionSuite.scala | 76 +++++++++++++++--- .../sql/execution/UnionCodegenSuite.scala | 80 +++++++++++++++++-- .../adaptive/AdaptiveQueryExecSuite.scala | 4 +- 8 files changed, 221 insertions(+), 65 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index f9be03cc1e1b9..e238ebac222bf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2938,9 +2938,9 @@ object SQLConf { .doc("When both this conf and `spark.sql.codegen.wholeStage` are true, an eligible " + "UnionExec on its non-partitioning-aware path takes part in whole-stage codegen. " + "The union's other eligibility checks still apply, and a child that does not support " + - "codegen still ends the stage at an InputAdapter. The value is read when a UnionExec's " + - "decision is fixed during physical preparation, so a change does not reach a " + - "decision already taken.") + "codegen still ends the stage at an InputAdapter. The value is captured once for the " + + "plan a physical preparation prepares, before its rule sequence, and every UnionExec " + + "decision taken in that sequence comes from that capture.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .booleanConf @@ -2955,9 +2955,9 @@ object SQLConf { "bytecode size, constant pool growth, JIT compilation time) rather " + "than the JVM per-method bytecode limit. Unions with more children " + "fall back to per-child codegen stages. Only effective when " + - s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true. The value is read when a " + - "UnionExec's decision is fixed during physical preparation, so a change does not " + - "reach a decision already taken.") + s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true. The value is captured once for the " + + "plan a physical preparation prepares, before its rule sequence, and every UnionExec " + + "decision taken in that sequence comes from that capture.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .intConf @@ -8182,8 +8182,9 @@ object SQLConf { .internal() .doc("When set to true, the output partitioning of UnionExec will be the same as the " + "input partitioning if its children have same partitioning. Otherwise, it will be a " + - "default partitioning. The value is read during physical preparation, so a change does " + - "not reach a decision already taken.") + "default partitioning. The value is captured once for the plan a physical preparation " + + "prepares, before its rule sequence, and every UnionExec decision taken in that sequence " + + "comes from that capture.") .version("4.1.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index fe0220971b993..ebc9acccb3b77 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -832,7 +832,7 @@ object QueryExecution { // Must run before `EnsureRequirements`, which asks a `UnionExec` what it reports: it // records the conf that answer depends on, so the following `StampUnionDecisions` freezes the // decision under the same value the exchanges were planned against. - new SnapshotUnionOutputPartitioningConf(unionConf), + new SnapshotUnionPreparationConf(unionConf), EnsureRequirements(), // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, and // the answer to fix is the one the exchanges around it were planned against. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index ae8570ab2d82c..0c5bc21089be6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -62,32 +62,51 @@ class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { } /** - * Records on each [[UnionExec]] the `UNION_OUTPUT_PARTITIONING` value it answers from, before - * `EnsureRequirements` asks it what it reports. + * Records the preparation's [[UnionConfSnapshot]] on each [[UnionExec]], so that every reader ahead + * of a barrier answers from the same values the barrier will stamp. * - * Without this, the two phases sample the conf separately: `EnsureRequirements` reads what the - * union reports under the value then, and [[StampUnionDecisions]] freezes the decision under the - * value one rule later. `conf` is the live session conf, so another thread turning it off in that - * window would let a parent drop an exchange over a concrete partitioning and then have the union - * concatenate, which puts one group in two partitions. + * Without this, each phase samples the confs separately: `EnsureRequirements` reads what a union + * reports under the value then, and [[StampUnionDecisions]] freezes the decision under the value + * one rule later. `conf` is the live session conf, so another thread turning + * `UNION_OUTPUT_PARTITIONING` off in that window would let a parent drop an exchange over a + * concrete partitioning and then have the union concatenate, which puts one group in two + * partitions. * - * Only the conf is recorded, never a partitioning. `EnsureRequirements` has not inserted the + * The same gap opens between two injected rules: one can return a `UnionExec` of its own, which + * carries no record yet, and a later one can plan requirements over it or ask its codegen gate. + * [[SnapshotUnionPreparationConf.before]] closes that for the injected query-stage preparation + * rules, the list where a consumer can still add or drop an exchange. Two windows stay open, both + * behind a barrier that decides before execution: the injected columnar rules, which share one + * `ApplyColumnarRulesAndInsertTransitions` and so cannot be interleaved from outside it, and the + * injected query-stage optimizer rules, which run on a plan whose exchanges are already fixed. + * + * Only the confs are recorded, never a partitioning. `EnsureRequirements` has not inserted the * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only * become co-partitioned there, which is why the decision itself waits for the barrier behind it. * - * The value comes from the preparation's [[UnionConfSnapshot]], so every union in it answers from - * the same read. Writing the tag in place is safe for the reason given on [[StampUnionDecisions]]. + * Writing the tag in place is safe for the reason given on [[StampUnionDecisions]]. */ -class SnapshotUnionOutputPartitioningConf(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { +class SnapshotUnionPreparationConf(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { plan.foreach { - case u: UnionExec => u.snapshotOutputPartitioningConf(snapshot.outputPartitioning) + case u: UnionExec => u.recordPreparationConf(snapshot) case _ => } plan } } +object SnapshotUnionPreparationConf { + /** + * `rules` with a snapshot pass ahead of each of them, so a `UnionExec` one rule creates carries + * the preparation's confs before the next rule reads them. Ahead of rather than behind, so that a + * union left by the rules listed before `rules` is covered too. `rules` is empty unless an + * extension injected something, so this adds no pass to an ordinary preparation. + */ + def before(snapshot: UnionConfSnapshot, rules: Seq[Rule[SparkPlan]]): Seq[Rule[SparkPlan]] = + rules.flatMap(rule => Seq(new SnapshotUnionPreparationConf(snapshot), rule)) +} + /** * The union confs one preparation answers from, read once and shared by every barrier in it. * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index 2c7b66e83e0c0..b7d42d1fbd869 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -143,7 +143,7 @@ case class AdaptiveSparkPlanExec( // Must run before `ensureRequirements`, which asks a `UnionExec` what it reports: it // records the conf that answer depends on, so the following `StampUnionDecisions` freezes the // decision under the same value the exchanges were planned against. - new SnapshotUnionOutputPartitioningConf(unionConf), + new SnapshotUnionPreparationConf(unionConf), ensureRequirements, // Must run after `EnsureRequirements`: it fixes each `UnionExec`'s partitioning decision, so // every rule below and the execution itself read the answer the exchanges above it were @@ -177,7 +177,8 @@ case class AdaptiveSparkPlanExec( // channel, opt-in). Runs last so skew handling and sort cleanup have settled // before placement is decided. AQEEnablePipelinedShuffle - ) ++ context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules :+ + ) ++ SnapshotUnionPreparationConf.before( + unionConf, context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules) :+ // A barrier for a `UnionExec` an injected prep rule just created. Decisions already stamped // above are kept. new StampUnionDecisions(unionConf) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index e67ede0694a44..80fc26d563162 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -981,7 +981,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * The SPARK-52921 candidate partitioning derived from the children, which `outputPartitioning` * reports when the decision says this union is not a plain concatenation. That decision comes out * plain on either of two grounds: `UNION_OUTPUT_PARTITIONING` being off, which is taken from the - * record `SnapshotUnionOutputPartitioningConf` writes rather than read here, or this coming back + * record `SnapshotUnionPreparationConf` writes rather than read here, or this coming back * `UnknownPartitioning`. Under that conf the candidate can still be concrete while the union * reports unknown. */ @@ -1057,7 +1057,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * for `numOutputRows`. A fresh copy inherits the answer instead, since `withNewChildren` ends in * `copyTagsFrom`. * - * `UNION_OUTPUT_PARTITIONING` is taken from `snapshotOutputPartitioningConf`, recorded before + * `UNION_OUTPUT_PARTITIONING` is taken from `recordPreparationConf`, recorded before * `EnsureRequirements`, so the value the exchanges are planned against is the value execution * uses. Reading it live here would leave one rule between the two: `conf` is live, and another * thread setting it in that window would let a parent drop an exchange over a concrete @@ -1083,21 +1083,25 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup private def stampedDecisions: Option[UnionExec.Decisions] = getTagValue(UnionExec.DECISIONS) + private def preparationConf: Option[UnionConfSnapshot] = + getTagValue(UnionExec.PREPARATION_CONF) + private def outputPartitioningEnabled: Boolean = - getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF) + preparationConf.map(_.outputPartitioning) .getOrElse(conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) /** - * Records the conf `isPlainUnion` answers from, read once per preparation into a - * `UnionConfSnapshot` and passed in here. `SnapshotUnionOutputPartitioningConf` does it ahead of - * `EnsureRequirements`, whose reads the following stamp has to agree with, and `stampDecisions` - * does it for a node that pass never saw. Only the conf, never a partitioning: the exchanges - * `EnsureRequirements` adds are not there yet, so a decision taken there would freeze plain on a - * union whose children only become co-partitioned in it. + * Records the confs this node answers from until its decisions are stamped, read once per + * preparation into a `UnionConfSnapshot` and passed in here. `SnapshotUnionPreparationConf` does + * it ahead of `EnsureRequirements`, whose reads the following stamp has to agree with, and ahead + * of each injected rule, so a union an earlier one created is not read live by a later one; + * `stampDecisions` does it for a node no such pass saw. Only the confs, never a partitioning: the + * exchanges `EnsureRequirements` adds are not there yet, so a decision taken there would freeze + * plain on a union whose children only become co-partitioned in it. */ - private[execution] def snapshotOutputPartitioningConf(enabled: Boolean): Unit = - if (getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF).isEmpty) { - setTagValue(UnionExec.OUTPUT_PARTITIONING_CONF, enabled) + private[execution] def recordPreparationConf(snapshot: UnionConfSnapshot): Unit = + if (preparationConf.isEmpty) { + setTagValue(UnionExec.PREPARATION_CONF, snapshot) } /** @@ -1115,11 +1119,15 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // than the live conf, so a late barrier cannot decide from a conf changed since. Reachable // for a union an injected rule added, and for one it rebuilt carrying tags of its own, which // is enough to stop `copyTagsFrom` from bringing this tag across. - snapshotOutputPartitioningConf(snapshot.outputPartitioning) + recordPreparationConf(snapshot) + // Every field from the record this node now carries, which is `snapshot` unless a pass had + // already recorded one. Reading some fields from the record and others from the argument + // would stamp half from each if the two ever differed. + val recorded = preparationConf.getOrElse(snapshot) setTagValue(UnionExec.DECISIONS, UnionExec.Decisions( plainUnion = isPlainUnion, - unionCodegenEnabled = snapshot.codegenEnabled, - maxChildren = snapshot.maxChildren)) + unionCodegenEnabled = recorded.codegenEnabled, + maxChildren = recorded.maxChildren)) } /** @@ -1171,13 +1179,17 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // so the gate, `metrics` and the copy `insertInputAdapter` puts inside the codegen shell would // otherwise be free to read different values. When a child is not `CodegenSupport` that copy is // real and its first evaluation lands at execution, which is where reading the conf produced the - // failure described on `isPlainUnion`. + // failure described on `isPlainUnion`. Before the stamp they come from the preparation's snapshot + // where a pass recorded one, so a rule reading the gate on a union an earlier rule created gets + // the value that preparation will stamp rather than the live conf. private def unionCodegenEnabled: Boolean = stampedDecisions.map(_.unionCodegenEnabled) + .orElse(preparationConf.map(_.codegenEnabled)) .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) private def maxCodegenChildren: Int = stampedDecisions.map(_.maxChildren) + .orElse(preparationConf.map(_.maxChildren)) .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) // Memoized per instance rather than stamped on the tag. Every term here reads the children, and a @@ -1446,14 +1458,15 @@ object UnionExec { private val DECISIONS = TreeNodeTag[Decisions]("unionDecisions") /** - * The `UNION_OUTPUT_PARTITIONING` value `isPlainUnion` answers from until the decision is - * stamped. See `snapshotOutputPartitioningConf`. `SnapshotUnionOutputPartitioningConf` writes the - * value before `EnsureRequirements`, and the stamp after it reads what that rule wrote, so both - * phases use one value; `stampDecisions` writes the value itself for a node that pass never saw. - * The tag travels onto rebuilt nodes the same way `DECISIONS` does, which is what carries it - * across the copies `EnsureRequirements` makes. + * The [[UnionConfSnapshot]] this node answers from until its decisions are stamped. See + * `recordPreparationConf`. `SnapshotUnionPreparationConf` writes it before `EnsureRequirements`, + * and the stamp after it reads what that rule wrote, so both phases use one value; the same rule + * runs ahead of each injected rule, so a union one of them creates is not read live by a later + * one, and `stampDecisions` writes it itself for a node no such pass saw. The tag travels onto + * rebuilt nodes the same way `DECISIONS` does, which is what carries it across the copies + * `EnsureRequirements` makes. */ - private val OUTPUT_PARTITIONING_CONF = TreeNodeTag[Boolean]("unionOutputPartitioningConf") + private val PREPARATION_CONF = TreeNodeTag[UnionConfSnapshot]("unionPreparationConf") /** * Codegen operators that return more than one RDD from `inputRDDs()`. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index a21ec970f80fe..a6b0a423f2537 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -623,10 +623,10 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { * Prepares a plan whose `UnionExec` was made by an injected rule, then turns * `UNION_OUTPUT_PARTITIONING` off and reads that node again. A union the `StampUnionDecisions` * barrier following the hook reached keeps answering from the decision it was prepared with; one - * no barrier reached has no decision to answer from, so this read derives one from the conf as it - * is now and comes back `UnknownPartitioning`. `UnionCodegenSuite` covers the stamping itself by - * calling the rule directly, so it stays green if one of the post-hook listings is dropped; these - * pin the post-hook stamping in each pipeline. + * that reached no barrier and no snapshot pass has nothing to answer from, so this read derives + * an answer from the conf as it is now and comes back `UnknownPartitioning`. `UnionCodegenSuite` + * covers the stamping itself by calling the rule directly, so it stays green if one of the + * post-hook listings is dropped; these pin the post-hook stamping in each pipeline. */ private def checkInjectedUnionIsStamped( extensions: Seq[SparkSessionExtensionsProvider], aqeEnabled: Boolean): Unit = { @@ -688,6 +688,44 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { create(_.injectColumnar(_ => WrapRootInUnionColumnarRule)), aqeEnabled = true) } + test("SPARK-59122: a union one injected prep rule adds is recorded before the next one reads") { + // The barrier sits after the whole injected list, so between two of them a fresh union used to + // have no record and answered live. A rule planning requirements over it could then elide an + // exchange over a concrete partitioning while the barrier stamped the preparation's value and + // execution concatenated, which puts one group in two partitions. + // `SnapshotUnionPreparationConf.before` lists a snapshot pass ahead of each injected rule. + val seen = ListBuffer.empty[Partitioning] + checkInjectedUnionIsStamped( + create { extensions => + extensions.injectQueryStagePrepRule(_ => WrapRootInUnion) + extensions.injectQueryStagePrepRule(_ => ObserveUnionPartitioning(seen)) + }, aqeEnabled = true) + assert(seen.nonEmpty, "the second prep rule must have seen the union") + assert(!seen.exists(_.isInstanceOf[UnknownPartitioning]), + s"the snapshot pass between the two rules must have recorded the conf: $seen") + } + + test("SPARK-59122: the codegen conf is recorded for a prep rule after the one that added it") { + // The same window, read through the codegen gate rather than the partitioning. The repartition + // is what makes AQE engage at all; it is round-robin, so the union above it has nothing to pass + // through and stays plain, which leaves its gate turning on the codegen confs. Only enablement + // is observable here: a one-child union is under any legal `maxChildren`, which cannot go below + // two, so `UnionCodegenSuite` pins that field on a union of three. + val seen = ListBuffer.empty[Boolean] + withSession(create { extensions => + extensions.injectQueryStagePrepRule(_ => WrapRootInUnion) + extensions.injectQueryStagePrepRule(_ => ObserveUnionSupportCodegen(seen)) + }) { session => + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) + val df = session.range(0, 20, 1, 2).repartition(2).selectExpr("id + 1 AS w") + assert(df.collect().map(_.getLong(0)).sorted.toSeq == (1L to 20L).toSeq, + "the injected union must not drop or duplicate rows") + assert(seen.nonEmpty, "the second prep rule must have seen the union") + assert(seen.forall(identity), + s"the snapshot pass between the two rules must have recorded the codegen conf: $seen") + } + } + test("custom aggregate hint") { // The custom hint allows us to replace the aggregate (without grouping keys) with just // Literal. @@ -1476,11 +1514,11 @@ object WrapRootInUnionColumnarRule extends ColumnarRule { } /** - * Records what each `UnionExec` reports while the AQE stage optimizers run, which is after the - * barrier at the end of the query stage preparation rules and before the one in - * `postStageCreationRules`. Reads it with `UNION_OUTPUT_PARTITIONING` turned off: the union an - * injected prep rule adds carries no recorded conf of its own, so a concrete answer can only come - * from a decision stamped earlier. Puts the conf back, so nothing downstream sees the flip. + * Records what each `UnionExec` reports, with `UNION_OUTPUT_PARTITIONING` turned off: a concrete + * answer can then only come from a decision stamped earlier, or from the conf a snapshot pass + * recorded. Injected as a stage-optimizer rule it runs after the barrier at the end of the query + * stage preparation rules; injected as a prep rule it runs among them. Puts the conf back, so + * nothing downstream sees the flip. */ case class ObserveUnionPartitioning(seen: ListBuffer[Partitioning]) extends Rule[SparkPlan] { override def apply(plan: SparkPlan): SparkPlan = { @@ -1496,6 +1534,26 @@ case class ObserveUnionPartitioning(seen: ListBuffer[Partitioning]) extends Rule } } +/** + * Records whether each `UnionExec` says it supports codegen, read with + * `WHOLESTAGE_UNION_CODEGEN_ENABLED` turned off: a `true` can then only come from the conf a + * snapshot pass recorded for it, or from a decision stamped earlier. Puts the conf back, so nothing + * downstream sees the flip. + */ +case class ObserveUnionSupportCodegen(seen: ListBuffer[Boolean]) extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + plan.foreach { + case u: UnionExec => + val enabled = u.conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED) + u.conf.setConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED, false) + try seen += u.supportCodegen + finally u.conf.setConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED, enabled) + case _ => + } + plan + } +} + // Example of an Aggregate hint that tells that 'attribute' values are no larger than 'max'. // We will use them to rewrite MAX(attribute) with 'max' constant. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index b835a73883c8b..ef5a502b79132 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -22,6 +22,7 @@ import java.util.concurrent.{CountDownLatch, Executors, TimeUnit} import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioningLike, UnknownPartitioning} +import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper} import org.apache.spark.sql.execution.columnar.InMemoryTableScanLike @@ -63,7 +64,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper * `SparkPlan.collect` stops at them, since both are `LeafExecNode`s. * * Stricter than `codegenUnions` on purpose: this matches only a union that is the root of its own - * codegen stage, which is the node the callers here reach for its tags and metrics. + * codegen stage, which is the node whose tags and metrics the callers here inspect. */ private def fusedUnions(df: DataFrame): Seq[UnionExec] = collect(df.queryExecution.executedPlan) { @@ -116,8 +117,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper private def stampUnionDecisions(plan: SparkPlan): SparkPlan = new StampUnionDecisions(UnionConfSnapshot(SQLConf.get))(plan) - private def snapshotUnionOutputPartitioningConf(plan: SparkPlan): SparkPlan = - new SnapshotUnionOutputPartitioningConf(UnionConfSnapshot(SQLConf.get))(plan) + private def snapshotUnionPreparationConf(plan: SparkPlan): SparkPlan = + new SnapshotUnionPreparationConf(UnionConfSnapshot(SQLConf.get))(plan) /** * Run `buildDf()` with union codegen on, then again with it off, and assert the two agree. @@ -800,7 +801,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper "planning and execution") { // `supportCodegenFailureReason` used to read `WHOLESTAGE_UNION_CODEGEN_ENABLED` live, and the // copy that `insertInputAdapter` puts inside the codegen shell evaluated it for the first time - // at execution. Planned with the conf on the union is fused, so the generated code increments + // at execution. Planned with the conf on, the union is fused, so the generated code increments // `numOutputRows`; if the copy re-derives the reason with the conf off, `metrics` comes back // empty and `doProduce` throws `key not found: numOutputRows`. withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { @@ -1043,7 +1044,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // `EnsureRequirements` asks the union what it reports, and the barrier behind it freezes that // answer one rule later. `conf` is live, so another thread turning `UNION_OUTPUT_PARTITIONING` // off in between would leave the parent's elided exchange standing over a union that then - // concatenates. `SnapshotUnionOutputPartitioningConf` records the value ahead of + // concatenates. `SnapshotUnionPreparationConf` records the value ahead of // `EnsureRequirements` for both to use. Driven rule by rule, because the two sit next to each // other in the pipeline and no injected rule can run in the window. withSQLConf( @@ -1053,7 +1054,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper // cannot land between them, but an edit to the list can. AQE builds its own list, and // `AdaptiveQueryExecSuite` asserts the same order there. val rules = QueryExecution.preparations(spark, subquery = false) - val snapshot = rules.indexWhere(_.isInstanceOf[SnapshotUnionOutputPartitioningConf]) + val snapshot = rules.indexWhere(_.isInstanceOf[SnapshotUnionPreparationConf]) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) assert(snapshot >= 0 && snapshot == ensureRequirements - 1, s"expected the conf snapshot right before EnsureRequirements at $ensureRequirements, " + @@ -1063,7 +1064,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper .union(spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k"))) .groupBy("k").count() val required = EnsureRequirements()( - snapshotUnionOutputPartitioningConf(df.queryExecution.sparkPlan.clone())) + snapshotUnionPreparationConf(df.queryExecution.sparkPlan.clone())) assert(required.collect { case s: ShuffleExchangeExec => s.shuffleOrigin } == Seq(REPARTITION_BY_NUM, REPARTITION_BY_NUM), "the aggregate's exchange must have been elided, or the window has nothing at stake") @@ -1092,7 +1093,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper .union(spark.range(20, 40, 1, 2).selectExpr("id % 5 AS k").repartition(4, col("k"))) .groupBy("k").count() val prepared = stampUnionDecisions(EnsureRequirements()( - snapshotUnionOutputPartitioningConf(df.queryExecution.sparkPlan.clone()))) + snapshotUnionPreparationConf(df.queryExecution.sparkPlan.clone()))) val union = prepared.collect { case u: UnionExec => u } assert(union.size == 1) // The value a preparation would carry to its late barrier, taken while the conf still says @@ -1186,6 +1187,69 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-59122: a snapshot pass ahead of each injected rule fills in what one created") { + // The barrier sits after the whole injected list, so between two injected rules a fresh union + // used to carry no record and answer live. `SnapshotUnionPreparationConf.before` lists a pass + // ahead of each of them. Driven by hand here; `SparkSessionExtensionSuite` runs it through real + // injected rules. + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") { + 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() + val prepared = EnsureRequirements()( + snapshotUnionPreparationConf(df.queryExecution.sparkPlan.clone())) + val kids = prepared.collect { case u: UnionExec => u }.head.children + // What a preparation carries to the rules below it, taken while the conf still says what + // `EnsureRequirements` above read. + val prepConf = UnionConfSnapshot(SQLConf.get) + + // Stands for an injected rule that returns a `UnionExec` of its own, twice over: the pass + // between the two is what the rule listed second would otherwise be missing. + val fresh = UnionExec(kids) + val injected = new Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = fresh + } + val listed = SnapshotUnionPreparationConf.before(prepConf, Seq(injected, injected)) + assert(listed.size == 4 && + listed.head.isInstanceOf[SnapshotUnionPreparationConf] && + listed(2).isInstanceOf[SnapshotUnionPreparationConf], + s"a snapshot pass must sit ahead of each injected rule, got ${listed.map(_.getClass)}") + + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + assert(fresh.isPlainUnion, "the fresh union must reach the list with no record of its own") + listed.foldLeft(prepared)((plan, rule) => rule(plan)) + assert(!fresh.outputPartitioning.isInstanceOf[UnknownPartitioning], + "the pass after the rule that created it must have recorded the preparation's conf, " + + s"got ${fresh.outputPartitioning}") + } + } + } + + test("SPARK-59122: a recorded conf answers the codegen gate before a decision is stamped") { + // The gate's two codegen confs fall back to the record the way the partitioning decision does, + // so a rule reading the gate on a union an earlier rule created gets the preparation's values. + // Three children, because the cap cannot go below two and so says nothing about a union of one. + val df = rangeDF(100).union(rangeDF(100)).union(rangeDF(100)) + val kids = df.queryExecution.sparkPlan.collect { case u: UnionExec => u }.head.children + assert(kids.size == 3, s"expected one union of three children, got ${kids.size}") + val prepConf = withSQLConf( + SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "3") { + UnionConfSnapshot(SQLConf.get) + } + + Seq( + "enablement" -> (SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false"), + "the child cap" -> (SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "2") + ).foreach { case (what, flipped) => + val union = UnionExec(kids) + new SnapshotUnionPreparationConf(prepConf)(union) + withSQLConf(flipped) { + assert(union.supportCodegen, s"the record must answer $what, not the value read now") + } + } + } + test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { // `InputFileName` is `Nondeterministic` but reads from `InputFileBlockHolder` // (a per-task thread-local) and does not embed `partitionIndex`. The gate's diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala index 4a97a5f962fe7..138f27a4fd9e1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -5327,7 +5327,7 @@ class AdaptiveQueryExecSuite test("SPARK-59122: query stage preparation keeps the union barriers around EnsureRequirements") { // `EnsureRequirements` asks a `UnionExec` what it reports, and `StampUnionDecisions` freezes // that answer so every rule below it and the execution read what the exchanges were planned - // against. `SnapshotUnionOutputPartitioningConf` has to run first, or the value the stamp reads + // against. `SnapshotUnionPreparationConf` has to run first, or the value the stamp reads // is whatever `conf` says by then rather than the one `EnsureRequirements` saw. The two sit // next to `EnsureRequirements` with nothing in between, and nothing at the list itself says the // three have to stay contiguous: an injected rule cannot land between them, since those are @@ -5345,7 +5345,7 @@ class AdaptiveQueryExecSuite } assert(aqe.isDefined, s"expected an AdaptiveSparkPlanExec:\n${df.queryExecution}") val rules = aqe.get.queryStagePreparationRules - val snapshot = rules.indexWhere(_.isInstanceOf[SnapshotUnionOutputPartitioningConf]) + val snapshot = rules.indexWhere(_.isInstanceOf[SnapshotUnionPreparationConf]) val ensureRequirements = rules.indexWhere(_.isInstanceOf[EnsureRequirements]) // Both barriers, not the first one: the list ends with a second `StampUnionDecisions` for a // union an injected prep rule created, and asking only for the first index would let that one From ef1eb31d2bdedfeb46145ada05196e452cef26ab Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 22 Sep 2026 15:59:37 +0800 Subject: [PATCH 20/24] Close the post planner strategy window and drop the dead stamped confs --- .../apache/spark/sql/internal/SQLConf.scala | 17 ++-- .../sql/execution/StampUnionDecisions.scala | 14 ++- .../adaptive/AdaptiveSparkPlanExec.scala | 6 +- .../execution/basicPhysicalOperators.scala | 99 +++++++++---------- .../sql/SparkSessionExtensionSuite.scala | 15 +++ .../sql/execution/UnionCodegenSuite.scala | 42 +++++--- 6 files changed, 109 insertions(+), 84 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index e238ebac222bf..b9b6f6058b7d5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2938,9 +2938,9 @@ object SQLConf { .doc("When both this conf and `spark.sql.codegen.wholeStage` are true, an eligible " + "UnionExec on its non-partitioning-aware path takes part in whole-stage codegen. " + "The union's other eligibility checks still apply, and a child that does not support " + - "codegen still ends the stage at an InputAdapter. The value is captured once for the " + - "plan a physical preparation prepares, before its rule sequence, and every UnionExec " + - "decision taken in that sequence comes from that capture.") + "codegen still ends the stage at an InputAdapter. The value is read once per physical " + + "preparation, so a union's codegen gate and the copy of it inside the generated stage " + + "agree.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .booleanConf @@ -2955,9 +2955,9 @@ object SQLConf { "bytecode size, constant pool growth, JIT compilation time) rather " + "than the JVM per-method bytecode limit. Unions with more children " + "fall back to per-child codegen stages. Only effective when " + - s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true. The value is captured once for the " + - "plan a physical preparation prepares, before its rule sequence, and every UnionExec " + - "decision taken in that sequence comes from that capture.") + s"`${WHOLESTAGE_UNION_CODEGEN_ENABLED.key}` is true. The value is read once per physical " + + "preparation, so a union's codegen gate and the copy of it inside the generated stage " + + "agree.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .intConf @@ -8182,9 +8182,8 @@ object SQLConf { .internal() .doc("When set to true, the output partitioning of UnionExec will be the same as the " + "input partitioning if its children have same partitioning. Otherwise, it will be a " + - "default partitioning. The value is captured once for the plan a physical preparation " + - "prepares, before its rule sequence, and every UnionExec decision taken in that sequence " + - "comes from that capture.") + "default partitioning. The value is read once per physical preparation, so the exchanges " + + "planned around a UnionExec and the decision it executes under agree.") .version("4.1.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index 0c5bc21089be6..11dda8178d229 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -74,11 +74,15 @@ class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { * * The same gap opens between two injected rules: one can return a `UnionExec` of its own, which * carries no record yet, and a later one can plan requirements over it or ask its codegen gate. - * [[SnapshotUnionPreparationConf.before]] closes that for the injected query-stage preparation - * rules, the list where a consumer can still add or drop an exchange. Two windows stay open, both - * behind a barrier that decides before execution: the injected columnar rules, which share one - * `ApplyColumnarRulesAndInsertTransitions` and so cannot be interleaved from outside it, and the - * injected query-stage optimizer rules, which run on a plan whose exchanges are already fixed. + * [[SnapshotUnionPreparationConf.before]] closes that for the two injected lists whose rules can + * still add or drop an exchange, the AQE post-planner-strategy rules and the query-stage + * preparation rules. Two windows stay open, and in neither can a reader add or drop an exchange. + * The injected columnar rules share one `ApplyColumnarRulesAndInsertTransitions`, so a pass cannot + * be listed between them from outside it. The injected query-stage optimizer rules run on a plan + * whose exchanges are fixed, and they are listed after the built-in ones, so what a live read there + * can still move is what the codegen gate answers, whose own barrier lands before + * `CollapseCodegenStages`, and, for an injected rule that is itself an `AQEShuffleReadRule`, + * whether `ValidateRequirements` keeps its rewrite. * * Only the confs are recorded, never a partitioning. `EnsureRequirements` has not inserted the * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index b7d42d1fbd869..ccf1c112714d6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -251,7 +251,11 @@ case class AdaptiveSparkPlanExec( private def applyQueryPostPlannerStrategyRules(plan: SparkPlan): SparkPlan = { applyPhysicalRules( plan, - context.session.sessionState.adaptiveRulesHolder.queryPostPlannerStrategyRules, + // These rules run before `ensureRequirements`, so one of them can still add or drop an + // exchange over a `UnionExec` another just created. A snapshot pass ahead of each is what + // keeps that read off the live conf. + SnapshotUnionPreparationConf.before( + unionConf, context.session.sessionState.adaptiveRulesHolder.queryPostPlannerStrategyRules), "AQE Query Post Planner Strategy Rules" ) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 80fc26d563162..8ed2bee76d79d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1094,10 +1094,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * Records the confs this node answers from until its decisions are stamped, read once per * preparation into a `UnionConfSnapshot` and passed in here. `SnapshotUnionPreparationConf` does * it ahead of `EnsureRequirements`, whose reads the following stamp has to agree with, and ahead - * of each injected rule, so a union an earlier one created is not read live by a later one; - * `stampDecisions` does it for a node no such pass saw. Only the confs, never a partitioning: the - * exchanges `EnsureRequirements` adds are not there yet, so a decision taken there would freeze - * plain on a union whose children only become co-partitioned in it. + * of the two injected lists that can still change an exchange; see that rule for the two windows + * it does not cover. `stampDecisions` does it for a node no such pass saw. Only the confs, never + * a partitioning: the exchanges `EnsureRequirements` adds are not there yet, so a decision taken + * there would freeze plain on a union whose children only become co-partitioned in it. */ private[execution] def recordPreparationConf(snapshot: UnionConfSnapshot): Unit = if (preparationConf.isEmpty) { @@ -1120,14 +1120,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup // for a union an injected rule added, and for one it rebuilt carrying tags of its own, which // is enough to stop `copyTagsFrom` from bringing this tag across. recordPreparationConf(snapshot) - // Every field from the record this node now carries, which is `snapshot` unless a pass had - // already recorded one. Reading some fields from the record and others from the argument - // would stamp half from each if the two ever differed. - val recorded = preparationConf.getOrElse(snapshot) - setTagValue(UnionExec.DECISIONS, UnionExec.Decisions( - plainUnion = isPlainUnion, - unionCodegenEnabled = recorded.codegenEnabled, - maxChildren = recorded.maxChildren)) + setTagValue(UnionExec.DECISIONS, UnionExec.Decisions(plainUnion = isPlainUnion)) } /** @@ -1175,30 +1168,30 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // The confs the gate reads, stamped for the reason the plain-union decision is: `conf` is live, - // so the gate, `metrics` and the copy `insertInputAdapter` puts inside the codegen shell would - // otherwise be free to read different values. When a child is not `CodegenSupport` that copy is - // real and its first evaluation lands at execution, which is where reading the conf produced the - // failure described on `isPlainUnion`. Before the stamp they come from the preparation's snapshot - // where a pass recorded one, so a rule reading the gate on a union an earlier rule created gets - // the value that preparation will stamp rather than the live conf. + // The confs the gate reads, taken from the record rather than live for the reason the plain-union + // decision is stamped: `conf` is live, so the gate, `metrics` and the copy `insertInputAdapter` + // puts inside the codegen shell would otherwise be free to read different values. When a child is + // not `CodegenSupport` that copy is real and its first evaluation lands at execution, which is + // where reading the conf produced the failure described on `isPlainUnion`. A stamped node always + // carries a record, so these do not need a branch for the stamp. private def unionCodegenEnabled: Boolean = - stampedDecisions.map(_.unionCodegenEnabled) - .orElse(preparationConf.map(_.codegenEnabled)) + preparationConf.map(_.codegenEnabled) .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) private def maxCodegenChildren: Int = - stampedDecisions.map(_.maxChildren) - .orElse(preparationConf.map(_.maxChildren)) + preparationConf.map(_.maxChildren) .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) // Memoized per instance rather than stamped on the tag. Every term here reads the children, and a // tag outlives them: `SQLExecution` builds the initial `SparkPlanInfo` before execution, forcing // `metrics` on every node it visits, so a rule that replaces a child after that would inherit an // allowing answer and fuse a topology that `hasPartitionIndexDependentCodegen` or - // `supportsColumnar` rejects. The copy in the codegen shell still agrees with the gate: - // `InputAdapter` delegates `output` and `supportsColumnar` to its child, the other terms walk the - // subtree through it, and each of those is fixed for a given set of children. + // `supportsColumnar` rejects. That is a trade rather than a free win: a rule placed after + // `CollapseCodegenStages` that weakened a child would instead leave a stamped union with empty + // `metrics` inside a shell that still fuses, which `doProduce` reports as a missing + // `numOutputRows`. Nothing in the tree does that. The copy in the codegen shell agrees with the + // gate either way: `InputAdapter` delegates `output` and `supportsColumnar` to its child, the + // other terms walk the subtree through it, and each is fixed for a given set of children. @transient private lazy val childTopologyFailureReason: Option[String] = { if (children.exists(_.exists(_.isInstanceOf[UnionExec]))) { Some("nested-union") @@ -1216,13 +1209,14 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // The three preparation-scoped terms are recomputed per call, because memoizing them would let a - // read arriving before `stampDecisions` settle the gate on the live conf. The stamp can install - // the opposite value, and the gate would then keep the memoized answer while the copy - // `insertInputAdapter` builds, a fresh instance carrying the stamped tag, derives the other one. - // Nothing on Spark's own path reads a union that early: every barrier runs inside `preparations`, - // ahead of the `SparkPlanInfo` that forces `metrics`. A late extension hook can, and so can a - // caller inspecting `sparkPlan`. + // The three preparation-scoped terms are recomputed per call. `isPlainUnion` has to be: it + // derives `rawPartitioning` from the children, so a read before `stampDecisions` would otherwise + // settle the gate on children that move. The two conf terms have to be for a different reason: a + // node with no record answers live, and a pass can write the record afterwards, so a memoized + // answer would keep the live value while the copy `insertInputAdapter` builds, a fresh instance + // carrying that record, reads it. No Spark rule reads this gate before `CollapseCodegenStages`, + // which every barrier precedes; a late extension hook can, and so can a caller inspecting + // `sparkPlan`. private def supportCodegenFailureReason: Option[String] = { if (!unionCodegenEnabled) { Some("union-codegen-disabled") @@ -1247,13 +1241,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } } - // Registered only when fusion will actually run, so plans that fall back to `doExecute` (which - // never updates the metric) do not surface a 0-valued row count in the SQL UI. `doConsume` is the - // sole incrementer. An unstamped node is the exception: its gate is still provisional, the stamp - // can land either way, and a map built without the metric would leave `doProduce` asking - // `metricTerm` for one that is not there. Registering it then costs an unused metric on a union - // an extension inspected and the stamp went on to reject; Spark's own force, the `SparkPlanInfo` - // `SQLExecution` builds, runs after every barrier, so ordinary fallback unions still omit it. + // Registered only when this union's own codegen gate allows fusion, so a union the gate rejects + // does not surface a 0-valued row count in the SQL UI; `doConsume` is the sole incrementer, and + // `doExecute` never touches it. An unstamped node is the exception: its gate is provisional, + // because `isPlainUnion` derives from children that can still move, so the stamp can land either + // way and a map built without the metric would leave `doProduce` asking `metricTerm` for one that + // is not there. Registering it then costs an unused metric on a union an extension inspected and + // the stamp went on to reject. override lazy val metrics: Map[String, SQLMetric] = if (stampedDecisions.isEmpty || supportCodegenFailureReason.isEmpty) { Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) @@ -1438,14 +1432,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup object UnionExec { /** - * What `StampUnionDecisions` fixes on a `UnionExec`: whether it is a plain concatenation, and the - * two confs the codegen gate reads. Everything else the gate asks is derived per instance, so a - * rule replacing a child cannot inherit an answer taken from the topology it replaced. + * What `StampUnionDecisions` fixes on a `UnionExec`: whether it is a plain concatenation. That is + * the one answer the confs alone do not give, since it also depends on the children. The confs + * themselves stay on `PREPARATION_CONF`, which a stamped node always carries. Everything else the + * gate asks is derived per instance, so a rule replacing a child cannot inherit an answer taken + * from the topology it replaced. */ - private case class Decisions( - plainUnion: Boolean, - unionCodegenEnabled: Boolean, - maxChildren: Int) + private case class Decisions(plainUnion: Boolean) /** * The stamped decisions. See `isPlainUnion` and `stampDecisions`. @@ -1458,12 +1451,12 @@ object UnionExec { private val DECISIONS = TreeNodeTag[Decisions]("unionDecisions") /** - * The [[UnionConfSnapshot]] this node answers from until its decisions are stamped. See - * `recordPreparationConf`. `SnapshotUnionPreparationConf` writes it before `EnsureRequirements`, - * and the stamp after it reads what that rule wrote, so both phases use one value; the same rule - * runs ahead of each injected rule, so a union one of them creates is not read live by a later - * one, and `stampDecisions` writes it itself for a node no such pass saw. The tag travels onto - * rebuilt nodes the same way `DECISIONS` does, which is what carries it across the copies + * The [[UnionConfSnapshot]] this node answers from, and the one a stamped node's codegen gate + * keeps answering from. See `recordPreparationConf`. `SnapshotUnionPreparationConf` writes it + * before `EnsureRequirements`, and the stamp after it reads what that rule wrote, so both phases + * use one value; the same rule runs ahead of the two injected lists that can still change an + * exchange, and `stampDecisions` writes it itself for a node no such pass saw. The tag travels + * onto rebuilt nodes the same way `DECISIONS` does, which is what carries it across the copies * `EnsureRequirements` makes. */ private val PREPARATION_CONF = TreeNodeTag[UnionConfSnapshot]("unionPreparationConf") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index a6b0a423f2537..8adaba2e306d5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -705,6 +705,21 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { s"the snapshot pass between the two rules must have recorded the conf: $seen") } + test("SPARK-59122: a union an injected post planner strategy rule adds is recorded next") { + // These rules run before `ensureRequirements`, so a reader among them can still drop an + // exchange over a union another of them created, which makes this the injected list where the + // window costs wrong rows rather than a lost fusion. + val seen = ListBuffer.empty[Partitioning] + checkInjectedUnionIsStamped( + create { extensions => + extensions.injectQueryPostPlannerStrategyRule(_ => WrapRootInUnion) + extensions.injectQueryPostPlannerStrategyRule(_ => ObserveUnionPartitioning(seen)) + }, aqeEnabled = true) + assert(seen.nonEmpty, "the second post planner strategy rule must have seen the union") + assert(!seen.exists(_.isInstanceOf[UnknownPartitioning]), + s"the snapshot pass between the two rules must have recorded the conf: $seen") + } + test("SPARK-59122: the codegen conf is recorded for a prep rule after the one that added it") { // The same window, read through the codegen gate rather than the partitioning. The repartition // is what makes AQE engage at all; it is round-robin, so the union above it has nothing to pass diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala index ef5a502b79132..4636a3ae09554 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala @@ -38,7 +38,8 @@ import org.apache.spark.sql.types._ */ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { - // Union codegen fusion is off by default; turn it on for this suite. + // Pinned rather than inherited: these cases turn on which value the union was prepared with, so a + // change to the conf's default must not silently change what they exercise. override protected def sparkConf: SparkConf = super.sparkConf.set(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key, "true") @@ -824,9 +825,9 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val copy = fusedUnions(planned) assert(copy.size == 1) // The copy this test needs: `insertInputAdapter` wrapped both children, so the shell holds - // a copy rather than the instance the gate answered on. This copy's reason is first forced - // by the `SparkPlanInfo` that `collect()` above builds, with the conf already off, so what - // it answers can only come from the stamp. + // a copy rather than the instance the gate answered on. The copy carries the record the + // preparation wrote, which is the only place its gate can get the conf from now that the + // live one says the opposite. assert(copy.head.children.forall(_.isInstanceOf[InputAdapter])) assert(copy.head.supportCodegen, "the copy in the shell must keep the decision it was planned with") @@ -949,7 +950,8 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> (!live).toString) { stampUnionDecisions(union) assert(union.supportCodegen == !live, - s"the gate must answer from the stamp, not from the read taken with the conf $live") + s"the gate must answer from the barrier's snapshot, not from the read taken with the " + + s"conf $live") } } } @@ -1023,10 +1025,10 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(fresh.isPlainUnion, "the barrier must have decided the fresh node") } - // The other half. The conf a decision was stamped with is the part a second pass could move, - // so the node to watch is one whose gate the conf still answers: plain, and with its reason - // not yet forced. `fusedUnions` returns the copy inside the codegen shell, whose reason no - // preparation rule has asked for, so what it answers below comes from the stamp alone. + // The other half. The conf a barrier records is the part a second pass could move, so the + // node to watch is one whose gate the conf still answers: plain, and with its reason not yet + // forced. `fusedUnions` returns the copy inside the codegen shell, whose reason no + // preparation rule has asked for, so what it answers below comes from the recorded conf. val fused = withSQLConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "true") { val df = rangeDF(100).repartition(2).union(rangeDF(100).repartition(2)) val union = fusedUnions(df) @@ -1118,12 +1120,11 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper } } - test("SPARK-59122: the stamp takes the codegen decisions from the preparation's snapshot") { - // The case above pins this for `UNION_OUTPUT_PARTITIONING`. The two codegen fields ride in the - // same `UnionConfSnapshot` and need the same pin: a stamp that read them live would agree with - // the snapshot everywhere else in this suite, because the other cases flip the conf after - // stamping and so pin the gate rather than the stamp. Nothing here is executed, so AQE does not - // enter into it. + test("SPARK-59122: the barrier records the preparation's codegen confs for a fresh node") { + // The case above pins the same for `UNION_OUTPUT_PARTITIONING`. The codegen confs live on + // the record rather than on the stamped decision, so what needs pinning here is that a barrier + // reaching a union no snapshot pass saw records that preparation's values and not the ones the + // conf says by then. Nothing here is executed, so AQE does not enter into it. val df = rangeDF(100).union(rangeDF(100)).union(rangeDF(100)) val planned = df.queryExecution.sparkPlan.collect { case u: UnionExec => u } assert(planned.size == 1 && planned.head.children.size == 3, @@ -1144,7 +1145,7 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper val union = UnionExec(kids) withSQLConf(flipped) { new StampUnionDecisions(prepConf)(union) - assert(union.supportCodegen, s"the stamp must take $what from the snapshot") + assert(union.supportCodegen, s"the record must hold $what from the snapshot") } } } @@ -1248,6 +1249,15 @@ class UnionCodegenSuite extends SharedSparkSession with AdaptiveSparkPlanHelper assert(union.supportCodegen, s"the record must answer $what, not the value read now") } } + + // A record is not a decision: `isPlainUnion` still derives from the children, so an unstamped + // gate stays provisional and the metric has to be there whichever way the stamp lands. Recorded + // with fusion off, so the gate denies and only the unstamped term can register it. + val denied = UnionExec(kids) + new SnapshotUnionPreparationConf(prepConf.copy(codegenEnabled = false))(denied) + assert(!denied.supportCodegen, "this half needs the gate to deny, or it pins nothing") + assert(denied.metrics.contains("numOutputRows"), + "a recorded but unstamped union must register the metric even where its gate denies") } test("SPARK-56482: input_file_name child fuses (Nondeterministic but partition-index-free)") { From 8a64d12f45867fb547f7c4cd0f9f6e949163ceb3 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Thu, 24 Sep 2026 16:32:26 +0800 Subject: [PATCH 21/24] Pin the late AQE barrier's snapshot in tests and qualify the conf doc --- .../apache/spark/sql/internal/SQLConf.scala | 6 +- .../sql/SparkSessionExtensionSuite.scala | 82 +++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index b9b6f6058b7d5..1d10b5cf112da 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -8182,8 +8182,10 @@ object SQLConf { .internal() .doc("When set to true, the output partitioning of UnionExec will be the same as the " + "input partitioning if its children have same partitioning. Otherwise, it will be a " + - "default partitioning. The value is read once per physical preparation, so the exchanges " + - "planned around a UnionExec and the decision it executes under agree.") + "default partitioning. The value is read once per physical preparation, and the decision " + + "taken with it, so the exchanges planned around a UnionExec and the decision it executes " + + "under agree. One decided to concatenate keeps reporting the default partitioning if its " + + "children come to share one afterwards.") .version("4.1.0") .booleanConf .createWithDefault(true) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index 8adaba2e306d5..20a45b375ca0e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -741,6 +741,70 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { } } + test("SPARK-59122: the barrier in AQE post stage creation stamps from the adaptive snapshot") { + // The cases above leave the conf where it is until their barriers have run, so one answering + // from a read of its own rather than the snapshot `AdaptiveSparkPlanExec` took at construction + // would pass them too. Here the two differ: the flip lands after the wrapper is built and + // before the stage carrying the injected union is created. This is where it can be seen, since + // the rules listed per stage are the ones built anew each time; the lists that run per + // re-planning round are built once, with the snapshot, so a barrier in them cannot take a + // later value to begin with. + withSession(create(_.injectColumnar(_ => WrapRootInUnionColumnarRule))) { session => + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) + session.conf.set(SQLConf.UNION_OUTPUT_PARTITIONING.key, true) + val df = session.range(0, 20, 1, 2).selectExpr("id % 5 AS k", "id AS v") + .repartition(4, col("k")).selectExpr("k", "v + 1 AS w") + // Constructing the wrapper is what reads AQE's snapshot. + assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) + + session.conf.set(SQLConf.UNION_OUTPUT_PARTITIONING.key, false) + // The stages, the columnar rules and the barrier behind them all run here. + assert(df.collect().map(_.getLong(1)).sorted.toSeq == (1L to 20L).toSeq, + "the injected union must not drop or duplicate rows") + + val unions = collect(df.queryExecution.executedPlan) { case u: UnionExec => u } + assert(unions.size == 1, s"expected the one union the rule adds, got ${unions.size}") + assert(!unions.head.outputPartitioning.isInstanceOf[UnknownPartitioning], + "the barrier must stamp from the snapshot, not from the value the conf holds when it " + + s"runs, got ${unions.head.outputPartitioning}") + } + } + + test("SPARK-59122: a late AQE barrier records the codegen confs the wrapper was built with") { + // The same divergence read through the codegen gate. A barrier taking the value the conf holds + // when it runs would record fusion as out, so nothing would fuse and the copy of the union + // inside the shell would carry no `numOutputRows`. Round-robin children, so the union has + // nothing to pass through and stays plain, which is what leaves its gate on these two confs. + withSession(create(_.injectColumnar(_ => RebuildRootUnionColumnarRule))) { session => + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) + Seq( + "enablement" -> (SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key -> "false"), + "the child cap" -> (SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key -> "2") + ).foreach { case (what, (key, flipped)) => + session.conf.set(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED.key, "true") + // Three children, so that a cap of two excludes this union; the cap cannot go below two. + session.conf.set(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN.key, "3") + val df = session.range(0, 20, 1, 2).repartition(2) + .union(session.range(20, 40, 1, 2).repartition(2)) + .union(session.range(40, 60, 1, 2).repartition(2)) + assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) + + session.conf.set(key, flipped) + assert(df.collect().map(_.longValue()).sorted.toSeq == (0L until 60L).toSeq, + "the rebuilt union must not drop or duplicate rows") + + val fused = collect(df.queryExecution.executedPlan) { case w: WholeStageCodegenExec => w } + .flatMap(_.collect { case u: UnionExec => u }) + assert(fused.size == 1, s"the rebuilt union must fuse with $what flipped, got\n" + + df.queryExecution.executedPlan) + assert(fused.head.children.forall(_.isInstanceOf[InputAdapter]), + s"expected the copy inside the shell, got ${fused.head.children.map(_.getClass)}") + assert(fused.head.metrics("numOutputRows").value == 60, + s"the copy inside the shell must count the rows with $what flipped") + } + } + } + test("custom aggregate hint") { // The custom hint allows us to replace the aggregate (without grouping keys) with just // Literal. @@ -1528,6 +1592,24 @@ object WrapRootInUnionColumnarRule extends ColumnarRule { override def postColumnarTransitions: Rule[SparkPlan] = WrapRootInUnion } +/** + * Stands for an extension that returns a `UnionExec` of its own in place of one already in the + * plan: a fresh instance over the same children, so it carries no stamped decision, and the rows + * are the same ones. Replaces the root only, which is where `postStageCreationRules` hands it a + * union; a stage's own plan is rooted at the exchange, so those applications leave it alone. + */ +object RebuildRootUnion extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = plan match { + case u: UnionExec => UnionExec(u.children) + case other => other + } +} + +/** The columnar-rule wrapper for `RebuildRootUnion`. */ +object RebuildRootUnionColumnarRule extends ColumnarRule { + override def postColumnarTransitions: Rule[SparkPlan] = RebuildRootUnion +} + /** * Records what each `UnionExec` reports, with `UNION_OUTPUT_PARTITIONING` turned off: a concrete * answer can then only come from a decision stamped earlier, or from the conf a snapshot pass From c9d9095237b1913938514eab2eea971970832519 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 25 Sep 2026 19:32:20 +0800 Subject: [PATCH 22/24] Record the preparation conf between the folded extension rules --- .../spark/sql/execution/QueryExecution.scala | 3 +- .../sql/execution/StampUnionDecisions.scala | 47 +++++++++++++++---- .../adaptive/AdaptiveSparkPlanExec.scala | 14 +++++- .../sql/SparkSessionExtensionSuite.scala | 45 ++++++++++++++++++ 4 files changed, 97 insertions(+), 12 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index ebc9acccb3b77..76a5eb6921525 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -861,7 +861,8 @@ object QueryExecution { // see `AdaptiveSparkPlanExec.queryStagePreparationRules`.) RemoveRedundantSorts, ApplyColumnarRulesAndInsertTransitions( - sparkSession.sessionState.columnarRules, outputsColumnar = false), + SnapshotUnionPreparationConf.after(unionConf, sparkSession.sessionState.columnarRules), + outputsColumnar = false), // A barrier for a `UnionExec` an injected columnar rule just created, which has no decision // yet and would otherwise take one wherever it is first asked. A decision already stamped on // a node is kept. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index 11dda8178d229..d4f60094803d6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -73,16 +73,14 @@ class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { * partitions. * * The same gap opens between two injected rules: one can return a `UnionExec` of its own, which - * carries no record yet, and a later one can plan requirements over it or ask its codegen gate. - * [[SnapshotUnionPreparationConf.before]] closes that for the two injected lists whose rules can - * still add or drop an exchange, the AQE post-planner-strategy rules and the query-stage - * preparation rules. Two windows stay open, and in neither can a reader add or drop an exchange. - * The injected columnar rules share one `ApplyColumnarRulesAndInsertTransitions`, so a pass cannot - * be listed between them from outside it. The injected query-stage optimizer rules run on a plan - * whose exchanges are fixed, and they are listed after the built-in ones, so what a live read there - * can still move is what the codegen gate answers, whose own barrier lands before - * `CollapseCodegenStages`, and, for an injected rule that is itself an `AQEShuffleReadRule`, - * whether `ValidateRequirements` keeps its rewrite. + * carries no record yet, and a later one can plan requirements over it, build a parent over it, or + * ask its codegen gate. [[SnapshotUnionPreparationConf.before]] closes that for the two lists whose + * rules are listed as plain `Rule[SparkPlan]`, the AQE post-planner-strategy rules and the + * query-stage preparation rules. [[SnapshotUnionPreparationConf.after]] closes it for the injected + * columnar rules, which share one `ApplyColumnarRulesAndInsertTransitions` and so have no place + * between them for a listed pass: the record rides on each rule's own transitions instead. + * `AdaptiveSparkPlanExec.optimizeQueryStage` writes it on each rule result it folds over, which + * covers the next rule and the `ValidateRequirements` check on an `AQEShuffleReadRule`'s rewrite. * * Only the confs are recorded, never a partitioning. `EnsureRequirements` has not inserted the * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only @@ -109,6 +107,35 @@ object SnapshotUnionPreparationConf { */ def before(snapshot: UnionConfSnapshot, rules: Seq[Rule[SparkPlan]]): Seq[Rule[SparkPlan]] = rules.flatMap(rule => Seq(new SnapshotUnionPreparationConf(snapshot), rule)) + + /** + * `rules` with the same record written behind each of their transitions. + * `ApplyColumnarRulesAndInsertTransitions` applies these itself, every + * `preColumnarTransitions` in order and then every `postColumnarTransitions` in reverse, so a + * pass cannot be listed between two of them from outside; wrapping the transitions puts it there. + * Behind rather than ahead, because the plan reaching the wrapper has passed a barrier already, + * and it is what a rule returns that can hold a union no pass has seen. `rules` is empty unless + * an extension injected something. + */ + def after(snapshot: UnionConfSnapshot, rules: Seq[ColumnarRule]): Seq[ColumnarRule] = + rules.map(new RecordUnionPreparationConf(_, snapshot)) +} + +/** + * Records `snapshot` on each `UnionExec` in what `inner` returns, so a union `inner` created is + * read from the preparation's confs by the next rule in the same + * `ApplyColumnarRulesAndInsertTransitions`. See [[SnapshotUnionPreparationConf.after]]. + */ +private class RecordUnionPreparationConf(inner: ColumnarRule, snapshot: UnionConfSnapshot) + extends ColumnarRule { + + private val record = new SnapshotUnionPreparationConf(snapshot) + + override def preColumnarTransitions: Rule[SparkPlan] = + plan => record(inner.preColumnarTransitions(plan)) + + override def postColumnarTransitions: Rule[SparkPlan] = + plan => record(inner.postColumnarTransitions(plan)) } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index ccf1c112714d6..2245697f5ff10 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -124,6 +124,10 @@ case class AdaptiveSparkPlanExec( // same session conf this query's `QueryExecution.preparations` reads. @transient private val unionConf = UnionConfSnapshot(context.session.sessionState.conf) + // The same record the listed passes write, for the two places that need it between two rules + // rather than as a pass of their own. See `SnapshotUnionPreparationConf`. + @transient private val recordUnionConf = new SnapshotUnionPreparationConf(unionConf) + // A list of physical plan rules to be applied before creation of query stages. The physical // plan should reach a final status of query stages (i.e., no more addition or removal of // Exchange nodes) after running these rules. @@ -205,7 +209,8 @@ case class AdaptiveSparkPlanExec( // plan to these rules has exchange as its root node. private def postStageCreationRules(outputsColumnar: Boolean) = Seq( ApplyColumnarRulesAndInsertTransitions( - context.session.sessionState.columnarRules, outputsColumnar), + SnapshotUnionPreparationConf.after(unionConf, context.session.sessionState.columnarRules), + outputsColumnar), // A barrier for a `UnionExec` an injected stage-optimizer or columnar rule just created, which // has no decision yet and would otherwise take one wherever it is first asked. A decision // already stamped on a node is kept, so this pass cannot move one. @@ -222,6 +227,13 @@ case class AdaptiveSparkPlanExec( } val optimized = rules.foldLeft(plan) { case (latestPlan, rule) => val applied = rule.apply(latestPlan) + if (applied ne latestPlan) { + // A `UnionExec` this rule just created carries no record of the confs this execution + // answers from, and the `ValidateRequirements` check below and the next rule both read the + // plan before the barrier in `postStageCreationRules` stamps it. A rule that returned its + // input added nothing. + recordUnionConf(applied) + } val result = rule match { case _: AQEShuffleReadRule if !applied.fastEquals(latestPlan) => val distribution = if (isFinalStage) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index 20a45b375ca0e..958320727431f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -720,6 +720,45 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { s"the snapshot pass between the two rules must have recorded the conf: $seen") } + test("SPARK-59122: a union one injected stage optimizer rule adds is recorded before the next") { + // These rules are folded over inside `optimizeQueryStage`, which validates an + // `AQEShuffleReadRule`'s rewrite and then folds on, both before the barrier in + // `postStageCreationRules`. The record is written on each rule's result there, there being no + // list to put a pass into. + val seen = ListBuffer.empty[Partitioning] + withSession(create { extensions => + extensions.injectQueryStageOptimizerRule(_ => WrapRootInUnion) + extensions.injectQueryStageOptimizerRule(_ => ObserveUnionPartitioning(seen)) + }) { session => + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) + val df = session.range(0, 20, 1, 2).selectExpr("id % 5 AS k", "id AS v") + .repartition(4, col("k")).selectExpr("k", "v + 1 AS w") + assert(df.collect().map(_.getLong(1)).sorted.toSeq == (1L to 20L).toSeq, + "the injected union must not drop or duplicate rows") + // These rules run on each stage's plan as well, and the union the first one adds inside the + // shuffle stands over a range, which has nothing to pass through either way. + assert(seen.exists(!_.isInstanceOf[UnknownPartitioning]), + s"the union over the shuffle read must answer from the recorded conf: $seen") + } + } + + test("SPARK-59122: a union one injected columnar rule adds is recorded before the next") { + // The injected columnar rules share one `ApplyColumnarRulesAndInsertTransitions`, so the record + // rides on each rule's own transitions. Those run in reverse order on the way out, so the rule + // that adds the union is the one injected second. + Seq(false, true).foreach { aqeEnabled => + val seen = ListBuffer.empty[Partitioning] + checkInjectedUnionIsStamped( + create { extensions => + extensions.injectColumnar(_ => ObserveUnionPartitioningColumnarRule(seen)) + extensions.injectColumnar(_ => WrapRootInUnionColumnarRule) + }, aqeEnabled) + assert(seen.nonEmpty, s"the second columnar rule must have seen the union, aqe=$aqeEnabled") + assert(!seen.exists(_.isInstanceOf[UnknownPartitioning]), + s"the record must be written behind the rule that added it, aqe=$aqeEnabled: $seen") + } + } + test("SPARK-59122: the codegen conf is recorded for a prep rule after the one that added it") { // The same window, read through the codegen gate rather than the partitioning. The repartition // is what makes AQE engage at all; it is round-robin, so the union above it has nothing to pass @@ -1631,6 +1670,12 @@ case class ObserveUnionPartitioning(seen: ListBuffer[Partitioning]) extends Rule } } +/** The columnar-rule wrapper for `ObserveUnionPartitioning`. */ +case class ObserveUnionPartitioningColumnarRule(seen: ListBuffer[Partitioning]) + extends ColumnarRule { + override def postColumnarTransitions: Rule[SparkPlan] = ObserveUnionPartitioning(seen) +} + /** * Records whether each `UnionExec` says it supports codegen, read with * `WHOLESTAGE_UNION_CODEGEN_ENABLED` turned off: a `true` can then only come from the conf a From b49df28e34c2ce9ecf52ecb372718c1e2cb2b057 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 25 Sep 2026 20:28:01 +0800 Subject: [PATCH 23/24] Make the columnar record comparable and cover its pre phase --- .../sql/execution/StampUnionDecisions.scala | 10 +++-- .../adaptive/AdaptiveSparkPlanExec.scala | 7 +-- .../execution/basicPhysicalOperators.scala | 13 +++--- .../sql/SparkSessionExtensionSuite.scala | 44 ++++++++++++++----- 4 files changed, 51 insertions(+), 23 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index d4f60094803d6..f3d43842d3df5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -79,8 +79,10 @@ class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { * query-stage preparation rules. [[SnapshotUnionPreparationConf.after]] closes it for the injected * columnar rules, which share one `ApplyColumnarRulesAndInsertTransitions` and so have no place * between them for a listed pass: the record rides on each rule's own transitions instead. - * `AdaptiveSparkPlanExec.optimizeQueryStage` writes it on each rule result it folds over, which - * covers the next rule and the `ValidateRequirements` check on an `AQEShuffleReadRule`'s rewrite. + * `AdaptiveSparkPlanExec.optimizeQueryStage` writes it on every rule result that changed the plan, + * which covers the next rule, and `ValidateRequirements` for an injected rule that is itself an + * `AQEShuffleReadRule`. The built-in ones are listed ahead of the injected list, so what they + * validate holds no union without a record. * * Only the confs are recorded, never a partitioning. `EnsureRequirements` has not inserted the * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only @@ -118,7 +120,7 @@ object SnapshotUnionPreparationConf { * an extension injected something. */ def after(snapshot: UnionConfSnapshot, rules: Seq[ColumnarRule]): Seq[ColumnarRule] = - rules.map(new RecordUnionPreparationConf(_, snapshot)) + rules.map(RecordUnionPreparationConf(_, snapshot)) } /** @@ -126,7 +128,7 @@ object SnapshotUnionPreparationConf { * read from the preparation's confs by the next rule in the same * `ApplyColumnarRulesAndInsertTransitions`. See [[SnapshotUnionPreparationConf.after]]. */ -private class RecordUnionPreparationConf(inner: ColumnarRule, snapshot: UnionConfSnapshot) +private case class RecordUnionPreparationConf(inner: ColumnarRule, snapshot: UnionConfSnapshot) extends ColumnarRule { private val record = new SnapshotUnionPreparationConf(snapshot) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index 2245697f5ff10..6ce04b5cdf565 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -229,9 +229,10 @@ case class AdaptiveSparkPlanExec( val applied = rule.apply(latestPlan) if (applied ne latestPlan) { // A `UnionExec` this rule just created carries no record of the confs this execution - // answers from, and the `ValidateRequirements` check below and the next rule both read the - // plan before the barrier in `postStageCreationRules` stamps it. A rule that returned its - // input added nothing. + // answers from, and the next rule reads the plan before the barrier in + // `postStageCreationRules` stamps it. So does the `ValidateRequirements` check below, for a + // rule that is itself an `AQEShuffleReadRule`. A rule that returned its input added + // nothing. recordUnionConf(applied) } val result = rule match { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index 8ed2bee76d79d..070d064513121 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -1094,8 +1094,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup * Records the confs this node answers from until its decisions are stamped, read once per * preparation into a `UnionConfSnapshot` and passed in here. `SnapshotUnionPreparationConf` does * it ahead of `EnsureRequirements`, whose reads the following stamp has to agree with, and ahead - * of the two injected lists that can still change an exchange; see that rule for the two windows - * it does not cover. `stampDecisions` does it for a node no such pass saw. Only the confs, never + * of the two injected lists that can still change an exchange. Two more writers cover the rules + * that are folded over rather than listed: each injected columnar rule's transitions carry the + * record behind them, and `optimizeQueryStage` writes it on every rule result that changed the + * plan. `stampDecisions` does it for a node no such pass saw. Only the confs, never * a partitioning: the exchanges `EnsureRequirements` adds are not there yet, so a decision taken * there would freeze plain on a union whose children only become co-partitioned in it. */ @@ -1455,9 +1457,10 @@ object UnionExec { * keeps answering from. See `recordPreparationConf`. `SnapshotUnionPreparationConf` writes it * before `EnsureRequirements`, and the stamp after it reads what that rule wrote, so both phases * use one value; the same rule runs ahead of the two injected lists that can still change an - * exchange, and `stampDecisions` writes it itself for a node no such pass saw. The tag travels - * onto rebuilt nodes the same way `DECISIONS` does, which is what carries it across the copies - * `EnsureRequirements` makes. + * exchange, behind each injected columnar rule's transitions, and on every rule result + * `optimizeQueryStage` sees changed, and `stampDecisions` writes it itself for a node no such + * pass saw. The tag travels onto rebuilt nodes the same way `DECISIONS` does, which is what + * carries it across the copies `EnsureRequirements` makes. */ private val PREPARATION_CONF = TreeNodeTag[UnionConfSnapshot]("unionPreparationConf") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index 958320727431f..c17aedb5c3f43 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -721,10 +721,11 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { } test("SPARK-59122: a union one injected stage optimizer rule adds is recorded before the next") { - // These rules are folded over inside `optimizeQueryStage`, which validates an - // `AQEShuffleReadRule`'s rewrite and then folds on, both before the barrier in - // `postStageCreationRules`. The record is written on each rule's result there, there being no - // list to put a pass into. + // These rules are folded over inside `optimizeQueryStage`, which has no list to put a pass + // into, so the record is written on each rule result that changed the plan. What this case + // reads is the next rule; the `ValidateRequirements` check in the same fold answers from the + // same write, and reaching it takes an injected rule that is itself an `AQEShuffleReadRule`, + // since the built-in ones are listed ahead of the injected list. val seen = ListBuffer.empty[Partitioning] withSession(create { extensions => extensions.injectQueryStageOptimizerRule(_ => WrapRootInUnion) @@ -744,18 +745,28 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { test("SPARK-59122: a union one injected columnar rule adds is recorded before the next") { // The injected columnar rules share one `ApplyColumnarRulesAndInsertTransitions`, so the record - // rides on each rule's own transitions. Those run in reverse order on the way out, so the rule - // that adds the union is the one injected second. + // rides on each rule's own transitions, one override per phase. The post transitions run in + // reverse list order and the pre transitions in list order, so each phase needs the rule that + // adds the union injected at the other end. Seq(false, true).foreach { aqeEnabled => - val seen = ListBuffer.empty[Partitioning] + val post = ListBuffer.empty[Partitioning] checkInjectedUnionIsStamped( create { extensions => - extensions.injectColumnar(_ => ObserveUnionPartitioningColumnarRule(seen)) + extensions.injectColumnar(_ => ObserveUnionPartitioningColumnarRule(post)) extensions.injectColumnar(_ => WrapRootInUnionColumnarRule) }, aqeEnabled) - assert(seen.nonEmpty, s"the second columnar rule must have seen the union, aqe=$aqeEnabled") - assert(!seen.exists(_.isInstanceOf[UnknownPartitioning]), - s"the record must be written behind the rule that added it, aqe=$aqeEnabled: $seen") + val pre = ListBuffer.empty[Partitioning] + checkInjectedUnionIsStamped( + create { extensions => + extensions.injectColumnar(_ => WrapRootInUnionPreColumnarRule) + extensions.injectColumnar(_ => ObserveUnionPartitioningPreColumnarRule(pre)) + }, aqeEnabled) + Seq("post" -> post, "pre" -> pre).foreach { case (phase, seen) => + assert(seen.nonEmpty, + s"the second columnar rule must have seen the union, $phase, aqe=$aqeEnabled") + assert(!seen.exists(_.isInstanceOf[UnknownPartitioning]), + s"the record must ride behind the rule that added it, $phase, aqe=$aqeEnabled: $seen") + } } } @@ -1631,6 +1642,11 @@ object WrapRootInUnionColumnarRule extends ColumnarRule { override def postColumnarTransitions: Rule[SparkPlan] = WrapRootInUnion } +/** `WrapRootInUnion` in the other phase, where the injected rules run in list order. */ +object WrapRootInUnionPreColumnarRule extends ColumnarRule { + override def preColumnarTransitions: Rule[SparkPlan] = WrapRootInUnion +} + /** * Stands for an extension that returns a `UnionExec` of its own in place of one already in the * plan: a fresh instance over the same children, so it carries no stamped decision, and the rows @@ -1676,6 +1692,12 @@ case class ObserveUnionPartitioningColumnarRule(seen: ListBuffer[Partitioning]) override def postColumnarTransitions: Rule[SparkPlan] = ObserveUnionPartitioning(seen) } +/** `ObserveUnionPartitioning` in the other phase. */ +case class ObserveUnionPartitioningPreColumnarRule(seen: ListBuffer[Partitioning]) + extends ColumnarRule { + override def preColumnarTransitions: Rule[SparkPlan] = ObserveUnionPartitioning(seen) +} + /** * Records whether each `UnionExec` says it supports codegen, read with * `WHOLESTAGE_UNION_CODEGEN_ENABLED` turned off: a `true` can then only come from the conf a From 7d091309356cc3ddea8c5b7fc16d1c409d5126ac Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 25 Sep 2026 22:56:27 +0800 Subject: [PATCH 24/24] Cover the validation consumer and name the cached-plan barrier --- .../sql/execution/StampUnionDecisions.scala | 15 +++--- .../sql/SparkSessionExtensionSuite.scala | 46 ++++++++++++++++++- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala index f3d43842d3df5..bd86fe1659644 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala @@ -25,12 +25,15 @@ import org.apache.spark.sql.internal.SQLConf * defined point. * * `UnionExec` derives both from state that moves: its children's `outputPartitioning` sharpens as - * AQE finalises the plans behind them, and `conf` is the live session conf. Every reader used to - * derive its own answer, so the answer depended on when it was read: the codegen gate could fuse a - * union whose copy in the shell then answered the other way, so `metrics` came back empty and - * `doProduce` failed asking `metricTerm` for `numOutputRows`. This rule asks right after - * `EnsureRequirements`, so the decision the exchanges around a union were planned against is the - * one `unionRDDs` and the codegen gate use. + * AQE finalises the plans behind them, and `conf` is the live session conf. That sharpening does + * not run backwards: a cached plan is executed at most once, through `SparkPlan`'s memoized + * `execute` or `executeColumnar`, so a node planned against it cannot see the window where it + * reports no final plan reopen, and a rebuild replaces the relation rather than re-entering that + * plan. Every reader used to derive its own answer, so the answer depended on when it was read: the + * codegen gate could fuse a union whose copy in the shell then answered the other way, so `metrics` + * came back empty and `doProduce` failed asking `metricTerm` for `numOutputRows`. This rule asks + * right after `EnsureRequirements`, so the decision the exchanges around a union were planned + * against is the one `unionRDDs` and the codegen gate use. * * It is listed again after the injected columnar and query-stage rules, the hooks that can add a * `UnionExec` of their own. One created there has no decision yet, and would otherwise take one diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala index c17aedb5c3f43..a54acd00eb77f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -42,10 +42,10 @@ import org.apache.spark.sql.classic.Dataset import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.connector.write.WriterCommitMessage import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, AQEShuffleReadRule, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.HashAggregateExec import org.apache.spark.sql.execution.datasources.{FileFormat, WriteFilesExec, WriteFilesExecBase, WriteFilesSpec} -import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ENSURE_REQUIREMENTS, ShuffleExchangeExec, ShuffleExchangeLike, ShuffleOrigin} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -791,6 +791,29 @@ class SparkSessionExtensionSuite extends PlanTest with AdaptiveSparkPlanHelper { } } + test("SPARK-59122: an injected shuffle read rule's union is recorded before validation") { + // The other consumer inside `optimizeQueryStage`. For an `AQEShuffleReadRule` the fold checks + // the rewrite with `ValidateRequirements` and keeps the plan from before the rule when it + // fails, so a union the rule put under the aggregate is present afterwards only if the check + // saw the recorded conf. The flip lands after the wrapper read its snapshot, so the live value + // says the union concatenates while the record says it passes the partitioning through. + withSession(create(_.injectQueryStageOptimizerRule(_ => WrapAggChildInUnion))) { session => + session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, true) + session.conf.set(SQLConf.UNION_OUTPUT_PARTITIONING.key, true) + val df = session.range(0, 20, 1, 2).selectExpr("id % 5 AS k").groupBy("k").count() + assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) + + session.conf.set(SQLConf.UNION_OUTPUT_PARTITIONING.key, false) + assert(df.collect().map(r => (r.getLong(0), r.getLong(1))).sortBy(_._1).toSeq == + (0L until 5L).map((_, 4L)), "the injected union must not drop or duplicate rows") + + val unions = collect(df.queryExecution.executedPlan) { case u: UnionExec => u } + assert(unions.size == 1, + "the aggregate's requirement must have been judged against the recorded conf, or the " + + s"rewrite was dropped, got\n${df.queryExecution.executedPlan}") + } + } + test("SPARK-59122: the barrier in AQE post stage creation stamps from the adaptive snapshot") { // The cases above leave the conf where it is until their barriers have run, so one answering // from a read of its own rather than the snapshot `AdaptiveSparkPlanExec` took at construction @@ -1637,6 +1660,25 @@ object WrapRootInUnion extends Rule[SparkPlan] { } } +/** + * Stands for an injected shuffle-read optimizer that builds a parent over a `UnionExec` of its own: + * a fresh one-child union under the final aggregate, which leaves the rows alone and makes + * `ValidateRequirements` judge the aggregate's clustering requirement against what that union + * reports. `optimizeQueryStage` drops the whole rewrite when that check fails, so the union + * survives only if it answered from the recorded conf. Matching `Final` mode keeps it off the + * partial aggregate inside the shuffle stage, whose requirement is unspecified either way. + */ +object WrapAggChildInUnion extends Rule[SparkPlan] with AQEShuffleReadRule { + override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] = Seq(ENSURE_REQUIREMENTS) + + override def apply(plan: SparkPlan): SparkPlan = plan.transformUp { + case agg: HashAggregateExec + if agg.aggregateExpressions.exists(_.mode == Final) && + !agg.child.isInstanceOf[UnionExec] => + agg.withNewChildren(Seq(UnionExec(Seq(agg.child)))) + } +} + /** The columnar-rule wrapper for `WrapRootInUnion`. */ object WrapRootInUnionColumnarRule extends ColumnarRule { override def postColumnarTransitions: Rule[SparkPlan] = WrapRootInUnion