Skip to content

Commit cdc2fc0

Browse files
adriangbclaude
andauthored
feat: support WITH ... AS [NOT] MATERIALIZED CTEs (#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 #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]>
1 parent b645de0 commit cdc2fc0

13 files changed

Lines changed: 2738 additions & 26 deletions

File tree

‎datafusion/core/src/datasource/view_test.rs‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,4 +456,75 @@ mod tests {
456456

457457
Ok(())
458458
}
459+
460+
/// A [`ViewTable`] that is not inlined: each read plans the saved query in
461+
/// `scan`.
462+
///
463+
/// [`ViewTable`]: crate::datasource::ViewTable
464+
#[derive(Debug)]
465+
struct NotInlinedView(crate::datasource::ViewTable);
466+
467+
#[async_trait::async_trait]
468+
impl crate::datasource::TableProvider for NotInlinedView {
469+
fn schema(&self) -> arrow::datatypes::SchemaRef {
470+
self.0.schema()
471+
}
472+
473+
fn table_type(&self) -> crate::datasource::TableType {
474+
self.0.table_type()
475+
}
476+
477+
async fn scan(
478+
&self,
479+
state: &dyn crate::catalog::Session,
480+
projection: Option<&[usize]>,
481+
filters: &[datafusion_expr::Expr],
482+
limit: Option<usize>,
483+
) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
484+
self.0.scan(state, projection, filters, limit).await
485+
}
486+
}
487+
488+
#[tokio::test]
489+
async fn query_not_inlined_view_with_materialized_cte_twice() -> Result<()> {
490+
// sqlparser only parses `[NOT] MATERIALIZED` for the PostgreSQL dialect.
491+
let session_ctx = SessionContext::new_with_config(
492+
SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL"),
493+
);
494+
session_ctx
495+
.sql("CREATE TABLE t AS VALUES (1), (2)")
496+
.await?
497+
.collect()
498+
.await?;
499+
let plan = session_ctx
500+
.sql(
501+
"WITH r AS MATERIALIZED (SELECT column1 AS a FROM t) \
502+
SELECT a FROM r",
503+
)
504+
.await?
505+
.into_optimized_plan()?;
506+
let view = crate::datasource::ViewTable::new(plan, None);
507+
session_ctx.register_table("w", std::sync::Arc::new(NotInlinedView(view)))?;
508+
509+
// Each read of `w` plans the saved query, which has the same CTE id
510+
// each time.
511+
let results = session_ctx
512+
.sql("SELECT x.a AS x, y.a AS y FROM w x, w y ORDER BY x, y")
513+
.await?
514+
.collect()
515+
.await?;
516+
517+
insta::assert_snapshot!(batches_to_string(&results),@r"
518+
+---+---+
519+
| x | y |
520+
+---+---+
521+
| 1 | 1 |
522+
| 1 | 2 |
523+
| 2 | 1 |
524+
| 2 | 2 |
525+
+---+---+
526+
");
527+
528+
Ok(())
529+
}
459530
}

‎datafusion/core/src/physical_planner.rs‎

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ use crate::physical_plan::joins::{
4545
PartitionMode, SortMergeJoinExec,
4646
};
4747
use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
48+
use crate::physical_plan::materialized_cte::{
49+
MaterializedCteBuffer, MaterializedCteExec, MaterializedCteScanExec,
50+
};
4851
use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr};
4952
use crate::physical_plan::repartition::RepartitionExec;
5053
use crate::physical_plan::sorts::sort::SortExec;
@@ -230,14 +233,160 @@ impl DefaultPhysicalPlanner {
230233
{
231234
return Ok(plan);
232235
}
236+
let renumbered = renumber_materialized_ctes(logical_plan)?;
237+
let logical_plan = renumbered.as_ref().unwrap_or(logical_plan);
233238
let plan = self
234239
.create_initial_plan(logical_plan, session_state)
235240
.await?;
241+
let plan = bind_materialized_cte_scans(plan)?;
236242

237243
self.optimize_physical_plan(plan, session_state, |_, _| {})
238244
}
239245
}
240246

247+
/// Give each occurrence of a [`MaterializedCte`] in `plan` a new id, and
248+
/// update the scans in its continuation, including the scans in subqueries.
249+
///
250+
/// The id is copied when a logical plan is cloned, so one CTE can occur more
251+
/// than once, for example when a view that declares it is referenced two
252+
/// times, or when a DataFrame is joined to itself. Each occurrence is a
253+
/// separate CTE, as in PostgreSQL.
254+
///
255+
/// The ids must also be new for each planning run, not only unique within
256+
/// `plan`. A table provider can plan a saved logical plan in its `scan`, as a
257+
/// [`ViewTable`] that is not inlined does. Each such scan is a separate
258+
/// planning run of the same saved plan, and its CTEs end up in the same
259+
/// physical plan as the CTEs of the outer run.
260+
///
261+
/// Returns `None` when `plan` has no materialized CTE.
262+
///
263+
/// [`MaterializedCte`]: datafusion_expr::MaterializedCte
264+
/// [`ViewTable`]: crate::datasource::ViewTable
265+
fn renumber_materialized_ctes(plan: &LogicalPlan) -> Result<Option<LogicalPlan>> {
266+
fn as_cte(plan: &LogicalPlan) -> Option<&datafusion_expr::MaterializedCte> {
267+
match plan {
268+
LogicalPlan::Extension(Extension { node }) => node.as_any().downcast_ref(),
269+
_ => None,
270+
}
271+
}
272+
273+
let mut has_cte = false;
274+
plan.apply_with_subqueries(|node| {
275+
if as_cte(node).is_some() {
276+
has_cte = true;
277+
return Ok(TreeNodeRecursion::Stop);
278+
}
279+
Ok(TreeNodeRecursion::Continue)
280+
})?;
281+
if !has_cte {
282+
return Ok(None);
283+
}
284+
285+
plan.clone()
286+
.transform_down_with_subqueries(|node| {
287+
let Some(cte) = as_cte(&node) else {
288+
return Ok(Transformed::no(node));
289+
};
290+
let old_id = cte.id;
291+
let new_id = datafusion_expr::MaterializedCteId::next();
292+
let continuation = cte
293+
.continuation
294+
.clone()
295+
.transform_down_with_subqueries(|node| {
296+
// A nested occurrence of the same CTE owns the scans below
297+
// it, and is renumbered on its own.
298+
if as_cte(&node).is_some_and(|nested| nested.id == old_id) {
299+
return Ok(Transformed::new(
300+
node,
301+
false,
302+
TreeNodeRecursion::Jump,
303+
));
304+
}
305+
let LogicalPlan::Extension(Extension { node: scan }) = &node else {
306+
return Ok(Transformed::no(node));
307+
};
308+
match scan
309+
.as_any()
310+
.downcast_ref::<datafusion_expr::MaterializedCteScan>()
311+
{
312+
Some(scan) if scan.id == old_id => {
313+
let scan = datafusion_expr::MaterializedCteScan {
314+
id: new_id,
315+
..scan.clone()
316+
};
317+
Ok(Transformed::yes(LogicalPlan::Extension(Extension {
318+
node: Arc::new(scan),
319+
})))
320+
}
321+
_ => Ok(Transformed::no(node)),
322+
}
323+
})?
324+
.data;
325+
let cte = datafusion_expr::MaterializedCte {
326+
id: new_id,
327+
name: cte.name.clone(),
328+
cte: cte.cte.clone(),
329+
continuation,
330+
};
331+
Ok(Transformed::yes(LogicalPlan::Extension(Extension {
332+
node: Arc::new(cte),
333+
})))
334+
})
335+
.map(|t| Some(t.data))
336+
}
337+
338+
/// Point every [`MaterializedCteScanExec`] at the buffer of the
339+
/// [`MaterializedCteExec`] with the same id.
340+
///
341+
/// [`renumber_materialized_ctes`] makes the ids unique before the physical
342+
/// plan is created.
343+
///
344+
/// A scan that is already bound was planned and bound by a separate planning
345+
/// run, such as the `scan` of a [`ViewTable`], and is left as is. Binding it
346+
/// again would count its partitions twice in the buffer.
347+
///
348+
/// Scans are planned before the node that owns their CTE, and a scan inside a
349+
/// scalar subquery is planned in a separate subtree, so the binding is done
350+
/// once on the whole initial plan.
351+
///
352+
/// [`ViewTable`]: crate::datasource::ViewTable
353+
fn bind_materialized_cte_scans(
354+
plan: Arc<dyn ExecutionPlan>,
355+
) -> Result<Arc<dyn ExecutionPlan>> {
356+
let mut buffers = HashMap::new();
357+
plan.apply(|node| {
358+
if let Some(cte) = node.downcast_ref::<MaterializedCteExec>()
359+
&& buffers
360+
.insert(cte.buffer().id(), Arc::clone(cte.buffer()))
361+
.is_some()
362+
{
363+
return internal_err!(
364+
"Two MaterializedCteExec nodes have the id {}",
365+
cte.buffer().id()
366+
);
367+
}
368+
Ok(TreeNodeRecursion::Continue)
369+
})?;
370+
if buffers.is_empty() {
371+
return Ok(plan);
372+
}
373+
plan.transform_up(|node| {
374+
let Some(scan) = node.downcast_ref::<MaterializedCteScanExec>() else {
375+
return Ok(Transformed::no(node));
376+
};
377+
if scan.is_bound() {
378+
return Ok(Transformed::no(node));
379+
}
380+
let Some(buffer) = buffers.get(&scan.id()) else {
381+
return internal_err!("MaterializedCteScanExec {} has no CTE", scan.id());
382+
};
383+
Ok(Transformed::yes(
384+
Arc::new(scan.bind(Arc::clone(buffer))) as Arc<dyn ExecutionPlan>
385+
))
386+
})
387+
.map(|t| t.data)
388+
}
389+
241390
#[derive(Debug)]
242391
struct ExecutionPlanChild {
243392
/// Index needed to order children of parent to ensure consistency with original
@@ -1904,6 +2053,58 @@ impl DefaultPhysicalPlanner {
19042053

19052054
// N Children
19062055
LogicalPlan::Union(_) => UnionExec::try_new(children.vec())?,
2056+
LogicalPlan::Extension(Extension { node })
2057+
if node.as_any().is::<datafusion_expr::MaterializedCte>() =>
2058+
{
2059+
let cte = node
2060+
.as_any()
2061+
.downcast_ref::<datafusion_expr::MaterializedCte>()
2062+
.unwrap();
2063+
let [body, continuation] = children.two()?;
2064+
// A CTE that nothing reads is never evaluated.
2065+
let mut is_referenced = false;
2066+
cte.continuation.apply_with_subqueries(|node| {
2067+
if let LogicalPlan::Extension(Extension { node }) = node
2068+
&& let Some(scan) =
2069+
node.as_any()
2070+
.downcast_ref::<datafusion_expr::MaterializedCteScan>()
2071+
&& scan.id == cte.id
2072+
{
2073+
is_referenced = true;
2074+
return Ok(TreeNodeRecursion::Stop);
2075+
}
2076+
Ok(TreeNodeRecursion::Continue)
2077+
})?;
2078+
if !is_referenced {
2079+
return Ok(continuation);
2080+
}
2081+
// The body is fully buffered before a scan yields a row, so an
2082+
// unbounded body would never produce output.
2083+
if body.boundedness().is_unbounded() {
2084+
return plan_err!(
2085+
"MATERIALIZED CTE {} reads an unbounded source, which cannot be materialized",
2086+
cte.name
2087+
);
2088+
}
2089+
let buffer =
2090+
Arc::new(MaterializedCteBuffer::new(cte.id.as_u64(), &cte.name));
2091+
Arc::new(MaterializedCteExec::new(body, continuation, buffer))
2092+
}
2093+
LogicalPlan::Extension(Extension { node })
2094+
if node.as_any().is::<datafusion_expr::MaterializedCteScan>() =>
2095+
{
2096+
let scan = node
2097+
.as_any()
2098+
.downcast_ref::<datafusion_expr::MaterializedCteScan>()
2099+
.unwrap();
2100+
// Unbound until `bind_materialized_cte_scans` runs on the whole plan.
2101+
Arc::new(MaterializedCteScanExec::new(
2102+
scan.id.as_u64(),
2103+
&scan.name,
2104+
Arc::clone(scan.schema.inner()),
2105+
session_state.config().target_partitions(),
2106+
))
2107+
}
19072108
LogicalPlan::Extension(Extension { node }) => {
19082109
let mut maybe_plan = None;
19092110
let children = children.vec();

0 commit comments

Comments
 (0)