adriangb commented on code in PR #25580:
URL: https://github.com/apache/datafusion/pull/25580#discussion_r4066954557
##########
datafusion/sql/src/query.rs:
##########
@@ -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)))
Review Comment:
When a `MATERIALIZED` CTE is declared inside an `IN` or `NOT IN` subquery,
planning fails. Before this PR the directive was ignored, so the same query
worked:
```sql
set datafusion.sql_parser.dialect = 'PostgreSQL';
CREATE TABLE t(k INT, v INT) AS VALUES (1, 10), (2, 20), (3, 30), (4, 40);
SELECT k FROM t WHERE k IN (WITH r AS MATERIALIZED (SELECT k FROM t WHERE k
> 2) SELECT k FROM r);
-- Optimizer rule 'decorrelate_predicate_subquery' failed
-- caused by
-- Error during planning: single expression required.
```
`NOT IN` gives the same error. With `NOT MATERIALIZED` the query returns `3,
4`.
The cause is `LogicalPlan::head_output_expr`, which returns `None` for an
`Extension` node. I added an arm that forwards to
`continuation.head_output_expr()`, and then `IN` and `NOT IN` returned the
correct rows.
A correlated `EXISTS` also fails, because the decorrelation does not pull
the correlated filter up through the extension node:
```sql
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);
-- This feature is not implemented: Physical plan does not support logical
expression Exists(...)
```
With `NOT MATERIALIZED` it returns `1, 2, 3`. Uncorrelated `EXISTS` and
uncorrelated scalar subqueries work.
I think this is related to open question 1: the rules that match on
`LogicalPlan` variants treat an `Extension` as opaque, so each of them needs a
special case for `MaterializedCte`. Could you add these queries to
`cte_materialized.slt`, and make them work or fail with a clear `not_impl_err`?
##########
datafusion/sql/src/cte.rs:
##########
@@ -64,10 +92,29 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
} else {
self.apply_table_alias(cte_plan, cte.alias)?
};
+ if is_materialized && !is_recursive {
Review Comment:
In a `WITH RECURSIVE` block, `is_recursive` is true for every CTE, so this
condition ignores `MATERIALIZED` on all CTEs of the block, also on CTEs that do
not refer to themselves. There is no error:
```sql
WITH RECURSIVE r AS MATERIALIZED (SELECT random() AS x)
SELECT count(DISTINCT x) FROM (SELECT x FROM r UNION ALL SELECT x FROM r);
-- 2, expected 1
```
I removed `&& !is_recursive` locally. The scan goes into the planner context
only after the body is planned, so a self-reference in the body is still
planned as `WorkTableExec`. With this change:
- the query above returns `1`;
- a recursive CTE that is referenced two times (`SELECT count(*) FROM r a, r
b` over `SELECT 1 AS n UNION ALL SELECT n + 1 FROM r WHERE n < 3`) runs
`RecursiveQueryExec` one time under `MaterializedCteExec` and returns `9`;
- a materialized CTE that a recursive CTE of the same block reads also
returns the correct rows.
```suggestion
if is_materialized {
```
##########
datafusion/physical-plan/src/materialized_cte.rs:
##########
@@ -0,0 +1,534 @@
+// 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 file, so row order within a
+//! partition is kept.
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+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::{StreamExt, TryStreamExt};
+use parking_lot::Mutex;
+
+use crate::execution_plan::{
+ Boundedness, CardinalityEffect, EmissionType, ExecutionPlan,
ExecutionPlanProperties,
+ PlanProperties, SchedulingType,
+};
+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
+/// an optional spill file with the remaining batches.
+struct BufferedPartition {
+ batches: Vec<RecordBatch>,
+ spill_file: Option<Arc<dyn SpillFile>>,
+ /// Released when the buffer is dropped.
+ _reservation: MemoryReservation,
+}
+
+/// The materialized output of a CTE body.
+struct MaterializedOutput {
+ partitions: Vec<BufferedPartition>,
+ spill_manager: SpillManager,
+}
+
+/// 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>>>,
+ 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()),
+ 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 metrics = self.metrics.clone();
+ once.try_once(move || Ok(buffer_body(name, body, context, metrics)))
+ }
+}
+
+async fn buffer_body(
+ name: String,
+ body: Arc<dyn ExecutionPlan>,
+ 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);
+ Ok(MaterializedOutput {
+ partitions: partitions.into_iter().map(|(_, p)| p).collect(),
+ spill_manager,
+ })
+}
+
+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;
+ 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;
+ }
+ spill
+
.get_or_insert(spill_manager.create_in_progress_file("MaterializedCte")?)
+ .append_batch(&batch)?;
Review Comment:
`Option::get_or_insert` evaluates its argument before the call, also when
`spill` is already `Some`. Thus each spilled batch creates a new temporary file
through the `DiskManager`, and then drops it. In the 200k-row slt test that is
approximately 20 unnecessary files. Also, if the `DiskManager` cannot create a
file, a batch that goes to an existing file fails.
```suggestion
if spill.is_none() {
spill =
Some(spill_manager.create_in_progress_file("MaterializedCte")?);
}
if let Some(file) = spill.as_mut() {
file.append_batch(&batch)?;
}
```
##########
datafusion/physical-plan/src/materialized_cte.rs:
##########
@@ -0,0 +1,534 @@
+// 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 file, so row order within a
+//! partition is kept.
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+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::{StreamExt, TryStreamExt};
+use parking_lot::Mutex;
+
+use crate::execution_plan::{
+ Boundedness, CardinalityEffect, EmissionType, ExecutionPlan,
ExecutionPlanProperties,
+ PlanProperties, SchedulingType,
+};
+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
+/// an optional spill file with the remaining batches.
+struct BufferedPartition {
+ batches: Vec<RecordBatch>,
+ spill_file: Option<Arc<dyn SpillFile>>,
+ /// Released when the buffer is dropped.
+ _reservation: MemoryReservation,
+}
+
+/// The materialized output of a CTE body.
+struct MaterializedOutput {
+ partitions: Vec<BufferedPartition>,
+ spill_manager: SpillManager,
+}
+
+/// 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>>>,
+ 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()),
+ 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 metrics = self.metrics.clone();
+ once.try_once(move || Ok(buffer_body(name, body, context, metrics)))
+ }
+}
+
+async fn buffer_body(
+ name: String,
+ body: Arc<dyn ExecutionPlan>,
+ 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);
+ Ok(MaterializedOutput {
+ partitions: partitions.into_iter().map(|(_, p)| p).collect(),
+ spill_manager,
+ })
+}
+
+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;
+ 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;
+ }
+ spill
+
.get_or_insert(spill_manager.create_in_progress_file("MaterializedCte")?)
+ .append_batch(&batch)?;
+ }
+ let spill_file = match spill {
+ Some(mut file) => file.finish()?,
+ None => None,
+ };
+ Ok(BufferedPartition {
+ batches,
+ spill_file,
+ _reservation: reservation,
+ })
+}
+
+/// Stream the buffered partitions `partition, partition + n, ...` of `output`,
+/// where `n` is the number of output partitions of the scan.
+fn replay(
+ output: &MaterializedOutput,
+ partition: usize,
+ output_partitions: usize,
+) -> Result<Vec<SendableRecordBatchStream>> {
+ 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);
+ if let Some(file) = &buffered.spill_file {
+ streams.push(
+ output
+ .spill_manager
+ .read_spill_as_stream(Arc::clone(file), None)?,
+ );
+ }
+ }
+ Ok(streams)
+}
+
+/// 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),
+ )
+ }
+
+ 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
+ }
+}
+
+/// 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 {
+ Self {
+ id: self.id,
+ name: self.name.clone(),
+ schema: Arc::clone(&self.schema),
+ buffer: Some(buffer),
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache: Arc::clone(&self.cache),
+ }
+ }
+
+ pub fn is_bound(&self) -> bool {
+ self.buffer.is_some()
+ }
+
+ /// The id of the [`MaterializedCteBuffer`] this scan reads.
+ pub fn id(&self) -> u64 {
+ self.id
+ }
+}
+
+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?;
+ let streams = replay(&output, partition, output_partitions)?;
+ Ok::<_, datafusion_common::DataFusionError>(
+ futures::stream::iter(streams).flatten(),
+ )
+ })
+ .try_flatten()
+ .inspect_ok(move |batch| baseline.record_output(batch.num_rows()));
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&self.schema),
+ stream,
+ )))
Review Comment:
The scan declares `SchedulingType::Cooperative`, so `EnsureCooperative` does
not add a `CooperativeExec` above it. But the replay of the in-memory batches
is `futures::stream::iter`, which never returns `Pending` and does not use the
Tokio budget. `WorkTableExec` and `MemorySourceConfig` declare the same
scheduling type and wrap their streams in `cooperative(...)`. Without the
wrapper, an operator that drains the scan in one poll can keep the worker
thread for the full replay, and cancellation must wait for it.
`cooperative` requires `Unpin`, so the inner stream must be boxed. This
compiles with `use crate::coop::cooperative;`:
```suggestion
Ok(Box::pin(cooperative(RecordBatchStreamAdapter::new(
Arc::clone(&self.schema),
Box::pin(stream),
))))
```
##########
datafusion/core/src/physical_planner.rs:
##########
@@ -216,11 +219,45 @@ impl DefaultPhysicalPlanner {
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, |_, _| {})
}
}
+/// Point every [`MaterializedCteScanExec`] at the buffer of the
+/// [`MaterializedCteExec`] with the same id.
+///
+/// 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.
+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()));
Review Comment:
The id is copied when a logical plan is cloned. When the same logical plan
occurs two times in one query, for example a view that is referenced two times
or a DataFrame that is joined to itself, two `MaterializedCteExec` nodes have
the same id. `insert` keeps the last one. The scans of both copies bind to that
buffer, and the body of the other copy never runs:
```sql
CREATE VIEW vv AS WITH r AS MATERIALIZED (SELECT random() AS x) SELECT x
FROM r;
SELECT a.x = b.x FROM vv a, vv b;
-- true
```
In PostgreSQL each reference to a view is a separate subquery with its own
CTE, so the result is `false`.
Could the binding be scoped? For example, walk the plan top-down and bind
the scans in the continuation of each `MaterializedCteExec` to the buffer of
that node. A scan in a scalar subquery is in the same subtree, because
`ScalarSubqueryExec` is inside the continuation. If that is too much for this
PR, an internal error for a duplicate id is better than a silent change of
results.
##########
datafusion/physical-plan/src/materialized_cte.rs:
##########
@@ -0,0 +1,534 @@
+// 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 file, so row order within a
+//! partition is kept.
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+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::{StreamExt, TryStreamExt};
+use parking_lot::Mutex;
+
+use crate::execution_plan::{
+ Boundedness, CardinalityEffect, EmissionType, ExecutionPlan,
ExecutionPlanProperties,
+ PlanProperties, SchedulingType,
+};
+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
+/// an optional spill file with the remaining batches.
+struct BufferedPartition {
+ batches: Vec<RecordBatch>,
+ spill_file: Option<Arc<dyn SpillFile>>,
+ /// Released when the buffer is dropped.
+ _reservation: MemoryReservation,
+}
+
+/// The materialized output of a CTE body.
+struct MaterializedOutput {
+ partitions: Vec<BufferedPartition>,
+ spill_manager: SpillManager,
+}
+
+/// 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>>>,
+ 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()),
+ 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 metrics = self.metrics.clone();
+ once.try_once(move || Ok(buffer_body(name, body, context, metrics)))
+ }
+}
+
+async fn buffer_body(
+ name: String,
+ body: Arc<dyn ExecutionPlan>,
+ 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);
+ Ok(MaterializedOutput {
+ partitions: partitions.into_iter().map(|(_, p)| p).collect(),
+ spill_manager,
+ })
+}
+
+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;
+ 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;
+ }
+ spill
+
.get_or_insert(spill_manager.create_in_progress_file("MaterializedCte")?)
+ .append_batch(&batch)?;
+ }
+ let spill_file = match spill {
+ Some(mut file) => file.finish()?,
+ None => None,
+ };
+ Ok(BufferedPartition {
+ batches,
+ spill_file,
+ _reservation: reservation,
+ })
+}
+
+/// Stream the buffered partitions `partition, partition + n, ...` of `output`,
+/// where `n` is the number of output partitions of the scan.
+fn replay(
+ output: &MaterializedOutput,
+ partition: usize,
+ output_partitions: usize,
+) -> Result<Vec<SendableRecordBatchStream>> {
+ 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);
+ if let Some(file) = &buffered.spill_file {
+ streams.push(
+ output
+ .spill_manager
+ .read_spill_as_stream(Arc::clone(file), None)?,
+ );
+ }
+ }
+ Ok(streams)
+}
+
+/// 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),
+ )
+ }
+
+ 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
+ }
+}
+
+/// 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 {
+ Self {
+ id: self.id,
+ name: self.name.clone(),
+ schema: Arc::clone(&self.schema),
+ buffer: Some(buffer),
+ metrics: ExecutionPlanMetricsSet::new(),
+ cache: Arc::clone(&self.cache),
+ }
+ }
+
+ pub fn is_bound(&self) -> bool {
Review Comment:
nit: `is_bound` has no caller. Could you remove it until something needs it?
##########
datafusion/sqllogictest/test_files/cte_materialized.slt:
##########
@@ -0,0 +1,233 @@
+# 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
+
+# Spill: the body does not fit in the memory limit, so it is written to disk
+# and replayed from there by every reference.
+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=1,
<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>]
+
+statement ok
+RESET datafusion.runtime.memory_limit
+
+statement ok
+RESET datafusion.catalog.create_default_catalog_and_schema
+
Review Comment:
nit: this file does not set `create_default_catalog_and_schema`, so this
`RESET` is not necessary.
```suggestion
```
##########
datafusion/expr/src/logical_plan/materialized_cte.rs:
##########
@@ -0,0 +1,186 @@
+// 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,
+}
+
+impl PartialOrd for MaterializedCte {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ match self.id.cmp(&other.id) {
+ Ordering::Equal => self.name.partial_cmp(&other.name),
+ ord => Some(ord),
+ }
+ }
+}
+
+impl UserDefinedLogicalNodeCore for MaterializedCte {
Review Comment:
`MaterializedCte` does not override `prevent_predicate_push_down_columns`,
so `PushDownFilter` keeps every filter above this node. This node is the root
of the query that declares the CTE, so a filter from an outer query is not
pushed into the continuation. Only the scans must be a fence. On the physical
side, `MaterializedCteExec` uses the default `gather_filters_for_pushdown`, so
parent filters and dynamic filters also stop there:
```sql
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;
-- FilterExec: v = 10 stays above MaterializedCteExec and is not pushed to
the scan of t
```
Without `MATERIALIZED` the filter is pushed to the scan of `t`. This does
not give incorrect results, but a view that uses `MATERIALIZED` can become much
slower. `PushDownFilter` sends the pushed predicates to all inputs of an
extension node, so the logical side needs a special case (or a `LogicalPlan`
variant, see open question 1). On the physical side,
`gather_filters_for_pushdown` can send the parent filters to child 1 only. If
you prefer to do this in a follow-up, could you add it to the list in the PR
description?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
