1use 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
34pub struct DefaultProvider<U, V = NoStore, D = NoDeletes, Ctx = DefaultContext> {
232 pub base_vectors: U,
234
235 pub aux_vectors: V,
237
238 pub(crate) neighbor_provider: SimpleNeighborProviderAsync<u32>,
240
241 pub(super) deleted: D,
245
246 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 pub max_points: usize,
258
259 pub frozen_points: NonZeroUsize,
262
263 pub dim: usize,
265
266 pub metric: Metric,
268
269 pub prefetch_lookahead: Option<usize>,
274
275 pub prefetch_cache_line_level: Option<PrefetchCacheLineLevel>,
276
277 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 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 #[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 pub fn starting_points(&self) -> ANNResult<Vec<u32>> {
338 Ok(self.start_points.range().collect())
339 }
340
341 pub fn iter(&self) -> std::ops::Range<u32> {
343 0..self.start_points.end()
344 }
345
346 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 pub fn capacity(&self) -> usize {
357 self.start_points.start().into_usize()
358 }
359
360 pub fn total_points(&self) -> usize {
362 self.start_points.end().into_usize()
363 }
364}
365
366impl<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 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 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
433impl<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
490impl<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 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
560impl<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 type InternalId = u32;
574 type ExternalId = u32;
576 type Error = ANNError;
578 type Guard = NoopGuard<u32>;
580
581 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 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
600impl<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 #[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 #[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 self.status_by_internal_id(context, *gid)
640 }
641
642 #[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
694impl<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 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 if let Err(err) = self.aux_vectors.set_element(id, element) {
718 return std::future::ready(Err(err));
719 }
720
721 if let Err(err) = self.base_vectors.set_element(id, element) {
723 return std::future::ready(Err(err));
724 }
725
726 std::future::ready(Ok(NoopGuard::new(*id)))
728 }
729}
730
731#[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 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 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 for i in iter.clone() {
792 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 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 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 assert!(
834 provider
835 .set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0])
836 .await
837 .is_err()
838 );
839 }
840}