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