apache / apache/datafusion-ballista
AQE: immediate intermediate-shuffle reclaim on job success is a no-op under the adaptive planner
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
### Background
#1982 reclaims a job's intermediate shuffle data immediately on success (instead of waiting for `finished_job_data_clean_up_interval_seconds`), to relieve executor disk pressure (#1981). It identifies intermediate stages via `ExecutionGraph::intermediate_stage_ids()`, which selects stages whose `output_links` is non-empty (the final stage has empty `output_links`).
### Problem
This works for the **static** planner (`StaticExecutionGraph`), where `ExecutionStageBuilder` populates `output_links` up front. It does **not** work under **AQE** (`AdaptiveExecutionGraph`): stages are created incrementally by `create_resolved_stage`, which always sets `output_links` to `vec![]` (`ballista/scheduler/src/state/aqe/mod.rs:212`, *"we do not know output links at this moment"*), and no AQE code path ever back-fills it.
Consequently, under AQE every stage has empty `output_links`, so `intermediate_stage_ids()` returns an empty set and the immediate reclaim does nothing — the job falls back to the existing delayed whole-job cleanup. This is **safe** (the "never delete a final stage" invariant holds trivially), but delivers **no disk benefit under AQE**. That matters because the motivating workload (the SF10 CI job in #1981) runs all 22 queries **both AQE off and on**, so the AQE-on half gets no relief.
### Proposed fix
Override `intermediate_stage_ids()` for `AdaptiveExecutionGraph` to derive final stages from `output_locations` rather than `output_links`: each `PartitionLocation` in `output_locations` carries its terminal `stage_id`, so `final = {stage_ids in output_locations}` and `intermediate = all_stages − final`. Guard: if `output_locations` is empty, return an empty set (reclaim nothing) to preserve the safety invariant.
### Tests
Add an AQE test that drives a multi-stage adaptive job to success and asserts `intermediate_stage_ids()` returns exactly the non-terminal stages (and never a stage present in `output_locations`).
Follow-up to #1982. Related: #1981.
Contributor guide
Research direction
Start in ballista/scheduler/src/state/aqe/mod.rs around AdaptiveExecutionGraph::create_resolved_stage and its empty output_links. Trace output_locations and compare terminal stage IDs with all stages, preserving the empty-output safety guard. Add an AQE multi-stage success test that verifies intermediate_stage_ids() contains exactly non-terminal stages and excludes every stage in output_locations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100