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

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