ulysses-you commented on PR #59630: URL: https://github.com/apache/spark/pull/59630#issuecomment-6519117017
Thank you @cloud-fan for the follow-up review. Addressed the six documentation and test findings in 0a98c1e6a70. **Staged-union cost:** Preparation now restores the union cost mark using `hasPushedShuffle`. This relies on the current rewrite producing two direct branch shapes: an `ENSURE_REQUIREMENTS` shuffle, or a `REPARTITION_BY_COL` shuffle whose count differs from the effective default after retargeting. Column retargeting preserves keys and their order; query stages retain the shuffle origin and count. The check does not look through intermediate operators and must be updated if future rewrites change these shapes. It covers partially staged unions as well as fully staged ones. The Boolean cost mark remains, and plain unions are not newly marked. An alternative is a dedicated push-down `ShuffleOrigin` carrying the original origin, which would make provenance explicit rather than inferred. That would also require adapting origin-based optimization guards and preserving exchange canonicalization/reuse behavior. We keep the narrower shape check here; explicit origin tracking can be considered if these assumptions become insufficient. A runtime SQL test records the actual current/proposal evaluations while delegating to `SimpleCostEvaluator`. It checks every evaluation, including the restored mark and the union cost of one, allowing either one or two comparison rounds because AQE may batch stage completion events. Disabling stage-based marking makes the test fail. **Forced skew:** We retain the existing opt-in behavior. The documentation now explicitly states that restoring the union exchange after moved branches have shuffled can shuffle those branches a second time. **Reducing operators:** We also retain the look-through opt-in. The documentation now states that reducers remain above the union, moved branches shuffle before reduction, and target ranking does not estimate that loss. Pushing partial aggregation through union branches can be explored in a follow-up PR to preserve reduction before shuffling moved branches. The change also restricts column-repartition retargeting to count changes, preserving keys and their order, with end-to-end SQL coverage. Rejected zero-move candidates now keep unary wrappers isolated from origin-tag inheritance. -- 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]
