[SPARK-59122][SQL] Take UnionExec's plain-union decision once instead of re-deriving it per caller #58419
+1,558
−108
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,165 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one or more | ||
| * contributor license agreements. See the NOTICE file distributed with | ||
| * this work for additional information regarding copyright ownership. | ||
| * The ASF licenses this file to You under the Apache License, Version 2.0 | ||
| * (the "License"); you may not use this file except in compliance with | ||
| * the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| package org.apache.spark.sql.execution | ||
|
|
||
| import org.apache.spark.sql.catalyst.rules.Rule | ||
| import org.apache.spark.sql.internal.SQLConf | ||
|
|
||
| /** | ||
| * Fixes each [[UnionExec]]'s partitioning decision and the confs its codegen gate reads, at one | ||
| * defined point. | ||
| * | ||
| * `UnionExec` derives both from state that moves: its children's `outputPartitioning` sharpens as | ||
| * AQE finalises the plans behind them, and `conf` is the live session conf. That sharpening does | ||
| * not run backwards: a cached plan is executed at most once, through `SparkPlan`'s memoized | ||
| * `execute` or `executeColumnar`, so a node planned against it cannot see the window where it | ||
| * reports no final plan reopen, and a rebuild replaces the relation rather than re-entering that | ||
| * plan. Every reader used to derive its own answer, so the answer depended on when it was read: the | ||
| * codegen gate could fuse a union whose copy in the shell then answered the other way, so `metrics` | ||
| * came back empty and `doProduce` failed asking `metricTerm` for `numOutputRows`. This rule asks | ||
| * right after `EnsureRequirements`, so the decision the exchanges around a union were planned | ||
| * against is the one `unionRDDs` and the codegen gate use. | ||
| * | ||
| * It is listed again after the injected columnar and query-stage rules, the hooks that can add a | ||
| * `UnionExec` of their own. One created there has no decision yet, and would otherwise take one | ||
| * wherever it is first asked, where the copy in the codegen shell can disagree with the gate. The | ||
| * cached-scan branch of stage creation needs no barrier: it rejects a result that is no longer an | ||
| * `InMemoryTableScanLike`, which is a leaf. | ||
| * | ||
| * A tag rather than a constructor field, because a field would land in `argString` and so in | ||
| * every `explain` and `PlanStability` golden holding a `Union`, and in `canonicalized`, which | ||
| * exchange and cached-plan reuse key on. Writing it in place is safe here, unlike in | ||
| * `MarkSingleTaskExecution`, because preparation runs on `sparkPlan.clone()` and every | ||
| * shared-subtree boundary is a leaf, so `foreach` cannot reach a node another query owns. | ||
| * | ||
| * It only writes what is not there yet, so a later pass cannot move a decision already stamped on | ||
| * a node, a second pass over the same nodes keeps the first answer, and a node rebuilt from a | ||
| * stamped one keeps the tag `copyTagsFrom` gave it. AQE re-plans between rounds, so a union above | ||
| * the stages already created is stamped again, from what that round sees and against that round's | ||
| * exchanges; a round whose plan loses on cost is discarded whole, stamps included. A union inside a | ||
| * stage is not revisited, since `foreach` stops at `QueryStageExec`. | ||
| */ | ||
| class StampUnionDecisions(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { | ||
| override def apply(plan: SparkPlan): SparkPlan = { | ||
| plan.foreach { | ||
| case u: UnionExec => u.stampDecisions(snapshot) | ||
| case _ => | ||
| } | ||
| plan | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Records the preparation's [[UnionConfSnapshot]] on each [[UnionExec]], so that every reader ahead | ||
| * of a barrier answers from the same values the barrier will stamp. | ||
| * | ||
| * Without this, each phase samples the confs separately: `EnsureRequirements` reads what a union | ||
| * reports under the value then, and [[StampUnionDecisions]] freezes the decision under the value | ||
| * one rule later. `conf` is the live session conf, so another thread turning | ||
| * `UNION_OUTPUT_PARTITIONING` off in that window would let a parent drop an exchange over a | ||
| * concrete partitioning and then have the union concatenate, which puts one group in two | ||
| * partitions. | ||
| * | ||
| * The same gap opens between two injected rules: one can return a `UnionExec` of its own, which | ||
| * carries no record yet, and a later one can plan requirements over it, build a parent over it, or | ||
| * ask its codegen gate. [[SnapshotUnionPreparationConf.before]] closes that for the two lists whose | ||
| * rules are listed as plain `Rule[SparkPlan]`, the AQE post-planner-strategy rules and the | ||
| * query-stage preparation rules. [[SnapshotUnionPreparationConf.after]] closes it for the injected | ||
| * columnar rules, which share one `ApplyColumnarRulesAndInsertTransitions` and so have no place | ||
| * between them for a listed pass: the record rides on each rule's own transitions instead. | ||
| * `AdaptiveSparkPlanExec.optimizeQueryStage` writes it on every rule result that changed the plan, | ||
| * which covers the next rule, and `ValidateRequirements` for an injected rule that is itself an | ||
| * `AQEShuffleReadRule`. The built-in ones are listed ahead of the injected list, so what they | ||
| * validate holds no union without a record. | ||
| * | ||
| * Only the confs are recorded, never a partitioning. `EnsureRequirements` has not inserted the | ||
| * exchanges it adds yet, so a decision taken now would freeze plain on a union whose children only | ||
| * become co-partitioned there, which is why the decision itself waits for the barrier behind it. | ||
| * | ||
| * Writing the tag in place is safe for the reason given on [[StampUnionDecisions]]. | ||
| */ | ||
| class SnapshotUnionPreparationConf(snapshot: UnionConfSnapshot) extends Rule[SparkPlan] { | ||
| override def apply(plan: SparkPlan): SparkPlan = { | ||
| plan.foreach { | ||
| case u: UnionExec => u.recordPreparationConf(snapshot) | ||
| case _ => | ||
| } | ||
| plan | ||
| } | ||
| } | ||
|
|
||
| object SnapshotUnionPreparationConf { | ||
| /** | ||
| * `rules` with a snapshot pass ahead of each of them, so a `UnionExec` one rule creates carries | ||
| * the preparation's confs before the next rule reads them. Ahead of rather than behind, so that a | ||
| * union left by the rules listed before `rules` is covered too. `rules` is empty unless an | ||
| * extension injected something, so this adds no pass to an ordinary preparation. | ||
| */ | ||
| def before(snapshot: UnionConfSnapshot, rules: Seq[Rule[SparkPlan]]): Seq[Rule[SparkPlan]] = | ||
| rules.flatMap(rule => Seq(new SnapshotUnionPreparationConf(snapshot), rule)) | ||
|
|
||
| /** | ||
| * `rules` with the same record written behind each of their transitions. | ||
| * `ApplyColumnarRulesAndInsertTransitions` applies these itself, every | ||
| * `preColumnarTransitions` in order and then every `postColumnarTransitions` in reverse, so a | ||
| * pass cannot be listed between two of them from outside; wrapping the transitions puts it there. | ||
| * Behind rather than ahead, because the plan reaching the wrapper has passed a barrier already, | ||
| * and it is what a rule returns that can hold a union no pass has seen. `rules` is empty unless | ||
| * an extension injected something. | ||
| */ | ||
| def after(snapshot: UnionConfSnapshot, rules: Seq[ColumnarRule]): Seq[ColumnarRule] = | ||
| rules.map(RecordUnionPreparationConf(_, snapshot)) | ||
| } | ||
|
|
||
| /** | ||
| * Records `snapshot` on each `UnionExec` in what `inner` returns, so a union `inner` created is | ||
| * read from the preparation's confs by the next rule in the same | ||
| * `ApplyColumnarRulesAndInsertTransitions`. See [[SnapshotUnionPreparationConf.after]]. | ||
| */ | ||
| private case class RecordUnionPreparationConf(inner: ColumnarRule, snapshot: UnionConfSnapshot) | ||
| extends ColumnarRule { | ||
|
|
||
| private val record = new SnapshotUnionPreparationConf(snapshot) | ||
|
|
||
| override def preColumnarTransitions: Rule[SparkPlan] = | ||
| plan => record(inner.preColumnarTransitions(plan)) | ||
|
|
||
| override def postColumnarTransitions: Rule[SparkPlan] = | ||
| plan => record(inner.postColumnarTransitions(plan)) | ||
| } | ||
|
|
||
| /** | ||
| * The union confs one preparation answers from, read once and shared by every barrier in it. | ||
| * | ||
| * A barrier that read the live conf instead would let two of them disagree: an injected rule can | ||
| * return an equivalent `UnionExec` carrying tags it set itself, and `copyTagsFrom` adds nothing to | ||
| * a node that already has one, so such a replacement reaches the late barrier with no record of its | ||
| * own and would be stamped from whatever the conf says by then, rather than from what the exchanges | ||
| * above it were planned against. | ||
| */ | ||
| case class UnionConfSnapshot( | ||
| outputPartitioning: Boolean, | ||
| codegenEnabled: Boolean, | ||
| maxChildren: Int) | ||
|
|
||
| object UnionConfSnapshot { | ||
| def apply(conf: SQLConf): UnionConfSnapshot = UnionConfSnapshot( | ||
| outputPartitioning = conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING), | ||
| codegenEnabled = conf.getConf(SQLConf.WHOLESTAGE_UNION_CODEGEN_ENABLED), | ||
| maxChildren = conf.getConf(SQLConf.WHOLESTAGE_UNION_MAX_CHILDREN)) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Non-blocking (P2): The AQE configuration-flip case exercises a union already stamped while
initialPlanis built, and the extension cases keep the initial value stable until their late barriers have run. Replacing a later barrier's storedunionConfwith a fresh SQLConf read would therefore leave those assertions green. Please add exchange-backed AQE cases that flip output partitioning, codegen enablement, andmaxChildrenafter this snapshot, then create a fresh union at a later stage/extension barrier and assert its layout, fusion, shell copy, andnumOutputRows.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added two cases in
SparkSessionExtensionSuite, both flipping a conf after theAdaptiveSparkPlanExecwrapper has read its snapshot and before the stage that carries the injected union is created.the barrier in AQE post stage creation stamps from the adaptive snapshot: an injected columnar rule adds the union,spark.sql.unionOutputPartitioninggoes off onceexecutedPlanis built, and the union still reports its children'sHashPartitioning.a late AQE barrier records the codegen confs the wrapper was built with: an injected columnar rule rebuilds a three-child union, withspark.sql.codegen.wholeStage.union.enabledoff in one iteration andmaxChildrenat 2 in the other. It asserts the union still fuses, its children inside the shell areInputAdapters, and the shell copy'snumOutputRowsreads 60.Replacing that barrier's
unionConfwithUnionConfSnapshot(conf)turns both red:UnknownPartitioning(0)for the first, no fused union for the second.The query-stage-preparation barrier has no case of its own. It cannot show the difference:
queryStagePreparationRulesis aval, so the barrier in it holds the construction snapshot whichever round applies it, and the union an injected prep rule adds is stamped before a flip can land. With that barrier changed to readSQLConf.getper application, a case of the same shape stays green.postStageCreationRulesis the list rebuilt per stage, so it is the one whose value could come from a later read.