mongreldb_core/
memtable.rs

1//! In-memory write buffer (the "memtable").
2//!
3//! Phase 11.2 wires the buffered [`crate::be_tree::BeTree`] (a Bε-tree over the
4//! composite `(RowId, Epoch)` version key) in as the live memtable, replacing
5//! the prototype skip list. A Bε-tree buffers many pending mutations per
6//! internal node and flushes them to one child in bulk, so write amplification
7//! approaches O(1) — the update-amplification win the design calls for. The
8//! composite key keeps multiple versions of a logical row coexisting so a
9//! snapshot read returns the newest version with `committed_epoch <= snapshot`,
10//! which is what makes MVCC correct within the live memtable.
11
12use crate::be_tree::BeTree;
13use crate::epoch::Epoch;
14use crate::rowid::RowId;
15use serde::{Deserialize, Serialize};
16use std::collections::{BTreeMap, HashMap};
17use std::sync::Arc;
18
19/// A cell value in the in-memory path. The flush path re-encodes these into
20/// columnar pages; it is intentionally simple for the prototype.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22pub enum Value {
23    Null,
24    Bool(bool),
25    Int64(i64),
26    Float64(f64),
27    Bytes(Vec<u8>),
28    Embedding(Vec<f32>),
29    /// Unscaled decimal value (i128). The column's `TypeId::Decimal128`
30    /// carries the precision/scale for formatting.
31    Decimal(i128),
32    /// SQL INTERVAL value: months, days, nanoseconds.
33    Interval {
34        months: i64,
35        days: i32,
36        nanos: i64,
37    },
38    /// RFC 4122 UUID (16 bytes, big-endian for sort order).
39    Uuid([u8; 16]),
40    /// JSON value stored as a UTF-8 byte sequence.
41    Json(Vec<u8>),
42    /// Generated embedding with durable source and model provenance.
43    ///
44    /// Kept last so existing bincode enum discriminants remain stable.
45    GeneratedEmbedding(Box<crate::embedding::GeneratedEmbeddingValue>),
46}
47
48impl Value {
49    pub fn as_embedding(&self) -> Option<&[f32]> {
50        match self {
51            Self::Embedding(values) => Some(values),
52            Self::GeneratedEmbedding(value) => Some(&value.vector),
53            _ => None,
54        }
55    }
56
57    pub fn generated_embedding_metadata(
58        &self,
59    ) -> Option<&crate::embedding::GeneratedEmbeddingMetadata> {
60        match self {
61            Self::GeneratedEmbedding(value) => Some(&value.metadata),
62            _ => None,
63        }
64    }
65
66    /// Lexicographically-comparable byte encoding for index keys (PK HOT,
67    /// bitmaps). Big-endian for integers so byte order matches value order.
68    pub fn encode_key(&self) -> Vec<u8> {
69        match self {
70            Value::Null => Vec::new(),
71            Value::Bool(b) => vec![*b as u8],
72            Value::Int64(n) => n.to_be_bytes().to_vec(),
73            Value::Float64(f) => f.to_bits().to_be_bytes().to_vec(),
74            Value::Bytes(b) => b.clone(),
75            Value::Embedding(v) => {
76                let mut out = Vec::with_capacity(v.len() * 4);
77                for x in v {
78                    out.extend_from_slice(&x.to_bits().to_be_bytes());
79                }
80                out
81            }
82            Value::GeneratedEmbedding(value) => {
83                let mut out = Vec::with_capacity(value.vector.len() * 4);
84                for x in &value.vector {
85                    out.extend_from_slice(&x.to_bits().to_be_bytes());
86                }
87                out
88            }
89            Value::Decimal(d) => d.to_be_bytes().to_vec(),
90            Value::Interval {
91                months,
92                days,
93                nanos,
94            } => {
95                let mut out = Vec::with_capacity(20);
96                out.extend_from_slice(&months.to_be_bytes());
97                out.extend_from_slice(&days.to_be_bytes());
98                out.extend_from_slice(&nanos.to_be_bytes());
99                out
100            }
101            Value::Uuid(b) => b.to_vec(),
102            Value::Json(b) => b.clone(),
103        }
104    }
105
106    pub(crate) fn estimated_bytes(&self) -> u64 {
107        match self {
108            Value::Null => 1,
109            Value::Bool(_) => 1,
110            Value::Int64(_) | Value::Float64(_) => 8,
111            Value::Bytes(bytes) | Value::Json(bytes) => 16 + bytes.len() as u64,
112            Value::Embedding(values) => 16 + (values.len() as u64) * 4,
113            Value::GeneratedEmbedding(value) => {
114                16 + (value.vector.len() as u64) * 4
115                    + value.metadata.provider_id.len() as u64
116                    + value.metadata.model_id.len() as u64
117                    + value.metadata.model_version.len() as u64
118                    + value.metadata.preprocessing_version.len() as u64
119                    + 48
120            }
121            Value::Decimal(_) | Value::Uuid(_) => 16,
122            Value::Interval { .. } => 20,
123        }
124    }
125}
126
127/// One logical row held in the memtable. A `deleted` row is a tombstone.
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct Row {
130    pub row_id: RowId,
131    pub committed_epoch: Epoch,
132    pub columns: HashMap<u16, Value>,
133    pub deleted: bool,
134}
135
136impl Row {
137    pub fn new(row_id: RowId, committed_epoch: Epoch) -> Self {
138        Self {
139            row_id,
140            committed_epoch,
141            columns: HashMap::new(),
142            deleted: false,
143        }
144    }
145
146    pub fn with_column(mut self, column_id: u16, value: Value) -> Self {
147        self.columns.insert(column_id, value);
148        self
149    }
150
151    /// Rough byte estimate for flush-threshold decisions.
152    pub fn estimated_bytes(&self) -> u64 {
153        self.columns
154            .values()
155            .fold(32, |bytes, value| bytes + value.estimated_bytes())
156    }
157}
158
159/// Bε-tree-backed memtable, ordered by `(RowId, Epoch)`. A drop-in replacement
160/// for the prototype skip list: the same MVCC semantics with lower write
161/// amplification (buffered messages flush to children in bulk).
162#[derive(Clone)]
163struct MemtableSegment {
164    tree: BeTree,
165    byte_size: u64,
166}
167
168/// Structurally shared committed overlays plus one small mutable write delta.
169#[derive(Clone)]
170pub struct Memtable {
171    frozen: Arc<Vec<Arc<MemtableSegment>>>,
172    active: MemtableSegment,
173    byte_size: u64,
174}
175
176impl Default for Memtable {
177    fn default() -> Self {
178        Self::new()
179    }
180}
181
182impl Memtable {
183    pub fn new() -> Self {
184        Self {
185            frozen: Arc::new(Vec::new()),
186            active: MemtableSegment {
187                tree: BeTree::new(),
188                byte_size: 0,
189            },
190            byte_size: 0,
191        }
192    }
193
194    /// Append a row version (keyed by `(row_id, committed_epoch)`). Versions are
195    /// never overwritten; the newest visible one wins at read time.
196    pub fn upsert(&mut self, row: Row) {
197        let bytes = row.estimated_bytes();
198        self.byte_size = self.byte_size.saturating_add(bytes);
199        self.active.byte_size = self.active.byte_size.saturating_add(bytes);
200        self.active.tree.insert_row(row);
201    }
202
203    /// Append a tombstone version for `row_id` at `epoch`. The tombstone copies
204    /// the columns from the newest live version so that engine-level HOT cleanup
205    /// can recover the primary-key value during WAL replay.
206    pub fn tombstone(&mut self, row_id: RowId, epoch: Epoch) {
207        let mut columns = HashMap::new();
208        if let Some(live) = self.get(row_id, Epoch(epoch.0.saturating_sub(1))) {
209            columns = live.columns;
210        }
211        let row = Row {
212            row_id,
213            committed_epoch: epoch,
214            columns,
215            deleted: true,
216        };
217        self.upsert(row);
218    }
219
220    /// Read the row at `row_id` visible to `snapshot`: the newest version with
221    /// `epoch <= snapshot`. Returns `None` if that version is a tombstone (or no
222    /// such version exists).
223    pub fn get(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<Row> {
224        self.get_version(row_id, snapshot_epoch)
225            .and_then(|(_, row)| (!row.deleted).then_some(row))
226    }
227
228    /// Newest version of `row_id` with `epoch <= snapshot`, **including
229    /// tombstones** (as a `Row` with `deleted=true`). Used by the engine to
230    /// merge versions across the memtable and sorted runs.
231    pub fn get_version(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<(Epoch, Row)> {
232        let mut best = self.active.tree.get_version(row_id, snapshot_epoch);
233        for segment in self.frozen.iter().rev() {
234            let Some(candidate) = segment.tree.get_version(row_id, snapshot_epoch) else {
235                continue;
236            };
237            if best.as_ref().is_none_or(|(epoch, _)| candidate.0 > *epoch) {
238                best = Some(candidate);
239            }
240        }
241        best
242    }
243
244    /// Number of stored versions.
245    pub fn len(&self) -> usize {
246        self.active.tree.mutations()
247            + self
248                .frozen
249                .iter()
250                .map(|segment| segment.tree.mutations())
251                .sum::<usize>()
252    }
253
254    pub fn is_empty(&self) -> bool {
255        self.active.tree.is_empty() && self.frozen.is_empty()
256    }
257
258    pub fn approx_bytes(&self) -> u64 {
259        self.byte_size
260    }
261
262    /// Visible rows at `snapshot`, deduplicated to the newest version per
263    /// `RowId` (tombstones drop their row). Returned in ascending `RowId` order.
264    pub fn visible_rows(&self, snapshot_epoch: Epoch) -> Vec<Row> {
265        self.visible_versions(snapshot_epoch)
266            .into_iter()
267            .filter(|r| !r.deleted)
268            .collect()
269    }
270
271    /// Newest visible version per `RowId` at `snapshot`, **including
272    /// tombstones** (as `Row`s with `deleted=true`). Used by the engine to merge
273    /// versions across the memtable and sorted runs.
274    pub fn visible_versions(&self, snapshot_epoch: Epoch) -> Vec<Row> {
275        let mut by_row: BTreeMap<RowId, Row> = BTreeMap::new();
276        for segment in self
277            .frozen
278            .iter()
279            .map(|segment| &segment.tree)
280            .chain(std::iter::once(&self.active.tree))
281        {
282            for row in segment.versions() {
283                if row.committed_epoch <= snapshot_epoch {
284                    by_row
285                        .entry(row.row_id)
286                        .and_modify(|existing| {
287                            if row.committed_epoch > existing.committed_epoch {
288                                *existing = row.clone();
289                            }
290                        })
291                        .or_insert(row);
292                }
293            }
294        }
295        by_row.into_values().collect()
296    }
297
298    /// Freeze the current write delta so future clones share it by `Arc`.
299    pub(crate) fn seal(&mut self) {
300        if self.active.tree.is_empty() {
301            return;
302        }
303        let active = std::mem::replace(
304            &mut self.active,
305            MemtableSegment {
306                tree: BeTree::new(),
307                byte_size: 0,
308            },
309        );
310        Arc::make_mut(&mut self.frozen).push(Arc::new(active));
311        if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
312            self.consolidate();
313        }
314    }
315
316    fn consolidate(&mut self) {
317        let mut tree = BeTree::new();
318        for row in self
319            .frozen
320            .iter()
321            .flat_map(|segment| segment.tree.versions())
322        {
323            tree.insert_row(row);
324        }
325        self.frozen = Arc::new(vec![Arc::new(MemtableSegment {
326            tree,
327            byte_size: self.byte_size,
328        })]);
329    }
330
331    #[cfg(test)]
332    pub(crate) fn frozen_layer_count(&self) -> usize {
333        self.frozen.len()
334    }
335
336    /// Drain all versions (for a memtable-to-run flush). Returns them in
337    /// ascending `(RowId, Epoch)` order.
338    pub fn drain_sorted(&mut self) -> Vec<Row> {
339        let mut out = self
340            .frozen
341            .iter()
342            .flat_map(|segment| segment.tree.versions())
343            .chain(self.active.tree.versions())
344            .collect::<Vec<_>>();
345        out.sort_by_key(|row| (row.row_id, row.committed_epoch));
346        self.frozen = Arc::new(Vec::new());
347        self.active = MemtableSegment {
348            tree: BeTree::new(),
349            byte_size: 0,
350        };
351        self.byte_size = 0;
352        out
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359
360    fn row(id: u64, epoch: u64) -> Row {
361        Row::new(RowId(id), Epoch(epoch)).with_column(1, Value::Int64(id as i64 * 10))
362    }
363
364    #[test]
365    fn upsert_get_and_visibility() {
366        let mut m = Memtable::new();
367        m.upsert(row(1, 5));
368        assert_eq!(m.len(), 1);
369        assert!(m.get(RowId(1), Epoch(5)).is_some());
370        assert!(m.get(RowId(1), Epoch(4)).is_none()); // not yet visible
371        assert!(m.get(RowId(2), Epoch(9)).is_none()); // missing
372    }
373
374    #[test]
375    fn tombstone_supersedes_at_its_epoch() {
376        let mut m = Memtable::new();
377        m.upsert(row(1, 1));
378        // Before the tombstone: the live version is visible.
379        assert!(m.get(RowId(1), Epoch(1)).is_some());
380        m.tombstone(RowId(1), Epoch(2));
381        // At/after the tombstone: hidden.
382        assert!(m.get(RowId(1), Epoch(2)).is_none());
383        assert!(m.get(RowId(1), Epoch(9)).is_none());
384        // A snapshot before the tombstone still sees the live version.
385        assert!(m.get(RowId(1), Epoch(1)).is_some());
386    }
387
388    #[test]
389    fn sealed_generations_share_rows_and_consolidate() {
390        let mut writer = Memtable::new();
391        for id in 0..crate::MAX_READ_GENERATION_LAYERS as u64 + 2 {
392            writer.upsert(row(id, id + 1));
393            writer.seal();
394        }
395        assert!(writer.frozen_layer_count() < crate::MAX_READ_GENERATION_LAYERS);
396        let generation = writer.clone();
397        writer.upsert(row(99, 99));
398        assert!(generation.get(RowId(99), Epoch(99)).is_none());
399        assert!(writer.get(RowId(99), Epoch(99)).is_some());
400    }
401
402    #[test]
403    fn drain_sorted_is_ascending_and_empties() {
404        let mut m = Memtable::new();
405        m.upsert(row(3, 1));
406        m.upsert(row(1, 1));
407        m.upsert(row(2, 1));
408        let out = m.drain_sorted();
409        let ids: Vec<u64> = out.iter().map(|r| r.row_id.0).collect();
410        assert_eq!(ids, vec![1, 2, 3]);
411        assert!(m.is_empty());
412        assert_eq!(m.approx_bytes(), 0);
413    }
414
415    #[test]
416    fn visible_rows_dedups_to_newest_version() {
417        let mut m = Memtable::new();
418        m.upsert(row(1, 1));
419        m.upsert(row(2, 9)); // future relative to snapshot 5
420        m.upsert(row(3, 1));
421        m.upsert(row(1, 3)); // newer version of row 1
422        let ids: Vec<u64> = m
423            .visible_rows(Epoch(5))
424            .iter()
425            .map(|r| r.row_id.0)
426            .collect();
427        assert_eq!(ids, vec![1, 3]);
428    }
429}
mongreldb_core/memtable.rs

mongreldb_core/
memtable.rs
