laminar_sql/datafusion/
live_source.rs

1//! Swappable table provider that eliminates per-cycle catalog churn and
2//! enables physical plan caching.
3//!
4//! Register a [`LiveSourceProvider`] once at pipeline startup. Each cycle,
5//! swap batches via [`LiveSourceHandle`], then execute the cached physical
6//! plan. The internal `LiveSourceExec` reads from the shared slot at `execute()` time,
7//! so the cached plan always sees fresh data.
8
9use std::sync::Arc;
10
11use arrow::array::RecordBatch;
12use arrow::datatypes::SchemaRef;
13use async_trait::async_trait;
14use datafusion::catalog::Session;
15use datafusion::datasource::TableProvider;
16use datafusion::error::DataFusionError;
17use datafusion::execution::{SendableRecordBatchStream, TaskContext};
18use datafusion::logical_expr::Expr;
19use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
20use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
21use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
22use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
23use datafusion_expr::TableType;
24use parking_lot::Mutex;
25
26/// Shared batch slot used by the provider and every `LiveSourceExec` scan.
27///
28/// Wrapping the batch list in `Arc` means the hot-path `execute()` clone is
29/// an O(1) Arc bump, not a full `Vec<RecordBatch>` clone (which would
30/// Arc-bump every column of every batch). `swap()` replaces the Arc
31/// wholesale while any in-flight reader keeps its prior snapshot alive.
32type BatchSlot = Arc<Mutex<Arc<Vec<RecordBatch>>>>;
33
34fn new_slot() -> BatchSlot {
35    Arc::new(Mutex::new(Arc::new(Vec::new())))
36}
37
38// ── TableProvider ────────────────────────────────────────────────────
39
40/// Swappable `TableProvider` for streaming micro-batch execution.
41///
42/// `scan()` returns an internal execution plan that reads from the shared
43/// batch slot at `execute()` time — enabling physical plan caching.
44pub struct LiveSourceProvider {
45    current: BatchSlot,
46    schema: SchemaRef,
47}
48
49impl LiveSourceProvider {
50    /// Creates a provider with the given schema and an empty batch slot.
51    #[must_use]
52    pub fn new(schema: SchemaRef) -> Self {
53        Self {
54            current: new_slot(),
55            schema,
56        }
57    }
58
59    /// Returns a handle for swapping batches into this provider.
60    #[must_use]
61    pub fn handle(&self) -> LiveSourceHandle {
62        LiveSourceHandle {
63            slot: Arc::clone(&self.current),
64        }
65    }
66}
67
68impl std::fmt::Debug for LiveSourceProvider {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        f.debug_struct("LiveSourceProvider")
71            .field("schema_fields", &self.schema.fields().len())
72            .finish_non_exhaustive()
73    }
74}
75
76#[async_trait]
77impl TableProvider for LiveSourceProvider {
78    fn as_any(&self) -> &dyn std::any::Any {
79        self
80    }
81
82    fn schema(&self) -> SchemaRef {
83        self.schema.clone()
84    }
85
86    fn table_type(&self) -> TableType {
87        TableType::Base
88    }
89
90    async fn scan(
91        &self,
92        _state: &dyn Session,
93        projection: Option<&Vec<usize>>,
94        _filters: &[Expr],
95        _limit: Option<usize>,
96    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
97        Ok(Arc::new(LiveSourceExec::new(
98            Arc::clone(&self.current),
99            self.schema.clone(),
100            projection.cloned(),
101        )))
102    }
103}
104
105// ── ExecutionPlan ────────────────────────────────────────────────────
106
107/// Leaf `ExecutionPlan` that reads from a shared batch slot at `execute()`
108/// time, not at construction time. This enables physical plan caching:
109/// the plan tree is built once, and each `execute()` call sees fresh data.
110pub(crate) struct LiveSourceExec {
111    slot: BatchSlot,
112    schema: SchemaRef,
113    projection: Option<Vec<usize>>,
114    properties: Arc<PlanProperties>,
115}
116
117impl LiveSourceExec {
118    fn new(slot: BatchSlot, source_schema: SchemaRef, projection: Option<Vec<usize>>) -> Self {
119        let schema = match &projection {
120            Some(indices) => {
121                let fields: Vec<_> = indices
122                    .iter()
123                    .map(|&i| source_schema.field(i).clone())
124                    .collect();
125                Arc::new(arrow::datatypes::Schema::new(fields))
126            }
127            None => source_schema,
128        };
129        let properties = Arc::new(PlanProperties::new(
130            EquivalenceProperties::new(schema.clone()),
131            Partitioning::UnknownPartitioning(1),
132            EmissionType::Final,
133            Boundedness::Bounded,
134        ));
135        Self {
136            slot,
137            schema,
138            projection,
139            properties,
140        }
141    }
142}
143
144impl std::fmt::Debug for LiveSourceExec {
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        f.debug_struct("LiveSourceExec")
147            .field("schema_fields", &self.schema.fields().len())
148            .finish_non_exhaustive()
149    }
150}
151
152impl DisplayAs for LiveSourceExec {
153    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154        match t {
155            DisplayFormatType::Default | DisplayFormatType::Verbose => {
156                write!(f, "LiveSourceExec: schema={}", self.schema.fields().len())
157            }
158            DisplayFormatType::TreeRender => write!(f, "LiveSourceExec"),
159        }
160    }
161}
162
163impl ExecutionPlan for LiveSourceExec {
164    fn as_any(&self) -> &dyn std::any::Any {
165        self
166    }
167
168    fn name(&self) -> &'static str {
169        "LiveSourceExec"
170    }
171
172    fn schema(&self) -> SchemaRef {
173        self.schema.clone()
174    }
175
176    fn properties(&self) -> &Arc<PlanProperties> {
177        &self.properties
178    }
179
180    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
181        vec![]
182    }
183
184    fn with_new_children(
185        self: Arc<Self>,
186        children: Vec<Arc<dyn ExecutionPlan>>,
187    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
188        if children.is_empty() {
189            Ok(self)
190        } else {
191            Err(DataFusionError::Plan(
192                "LiveSourceExec is a leaf node".to_string(),
193            ))
194        }
195    }
196
197    fn execute(
198        &self,
199        partition: usize,
200        _context: Arc<TaskContext>,
201    ) -> Result<SendableRecordBatchStream, DataFusionError> {
202        if partition != 0 {
203            return Err(DataFusionError::Plan(format!(
204                "LiveSourceExec only supports partition 0, got {partition}"
205            )));
206        }
207
208        // O(1) snapshot: bump the Arc under the lock and release.
209        // Any concurrent `swap()` installs a new Arc; this reader keeps
210        // its prior snapshot alive until the returned stream is dropped.
211        let batches_arc: Arc<Vec<RecordBatch>> = Arc::clone(&self.slot.lock());
212        let schema = self.schema.clone();
213        let projection = self.projection.clone();
214
215        let output = futures::stream::iter(if batches_arc.is_empty() {
216            vec![Ok(RecordBatch::new_empty(schema))]
217        } else if let Some(indices) = projection {
218            batches_arc
219                .iter()
220                .map(|batch| batch.project(&indices).map_err(DataFusionError::from))
221                .collect()
222        } else {
223            batches_arc.iter().cloned().map(Ok).collect()
224        });
225
226        Ok(Box::pin(RecordBatchStreamAdapter::new(
227            self.schema.clone(),
228            output,
229        )))
230    }
231}
232
233impl datafusion::physical_plan::ExecutionPlanProperties for LiveSourceExec {
234    fn output_partitioning(&self) -> &Partitioning {
235        self.properties.output_partitioning()
236    }
237
238    fn output_ordering(&self) -> Option<&datafusion::physical_expr::LexOrdering> {
239        self.properties.output_ordering()
240    }
241
242    fn boundedness(&self) -> Boundedness {
243        Boundedness::Bounded
244    }
245
246    fn pipeline_behavior(&self) -> EmissionType {
247        EmissionType::Final
248    }
249
250    fn equivalence_properties(&self) -> &EquivalenceProperties {
251        self.properties.equivalence_properties()
252    }
253}
254
255// ── Handle ───────────────────────────────────────────────────────────
256
257/// Handle for swapping batches into a [`LiveSourceProvider`].
258#[derive(Clone)]
259pub struct LiveSourceHandle {
260    slot: BatchSlot,
261}
262
263impl LiveSourceHandle {
264    /// Replace current batches. In-flight readers that captured the
265    /// prior snapshot continue to see the prior data; the next scan
266    /// picks up the new one.
267    pub fn swap(&self, batches: Vec<RecordBatch>) {
268        *self.slot.lock() = Arc::new(batches);
269    }
270
271    /// Clear all pending batches.
272    pub fn clear(&self) {
273        *self.slot.lock() = Arc::new(Vec::new());
274    }
275}
276
277impl std::fmt::Debug for LiveSourceHandle {
278    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
279        f.debug_struct("LiveSourceHandle").finish()
280    }
281}
282
283// ── Tests ────────────────────────────────────────────────────────────
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use arrow::array::{Float64Array, Int64Array, StringArray};
289    use arrow::datatypes::{DataType, Field, Schema};
290
291    fn test_schema() -> SchemaRef {
292        Arc::new(Schema::new(vec![
293            Field::new("id", DataType::Int64, false),
294            Field::new("name", DataType::Utf8, true),
295            Field::new("price", DataType::Float64, true),
296        ]))
297    }
298
299    fn make_batch(ids: &[i64], names: &[&str], prices: &[f64]) -> RecordBatch {
300        RecordBatch::try_new(
301            test_schema(),
302            vec![
303                Arc::new(Int64Array::from(ids.to_vec())),
304                Arc::new(StringArray::from(
305                    names.iter().map(|s| Some(*s)).collect::<Vec<_>>(),
306                )),
307                Arc::new(Float64Array::from(prices.to_vec())),
308            ],
309        )
310        .unwrap()
311    }
312
313    fn test_ctx() -> datafusion::prelude::SessionContext {
314        // Use a plain context (no streaming validator) for unit tests.
315        datafusion::prelude::SessionContext::new()
316    }
317
318    async fn count_rows(ctx: &datafusion::prelude::SessionContext, sql: &str) -> usize {
319        let df = ctx.sql(sql).await.unwrap();
320        df.collect()
321            .await
322            .unwrap()
323            .iter()
324            .map(RecordBatch::num_rows)
325            .sum()
326    }
327
328    #[test]
329    fn test_handle_swap_and_clear() {
330        let provider = LiveSourceProvider::new(test_schema());
331        let h1 = provider.handle();
332        let h2 = h1.clone();
333
334        h1.swap(vec![make_batch(&[1, 2], &["A", "B"], &[1.0, 2.0])]);
335        assert_eq!(h2.slot.lock().len(), 1);
336
337        h2.clear();
338        assert_eq!(h1.slot.lock().len(), 0);
339    }
340
341    #[tokio::test]
342    async fn test_scan_reads_fresh_data_each_execute() {
343        let provider = Arc::new(LiveSourceProvider::new(test_schema()));
344        let handle = provider.handle();
345        let ctx = test_ctx();
346        ctx.register_table("t", provider).unwrap();
347
348        handle.swap(vec![make_batch(
349            &[1, 2, 3],
350            &["A", "B", "C"],
351            &[10.0, 20.0, 30.0],
352        )]);
353        assert_eq!(count_rows(&ctx, "SELECT * FROM t").await, 3);
354        assert_eq!(count_rows(&ctx, "SELECT * FROM t").await, 3);
355    }
356
357    #[tokio::test]
358    async fn test_scan_empty() {
359        let provider = Arc::new(LiveSourceProvider::new(test_schema()));
360        let ctx = test_ctx();
361        ctx.register_table("t", provider).unwrap();
362        assert_eq!(count_rows(&ctx, "SELECT * FROM t").await, 0);
363    }
364
365    #[tokio::test]
366    async fn test_projection() {
367        let provider = Arc::new(LiveSourceProvider::new(test_schema()));
368        let handle = provider.handle();
369        let ctx = test_ctx();
370        ctx.register_table("t", provider).unwrap();
371
372        handle.swap(vec![make_batch(
373            &[1, 2, 3],
374            &["A", "B", "C"],
375            &[10.0, 20.0, 30.0],
376        )]);
377
378        let df = ctx.sql("SELECT id, price FROM t").await.unwrap();
379        let result = df.collect().await.unwrap();
380        assert_eq!(result.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
381        assert_eq!(result[0].schema().fields().len(), 2);
382        assert_eq!(result[0].schema().field(0).name(), "id");
383        assert_eq!(result[0].schema().field(1).name(), "price");
384    }
385
386    #[tokio::test]
387    async fn test_multi_cycle() {
388        let provider = Arc::new(LiveSourceProvider::new(test_schema()));
389        let handle = provider.handle();
390        let ctx = test_ctx();
391        ctx.register_table("t", provider).unwrap();
392
393        handle.swap(vec![make_batch(&[1], &["A"], &[10.0])]);
394        assert_eq!(count_rows(&ctx, "SELECT * FROM t").await, 1);
395
396        handle.swap(vec![make_batch(&[2, 3], &["B", "C"], &[20.0, 30.0])]);
397        assert_eq!(count_rows(&ctx, "SELECT * FROM t").await, 2);
398
399        handle.clear();
400        assert_eq!(count_rows(&ctx, "SELECT * FROM t").await, 0);
401    }
402
403    #[tokio::test]
404    async fn test_cached_plan_sees_fresh_data() {
405        use datafusion::physical_plan::ExecutionPlanProperties as _;
406
407        let provider = Arc::new(LiveSourceProvider::new(test_schema()));
408        let handle = provider.handle();
409        let ctx = test_ctx();
410        ctx.register_table("t", provider).unwrap();
411
412        handle.swap(vec![make_batch(&[1], &["A"], &[10.0])]);
413        let logical = ctx
414            .state()
415            .create_logical_plan("SELECT * FROM t")
416            .await
417            .unwrap();
418        let physical = ctx.state().create_physical_plan(&logical).await.unwrap();
419        assert_eq!(physical.output_partitioning().partition_count(), 1);
420
421        let task_ctx = ctx.task_ctx();
422        let r1 = datafusion::physical_plan::collect(physical.clone(), task_ctx.clone())
423            .await
424            .unwrap();
425        assert_eq!(r1.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
426
427        handle.swap(vec![make_batch(
428            &[2, 3, 4],
429            &["B", "C", "D"],
430            &[20.0, 30.0, 40.0],
431        )]);
432        let r2 = datafusion::physical_plan::collect(physical.clone(), task_ctx.clone())
433            .await
434            .unwrap();
435        assert_eq!(r2.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
436
437        handle.clear();
438        let r3 = datafusion::physical_plan::collect(physical, task_ctx)
439            .await
440            .unwrap();
441        assert_eq!(r3.iter().map(RecordBatch::num_rows).sum::<usize>(), 0);
442    }
443}
laminar_sql/datafusion/live_source.rs

laminar_sql/datafusion/
live_source.rs
