jayzhan211 opened a new pull request, #26984:
URL: https://github.com/apache/datafusion/pull/26984
## Which issue does this PR close?
- Part of #26711.
## Rationale for this change
A spilled sort or aggregate fails with `ResourcesExhausted` when its rows
are wide, even though each row fits the pool. Two examples, both failing on
main (5db4b233ef) and passing with this PR:
- The reproducer in #25851 (`array_agg` over 16 groups, `datafusion-cli -m
128M`). Main fails with `Failed to allocate additional 179.9 MB for
FinalHashAggregateStream[0]`. This PR returns 16 groups and 3,000,000 values.
- The reproducer in #25910 (sort over 64 KiB rows, `-m 12m --mem-pool-type
fair`). Main fails with `Failed to allocate additional 20.0 MB for
ExternalSorterMerge[0]`. This PR finishes with 7 spills.
Both have the same cause. Sorted spill runs are written in batches of
`batch_size` rows, whatever their size in bytes, and the multi-level merge
reserves twice the largest batch of each run. The merge copes with this through
several mechanisms added one at a time: re-spilling a run with halved batches
and a per-run batch-size limit (#22945, #23286), a scan for single-row runs and
a retry without headroom (#25383), and two options that only the aggregate
sets, replay headroom (#25383) and intermediate merge sizing (#25428). #26711
asks to revisit these together. The open PRs #25910, #25853 and #26264 each add
another option to the same layer, for the sort writer, the aggregate, and sort
headroom respectively.
This draft tries the combined version: one byte bound for every sorted run,
and one admission rule for every caller. With those, the per-caller options and
the halving re-spill are no longer needed. It takes the batch splitter from
#25910, the output cap approach from #25853 and the admission formula from
#26264. I'm opening it as a draft to compare with those PRs.
## What changes are included in this PR?
- `SpillManager::with_max_batch_bytes`: every batch written through the
manager is at most this many bytes unless it holds one row. Larger batches are
split with `SpillPieces` (from #25910). The sort and both aggregate paths set
it, and the multi-level merge writes its intermediate runs through the same
manager.
- The bound, `max_spill_batch_bytes`:
- For the sort, `sort_spill_reservation_bytes / 4`, so two runs with one
batch read ahead fit the workspace it already retains (the formula from #25910).
- The aggregate retains no workspace, so its bound is also at most an
eighth of the state it held when it spilled. The smallest merge, plus as much
again for the replay, then fits in what it held.
- A merge of spilled runs ends an output batch at `batch_size` rows or at
about the bound in bytes, estimated from the average row size of each input
batch.
- One admission rule for sort and aggregate (the formula from #26264): each
run must leave as much again as the merge takes beyond the workspace it already
holds. Only the two runs a pass needs, with one batch read ahead, may use the
whole pool.
- Intermediate merge sizing (#25428) now applies to every caller, estimated
with the widest run and capped by `max_spill_merge_fan_in`.
- Removed:
- `with_replay_headroom` and `with_intermediate_merge_sizing` on both
merge builders
- the per-run batch-size limit
- `split_spill_file_in_half`, `split_batch_in_half`, and the single-row
scan with its retry flags
- `MergeStep` and `SpillFilesToMerge`
- `AggregateSpill`'s minimum batch-row tracking
## What is the testing strategy for this PR?
- **Unit tests:**
- Tests of the removed re-spill mechanism are removed.
- Tests of the merge's outcomes (output order, rewrite counts, fan-in,
headroom release) are kept, with the same expected values.
- New: `bounded_runs_merge_in_a_pool_smaller_than_two_whole_batches` and
`read_ahead_needs_headroom`.
- **Fuzz assertion:** in `spilling_fuzz_in_memory_constrained_env`, the
output batch count of a spilled sort is now checked with `>=` instead of `==`.
With `sort_spill_reservation_bytes = 1` the bound is 64 KiB, so the output
comes in more, smaller batches.
- **Local runs (all pass):**
- physical-plan lib tests
- core `memory_limit` tests
- extended fuzz tests (spilling, sort, aggregate,
`memory_limit_validation`)
- the full sqllogictest suite
- **Aggregate under small pools,** `datafusion-cli`, main vs this PR:
| Values per group | Memory limit | main | This PR |
|---|---|---|---|
| 4 × 40 B | 3, 4, 6, 8, 12 MB | ok | ok |
| 4 × 200 B | 3 MB | `ResourcesExhausted` | ok |
| 4 × 200 B | 4, 6, 8, 12 MB | ok | ok |
| 4 × 600 B | 3, 4, 6, 8 MB | `ResourcesExhausted` | ok |
| 4 × 600 B | 12 MB | ok | ok |
<details><summary>Query</summary>
```sql
SET datafusion.execution.target_partitions = 1;
-- 4 × 200 B: 40000 groups, repeat('x', 200), 160000 rows
SELECT count(*), sum(length(s)) FROM (
SELECT k, string_agg(v, ',') AS s
FROM (SELECT value % 40000 AS k, repeat('x', 200) AS v FROM
generate_series(1, 160000))
GROUP BY k
);
```
The 4 × 40 B row uses `value % 100000`, `repeat('x', 40)` and 400000 rows.
</details>
- **Unsplittable columns:** a sort over 64 KiB rows held in a `Struct` of a
`Dictionary` column, which `SpillPieces` writes whole, behaves as on main. Both
fail at 12-64 MB and finish without spilling at 96 MB and above.
- **Not done yet:**
- Benchmarks: `sort_tpch` and `external_aggr` under memory limits,
`sort_preserving_merge`, `spill_io`.
- Spilled sorts with union, run-end encoded or list view columns.
## Are there any user-facing changes?
No query results change. Memory behavior does:
- A spilled sort's merge now leaves headroom beyond its workspace, as #26264
proposes. Under memory pressure it may merge fewer runs per pass.
- Spilled sorts and aggregates with wide rows, or with a small
`sort_spill_reservation_bytes`, write and emit smaller batches.
- New public `SpillManager::with_max_batch_bytes`. Callers of
`StreamingMergeBuilder::with_sorted_spill_files` that write their own runs no
longer get the halving re-spill; they can set this bound instead.
--
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]
