DeviousCardi opened a new pull request, #2613: URL: https://github.com/apache/datafusion-ballista/pull/2613
# Which issue does this PR close? Closes #2594. # Rationale for this change `AdaptiveExecutionGraph::resolve_stage`, `rollback_running_stage` and `rollback_resolved_stage` removed the stage from `stages` before running the fallible plan rewrite (`to_resolved` / `to_unresolved`). If the rewrite failed, `?` returned early and the stage was never put back, so later lookups failed with "stage not found" instead of the original error. `StaticExecutionGraph` already avoids this. # What changes are included in this PR? - The three AQE methods now `get` the stage, build the new one, and `insert` it only once that succeeds, the same way `StaticExecutionGraph` does it. `succeed_stage`, `fail_stage` and `rerun_successful_stage` keep `remove` because their conversions can't fail. - `FailingPlanRewriteExec` / `fail_plan_rewrites` move from the static graph's test module to `test_utils.rs` so both graphs can use them. This is a move only; the static graph's tests are otherwise unchanged. - New `state/aqe/test/stage_transitions.rs` with three tests that mirror the static graph's tests. Each one forces the plan rewrite to fail and checks that the stage is still in the graph in its original state (same plan, attempt number, and pending and running tasks). All three fail on `main` and pass with this change. `cargo test -p ballista-scheduler` passes, and so does `cargo clippy -p ballista-scheduler --all-targets --all-features -- -D warnings`. # Are there any user-facing changes? No. --- AI disclosure: Claude Code prepared the implementation and tests. I reviewed them and will reply to review comments myself. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01Y1b55sEL1gA1c3FCMcQCbm -- 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]
