LuciferYang commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r4001335059
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1035,13 +1039,79 @@ 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 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[execution] def isPlainUnion: Boolean =
stampedDecisions.map(_.plainUnion).getOrElse {
+ !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) ||
Review Comment:
Real, and the window is one rule wide. Closed in `65e3f96ee6c`: a new
`SnapshotUnionOutputPartitioningConf` runs before `EnsureRequirements` in both
pipelines and records `spark.sql.unionOutputPartitioning` on each unstamped
`UnionExec`; `isPlainUnion` answers from that record until the decision is
stamped, and only a node created after that pass reads the conf live. The rule
reads the conf once for the whole plan, so every union in one plan answers from
the same value.
Only the conf, never a partitioning: `EnsureRequirements` has not inserted
its exchanges yet, so a decision taken there would freeze plain on a union
whose children only become co-partitioned after it. That is your first Risk,
and it is why the decision itself still waits for the barrier behind
`EnsureRequirements`.
Coverage is `UnionCodegenSuite`'s "the stamp uses the conf the exchanges
were planned against". It drives the three rules in order and flips the conf
between the second and the third, since the two sit next to each other in the
pipeline and no injected rule can run in the window. It asserts both halves:
the aggregate's exchange was elided (both remaining exchanges are
`REPARTITION_BY_NUM`), and the union still reports a concrete partitioning
after the flip. Reverting the snapshot read to a live one fails it with
`UnknownPartitioning(0)`.
One thing I checked before agreeing: whether `EnsureRequirements` has the
same hazard internally, since it reads `conf.numShufflePartitions` per child at
two sites. It does not, because it reconciles the children against each other's
actual partitioning afterwards rather than against the conf.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1044,11 +1091,21 @@ 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] = {
+ // Latched for the same reason `isPlainUnion` is: `supportCodegen` and
`metrics` must see one
+ // answer, and `conf` is live. When a child is not `CodegenSupport`,
`insertInputAdapter` wraps
+ // it, so `withNewChildren` returns a real copy whose first evaluation of
this would land at
+ // execution; re-deriving there left `metrics` empty while `doProduce` asked
`metricTerm` for
+ // `numOutputRows`. The first force is not always the gate: under AQE it is
a plan-update event
+ // on the pre-stage-creation tree, so a term added here sees more of the
plan than the gate does.
+ private def supportCodegenFailureReason: Option[String] =
decisionLock.synchronized {
+ getTagValue(UnionExec.CODEGEN_FAILURE_REASON).getOrElse {
Review Comment:
Done in `65e3f96ee6c`. The case now injects a stage-optimizer rule alongside
the prep rule; it reads the union with `spark.sql.unionOutputPartitioning`
turned off and puts the conf back, so a concrete answer can only come from a
decision stamped before the stage optimizers ran, and the union an injected
prep rule adds carries no recorded conf of its own to answer from.
Removing the barrier after the injected prep rules now fails that case with
`ListBuffer(UnknownPartitioning(0))`, and it is the only case that fails, so
each of the three listings has its own signal.
--
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]
