LuciferYang commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r3939227491
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1023,13 +1026,53 @@ 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 `doExecute`.
- // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in
`doExecute`, 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]
+ // Serializes the latch below so concurrent first readers agree on one
answer. Not this node's own
+ // monitor, which `unionedInputRDD`'s `lazy val` holds while it drives
`child.execute()`.
+ // Driver-only, hence `@transient`.
+ @transient private val decisionLock = new Object()
+
+ /**
+ * True when this union behaves as a plain concatenation, so
`unionedInputRDD` matches
+ * `sparkContext.union(...)` in `doExecute` and the codegen path applies. A
`KeyedPartitioning`
+ * union also concatenates, but codegen stays off for it:
`supportCodegenFailureReason` reports
+ * "partitioning-aware", because a downstream `GroupPartitionsExec` consumes
its key descriptor.
+ *
+ * Latched, 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-derives comes back
+ * with empty `metrics` while `doProduce` asks `metricTerm` for
`numOutputRows`. A `TreeNodeTag`
+ * survives that rebuild where a field would not, since `withNewChildren`
ends in `copyTagsFrom`.
+ *
+ * `UNION_OUTPUT_PARTITIONING` is read here rather than in `rawPartitioning`
so it is latched too:
+ * `conf` is live, and re-reading it let a plan made with the conf on
execute with it off.
+ */
+ private[sql] def isPlainUnion: Boolean = {
+
decisionLock.synchronized(getTagValue(UnionExec.PLAIN_UNION_DECISION)).getOrElse
{
+ val plain = !conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING) ||
+ rawPartitioning.isInstanceOf[UnknownPartitioning]
+ decisionLock.synchronized {
+ getTagValue(UnionExec.PLAIN_UNION_DECISION).getOrElse {
+ setTagValue(UnionExec.PLAIN_UNION_DECISION, plain)
+ plain
+ }
+ }
+ }
+ }
+
+ /**
+ * A node latched 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.
+ *
+ * The other branch is derived per call and can come back
`UnknownPartitioning` later -- AQE skew
+ * splitting through a union leaves the children's partition counts
divergent. Failing there was
+ * tried and reverted: nothing in those plans required the reported
partitioning, and this node
Review Comment:
Good catch, and done in `d60079ebca4`. Two later passes tightened the same
paragraph again after I found it was still claiming more than the code shows:
it now states the mechanism and names what reconciles a change, with no
tried-and-reverted history left in it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:
##########
@@ -628,6 +654,89 @@ class UnionCodegenSuite extends SharedSparkSession {
}
}
+ test("SPARK-59122: a fused union keeps numOutputRows when a child's
partitioning firms up") {
Review Comment:
Agreed, merged in `9da6b20394e`. One test asserts both halves now, and a
later pass also pinned `spark.sql.unionOutputPartitioning` inside it, so it
cannot pass vacuously if that default ever flips.
--
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]
