adriangb opened a new pull request, #26456: URL: https://github.com/apache/datafusion/pull/26456
## Which issue does this PR close? Closes nothing. Part of #22676. Stacked on #26270. Only the last commit is new. This PR supersedes #26272: it includes the release of the output when the last scan finishes. ## Rationale for this change The body of a materialized CTE ran to its end into a buffer, and each scan then replayed the buffer. A `LIMIT 5` over the CTE computed and kept the whole body. Now the body streams to its scans. Memory holds a few batches ahead of the fastest scan, and the batches that a slower scan has not read yet. `datafusion-cli` (release-nonlto, 12 cores, default settings), max RSS. Before is #26272. The body is `SELECT value AS v, value * 2 AS w FROM generate_series(1, N)`. | Query | Before | After | |---|---|---| | `SELECT * FROM c LIMIT 5`, N = 30M | 498 MB | 38 MB | | `UNION ALL` of two references, N = 30M | 503 MB | 44 MB | | Self hash join `c a JOIN c b ON a.v = b.w`, N = 5M | 305 MB | 306 MB | | `UNION ALL` body of two 15M-row inputs, one scan partition (`target_partitions = 1`) | 267 MB | 37 MB | A scan below a `LIMIT` stops the body. The hash join builds its whole build side before it reads the probe side, so it keeps the same data as before. Timing on a GCP c3-standard-22 VM (10 cores pinned, 5 repetitions, median ms), measured on an earlier revision of this branch with the same algorithm. B is #26272 (full buffer). "Inlined" is `NOT MATERIALIZED` on the merge base. | Case | Partitions | Inlined | B | This PR | |---|---|---|---|---| | `LIMIT 10` over 30M rows | 10 | 2 | 95 | 3 | | `UNION ALL` of two references, 30M rows | 10 | 58 | 131 | 60 | | Three scalar subqueries over one aggregate | 10 | 281 | 172 | 174 | | Self hash join, 5M rows | 10 | 89 | 79 | 71 | | Self hash join of wide rows, 512 MB pool | 10 | 98 | fails | 199 | | 64-row batches, two references | 10 | 1091 | 1588 | 1491 | | Wide strings, 512 MB pool | 10 | 130 | 266 | 99 | With a 512 MB pool and wide strings, B spills 564 MB and this PR spills none. Peak RSS for 30M rows is about 100 MB against about 750 MB for B. TPC-H and TPC-DS (SF1 and SF10, without `MATERIALIZED`) have the same geomean as the merge base, within 1.1%. **Where materialization costs.** A materialized CTE is slower than the inlined CTE in three cases. These are not regressions against B. - One reference reads the whole body before another reference starts, for example the build and probe sides of one hash join. The log keeps the whole body (TPC-DS q39 and q95). - Very small batches at a low partition count. Each batch goes through the log. - References with different selective filters. Inlined, each filter goes into its own scan. Materialized, the body is computed without them (TPC-DS q31). ## What changes are included in this PR? - Each body partition has a fan-out: a log of batches, and a cursor for each scan that reads it. A spawned task pulls the body partition and appends to the log. The log drops the batches that every scan has read. - The task appends while the fastest active scan is less than 4 batches behind, and the log fits under its memory cap. When the fastest scan waits at the end of the log, the task always appends. A gate on the fastest scan, not the slowest, cannot deadlock when two references are the two sides of one hash join. - The logs of one CTE keep at most half of a bounded memory pool. Past that, batches go to spill files. A reader reads a spill file one piece at a time. - Partition `p` of a scan with `n` partitions reads the body partitions `p, p + n, ...` at the same time, and interleaves their batches. The scan claims no ordering. A `UNION ALL` body has more partitions than the scan, and its partitions stream in step. - A body can couple its partitions (for example, outputs that share one input through bounded channels). Because a scan partition polls all of its body partitions, each of them has an active reader, and the body completes. Known limit: such a body can wait forever if one of its partitions has no polled reader, for example a scan partition that is never executed, or a re-execution that computes only some body partitions again. A fan-out whose readers have all finished drops its body stream, also one that no reader started. - The output of a run is released when the last scan partition finishes or is dropped. `RecursiveQueryExec` holds the CTEs that its recursive term reads and does not declare. **Comparison with PostgreSQL and DuckDB.** PostgreSQL fills the CTE tuplestore lazily when a `CteScan` needs a row, and never trims it while the query runs ([nodeCtescan.c](https://github.com/postgres/postgres/blob/af090508018be5ace443ba6e2814500c10842a4d/src/backend/executor/nodeCtescan.c)). This PR is the parallel version of that lazy fill, and it also drops the batches that every scan has read. DuckDB broadcasts the CTE pipeline to its readers ([duckdb/duckdb#23750](https://github.com/duckdb/duckdb/pull/23750)). When it has more than one reader and the buffer passes a high-water mark, it moves the rows to a shared spool that can spill to disk, and it stops throttling the producer ([pipeline_broadcast_exchange.cpp](https://github.com/duckdb/duckdb/blob/1b4c3e54ba3625bac56ad184f0ce642a8d2bfbd5/src/parallel/pipeline_broadcast_exchange.cpp#L1843-L1861)). It releases the buffer when the last reader unregisters ([pipeline_broadcast_exchange.cpp](https://github.com/duckdb/duckdb/blob/ 1b4c3e54ba3625bac56ad184f0ce642a8d2bfbd5/src/parallel/pipeline_broadcast_exchange.cpp#L1708)). This PR uses the same approach. The difference is that the producer here stays gated on the fastest reader, so readers that keep pace never spill. ## Are these changes tested? Yes. Unit tests in `materialized_cte.rs` cover both sides of a hash join, scans with fewer partitions than the body, spill and the release of spill files, readers in step that keep a few batches, a `LIMIT` that stops the body, dropped scans, body errors, a panic in `execute` and a panic of the body stream, re-execution, holds, and the release of the run. A test body with two outputs that share one input through channels of capacity 1 completes with a 1-partition scan, also with a small memory pool. A scan with fewer partitions than the body does not spill and keeps a few batches, also through SQL with a `UNION ALL` body. A scan partition that is dropped before its first poll does not stop a coupled body. Two more tests cover a late scan that reads a body partition, and the drop of a finished reader whose slot a new reader reuses. Other tests cover the retirement of each batch, the drop of a hold, and a batch larger than 2 MiB that is spilled in pieces. `cte_materialized.slt` cover s the SQL cases. ## Are there any user-facing changes? New `EXPLAIN ANALYZE` metrics for `MaterializedCteExec`: `retired_batches`, `peak_buffered_bytes` and `producer_wait_time`. The spill metrics existed before and are now per body partition. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
