apache / apache/datafusion-ballista

AQE: immediate intermediate-shuffle reclaim on job success is a no-op under the adaptive planner

Open
#1,996 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.