diskann_providers/model/graph/provider/async_/
common.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{cell::UnsafeCell, mem, num::NonZeroUsize, ops::Deref, slice, sync::Arc};
7
8use crate::storage::{StorageReadProvider, StorageWriteProvider};
9use arc_swap::Guard;
10use diskann::{ANNError, ANNResult, always_escalate, utils::IntoUsize};
11use diskann_utils::future::AsyncFriendly;
12use diskann_vector::distance::Metric;
13
14use crate::{
15    model::graph::provider::async_::{TableDeleteProviderAsync, postprocess},
16    storage::{AsyncIndexMetadata, AsyncQuantLoadContext, LoadWith, SaveWith},
17};
18
19pub use diskann::graph::strategy::{FullPrecision, Quantized};
20
21/// Represents a range of start points for an index.
22/// The range includes `start` and excludes `end`.
23/// `start` is the first valid point, and `end - 1` is the last valid point.
24pub struct StartPoints {
25    start: u32,
26    end: u32,
27}
28
29impl StartPoints {
30    pub fn new(valid_points: u32, frozen_points: NonZeroUsize) -> ANNResult<Self> {
31        Ok(Self {
32            start: valid_points,
33            end: match valid_points.checked_add(frozen_points.get() as u32) {
34                Some(end) => end,
35                None => {
36                    return Err(ANNError::log_index_error(
37                        "Sum of valid points and frozen points exceeds u32::MAX",
38                    ));
39                }
40            },
41        })
42    }
43
44    pub fn range(&self) -> std::ops::Range<u32> {
45        self.start..self.end
46    }
47
48    pub fn len(&self) -> usize {
49        (self.end - self.start).into_usize()
50    }
51
52    pub fn is_empty(&self) -> bool {
53        self.len() == 0
54    }
55
56    pub fn start(&self) -> u32 {
57        self.start
58    }
59
60    pub fn end(&self) -> u32 {
61        self.end
62    }
63}
64
65pub struct VectorGuard<T> {
66    inner: Guard<Arc<Vec<T>>>,
67}
68
69impl<T> VectorGuard<T> {
70    pub(crate) fn from_guard(guard: Guard<Arc<Vec<T>>>) -> Self {
71        Self { inner: guard }
72    }
73}
74
75impl<T> Deref for VectorGuard<T> {
76    type Target = [T];
77
78    fn deref(&self) -> &Self::Target {
79        self.inner.deref()
80    }
81}
82
83/// Memory-backed vector storage that aligns vectors to cache lines.
84///
85/// This stores vectors in a one giant allocation and guarantees that
86/// each vector starts on a cache-aligned boundary (64 byte aligned).
87/// To achieve this, vectors may not be densely packed into the underlying
88/// buffer.
89pub struct AlignedMemoryVectorStore<T: bytemuck::Pod> {
90    store: UnsafeCell<Vec<T>>,
91    max_vectors: usize,
92    start_index: usize,
93    dim: usize,
94    padded_vector_dim: usize,
95}
96
97// SAFETY: It's not really, but the `bytemuck::Pod` bound helps mitigate the fallout.
98unsafe impl<T: bytemuck::Pod + Sync> Sync for AlignedMemoryVectorStore<T> {}
99
100// SAFETY: It's not really, but the `bytemuck::Pod` bound helps mitigate the fallout.
101unsafe impl<T: bytemuck::Pod + Send> Send for AlignedMemoryVectorStore<T> {}
102
103impl<T: bytemuck::Pod> AlignedMemoryVectorStore<T> {
104    pub fn with_capacity(max_vectors: usize, dim: usize) -> Self {
105        let elem_size = mem::size_of::<T>();
106        assert!(64 % elem_size == 0);
107        let vector_size = elem_size * dim;
108        let extra_size = vector_size % 64;
109        let padded_vector_dim = if extra_size == 0 {
110            // vectors will be naturally aligned when packed
111            dim
112        } else {
113            let padding_needed_size = 64 - extra_size;
114            assert!(padding_needed_size.is_multiple_of(elem_size));
115            let extra_elems = padding_needed_size / elem_size;
116            dim + extra_elems
117        };
118
119        assert!((padded_vector_dim * elem_size).is_multiple_of(64));
120
121        // Our allocation may start unaligned, so we will offset the first vector to ensure
122        // correct alignment. This means we need some extra elements at the end to compensate.
123        let last_elems: usize = 64 / elem_size - 1;
124
125        let count = max_vectors * padded_vector_dim + last_elems;
126        let mut store: UnsafeCell<Vec<T>> =
127            UnsafeCell::new(vec![<T as bytemuck::Zeroable>::zeroed(); count]);
128
129        let start_index = store.get_mut().as_ptr().align_offset(64);
130
131        Self {
132            store,
133            max_vectors,
134            start_index,
135            dim,
136            padded_vector_dim,
137        }
138    }
139
140    pub fn max_vectors(&self) -> usize {
141        self.max_vectors
142    }
143
144    pub fn dim(&self) -> usize {
145        self.dim
146    }
147
148    /// Return a vector as a slice.
149    ///
150    /// # Safety
151    ///
152    /// This function will not read out of bounds, but it may observe a torn read if reading a vector at the same time
153    /// as it is being written. It is up to the caller to deal with torn reads, but it is expected that clients wanting maximum
154    /// performance will be okay with that tradeoff.
155    ///
156    /// Note that as vector elements are plain data, the impact of memory races is limited.
157    #[inline(always)]
158    pub unsafe fn get_slice(&self, index: usize) -> &[T] {
159        assert!(
160            index < self.max_vectors,
161            "index ({}) exceeded max_vectors ({})",
162            index,
163            self.max_vectors
164        );
165        let index = index * self.padded_vector_dim + self.start_index;
166
167        // SAFETY: Constructing a slice to the inside of the allocation. We know this is
168        // valid memory because the allocation is sized so that all vectors fit, and we know
169        // that the `index < max_vectors`.
170        let buf = unsafe { (*self.store.get()).as_ptr() };
171
172        // SAFETY: See comment above.
173        unsafe { slice::from_raw_parts(buf.add(index), self.dim) }
174    }
175
176    /// Return a vector as a mutable slice.
177    ///
178    /// # Safety
179    ///
180    /// This function will not synchronize access, but the memory is guaranteed to be valid. Callers must synchronize
181    /// themselves if they require consistency.
182    #[allow(clippy::mut_from_ref)]
183    pub unsafe fn get_mut_slice(&self, index: usize) -> &mut [T] {
184        assert!(index < self.max_vectors);
185        let index = index * self.padded_vector_dim + self.start_index;
186
187        // SAFETY: Constructing a mutable slice to the inside of the allocation. We know
188        // this is valid memory because the allocation is sized so that all vectors fit,
189        // and we know that the `index < max_vectors`.
190        unsafe {
191            let buf = (*self.store.get()).as_mut_ptr();
192            slice::from_raw_parts_mut(buf.add(index), self.dim)
193        }
194    }
195}
196
197/// Prefetch cache line level.
198#[derive(Debug, Default, Clone, Copy, Eq, PartialEq)]
199pub enum PrefetchCacheLineLevel {
200    /// 4 cache lines
201    CacheLine4,
202    /// 8 cache lines
203    CacheLine8,
204    /// 16 cache lines
205    #[default]
206    CacheLine16,
207    /// prefetch all cache lines
208    All,
209}
210
211//////////////////////////////////////////////////////////
212// Common data structure and traits for async providers //
213//////////////////////////////////////////////////////////
214
215/// A helper trait to set element in the Async index.
216pub trait SetElementHelper<T> {
217    /// Set an element in the index.
218    fn set_element(&self, index: &u32, element: &[T]) -> ANNResult<()>;
219}
220
221/// A helper trait to select the quant vector store.
222///
223/// This is also implemented for [`NoStore`], which explicitly disables deletion
224/// related functionality.
225pub trait CreateVectorStore {
226    /// The type of the created vector store.
227    type Target: VectorStore;
228
229    /// Create a quant provider capable of tracking `max_points`.
230    fn create(
231        self,
232        max_points: usize,
233        metric: Metric,
234        prefetch_lookahead: Option<usize>,
235    ) -> Self::Target;
236}
237
238/// A helper trait to select the delete provider.
239///
240/// This is also implemented for [`NoDeletes`], which explicitly disables deletion
241/// related functionality.
242pub trait CreateDeleteProvider {
243    /// The type of the created delete provider.
244    type Target;
245
246    /// Create a delete provider capable of tracking `total_points` number of deletes
247    /// (or disabling deletion check all together).
248    ///
249    /// NOTE: The value `total_points` consists of the sum of `max_points` and
250    /// `frozen_points`.
251    fn create(self, total_points: usize) -> Self::Target;
252}
253
254pub trait VectorStore: AsyncFriendly {
255    /// Total number of vectors in the store.
256    fn total(&self) -> usize;
257
258    /// Return the number of vector reads for a vector store.
259    fn count_for_get_vector(&self) -> usize;
260}
261
262/// A tag type indicating that a method fails via panic instead of returning an error.
263///
264/// This is an enum with no alternatives, so is impossible to construct. Therefore, we know
265/// that there can never be an actual value with this type.
266///
267#[derive(Debug)]
268pub enum Panics {}
269
270impl std::fmt::Display for Panics {
271    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        write!(f, "panics")
273    }
274}
275
276impl std::error::Error for Panics {}
277impl From<Panics> for ANNError {
278    #[cold]
279    fn from(_: Panics) -> ANNError {
280        ANNError::log_async_error("unreachable")
281    }
282}
283
284always_escalate!(Panics);
285
286/// A tag type used to indicate that no store should be used.
287///
288/// Typically this would be for full precision only or quant only setups.
289#[derive(Debug, Clone, Copy)]
290pub struct NoStore;
291
292impl CreateVectorStore for NoStore {
293    type Target = NoStore;
294    fn create(
295        self,
296        _max_points: usize,
297        _metric: Metric,
298        _prefetch_lookahead: Option<usize>,
299    ) -> Self::Target {
300        self
301    }
302}
303
304impl VectorStore for NoStore {
305    fn total(&self) -> usize {
306        0
307    }
308
309    fn count_for_get_vector(&self) -> usize {
310        0
311    }
312}
313
314impl LoadWith<AsyncQuantLoadContext> for NoStore {
315    type Error = ANNError;
316    async fn load_with<P>(_: &P, _: &AsyncQuantLoadContext) -> ANNResult<Self>
317    where
318        P: StorageReadProvider,
319    {
320        Ok(Self)
321    }
322}
323
324impl SaveWith<AsyncIndexMetadata> for NoStore {
325    type Ok = usize;
326    type Error = ANNError;
327    async fn save_with<P>(&self, _provider: &P, _auxiliary: &AsyncIndexMetadata) -> ANNResult<usize>
328    where
329        P: StorageWriteProvider,
330    {
331        Ok(0)
332    }
333}
334
335impl<T> SetElementHelper<T> for NoStore {
336    fn set_element(&self, _index: &u32, _element: &[T]) -> ANNResult<()> {
337        Ok(())
338    }
339}
340
341/// A tag type to indicate that no deletes are allowed for this provider.
342///
343/// This effectively disables deletion support at compile-time.
344#[derive(Debug, Clone, Copy)]
345pub struct NoDeletes;
346
347impl postprocess::DeletionCheck for NoDeletes {
348    /// Always mark IDs as not deleted.
349    ///
350    /// We rely on constant propagation and dead-code elimination to optimize call-sites
351    /// accordingly.
352    #[inline(always)]
353    fn deletion_check(&self, _: u32) -> bool {
354        false
355    }
356}
357
358impl CreateDeleteProvider for NoDeletes {
359    type Target = Self;
360    fn create(self, _: usize) -> Self {
361        Self
362    }
363}
364
365/// A tag type used to indicate that the `TableDeleteProviderAsync` should be used.
366#[derive(Debug, Clone, Copy)]
367pub struct TableBasedDeletes;
368
369impl CreateDeleteProvider for TableBasedDeletes {
370    type Target = TableDeleteProviderAsync;
371    fn create(self, total_points: usize) -> Self::Target {
372        TableDeleteProviderAsync::new(total_points)
373    }
374}
375
376/// Operates primarily in the quantized space with selective use of full precision.
377///
378/// # Search
379/// Search is performed in the quantized space. Full-precision vectors are used only
380/// to rerank the final candidate set.
381///
382/// # Insert and Prune
383/// During insert operations, the search step uses quantized vectors.
384/// Pruning then combines quantized vectors with a limited number of full-precision vectors.
385///
386/// The number of full-precision vectors used in pruning can be configured with
387/// the `max_fp_vecs_per_prune` option when constructing a `BfTreeProvider`.
388#[derive(Debug, Clone, Copy)]
389pub struct Hybrid {
390    /// Maximum number of full-precision vectors to use during pruning.
391    /// This field is ignored during search, where full-precision vectors are never used.
392    /// `None` defaults to use all full-precision vectors.
393    pub max_fp_vecs_per_prune: Option<usize>,
394}
395
396impl Hybrid {
397    /// Create a new `Hybrid` strategy with the specified maximum number of full-precision vectors
398    /// to use during pruning.
399    ///
400    /// If `max_fp_vecs_per_prune` is `None`, use all full-precision vectors.
401    pub fn new(max_fp_vecs_per_prune: Option<usize>) -> Self {
402        Self {
403            max_fp_vecs_per_prune,
404        }
405    }
406}
407
408#[cfg(test)]
409pub struct TestCallCount {
410    count: std::sync::atomic::AtomicUsize,
411}
412
413#[cfg(test)]
414impl TestCallCount {
415    pub fn new() -> Self {
416        TestCallCount {
417            count: std::sync::atomic::AtomicUsize::new(0),
418        }
419    }
420
421    pub fn enabled() -> bool {
422        true
423    }
424
425    pub fn increment(&self) {
426        self.count
427            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
428    }
429
430    pub fn get(&self) -> usize {
431        self.count.load(std::sync::atomic::Ordering::Relaxed)
432    }
433}
434
435#[cfg(not(test))]
436pub struct TestCallCount {}
437
438#[cfg(not(test))]
439impl TestCallCount {
440    pub fn new() -> Self {
441        TestCallCount {}
442    }
443
444    pub fn enabled() -> bool {
445        false
446    }
447
448    pub fn increment(&self) {}
449
450    pub fn get(&self) -> usize {
451        0
452    }
453}
454
455impl Default for TestCallCount {
456    fn default() -> Self {
457        Self::new()
458    }
459}
460
461#[cfg(test)]
462mod tests {
463    use std::num::NonZeroUsize;
464
465    use super::*;
466
467    #[test]
468    fn new_creates_correct_range() {
469        // valid_points of ten with five frozen points gives range 10..15
470        let sp = StartPoints::new(10, NonZeroUsize::new(5).unwrap())
471            .expect("should construct without overflow");
472        let r = sp.range().collect::<Vec<_>>();
473        assert_eq!(r, vec![10, 11, 12, 13, 14]);
474        assert_eq!(sp.end(), 15);
475    }
476
477    #[test]
478    fn new_returns_error_on_overflow() {
479        // valid_points at u32::MAX plus one frozen point must overflow
480        let max = u32::MAX;
481        let res = StartPoints::new(max, NonZeroUsize::new(1).unwrap());
482        assert!(res.is_err(), "expected an error when sum exceeds u32::MAX");
483        if let Err(err) = res {
484            let msg = err.to_string();
485            assert!(
486                msg.contains("Sum of valid points and frozen points exceeds u32::MAX"),
487                "unexpected error message: {}",
488                msg
489            );
490        }
491    }
492}
diskann_providers/model/graph/provider/async_/common.rs

diskann_providers/model/graph/provider/async_/
common.rs
