Repository navigation
[EPIC] Share execution graph logic between StaticExecutionGraph and AdaptiveExecutionGraph #2595
Description
Is your feature request related to a problem or challenge? Please describe what you are trying to do.
StaticExecutionGraph (state/execution_graph.rs) and AdaptiveExecutionGraph (state/aqe/mod.rs) carry near-identical copies of most of their logic. Both files are among the most frequently changed in the scheduler (29 and 25 commits in the last 12 months), and the copies have already drifted:
resolve_stage/rollback_*_stagelose the stage on error in the AQE copy only (AQE: resolve_stage and rollback_*_stage drop the stage from the graph when the conversion fails #2594).- The AQE
update_task_statushas deadreset_running_stageshandling, lacks the UnResolved late fetch-failure path, and reports a less useful max-failures message (details in AQE: handle executor failure on the adaptive execution path #1986).
Duplicated code:
| What | Static graph | Adaptive graph |
|---|---|---|
update_task_status |
L760-L1195, 376 lines, cognitive complexity 164, nesting 9 | L781-L1142, 361 lines, cognitive complexity 121 |
reset_stages_internal |
L532-L655 | L548-L676 (adds planner.remove_exchange_locations) |
accessors, revive, is_successful, completed_stages, lost-executor reset tail |
L609-L766 | L630-L787 (identical apart from type names) |
stage transitions (resolve_stage, succeed_stage, fail_stage, rollback_*, rerun_successful_stage, fail_job, succeed_job) |
L1307-L1438 | L1254-L1391 |
processing_stages_update, get_running_stage_id |
L390-L455, L504-L530 | copied |
pop_next_task (test only) |
L1484+ | L1440+ |
Describe the solution you'd like
Proposed sub-tasks, each a behaviour-preserving PR (apart from the bug fixes) relying on the existing execution-graph tests:
- Fix the AQE stage-loss bug (AQE: resolve_stage and rollback_*_stage drop the stage from the graph when the conversion fails #2594) so both copies behave the same before they are merged.
- Extract the shared state (job/session ids, status, timestamps,
stages,failed_stage_attempts, ...) into a common struct embedded by both graphs, and implement the accessors,revive,is_successful,completed_stagesand lost-executor reset once. - Share the stage-transition methods and
processing_stages_update/get_running_stage_idas default trait methods or free functions over the stages map. - Share
reset_stages_internalas a function returning a reset plan; the AQE graph adds itsremove_exchange_locationsstep afterwards. Replace the nestedfor_eachclosures with plain loops. - Break
update_task_statusintohandle_failed_task,handle_successful_task,handle_late_task_for_unresolved_stageand the resubmit/reset post-passes, then share those between the graphs. This also resolves the AQE gaps described in AQE: handle executor failure on the adaptive execution path #1986. - Share the test-only
pop_next_task.
Describe alternatives you've considered
Keep the copies and port each fix by hand. The drift above shows this is already failing.
Additional context
Related: #1986, #2539, #2150. Line references are pinned to commit ea1126c.
Activity
milenkovicm commentedon Oct 14, 2026
i will take blame for this, meant to clean them up when AQE get mature but can't really catch up with ballista this days