Repository navigation
Commit cdc2fc0
feat: support
## Which issue does this PR close?
- Closes #17737.
- Related: #8777, #22676, #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 (#22551, #22675) bundled automatic materialization,
heuristics and filter pushdown. This PR is limited to the explicit
directive, as proposed in
#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 #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.
#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 (#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 #22676:
- Protobuf serialization. These nodes have no proto encoding yet, so
serializing a plan that contains them fails in the default extension
codec (#26231).
- Statistics, ordering and partitioning on the scans (#26230).
- Release the buffer when the last scan finishes rather than when the
plan is dropped (#26233).
- Byte-bounded spill files (#26232).
- Keep the output per run, so that `reset_state` follows its contract
(#26229).
- Deciding automatically when to materialize (#26237).
- First-class `LogicalPlan` variants (#26238).
- Pushing filters that all references share into the CTE body (#26239).
- Pruning unused columns from the CTE body (#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 <[email protected]>WITH ... AS [NOT] MATERIALIZED CTEs (#25580)1 parent b645de0 commit cdc2fc0
13 files changed
Lines changed: 2738 additions & 26 deletions
File tree
- datafusion
- core/src
- datasource
- expr/src/logical_plan
- optimizer/src
- analyzer
- physical-plan/src
- sqllogictest/test_files
- sql/src
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
456 | 456 | | |
457 | 457 | | |
458 | 458 | | |
| 459 | + | |
| 460 | + | |
| 461 | + | |
| 462 | + | |
| 463 | + | |
| 464 | + | |
| 465 | + | |
| 466 | + | |
| 467 | + | |
| 468 | + | |
| 469 | + | |
| 470 | + | |
| 471 | + | |
| 472 | + | |
| 473 | + | |
| 474 | + | |
| 475 | + | |
| 476 | + | |
| 477 | + | |
| 478 | + | |
| 479 | + | |
| 480 | + | |
| 481 | + | |
| 482 | + | |
| 483 | + | |
| 484 | + | |
| 485 | + | |
| 486 | + | |
| 487 | + | |
| 488 | + | |
| 489 | + | |
| 490 | + | |
| 491 | + | |
| 492 | + | |
| 493 | + | |
| 494 | + | |
| 495 | + | |
| 496 | + | |
| 497 | + | |
| 498 | + | |
| 499 | + | |
| 500 | + | |
| 501 | + | |
| 502 | + | |
| 503 | + | |
| 504 | + | |
| 505 | + | |
| 506 | + | |
| 507 | + | |
| 508 | + | |
| 509 | + | |
| 510 | + | |
| 511 | + | |
| 512 | + | |
| 513 | + | |
| 514 | + | |
| 515 | + | |
| 516 | + | |
| 517 | + | |
| 518 | + | |
| 519 | + | |
| 520 | + | |
| 521 | + | |
| 522 | + | |
| 523 | + | |
| 524 | + | |
| 525 | + | |
| 526 | + | |
| 527 | + | |
| 528 | + | |
| 529 | + | |
459 | 530 | | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
45 | 45 | | |
46 | 46 | | |
47 | 47 | | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
48 | 51 | | |
49 | 52 | | |
50 | 53 | | |
| |||
230 | 233 | | |
231 | 234 | | |
232 | 235 | | |
| 236 | + | |
| 237 | + | |
233 | 238 | | |
234 | 239 | | |
235 | 240 | | |
| 241 | + | |
236 | 242 | | |
237 | 243 | | |
238 | 244 | | |
239 | 245 | | |
240 | 246 | | |
| 247 | + | |
| 248 | + | |
| 249 | + | |
| 250 | + | |
| 251 | + | |
| 252 | + | |
| 253 | + | |
| 254 | + | |
| 255 | + | |
| 256 | + | |
| 257 | + | |
| 258 | + | |
| 259 | + | |
| 260 | + | |
| 261 | + | |
| 262 | + | |
| 263 | + | |
| 264 | + | |
| 265 | + | |
| 266 | + | |
| 267 | + | |
| 268 | + | |
| 269 | + | |
| 270 | + | |
| 271 | + | |
| 272 | + | |
| 273 | + | |
| 274 | + | |
| 275 | + | |
| 276 | + | |
| 277 | + | |
| 278 | + | |
| 279 | + | |
| 280 | + | |
| 281 | + | |
| 282 | + | |
| 283 | + | |
| 284 | + | |
| 285 | + | |
| 286 | + | |
| 287 | + | |
| 288 | + | |
| 289 | + | |
| 290 | + | |
| 291 | + | |
| 292 | + | |
| 293 | + | |
| 294 | + | |
| 295 | + | |
| 296 | + | |
| 297 | + | |
| 298 | + | |
| 299 | + | |
| 300 | + | |
| 301 | + | |
| 302 | + | |
| 303 | + | |
| 304 | + | |
| 305 | + | |
| 306 | + | |
| 307 | + | |
| 308 | + | |
| 309 | + | |
| 310 | + | |
| 311 | + | |
| 312 | + | |
| 313 | + | |
| 314 | + | |
| 315 | + | |
| 316 | + | |
| 317 | + | |
| 318 | + | |
| 319 | + | |
| 320 | + | |
| 321 | + | |
| 322 | + | |
| 323 | + | |
| 324 | + | |
| 325 | + | |
| 326 | + | |
| 327 | + | |
| 328 | + | |
| 329 | + | |
| 330 | + | |
| 331 | + | |
| 332 | + | |
| 333 | + | |
| 334 | + | |
| 335 | + | |
| 336 | + | |
| 337 | + | |
| 338 | + | |
| 339 | + | |
| 340 | + | |
| 341 | + | |
| 342 | + | |
| 343 | + | |
| 344 | + | |
| 345 | + | |
| 346 | + | |
| 347 | + | |
| 348 | + | |
| 349 | + | |
| 350 | + | |
| 351 | + | |
| 352 | + | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
| 358 | + | |
| 359 | + | |
| 360 | + | |
| 361 | + | |
| 362 | + | |
| 363 | + | |
| 364 | + | |
| 365 | + | |
| 366 | + | |
| 367 | + | |
| 368 | + | |
| 369 | + | |
| 370 | + | |
| 371 | + | |
| 372 | + | |
| 373 | + | |
| 374 | + | |
| 375 | + | |
| 376 | + | |
| 377 | + | |
| 378 | + | |
| 379 | + | |
| 380 | + | |
| 381 | + | |
| 382 | + | |
| 383 | + | |
| 384 | + | |
| 385 | + | |
| 386 | + | |
| 387 | + | |
| 388 | + | |
| 389 | + | |
241 | 390 | | |
242 | 391 | | |
243 | 392 | | |
| |||
1904 | 2053 | | |
1905 | 2054 | | |
1906 | 2055 | | |
| 2056 | + | |
| 2057 | + | |
| 2058 | + | |
| 2059 | + | |
| 2060 | + | |
| 2061 | + | |
| 2062 | + | |
| 2063 | + | |
| 2064 | + | |
| 2065 | + | |
| 2066 | + | |
| 2067 | + | |
| 2068 | + | |
| 2069 | + | |
| 2070 | + | |
| 2071 | + | |
| 2072 | + | |
| 2073 | + | |
| 2074 | + | |
| 2075 | + | |
| 2076 | + | |
| 2077 | + | |
| 2078 | + | |
| 2079 | + | |
| 2080 | + | |
| 2081 | + | |
| 2082 | + | |
| 2083 | + | |
| 2084 | + | |
| 2085 | + | |
| 2086 | + | |
| 2087 | + | |
| 2088 | + | |
| 2089 | + | |
| 2090 | + | |
| 2091 | + | |
| 2092 | + | |
| 2093 | + | |
| 2094 | + | |
| 2095 | + | |
| 2096 | + | |
| 2097 | + | |
| 2098 | + | |
| 2099 | + | |
| 2100 | + | |
| 2101 | + | |
| 2102 | + | |
| 2103 | + | |
| 2104 | + | |
| 2105 | + | |
| 2106 | + | |
| 2107 | + | |
1907 | 2108 | | |
1908 | 2109 | | |
1909 | 2110 | | |
| |||
0 commit comments