LuciferYang commented on PR #58419: URL: https://github.com/apache/spark/pull/58419#issuecomment-5832365068
You are right that the phase argument leans on a convention the APIs do not enforce, and the case that settles it is a rule that builds a parent over a union it just made: nothing preparation did protects a node preparation never saw. Both windows are closed now. `optimizeQueryStage` records the snapshot on each rule result that changed the plan, before the `ValidateRequirements` check and before the fold continues. For the columnar rules, `SnapshotUnionPreparationConf.after` wraps each injected `ColumnarRule` so the record is written behind each of its two transitions: `ApplyColumnarRulesAndInsertTransitions` keeps its signature, and what a rule returns is recorded whether it came from the pre pass or the reverse post pass. `StampUnionDecisions` still lands where it did, so `EnsureRequirements` and the children decide the plain-union answer, and no partition count or topology term is frozen. Two cases in `SparkSessionExtensionSuite`: two injected query stage optimizer rules, the first adding a union and the second reading its partitioning with the conf off, and the same pair as injected columnar rules, in each transition phase and with AQE off and on. Dropping the record from the fold fails the first, with `ListBuffer(UnknownPartitioning(0), UnknownPartitioning(0))`. Leaving either columnar transition unrecorded fails the second, one call site and one phase at a time. One thing the tests do not reach: with `ValidateRequirements` the consumer is only reachable through an injected rule that is itself an `AQEShuffleReadRule`, since the three built-in ones are listed ahead of the injected list, so the plan they validate holds no union without a record. That consumer answers from the same write, and the comments now say so rather than implying a case covers it. The description is updated for both, and its mechanism paragraph now names `SnapshotUnionPreparationConf` rather than the class name it carried before the rule recorded all three confs. -- 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]
