nodedb_lite/query/
columnar_provider.rs

1//! DataFusion `TableProvider` for columnar collections.
2//!
3//! Reads compressed segments from storage, decodes requested columns via
4//! SegmentReader with projection pushdown and delete bitmap masking,
5//! and returns Arrow RecordBatches to DataFusion.
6
7use std::any::Any;
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use datafusion::arrow::array::{
12    Array, BinaryArray, BooleanArray, Float64Array, Int64Array, RecordBatch, StringArray,
13    TimestampMicrosecondArray,
14};
15use datafusion::arrow::buffer::{BooleanBuffer, NullBuffer};
16use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
17use datafusion::catalog::{Session, TableProvider};
18use datafusion::error::DataFusionError;
19use datafusion::logical_expr::{Expr, TableType};
20use datafusion::physical_plan::ExecutionPlan;
21use nodedb_columnar::delete_bitmap::DeleteBitmap;
22use nodedb_columnar::reader::{DecodedColumn, SegmentReader};
23use nodedb_types::Namespace;
24use nodedb_types::columnar::{ColumnType, ColumnarSchema};
25
26use crate::storage::engine::StorageEngine;
27
28/// A DataFusion `TableProvider` that reads from columnar segments.
29///
30/// Supports column projection pushdown: only requested columns are decoded.
31/// Delete bitmaps are applied to mask deleted rows.
32pub struct ColumnarTableProvider<S: StorageEngine> {
33    collection: String,
34    arrow_schema: SchemaRef,
35    columnar_schema: ColumnarSchema,
36    storage: Arc<S>,
37    /// Segment IDs to read.
38    segment_ids: Vec<u32>,
39    /// Per-segment delete bitmaps (cloned from engine state).
40    delete_bitmaps: Vec<(u32, DeleteBitmap)>,
41}
42
43impl<S: StorageEngine> std::fmt::Debug for ColumnarTableProvider<S> {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("ColumnarTableProvider")
46            .field("collection", &self.collection)
47            .field("segments", &self.segment_ids.len())
48            .finish()
49    }
50}
51
52impl<S: StorageEngine> ColumnarTableProvider<S> {
53    /// Create a table provider for a columnar collection.
54    pub fn new(
55        collection: String,
56        schema: &ColumnarSchema,
57        storage: Arc<S>,
58        segment_ids: Vec<u32>,
59        delete_bitmaps: Vec<(u32, DeleteBitmap)>,
60    ) -> Self {
61        let arrow_schema = columnar_schema_to_arrow(schema);
62        Self {
63            collection,
64            arrow_schema,
65            columnar_schema: schema.clone(),
66            storage,
67            segment_ids,
68            delete_bitmaps,
69        }
70    }
71
72    /// Scan segments and build RecordBatches.
73    fn scan_to_batches(
74        &self,
75        projection: Option<&Vec<usize>>,
76        limit: Option<usize>,
77    ) -> Result<Vec<RecordBatch>, DataFusionError> {
78        let col_indices: Vec<usize> = match projection {
79            Some(proj) => proj.clone(),
80            None => (0..self.columnar_schema.columns.len()).collect(),
81        };
82
83        let projected_schema = if projection.is_some() {
84            Arc::new(
85                self.arrow_schema
86                    .project(&col_indices)
87                    .map_err(|e| DataFusionError::Execution(format!("project schema: {e}")))?,
88            )
89        } else {
90            self.arrow_schema.clone()
91        };
92
93        // Collect decoded columns across all segments.
94        let mut batches: Vec<RecordBatch> = Vec::new();
95        let mut total_rows = 0usize;
96
97        for &seg_id in &self.segment_ids {
98            let seg_key = format!("{}:seg:{}", self.collection, seg_id);
99            let seg_bytes = tokio::task::block_in_place(|| {
100                let handle = tokio::runtime::Handle::current();
101                handle.block_on(async {
102                    self.storage
103                        .get(Namespace::Columnar, seg_key.as_bytes())
104                        .await
105                })
106            })
107            .map_err(|e| DataFusionError::Execution(format!("storage read: {e}")))?;
108
109            let Some(seg_bytes) = seg_bytes else {
110                continue;
111            };
112
113            let reader = SegmentReader::open(&seg_bytes)
114                .map_err(|e| DataFusionError::Execution(format!("open segment: {e}")))?;
115
116            // Find delete bitmap for this segment.
117            let empty_bitmap = DeleteBitmap::new();
118            let bitmap = self
119                .delete_bitmaps
120                .iter()
121                .find(|(id, _)| *id == seg_id)
122                .map(|(_, bm)| bm)
123                .unwrap_or(&empty_bitmap);
124
125            // Read requested columns with delete masking (one batch per segment).
126            let seg_columns: Vec<DecodedColumn> = col_indices
127                .iter()
128                .map(|&idx| reader.read_column_with_deletes(idx, &[], bitmap))
129                .collect::<Result<Vec<_>, _>>()
130                .map_err(|e| DataFusionError::Execution(format!("read columns: {e}")))?;
131
132            let seg_row_count = reader.row_count() as usize;
133
134            // Convert decoded columns to Arrow arrays.
135            let mut arrow_arrays: Vec<Arc<dyn Array>> = Vec::with_capacity(col_indices.len());
136            for (i, decoded) in seg_columns.into_iter().enumerate() {
137                let col_schema_idx = col_indices[i];
138                let col_type = &self.columnar_schema.columns[col_schema_idx].column_type;
139                let arr = decoded_to_arrow(decoded, col_type, seg_row_count)?;
140                arrow_arrays.push(arr);
141            }
142
143            let batch = RecordBatch::try_new(projected_schema.clone(), arrow_arrays)
144                .map_err(|e| DataFusionError::Execution(format!("build batch: {e}")))?;
145
146            batches.push(batch);
147            total_rows += batches.last().map_or(0, |b| b.num_rows());
148            if let Some(lim) = limit
149                && total_rows >= lim
150            {
151                break;
152            }
153        }
154
155        if batches.is_empty() {
156            return Ok(vec![RecordBatch::new_empty(projected_schema)]);
157        }
158
159        Ok(batches)
160    }
161}
162
163#[async_trait]
164impl<S: StorageEngine> TableProvider for ColumnarTableProvider<S> {
165    fn as_any(&self) -> &dyn Any {
166        self
167    }
168
169    fn schema(&self) -> SchemaRef {
170        Arc::clone(&self.arrow_schema)
171    }
172
173    fn table_type(&self) -> TableType {
174        TableType::Base
175    }
176
177    async fn scan(
178        &self,
179        state: &dyn Session,
180        projection: Option<&Vec<usize>>,
181        _filters: &[Expr],
182        limit: Option<usize>,
183    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
184        let batches = self.scan_to_batches(projection, limit)?;
185        let schema = if let Some(proj) = projection {
186            Arc::new(self.arrow_schema.project(proj)?)
187        } else {
188            self.arrow_schema.clone()
189        };
190        let mem_table = datafusion::datasource::MemTable::try_new(schema, vec![batches])?;
191        mem_table.scan(state, None, &[], limit).await
192    }
193}
194
195/// Convert a ColumnarSchema to Arrow Schema.
196fn columnar_schema_to_arrow(schema: &ColumnarSchema) -> SchemaRef {
197    let fields: Vec<Field> = schema
198        .columns
199        .iter()
200        .map(|col| {
201            let dt = match &col.column_type {
202                ColumnType::Int64 => DataType::Int64,
203                ColumnType::Float64 => DataType::Float64,
204                ColumnType::String => DataType::Utf8,
205                ColumnType::Bool => DataType::Boolean,
206                ColumnType::Bytes | ColumnType::Geometry => DataType::Binary,
207                ColumnType::Timestamp => DataType::Timestamp(TimeUnit::Microsecond, None),
208                ColumnType::Decimal | ColumnType::Uuid => DataType::Utf8,
209                ColumnType::Vector(_) => DataType::Binary,
210            };
211            Field::new(&col.name, dt, col.nullable)
212        })
213        .collect();
214    Arc::new(Schema::new(fields))
215}
216
217/// Convert a DecodedColumn to an Arrow ArrayRef.
218fn decoded_to_arrow(
219    decoded: DecodedColumn,
220    col_type: &ColumnType,
221    _row_count: usize,
222) -> Result<Arc<dyn Array>, DataFusionError> {
223    match decoded {
224        DecodedColumn::Int64 { values, valid } => {
225            let null_buf = build_null_buffer(&valid);
226            let arr = Int64Array::new(values.into(), null_buf);
227            Ok(Arc::new(arr))
228        }
229        DecodedColumn::Float64 { values, valid } => {
230            let null_buf = build_null_buffer(&valid);
231            let arr = Float64Array::new(values.into(), null_buf);
232            Ok(Arc::new(arr))
233        }
234        DecodedColumn::Timestamp { values, valid } => {
235            let null_buf = build_null_buffer(&valid);
236            let arr = TimestampMicrosecondArray::new(values.into(), null_buf);
237            Ok(Arc::new(arr))
238        }
239        DecodedColumn::Bool { values, valid } => {
240            let null_buf = build_null_buffer(&valid);
241            let bool_buf = BooleanBuffer::from(values);
242            let arr = BooleanArray::new(bool_buf, null_buf);
243            Ok(Arc::new(arr))
244        }
245        DecodedColumn::Binary {
246            data,
247            offsets,
248            valid,
249        } => {
250            match col_type {
251                ColumnType::String | ColumnType::Uuid | ColumnType::Decimal => {
252                    // Build StringArray from offsets + data.
253                    let null_buf = build_null_buffer(&valid);
254                    let mut strs: Vec<Option<&str>> = Vec::with_capacity(valid.len());
255                    for (i, &is_valid) in valid.iter().enumerate() {
256                        if is_valid && i + 1 < offsets.len() {
257                            let start = offsets[i] as usize;
258                            let end = offsets[i + 1] as usize;
259                            strs.push(Some(std::str::from_utf8(&data[start..end]).unwrap_or("")));
260                        } else {
261                            strs.push(None);
262                        }
263                    }
264                    let arr = StringArray::from(strs);
265                    if null_buf.is_some() {
266                        let data = arr
267                            .into_data()
268                            .into_builder()
269                            .null_bit_buffer(null_buf.map(|nb| nb.into_inner().into_inner()))
270                            .build()
271                            .map_err(|e| {
272                                DataFusionError::Execution(format!("build string array: {e}"))
273                            })?;
274                        Ok(Arc::new(StringArray::from(data)))
275                    } else {
276                        Ok(Arc::new(arr))
277                    }
278                }
279                _ => {
280                    // Binary types (Bytes, Geometry, Vector).
281                    let mut blobs: Vec<Option<&[u8]>> = Vec::with_capacity(valid.len());
282                    for (i, &is_valid) in valid.iter().enumerate() {
283                        if is_valid && i + 1 < offsets.len() {
284                            let start = offsets[i] as usize;
285                            let end = offsets[i + 1] as usize;
286                            blobs.push(Some(&data[start..end]));
287                        } else {
288                            blobs.push(None);
289                        }
290                    }
291                    let arr = BinaryArray::from(blobs);
292                    Ok(Arc::new(arr))
293                }
294            }
295        }
296    }
297}
298
299/// Build an Arrow NullBuffer from a validity vector.
300fn build_null_buffer(valid: &[bool]) -> Option<NullBuffer> {
301    let null_count = valid.iter().filter(|&&v| !v).count();
302    if null_count == 0 {
303        return None;
304    }
305    Some(NullBuffer::new(BooleanBuffer::from(valid.to_vec())))
306}
nodedb_lite/query/columnar_provider.rs

nodedb_lite/query/
columnar_provider.rs
