cloud-fan commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r4068257893
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -163,7 +177,10 @@ case class AdaptiveSparkPlanExec(
// channel, opt-in). Runs last so skew handling and sort cleanup have
settled
// before placement is decided.
AQEEnablePipelinedShuffle
- ) ++ context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules
+ ) ++ context.session.sessionState.adaptiveRulesHolder.queryStagePrepRules
:+
Review Comment:
**Non-blocking (P2):** This stamp runs only after the complete injected
prep-rule list. If one rule creates a fresh co-partitioned `UnionExec` and the
next applies requirements planning, that consumer still reads live
`UNION_OUTPUT_PARTITIONING`; the suffix can then stamp the older preparation
value. The parent may omit its exchange from the concrete provisional answer,
while execution later concatenates under the stamped plain decision, so the
same aggregate key can be emitted from multiple partitions. The preparation
snapshot needs to be available before the next supported rule can consume the
new union.
**Recommended change:** Carry the full preparation UnionConfSnapshot as
provisional UnionExec state and refresh that provisional state after each
injected physical or columnar rule, before the next rule runs; retain the
existing phase-end completed-decision barriers and write-once semantics.
**Why this works:** Store or propagate the immutable three-field
UnionConfSnapshot separately from the completed child-derived decision. After
each supported external rule invocation, traverse only to fill missing
provisional snapshot state on newly created unions. Later rules then read the
preparation value for partitioning and codegen gates, while the existing suffix
stamp combines that same snapshot with the topology present at its defined
barrier.
**Scope:** Make one preparation snapshot visible throughout ordered
extension-rule composition without freezing child-derived partitioning before
the established completed-decision points.
**Compatibility:** Preserve live, non-mutating answers on genuinely
unprepared plans; write-once completed decisions; dynamic AQE partition counts;
and the documented optimization tradeoffs for already prepared unions.
**Risks:** Missing either the forward or reverse ColumnarRule fold would
leave a same-list gap. Propagating a completed decision instead of only
configuration could freeze topology too early and lose valid exchange
elimination. A traversal inserted after rather than between external rules
would reproduce the current defect.
**Constraints:** Do not derive or retain a Partitioning object before
EnsureRequirements; AQE partition counts and child-derived raw partitioning
must remain dynamic. Unprepared plan inspection remains non-mutating and may
reflect live configuration. Already completed decisions remain write-once and
survive intended tagless copies. Do not add preparation state to case-class
product identity, canonicalization, or explain output.
**Success:** A fresh UnionExec created by one injected rule exposes the
preparation snapshot to every later rule in that ordered extension phase.
Requirements planning and eventual union execution cannot observe opposing
UNION_OUTPUT_PARTITIONING values for the same fresh node. Pre-stamp codegen
inspection and eventual stamped codegen decisions use the same codegen-enabled
and maxChildren values. Child-derived partitioning is still decided only at the
established post-requirements or post-extension barrier.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -2935,10 +2935,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
when a UnionExec's " +
+ "decision is fixed during physical preparation, so a change does not
reach a " +
Review Comment:
**Nit (P3):** The value is not read when each union's decision is fixed: it
is captured once before the physical preparation/adaptive rule sequence and
reused by later barriers. Please describe that preparation-wide capture point
here and on `wholeStage.union.enabled`, so a mid-preparation change is not
documented as affecting a later-created union.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:
##########
@@ -53,16 +58,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 the callers here reach for its tags and
metrics.
Review Comment:
**Nit (P3):** `the node the callers here reach for its tags and metrics` is
ungrammatical; this can say `the node whose tags and metrics the callers here
inspect`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:
##########
@@ -628,6 +706,486 @@ 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
Review Comment:
**Nit (P3):** This should read `Planned with the conf on, the union is
fused`; the comma separates the planning condition from the subject.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -119,10 +119,16 @@ 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)
Review Comment:
**Non-blocking (P2):** The AQE configuration-flip case exercises a union
already stamped while `initialPlan` is built, and the extension cases keep the
initial value stable until their late barriers have run. Replacing a later
barrier's stored `unionConf` with a fresh SQLConf read would therefore leave
those assertions green. Please add exchange-backed AQE cases that flip output
partitioning, codegen enablement, and `maxChildren` after this snapshot, then
create a fresh union at a later stage/extension barrier and assert its layout,
fusion, shell copy, and `numOutputRows`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
