jayzhan211 opened a new pull request, #26987:
URL: https://github.com/apache/datafusion/pull/26987
## 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. When two runs don't fit, the
merge re-spills a run with every batch halved, tracks a smaller batch size for
each re-spilled run, and scans for runs of single rows (#22945, #23286,
#25383). Bounding the batches when the runs are written makes all of that
unnecessary.
#25910 and #25853 each bound spill batches for one operator, the sort and
the aggregate respectively. This PR applies one bound to every sorted run. It
takes the batch splitter from #25910 and the output cap approach from #25853,
and removes the halving re-spill. The merge's admission policy is unchanged.
## What changes are included in this PR?
The three commits can be reviewed separately.
1. `SpillManager::with_max_batch_bytes`: every batch written through the
manager is at most this many bytes. The exceptions are a batch of one row, and
a batch with a column that cannot be split by rows: dictionaries or run-end
encoded arrays below the top level, unions and list views. Larger batches are
split with `SpillPieces` (from #25910), which here also splits top-level
run-end encoded columns.
2. The sort and both aggregate streams set the bound, and the multi-level
merge writes its intermediate runs through the same manager:
- For the sort, the bound is `sort_spill_reservation_bytes / 4` (at least
64 KiB), so two runs with one batch read ahead fit the workspace the sort
already retains. This is the formula from #25910.
- The aggregate retains no workspace, so its bound is also at most an
eighth of the state of its first spill. The smallest merge, plus as much again
for the replay, then fits in what it held when memory first ran out. Later
spills, such as the remainder spilled at the end of the input, don't lower it.
- 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. Merges without a bound skip this estimate.
3. The merge no longer re-spills runs. Removed: `split_spill_file_in_half`,
`split_batch_in_half`, the per-run batch-size limit, the single-row scan with
its retry flags, `MergeStep` and `SpillFilesToMerge`. If two runs don't fit
with one batch read ahead, the merge returns `ResourcesExhausted`, with the
size it needed in the message. The aggregate's replay headroom keeps one
change: with one batch read ahead, the two runs a pass needs are admitted
against the pool alone, which on main happened only after re-spilling failed.
Also includes an upgrade guide entry.
## What is the testing strategy for this PR?
- **`spill_batch_bound.slt`:** spilled sorts and aggregates over wide rows,
checking results. On main, three of its five queries fail with
`ResourcesExhausted`: a sort of 64 KiB rows at 16 MB, the #25851 query scaled
to 4 groups at 12 MB, and a `string_agg` over 600-byte values at 4 MB. The
other two pass on both: a run-end encoded sort at 32 MB, and the legacy
aggregate stream at 4 MB.
- **Unit tests:**
- `SpillPieces`: row order and the size bound for strings, string views,
lists, skewed rows, top-level dictionaries and run-end encoded columns; one
oversized row; a nested dictionary written whole.
- Merge tests of the removed re-spill are removed. The others are kept;
those that called the re-spill directly now check the same outcomes without it.
- 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, 5, 6, 8, 12 MB | ok | ok |
| 4 × 200 B | 3 MB | `ResourcesExhausted` | ok |
| 4 × 200 B | 4, 5, 6, 8, 12 MB | ok | ok |
| 4 × 600 B | 3, 4, 5, 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>
With `enable_migration_aggregate = false` (the legacy stream), each of
these queries succeeds or fails exactly as on main.
- **Columns that are hard to split:** 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 and 64 MB and finish without spilling at 96 MB. The
same rows as run-end encoded strings succeed on both at 32, 48 and 64 MB.
- **Not done yet:**
- Benchmarks: `sort_tpch` and `external_aggr` under memory limits,
`sort_preserving_merge`, `spill_io`.
- Spilled sorts with union or list view columns.
## Are there any user-facing changes?
Query results don't change, and in every comparison above, a query that
succeeds on main also succeeds here. Memory behavior changes in two ways:
- 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` and `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 the bound instead
(see the upgrade guide).
--
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]
