fix(vector): preserve tombstones through checkpoint and compact doc_id_map

Checkpoint serialization previously skipped deleted vectors entirely,
so on restore those local ids were missing and all subsequent ids shifted
by one — corrupting the HNSW graph's neighbor adjacency. Fix by capturing
the deleted flag for every vector (growing and building segments) and
replaying tombstones via insert_tombstoned / index.delete() on restore.

Separately, compact() previously discarded doc_id_map and multi_doc_map
entries for the segment being compacted. Use compact_with_map() to obtain
the old→new local id remapping and rewrite both maps so that global ids
continue to resolve to the correct document strings after compaction.

Author: farhan-syah

nodedb-vector/src/collection/checkpoint.rs

@@ -7,6 +7,7 @@ use crate::collection::tier::StorageTier;
 use crate::distance::DistanceMetric;
 use crate::flat::FlatIndex;
 use crate::hnsw::{HnswIndex, HnswParams};
+use crate::quantize::pq::PqCodec;
 
 use super::lifecycle::VectorCollection;
 
@@ -33,12 +34,18 @@ pub(crate) struct CollectionSnapshot {
 pub(crate) struct SealedSnapshot {
     pub base_id: u32,
     pub hnsw_bytes: Vec<u8>,
+    #[serde(default)]
+    pub pq_bytes: Option<Vec<u8>>,
+    #[serde(default)]
+    pub pq_codes: Option<Vec<u8>>,
 }
 
 #[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
 pub(crate) struct BuildingSnapshot {
     pub base_id: u32,
     pub vectors: Vec<Vec<f32>>,
+    #[serde(default)]
+    pub deleted: Vec<bool>,
 }
 
 impl VectorCollection {
@@ -53,17 +60,27 @@ impl VectorCollection {
             next_id: self.next_id,
             growing_base_id: self.growing_base_id,
             growing_vectors: (0..self.growing.len() as u32)
-                .filter_map(|i| self.growing.get_vector(i).map(|v| v.to_vec()))
+                .filter_map(|i| self.growing.get_vector_raw(i).map(|v| v.to_vec()))
                 .collect(),
             growing_deleted: (0..self.growing.len() as u32)
-                .map(|i| self.growing.get_vector(i).is_none())
+                .map(|i| self.growing.is_deleted(i))
                 .collect(),
             sealed_segments: self
                 .sealed
                 .iter()
-                .map(|s| SealedSnapshot {
-                    base_id: s.base_id,
-                    hnsw_bytes: s.index.checkpoint_to_bytes(),
+                .map(|s| {
+                    let (pq_bytes, pq_codes) = match &s.pq {
+                        Some((codec, codes)) => {
+                            (zerompk::to_msgpack_vec(codec).ok(), Some(codes.clone()))
+                        }
+                        None => (None, None),
+                    };
+                    SealedSnapshot {
+                        base_id: s.base_id,
+                        hnsw_bytes: s.index.checkpoint_to_bytes(),
+                        pq_bytes,
+                        pq_codes,
+                    }
                 })
                 .collect(),
             building_segments: self
@@ -72,7 +89,10 @@ impl VectorCollection {
                 .map(|b| BuildingSnapshot {
                     base_id: b.base_id,
                     vectors: (0..b.flat.len() as u32)
-                        .filter_map(|i| b.flat.get_vector(i).map(|v| v.to_vec()))
+                        .filter_map(|i| b.flat.get_vector_raw(i).map(|v| v.to_vec()))
+                        .collect(),
+                    deleted: (0..b.flat.len() as u32)
+                        .map(|i| b.flat.is_deleted(i))
                         .collect(),
                 })
                 .collect(),
@@ -118,18 +138,35 @@ impl VectorCollection {
         };
 
         let mut growing = FlatIndex::new(snap.dim, metric);
-        for v in &snap.growing_vectors {
-            growing.insert(v.clone());
+        for (i, v) in snap.growing_vectors.iter().enumerate() {
+            let deleted = snap.growing_deleted.get(i).copied().unwrap_or(false);
+            if deleted {
+                growing.insert_tombstoned(v.clone());
+            } else {
+                growing.insert(v.clone());
+            }
         }
 
         let mut sealed = Vec::with_capacity(snap.sealed_segments.len());
         for ss in &snap.sealed_segments {
             if let Some(index) = HnswIndex::from_checkpoint(&ss.hnsw_bytes) {
-                let sq8 = VectorCollection::build_sq8_for_index(&index);
+                let pq = match (&ss.pq_bytes, &ss.pq_codes) {
+                    (Some(bytes), Some(codes)) => zerompk::from_msgpack::<PqCodec>(bytes)
+                        .ok()
+                        .map(|codec| (codec, codes.clone())),
+                    _ => None,
+                };
+                // Only train SQ8 when PQ isn't present — a segment never carries both.
+                let sq8 = if pq.is_some() {
+                    None
+                } else {
+                    VectorCollection::build_sq8_for_index(&index)
+                };
                 sealed.push(SealedSegment {
                     index,
                     base_id: ss.base_id,
                     sq8,
+                    pq,
                     tier: StorageTier::L0Ram,
                     mmap_vectors: None,
                 });
@@ -143,18 +180,29 @@ impl VectorCollection {
                     .insert(v.clone())
                     .expect("dimension guaranteed by checkpoint");
             }
+            // Replay building-segment tombstones onto the HNSW index.
+            for (i, &dead) in bs.deleted.iter().enumerate() {
+                if dead {
+                    index.delete(i as u32);
+                }
+            }
             let sq8 = VectorCollection::build_sq8_for_index(&index);
             sealed.push(SealedSegment {
                 index,
                 base_id: bs.base_id,
                 sq8,
+                pq: None,
                 tier: StorageTier::L0Ram,
                 mmap_vectors: None,
             });
         }
 
         let next_segment_id = (sealed.len() + 1) as u32;
 
+        let index_config = crate::index_config::IndexConfig {
+            hnsw: params.clone(),
+            ..crate::index_config::IndexConfig::default()
+        };
         Some(Self {
             growing,
             growing_base_id: snap.growing_base_id,
@@ -171,6 +219,7 @@ impl VectorCollection {
             doc_id_map: snap.doc_id_map.into_iter().collect(),
             multi_doc_map: snap.multi_doc_map.into_iter().collect(),
             seal_threshold: DEFAULT_SEAL_THRESHOLD,
+            index_config,
         })
     }
 }

nodedb-vector/src/collection/segment.rs

@@ -4,6 +4,7 @@ use crate::collection::tier::StorageTier;
 use crate::flat::FlatIndex;
 use crate::hnsw::{HnswIndex, HnswParams};
 use crate::mmap_segment::MmapVectorSegment;
+use crate::quantize::pq::PqCodec;
 use crate::quantize::sq8::Sq8Codec;
 
 /// Default threshold for sealing the growing segment.
@@ -44,6 +45,8 @@ pub struct SealedSegment {
     pub base_id: u32,
     /// Optional SQ8 quantized vectors for accelerated traversal.
     pub sq8: Option<(Sq8Codec, Vec<u8>)>,
+    /// Optional PQ-compressed codes (for HnswPq-configured indexes).
+    pub pq: Option<(PqCodec, Vec<u8>)>,
     /// Storage tier: L0Ram = FP32 in HNSW nodes, L1Nvme = FP32 in mmap file.
     pub tier: StorageTier,
     /// mmap-backed vector segment for L1 NVMe tier.

nodedb-vector/tests/collection_checkpoint_tombstones.rs

@@ -0,0 +1,98 @@
+//! Soft-deletes in growing / building segments must survive checkpoint restore.
+//!
+//! Spec: `delete(id)` on a vector in the growing segment (or a building
+//! segment awaiting HNSW completion) tombstones the vector. Checkpoints
+//! MUST serialize that tombstone, and `from_checkpoint` MUST apply it so
+//! the restored collection reports the same `live_count()` and excludes
+//! the deleted vector from `search()` results.
+//!
+//! Today:
+//!   - `FlatIndex::get_vector` returns `Some(..)` even for tombstoned
+//!     slots, so `growing_deleted` is serialized as all-false.
+//!   - `from_checkpoint` ignores the `growing_deleted` field entirely and
+//!     re-inserts every vector as live.
+//!
+//! Result: crash recovery silently resurrects soft-deleted rows — a
+//! correctness regression for any workflow using `valid_until` deletes.
+
+#![cfg(feature = "collection")]
+
+use nodedb_vector::DistanceMetric;
+use nodedb_vector::collection::VectorCollection;
+use nodedb_vector::hnsw::HnswParams;
+
+fn params() -> HnswParams {
+    HnswParams {
+        metric: DistanceMetric::L2,
+        ..HnswParams::default()
+    }
+}
+
+#[test]
+fn growing_segment_tombstones_survive_checkpoint_roundtrip() {
+    let mut coll = VectorCollection::new(2, params());
+    for i in 0..10u32 {
+        coll.insert(vec![i as f32, 0.0]);
+    }
+    assert!(coll.delete(3), "delete on live growing vector must succeed");
+    assert!(coll.delete(7), "delete on live growing vector must succeed");
+    let live_before = coll.live_count();
+    assert_eq!(live_before, 8);
+
+    let bytes = coll.checkpoint_to_bytes();
+    let restored = VectorCollection::from_checkpoint(&bytes).expect("checkpoint deserializes");
+
+    assert_eq!(
+        restored.live_count(),
+        live_before,
+        "tombstoned growing-segment vectors resurrected on restore"
+    );
+
+    let results = restored.search(&[3.0, 0.0], 10, 64);
+    let ids: std::collections::HashSet<u32> = results.iter().map(|r| r.id).collect();
+    assert!(
+        !ids.contains(&3),
+        "soft-deleted id=3 reappeared in search after restore"
+    );
+    assert!(
+        !ids.contains(&7),
+        "soft-deleted id=7 reappeared in search after restore"
+    );
+}
+
+#[test]
+fn building_segment_tombstones_survive_checkpoint_roundtrip() {
+    // Force a seal so the deleted rows live in a building segment at
+    // snapshot time, exercising the `building_segments` encode path.
+    let mut coll = VectorCollection::with_seal_threshold(2, params(), 20);
+    for i in 0..20u32 {
+        coll.insert(vec![i as f32, 0.0]);
+    }
+    let _req = coll.seal("k").expect("seal produced request");
+    // Intentionally do NOT complete the build — vectors now sit in the
+    // building segment as a FlatIndex.
+    assert!(coll.delete(5), "delete on building vector must succeed");
+    assert!(coll.delete(15), "delete on building vector must succeed");
+    let live_before = coll.live_count();
+    assert_eq!(live_before, 18);
+
+    let bytes = coll.checkpoint_to_bytes();
+    let restored = VectorCollection::from_checkpoint(&bytes).expect("checkpoint deserializes");
+
+    assert_eq!(
+        restored.live_count(),
+        live_before,
+        "tombstoned building-segment vectors resurrected on restore"
+    );
+
+    let results = restored.search(&[5.0, 0.0], 20, 64);
+    let ids: std::collections::HashSet<u32> = results.iter().map(|r| r.id).collect();
+    assert!(
+        !ids.contains(&5),
+        "soft-deleted id=5 reappeared after restore"
+    );
+    assert!(
+        !ids.contains(&15),
+        "soft-deleted id=15 reappeared after restore"
+    );
+}

nodedb-vector/tests/collection_compact_doc_map.rs

@@ -0,0 +1,122 @@
+//! `compact()` must keep `doc_id_map` / `multi_doc_map` consistent with the
+//! renumbered HNSW local IDs.
+//!
+//! Spec: `HnswIndex::compact()` removes tombstoned nodes and renumbers
+//! surviving local node ids. The collection stores `doc_id_map` and
+//! `multi_doc_map` keyed on GLOBAL ids (`seg.base_id + local`). After
+//! compaction those globals shift too — the collection MUST walk both
+//! maps and rewrite every entry for the compacted segment to the new
+//! `(seg.base_id + new_local)` globals. Without the rewrite,
+//! `get_doc_id(vid)` and `delete_multi_vector(doc)` point at stale or
+//! wrong vectors.
+
+#![cfg(feature = "collection")]
+
+use nodedb_vector::DistanceMetric;
+use nodedb_vector::collection::VectorCollection;
+use nodedb_vector::hnsw::{HnswIndex, HnswParams};
+
+fn params() -> HnswParams {
+    HnswParams {
+        metric: DistanceMetric::L2,
+        ..HnswParams::default()
+    }
+}
+
+fn build_collection_with_docs() -> VectorCollection {
+    let mut coll = VectorCollection::with_seal_threshold(2, params(), 6);
+    // Six docs, one vector each. Global ids 0..6.
+    for i in 0..6u32 {
+        coll.insert_with_doc_id(vec![i as f32, 0.0], format!("doc_{i}"));
+    }
+    // Seal + complete → sealed segment with base_id=0, local ids 0..6.
+    let req = coll.seal("k").expect("seal produced request");
+    let mut idx = HnswIndex::new(req.dim, req.params.clone());
+    for v in &req.vectors {
+        idx.insert(v.clone()).unwrap();
+    }
+    coll.complete_build(req.segment_id, idx);
+    coll
+}
+
+#[test]
+fn doc_id_map_stays_correct_after_compact() {
+    let mut coll = build_collection_with_docs();
+
+    // Tombstone two vectors in the middle of the sealed segment.
+    assert!(coll.delete(1));
+    assert!(coll.delete(3));
+
+    // Sanity: pre-compact, the surviving doc mapping still resolves.
+    assert_eq!(coll.get_doc_id(0), Some("doc_0"));
+    assert_eq!(coll.get_doc_id(5), Some("doc_5"));
+
+    let removed = coll.compact();
+    assert_eq!(removed, 2, "compact should remove 2 tombstoned nodes");
+
+    // Spec: the search results (identified by renumbered global ids) still
+    // resolve to the original doc strings. For the surviving vectors
+    // {0, 2, 4, 5} post-compact globals become {0, 1, 2, 3}. `get_doc_id`
+    // MUST map those new globals to "doc_0", "doc_2", "doc_4", "doc_5".
+    let results = coll.search(&[0.0, 0.0], 4, 64);
+    let ids: Vec<u32> = results.iter().map(|r| r.id).collect();
+    assert_eq!(ids.len(), 4, "expected 4 live vectors post-compact");
+
+    let observed_docs: std::collections::HashSet<String> = ids
+        .iter()
+        .filter_map(|id| coll.get_doc_id(*id).map(|s| s.to_string()))
+        .collect();
+    let expected_docs: std::collections::HashSet<String> = ["doc_0", "doc_2", "doc_4", "doc_5"]
+        .into_iter()
+        .map(String::from)
+        .collect();
+
+    assert_eq!(
+        observed_docs, expected_docs,
+        "doc_id_map was not rewritten after compact — globals shifted but the map did not"
+    );
+}
+
+#[test]
+fn multi_doc_map_stays_correct_after_compact() {
+    let mut coll = VectorCollection::with_seal_threshold(2, params(), 6);
+
+    // Two multi-vector docs: doc_a owns globals 0,1,2; doc_b owns 3,4,5.
+    let a_vecs: Vec<Vec<f32>> = (0..3u32).map(|i| vec![i as f32, 0.0]).collect();
+    let a_refs: Vec<&[f32]> = a_vecs.iter().map(|v| v.as_slice()).collect();
+    let a_ids = coll.insert_multi_vector(&a_refs, "doc_a".to_string());
+    assert_eq!(a_ids, vec![0, 1, 2]);
+
+    let b_vecs: Vec<Vec<f32>> = (3..6u32).map(|i| vec![i as f32, 0.0]).collect();
+    let b_refs: Vec<&[f32]> = b_vecs.iter().map(|v| v.as_slice()).collect();
+    let b_ids = coll.insert_multi_vector(&b_refs, "doc_b".to_string());
+    assert_eq!(b_ids, vec![3, 4, 5]);
+
+    let req = coll.seal("k").expect("seal produced request");
+    let mut idx = HnswIndex::new(req.dim, req.params.clone());
+    for v in &req.vectors {
+        idx.insert(v.clone()).unwrap();
+    }
+    coll.complete_build(req.segment_id, idx);
+
+    // Tombstone one vector from each doc (middle of each group).
+    assert!(coll.delete(1));
+    assert!(coll.delete(4));
+
+    coll.compact();
+
+    // Spec: `delete_multi_vector("doc_a")` must reach the two remaining
+    // vectors that originally belonged to doc_a, regardless of the local
+    // id renumbering performed by HnswIndex::compact.
+    let deleted_a = coll.delete_multi_vector("doc_a");
+    assert_eq!(
+        deleted_a, 2,
+        "delete_multi_vector(doc_a) must find its 2 remaining vectors after compact"
+    );
+
+    let live_after = coll.live_count();
+    assert_eq!(
+        live_after, 2,
+        "post-compact + doc_a delete: only doc_b's 2 remaining vectors survive"
+    );
+}