PDGGK opened a new pull request, #2753:
URL: https://github.com/apache/datafusion-ballista/pull/2753
# Which issue does this PR close?
Closes #1996.
# Rationale for this change
#1982 frees a job's intermediate shuffle files as soon as the job succeeds.
It picks them with `ExecutionGraph::intermediate_stage_ids()`, whose default
implementation treats a stage with no `output_links` as final.
The adaptive graph creates every stage in `create_resolved_stage` with
`output_links` set to `vec![]` ("we do not know output links at this moment")
and never fills it in. Under AQE every stage therefore looks final,
`intermediate_stage_ids()` returns an empty list, and
`clean_up_intermediate_job_data` returns early. The intermediate shuffle files
stay on the executors until the delayed whole-job cleanup,
`finished_job_data_clean_up_interval_seconds` later (300s by default). AQE is
the default planner now, so this is what a default deployment does. When
queries run back to back, as in the TPC-H SF10 CI job described in the issue,
shuffle data piles up until the disk is full (`No space left on device` on
query 21 with AQE on, while the AQE-off leg passed).
Reproduction: a standalone cluster (in-process scheduler and pull-staged
executor), `target_partitions = 4`, running
```sql
SELECT c1, MIN(c12), MAX(c12) FROM test WHERE c11 > 0.1 AND c11 < 0.9 GROUP
BY c1
```
on `examples/testdata/aggregate_test_100.csv`, then listing the stage
directories left in the executor work dir 5 seconds after the job succeeds.
Before, on `main`:
```
adaptive planner = true, result rows = 5
job=038D1KKB04000 stage=0
job=038D1KKB04000 stage=1
count = 2
adaptive planner = false, result rows = 5
job=038D1KKN7Y000 stage=2
count = 1
```
After:
```
adaptive planner = true, result rows = 5
job=038D1KGY5A000 stage=1
count = 1
adaptive planner = false, result rows = 5
job=038D1KH8CT000 stage=2
count = 1
```
With the static planner only the final stage is left; with AQE the
intermediate stage was left as well. After the change both planners leave only
the final stage, whose files the client reads.
# What changes are included in this PR?
- `AdaptiveExecutionGraph` overrides `intermediate_stage_ids()`. A stage is
final when the job's `output_locations` point at it (each `PartitionLocation`
carries the id of the stage that wrote it); every other stage in the graph is
intermediate. These are the same locations the client fetches and that
`DataFrame::cache()` records and pins, so the stage they point at is never
reclaimed.
- `output_locations` is only set once the planner has no more stages to run.
Until then the override returns an empty list, so a final stage is never
returned and the "never delete a final stage" rule from #1982 still holds.
- Two tests in `state/aqe/test/job_success.rs` drive an adaptive job to
success, a two-stage aggregation and a three-stage nested aggregation, and
check that `intermediate_stage_ids()` is empty while the job runs and
afterwards returns exactly the stages that are not in `output_locations`.
Without the change both fail with `left: []`, `right: [0]` and `left: []`,
`right: [0, 1]`.
This is the approach proposed in the issue. #2149 implemented it earlier and
was closed as stale without review.
Not changed here: `result_schema` in `flight_sql.rs` also finds the final
stage as the first stage with empty `output_links`, so under AQE it can pick a
non-final stage. It is behind the `flight-sql` feature and needs its own test,
so I left it for a separate change.
# Are there any user-facing changes?
No API or configuration changes. Jobs planned by the adaptive planner now
free their intermediate shuffle files when they succeed, as jobs planned by the
static planner already did.
# Testing
Run locally:
- `cargo test -p ballista-scheduler --lib job_success`: 2 passed. With the
`aqe/mod.rs` change reverted, both fail as shown above.
- `cargo test -p ballista-scheduler --lib -- state::aqe
state::execution_graph scheduler_server::query_stage_scheduler
state::task_manager state::cache`: 327 passed.
- The standalone reproduction above, before and after the change.
- `cargo fmt --all -- --check` and `cargo clippy -p ballista-scheduler
--all-targets -- -D warnings`.
Left to CI: the full workspace test suite, clippy with `--all-features`, and
the TPC-H/TPC-DS suites.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
