datafusion_catalog/
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//! A simplified [`TableProvider`] for streaming partitioned datasets
19
20use std::sync::Arc;
21
22use arrow::datatypes::SchemaRef;
23use async_trait::async_trait;
24use datafusion_common::{DFSchema, Result, plan_err};
25use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
26use datafusion_expr::{Expr, SortExpr, TableType};
27use datafusion_physical_expr::equivalence::project_ordering;
28use datafusion_physical_expr::projection::ProjectionMapping;
29use datafusion_physical_expr::{
30    EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs,
31};
32use datafusion_physical_plan::ExecutionPlan;
33use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec};
34use log::debug;
35
36use crate::{Session, TableProvider};
37
38/// A [`TableProvider`] that streams a set of [`PartitionStream`]
39#[derive(Debug)]
40pub struct StreamingTable {
41    schema: SchemaRef,
42    partitions: Vec<Arc<dyn PartitionStream>>,
43    infinite: bool,
44    sort_order: Vec<SortExpr>,
45    output_partitioning: Option<Partitioning>,
46}
47
48impl StreamingTable {
49    /// Try to create a new [`StreamingTable`] returning an error if the schema is incorrect
50    pub fn try_new(
51        schema: SchemaRef,
52        partitions: Vec<Arc<dyn PartitionStream>>,
53    ) -> Result<Self> {
54        for x in partitions.iter() {
55            let partition_schema = x.schema();
56            if !schema.contains(partition_schema) {
57                debug!(
58                    "target schema does not contain partition schema. \
59                        Target_schema: {schema:?}. Partition Schema: {partition_schema:?}"
60                );
61                return plan_err!("Mismatch between schema and batches");
62            }
63        }
64
65        Ok(Self {
66            schema,
67            partitions,
68            infinite: false,
69            sort_order: vec![],
70            output_partitioning: None,
71        })
72    }
73
74    /// Sets streaming table can be infinite.
75    pub fn with_infinite_table(mut self, infinite: bool) -> Self {
76        self.infinite = infinite;
77        self
78    }
79
80    /// Sets the existing ordering of streaming table.
81    pub fn with_sort_order(mut self, sort_order: Vec<SortExpr>) -> Self {
82        self.sort_order = sort_order;
83        self
84    }
85
86    /// Declares the output partitioning of this streaming table.
87    ///
88    /// The partitioning expressions refer to the table schema before scan
89    /// projection. If a scan projection removes a partitioning expression, the
90    /// physical plan reports unknown partitioning.
91    pub fn with_output_partitioning(mut self, output_partitioning: Partitioning) -> Self {
92        self.output_partitioning = Some(output_partitioning);
93        self
94    }
95
96    fn output_partitioning(
97        &self,
98        projection: Option<&Vec<usize>>,
99    ) -> Result<Partitioning> {
100        let Some(output_partitioning) = &self.output_partitioning else {
101            return Ok(Partitioning::UnknownPartitioning(self.partitions.len()));
102        };
103        let Some(projection) = projection else {
104            return Ok(output_partitioning.clone());
105        };
106
107        let projection_mapping =
108            ProjectionMapping::from_indices(projection, &self.schema)?;
109        let eq_properties = EquivalenceProperties::new(Arc::clone(&self.schema));
110        Ok(output_partitioning.project(&projection_mapping, &eq_properties))
111    }
112}
113
114#[async_trait]
115impl TableProvider for StreamingTable {
116    fn schema(&self) -> SchemaRef {
117        Arc::clone(&self.schema)
118    }
119
120    fn table_type(&self) -> TableType {
121        TableType::View
122    }
123
124    async fn scan(
125        &self,
126        state: &dyn Session,
127        projection: Option<&Vec<usize>>,
128        _filters: &[Expr],
129        limit: Option<usize>,
130    ) -> Result<Arc<dyn ExecutionPlan>> {
131        let physical_sort = if !self.sort_order.is_empty() {
132            let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
133            let eqp = state.execution_props();
134
135            let original_sort_exprs = create_physical_sort_exprs(
136                &self.sort_order,
137                &df_schema,
138                eqp,
139                &PhysicalPlanningContext::default(),
140            )?;
141
142            if let Some(p) = projection {
143                // When performing a projection, the output columns will not match
144                // the original physical sort expression indices. Also the sort columns
145                // may not be in the output projection. To correct for these issues
146                // we need to project the ordering based on the output schema.
147                let schema = Arc::new(self.schema.project(p)?);
148                LexOrdering::new(original_sort_exprs)
149                    .and_then(|lex_ordering| project_ordering(&lex_ordering, &schema))
150                    .map(|lex_ordering| lex_ordering.to_vec())
151                    .unwrap_or_default()
152            } else {
153                original_sort_exprs
154            }
155        } else {
156            vec![]
157        };
158
159        let exec = StreamingTableExec::try_new(
160            Arc::clone(&self.schema),
161            self.partitions.clone(),
162            projection,
163            LexOrdering::new(physical_sort),
164            self.infinite,
165            limit,
166        )?
167        .with_output_partitioning(self.output_partitioning(projection)?)?;
168
169        Ok(Arc::new(exec))
170    }
171}
datafusion_catalog/streaming.rs

datafusion_catalog/
streaming.rs
