diskann_bftree/
vectors.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Bf-Tree vector provider
7
8use std::marker::PhantomData;
9
10use crate::{AccessError, VectorError, VectorUnavailable};
11use bf_tree::{BfTree, Config};
12use bytemuck::cast_slice;
13use diskann::{error::RankedError, utils::VectorRepr, ANNError, ANNErrorKind, ANNResult};
14use thiserror::Error;
15
16use super::ConfigError;
17use crate::{bftree_insert, TestCallCount};
18
19pub struct VectorProvider<T: VectorRepr> {
20    dim: usize,
21    pub max_vectors: usize,
22    pub num_start_points: usize,
23    vector_index: BfTree,
24    pub(super) num_get_calls: TestCallCount,
25    _phantom: PhantomData<T>,
26}
27
28impl<T: VectorRepr> VectorProvider<T> {
29    /// Create a new instance based on bf-tree Config directly
30    pub fn new_with_config(
31        max_vectors: usize,
32        dim: usize,
33        num_start_points: usize,
34        config: Config,
35    ) -> ANNResult<Self> {
36        crate::validate_record_size(
37            "vector_provider",
38            &config,
39            std::mem::size_of::<usize>(),
40            dim * std::mem::size_of::<T>(),
41        )?;
42
43        let vector_index = BfTree::with_config(config, None).map_err(ConfigError)?;
44
45        Ok(Self {
46            dim,
47            max_vectors,
48            num_start_points,
49            vector_index,
50            num_get_calls: TestCallCount::default(),
51            _phantom: PhantomData,
52        })
53    }
54
55    /// Create a new instance from an existing BfTree (for loading from snapshot)
56    ///
57    #[inline(always)]
58    pub fn new_from_bftree(
59        max_vectors: usize,
60        dim: usize,
61        num_start_points: usize,
62        vector_index: BfTree,
63    ) -> Self {
64        Self {
65            dim,
66            max_vectors,
67            num_start_points,
68            vector_index,
69            num_get_calls: TestCallCount::default(),
70            _phantom: PhantomData,
71        }
72    }
73
74    /// Return the total number of points including starting points
75    ///
76    #[inline(always)]
77    pub fn total(&self) -> usize {
78        self.max_vectors + self.num_start_points
79    }
80
81    // /// Return vector Id range for the starting points
82    // ///
83    // #[inline(always)]
84    // pub fn start_point_range(&self) -> std::ops::Range<usize> {
85    //     self.max_vectors..self.total()
86    // }
87
88    /// Return the vector dimension
89    ///
90    #[inline(always)]
91    pub fn dim(&self) -> usize {
92        self.dim
93    }
94
95    /// Return a vector of vector Ids of the starting points
96    ///
97    #[inline(always)]
98    pub fn starting_points(&self) -> ANNResult<Vec<u32>> {
99        (self.max_vectors..self.total())
100            .map(|i| {
101                u32::try_from(i).map_err(|_| {
102                    ANNError::log_index_error(format_args!("start point id {i} exceeds u32::MAX"))
103                })
104            })
105            .collect()
106    }
107
108    /// Access the BfTree config
109    pub(crate) fn config(&self) -> &Config {
110        self.vector_index.config()
111    }
112
113    /// Access the underlying BfTree
114    pub(crate) fn bftree(&self) -> &BfTree {
115        &self.vector_index
116    }
117
118    /// Set vector with Id, `i``, to `v`
119    ///
120    /// Errors if:
121    ///
122    /// * `i > self.total()`: `i` must be in bounds.
123    /// * `v.dim() != self.dim()`: The slice must have the proper length
124    #[inline(always)]
125    pub(crate) fn set_vector_sync(&self, i: usize, v: &[T]) -> ANNResult<()> {
126        if v.len() != self.dim {
127            return Err(ANNError::log_index_error(
128                "Vector dimension is not equal to the expected dimension.",
129            ));
130        }
131        if i >= self.total() {
132            return Err(ANNError::log_index_error(
133                "Vector id is out of boundary in the dataset.",
134            ));
135        }
136
137        // Serialize the key, vector_id, into a byte string, &[u8]
138        let key = bytemuck::bytes_of(&i);
139        let value = cast_slice::<T, u8>(v);
140
141        bftree_insert(&self.vector_index, key, value)?;
142
143        Ok(())
144    }
145
146    pub(crate) fn get_vector_into(&self, i: usize, buffer: &mut [T]) -> Result<(), AccessError> {
147        if buffer.len() != self.dim {
148            #[derive(Debug, Error)]
149            #[error("expected a buffer with dim {0}, instead got {1}")]
150            struct WrongDim(usize, usize);
151
152            return Err(RankedError::Error(ANNError::new(
153                ANNErrorKind::IndexError,
154                WrongDim(self.dim(), buffer.len()),
155            )));
156        }
157
158        self.num_get_calls.increment();
159        match self.vector_index.read(
160            bytemuck::bytes_of(&i),
161            bytemuck::must_cast_slice_mut::<_, u8>(buffer),
162        ) {
163            bf_tree::LeafReadResult::Found(read_size) => {
164                let vector_size = std::mem::size_of::<T>() * self.dim;
165                if read_size as usize != vector_size {
166                    return Err(RankedError::Error(ANNError::log_index_error(format!(
167                        "The bf-tree entry for vector id {} is marked as found but has size {} instead of the expected size {}",
168                        i, read_size, vector_size,
169                    ))));
170                }
171            }
172            bf_tree::LeafReadResult::Deleted => {
173                return Err(RankedError::Transient(VectorUnavailable {
174                    id: i,
175                    err: VectorError::Deleted,
176                }));
177            }
178            bf_tree::LeafReadResult::InvalidKey => {
179                return Err(RankedError::Error(ANNError::log_index_error(format!(
180                    "The bf-tree entry for vector id {} is marked as invalid",
181                    i
182                ))));
183            }
184            bf_tree::LeafReadResult::NotFound => {
185                return Err(RankedError::Transient(VectorUnavailable {
186                    id: i,
187                    err: VectorError::NotFound,
188                }));
189            }
190        };
191
192        Ok(())
193    }
194
195    /// Return the vector at index `i`
196    #[inline(always)]
197    pub(crate) fn get_vector_sync(&self, i: usize) -> Result<Vec<T>, AccessError> {
198        // Search for the corresponding vector
199        let mut vector = vec![T::default(); self.dim];
200        self.get_vector_into(i, &mut vector)?;
201        Ok(vector)
202    }
203
204    pub(crate) fn delete_vector(&self, i: usize) {
205        let key = bytemuck::bytes_of(&i);
206        self.vector_index.delete(key);
207    }
208}
209
210///////////
211// Tests //
212///////////
213
214/// These unit tests target the functionality of Bf-Tree vector provider alone
215///
216#[cfg(test)]
217mod tests {
218    use std::sync::Arc;
219
220    use tokio::task::JoinSet;
221
222    use super::*;
223
224    /// Test the interleaved and parallell traversal of the Bf-Tree
225    /// by invoking the async accessors of the vector provider
226    ///
227    #[tokio::test(flavor = "multi_thread", worker_threads = 5)]
228    async fn test_parallel_tree_traversal() {
229        let num_points = 100;
230        let bf_tree_config = Config::default();
231        let vector_provider = Arc::new(
232            VectorProvider::<f32>::new_with_config(num_points, 3, 2, bf_tree_config).unwrap(),
233        );
234        let mut set = JoinSet::new();
235        for i in 0..num_points {
236            let vector = vec![i as f32, (i + 1) as f32, (i + 2) as f32];
237            let vector_provider_clone = Arc::clone(&vector_provider);
238            set.spawn(async move {
239                // One tokio task per vector insertion
240                vector_provider_clone.set_vector_sync(i, &vector).unwrap()
241            });
242        }
243
244        while let Some(res) = set.join_next().await {
245            res.unwrap();
246        }
247
248        for i in 0..num_points {
249            // SAFETY: We're only accessing one at a time.
250            let vector = vector_provider.get_vector_sync(i).unwrap();
251            assert_eq!(&vector, &vec![(i as f32), (i + 1) as f32, (i + 2) as f32]);
252        }
253        if TestCallCount::enabled() {
254            assert_eq!(vector_provider.num_get_calls.get(), num_points);
255        }
256    }
257
258    /// Test other methods and edge cases of the vector provider and sycrhnoization mechanism of Bf-Tree
259    ///
260    #[tokio::test(flavor = "multi_thread", worker_threads = 5)]
261    async fn test_parallel_vector_access() {
262        let num_points = 3;
263        let frozen_points = 2;
264        let dim = 3;
265        let bf_tree_config = Config::default();
266
267        let provider = Arc::new(
268            VectorProvider::<f32>::new_with_config(num_points, dim, frozen_points, bf_tree_config)
269                .unwrap(),
270        );
271
272        let mut set = JoinSet::new();
273        for _ in 0..5 {
274            let provider_ref = Arc::clone(&provider);
275            set.spawn(async move {
276                provider_ref.set_vector_sync(0, &[0.0, 0.0, 0.0]).unwrap();
277                provider_ref.set_vector_sync(1, &[1.0, 1.0, 1.0]).unwrap();
278                provider_ref.set_vector_sync(2, &[2.0, 2.0, 2.0]).unwrap();
279                provider_ref.set_vector_sync(3, &[3.0, 3.0, 3.0]).unwrap();
280                provider_ref.set_vector_sync(4, &[4.0, 4.0, 4.0]).unwrap();
281
282                assert_eq!(provider_ref.get_vector_sync(4).unwrap(), &[4.0, 4.0, 4.0]);
283                assert_eq!(provider_ref.get_vector_sync(3).unwrap(), &[3.0, 3.0, 3.0]);
284                assert_eq!(provider_ref.get_vector_sync(2).unwrap(), &[2.0, 2.0, 2.0]);
285                assert_eq!(provider_ref.get_vector_sync(1).unwrap(), &[1.0, 1.0, 1.0]);
286                assert_eq!(provider_ref.get_vector_sync(0).unwrap(), &[0.0, 0.0, 0.0]);
287
288                // Error checking.
289                assert!(provider_ref.set_vector_sync(5, &[0.0, 0.0, 0.0]).is_err());
290                assert!(provider_ref.set_vector_sync(2, &[0.0, 0.0]).is_err());
291            });
292        }
293
294        while let Some(res) = set.join_next().await {
295            res.unwrap();
296        }
297    }
298
299    /// Test new_from_bftree constructor
300    #[tokio::test]
301    async fn test_new_from_bftree() {
302        let bftree = BfTree::with_config(Config::default(), None).expect("Failed to create BfTree");
303        let provider = VectorProvider::<f32>::new_from_bftree(100, 3, 2, bftree);
304
305        // Verify fields are set correctly
306        assert_eq!(provider.dim(), 3);
307        assert_eq!(provider.max_vectors, 100);
308        assert_eq!(provider.num_start_points, 2);
309        assert_eq!(provider.total(), 102);
310
311        // Verify the provider is functional
312        provider.set_vector_sync(0, &[1.0, 2.0, 3.0]).unwrap();
313        let result = provider.get_vector_sync(0).unwrap();
314        assert_eq!(result, vec![1.0, 2.0, 3.0]);
315    }
316}
diskann_bftree/vectors.rs

diskann_bftree/
vectors.rs
