diskann_disk/data_model/
cache.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use crate::data_model::GraphDataType;
7use diskann::{graph::AdjacencyList, ANNError, ANNResult};
8use hashbrown::{hash_map::Entry::Occupied, HashMap};
9
10pub struct Cache<Data: GraphDataType<VectorIdType = u32>> {
11    // Maintains the mapping of vector_id to index in the global cached nodes list.
12    mapping: HashMap<Data::VectorIdType, usize>,
13
14    // Flat buffer holding `capacity * dimension` vector elements, laid out row-major.
15    vectors: Vec<Data::VectorDataType>,
16
17    // The cached adjacency lists.
18    adjacency_lists: Vec<AdjacencyList<Data::VectorIdType>>,
19
20    // The cached associated data list.
21    associated_data: Vec<Data::AssociatedDataType>,
22
23    // The dimension of the vectors in the cache.
24    dimension: usize,
25
26    // The capacity of the cache.
27    capacity: usize,
28}
29
30impl<Data> Cache<Data>
31where
32    Data: GraphDataType<VectorIdType = u32>,
33{
34    // Creates a new cache with the specified dimension and capacity.
35    pub fn new(dimension: usize, capacity: usize) -> ANNResult<Self> {
36        Ok(Self {
37            mapping: HashMap::new(),
38            vectors: vec![Data::VectorDataType::default(); capacity * dimension],
39            adjacency_lists: Vec::with_capacity(capacity),
40            associated_data: Vec::with_capacity(capacity),
41            dimension,
42            capacity,
43        })
44    }
45
46    // Returns `true` if the cache contains the `vector_id`, otherwise `false`.
47    pub fn contains(&self, vector_id: &Data::VectorIdType) -> bool {
48        self.mapping.contains_key(vector_id)
49    }
50
51    // Returns the vector associated with the `vector_id`, if it exists in the cache otherwise `Option::None`.
52    pub fn get_vector(&self, vector_id: &Data::VectorIdType) -> Option<&[Data::VectorDataType]> {
53        if let Some(idx) = self.mapping.get(vector_id) {
54            Some(&self.vectors[idx * self.dimension..(idx + 1) * self.dimension])
55        } else {
56            Option::None
57        }
58    }
59
60    // Returns the adjacency list associated with the `vector_id``, if it exists in the cache otherwise `Option::None`.
61    pub fn get_adjacency_list(
62        &self,
63        vector_id: &Data::VectorIdType,
64    ) -> Option<&AdjacencyList<Data::VectorIdType>> {
65        if let Some(idx) = self.mapping.get(vector_id) {
66            Some(&self.adjacency_lists[*idx])
67        } else {
68            Option::None
69        }
70    }
71
72    // Returns the associated data associated with the `vector_id`, if it exists in the cache otherwise `Option::None`.
73    pub fn get_associated_data(
74        &self,
75        vector_id: &Data::VectorIdType,
76    ) -> Option<&Data::AssociatedDataType> {
77        if let Some(idx) = self.mapping.get(vector_id) {
78            Some(&self.associated_data[*idx])
79        } else {
80            Option::None
81        }
82    }
83
84    // Inserts a new node in the cache, if the node already exists in the cache, it updates the node.
85    // If the cache is full, it returns an error.
86    pub fn insert(
87        &mut self,
88        vector_id: &Data::VectorIdType,
89        vector: &[Data::VectorDataType],
90        adjacency_list: AdjacencyList<Data::VectorIdType>,
91        associated_data: Data::AssociatedDataType,
92    ) -> ANNResult<()> {
93        if self.dimension != vector.len() {
94            return ANNResult::Err(ANNError::log_index_error(
95                "Vector dimension does not match the dimension set in cache.",
96            ));
97        }
98
99        if let Occupied(occupied_entry) = self.mapping.entry(*vector_id) {
100            let idx = *occupied_entry.get();
101            self.copy_to_cache(idx, vector, adjacency_list, associated_data);
102            return ANNResult::Ok(());
103        }
104
105        if self.mapping.len() >= self.capacity {
106            return ANNResult::Err(ANNError::log_index_error(
107                "Cache is full, cannot insert more nodes",
108            ));
109        }
110
111        let idx = self.mapping.len();
112        self.mapping.insert(*vector_id, idx);
113        self.copy_to_cache(idx, vector, adjacency_list, associated_data);
114        ANNResult::Ok(())
115    }
116
117    // Returns `true` if the cache is empty, otherwise `false`.
118    pub fn is_empty(&self) -> bool {
119        self.mapping.is_empty()
120    }
121
122    // Returns the number of nodes in the cache.
123    pub fn len(&self) -> usize {
124        self.mapping.len()
125    }
126
127    fn copy_to_cache(
128        &mut self,
129        idx: usize,
130        vector: &[Data::VectorDataType],
131        adjacency_list: AdjacencyList<Data::VectorIdType>,
132        associated_data: Data::AssociatedDataType,
133    ) {
134        self.vectors[idx * self.dimension..(idx + 1) * self.dimension].copy_from_slice(vector);
135        self.adjacency_lists.push(adjacency_list);
136        self.associated_data.push(associated_data);
137    }
138}
139
140#[derive(PartialEq)]
141pub enum CachingStrategy {
142    None,
143    StaticCacheWithBfsNodes(usize),
144}
145
146#[cfg(test)]
147mod tests {
148    use crate::test_utils::GraphDataF32VectorUnitData;
149    use diskann::graph::AdjacencyList;
150    use rstest::rstest;
151
152    use crate::data_model::Cache;
153
154    #[rstest]
155    fn test_contains() {
156        let mut cache =
157            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 2).unwrap();
158        insert_a_random_node(&mut cache);
159        let vector_id = 1;
160        let vector = vec![1.0; 10];
161        let adjacency_list = AdjacencyList::from_iter_untrusted([2, 3, 4]);
162        cache
163            .insert(&vector_id, &vector, adjacency_list, ())
164            .unwrap();
165
166        assert!(cache.contains(&vector_id));
167
168        let not_exist_vector_id = 2;
169        assert!(!cache.contains(&not_exist_vector_id));
170    }
171
172    #[rstest]
173    fn test_get_vector() {
174        let mut cache =
175            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 2).unwrap();
176        insert_a_random_node(&mut cache);
177        let vector_id = 1;
178        let vector = vec![1.0; 10];
179        let adjacency_list = AdjacencyList::from_iter_untrusted([2, 3, 4]);
180        cache
181            .insert(&vector_id, &vector, adjacency_list, ())
182            .unwrap();
183
184        let result = cache.get_vector(&vector_id).unwrap();
185        assert_eq!(result, vector.as_slice());
186
187        let not_exist_vector_id = 2;
188        assert!(cache.get_vector(&not_exist_vector_id).is_none());
189    }
190
191    #[rstest]
192    fn test_get_adjacency_list() {
193        let mut cache =
194            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 2).unwrap();
195        insert_a_random_node(&mut cache);
196        let vector_id = 1;
197        let vector = vec![1.0; 10];
198        let adjacency_list = AdjacencyList::from_iter_untrusted([2, 3, 4]);
199        cache
200            .insert(&vector_id, &vector, adjacency_list.clone(), ())
201            .unwrap();
202
203        let result = cache.get_adjacency_list(&vector_id).unwrap();
204        assert_eq!(*result, adjacency_list);
205
206        let not_exist_vector_id = 2;
207        assert!(cache.get_adjacency_list(&not_exist_vector_id).is_none());
208    }
209
210    #[rstest]
211    fn test_get_associated_data() {
212        let mut cache =
213            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 2).unwrap();
214        insert_a_random_node(&mut cache);
215        let vector_id = 1;
216        let vector = vec![1.0; 10];
217        let adjacency_list = AdjacencyList::from_iter_untrusted([2, 3, 4]);
218        let associated_data = ();
219        cache
220            .insert(&vector_id, &vector, adjacency_list, associated_data)
221            .unwrap();
222
223        let result = cache.get_associated_data(&vector_id);
224        assert!(result.is_some());
225
226        let not_exist_vector_id = 2;
227        assert!(cache.get_associated_data(&not_exist_vector_id).is_none());
228    }
229
230    #[rstest]
231    fn test_insert() {
232        let mut cache =
233            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 2).unwrap();
234        insert_a_random_node(&mut cache);
235        let vector_id = 1;
236        let vector = vec![1.0; 10];
237        let adjacency_list = AdjacencyList::from_iter_untrusted([2, 3, 4]);
238
239        // Insert in cache
240        cache
241            .insert(&vector_id, &vector, adjacency_list.clone(), ())
242            .unwrap();
243        assert!(cache.contains(&vector_id));
244
245        // Update in cache
246        let updated_vector = vec![2.0; 10];
247        cache
248            .insert(&vector_id, &updated_vector, adjacency_list.clone(), ())
249            .unwrap();
250        assert_eq!(
251            cache.get_vector(&vector_id).unwrap(),
252            updated_vector.as_slice()
253        );
254
255        // Cache is Full
256        let vector_id_2 = 2;
257        let result = cache.insert(&vector_id_2, &vector, adjacency_list.clone(), ());
258        assert!(result.is_err());
259
260        // Wrong dimention Insert fails.
261        let wrong_dimentions_vector = vec![1.0; 11];
262        assert!(cache
263            .insert(&vector_id, &wrong_dimentions_vector, adjacency_list, ())
264            .is_err());
265    }
266
267    #[rstest]
268    fn test_is_empty() {
269        let mut cache =
270            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 1).unwrap();
271
272        assert!(cache.is_empty());
273
274        insert_a_random_node(&mut cache);
275
276        assert!(!cache.is_empty());
277    }
278
279    #[rstest]
280    fn test_len() {
281        let mut cache =
282            Cache::<GraphDataF32VectorUnitData>::new(/*dimention=*/ 10, /*capacity=*/ 5).unwrap();
283
284        assert_eq!(cache.len(), 0);
285
286        let vector_id = 1;
287        let vector = vec![1.0; 10];
288        let adjacency_list = AdjacencyList::from_iter_untrusted([2, 3, 4]);
289        cache
290            .insert(&vector_id, &vector, adjacency_list.clone(), ())
291            .unwrap();
292        let vector_id_2 = 2;
293        cache
294            .insert(&vector_id_2, &vector, adjacency_list.clone(), ())
295            .unwrap();
296        let vector_id_3 = 3;
297        cache
298            .insert(&vector_id_3, &vector, adjacency_list, ())
299            .unwrap();
300
301        assert_eq!(cache.len(), 3);
302    }
303
304    fn insert_a_random_node(cache: &mut Cache<GraphDataF32VectorUnitData>) {
305        let vector_id = 99;
306        let vector = vec![9.0; 10];
307        cache
308            .insert(
309                &vector_id,
310                &vector,
311                AdjacencyList::from_iter_untrusted([20, 30, 40]),
312                (),
313            )
314            .unwrap();
315    }
316}
diskann_disk/data_model/cache.rs

diskann_disk/data_model/
cache.rs
