dongjoon-hyun commented on code in PR #58419:
URL: https://github.com/apache/spark/pull/58419#discussion_r3914971055
##########
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.
Review Comment:
This invariant only holds on the row path. `doExecuteColumnar` (L1307) still
does `sparkContext.union(children.map(_.executeColumnar()))` regardless of
`outputPartitioning`, so a non-plain union that stays columnar concatenates
while advertising its children's `HashPartitioning`.
I reproduced wrong results on `master` with two bucketed Parquet tables
(`bucketBy(4, "k")`):
```sql
SELECT k, count(*) FROM (SELECT * FROM t1 UNION ALL SELECT * FROM t2) GROUP
BY k
```
plans as `HashAggregate <- ColumnarToRow <- Union(FileScan bucketed,
FileScan bucketed)` with no exchange and returns 10 rows instead of 5
(`spark.sql.unionOutputPartitioning=false` gives 5). This predates this PR
(SPARK-52921), but since the doc here states the invariant, could we either
mirror `doExecute` in `doExecuteColumnar` or make `supportsColumnar` false when
`!isPlainUnion`? A separate JIRA is fine too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1047,10 +1090,10 @@ 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.
+ // Memoized so `supportCodegen` (called repeatedly by
`CollapseCodegenStages`)
+ // and `metrics` see one reason on one instance; `conf` is live, so
re-deriving
+ // could answer differently. Agreeing with the `withNewChildren` copy is the
+ // `isPlainUnion` tag's job, not this memo's.
@transient private lazy val supportCodegenFailureReason: Option[String] = {
Review Comment:
The comment above says agreeing with the `withNewChildren` copy is the tag's
job, but the tag only covers the `isPlainUnion` input. The copy that
`insertInputAdapter` puts inside the shell still evaluates this lazy val for
the first time at execution, so `WHOLESTAGE_UNION_CODEGEN_ENABLED` /
`WHOLESTAGE_UNION_MAX_CHILDREN` flipped between `executedPlan` and `collect()`
still gives empty `metrics` and the same `key not found: numOutputRows` from
`doProduce` (AQE off, children with an exchange so the copy is fresh).
That's pre-existing, but it suggests the cleaner fix is registering
`numOutputRows` unconditionally like the other `CodegenSupport` operators,
rather than tying `metrics` to this memo.
##########
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:
`rawPartitioning` only walks the children (no `execute`, no lock another
thread could hold while waiting on this one), so this can be a single
`decisionLock.synchronized { getTagValue(...).getOrElse { ...;
setTagValue(...); plain } }`, like `SparkPlan.prepare()`. Concurrent first
readers then wait behind one derivation instead of each deriving and discarding.
##########
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:
This is `null` after deserialization, so `isPlainUnion` /
`outputPartitioning` / `doExecute` NPE on a deserialized `UnionExec` (e.g.
inside a `ScalarSubquery` / `InSubqueryExec` plan captured in a task closure),
where the previous `outputPartitioning` still worked. `@transient private lazy
val` keeps it driver-only and survives the round trip, as SPARK-23731 did for
`FileSourceScanExec`.
##########
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:
nit: this paragraph (and "not this memo's job" in the memo comment below)
reads as PR history rather than a description of the code. I'd keep the first
paragraph plus a one-liner that the non-plain branch may become
`UnknownPartitioning` after AQE skew splitting and is tolerated, and drop the
tried-and-reverted sentence.
##########
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:
This and the next test have identical setup and assert two halves of the
same decision. One test asserting both `metrics.contains("numOutputRows")` and
`outputPartitioning.isInstanceOf[UnknownPartitioning]` would do.
##########
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:
Since this write happens on first read, reading `outputPartitioning` on the
un-prepared `queryExecution.sparkPlan` now latches `plain` (children haven't
got their exchanges yet), and `executedPlan` inherits it through
`sparkPlan.clone()` -> `makeCopy` -> `copyTagsFrom`
(`QueryExecution.scala:395`). After `EnsureRequirements` inserts matching
`HashPartitioning` under both children, the union still reports
`UnknownPartitioning` and the parent's exchange elimination is lost.
`QueryTest.checkAnswer(_, planFunction, _)` and `PlannerSuite` inspect
`sparkPlan` this way, and listeners/extensions can too. Before the PR
`outputPartitioning` was side-effect free.
##########
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:
One consumer of `isPlainUnion` whose behavior this latch changes is
`CoalesceShufflePartitions.childrenNeedCompatiblePartitioning` (L189).
With AQE on, `SparkPlanInfo.fromSparkPlan` / `EnsureRequirements` latch the
union non-plain while both children are identical rebalance exchanges. If
`OptimizeSkewInRebalancePartitions` later splits only one child, that child
reports `UnknownPartitioning`, but the tag still says non-plain, so both
children land in one coalesce group with mixed specs and
`coalescePartitionsWithSkew` bails out for both (`Could not apply partition
coalescing ...`). Before this PR the re-derived answer was plain and each child
was coalesced independently.
e.g. `df1.hint("rebalance", "k").union(df2.hint("rebalance", "k")).count()`
with skew only in `df1`. Results are still correct, but this is a behavior
change that isn't mentioned or tested.
##########
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:
This guard is unreachable: both callers are inside `withTempView("v")`,
which already uncaches via `Catalog.dropTempView -> uncacheView` at the end of
each test.
##########
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:
Do we need a tag + lock for this? `UnionExec` already overrides
`withNewChildrenInternal`, which is the only copy path between
`CollapseCodegenStages` and execution, so an explicit per-node field forwarded
there would survive the same rebuilds without depending on `copyTagsFrom`'s
only-into-a-tagless-node rule (`QueryStageExec._resultOption` is the precedent).
And the `UNION_OUTPUT_PARTITIONING` half could be fixed at planning as a
constructor field set in `SparkStrategies`, like `ExpandExec.useSingleTask` /
`ShuffleExchangeExec.shuffleOrigin`. That also pins it at planning for a root
union with codegen off, where the first read is currently in `doExecute`.
##########
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:
The answer is fully determined here (4 ids per `k` from each range), so
`checkAnswer(planned, (0L until 5L).map(k => Row(k, 8L)))` is stronger and
avoids a second plan/execute. The parity oracle would still pass if both paths
regressed to 10 rows of 4.
--
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]
