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