Skip to main content

nodedb_lite/engine/array/
memtable.rs

1//! In-memory write buffer for one array — thin re-implementation of the
2//! Origin memtable adapted for Lite's sync context.
3//!
4//! Cells are bucketed by `TileId` (`hilbert_prefix`, `system_from_ms`).
5//! Each bucket is a `HashMap<encoded_coord, payload_bytes>`. Flush drains
6//! the buckets through `SparseTileBuilder` to produce `SparseTile`s ready
7//! for `SegmentWriter`.
8
9use std::collections::{BTreeMap, HashMap};
10
11use nodedb_array::schema::ArraySchema;
12use nodedb_array::tile::cell_payload::{
13    CELL_GDPR_ERASURE_SENTINEL, CELL_TOMBSTONE_SENTINEL, CellPayload, is_cell_gdpr_erasure,
14    is_cell_sentinel, is_cell_tombstone,
15};
16use nodedb_array::tile::sparse_tile::{RowKind, SparseRow, SparseTile, SparseTileBuilder};
17use nodedb_array::tile::tile_id_for_cell;
18use nodedb_array::types::TileId;
19use nodedb_array::types::cell_value::value::CellValue;
20use nodedb_array::types::coord::value::CoordValue;
21use nodedb_types::{OPEN_UPPER, Surrogate};
22
23use crate::error::LiteError;
24
25/// The zerompk coordinate key. The memtable and the surrogate binder share it.
26pub(crate) fn encode_coord(coord: &[CoordValue]) -> Result<Vec<u8>, LiteError> {
27    zerompk::to_msgpack_vec(&coord.to_vec()).map_err(|e| LiteError::Serialization {
28        detail: format!("encode coord: {e}"),
29    })
30}
31
32fn decode_coord(key: &[u8]) -> Result<Vec<CoordValue>, LiteError> {
33    zerompk::from_msgpack(key).map_err(|e| LiteError::Serialization {
34        detail: format!("decode coord: {e}"),
35    })
36}
37
38/// Raw write buffer for a single tile version.
39#[derive(Debug, Default, Clone)]
40struct TileBuffer {
41    entries: HashMap<Vec<u8>, Vec<u8>>,
42}
43
44impl TileBuffer {
45    fn entry_count(&self) -> usize {
46        self.entries.len()
47    }
48
49    fn live_cell_count(&self) -> usize {
50        self.entries
51            .values()
52            .filter(|v| !is_cell_sentinel(v))
53            .count()
54    }
55
56    fn get(&self, coord: &[CoordValue]) -> Option<&[u8]> {
57        let key = encode_coord(coord).ok()?;
58        self.entries.get(&key).map(|v| v.as_slice())
59    }
60
61    fn materialise(&self, schema: &ArraySchema) -> Result<SparseTile, LiteError> {
62        let mut b = SparseTileBuilder::new(schema);
63        for (coord_key, bytes) in &self.entries {
64            let coord = decode_coord(coord_key)?;
65            if is_cell_tombstone(bytes) {
66                b.push_row(SparseRow {
67                    coord: &coord,
68                    attrs: &[],
69                    surrogate: None,
70                    valid_from_ms: 0,
71                    valid_until_ms: OPEN_UPPER,
72                    kind: RowKind::Tombstone,
73                })
74                .map_err(|e| LiteError::Storage {
75                    detail: format!("push tombstone row: {e}"),
76                })?;
77            } else if is_cell_gdpr_erasure(bytes) {
78                b.push_row(SparseRow {
79                    coord: &coord,
80                    attrs: &[],
81                    surrogate: None,
82                    valid_from_ms: 0,
83                    valid_until_ms: OPEN_UPPER,
84                    kind: RowKind::GdprErased,
85                })
86                .map_err(|e| LiteError::Storage {
87                    detail: format!("push erasure row: {e}"),
88                })?;
89            } else {
90                let payload = CellPayload::decode(bytes).map_err(|e| LiteError::Storage {
91                    detail: format!("decode CellPayload: {e}"),
92                })?;
93                b.push_row(SparseRow {
94                    coord: &coord,
95                    attrs: &payload.attrs,
96                    surrogate: Some(payload.surrogate),
97                    valid_from_ms: payload.valid_from_ms,
98                    valid_until_ms: payload.valid_until_ms,
99                    kind: RowKind::Live,
100                })
101                .map_err(|e| LiteError::Storage {
102                    detail: format!("push live row: {e}"),
103                })?;
104            }
105        }
106        Ok(b.build())
107    }
108}
109
110/// Statistics about the current memtable state.
111#[derive(Debug, Default)]
112pub struct MemtableStats {
113    pub tile_count: usize,
114    pub cell_count: usize,
115}
116
117/// Arguments for `ArrayMemtable::put_cell`.
118pub struct PutCellArgs {
119    pub coord: Vec<CoordValue>,
120    pub attrs: Vec<CellValue>,
121    pub surrogate: Surrogate,
122    pub system_from_ms: i64,
123    pub valid_from_ms: i64,
124    pub valid_until_ms: i64,
125}
126
127/// In-memory write buffer for a single array.
128#[derive(Debug, Default)]
129pub struct ArrayMemtable {
130    tiles: BTreeMap<TileId, TileBuffer>,
131}
132
133impl ArrayMemtable {
134    pub fn new() -> Self {
135        Self::default()
136    }
137
138    /// Insert a cell. Creates a new tile-version for the given `system_from_ms`.
139    pub fn put_cell(
140        &mut self,
141        schema: &ArraySchema,
142        args: PutCellArgs,
143    ) -> Result<TileId, LiteError> {
144        let PutCellArgs {
145            coord,
146            attrs,
147            surrogate,
148            system_from_ms,
149            valid_from_ms,
150            valid_until_ms,
151        } = args;
152        let tile_id =
153            tile_id_for_cell(schema, &coord, system_from_ms).map_err(|e| LiteError::Storage {
154                detail: format!("tile_id_for_cell: {e}"),
155            })?;
156        let payload = CellPayload {
157            valid_from_ms,
158            valid_until_ms,
159            attrs,
160            surrogate,
161        };
162        let bytes = payload.encode().map_err(|e| LiteError::Storage {
163            detail: format!("encode CellPayload: {e}"),
164        })?;
165        let key = encode_coord(&coord)?;
166        self.tiles
167            .entry(tile_id)
168            .or_default()
169            .entries
170            .insert(key, bytes);
171        Ok(tile_id)
172    }
173
174    /// Append a tombstone version for `(coord, system_from_ms)`.
175    pub fn delete_cell(
176        &mut self,
177        schema: &ArraySchema,
178        coord: Vec<CoordValue>,
179        system_from_ms: i64,
180    ) -> Result<TileId, LiteError> {
181        let tile_id =
182            tile_id_for_cell(schema, &coord, system_from_ms).map_err(|e| LiteError::Storage {
183                detail: format!("tile_id_for_cell: {e}"),
184            })?;
185        let key = encode_coord(&coord)?;
186        self.tiles
187            .entry(tile_id)
188            .or_default()
189            .entries
190            .insert(key, CELL_TOMBSTONE_SENTINEL.to_vec());
191        Ok(tile_id)
192    }
193
194    /// Append a GDPR erasure sentinel for `(coord, system_from_ms)`.
195    pub fn erase_cell(
196        &mut self,
197        schema: &ArraySchema,
198        coord: Vec<CoordValue>,
199        system_from_ms: i64,
200    ) -> Result<TileId, LiteError> {
201        let tile_id =
202            tile_id_for_cell(schema, &coord, system_from_ms).map_err(|e| LiteError::Storage {
203                detail: format!("tile_id_for_cell: {e}"),
204            })?;
205        let key = encode_coord(&coord)?;
206        self.tiles
207            .entry(tile_id)
208            .or_default()
209            .entries
210            .insert(key, CELL_GDPR_ERASURE_SENTINEL.to_vec());
211        Ok(tile_id)
212    }
213
214    /// Total cell count (live cells only, sentinels excluded).
215    pub fn stats(&self) -> MemtableStats {
216        MemtableStats {
217            tile_count: self.tiles.len(),
218            cell_count: self.tiles.values().map(|b| b.live_cell_count()).sum(),
219        }
220    }
221
222    pub fn is_empty(&self) -> bool {
223        self.tiles.values().all(|b| b.entry_count() == 0)
224    }
225
226    /// Iterate all tile IDs currently stored (used for prefix enumeration).
227    pub fn tiles_iter(&self) -> impl Iterator<Item = TileId> + '_ {
228        self.tiles.keys().copied()
229    }
230
231    /// Drain tiles sorted by `TileId` and materialise each into a `SparseTile`.
232    pub fn drain_and_materialise(
233        &mut self,
234        schema: &ArraySchema,
235    ) -> Result<Vec<(TileId, SparseTile)>, LiteError> {
236        let tiles: Vec<_> = std::mem::take(&mut self.tiles).into_iter().collect();
237        tiles
238            .into_iter()
239            .map(|(id, buf)| buf.materialise(schema).map(|tile| (id, tile)))
240            .collect()
241    }
242
243    /// Iterate memtable tile versions for `hilbert_prefix` at or before
244    /// `system_as_of`, newest-first. Returns `(TileId, raw_cell_bytes)` for
245    /// the requested `coord`.
246    pub fn iter_cell_versions(
247        &self,
248        hilbert_prefix: u64,
249        system_as_of: i64,
250        coord: &[CoordValue],
251    ) -> Vec<(TileId, Vec<u8>)> {
252        let lo = TileId::new(hilbert_prefix, i64::MIN);
253        let hi = TileId::new(hilbert_prefix, system_as_of);
254        self.tiles
255            .range(lo..=hi)
256            .rev()
257            .filter_map(|(id, buf)| buf.get(coord).map(|b| (*id, b.to_vec())))
258            .collect()
259    }
260
261    /// Iterate all tile versions for `hilbert_prefix` at or before `system_as_of`
262    /// returning `(TileId, SparseTile)` — used by slice queries.
263    pub fn iter_tiles_for_prefix(
264        &self,
265        hilbert_prefix: u64,
266        system_as_of: i64,
267        schema: &ArraySchema,
268    ) -> Result<Vec<(TileId, SparseTile)>, LiteError> {
269        let lo = TileId::new(hilbert_prefix, i64::MIN);
270        let hi = TileId::new(hilbert_prefix, system_as_of);
271        self.tiles
272            .range(lo..=hi)
273            .map(|(id, buf)| buf.materialise(schema).map(|tile| (*id, tile)))
274            .collect()
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281    use nodedb_array::schema::ArraySchemaBuilder;
282    use nodedb_array::schema::attr_spec::{AttrSpec, AttrType};
283    use nodedb_array::schema::dim_spec::{DimSpec, DimType};
284    use nodedb_array::types::domain::{Domain, DomainBound};
285
286    fn schema() -> ArraySchema {
287        ArraySchemaBuilder::new("a")
288            .dim(DimSpec::new(
289                "x",
290                DimType::Int64,
291                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
292            ))
293            .attr(AttrSpec::new("v", AttrType::Int64, true))
294            .tile_extents(vec![4])
295            .build()
296            .unwrap()
297    }
298
299    #[test]
300    fn put_and_stats() {
301        let s = schema();
302        let mut m = ArrayMemtable::new();
303        m.put_cell(
304            &s,
305            PutCellArgs {
306                coord: vec![CoordValue::Int64(1)],
307                attrs: vec![CellValue::Int64(10)],
308                surrogate: Surrogate::new(1),
309                system_from_ms: 100,
310                valid_from_ms: 0,
311                valid_until_ms: OPEN_UPPER,
312            },
313        )
314        .unwrap();
315        assert_eq!(m.stats().cell_count, 1);
316    }
317
318    #[test]
319    fn delete_adds_tombstone_version() {
320        let s = schema();
321        let mut m = ArrayMemtable::new();
322        m.put_cell(
323            &s,
324            PutCellArgs {
325                coord: vec![CoordValue::Int64(1)],
326                attrs: vec![CellValue::Int64(1)],
327                surrogate: Surrogate::new(1),
328                system_from_ms: 100,
329                valid_from_ms: 0,
330                valid_until_ms: OPEN_UPPER,
331            },
332        )
333        .unwrap();
334        m.delete_cell(&s, vec![CoordValue::Int64(1)], 200).unwrap();
335        assert_eq!(m.tiles.len(), 2);
336        assert_eq!(m.stats().cell_count, 1);
337    }
338
339    #[test]
340    fn erase_adds_gdpr_sentinel() {
341        let s = schema();
342        let mut m = ArrayMemtable::new();
343        m.put_cell(
344            &s,
345            PutCellArgs {
346                coord: vec![CoordValue::Int64(2)],
347                attrs: vec![CellValue::Int64(99)],
348                surrogate: Surrogate::new(1),
349                system_from_ms: 10,
350                valid_from_ms: 0,
351                valid_until_ms: OPEN_UPPER,
352            },
353        )
354        .unwrap();
355        m.erase_cell(&s, vec![CoordValue::Int64(2)], 20).unwrap();
356        assert_eq!(m.tiles.len(), 2);
357    }
358
359    #[test]
360    fn drain_and_materialise_empties() {
361        let s = schema();
362        let mut m = ArrayMemtable::new();
363        m.put_cell(
364            &s,
365            PutCellArgs {
366                coord: vec![CoordValue::Int64(3)],
367                attrs: vec![CellValue::Int64(7)],
368                surrogate: Surrogate::new(1),
369                system_from_ms: 1,
370                valid_from_ms: 0,
371                valid_until_ms: OPEN_UPPER,
372            },
373        )
374        .unwrap();
375        let tiles = m.drain_and_materialise(&s).unwrap();
376        assert_eq!(tiles.len(), 1);
377        assert!(m.is_empty());
378    }
379
380    #[test]
381    fn iter_cell_versions_newest_first() {
382        let s = schema();
383        let mut m = ArrayMemtable::new();
384        let coord = vec![CoordValue::Int64(0)];
385        for sys in [100i64, 200, 300] {
386            m.put_cell(
387                &s,
388                PutCellArgs {
389                    coord: coord.clone(),
390                    attrs: vec![CellValue::Int64(sys)],
391                    surrogate: Surrogate::new(1),
392                    system_from_ms: sys,
393                    valid_from_ms: 0,
394                    valid_until_ms: OPEN_UPPER,
395                },
396            )
397            .unwrap();
398        }
399        let tile_id = nodedb_array::tile::tile_id_for_cell(&s, &coord, 100).unwrap();
400        let versions = m.iter_cell_versions(tile_id.hilbert_prefix, 300, &coord);
401        assert_eq!(versions.len(), 3);
402        assert!(versions[0].0.system_from_ms >= versions[1].0.system_from_ms);
403    }
404}