zhuqi-lucas commented on code in PR #25580: URL: https://github.com/apache/datafusion/pull/25580#discussion_r4244985567
########## datafusion/physical-plan/src/materialized_cte.rs: ########## @@ -0,0 +1,834 @@ +// 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. + +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, +} + +impl BufferedPartition { + /// Move in-memory batches from the end of `batches` to a new spill file, + /// until at least `bytes` are released or no in-memory batch is left. + /// The new file is read before the existing spill files, so the order of + /// the partition is kept. + fn spill_in_memory_suffix( + &mut self, + bytes: usize, + spill_manager: &SpillManager, + ) -> Result<()> { + let mut split = self.batches.len(); + let mut released = 0; + while split > 0 && released < bytes { + split -= 1; + released += self.batches[split].get_array_memory_size(); + } + let suffix = self.batches.split_off(split); + let mut file = spill_manager.create_in_progress_file("MaterializedCte")?; + for batch in &suffix { + let size = file.append_batch(batch)?; + 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 `partitions`. +/// +/// 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(partitions: &[BufferedPartition], scan_partitions: &[usize]) -> usize { + scan_partitions + .iter() + .map(|&n| { + (0..n) + .map(|p| { + partitions + .iter() + .skip(p) + .step_by(n) + .map(|b| b.max_spilled_batch_size) + .max() + .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, from +/// the partition that holds the most memory first, until it can. +fn reserve_replay_memory( Review Comment: Spilling here can raise the requirement it is trying to satisfy, so this can fail a query that had enough memory. `spill_in_memory_suffix` only ever grows `max_spilled_batch_size`, and `replay_memory` sums exactly that value across scan partitions. So moving an in-memory batch that is larger than anything already spilled *increases* `required` on the next iteration. Reproducible as SQL (two references, 18 MB limit). The first batch holds large strings and stays in memory; later batches are smaller and spill, so `max_spilled_batch_size` is small: ```sql SET datafusion.runtime.memory_limit = '18M'; 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); ``` ``` Resources exhausted: Failed to allocate additional 49.1 MB for MaterializedCte[c] replay with 0.0 B already allocated - 18.0 MB remain available for the total memory pool: greedy(used: 72.0 B, pool_size: 18.0 MB) ``` Note `used: 72.0 B` — the whole buffer has been spilled, the pool is empty, and it still fails, because `required` grew to 2.7x the entire pool. Before the loop spilled the in-memory prefix, `required` was 2 x ~3 MB, which fits comfortably. I isolated the cause: suppressing the `max_spilled_batch_size` update inside `spill_in_memory_suffix` makes the same case succeed (that is not the fix — it would then under-reserve for the large spilled batch — it just pins the mechanism). `spill_replay_memory_is_reserved` cannot catch this because `generate_series` produces uniform batches, so `max_spilled_batch_size` never changes as more batches spill. Possible directions: prefer spilling batches no larger than the current `max_spilled_batch_size`, or compute the post-spill `required` before committing and only spill when it actually improves. ########## datafusion/physical-plan/src/materialized_cte.rs: ########## @@ -0,0 +1,834 @@ +// 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. + +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, +} + +impl BufferedPartition { + /// Move in-memory batches from the end of `batches` to a new spill file, + /// until at least `bytes` are released or no in-memory batch is left. + /// The new file is read before the existing spill files, so the order of + /// the partition is kept. + fn spill_in_memory_suffix( + &mut self, + bytes: usize, + spill_manager: &SpillManager, + ) -> Result<()> { + let mut split = self.batches.len(); + let mut released = 0; + while split > 0 && released < bytes { + split -= 1; + released += self.batches[split].get_array_memory_size(); + } + let suffix = self.batches.split_off(split); + let mut file = spill_manager.create_in_progress_file("MaterializedCte")?; + for batch in &suffix { + let size = file.append_batch(batch)?; + 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 `partitions`. +/// +/// 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(partitions: &[BufferedPartition], scan_partitions: &[usize]) -> usize { + scan_partitions + .iter() + .map(|&n| { + (0..n) + .map(|p| { + partitions + .iter() + .skip(p) + .step_by(n) + .map(|b| b.max_spilled_batch_size) + .max() + .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, from +/// the partition that holds the most memory first, until it can. +fn reserve_replay_memory( + partitions: &mut [BufferedPartition], + scan_partitions: &[usize], + spill_manager: &SpillManager, + reservation: &mut MemoryReservation, +) -> Result<()> { + loop { + let required = replay_memory(partitions, scan_partitions); + let Err(e) = reservation.try_resize(required) else { + return Ok(()); + }; + let Some(partition) = partitions + .iter_mut() + .filter(|p| !p.batches.is_empty()) + .max_by_key(|p| p.reservation.size()) + else { + return Err(e); + }; + partition.spill_in_memory_suffix( + required.saturating_sub(reservation.size()), + 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), + ) + } + + 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()) + } + + fn statistics_from_inputs( Review Comment: Worth recording that unknown statistics here are not only a missed optimization — they can make a query slower, which is the opposite of what `MATERIALIZED` is asked for. ```sql CREATE TABLE small AS SELECT value AS k FROM generate_series(1, 5); CREATE TABLE huge AS SELECT value AS k FROM generate_series(1, 100000); EXPLAIN WITH c AS MATERIALIZED (SELECT k FROM small) SELECT count(*) FROM huge JOIN c ON huge.k = c.k; ``` Inlined, `JoinSelection` puts the 5-row side on the build side. Materialized, it swaps: ``` 07)----------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)] 08)------------CoalescePartitionsExec 09)--------------DataSourceExec: partitions=4 <- 100k rows, build side 10)------------MaterializedCteScanExec: name=c <- 5 rows, probe side ``` So a 100k-row hash table is built instead of a 5-row one, and the probe side loses its partitioning to `CoalescePartitionsExec`. Writing the CTE first in the join gives the same plan, so it is statistics-driven rather than textual. I tried the obvious shortcut before writing this: the scan holds `Arc<MaterializedCteBuffer>`, which holds the body, so it looked like it could read the body's statistics without `StatisticsContext`. It does not work — `ExecutionPlan::partition_statistics` is deprecated and returns unknown for `DataSourceExec` (which overrides `statistics_from_inputs` instead), and `statistics_from_inputs` is documented as not allowed to call `StatisticsContext::compute`. Your note about a leaf not being able to ask its producer is accurate, and the gap is in the statistics API rather than here. Not asking for it in this PR — but given it can regress plans rather than just leave them flat, it may deserve a mention as a known limitation alongside the follow-up. -- 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]
