Repository navigation
Conversation
github-actions
Bot
added
core
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
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_enforcementfails 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_reservationgoes from 2.9 s to 198 s, andtest_memory_limit_with_spillfrom 4.9 s to 192 s.test_sort_skewed_batches_spill(reservation 1) still passes but no longer reachessplit_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'sGenericByteViewArray::total_buffer_bytes_used(). It could bearray.len() * VIEW_SIZE_BYTES + array.total_buffer_bytes_used().- The
compact_rangedoc (in_progress_spill_file.rs:296) says the second pass exists because view builders allocate with spare capacity, but arrow 60'sgc()allocates exactlytotal_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 Report❌ Patch coverage is 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. 🚀 New features to boost your workflow:
|
|
Thanks @andygrove! Addressed in 0bc3b6c and 7dec81f (description updated):
|
|
What's the difference between this PR and |
|
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 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() |
There was a problem hiding this comment.
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.
## 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>
- 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)
## 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
left a comment
andygrove
left a comment
There was a problem hiding this comment.
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
- 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 10MBsort_spill_reservation_bytesthe 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. Becausemax_batch_rowscan now be belowbatch_size_limitin the intermediate-spill bookkeeping,size_intermediate_mergesalso looks like it gets turned off for those sorts. Is that intended? If not, could the bound apply only to intermediate spill files? - Per-row byte estimate in
sorts/builder.rs. It isget_record_batch_memory_size(batch) / num_rowsof 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 hitmax_batch_bytesafter a few rows and emit very small batches. Rows within a batch also vary, which the average ignores. - The
sort_spill_reservation_bytes / 4bound insorts/sort.rsrelies onget_reserved_bytes_for_record_batch_sizebeingsize + sliced_sizeand 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 hitResourcesExhausted, the case this PR is fixing. A value of 0 silently disables the bound. Only the unit test would catch a change toget_reserved_bytes_for_record_batch_size. - Zero-width
FixedSizeList(reproduced with a probe).compact_arrayrebuilds withFixedSizeListArray::try_new, which takes the length fromvalues.len() / size. For aFixedSizeList(_, 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
- 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
interleavecould reuse it and copy only keys. Now each batch carries its own dictionary, sointerleavehas 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? list_view_sizecallsreferenced_sizeon a one-row child slice for every row, andsplit_rangesrepeats that at each recursion level. For a list-view with a dictionary child, each call goes throughdictionary_values_size, which allocates and scans a bitmap of the full dictionary length. That is O(rows × dict_len) per level.- Repeated size measurement.
split_rangesmeasures both halves at every level, thencompact_piececompacts and measures each piece again, so each batch is traversed about log2(rows) + 2 times.SpillManager::spilled_batch_sizecompacts a whole batch only to get a number that is then discarded.
Design
- About 450 lines of hand-written size model and compaction partly duplicate
gc_view_arraysand Arrow's own sizing (get_slice_memory_size).compact_viewre-implements view gc thatgc_view_arraysalready 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? SpillManager::with_max_batch_byteschanges the one-append-one-batch contract, and only a doc comment protects other consumers. The constructor ispub(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 onappend_batchwould 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.
Adds integration tests that sort 512 rows of 64 KiB (Utf8, Utf8View and List) with a 12 MiB pool. Each input batch fits in the pool, but the sort writes batch_size-row spill batches that do not, so the final merge fails with ResourcesExhausted. These tests fail until the next commits. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SpillManager recorded the row count of the appended batch as the largest batch of the file. The writer now returns the row count of what it wrote, so the next commit can write one appended batch as several smaller ones. No change in behavior. 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>
…ws can merge The sort wrote spill batches of batch_size rows whatever their size. The final merge reserves twice the largest batch of each spill file, so with wide rows it could not start. InProgressSpillFile::with_max_batch_bytes now splits a batch larger than the bound into row ranges. Each range is copied so that it holds only its own rows, and the ranges are written one at a time. A batch within the bound is written as before, so its dictionaries stay shared. Batches with list views, unions, run-end encoded arrays or nested dictionaries are written whole. The bound is set per writer, not on SpillManager, so readers that expect one batch per append (spill_pool, sort-merge join) cannot get a bounded file. The sort sets it for its spill files and for the intermediate spill files of its merge. The bound is the largest batch for which the minimum merge (two files, get_reserved_bytes_for_record_batch_size each) fits in sort_spill_reservation_bytes, and at least 64 KiB. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The sort now bounds its spill batches, so the merge can seat two streams without the fallback that re-spills a run at half the batch size. The test now checks that the sort spills and finishes. The fallback keeps its unit tests in multi_level_merge.rs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Covers view, list, map, struct and dictionary columns, and the cases that are written whole: one wide row, batches within the bound (dictionary kept), list views with shared children and nested dictionaries. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
7dec81f to
782ef12
Compare
Which issue does this PR close?
Rationale for this change
An
ORDER BYover wide rows fails withResourcesExhaustedunder a memory limit, even when each input batch fits in the pool and the sort spills.Run with
datafusion-cli -m 12m --mem-pool-type fair -f repro.sql:On
main:With this PR the sort finishes.
The sort writes spill batches of up to
batch_sizerows, whatever their size in bytes. The final merge reserves twice the largest batch of each spill file, and it reads at least two files. With wide rows, that reservation is larger than the pool, so the merge cannot start. Smaller input batches do not help, because the sort buildsbatch_size-row batches again before it writes them.What changes are included in this PR?
A spill writer can now have a byte bound. The sort sets one.
SpillPieces(new filespill/spill_pieces.rs) decides the pieces:sizeis the existingget_sliced_sizemeasure. A schema "cannot be split" when it has a list view, a union, a run-end encoded array, or a dictionary below the top level. A row range does not bound what these columns hold, so the batch is written whole.InProgressSpillFile::with_max_batch_bytesSpillManager. Sospill_pooland the sort-merge join, which expect one batch per append, cannot get a bounded file from a shared manager.ExternalSorterspill_max_batch_bytesMIN_SPILL_MERGE_FAN_INfiles ×get_reserved_bytes_for_record_batch_size) fits insort_spill_reservation_bytes, found from that formula. The bound is at least 64 KiB, also when the reservation is 0.append_batch_async_with_statsThe commits are split for review:
test:the wide-row integration tests (fail onmain).refactor:the writer returns the row count of what it wrote (no change in behavior).bench:asort_spillbenchmark, so that you can run it before and after the fix.fix:the change above (the tests from commit 1 pass).test:one fuzz test checked for the old re-spill fallback. It now checks the outcome (see the testing section).test:unit tests for the split.What is the testing strategy for this PR?
Before:
wide_row_sort_spill.rs(Utf8, Utf8View and List payloads) fails onmain:After: the three tests pass. Each checks that the sort spilled and that the rows come back in order.
Unit tests in
in_progress_spill_file.rswrite through a bounded file, read it back, and check that each batch read back is within the recordedmax_record_batch_memory(plus the reader'sSPILL_BATCH_MEMORY_MARGIN):Utf8Viewarray, repeated view values, a dictionary with large values, slices of list, map and struct-with-list columns, keys that are all null;test_spill_max_batch_bytes_fits_minimum_mergechecks that the bound is the largest one whose minimum merge fits in the headroom, and the floor for a small or zero headroom.test_sort_with_limited_memory_and_oversized_record_batch(fuzz) used to check that the merge re-spilled a run at half the batch size. With bounded spill batches the merge does not need that fallback, so the test now checks that the sort spills and finishes. The fallback keeps its unit tests inmulti_level_merge.rs(skewed_runs_are_respilled_so_the_merge_fitsand others).Benchmark:
cargo bench -p datafusion-physical-plan --features test_utils --bench sort_spill. 16 batches of 8,192 rows. Each case spills. Non-LTO bench profile, Apple M-series laptop, two runs on each side:mainThe dictionary and short-string cases show that sorts that do not split their batches do not get slower. I did not profile why the dictionary case is faster. The 1 KiB case is noisy.
Are there any user-facing changes?
External sorts over wide rows can finish under memory limits where they failed before. A spilling sort can write more, smaller spill batches when rows are wide. No public API or configuration changes.
Known limitations:
main.sort_spill_reservation_bytesbelow 256 KiB (or 0), the 64 KiB floor is larger than what fits. The merge then takes the rest from the pool, as it does onmain.🤖 Generated with Claude Code