Skip to main content

nodedb_vector/collection/
checkpoint.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Checkpoint serialization and deserialization for `VectorCollection`.
4//!
5//! ## On-disk framing
6//!
7//! **Plaintext** checkpoints are raw MessagePack bytes (existing format).
8//! The first 4 bytes are never `SEGV`, so detection is unambiguous.
9//!
10//! **Encrypted** checkpoints use the following layout:
11//!
12//! ```text
13//! [SEGV (4B)] [version_u16_le (2B)] [cipher_alg_u8 (1B)] [kid_u8 (1B)]
14//! [epoch (4B)] [reserved (4B)] [AES-256-GCM ciphertext of msgpack payload]
15//! ```
16//!
17//! The first 16 bytes form a `SegmentPreamble` (reusing the existing preamble
18//! layout with a distinct `SEGV` magic). These 16 bytes are included as AAD,
19//! preventing preamble-swap attacks. The nonce is `(epoch, lsn=0)` — epoch
20//! provides per-write uniqueness even without an LSN.
21
22use 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
39/// Magic bytes identifying an encrypted vector checkpoint. Shared with
40/// `nodedb-spatial`'s SEGV checkpoint format.
41const SEGV_MAGIC: [u8; 4] = *b"SEGV";
42
43/// Encrypt `plaintext` into the SEGV envelope from [`nodedb_wal::crypto`].
44fn 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
55/// Decrypt an encrypted checkpoint blob (starting at byte 0, which is `SEGV`).
56fn 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    /// `(global_vector_id, surrogate_u32)` pairs.
81    #[serde(default)]
82    pub surrogate_map: Vec<(u32, u32)>,
83    /// `(document_surrogate_u32, [global_vector_ids])` pairs.
84    #[serde(default)]
85    pub multi_doc_map: Vec<(u32, Vec<u32>)>,
86    /// Quantization mode for the collection-level codec-dispatch index.
87    /// Serialised as a u8 matching `VectorQuantization` discriminants.
88    /// 0 = None (default, backward-compatible).
89    #[serde(default)]
90    pub quantization_tag: u8,
91    /// Serialised `PayloadIndexSetSnapshot` (msgpack bytes).
92    /// Empty vec = no payload indexes (default, backward-compatible).
93    #[serde(default)]
94    pub payload_index_bytes: Vec<u8>,
95    /// Highest WAL LSN whose write is included in this checkpoint. Startup WAL
96    /// replay skips records at or below it for this collection so a
97    /// straddling-segment record already absorbed here is not re-applied.
98    /// `0` (default) on legacy checkpoints predating this field — those gate
99    /// nothing and replay everything, exactly as before.
100    #[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    /// Serialized [`Sq8Codec`] bytes (magic + version + msgpack).
113    /// Present when SQ8 quantization is active and PQ is absent.
114    /// `None` means no SQ8 quantization for this segment.
115    #[serde(default)]
116    pub sq8_bytes: Option<Vec<u8>>,
117    /// Pre-quantized SQ8 codes for all vectors in this segment.
118    /// Layout: `[v0_d0, v0_d1, ..., v1_d0, ...]` (dim bytes per vector).
119    /// `None` when SQ8 is not configured.
120    #[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    /// Serialize all segments for checkpointing.
134    ///
135    /// When `kek` is `Some`, the MessagePack payload is wrapped in an
136    /// AES-256-GCM encrypted envelope with a `SEGV` preamble. When `None`,
137    /// raw MessagePack bytes are returned (existing plaintext format).
138    ///
139    /// Returns `Err` on serialization or encryption failure — callers must
140    /// propagate rather than treat a failed encode as an empty checkpoint.
141    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                    // Only serialize SQ8 when PQ is absent — a segment never carries both.
175                    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    /// Restore a collection from checkpoint bytes.
247    ///
248    /// `kek` controls the expected framing:
249    /// - `None` → the file must be plaintext MessagePack (starting with bytes
250    ///   that are NOT `SEGV`). If the file starts with `SEGV` and no key is
251    ///   provided, returns `Err(CheckpointEncryptedNoKey)`.
252    /// - `Some(key)` → encryption is **required**. If the file starts with
253    ///   `SEGV`, it is decrypted with `key`. If the file is plaintext, returns
254    ///   `Err(CheckpointPlaintextKeyRequired)` — refuse to silently load
255    ///   unencrypted data when the operator has enabled at-rest encryption.
256    ///
257    /// Returns `Err(CheckpointDeserializationError)` if the (decrypted)
258    /// MessagePack payload, an embedded sealed-segment HNSW checkpoint, or an
259    /// embedded PQ/SQ8 codec is structurally corrupt.
260    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            // Sealed segments always carry the `RKHNS\0` magic — we wrote
322            // them ourselves. Both a decode `Err` and an `Ok(None)` (magic
323            // missing) mean the checkpoint is structurally corrupt; either
324            // way the whole restore must fail rather than silently drop the
325            // segment's vectors.
326            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            // Restore SQ8 from persisted bytes — never recompute on load.
343            // A segment never carries both PQ and SQ8.
344            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            // Replay building-segment tombstones onto the HNSW index.
379            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
447/// Encode a `VectorQuantization` to a u8 tag for storage.
448fn 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
462/// Decode a u8 tag back to `VectorQuantization`.
463fn 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    /// SQ8 calibration data must survive a checkpoint round-trip without
484    /// being recomputed. Verifies that the O(N*dim) rebuild-on-restart bug
485    /// is eliminated: `sq8` on the restored sealed segment is `Some` and
486    /// its `inv_scales` match the original exactly.
487    #[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        // 1024 vectors of dim=8 — enough to pass the ≥1000 threshold in
497        // `build_sq8_for_index`, so sq8 will be Some after complete_build.
498        // Uses plain HNSW (default IndexType::Hnsw) so SQ8 is selected.
499        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        // Capture serialized form as ground truth.
522        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        // Byte-level equality guarantees calibration data is persisted, not recomputed.
536        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    /// The checkpoint watermark (`checkpoint_wal_lsn`) must survive a
565    /// checkpoint round-trip. `note_checkpoint_lsn` only feeds the running
566    /// `applied_wal_lsn` max — it does NOT move the replay gate. The gate
567    /// moves only at checkpoint save (which folds `applied_wal_lsn` into the
568    /// persisted watermark via `max`) and at checkpoint load (which exposes
569    /// that persisted watermark as the new gate). A legacy checkpoint
570    /// predating the field decodes to `0` (via `#[serde(default)]`), which
571    /// gates nothing.
572    #[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    /// `note_checkpoint_lsn` is monotonic and ignores `0` (an unassigned LSN
600    /// must never move the watermark, or a stale-low checkpoint would fail to
601    /// gate a straddling replay). This behavior lives on `applied_wal_lsn`,
602    /// the running max `note_checkpoint_lsn` advances.
603    #[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); // lower — ignored
609        assert_eq!(coll.applied_wal_lsn(), 10);
610        coll.note_checkpoint_lsn(0); // unassigned — ignored
611        assert_eq!(coll.applied_wal_lsn(), 10);
612        coll.note_checkpoint_lsn(20); // higher — advances
613        assert_eq!(coll.applied_wal_lsn(), 20);
614    }
615
616    /// Payload bitmap indexes registered on a vector-primary collection
617    /// must survive a checkpoint round-trip — otherwise `WHERE` filters
618    /// would silently return zero rows after a node restart.
619    #[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}