laminar_sql/datafusion/
table_provider.rs

1//! Streaming table provider for `DataFusion` integration
2//!
3//! This module provides `StreamingTableProvider` which implements `DataFusion`'s
4//! `TableProvider` trait, allowing streaming sources to be registered as
5//! tables in a `SessionContext` and queried with SQL.
6
7use std::sync::Arc;
8
9use arrow_schema::SchemaRef;
10use async_trait::async_trait;
11use datafusion::catalog::Session;
12use datafusion::datasource::TableProvider;
13use datafusion::physical_plan::ExecutionPlan;
14use datafusion_common::DataFusionError;
15use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType};
16
17use super::exec::StreamingScanExec;
18use super::source::StreamSourceRef;
19
20/// A `DataFusion` table provider backed by a streaming source.
21///
22/// This allows streaming sources to be registered as tables in `DataFusion`'s
23/// `SessionContext` and queried using SQL. The provider handles:
24///
25/// - Schema exposure to `DataFusion`'s catalog
26/// - Projection pushdown to the source
27/// - Filter pushdown when supported by the source
28///
29/// # Usage
30///
31/// ```rust,ignore
32/// let source = Arc::new(ChannelStreamSource::new(schema));
33/// let provider = StreamingTableProvider::new("events", source);
34/// ctx.register_table("events", Arc::new(provider))?;
35///
36/// let df = ctx.sql("SELECT * FROM events WHERE id > 100").await?;
37/// ```
38#[derive(Debug)]
39pub struct StreamingTableProvider {
40    /// Table name
41    name: String,
42    /// The underlying streaming source
43    source: StreamSourceRef,
44}
45
46impl StreamingTableProvider {
47    /// Creates a new streaming table provider.
48    #[must_use]
49    pub fn new(name: impl Into<String>, source: StreamSourceRef) -> Self {
50        Self {
51            name: name.into(),
52            source,
53        }
54    }
55
56    /// Returns the table name.
57    #[must_use]
58    pub fn name(&self) -> &str {
59        &self.name
60    }
61
62    /// Returns the underlying streaming source.
63    #[must_use]
64    pub fn source(&self) -> &StreamSourceRef {
65        &self.source
66    }
67}
68
69#[async_trait]
70impl TableProvider for StreamingTableProvider {
71    fn as_any(&self) -> &dyn std::any::Any {
72        self
73    }
74
75    fn schema(&self) -> SchemaRef {
76        self.source.schema()
77    }
78
79    fn table_type(&self) -> TableType {
80        // Streaming tables behave like base tables but are read-only
81        TableType::Base
82    }
83
84    fn supports_filters_pushdown(
85        &self,
86        filters: &[&Expr],
87    ) -> Result<Vec<TableProviderFilterPushDown>, DataFusionError> {
88        // Ask the source which filters it can handle
89        let expr_refs: Vec<Expr> = filters.iter().map(|e| (*e).clone()).collect();
90        let supported = self.source.supports_filters(&expr_refs);
91
92        Ok(supported
93            .into_iter()
94            .map(|s| {
95                if s {
96                    TableProviderFilterPushDown::Exact
97                } else {
98                    TableProviderFilterPushDown::Unsupported
99                }
100            })
101            .collect())
102    }
103
104    async fn scan(
105        &self,
106        _state: &dyn Session,
107        projection: Option<&Vec<usize>>,
108        filters: &[Expr],
109        _limit: Option<usize>,
110    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
111        // DataFusion only passes filters here that `supports_filters_pushdown`
112        // already claimed as `Exact`/`Inexact` — no need to re-ask the source
113        // which ones it can handle. Forwarding all of them was the old path's
114        // cost (two Vec<Expr> clones per scan for nothing).
115        Ok(Arc::new(StreamingScanExec::new(
116            Arc::clone(&self.source),
117            projection.cloned(),
118            filters.to_vec(),
119        )))
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::datafusion::source::StreamSource;
127    use arrow_schema::{DataType, Field, Schema};
128    use datafusion::execution::SendableRecordBatchStream;
129
130    #[derive(Debug)]
131    struct MockSource {
132        schema: SchemaRef,
133        supports_eq_filter: bool,
134    }
135
136    #[async_trait]
137    impl StreamSource for MockSource {
138        fn schema(&self) -> SchemaRef {
139            Arc::clone(&self.schema)
140        }
141
142        fn stream(
143            &self,
144            _projection: Option<Vec<usize>>,
145            _filters: Vec<Expr>,
146        ) -> Result<SendableRecordBatchStream, DataFusionError> {
147            Err(DataFusionError::NotImplemented("mock".to_string()))
148        }
149
150        fn supports_filters(&self, filters: &[Expr]) -> Vec<bool> {
151            filters
152                .iter()
153                .map(|f| {
154                    if self.supports_eq_filter {
155                        // Only support equality filters for testing
156                        matches!(f, Expr::BinaryExpr(e) if e.op == datafusion_expr::Operator::Eq)
157                    } else {
158                        false
159                    }
160                })
161                .collect()
162        }
163    }
164
165    fn test_schema() -> SchemaRef {
166        Arc::new(Schema::new(vec![
167            Field::new("id", DataType::Int64, false),
168            Field::new("name", DataType::Utf8, true),
169        ]))
170    }
171
172    #[test]
173    fn test_table_provider_schema() {
174        let schema = test_schema();
175        let source: StreamSourceRef = Arc::new(MockSource {
176            schema: Arc::clone(&schema),
177            supports_eq_filter: false,
178        });
179        let provider = StreamingTableProvider::new("test_table", source);
180
181        assert_eq!(provider.schema(), schema);
182        assert_eq!(provider.name(), "test_table");
183    }
184
185    #[test]
186    fn test_table_provider_type() {
187        let schema = test_schema();
188        let source: StreamSourceRef = Arc::new(MockSource {
189            schema,
190            supports_eq_filter: false,
191        });
192        let provider = StreamingTableProvider::new("test", source);
193
194        assert_eq!(provider.table_type(), TableType::Base);
195    }
196
197    #[test]
198    fn test_filter_pushdown_unsupported() {
199        let schema = test_schema();
200        let source: StreamSourceRef = Arc::new(MockSource {
201            schema,
202            supports_eq_filter: false,
203        });
204        let provider = StreamingTableProvider::new("test", source);
205
206        let filter = Expr::Literal(datafusion_common::ScalarValue::Int64(Some(1)), None);
207        let result = provider.supports_filters_pushdown(&[&filter]).unwrap();
208
209        assert_eq!(result.len(), 1);
210        assert!(matches!(
211            result[0],
212            TableProviderFilterPushDown::Unsupported
213        ));
214    }
215
216    #[test]
217    fn test_filter_pushdown_supported() {
218        let schema = test_schema();
219        let source: StreamSourceRef = Arc::new(MockSource {
220            schema,
221            supports_eq_filter: true,
222        });
223        let provider = StreamingTableProvider::new("test", source);
224
225        // Create an equality filter: id = 1
226        let filter = Expr::BinaryExpr(datafusion_expr::BinaryExpr {
227            left: Box::new(Expr::Column(datafusion_common::Column::new_unqualified(
228                "id",
229            ))),
230            op: datafusion_expr::Operator::Eq,
231            right: Box::new(Expr::Literal(
232                datafusion_common::ScalarValue::Int64(Some(1)),
233                None,
234            )),
235        });
236        let result = provider.supports_filters_pushdown(&[&filter]).unwrap();
237
238        assert_eq!(result.len(), 1);
239        assert!(matches!(result[0], TableProviderFilterPushDown::Exact));
240    }
241
242    #[tokio::test]
243    async fn test_scan_creates_exec() {
244        use crate::datafusion::create_session_context;
245
246        let schema = test_schema();
247        let source: StreamSourceRef = Arc::new(MockSource {
248            schema: Arc::clone(&schema),
249            supports_eq_filter: false,
250        });
251        let provider = StreamingTableProvider::new("test", source);
252
253        let ctx = create_session_context();
254        let session_state = ctx.state();
255
256        let exec = provider
257            .scan(&session_state, None, &[], None)
258            .await
259            .unwrap();
260
261        // Verify it's a StreamingScanExec
262        assert!(exec.as_any().is::<StreamingScanExec>());
263        assert_eq!(exec.schema(), schema);
264    }
265
266    #[tokio::test]
267    async fn test_scan_with_projection() {
268        use crate::datafusion::create_session_context;
269
270        let schema = test_schema();
271        let source: StreamSourceRef = Arc::new(MockSource {
272            schema,
273            supports_eq_filter: false,
274        });
275        let provider = StreamingTableProvider::new("test", source);
276
277        let ctx = create_session_context();
278        let session_state = ctx.state();
279
280        let projection = vec![0]; // Only id column
281        let exec = provider
282            .scan(&session_state, Some(&projection), &[], None)
283            .await
284            .unwrap();
285
286        let output_schema = exec.schema();
287        assert_eq!(output_schema.fields().len(), 1);
288        assert_eq!(output_schema.field(0).name(), "id");
289    }
290}
laminar_sql/datafusion/table_provider.rs

laminar_sql/datafusion/
table_provider.rs
