LuciferYang commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r3939227427
##########
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
{
Review Comment:
Done in `d60079ebca4`. It reads much better as one block, and `prepare()`
was the right precedent to point me at.
##########
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 (`825c1205952`, `36c62b8f542`) 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. A later
pass (`825c1205952`) also pinned `spark.sql.unionOutputPartitioning` inside it,
so it cannot pass vacuously if that default ever flips.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/UnionCodegenSuite.scala:
##########
@@ -58,6 +61,29 @@ 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.
+ */
+ private def fusedUnions(df: DataFrame): Seq[UnionExec] =
+ collect(df.queryExecution.executedPlan) {
+ case w: WholeStageCodegenExec if w.child.isInstanceOf[UnionExec] =>
+ w.child.asInstanceOf[UnionExec]
+ }
+
+ /** A cached aggregate, so the union's children read an
`InMemoryTableScanExec`. */
+ private def cacheAggregateView(view: String): Unit = {
+ spark.range(0, 200, 1, 4)
+ .selectExpr("id % 10 AS k", "id AS v")
+ .groupBy("k")
+ .agg(sum("v").as("s"))
+ .createOrReplaceTempView(view)
+ // Both callers need the cache unmaterialized, and `CacheManager` no-ops
on an already-cached
+ // plan, so drop whatever an earlier test left for this one. `isCached`
matches by plan.
+ if (spark.catalog.isCached(view)) spark.catalog.uncacheTable(view)
Review Comment:
You are right, and it is gone in `9da6b20394e`. The helper's doc comment now
just says what `withTempView` does instead of promising more than
`dropTempView` actually gives.
##########
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") {
+ // 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
`HashPartitioning`, and
+ // re-deriving the decision at that point left `metrics` empty while the
generated code still
+ // incremented it, so `doProduce` threw `key not found: numOutputRows`.
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.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")
+ assert(fused.forall(_.metrics.contains("numOutputRows")),
+ "a fused union must register the metric its generated code
increments")
+ }
+ }
+ }
+
+ test("SPARK-59122: a fused union reports UnknownPartitioning, so no parent
skips an exchange") {
+ // The other half of the same decision. A fused union concatenates its
children's partitions,
+ // so if it went on claiming the children's `HashPartitioning` a parent
could satisfy a
+ // clustered distribution from an RDD that does not have it -- a wrong
answer rather than a
+ // crash, which is why the missing metric must not simply be registered
unconditionally.
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.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")
+ df.collect()
+ val fused = fusedUnions(df)
+ assert(fused.nonEmpty, "this shape must actually fuse")
+ fused.foreach { u =>
+ 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 where the plain-union
decision is latched, 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
`collect()` stays outside
+ // the block that planned the DataFrame on purpose: `executedPlan` is
memoized on first read,
+ // and moving it back inside makes both phases see the same conf.
+ 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")
+ def build(): DataFrame =
+ left.repartition(4, col("k")).union(right.repartition(4,
col("k"))).groupBy("k").count()
+
+ val expected = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key ->
"false") {
Review Comment:
Done in `9da6b20394e`: it asserts `(0L until 5L).map(k => Row(k, 8L))` now
and skips the second execution. Your point about both paths regressing together
is exactly why, and I made the same change to the SPARK-59141 test for the same
reason.
--
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]
