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 ZST for [`MultiInsertStrategy::seed`](diskann::graph::glue::MultiInsertStrategy::Seed)
216/// indicating no use of the input batch.
217///
218/// Inmem providers typically don't use a working set at all, instead passing through accesses
219/// directly to the underlying provider. As such, no seeding is needed.
220#[derive(Debug, Clone, Copy)]
221pub struct Unseeded;
222
223/// A helper trait to set element in the Async index.
224pub trait SetElementHelper<T> {
225    /// Set an element in the index.
226    fn set_element(&self, index: &u32, element: &[T]) -> ANNResult<()>;
227}
228
229/// A helper trait to select the quant vector store.
230///
231/// This is also implemented for [`NoStore`], which explicitly disables deletion
232/// related functionality.
233pub trait CreateVectorStore {
234    /// The type of the created vector store.
235    type Target: VectorStore;
236
237    /// Create a quant provider capable of tracking `max_points`.
238    fn create(
239        self,
240        max_points: usize,
241        metric: Metric,
242        prefetch_lookahead: Option<usize>,
243    ) -> Self::Target;
244}
245
246/// A helper trait to select the delete provider.
247///
248/// This is also implemented for [`NoDeletes`], which explicitly disables deletion
249/// related functionality.
250pub trait CreateDeleteProvider {
251    /// The type of the created delete provider.
252    type Target;
253
254    /// Create a delete provider capable of tracking `total_points` number of deletes
255    /// (or disabling deletion check all together).
256    ///
257    /// NOTE: The value `total_points` consists of the sum of `max_points` and
258    /// `frozen_points`.
259    fn create(self, total_points: usize) -> Self::Target;
260}
261
262pub trait VectorStore: AsyncFriendly {
263    /// Total number of vectors in the store.
264    fn total(&self) -> usize;
265
266    /// Return the number of vector reads for a vector store.
267    fn count_for_get_vector(&self) -> usize;
268}
269
270/// A tag type indicating that a method fails via panic instead of returning an error.
271///
272/// This is an enum with no alternatives, so is impossible to construct. Therefore, we know
273/// that there can never be an actual value with this type.
274///
275#[derive(Debug)]
276pub enum Panics {}
277
278impl std::fmt::Display for Panics {
279    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
280        write!(f, "panics")
281    }
282}
283
284impl std::error::Error for Panics {}
285impl From<Panics> for ANNError {
286    #[cold]
287    fn from(_: Panics) -> ANNError {
288        ANNError::log_async_error("unreachable")
289    }
290}
291
292always_escalate!(Panics);
293
294/// A tag type used to indicate that no store should be used.
295///
296/// Typically this would be for full precision only or quant only setups.
297#[derive(Debug, Clone, Copy)]
298pub struct NoStore;
299
300impl CreateVectorStore for NoStore {
301    type Target = NoStore;
302    fn create(
303        self,
304        _max_points: usize,
305        _metric: Metric,
306        _prefetch_lookahead: Option<usize>,
307    ) -> Self::Target {
308        self
309    }
310}
311
312impl VectorStore for NoStore {
313    fn total(&self) -> usize {
314        0
315    }
316
317    fn count_for_get_vector(&self) -> usize {
318        0
319    }
320}
321
322impl LoadWith<AsyncQuantLoadContext> for NoStore {
323    type Error = ANNError;
324    async fn load_with<P>(_: &P, _: &AsyncQuantLoadContext) -> ANNResult<Self>
325    where
326        P: StorageReadProvider,
327    {
328        Ok(Self)
329    }
330}
331
332impl SaveWith<AsyncIndexMetadata> for NoStore {
333    type Ok = usize;
334    type Error = ANNError;
335    async fn save_with<P>(&self, _provider: &P, _auxiliary: &AsyncIndexMetadata) -> ANNResult<usize>
336    where
337        P: StorageWriteProvider,
338    {
339        Ok(0)
340    }
341}
342
343impl<T> SetElementHelper<T> for NoStore {
344    fn set_element(&self, _index: &u32, _element: &[T]) -> ANNResult<()> {
345        Ok(())
346    }
347}
348
349/// A tag type to indicate that no deletes are allowed for this provider.
350///
351/// This effectively disables deletion support at compile-time.
352#[derive(Debug, Clone, Copy)]
353pub struct NoDeletes;
354
355impl postprocess::DeletionCheck for NoDeletes {
356    /// Always mark IDs as not deleted.
357    ///
358    /// We rely on constant propagation and dead-code elimination to optimize call-sites
359    /// accordingly.
360    #[inline(always)]
361    fn deletion_check(&self, _: u32) -> bool {
362        false
363    }
364}
365
366impl CreateDeleteProvider for NoDeletes {
367    type Target = Self;
368    fn create(self, _: usize) -> Self {
369        Self
370    }
371}
372
373/// A tag type used to indicate that the `TableDeleteProviderAsync` should be used.
374#[derive(Debug, Clone, Copy)]
375pub struct TableBasedDeletes;
376
377impl CreateDeleteProvider for TableBasedDeletes {
378    type Target = TableDeleteProviderAsync;
379    fn create(self, total_points: usize) -> Self::Target {
380        TableDeleteProviderAsync::new(total_points)
381    }
382}
383
384/// Operates primarily in the quantized space with selective use of full precision.
385///
386/// # Search
387/// Search is performed in the quantized space. Full-precision vectors are used only
388/// to rerank the final candidate set.
389///
390/// # Insert and Prune
391/// During insert operations, the search step uses quantized vectors.
392/// Pruning then combines quantized vectors with a limited number of full-precision vectors.
393///
394/// The number of full-precision vectors used in pruning can be configured with
395/// the `max_fp_vecs_per_prune` option when constructing a `BfTreeProvider`.
396#[derive(Debug, Clone, Copy)]
397pub struct Hybrid {
398    /// Maximum number of full-precision vectors to use during pruning.
399    /// This field is ignored during search, where full-precision vectors are never used.
400    /// `None` defaults to use all full-precision vectors.
401    pub max_fp_vecs_per_prune: Option<usize>,
402}
403
404impl Hybrid {
405    /// Create a new `Hybrid` strategy with the specified maximum number of full-precision vectors
406    /// to use during pruning.
407    ///
408    /// If `max_fp_vecs_per_prune` is `None`, use all full-precision vectors.
409    pub fn new(max_fp_vecs_per_prune: Option<usize>) -> Self {
410        Self {
411            max_fp_vecs_per_prune,
412        }
413    }
414}
415
416#[cfg(test)]
417pub struct TestCallCount {
418    count: std::sync::atomic::AtomicUsize,
419}
420
421#[cfg(test)]
422impl TestCallCount {
423    pub fn new() -> Self {
424        TestCallCount {
425            count: std::sync::atomic::AtomicUsize::new(0),
426        }
427    }
428
429    pub fn enabled() -> bool {
430        true
431    }
432
433    pub fn increment(&self) {
434        self.count
435            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
436    }
437
438    pub fn get(&self) -> usize {
439        self.count.load(std::sync::atomic::Ordering::Relaxed)
440    }
441}
442
443#[cfg(not(test))]
444pub struct TestCallCount {}
445
446#[cfg(not(test))]
447impl TestCallCount {
448    pub fn new() -> Self {
449        TestCallCount {}
450    }
451
452    pub fn enabled() -> bool {
453        false
454    }
455
456    pub fn increment(&self) {}
457
458    pub fn get(&self) -> usize {
459        0
460    }
461}
462
463impl Default for TestCallCount {
464    fn default() -> Self {
465        Self::new()
466    }
467}
468
469#[cfg(test)]
470mod tests {
471    use std::num::NonZeroUsize;
472
473    use super::*;
474
475    #[test]
476    fn new_creates_correct_range() {
477        // valid_points of ten with five frozen points gives range 10..15
478        let sp = StartPoints::new(10, NonZeroUsize::new(5).unwrap())
479            .expect("should construct without overflow");
480        let r = sp.range().collect::<Vec<_>>();
481        assert_eq!(r, vec![10, 11, 12, 13, 14]);
482        assert_eq!(sp.end(), 15);
483    }
484
485    #[test]
486    fn new_returns_error_on_overflow() {
487        // valid_points at u32::MAX plus one frozen point must overflow
488        let max = u32::MAX;
489        let res = StartPoints::new(max, NonZeroUsize::new(1).unwrap());
490        assert!(res.is_err(), "expected an error when sum exceeds u32::MAX");
491        if let Err(err) = res {
492            let msg = err.to_string();
493            assert!(
494                msg.contains("Sum of valid points and frozen points exceeds u32::MAX"),
495                "unexpected error message: {}",
496                msg
497            );
498        }
499    }
500}
diskann_providers/model/graph/provider/async_/common.rs

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