feox_ann/
insert.rs

1use crate::graph::{compare_scored_desc, Node, ScoredNode, SearchScratch};
2use crate::math::{dot, normalize_in_place};
3use crate::{AnnIndex, Result};
4
5pub struct AnnInsertCursor<'a> {
6    index: &'a mut AnnIndex,
7    scratch: SearchScratch,
8}
9
10impl AnnIndex {
11    pub fn insert_cursor(&mut self) -> AnnInsertCursor<'_> {
12        AnnInsertCursor {
13            index: self,
14            scratch: SearchScratch::default(),
15        }
16    }
17
18    pub fn upsert(&mut self, id: String, vector: &[f32]) -> Result<()> {
19        let mut scratch = SearchScratch::default();
20        self.insert_vector(id, vector.to_vec(), &mut scratch)
21    }
22
23    pub fn upsert_owned(&mut self, id: String, vector: Vec<f32>) -> Result<()> {
24        let mut scratch = SearchScratch::default();
25        self.insert_vector(id, vector, &mut scratch)
26    }
27
28    pub(crate) fn insert_vector(
29        &mut self,
30        id: String,
31        mut vector: Vec<f32>,
32        scratch: &mut SearchScratch,
33    ) -> Result<()> {
34        self.validate_vector(&vector)?;
35        if let Some(index) = self.ids.remove(&id) {
36            if !self.nodes[index].deleted {
37                self.nodes[index].deleted = true;
38                self.active = self.active.saturating_sub(1);
39            }
40        }
41
42        normalize_in_place(&mut vector);
43        let level = self.level_for_id(&id);
44        let index = self.nodes.len();
45        self.nodes.push(Node {
46            id: id.clone(),
47            vector,
48            deleted: false,
49            neighbors: vec![Vec::new(); level + 1],
50        });
51        self.ids.insert(id, index);
52        self.add_entry_point(index);
53        self.active += 1;
54
55        let Some(mut entry) = self.entry else {
56            self.entry = Some(index);
57            self.max_level = level;
58            return Ok(());
59        };
60
61        for layer in ((level + 1)..=self.max_level).rev() {
62            entry = self.greedy_search(index, entry, layer);
63        }
64
65        let top_layer = level.min(self.max_level);
66        for layer in (0..=top_layer).rev() {
67            self.search_layer(index, entry, layer, self.config.ef_construction, scratch);
68            let selected = self.select_neighbors(index, layer, &mut scratch.found);
69            for &neighbor in &selected {
70                self.nodes[neighbor].neighbors[layer].push(index);
71                self.prune_neighbors(neighbor, layer);
72            }
73            self.nodes[index].neighbors[layer] = selected;
74            if let Some(best) = self.best_candidate(index, layer) {
75                entry = best;
76            }
77        }
78
79        if level > self.max_level {
80            self.entry = Some(index);
81            self.max_level = level;
82        }
83
84        Ok(())
85    }
86
87    pub(crate) fn select_neighbors(
88        &self,
89        query_index: usize,
90        layer: usize,
91        candidates: &mut Vec<ScoredNode>,
92    ) -> Vec<usize> {
93        let limit = self.neighbor_limit(layer);
94        candidates.retain(|candidate| candidate.index != query_index);
95        candidates.sort_by(compare_scored_desc);
96        candidates.dedup_by_key(|candidate| candidate.index);
97
98        let mut selected: Vec<usize> = Vec::with_capacity(limit);
99        let mut pruned: Vec<usize> = Vec::new();
100        for candidate in candidates.iter() {
101            if selected.len() >= limit {
102                break;
103            }
104            let vector = &self.nodes[candidate.index].vector;
105            let diverse = selected
106                .iter()
107                .all(|&kept| candidate.score >= dot(vector, &self.nodes[kept].vector));
108            if diverse {
109                selected.push(candidate.index);
110            } else {
111                pruned.push(candidate.index);
112            }
113        }
114        for index in pruned {
115            if selected.len() >= limit {
116                break;
117            }
118            selected.push(index);
119        }
120        selected
121    }
122
123    pub(crate) fn prune_neighbors(&mut self, index: usize, layer: usize) {
124        let limit = self.neighbor_limit(layer);
125        if self.nodes[index].neighbors[layer].len() <= limit {
126            return;
127        }
128        let mut neighbors = self.nodes[index].neighbors[layer].clone();
129        neighbors.sort_unstable();
130        neighbors.dedup();
131        if neighbors.len() <= limit {
132            self.nodes[index].neighbors[layer] = neighbors;
133            return;
134        }
135
136        let source = &self.nodes[index].vector;
137        let mut candidates = Vec::with_capacity(neighbors.len());
138        for neighbor in neighbors {
139            candidates.push(ScoredNode {
140                index: neighbor,
141                score: dot(source, &self.nodes[neighbor].vector),
142            });
143        }
144        self.nodes[index].neighbors[layer] = self.select_neighbors(index, layer, &mut candidates);
145    }
146}
147
148impl AnnInsertCursor<'_> {
149    pub fn reserve(&mut self, additional: usize) {
150        self.index.reserve(additional);
151    }
152
153    pub fn upsert_owned(&mut self, id: String, vector: Vec<f32>) -> Result<()> {
154        self.index.insert_vector(id, vector, &mut self.scratch)
155    }
156}
feox_ann/insert.rs

feox_ann/
insert.rs
