Skip to content

fix: write sorted spill runs in byte-bounded batches and drop the halving re-spill - #26987

Draft
jayzhan211 wants to merge 3 commits into
apache:mainfrom
jayzhan211:spill-batch-bound-a
Draft

jayzhan211 wants to merge 3 commits into
apache:mainfrom
jayzhan211:spill-batch-bound-a

Conversation

@jayzhan211

@jayzhan211 jayzhan211 commented Nov 1, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

In short: sorted spill runs are now written in batches bounded by bytes, not only by rows, so the multi-level merge always knows how much memory two runs need. That fixes two ResourcesExhausted failures on wide rows, and lets the merge drop its halving re-spill.

A spilled sort or aggregate fails when its rows are wide, even though each row fits the pool. Both of these fail on main (5db4b23) and pass with this PR:

Reproducer main This PR
#25851: array_agg over 16 groups, datafusion-cli -m 128M Failed to allocate additional 179.9 MB for FinalHashAggregateStream[0] 16 groups, 3,000,000 values
#25910: sort over 64 KiB rows, -m 12m --mem-pool-type fair Failed to allocate additional 20.0 MB for ExternalSorterMerge[0] finishes, 7 spills

Cause: sorted spill runs are written in batches of batch_size rows, whatever their size in bytes. The multi-level merge reserves twice the largest batch of each run it reads. 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 (sort) and #25853 (aggregate) each bound spill batches for one operator. This PR applies one bound to every sorted run: it takes the batch splitter from #25910 and the output cap approach from #25853. The merge's admission policy is otherwise unchanged; #26984 is a follow-up that unifies it.

What changes are included in this PR?

How to review

The three commits are meant to be read in order.

1. feat: let SpillManager bound the size of the batches it writes (+428/−15, of which 155 are tests)

  • spill/spill_pieces.rs: SpillPieces is fix: bound the spill batches of an external sort by bytes #25910's splitter, plus top-level run-end encoded columns, which compact rebuilds with take.
  • spill/in_progress_spill_file.rs: append_batch and append_batch_async write the pieces and return the size of the largest one.
  • spill/spill_manager.rs: the bound lives on the manager (with_max_batch_bytes), so every writer it creates is bounded.
  • Worth checking:
    • Pieces keep row order.
    • The returned size is what a reader decodes.
    • Columns that can't be split by rows are written whole: dictionaries or run-end encoded arrays below the top level, unions and list views.

2. fix: write sorted spill runs in byte-bounded batches (+280/−18)

  • sorts/multi_level_merge.rs: max_spill_batch_bytes(merge_memory) is merge_memory / 4, at least 64 KiB. Two runs with one batch read ahead, each batch charged twice its size, fit merge_memory.
  • sorts/sort.rs, ExternalSorter::new: the sort's bound is max_spill_batch_bytes(sort_spill_reservation_bytes), the workspace it already retains (fix: bound the spill batches of an external sort by bytes #25910's formula).
  • aggregates/spill.rs, bound_spill_batches_by_first_spill, called from both aggregate streams:
    • The aggregate retains no workspace, so its bound is also max_spill_batch_bytes(first_spill / 2): the smallest merge, plus as much again for the replay, fits what it held when memory first ran out.
    • Later spills, such as the remainder at the end of the input, don't lower it.
  • sorts/builder.rs and sorts/merge.rs: a merge of spilled runs ends an output batch at batch_size rows or at about the bound in bytes. The size is estimated from each input batch's average row size, only when a bound is set.

3. refactor: remove the spill merge's halving re-spill (+225/−997, mostly sorts/multi_level_merge.rs)

  • 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.
  • Worth checking:
    • In get_sorted_spill_files_to_merge, if two runs don't fit with one batch read ahead, the merge now returns ResourcesExhausted, with the size it needed in the message.
    • The aggregate's replay headroom changes in one place: with one batch read ahead, the two runs a pass needs are admitted against the pool alone. On main that happened only after re-spilling failed.
  • Also adds 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:

    The other two, a run-end encoded sort at 32 MB and the legacy aggregate stream at 4 MB, pass on both builds.

  • 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. Tests that called it 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: in spilling_fuzz_in_memory_constrained_env, a spilled sort's output batch count 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 passing: physical-plan lib tests, core memory_limit tests, extended fuzz tests (spilling, sort, aggregate, memory_limit_validation), the full sqllogictest suite, and clippy.

  • Main vs this PR with datafusion-cli. No query that succeeds on main fails here:

    Query Memory limit main This PR
    string_agg, 4 × 40 B per group 3, 4, 5, 6, 8, 12 MB ok ok
    string_agg, 4 × 200 B per group 3 MB ResourcesExhausted ok
    same 4, 5, 6, 8, 12 MB ok ok
    string_agg, 4 × 600 B per group 3, 4, 5, 6, 8 MB ResourcesExhausted ok
    same 12 MB ok ok
    sort, 64 KiB run-end encoded strings 32, 48, 64 MB ok ok
    sort, 64 KiB strings in a Struct of a Dictionary (written whole) 12, 64 MB ResourcesExhausted ResourcesExhausted
    same 96 MB ok, no spill ok, no spill

    On the legacy aggregate stream (enable_migration_aggregate = false), each string_agg query succeeds or fails exactly as on main at every limit above.

    Aggregate query
    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
    );

    4 × 40 B uses value % 100000, repeat('x', 40) and 400000 rows. 4 × 600 B uses repeat('x', 600).

  • 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. 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).

@github-actions github-actions Bot added documentation Improvements or additions to documentation core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) physical-plan Changes to the physical-plan crate labels Nov 1, 2026
@adriangb

adriangb commented Nov 2, 2026

Copy link
Copy Markdown
Contributor

@jayzhan211 this seems to overlap heavily with #25910. This seems to have taken some commits from #25910 but has not kept up with fixes. In #25910 (comment) I proposed basically what you are doing here, but I wonder if we shouldn't seek to e.g. land #25910 first and then do followups instead of trying to roll it all (including the deletion of spill merge's halving re-spill) into 1 PR? Thoughts on how to coordinate this PR with the other open PRs it overlaps / re-implements?

@jayzhan211

jayzhan211 commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor Author

@jayzhan211 this seems to overlap heavily with #25910. This seems to have taken some commits from #25910 but has not kept up with fixes. In #25910 (comment) I proposed basically what you are doing here, but I wonder if we shouldn't seek to e.g. land #25910 first and then do followups instead of trying to roll it all (including the deletion of spill merge's halving re-spill) into 1 PR? Thoughts on how to coordinate this PR with the other open PRs it overlaps / re-implements?

Sure! I plan to review #25910 and other overlapping PRs, then check what is left for #26711

adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 10, 2026
An aggregate spills sorted runs and merges them for replay. Its runs are
bounded only by the producer target of 1 MiB, so in a small pool two runs
can need more than the aggregate held, and the merge has to re-spill.

This commit sets the aggregate's bound on its `SpillManager`, so each run
is written as byte-bounded pieces. On the first spill it lowers the bound
so that the smallest merge, plus as much again for the replay, fits in
half of what the aggregate held (the approach of apache#26987).
The aggregate writes through `append_batch_with_stats`, so the merge
sizes its intermediate batches from the exact rows of each piece.

The replay test allows 1 KiB over the pool limit. The merge grows a second
reservation of the aggregate's consumer while the first one holds a few
bytes, and `FairSpillPool` checks each reservation against the fair share
on its own. Smaller batches let the merge reach that share exactly.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 10, 2026
…pill merge

The multi-level merge re-spills a run at half its batch size when two runs
do not fit in memory, and it tracks a limit per run to drive that. The sort
and the aggregate now bound their spill batches so that the smallest merge
fits, so the fallback is no longer needed.

This commit deletes the re-spill, the per-run limits and their tests. The
merge keeps the output cap of apache#25853. It takes the change
from apache#26987, adapted to keep that cap. The upgrade guide
says to keep spilled batches to at most a quarter of the merge memory.

Two aggregate memory call-sequence snapshots lose one `shrink`: the merge
no longer releases its reservation to try a re-spill.

Co-Authored-By: Jay Zhan <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 11, 2026
An ORDER BY over 64 KiB rows fails with ResourcesExhausted under a 16 MB
limit, although each input batch fits in the pool and the sort spills.
The sort writes spill batches of batch_size rows, and the final merge
reserves twice the largest batch of each of at least two runs.

This commit adds spill_batch_bound.slt. It records the error of the wide
sort on main, so the fix shows the error becoming a result. A run-end
encoded sort at 32 MB passes on main and guards that the fix keeps it
passing. The queries come from apache#26987.

Co-authored-by: Jay Zhan <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 11, 2026
An ORDER BY over wide rows fails with ResourcesExhausted under a memory
limit, even when each input batch fits and the sort spills. The sort
writes spill batches of batch_size rows, whatever their size in bytes.
The final merge reserves twice the largest batch of each run and reads
at least two runs, so with wide rows that reservation exceeds the pool.

ExternalSorter now sets a bound on its SpillManager: the largest batch
for which the smallest merge (MIN_SPILL_MERGE_FAN_IN runs, each charged
twice its batch) fits in sort_spill_reservation_bytes, at least 64 KiB.
It also passes the bound to StreamingMergeBuilder::with_max_batch_bytes,
so every merge pass ends its output batches at the same bound. The
multi-level merge keeps its halving re-spill for batches that cannot be
split yet, such as run-end encoded columns.

The wide sort in spill_batch_bound.slt returns its result instead of the
error. A spilled sort can emit more, smaller batches, so the fuzz check
of the output batch count becomes a lower bound.

The bound formula follows a review suggestion on apache#25910,
and the merge output cap and the fuzz change follow
apache#26987.

Co-authored-by: Jay Zhan <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 12, 2026
An ORDER BY over wide rows fails with ResourcesExhausted under a memory
limit, even when each input batch fits and the sort spills. The sort
writes spill batches of batch_size rows, whatever their size in bytes.
The final merge reserves twice the largest batch of each run and reads
at least two runs, so with wide rows that reservation exceeds the pool.

ExternalSorter now sets a bound on its SpillManager: the largest batch
for which the smallest merge (MIN_SPILL_MERGE_FAN_IN runs, each charged
twice its batch) fits in sort_spill_reservation_bytes, at least 64 KiB.
It also passes the bound to StreamingMergeBuilder::with_max_batch_bytes,
so every merge pass ends its output batches at the same bound. The
multi-level merge keeps its halving re-spill, for single rows and rows
that share data larger than the bound.

The wide sort in spill_batch_bound.slt returns its result instead of the
error, and the run-end encoded sort now fits the same 16 MB limit. A
spilled sort can emit more, smaller batches, so the fuzz check
of the output batch count becomes a lower bound.

The bound formula follows a review suggestion on apache#25910,
and the merge output cap and the fuzz change follow
apache#26987.

Co-authored-by: Jay Zhan <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 14, 2026
A string_agg over 40000 groups of 200 or 600 byte values fails with
ResourcesExhausted under a 2 MB limit. The aggregate spills runs of
1 MiB batches, and the smallest merge of two runs does not fit in what
the aggregate held.

This commit adds the aggregate queries of apache#26987 to
spill_batch_bound.slt. They pass on main and guard against regressions:
array_agg over a few large groups (apache#25851), string_agg at 4 MB, and the
legacy aggregate stream. It also adds the two 2 MB queries, recorded
with the error of main.

Co-authored-by: Jay Zhan <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
adriangb added a commit to adriangb/datafusion that referenced this pull request Nov 14, 2026
A string_agg over many groups fails with ResourcesExhausted under a 2 MB
limit. The aggregate spills runs of up to 1 MiB batches, and the replay
merge reserves twice the largest batch of each of at least two runs.
That is more than the aggregate held when it first spilled.

This commit sets a bound on the SpillManager of both aggregate streams:
the 1 MiB replay target, or less when sort_spill_reservation_bytes fits
only smaller batches (the config bound of apache#26987). On
the first spill, the aggregate lowers the bound so that the smallest
merge, plus as much again for the replay, fits in half of what it held.
Later spills, such as the remainder at the end of the input, do not lower
it. The aggregate writes through append_batch_with_stats, so the merge
sizes its intermediate batches from the exact rows of each piece. The
legacy stream also passes the bound to its merge.

The two 2 MB queries in spill_batch_bound.slt return their results.

The replay test allows 1 KiB over the pool limit. The merge grows a
second reservation of the aggregate's consumer while the first one
holds a few bytes, and FairSpillPool checks each reservation against
the fair share on its own. Smaller batches let the merge reach that
share exactly. The pool behavior exists on main.

Co-authored-by: Jay Zhan <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate documentation Improvements or additions to documentation physical-plan Changes to the physical-plan crate sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants