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

[andygrove]

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.

[andygrove]

This just bit PR #2064 in CI, so here is a concrete reproduction to go with the code analysis above.

Run 29923718090, leg TPC-H SF10 (AQE on, multi-partition tasks). Query 21 failed about 2 seconds in with No space left on device (os error 28) while writing shuffle output. Queries 1 through 20 all passed and verified. The AQE-off leg passed, which is exactly the asymmetry this issue predicts.

Task 1 failed: DataFusion error: Arrow error: External error: Execution error:
Execution("ArrowError(IoError(\"No space left on device (os error 28)\",
Os { code: 28, kind: StorageFull, message: \"No space left on device\" }), Some(\"\"))")

The scheduler log shows 20 Job finished successfully events between 13:57:06 and 13:59:02. The executor log tail covers three of those completions (jobs ypJm7Ol, qPMmk2t, V3bztez, queries 18 through 20) and contains zero Remove data for job and zero Remove intermediate data for job lines. The only reclaim anywhere in the window is for the failed job, which goes through clean_up_failed_job and so calls the undelayed clean_up_job_data.

What makes this unrecoverable rather than merely suboptimal is the timing. finished_job_data_clean_up_interval_seconds defaults to 300 (ballista/scheduler/src/config.rs:302), and the full 21 query suite finishes in 116 seconds. With the immediate reclaim no-oping under AQE, the delayed whole-job cleanup never fires either, so shuffle output grows monotonically from query 1 through query 21 and nothing is ever freed during the run. For scale, one q18 stage was spilling around 235 MB per event across 4 partitions repeatedly, and the executor reported 69.7 GB available at registration.

Also worth flagging that #1981 is currently closed as completed, but #1982 only fixed the static planner path. The flake it describes is still reachable through this issue, and it is showing up as red X's on unrelated PRs again.