This is an automated email from the ASF dual-hosted git repository.
LuciferYang pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 13e88a007e43 [SPARK-59122][4.3][SQL] Take UnionExec's plain-union
decision once instead of re-deriving it per caller
13e88a007e43 is described below
commit 13e88a007e43867fdc95da334ea9f8f75073caf9
Author: YangJie <[email protected]>
AuthorDate: Fri Oct 2 21:54:40 2026 +0800
[SPARK-59122][4.3][SQL] Take UnionExec's plain-union decision once instead
of re-deriving it per caller
### What changes were proposed in this pull request?
This backports SPARK-59122 (#58419, already in master and branch-4.x) to
branch-4.3. The change makes `UnionExec` decide once, right after
`EnsureRequirements`, whether it is a plain concatenation, and freeze that
decision plus the codegen confs in a `TreeNodeTag`, so the fusion gate,
`metrics`, `unionRDDs` and the copy of the union inside the codegen shell can
no longer give different answers depending on when each is read. See #58419 for
the full design and the two failures it fixes.
Two conflicts came up against branch-4.3, both because the branch predates
code the master change was built on. Neither changes the fix:
- branch-4.3 does not have SPARK-57399, so `AQEEnablePipelinedShuffle` is
not in the AQE query-stage preparation rule list. The two union barriers are
appended after `RemoveRedundantSorts` exactly as on master, without that rule.
- branch-4.3 does not have SPARK-58819 or SPARK-59285, so `UnionExec` still
derives its candidate partitioning by remapping to the first child's attributes
and lifting with a local `toUnionOutput`, and `comparePartitioning` does not
compare `keyDataTypes`. The backport keeps that body. The SPARK-59122 edit to
it is only the rename of the old `outputPartitioning` body to a private
`rawPartitioning`, moving the `spark.sql.unionOutputPartitioning` gate out of
an early return and into the [...]
### Why are the changes needed?
SPARK-59122 is a Blocker correctness fix. On default settings a union whose
children share a partitioning could concatenate at execution while a parent had
already dropped the exchange it needed, giving wrong results; a fused union
could also lose the `numOutputRows` metric its generated code asks for and
crash. branch-4.3 is exposed to both, since `spark.sql.unionOutputPartitioning`
defaults to true there.
### Does this PR introduce _any_ user-facing change?
Yes, the same change as #58419: the wrong-result and codegen-crash cases
above now behave correctly. The three confs the decision reads no longer reach
a plan that has already been prepared, only plans prepared afterwards.
### How was this patch tested?
The tests come with the cherry-pick, unchanged from #58419; CI runs them
against branch-4.3. The AQE rule-order case asserts positions relative to
`EnsureRequirements` and the end of the list, so it holds with
`AQEEnablePipelinedShuffle` absent without editing.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5
Closes #59151 from LuciferYang/SPARK-59122-4.3.
Authored-by: YangJie <[email protected]>
Signed-off-by: yangjie01 <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 19 +-
.../spark/sql/execution/QueryExecution.scala | 17 +-
.../spark/sql/execution/StampUnionDecisions.scala | 165 +++++
.../execution/adaptive/AdaptiveSparkPlanExec.scala | 47 +-
.../sql/execution/basicPhysicalOperators.scala | 240 ++++++-
.../spark/sql/SparkSessionExtensionSuite.scala | 379 ++++++++++-
.../spark/sql/execution/UnionCodegenSuite.scala | 758 +++++++++++++++++++--
.../adaptive/AdaptiveQueryExecSuite.scala | 41 +-
8 files changed, 1558 insertions(+), 108 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 04f778aca33c..2a42778e283e 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
@@ -2865,10 +2865,12 @@ 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.")
+ .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
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
@@ -2883,7 +2885,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
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
@@ -7928,7 +7932,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.")
+ "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/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index 586cb5226173..aea928107900 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
@@ -772,6 +772,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 ++
@@ -780,7 +783,14 @@ 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.
+ 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.
+ 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
@@ -805,7 +815,12 @@ 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.
+ 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
new file mode 100644
index 000000000000..bd86fe165964
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/StampUnionDecisions.scala
@@ -0,0 +1,165 @@
+/*
+ * 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
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * 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.
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
+ * 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 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`.
+ */
+class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan]
{
+ override def apply(plan: SparkPlan): SparkPlan = {
+ plan.foreach {
+ case u: UnionExec => u.stampDecisions(snapshot)
+ case _ =>
+ }
+ plan
+ }
+}
+
+/**
+ * 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, 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.
+ *
+ * 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, 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 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
+ * become co-partitioned there, which is why the decision itself waits for the
barrier behind it.
+ *
+ * Writing the tag in place is safe for the reason given on
[[StampUnionDecisions]].
+ */
+class SnapshotUnionPreparationConf(snapshot: UnionConfSnapshot) extends
Rule[SparkPlan] {
+ override def apply(plan: SparkPlan): SparkPlan = {
+ plan.foreach {
+ 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))
+
+ /**
+ * `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(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 case 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))
+}
+
+/**
+ * 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 bfe6a9a3f633..30af829c324b 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
@@ -107,10 +107,20 @@ 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)
+
+ // 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.
- @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
@@ -122,7 +132,15 @@ 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.
+ 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
+ // planned against.
+ new StampUnionDecisions(unionConf),
// This rule must be run after `EnsureRequirements`.
InsertSortForLimitAndOffset,
AdjustShuffleExchangePosition,
@@ -147,7 +165,11 @@ 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
+ ) ++ 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)
}
// A list of physical optimizer rules to be applied to a new stage before
its execution. These
@@ -171,7 +193,12 @@ 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.
+ new StampUnionDecisions(unionConf),
collapseCodegenStagesRule
)
@@ -184,6 +211,14 @@ 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 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 {
case _: AQEShuffleReadRule if !applied.fastEquals(latestPlan) =>
val distribution = if (isFinalStage) {
@@ -213,7 +248,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 9a417876b830..9677243b0770 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
@@ -35,6 +35,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}
@@ -923,7 +924,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
@@ -931,7 +932,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
// For `KeyedPartitioning`, only the partition expressions must match
(the other child's
// expressions have already been remapped to the first child's
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) }
@@ -941,11 +942,15 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
}
- override def outputPartitioning: Partitioning = {
- if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) {
- return super.outputPartitioning
- }
-
+ /**
+ * 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 `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.
+ */
+ private def rawPartitioning: Partitioning = {
// Children's partitionings with attributes remapped to the first child's
attributes.
val partitionings = prepareOutputPartitioning()
// Map from the first child's attributes to this union's own output
attributes.
@@ -1013,13 +1018,119 @@ 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 `unionRDDs`.
- // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in
`unionRDDs`, 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 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. 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.
+ *
+ * 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
+ * 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 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
+ * 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 = 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)
+
+ private def preparationConf: Option[UnionConfSnapshot] =
+ getTagValue(UnionExec.PREPARATION_CONF)
+
+ private def outputPartitioningEnabled: Boolean =
+ preparationConf.map(_.outputPartitioning)
+ .getOrElse(conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING))
+
+ /**
+ * 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. 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.
+ */
+ private[execution] def recordPreparationConf(snapshot: UnionConfSnapshot):
Unit =
+ if (preparationConf.isEmpty) {
+ setTagValue(UnionExec.PREPARATION_CONF, snapshot)
+ }
+
+ /**
+ * 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 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(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.
+ recordPreparationConf(snapshot)
+ setTagValue(UnionExec.DECISIONS, UnionExec.Decisions(plainUnion =
isPlainUnion))
+ }
+
+ /**
+ * 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.
+ *
+ * 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, whereas re-deriving at the gate would
have fused it.
+ * `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.
+ *
+ * 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 = {
+ // 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
@@ -1037,23 +1148,37 @@ 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.
- @transient private lazy val supportCodegenFailureReason: Option[String] = {
- if (!conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED)) {
- Some("union-codegen-disabled")
- } else if (!isPlainUnion) {
- Some("partitioning-aware")
- } else if (children.exists(_.exists(_.isInstanceOf[UnionExec]))) {
+ // 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 =
+ preparationConf.map(_.codegenEnabled)
+ .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED))
+
+ private def maxCodegenChildren: Int =
+ 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. 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")
} 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 >
conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) {
- Some("max-children-exceeded")
} else if (supportsColumnar) {
Some("columnar")
} else if (children.exists(c =>
@@ -1064,6 +1189,26 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
}
+ // 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")
+ } 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
@@ -1076,11 +1221,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 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 (supportCodegenFailureReason.isEmpty) {
+ if (stampedDecisions.isEmpty || supportCodegenFailureReason.isEmpty) {
Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of
output rows"))
} else {
Map.empty
@@ -1262,6 +1411,37 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
object UnionExec {
+ /**
+ * 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)
+
+ /**
+ * 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. 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 [[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, 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")
+
/**
* 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 05619efa285f..a54acd00eb77 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._
@@ -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
@@ -619,6 +619,265 @@ 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
+ * 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 = {
+ 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") {
+ // 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 { 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") {
+ // 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("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: 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: a union one injected stage optimizer rule adds is
recorded before the next") {
+ // 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)
+ 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, 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 post = ListBuffer.empty[Partitioning]
+ checkInjectedUnionIsStamped(
+ create { extensions =>
+ extensions.injectColumnar(_ =>
ObserveUnionPartitioningColumnarRule(post))
+ extensions.injectColumnar(_ => WrapRootInUnionColumnarRule)
+ }, aqeEnabled)
+ 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")
+ }
+ }
+ }
+
+ 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("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
+ // 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.
@@ -1387,6 +1646,120 @@ 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
+ }
+}
+
+/**
+ * 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
+}
+
+/** `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
+ * 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
+ * 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 = {
+ 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
+ }
+}
+
+/** The columnar-rule wrapper for `ObserveUnionPartitioning`. */
+case class ObserveUnionPartitioningColumnarRule(seen: ListBuffer[Partitioning])
+ extends ColumnarRule {
+ 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
+ * 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 6af286b610e2..4636a3ae0955 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,12 @@ 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
+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
@@ -30,9 +36,10 @@ 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.
+ // 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")
@@ -53,16 +60,86 @@ class UnionCodegenSuite extends SharedSparkSession {
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 `codegenUnions` on purpose: this matches only a union that
is the root of its own
+ * 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) {
+ case w: WholeStageCodegenExec if w.child.isInstanceOf[UnionExec] =>
+ w.child.asInstanceOf[UnionExec]
+ }
- /** Run query with flag on, then flag off, assert results match. */
- protected def assertFlagParity(buildDf: () => DataFrame): Unit = {
- val onRows = buildDf().collect().toSeq
+ /**
+ * 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`
+ * uncaches this view's plan.
+ */
+ 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)
+ 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 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.
+ *
+ * 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)
}
}
@@ -86,14 +163,14 @@ class UnionCodegenSuite extends SharedSparkSession {
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)
}
}
@@ -102,7 +179,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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)
}
}
@@ -200,7 +277,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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"))
}
@@ -223,7 +300,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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)))
}
@@ -244,7 +321,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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
@@ -267,7 +344,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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
@@ -401,16 +478,15 @@ class UnionCodegenSuite extends SharedSparkSession {
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}")
}
}
@@ -418,17 +494,19 @@ class UnionCodegenSuite extends SharedSparkSession {
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}")
}
}
@@ -442,14 +520,16 @@ class UnionCodegenSuite extends SharedSparkSession {
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")))
}
+ // `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}")
}
}
@@ -480,11 +560,12 @@ class UnionCodegenSuite extends SharedSparkSession {
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()
}
@@ -538,7 +619,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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)
}
}
@@ -565,7 +646,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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"))
@@ -575,19 +656,19 @@ class UnionCodegenSuite extends SharedSparkSession {
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") {
@@ -608,16 +689,15 @@ class UnionCodegenSuite extends SharedSparkSession {
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
@@ -628,6 +708,558 @@ class UnionCodegenSuite extends SharedSparkSession {
}
}
+ 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
+ // 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
concrete layout, and
+ // re-deriving the decision at that point answered "partitioning-aware",
so `metrics` came back
+ // empty while `doProduce` asked `metricTerm` for `numOutputRows`.
+ //
+ // 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(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.WHOLESTAGE_UNION_CODEGEN_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")
+ // 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")
+ 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("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}")
+ }
+ }
+ }
+ }
+
+ test("SPARK-59122: a partitioning-aware union keeps its layout when the conf
changes between " +
+ "planning and execution") {
+ // `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")
+
+ val planned = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key ->
"true") {
+ 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)
+ // 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),
+ s"expected only the two repartitions, got
${shuffles.map(_.shuffleOrigin)}")
+ df
+ }
+
+ withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") {
+ // 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)))
+ }
+ }
+ }
+
+ 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 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, 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
+ }
+ 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. 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. 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")
+ }
+ }
+ }
+
+ 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.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
+ // 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.
+ // `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") {
+ 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
+ // 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")
+ // 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)
+ }
+ }
+
+ 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 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 barrier's snapshot, not from the
read taken with the " +
+ s"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
+ // `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 -> "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}")
+ }
+ }
+
+ 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, 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 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(_.isInstanceOf[StampUnionDecisions])
+ val ensureRequirements =
rules.indexWhere(_.isInstanceOf[EnsureRequirements])
+ val columnarRules =
+
rules.indexWhere(_.isInstanceOf[ApplyColumnarRulesAndInsertTransitions])
+ assert(ensureRequirements >= 0 && firstStamp == ensureRequirements + 1 &&
+ 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"))
+ .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 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)
+ 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-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. `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(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ // 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(_.isInstanceOf[SnapshotUnionPreparationConf])
+ val ensureRequirements =
rules.indexWhere(_.isInstanceOf[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")))
+ .groupBy("k").count()
+ val required = EnsureRequirements()(
+ 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")
+
+ 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-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()(
+ 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
+ // 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-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,
+ 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 record must hold $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-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")
+ }
+ }
+
+ // 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)") {
// `InputFileName` is `Nondeterministic` but reads from
`InputFileBlockHolder`
// (a per-task thread-local) and does not embed `partitionIndex`. The
gate's
@@ -638,7 +1270,7 @@ class UnionCodegenSuite extends SharedSparkSession {
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"))
}
@@ -667,7 +1299,7 @@ class UnionCodegenSuite extends SharedSparkSession {
// 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 9523ce897ecf..38832ad3723e 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
@@ -43,7 +43,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}
@@ -4241,6 +4241,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. `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
+ // 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(_.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
+ // stand in for the barrier behind `EnsureRequirements`.
+ val stamps = rules.zipWithIndex.collect {
+ case (rule, i) if rule.isInstanceOf[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",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
