apache / apache/datafusion-ballista
Adaptive (AQE) execution hangs when a replan cancels an in-flight stage
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
## Describe the bug
Under adaptive execution, when a replan collapses the residual plan and removes an in-flight stage's `ExchangeExec`, the scheduler cancels that stage but its already-running tasks keep going. When those tasks complete, the late status update fails and the job never finishes — the client hangs indefinitely.
The scheduler logs:
```
WARN ballista_scheduler::state::aqe: there are stages to be cancelled but its not implemented. stages to cancel: {2}
ERROR ballista_scheduler::scheduler_server::query_stage_scheduler: Failed to update 3 task statuses
for Executor : DataFusionError(Execution("Can't find active cache resolve"))
```
Mechanism:
1. `AdaptivePlanner::actionable_stages` finds stage 2's exchange is gone from the replanned plan and returns it in `stages_to_cancel`.
2. `AdaptiveExecutionGraph::update_stage_progress` calls `cancel_stage(2)`, which removes the stage from `runnable_stage_cache` (`aqe/planner.rs`, `cancel_stage` — "TODO consider marking cancelled stage").
3. Stage 2's tasks were already launched and run to completion regardless — cancellation is not implemented on the executor side.
4. Their status update reaches `finalise_stage(2)`, whose `runnable_stage_cache.get(&stage_id)` now misses and returns `Execution("Can't find active cache resolve")`.
5. The error propagates out of the task-status batch, so the job's completion event is never emitted and the query hangs.
Note the final stage does complete successfully before this happens — the result is computed, it just never gets delivered.
## To Reproduce
Generate SF1 TPC-DS Parquet, start a scheduler + executor, then, with #2047 fixed (on `main` today the `EmptyExec` assertion aborts these queries first and masks this):
```
tpcds --host localhost --port 50050 --path \
--query 1 -c datafusion.optimizer.prefer_hash_join=false \
-c ballista.planner.adaptive.enabled=true
# -> hangs; scheduler logs "Can't find active cache resolve"
```
Reproduced on q1 and q44 at SF1 (1 scheduler / 1 executor, 8 slots). The other 13 queries listed in #2047 complete once the assertion is fixed. The same queries succeed under the static planner.
## Expected behavior
A completion update for a stage that a replan has cancelled is ignored rather than treated as an error, and the job runs to completion.
## Additional context
Reproduced on `apache/main` @ e843c8e8 plus the fix in #2062. The `Can't find active cache resolve` error is already visible in the scheduler log on unmodified `main` for these queries — it is pre-existing and independent of #2047, which only removed the assertion that was terminating the job first.
Part of the adaptive-execution epic #1359. Distinct from the q23/q77 hang noted in #2047 (a suspected cross-stage dynamic-filter deadlock) and from the AQE correctness divergences in #2046.
Contributor guide
Research direction
Start with AdaptivePlanner::actionable_stages and AdaptiveExecutionGraph::update_stage_progress, then inspect cancel_stage and finalise_stage in aqe/planner.rs. Reproduce with the TPC-DS q1 or q44 command after #2047 is fixed. Done means completion updates from replanned-away stages are ignored and the job finishes instead of hanging.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100