nodedb_lite/query/
strict_provider.rs

1//! DataFusion `TableProvider` for strict document collections.
2//!
3//! Reads Binary Tuples from the StrictEngine, extracts columns into Arrow
4//! arrays via `nodedb-strict`'s vectorized extraction, and feeds them to
5//! DataFusion as RecordBatches with projection pushdown.
6
7use std::any::Any;
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use datafusion::arrow::array::RecordBatch;
12use datafusion::arrow::datatypes::SchemaRef;
13use datafusion::catalog::{Session, TableProvider};
14use datafusion::error::DataFusionError;
15use datafusion::logical_expr::{Expr, TableType};
16use datafusion::physical_plan::ExecutionPlan;
17use nodedb_strict::TupleDecoder;
18use nodedb_strict::arrow_extract::extract_column_to_arrow;
19use nodedb_types::Namespace;
20use nodedb_types::columnar::StrictSchema;
21
22use crate::engine::strict::strict_schema_to_arrow;
23use crate::storage::engine::StorageEngine;
24
25/// A DataFusion `TableProvider` that reads from a strict document collection.
26///
27/// Supports column projection pushdown: only the requested columns are
28/// decoded from the Binary Tuples. Unneeded columns are never touched.
29pub struct StrictTableProvider<S: StorageEngine> {
30    collection: String,
31    arrow_schema: SchemaRef,
32    strict_schema: StrictSchema,
33    decoder: TupleDecoder,
34    storage: Arc<S>,
35}
36
37impl<S: StorageEngine> std::fmt::Debug for StrictTableProvider<S> {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("StrictTableProvider")
40            .field("collection", &self.collection)
41            .finish()
42    }
43}
44
45impl<S: StorageEngine> StrictTableProvider<S> {
46    /// Create a table provider for a strict collection.
47    pub fn new(collection: String, schema: &StrictSchema, storage: Arc<S>) -> Self {
48        let arrow_schema = strict_schema_to_arrow(schema);
49        let decoder = TupleDecoder::new(schema);
50        Self {
51            collection,
52            arrow_schema,
53            strict_schema: schema.clone(),
54            decoder,
55            storage,
56        }
57    }
58
59    /// Scan tuples and build Arrow RecordBatches.
60    ///
61    /// `projection` specifies which column indices to decode. If `None`, all
62    /// columns are decoded. `limit` caps the number of rows.
63    fn scan_to_batches(
64        &self,
65        projection: Option<&Vec<usize>>,
66        limit: Option<usize>,
67    ) -> Result<Vec<RecordBatch>, DataFusionError> {
68        // Read raw tuples from storage.
69        let prefix = format!("{}:", self.collection);
70
71        // StorageEngine is async, but scan_to_batches is sync (called from
72        // DataFusion's async scan). We use tokio::runtime::Handle to block.
73        let tuples = tokio::task::block_in_place(|| {
74            let handle = tokio::runtime::Handle::current();
75            handle.block_on(async {
76                self.storage
77                    .scan_prefix(Namespace::Strict, prefix.as_bytes())
78                    .await
79            })
80        })
81        .map_err(|e| DataFusionError::Execution(format!("storage scan: {e}")))?;
82
83        // Apply limit.
84        let tuple_bytes: Vec<Vec<u8>> = if let Some(n) = limit {
85            tuples.into_iter().take(n).map(|(_, v)| v).collect()
86        } else {
87            tuples.into_iter().map(|(_, v)| v).collect()
88        };
89
90        if tuple_bytes.is_empty() {
91            let batch = RecordBatch::new_empty(self.arrow_schema.clone());
92            return Ok(vec![batch]);
93        }
94
95        let refs: Vec<&[u8]> = tuple_bytes.iter().map(|t| t.as_slice()).collect();
96
97        // Determine which columns to extract.
98        let col_indices: Vec<usize> = match projection {
99            Some(proj) => proj.to_vec(),
100            None => (0..self.strict_schema.columns.len()).collect(),
101        };
102
103        // Build the projected Arrow schema.
104        let projected_schema = if projection.is_some() {
105            Arc::new(
106                self.arrow_schema
107                    .project(&col_indices)
108                    .map_err(|e| DataFusionError::Execution(format!("schema projection: {e}")))?,
109            )
110        } else {
111            self.arrow_schema.clone()
112        };
113
114        // Extract each column into an Arrow array.
115        let mut arrays = Vec::with_capacity(col_indices.len());
116        for &idx in &col_indices {
117            let arr = extract_column_to_arrow(&self.strict_schema, &self.decoder, &refs, idx)
118                .map_err(|e| DataFusionError::Execution(format!("extract column: {e}")))?;
119            arrays.push(arr);
120        }
121
122        let batch = RecordBatch::try_new(projected_schema, arrays)
123            .map_err(|e| DataFusionError::Execution(format!("build batch: {e}")))?;
124
125        Ok(vec![batch])
126    }
127}
128
129#[async_trait]
130impl<S: StorageEngine> TableProvider for StrictTableProvider<S> {
131    fn as_any(&self) -> &dyn Any {
132        self
133    }
134
135    fn schema(&self) -> SchemaRef {
136        Arc::clone(&self.arrow_schema)
137    }
138
139    fn table_type(&self) -> TableType {
140        TableType::Base
141    }
142
143    async fn scan(
144        &self,
145        state: &dyn Session,
146        projection: Option<&Vec<usize>>,
147        _filters: &[Expr],
148        limit: Option<usize>,
149    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
150        let batches = self.scan_to_batches(projection, limit)?;
151        let schema = if let Some(proj) = projection {
152            Arc::new(self.arrow_schema.project(proj)?)
153        } else {
154            self.arrow_schema.clone()
155        };
156        let mem_table = datafusion::datasource::MemTable::try_new(schema, vec![batches])?;
157        mem_table.scan(state, None, &[], limit).await
158    }
159}
nodedb_lite/query/strict_provider.rs

nodedb_lite/query/
strict_provider.rs
