scirs2_io/
datafusion_provider.rs

1//! DataFusion-compatible table provider interface.
2//!
3//! Provides a pure-Rust table abstraction that mirrors the Apache DataFusion
4//! `TableProvider` trait without pulling in the `datafusion` or `arrow` crates.
5//!
6//! # Components
7//!
8//! - `DataType` / `ColumnDef` / `TableSchema` — schema definitions
9//! - `ColumnData` / `RecordBatch` — columnar in-memory batches
10//! - `Expr` / `BinaryOperator` / `LiteralValue` — filter expression trees
11//! - `TableProvider` trait — uniform scan interface
12//! - `MemTableProvider` — simple in-memory implementation
13
14use std::sync::Arc;
15
16use scirs2_core::ndarray::Array2;
17
18// ──────────────────────────────────────────────────────────────────────────────
19// Error type
20// ──────────────────────────────────────────────────────────────────────────────
21
22/// Errors returned from table provider operations.
23#[derive(Debug, thiserror::Error)]
24pub enum TableProviderError {
25    /// The requested column was not found in the schema.
26    #[error("Column not found: {0}")]
27    ColumnNotFound(std::string::String),
28    /// A type mismatch was encountered during a schema or value operation.
29    #[error("Type error: {0}")]
30    TypeError(std::string::String),
31    /// An error occurred during a table scan.
32    #[error("Scan error: {0}")]
33    ScanError(std::string::String),
34}
35
36// ──────────────────────────────────────────────────────────────────────────────
37// Schema types
38// ──────────────────────────────────────────────────────────────────────────────
39
40/// Supported column data types.
41#[derive(Debug, Clone, PartialEq)]
42pub enum DataType {
43    /// 32-bit signed integer.
44    Int32,
45    /// 64-bit signed integer.
46    Int64,
47    /// 32-bit IEEE 754 floating-point number.
48    Float32,
49    /// 64-bit IEEE 754 floating-point number.
50    Float64,
51    /// Boolean.
52    Boolean,
53    /// UTF-8 string.
54    Utf8,
55    /// Opaque binary data.
56    Binary,
57    /// Variable-length list.
58    List(Box<DataType>),
59}
60
61/// A named column descriptor.
62#[derive(Debug, Clone)]
63pub struct ColumnDef {
64    /// Column name.
65    pub name: std::string::String,
66    /// Column data type.
67    pub data_type: DataType,
68    /// Whether the column may contain null values.
69    pub nullable: bool,
70}
71
72/// An ordered collection of `ColumnDef` descriptors forming a table schema.
73#[derive(Debug, Clone)]
74pub struct TableSchema {
75    /// Ordered column descriptors.
76    pub columns: Vec<ColumnDef>,
77}
78
79impl TableSchema {
80    /// Create a new schema from a list of column definitions.
81    pub fn new(columns: Vec<ColumnDef>) -> Self {
82        Self { columns }
83    }
84
85    /// Find a column by name (case-sensitive).
86    pub fn find_column(&self, name: &str) -> Option<&ColumnDef> {
87        self.columns.iter().find(|c| c.name == name)
88    }
89
90    /// Return the zero-based index of a column by name, or `None` if not found.
91    pub fn field_index(&self, name: &str) -> Option<usize> {
92        self.columns.iter().position(|c| c.name == name)
93    }
94}
95
96// ──────────────────────────────────────────────────────────────────────────────
97// Columnar data
98// ──────────────────────────────────────────────────────────────────────────────
99
100/// The typed contents of one column inside a `RecordBatch`.
101#[derive(Debug, Clone)]
102pub enum ColumnData {
103    /// 32-bit signed integers (non-nullable).
104    Int32(Vec<i32>),
105    /// 64-bit signed integers (non-nullable).
106    Int64(Vec<i64>),
107    /// 32-bit floats (non-nullable).
108    Float32(Vec<f32>),
109    /// 64-bit floats (non-nullable).
110    Float64(Vec<f64>),
111    /// Booleans (non-nullable).
112    Boolean(Vec<bool>),
113    /// UTF-8 strings (non-nullable).
114    Utf8(Vec<std::string::String>),
115    /// All-null column of the given length.
116    Null(usize),
117}
118
119impl ColumnData {
120    /// Return the number of rows in this column.
121    pub fn len(&self) -> usize {
122        match self {
123            ColumnData::Int32(v) => v.len(),
124            ColumnData::Int64(v) => v.len(),
125            ColumnData::Float32(v) => v.len(),
126            ColumnData::Float64(v) => v.len(),
127            ColumnData::Boolean(v) => v.len(),
128            ColumnData::Utf8(v) => v.len(),
129            ColumnData::Null(n) => *n,
130        }
131    }
132
133    /// Return `true` when the column contains no rows.
134    pub fn is_empty(&self) -> bool {
135        self.len() == 0
136    }
137
138    /// Filter rows by a boolean mask of the same length.
139    pub fn filter_by_mask(&self, mask: &[bool]) -> ColumnData {
140        match self {
141            ColumnData::Int32(v) => ColumnData::Int32(
142                v.iter()
143                    .zip(mask)
144                    .filter_map(|(&val, &m)| if m { Some(val) } else { None })
145                    .collect(),
146            ),
147            ColumnData::Int64(v) => ColumnData::Int64(
148                v.iter()
149                    .zip(mask)
150                    .filter_map(|(&val, &m)| if m { Some(val) } else { None })
151                    .collect(),
152            ),
153            ColumnData::Float32(v) => ColumnData::Float32(
154                v.iter()
155                    .zip(mask)
156                    .filter_map(|(&val, &m)| if m { Some(val) } else { None })
157                    .collect(),
158            ),
159            ColumnData::Float64(v) => ColumnData::Float64(
160                v.iter()
161                    .zip(mask)
162                    .filter_map(|(&val, &m)| if m { Some(val) } else { None })
163                    .collect(),
164            ),
165            ColumnData::Boolean(v) => ColumnData::Boolean(
166                v.iter()
167                    .zip(mask)
168                    .filter_map(|(&val, &m)| if m { Some(val) } else { None })
169                    .collect(),
170            ),
171            ColumnData::Utf8(v) => ColumnData::Utf8(
172                v.iter()
173                    .zip(mask)
174                    .filter_map(|(val, &m)| if m { Some(val.clone()) } else { None })
175                    .collect(),
176            ),
177            ColumnData::Null(_) => {
178                let count = mask.iter().filter(|&&m| m).count();
179                ColumnData::Null(count)
180            }
181        }
182    }
183
184    /// Select rows by a vec of indices.
185    pub fn select_rows(&self, indices: &[usize]) -> ColumnData {
186        match self {
187            ColumnData::Int32(v) => {
188                ColumnData::Int32(indices.iter().filter_map(|&i| v.get(i).copied()).collect())
189            }
190            ColumnData::Int64(v) => {
191                ColumnData::Int64(indices.iter().filter_map(|&i| v.get(i).copied()).collect())
192            }
193            ColumnData::Float32(v) => {
194                ColumnData::Float32(indices.iter().filter_map(|&i| v.get(i).copied()).collect())
195            }
196            ColumnData::Float64(v) => {
197                ColumnData::Float64(indices.iter().filter_map(|&i| v.get(i).copied()).collect())
198            }
199            ColumnData::Boolean(v) => {
200                ColumnData::Boolean(indices.iter().filter_map(|&i| v.get(i).copied()).collect())
201            }
202            ColumnData::Utf8(v) => {
203                ColumnData::Utf8(indices.iter().filter_map(|&i| v.get(i).cloned()).collect())
204            }
205            ColumnData::Null(_) => ColumnData::Null(indices.len()),
206        }
207    }
208}
209
210/// A batch of rows stored in columnar format.
211#[derive(Debug, Clone)]
212pub struct RecordBatch {
213    /// Schema describing the columns.
214    pub schema: Arc<TableSchema>,
215    /// One `ColumnData` entry per column in the schema.
216    pub columns: Vec<ColumnData>,
217    /// Number of rows represented in this batch.
218    pub num_rows: usize,
219}
220
221impl RecordBatch {
222    /// Create a new `RecordBatch` from a schema and column data.
223    pub fn new(schema: Arc<TableSchema>, columns: Vec<ColumnData>) -> Self {
224        let num_rows = columns.first().map(|c| c.len()).unwrap_or(0);
225        Self {
226            schema,
227            columns,
228            num_rows,
229        }
230    }
231
232    /// Return a reference to the column at `index`.
233    pub fn column(&self, index: usize) -> Option<&ColumnData> {
234        self.columns.get(index)
235    }
236
237    /// Return a reference to the column with the given name.
238    pub fn column_by_name(&self, name: &str) -> Option<&ColumnData> {
239        self.schema
240            .field_index(name)
241            .and_then(|i| self.columns.get(i))
242    }
243}
244
245// ──────────────────────────────────────────────────────────────────────────────
246// Expression tree
247// ──────────────────────────────────────────────────────────────────────────────
248
249/// A scalar literal value used in filter expressions.
250#[derive(Debug, Clone)]
251pub enum LiteralValue {
252    /// 64-bit integer.
253    Int64(i64),
254    /// 64-bit float.
255    Float64(f64),
256    /// Boolean.
257    Boolean(bool),
258    /// UTF-8 string.
259    Utf8(std::string::String),
260    /// SQL NULL.
261    Null,
262}
263
264/// Binary arithmetic or comparison operator.
265#[derive(Debug, Clone, Copy, PartialEq, Eq)]
266pub enum BinaryOperator {
267    /// Equality (`=`).
268    Eq,
269    /// Inequality (`!=`).
270    NotEq,
271    /// Less-than (`<`).
272    Lt,
273    /// Less-than-or-equal (`<=`).
274    LtEq,
275    /// Greater-than (`>`).
276    Gt,
277    /// Greater-than-or-equal (`>=`).
278    GtEq,
279    /// Logical AND.
280    And,
281    /// Logical OR.
282    Or,
283    /// Arithmetic addition.
284    Plus,
285    /// Arithmetic subtraction.
286    Minus,
287    /// Arithmetic multiplication.
288    Multiply,
289    /// Arithmetic division.
290    Divide,
291}
292
293/// A filter or projection expression.
294#[derive(Debug, Clone)]
295pub enum Expr {
296    /// Reference to a column by name.
297    Column(std::string::String),
298    /// A scalar literal constant.
299    Literal(LiteralValue),
300    /// Binary operation on two sub-expressions.
301    BinaryOp {
302        /// Left-hand operand.
303        left: Box<Expr>,
304        /// Operator.
305        op: BinaryOperator,
306        /// Right-hand operand.
307        right: Box<Expr>,
308    },
309    /// Test whether a column value is SQL NULL.
310    IsNull(Box<Expr>),
311    /// Test whether a column value is not SQL NULL.
312    IsNotNull(Box<Expr>),
313    /// Logical negation.
314    Not(Box<Expr>),
315}
316
317// ──────────────────────────────────────────────────────────────────────────────
318// TableProvider trait
319// ──────────────────────────────────────────────────────────────────────────────
320
321/// Uniform table scan interface compatible with DataFusion's `TableProvider`.
322pub trait TableProvider: Send + Sync {
323    /// Return the table schema.
324    fn schema(&self) -> Arc<TableSchema>;
325
326    /// Scan the table.
327    ///
328    /// # Parameters
329    /// - `projection`: Optional list of column indices to return. `None` returns all columns.
330    /// - `filters`: Filter expressions to push down (best-effort; implementations may ignore).
331    /// - `limit`: Optional maximum number of rows to return.
332    fn scan(
333        &self,
334        projection: Option<&[usize]>,
335        filters: &[Expr],
336        limit: Option<usize>,
337    ) -> Result<Vec<RecordBatch>, TableProviderError>;
338}
339
340// ──────────────────────────────────────────────────────────────────────────────
341// In-memory table provider
342// ──────────────────────────────────────────────────────────────────────────────
343
344/// An in-memory table that stores data as a vector of `RecordBatch` slices.
345pub struct MemTableProvider {
346    schema: Arc<TableSchema>,
347    batches: Vec<RecordBatch>,
348}
349
350impl MemTableProvider {
351    /// Create a new provider from a schema and a pre-built list of batches.
352    pub fn new(schema: TableSchema, batches: Vec<RecordBatch>) -> Self {
353        Self {
354            schema: Arc::new(schema),
355            batches,
356        }
357    }
358
359    /// Construct a `MemTableProvider` from a 2-D `f64` matrix.
360    ///
361    /// Each column in `matrix` becomes a `Float64` column. `column_names` must
362    /// have exactly as many entries as there are columns in `matrix`.
363    pub fn from_f64_matrix(
364        matrix: &Array2<f64>,
365        column_names: &[&str],
366    ) -> Result<Self, TableProviderError> {
367        let ncols = matrix.ncols();
368        if column_names.len() != ncols {
369            return Err(TableProviderError::TypeError(format!(
370                "matrix has {ncols} columns but {} names were supplied",
371                column_names.len()
372            )));
373        }
374
375        let columns_def: Vec<ColumnDef> = column_names
376            .iter()
377            .map(|&name| ColumnDef {
378                name: name.to_string(),
379                data_type: DataType::Float64,
380                nullable: false,
381            })
382            .collect();
383        let schema = Arc::new(TableSchema::new(columns_def));
384
385        let columns: Vec<ColumnData> = (0..ncols)
386            .map(|col_idx| {
387                let col_vec: Vec<f64> = matrix.column(col_idx).iter().copied().collect();
388                ColumnData::Float64(col_vec)
389            })
390            .collect();
391
392        let num_rows = matrix.nrows();
393        let batch = RecordBatch {
394            schema: Arc::clone(&schema),
395            columns,
396            num_rows,
397        };
398
399        Ok(Self {
400            schema,
401            batches: vec![batch],
402        })
403    }
404}
405
406impl TableProvider for MemTableProvider {
407    fn schema(&self) -> Arc<TableSchema> {
408        Arc::clone(&self.schema)
409    }
410
411    fn scan(
412        &self,
413        projection: Option<&[usize]>,
414        _filters: &[Expr],
415        limit: Option<usize>,
416    ) -> Result<Vec<RecordBatch>, TableProviderError> {
417        let mut result_batches: Vec<RecordBatch> = Vec::new();
418        let mut rows_remaining = limit;
419
420        for batch in &self.batches {
421            // Determine how many rows to take from this batch.
422            let take_rows = match rows_remaining {
423                None => batch.num_rows,
424                Some(0) => break,
425                Some(rem) => rem.min(batch.num_rows),
426            };
427
428            let projected_schema: Arc<TableSchema>;
429            let projected_cols: Vec<ColumnData>;
430
431            match projection {
432                None => {
433                    // Return all columns, sliced to `take_rows`.
434                    projected_schema = Arc::clone(&batch.schema);
435                    projected_cols = batch
436                        .columns
437                        .iter()
438                        .map(|c| slice_column(c, 0, take_rows))
439                        .collect();
440                }
441                Some(indices) => {
442                    // Return only the projected columns.
443                    let proj_defs: Vec<ColumnDef> = indices
444                        .iter()
445                        .map(|&i| {
446                            batch.schema.columns.get(i).cloned().ok_or_else(|| {
447                                TableProviderError::ColumnNotFound(format!(
448                                    "projection index {i} out of range"
449                                ))
450                            })
451                        })
452                        .collect::<Result<Vec<_>, _>>()?;
453
454                    projected_schema = Arc::new(TableSchema::new(proj_defs));
455
456                    projected_cols = indices
457                        .iter()
458                        .map(|&i| {
459                            batch
460                                .columns
461                                .get(i)
462                                .map(|c| slice_column(c, 0, take_rows))
463                                .ok_or_else(|| {
464                                    TableProviderError::ColumnNotFound(format!(
465                                        "projection index {i} out of range"
466                                    ))
467                                })
468                        })
469                        .collect::<Result<Vec<_>, _>>()?;
470                }
471            }
472
473            result_batches.push(RecordBatch {
474                schema: projected_schema,
475                columns: projected_cols,
476                num_rows: take_rows,
477            });
478
479            if let Some(ref mut rem) = rows_remaining {
480                *rem -= take_rows;
481            }
482        }
483
484        Ok(result_batches)
485    }
486}
487
488// ──────────────────────────────────────────────────────────────────────────────
489// Helpers
490// ──────────────────────────────────────────────────────────────────────────────
491
492/// Slice `col` to `[offset, offset + len)`.
493pub(crate) fn slice_column(col: &ColumnData, offset: usize, len: usize) -> ColumnData {
494    let end = (offset + len).min(col.len());
495    match col {
496        ColumnData::Int32(v) => ColumnData::Int32(v[offset..end].to_vec()),
497        ColumnData::Int64(v) => ColumnData::Int64(v[offset..end].to_vec()),
498        ColumnData::Float32(v) => ColumnData::Float32(v[offset..end].to_vec()),
499        ColumnData::Float64(v) => ColumnData::Float64(v[offset..end].to_vec()),
500        ColumnData::Boolean(v) => ColumnData::Boolean(v[offset..end].to_vec()),
501        ColumnData::Utf8(v) => ColumnData::Utf8(v[offset..end].to_vec()),
502        ColumnData::Null(n) => ColumnData::Null((end - offset).min(*n)),
503    }
504}
505
506// ──────────────────────────────────────────────────────────────────────────────
507// Tests
508// ──────────────────────────────────────────────────────────────────────────────
509
510#[cfg(test)]
511mod tests {
512    use super::*;
513    use scirs2_core::ndarray::array;
514
515    fn make_batch() -> RecordBatch {
516        let schema = Arc::new(TableSchema::new(vec![
517            ColumnDef {
518                name: "id".to_string(),
519                data_type: DataType::Int32,
520                nullable: false,
521            },
522            ColumnDef {
523                name: "score".to_string(),
524                data_type: DataType::Float64,
525                nullable: false,
526            },
527            ColumnDef {
528                name: "label".to_string(),
529                data_type: DataType::Utf8,
530                nullable: true,
531            },
532        ]));
533        let columns = vec![
534            ColumnData::Int32(vec![1, 2, 3, 4, 5]),
535            ColumnData::Float64(vec![1.1, 2.2, 3.3, 4.4, 5.5]),
536            ColumnData::Utf8(vec![
537                "a".to_string(),
538                "b".to_string(),
539                "c".to_string(),
540                "d".to_string(),
541                "e".to_string(),
542            ]),
543        ];
544        RecordBatch::new(schema, columns)
545    }
546
547    #[test]
548    fn test_mem_table_scan_all() {
549        let batch = make_batch();
550        let schema = (*batch.schema).clone();
551        let provider = MemTableProvider::new(schema, vec![batch]);
552
553        let result = provider.scan(None, &[], None).expect("scan failed");
554        assert_eq!(result.len(), 1);
555        assert_eq!(result[0].num_rows, 5);
556        assert_eq!(result[0].columns.len(), 3);
557    }
558
559    #[test]
560    fn test_mem_table_projection() {
561        let batch = make_batch();
562        let schema = (*batch.schema).clone();
563        let provider = MemTableProvider::new(schema, vec![batch]);
564
565        // Project only column 0 (id) and column 2 (label).
566        let result = provider
567            .scan(Some(&[0, 2]), &[], None)
568            .expect("scan failed");
569        assert_eq!(result.len(), 1);
570        let rb = &result[0];
571        assert_eq!(rb.columns.len(), 2);
572        assert_eq!(rb.schema.columns[0].name, "id");
573        assert_eq!(rb.schema.columns[1].name, "label");
574    }
575
576    #[test]
577    fn test_mem_table_from_matrix() {
578        let mat = array![[1.0_f64, 2.0], [3.0, 4.0], [5.0, 6.0]];
579        let provider =
580            MemTableProvider::from_f64_matrix(&mat, &["x", "y"]).expect("from_f64_matrix failed");
581
582        let result = provider.scan(None, &[], None).expect("scan failed");
583        assert_eq!(result.len(), 1);
584        assert_eq!(result[0].num_rows, 3);
585
586        if let ColumnData::Float64(vals) = &result[0].columns[0] {
587            assert!((vals[0] - 1.0).abs() < 1e-12);
588            assert!((vals[2] - 5.0).abs() < 1e-12);
589        } else {
590            panic!("Expected Float64 column");
591        }
592    }
593
594    #[test]
595    fn test_table_schema_find() {
596        let batch = make_batch();
597        let schema = (*batch.schema).clone();
598
599        let col = schema.find_column("score");
600        assert!(col.is_some());
601        assert_eq!(col.unwrap().data_type, DataType::Float64);
602
603        let missing = schema.find_column("nonexistent");
604        assert!(missing.is_none());
605
606        let idx = schema.field_index("label");
607        assert_eq!(idx, Some(2));
608    }
609}
scirs2_io/datafusion_provider.rs

scirs2_io/
datafusion_provider.rs
