adriangb opened a new pull request, #25580: URL: https://github.com/apache/datafusion/pull/25580
## 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 https://github.com/apache/datafusion/issues/17737#issuecomment-5767781260: - 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. - A recursive CTE is still planned as `RecursiveQuery`. - **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 one 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 file. ### 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. ### Out of scope, possible follow-ups - Protobuf serialization. These nodes have no proto encoding yet, so serializing a plan that contains them fails in the default extension codec. - Statistics, ordering and partitioning on the scans, and releasing the buffer when the last scan finishes rather than when the plan is dropped. - Deciding automatically when to materialize. - A streaming fan-out for consumers that are known to be independent. - Pushing filters into the CTE body. - `sqlparser` parses `[NOT] MATERIALIZED` only in `PostgreSqlDialect`, so the tests set that dialect. If support in the generic dialect is wanted, that is a small `sqlparser-rs` issue we can open. ## 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; - column aliases; - a spill under `memory_limit`: 200k rows, about 167k spilled, and both scans still return 200k rows. 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? -- 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]
