nodedb_columnar/memtable/
mutation.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Row append, truncate, drain, and dictionary-encoding mutations.
4
5use nodedb_types::columnar::ColumnarSchema;
6use nodedb_types::value::Value;
7
8use crate::error::ColumnarError;
9
10use super::column_data::{ColumnData, DICT_ENCODE_MAX_CARDINALITY};
11use super::core::ColumnarMemtable;
12use super::ingest_value::IngestValue;
13
14impl ColumnarMemtable {
15    /// Append a row of values. Validates types and nullability.
16    ///
17    /// The append is all-or-nothing. On an error, every column is cut back
18    /// to `row_count`, so the memtable holds exactly the rows it held before
19    /// the call. A rolled-back push to a `DictEncoded` column can leave one
20    /// dictionary entry that no row references. Every id stays valid.
21    pub fn append_row(&mut self, values: &[Value]) -> Result<(), ColumnarError> {
22        if values.len() != self.schema.columns.len() {
23            return Err(ColumnarError::SchemaMismatch {
24                expected: self.schema.columns.len(),
25                got: values.len(),
26            });
27        }
28
29        let pushed = self
30            .schema
31            .columns
32            .iter()
33            .zip(values.iter())
34            .zip(self.columns.iter_mut())
35            .try_for_each(|((col_def, value), column)| {
36                if matches!(value, Value::Null) && !col_def.nullable {
37                    return Err(ColumnarError::NullViolation(col_def.name.clone()));
38                }
39                column.push(value, &col_def.name, &col_def.column_type)
40            });
41        if let Err(e) = pushed {
42            self.rollback_partial_row();
43            return Err(e);
44        }
45
46        self.row_count += 1;
47        debug_assert!(
48            self.columns.iter().all(|c| c.len() == self.row_count),
49            "column lengths must stay aligned with row_count"
50        );
51        Ok(())
52    }
53
54    /// Cut every column back to `row_count` after a failed row push.
55    ///
56    /// The columns before the failed one hold one extra value. The failed
57    /// column and the columns after it hold `row_count` values already.
58    fn rollback_partial_row(&mut self) {
59        let n = self.row_count;
60        for col in &mut self.columns {
61            col.truncate(n);
62        }
63        debug_assert!(
64            self.columns.iter().all(|c| c.len() == self.row_count),
65            "column lengths must stay aligned with row_count after rollback"
66        );
67    }
68
69    /// Convert low-cardinality `String` columns to `DictEncoded` in-place.
70    pub fn try_dict_encode_columns(&mut self, max_cardinality: u32) {
71        for col in &mut self.columns {
72            if let ColumnData::String { .. } = col
73                && let Some(encoded) = ColumnData::try_dict_encode(col, max_cardinality)
74            {
75                *col = encoded;
76            }
77        }
78    }
79
80    /// Truncate the memtable to `n` rows, discarding all rows beyond that point.
81    ///
82    /// Used by transaction rollback to restore the memtable to its pre-write state.
83    /// No-op if `n >= self.row_count`.
84    pub fn truncate_to(&mut self, n: usize) {
85        if n >= self.row_count {
86            return;
87        }
88        for col in &mut self.columns {
89            col.truncate(n);
90        }
91        self.row_count = n;
92        debug_assert!(
93            self.columns.iter().all(|c| c.len() == self.row_count),
94            "column lengths misaligned after truncate_to"
95        );
96    }
97
98    /// Drain the memtable: return all column data and reset to empty.
99    pub fn drain(&mut self) -> (ColumnarSchema, Vec<ColumnData>, usize) {
100        let columns = std::mem::replace(
101            &mut self.columns,
102            self.schema
103                .columns
104                .iter()
105                .map(|col| ColumnData::new(&col.column_type, col.nullable))
106                .collect(),
107        );
108        let row_count = self.row_count;
109        self.row_count = 0;
110        (self.schema.clone(), columns, row_count)
111    }
112
113    /// Drain with automatic dictionary encoding for low-cardinality String columns.
114    pub fn drain_optimized(&mut self) -> (ColumnarSchema, Vec<ColumnData>, usize) {
115        self.try_dict_encode_columns(DICT_ENCODE_MAX_CARDINALITY);
116        self.drain()
117    }
118
119    /// Zero-copy row ingest for timeseries and high-throughput paths.
120    ///
121    /// Accepts borrowed values via `IngestValue<'_>`, avoiding string cloning
122    /// for tag columns that are already interned in the `DictEncoded` dictionary.
123    /// The ingest is all-or-nothing, as in [`Self::append_row`].
124    pub fn ingest_row_refs(&mut self, values: &[IngestValue<'_>]) -> Result<(), ColumnarError> {
125        if values.len() != self.schema.columns.len() {
126            return Err(ColumnarError::SchemaMismatch {
127                expected: self.schema.columns.len(),
128                got: values.len(),
129            });
130        }
131
132        let pushed = self
133            .schema
134            .columns
135            .iter()
136            .zip(values.iter())
137            .zip(self.columns.iter_mut())
138            .try_for_each(|((col_def, value), column)| {
139                if matches!(value, IngestValue::Null) && !col_def.nullable {
140                    return Err(ColumnarError::NullViolation(col_def.name.clone()));
141                }
142                column.push_ref(value, &col_def.name)
143            });
144        if let Err(e) = pushed {
145            self.rollback_partial_row();
146            return Err(e);
147        }
148
149        self.row_count += 1;
150        Ok(())
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};
157
158    use super::*;
159
160    fn test_schema() -> ColumnarSchema {
161        ColumnarSchema::new(vec![
162            ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
163            ColumnDef::required("name", ColumnType::String),
164            ColumnDef::nullable("score", ColumnType::Float64),
165        ])
166        .expect("valid schema")
167    }
168
169    #[test]
170    fn append_and_count() {
171        let schema = test_schema();
172        let mut mt = ColumnarMemtable::new(&schema);
173
174        mt.append_row(&[
175            Value::Integer(1),
176            Value::String("Alice".into()),
177            Value::Float(0.75),
178        ])
179        .expect("append");
180
181        mt.append_row(&[Value::Integer(2), Value::String("Bob".into()), Value::Null])
182            .expect("append");
183
184        assert_eq!(mt.row_count(), 2);
185        assert!(!mt.is_empty());
186    }
187
188    #[test]
189    fn null_violation_rejected() {
190        let schema = test_schema();
191        let mut mt = ColumnarMemtable::new(&schema);
192
193        let err = mt
194            .append_row(&[Value::Null, Value::String("x".into()), Value::Null])
195            .unwrap_err();
196        assert!(matches!(err, ColumnarError::NullViolation(ref s) if s == "id"));
197    }
198
199    #[test]
200    fn null_violation_mid_row_leaves_columns_aligned() {
201        let schema = test_schema();
202        let mut mt = ColumnarMemtable::new(&schema);
203
204        let err = mt
205            .append_row(&[Value::Integer(1), Value::Null, Value::Null])
206            .unwrap_err();
207        assert!(matches!(err, ColumnarError::NullViolation(ref s) if s == "name"));
208        assert_eq!(mt.row_count(), 0);
209        assert!(mt.columns().iter().all(|c| c.len() == 0));
210
211        mt.append_row(&[Value::Integer(2), Value::String("b".into()), Value::Null])
212            .expect("append after rejected row");
213        assert_eq!(mt.row_count(), 1);
214        assert!(mt.columns().iter().all(|c| c.len() == 1));
215        assert_eq!(
216            mt.get_row(0).expect("read"),
217            Some(vec![
218                Value::Integer(2),
219                Value::String("b".into()),
220                Value::Null
221            ])
222        );
223    }
224
225    #[test]
226    fn type_mismatch_mid_row_leaves_columns_aligned() {
227        let schema = test_schema();
228        let mut mt = ColumnarMemtable::new(&schema);
229        mt.append_row(&[
230            Value::Integer(1),
231            Value::String("a".into()),
232            Value::Float(0.5),
233        ])
234        .expect("append");
235
236        let err = mt
237            .append_row(&[
238                Value::Integer(2),
239                Value::String("b".into()),
240                Value::Bool(true),
241            ])
242            .unwrap_err();
243        assert!(matches!(err, ColumnarError::TypeMismatch { ref column, .. } if column == "score"));
244        assert_eq!(mt.row_count(), 1);
245        assert!(mt.columns().iter().all(|c| c.len() == 1));
246
247        mt.append_row(&[
248            Value::Integer(3),
249            Value::String("c".into()),
250            Value::Float(0.25),
251        ])
252        .expect("append after rejected row");
253        assert_eq!(
254            mt.get_row(1).expect("read"),
255            Some(vec![
256                Value::Integer(3),
257                Value::String("c".into()),
258                Value::Float(0.25)
259            ])
260        );
261    }
262
263    #[test]
264    fn rejected_ref_ingest_leaves_columns_aligned() {
265        let schema = test_schema();
266        let mut mt = ColumnarMemtable::new(&schema);
267
268        let err = mt
269            .ingest_row_refs(&[
270                IngestValue::Int64(1),
271                IngestValue::Str("a"),
272                IngestValue::Bool(true),
273            ])
274            .unwrap_err();
275        assert!(matches!(err, ColumnarError::TypeMismatch { .. }));
276        assert_eq!(mt.row_count(), 0);
277        assert!(mt.columns().iter().all(|c| c.len() == 0));
278    }
279
280    #[test]
281    fn schema_mismatch_rejected() {
282        let schema = test_schema();
283        let mut mt = ColumnarMemtable::new(&schema);
284
285        let err = mt.append_row(&[Value::Integer(1)]).unwrap_err();
286        assert!(matches!(err, ColumnarError::SchemaMismatch { .. }));
287    }
288
289    #[test]
290    fn flush_threshold() {
291        let schema = test_schema();
292        let mut mt = ColumnarMemtable::with_threshold(&schema, 3);
293
294        for i in 0..2 {
295            mt.append_row(&[
296                Value::Integer(i),
297                Value::String(format!("u{i}")),
298                Value::Null,
299            ])
300            .expect("append");
301        }
302        assert!(!mt.should_flush());
303
304        mt.append_row(&[Value::Integer(2), Value::String("u2".into()), Value::Null])
305            .expect("append");
306        assert!(mt.should_flush());
307    }
308
309    #[test]
310    fn drain_resets() {
311        let schema = test_schema();
312        let mut mt = ColumnarMemtable::new(&schema);
313
314        mt.append_row(&[
315            Value::Integer(1),
316            Value::String("x".into()),
317            Value::Float(0.5),
318        ])
319        .expect("append");
320
321        let (_schema, columns, row_count) = mt.drain();
322        assert_eq!(row_count, 1);
323        assert_eq!(columns.len(), 3);
324        assert_eq!(mt.row_count(), 0);
325        assert!(mt.is_empty());
326
327        match &columns[0] {
328            ColumnData::Int64 { values, valid } => {
329                assert_eq!(values, &[1]);
330                assert!(valid.is_none());
331            }
332            _ => panic!("expected Int64"),
333        }
334        match &columns[1] {
335            ColumnData::String {
336                data,
337                offsets,
338                valid,
339            } => {
340                assert_eq!(std::str::from_utf8(data).unwrap(), "x");
341                assert_eq!(offsets, &[0, 1]);
342                assert!(valid.is_none());
343            }
344            _ => panic!("expected String"),
345        }
346    }
347
348    #[test]
349    fn all_types() {
350        let schema = ColumnarSchema::new(vec![
351            ColumnDef::required("i", ColumnType::Int64),
352            ColumnDef::required("f", ColumnType::Float64),
353            ColumnDef::required("b", ColumnType::Bool),
354            ColumnDef::required("ts", ColumnType::Timestamp),
355            ColumnDef::required("s", ColumnType::String),
356            ColumnDef::required("raw", ColumnType::Bytes),
357            ColumnDef::required("vec", ColumnType::Vector(3)),
358        ])
359        .expect("valid");
360
361        let mut mt = ColumnarMemtable::new(&schema);
362        mt.append_row(&[
363            Value::Integer(42),
364            Value::Float(0.25),
365            Value::Bool(true),
366            Value::Integer(1_700_000_000),
367            Value::String("hello".into()),
368            Value::Bytes(vec![0xDE, 0xAD]),
369            Value::Array(vec![
370                Value::Float(1.0),
371                Value::Float(2.0),
372                Value::Float(3.0),
373            ]),
374        ])
375        .expect("append all types");
376
377        assert_eq!(mt.row_count(), 1);
378    }
379
380    #[test]
381    fn dict_encode_low_cardinality() {
382        let schema = ColumnarSchema::new(vec![ColumnDef::required("qtype", ColumnType::String)])
383            .expect("valid");
384
385        let mut mt = ColumnarMemtable::new(&schema);
386        let qtypes = ["A", "B", "AAAA", "NS", "MX", "SOA", "CNAME", "PTR"];
387        for _ in 0..10 {
388            for &q in &qtypes {
389                mt.append_row(&[Value::String(q.into())]).expect("append");
390            }
391        }
392        assert_eq!(mt.row_count(), 80);
393
394        mt.try_dict_encode_columns(DICT_ENCODE_MAX_CARDINALITY);
395
396        let (_schema, columns, _row_count) = mt.drain();
397        match &columns[0] {
398            ColumnData::DictEncoded {
399                ids,
400                dictionary,
401                valid,
402                ..
403            } => {
404                assert_eq!(ids.len(), 80);
405                assert!(valid.is_none());
406                assert_eq!(dictionary.len(), 8);
407                for &id in ids {
408                    assert!((id as usize) < dictionary.len());
409                }
410                for (i, &q) in qtypes.iter().enumerate().take(8) {
411                    let expected_id = dictionary.iter().position(|s| s == q).expect("in dict");
412                    assert_eq!(ids[i], expected_id as u32);
413                }
414            }
415            _ => panic!("expected DictEncoded after try_dict_encode_columns"),
416        }
417    }
418
419    #[test]
420    fn dict_encode_exceeds_cardinality_stays_string() {
421        let schema = ColumnarSchema::new(vec![ColumnDef::required("name", ColumnType::String)])
422            .expect("valid");
423
424        let mut mt = ColumnarMemtable::new(&schema);
425        let max: u32 = 4;
426        for i in 0..=max {
427            mt.append_row(&[Value::String(format!("val_{i}"))])
428                .expect("append");
429        }
430
431        mt.try_dict_encode_columns(max);
432
433        let (_schema, columns, _row_count) = mt.drain();
434        assert!(matches!(columns[0], ColumnData::String { .. }));
435    }
436
437    #[test]
438    fn dict_encode_with_nulls() {
439        let schema = ColumnarSchema::new(vec![ColumnDef::nullable("tag", ColumnType::String)])
440            .expect("valid");
441
442        let mut mt = ColumnarMemtable::new(&schema);
443        mt.append_row(&[Value::String("foo".into())])
444            .expect("append");
445        mt.append_row(&[Value::Null]).expect("append null");
446        mt.append_row(&[Value::String("bar".into())])
447            .expect("append");
448        mt.append_row(&[Value::Null]).expect("append null");
449
450        mt.try_dict_encode_columns(DICT_ENCODE_MAX_CARDINALITY);
451
452        let (_schema, columns, _row_count) = mt.drain();
453        match &columns[0] {
454            ColumnData::DictEncoded {
455                ids,
456                valid,
457                dictionary,
458                ..
459            } => {
460                assert_eq!(ids.len(), 4);
461                let v = valid.as_ref().expect("nullable column has validity bitmap");
462                assert_eq!(v.len(), 4);
463                assert!(v[0]);
464                assert!(!v[1]);
465                assert!(v[2]);
466                assert!(!v[3]);
467                assert_eq!(dictionary.len(), 2);
468            }
469            _ => panic!("expected DictEncoded"),
470        }
471    }
472}
nodedb_columnar/memtable/mutation.rs

nodedb_columnar/memtable/
mutation.rs
