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        // no-governor: VectorCollection is !Send and has no governor field; budget is enforced by the Data Plane core's arena before multi-vector insert is dispatched
38        let mut ids = Vec::with_capacity(vectors.len());
39        for &v in vectors {
40            let id = self.insert(v.to_vec());
41            if document_surrogate != Surrogate::ZERO {
42                self.surrogate_map.insert(id, document_surrogate);
43            }
44            ids.push(id);
45        }
46        if document_surrogate != Surrogate::ZERO {
47            self.multi_doc_map.insert(document_surrogate, ids.clone());
48        }
49        ids
50    }
51
52    /// Delete all vectors belonging to a multi-vector document.
53    pub fn delete_multi_vector(&mut self, document_surrogate: Surrogate) -> usize {
54        let Some(ids) = self.multi_doc_map.remove(&document_surrogate) else {
55            return 0;
56        };
57        let mut deleted = 0;
58        for id in &ids {
59            if self.delete(*id) {
60                deleted += 1;
61            }
62            self.surrogate_map.remove(id);
63        }
64        self.surrogate_to_local.remove(&document_surrogate);
65        deleted
66    }
67
68    /// Look up the surrogate for a global vector ID.
69    pub fn get_surrogate(&self, vector_id: u32) -> Option<Surrogate> {
70        self.surrogate_map.get(&vector_id).copied()
71    }
72
73    /// Resolve a surrogate back to its global vector ID, if bound.
74    pub fn local_for_surrogate(&self, surrogate: Surrogate) -> Option<u32> {
75        self.surrogate_to_local.get(&surrogate).copied()
76    }
77
78    /// Soft-delete a vector by global ID.
79    pub fn delete(&mut self, id: u32) -> bool {
80        let ok = self.delete_inner(id);
81        if ok && let Some(s) = self.surrogate_map.remove(&id) {
82            self.surrogate_to_local.remove(&s);
83        }
84        ok
85    }
86
87    pub(super) fn delete_inner(&mut self, id: u32) -> bool {
88        if id >= self.growing_base_id {
89            let local = id - self.growing_base_id;
90            if (local as usize) < self.growing.len() {
91                return self.growing.delete(local);
92            }
93        }
94        for seg in &mut self.sealed {
95            if id >= seg.base_id {
96                let local = id - seg.base_id;
97                if (local as usize) < seg.index.len() {
98                    return seg.index.delete(local);
99                }
100            }
101        }
102        for seg in &mut self.building {
103            if id >= seg.base_id {
104                let local = id - seg.base_id;
105                if (local as usize) < seg.flat.len() {
106                    return seg.flat.delete(local);
107                }
108            }
109        }
110        false
111    }
112
113    /// Soft-delete a vector by surrogate.
114    pub fn delete_by_surrogate(&mut self, surrogate: Surrogate) -> bool {
115        let Some(global_id) = self.surrogate_to_local.get(&surrogate).copied() else {
116            return false;
117        };
118        self.delete(global_id)
119    }
120
121    /// Un-delete a previously soft-deleted vector (for transaction rollback).
122    pub fn undelete(&mut self, id: u32) -> bool {
123        for seg in &mut self.sealed {
124            if id >= seg.base_id {
125                let local = id - seg.base_id;
126                if (local as usize) < seg.index.len() {
127                    return seg.index.undelete(local);
128                }
129            }
130        }
131        false
132    }
133}