Skip to main content

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};
17
18/// A cell value in the in-memory path. The flush path re-encodes these into
19/// columnar pages; it is intentionally simple for the prototype.
20#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
21pub enum Value {
22    Null,
23    Bool(bool),
24    Int64(i64),
25    Float64(f64),
26    Bytes(Vec<u8>),
27    Embedding(Vec<f32>),
28    /// Unscaled decimal value (i128). The column's `TypeId::Decimal128`
29    /// carries the precision/scale for formatting.
30    Decimal(i128),
31    /// SQL INTERVAL value: months, days, nanoseconds.
32    Interval {
33        months: i64,
34        days: i32,
35        nanos: i64,
36    },
37    /// RFC 4122 UUID (16 bytes, big-endian for sort order).
38    Uuid([u8; 16]),
39    /// JSON value stored as a UTF-8 byte sequence.
40    Json(Vec<u8>),
41}
42
43impl Value {
44    /// Lexicographically-comparable byte encoding for index keys (PK HOT,
45    /// bitmaps). Big-endian for integers so byte order matches value order.
46    pub fn encode_key(&self) -> Vec<u8> {
47        match self {
48            Value::Null => Vec::new(),
49            Value::Bool(b) => vec![*b as u8],
50            Value::Int64(n) => n.to_be_bytes().to_vec(),
51            Value::Float64(f) => f.to_bits().to_be_bytes().to_vec(),
52            Value::Bytes(b) => b.clone(),
53            Value::Embedding(v) => {
54                let mut out = Vec::with_capacity(v.len() * 4);
55                for x in v {
56                    out.extend_from_slice(&x.to_bits().to_be_bytes());
57                }
58                out
59            }
60            Value::Decimal(d) => d.to_be_bytes().to_vec(),
61            Value::Interval {
62                months,
63                days,
64                nanos,
65            } => {
66                let mut out = Vec::with_capacity(20);
67                out.extend_from_slice(&months.to_be_bytes());
68                out.extend_from_slice(&days.to_be_bytes());
69                out.extend_from_slice(&nanos.to_be_bytes());
70                out
71            }
72            Value::Uuid(b) => b.to_vec(),
73            Value::Json(b) => b.clone(),
74        }
75    }
76}
77
78/// One logical row held in the memtable. A `deleted` row is a tombstone.
79#[derive(Debug, Clone, Serialize, Deserialize)]
80pub struct Row {
81    pub row_id: RowId,
82    pub committed_epoch: Epoch,
83    pub columns: HashMap<u16, Value>,
84    pub deleted: bool,
85}
86
87impl Row {
88    pub fn new(row_id: RowId, committed_epoch: Epoch) -> Self {
89        Self {
90            row_id,
91            committed_epoch,
92            columns: HashMap::new(),
93            deleted: false,
94        }
95    }
96
97    pub fn with_column(mut self, column_id: u16, value: Value) -> Self {
98        self.columns.insert(column_id, value);
99        self
100    }
101
102    /// Rough byte estimate for flush-threshold decisions.
103    pub fn estimated_bytes(&self) -> u64 {
104        let mut n = 32; // header overhead
105        for v in self.columns.values() {
106            n += match v {
107                Value::Null => 1,
108                Value::Bool(_) => 1,
109                Value::Int64(_) => 8,
110                Value::Float64(_) => 8,
111                Value::Bytes(b) => 16 + b.len() as u64,
112                Value::Embedding(v) => 16 + (v.len() as u64) * 4,
113                Value::Decimal(_) => 16,
114                Value::Interval { .. } => 20,
115                Value::Uuid(_) => 16,
116                Value::Json(b) => 16 + b.len() as u64,
117            };
118        }
119        n
120    }
121}
122
123/// Bε-tree-backed memtable, ordered by `(RowId, Epoch)`. A drop-in replacement
124/// for the prototype skip list: the same MVCC semantics with lower write
125/// amplification (buffered messages flush to children in bulk).
126#[derive(Clone)]
127pub struct Memtable {
128    tree: BeTree,
129    byte_size: u64,
130}
131
132impl Default for Memtable {
133    fn default() -> Self {
134        Self::new()
135    }
136}
137
138impl Memtable {
139    pub fn new() -> Self {
140        Self {
141            tree: BeTree::new(),
142            byte_size: 0,
143        }
144    }
145
146    /// Append a row version (keyed by `(row_id, committed_epoch)`). Versions are
147    /// never overwritten; the newest visible one wins at read time.
148    pub fn upsert(&mut self, row: Row) {
149        self.byte_size += row.estimated_bytes();
150        self.tree.insert_row(row);
151    }
152
153    /// Append a tombstone version for `row_id` at `epoch`. The tombstone copies
154    /// the columns from the newest live version so that engine-level HOT cleanup
155    /// can recover the primary-key value during WAL replay.
156    pub fn tombstone(&mut self, row_id: RowId, epoch: Epoch) {
157        let mut columns = HashMap::new();
158        if let Some(live) = self.get(row_id, Epoch(epoch.0.saturating_sub(1))) {
159            columns = live.columns;
160        }
161        let row = Row {
162            row_id,
163            committed_epoch: epoch,
164            columns,
165            deleted: true,
166        };
167        self.upsert(row);
168    }
169
170    /// Read the row at `row_id` visible to `snapshot`: the newest version with
171    /// `epoch <= snapshot`. Returns `None` if that version is a tombstone (or no
172    /// such version exists).
173    pub fn get(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<Row> {
174        self.tree.get_visible(row_id, snapshot_epoch)
175    }
176
177    /// Newest version of `row_id` with `epoch <= snapshot`, **including
178    /// tombstones** (as a `Row` with `deleted=true`). Used by the engine to
179    /// merge versions across the memtable and sorted runs.
180    pub fn get_version(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<(Epoch, Row)> {
181        self.tree.get_version(row_id, snapshot_epoch)
182    }
183
184    /// Number of stored versions.
185    pub fn len(&self) -> usize {
186        self.tree.mutations()
187    }
188
189    pub fn is_empty(&self) -> bool {
190        self.tree.is_empty()
191    }
192
193    pub fn approx_bytes(&self) -> u64 {
194        self.byte_size
195    }
196
197    /// Visible rows at `snapshot`, deduplicated to the newest version per
198    /// `RowId` (tombstones drop their row). Returned in ascending `RowId` order.
199    pub fn visible_rows(&self, snapshot_epoch: Epoch) -> Vec<Row> {
200        self.visible_versions(snapshot_epoch)
201            .into_iter()
202            .filter(|r| !r.deleted)
203            .collect()
204    }
205
206    /// Newest visible version per `RowId` at `snapshot`, **including
207    /// tombstones** (as `Row`s with `deleted=true`). Used by the engine to merge
208    /// versions across the memtable and sorted runs.
209    pub fn visible_versions(&self, snapshot_epoch: Epoch) -> Vec<Row> {
210        let mut by_row: BTreeMap<RowId, Row> = BTreeMap::new();
211        for row in self.tree.versions() {
212            if row.committed_epoch <= snapshot_epoch {
213                by_row
214                    .entry(row.row_id)
215                    .and_modify(|e| {
216                        if row.committed_epoch > e.committed_epoch {
217                            *e = row.clone();
218                        }
219                    })
220                    .or_insert(row);
221            }
222        }
223        by_row.into_values().collect()
224    }
225
226    /// Drain all versions (for a memtable-to-run flush). Returns them in
227    /// ascending `(RowId, Epoch)` order.
228    pub fn drain_sorted(&mut self) -> Vec<Row> {
229        let out = std::mem::take(&mut self.tree).into_sorted_rows();
230        self.byte_size = 0;
231        out
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    fn row(id: u64, epoch: u64) -> Row {
240        Row::new(RowId(id), Epoch(epoch)).with_column(1, Value::Int64(id as i64 * 10))
241    }
242
243    #[test]
244    fn upsert_get_and_visibility() {
245        let mut m = Memtable::new();
246        m.upsert(row(1, 5));
247        assert_eq!(m.len(), 1);
248        assert!(m.get(RowId(1), Epoch(5)).is_some());
249        assert!(m.get(RowId(1), Epoch(4)).is_none()); // not yet visible
250        assert!(m.get(RowId(2), Epoch(9)).is_none()); // missing
251    }
252
253    #[test]
254    fn tombstone_supersedes_at_its_epoch() {
255        let mut m = Memtable::new();
256        m.upsert(row(1, 1));
257        // Before the tombstone: the live version is visible.
258        assert!(m.get(RowId(1), Epoch(1)).is_some());
259        m.tombstone(RowId(1), Epoch(2));
260        // At/after the tombstone: hidden.
261        assert!(m.get(RowId(1), Epoch(2)).is_none());
262        assert!(m.get(RowId(1), Epoch(9)).is_none());
263        // A snapshot before the tombstone still sees the live version.
264        assert!(m.get(RowId(1), Epoch(1)).is_some());
265    }
266
267    #[test]
268    fn drain_sorted_is_ascending_and_empties() {
269        let mut m = Memtable::new();
270        m.upsert(row(3, 1));
271        m.upsert(row(1, 1));
272        m.upsert(row(2, 1));
273        let out = m.drain_sorted();
274        let ids: Vec<u64> = out.iter().map(|r| r.row_id.0).collect();
275        assert_eq!(ids, vec![1, 2, 3]);
276        assert!(m.is_empty());
277        assert_eq!(m.approx_bytes(), 0);
278    }
279
280    #[test]
281    fn visible_rows_dedups_to_newest_version() {
282        let mut m = Memtable::new();
283        m.upsert(row(1, 1));
284        m.upsert(row(2, 9)); // future relative to snapshot 5
285        m.upsert(row(3, 1));
286        m.upsert(row(1, 3)); // newer version of row 1
287        let ids: Vec<u64> = m
288            .visible_rows(Epoch(5))
289            .iter()
290            .map(|r| r.row_id.0)
291            .collect();
292        assert_eq!(ids, vec![1, 3]);
293    }
294}