LuciferYang commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r4038309976
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1035,13 +1040,96 @@ 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
`snapshotOutputPartitioningConf`, recorded before
+ * `EnsureRequirements`, so the value the exchanges are planned against is
the value execution
+ * uses; a node created after that pass carries no record and reads the live
conf. Reading it live
+ * here would leave one rule between the two: `conf` is live, and another
thread setting it in
+ * that window would let a parent drop an exchange over a concrete
partitioning and then have the
+ * stamp freeze plain concatenation under it.
+ *
+ * A read before `StampUnionDecisions` answers from the children as they are
then, and does not
+ * write, so observing an unprepared plan cannot decide anything for the
prepared one.
+ */
+ private[execution] def isPlainUnion: Boolean =
stampedDecisions.map(_.plainUnion).getOrElse {
+ !outputPartitioningEnabled ||
rawPartitioning.isInstanceOf[UnknownPartitioning]
+ }
+
+ private def stampedDecisions: Option[UnionExec.Decisions] =
+ getTagValue(UnionExec.DECISIONS)
+
+ private def outputPartitioningEnabled: Boolean =
+ getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF)
+ .getOrElse(conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING))
+
+ /**
+ * Records the conf `isPlainUnion` answers from, read once for the whole
plan by
+ * `SnapshotUnionOutputPartitioningConf` and passed in here, ahead of
`EnsureRequirements`, whose
+ * reads the following stamp has to agree with. Only the conf, never a
partitioning: the exchanges
+ * `EnsureRequirements` adds are not there yet, so a decision taken here
would freeze plain on a
+ * union whose children only become co-partitioned there.
+ */
+ private[execution] def snapshotOutputPartitioningConf(enabled: Boolean):
Unit =
+ if (getTagValue(UnionExec.OUTPUT_PARTITIONING_CONF).isEmpty) {
+ setTagValue(UnionExec.OUTPUT_PARTITIONING_CONF, enabled)
+ }
+
+ /**
+ * Fixes this node's decisions for the rest of the plan's life. Called by
`StampUnionDecisions`,
+ * first right after `EnsureRequirements`, so what the exchanges around this
union were planned
+ * against is what execution uses; the two confs come from one read per plan
there. Nothing else
+ * writes this tag on an existing node, and the nodes the rule writes are
freshly planned and not
+ * yet published, so no reader can be looking at one; `metrics` and the
codegen gate read it
+ * later, and a node that already carries it keeps it, which is how the copy
in the codegen shell
+ * stays in step with the gate.
+ */
+ private[execution] def stampDecisions(codegenEnabled: Boolean, maxChildren:
Int): Unit =
+ if (stampedDecisions.isEmpty) {
+ setTagValue(UnionExec.DECISIONS, UnionExec.Decisions(
+ plainUnion = isPlainUnion,
+ unionCodegenEnabled = codegenEnabled,
+ maxChildren = maxChildren))
+ }
+
+ /**
+ * A node stamped plain reports `UnknownPartitioning` even once its children
agree on a concrete
+ * 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 =
+ if (isPlainUnion) super.outputPartitioning else rawPartitioning
Review Comment:
Done. `outputPartitioning` holds a method-local `lazy val`, and the decision
takes the raw partitioning by name, so a stamped node still never derives it
and nothing is cached across calls.
--
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]
