1use std::collections::HashMap;
23
24use nodedb_types::{Surrogate, VectorQuantization};
25use serde::{Deserialize, Serialize};
26
27use crate::collection::payload_index::PayloadIndexSetSnapshot;
28use crate::collection::segment::{DEFAULT_SEAL_THRESHOLD, SealedSegment};
29use crate::collection::tier::StorageTier;
30use crate::distance::DistanceMetric;
31use crate::error::VectorError;
32use crate::flat::FlatIndex;
33use crate::hnsw::{HnswIndex, HnswParams};
34use crate::quantize::pq::PqCodec;
35use crate::quantize::sq8::Sq8Codec;
36
37use super::lifecycle::VectorCollection;
38
39const SEGV_MAGIC: [u8; 4] = *b"SEGV";
42
43fn encrypt_checkpoint(
45 key: &nodedb_wal::crypto::WalEncryptionKey,
46 plaintext: &[u8],
47) -> Result<Vec<u8>, VectorError> {
48 nodedb_wal::crypto::encrypt_segment_envelope(key, &SEGV_MAGIC, plaintext).map_err(|e| {
49 VectorError::CheckpointEncryptionError {
50 detail: e.to_string(),
51 }
52 })
53}
54
55fn decrypt_checkpoint(
57 key: &nodedb_wal::crypto::WalEncryptionKey,
58 blob: &[u8],
59) -> Result<Vec<u8>, VectorError> {
60 nodedb_wal::crypto::decrypt_segment_envelope(key, &SEGV_MAGIC, blob).map_err(|e| {
61 VectorError::CheckpointEncryptionError {
62 detail: e.to_string(),
63 }
64 })
65}
66
67#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
68pub(crate) struct CollectionSnapshot {
69 pub dim: usize,
70 pub params_m: usize,
71 pub params_m0: usize,
72 pub params_ef_construction: usize,
73 pub params_metric: u8,
74 pub next_id: u32,
75 pub growing_base_id: u32,
76 pub growing_vectors: Vec<Vec<f32>>,
77 pub growing_deleted: Vec<bool>,
78 pub sealed_segments: Vec<SealedSnapshot>,
79 pub building_segments: Vec<BuildingSnapshot>,
80 #[serde(default)]
82 pub surrogate_map: Vec<(u32, u32)>,
83 #[serde(default)]
85 pub multi_doc_map: Vec<(u32, Vec<u32>)>,
86 #[serde(default)]
90 pub quantization_tag: u8,
91 #[serde(default)]
94 pub payload_index_bytes: Vec<u8>,
95 #[serde(default)]
101 pub checkpoint_wal_lsn: u64,
102}
103
104#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
105pub(crate) struct SealedSnapshot {
106 pub base_id: u32,
107 pub hnsw_bytes: Vec<u8>,
108 #[serde(default)]
109 pub pq_bytes: Option<Vec<u8>>,
110 #[serde(default)]
111 pub pq_codes: Option<Vec<u8>>,
112 #[serde(default)]
116 pub sq8_bytes: Option<Vec<u8>>,
117 #[serde(default)]
121 pub sq8_codes: Option<Vec<u8>>,
122}
123
124#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
125pub(crate) struct BuildingSnapshot {
126 pub base_id: u32,
127 pub vectors: Vec<Vec<f32>>,
128 #[serde(default)]
129 pub deleted: Vec<bool>,
130}
131
132impl VectorCollection {
133 pub fn checkpoint_to_bytes(
142 &self,
143 kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
144 ) -> Result<Vec<u8>, VectorError> {
145 let snapshot = CollectionSnapshot {
146 dim: self.dim,
147 params_m: self.params.m,
148 params_m0: self.params.m0,
149 params_ef_construction: self.params.ef_construction,
150 params_metric: self.params.metric as u8,
151 next_id: self.next_id,
152 growing_base_id: self.growing_base_id,
153 growing_vectors: (0..self.growing.len() as u32)
154 .filter_map(|i| self.growing.get_vector_raw(i).map(|v| v.to_vec()))
155 .collect(),
156 growing_deleted: (0..self.growing.len() as u32)
157 .map(|i| self.growing.is_deleted(i))
158 .collect(),
159 sealed_segments: self
160 .sealed
161 .iter()
162 .map(|s| {
163 let (pq_bytes, pq_codes) = match &s.pq {
164 Some((codec, codes)) => (
165 Some(codec.to_bytes().map_err(|e| {
166 VectorError::CheckpointSerializationError {
167 detail: format!("PQ codec encode: {e}"),
168 }
169 })?),
170 Some(codes.clone()),
171 ),
172 None => (None, None),
173 };
174 let (sq8_bytes, sq8_codes) = if pq_bytes.is_none() {
176 match &s.sq8 {
177 Some((codec, codes)) => (Some(codec.to_bytes()), Some(codes.clone())),
178 None => (None, None),
179 }
180 } else {
181 (None, None)
182 };
183 Ok(SealedSnapshot {
184 base_id: s.base_id,
185 hnsw_bytes: s.index.checkpoint_to_bytes()?,
186 pq_bytes,
187 pq_codes,
188 sq8_bytes,
189 sq8_codes,
190 })
191 })
192 .collect::<Result<Vec<_>, VectorError>>()?,
193 building_segments: self
194 .building
195 .iter()
196 .map(|b| BuildingSnapshot {
197 base_id: b.base_id,
198 vectors: (0..b.flat.len() as u32)
199 .filter_map(|i| b.flat.get_vector_raw(i).map(|v| v.to_vec()))
200 .collect(),
201 deleted: (0..b.flat.len() as u32)
202 .map(|i| b.flat.is_deleted(i))
203 .collect(),
204 })
205 .collect(),
206 surrogate_map: self
207 .surrogate_map
208 .iter()
209 .map(|(&k, s)| (k, s.as_u32()))
210 .collect(),
211 multi_doc_map: self
212 .multi_doc_map
213 .iter()
214 .map(|(k, v)| (k.as_u32(), v.clone()))
215 .collect(),
216 quantization_tag: quantization_to_tag(self.quantization),
217 payload_index_bytes: {
218 let snap = self.payload.to_snapshot();
219 match zerompk::to_msgpack_vec(&snap) {
220 Ok(bytes) => bytes,
221 Err(e) => {
222 return Err(VectorError::CheckpointSerializationError {
223 detail: e.to_string(),
224 });
225 }
226 }
227 },
228 checkpoint_wal_lsn: self.checkpoint_wal_lsn.max(self.applied_wal_lsn),
229 };
230 let msgpack = match zerompk::to_msgpack_vec(&snapshot) {
231 Ok(bytes) => bytes,
232 Err(e) => {
233 return Err(VectorError::CheckpointSerializationError {
234 detail: e.to_string(),
235 });
236 }
237 };
238
239 if let Some(key) = kek {
240 encrypt_checkpoint(key, &msgpack)
241 } else {
242 Ok(msgpack)
243 }
244 }
245
246 pub fn from_checkpoint(
261 bytes: &[u8],
262 kek: Option<&nodedb_wal::crypto::WalEncryptionKey>,
263 ) -> Result<Self, VectorError> {
264 let is_encrypted = bytes.len() >= 4 && bytes[0..4] == SEGV_MAGIC;
265
266 let msgpack: Vec<u8>;
267 let msgpack_ref: &[u8];
268
269 if is_encrypted {
270 if let Some(key) = kek {
271 msgpack = decrypt_checkpoint(key, bytes)?;
272 msgpack_ref = &msgpack;
273 } else {
274 return Err(VectorError::CheckpointEncryptedNoKey);
275 }
276 } else if kek.is_some() {
277 return Err(VectorError::CheckpointPlaintextKeyRequired);
278 } else {
279 msgpack_ref = bytes;
280 }
281
282 let snap: CollectionSnapshot = match zerompk::from_msgpack(msgpack_ref) {
283 Ok(s) => s,
284 Err(e) => {
285 return Err(VectorError::CheckpointDeserializationError {
286 detail: format!("collection msgpack decode: {e}"),
287 });
288 }
289 };
290 let metric = match snap.params_metric {
291 0 => DistanceMetric::L2,
292 1 => DistanceMetric::Cosine,
293 2 => DistanceMetric::InnerProduct,
294 3 => DistanceMetric::Manhattan,
295 4 => DistanceMetric::Chebyshev,
296 5 => DistanceMetric::Hamming,
297 6 => DistanceMetric::Jaccard,
298 7 => DistanceMetric::Pearson,
299 _ => DistanceMetric::Cosine,
300 };
301 let params = HnswParams {
302 m: snap.params_m,
303 m0: snap.params_m0,
304 ef_construction: snap.params_ef_construction,
305 metric,
306 dtype: nodedb_types::vector_dtype::VectorStorageDtype::F32,
307 };
308
309 let mut growing = FlatIndex::new(snap.dim, metric);
310 for (i, v) in snap.growing_vectors.iter().enumerate() {
311 let deleted = snap.growing_deleted.get(i).copied().unwrap_or(false);
312 if deleted {
313 growing.insert_tombstoned(v.clone());
314 } else {
315 growing.insert(v.clone());
316 }
317 }
318
319 let mut sealed = Vec::with_capacity(snap.sealed_segments.len());
320 for ss in &snap.sealed_segments {
321 let index = HnswIndex::from_checkpoint(&ss.hnsw_bytes)?.ok_or_else(|| {
327 VectorError::CheckpointDeserializationError {
328 detail: "sealed segment hnsw bytes missing RKHNS magic".into(),
329 }
330 })?;
331 let pq = match (&ss.pq_bytes, &ss.pq_codes) {
332 (Some(bytes), Some(codes)) => Some((
333 PqCodec::from_bytes(bytes).map_err(|e| {
334 VectorError::CheckpointDeserializationError {
335 detail: format!("PQ codec decode: {e}"),
336 }
337 })?,
338 codes.clone(),
339 )),
340 _ => None,
341 };
342 let sq8 = if pq.is_some() {
345 None
346 } else {
347 match (&ss.sq8_bytes, &ss.sq8_codes) {
348 (Some(codec_bytes), Some(codes)) => Some((
349 Sq8Codec::from_bytes(codec_bytes).map_err(|e| {
350 VectorError::CheckpointDeserializationError {
351 detail: format!("SQ8 codec decode: {e}"),
352 }
353 })?,
354 codes.clone(),
355 )),
356 _ => None,
357 }
358 };
359 sealed.push(SealedSegment {
360 index,
361 base_id: ss.base_id,
362 sq8,
363 pq,
364 tier: StorageTier::L0Ram,
365 mmap_vectors: None,
366 });
367 }
368
369 for bs in &snap.building_segments {
370 let mut index = HnswIndex::new(snap.dim, params.clone());
371 for v in &bs.vectors {
372 index.insert(v.clone()).map_err(|e| {
373 VectorError::CheckpointDeserializationError {
374 detail: format!("building-segment replay insert: {e}"),
375 }
376 })?;
377 }
378 for (i, &dead) in bs.deleted.iter().enumerate() {
380 if dead {
381 index.delete(i as u32);
382 }
383 }
384 let sq8 = VectorCollection::build_sq8_for_index(&index);
385 sealed.push(SealedSegment {
386 index,
387 base_id: bs.base_id,
388 sq8,
389 pq: None,
390 tier: StorageTier::L0Ram,
391 mmap_vectors: None,
392 });
393 }
394
395 let next_segment_id = (sealed.len() + 1) as u32;
396
397 let index_config = crate::index_config::IndexConfig {
398 hnsw: params.clone(),
399 ..crate::index_config::IndexConfig::default()
400 };
401 Ok(Self {
402 growing,
403 growing_base_id: snap.growing_base_id,
404 sealed,
405 building: Vec::new(),
406 params,
407 next_id: snap.next_id,
408 next_segment_id,
409 dim: snap.dim,
410 data_dir: None,
411 ram_budget_bytes: 0,
412 mmap_fallback_count: 0,
413 mmap_segment_count: 0,
414 surrogate_map: snap
415 .surrogate_map
416 .iter()
417 .map(|&(k, s)| (k, Surrogate::new(s)))
418 .collect(),
419 surrogate_to_local: snap
420 .surrogate_map
421 .iter()
422 .map(|&(k, s)| (Surrogate::new(s), k))
423 .collect(),
424 multi_doc_map: snap
425 .multi_doc_map
426 .into_iter()
427 .map(|(k, v)| (Surrogate::new(k), v))
428 .collect::<HashMap<_, _>>(),
429 seal_threshold: DEFAULT_SEAL_THRESHOLD,
430 index_config,
431 codec_dispatch: None,
432 quantization: quantization_from_tag(snap.quantization_tag),
433 payload: if snap.payload_index_bytes.is_empty() {
434 super::payload_index::PayloadIndexSet::default()
435 } else {
436 zerompk::from_msgpack::<PayloadIndexSetSnapshot>(&snap.payload_index_bytes)
437 .map(super::payload_index::PayloadIndexSet::from_snapshot)
438 .unwrap_or_default()
439 },
440 arena_index: None,
441 checkpoint_wal_lsn: snap.checkpoint_wal_lsn,
442 applied_wal_lsn: snap.checkpoint_wal_lsn,
443 })
444 }
445}
446
447fn quantization_to_tag(q: VectorQuantization) -> u8 {
449 match q {
450 VectorQuantization::None => 0,
451 VectorQuantization::Sq8 => 1,
452 VectorQuantization::Pq => 2,
453 VectorQuantization::RaBitQ => 3,
454 VectorQuantization::Bbq => 4,
455 VectorQuantization::Binary => 5,
456 VectorQuantization::Ternary => 6,
457 VectorQuantization::Opq => 7,
458 _ => 0,
459 }
460}
461
462fn quantization_from_tag(tag: u8) -> VectorQuantization {
464 match tag {
465 0 => VectorQuantization::None,
466 1 => VectorQuantization::Sq8,
467 2 => VectorQuantization::Pq,
468 3 => VectorQuantization::RaBitQ,
469 4 => VectorQuantization::Bbq,
470 5 => VectorQuantization::Binary,
471 6 => VectorQuantization::Ternary,
472 7 => VectorQuantization::Opq,
473 _ => VectorQuantization::None,
474 }
475}
476
477#[cfg(test)]
478mod tests {
479 use crate::collection::lifecycle::VectorCollection;
480 use crate::distance::DistanceMetric;
481 use crate::hnsw::HnswParams;
482
483 #[test]
488 fn checkpoint_roundtrip_preserves_sq8() {
489 use crate::collection::lifecycle::VectorCollection;
490 use crate::hnsw::{HnswIndex, HnswParams};
491
492 let params = HnswParams {
493 metric: crate::distance::DistanceMetric::L2,
494 ..HnswParams::default()
495 };
496 let mut coll = VectorCollection::with_seal_threshold(8, params, 1024);
500 for i in 0..1024u32 {
501 let mut v = vec![0.0f32; 8];
502 for (d, slot) in v.iter_mut().enumerate() {
503 *slot = ((i as f32) * 0.01 + (d as f32) * 0.1).sin();
504 }
505 coll.insert(v);
506 }
507 let req = coll.seal("sq8_test").expect("seal produced request");
508 let mut idx = HnswIndex::new(req.dim, req.params.clone());
509 for v in &req.vectors {
510 idx.insert(v.clone()).unwrap();
511 }
512 coll.complete_build(req.segment_id, idx);
513
514 let sealed = coll.sealed_segments();
515 assert!(!sealed.is_empty(), "expected at least one sealed segment");
516 let orig_sq8 = sealed[0]
517 .sq8
518 .as_ref()
519 .expect("sq8 must be Some after complete_build with ≥1000 vectors");
520 let orig_dim = orig_sq8.0.dim();
521 let orig_bytes = orig_sq8.0.to_bytes();
523
524 let checkpoint = coll.checkpoint_to_bytes(None).unwrap();
525 let restored = VectorCollection::from_checkpoint(&checkpoint, None).unwrap();
526
527 let restored_sealed = restored.sealed_segments();
528 assert!(!restored_sealed.is_empty());
529 let restored_sq8 = restored_sealed[0]
530 .sq8
531 .as_ref()
532 .expect("sq8 must be Some after restoring checkpoint — never recomputed");
533
534 assert_eq!(restored_sq8.0.dim(), orig_dim, "dim mismatch after restore");
535 assert_eq!(
537 restored_sq8.0.to_bytes(),
538 orig_bytes,
539 "sq8 codec bytes differ — calibration data was recomputed rather than persisted"
540 );
541 }
542
543 #[test]
544 fn checkpoint_roundtrip() {
545 let mut coll = VectorCollection::new(
546 3,
547 HnswParams {
548 metric: DistanceMetric::L2,
549 ..HnswParams::default()
550 },
551 );
552 for i in 0..50u32 {
553 coll.insert(vec![i as f32, 0.0, 0.0]);
554 }
555 let bytes = coll.checkpoint_to_bytes(None).unwrap();
556 let restored = VectorCollection::from_checkpoint(&bytes, None).unwrap();
557 assert_eq!(restored.len(), 50);
558 assert_eq!(restored.dim(), 3);
559
560 let results = restored.search(&[25.0, 0.0, 0.0], 1, 64);
561 assert_eq!(results[0].id, 25);
562 }
563
564 #[test]
573 fn checkpoint_roundtrip_preserves_wal_lsn() {
574 let mut coll = VectorCollection::new(
575 3,
576 HnswParams {
577 metric: DistanceMetric::L2,
578 ..HnswParams::default()
579 },
580 );
581 coll.insert(vec![1.0, 0.0, 0.0]);
582 coll.note_checkpoint_lsn(42);
583 assert_eq!(coll.applied_wal_lsn(), 42);
584 assert_eq!(
585 coll.checkpoint_wal_lsn(),
586 0,
587 "gate is not advanced by note; only by load/save"
588 );
589
590 let bytes = coll.checkpoint_to_bytes(None).unwrap();
591 let restored = VectorCollection::from_checkpoint(&bytes, None).unwrap();
592 assert_eq!(
593 restored.checkpoint_wal_lsn(),
594 42,
595 "save folds applied_wal_lsn into the persisted gate, so restore exposes it"
596 );
597 }
598
599 #[test]
604 fn note_checkpoint_lsn_is_monotonic_and_ignores_zero() {
605 let mut coll = VectorCollection::new(3, HnswParams::default());
606 coll.note_checkpoint_lsn(10);
607 assert_eq!(coll.applied_wal_lsn(), 10);
608 coll.note_checkpoint_lsn(5); assert_eq!(coll.applied_wal_lsn(), 10);
610 coll.note_checkpoint_lsn(0); assert_eq!(coll.applied_wal_lsn(), 10);
612 coll.note_checkpoint_lsn(20); assert_eq!(coll.applied_wal_lsn(), 20);
614 }
615
616 #[test]
620 fn checkpoint_roundtrip_preserves_payload_bitmap() {
621 use crate::collection::PayloadIndexKind;
622 use crate::collection::payload_index::FilterPredicate;
623 use nodedb_types::Value;
624 use std::collections::HashMap;
625
626 let mut coll = VectorCollection::new(
627 3,
628 HnswParams {
629 metric: DistanceMetric::L2,
630 ..HnswParams::default()
631 },
632 );
633 coll.payload
634 .add_index("category".to_string(), PayloadIndexKind::Equality);
635 for i in 0u32..10 {
636 let node_id = coll.insert(vec![i as f32, 0.0, 0.0]);
637 let mut fields = HashMap::new();
638 let cat = if i % 2 == 0 { "A" } else { "B" };
639 fields.insert("category".to_string(), Value::String(cat.to_string()));
640 coll.payload.insert_row(node_id, &fields);
641 }
642
643 let bytes = coll.checkpoint_to_bytes(None).unwrap();
644 let restored = VectorCollection::from_checkpoint(&bytes, None).unwrap();
645
646 let pred = FilterPredicate::Eq {
647 field: "category".to_string(),
648 value: Value::String("A".to_string()),
649 };
650 let bm = restored
651 .payload
652 .pre_filter(&pred)
653 .expect("payload index 'category' must be present after restore");
654 assert_eq!(
655 bm.len(),
656 5,
657 "5 rows of category=A must survive checkpoint round-trip"
658 );
659 }
660}