LuciferYang commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r3931513453
##########
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 = {
Review Comment:
Keeping the latch, and I would rather not have the consumer read
`rawPartitioning`.
The reason is narrower than "the latch is better". `unionRDDs` branches on
the latched decision through `outputPartitioning`, so if
`childrenNeedCompatiblePartitioning` keyed off a freshly derived value, the
grouping decision and the arm the union actually takes could come from two
different reads. Today they cannot. There is also a functional difference:
reading `isPlainUnion` latches an unlatched union at the earliest AQE consumer,
and reading `rawPartitioning` would not.
I could not construct a wrong result from the fresh read, for what it is
worth. Wherever a parent relies on the union's partitioning, `rawPartitioning`
has stayed concrete and the two values agree; where they diverge, nothing
relies on it and both arms produce the same rows in a different layout. So this
is about keeping the decision locally evident rather than contingent on that
analysis.
On the behavior change itself you are right that it is one, and that it was
neither documented nor tested. I have added it to the description: after a skew
split under a non-plain union both children lose coalescing, where before the
re-derived answer was plain and each child was coalesced independently.
One correction to the repro. `df1.hint("rebalance",
"k").union(df2.hint("rebalance", "k")).count()` does not reach it: the
`ProjectExec` that `count()` puts above the union drops `k`, so the union
latches plain and takes the independent-group path. It needs an aggregate that
keeps the key, for example `groupBy("k").max("v")`.
On a test, the cheapest shape I found pins the timing rather than the skew
path: cache a child so the union latches while the inner AQE plan is non-final,
turn coalescing on, and assert that the two children's read specs may differ. I
can add that if you want it, though it does not cover the skew axis, and
building a stable one-side-skewed rebalance pair looked more expensive than
this behavior change warrants. Your call.
##########
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)
Review Comment:
Not acceptable as is, agreed. I do not think "latch only once the plan is
prepared" can be expressed from inside the node either, since it has no way to
know whether it is being read before or after preparation.
What I would do instead is stamp the decision from a rule appended after
`EnsureRequirements`, in both `QueryExecution.preparations` and AQE's
stage-prep list. That turns "whoever reads first" into a defined point and
makes `outputPartitioning` side-effect free again.
That also answers half of your design comment below: the conf can be a
planner-set field, but the partitioning half cannot be decided in
`SparkStrategies`, because the children there are `PlanLater` placeholders. A
decision taken at that point comes out plain for every union, including ones
whose children are co-partitioned without any exchange.
Would you rather I do that in this PR, or land the current read-on-first-use
and follow up?
##########
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()
Review Comment:
I would rather keep the `val`, because `lazy val` trades this NPE for a
lock-ordering problem.
A `lazy val`'s initializer takes the enclosing instance's monitor in Scala
2.13, and that is the same monitor `unionedInputRDD`'s `lazy val` holds while
it builds the children's RDDs. I checked the bytecode on this branch rather
than trusting the reference: `unionedInputRDD$lzycompute` does `monitorenter`
on `this` and runs `children.map(...)` and `new UnionRDD(...)` inside it.
`supportCodegenFailureReason` is another `lazy val` that calls `isPlainUnion`
from inside its own initializer. So a lazy lock puts its own initialization
behind the monitor that the separate lock exists to stay out of, and
`CoalesceShufflePartitions` reads `isPlainUnion` while holding the AQE lock.
On the NPE I read it as parity with what `SparkPlan` already does rather
than as unreachable. `SparkPlan` has `@transient private val prepareLock = new
Object()`, taken in `prepare()` and `waitForSubqueries()`, which are on the
path of every `execute*`, and `@transient val session ... orNull` is on the
path of `conf`, `sparkContext` and `metrics`. A deserialized plan that anyone
uses as a plan has been failing on those long before this, so `decisionLock`
adds no exposure that was not already there. I would not claim more than that:
I did not sweep every path that can put a plan in a closure.
If you would still rather not add another one, the shape that avoids the
monitor without the NPE is an explicit field forwarded in
`withNewChildrenInternal`, which is your next comment.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1267,6 +1313,9 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
object UnionExec {
+ /** The latched "is this a plain concatenation" decision. See
`isPlainUnion`. */
+ private val PLAIN_UNION_DECISION = TreeNodeTag[Boolean]("plainUnionDecision")
Review Comment:
Reasoning on record, and it is not a strong preference.
The tag does one thing a field does not: it survives the `withNewChildren`
copy that `insertInputAdapter` puts inside the codegen shell without each
future rebuild path having to remember to forward it. A field forwarded in
`withNewChildrenInternal` covers that path, and that override is the only copy
between `CollapseCodegenStages` and execution today, so it would work. It moves
the invariant from "tags are copied" to "this override forwards it", and
`EnsureRequirements` already has two `g.copy(...)` sites that call
`copyTagsFrom` by hand, which is the same kind of thing being forgotten.
The conf as a planner-set constructor field I like independently of the
rest, since it removes a live read at execution.
The part I do not think works is deciding the partitioning half in
`SparkStrategies`: the children there are `PlanLater` placeholders, so the
answer comes out plain for everything, including unions that are co-partitioned
without an exchange. If you want the decision at a defined point rather than on
first read, the place that has real children and still runs before execution is
a rule after `EnsureRequirements`, which is what I proposed on the `sparkPlan`
thread.
--
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]
