Skip to main content

nodedb_vector/collection/
lifecycle_insert_ops.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Insert, delete, and surrogate-map operations for `VectorCollection`.
4
5use nodedb_types::Surrogate;
6
7use super::lifecycle::VectorCollection;
8
9impl VectorCollection {
10    /// Insert a vector. Returns the global vector ID.
11    pub fn insert(&mut self, vector: Vec<f32>) -> u32 {
12        let id = self.next_id;
13        self.growing.insert(vector);
14        self.next_id += 1;
15        id
16    }
17
18    /// Insert a vector with an associated surrogate. The surrogate is
19    /// allocated by the Control Plane before the call; the engine only
20    /// stores the binding.
21    pub fn insert_with_surrogate(&mut self, vector: Vec<f32>, surrogate: Surrogate) -> u32 {
22        let id = self.insert(vector);
23        if surrogate != Surrogate::ZERO {
24            self.surrogate_map.insert(id, surrogate);
25            self.surrogate_to_local.insert(surrogate, id);
26        }
27        id
28    }
29
30    /// Insert multiple vectors for a single document (ColBERT-style).
31    /// All N vectors are bound to the same `document_surrogate`.
32    pub fn insert_multi_vector(
33        &mut self,
34        vectors: &[&[f32]],
35        document_surrogate: Surrogate,
36    ) -> Vec<u32> {
37        let mut ids = Vec::with_capacity(vectors.len());
38        for &v in vectors {
39            let id = self.insert(v.to_vec());
40            if document_surrogate != Surrogate::ZERO {
41                self.surrogate_map.insert(id, document_surrogate);
42            }
43            ids.push(id);
44        }
45        if document_surrogate != Surrogate::ZERO {
46            self.multi_doc_map.insert(document_surrogate, ids.clone());
47        }
48        ids
49    }
50
51    /// Delete all vectors belonging to a multi-vector document.
52    pub fn delete_multi_vector(&mut self, document_surrogate: Surrogate) -> usize {
53        let Some(ids) = self.multi_doc_map.remove(&document_surrogate) else {
54            return 0;
55        };
56        let mut deleted = 0;
57        for id in &ids {
58            if self.delete(*id) {
59                deleted += 1;
60            }
61            self.surrogate_map.remove(id);
62        }
63        self.surrogate_to_local.remove(&document_surrogate);
64        deleted
65    }
66
67    /// Look up the surrogate for a global vector ID.
68    pub fn get_surrogate(&self, vector_id: u32) -> Option<Surrogate> {
69        self.surrogate_map.get(&vector_id).copied()
70    }
71
72    /// Resolve a surrogate back to its global vector ID, if bound.
73    pub fn local_for_surrogate(&self, surrogate: Surrogate) -> Option<u32> {
74        self.surrogate_to_local.get(&surrogate).copied()
75    }
76
77    /// Soft-delete a vector by global ID.
78    pub fn delete(&mut self, id: u32) -> bool {
79        let ok = self.delete_inner(id);
80        if ok && let Some(s) = self.surrogate_map.remove(&id) {
81            self.surrogate_to_local.remove(&s);
82        }
83        ok
84    }
85
86    pub(super) fn delete_inner(&mut self, id: u32) -> bool {
87        if id >= self.growing_base_id {
88            let local = id - self.growing_base_id;
89            if (local as usize) < self.growing.len() {
90                return self.growing.delete(local);
91            }
92        }
93        for seg in &mut self.sealed {
94            if id >= seg.base_id {
95                let local = id - seg.base_id;
96                if (local as usize) < seg.index.len() {
97                    return seg.index.delete(local);
98                }
99            }
100        }
101        for seg in &mut self.building {
102            if id >= seg.base_id {
103                let local = id - seg.base_id;
104                if (local as usize) < seg.flat.len() {
105                    return seg.flat.delete(local);
106                }
107            }
108        }
109        false
110    }
111
112    /// Soft-delete a vector by surrogate.
113    pub fn delete_by_surrogate(&mut self, surrogate: Surrogate) -> bool {
114        let Some(global_id) = self.surrogate_to_local.get(&surrogate).copied() else {
115            return false;
116        };
117        self.delete(global_id)
118    }
119
120    /// Un-delete a previously soft-deleted vector (for transaction rollback).
121    pub fn undelete(&mut self, id: u32) -> bool {
122        for seg in &mut self.sealed {
123            if id >= seg.base_id {
124                let local = id - seg.base_id;
125                if (local as usize) < seg.index.len() {
126                    return seg.index.undelete(local);
127                }
128            }
129        }
130        false
131    }
132}