andygrove opened a new issue, #2594: URL: https://github.com/apache/datafusion-ballista/issues/2594
**Describe the bug** In the adaptive execution graph, `resolve_stage`, `rollback_running_stage` and `rollback_resolved_stage` remove the stage from `self.stages` *before* the fallible conversion. If the conversion returns an error, the `?` exits and the stage is never put back, so it disappears from the graph. [aqe/mod.rs#L1254-L1262](https://github.com/apache/datafusion-ballista/blob/ea1126c0649b2bc08e1e233b1972b64318821c16/ballista/scheduler/src/state/aqe/mod.rs#L1254-L1262): ```rust if let Some(ExecutionStage::UnResolved(stage)) = self.stages.remove(&stage_id) { self.stages.insert( stage_id, ExecutionStage::Resolved(stage.to_resolved(self.session_config.options())?), ); ``` The same `remove`-then-`?` pattern is used in `rollback_running_stage` (L1313) and `rollback_resolved_stage` (L1344) with `to_unresolved(..)?`. The static graph's copies of these methods do this safely: they `get` the stage, compute the new one, and only then `insert`, so an error leaves the original stage in place ([execution_graph.rs#L1307-L1311](https://github.com/apache/datafusion-ballista/blob/ea1126c0649b2bc08e1e233b1972b64318821c16/ballista/scheduler/src/state/execution_graph.rs#L1307-L1311), and L1363, L1390). It looks like a fix landed in one copy but not the other. **To Reproduce** Any error from `UnresolvedStage::to_resolved` or `RunningStage/ResolvedStage::to_unresolved` on the AQE path. After the error, `stages` no longer contains the stage id, and later lookups fail with "stage not found" style errors rather than the original error. **Expected behavior** On error the stage stays in the graph, matching `StaticExecutionGraph`. The minimal fix is to port the `get`-compute-`insert` pattern to the three AQE methods. **Additional context** These state-transition methods are duplicated between the two graphs, which is how they drifted apart. Sharing them is tracked in the execution-graph deduplication epic. Line references are pinned to commit ea1126c0649b. -- 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]
