[jayzhan211]

Which issue does this PR close?

• Part of [EPIC] Simplify memory-limited (spilling) execution #26711.

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 spill batches by bytes and remove the halving re-spill from the spill merge #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 spill batches by bytes and remove the halving re-spill from the spill merge #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:

  • a sort of 64 KiB rows at 16 MB
  • the Hash aggregate with a few large groups fails after spilling: spill batches are sized by rows, not bytes #25851 query scaled to 4 groups at 12 MB
  • a string_agg over 600-byte values at 4 MB

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

[adriangb]

@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 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