1use std::collections::HashMap;
15
16use nodedb_types::columnar::ColumnarSchema;
17use nodedb_types::surrogate::Surrogate;
18use serde::{Deserialize, Serialize};
19use zerompk::{FromMessagePack, ToMessagePack};
20
21use crate::delete_bitmap::DeleteBitmap;
22use crate::error::ColumnarError;
23use crate::memtable::ColumnData;
24use crate::pk_index::PkIndex;
25
26use super::engine::MutationEngine;
27
28pub type FlushedSurrogateTable = Vec<Vec<Option<Surrogate>>>;
32
33#[derive(Debug, Clone, Serialize, Deserialize, ToMessagePack, FromMessagePack)]
41pub enum ColumnDataSnapshot {
42 Int64 {
43 values: Vec<i64>,
44 valid: Option<Vec<bool>>,
45 },
46 Float64 {
47 values: Vec<f64>,
48 valid: Option<Vec<bool>>,
49 },
50 Bool {
51 values: Vec<bool>,
52 valid: Option<Vec<bool>>,
53 },
54 Timestamp {
55 values: Vec<i64>,
56 valid: Option<Vec<bool>>,
57 },
58 Decimal {
59 values: Vec<[u8; 16]>,
60 valid: Option<Vec<bool>>,
61 },
62 Uuid {
63 values: Vec<[u8; 16]>,
64 valid: Option<Vec<bool>>,
65 },
66 String {
67 data: Vec<u8>,
68 offsets: Vec<u32>,
69 valid: Option<Vec<bool>>,
70 },
71 Bytes {
72 data: Vec<u8>,
73 offsets: Vec<u32>,
74 valid: Option<Vec<bool>>,
75 },
76 Json {
77 data: Vec<u8>,
78 offsets: Vec<u32>,
79 valid: Option<Vec<bool>>,
80 },
81 Geometry {
82 data: Vec<u8>,
83 offsets: Vec<u32>,
84 valid: Option<Vec<bool>>,
85 },
86 Vector {
87 data: Vec<f32>,
88 dim: u32,
89 valid: Option<Vec<bool>>,
90 },
91 DictEncoded {
96 ids: Vec<u32>,
97 dictionary: Vec<String>,
98 valid: Option<Vec<bool>>,
99 },
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize, ToMessagePack, FromMessagePack)]
110#[msgpack(map)]
111pub struct ColumnarEngineSnapshot {
112 pub collection: String,
114 pub schema: ColumnarSchema,
116 pub memtable_columns: Vec<ColumnDataSnapshot>,
118 pub memtable_surrogates: Vec<Option<Surrogate>>,
120 pub pk_index_bytes: Vec<u8>,
122 pub delete_bitmaps: Vec<(u64, Vec<u8>)>,
127 pub memtable_delete_bitmap_bytes: Vec<u8>,
130 pub flushed_segments: Vec<Vec<u8>>,
133 #[msgpack(default)]
137 #[serde(default)]
138 pub flushed_surrogates: FlushedSurrogateTable,
139 pub next_segment_id: u64,
141 pub memtable_segment_id: u64,
143 pub memtable_row_counter: u32,
145}
146
147impl MutationEngine {
150 pub fn export_snapshot(
161 &self,
162 flushed_segments: &[Vec<u8>],
163 flushed_surrogates: &[Vec<Option<Surrogate>>],
164 ) -> Result<ColumnarEngineSnapshot, ColumnarError> {
165 let memtable_columns = self
167 .memtable
168 .columns()
169 .iter()
170 .map(column_to_snapshot)
171 .collect::<Vec<_>>();
172
173 let pk_index_bytes = self.pk_index.to_bytes()?;
175
176 let mut delete_bitmaps: Vec<(u64, Vec<u8>)> = Vec::new();
179 let mut memtable_delete_bitmap_bytes = Vec::new();
180
181 for (&seg_id, bitmap) in &self.delete_bitmaps {
182 let bytes = bitmap.to_bytes()?;
183 if seg_id == self.memtable_segment_id {
184 memtable_delete_bitmap_bytes = bytes;
185 } else {
186 delete_bitmaps.push((seg_id, bytes));
187 }
188 }
189
190 Ok(ColumnarEngineSnapshot {
191 collection: self.collection.clone(),
192 schema: self.schema.clone(),
193 memtable_columns,
194 memtable_surrogates: self.memtable_surrogates.clone(),
195 pk_index_bytes,
196 delete_bitmaps,
197 memtable_delete_bitmap_bytes,
198 flushed_segments: flushed_segments.to_vec(),
199 flushed_surrogates: flushed_surrogates.to_vec(),
200 next_segment_id: self.next_segment_id,
201 memtable_segment_id: self.memtable_segment_id,
202 memtable_row_counter: self.memtable_row_counter,
203 })
204 }
205
206 pub fn from_snapshot(
224 snap: ColumnarEngineSnapshot,
225 ) -> Result<(MutationEngine, Vec<Vec<u8>>, FlushedSurrogateTable), ColumnarError> {
226 let col_count = snap.schema.columns.len();
227 if snap.memtable_columns.len() != col_count {
228 return Err(ColumnarError::SchemaMismatch {
229 expected: col_count,
230 got: snap.memtable_columns.len(),
231 });
232 }
233
234 let pk_index = PkIndex::from_bytes(&snap.pk_index_bytes)?;
236
237 let mut delete_bitmaps: HashMap<u64, DeleteBitmap> = HashMap::new();
239 for (seg_id, bytes) in snap.delete_bitmaps {
240 let bm = DeleteBitmap::from_bytes(&bytes)?;
241 delete_bitmaps.insert(seg_id, bm);
242 }
243
244 if !snap.memtable_delete_bitmap_bytes.is_empty() {
246 let bm = DeleteBitmap::from_bytes(&snap.memtable_delete_bitmap_bytes)?;
247 delete_bitmaps.insert(snap.memtable_segment_id, bm);
248 }
249
250 let columns: Vec<ColumnData> = snap
252 .memtable_columns
253 .into_iter()
254 .map(snapshot_to_column)
255 .collect();
256
257 let pk_col_indices: Vec<usize> = snap
259 .schema
260 .columns
261 .iter()
262 .enumerate()
263 .filter(|(_, c)| c.primary_key)
264 .map(|(i, _)| i)
265 .collect();
266
267 let memtable = crate::memtable::ColumnarMemtable::from_raw_columns(
269 &snap.schema,
270 columns,
271 snap.memtable_row_counter as usize,
272 );
273
274 let engine = MutationEngine {
275 collection: snap.collection,
276 schema: snap.schema,
277 memtable,
278 pk_index,
279 delete_bitmaps,
280 pk_col_indices,
281 next_segment_id: snap.next_segment_id,
282 memtable_segment_id: snap.memtable_segment_id,
283 memtable_row_counter: snap.memtable_row_counter,
284 memtable_surrogates: snap.memtable_surrogates,
285 };
286
287 Ok((engine, snap.flushed_segments, snap.flushed_surrogates))
288 }
289}
290
291fn column_to_snapshot(col: &ColumnData) -> ColumnDataSnapshot {
295 match col {
296 ColumnData::Int64 { values, valid } => ColumnDataSnapshot::Int64 {
297 values: values.clone(),
298 valid: valid.clone(),
299 },
300 ColumnData::Float64 { values, valid } => ColumnDataSnapshot::Float64 {
301 values: values.clone(),
302 valid: valid.clone(),
303 },
304 ColumnData::Bool { values, valid } => ColumnDataSnapshot::Bool {
305 values: values.clone(),
306 valid: valid.clone(),
307 },
308 ColumnData::Timestamp { values, valid } => ColumnDataSnapshot::Timestamp {
309 values: values.clone(),
310 valid: valid.clone(),
311 },
312 ColumnData::Decimal { values, valid } => ColumnDataSnapshot::Decimal {
313 values: values.clone(),
314 valid: valid.clone(),
315 },
316 ColumnData::Uuid { values, valid } => ColumnDataSnapshot::Uuid {
317 values: values.clone(),
318 valid: valid.clone(),
319 },
320 ColumnData::String {
321 data,
322 offsets,
323 valid,
324 } => ColumnDataSnapshot::String {
325 data: data.clone(),
326 offsets: offsets.clone(),
327 valid: valid.clone(),
328 },
329 ColumnData::Bytes {
330 data,
331 offsets,
332 valid,
333 } => ColumnDataSnapshot::Bytes {
334 data: data.clone(),
335 offsets: offsets.clone(),
336 valid: valid.clone(),
337 },
338 ColumnData::Json {
339 data,
340 offsets,
341 valid,
342 } => ColumnDataSnapshot::Json {
343 data: data.clone(),
344 offsets: offsets.clone(),
345 valid: valid.clone(),
346 },
347 ColumnData::Geometry {
348 data,
349 offsets,
350 valid,
351 } => ColumnDataSnapshot::Geometry {
352 data: data.clone(),
353 offsets: offsets.clone(),
354 valid: valid.clone(),
355 },
356 ColumnData::Vector { data, dim, valid } => ColumnDataSnapshot::Vector {
357 data: data.clone(),
358 dim: *dim,
359 valid: valid.clone(),
360 },
361 ColumnData::DictEncoded {
362 ids,
363 dictionary,
364 valid,
365 ..
367 } => ColumnDataSnapshot::DictEncoded {
368 ids: ids.clone(),
369 dictionary: dictionary.clone(),
370 valid: valid.clone(),
371 },
372 }
373}
374
375fn snapshot_to_column(snap: ColumnDataSnapshot) -> ColumnData {
379 match snap {
380 ColumnDataSnapshot::Int64 { values, valid } => ColumnData::Int64 { values, valid },
381 ColumnDataSnapshot::Float64 { values, valid } => ColumnData::Float64 { values, valid },
382 ColumnDataSnapshot::Bool { values, valid } => ColumnData::Bool { values, valid },
383 ColumnDataSnapshot::Timestamp { values, valid } => ColumnData::Timestamp { values, valid },
384 ColumnDataSnapshot::Decimal { values, valid } => ColumnData::Decimal { values, valid },
385 ColumnDataSnapshot::Uuid { values, valid } => ColumnData::Uuid { values, valid },
386 ColumnDataSnapshot::String {
387 data,
388 offsets,
389 valid,
390 } => ColumnData::String {
391 data,
392 offsets,
393 valid,
394 },
395 ColumnDataSnapshot::Bytes {
396 data,
397 offsets,
398 valid,
399 } => ColumnData::Bytes {
400 data,
401 offsets,
402 valid,
403 },
404 ColumnDataSnapshot::Json {
405 data,
406 offsets,
407 valid,
408 } => ColumnData::Json {
409 data,
410 offsets,
411 valid,
412 },
413 ColumnDataSnapshot::Geometry {
414 data,
415 offsets,
416 valid,
417 } => ColumnData::Geometry {
418 data,
419 offsets,
420 valid,
421 },
422 ColumnDataSnapshot::Vector { data, dim, valid } => ColumnData::Vector { data, dim, valid },
423 ColumnDataSnapshot::DictEncoded {
424 ids,
425 dictionary,
426 valid,
427 } => {
428 let reverse: HashMap<String, u32> = dictionary
430 .iter()
431 .enumerate()
432 .map(|(i, s)| (s.clone(), i as u32))
433 .collect();
434 ColumnData::DictEncoded {
435 ids,
436 dictionary,
437 reverse,
438 valid,
439 }
440 }
441 }
442}
443
444#[cfg(test)]
447mod tests {
448 use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};
449 use nodedb_types::value::Value;
450
451 use super::*;
452
453 fn simple_schema() -> ColumnarSchema {
454 ColumnarSchema {
455 columns: vec![
456 ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
457 ColumnDef::required("name", ColumnType::String),
458 ColumnDef::nullable("score", ColumnType::Float64),
459 ],
460 version: 1,
461 }
462 }
463
464 fn insert_row(engine: &mut MutationEngine, id: i64, name: &str, score: Option<f64>) {
465 let score_val = score.map(Value::Float).unwrap_or(Value::Null);
466 engine
467 .insert(&[Value::Integer(id), Value::String(name.into()), score_val])
468 .expect("insert");
469 }
470
471 #[test]
472 fn round_trip_memtable_rows() {
473 let schema = simple_schema();
474 let mut engine = MutationEngine::new("test_col".to_string(), schema);
475
476 insert_row(&mut engine, 1, "Alice", Some(0.9));
477 insert_row(&mut engine, 2, "Bob", None);
478 insert_row(&mut engine, 3, "Carol", Some(0.5));
479
480 let snap = engine.export_snapshot(&[], &[]).expect("export");
482
483 assert_eq!(snap.collection, "test_col");
485 assert_eq!(snap.memtable_row_counter, 3);
486 assert_eq!(snap.next_segment_id, 1);
487 assert_eq!(snap.memtable_segment_id, 0);
488 assert_eq!(snap.flushed_segments.len(), 0);
489
490 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
492 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
493
494 let (restored, flushed, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
496 assert!(flushed.is_empty());
497 assert_eq!(restored.next_segment_id(), 1);
498 assert_eq!(restored.memtable_segment_id(), 0);
499 assert_eq!(restored.memtable_row_counter, 3);
500
501 assert_eq!(restored.pk_index().len(), 3);
503
504 let rows: Vec<Vec<Value>> = restored.scan_memtable_rows().collect();
506 assert_eq!(rows.len(), 3);
507 assert_eq!(rows[0][0], Value::Integer(1));
508 assert_eq!(rows[1][1], Value::String("Bob".into()));
509 assert_eq!(rows[2][2], Value::Float(0.5));
510 }
511
512 #[test]
513 fn round_trip_with_tombstone() {
514 let schema = simple_schema();
515 let mut engine = MutationEngine::new("tombstone_test".to_string(), schema);
516
517 insert_row(&mut engine, 10, "Xena", Some(1.0));
518 insert_row(&mut engine, 20, "Yara", Some(2.0));
519
520 engine
523 .delete_bitmap_mut(engine.memtable_segment_id)
524 .mark_deleted(1); let snap = engine.export_snapshot(&[], &[]).expect("export");
527
528 assert!(!snap.memtable_delete_bitmap_bytes.is_empty());
530
531 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
532 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
533 let (restored, _, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
534
535 let rows: Vec<Vec<Value>> = restored.scan_memtable_rows().collect();
537 assert_eq!(rows.len(), 1);
538 assert_eq!(rows[0][0], Value::Integer(10));
539
540 assert!(
542 restored
543 .delete_bitmap(restored.memtable_segment_id())
544 .is_some_and(|bm| bm.is_deleted(1))
545 );
546 }
547
548 #[test]
549 fn round_trip_flushed_segment_blob() {
550 let schema = simple_schema();
551 let engine = MutationEngine::new("flushed_test".to_string(), schema);
552
553 let fake_blob: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0x01, 0x02, 0x03]; let snap = engine
557 .export_snapshot(std::slice::from_ref(&fake_blob), &[])
558 .expect("export");
559 assert_eq!(snap.flushed_segments.len(), 1);
560
561 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
562 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
563 let (_, flushed, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
564
565 assert_eq!(flushed.len(), 1);
566 assert_eq!(flushed[0], fake_blob);
567 }
568
569 #[test]
570 fn schema_mismatch_rejected() {
571 let schema = simple_schema();
572 let engine = MutationEngine::new("mismatch".to_string(), schema);
573 let mut snap = engine.export_snapshot(&[], &[]).expect("export");
574
575 snap.memtable_columns.push(ColumnDataSnapshot::Int64 {
577 values: vec![],
578 valid: None,
579 });
580
581 let result = MutationEngine::from_snapshot(snap);
582 assert!(
583 matches!(
584 result,
585 Err(ColumnarError::SchemaMismatch {
586 expected: 3,
587 got: 4
588 })
589 ),
590 "expected SchemaMismatch error on extra column",
591 );
592 }
593
594 #[test]
595 fn pk_index_survives_round_trip() {
596 let schema = simple_schema();
597 let mut engine = MutationEngine::new("pk_test".to_string(), schema);
598
599 for i in 0..5i64 {
600 insert_row(&mut engine, i, &format!("u{i}"), None);
601 }
602
603 let snap = engine.export_snapshot(&[], &[]).expect("export");
604 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
605 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
606 let (restored, _, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
607
608 assert_eq!(restored.pk_index().len(), 5);
609 for i in 0..5i64 {
610 let pk = crate::pk_index::encode_pk(&Value::Integer(i));
611 assert!(restored.pk_index().contains(&pk), "missing pk {i}");
612 }
613 }
614
615 #[test]
616 fn counters_preserved() {
617 let schema = simple_schema();
618 let mut engine = MutationEngine::new("counters".to_string(), schema);
619
620 insert_row(&mut engine, 99, "Z", Some(2.5));
621
622 engine.next_segment_id = 7;
624 engine.memtable_segment_id = 6;
625
626 let snap = engine.export_snapshot(&[], &[]).expect("export");
627 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
628 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
629 let (restored, _, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
630
631 assert_eq!(restored.next_segment_id, 7);
632 assert_eq!(restored.memtable_segment_id, 6);
633 assert_eq!(restored.memtable_row_counter, 1);
634 }
635
636 #[test]
637 fn flushed_surrogates_survive_round_trip() {
638 let schema = simple_schema();
639 let engine = MutationEngine::new("surr_test".to_string(), schema);
640
641 let blob0: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0xAA];
643 let blob1: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0xBB];
644 let surrogates: FlushedSurrogateTable = vec![
645 vec![Some(Surrogate::new(10)), None, Some(Surrogate::new(12))],
646 vec![Some(Surrogate::new(20))],
647 ];
648
649 let snap = engine
650 .export_snapshot(&[blob0.clone(), blob1.clone()], &surrogates)
651 .expect("export");
652 assert_eq!(snap.flushed_segments.len(), 2);
653 assert_eq!(snap.flushed_surrogates.len(), 2);
654
655 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
656 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
657 let (_, flushed, flushed_surrogates) =
658 MutationEngine::from_snapshot(snap2).expect("from_snapshot");
659
660 assert_eq!(flushed, vec![blob0, blob1]);
661 assert_eq!(flushed_surrogates, surrogates);
662 }
663
664 #[test]
665 fn missing_flushed_surrogates_decode_empty() {
666 let schema = simple_schema();
671 let engine = MutationEngine::new("compat_test".to_string(), schema);
672
673 let blob: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0x01];
674
675 let mut snap = engine
678 .export_snapshot(std::slice::from_ref(&blob), &[])
679 .expect("export");
680 snap.flushed_surrogates.clear();
681
682 let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
683 let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
684
685 assert!(snap2.flushed_surrogates.is_empty());
687 assert_eq!(snap2.flushed_segments.len(), 1);
688
689 let (_, flushed, flushed_surrogates) =
690 MutationEngine::from_snapshot(snap2).expect("from_snapshot");
691 assert_eq!(flushed.len(), 1);
692 assert!(flushed_surrogates.is_empty());
693 }
694}