Skip to main content

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

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{fmt::Debug, future::Future, num::NonZeroUsize};
7
8use crate::storage::{StorageReadProvider, StorageWriteProvider};
9use diskann::{
10    ANNError, ANNResult,
11    graph::AdjacencyList,
12    provider::{
13        DataProvider, DefaultAccessor, DefaultContext, Delete, ElementStatus, ExecutionContext,
14        NeighborAccessor, NeighborAccessorMut, NoopGuard, SetElement,
15    },
16    utils::{IntoUsize, ONE, VectorRepr},
17};
18use diskann_utils::future::AsyncFriendly;
19use diskann_vector::distance::Metric;
20
21use crate::{
22    model::graph::provider::async_::{
23        SimpleNeighborProviderAsync, StartPoints, TableDeleteProviderAsync,
24        common::{
25            CreateDeleteProvider, CreateVectorStore, NoDeletes, NoStore, PrefetchCacheLineLevel,
26            SetElementHelper, VectorStore,
27        },
28    },
29    storage::{AsyncIndexMetadata, AsyncQuantLoadContext, DiskGraphOnly, LoadWith, SaveWith},
30};
31
32/////////////////////
33// DefaultProvider //
34/////////////////////
35
36/// An in-memory implementation of a [`DataProvider`] built around the idea of having up to
37/// two layers of vector stores: a base store and an auxiliary store.
38///
39/// This provider must be pre-configured with the number of elements it is going to contain
40/// and uses the identity mapping between external and internal vector IDs.
41///
42/// In addition to a pre-configured number of points, this struct is also parameterized by
43/// the concept of "frozen" point, which serve as the entry points for graph search.
44/// Internally, these are stored consecutively just after the "max_points" slot.
45///
46/// # Type Parameters:
47///
48/// * `U`: The primary vector store that holds the main representation of vectors.
49///   Typical use cases:
50///   - Full precision vectors (e.g., [`FullPrecisionStore`])
51///   - Quantized vectors when no higher fidelity representation is required
52///   - May be `NoStore` if no base representation is required
53///
54/// * `V`: The auxiliary vector store that complements `base_vectors`.
55///   Typical use cases:
56///   - Quantized vectors when `base_vectors` holds full precision
57///   - Alternative compressed formats (e.g., 1-bit scalar with 8-bit scalar)
58///   - May be `NoStore` if no auxiliary representation is required
59///
60/// * `D`: The type of the deleted vector store. Like the quantized store, this is also
61///   not constrained by a trait and rather relies on implementation for concrete types.
62///   These are:
63///
64///   - [`NoDeletes`]: Do not support deletion at all (this disables implementation of
65///     the [`Delete`] trait.
66///   - [`TableDeleteProviderAsync`]: A bitmap storing deletion information.
67///
68/// * `Ctx`: A parameter controlling the [`ExecutionContext`] to be associated with this
69///   provider. For the majority of cases, this is [`DefaultContext`], but is left as
70///   a parameter to allow extension.
71///
72/// # Indexing Strategies
73///
74/// * [`FullPrecision`]: The strategies implemented by [`FullPrecision`] only retrieve data
75///   from the full-precision portion of the index. No quantized vectors are used.
76///
77///   During search, start points are filtered from the final results.
78///
79/// * [`Quantized`]: The strategies implemented by [`Quantized`] can use a mix of quantized
80///   and full-precision vectors.
81///
82///   - Search: During search, quantized vectors are used with reranking applied to the
83///     results before returning.
84///
85///   - Insertion: Quantized vectors are used during the search phase. During the pruning
86///     phase, a hybrid of quantized and full-precision vectors are used.
87///
88/// # Examples
89///
90/// The following code demonstrates how to instantiate and use the `DefaultProvider` in
91/// a number of different scenarios.
92///
93/// ## Full-Precision Only - No Deletes
94///
95/// This example demonstrates how to create a `DefaultProvider` that only supports
96/// full-precision vectors.
97/// ```
98/// use std::num::NonZeroUsize;
99///
100/// use diskann::provider::DefaultContext;
101/// use diskann_providers::model::graph::provider::async_::{
102///     inmem::{
103///         DefaultProvider, DefaultProviderParameters,
104///         CreateFullPrecision,
105///     },
106///     common::{NoStore, NoDeletes},
107/// };
108/// use diskann_vector::distance::Metric;
109///
110/// let dim = 4;
111/// let prefetch_cache_line_level = None;
112/// let parameters = DefaultProviderParameters {
113///     max_points: 5,
114///     frozen_points: NonZeroUsize::new(1).unwrap(),
115///     dim,
116///     metric: Metric::L2,
117///     prefetch_lookahead: None,
118///     max_degree: 40,
119///     prefetch_cache_line_level,
120/// };
121///
122/// // Create a table that supports 5 points and 1 "frozen" point.
123/// let provider = DefaultProvider::<_, _, _, DefaultContext>::new_empty(
124///     parameters,
125///     CreateFullPrecision::<f32>::new(dim, prefetch_cache_line_level),
126///     NoStore,
127///     NoDeletes,
128/// );
129/// ```
130///
131/// ## Full-Precision and PQ - No Deletes
132///
133/// To create a two-level provider with a PQ-based quant vector store, a
134/// [`FixedChunkPQTable`] can be supplied for the `quant_precursor` argument, as this
135/// implements the [`CreateQuantProvider`] trait.
136/// ```
137/// use std::num::NonZeroUsize;
138///
139/// use diskann::provider::DefaultContext;
140/// use diskann_providers::model::{
141///     pq::FixedChunkPQTable,
142///     graph::provider::async_::{
143///         inmem::{
144///             DefaultProvider, DefaultProviderParameters,
145///             CreateFullPrecision,
146///         },
147///         common::NoDeletes,
148///     },
149/// };
150/// use diskann_vector::distance::Metric;
151///
152/// // An example PQ table.
153/// let dim = 4;
154/// let table = FixedChunkPQTable::new(
155///     dim,
156///     Box::new([0.0, 0.0, 0.0, 0.0]),
157///     Box::new([0, dim]),
158/// ).unwrap();
159///
160/// let prefetch_cache_line_level = None;
161/// let parameters = DefaultProviderParameters {
162///     max_points: 5,
163///     frozen_points: NonZeroUsize::new(1).unwrap(),
164///     dim,
165///     metric: Metric::L2,
166///     prefetch_lookahead: None,
167///     max_degree: 40,
168///     prefetch_cache_line_level,
169/// };
170///
171/// // Create a table that supports 5 points and 1 "frozen" point.
172/// let provider = DefaultProvider::<_, _, _, DefaultContext>::new_empty(
173///     parameters,
174///     CreateFullPrecision::<f32>::new(dim, prefetch_cache_line_level),
175///     table,
176///     NoDeletes,
177/// );
178/// ```
179///
180/// ## Full-Precision and PQ - With Deletes.
181///
182/// If deletes are desired, than the type [`TableBasedDeletes`] can be passed to the
183/// constructor.
184/// ```
185/// use std::num::NonZeroUsize;
186///
187/// use diskann::provider::DefaultContext;
188/// use diskann_providers::model::{
189///     pq::FixedChunkPQTable,
190///     graph::provider::async_::{
191///         inmem::{
192///             DefaultProvider, DefaultProviderParameters,
193///             CreateFullPrecision,
194///         },
195///         common::TableBasedDeletes,
196///     },
197/// };
198/// use diskann_vector::distance::Metric;
199///
200/// // An example PQ table.
201/// let dim = 4;
202/// let table = FixedChunkPQTable::new(
203///     dim,
204///     Box::new([0.0, 0.0, 0.0, 0.0]),
205///     Box::new([0, dim]),
206/// ).unwrap();
207///
208/// let prefetch_cache_line_level = None;
209/// let parameters = DefaultProviderParameters {
210///     max_points: 5,
211///     frozen_points: NonZeroUsize::new(1).unwrap(),
212///     dim,
213///     metric: Metric::L2,
214///     prefetch_lookahead: None,
215///     max_degree: 40,
216///     prefetch_cache_line_level,
217/// };
218///
219/// // Create a table that supports 5 points and 1 "frozen" point.
220/// let provider = DefaultProvider::<_, _, _, DefaultContext>::new_empty(
221///     parameters,
222///     CreateFullPrecision::<f32>::new(dim, prefetch_cache_line_level),
223///     table,
224///     TableBasedDeletes,
225/// );
226/// ```
227pub struct DefaultProvider<U, V = NoStore, D = NoDeletes, Ctx = DefaultContext> {
228    /// The primary vector store that holds the main representation of vectors.
229    pub base_vectors: U,
230
231    /// The auxiliary vector store that complements `base_vectors`.
232    pub aux_vectors: V,
233
234    // Provider that holds the graph structure as neighbors of vectors.
235    pub(crate) neighbor_provider: SimpleNeighborProviderAsync,
236
237    /// The delete provider. If `D == NoDeletes`, then delete related operations are disabled.
238    ///
239    /// The size of this store must be kept in-sync with `quant_vectors` and `full-vectors`.
240    pub(super) deleted: D,
241
242    /// The metric to use for distances.
243    pub(super) metric: Metric,
244
245    pub(super) start_points: StartPoints,
246
247    context: std::marker::PhantomData<Ctx>,
248}
249
250#[derive(Debug, Clone)]
251pub struct DefaultProviderParameters {
252    /// The maximum number of valid points that provider can hold.
253    pub max_points: usize,
254
255    /// The number of frozen-points (start points) to store. The two level provider
256    /// stores these points starting at the linear index just after `max_points`.
257    pub frozen_points: NonZeroUsize,
258
259    /// The logical dimension of the full-precision data.
260    pub dim: usize,
261
262    /// The metric to use for distance computations.
263    pub metric: Metric,
264
265    /// The prefetch amount to use when performing bulk retrievals.
266    ///
267    /// Careful selection of this parameter can have a dramatic improvement on search
268    /// performance.
269    pub prefetch_lookahead: Option<usize>,
270
271    pub prefetch_cache_line_level: Option<PrefetchCacheLineLevel>,
272
273    /// The **actual** maximum number of neighbors to store for each vector.
274    pub max_degree: u32,
275}
276
277impl DefaultProviderParameters {
278    pub fn simple(max_points: usize, dim: usize, metric: Metric, max_degree: u32) -> Self {
279        Self {
280            max_points,
281            frozen_points: ONE,
282            metric,
283            dim,
284            prefetch_lookahead: None,
285            prefetch_cache_line_level: None,
286            max_degree,
287        }
288    }
289}
290
291impl<U, V, D, Ctx> DefaultProvider<U, V, D, Ctx> {
292    /// Construct a new, unpopulated data provider.
293    ///
294    /// # Arguments
295    /// * `params`: An instance of [`DefaultProviderParameters`] collecting shared
296    ///   configuration information.
297    /// * `base_precursor`: A precursor type for the base layer.
298    /// * `aux_precursor`: A precursor type for the auxiliary layer.
299    /// * `delete_precursor`: A precursor type for the delete layer.
300    /// * `neighbor_precursor`: A precursor type for the neighbor layer.
301    pub fn new_empty<CU, CV, CD>(
302        params: DefaultProviderParameters,
303        base_precursor: CU,
304        aux_precursor: CV,
305        delete_precursor: CD,
306    ) -> ANNResult<Self>
307    where
308        CU: CreateVectorStore<Target = U>,
309        CV: CreateVectorStore<Target = V>,
310        CD: CreateDeleteProvider<Target = D>,
311    {
312        let npts = params.max_points + params.frozen_points.get();
313        Ok(Self {
314            base_vectors: base_precursor.create(npts, params.metric, params.prefetch_lookahead),
315            aux_vectors: aux_precursor.create(npts, params.metric, params.prefetch_lookahead),
316            neighbor_provider: SimpleNeighborProviderAsync::new(npts, 1, params.max_degree, 1.0),
317            deleted: delete_precursor.create(npts),
318            metric: params.metric,
319            start_points: StartPoints::new(params.max_points as u32, params.frozen_points)?,
320            context: std::marker::PhantomData,
321        })
322    }
323
324    /// Return a vector of starting points.
325    pub fn starting_points(&self) -> ANNResult<Vec<u32>> {
326        Ok(self.start_points.range().collect())
327    }
328
329    /// An iterator over all ids including start points (even if they are deleted).
330    pub fn iter(&self) -> std::ops::Range<u32> {
331        0..self.start_points.end()
332    }
333
334    /// Return a reference to the neighbor provider.
335    pub fn neighbors(&self) -> &SimpleNeighborProviderAsync {
336        &self.neighbor_provider
337    }
338
339    pub fn num_start_points(&self) -> usize {
340        self.start_points.len()
341    }
342
343    /// Return the total capacity of the provider, **excluding** start points.
344    pub fn capacity(&self) -> usize {
345        self.start_points.start().into_usize()
346    }
347
348    /// Return the total capacity of the provider, **including** start points.
349    pub fn total_points(&self) -> usize {
350        self.start_points.end().into_usize()
351    }
352}
353
354/// Allow `&DefaultProvider` to implement `IntoIter`.
355impl<U, V, D, Ctx> IntoIterator for &DefaultProvider<U, V, D, Ctx> {
356    type Item = u32;
357    type IntoIter = std::ops::Range<u32>;
358    fn into_iter(self) -> Self::IntoIter {
359        self.iter()
360    }
361}
362
363impl<U, V, Ctx> DefaultProvider<U, V, TableDeleteProviderAsync, Ctx> {
364    /// A temporary method while development of deletion is in progress.
365    pub fn clear_delete_set(&self) {
366        self.deleted.clear();
367    }
368}
369
370impl<U, V, D, Ctx> DefaultProvider<U, V, D, Ctx>
371where
372    U: VectorStore,
373    V: VectorStore,
374{
375    /// Return the number of vector reads for base vector and aux vector stores respectively.
376    pub fn counts_for_get_vector(&self) -> (usize, usize) {
377        (
378            self.base_vectors.count_for_get_vector(),
379            self.aux_vectors.count_for_get_vector(),
380        )
381    }
382}
383
384pub trait SetStartPoints<T>
385where
386    T: ?Sized + 'static,
387{
388    fn set_start_points<'a, Itr>(&self, itr: Itr) -> ANNResult<()>
389    where
390        Itr: ExactSizeIterator<Item = &'a T> + 'a;
391}
392
393impl<T, U, V, D> SetStartPoints<[T]> for DefaultProvider<U, V, D>
394where
395    U: SetElementHelper<T>,
396    V: SetElementHelper<T>,
397    T: std::fmt::Debug + 'static,
398{
399    fn set_start_points<'a, Itr>(&self, itr: Itr) -> ANNResult<()>
400    where
401        Itr: ExactSizeIterator<Item = &'a [T]> + 'a,
402    {
403        let start_points = self.start_points.range();
404        if itr.len() != start_points.len() {
405            return Err(ANNError::log_async_index_error(format!(
406                "expected `itr` to contain `{}` items, instead it has {}",
407                start_points.len(),
408                itr.len(),
409            )));
410        }
411
412        for (i, v) in std::iter::zip(start_points, itr) {
413            self.aux_vectors.set_element(&i, v)?;
414            self.base_vectors.set_element(&i, v)?;
415        }
416
417        Ok(())
418    }
419}
420
421////////////
422// Saving //
423////////////
424
425impl<U, V, D, Ctx> SaveWith<(u32, AsyncIndexMetadata)> for DefaultProvider<U, V, D, Ctx>
426where
427    U: AsyncFriendly + SaveWith<AsyncIndexMetadata>,
428    V: AsyncFriendly + SaveWith<AsyncIndexMetadata>,
429    D: AsyncFriendly,
430    ANNError: From<U::Error> + From<V::Error>,
431    Ctx: ExecutionContext,
432{
433    type Ok = ();
434    type Error = ANNError;
435
436    async fn save_with<P>(
437        &self,
438        provider: &P,
439        auxiliary: &(u32, AsyncIndexMetadata),
440    ) -> Result<Self::Ok, Self::Error>
441    where
442        P: StorageWriteProvider,
443    {
444        self.base_vectors.save_with(provider, &auxiliary.1).await?;
445        self.aux_vectors.save_with(provider, &auxiliary.1).await?;
446        self.neighbor_provider
447            .save_with(provider, auxiliary)
448            .await?;
449        Ok(())
450    }
451}
452
453impl<U, V, D, Ctx> SaveWith<(u32, u32, DiskGraphOnly)> for DefaultProvider<U, V, D, Ctx>
454where
455    U: AsyncFriendly,
456    V: AsyncFriendly,
457    D: AsyncFriendly,
458    Ctx: ExecutionContext,
459{
460    type Ok = ();
461    type Error = ANNError;
462
463    async fn save_with<P>(
464        &self,
465        provider: &P,
466        auxiliary: &(u32, u32, DiskGraphOnly),
467    ) -> Result<Self::Ok, Self::Error>
468    where
469        P: StorageWriteProvider,
470    {
471        self.neighbor_provider
472            .save_with(provider, auxiliary)
473            .await?;
474        Ok(())
475    }
476}
477
478/////////////
479// Loading //
480/////////////
481
482impl<U, V, D, Ctx> LoadWith<AsyncQuantLoadContext> for DefaultProvider<U, V, D, Ctx>
483where
484    U: VectorStore + LoadWith<AsyncQuantLoadContext>,
485    V: VectorStore + AsyncFriendly + LoadWith<AsyncQuantLoadContext>,
486    D: AsyncFriendly + LoadWith<usize>,
487    ANNError: From<U::Error> + From<V::Error> + From<D::Error>,
488    Ctx: ExecutionContext,
489{
490    type Error = ANNError;
491
492    async fn load_with<P>(provider: &P, ctx: &AsyncQuantLoadContext) -> ANNResult<Self>
493    where
494        P: StorageReadProvider,
495    {
496        let base_vectors = U::load_with(provider, ctx).await?;
497        let aux_vectors = V::load_with(provider, ctx).await?;
498        let deleted = D::load_with(provider, &base_vectors.total()).await?;
499
500        // Take the maximum of the two totals so that if either store is `NoStore`,
501        // we still compute the correct overall number of points.
502        let npts = std::cmp::max(base_vectors.total(), aux_vectors.total());
503
504        let valid_points = npts
505            .checked_sub(ctx.num_frozen_points.get())
506            .ok_or_else(|| {
507                ANNError::log_index_error(format_args!(
508                    "Expected {} start points but the stored index only has {} total points",
509                    ctx.num_frozen_points.get(),
510                    base_vectors.total(),
511                ))
512            })?;
513        let start_points = StartPoints::new(valid_points as u32, ctx.num_frozen_points)?;
514        Ok(Self {
515            base_vectors,
516            aux_vectors,
517            neighbor_provider: SimpleNeighborProviderAsync::load_with(provider, ctx).await?,
518            deleted,
519            metric: ctx.metric,
520            start_points,
521            context: std::marker::PhantomData,
522        })
523    }
524}
525
526impl LoadWith<usize> for NoDeletes {
527    type Error = ANNError;
528
529    async fn load_with<P>(_: &P, _num_points: &usize) -> ANNResult<Self>
530    where
531        P: StorageReadProvider,
532    {
533        Ok(NoDeletes)
534    }
535}
536
537impl LoadWith<usize> for TableDeleteProviderAsync {
538    type Error = ANNError;
539
540    async fn load_with<P>(_: &P, num_points: &usize) -> ANNResult<Self>
541    where
542        P: StorageReadProvider,
543    {
544        Ok(TableDeleteProviderAsync::new(*num_points))
545    }
546}
547
548///////////////////
549// Data Provider //
550///////////////////
551
552impl<U, V, D, Ctx> DataProvider for DefaultProvider<U, V, D, Ctx>
553where
554    U: AsyncFriendly,
555    V: AsyncFriendly,
556    D: AsyncFriendly,
557    Ctx: ExecutionContext,
558{
559    type Context = Ctx;
560    /// The `DefaultProvider` uses the identity map for IDs.
561    type InternalId = u32;
562    /// The `DefaultProvider` uses the identity map for IDs.
563    type ExternalId = u32;
564    /// Use a general error type for now.
565    type Error = ANNError;
566    /// The guard to (not) roll back pending changes.
567    type Guard = NoopGuard<u32>;
568
569    /// Translate an external id to its corresponding internal id.
570    fn to_internal_id(
571        &self,
572        _context: &Self::Context,
573        gid: &Self::ExternalId,
574    ) -> Result<Self::InternalId, Self::Error> {
575        Ok(*gid)
576    }
577
578    /// Translate an internal id its corresponding external id.
579    fn to_external_id(
580        &self,
581        _context: &Self::Context,
582        id: Self::InternalId,
583    ) -> Result<Self::ExternalId, Self::Error> {
584        Ok(id)
585    }
586}
587
588/// Support deletes when we have a valid delete provider.
589impl<U, V, Ctx> Delete for DefaultProvider<U, V, TableDeleteProviderAsync, Ctx>
590where
591    U: AsyncFriendly,
592    V: AsyncFriendly,
593    Ctx: ExecutionContext,
594{
595    fn release(
596        &self,
597        _context: &Ctx,
598        id: Self::InternalId,
599    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
600        self.deleted.undelete(id.into_usize());
601        let res = self
602            .neighbor_provider
603            .set_neighbors_sync(id.into_usize(), &[])
604            .map_err(|err| err.context(format!("resetting neighbors for undeleted id {}", id)));
605        std::future::ready(res)
606    }
607
608    /// Delete an item by external ID.
609    #[inline]
610    fn delete(
611        &self,
612        _context: &Ctx,
613        gid: &Self::ExternalId,
614    ) -> impl Future<Output = Result<(), Self::Error>> + Send {
615        self.deleted.delete(gid.into_usize());
616        std::future::ready(Ok(()))
617    }
618
619    /// Check the status via external ID.
620    #[inline]
621    fn status_by_external_id(
622        &self,
623        context: &Ctx,
624        gid: &Self::ExternalId,
625    ) -> impl Future<Output = Result<ElementStatus, Self::Error>> + Send {
626        // NOTE: ID translation is the identity, so we can refer to `status_by_internal_id`.
627        self.status_by_internal_id(context, *gid)
628    }
629
630    /// Check the status via internal ID.
631    #[inline]
632    fn status_by_internal_id(
633        &self,
634        _context: &Ctx,
635        id: Self::InternalId,
636    ) -> impl Future<Output = Result<ElementStatus, Self::Error>> + Send {
637        let status = if self.deleted.is_deleted(id.into_usize()) {
638            ElementStatus::Deleted
639        } else {
640            ElementStatus::Valid
641        };
642        std::future::ready(Ok(status))
643    }
644}
645
646impl NeighborAccessor for &SimpleNeighborProviderAsync {
647    async fn get_neighbors(
648        &mut self,
649        id: Self::Id,
650        neighbors: &mut AdjacencyList<Self::Id>,
651    ) -> ANNResult<()> {
652        self.get_neighbors_sync(id.into_usize(), neighbors)?;
653        Ok(())
654    }
655}
656
657impl NeighborAccessorMut for &SimpleNeighborProviderAsync {
658    async fn set_neighbors(&mut self, id: u32, neighbors: &[u32]) -> ANNResult<()> {
659        self.set_neighbors_sync(id.into_usize(), neighbors)?;
660        Ok(())
661    }
662
663    async fn append_vector(&mut self, id: u32, new_neighbor_ids: &[u32]) -> ANNResult<()> {
664        self.append_vector_sync(id.into_usize(), new_neighbor_ids)?;
665        Ok(())
666    }
667}
668
669impl<U, V, D, Ctx> DefaultAccessor for DefaultProvider<U, V, D, Ctx>
670where
671    U: AsyncFriendly,
672    V: AsyncFriendly,
673    D: AsyncFriendly,
674    Ctx: ExecutionContext,
675{
676    type Accessor<'a> = &'a SimpleNeighborProviderAsync;
677    fn default_accessor(&self) -> Self::Accessor<'_> {
678        self.neighbors()
679    }
680}
681
682////////////////
683// SetElement //
684////////////////
685
686// Assign to both the base and aux vector stores.
687impl<U, V, D, Ctx, T> SetElement<&[T]> for DefaultProvider<U, V, D, Ctx>
688where
689    T: VectorRepr,
690    U: AsyncFriendly + SetElementHelper<T>,
691    V: AsyncFriendly + SetElementHelper<T>,
692    D: AsyncFriendly,
693    Ctx: ExecutionContext,
694{
695    type SetError = ANNError;
696
697    /// Store the provided element in just the full-precision vector stores.
698    fn set_element(
699        &self,
700        _context: &Self::Context,
701        id: &u32,
702        element: &[T],
703    ) -> impl Future<Output = Result<Self::Guard, Self::SetError>> + Send {
704        // First try adding to the aux vector store
705        if let Err(err) = self.aux_vectors.set_element(id, element) {
706            return std::future::ready(Err(err));
707        }
708
709        // Next, add to the base vector store.
710        if let Err(err) = self.base_vectors.set_element(id, element) {
711            return std::future::ready(Err(err));
712        }
713
714        // Success.
715        std::future::ready(Ok(NoopGuard::new(*id)))
716    }
717}
718
719///////////
720// Tests //
721///////////
722
723#[cfg(test)]
724mod tests {
725    use super::*;
726    use crate::model::graph::provider::async_::{
727        common::{NoStore, TableBasedDeletes},
728        inmem::CreateFullPrecision,
729    };
730
731    #[tokio::test]
732    async fn test_data_provider_and_delete_interface() {
733        let ctx = &DefaultContext;
734        let provider = DefaultProvider::new_empty(
735            DefaultProviderParameters {
736                max_points: 10,
737                frozen_points: NonZeroUsize::new(2).unwrap(),
738                dim: 5,
739                metric: Metric::L2,
740                prefetch_lookahead: None,
741                max_degree: (64.0 * 1.2) as u32,
742                prefetch_cache_line_level: None,
743            },
744            CreateFullPrecision::<f32>::new(5, None),
745            NoStore,
746            TableBasedDeletes,
747        )
748        .unwrap();
749
750        // Iterator
751        assert_eq!((&provider).into_iter(), 0..(10 + 2));
752
753        let iter = provider.iter();
754        for i in iter.clone() {
755            assert_eq!(provider.to_external_id(ctx, i).unwrap(), i);
756            assert_eq!(provider.to_internal_id(ctx, &i).unwrap(), i);
757            assert_eq!(
758                provider.status_by_internal_id(ctx, i).await.unwrap(),
759                ElementStatus::Valid
760            );
761            assert_eq!(
762                provider.status_by_external_id(ctx, &i).await.unwrap(),
763                ElementStatus::Valid
764            );
765
766            // Delete by external ID.
767            provider.delete(ctx, &i).await.unwrap();
768            assert_eq!(
769                provider.status_by_internal_id(ctx, i).await.unwrap(),
770                ElementStatus::Deleted
771            );
772            assert_eq!(
773                provider.status_by_external_id(ctx, &i).await.unwrap(),
774                ElementStatus::Deleted
775            );
776        }
777
778        // Call `release` to "undelete" it ID.
779        for i in iter.clone() {
780            // set adjacency list to non-empty before release
781            provider
782                .neighbors()
783                .set_neighbors(i, &[1, 2])
784                .await
785                .unwrap();
786            provider.release(ctx, i).await.unwrap();
787            assert_eq!(
788                provider.status_by_internal_id(ctx, i).await.unwrap(),
789                ElementStatus::Valid
790            );
791            assert_eq!(
792                provider.status_by_external_id(ctx, &i).await.unwrap(),
793                ElementStatus::Valid
794            );
795            // check that adjacency list was reset after release
796            let mut neighbors = AdjacencyList::new();
797            provider
798                .neighbors()
799                .get_neighbors(i, &mut neighbors)
800                .await
801                .unwrap();
802            assert!(neighbors.to_vec().is_empty());
803
804            // Put it back to "deleted" to test `clear`.
805            provider.delete(ctx, &i).await.unwrap();
806        }
807
808        provider.clear_delete_set();
809        for i in iter.clone() {
810            assert_eq!(
811                provider.status_by_internal_id(ctx, i).await.unwrap(),
812                ElementStatus::Valid
813            );
814            assert_eq!(
815                provider.status_by_external_id(ctx, &i).await.unwrap(),
816                ElementStatus::Valid
817            );
818        }
819
820        // out-of-bound set-element fails.
821        assert!(
822            provider
823                .set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0])
824                .await
825                .is_err()
826        );
827    }
828}