nodedb_lite/query/
table_provider.rs

1//! DataFusion `TableProvider` backed by Loro CRDT documents.
2//!
3//! Provides full SQL query capability over Lite's document store.
4//! Documents are scanned from the in-memory Loro state, converted
5//! to Arrow RecordBatches, and fed into DataFusion's execution engine.
6
7use std::any::Any;
8use std::sync::{Arc, Mutex};
9
10use async_trait::async_trait;
11use datafusion::arrow::array::{RecordBatch, StringArray};
12use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
13use datafusion::catalog::{Session, TableProvider};
14use datafusion::error::DataFusionError;
15use datafusion::logical_expr::{Expr, TableType};
16use datafusion::physical_plan::ExecutionPlan;
17
18use crate::engine::crdt::CrdtEngine;
19
20/// A DataFusion `TableProvider` that reads documents from a Loro collection.
21///
22/// Each document becomes a row. All fields are stored as JSON strings
23/// in a schemaless layout: `(id TEXT, document TEXT)`. DataFusion's
24/// JSON functions can extract fields for WHERE/ORDER BY/GROUP BY.
25///
26/// For typed collections (with known fields), a richer Arrow schema
27/// is generated with proper column types.
28pub struct LiteTableProvider {
29    collection: String,
30    schema: SchemaRef,
31    crdt: Arc<Mutex<CrdtEngine>>,
32}
33
34impl std::fmt::Debug for LiteTableProvider {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("LiteTableProvider")
37            .field("collection", &self.collection)
38            .finish()
39    }
40}
41
42impl LiteTableProvider {
43    /// Create a table provider for a schemaless collection.
44    ///
45    /// Schema: `(id TEXT NOT NULL, document TEXT)` — all fields stored
46    /// as a JSON blob in the `document` column. DataFusion JSON functions
47    /// extract individual fields.
48    pub fn new(collection: String, crdt: Arc<Mutex<CrdtEngine>>) -> Self {
49        let schema = Arc::new(Schema::new(vec![
50            Field::new("id", DataType::Utf8, false),
51            Field::new("document", DataType::Utf8, true),
52        ]));
53        Self {
54            collection,
55            schema,
56            crdt,
57        }
58    }
59
60    /// Create a table provider with a known schema (typed collection).
61    pub fn with_schema(
62        collection: String,
63        schema: SchemaRef,
64        crdt: Arc<Mutex<CrdtEngine>>,
65    ) -> Self {
66        Self {
67            collection,
68            schema,
69            crdt,
70        }
71    }
72
73    /// Scan documents from the Loro collection into Arrow RecordBatches.
74    ///
75    /// `limit` is pushed down from DataFusion — if set, only reads that
76    /// many documents instead of the entire collection.
77    fn scan_to_batches(&self, limit: Option<usize>) -> Result<Vec<RecordBatch>, DataFusionError> {
78        let crdt = self
79            .crdt
80            .lock()
81            .map_err(|e| DataFusionError::Execution(format!("crdt lock: {e}")))?;
82
83        let mut ids = crdt.list_ids(&self.collection);
84        // Apply limit pushdown: don't load more documents than needed.
85        if let Some(n) = limit {
86            ids.truncate(n);
87        }
88        if ids.is_empty() {
89            let batch = RecordBatch::new_empty(self.schema.clone());
90            return Ok(vec![batch]);
91        }
92
93        // For schemaless: serialize each document as JSON.
94        if self.schema.fields().len() == 2
95            && self.schema.field(0).name() == "id"
96            && self.schema.field(1).name() == "document"
97        {
98            return self.scan_schemaless(&crdt, &ids);
99        }
100
101        // For typed: extract fields into typed columns.
102        self.scan_typed(&crdt, &ids)
103    }
104
105    /// Schemaless scan: (id, document_json) pairs.
106    fn scan_schemaless(
107        &self,
108        crdt: &CrdtEngine,
109        ids: &[String],
110    ) -> Result<Vec<RecordBatch>, DataFusionError> {
111        let mut id_values = Vec::with_capacity(ids.len());
112        let mut doc_values = Vec::with_capacity(ids.len());
113
114        for id in ids {
115            if let Some(loro_val) = crdt.read(&self.collection, id) {
116                let doc = crate::nodedb::convert::loro_value_to_document(id, &loro_val);
117                let json = serde_json::to_string(&doc.fields).unwrap_or_else(|e| {
118                    tracing::warn!(id = %id, error = %e, "JSON serialization failed for document");
119                    "{}".to_string()
120                });
121                id_values.push(id.clone());
122                doc_values.push(json);
123            }
124        }
125
126        let id_array = StringArray::from(id_values);
127        let doc_array = StringArray::from(doc_values);
128        let batch = RecordBatch::try_new(
129            self.schema.clone(),
130            vec![Arc::new(id_array), Arc::new(doc_array)],
131        )
132        .map_err(|e| DataFusionError::Execution(format!("build batch: {e}")))?;
133
134        Ok(vec![batch])
135    }
136
137    /// Typed scan: extract known fields into proper Arrow columns.
138    fn scan_typed(
139        &self,
140        crdt: &CrdtEngine,
141        ids: &[String],
142    ) -> Result<Vec<RecordBatch>, DataFusionError> {
143        use datafusion::arrow::array::{BooleanArray, Float64Array, Int64Array};
144
145        let field_count = self.schema.fields().len();
146        let mut columns: Vec<Vec<Option<nodedb_types::Value>>> =
147            vec![Vec::with_capacity(ids.len()); field_count];
148
149        for id in ids {
150            if let Some(loro_val) = crdt.read(&self.collection, id) {
151                let doc = crate::nodedb::convert::loro_value_to_document(id, &loro_val);
152                for (i, field) in self.schema.fields().iter().enumerate() {
153                    let val = if field.name() == "id" {
154                        Some(nodedb_types::Value::String(id.clone()))
155                    } else {
156                        doc.fields.get(field.name()).cloned()
157                    };
158                    columns[i].push(val);
159                }
160            }
161        }
162
163        // Build Arrow arrays from Value columns.
164        let mut arrow_columns: Vec<Arc<dyn datafusion::arrow::array::Array>> =
165            Vec::with_capacity(field_count);
166
167        for (i, field) in self.schema.fields().iter().enumerate() {
168            let col = &columns[i];
169            let array: Arc<dyn datafusion::arrow::array::Array> = match field.data_type() {
170                DataType::Utf8 => {
171                    let vals: Vec<Option<String>> = col
172                        .iter()
173                        .map(|v| match v {
174                            Some(nodedb_types::Value::String(s)) => Some(s.clone()),
175                            Some(other) => Some(format!("{other:?}")),
176                            None => None,
177                        })
178                        .collect();
179                    Arc::new(StringArray::from(vals))
180                }
181                DataType::Int64 => {
182                    let vals: Vec<Option<i64>> = col
183                        .iter()
184                        .map(|v| match v {
185                            Some(nodedb_types::Value::Integer(i)) => Some(*i),
186                            Some(nodedb_types::Value::Float(f)) => Some(*f as i64),
187                            _ => None,
188                        })
189                        .collect();
190                    Arc::new(Int64Array::from(vals))
191                }
192                DataType::Float64 => {
193                    let vals: Vec<Option<f64>> = col
194                        .iter()
195                        .map(|v| match v {
196                            Some(nodedb_types::Value::Float(f)) => Some(*f),
197                            Some(nodedb_types::Value::Integer(i)) => Some(*i as f64),
198                            _ => None,
199                        })
200                        .collect();
201                    Arc::new(Float64Array::from(vals))
202                }
203                DataType::Boolean => {
204                    let vals: Vec<Option<bool>> = col
205                        .iter()
206                        .map(|v| match v {
207                            Some(nodedb_types::Value::Bool(b)) => Some(*b),
208                            _ => None,
209                        })
210                        .collect();
211                    Arc::new(BooleanArray::from(vals))
212                }
213                _ => {
214                    // Fallback: serialize as string.
215                    let vals: Vec<Option<String>> = col
216                        .iter()
217                        .map(|v| v.as_ref().map(|v| format!("{v:?}")))
218                        .collect();
219                    Arc::new(StringArray::from(vals))
220                }
221            };
222            arrow_columns.push(array);
223        }
224
225        let batch = RecordBatch::try_new(self.schema.clone(), arrow_columns)
226            .map_err(|e| DataFusionError::Execution(format!("build typed batch: {e}")))?;
227
228        Ok(vec![batch])
229    }
230}
231
232#[async_trait]
233impl TableProvider for LiteTableProvider {
234    fn as_any(&self) -> &dyn Any {
235        self
236    }
237
238    fn schema(&self) -> SchemaRef {
239        Arc::clone(&self.schema)
240    }
241
242    fn table_type(&self) -> TableType {
243        TableType::Base
244    }
245
246    async fn scan(
247        &self,
248        state: &dyn Session,
249        projection: Option<&Vec<usize>>,
250        _filters: &[Expr],
251        limit: Option<usize>,
252    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
253        let batches = self.scan_to_batches(limit)?;
254        // Use MemTable to create a physical plan from in-memory batches.
255        let mem_table =
256            datafusion::datasource::MemTable::try_new(self.schema.clone(), vec![batches])?;
257        mem_table.scan(state, projection, &[], limit).await
258    }
259}
nodedb_lite/query/table_provider.rs

nodedb_lite/query/
table_provider.rs
