nodedb_columnar/mutation/
truncate.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Whole-collection `TRUNCATE` on a `MutationEngine`, reversible for a
4//! transaction batch that aborts after it.
5
6use std::collections::HashMap;
7
8use nodedb_types::surrogate::Surrogate;
9
10use crate::delete_bitmap::DeleteBitmap;
11use crate::memtable::ColumnarMemtable;
12use crate::pk_index::PkIndex;
13
14use super::engine::MutationEngine;
15
16/// Every row-bearing part of a `MutationEngine`, taken out by
17/// [`MutationEngine::truncate`] and put back by
18/// [`MutationEngine::restore_truncated`]. Holds every bitemporal version and
19/// every tombstone exactly as they were, so a restore is exact.
20pub struct TruncatedRows {
21    memtable: ColumnarMemtable,
22    pk_index: PkIndex,
23    delete_bitmaps: HashMap<u64, DeleteBitmap>,
24    memtable_row_counter: u32,
25    memtable_surrogates: Vec<Option<Surrogate>>,
26}
27
28impl MutationEngine {
29    /// Number of live rows: one per bound primary key. Historical bitemporal
30    /// versions and tombstoned rows are not counted.
31    pub fn live_row_count(&self) -> usize {
32        self.pk_index.len()
33    }
34
35    /// Remove every row, every bitemporal version, and every tombstone, and
36    /// hand the pre-image back. Segment ids stay monotonic: the caller drops
37    /// its flushed segment bytes for the same collection, and the next flush
38    /// takes the next unused id rather than reusing one a stale reader could
39    /// still name.
40    pub fn truncate(&mut self) -> TruncatedRows {
41        let threshold = self.memtable.flush_threshold();
42        let memtable = std::mem::replace(
43            &mut self.memtable,
44            ColumnarMemtable::with_threshold(&self.schema, threshold),
45        );
46        TruncatedRows {
47            memtable,
48            pk_index: std::mem::take(&mut self.pk_index),
49            delete_bitmaps: std::mem::take(&mut self.delete_bitmaps),
50            memtable_row_counter: std::mem::replace(&mut self.memtable_row_counter, 0),
51            memtable_surrogates: std::mem::take(&mut self.memtable_surrogates),
52        }
53    }
54
55    /// Put back what [`Self::truncate`] took out. Used exclusively by the
56    /// transaction undo log; never called on the normal write path.
57    pub fn restore_truncated(&mut self, rows: TruncatedRows) {
58        let TruncatedRows {
59            memtable,
60            pk_index,
61            delete_bitmaps,
62            memtable_row_counter,
63            memtable_surrogates,
64        } = rows;
65        self.memtable = memtable;
66        self.pk_index = pk_index;
67        self.delete_bitmaps = delete_bitmaps;
68        self.memtable_row_counter = memtable_row_counter;
69        self.memtable_surrogates = memtable_surrogates;
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};
76    use nodedb_types::value::Value;
77
78    use super::*;
79
80    fn schema() -> ColumnarSchema {
81        ColumnarSchema::new(vec![
82            ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
83            ColumnDef::required("name", ColumnType::String),
84        ])
85        .expect("valid")
86    }
87
88    fn seeded() -> MutationEngine {
89        let mut engine = MutationEngine::with_flush_threshold("t".into(), schema(), 7);
90        for i in 0..4 {
91            engine
92                .insert_with_surrogate(
93                    &[Value::Integer(i), Value::String(format!("r{i}"))],
94                    Surrogate(100 + i as u32),
95                )
96                .expect("insert");
97        }
98        engine.delete(&Value::Integer(1)).expect("delete");
99        engine
100    }
101
102    #[test]
103    fn truncate_empties_every_row_bearing_part_and_keeps_the_threshold() {
104        let mut engine = seeded();
105        assert_eq!(engine.live_row_count(), 3);
106
107        let _pre = engine.truncate();
108
109        assert_eq!(engine.live_row_count(), 0);
110        assert_eq!(engine.memtable().row_count(), 0);
111        assert!(engine.delete_bitmaps().is_empty());
112        assert!(engine.memtable_surrogates().is_empty());
113        assert_eq!(engine.scan_memtable_rows().count(), 0);
114        assert_eq!(engine.memtable().flush_threshold(), 7);
115        assert_eq!(engine.next_segment_id(), 1);
116
117        engine
118            .insert(&[Value::Integer(9), Value::String("z".into())])
119            .expect("insert after truncate");
120        assert_eq!(engine.live_row_count(), 1);
121        assert_eq!(engine.scan_memtable_rows().count(), 1);
122    }
123
124    #[test]
125    fn restore_truncated_brings_back_rows_and_tombstones_exactly() {
126        let mut engine = seeded();
127        let before: Vec<Vec<Value>> = engine
128            .scan_memtable_rows()
129            .collect::<Result<_, _>>()
130            .expect("read");
131
132        let pre = engine.truncate();
133        engine
134            .insert(&[Value::Integer(9), Value::String("z".into())])
135            .expect("insert after truncate");
136        engine.restore_truncated(pre);
137
138        let after: Vec<Vec<Value>> = engine
139            .scan_memtable_rows()
140            .collect::<Result<_, _>>()
141            .expect("read");
142        assert_eq!(
143            after, before,
144            "restore must reproduce the pre-truncate rows"
145        );
146        assert_eq!(engine.live_row_count(), 3);
147        assert!(
148            engine
149                .pk_index()
150                .get(&crate::pk_index::encode_pk(&Value::Integer(1)))
151                .is_none(),
152            "the tombstone deleted before the truncate stays deleted"
153        );
154        assert_eq!(engine.memtable_surrogates().len(), 4);
155    }
156
157    #[test]
158    fn truncate_keeps_segment_ids_monotonic_after_a_flush() {
159        let mut engine = seeded();
160        engine.on_memtable_flushed(1).expect("flush");
161        assert_eq!(engine.next_segment_id(), 2);
162        let _pre = engine.truncate();
163        assert_eq!(engine.next_segment_id(), 2);
164    }
165}
nodedb_columnar/mutation/truncate.rs

nodedb_columnar/mutation/
truncate.rs
