1use crate::be_tree::BeTree;
13use crate::epoch::Epoch;
14use crate::rowid::RowId;
15use serde::{Deserialize, Serialize};
16use std::collections::{BTreeMap, HashMap};
17
18#[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 Decimal(i128),
31 Interval {
33 months: i64,
34 days: i32,
35 nanos: i64,
36 },
37 Uuid([u8; 16]),
39 Json(Vec<u8>),
41}
42
43impl Value {
44 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#[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 pub fn estimated_bytes(&self) -> u64 {
104 let mut n = 32; 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#[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 pub fn upsert(&mut self, row: Row) {
149 self.byte_size += row.estimated_bytes();
150 self.tree.insert_row(row);
151 }
152
153 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 pub fn get(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<Row> {
174 self.tree.get_visible(row_id, snapshot_epoch)
175 }
176
177 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 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 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 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 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()); assert!(m.get(RowId(2), Epoch(9)).is_none()); }
252
253 #[test]
254 fn tombstone_supersedes_at_its_epoch() {
255 let mut m = Memtable::new();
256 m.upsert(row(1, 1));
257 assert!(m.get(RowId(1), Epoch(1)).is_some());
259 m.tombstone(RowId(1), Epoch(2));
260 assert!(m.get(RowId(1), Epoch(2)).is_none());
262 assert!(m.get(RowId(1), Epoch(9)).is_none());
263 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)); m.upsert(row(3, 1));
286 m.upsert(row(1, 3)); 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}