You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/15 22:00:29 UTC
(datafusion) branch main updated: feat: support `WITH ... AS [NOT] MATERIALIZED` CTEs (#25580)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new cdc2fc041b feat: support `WITH ... AS [NOT] MATERIALIZED` CTEs (#25580)
cdc2fc041b is described below
commit cdc2fc041b6a812a24933abc40da6475dbbf5729
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Thu Oct 15 21:30:04 2026 +0000
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
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.
- `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.
https://github.com/apache/datafusion/pull/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 <no...@anthropic.com>
---
datafusion/core/src/datasource/view_test.rs | 71 ++
datafusion/core/src/physical_planner.rs | 201 +++++
.../expr/src/logical_plan/materialized_cte.rs | 236 +++++
datafusion/expr/src/logical_plan/mod.rs | 2 +
datafusion/expr/src/logical_plan/plan.rs | 9 +-
datafusion/optimizer/src/analyzer/type_coercion.rs | 121 ++-
datafusion/optimizer/src/decorrelate.rs | 37 +-
datafusion/optimizer/src/push_down_filter.rs | 22 +-
datafusion/physical-plan/src/lib.rs | 1 +
datafusion/physical-plan/src/materialized_cte.rs | 989 ++++++++++++++++++++
datafusion/sql/src/cte.rs | 60 +-
datafusion/sql/src/query.rs | 17 +-
.../sqllogictest/test_files/cte_materialized.slt | 998 +++++++++++++++++++++
13 files changed, 2738 insertions(+), 26 deletions(-)
diff --git a/datafusion/core/src/datasource/view_test.rs b/datafusion/core/src/datasource/view_test.rs
index 35418d6dea..da4745d6d3 100644
--- a/datafusion/core/src/datasource/view_test.rs
+++ b/datafusion/core/src/datasource/view_test.rs
@@ -456,4 +456,75 @@ mod tests {
Ok(())
}
+
+ /// A [`ViewTable`] that is not inlined: each read plans the saved query in
+ /// `scan`.
+ ///
+ /// [`ViewTable`]: crate::datasource::ViewTable
+ #[derive(Debug)]
+ struct NotInlinedView(crate::datasource::ViewTable);
+
+ #[async_trait::async_trait]
+ impl crate::datasource::TableProvider for NotInlinedView {
+ fn schema(&self) -> arrow::datatypes::SchemaRef {
+ self.0.schema()
+ }
+
+ fn table_type(&self) -> crate::datasource::TableType {
+ self.0.table_type()
+ }
+
+ async fn scan(
+ &self,
+ state: &dyn crate::catalog::Session,
+ projection: Option<&[usize]>,
+ filters: &[datafusion_expr::Expr],
+ limit: Option<usize>,
+ ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
+ self.0.scan(state, projection, filters, limit).await
+ }
+ }
+
+ #[tokio::test]
+ async fn query_not_inlined_view_with_materialized_cte_twice() -> Result<()> {
+ // sqlparser only parses `[NOT] MATERIALIZED` for the PostgreSQL dialect.
+ let session_ctx = SessionContext::new_with_config(
+ SessionConfig::new().set_str("datafusion.sql_parser.dialect", "PostgreSQL"),
+ );
+ session_ctx
+ .sql("CREATE TABLE t AS VALUES (1), (2)")
+ .await?
+ .collect()
+ .await?;
+ let plan = session_ctx
+ .sql(
+ "WITH r AS MATERIALIZED (SELECT column1 AS a FROM t) \
+ SELECT a FROM r",
+ )
+ .await?
+ .into_optimized_plan()?;
+ let view = crate::datasource::ViewTable::new(plan, None);
+ session_ctx.register_table("w", std::sync::Arc::new(NotInlinedView(view)))?;
+
+ // Each read of `w` plans the saved query, which has the same CTE id
+ // each time.
+ let results = session_ctx
+ .sql("SELECT x.a AS x, y.a AS y FROM w x, w y ORDER BY x, y")
+ .await?
+ .collect()
+ .await?;
+
+ insta::assert_snapshot!(batches_to_string(&results),@r"
+ +---+---+
+ | x | y |
+ +---+---+
+ | 1 | 1 |
+ | 1 | 2 |
+ | 2 | 1 |
+ | 2 | 2 |
+ +---+---+
+ ");
+
+ Ok(())
+ }
}
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index f1bc734d72..96c8235d4a 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -45,6 +45,9 @@ use crate::physical_plan::joins::{
PartitionMode, SortMergeJoinExec,
};
use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
+use crate::physical_plan::materialized_cte::{
+ MaterializedCteBuffer, MaterializedCteExec, MaterializedCteScanExec,
+};
use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr};
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort::SortExec;
@@ -230,14 +233,160 @@ impl DefaultPhysicalPlanner {
{
return Ok(plan);
}
+ let renumbered = renumber_materialized_ctes(logical_plan)?;
+ let logical_plan = renumbered.as_ref().unwrap_or(logical_plan);
let plan = self
.create_initial_plan(logical_plan, session_state)
.await?;
+ let plan = bind_materialized_cte_scans(plan)?;
self.optimize_physical_plan(plan, session_state, |_, _| {})
}
}
+/// Give each occurrence of a [`MaterializedCte`] in `plan` a new id, and
+/// update the scans in its continuation, including the scans in subqueries.
+///
+/// The id is copied when a logical plan is cloned, so one CTE can occur more
+/// than once, for example when a view that declares it is referenced two
+/// times, or when a DataFrame is joined to itself. Each occurrence is a
+/// separate CTE, as in PostgreSQL.
+///
+/// The ids must also be new for each planning run, not only unique within
+/// `plan`. A table provider can plan a saved logical plan in its `scan`, as a
+/// [`ViewTable`] that is not inlined does. Each such scan is a separate
+/// planning run of the same saved plan, and its CTEs end up in the same
+/// physical plan as the CTEs of the outer run.
+///
+/// Returns `None` when `plan` has no materialized CTE.
+///
+/// [`MaterializedCte`]: datafusion_expr::MaterializedCte
+/// [`ViewTable`]: crate::datasource::ViewTable
+fn renumber_materialized_ctes(plan: &LogicalPlan) -> Result<Option<LogicalPlan>> {
+ fn as_cte(plan: &LogicalPlan) -> Option<&datafusion_expr::MaterializedCte> {
+ match plan {
+ LogicalPlan::Extension(Extension { node }) => node.as_any().downcast_ref(),
+ _ => None,
+ }
+ }
+
+ let mut has_cte = false;
+ plan.apply_with_subqueries(|node| {
+ if as_cte(node).is_some() {
+ has_cte = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ if !has_cte {
+ return Ok(None);
+ }
+
+ plan.clone()
+ .transform_down_with_subqueries(|node| {
+ let Some(cte) = as_cte(&node) else {
+ return Ok(Transformed::no(node));
+ };
+ let old_id = cte.id;
+ let new_id = datafusion_expr::MaterializedCteId::next();
+ let continuation = cte
+ .continuation
+ .clone()
+ .transform_down_with_subqueries(|node| {
+ // A nested occurrence of the same CTE owns the scans below
+ // it, and is renumbered on its own.
+ if as_cte(&node).is_some_and(|nested| nested.id == old_id) {
+ return Ok(Transformed::new(
+ node,
+ false,
+ TreeNodeRecursion::Jump,
+ ));
+ }
+ let LogicalPlan::Extension(Extension { node: scan }) = &node else {
+ return Ok(Transformed::no(node));
+ };
+ match scan
+ .as_any()
+ .downcast_ref::<datafusion_expr::MaterializedCteScan>()
+ {
+ Some(scan) if scan.id == old_id => {
+ let scan = datafusion_expr::MaterializedCteScan {
+ id: new_id,
+ ..scan.clone()
+ };
+ Ok(Transformed::yes(LogicalPlan::Extension(Extension {
+ node: Arc::new(scan),
+ })))
+ }
+ _ => Ok(Transformed::no(node)),
+ }
+ })?
+ .data;
+ let cte = datafusion_expr::MaterializedCte {
+ id: new_id,
+ name: cte.name.clone(),
+ cte: cte.cte.clone(),
+ continuation,
+ };
+ Ok(Transformed::yes(LogicalPlan::Extension(Extension {
+ node: Arc::new(cte),
+ })))
+ })
+ .map(|t| Some(t.data))
+}
+
+/// Point every [`MaterializedCteScanExec`] at the buffer of the
+/// [`MaterializedCteExec`] with the same id.
+///
+/// [`renumber_materialized_ctes`] makes the ids unique before the physical
+/// plan is created.
+///
+/// A scan that is already bound was planned and bound by a separate planning
+/// run, such as the `scan` of a [`ViewTable`], and is left as is. Binding it
+/// again would count its partitions twice in the buffer.
+///
+/// Scans are planned before the node that owns their CTE, and a scan inside a
+/// scalar subquery is planned in a separate subtree, so the binding is done
+/// once on the whole initial plan.
+///
+/// [`ViewTable`]: crate::datasource::ViewTable
+fn bind_materialized_cte_scans(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let mut buffers = HashMap::new();
+ plan.apply(|node| {
+ if let Some(cte) = node.downcast_ref::<MaterializedCteExec>()
+ && buffers
+ .insert(cte.buffer().id(), Arc::clone(cte.buffer()))
+ .is_some()
+ {
+ return internal_err!(
+ "Two MaterializedCteExec nodes have the id {}",
+ cte.buffer().id()
+ );
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ if buffers.is_empty() {
+ return Ok(plan);
+ }
+ plan.transform_up(|node| {
+ let Some(scan) = node.downcast_ref::<MaterializedCteScanExec>() else {
+ return Ok(Transformed::no(node));
+ };
+ if scan.is_bound() {
+ return Ok(Transformed::no(node));
+ }
+ let Some(buffer) = buffers.get(&scan.id()) else {
+ return internal_err!("MaterializedCteScanExec {} has no CTE", scan.id());
+ };
+ Ok(Transformed::yes(
+ Arc::new(scan.bind(Arc::clone(buffer))) as Arc<dyn ExecutionPlan>
+ ))
+ })
+ .map(|t| t.data)
+}
+
#[derive(Debug)]
struct ExecutionPlanChild {
/// Index needed to order children of parent to ensure consistency with original
@@ -1904,6 +2053,58 @@ impl DefaultPhysicalPlanner {
// N Children
LogicalPlan::Union(_) => UnionExec::try_new(children.vec())?,
+ LogicalPlan::Extension(Extension { node })
+ if node.as_any().is::<datafusion_expr::MaterializedCte>() =>
+ {
+ let cte = node
+ .as_any()
+ .downcast_ref::<datafusion_expr::MaterializedCte>()
+ .unwrap();
+ let [body, continuation] = children.two()?;
+ // A CTE that nothing reads is never evaluated.
+ let mut is_referenced = false;
+ cte.continuation.apply_with_subqueries(|node| {
+ if let LogicalPlan::Extension(Extension { node }) = node
+ && let Some(scan) =
+ node.as_any()
+ .downcast_ref::<datafusion_expr::MaterializedCteScan>()
+ && scan.id == cte.id
+ {
+ is_referenced = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ if !is_referenced {
+ return Ok(continuation);
+ }
+ // The body is fully buffered before a scan yields a row, so an
+ // unbounded body would never produce output.
+ if body.boundedness().is_unbounded() {
+ return plan_err!(
+ "MATERIALIZED CTE {} reads an unbounded source, which cannot be materialized",
+ cte.name
+ );
+ }
+ let buffer =
+ Arc::new(MaterializedCteBuffer::new(cte.id.as_u64(), &cte.name));
+ Arc::new(MaterializedCteExec::new(body, continuation, buffer))
+ }
+ LogicalPlan::Extension(Extension { node })
+ if node.as_any().is::<datafusion_expr::MaterializedCteScan>() =>
+ {
+ let scan = node
+ .as_any()
+ .downcast_ref::<datafusion_expr::MaterializedCteScan>()
+ .unwrap();
+ // Unbound until `bind_materialized_cte_scans` runs on the whole plan.
+ Arc::new(MaterializedCteScanExec::new(
+ scan.id.as_u64(),
+ &scan.name,
+ Arc::clone(scan.schema.inner()),
+ session_state.config().target_partitions(),
+ ))
+ }
LogicalPlan::Extension(Extension { node }) => {
let mut maybe_plan = None;
let children = children.vec();
diff --git a/datafusion/expr/src/logical_plan/materialized_cte.rs b/datafusion/expr/src/logical_plan/materialized_cte.rs
new file mode 100644
index 0000000000..aaf8f46dc6
--- /dev/null
+++ b/datafusion/expr/src/logical_plan/materialized_cte.rs
@@ -0,0 +1,236 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Logical nodes for `WITH x AS MATERIALIZED (...)`.
+//!
+//! A materialized CTE is planned as one [`MaterializedCte`] node, which owns the
+//! CTE body and the rest of the query (the "continuation"), and one
+//! [`MaterializedCteScan`] leaf for each reference to the CTE inside the
+//! continuation. The body is therefore optimized and executed once, and the
+//! scans read its buffered output.
+//!
+//! The scans are leaves, so filters and projections of one reference are not
+//! pushed into the shared body. This is the same optimization fence that
+//! PostgreSQL applies to `MATERIALIZED` CTEs.
+
+use std::cmp::Ordering;
+use std::fmt;
+use std::hash::Hash;
+use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
+
+use datafusion_common::{DFSchemaRef, Result, assert_eq_or_internal_err};
+
+use crate::{Expr, LogicalPlan, UserDefinedLogicalNodeCore};
+
+/// Identifies one materialized CTE and binds its scans to it.
+///
+/// CTE names are not unique (two sibling subqueries can each declare `WITH t`),
+/// so the planner allocates a process-unique id.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct MaterializedCteId(u64);
+
+impl MaterializedCteId {
+ /// Allocate a new unique id.
+ pub fn next() -> Self {
+ static NEXT: AtomicU64 = AtomicU64::new(0);
+ Self(NEXT.fetch_add(1, AtomicOrdering::Relaxed))
+ }
+
+ /// The numeric value of this id.
+ pub fn as_u64(self) -> u64 {
+ self.0
+ }
+}
+
+impl fmt::Display for MaterializedCteId {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ write!(f, "{}", self.0)
+ }
+}
+
+/// Computes `cte` once and evaluates `continuation`, in which every
+/// [`MaterializedCteScan`] with the same `id` reads the buffered output of
+/// `cte`. The output of this node is the output of `continuation`.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct MaterializedCte {
+ pub id: MaterializedCteId,
+ pub name: String,
+ pub cte: LogicalPlan,
+ pub continuation: LogicalPlan,
+}
+
+/// Orders by `id` and `name`. Nodes with the same `id` and `name` but a
+/// different `cte` or `continuation` are not comparable, so `Some(Equal)` is
+/// returned only for equal nodes.
+impl PartialOrd for MaterializedCte {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ match self
+ .id
+ .cmp(&other.id)
+ .then_with(|| self.name.cmp(&other.name))
+ {
+ Ordering::Equal if self != other => None,
+ ord => Some(ord),
+ }
+ }
+}
+
+impl UserDefinedLogicalNodeCore for MaterializedCte {
+ fn name(&self) -> &str {
+ "MaterializedCte"
+ }
+
+ fn inputs(&self) -> Vec<&LogicalPlan> {
+ vec![&self.cte, &self.continuation]
+ }
+
+ fn schema(&self) -> &DFSchemaRef {
+ self.continuation.schema()
+ }
+
+ fn expressions(&self) -> Vec<Expr> {
+ vec![]
+ }
+
+ fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "MaterializedCte: name={}", self.name)
+ }
+
+ fn with_exprs_and_inputs(
+ &self,
+ exprs: Vec<Expr>,
+ mut inputs: Vec<LogicalPlan>,
+ ) -> Result<Self> {
+ assert_eq_or_internal_err!(exprs.len(), 0, "MaterializedCte takes no exprs");
+ assert_eq_or_internal_err!(inputs.len(), 2, "MaterializedCte takes 2 inputs");
+ let continuation = inputs.pop().unwrap();
+ let cte = inputs.pop().unwrap();
+ Ok(Self {
+ id: self.id,
+ name: self.name.clone(),
+ cte,
+ continuation,
+ })
+ }
+
+ fn necessary_children_exprs(
+ &self,
+ output_columns: &[usize],
+ ) -> Option<Vec<Vec<usize>>> {
+ // The scans need every column of the body, since they are leaves whose
+ // schema is fixed at planning time. The continuation only needs the
+ // columns the parent asks for.
+ let all_cte_columns = (0..self.cte.schema().fields().len()).collect();
+ Some(vec![all_cte_columns, output_columns.to_vec()])
+ }
+}
+
+/// Reads the buffered output of the [`MaterializedCte`] with the same `id`.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct MaterializedCteScan {
+ pub id: MaterializedCteId,
+ pub name: String,
+ pub schema: DFSchemaRef,
+}
+
+/// Orders by `id` and `name`. Scans with the same `id` and `name` but a
+/// different `schema` are not comparable.
+impl PartialOrd for MaterializedCteScan {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ match self
+ .id
+ .cmp(&other.id)
+ .then_with(|| self.name.cmp(&other.name))
+ {
+ Ordering::Equal if self != other => None,
+ ord => Some(ord),
+ }
+ }
+}
+
+impl UserDefinedLogicalNodeCore for MaterializedCteScan {
+ fn name(&self) -> &str {
+ "MaterializedCteScan"
+ }
+
+ fn inputs(&self) -> Vec<&LogicalPlan> {
+ vec![]
+ }
+
+ fn schema(&self) -> &DFSchemaRef {
+ &self.schema
+ }
+
+ fn expressions(&self) -> Vec<Expr> {
+ vec![]
+ }
+
+ fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "MaterializedCteScan: name={}", self.name)
+ }
+
+ fn with_exprs_and_inputs(
+ &self,
+ exprs: Vec<Expr>,
+ inputs: Vec<LogicalPlan>,
+ ) -> Result<Self> {
+ assert_eq_or_internal_err!(exprs.len(), 0, "MaterializedCteScan takes no exprs");
+ assert_eq_or_internal_err!(
+ inputs.len(),
+ 0,
+ "MaterializedCteScan takes no inputs"
+ );
+ Ok(self.clone())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::EmptyRelation;
+ use std::sync::Arc;
+
+ fn empty(produce_one_row: bool) -> LogicalPlan {
+ LogicalPlan::EmptyRelation(EmptyRelation {
+ produce_one_row,
+ schema: Arc::new(datafusion_common::DFSchema::empty()),
+ })
+ }
+
+ #[test]
+ fn partial_cmp_is_equal_only_for_equal_nodes() {
+ let id = MaterializedCteId::next();
+ let a = MaterializedCte {
+ id,
+ name: "c".to_string(),
+ cte: empty(false),
+ continuation: empty(false),
+ };
+ let mut b = a.clone();
+ assert_eq!(a.partial_cmp(&b), Some(Ordering::Equal));
+
+ b.continuation = empty(true);
+ assert_ne!(a, b);
+ assert_eq!(a.partial_cmp(&b), None);
+
+ let c = MaterializedCte {
+ id: MaterializedCteId::next(),
+ ..a.clone()
+ };
+ assert_eq!(a.partial_cmp(&c), Some(Ordering::Less));
+ }
+}
diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs
index 7f2208ae12..81483dae7f 100644
--- a/datafusion/expr/src/logical_plan/mod.rs
+++ b/datafusion/expr/src/logical_plan/mod.rs
@@ -21,6 +21,7 @@ pub mod display;
pub mod dml;
mod extension;
pub(crate) mod invariants;
+pub mod materialized_cte;
pub use invariants::{InvariantLevel, assert_expected_schema, check_subquery_expr};
mod plan;
mod statement;
@@ -60,3 +61,4 @@ pub use datafusion_common::format::ExplainFormat;
pub use display::display_schema;
pub use extension::{UserDefinedLogicalNode, UserDefinedLogicalNodeCore};
+pub use materialized_cte::{MaterializedCte, MaterializedCteId, MaterializedCteScan};
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index 0213f78283..1844966ec6 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -39,6 +39,7 @@ use crate::expr_rewriter::{
};
use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor};
use crate::logical_plan::extension::UserDefinedLogicalNode;
+use crate::logical_plan::materialized_cte::MaterializedCte;
use crate::logical_plan::{DmlStatement, Statement, WriteOp};
use crate::utils::{
check_aggregate_and_window_nesting, check_no_window_functions,
@@ -610,13 +611,19 @@ impl LogicalPlan {
})
.map_or(Ok(None), |v| v.map(Some))
}
+ // A materialized CTE outputs the output of its continuation.
+ LogicalPlan::Extension(Extension { node }) => {
+ match node.as_any().downcast_ref::<MaterializedCte>() {
+ Some(cte) => cte.continuation.head_output_expr(),
+ None => Ok(None),
+ }
+ }
LogicalPlan::Subquery(_) => Ok(None),
LogicalPlan::EmptyRelation(_)
| LogicalPlan::Statement(_)
| LogicalPlan::Values(_)
| LogicalPlan::Explain(_)
| LogicalPlan::Analyze(_)
- | LogicalPlan::Extension(_)
| LogicalPlan::Dml(_)
| LogicalPlan::Copy(_)
| LogicalPlan::Ddl(_)
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
index 47aba9e715..e31dc65173 100644
--- a/datafusion/optimizer/src/analyzer/type_coercion.rs
+++ b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -28,7 +28,9 @@ use crate::utils::{NamePreserver, merge_into_schema};
use arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit};
use arrow::temporal_conversions::SECONDS_IN_DAY;
use datafusion_common::config::ConfigOptions;
-use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
+use datafusion_common::tree_node::{
+ Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter,
+};
use datafusion_common::{
Column, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference,
exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err,
@@ -41,7 +43,9 @@ use datafusion_expr::expr::{
};
use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
use datafusion_expr::expr_schema::{cast_subquery, cast_subquery_columns};
-use datafusion_expr::logical_plan::Subquery;
+use datafusion_expr::logical_plan::{
+ Extension, MaterializedCte, MaterializedCteId, MaterializedCteScan, Subquery,
+};
use datafusion_expr::type_coercion::binary::{
comparison_coercion_with_session_timezone, like_coercion, regex_coercion,
type_union_coercion,
@@ -96,8 +100,6 @@ impl AnalyzerRule for TypeCoercion {
}
fn analyze(&self, plan: LogicalPlan, config: &ConfigOptions) -> Result<LogicalPlan> {
- static EMPTY_SCHEMA: LazyLock<DFSchema> = LazyLock::new(DFSchema::empty);
-
// Record bare row-count parameter types before coercion makes them indistinguishable
// from explicit casts. A type inferred elsewhere takes precedence over the default.
let parameter_fields = plan.get_parameter_fields()?;
@@ -118,21 +120,114 @@ impl AnalyzerRule for TypeCoercion {
.data;
// recurse
- let transformed_plan = plan
- .transform_up_with_subqueries(|plan| {
- analyze_internal(
- &EMPTY_SCHEMA,
- plan,
- config.execution.time_zone.as_ref().map(|tz| tz.as_str()),
- )
- })?
- .data;
+ let transformed_plan = coerce_plan_and_subqueries(plan, config)?;
// finish
coerce_output(transformed_plan, config)
}
}
+/// Coerce each node of `plan` and of its subqueries, from the leaves up.
+fn coerce_plan_and_subqueries(
+ plan: LogicalPlan,
+ config: &ConfigOptions,
+) -> Result<LogicalPlan> {
+ static EMPTY_SCHEMA: LazyLock<DFSchema> = LazyLock::new(DFSchema::empty);
+
+ plan.transform_down_up_with_subqueries(
+ |plan| {
+ let cte = match &plan {
+ LogicalPlan::Extension(Extension { node }) => {
+ node.as_any().downcast_ref::<MaterializedCte>().cloned()
+ }
+ _ => None,
+ };
+ let Some(cte) = cte else {
+ return Ok(Transformed::no(plan));
+ };
+ // Coercion can change the types of the body, for example the type
+ // of a `NULL` column of a `UNION`. Coerce the body first, and give
+ // its new schema to the scans before the continuation is coerced.
+ let body = coerce_plan_and_subqueries(cte.cte, config)?;
+ let continuation = with_materialized_cte_scan_schema(
+ cte.continuation,
+ cte.id,
+ body.schema(),
+ )?;
+ let continuation = coerce_plan_and_subqueries(continuation, config)?;
+ Ok(Transformed::new(
+ LogicalPlan::Extension(Extension {
+ node: Arc::new(MaterializedCte {
+ id: cte.id,
+ name: cte.name,
+ cte: body,
+ continuation,
+ }),
+ }),
+ true,
+ TreeNodeRecursion::Jump,
+ ))
+ },
+ |plan| {
+ analyze_internal(
+ &EMPTY_SCHEMA,
+ plan,
+ config.execution.time_zone.as_ref().map(|tz| tz.as_str()),
+ )
+ },
+ )
+ .map(|t| t.data)
+}
+
+/// Give the fields of `body_schema` to each scan of the materialized CTE `id`
+/// in `continuation`, also in subquery expressions. A scan keeps its
+/// qualifiers.
+fn with_materialized_cte_scan_schema(
+ continuation: LogicalPlan,
+ id: MaterializedCteId,
+ body_schema: &DFSchemaRef,
+) -> Result<LogicalPlan> {
+ continuation
+ .transform_down_with_subqueries(|plan| {
+ let LogicalPlan::Extension(Extension { node }) = &plan else {
+ return Ok(Transformed::no(plan));
+ };
+ // A nested occurrence of the same CTE owns the scans below it.
+ if node
+ .as_any()
+ .downcast_ref::<MaterializedCte>()
+ .is_some_and(|nested| nested.id == id)
+ {
+ return Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump));
+ }
+ let Some(scan) =
+ node.as_any()
+ .downcast_ref::<MaterializedCteScan>()
+ .filter(|scan| {
+ scan.id == id && scan.schema.inner() != body_schema.inner()
+ })
+ else {
+ return Ok(Transformed::no(plan));
+ };
+ let qualifiers = scan.schema.iter().map(|(q, _)| q.cloned()).collect();
+ let schema = DFSchema::from_field_specific_qualified_schema(
+ qualifiers,
+ body_schema.inner(),
+ )?
+ .with_functional_dependencies(
+ body_schema.functional_dependencies().clone(),
+ )?;
+ Ok(Transformed::yes(LogicalPlan::Extension(Extension {
+ node: Arc::new(MaterializedCteScan {
+ id: scan.id,
+ name: scan.name.clone(),
+ schema: Arc::new(schema),
+ }),
+ })))
+ })
+ .map(|t| t.data)
+}
+
/// use the external schema to handle the correlated subqueries case
///
/// Assumes that children have already been optimized
diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs
index 43037b258f..ca2f56c417 100644
--- a/datafusion/optimizer/src/decorrelate.rs
+++ b/datafusion/optimizer/src/decorrelate.rs
@@ -36,8 +36,9 @@ use datafusion_expr::utils::{
collect_subquery_cols, conjunction, find_join_exprs, split_conjunction,
};
use datafusion_expr::{
- BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, FetchType,
- LogicalPlan, LogicalPlanBuilder, Operator, SkipType, expr, lit,
+ BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, Extension, FetchType,
+ LogicalPlan, LogicalPlanBuilder, MaterializedCte, Operator, SkipType,
+ UserDefinedLogicalNodeCore, expr, lit,
};
/// This struct rewrite the sub query plan by pull up the correlated
@@ -163,6 +164,38 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
LogicalPlan::Subquery(_) => {
Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
}
+ LogicalPlan::Extension(Extension { node })
+ if node.as_any().is::<MaterializedCte>() =>
+ {
+ let cte = node.as_any().downcast_ref::<MaterializedCte>().unwrap();
+ if !cte.cte.all_out_ref_exprs().is_empty() {
+ // A correlated body would have to run once per outer row.
+ return Ok(
+ self.stop_pull_up(LogicalPlan::Extension(Extension { node }))
+ );
+ }
+ // Only the continuation can hold correlated expressions. Rewrite
+ // it alone, so that the body shared by every scan keeps its schema.
+ let continuation = cte.continuation.clone().rewrite(self)?;
+ if !continuation.transformed {
+ return Ok(Transformed::new(
+ LogicalPlan::Extension(Extension { node }),
+ false,
+ TreeNodeRecursion::Jump,
+ ));
+ }
+ let cte = cte.with_exprs_and_inputs(
+ vec![],
+ vec![cte.cte.clone(), continuation.data],
+ )?;
+ Ok(Transformed::new(
+ LogicalPlan::Extension(Extension {
+ node: Arc::new(cte),
+ }),
+ true,
+ TreeNodeRecursion::Jump,
+ ))
+ }
// A correlated filter can only move above a join from a side whose
// rows the join preserves. Below the side an outer join fills with
// NULLs, the filter decides which rows are unmatched, so pulling it
diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs
index 2fb2c953fe..becc789d81 100644
--- a/datafusion/optimizer/src/push_down_filter.rs
+++ b/datafusion/optimizer/src/push_down_filter.rs
@@ -62,8 +62,9 @@ use datafusion_expr::utils::{
conjunction, expr_to_columns, split_conjunction, split_conjunction_owned,
};
use datafusion_expr::{
- BinaryExpr, Distinct, EmptyRelation, Expr, Filter, Operator, Projection,
- TableProviderFilterPushDown, and, or,
+ BinaryExpr, Distinct, EmptyRelation, Expr, Extension, Filter, MaterializedCte,
+ Operator, Projection, TableProviderFilterPushDown, UserDefinedLogicalNodeCore, and,
+ or,
};
use crate::extract_leaf_expressions::is_pure_extraction_projection;
@@ -1398,6 +1399,23 @@ impl OptimizerRule for PushDownFilter {
LogicalPlan::TableScan(scan),
)))
}
+ LogicalPlan::Extension(Extension { node })
+ if node.as_any().is::<MaterializedCte>() =>
+ {
+ // The output is the output of the continuation, so the filter
+ // moves onto the continuation. The body is shared by every scan
+ // of the CTE and is never filtered for one of them.
+ let cte = node.as_any().downcast_ref::<MaterializedCte>().unwrap();
+ let continuation = LogicalPlan::Filter(Filter::new(
+ filter.predicate,
+ Arc::new(cte.continuation.clone()),
+ ));
+ let cte = cte
+ .with_exprs_and_inputs(vec![], vec![cte.cte.clone(), continuation])?;
+ Ok(Transformed::yes(LogicalPlan::Extension(Extension {
+ node: Arc::new(cte),
+ })))
+ }
LogicalPlan::Extension(extension_plan) => {
// This check prevents the Filter from being removed when the extension node has no children,
// so we return the original Filter unchanged.
diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs
index 77780c9f8e..cce3f08572 100644
--- a/datafusion/physical-plan/src/lib.rs
+++ b/datafusion/physical-plan/src/lib.rs
@@ -86,6 +86,7 @@ pub mod filter;
pub mod filter_pushdown;
pub mod joins;
pub mod limit;
+pub mod materialized_cte;
pub mod memory;
pub mod metrics;
pub mod operator_statistics;
diff --git a/datafusion/physical-plan/src/materialized_cte.rs b/datafusion/physical-plan/src/materialized_cte.rs
new file mode 100644
index 0000000000..194cd6825b
--- /dev/null
+++ b/datafusion/physical-plan/src/materialized_cte.rs
@@ -0,0 +1,989 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Execution plans for `WITH x AS MATERIALIZED (...)`.
+//!
+//! [`MaterializedCteExec`] owns the CTE body and the continuation (the rest of
+//! the query). Every [`MaterializedCteScanExec`] in the continuation shares one
+//! [`MaterializedCteBuffer`] with it. The first consumer to execute runs every
+//! partition of the body to completion and buffers the output; every scan then
+//! replays that buffer.
+//!
+//! The body is fully materialized before any scan yields a row. This is a
+//! pipeline break, but it cannot deadlock when two consumers of the same CTE
+//! are read at different rates (for example the build and probe side of one
+//! hash join), which a bounded fan-out channel can.
+//!
+//! Buffered batches are accounted in the memory pool. When a reservation
+//! fails, the rest of that partition is written to a spill file, and the scans
+//! read the in-memory prefix and then the spill files, so row order within a
+//! partition is kept.
+//!
+//! A scan reads a spill file one batch at a time, without read-ahead. After
+//! the body is buffered, memory for one decoded batch per concurrent reader is
+//! reserved. If the pool cannot grant it, in-memory batches are moved to spill
+//! files until it can. An in-memory batch that is larger than the batches on
+//! disk is written in pieces, so that the move does not raise the memory that
+//! each reader needs.
+
+use std::fmt;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::TreeNodeRecursion;
+use datafusion_common::{Result, Statistics, internal_err};
+use datafusion_common_runtime::JoinSet;
+use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion_execution::{SendableRecordBatchStream, SpillFile, TaskContext};
+use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
+use futures::stream::BoxStream;
+use futures::{Stream, StreamExt, TryStreamExt};
+use parking_lot::Mutex;
+
+use crate::coop::cooperative;
+use crate::execution_plan::{
+ Boundedness, CardinalityEffect, EmissionType, ExecutionPlan, ExecutionPlanProperties,
+ PlanProperties, SchedulingType,
+};
+use crate::filter_pushdown::{
+ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase,
+ FilterPushdownPropagation,
+};
+use crate::joins::utils::{OnceAsync, OnceFut};
+use crate::metrics::{
+ BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, SpillMetrics,
+};
+use crate::spill::spill_manager::SpillManager;
+use crate::statistics::{ChildStats, StatisticsArgs};
+use crate::stream::RecordBatchStreamAdapter;
+use crate::{
+ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ReplaceChildrenOptions,
+};
+
+/// The buffered output of one body partition: an in-memory prefix followed by
+/// spill files with the remaining batches, in order.
+struct BufferedPartition {
+ batches: Vec<RecordBatch>,
+ spill_files: Vec<Arc<dyn SpillFile>>,
+ /// The memory size of the largest batch in `spill_files` when decoded.
+ max_spilled_batch_size: usize,
+ /// Holds `batches`. Released when the buffer is dropped.
+ reservation: MemoryReservation,
+}
+
+/// A plan to move the in-memory batches `split..` of a [`BufferedPartition`]
+/// to a spill file. See [`BufferedPartition::plan_spill`].
+struct SpillPlan {
+ split: usize,
+ /// A batch larger than this is written in pieces of about this size.
+ piece_size: usize,
+ /// The memory that the move releases, less the memory that it adds to
+ /// the replay of the spill files.
+ gain: usize,
+}
+
+impl BufferedPartition {
+ /// Plan a move of in-memory batches from the end of `batches` to a spill
+ /// file, to gain at least `bytes`.
+ ///
+ /// A spilled batch that is larger than every batch already on disk makes
+ /// each scan reserve more memory to decode it. `added_replay_memory`
+ /// returns that increase for a new `max_spilled_batch_size`. To keep it
+ /// small, a batch larger than `piece_size` is written in pieces.
+ ///
+ /// Returns `None` when no move releases more memory than it adds.
+ fn plan_spill(
+ &self,
+ bytes: usize,
+ piece_size: usize,
+ added_replay_memory: impl Fn(usize) -> usize,
+ ) -> Option<SpillPlan> {
+ let mut best: Option<SpillPlan> = None;
+ let mut released = 0;
+ let mut max_size = self.max_spilled_batch_size;
+ let mut added = 0;
+ for (split, batch) in self.batches.iter().enumerate().rev() {
+ let size = batch.get_array_memory_size();
+ released += size;
+ // A batch of one row cannot be written in pieces.
+ let spilled_size = match batch.num_rows() {
+ 0 | 1 => size,
+ _ => size.min(piece_size),
+ };
+ if spilled_size > max_size {
+ max_size = spilled_size;
+ added = added_replay_memory(max_size);
+ }
+ let gain = released.saturating_sub(added);
+ if gain > best.as_ref().map_or(0, |b| b.gain) {
+ best = Some(SpillPlan {
+ split,
+ piece_size,
+ gain,
+ });
+ }
+ if gain >= bytes {
+ break;
+ }
+ }
+ best
+ }
+
+ /// Move the in-memory batches of `plan` to a new spill file. The new file
+ /// is read before the existing spill files, so the order of the partition
+ /// is kept.
+ fn spill(&mut self, plan: &SpillPlan, spill_manager: &SpillManager) -> Result<()> {
+ let suffix = self.batches.split_off(plan.split);
+ let mut file = spill_manager.create_in_progress_file("MaterializedCte")?;
+ let mut released = 0;
+ for batch in &suffix {
+ let size = batch.get_array_memory_size();
+ released += size;
+ let rows = batch.num_rows();
+ // Rows per piece, so that a piece is about `piece_size`.
+ let piece_rows = rows.div_ceil(size.div_ceil(plan.piece_size).max(1)).max(1);
+ for offset in (0..rows.max(1)).step_by(piece_rows) {
+ let piece = batch.slice(offset, piece_rows.min(rows - offset));
+ let size = file.append_batch(&piece)?;
+ self.max_spilled_batch_size = self.max_spilled_batch_size.max(size);
+ }
+ }
+ if let Some(file) = file.finish()? {
+ self.spill_files.insert(0, file);
+ }
+ self.reservation.shrink(released);
+ Ok(())
+ }
+}
+
+/// The materialized output of a CTE body.
+struct MaterializedOutput {
+ partitions: Vec<BufferedPartition>,
+ spill_manager: SpillManager,
+ /// Memory for the batches that the scans decode from the spill files.
+ /// Released when the buffer is dropped.
+ _replay_reservation: MemoryReservation,
+}
+
+/// The memory that the scans need to read the spill files of the body
+/// partitions. `max_spilled_batch_sizes` has one entry per body partition.
+///
+/// Partition `p` of a scan with `n` partitions reads the body partitions
+/// `p, p + n, ...` one after the other, with at most one decoded batch in
+/// memory. All scan partitions can read at the same time.
+fn replay_memory(max_spilled_batch_sizes: &[usize], scan_partitions: &[usize]) -> usize {
+ scan_partitions
+ .iter()
+ .map(|&n| {
+ (0..n)
+ .map(|p| {
+ max_spilled_batch_sizes
+ .iter()
+ .skip(p)
+ .step_by(n)
+ .max()
+ .copied()
+ .unwrap_or(0)
+ })
+ .sum::<usize>()
+ })
+ .sum()
+}
+
+/// Grow `reservation` to the memory the scans need to replay `partitions`.
+/// If the pool cannot grant it, move in-memory batches to spill files until
+/// it can.
+///
+/// A move can also raise the memory that the replay needs, because a scan
+/// reserves memory for the largest batch that it reads from disk. Each step
+/// makes the move that gains the most memory, and a batch that is larger than
+/// the batches on disk is written in pieces. If no move gains memory, the
+/// allocation error is returned.
+fn reserve_replay_memory(
+ partitions: &mut [BufferedPartition],
+ scan_partitions: &[usize],
+ spill_manager: &SpillManager,
+ reservation: &mut MemoryReservation,
+) -> Result<()> {
+ loop {
+ let sizes: Vec<_> = partitions
+ .iter()
+ .map(|p| p.max_spilled_batch_size)
+ .collect();
+ let required = replay_memory(&sizes, scan_partitions);
+ let Err(e) = reservation.try_resize(required) else {
+ return Ok(());
+ };
+ // `required` is not zero, so at least one partition has a spill file.
+ let largest = sizes.iter().copied().max().unwrap_or(0).max(1);
+ let plan = partitions
+ .iter()
+ .enumerate()
+ .filter_map(|(i, partition)| {
+ let piece_size = match sizes[i] {
+ 0 => largest,
+ size => size,
+ };
+ let added_replay_memory = |size| {
+ let mut sizes = sizes.clone();
+ sizes[i] = size;
+ replay_memory(&sizes, scan_partitions).saturating_sub(required)
+ };
+ let plan =
+ partition.plan_spill(required, piece_size, added_replay_memory)?;
+ Some((i, plan))
+ })
+ .max_by_key(|(_, plan)| plan.gain);
+ let Some((i, plan)) = plan else {
+ return Err(e);
+ };
+ partitions[i].spill(&plan, spill_manager)?;
+ }
+}
+
+/// State shared by one [`MaterializedCteExec`] and all its
+/// [`MaterializedCteScanExec`]s.
+pub struct MaterializedCteBuffer {
+ id: u64,
+ name: String,
+ /// The body to run. [`MaterializedCteExec`] updates it every time it is
+ /// rebuilt (for example by a physical optimizer rule), so it is the body of
+ /// the final plan. A scan can run the body before the owning
+ /// `MaterializedCteExec` executes, for example from a scalar subquery.
+ body: Mutex<Option<Arc<dyn ExecutionPlan>>>,
+ output: Mutex<Arc<OnceAsync<MaterializedOutput>>>,
+ /// The partition count of every scan bound to this buffer.
+ scan_partitions: Mutex<Vec<usize>>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl fmt::Debug for MaterializedCteBuffer {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("MaterializedCteBuffer")
+ .field("name", &self.name)
+ .finish_non_exhaustive()
+ }
+}
+
+impl MaterializedCteBuffer {
+ pub fn new(id: u64, name: impl Into<String>) -> Self {
+ Self {
+ id,
+ name: name.into(),
+ body: Mutex::new(None),
+ output: Mutex::new(Arc::default()),
+ scan_partitions: Mutex::new(vec![]),
+ metrics: ExecutionPlanMetricsSet::new(),
+ }
+ }
+
+ /// The id that binds scans to this buffer.
+ pub fn id(&self) -> u64 {
+ self.id
+ }
+
+ fn set_body(&self, body: Arc<dyn ExecutionPlan>) {
+ *self.body.lock() = Some(body);
+ }
+
+ fn reset(&self) {
+ *self.output.lock() = Arc::default();
+ }
+
+ /// Return a future that resolves once the body has been fully buffered.
+ /// The body runs at most once, whichever consumer asks first.
+ fn materialize(
+ self: &Arc<Self>,
+ context: Arc<TaskContext>,
+ ) -> Result<OnceFut<MaterializedOutput>> {
+ let Some(body) = self.body.lock().clone() else {
+ return internal_err!("MaterializedCte {} has no body", self.name);
+ };
+ let once = Arc::clone(&self.output.lock());
+ let name = self.name.clone();
+ let scan_partitions = self.scan_partitions.lock().clone();
+ let metrics = self.metrics.clone();
+ once.try_once(move || {
+ Ok(buffer_body(name, body, scan_partitions, context, metrics))
+ })
+ }
+}
+
+async fn buffer_body(
+ name: String,
+ body: Arc<dyn ExecutionPlan>,
+ scan_partitions: Vec<usize>,
+ context: Arc<TaskContext>,
+ metrics: ExecutionPlanMetricsSet,
+) -> Result<MaterializedOutput> {
+ let spill_manager = SpillManager::new(
+ context.runtime_env(),
+ SpillMetrics::new(&metrics, 0),
+ body.schema(),
+ )
+ .with_compression_type(context.session_config().spill_compression());
+ let buffered_rows = MetricBuilder::new(&metrics).global_counter("buffered_rows");
+
+ let mut join_set = JoinSet::new();
+ for partition in 0..body.output_partitioning().partition_count() {
+ let stream = body.execute(partition, Arc::clone(&context))?;
+ let reservation =
+ MemoryConsumer::new(format!("MaterializedCte[{name}][{partition}]"))
+ .with_can_spill(true)
+ .register(context.memory_pool());
+ let spill_manager = spill_manager.clone();
+ let buffered_rows = buffered_rows.clone();
+ join_set.spawn(async move {
+ let result =
+ buffer_partition(stream, reservation, &spill_manager, &buffered_rows)
+ .await;
+ (partition, result)
+ });
+ }
+
+ let mut partitions = Vec::with_capacity(join_set.len());
+ while let Some(joined) = join_set.join_next().await {
+ match joined {
+ Ok((partition, result)) => partitions.push((partition, result?)),
+ Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()),
+ Err(e) => return internal_err!("MaterializedCte task failed: {e}"),
+ }
+ }
+ partitions.sort_by_key(|(partition, _)| *partition);
+ let mut partitions: Vec<_> = partitions.into_iter().map(|(_, p)| p).collect();
+
+ let mut replay_reservation =
+ MemoryConsumer::new(format!("MaterializedCte[{name}] replay"))
+ .register(context.memory_pool());
+ reserve_replay_memory(
+ &mut partitions,
+ &scan_partitions,
+ &spill_manager,
+ &mut replay_reservation,
+ )?;
+ Ok(MaterializedOutput {
+ partitions,
+ spill_manager,
+ _replay_reservation: replay_reservation,
+ })
+}
+
+async fn buffer_partition(
+ mut stream: SendableRecordBatchStream,
+ reservation: MemoryReservation,
+ spill_manager: &SpillManager,
+ buffered_rows: &crate::metrics::Count,
+) -> Result<BufferedPartition> {
+ let mut batches = vec![];
+ let mut spill = None;
+ let mut max_spilled_batch_size = 0;
+ while let Some(batch) = stream.next().await.transpose()? {
+ buffered_rows.add(batch.num_rows());
+ // Once a partition spills, every later batch goes to the same file,
+ // so replay keeps the order of the partition.
+ if spill.is_none() && reservation.try_grow(batch.get_array_memory_size()).is_ok()
+ {
+ batches.push(batch);
+ continue;
+ }
+ if spill.is_none() {
+ spill = Some(spill_manager.create_in_progress_file("MaterializedCte")?);
+ }
+ if let Some(file) = spill.as_mut() {
+ let size = file.append_batch(&batch)?;
+ max_spilled_batch_size = max_spilled_batch_size.max(size);
+ }
+ }
+ let spill_file = match spill {
+ Some(mut file) => file.finish()?,
+ None => None,
+ };
+ Ok(BufferedPartition {
+ batches,
+ spill_files: spill_file.into_iter().collect(),
+ max_spilled_batch_size,
+ reservation,
+ })
+}
+
+/// Stream the buffered partitions `partition, partition + n, ...` of `output`,
+/// where `n` is the number of output partitions of the scan.
+fn replay(
+ output: Arc<MaterializedOutput>,
+ partition: usize,
+ output_partitions: usize,
+) -> Result<ReplayStream> {
+ let mut streams = vec![];
+ for buffered in output
+ .partitions
+ .iter()
+ .skip(partition)
+ .step_by(output_partitions)
+ {
+ let schema = Arc::clone(output.spill_manager.schema());
+ let batches = buffered.batches.clone();
+ streams.push(Box::pin(RecordBatchStreamAdapter::new(
+ schema,
+ futures::stream::iter(batches.into_iter().map(Ok)),
+ )) as SendableRecordBatchStream);
+ // Read without read-ahead, so that the decoded batch fits in the
+ // replay reservation.
+ for file in &buffered.spill_files {
+ streams.push(output.spill_manager.read_spill_as_stream_unbuffered(
+ Arc::clone(file),
+ Some(buffered.max_spilled_batch_size),
+ )?);
+ }
+ }
+ Ok(ReplayStream {
+ stream: futures::stream::iter(streams).flatten().boxed(),
+ _output: output,
+ })
+}
+
+/// The replay of one scan partition.
+///
+/// Holds the [`MaterializedOutput`], and with it the memory reservations of
+/// the buffered batches and of the spill reads, until the replay is dropped.
+/// The scan can outlive its plan, for example when the plan is dropped after
+/// `execute`, so the output must not be released earlier.
+struct ReplayStream {
+ stream: BoxStream<'static, Result<RecordBatch>>,
+ _output: Arc<MaterializedOutput>,
+}
+
+impl Stream for ReplayStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.stream.poll_next_unpin(cx)
+ }
+}
+
+/// Computes a CTE body once and runs the continuation, whose
+/// [`MaterializedCteScanExec`]s read the buffered body output.
+///
+/// Children: `[body, continuation]`. The output is the output of the
+/// continuation.
+#[derive(Debug)]
+pub struct MaterializedCteExec {
+ body: Arc<dyn ExecutionPlan>,
+ continuation: Arc<dyn ExecutionPlan>,
+ buffer: Arc<MaterializedCteBuffer>,
+ cache: Arc<PlanProperties>,
+}
+
+impl MaterializedCteExec {
+ pub fn new(
+ body: Arc<dyn ExecutionPlan>,
+ continuation: Arc<dyn ExecutionPlan>,
+ buffer: Arc<MaterializedCteBuffer>,
+ ) -> Self {
+ buffer.set_body(Arc::clone(&body));
+ let cache = Arc::clone(continuation.properties());
+ Self {
+ body,
+ continuation,
+ buffer,
+ cache,
+ }
+ }
+
+ pub fn buffer(&self) -> &Arc<MaterializedCteBuffer> {
+ &self.buffer
+ }
+}
+
+impl DisplayAs for MaterializedCteExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "MaterializedCteExec: name={}", self.buffer.name)
+ }
+ DisplayFormatType::TreeRender => write!(f, "name={}", self.buffer.name),
+ }
+ }
+}
+
+impl ExecutionPlan for MaterializedCteExec {
+ fn name(&self) -> &'static str {
+ "MaterializedCteExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.body, &self.continuation]
+ }
+
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 2 {
+ return internal_err!("MaterializedCteExec takes 2 children");
+ }
+ let continuation = children.pop().unwrap();
+ let body = children.pop().unwrap();
+ Ok(Arc::new(Self::new(
+ body,
+ continuation,
+ Arc::clone(&self.buffer),
+ )))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ /// Clears the shared buffer in place and returns `self`.
+ ///
+ /// This is an exception to the `reset_state` contract, which expects a
+ /// new node and leaves the old node unchanged. `RecursiveQueryExec` resets
+ /// the recursive term before each iteration, and a `MATERIALIZED` CTE in
+ /// that term must be computed again. A new buffer would also need new
+ /// scans, and a scan in a scalar subquery is not below this node, so this
+ /// node cannot rebind it. `ScalarSubqueryExec::reset_state` does the same.
+ ///
+ /// Two overlapping runs of one plan share this buffer. Nothing in
+ /// DataFusion does that today. See
+ /// <https://github.com/apache/datafusion/issues/26229>.
+ fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> {
+ self.buffer.reset();
+ Ok(self)
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![false, true]
+ }
+
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ vec![true, false]
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ // The body runs lazily, when the first scan polls it.
+ self.buffer.set_body(Arc::clone(&self.body));
+ self.continuation.execute(partition, context)
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.buffer.metrics.clone_inner())
+ }
+
+ fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+ vec![ChildStats::Skip, ChildStats::At(partition)]
+ }
+
+ fn statistics_from_inputs(
+ &self,
+ input_stats: &[Arc<Statistics>],
+ _args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ Ok(Arc::clone(&input_stats[1]))
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ CardinalityEffect::Equal
+ }
+
+ /// The output is the output of the continuation, so parent filters go to
+ /// the continuation only. The body is shared by every scan and must not
+ /// be filtered for one consumer.
+ fn gather_filters_for_pushdown(
+ &self,
+ _phase: FilterPushdownPhase,
+ parent_filters: Vec<Arc<dyn PhysicalExpr>>,
+ _config: &ConfigOptions,
+ ) -> Result<FilterDescription> {
+ Ok(FilterDescription::new()
+ .with_child(ChildFilterDescription::all_unsupported(&parent_filters))
+ .with_child(ChildFilterDescription::from_child(
+ &parent_filters,
+ &self.continuation,
+ )?))
+ }
+
+ fn handle_child_pushdown_result(
+ &self,
+ _phase: FilterPushdownPhase,
+ child_pushdown_result: ChildPushdownResult,
+ _config: &ConfigOptions,
+ ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
+ // The body always reports `No`, so a filter is handled when the
+ // continuation handles it.
+ Ok(FilterPushdownPropagation::if_any(child_pushdown_result))
+ }
+}
+
+/// Reads the buffered output of a [`MaterializedCteExec`] body.
+#[derive(Debug)]
+pub struct MaterializedCteScanExec {
+ id: u64,
+ name: String,
+ schema: SchemaRef,
+ buffer: Option<Arc<MaterializedCteBuffer>>,
+ metrics: ExecutionPlanMetricsSet,
+ cache: Arc<PlanProperties>,
+}
+
+impl MaterializedCteScanExec {
+ /// Create a scan that is not yet bound to a buffer. See [`Self::bind`].
+ pub fn new(
+ id: u64,
+ name: impl Into<String>,
+ schema: SchemaRef,
+ partitions: usize,
+ ) -> Self {
+ let cache = PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(partitions.max(1)),
+ EmissionType::Final,
+ Boundedness::Bounded,
+ )
+ .with_scheduling_type(SchedulingType::Cooperative);
+ Self {
+ id,
+ name: name.into(),
+ schema,
+ buffer: None,
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache: Arc::new(cache),
+ }
+ }
+
+ /// Bind this scan to the buffer of its [`MaterializedCteExec`].
+ pub fn bind(&self, buffer: Arc<MaterializedCteBuffer>) -> Self {
+ buffer
+ .scan_partitions
+ .lock()
+ .push(self.cache.partitioning.partition_count());
+ Self {
+ id: self.id,
+ name: self.name.clone(),
+ schema: Arc::clone(&self.schema),
+ buffer: Some(buffer),
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache: Arc::clone(&self.cache),
+ }
+ }
+
+ /// The id of the [`MaterializedCteBuffer`] this scan reads.
+ pub fn id(&self) -> u64 {
+ self.id
+ }
+
+ /// Whether [`Self::bind`] has been called on this scan.
+ pub fn is_bound(&self) -> bool {
+ self.buffer.is_some()
+ }
+}
+
+impl DisplayAs for MaterializedCteScanExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "MaterializedCteScanExec: name={}", self.name)
+ }
+ DisplayFormatType::TreeRender => write!(f, "name={}", self.name),
+ }
+ }
+}
+
+impl ExecutionPlan for MaterializedCteScanExec {
+ fn name(&self) -> &'static str {
+ "MaterializedCteScanExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ _: Vec<Arc<dyn ExecutionPlan>>,
+ _: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let Some(buffer) = &self.buffer else {
+ return internal_err!(
+ "MaterializedCteScanExec {} is not bound to its CTE",
+ self.name
+ );
+ };
+ let mut output = buffer.materialize(context)?;
+ let output_partitions = self.cache.partitioning.partition_count();
+ let baseline = BaselineMetrics::new(&self.metrics, partition);
+ let stream = futures::stream::once(async move {
+ let output = std::future::poll_fn(|cx| output.get_shared(cx)).await?;
+ replay(output, partition, output_partitions)
+ })
+ .try_flatten()
+ .inspect_ok(move |batch| baseline.record_output(batch.num_rows()));
+ // The replay of in-memory batches never returns `Pending`, so it must
+ // consume the Tokio budget explicitly.
+ Ok(Box::pin(cooperative(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.schema),
+ Box::pin(stream),
+ ))))
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+
+ /// The statistics are unknown: a leaf cannot ask the body for its
+ /// statistics through [`StatisticsArgs`].
+ ///
+ /// This is a known limitation, and it can make a plan worse than the
+ /// plan of the inlined CTE. For example, `JoinSelection` does not know
+ /// that a scan is small, and can put the other, larger input of a hash
+ /// join on the build side.
+ fn statistics_from_inputs(
+ &self,
+ _input_stats: &[Arc<Statistics>],
+ _args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ Ok(Arc::new(Statistics::new_unknown(&self.schema)))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::common::collect;
+ use crate::empty::EmptyExec;
+ use crate::filter_pushdown::PushedDown;
+ use crate::test::TestMemoryExec;
+ use arrow::array::{Array, Int32Array};
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+ use datafusion_physical_expr::expressions::{col, lit};
+
+ #[test]
+ fn parent_filters_go_to_the_continuation_only() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
+ let body = Arc::new(EmptyExec::new(Arc::clone(&schema)));
+ let continuation = Arc::new(EmptyExec::new(Arc::clone(&schema)));
+ let exec = MaterializedCteExec::new(
+ body,
+ continuation,
+ Arc::new(MaterializedCteBuffer::new(0, "c")),
+ );
+ let filter = Arc::new(datafusion_physical_expr::expressions::BinaryExpr::new(
+ col("a", &schema)?,
+ datafusion_expr::Operator::Eq,
+ lit(1i32),
+ )) as Arc<dyn PhysicalExpr>;
+
+ let description = exec.gather_filters_for_pushdown(
+ FilterPushdownPhase::Pre,
+ vec![filter],
+ &ConfigOptions::default(),
+ )?;
+ let parent_filters = description.parent_filters();
+ assert_eq!(parent_filters.len(), 2);
+ assert!(matches!(parent_filters[0][0].discriminant, PushedDown::No));
+ assert!(matches!(parent_filters[1][0].discriminant, PushedDown::Yes));
+ Ok(())
+ }
+
+ /// The in-memory prefix takes all the memory that the pool can grant.
+ /// Some of it is then moved to disk, so that the two scans have memory to
+ /// decode the spill files, and both scans still return every row in order.
+ #[tokio::test]
+ async fn spill_replay_memory_is_reserved() -> Result<()> {
+ let memory_limit = 20_000;
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+ let batches = (0..10)
+ .map(|i| {
+ let values = Int32Array::from_iter_values(i * 1000..(i + 1) * 1000);
+ RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(values)])
+ })
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ let body = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?;
+ let buffer = Arc::new(MaterializedCteBuffer::new(0, "c"));
+ let exec = MaterializedCteExec::new(
+ body,
+ Arc::new(EmptyExec::new(Arc::clone(&schema))),
+ Arc::clone(&buffer),
+ );
+ let scan = MaterializedCteScanExec::new(0, "c", Arc::clone(&schema), 1);
+ let scans = [
+ scan.bind(Arc::clone(&buffer)),
+ scan.bind(Arc::clone(&buffer)),
+ ];
+
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(memory_limit, 1.0)
+ .build_arc()?;
+ let context = Arc::new(TaskContext::default().with_runtime(runtime));
+
+ let mut once = buffer.materialize(Arc::clone(&context))?;
+ let output = std::future::poll_fn(|cx| once.get_shared(cx)).await?;
+ let partition = &output.partitions[0];
+ assert!(
+ partition.spill_files.len() > 1,
+ "the prefix was not spilled"
+ );
+ assert!(partition.max_spilled_batch_size > 0);
+ // The pool holds the in-memory prefix and one decoded batch per scan.
+ let reserved =
+ partition.reservation.size() + 2 * partition.max_spilled_batch_size;
+ assert_eq!(context.memory_pool().reserved(), reserved);
+ assert!(reserved <= memory_limit);
+
+ // The replay streams keep the reservations after the plan is dropped.
+ let mut streams = scans
+ .iter()
+ .map(|scan| scan.execute(0, Arc::clone(&context)))
+ .collect::<Result<Vec<_>>>()?;
+ drop((once, output, scans, exec, buffer));
+ // Once every stream has returned a batch, only the streams hold the
+ // buffered output.
+ let mut firsts = vec![];
+ for stream in &mut streams {
+ firsts.push(stream.next().await.transpose()?.unwrap());
+ }
+ assert_eq!(context.memory_pool().reserved(), reserved);
+
+ for (stream, first) in streams.into_iter().zip(firsts) {
+ let mut batches = vec![first];
+ batches.extend(collect(stream).await?);
+ let values: Vec<i32> = batches
+ .iter()
+ .flat_map(|b| {
+ let array = b.column(0).as_any().downcast_ref::<Int32Array>();
+ array.unwrap().values().to_vec()
+ })
+ .collect();
+ assert_eq!(values, (0..10_000).collect::<Vec<_>>());
+ }
+ assert_eq!(context.memory_pool().reserved(), 0);
+ Ok(())
+ }
+
+ /// The first batch is larger than the batches on disk. To move it to disk
+ /// whole would raise the memory that each scan needs to decode a batch
+ /// above the memory limit, so it is written in pieces.
+ #[tokio::test]
+ async fn spill_of_large_in_memory_batch_does_not_raise_replay_memory() -> Result<()> {
+ let memory_limit = 17_500;
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+ let batch = |range: std::ops::Range<i32>| {
+ let values = Int32Array::from_iter_values(range);
+ RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(values)])
+ };
+ let mut batches = vec![batch(0..4000)?];
+ for i in 0..10 {
+ batches.push(batch(4000 + i * 250..4000 + (i + 1) * 250)?);
+ }
+ let small_batch_size = batches[1].get_array_memory_size();
+ let body = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?;
+ let buffer = Arc::new(MaterializedCteBuffer::new(0, "c"));
+ let _exec = MaterializedCteExec::new(
+ body,
+ Arc::new(EmptyExec::new(Arc::clone(&schema))),
+ Arc::clone(&buffer),
+ );
+ let scan = MaterializedCteScanExec::new(0, "c", Arc::clone(&schema), 1);
+ let scans = [
+ scan.bind(Arc::clone(&buffer)),
+ scan.bind(Arc::clone(&buffer)),
+ ];
+
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(memory_limit, 1.0)
+ .build_arc()?;
+ let context = Arc::new(TaskContext::default().with_runtime(runtime));
+
+ let mut once = buffer.materialize(Arc::clone(&context))?;
+ let output = std::future::poll_fn(|cx| once.get_shared(cx)).await?;
+ let partition = &output.partitions[0];
+ assert!(partition.batches.is_empty(), "the large batch is in memory");
+ assert!(partition.max_spilled_batch_size <= small_batch_size);
+ assert!(context.memory_pool().reserved() <= memory_limit);
+
+ for scan in &scans {
+ let batches = collect(scan.execute(0, Arc::clone(&context))?).await?;
+ let values: Vec<i32> = batches
+ .iter()
+ .flat_map(|b| {
+ let array = b.column(0).as_any().downcast_ref::<Int32Array>();
+ array.unwrap().values().to_vec()
+ })
+ .collect();
+ assert_eq!(values, (0..6500).collect::<Vec<_>>());
+ }
+ Ok(())
+ }
+}
diff --git a/datafusion/sql/src/cte.rs b/datafusion/sql/src/cte.rs
index f735b33601..b6fbd21822 100644
--- a/datafusion/sql/src/cte.rs
+++ b/datafusion/sql/src/cte.rs
@@ -24,16 +24,42 @@ use datafusion_common::{
Result, TableReference, not_impl_err, plan_err,
tree_node::{TreeNode, TreeNodeRecursion},
};
-use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, TableSource};
-use sqlparser::ast::{Ident, Query, SetExpr, SetOperator, With};
+use datafusion_expr::{
+ Extension, LogicalPlan, LogicalPlanBuilder, MaterializedCte, MaterializedCteId,
+ MaterializedCteScan, TableSource,
+};
+use sqlparser::ast::{CteAsMaterialized, Ident, Query, SetExpr, SetOperator, With};
+
+/// A `WITH x AS MATERIALIZED (...)` CTE whose body must wrap the query that
+/// declared it. See [`MaterializedCte`].
+pub(crate) struct PendingMaterializedCte {
+ id: MaterializedCteId,
+ name: String,
+ cte: LogicalPlan,
+}
+
+impl PendingMaterializedCte {
+ /// Wrap `continuation` so that its scans of this CTE read one shared result.
+ pub(crate) fn wrap(self, continuation: LogicalPlan) -> LogicalPlan {
+ LogicalPlan::Extension(Extension {
+ node: Arc::new(MaterializedCte {
+ id: self.id,
+ name: self.name,
+ cte: self.cte,
+ continuation,
+ }),
+ })
+ }
+}
impl<S: ContextProvider> SqlToRel<'_, S> {
pub(super) fn plan_with_clause(
&self,
with: With,
planner_context: &mut PlannerContext,
- ) -> Result<()> {
+ ) -> Result<Vec<PendingMaterializedCte>> {
let is_recursive = with.recursive;
+ let mut materialized = vec![];
// Process CTEs from top to bottom
for cte in with.cte_tables {
// A `WITH` block can't use the same name more than once
@@ -55,6 +81,12 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
// Each `WITH` block can change the column names in the last projection
// (e.g. "WITH table(t1, t2) AS SELECT 1, 2"). Recursive CTEs apply those
// to the static term in recursive_cte(), so only the relation name here.
+ // A body that refers to an outer query would have to run once
+ // per outer row, which is not supported. Such a CTE is inlined,
+ // as if MATERIALIZED was not given.
+ let is_materialized = cte.materialized
+ == Some(CteAsMaterialized::Materialized)
+ && cte_plan.all_out_ref_exprs().is_empty();
let final_plan = if is_recursive {
LogicalPlanBuilder::from(cte_plan)
.alias(TableReference::bare(
@@ -64,10 +96,30 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
} else {
self.apply_table_alias(cte_plan, cte.alias)?
};
+ if is_materialized {
+ // Every reference reads the shared result instead of a copy
+ // of the body. The body is already planned, so a reference
+ // of a recursive CTE to itself stays a work table scan.
+ let id = MaterializedCteId::next();
+ let scan = LogicalPlan::Extension(Extension {
+ node: Arc::new(MaterializedCteScan {
+ id,
+ name: cte_name.clone(),
+ schema: Arc::clone(final_plan.schema()),
+ }),
+ });
+ materialized.push(PendingMaterializedCte {
+ id,
+ name: cte_name.clone(),
+ cte: final_plan,
+ });
+ planner_context.insert_cte(cte_name, scan);
+ continue;
+ }
// Export the CTE to the outer query
planner_context.insert_cte(cte_name, final_plan);
}
- Ok(())
+ Ok(materialized)
}
fn non_recursive_cte(
diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs
index f26411cd60..a943abe2f8 100644
--- a/datafusion/sql/src/query.rs
+++ b/datafusion/sql/src/query.rs
@@ -61,9 +61,10 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
return not_impl_err!("FETCH clause is not supported yet");
}
- if let Some(with) = with {
- self.plan_with_clause(with, planner_context)?;
- }
+ let materialized_ctes = match with {
+ Some(with) => self.plan_with_clause(with, planner_context)?,
+ None => vec![],
+ };
let set_expr = *body;
let plan = match set_expr {
@@ -96,7 +97,15 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
}
}?;
- self.pipe_operators(plan, pipe_operators, planner_context)
+ let plan = self.pipe_operators(plan, pipe_operators, planner_context)?;
+ if !materialized_ctes.is_empty() && matches!(plan, LogicalPlan::Ddl(_)) {
+ return not_impl_err!("MATERIALIZED CTEs are not supported with SELECT INTO");
+ }
+ // A later CTE may read an earlier one, so the first CTE is the outermost.
+ Ok(materialized_ctes
+ .into_iter()
+ .rev()
+ .fold(plan, |plan, cte| cte.wrap(plan)))
}
/// Apply pipe operators to a plan
diff --git a/datafusion/sqllogictest/test_files/cte_materialized.slt b/datafusion/sqllogictest/test_files/cte_materialized.slt
new file mode 100644
index 0000000000..74482b01da
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/cte_materialized.slt
@@ -0,0 +1,998 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Tests for `WITH x AS MATERIALIZED (...)`: the CTE body runs once and every
+# reference reads its buffered output.
+
+# sqlparser only parses `[NOT] MATERIALIZED` for the PostgreSQL dialect.
+statement ok
+set datafusion.sql_parser.dialect = 'PostgreSQL';
+
+statement ok
+set datafusion.explain.logical_plan_only = false;
+
+statement ok
+CREATE TABLE t(k INT, v INT) AS VALUES (1, 10), (2, 20), (3, 30), (4, 40);
+
+# A volatile body is evaluated once, so both references see the same value.
+query I
+WITH r AS MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+1
+
+# Without MATERIALIZED the body is inlined and evaluated once per reference.
+query I
+WITH r AS NOT MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+2
+
+query I
+WITH r AS (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+2
+
+# The table is scanned once, below MaterializedCteExec, and each reference is
+# a MaterializedCteScanExec. The filter of each reference stays above its scan.
+query TT
+EXPLAIN WITH c AS MATERIALIZED (SELECT k, v * 2 AS v2 FROM t)
+SELECT k FROM c WHERE v2 > 30
+UNION ALL
+SELECT k FROM c WHERE v2 < 30
+----
+logical_plan
+01)MaterializedCte: name=c
+02)--SubqueryAlias: c
+03)----Projection: t.k, CAST(t.v AS Int64) * Int64(2) AS v2
+04)------TableScan: t projection=[k, v]
+05)--Union
+06)----Projection: c.k
+07)------Filter: c.v2 > Int64(30)
+08)--------MaterializedCteScan: name=c
+09)----Projection: c.k
+10)------Filter: c.v2 < Int64(30)
+11)--------MaterializedCteScan: name=c
+physical_plan
+01)MaterializedCteExec: name=c
+02)--ProjectionExec: expr=[k@0 as k, CAST(v@1 AS Int64) * 2 as v2]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+04)--UnionExec
+05)----FilterExec: v2@1 > 30, projection=[k@0]
+06)------MaterializedCteScanExec: name=c
+07)----FilterExec: v2@1 < 30, projection=[k@0]
+08)------MaterializedCteScanExec: name=c
+
+query I rowsort
+WITH c AS MATERIALIZED (SELECT k, v * 2 AS v2 FROM t)
+SELECT k FROM c WHERE v2 > 30
+UNION ALL
+SELECT k FROM c WHERE v2 < 30
+----
+1
+2
+3
+4
+
+# EXPLAIN ANALYZE: the source is read once (4 rows) although the CTE is
+# referenced three times, and each scan replays all 4 rows.
+query TT
+EXPLAIN ANALYZE WITH c AS MATERIALIZED (SELECT k, v FROM t)
+SELECT count(*) FROM (
+ SELECT k FROM c UNION ALL SELECT k FROM c UNION ALL SELECT k FROM c
+)
+----
+Plan with Metrics
+01)MaterializedCteExec: name=c, metrics=[spill_count=0, <slt:ignore>buffered_rows=4]
+02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[<slt:ignore>]
+03)--ProjectionExec: expr=[count(Int64(1))@0 as count(*)], metrics=[<slt:ignore>]
+04)----AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))], metrics=[<slt:ignore>]
+05)------CoalescePartitionsExec, metrics=[<slt:ignore>]
+06)--------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))], metrics=[<slt:ignore>]
+07)----------UnionExec, metrics=[<slt:ignore>]
+08)------------ProjectionExec: expr=[], metrics=[<slt:ignore>]
+09)--------------MaterializedCteScanExec: name=c, metrics=[output_rows=4, <slt:ignore>]
+10)------------ProjectionExec: expr=[], metrics=[<slt:ignore>]
+11)--------------MaterializedCteScanExec: name=c, metrics=[output_rows=4, <slt:ignore>]
+12)------------ProjectionExec: expr=[], metrics=[<slt:ignore>]
+13)--------------MaterializedCteScanExec: name=c, metrics=[output_rows=4, <slt:ignore>]
+
+# For comparison, the inlined CTE scans the source once per reference.
+query TT
+EXPLAIN WITH c AS (SELECT k, v FROM t)
+SELECT sum(k) FROM (
+ SELECT k FROM c UNION ALL SELECT k FROM c UNION ALL SELECT k FROM c
+)
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[sum(CAST(k AS Int64))]]
+02)--Union
+03)----SubqueryAlias: c
+04)------TableScan: t projection=[k]
+05)----SubqueryAlias: c
+06)------TableScan: t projection=[k]
+07)----SubqueryAlias: c
+08)------TableScan: t projection=[k]
+physical_plan
+01)AggregateExec: mode=Final, gby=[], aggr=[sum(k)]
+02)--CoalescePartitionsExec
+03)----AggregateExec: mode=Partial, gby=[], aggr=[sum(k)]
+04)------UnionExec
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)--------DataSourceExec: partitions=1, partition_sizes=[1]
+07)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Both sides of one hash join read the same CTE. The build side must be read
+# completely before the probe side starts, which full materialization allows.
+query III rowsort
+WITH c AS MATERIALIZED (SELECT k, v FROM t)
+SELECT a.k, a.v, b.v FROM c a JOIN c b ON a.k = b.k
+----
+1 10 10
+2 20 20
+3 30 30
+4 40 40
+
+# A reference inside a scalar subquery. The scalar subquery runs before the
+# main plan, so the scan must be able to start the body itself.
+query II
+WITH c AS MATERIALIZED (SELECT k, v FROM t)
+SELECT k, v FROM c WHERE v = (SELECT max(v) FROM c)
+----
+4 40
+
+# A later CTE reads an earlier one.
+query II rowsort
+WITH a AS MATERIALIZED (SELECT k, v FROM t WHERE k > 1),
+ b AS MATERIALIZED (SELECT k, v + 1 AS v FROM a)
+SELECT a.k, b.v FROM a JOIN b ON a.k = b.k
+----
+2 21
+3 31
+4 41
+
+# Sibling subqueries declare CTEs with the same name.
+query II
+SELECT x.cnt, y.cnt FROM
+ (WITH s AS MATERIALIZED (SELECT k FROM t WHERE k <= 2)
+ SELECT count(*) AS cnt FROM s p JOIN s q ON p.k = q.k) x,
+ (WITH s AS MATERIALIZED (SELECT k FROM t WHERE k >= 2)
+ SELECT count(*) AS cnt FROM s p JOIN s q ON p.k = q.k) y
+----
+2 3
+
+# A column alias list on a materialized CTE.
+query I rowsort
+WITH c(a, b) AS MATERIALIZED (SELECT k, v FROM t)
+SELECT a FROM c WHERE b > 20
+----
+3
+4
+
+# IN and NOT IN subqueries that declare a materialized CTE.
+query I rowsort
+SELECT k FROM t WHERE k IN (
+ WITH r AS MATERIALIZED (SELECT k FROM t WHERE k > 2) SELECT k FROM r)
+----
+3
+4
+
+query I rowsort
+SELECT k FROM t WHERE k NOT IN (
+ WITH r AS MATERIALIZED (SELECT k FROM t WHERE k > 2) SELECT k FROM r)
+----
+1
+2
+
+# Correlated subqueries whose continuation refers to the outer query.
+query I rowsort
+SELECT k FROM t WHERE EXISTS (
+ WITH r AS MATERIALIZED (SELECT k AS rk FROM t)
+ SELECT 1 FROM r WHERE r.rk = t.k + 1)
+----
+1
+2
+3
+
+query I
+SELECT k FROM t WHERE NOT EXISTS (
+ WITH r AS MATERIALIZED (SELECT k AS rk FROM t)
+ SELECT 1 FROM r WHERE r.rk = t.k + 1)
+----
+4
+
+query I rowsort
+SELECT k FROM t WHERE v IN (
+ WITH r AS MATERIALIZED (SELECT k AS rk, v AS rv FROM t)
+ SELECT rv FROM r WHERE r.rk = t.k AND r.rk > 2)
+----
+3
+4
+
+# A body that refers to the outer query would have to run once per outer row.
+# Such a CTE is inlined, as if MATERIALIZED was not given.
+query I rowsort
+SELECT k FROM t WHERE EXISTS (
+ WITH r AS MATERIALIZED (SELECT k AS rk FROM t t2 WHERE t2.k = t.k + 1)
+ SELECT 1 FROM r)
+----
+1
+2
+3
+
+query I
+SELECT k FROM t WHERE NOT EXISTS (
+ WITH r AS MATERIALIZED (SELECT k AS rk FROM t t2 WHERE t2.k = t.k + 1)
+ SELECT 1 FROM r)
+----
+4
+
+query II rowsort
+SELECT k, (
+ WITH r AS MATERIALIZED (SELECT v FROM t t2 WHERE t2.k = t.k)
+ SELECT sum(v) * 2 FROM r)
+FROM t
+----
+1 20
+2 40
+3 60
+4 80
+
+# The correlated CTE is inlined, so the plan has no MaterializedCteExec.
+query TT
+EXPLAIN SELECT k FROM t WHERE EXISTS (
+ WITH r AS MATERIALIZED (SELECT k AS rk FROM t t2 WHERE t2.k = t.k + 1)
+ SELECT 1 FROM r)
+----
+logical_plan
+01)LeftSemi Join: CAST(t.k AS Int64) + Int64(1) = CAST(__correlated_sq_1.k AS Int64)
+02)--TableScan: t projection=[k]
+03)--SubqueryAlias: __correlated_sq_1
+04)----SubqueryAlias: r
+05)------SubqueryAlias: t2
+06)--------TableScan: t projection=[k]
+physical_plan
+01)HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(t.k + Int64(1)@1, CAST(__correlated_sq_1.k AS Int64)@1)], projection=[k@0]
+02)--ProjectionExec: expr=[k@0 as k, CAST(k@0 AS Int64) + 1 as t.k + Int64(1)]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+04)--ProjectionExec: expr=[k@0 as k, CAST(k@0 AS Int64) as CAST(__correlated_sq_1.k AS Int64)]
+05)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+# In a WITH RECURSIVE block, MATERIALIZED applies to a CTE that does not refer
+# to itself.
+query I
+WITH RECURSIVE r AS MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+1
+
+# A recursive CTE is also materialized: RecursiveQueryExec runs once, and its
+# reference to itself is still a work table.
+query TT
+EXPLAIN WITH RECURSIVE r AS MATERIALIZED (
+ SELECT 1 AS n UNION ALL SELECT n + 1 FROM r WHERE n < 3
+)
+SELECT count(*) FROM r a, r b
+----
+logical_plan
+01)MaterializedCte: name=r
+02)--SubqueryAlias: r
+03)----RecursiveQuery: is_distinct=false
+04)------Projection: Int64(1) AS n
+05)--------EmptyRelation: rows=1
+06)------Projection: r.n + Int64(1)
+07)--------Filter: r.n < Int64(3)
+08)----------TableScan: r projection=[n]
+09)--Projection: count(Int64(1)) AS count(*)
+10)----Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]
+11)------Cross Join:
+12)--------SubqueryAlias: a
+13)----------Projection:
+14)------------MaterializedCteScan: name=r
+15)--------SubqueryAlias: b
+16)----------Projection:
+17)------------MaterializedCteScan: name=r
+physical_plan
+01)MaterializedCteExec: name=r
+02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+03)----RecursiveQueryExec: name=r, is_distinct=false
+04)------ProjectionExec: expr=[CAST(1 AS Int64) as n]
+05)--------PlaceholderRowExec
+06)------CoalescePartitionsExec
+07)--------ProjectionExec: expr=[n@0 + 1 as n]
+08)----------FilterExec: n@0 < 3
+09)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+10)--------------WorkTableExec: name=r
+11)--ProjectionExec: expr=[count(Int64(1))@0 as count(*)]
+12)----AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]
+13)------CoalescePartitionsExec
+14)--------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
+15)----------CrossJoinExec
+16)------------CoalescePartitionsExec
+17)--------------ProjectionExec: expr=[]
+18)----------------MaterializedCteScanExec: name=r
+19)------------ProjectionExec: expr=[]
+20)--------------MaterializedCteScanExec: name=r
+
+query I
+WITH RECURSIVE r AS MATERIALIZED (
+ SELECT 1 AS n UNION ALL SELECT n + 1 FROM r WHERE n < 3
+)
+SELECT count(*) FROM r a, r b
+----
+9
+
+# A recursive CTE reads a materialized CTE of the same block.
+query I rowsort
+WITH RECURSIVE m AS MATERIALIZED (SELECT 3 AS lim),
+ r AS (SELECT 1 AS n UNION ALL SELECT n + 1 FROM r, m WHERE n < m.lim)
+SELECT n FROM r
+----
+1
+2
+3
+
+# Each reference to a view is a separate query, so each has its own
+# materialized CTE, as in PostgreSQL.
+statement ok
+CREATE VIEW vv AS WITH r AS MATERIALIZED (SELECT random() AS x) SELECT x FROM r;
+
+query B
+SELECT a.x = b.x FROM vv a, vv b
+----
+false
+
+query I
+SELECT count(DISTINCT x) FROM (
+ SELECT x FROM vv UNION ALL SELECT x FROM vv UNION ALL SELECT (SELECT max(x) FROM vv)
+)
+----
+3
+
+# Inside one reference to a view, the CTE is still shared.
+statement ok
+CREATE VIEW vv2 AS WITH r AS MATERIALIZED (SELECT random() AS x)
+SELECT a.x = b.x AS same FROM r a, r b;
+
+query B
+SELECT * FROM vv2 UNION ALL SELECT * FROM vv2
+----
+true
+true
+
+statement ok
+DROP VIEW vv;
+
+statement ok
+DROP VIEW vv2;
+
+# A filter above the query that declares the CTE goes into the continuation,
+# not into the shared body.
+query TT
+EXPLAIN SELECT * FROM (
+ WITH r AS MATERIALIZED (SELECT k FROM t)
+ SELECT t.k, t.v FROM t JOIN r ON t.k = r.k
+) s WHERE s.v = 10
+----
+logical_plan
+01)SubqueryAlias: s
+02)--MaterializedCte: name=r
+03)----SubqueryAlias: r
+04)------TableScan: t projection=[k]
+05)----Projection: t.k, t.v
+06)------Inner Join: t.k = r.k
+07)--------Filter: t.v = Int32(10)
+08)----------TableScan: t projection=[k, v]
+09)--------MaterializedCteScan: name=r
+physical_plan
+01)MaterializedCteExec: name=r
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+03)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)], projection=[k@0, v@1]
+04)----FilterExec: v@1 = 10
+05)------DataSourceExec: partitions=1, partition_sizes=[1]
+06)----MaterializedCteScanExec: name=r
+
+query II
+SELECT * FROM (
+ WITH r AS MATERIALIZED (SELECT k FROM t)
+ SELECT t.k, t.v FROM t JOIN r ON t.k = r.k
+) s WHERE s.v = 10
+----
+1 10
+
+# Spill: the body does not fit in the memory limit, so it is written to disk
+# and replayed from there by every reference. Part of the in-memory prefix is
+# also moved to disk, so that the two references have memory to decode the
+# spill files (spill_count=2).
+statement ok
+SET datafusion.execution.target_partitions = 1
+
+statement ok
+SET datafusion.runtime.memory_limit = '300K'
+
+query II
+WITH g AS MATERIALIZED (SELECT value AS x FROM generate_series(1, 200000))
+SELECT (SELECT count(*) FROM g), (SELECT sum(x) FROM g)
+----
+200000 20000100000
+
+query TT
+EXPLAIN ANALYZE WITH g AS MATERIALIZED (SELECT value AS x FROM generate_series(1, 200000))
+SELECT count(*) FROM (SELECT x FROM g UNION ALL SELECT x FROM g)
+----
+Plan with Metrics
+01)MaterializedCteExec: name=g, metrics=[spill_count=2, <slt:ignore>buffered_rows=200.0 K]
+02)--ProjectionExec: expr=[value@0 as x], metrics=[output_rows=200.0 K, <slt:ignore>]
+03)----LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=200000, batch_size=8192], metrics=[output_rows=200.0 K, <slt:ignore>]
+04)--ProjectionExec: expr=[count(Int64(1))@0 as count(*)], metrics=[<slt:ignore>]
+05)----AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))], metrics=[<slt:ignore>]
+06)------CoalescePartitionsExec, metrics=[<slt:ignore>]
+07)--------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))], metrics=[<slt:ignore>]
+08)----------UnionExec, metrics=[<slt:ignore>]
+09)------------ProjectionExec: expr=[], metrics=[<slt:ignore>]
+10)--------------MaterializedCteScanExec: name=g, metrics=[output_rows=200.0 K, <slt:ignore>]
+11)------------ProjectionExec: expr=[], metrics=[<slt:ignore>]
+12)--------------MaterializedCteScanExec: name=g, metrics=[output_rows=200.0 K, <slt:ignore>]
+
+# Spill: the first batch has long strings and stays in memory. The next
+# batches are smaller and go to disk. The two references then need memory to
+# decode the spill files, and the first batch must move to disk. It is
+# written in pieces, because each reference reserves memory for the largest
+# batch on disk.
+statement ok
+SET datafusion.runtime.memory_limit = '18M'
+
+query I
+WITH c AS MATERIALIZED (
+ SELECT repeat('x', CASE WHEN value <= 8192 THEN 2000 ELSE 375 END) AS s
+ FROM generate_series(1, 50000)
+)
+SELECT count(*) FROM (SELECT s FROM c UNION ALL SELECT s FROM c)
+----
+100000
+
+statement ok
+RESET datafusion.runtime.memory_limit
+
+# Known limitation: a reference to a materialized CTE has unknown statistics.
+# The join below builds the hash table on the 100000 rows of `huge`. When the
+# CTE is inlined, the hash table is built on the 5 rows of `small`.
+statement ok
+CREATE TABLE small AS SELECT value AS k FROM generate_series(1, 5)
+
+statement ok
+CREATE TABLE huge AS SELECT value AS k FROM generate_series(1, 100000)
+
+query TT
+EXPLAIN WITH c AS MATERIALIZED (SELECT k FROM small)
+SELECT count(*) FROM huge JOIN c ON huge.k = c.k
+----
+logical_plan
+01)MaterializedCte: name=c
+02)--SubqueryAlias: c
+03)----TableScan: small projection=[k]
+04)--Projection: count(Int64(1)) AS count(*)
+05)----Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]
+06)------Projection:
+07)--------Inner Join: huge.k = c.k
+08)----------TableScan: huge projection=[k]
+09)----------MaterializedCteScan: name=c
+physical_plan
+01)MaterializedCteExec: name=c
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+03)--ProjectionExec: expr=[count(Int64(1))@0 as count(*)]
+04)----AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1))]
+05)------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)], projection=[]
+06)--------DataSourceExec: partitions=1, partition_sizes=[13]
+07)--------MaterializedCteScanExec: name=c
+
+query TT
+EXPLAIN WITH c AS NOT MATERIALIZED (SELECT k FROM small)
+SELECT count(*) FROM huge JOIN c ON huge.k = c.k
+----
+logical_plan
+01)Projection: count(Int64(1)) AS count(*)
+02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]
+03)----Projection:
+04)------Inner Join: huge.k = c.k
+05)--------TableScan: huge projection=[k]
+06)--------SubqueryAlias: c
+07)----------TableScan: small projection=[k]
+physical_plan
+01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)]
+02)--AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1))]
+03)----HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)], projection=[]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+05)------DataSourceExec: partitions=1, partition_sizes=[13]
+
+statement ok
+DROP TABLE small
+
+statement ok
+DROP TABLE huge
+
+# Setting a runtime option also sets this option to false.
+statement ok
+RESET datafusion.catalog.create_default_catalog_and_schema
+
+# Restore the slt runner default.
+statement ok
+SET datafusion.execution.target_partitions = 4
+
+# Type coercion changes the type of `a` in the body from Null to Int64, and
+# the scans get the new type.
+query II
+WITH c AS MATERIALIZED (SELECT NULL a, 12 b UNION ALL SELECT 11 a, 13 b)
+SELECT * FROM c ORDER BY 1, 2
+----
+11 13
+NULL 12
+
+query II
+WITH c AS MATERIALIZED (SELECT NULL a, 12 b UNION ALL SELECT 11 a, 13 b)
+SELECT x.a, x.b FROM c x WHERE EXISTS (SELECT 1 FROM c y WHERE y.a = x.a)
+----
+11 13
+
+# An unbounded body is rejected: it is fully buffered before any row is
+# returned, so the query would never produce output.
+statement ok
+CREATE UNBOUNDED EXTERNAL TABLE unbounded_t(c1 integer NOT NULL)
+STORED AS CSV
+LOCATION '../core/tests/data/empty.csv'
+OPTIONS ('format.has_header' 'true');
+
+query error DataFusion error: Error during planning: MATERIALIZED CTE u reads an unbounded source, which cannot be materialized
+WITH u AS MATERIALIZED (SELECT c1 FROM unbounded_t)
+SELECT count(*) FROM u
+
+# Without MATERIALIZED the CTE is inlined as before.
+statement ok
+EXPLAIN WITH u AS (SELECT c1 FROM unbounded_t) SELECT c1 FROM u
+
+# A MATERIALIZED CTE that nothing reads is not evaluated, so an unbounded body
+# is not an error.
+query I
+WITH u AS MATERIALIZED (SELECT c1 FROM unbounded_t) SELECT 1
+----
+1
+
+statement ok
+DROP TABLE unbounded_t;
+
+# Behaviour matrix. Each case runs with MATERIALIZED, NOT MATERIALIZED and no
+# keyword, and pins the current result. For a volatile body, count(DISTINCT x)
+# over two references is 1 when the references share one evaluation and 2 when
+# each reference evaluates the body.
+#
+# The PostgreSQL and DuckDB results are from PostgreSQL 17.11 and DuckDB 1.5.2,
+# with the same table t. Results are listed as sorted rows.
+# "(differs)" marks a result that is not the same as DataFusion. The
+# differences have two causes:
+# - PostgreSQL does not inline a CTE with a volatile body, also with NOT
+# MATERIALIZED, and it does not evaluate an uncorrelated subquery in a
+# recursive term again in each iteration.
+# - DuckDB materializes a CTE with no keyword when it is referenced more than
+# one time.
+# DataFusion inlines every CTE that does not have MATERIALIZED.
+#
+# Two references in the same query.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH m AS MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM m UNION ALL SELECT x FROM m)
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 2.
+query I rowsort
+WITH m AS NOT MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM m UNION ALL SELECT x FROM m)
+----
+2
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 1 (differs).
+query I rowsort
+WITH m AS (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM m UNION ALL SELECT x FROM m)
+----
+2
+
+# One reference in a scalar subquery, one in the main query.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH m AS MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT (SELECT x FROM m) AS x UNION ALL SELECT x FROM m)
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 2.
+query I rowsort
+WITH m AS NOT MATERIALIZED (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT (SELECT x FROM m) AS x UNION ALL SELECT x FROM m)
+----
+2
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 1 (differs).
+query I rowsort
+WITH m AS (SELECT random() AS x)
+SELECT count(DISTINCT x) FROM (SELECT (SELECT x FROM m) AS x UNION ALL SELECT x FROM m)
+----
+2
+
+# One reference in an IN subquery, one in the main query.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH m AS MATERIALIZED (SELECT random() AS x)
+SELECT count(*) FROM m WHERE x IN (SELECT x FROM m)
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 0.
+query I rowsort
+WITH m AS NOT MATERIALIZED (SELECT random() AS x)
+SELECT count(*) FROM m WHERE x IN (SELECT x FROM m)
+----
+0
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 1 (differs).
+query I rowsort
+WITH m AS (SELECT random() AS x)
+SELECT count(*) FROM m WHERE x IN (SELECT x FROM m)
+----
+0
+
+# A volatile scalar subquery inside the body. An uncorrelated scalar subquery
+# is evaluated once per query, so the result is 1 in every mode.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH m AS MATERIALIZED (SELECT (SELECT random()) AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM m UNION ALL SELECT x FROM m)
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1. DuckDB: 2 (differs).
+query I rowsort
+WITH m AS NOT MATERIALIZED (SELECT (SELECT random()) AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM m UNION ALL SELECT x FROM m)
+----
+1
+
+# no keyword
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH m AS (SELECT (SELECT random()) AS x)
+SELECT count(DISTINCT x) FROM (SELECT x FROM m UNION ALL SELECT x FROM m)
+----
+1
+
+# A CTE whose body declares another CTE.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH a AS MATERIALIZED (
+ WITH b AS MATERIALIZED (SELECT random() AS x) SELECT x FROM b UNION ALL SELECT x FROM b)
+SELECT count(DISTINCT x) FROM (SELECT x FROM a UNION ALL SELECT x FROM a)
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 4.
+query I rowsort
+WITH a AS NOT MATERIALIZED (
+ WITH b AS NOT MATERIALIZED (SELECT random() AS x) SELECT x FROM b UNION ALL SELECT x FROM b)
+SELECT count(DISTINCT x) FROM (SELECT x FROM a UNION ALL SELECT x FROM a)
+----
+4
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 1 (differs).
+query I rowsort
+WITH a AS (
+ WITH b AS (SELECT random() AS x) SELECT x FROM b UNION ALL SELECT x FROM b)
+SELECT count(DISTINCT x) FROM (SELECT x FROM a UNION ALL SELECT x FROM a)
+----
+4
+
+# Uncorrelated EXISTS, IN and scalar subqueries inside the body.
+# MATERIALIZED
+# PostgreSQL: 2 3. DuckDB: 2 3.
+query I rowsort
+WITH m AS MATERIALIZED (
+ SELECT k FROM t
+ WHERE EXISTS (SELECT 1 FROM t t2 WHERE t2.k > 3)
+ AND k IN (SELECT k FROM t WHERE v > 10)
+ AND v < (SELECT max(v) FROM t))
+SELECT k FROM m
+----
+2
+3
+
+# NOT MATERIALIZED
+# PostgreSQL: 2 3. DuckDB: 2 3.
+query I rowsort
+WITH m AS NOT MATERIALIZED (
+ SELECT k FROM t
+ WHERE EXISTS (SELECT 1 FROM t t2 WHERE t2.k > 3)
+ AND k IN (SELECT k FROM t WHERE v > 10)
+ AND v < (SELECT max(v) FROM t))
+SELECT k FROM m
+----
+2
+3
+
+# no keyword
+# PostgreSQL: 2 3. DuckDB: 2 3.
+query I rowsort
+WITH m AS (
+ SELECT k FROM t
+ WHERE EXISTS (SELECT 1 FROM t t2 WHERE t2.k > 3)
+ AND k IN (SELECT k FROM t WHERE v > 10)
+ AND v < (SELECT max(v) FROM t))
+SELECT k FROM m
+----
+2
+3
+
+# A correlated body in an EXISTS subquery (the MATERIALIZED case is above).
+# NOT MATERIALIZED
+# PostgreSQL: 1 2 3. DuckDB: 1 2 3.
+query I rowsort
+SELECT k FROM t WHERE EXISTS (
+ WITH r AS NOT MATERIALIZED (SELECT k AS rk FROM t t2 WHERE t2.k = t.k + 1)
+ SELECT 1 FROM r)
+----
+1
+2
+3
+
+# no keyword
+# PostgreSQL: 1 2 3. DuckDB: 1 2 3.
+query I rowsort
+SELECT k FROM t WHERE EXISTS (
+ WITH r AS (SELECT k AS rk FROM t t2 WHERE t2.k = t.k + 1)
+ SELECT 1 FROM r)
+----
+1
+2
+3
+
+# Recursive queries.
+#
+# A volatile CTE declared outside a recursive CTE and read by its recursive
+# term. Materialized, it is evaluated once for the whole query.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH RECURSIVE
+ m AS MATERIALIZED (SELECT random() AS x),
+ r AS (
+ SELECT 1 AS n, x FROM m
+ UNION ALL
+ SELECT n + 1, m.x FROM r, m WHERE n < 4)
+SELECT count(DISTINCT x) FROM r
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 4.
+query I rowsort
+WITH RECURSIVE
+ m AS NOT MATERIALIZED (SELECT random() AS x),
+ r AS (
+ SELECT 1 AS n, x FROM m
+ UNION ALL
+ SELECT n + 1, m.x FROM r, m WHERE n < 4)
+SELECT count(DISTINCT x) FROM r
+----
+4
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 1 (differs).
+query I rowsort
+WITH RECURSIVE
+ m AS (SELECT random() AS x),
+ r AS (
+ SELECT 1 AS n, x FROM m
+ UNION ALL
+ SELECT n + 1, m.x FROM r, m WHERE n < 4)
+SELECT count(DISTINCT x) FROM r
+----
+4
+
+# A volatile CTE declared inside the recursive term. The recursive term is a
+# new query in each iteration, so the CTE is evaluated again in each iteration.
+# MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 3.
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n, CAST(0 AS DOUBLE PRECISION) AS x
+ UNION ALL
+ SELECT n + 1, s.x FROM r, (WITH m AS MATERIALIZED (SELECT random() AS x) SELECT x FROM m) s
+ WHERE n < 4)
+SELECT count(DISTINCT x) FROM r WHERE n > 1
+----
+3
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 3.
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n, CAST(0 AS DOUBLE PRECISION) AS x
+ UNION ALL
+ SELECT n + 1, s.x FROM r, (WITH m AS NOT MATERIALIZED (SELECT random() AS x) SELECT x FROM m) s
+ WHERE n < 4)
+SELECT count(DISTINCT x) FROM r WHERE n > 1
+----
+3
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 3.
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n, CAST(0 AS DOUBLE PRECISION) AS x
+ UNION ALL
+ SELECT n + 1, s.x FROM r, (WITH m AS (SELECT random() AS x) SELECT x FROM m) s
+ WHERE n < 4)
+SELECT count(DISTINCT x) FROM r WHERE n > 1
+----
+3
+
+# Inside one iteration, two references to a CTE declared in the recursive term.
+# MATERIALIZED
+# PostgreSQL: 1 true 2 true 3 true 4 true. DuckDB: 1 true 2 true 3 true 4 true.
+query IB rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n, true AS same
+ UNION ALL
+ SELECT n + 1, s.same FROM r,
+ (WITH m AS MATERIALIZED (SELECT random() AS x) SELECT a.x = b.x AS same FROM m a, m b) s
+ WHERE n < 4)
+SELECT n, same FROM r
+----
+1 true
+2 true
+3 true
+4 true
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 true 2 true 3 true 4 true (differs). DuckDB: 1 true 2 false 3 false 4 false.
+query IB rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n, true AS same
+ UNION ALL
+ SELECT n + 1, s.same FROM r,
+ (WITH m AS NOT MATERIALIZED (SELECT random() AS x) SELECT a.x = b.x AS same FROM m a, m b) s
+ WHERE n < 4)
+SELECT n, same FROM r
+----
+1 true
+2 false
+3 false
+4 false
+
+# no keyword
+# PostgreSQL: 1 true 2 true 3 true 4 true (differs). DuckDB: 1 true 2 true 3 true 4 true (differs).
+query IB rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n, true AS same
+ UNION ALL
+ SELECT n + 1, s.same FROM r,
+ (WITH m AS (SELECT random() AS x) SELECT a.x = b.x AS same FROM m a, m b) s
+ WHERE n < 4)
+SELECT n, same FROM r
+----
+1 true
+2 false
+3 false
+4 false
+
+# A CTE in the recursive term whose body reads the recursive CTE. The body
+# must be evaluated again in each iteration, or the query does not terminate.
+# MATERIALIZED
+# PostgreSQL: 1 2 3 4. DuckDB: 1 2 3 4.
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n
+ UNION ALL
+ SELECT n + 1 FROM (WITH m AS MATERIALIZED (SELECT n FROM r) SELECT n FROM m) WHERE n < 4)
+SELECT n FROM r
+----
+1
+2
+3
+4
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 2 3 4. DuckDB: 1 2 3 4.
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n
+ UNION ALL
+ SELECT n + 1 FROM (WITH m AS NOT MATERIALIZED (SELECT n FROM r) SELECT n FROM m) WHERE n < 4)
+SELECT n FROM r
+----
+1
+2
+3
+4
+
+# no keyword
+# PostgreSQL: 1 2 3 4. DuckDB: 1 2 3 4.
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT 1 AS n
+ UNION ALL
+ SELECT n + 1 FROM (WITH m AS (SELECT n FROM r) SELECT n FROM m) WHERE n < 4)
+SELECT n FROM r
+----
+1
+2
+3
+4
+
+# A recursive CTE referenced two times, with a volatile static term.
+# MATERIALIZED
+# PostgreSQL: 1. DuckDB: 1.
+query I rowsort
+WITH RECURSIVE r AS MATERIALIZED (
+ SELECT random() AS x, 1 AS n
+ UNION ALL
+ SELECT x, n + 1 FROM r WHERE n < 3)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+1
+
+# NOT MATERIALIZED
+# PostgreSQL: 1 (differs). DuckDB: 2.
+query I rowsort
+WITH RECURSIVE r AS NOT MATERIALIZED (
+ SELECT random() AS x, 1 AS n
+ UNION ALL
+ SELECT x, n + 1 FROM r WHERE n < 3)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+2
+
+# no keyword
+# PostgreSQL: 1 (differs). DuckDB: 1 (differs).
+query I rowsort
+WITH RECURSIVE r AS (
+ SELECT random() AS x, 1 AS n
+ UNION ALL
+ SELECT x, n + 1 FROM r WHERE n < 3)
+SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r)
+----
+2
+
+statement ok
+DROP TABLE t;
+
+statement ok
+set datafusion.sql_parser.dialect = generic;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org