Skip to main content

nodedb_columnar/memtable/
iter.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Row-oriented iteration and single-row lookup over the memtable.
4
5use nodedb_types::columnar::ColumnDef;
6use nodedb_types::value::Value;
7
8use super::column_data::ColumnData;
9use super::core::ColumnarMemtable;
10use crate::error::ColumnarError;
11
12impl ColumnarMemtable {
13    /// Iterate rows as `Vec<Value>`. For scan/read operations.
14    ///
15    /// Every cell is typed by its column's declared type, so a time cell
16    /// comes back as the instant variant the schema declares.
17    pub fn iter_rows(&self) -> MemtableRowIter<'_> {
18        MemtableRowIter {
19            columns: &self.columns,
20            column_defs: &self.schema.columns,
21            row_count: self.row_count,
22            current: 0,
23        }
24    }
25
26    /// Get a single row by index as `Vec<Value>`.
27    ///
28    /// `Ok(None)` when `row_idx` is past the last row. `Err` when a cell of
29    /// the row is corrupt.
30    pub fn get_row(&self, row_idx: usize) -> Result<Option<Vec<Value>>, ColumnarError> {
31        if row_idx >= self.row_count {
32            return Ok(None);
33        }
34        read_row(&self.columns, &self.schema.columns, row_idx).map(Some)
35    }
36}
37
38/// Read row `row_idx` of `columns`, typing each cell by its column def.
39fn read_row(
40    columns: &[ColumnData],
41    column_defs: &[ColumnDef],
42    row_idx: usize,
43) -> Result<Vec<Value>, ColumnarError> {
44    columns
45        .iter()
46        .zip(column_defs)
47        .map(|(col, def)| col.get_value(row_idx, &def.column_type, &def.name))
48        .collect()
49}
50
51/// Row iterator over a columnar memtable.
52pub struct MemtableRowIter<'a> {
53    columns: &'a [ColumnData],
54    column_defs: &'a [ColumnDef],
55    row_count: usize,
56    current: usize,
57}
58
59/// Yields `Err` for a row with a corrupt cell. The iterator still advances
60/// past that row.
61impl Iterator for MemtableRowIter<'_> {
62    type Item = Result<Vec<Value>, ColumnarError>;
63
64    fn next(&mut self) -> Option<Self::Item> {
65        if self.current >= self.row_count {
66            return None;
67        }
68        let row = read_row(self.columns, self.column_defs, self.current);
69        self.current += 1;
70        Some(row)
71    }
72
73    fn size_hint(&self) -> (usize, Option<usize>) {
74        let remaining = self.row_count - self.current;
75        (remaining, Some(remaining))
76    }
77}
78
79impl ExactSizeIterator for MemtableRowIter<'_> {}
80
81#[cfg(test)]
82mod tests {
83    use nodedb_types::NdbDateTime;
84    use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};
85    use nodedb_types::value::Value;
86
87    use super::super::core::ColumnarMemtable;
88
89    const MICROS: i64 = 1_583_402_400_000_000;
90
91    /// A row read back from the memtable types each time cell by its declared
92    /// column: `Timestamp` is a naive instant, `Timestamptz` a UTC instant,
93    /// `SystemTimestamp` the integer stored.
94    #[test]
95    fn a_time_cell_reads_back_as_its_declared_type() {
96        let schema = ColumnarSchema::new(vec![
97            ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
98            ColumnDef::required("at", ColumnType::Timestamp),
99            ColumnDef::required("at_tz", ColumnType::Timestamptz),
100            ColumnDef::required("sys", ColumnType::SystemTimestamp),
101        ])
102        .expect("valid schema");
103        let mut mt = ColumnarMemtable::new(&schema);
104        let dt = NdbDateTime::from_micros(MICROS);
105        mt.append_row(&[
106            Value::Integer(1),
107            Value::NaiveDateTime(dt),
108            Value::DateTime(dt),
109            Value::Integer(7),
110        ])
111        .expect("append");
112
113        let expected = vec![
114            Value::Integer(1),
115            Value::NaiveDateTime(dt),
116            Value::DateTime(dt),
117            Value::Integer(7),
118        ];
119        assert_eq!(mt.get_row(0).expect("read"), Some(expected.clone()));
120        let rows: Vec<Vec<Value>> = mt.iter_rows().collect::<Result<_, _>>().expect("read");
121        assert_eq!(rows, vec![expected]);
122        assert_eq!(mt.get_row(1).expect("read"), None);
123    }
124}