AQE: a replan that reassigns stage ids orphans the sibling leaf stage, and its task completion errors out

[andygrove]

Background

Under AdaptiveExecutionGraph, finalising a stage can trigger a replan that changes the set of stages. AdaptivePlanner::cancel_stage drops the stage from runnable_stage_cache (ballista/scheduler/src/state/aqe/planner.rs:197), but nothing on the graph side transitions that stage out of Running or cancels its in-flight tasks. update_stage_progress just logs and moves on:

if !stages_to_cancel.is_empty() {
    warn!(
        "there are stages to be cancelled but its not implemented. stages to cancel: {:?}",
        stages_to_cancel
    );
}

(ballista/scheduler/src/state/aqe/mod.rs:825)

Problem

When the orphaned stage's tasks later report success, update_stage_progress calls AdaptivePlanner::finalise_stage, which no longer finds the stage in runnable_stage_cache and fails:

Err(DataFusionError(Execution("Can't find active cache resolve")))

(ballista/scheduler/src/state/aqe/planner.rs:270)

Two consequences:

1. The job can never complete. The error propagates out of update_task_status before the stage is transitioned, so the orphaned stage is left Running with every one of its tasks already reported successful and zero available tasks. is_successful() requires all stages to be Successful, so it stays false forever. The job neither finishes nor fails.
2. Unrelated task updates are dropped. TaskManager::update_task_statuses groups the incoming batch by job and uses ? inside that loop (ballista/scheduler/src/state/task_manager.rs:568), so one bad job aborts the whole batch. query_stage_scheduler only logs Failed to update N task statuses for Executor X and discards the rest, including statuses belonging to other jobs and any events already accumulated.

Reproduction

A join over two leaf stages. Completing stage 0 triggers a replan that creates stages 2 and 3 and drops stage 1 from the planner cache, while stage 1 remains Running in the graph. Reporting stage 1's tasks as successful then errors on the task that completes the stage:

initial stages: [0, 1]
popped tasks for stages: [0, 0, 1, 1]
stages after stage 0 finished: [0, 1, 2, 3]
  stage 0: Successful
  stage 1: Running
  stage 2: Resolved
  stage 3: Resolved
stage 1 task completion result: Ok([])
stage 1 task completion result: Err(DataFusionError(Execution("Can't find active cache resolve")))
is_successful after error: false
  final stage 0: Successful
  final stage 1: Running
  final stage 2: Running
  final stage 3: Running
running_stages: [2, 3, 1]
available_tasks: 4

The scratch test that produced this builds an AdaptiveExecutionGraph for left JOIN right ... GROUP BY left.id with target_partitions = 2, pops every task of both leaf stages up front, completes only stage 0's tasks, then reports stage 1's tasks as successful.

Proposed fix

Handle stages_to_cancel on the graph side rather than logging it: transition the cancelled stage out of Running (or remove it from stages entirely so is_successful() is not blocked), and emit CancelTasks for its in-flight tasks. Late task statuses arriving for a cancelled stage should be ignored rather than routed into finalise_stage.

Separately, TaskManager::update_task_statuses should isolate per-job failures so one job's error cannot discard another job's task updates in the same batch.

Notes

Found while writing an AQE test for #1996. Not caused by that change, and #2149 sidesteps it by testing with an aggregation chain instead of a join.