use std::collections::HashMap;
use nodedb_types::columnar::ColumnarSchema;
use nodedb_types::surrogate::Surrogate;
use serde::{Deserialize, Serialize};
use zerompk::{FromMessagePack, ToMessagePack};
use crate::delete_bitmap::DeleteBitmap;
use crate::error::ColumnarError;
use crate::memtable::ColumnData;
use crate::pk_index::PkIndex;
use super::engine::MutationEngine;
pub type FlushedSurrogateTable = Vec<Vec<Option<Surrogate>>>;
#[derive(Debug, Clone, Serialize, Deserialize, ToMessagePack, FromMessagePack)]
pub enum ColumnDataSnapshot {
Int64 {
values: Vec<i64>,
valid: Option<Vec<bool>>,
},
Float64 {
values: Vec<f64>,
valid: Option<Vec<bool>>,
},
Bool {
values: Vec<bool>,
valid: Option<Vec<bool>>,
},
Timestamp {
values: Vec<i64>,
valid: Option<Vec<bool>>,
},
Decimal {
values: Vec<[u8; 16]>,
valid: Option<Vec<bool>>,
},
Uuid {
values: Vec<[u8; 16]>,
valid: Option<Vec<bool>>,
},
String {
data: Vec<u8>,
offsets: Vec<u32>,
valid: Option<Vec<bool>>,
},
Bytes {
data: Vec<u8>,
offsets: Vec<u32>,
valid: Option<Vec<bool>>,
},
Json {
data: Vec<u8>,
offsets: Vec<u32>,
valid: Option<Vec<bool>>,
},
Geometry {
data: Vec<u8>,
offsets: Vec<u32>,
valid: Option<Vec<bool>>,
},
Vector {
data: Vec<f32>,
dim: u32,
valid: Option<Vec<bool>>,
},
DictEncoded {
ids: Vec<u32>,
dictionary: Vec<String>,
valid: Option<Vec<bool>>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, ToMessagePack, FromMessagePack)]
#[msgpack(map)]
pub struct ColumnarEngineSnapshot {
pub collection: String,
pub schema: ColumnarSchema,
pub memtable_columns: Vec<ColumnDataSnapshot>,
pub memtable_surrogates: Vec<Option<Surrogate>>,
pub pk_index_bytes: Vec<u8>,
pub delete_bitmaps: Vec<(u64, Vec<u8>)>,
pub memtable_delete_bitmap_bytes: Vec<u8>,
pub flushed_segments: Vec<Vec<u8>>,
#[msgpack(default)]
#[serde(default)]
pub flushed_surrogates: FlushedSurrogateTable,
pub next_segment_id: u64,
pub memtable_segment_id: u64,
pub memtable_row_counter: u32,
}
impl MutationEngine {
pub fn export_snapshot(
&self,
flushed_segments: &[Vec<u8>],
flushed_surrogates: &[Vec<Option<Surrogate>>],
) -> Result<ColumnarEngineSnapshot, ColumnarError> {
let memtable_columns = self
.memtable
.columns()
.iter()
.map(column_to_snapshot)
.collect::<Vec<_>>();
let pk_index_bytes = self.pk_index.to_bytes()?;
let mut delete_bitmaps: Vec<(u64, Vec<u8>)> = Vec::new();
let mut memtable_delete_bitmap_bytes = Vec::new();
for (&seg_id, bitmap) in &self.delete_bitmaps {
let bytes = bitmap.to_bytes()?;
if seg_id == self.memtable_segment_id {
memtable_delete_bitmap_bytes = bytes;
} else {
delete_bitmaps.push((seg_id, bytes));
}
}
Ok(ColumnarEngineSnapshot {
collection: self.collection.clone(),
schema: self.schema.clone(),
memtable_columns,
memtable_surrogates: self.memtable_surrogates.clone(),
pk_index_bytes,
delete_bitmaps,
memtable_delete_bitmap_bytes,
flushed_segments: flushed_segments.to_vec(),
flushed_surrogates: flushed_surrogates.to_vec(),
next_segment_id: self.next_segment_id,
memtable_segment_id: self.memtable_segment_id,
memtable_row_counter: self.memtable_row_counter,
})
}
pub fn from_snapshot(
snap: ColumnarEngineSnapshot,
) -> Result<(MutationEngine, Vec<Vec<u8>>, FlushedSurrogateTable), ColumnarError> {
let col_count = snap.schema.columns.len();
if snap.memtable_columns.len() != col_count {
return Err(ColumnarError::SchemaMismatch {
expected: col_count,
got: snap.memtable_columns.len(),
});
}
let pk_index = PkIndex::from_bytes(&snap.pk_index_bytes)?;
let mut delete_bitmaps: HashMap<u64, DeleteBitmap> = HashMap::new();
for (seg_id, bytes) in snap.delete_bitmaps {
let bm = DeleteBitmap::from_bytes(&bytes)?;
delete_bitmaps.insert(seg_id, bm);
}
if !snap.memtable_delete_bitmap_bytes.is_empty() {
let bm = DeleteBitmap::from_bytes(&snap.memtable_delete_bitmap_bytes)?;
delete_bitmaps.insert(snap.memtable_segment_id, bm);
}
let columns: Vec<ColumnData> = snap
.memtable_columns
.into_iter()
.map(snapshot_to_column)
.collect();
let pk_col_indices: Vec<usize> = snap
.schema
.columns
.iter()
.enumerate()
.filter(|(_, c)| c.primary_key)
.map(|(i, _)| i)
.collect();
let memtable = crate::memtable::ColumnarMemtable::from_raw_columns(
&snap.schema,
columns,
snap.memtable_row_counter as usize,
);
let engine = MutationEngine {
collection: snap.collection,
schema: snap.schema,
memtable,
pk_index,
delete_bitmaps,
pk_col_indices,
next_segment_id: snap.next_segment_id,
memtable_segment_id: snap.memtable_segment_id,
memtable_row_counter: snap.memtable_row_counter,
memtable_surrogates: snap.memtable_surrogates,
};
Ok((engine, snap.flushed_segments, snap.flushed_surrogates))
}
}
fn column_to_snapshot(col: &ColumnData) -> ColumnDataSnapshot {
match col {
ColumnData::Int64 { values, valid } => ColumnDataSnapshot::Int64 {
values: values.clone(),
valid: valid.clone(),
},
ColumnData::Float64 { values, valid } => ColumnDataSnapshot::Float64 {
values: values.clone(),
valid: valid.clone(),
},
ColumnData::Bool { values, valid } => ColumnDataSnapshot::Bool {
values: values.clone(),
valid: valid.clone(),
},
ColumnData::Timestamp { values, valid } => ColumnDataSnapshot::Timestamp {
values: values.clone(),
valid: valid.clone(),
},
ColumnData::Decimal { values, valid } => ColumnDataSnapshot::Decimal {
values: values.clone(),
valid: valid.clone(),
},
ColumnData::Uuid { values, valid } => ColumnDataSnapshot::Uuid {
values: values.clone(),
valid: valid.clone(),
},
ColumnData::String {
data,
offsets,
valid,
} => ColumnDataSnapshot::String {
data: data.clone(),
offsets: offsets.clone(),
valid: valid.clone(),
},
ColumnData::Bytes {
data,
offsets,
valid,
} => ColumnDataSnapshot::Bytes {
data: data.clone(),
offsets: offsets.clone(),
valid: valid.clone(),
},
ColumnData::Json {
data,
offsets,
valid,
} => ColumnDataSnapshot::Json {
data: data.clone(),
offsets: offsets.clone(),
valid: valid.clone(),
},
ColumnData::Geometry {
data,
offsets,
valid,
} => ColumnDataSnapshot::Geometry {
data: data.clone(),
offsets: offsets.clone(),
valid: valid.clone(),
},
ColumnData::Vector { data, dim, valid } => ColumnDataSnapshot::Vector {
data: data.clone(),
dim: *dim,
valid: valid.clone(),
},
ColumnData::DictEncoded {
ids,
dictionary,
valid,
..
} => ColumnDataSnapshot::DictEncoded {
ids: ids.clone(),
dictionary: dictionary.clone(),
valid: valid.clone(),
},
}
}
fn snapshot_to_column(snap: ColumnDataSnapshot) -> ColumnData {
match snap {
ColumnDataSnapshot::Int64 { values, valid } => ColumnData::Int64 { values, valid },
ColumnDataSnapshot::Float64 { values, valid } => ColumnData::Float64 { values, valid },
ColumnDataSnapshot::Bool { values, valid } => ColumnData::Bool { values, valid },
ColumnDataSnapshot::Timestamp { values, valid } => ColumnData::Timestamp { values, valid },
ColumnDataSnapshot::Decimal { values, valid } => ColumnData::Decimal { values, valid },
ColumnDataSnapshot::Uuid { values, valid } => ColumnData::Uuid { values, valid },
ColumnDataSnapshot::String {
data,
offsets,
valid,
} => ColumnData::String {
data,
offsets,
valid,
},
ColumnDataSnapshot::Bytes {
data,
offsets,
valid,
} => ColumnData::Bytes {
data,
offsets,
valid,
},
ColumnDataSnapshot::Json {
data,
offsets,
valid,
} => ColumnData::Json {
data,
offsets,
valid,
},
ColumnDataSnapshot::Geometry {
data,
offsets,
valid,
} => ColumnData::Geometry {
data,
offsets,
valid,
},
ColumnDataSnapshot::Vector { data, dim, valid } => ColumnData::Vector { data, dim, valid },
ColumnDataSnapshot::DictEncoded {
ids,
dictionary,
valid,
} => {
let reverse: HashMap<String, u32> = dictionary
.iter()
.enumerate()
.map(|(i, s)| (s.clone(), i as u32))
.collect();
ColumnData::DictEncoded {
ids,
dictionary,
reverse,
valid,
}
}
}
}
#[cfg(test)]
mod tests {
use nodedb_types::columnar::{ColumnDef, ColumnType, ColumnarSchema};
use nodedb_types::value::Value;
use super::*;
fn simple_schema() -> ColumnarSchema {
ColumnarSchema {
columns: vec![
ColumnDef::required("id", ColumnType::Int64).with_primary_key(),
ColumnDef::required("name", ColumnType::String),
ColumnDef::nullable("score", ColumnType::Float64),
],
version: 1,
}
}
fn insert_row(engine: &mut MutationEngine, id: i64, name: &str, score: Option<f64>) {
let score_val = score.map(Value::Float).unwrap_or(Value::Null);
engine
.insert(&[Value::Integer(id), Value::String(name.into()), score_val])
.expect("insert");
}
#[test]
fn round_trip_memtable_rows() {
let schema = simple_schema();
let mut engine = MutationEngine::new("test_col".to_string(), schema);
insert_row(&mut engine, 1, "Alice", Some(0.9));
insert_row(&mut engine, 2, "Bob", None);
insert_row(&mut engine, 3, "Carol", Some(0.5));
let snap = engine.export_snapshot(&[], &[]).expect("export");
assert_eq!(snap.collection, "test_col");
assert_eq!(snap.memtable_row_counter, 3);
assert_eq!(snap.next_segment_id, 1);
assert_eq!(snap.memtable_segment_id, 0);
assert_eq!(snap.flushed_segments.len(), 0);
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
let (restored, flushed, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
assert!(flushed.is_empty());
assert_eq!(restored.next_segment_id(), 1);
assert_eq!(restored.memtable_segment_id(), 0);
assert_eq!(restored.memtable_row_counter, 3);
assert_eq!(restored.pk_index().len(), 3);
let rows: Vec<Vec<Value>> = restored.scan_memtable_rows().collect();
assert_eq!(rows.len(), 3);
assert_eq!(rows[0][0], Value::Integer(1));
assert_eq!(rows[1][1], Value::String("Bob".into()));
assert_eq!(rows[2][2], Value::Float(0.5));
}
#[test]
fn round_trip_with_tombstone() {
let schema = simple_schema();
let mut engine = MutationEngine::new("tombstone_test".to_string(), schema);
insert_row(&mut engine, 10, "Xena", Some(1.0));
insert_row(&mut engine, 20, "Yara", Some(2.0));
engine
.delete_bitmap_mut(engine.memtable_segment_id)
.mark_deleted(1);
let snap = engine.export_snapshot(&[], &[]).expect("export");
assert!(!snap.memtable_delete_bitmap_bytes.is_empty());
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
let (restored, _, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
let rows: Vec<Vec<Value>> = restored.scan_memtable_rows().collect();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0][0], Value::Integer(10));
assert!(
restored
.delete_bitmap(restored.memtable_segment_id())
.is_some_and(|bm| bm.is_deleted(1))
);
}
#[test]
fn round_trip_flushed_segment_blob() {
let schema = simple_schema();
let engine = MutationEngine::new("flushed_test".to_string(), schema);
let fake_blob: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0x01, 0x02, 0x03];
let snap = engine
.export_snapshot(std::slice::from_ref(&fake_blob), &[])
.expect("export");
assert_eq!(snap.flushed_segments.len(), 1);
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
let (_, flushed, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
assert_eq!(flushed.len(), 1);
assert_eq!(flushed[0], fake_blob);
}
#[test]
fn schema_mismatch_rejected() {
let schema = simple_schema();
let engine = MutationEngine::new("mismatch".to_string(), schema);
let mut snap = engine.export_snapshot(&[], &[]).expect("export");
snap.memtable_columns.push(ColumnDataSnapshot::Int64 {
values: vec![],
valid: None,
});
let result = MutationEngine::from_snapshot(snap);
assert!(
matches!(
result,
Err(ColumnarError::SchemaMismatch {
expected: 3,
got: 4
})
),
"expected SchemaMismatch error on extra column",
);
}
#[test]
fn pk_index_survives_round_trip() {
let schema = simple_schema();
let mut engine = MutationEngine::new("pk_test".to_string(), schema);
for i in 0..5i64 {
insert_row(&mut engine, i, &format!("u{i}"), None);
}
let snap = engine.export_snapshot(&[], &[]).expect("export");
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
let (restored, _, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
assert_eq!(restored.pk_index().len(), 5);
for i in 0..5i64 {
let pk = crate::pk_index::encode_pk(&Value::Integer(i));
assert!(restored.pk_index().contains(&pk), "missing pk {i}");
}
}
#[test]
fn counters_preserved() {
let schema = simple_schema();
let mut engine = MutationEngine::new("counters".to_string(), schema);
insert_row(&mut engine, 99, "Z", Some(2.5));
engine.next_segment_id = 7;
engine.memtable_segment_id = 6;
let snap = engine.export_snapshot(&[], &[]).expect("export");
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
let (restored, _, _) = MutationEngine::from_snapshot(snap2).expect("from_snapshot");
assert_eq!(restored.next_segment_id, 7);
assert_eq!(restored.memtable_segment_id, 6);
assert_eq!(restored.memtable_row_counter, 1);
}
#[test]
fn flushed_surrogates_survive_round_trip() {
let schema = simple_schema();
let engine = MutationEngine::new("surr_test".to_string(), schema);
let blob0: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0xAA];
let blob1: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0xBB];
let surrogates: FlushedSurrogateTable = vec![
vec![Some(Surrogate::new(10)), None, Some(Surrogate::new(12))],
vec![Some(Surrogate::new(20))],
];
let snap = engine
.export_snapshot(&[blob0.clone(), blob1.clone()], &surrogates)
.expect("export");
assert_eq!(snap.flushed_segments.len(), 2);
assert_eq!(snap.flushed_surrogates.len(), 2);
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
let (_, flushed, flushed_surrogates) =
MutationEngine::from_snapshot(snap2).expect("from_snapshot");
assert_eq!(flushed, vec![blob0, blob1]);
assert_eq!(flushed_surrogates, surrogates);
}
#[test]
fn missing_flushed_surrogates_decode_empty() {
let schema = simple_schema();
let engine = MutationEngine::new("compat_test".to_string(), schema);
let blob: Vec<u8> = vec![0x4E, 0x44, 0x42, 0x53, 0x01];
let mut snap = engine
.export_snapshot(std::slice::from_ref(&blob), &[])
.expect("export");
snap.flushed_surrogates.clear();
let bytes = zerompk::to_msgpack_vec(&snap).expect("serialize");
let snap2: ColumnarEngineSnapshot = zerompk::from_msgpack(&bytes).expect("deserialize");
assert!(snap2.flushed_surrogates.is_empty());
assert_eq!(snap2.flushed_segments.len(), 1);
let (_, flushed, flushed_surrogates) =
MutationEngine::from_snapshot(snap2).expect("from_snapshot");
assert_eq!(flushed.len(), 1);
assert!(flushed_surrogates.is_empty());
}
}