datafusion_physical_plan/
streaming.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Generic plans for deferred execution: [`StreamingTableExec`] and [`PartitionStream`]
19
20use std::fmt::Debug;
21use std::sync::Arc;
22
23use super::{DisplayAs, DisplayFormatType, PlanProperties};
24use crate::coop::make_cooperative;
25use crate::display::{ProjectSchemaDisplay, display_orderings};
26use crate::execution_plan::{Boundedness, EmissionType, SchedulingType};
27use crate::limit::LimitStream;
28use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
29use crate::projection::{
30    ProjectionExec, all_alias_free_columns, new_projections_for_columns, update_ordering,
31};
32use crate::stream::RecordBatchStreamAdapter;
33use crate::{
34    ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions,
35    SendableRecordBatchStream,
36};
37
38use arrow::datatypes::{Schema, SchemaRef};
39use datafusion_common::tree_node::TreeNodeRecursion;
40use datafusion_common::{Result, internal_err, plan_err};
41use datafusion_execution::TaskContext;
42use datafusion_physical_expr::PhysicalExpr;
43use datafusion_physical_expr::projection::ProjectionMapping;
44use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
45
46use async_trait::async_trait;
47use futures::stream::StreamExt;
48use log::debug;
49
50/// A partition that can be converted into a [`SendableRecordBatchStream`]
51///
52/// Combined with [`StreamingTableExec`], you can use this trait to implement
53/// [`ExecutionPlan`] for a custom source with less boiler plate than
54/// implementing `ExecutionPlan` directly for many use cases.
55pub trait PartitionStream: Debug + Send + Sync {
56    /// Returns the schema of this partition
57    fn schema(&self) -> &SchemaRef;
58
59    /// Returns a stream yielding this partitions values
60    fn execute(&self, ctx: Arc<TaskContext>) -> SendableRecordBatchStream;
61}
62
63/// An [`ExecutionPlan`] for one or more [`PartitionStream`]s.
64///
65/// If your source can be represented as one or more [`PartitionStream`]s, you can
66/// use this struct to implement [`ExecutionPlan`].
67#[derive(Clone)]
68pub struct StreamingTableExec {
69    partitions: Vec<Arc<dyn PartitionStream>>,
70    projection: Option<Arc<[usize]>>,
71    projected_schema: SchemaRef,
72    projected_output_ordering: Vec<LexOrdering>,
73    infinite: bool,
74    limit: Option<usize>,
75    cache: Arc<PlanProperties>,
76    metrics: ExecutionPlanMetricsSet,
77}
78
79impl StreamingTableExec {
80    /// Try to create a new [`StreamingTableExec`] returning an error if the schema is incorrect
81    pub fn try_new(
82        schema: SchemaRef,
83        partitions: Vec<Arc<dyn PartitionStream>>,
84        projection: Option<&Vec<usize>>,
85        projected_output_ordering: impl IntoIterator<Item = LexOrdering>,
86        infinite: bool,
87        limit: Option<usize>,
88    ) -> Result<Self> {
89        for x in partitions.iter() {
90            let partition_schema = x.schema();
91            if !schema.eq(partition_schema) {
92                debug!(
93                    "Target schema does not match with partition schema. \
94                        Target_schema: {schema:?}. Partition Schema: {partition_schema:?}"
95                );
96                return plan_err!("Mismatch between schema and batches");
97            }
98        }
99
100        let projected_schema = match projection {
101            Some(p) => Arc::new(schema.project(p)?),
102            None => schema,
103        };
104        let projected_output_ordering =
105            projected_output_ordering.into_iter().collect::<Vec<_>>();
106        let cache = Self::compute_properties(
107            Arc::clone(&projected_schema),
108            projected_output_ordering.clone(),
109            Partitioning::UnknownPartitioning(partitions.len()),
110            infinite,
111        );
112        Ok(Self {
113            partitions,
114            projected_schema,
115            projection: projection.cloned().map(Into::into),
116            projected_output_ordering,
117            infinite,
118            limit,
119            cache: Arc::new(cache),
120            metrics: ExecutionPlanMetricsSet::new(),
121        })
122    }
123
124    /// Declares the output partitioning of this stream.
125    ///
126    /// `output_partitioning` must describe this plan's current output and have
127    /// the same number of partitions as the stream.
128    pub fn with_output_partitioning(
129        mut self,
130        output_partitioning: Partitioning,
131    ) -> Result<Self> {
132        if output_partitioning.partition_count() != self.partitions.len() {
133            return plan_err!(
134                "Output partitioning has {} partitions but stream has {} partitions",
135                output_partitioning.partition_count(),
136                self.partitions.len()
137            );
138        }
139        Arc::make_mut(&mut self.cache).partitioning = output_partitioning;
140        Ok(self)
141    }
142
143    pub fn partitions(&self) -> &Vec<Arc<dyn PartitionStream>> {
144        &self.partitions
145    }
146
147    pub fn partition_schema(&self) -> &SchemaRef {
148        self.partitions[0].schema()
149    }
150
151    pub fn projection(&self) -> &Option<Arc<[usize]>> {
152        &self.projection
153    }
154
155    pub fn projected_schema(&self) -> &Schema {
156        &self.projected_schema
157    }
158
159    pub fn projected_output_ordering(&self) -> impl IntoIterator<Item = LexOrdering> {
160        self.projected_output_ordering.clone()
161    }
162
163    pub fn is_infinite(&self) -> bool {
164        self.infinite
165    }
166
167    pub fn limit(&self) -> Option<usize> {
168        self.limit
169    }
170
171    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
172    fn compute_properties(
173        schema: SchemaRef,
174        orderings: Vec<LexOrdering>,
175        output_partitioning: Partitioning,
176        infinite: bool,
177    ) -> PlanProperties {
178        // Calculate equivalence properties:
179        let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings);
180
181        let boundedness = if infinite {
182            Boundedness::Unbounded {
183                requires_infinite_memory: false,
184            }
185        } else {
186            Boundedness::Bounded
187        };
188        PlanProperties::new(
189            eq_properties,
190            output_partitioning,
191            EmissionType::Incremental,
192            boundedness,
193        )
194        .with_scheduling_type(SchedulingType::Cooperative)
195    }
196}
197
198impl Debug for StreamingTableExec {
199    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
200        f.debug_struct("LazyMemTableExec").finish_non_exhaustive()
201    }
202}
203
204impl DisplayAs for StreamingTableExec {
205    fn fmt_as(
206        &self,
207        t: DisplayFormatType,
208        f: &mut std::fmt::Formatter,
209    ) -> std::fmt::Result {
210        match t {
211            DisplayFormatType::Default | DisplayFormatType::Verbose => {
212                write!(
213                    f,
214                    "StreamingTableExec: partition_sizes={:?}",
215                    self.partitions.len(),
216                )?;
217                if !self.projected_schema.fields().is_empty() {
218                    write!(
219                        f,
220                        ", projection={}",
221                        ProjectSchemaDisplay(&self.projected_schema)
222                    )?;
223                }
224                if self.infinite {
225                    write!(f, ", infinite_source=true")?;
226                }
227                if let Some(fetch) = self.limit {
228                    write!(f, ", fetch={fetch}")?;
229                }
230                if !matches!(
231                    self.cache.output_partitioning(),
232                    Partitioning::UnknownPartitioning(_)
233                ) {
234                    write!(
235                        f,
236                        ", output_partitioning={}",
237                        self.cache.output_partitioning()
238                    )?;
239                }
240
241                display_orderings(f, &self.projected_output_ordering)?;
242
243                Ok(())
244            }
245            DisplayFormatType::TreeRender => {
246                if self.infinite {
247                    writeln!(f, "infinite={}", self.infinite)?;
248                }
249                if let Some(limit) = self.limit {
250                    write!(f, "limit={limit}")?;
251                } else {
252                    write!(f, "limit=None")?;
253                }
254
255                Ok(())
256            }
257        }
258    }
259}
260
261#[async_trait]
262impl ExecutionPlan for StreamingTableExec {
263    fn name(&self) -> &'static str {
264        "StreamingTableExec"
265    }
266
267    fn properties(&self) -> &Arc<PlanProperties> {
268        &self.cache
269    }
270
271    fn fetch(&self) -> Option<usize> {
272        self.limit
273    }
274
275    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
276        vec![]
277    }
278
279    fn apply_expressions(
280        &self,
281        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
282    ) -> Result<TreeNodeRecursion> {
283        Ok(TreeNodeRecursion::Continue)
284    }
285
286    fn replace_children(
287        self: Arc<Self>,
288        children: Vec<Arc<dyn ExecutionPlan>>,
289        _: ReplaceChildrenOptions,
290    ) -> Result<Arc<dyn ExecutionPlan>> {
291        if children.is_empty() {
292            Ok(self)
293        } else {
294            internal_err!("Children cannot be replaced in {self:?}")
295        }
296    }
297
298    fn with_new_children(
299        self: Arc<Self>,
300        children: Vec<Arc<dyn ExecutionPlan>>,
301    ) -> Result<Arc<dyn ExecutionPlan>> {
302        self.replace_children(
303            children,
304            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
305        )
306    }
307
308    fn execute(
309        &self,
310        partition: usize,
311        ctx: Arc<TaskContext>,
312    ) -> Result<SendableRecordBatchStream> {
313        let stream = self.partitions[partition].execute(Arc::clone(&ctx));
314        let projected_stream = match self.projection.clone() {
315            Some(projection) => Box::pin(RecordBatchStreamAdapter::new(
316                Arc::clone(&self.projected_schema),
317                stream.map(move |x| {
318                    x.and_then(|b| b.project(projection.as_ref()).map_err(Into::into))
319                }),
320            )),
321            None => stream,
322        };
323        let stream = make_cooperative(projected_stream);
324
325        Ok(match self.limit {
326            None => stream,
327            Some(fetch) => {
328                let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
329                Box::pin(LimitStream::new(stream, 0, Some(fetch), baseline_metrics))
330            }
331        })
332    }
333
334    /// Tries to embed `projection` to its input (`streaming table`).
335    /// If possible, returns [`StreamingTableExec`] as the top plan. Otherwise,
336    /// returns `None`.
337    fn try_swapping_with_projection(
338        &self,
339        projection: &ProjectionExec,
340    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
341        if !all_alias_free_columns(projection.expr()) {
342            return Ok(None);
343        }
344
345        let streaming_table_projections =
346            self.projection().as_ref().map(|i| i.as_ref().to_vec());
347        let new_projections = new_projections_for_columns(
348            projection.expr(),
349            &streaming_table_projections
350                .unwrap_or_else(|| (0..self.schema().fields().len()).collect()),
351        );
352
353        let mut lex_orderings = vec![];
354        for ordering in self.projected_output_ordering().into_iter() {
355            let Some(ordering) = update_ordering(ordering, projection.expr())? else {
356                return Ok(None);
357            };
358            lex_orderings.push(ordering);
359        }
360        let projection_mapping = ProjectionMapping::try_new(
361            projection
362                .expr()
363                .iter()
364                .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())),
365            &self.schema(),
366        )?;
367        let output_partitioning = self
368            .cache
369            .output_partitioning()
370            .project(&projection_mapping, self.cache.equivalence_properties());
371
372        StreamingTableExec::try_new(
373            Arc::clone(self.partition_schema()),
374            self.partitions().clone(),
375            Some(new_projections.as_ref()),
376            lex_orderings,
377            self.is_infinite(),
378            self.limit(),
379        )
380        .and_then(|exec| exec.with_output_partitioning(output_partitioning))
381        .map(|e| Some(Arc::new(e) as _))
382    }
383
384    fn metrics(&self) -> Option<MetricsSet> {
385        Some(self.metrics.clone_inner())
386    }
387
388    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
389        Some(Arc::new(StreamingTableExec {
390            partitions: self.partitions.clone(),
391            projection: self.projection.clone(),
392            projected_schema: Arc::clone(&self.projected_schema),
393            projected_output_ordering: self.projected_output_ordering.clone(),
394            infinite: self.infinite,
395            limit,
396            cache: Arc::clone(&self.cache),
397            metrics: self.metrics.clone(),
398        }))
399    }
400}
401
402#[cfg(test)]
403mod test {
404    use super::*;
405    use crate::collect_partitioned;
406    use crate::streaming::PartitionStream;
407    use crate::test::{TestPartitionStream, make_partition};
408    use arrow::record_batch::RecordBatch;
409
410    #[tokio::test]
411    async fn test_no_limit() {
412        let exec = TestBuilder::new()
413            // Make 2 batches, each with 100 rows
414            .with_batches(vec![make_partition(100), make_partition(100)])
415            .build();
416
417        let counts = collect_num_rows(Arc::new(exec)).await;
418        assert_eq!(counts, vec![200]);
419    }
420
421    #[tokio::test]
422    async fn test_limit() {
423        let exec = TestBuilder::new()
424            // Make 2 batches, each with 100 rows
425            .with_batches(vec![make_partition(100), make_partition(100)])
426            // Limit to only the first 75 rows back
427            .with_limit(Some(75))
428            .build();
429
430        let counts = collect_num_rows(Arc::new(exec)).await;
431        assert_eq!(counts, vec![75]);
432    }
433
434    /// Runs the provided execution plan and returns a vector of the number of
435    /// rows in each partition
436    async fn collect_num_rows(exec: Arc<dyn ExecutionPlan>) -> Vec<usize> {
437        let ctx = Arc::new(TaskContext::default());
438        let partition_batches = collect_partitioned(exec, ctx).await.unwrap();
439        partition_batches
440            .into_iter()
441            .map(|batches| batches.iter().map(|b| b.num_rows()).sum::<usize>())
442            .collect()
443    }
444
445    #[derive(Default)]
446    struct TestBuilder {
447        schema: Option<SchemaRef>,
448        partitions: Vec<Arc<dyn PartitionStream>>,
449        projection: Option<Vec<usize>>,
450        projected_output_ordering: Vec<LexOrdering>,
451        infinite: bool,
452        limit: Option<usize>,
453    }
454
455    impl TestBuilder {
456        fn new() -> Self {
457            Self::default()
458        }
459
460        /// Set the batches for the stream
461        fn with_batches(mut self, batches: Vec<RecordBatch>) -> Self {
462            let stream = TestPartitionStream::new_with_batches(batches);
463            self.schema = Some(Arc::clone(stream.schema()));
464            self.partitions = vec![Arc::new(stream)];
465            self
466        }
467
468        /// Set the limit for the stream
469        fn with_limit(mut self, limit: Option<usize>) -> Self {
470            self.limit = limit;
471            self
472        }
473
474        fn build(self) -> StreamingTableExec {
475            StreamingTableExec::try_new(
476                self.schema.unwrap(),
477                self.partitions,
478                self.projection.as_ref(),
479                self.projected_output_ordering,
480                self.infinite,
481                self.limit,
482            )
483            .unwrap()
484        }
485    }
486}
datafusion_physical_plan/streaming.rs

datafusion_physical_plan/
streaming.rs
