Skip to content

fix: bound the spill batches of an external sort by bytes - #25910

Draft
adriangb wants to merge 6 commits into
apache:mainfrom
adriangb:spill-max-batch-bytes
Draft

adriangb wants to merge 6 commits into
apache:mainfrom
adriangb:spill-max-batch-bytes

Conversation

@adriangb

@adriangb adriangb commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

An ORDER BY over wide rows fails with ResourcesExhausted under a memory limit, even when each input batch fits in the pool and the sort spills:

SET datafusion.execution.target_partitions = 1;
SET datafusion.execution.sort_spill_reservation_bytes = 1048576;
SET datafusion.execution.batch_size = 4;
CREATE TABLE wide AS
SELECT value AS id, concat(CAST(value AS VARCHAR), repeat('x', 65536)) AS payload
FROM generate_series(0, 511);
SET datafusion.execution.batch_size = 512;
SET datafusion.runtime.memory_limit = '16M';
SELECT count(*) FROM (SELECT * FROM wide ORDER BY id DESC);
-- main: Resources exhausted: Failed to allocate additional ... for ExternalSorterMerge[0]

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 it reads at least two runs. With wide rows, that reservation is larger than the pool.

What changes are included in this PR?

A SpillManager can bound the bytes of the batches it writes. The external sort sets that bound, and every merge pass of the sort ends its output batches at the same bound. Every column type can be split: each piece is copied into fresh buffers before it is measured. The halving re-spill of the merge stays as the fallback for single rows wider than the bound and for rows that share data larger than the bound.

Commits

Each commit builds and passes the datafusion-physical-plan tests. Review them in order.

  1. test: spill_batch_bound.slt with the sort queries of fix: write sorted spill runs in byte-bounded batches and drop the halving re-spill #26987. The wide sort records the error of main. The run-end encoded sort passes on main at 32 MB and guards against a regression.
  2. bench: a sort_spill benchmark.
  3. refactor: the writer returns WrittenBatchStats (memory size and rows) instead of a bare usize. No change in behavior.
  4. refactor: SpillManager::with_max_batch_bytes (pub(crate)), SpillPieces and compact_batch, with unit tests. No operator sets a bound yet. An unbounded writer records the same size as on main.
  5. fix: ExternalSorter sets the bound (sort_spill_reservation_bytes / 4, at least 64 KiB) and passes it to StreamingMergeBuilder::with_max_batch_bytes. The slt diff shows the error becoming a result, and the run-end encoded sort fits the same 16 MB limit. The fuzz check of the output batch count becomes >=. The wide-row integration tests (core/tests/memory_limit/wide_row_sort_spill.rs) are added here, because they fail on main.
  6. test: the oversized-batch fuzz test checks that the sort spills and finishes. New tests sort wide run-end encoded, dense union and list view rows.

How SpillPieces cuts a batch

gc = gc_view_arrays(batch)
if no bound                              -> yield gc, size = get_sliced_size (as on main)
if within bound or 1 row                 -> yield gc, size = recorded_size
else
  cut row ranges by SpilledRowSizes, within (bound - IPC padding of a piece)
  for each range
    piece = compact_batch(slice)   // fresh buffers that hold only its rows
    if recorded_size(piece) > bound, more than 1 row, and smaller than its parent:
      halve and retry
    else yield piece        // the caller writes it before the next one is built

recorded_size is get_sliced_size plus the IPC padding of each buffer, so a batch read back never exceeds the size recorded for the file. For a compacted piece, get_sliced_size is exact for every type.

compact_batch is arrow take plus fixes for what take keeps whole: view buffers (gc), unused dictionary values, and list view children. Rows of a list view can share a child range, so the copy keeps the covered range once, and halving stops when a piece does not get smaller. Without that stop, 1024 rows over one shared 64 KiB value became 1024 pieces of 64 KiB each.

Trade-offs:

  • Each piece of a split batch is a copy. On the sort_spill benchmark the cost was within noise (0 to 6%).
  • A dictionary that the rows of a split batch share is written once per piece. A batch within the bound is written as on main, so its dictionaries and views stay shared.

Alternatives that were prototyped and compared under one harness (13 types, 10 to 80 MB):

Design Result
Copy each piece (this PR) Every type passes at 20 MB and above
Measure slices by offsets, keep narrow_lists Fixes run-end encoded and sparse unions. Dense unions and list views still fail below 80 MB. The measure must match how the arrow IPC writer narrows each type.
Reserve after decode in the merge Needs the merge to keep with_bypass_mempool, and loses the guarantee that the merge keeps its admitted memory

Scope and follow-ups

Code Status Replaced by
SpillManager::with_max_batch_bytes, the sort bound, the merge output cap for the sort, WrittenBatchStats Long term
Cuts by SpilledRowSizes, the halving retry Temporary workaround arrow-rs compacted_size and row_sizes (apache/arrow-rs#12184)
compact_batch (take plus fixes) Temporary workaround arrow-rs minify (apache/arrow-rs#12184)
gc_view_arrays of an unsplit batch Temporary workaround the IPC option to write only the used view bytes (apache/arrow-rs#12039)
Halving re-spill and per-run limits of the multi-level merge Kept as the fallback for single rows and shared data wider than the bound Deleted in a later PR

Follow-ups, not in this PR:

  • Aggregate spill bound by the first spill (stacked PR, see below).
  • Public with_max_batch_bytes and an upgrade guide entry: only needed if the re-spill is deleted, because then custom operators that write their own runs must bound them.
  • Delete the halving re-spill and the per-run limits (about −1,000 lines), and add Jay's read_ahead_needs_headroom and bounded_runs_merge_in_a_pool_smaller_than_two_whole_batches tests with it. Every type can now be split, so this no longer waits on arrow-rs.
  • The no-regression matrix of fix: write sorted spill runs in byte-bounded batches and drop the halving re-spill #26987 (string_agg at 3 to 12 MB, struct of dictionary) belongs with the deletion.

What is the testing strategy for this PR?

  • spill_batch_bound.slt: the wide sort fails on main and returns its result here. The run-end encoded sort passes on main at 32 MB and here at 16 MB.
  • wide_row_sort_spill.rs: Utf8, Utf8View, BinaryView, List, dictionaries, nested views, and view slices that share buffers. Run-end encoded, dense union and list view columns. Each test checks that the sort spilled and that the result equals an unlimited sort.
  • Unit tests in in_progress_spill_file.rs write through a bounded manager, read the file back, and check every batch against max_record_batch_memory.
  • compacted_slice_holds_only_its_rows checks that a 10-row slice of each of 12 types compacts to about its rows and keeps its values. types_with_unsliced_children_split_within_bound and list_view_with_shared_child_copies_it_once_per_piece check the pieces.
  • test_spill_max_batch_bytes_fits_minimum_merge checks the bound, its floor, and usize::MAX.
  • The extended suite (--features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption) passes.

Are there any user-facing changes?

  • External sorts over wide rows finish under memory limits where they failed before.
  • Spilling sorts can write and emit more, smaller batches.
  • No public API or configuration changes. The merge keeps its re-spill, so no upgrade guide entry is needed.

🤖 Generated with Claude Code

@github-actions github-actions Bot added core Core DataFusion crate physical-plan Changes to the physical-plan crate labels Sep 30, 2026

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note

This is a first-pass review using an LLM. I plan on doing a more careful review later, so please treat these as things to check rather than settled conclusions. Results marked "measured" come from local runs at this PR's head (caa7814, debug builds), compared with the same code with the new bound turned off.

Summary: bounding spilled batches gets the final merge past its admission check, but the merge still builds batch_size-row batches outside the memory pool. So wide-row sorts that failed cleanly with ResourcesExhausted on main can now succeed while using several times the pool. Separately, .max(1) turns sort_spill_reservation_bytes = 0 into a 1-byte bound, which breaks an existing test.

Main concerns

1. The merge can now exceed the memory limit without failing (sort.rs:316)

Smaller pieces shrink what the multi-level merge reserves per run: 2 × max_record_batch_memory × buffer_len (multi_level_merge.rs:599). But each initial run is still registered with the full batch_size as its output limit (multi_level_merge.rs:195), and spill merges bypass the memory pool (multi_level_merge.rs:552). So BatchBuilder holds up to batch_size rows of pieces and then interleaves a batch_size-row output batch, none of which the pool sees.

Measured with the new wide_row_sort_spill setup (12 MiB FairSpillPool, 512 rows of 64 KiB, batch_size 512, 1 MiB reservation): main fails with Failed to allocate additional 20.0 MB for ExternalSorterMerge[0]. With this PR the query succeeds, but the final merge holds 32 MiB across 49 pieces plus a 32 MiB output batch, peaking at ~64 MiB allocated (about 5× the pool). By the same logic, the default batch_size of 8192 with ~1 MiB rows would build GiB-sized batches while reserving tens of MiB, turning a clean error into a possible OOM kill. On the sort path the re-spill limit (multi_level_merge.rs:843) comes from old_batch_size, so it doesn't cap this either. It seems the merge output (or each run's row limit) also needs a byte bound for the fix to hold.

2. sort_spill_reservation_bytes = 0 becomes a 1-byte bound (sort.rs:316)

.max(1) turns 0 (and 1) into a 1-byte bound, so split_ranges halves every spilled batch, including intermediate merge re-spills, down to one-row IPC messages. 0 is used across tests and examples, and the sort's OOM message suggests decreasing this setting (sort.rs:898). Measured:

  • sql::runtime_config::test_max_temp_directory_size_enforcement fails at this PR's head with "Should not fail due to max temp directory size limit" (runtime_config.rs:224), and passes with the bound off. One-row Int64 messages take ~160 bytes per row instead of ~8, so the spill no longer fits in 1M.
  • test_external_sort_zero_merge_reservation goes from 2.9 s to 198 s, and test_memory_limit_with_spill from 4.9 s to 192 s.
  • test_sort_skewed_batches_spill (reservation 1) still passes but no longer reaches split_spill_file_in_half (0 calls vs 16), so it silently loses the coverage it was written for.

Treating 0 as "no bound", or applying a floor, would avoid this.

3. Batches the estimate accepts are written uncompacted (in_progress_spill_file.rs:281)

The no-split paths (line 281, and ranges.len() == 1 at line 286) decide using referenced_batch_size, then write and record the batch as is, with no take and no dictionary GC. Sort chunks come from take, which keeps a dictionary's full values, so a chunk whose used values fit the bound is written and recorded with the whole dictionary. Measured with a 1 MiB bound and a Dictionary<Int32, Utf8> sharing 256 × 64 KiB values: a 15-row batch is recorded at 16 MiB, while a 16-row batch is split into 0.5 MiB pieces. Sliced lists and maps (e.g. the halves from split_batch_in_half) are likewise recorded with the parent's whole child.

4. Nested dictionaries get halved down to one-row pieces (in_progress_spill_file.rs:369)

referenced_size narrows dictionaries at any depth, but compact_piece only GCs top-level dictionary columns, take keeps a nested dictionary's values whole, and compact_range never checks that halving actually shrinks a piece. Measured: Struct{d: Dictionary<Int32, Utf8>} with 16 MiB of values and a 1 MiB bound is written as 256 one-row messages, each still recorded at ~16 MiB, at ~4× the write time. The same data as a top-level dictionary gives ~0.5 MiB pieces.

5. A bound equal to the headroom doesn't make the minimum merge fit (sort.rs:312)

The comment says the merge "reserves memory for the largest batch of each spill file", but it reserves 2 × max_record_batch_memory × buffer_len per stream and needs at least two streams. Halving leaves the largest piece anywhere in (S/2, S], so a two-stream merge can need up to 4S, where S is sort_spill_reservation_bytes. When other partitions hold the rest of the pool, the merge can still fail at multi_level_merge.rs:660 (from reading the code, not measured). A bound of S/4 would fit the minimum merge in the retained headroom.

Other issues

6. Each split batch is fully copied before any piece is written (in_progress_spill_file.rs:289). split_for_spill take-copies every column of every piece into a Vec that lives until the last write, so spilling holds an extra, untracked copy of the batch just when memory is exhausted. For non-view columns this is new: main encoded uncompressed buffers by reference. Only list and map columns need their child narrowed, and compacting and writing one piece at a time would cap the extra memory at one piece.

7. A bounded SpillManager loses rows in spill_pool (spill_manager.rs:74). with_max_batch_bytes is public, but spill_pool counts one batch per push and stops reading a finished file at that count (spill_pool.rs:258, spill_pool.rs:607). Measured: passing a bounded manager to spsc_channel returned 8 of 128 rows, with no error. SMJ's restore path also reads a single batch per spill file (materializing_stream.rs:1022). Nothing in-tree combines these today, but the doc comment should at least warn about it.

8. ListView, dense Union and RunEndEncoded columns aren't narrowed (in_progress_spill_file.rs:443). They fall through to get_slice_memory_size, which counts the whole child. When other columns force a split, each ListView piece writes the full child (take keeps it, and arrow-ipc writes ListView children unsliced): a 15 MB ListView child plus 200 MB of Utf8 would spill 440 MB instead of 215 MB (from reading the code). Batches dominated by a dense Union or REE column are never split, even though take would compact them.

9. A full dictionary GC runs just to measure size (in_progress_spill_file.rs:415). This happens for every batch an ExternalSorter spills (the bound is always set), for both halves at every split level, and again in compact_range. Each call allocates a remap vector the size of the whole dictionary and copies the used values, then discards them. occupancy() plus the value lengths would give the same number without allocating.

10. max_batch_rows is counted before splitting (spill_manager.rs:181). With a bounded manager the returned memory is per written piece, but the rows are per appended batch, so size_intermediate_merges (multi_level_merge.rs:305) would stay on for runs with short batches. This is latent until an aggregate opts in.

Tests

11. The new integration tests (wide_row_sort_spill.rs:106) only check that the query completes with ids in order. They don't assert that the sort spilled (plan_spill_count) or that memory stayed within the pool, so they pass today while the merge builds a 32 MiB output batch against a 12 MiB pool. Per the testing guide, this also looks like it could be a sqllogictest: the repro in the PR description is already SQL, and with target_partitions = 1 the pool from SET datafusion.runtime.memory_limit should behave like FairSpillPool here.

12. The unit tests (in_progress_spill_file.rs:573) say each piece "passes the reader's size check", but that check only logs at debug! (spill/mod.rs:395), and round_trip never compares read-back sizes with the recorded max_record_batch_memory. The assertion from spill/mod.rs:1430 would cover the property the merge relies on. Also, the comment on repeated_view_values_are_split_and_kept_in_order says values are "deduplicated within each piece", but nothing here deduplicates.

Smaller things

  • view_values_size (in_progress_spill_file.rs:458) reimplements arrow's GenericByteViewArray::total_buffer_bytes_used(). It could be array.len() * VIEW_SIZE_BYTES + array.total_buffer_bytes_used().
  • The compact_range doc (in_progress_spill_file.rs:296) says the second pass exists because view builders allocate with spare capacity, but arrow 60's gc() allocates exactly total_buffer_bytes_used(). In practice the pass only does real work for small view buffers that skip GC (10 KiB or less) and for nested dictionaries (see 4).

@codecov-commenter

codecov-commenter commented Sep 30, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 74.78873% with 179 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.59%. Comparing base (6924c6a) to head (7dec81f).
⚠️ Report is 378 commits behind head on main.

Files with missing lines Patch % Lines
.../physical-plan/src/spill/in_progress_spill_file.rs 72.84% 103 Missing and 67 partials ⚠️
...atafusion/physical-plan/src/spill/spill_manager.rs 80.00% 2 Missing and 2 partials ⚠️
...afusion/physical-plan/src/sorts/streaming_merge.rs 75.00% 3 Missing ⚠️
...usion/physical-plan/src/sorts/multi_level_merge.rs 33.33% 0 Missing and 2 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #25910      +/-   ##
==========================================
- Coverage   82.60%   82.59%   -0.02%     
==========================================
  Files        1144     1144              
  Lines      442793   443590     +797     
  Branches   442793   443590     +797     
==========================================
+ Hits       365779   366388     +609     
- Misses      54863    54981     +118     
- Partials    22151    22221      +70     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@adriangb

Copy link
Copy Markdown
Contributor Author

Thanks @andygrove! Addressed in 0bc3b6c and 7dec81f (description updated):

  1. The merge of spill files now also caps output batches by the spill bound, estimated from average row size. Still not reserved in the pool, so skewed rows can exceed it (noted under limitations).
  2. 0 now means no bound; the affected tests pass at their old speed.
  3. Every written piece is compacted, including unsplit batches.
  4. Compaction is recursive, and halving stops when it doesn't shrink pieces.
  5. Bound is now max(reservation / 4, 1 MiB).
  6. Only dictionaries, view arrays and list views are copied; other columns stay zero-copy slices. Pieces are still compacted before the first write; happy to switch to one-at-a-time if you think it matters.
  7. with_max_batch_bytes is now pub(crate), documented as only for readers that read to the end.
  8. ListView is now handled; dense Union and REE are still written whole (listed as a limitation).
  9. Used dictionary values are measured with a bitmap, no GC.
  10. Rows are counted per written piece.
  11. The integration tests now assert that the sort spilled and that merged batches stay within the bound. Kept in Rust since sqllogictest can't check those.
  12. round_trip now checks read-back sizes against max_record_batch_memory; stale comment fixed. Also switched to total_buffer_bytes_used() and rewrote the compaction docs.

@2010YOUY01

Copy link
Copy Markdown
Member

@adriangb

adriangb commented Oct 1, 2026

Copy link
Copy Markdown
Contributor Author

They overlap quite a bit. We should probably discuss them together. Each PR has two parts: what gets written to the spill file, and what the merge builds from it.

The merge-output parts of all three do the same thing, and #25565 does it best. I propose we remove the merge cap from this PR so it covers only the write side, with #25565 as the merge-output fix. We need both: this PR alone removes the error but leaves the merge output outside the pool, and #25565 alone cannot start the merge because the spill batches are too large. They touch different code, so they can be reviewed in parallel and merged in either order.

Then we can change #25853 so the aggregate uses the SpillManager bound and does not have its own mechanism. That one has to come after this PR, because it needs the new API.

Wdyt?

@Weijun-H

Copy link
Copy Markdown
Member

They overlap quite a bit. We should probably discuss them together. Each PR has two parts: what gets written to the spill file, and what the merge builds from it.

* **[Limit spill merge output batches by byte target #25565](https://github.com/apache/datafusion/pull/25565)** fixes only the merge output. It reserves the output memory in the pool before building each batch.

* **[fix: size aggregate spill batches by bytes as well as rows #25853](https://github.com/apache/datafusion/pull/25853)** bounds aggregate spill batches through the sort iterator, and also caps merge output batches.

* **This PR** bounds spill batches in `SpillManager`, so any operator can opt in. Only the sort does in this PR. It also caps merge output batches, but by an estimate outside the pool.

The merge-output parts of all three do the same thing, and #25565 does it best. I propose we remove the merge cap from this PR so it covers only the write side, with #25565 as the merge-output fix. We need both: this PR alone removes the error but leaves the merge output outside the pool, and #25565 alone cannot start the merge because the spill batches are too large. They touch different code, so they can be reviewed in parallel and merged in either order.

Then we can change #25853 so the aggregate uses the SpillManager bound and does not have its own mechanism. That one has to come after this PR, because it needs the new API.

Wdyt?

The split makes sense, but #25565 already passes the Utf8 repro and doesn't cover List/Utf8View output accounting. Before removing the other merge caps, we need that coverage; moving aggregate batching into SpillManager also needs to preserve the bound before take builds the batch and track the actual written batch sizes.

let size = referenced_batch_size(batch)?;
let mut ranges = Vec::new();
split_ranges(batch, size, max_batch_bytes, &mut ranges)?;
ranges.iter().map(compact_piece).collect()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Preserve shared ListView children when splitting

For 1,024 ListView rows sharing one 64 KiB child, a 1 MiB limit expands a 74 KB input into 68 MB of simultaneously live pieces before the first write, outside the memory pool; the spill file also grows from 74 KB to 67 MB. Preserve overlapping child ranges and compact/write pieces incrementally; incremental writing alone leaves the disk and merge-reservation amplification.

Theodus pushed a commit to Theodus/datafusion that referenced this pull request Oct 15, 2026
## Which issue does this PR close?

- Closes apache#17737.
- Related: apache#8777, apache#22676, apache#10337.

## Rationale for this change

DataFusion inlines every CTE reference. When a query references the same
CTE several times, the CTE body is planned and executed once for each
reference. Postgres lets users opt out of this with `WITH x AS
MATERIALIZED (...)`. `sqlparser` already parses the directive into
`Cte.materialized`, but the SQL planner ignores it.

Earlier attempts (apache#22551, apache#22675) bundled automatic materialization,
heuristics and filter pushdown. This PR is limited to the explicit
directive, as proposed in
apache#17737 (comment):

- No config flag.
- No heuristic.
- Queries without `MATERIALIZED` plan exactly as before.

Deciding automatically when to materialize can be discussed separately.

Example:

```sql
WITH r AS MATERIALIZED (SELECT random() AS v)
SELECT a.v = b.v FROM r a, r b;
-- true: the body runs once. With NOT MATERIALIZED (or no directive) each reference evaluates random() separately.
```

## What changes are included in this PR?

- **Planner:** `datafusion/sql/src/cte.rs` and `query.rs` plan a
`MATERIALIZED` CTE as a `MaterializedCte` node, whose children are the
CTE body and the rest of the query, plus one `MaterializedCteScan` leaf
for each reference. These are `UserDefinedLogicalNode`s in
`datafusion-expr`.
  - `NOT MATERIALIZED` and the default keep today's inlining.
- `MATERIALIZED` also applies in a `WITH RECURSIVE` block. A recursive
CTE is still planned as `RecursiveQuery`, and its reference to itself
stays a work table scan.
- `head_output_expr`, `PullUpCorrelatedExpr` and `PushDownFilter` look
through `MaterializedCte` to its continuation, so `IN` / `EXISTS`
subqueries and filters from an outer query work as with an inlined CTE.
The body is never filtered.
- Each occurrence of a CTE in a plan gets its own id before physical
planning, so a view that is referenced two times runs its CTE two times,
as in PostgreSQL.
- **Physical plan:** `MaterializedCteExec` and `MaterializedCteScanExec`
in `datafusion-physical-plan`, and a physical planner pass that binds
each scan to its producer.
- **Execution:** the body is run to completion once, all body partitions
in parallel, on the first poll of any scan. It is started by one shared
future and runs in a `JoinSet` that is aborted when dropped. Every scan
then replays the buffered batches.
- **Why full buffering:** a bounded streaming fan-out deadlocks when,
for example, both sides of one hash join read the same CTE (see apache#8777).
This design cannot deadlock.
- **Memory:** buffered batches are tracked by a `MemoryReservation`
created with `can_spill(true)`. When the reservation cannot grow, the
rest of that partition is written to a spill file through
`SpillManager`. `SpillFile::read_stream` can be called by any number of
readers at the same time, so every scan replays both the in-memory
batches and the spill files. Scans read spill files without read-ahead,
and after buffering one decoded batch per concurrent reader is reserved.
If the pool cannot grant this, in-memory batches are moved to spill
files until it can. Each scan reserves memory for the largest batch on
disk, so an in-memory batch that is larger than the batches on disk is
written in pieces, and a move is made only when it releases more memory
than it adds to the scans.
apache#25910 adds a byte bound to
spill files with a better size model; when it merges, this operator can
use it and drop its own slicing. A body that reads an unbounded source
is rejected at planning time, because it would be buffered forever.

### Notes on earlier review feedback (apache#22675)

- **CTE names are not unique:** each scan binds to its producer through
an id allocated by the planner, not through the CTE name. Sibling
subqueries that declare CTEs with the same name are tested.
- **Alignment with `ScalarSubqueryExec`:** `MaterializedCteExec` follows
the same shape. The rest of the query is its output and passes through
input order, partitioning and statistics. A scan inside a scalar
subquery can start the body itself (tested).
- **Why not `WorkTable`:** `WorkTable` gives its batches to exactly one
consumer (`take()`) and returns an error when memory runs out;
`RecursiveQueryExec` does not spill. A materialized CTE needs N readers
and spilling. The scan does follow the same pattern as `WorkTableExec`:
a leaf that reads a handle bound at planning time.
- **Statistics:** the scans report unknown statistics for now. A leaf
cannot ask its producer for statistics through `StatisticsContext`. I
would like to do this in a follow-up. This is a known limitation that
can make a plan worse than the plan of the inlined CTE, not only leave
it the same: `JoinSelection` does not know that a scan is small, and can
put the other, larger input of a hash join on the build side (pinned in
`cte_materialized.slt`).

### Out of scope, possible follow-ups

Tracked under the EPIC apache#22676:

- Protobuf serialization. These nodes have no proto encoding yet, so
serializing a plan that contains them fails in the default extension
codec (apache#26231).
- Statistics, ordering and partitioning on the scans (apache#26230).
- Release the buffer when the last scan finishes rather than when the
plan is dropped (apache#26233).
- Byte-bounded spill files (apache#26232).
- Keep the output per run, so that `reset_state` follows its contract
(apache#26229).
- Deciding automatically when to materialize (apache#26237).
- First-class `LogicalPlan` variants (apache#26238).
- Pushing filters that all references share into the CTE body (apache#26239).
- Pruning unused columns from the CTE body (apache#26240).
- `sqlparser` parses `[NOT] MATERIALIZED` only in `PostgreSqlDialect`,
so the tests set that dialect (apache/datafusion-sqlparser-rs#2622).
- A streaming fan-out for consumers that are known to be independent.

## Are these changes tested?

Yes. The new file
`datafusion/sqllogictest/test_files/cte_materialized.slt` covers:

- a volatile body, which returns one value when `MATERIALIZED` and
differs per reference when `NOT MATERIALIZED` or with no directive;
- `EXPLAIN`, which shows one `DataSourceExec` for three references;
- `EXPLAIN ANALYZE`, where the body buffers 4 rows once and each of the
three scans outputs 4 rows;
- a hash join whose two sides both read the CTE, with no deadlock;
- a scalar subquery reference;
- a CTE that reads an earlier CTE;
- sibling CTEs with the same name;
- `IN`, `NOT IN` and correlated `EXISTS` / `NOT EXISTS` / `IN`
subqueries that declare a materialized CTE;
- `WITH RECURSIVE` with recursive and non-recursive materialized CTEs;
- a view referenced two times (separate CTEs) and one reference with two
scans (one shared CTE);
- a filter from an outer query, which is pushed into the continuation;
- column aliases;
- a spill under `memory_limit`: 200k rows, about 184k spilled (part of
the in-memory prefix moves to disk to make room for the two readers),
and both scans still return 200k rows.
- a spill where the first batch is larger than the batches on disk: the
query finishes under an 18 MB limit, and failed before with an empty
pool.

The full sqllogictest suite passes.

## Are there any user-facing changes?

Yes. `WITH ... AS MATERIALIZED` is now honored when the Postgres dialect
is used. Queries without the directive are unchanged. There is no new
config option and no public API break. The new nodes are additions.

## Open questions for reviewers

1. Should these be a new `LogicalPlan` variant, like `RecursiveQuery`,
or stay as built-in `UserDefinedLogicalNode`s, as in this PR, which
keeps the change smaller?
2. Is it acceptable for protobuf serialization of these nodes to fail
until a follow-up PR adds it?

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
adriangb added a commit to pydantic/datafusion that referenced this pull request Oct 16, 2026
- Closes apache#17737.
- Related: apache#8777, apache#22676, apache#10337.

DataFusion inlines every CTE reference. When a query references the same
CTE several times, the CTE body is planned and executed once for each
reference. Postgres lets users opt out of this with `WITH x AS
MATERIALIZED (...)`. `sqlparser` already parses the directive into
`Cte.materialized`, but the SQL planner ignores it.

Earlier attempts (apache#22551, apache#22675) bundled automatic materialization,
heuristics and filter pushdown. This PR is limited to the explicit
directive, as proposed in
apache#17737 (comment):

- No config flag.
- No heuristic.
- Queries without `MATERIALIZED` plan exactly as before.

Deciding automatically when to materialize can be discussed separately.

Example:

```sql
WITH r AS MATERIALIZED (SELECT random() AS v)
SELECT a.v = b.v FROM r a, r b;
-- true: the body runs once. With NOT MATERIALIZED (or no directive) each reference evaluates random() separately.
```

- **Planner:** `datafusion/sql/src/cte.rs` and `query.rs` plan a
`MATERIALIZED` CTE as a `MaterializedCte` node, whose children are the
CTE body and the rest of the query, plus one `MaterializedCteScan` leaf
for each reference. These are `UserDefinedLogicalNode`s in
`datafusion-expr`.
  - `NOT MATERIALIZED` and the default keep today's inlining.
- `MATERIALIZED` also applies in a `WITH RECURSIVE` block. A recursive
CTE is still planned as `RecursiveQuery`, and its reference to itself
stays a work table scan.
- `head_output_expr`, `PullUpCorrelatedExpr` and `PushDownFilter` look
through `MaterializedCte` to its continuation, so `IN` / `EXISTS`
subqueries and filters from an outer query work as with an inlined CTE.
The body is never filtered.
- Each occurrence of a CTE in a plan gets its own id before physical
planning, so a view that is referenced two times runs its CTE two times,
as in PostgreSQL.
- **Physical plan:** `MaterializedCteExec` and `MaterializedCteScanExec`
in `datafusion-physical-plan`, and a physical planner pass that binds
each scan to its producer.
- **Execution:** the body is run to completion once, all body partitions
in parallel, on the first poll of any scan. It is started by one shared
future and runs in a `JoinSet` that is aborted when dropped. Every scan
then replays the buffered batches.
- **Why full buffering:** a bounded streaming fan-out deadlocks when,
for example, both sides of one hash join read the same CTE (see apache#8777).
This design cannot deadlock.
- **Memory:** buffered batches are tracked by a `MemoryReservation`
created with `can_spill(true)`. When the reservation cannot grow, the
rest of that partition is written to a spill file through
`SpillManager`. `SpillFile::read_stream` can be called by any number of
readers at the same time, so every scan replays both the in-memory
batches and the spill files. Scans read spill files without read-ahead,
and after buffering one decoded batch per concurrent reader is reserved.
If the pool cannot grant this, in-memory batches are moved to spill
files until it can. Each scan reserves memory for the largest batch on
disk, so an in-memory batch that is larger than the batches on disk is
written in pieces, and a move is made only when it releases more memory
than it adds to the scans.
apache#25910 adds a byte bound to
spill files with a better size model; when it merges, this operator can
use it and drop its own slicing. A body that reads an unbounded source
is rejected at planning time, because it would be buffered forever.

- **CTE names are not unique:** each scan binds to its producer through
an id allocated by the planner, not through the CTE name. Sibling
subqueries that declare CTEs with the same name are tested.
- **Alignment with `ScalarSubqueryExec`:** `MaterializedCteExec` follows
the same shape. The rest of the query is its output and passes through
input order, partitioning and statistics. A scan inside a scalar
subquery can start the body itself (tested).
- **Why not `WorkTable`:** `WorkTable` gives its batches to exactly one
consumer (`take()`) and returns an error when memory runs out;
`RecursiveQueryExec` does not spill. A materialized CTE needs N readers
and spilling. The scan does follow the same pattern as `WorkTableExec`:
a leaf that reads a handle bound at planning time.
- **Statistics:** the scans report unknown statistics for now. A leaf
cannot ask its producer for statistics through `StatisticsContext`. I
would like to do this in a follow-up. This is a known limitation that
can make a plan worse than the plan of the inlined CTE, not only leave
it the same: `JoinSelection` does not know that a scan is small, and can
put the other, larger input of a hash join on the build side (pinned in
`cte_materialized.slt`).

Tracked under the EPIC apache#22676:

- Protobuf serialization. These nodes have no proto encoding yet, so
serializing a plan that contains them fails in the default extension
codec (apache#26231).
- Statistics, ordering and partitioning on the scans (apache#26230).
- Release the buffer when the last scan finishes rather than when the
plan is dropped (apache#26233).
- Byte-bounded spill files (apache#26232).
- Keep the output per run, so that `reset_state` follows its contract
(apache#26229).
- Deciding automatically when to materialize (apache#26237).
- First-class `LogicalPlan` variants (apache#26238).
- Pushing filters that all references share into the CTE body (apache#26239).
- Pruning unused columns from the CTE body (apache#26240).
- `sqlparser` parses `[NOT] MATERIALIZED` only in `PostgreSqlDialect`,
so the tests set that dialect (apache/datafusion-sqlparser-rs#2622).
- A streaming fan-out for consumers that are known to be independent.

Yes. The new file
`datafusion/sqllogictest/test_files/cte_materialized.slt` covers:

- a volatile body, which returns one value when `MATERIALIZED` and
differs per reference when `NOT MATERIALIZED` or with no directive;
- `EXPLAIN`, which shows one `DataSourceExec` for three references;
- `EXPLAIN ANALYZE`, where the body buffers 4 rows once and each of the
three scans outputs 4 rows;
- a hash join whose two sides both read the CTE, with no deadlock;
- a scalar subquery reference;
- a CTE that reads an earlier CTE;
- sibling CTEs with the same name;
- `IN`, `NOT IN` and correlated `EXISTS` / `NOT EXISTS` / `IN`
subqueries that declare a materialized CTE;
- `WITH RECURSIVE` with recursive and non-recursive materialized CTEs;
- a view referenced two times (separate CTEs) and one reference with two
scans (one shared CTE);
- a filter from an outer query, which is pushed into the continuation;
- column aliases;
- a spill under `memory_limit`: 200k rows, about 184k spilled (part of
the in-memory prefix moves to disk to make room for the two readers),
and both scans still return 200k rows.
- a spill where the first batch is larger than the batches on disk: the
query finishes under an 18 MB limit, and failed before with an empty
pool.

The full sqllogictest suite passes.

Yes. `WITH ... AS MATERIALIZED` is now honored when the Postgres dialect
is used. Queries without the directive are unchanged. There is no new
config option and no public API break. The new nodes are additions.

1. Should these be a new `LogicalPlan` variant, like `RecursiveQuery`,
or stay as built-in `UserDefinedLogicalNode`s, as in this PR, which
keeps the change smaller?
2. Is it acceptable for protobuf serialization of these nodes to fail
until a follow-up PR adds it?

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
(cherry picked from commit cdc2fc0)
adriangb added a commit to pydantic/datafusion that referenced this pull request Oct 16, 2026
## Which issue does this PR close?

- Closes apache#17737.
- Related: apache#8777, apache#22676, apache#10337.

## Rationale for this change

DataFusion inlines every CTE reference. When a query references the same
CTE several times, the CTE body is planned and executed once for each
reference. Postgres lets users opt out of this with `WITH x AS
MATERIALIZED (...)`. `sqlparser` already parses the directive into
`Cte.materialized`, but the SQL planner ignores it.

Earlier attempts (apache#22551, apache#22675) bundled automatic materialization,
heuristics and filter pushdown. This PR is limited to the explicit
directive, as proposed in
apache#17737 (comment):

- No config flag.
- No heuristic.
- Queries without `MATERIALIZED` plan exactly as before.

Deciding automatically when to materialize can be discussed separately.

Example:

```sql
WITH r AS MATERIALIZED (SELECT random() AS v)
SELECT a.v = b.v FROM r a, r b;
-- true: the body runs once. With NOT MATERIALIZED (or no directive) each reference evaluates random() separately.
```

## What changes are included in this PR?

- **Planner:** `datafusion/sql/src/cte.rs` and `query.rs` plan a
`MATERIALIZED` CTE as a `MaterializedCte` node, whose children are the
CTE body and the rest of the query, plus one `MaterializedCteScan` leaf
for each reference. These are `UserDefinedLogicalNode`s in
`datafusion-expr`.
  - `NOT MATERIALIZED` and the default keep today's inlining.
- `MATERIALIZED` also applies in a `WITH RECURSIVE` block. A recursive
CTE is still planned as `RecursiveQuery`, and its reference to itself
stays a work table scan.
- `head_output_expr`, `PullUpCorrelatedExpr` and `PushDownFilter` look
through `MaterializedCte` to its continuation, so `IN` / `EXISTS`
subqueries and filters from an outer query work as with an inlined CTE.
The body is never filtered.
- Each occurrence of a CTE in a plan gets its own id before physical
planning, so a view that is referenced two times runs its CTE two times,
as in PostgreSQL.
- **Physical plan:** `MaterializedCteExec` and `MaterializedCteScanExec`
in `datafusion-physical-plan`, and a physical planner pass that binds
each scan to its producer.
- **Execution:** the body is run to completion once, all body partitions
in parallel, on the first poll of any scan. It is started by one shared
future and runs in a `JoinSet` that is aborted when dropped. Every scan
then replays the buffered batches.
- **Why full buffering:** a bounded streaming fan-out deadlocks when,
for example, both sides of one hash join read the same CTE (see apache#8777).
This design cannot deadlock.
- **Memory:** buffered batches are tracked by a `MemoryReservation`
created with `can_spill(true)`. When the reservation cannot grow, the
rest of that partition is written to a spill file through
`SpillManager`. `SpillFile::read_stream` can be called by any number of
readers at the same time, so every scan replays both the in-memory
batches and the spill files. Scans read spill files without read-ahead,
and after buffering one decoded batch per concurrent reader is reserved.
If the pool cannot grant this, in-memory batches are moved to spill
files until it can. Each scan reserves memory for the largest batch on
disk, so an in-memory batch that is larger than the batches on disk is
written in pieces, and a move is made only when it releases more memory
than it adds to the scans.
apache#25910 adds a byte bound to
spill files with a better size model; when it merges, this operator can
use it and drop its own slicing. A body that reads an unbounded source
is rejected at planning time, because it would be buffered forever.

### Notes on earlier review feedback (apache#22675)

- **CTE names are not unique:** each scan binds to its producer through
an id allocated by the planner, not through the CTE name. Sibling
subqueries that declare CTEs with the same name are tested.
- **Alignment with `ScalarSubqueryExec`:** `MaterializedCteExec` follows
the same shape. The rest of the query is its output and passes through
input order, partitioning and statistics. A scan inside a scalar
subquery can start the body itself (tested).
- **Why not `WorkTable`:** `WorkTable` gives its batches to exactly one
consumer (`take()`) and returns an error when memory runs out;
`RecursiveQueryExec` does not spill. A materialized CTE needs N readers
and spilling. The scan does follow the same pattern as `WorkTableExec`:
a leaf that reads a handle bound at planning time.
- **Statistics:** the scans report unknown statistics for now. A leaf
cannot ask its producer for statistics through `StatisticsContext`. I
would like to do this in a follow-up. This is a known limitation that
can make a plan worse than the plan of the inlined CTE, not only leave
it the same: `JoinSelection` does not know that a scan is small, and can
put the other, larger input of a hash join on the build side (pinned in
`cte_materialized.slt`).

### Out of scope, possible follow-ups

Tracked under the EPIC apache#22676:

- Protobuf serialization. These nodes have no proto encoding yet, so
serializing a plan that contains them fails in the default extension
codec (apache#26231).
- Statistics, ordering and partitioning on the scans (apache#26230).
- Release the buffer when the last scan finishes rather than when the
plan is dropped (apache#26233).
- Byte-bounded spill files (apache#26232).
- Keep the output per run, so that `reset_state` follows its contract
(apache#26229).
- Deciding automatically when to materialize (apache#26237).
- First-class `LogicalPlan` variants (apache#26238).
- Pushing filters that all references share into the CTE body (apache#26239).
- Pruning unused columns from the CTE body (apache#26240).
- `sqlparser` parses `[NOT] MATERIALIZED` only in `PostgreSqlDialect`,
so the tests set that dialect (apache/datafusion-sqlparser-rs#2622).
- A streaming fan-out for consumers that are known to be independent.

## Are these changes tested?

Yes. The new file
`datafusion/sqllogictest/test_files/cte_materialized.slt` covers:

- a volatile body, which returns one value when `MATERIALIZED` and
differs per reference when `NOT MATERIALIZED` or with no directive;
- `EXPLAIN`, which shows one `DataSourceExec` for three references;
- `EXPLAIN ANALYZE`, where the body buffers 4 rows once and each of the
three scans outputs 4 rows;
- a hash join whose two sides both read the CTE, with no deadlock;
- a scalar subquery reference;
- a CTE that reads an earlier CTE;
- sibling CTEs with the same name;
- `IN`, `NOT IN` and correlated `EXISTS` / `NOT EXISTS` / `IN`
subqueries that declare a materialized CTE;
- `WITH RECURSIVE` with recursive and non-recursive materialized CTEs;
- a view referenced two times (separate CTEs) and one reference with two
scans (one shared CTE);
- a filter from an outer query, which is pushed into the continuation;
- column aliases;
- a spill under `memory_limit`: 200k rows, about 184k spilled (part of
the in-memory prefix moves to disk to make room for the two readers),
and both scans still return 200k rows.
- a spill where the first batch is larger than the batches on disk: the
query finishes under an 18 MB limit, and failed before with an empty
pool.

The full sqllogictest suite passes.

## Are there any user-facing changes?

Yes. `WITH ... AS MATERIALIZED` is now honored when the Postgres dialect
is used. Queries without the directive are unchanged. There is no new
config option and no public API break. The new nodes are additions.

## Open questions for reviewers

1. Should these be a new `LogicalPlan` variant, like `RecursiveQuery`,
or stay as built-in `UserDefinedLogicalNode`s, as in this PR, which
keeps the change smaller?
2. Is it acceptable for protobuf serialization of these nodes to fail
until a follow-up PR adds it?

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @adriangb, bounding spill batches by bytes is the right direction for wide-row sorts. I read through the change and ran the existing spill:: and sorts:: unit tests on this branch (216 pass). I did not run benchmarks, so only the zero-width FixedSizeList item below was reproduced with a probe. The rest come from reading the code, so please treat them as questions.

Behavior changes

  1. Final sort output batches get smaller. The byte bound is applied to the final merge output in multi_level_merge.rs, so it also shrinks what the sort returns. With the default 10MB sort_spill_reservation_bytes the bound is 2.5MB, so 8192 rows of ~500 B would come out as ~5000-row batches. Downstream writers and coalescing see more, smaller batches. Because max_batch_rows can now be below batch_size_limit in the intermediate-spill bookkeeping, size_intermediate_merges also looks like it gets turned off for those sorts. Is that intended? If not, could the bound apply only to intermediate spill files?
  2. Per-row byte estimate in sorts/builder.rs. It is get_record_batch_memory_size(batch) / num_rows of the whole input batch. That counts full shared buffers, so sliced or shared-buffer inputs (e.g. Utf8View or Binary arrays sharing one allocation) overestimate. The merge could then think it hit max_batch_bytes after a few rows and emit very small batches. Rows within a batch also vary, which the average ignores.
  3. The sort_spill_reservation_bytes / 4 bound in sorts/sort.rs relies on get_reserved_bytes_for_record_batch_size being size + sliced_size and on a minimum fan-in of 2. It is floored at 1 MiB, so for reservations below 4 MiB (the test uses 1 MiB) four times the bound exceeds the headroom. The minimum merge could then take extra memory from the pool and hit ResourcesExhausted, the case this PR is fixing. A value of 0 silently disables the bound. Only the unit test would catch a change to get_reserved_bytes_for_record_batch_size.
  4. Zero-width FixedSizeList (reproduced with a probe). compact_array rebuilds with FixedSizeListArray::try_new, which takes the length from values.len() / size. For a FixedSizeList(_, 0) column with no null buffer and a child needing compaction (dictionary or view), the length comes back as 0. That produces a column that doesn't match the row count. It is rare, but the length should be passed explicitly.

Performance

  1. Dictionaries are no longer shared. Every spilled batch now has its dictionary columns pruned into a private dictionary, even when no split happens, and the default config always sets a bound. Before, all batches in a spill file shared one dictionary, which the IPC reader decoded once, so interleave could reuse it and copy only keys. Now each batch carries its own dictionary, so interleave has to merge or concatenate values and IPC resends the dictionary per message. I expect high-cardinality dictionary sorts that spill to get slower and use more memory. Could you add a benchmark or test for a spilled sort on a dictionary column, and consider pruning only when a split happens?
  2. list_view_size calls referenced_size on a one-row child slice for every row, and split_ranges repeats that at each recursion level. For a list-view with a dictionary child, each call goes through dictionary_values_size, which allocates and scans a bitmap of the full dictionary length. That is O(rows × dict_len) per level.
  3. Repeated size measurement. split_ranges measures both halves at every level, then compact_piece compacts and measures each piece again, so each batch is traversed about log2(rows) + 2 times. SpillManager::spilled_batch_size compacts a whole batch only to get a number that is then discarded.

Design

  1. About 450 lines of hand-written size model and compaction partly duplicate gc_view_arrays and Arrow's own sizing (get_slice_memory_size). compact_view re-implements view gc that gc_view_arrays already provides. That is two copies of the view-compaction rules plus a model of what the IPC writer encodes, which has to stay in sync with arrow-ipc. Can we reuse the existing helpers, or split the sizing work into a follow-up so the bugfix is easier to review?
  2. SpillManager::with_max_batch_bytes changes the one-append-one-batch contract, and only a doc comment protects other consumers. The constructor is pub(crate), so a reader that assumes one batch per append (spill_pool, SMJ buffered batches) would silently break if handed the sort's manager. A separate writer type or a parameter on append_batch would make that misuse impossible.

I'd consider 1 and 5 the main things to settle before merging. Happy to help test or benchmark any of these.

@adriangb
adriangb force-pushed the spill-max-batch-bytes branch from 7dec81f to 782ef12 Compare October 26, 2026 00:25
@adriangb

Copy link
Copy Markdown
Contributor Author

Reworked as discussed: this PR now changes only what is written to spill files. #25565 sizes the merge output. Production code is now about 260 lines (was about 950). The commits are split for review: failing test → refactor → bench → fix → test updates.

@andygrove:

# Status
1, 2 Merge output cap removed. Output batches and size_intermediate_merges are as on main.
3 Bound is computed from MIN_SPILL_MERGE_FAN_IN × get_reserved_bytes_for_record_batch_size. Floor is 64 KiB, also for 0.
4 Custom compaction removed. Pieces are copied with MutableArrayData.
5 A batch within the bound is written as on main, so its dictionary stays shared. New sort_spill bench, dictionary case: 35 ms (main) → 21–25 ms.
6, 7 Removed. Each piece is copied and measured once. It is halved again only if it is still too large.
8 Size model removed. Uses get_sliced_size, gc_view_arrays and garbage_collect_any_dictionary.
9 The bound is set on the writer, not on SpillManager. spill_pool and SMJ cannot get split batches.
spill_manager
    .create_in_progress_file("Sorting")?
    .with_max_batch_bytes(Some(bound))

@Weijun-H:

Test change: test_sort_with_limited_memory_and_oversized_record_batch now checks that the sort spills and finishes. The merge no longer needs the half-batch re-spill, which keeps its unit tests in multi_level_merge.rs.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Shared spill-writer compaction and memory accounting need human validation, with dictionary performance and regression coverage concerns still unresolved.

2 open issues
What changed in this PR

This PR bounds batches written to external-sort spill files by bytes, allowing wide-row sorts to merge under tight memory limits. It complements the separate merge-output sizing work in #25565.

Changes:

  • Adds writer-local byte limits with row-preserving splitting and compaction.
  • Applies limits to initial and intermediate sort spills, reporting actual written-batch statistics.
  • Adds wide-row regression tests, spill round-trip coverage, and benchmarks.
File Description
datafusion/​physical-plan/​src/​spill/​spill_pieces.rs Implements byte-bounded splitting and compaction.
datafusion/​physical-plan/​src/​spill/​spill_manager.rs Passes writer limits and collects piece statistics.
datafusion/​physical-plan/​src/​spill/​mod.rs Registers the splitting module.
datafusion/​physical-plan/​src/​spill/​in_progress_spill_file.rs Writes bounded pieces and adds round-trip tests.
datafusion/​physical-plan/​src/​sorts/​streaming_merge.rs Forwards limits to intermediate spill merges.
datafusion/​physical-plan/​src/​sorts/​sort.rs Derives and applies sort spill byte limits.
datafusion/​physical-plan/​src/​sorts/​multi_level_merge.rs Bounds intermediate and fallback spill writes.
datafusion/​physical-plan/​Cargo.toml Adds dictionary-compaction dependency and benchmark registration.
datafusion/​physical-plan/​benches/​sort_spill.rs Benchmarks dictionary, narrow-string, and wide-string spills.
datafusion/​core/​tests/​memory_limit/​wide_row_sort_spill.rs Tests wide-row sorts under constrained memory.
datafusion/​core/​tests/​memory_limit/​mod.rs Registers wide-row regression tests.
datafusion/​core/​tests/​fuzz_cases/​spilling_fuzz_in_memory_constrained_env.rs Checks successful spilling without requiring the old fallback.
Cargo.lock Records the added arrow-select dependency.

Review effort: Balanced

Ok(piece) => piece,
Err(e) => return Some(Err(e)),
};
let size = match piece.get_sliced_size() {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this record an upper bound on decoded size, including IPC padding? With 4,096 UInt8 columns × 17 rows and a 64 KiB bound, the recorded maximum is 36,864 B but decoding needs 98,304 B. Two runs reserve 147,456 B yet hold 196,608 B. The sizing gap predates this PR, but splitting can exhaust the existing 2× margin.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 1182cb9: the recorded size now counts per-buffer IPC padding (schema alignment, or 64 B when compressed). Added your example as a test, plus a check across 6 types × {none, LZ4, Zstd}.

<= sort_spill_reservation_bytes
};
// Binary search for the largest batch size that fits: `fits` is true up to it.
let (mut fitting, mut too_large) = (0, sort_spill_reservation_bytes + 1);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Large reservations overflow here (e.g. 9223372036854775807 on 64-bit targets). Could this use (reservation / (2 * MIN_SPILL_MERGE_FAN_IN)).max(MIN_SPILL_MAX_BATCH_BYTES) instead? Division avoids both the search and the overflow; please add a boundary test.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes sense to me. I'll convert this PR to draft for now. @adriangb, please mark it ready for review once this is addressed.

@adriangb adriangb Nov 2, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in d8a922f, using your formula. It gives the same result as the search. Added a boundary test for usize::MAX.

@jayzhan211
jayzhan211 marked this pull request as draft November 2, 2026 13:04
@adriangb
adriangb force-pushed the spill-max-batch-bytes branch 2 times, most recently from 91fc305 to 1182cb9 Compare November 2, 2026 17:55

@adriangb adriangb left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

self review

pub(crate) async fn append_batch_async_with_stats(
&mut self,
batch: &RecordBatch,
) -> Result<(usize, usize)> {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's use a named newtype struct here to avoid mixing up our usizes


impl Iterator for SpillPieces {
/// A piece and its size, as [`GetSlicedSize`] measures it.
type Item = Result<(RecordBatch, usize)>;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A named newtype struct here would be nice

/// Copies a slice so that it holds only its own rows: list children are narrowed to the
/// range the slice uses, view arrays are garbage collected, and top-level dictionaries drop
/// the values that no key uses.
fn compact(slice: &RecordBatch) -> Result<RecordBatch> {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Surely there must be other copies / implementations of this functionality or similar in datafusion and/or arrow-rs.

}

/// Drops the values of top-level dictionaries that no key uses.
fn prune_dictionaries(batch: &RecordBatch) -> Result<RecordBatch> {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly here, I'm surprised this functionality doesn't exist. IIRC it does, and I also remember us putting in place bounds like "only gc / compact if more than X% of the space is wasted". Maybe that matters less here since spilling by nature involves "slow down drastically to be able to actually finish", but still if we can avoid copying all of the data for a marginal gain the tradeoff of e.g. inflating the reservation on the other end might be worth it.

/// views and unions, whose children a row range does not bound, run-end encoded arrays, or
/// dictionaries below the top level, whose values it does not prune. A batch with such a
/// column is written whole.
fn can_split(schema: &Schema) -> bool {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we sure these are not handled? Again maybe some existing shared functionality would handle them. Even if not I'm not convinced this list is accurate.

Comment on lines +122 to +124
/// The capacity granularity of arrow's `MutableBuffer`, which a compressed buffer is
/// decompressed into.
const DECOMPRESSED_ALIGNMENT: usize = 64;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm also surprised there is no existing functionality in arrow-rs or datafusion to calculate this sort of padding / alignment of IPC files. We should double check, and consider if this is the right place to put it. It might be a better fit for arrow-rs.

@jayzhan211

Copy link
Copy Markdown
Contributor

Feel free to ping me when this is ready for review.

@adriangb
adriangb force-pushed the spill-max-batch-bytes branch 2 times, most recently from 79cb02e to 4253f74 Compare November 10, 2026 17:26
@adriangb adriangb changed the title fix: bound spill file batches by bytes so external sorts over wide rows can merge fix: bound spill batches by bytes and remove the halving re-spill from the spill merge Nov 10, 2026

See [issue #18235](https://github.com/apache/datafusion/issues/18235) for details.

### Spill merges no longer re-spill runs with halved batches

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this upgrade guide is necessary.

adriangb and others added 3 commits November 11, 2026 12:46
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 <jayzhan211@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adds a sort_spill benchmark with three cases: a dictionary column with
100,000 distinct values, short strings, and 1 KiB strings. Each case uses a
memory pool smaller than the input, so the sort spills and merges.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SpillManager records the row count of the appended batch as the largest
batch of the file. A later commit writes one appended batch as several
smaller ones, so the writer must report what it wrote. A named struct
keeps the memory size and the row count apart, where a tuple of two
usize values does not.

This commit adds WrittenBatchStats and the *_with_stats append methods.
No change in behavior.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@adriangb
adriangb force-pushed the spill-max-batch-bytes branch from 4253f74 to 1f63322 Compare November 11, 2026 20:48
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 <jayzhan211@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@adriangb adriangb changed the title fix: bound spill batches by bytes and remove the halving re-spill from the spill merge fix: bound the spill batches of an external sort by bytes Nov 11, 2026
adriangb and others added 3 commits November 12, 2026 00:24
A spill batch has at most batch_size rows, whatever its size in bytes. A
merge of spill files reserves memory for the largest batch of each file,
so wide rows make that reservation larger than the pool. An operator
needs a way to bound the batches it spills by bytes.

This commit adds SpillManager::with_max_batch_bytes (crate-private) and
SpillPieces. A bounded writer cuts each appended batch into row ranges
by SpilledRowSizes, copies each range, and writes it as its own IPC
message. The recorded size of a bounded writer counts the IPC padding of
each buffer, so a batch read back never exceeds it.

A slice of a list, run-end encoded array, union or list view keeps whole
children or values, so its measured size is far above what the IPC reader
allocates for it. compact_batch copies each range into fresh buffers
that hold only its rows, so the measure is exact for every type and no
type has to be written whole. It uses arrow take, plus fixes for what
take keeps whole: view buffers, unused dictionary values and list view
children. Rows of a list view can share a child range, so it copies the
covered range once, and the halving of a piece stops when the halves do
not get smaller. arrow-rs minify (apache/arrow-rs#12184) replaces the
copy.

A dictionary that the rows of a split batch share is written once per
piece. A batch within the bound is written as before.

No operator sets a bound yet, and an unbounded writer records the same
size as before, so this commit changes no behavior.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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 <jayzhan211@gmail.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ll column types

The oversized-batch fuzz test asserted that the merge re-spilled runs to
add output batches. The sort now bounds its spill batches, so the merge
seats two runs without a re-spill. The test now checks that the sort
spills and finishes. The re-spill keeps its unit tests in
multi_level_merge.rs.

A slice of a run-end encoded array, a dense union or a list view keeps
whole children or values. New integration tests sort wide rows of each
type under a memory limit and check the result against an unbounded
sort.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@adriangb
adriangb force-pushed the spill-max-batch-bytes branch from 1f63322 to e843b11 Compare November 12, 2026 07:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants