cloud-fan commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r3996975958
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1035,13 +1038,67 @@ 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, and a copy that
re-derived here came
+ * back with empty `metrics` while `doProduce` asked `metricTerm` for
`numOutputRows`. A fresh
+ * copy inherits the answer instead, since `withNewChildren` ends in
`copyTagsFrom`.
+ *
+ * `UNION_OUTPUT_PARTITIONING` is read where the decision is stamped rather
than in
+ * `rawPartitioning`, so it too is fixed once the plan is prepared: `conf`
is live, and a plan
+ * must execute by the partitioning it was planned against.
+ *
+ * A read before `StampUnionDecisions` answers from the children as they are
then, and does not
+ * write, so observing an unprepared plan cannot decide anything for the
prepared one.
+ */
+ private[sql] def isPlainUnion: Boolean =
stampedDecisions.map(_.plainUnion).getOrElse {
+ !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) ||
+ rawPartitioning.isInstanceOf[UnknownPartitioning]
+ }
+
+ private def stampedDecisions: Option[UnionExec.Decisions] =
+ getTagValue(UnionExec.DECISIONS)
+
+ /**
+ * Fixes this node's decisions for the rest of the plan's life. Called by
`StampUnionDecisions`
+ * right after `EnsureRequirements`, so what the exchanges around this union
were planned against
+ * is what execution uses. Nothing else writes this tag on an existing node,
and the nodes the
+ * rule writes are freshly planned and not yet published, so no reader can
be looking at one;
+ * `metrics` and the codegen gate read it later, and a node that already
carries it keeps it,
+ * which is how the copy in the codegen shell stays in step with the gate.
+ */
+ private[execution] def stampDecisions(): Unit = if
(stampedDecisions.isEmpty) {
+ setTagValue(UnionExec.DECISIONS, UnionExec.Decisions(
+ plainUnion = isPlainUnion,
+ unionCodegenEnabled =
conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED),
+ maxChildren = conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)))
Review Comment:
**Non-blocking (P2):** `maxChildren` is now part of the prepared-plan
snapshot, but no test distinguishes this from the old live read. The existing
cap tests set it before planning, and the same-plan flip covers only
`union.enabled`. Please add the exchange-backed equivalent: prepare a two-child
union at cap 2, lower the cap before `collect()`, and assert the shell copy
still supports codegen and has `numOutputRows`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1059,12 +1116,31 @@ 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.
+ // The confs the gate reads, stamped for the reason the plain-union decision
is: `conf` is live,
+ // so the gate, `metrics` and the copy `insertInputAdapter` puts inside the
codegen shell would
+ // otherwise be free to read different values. When a child is not
`CodegenSupport` that copy is
+ // real and its first evaluation lands at execution; reading the conf there
left `metrics` empty
+ // while `doProduce` asked `metricTerm` for `numOutputRows`. A read before
the stamp answers from
+ // the conf as it is then and writes nothing, so observing an unprepared
plan cannot pin this
+ // either.
+ private def unionCodegenEnabled: Boolean =
+ stampedDecisions.map(_.unionCodegenEnabled)
+ .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED))
+
+ private def maxCodegenChildren: Int =
+ stampedDecisions.map(_.maxChildren)
+ .getOrElse(conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN))
+
+ // Memoized per instance rather than stamped on the tag. Every term below
the confs except
+ // `isPlainUnion` reads the children, and a tag outlives them:
`SparkPlanInfo` forces `metrics` on
Review Comment:
**Nit (P3):** This event order is reversed: `newQueryStage` applies
`optimizeQueryStage` and `postStageCreationRules` (including
`CollapseCodegenStages`) before `withFinalPlanUpdate` calls `onUpdatePlan`. The
early `metrics` read comes from `SQLExecution` building the initial
`SparkPlanInfo` before execution. Please use that trigger to justify keeping
the child-derived gates per instance.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:
##########
@@ -58,6 +61,34 @@ class UnionCodegenSuite extends SharedSparkSession {
case w: WholeStageCodegenExec if
w.find(_.isInstanceOf[UnionExec]).isDefined => w
}.nonEmpty
+ /**
+ * `AdaptiveSparkPlanHelper.collect` descends through AQE wrappers and query
stages;
+ * `SparkPlan.collect` stops at them, since both are `LeafExecNode`s.
+ *
+ * Stricter than `unionInsideWSCG` on purpose: this matches only a union
that is the root of its
+ * own codegen stage, which is what "fused" means for the callers here,
while `w.find` also
+ * matches one an `InputAdapter` left inside the stage.
Review Comment:
**Nit (P3):** This should read `matches one that an InputAdapter left inside
the stage`; the relative clause is currently missing `that`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -976,11 +977,13 @@ 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 pass-through partitioning, derived from the children.
`isPlainUnion` answers on
Review Comment:
**Nit (P3):** `answers on whether` is not grammatical here. Please say that
`isPlainUnion` answers whether `rawPartitioning` is `UnknownPartitioning`, or
that its answer is based on whether this method returns unknown.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -2908,7 +2908,9 @@ object SQLConf {
.doc("When both this conf and `spark.sql.codegen.wholeStage` are true, "
+
"UnionExec participates in whole-stage codegen on its " +
"non-partitioning-aware path: the parent and all children fuse into " +
- "a single WholeStageCodegenExec stage.")
+ "a single WholeStageCodegenExec stage. The value is read when a
UnionExec " +
Review Comment:
**Non-blocking (P2):** The new boundary is physical preparation, not plan
construction. A `DataFrame`/`sparkPlan` can be built under one value, changed
before `executedPlan` is first requested, and then stamped with the new value.
Please update this and the other two captured union-conf descriptions to say
that changes affect plans *prepared* afterward, not plans built afterward.
--
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]
