LuciferYang opened a new pull request, #58419:
URL: https://github.com/apache/spark/pull/58419
### What changes were proposed in this pull request?
`UnionExec` reads its children's `outputPartitioning` to decide three
separate things: whether whole-stage codegen fusion applies, whether
`numOutputRows` is registered, and which RDD `doExecute` builds. This PR makes
that decision once per node, stores it in a `TreeNodeTag`, and derives
`outputPartitioning` from the stored value.
The three edits:
- the old body of `outputPartitioning` becomes a private `rawPartitioning`,
still recomputed from the children on every call;
- `isPlainUnion` evaluates
`rawPartitioning.isInstanceOf[UnknownPartitioning]` on first use and latches
the answer in `UnionExec.PLAIN_UNION_DECISION`;
- `outputPartitioning` returns `super.outputPartitioning` for a node that
latched plain, and `rawPartitioning` otherwise.
`TreeNode.withNewChildren` ends in `copyTagsFrom`, which only writes into a
tagless node, so the copy that `CollapseCodegenStages.insertInputAdapter`
places inside the codegen shell inherits the decision instead of making its own.
### Why are the changes needed?
A fused `UnionExec` can fail at execution:
```
java.util.NoSuchElementException: key not found: numOutputRows
at org.apache.spark.sql.execution.SparkPlan.longMetric(SparkPlan.scala:154)
at
org.apache.spark.sql.execution.CodegenSupport.metricTerm(WholeStageCodegenExec.scala:71)
at
org.apache.spark.sql.execution.UnionExec.doProduce(basicPhysicalOperators.scala:1148)
at
org.apache.spark.sql.execution.WholeStageCodegenExec.doCodeGen(WholeStageCodegenExec.scala:676)
```
Reproduced on current master with default configuration:
```scala
spark.range(0, 200, 1, 4).selectExpr("id % 10 AS k", "id AS v")
.groupBy("k").agg(sum("v").as("s")).createOrReplaceTempView("v")
spark.catalog.cacheTable("v")
spark.sql("SELECT k, abs(s) AS s FROM v UNION ALL SELECT k, s FROM
v").collect()
```
The children's partitioning is not stable while the plan is being prepared.
`InMemoryTableScanExec` reads `cachedPlan.outputPartitioning`, and the inner
`AdaptiveSparkPlanExec` answers `UnknownPartitioning` until its final plan
exists.
In the query above, the expression on one branch keeps a `ProjectExec` from
being collapsed away, so the two children are shaped differently,
`comparePartitioning` rejects the pair, and the union looks plain and is fused.
`insertInputAdapter` then rebuilds the node and puts the copy in the shell, and
that copy evaluates the gate for the first time after the cache stages have
finalised. Both children now report the same `HashPartitioning`, the gate
answers `partitioning-aware`, `metrics` comes back empty, and the generated
code still increments the metric.
Registering the metric unconditionally would trade the crash for a wrong
answer. A fused union concatenates its children's partitions, so a node that
went on claiming their `HashPartitioning` could let a parent skip an exchange
it needs. Latching the decision keeps `outputPartitioning` and the RDD shape in
agreement.
This affects the fusion added in SPARK-56482, so 4.2.0, branch-4.3 and
master all carry it.
### Does this PR introduce _any_ user-facing change?
Yes. A query of the shape above fails on 4.2.0 and returns its rows after
this change.
There is a second, narrower change: a union whose children only agree on
partitioning after planning now keeps reporting `UnknownPartitioning`, so
SPARK-52921's exchange elimination no longer applies to that shape. That is the
conservative direction. The alternative is a node that concatenates partitions
while advertising a partitioning it does not have.
### How was this patch tested?
Two new cases in `UnionCodegenSuite`, one for each half of the decision: the
fused union keeps `numOutputRows`, and it reports `UnknownPartitioning`. Both
fail on master, the first with the exception above.
`UnionCodegenSuite`, `UnionCodegenAqeSuite` and `UnionCodegenAnsiSuite` run
135 tests, all passing. `./dev/scalastyle` is clean. The existing case
`SPARK-56482: partitioning-aware union falls back to non-codegen` covers
children that agree up front, so it pins that the latch does not disable the
partitioning-aware path.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Ducc (Claude Code, Claude Opus 5)
--
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]
