1use std::{
7 fmt::Debug,
8 future::Future,
9 io::{Read, Write},
10 num::NonZeroUsize,
11 str::FromStr,
12};
13
14use diskann_quantization::{
15 alloc::{GlobalAllocator, Poly},
16 spherical::iface::{self as spherical_iface, try_deserialize, Opaque, Quantizer},
17};
18use serde::{Deserialize, Serialize};
19
20use bf_tree::{BfTree, Config};
21use diskann::{
22 default_post_processor,
23 error::{ErrorExt, Infallible, RankedError},
24 graph::{
25 glue::{
26 self, Batch, CopyIds, DefaultPostProcessor, InplaceDeleteStrategy, InsertStrategy,
27 MultiInsertStrategy, PruneStrategy, SearchStrategy,
28 },
29 strategy::{FullPrecision, Quantized},
30 workingset::map,
31 AdjacencyList, SearchOutputBuffer,
32 },
33 neighbor::Neighbor,
34 provider::{DataProvider, DefaultContext, Delete, ElementStatus, HasId, NoopGuard, SetElement},
35 utils::{IntoUsize, VectorRepr},
36 ANNError, ANNResult,
37};
38use diskann_utils::{
39 future::{AsyncFriendly, SendFuture},
40 views::MatrixView,
41};
42use diskann_vector::{distance::Metric, DistanceFunction, PreprocessedDistanceFunction};
43
44use super::{
45 neighbors::{NeighborAccessor, NeighborProvider},
46 quant::QuantVectorProvider,
47 vectors::VectorProvider,
48 AccessError, NoStore,
49};
50use crate::locks::StripedLocks;
51use diskann_providers::model::graph::provider::async_::distances::UnwrapErr;
52use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, StorageWriteProvider};
53
54pub struct BfTreeProvider<T, Q = QuantVectorProvider>
181where
182 T: VectorRepr,
183{
184 pub(super) quant_vectors: Q,
187
188 pub(super) full_vectors: VectorProvider<T>,
191
192 pub(crate) neighbor_provider: NeighborProvider<u32>,
195
196 pub(super) metric: Metric,
199
200 pub(crate) graph_params: Option<GraphParams>,
203
204 pub(crate) use_snapshot: bool,
206 pub(crate) locks: StripedLocks,
225}
226
227#[derive(Debug, Clone)]
228pub struct BfTreeProviderParameters {
229 pub max_points: usize,
231
232 pub num_start_points: NonZeroUsize,
234
235 pub dim: usize,
237
238 pub metric: Metric,
240
241 pub max_degree: u32,
245
246 pub vector_provider_config: Config,
252
253 pub quant_vector_provider_config: Config,
256
257 pub neighbor_list_provider_config: Config,
260
261 pub graph_params: Option<GraphParams>,
263
264 pub use_snapshot: bool,
266}
267
268impl<T, Q> BfTreeProvider<T, Q>
269where
270 T: VectorRepr,
271{
272 fn new_empty<TQ>(mut params: BfTreeProviderParameters, quant_precursor: TQ) -> ANNResult<Self>
286 where
287 Self: StartPoint<T>,
288 TQ: CreateQuantProvider<Target = Q>,
289 {
290 params
293 .vector_provider_config
294 .use_snapshot(params.use_snapshot);
295 params
296 .neighbor_list_provider_config
297 .use_snapshot(params.use_snapshot);
298 params
299 .quant_vector_provider_config
300 .use_snapshot(params.use_snapshot);
301
302 Ok(Self {
303 quant_vectors: quant_precursor.create(params.quant_vector_provider_config)?,
304 full_vectors: VectorProvider::new_with_config(
305 params.max_points,
306 params.dim,
307 params.num_start_points.get(),
308 params.vector_provider_config,
309 )?,
310 neighbor_provider: NeighborProvider::new_with_config(
311 params.max_degree,
312 params.neighbor_list_provider_config,
313 )?,
314 metric: params.metric,
315 graph_params: params.graph_params,
316 use_snapshot: params.use_snapshot,
317 locks: StripedLocks::new(),
318 })
319 }
320
321 pub fn new<TQ>(
336 params: BfTreeProviderParameters,
337 start_points: MatrixView<'_, T>,
338 quant_precursor: TQ,
339 ) -> ANNResult<Self>
340 where
341 Self: StartPoint<T>,
342 TQ: CreateQuantProvider<Target = Q>,
343 {
344 if start_points.nrows() != params.num_start_points.get() {
346 return Err(ANNError::log_async_index_error(format!(
347 "start_points matrix has {} rows, but params.num_start_points is {}",
348 start_points.nrows(),
349 params.num_start_points.get(),
350 )));
351 }
352
353 let provider = Self::new_empty(params.clone(), quant_precursor)?;
354 provider.set_start_points(Hidden(()), start_points)?;
355 {
356 let mut scratch = provider.neighbor_provider.scratch(&provider.locks);
362 for i in 0..params.max_points {
363 let vector_id = i as u32;
364 scratch.write_neighbors(vector_id, &[])?;
365 }
366 }
367 Ok(provider)
368 }
369
370 pub fn starting_points(&self) -> ANNResult<Vec<u32>> {
382 self.full_vectors.starting_points()
383 }
384
385 pub fn iter(&self) -> std::ops::Range<u32> {
387 0..(self.full_vectors.total() as u32)
388 }
389
390 pub fn num_start_points(&self) -> usize {
391 self.full_vectors.num_start_points
392 }
393
394 pub fn max_points(&self) -> usize {
396 self.full_vectors.max_vectors
397 }
398
399 pub fn dim(&self) -> usize {
401 self.full_vectors.dim()
402 }
403
404 pub fn metric(&self) -> Metric {
406 self.metric
407 }
408
409 pub fn max_degree(&self) -> u32 {
411 self.neighbor_provider.max_degree()
412 }
413}
414
415impl<T> BfTreeProvider<T, QuantVectorProvider>
416where
417 T: VectorRepr,
418{
419 pub fn counts_for_get_vector(&self) -> (usize, usize) {
422 (
423 self.full_vectors.num_get_calls.get(),
424 self.quant_vectors.num_get_calls.get(),
425 )
426 }
427}
428
429impl<T> BfTreeProvider<T, NoStore>
430where
431 T: VectorRepr,
432{
433 pub fn counts_for_get_vector(&self) -> (usize, usize) {
436 (self.full_vectors.num_get_calls.get(), 0)
437 }
438}
439
440pub(crate) trait DeleteQuant {
445 fn delete_vector(&self, id: usize);
446}
447
448impl DeleteQuant for QuantVectorProvider {
449 fn delete_vector(&self, id: usize) {
450 QuantVectorProvider::delete_vector(self, id);
451 }
452}
453
454impl DeleteQuant for NoStore {
455 fn delete_vector(&self, _id: usize) {}
456}
457
458impl<T, Q> Delete for BfTreeProvider<T, Q>
471where
472 T: VectorRepr,
473 Q: AsyncFriendly + DeleteQuant,
474{
475 fn release(
476 &self,
477 _context: &Self::Context,
478 _id: Self::InternalId,
479 ) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send {
480 std::future::ready(Ok(()))
481 }
482
483 fn delete(
484 &self,
485 _context: &Self::Context,
486 gid: &Self::ExternalId,
487 ) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send {
488 let id = *gid;
489 let _guard = self.locks.lock(id as usize);
490 self.full_vectors.delete_vector(id as usize);
494 self.quant_vectors.delete_vector(id as usize);
495
496 std::future::ready(Ok(()))
497 }
498
499 fn status_by_external_id(
500 &self,
501 context: &Self::Context,
502 gid: &Self::ExternalId,
503 ) -> impl std::future::Future<Output = Result<diskann::provider::ElementStatus, Self::Error>> + Send
504 {
505 self.status_by_internal_id(context, *gid)
506 }
507
508 fn status_by_internal_id(
509 &self,
510 _context: &Self::Context,
511 id: Self::InternalId,
512 ) -> impl std::future::Future<Output = Result<diskann::provider::ElementStatus, Self::Error>> + Send
513 {
514 let status = match self.full_vectors.get_vector_sync(id.into_usize()) {
515 Ok(_) => Ok(ElementStatus::Valid),
516 Err(RankedError::Transient(_)) => Ok(ElementStatus::Deleted),
517 Err(RankedError::Error(e)) => Err(e),
518 };
519 std::future::ready(status)
520 }
521}
522
523impl<T, Q> IntoIterator for &BfTreeProvider<T, Q>
526where
527 T: VectorRepr,
528{
529 type Item = u32;
530 type IntoIter = std::ops::Range<u32>;
531 fn into_iter(self) -> Self::IntoIter {
532 self.iter()
533 }
534}
535
536pub trait CreateQuantProvider {
542 type Target;
545
546 fn create(self, bf_tree_config: Config) -> ANNResult<Self::Target>;
550}
551
552impl CreateQuantProvider for NoStore {
553 type Target = NoStore;
554 fn create(self, _bf_tree_config: Config) -> ANNResult<Self::Target> {
555 Ok(self)
556 }
557}
558
559impl CreateQuantProvider for Poly<dyn Quantizer> {
562 type Target = QuantVectorProvider;
563 fn create(self, bf_tree_config: Config) -> ANNResult<Self::Target> {
564 QuantVectorProvider::new_with_config(self, bf_tree_config)
565 }
566}
567
568impl<T, Q> BfTreeProvider<T, Q>
569where
570 T: VectorRepr,
571 Q: AsyncFriendly,
572{
573 pub fn neighbors(&self) -> &NeighborProvider<u32> {
574 &self.neighbor_provider
575 }
576}
577
578impl<T, Q> DataProvider for BfTreeProvider<T, Q>
583where
584 T: VectorRepr,
585 Q: AsyncFriendly,
586{
587 type Context = DefaultContext;
588
589 type InternalId = u32;
592
593 type ExternalId = u32;
596
597 type Error = ANNError;
600
601 type Guard = NoopGuard<u32>;
603
604 fn to_internal_id(
607 &self,
608 _context: &DefaultContext,
609 gid: &Self::ExternalId,
610 ) -> Result<Self::InternalId, Self::Error> {
611 Ok(*gid)
612 }
613
614 fn to_external_id(
617 &self,
618 _context: &DefaultContext,
619 id: Self::InternalId,
620 ) -> Result<Self::ExternalId, Self::Error> {
621 Ok(id)
622 }
623}
624
625impl<T, Q> HasId for BfTreeProvider<T, Q>
626where
627 T: VectorRepr,
628 Q: AsyncFriendly,
629{
630 type Id = u32;
631}
632
633impl<T> SetElement<&[T]> for BfTreeProvider<T, QuantVectorProvider>
640where
641 T: VectorRepr,
642{
643 type SetError = ANNError;
644
645 fn set_element(
652 &self,
653 _context: &Self::Context,
654 id: &u32,
655 element: &[T],
656 ) -> impl Future<Output = Result<Self::Guard, Self::SetError>> + Send {
657 let _guard = self.locks.lock(id.into_usize());
658
659 if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) {
661 return std::future::ready(Err(err));
662 }
663
664 if let Err(err) = self.quant_vectors.set_vector_sync(id.into_usize(), element) {
666 debug_assert!(
667 false,
668 "quant write failed after full-precision success: {err}"
669 );
670 return std::future::ready(Err(err));
671 }
672
673 std::future::ready(Ok(NoopGuard::new(*id)))
674 }
675}
676
677impl<T> SetElement<&[T]> for BfTreeProvider<T, NoStore>
680where
681 T: VectorRepr,
682{
683 type SetError = ANNError;
684
685 fn set_element(
688 &self,
689 _context: &Self::Context,
690 id: &u32,
691 element: &[T],
692 ) -> impl Future<Output = Result<Self::Guard, Self::SetError>> + Send {
693 let _guard = self.locks.lock(id.into_usize());
694
695 if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) {
696 return std::future::ready(Err(err));
697 }
698
699 std::future::ready(Ok(NoopGuard::new(*id)))
700 }
701}
702
703pub struct Hidden(());
711
712pub trait StartPoint<T> {
718 #[doc(hidden)]
724 fn set_start_points(&self, hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()>;
725}
726
727impl<T> StartPoint<T> for BfTreeProvider<T, QuantVectorProvider>
736where
737 T: VectorRepr,
738{
739 fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> {
740 let start_point_ids = self.full_vectors.starting_points()?;
741 if start_points.nrows() != start_point_ids.len() {
742 return Err(ANNError::log_async_index_error(format!(
743 "expected start_points to contain `{}` rows, instead it has {}",
744 start_point_ids.len(),
745 start_points.nrows(),
746 )));
747 }
748
749 let mut scratch = self.neighbor_provider.scratch(&self.locks);
750 for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) {
751 self.full_vectors.set_vector_sync(id.into_usize(), v)?;
753 self.quant_vectors.set_vector_sync(id.into_usize(), v)?;
754 scratch.write_neighbors(id, &[])?;
756 }
757
758 Ok(())
759 }
760}
761
762impl<T> StartPoint<T> for BfTreeProvider<T, NoStore>
767where
768 T: VectorRepr,
769{
770 fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> {
771 let start_point_ids = self.full_vectors.starting_points()?;
772 if start_points.nrows() != start_point_ids.len() {
773 return Err(ANNError::log_async_index_error(format!(
774 "expected start_points to contain `{}` rows, instead it has {}",
775 start_point_ids.len(),
776 start_points.nrows(),
777 )));
778 }
779
780 let mut scratch = self.neighbor_provider.scratch(&self.locks);
781 for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) {
782 self.full_vectors.set_vector_sync(id.into_usize(), v)?;
784 scratch.write_neighbors(id, &[])?;
786 }
787
788 Ok(())
789 }
790}
791
792pub struct FullAccessor<'a, T, Q>
804where
805 T: VectorRepr,
806 Q: AsyncFriendly,
807{
808 provider: &'a BfTreeProvider<T, Q>,
810 computer: T::QueryDistance,
812 element: Box<[T]>,
814}
815
816impl<'a, T, Q> FullAccessor<'a, T, Q>
817where
818 T: VectorRepr,
819 Q: AsyncFriendly,
820{
821 pub(crate) fn new(provider: &'a BfTreeProvider<T, Q>, query: &[T]) -> Self {
822 Self {
823 provider,
824 computer: T::query_distance(query, provider.metric),
825 element: (0..provider.full_vectors.dim())
826 .map(|_| T::default())
827 .collect(),
828 }
829 }
830
831 fn get_distance(&mut self, id: u32) -> Result<f32, AccessError> {
832 self.provider
833 .full_vectors
834 .get_vector_into(id.into_usize(), &mut self.element)
835 .map(|_: ()| self.computer.evaluate_similarity(&self.element))
836 }
837}
838
839impl<T, Q> HasId for FullAccessor<'_, T, Q>
840where
841 T: VectorRepr,
842 Q: AsyncFriendly,
843{
844 type Id = u32;
845}
846
847impl<T, Q> glue::SearchAccessor for FullAccessor<'_, T, Q>
848where
849 T: VectorRepr,
850 Q: AsyncFriendly,
851{
852 fn starting_points(&self) -> impl Future<Output = ANNResult<Vec<u32>>> {
853 std::future::ready(self.provider.starting_points())
854 }
855
856 async fn start_point_distances<F>(&mut self, mut f: F) -> ANNResult<()>
857 where
858 F: FnMut(Self::Id, f32) + Send,
859 {
860 for i in self.provider.starting_points()? {
861 f(
862 i,
863 self.get_distance(i)
864 .escalate("starting point retrieval must succeed")?,
865 )
866 }
867 Ok(())
868 }
869
870 async fn expand_beam<Itr, P, F>(
871 &mut self,
872 ids: Itr,
873 mut pred: P,
874 mut on_neighbors: F,
875 ) -> ANNResult<()>
876 where
877 Itr: Iterator<Item = Self::Id> + Send,
878 P: glue::HybridPredicate<Self::Id> + Send + Sync,
879 F: FnMut(Self::Id, f32) + Send,
880 {
881 let mut neighbors = AdjacencyList::new();
882 for n in ids {
883 self.provider.neighbors().get_neighbors(n, &mut neighbors)?;
884 for &id in neighbors.iter().filter(|i| pred.eval_mut(i)) {
885 if let Some(distance) = self
886 .get_distance(id)
887 .allow_transient("skipping deleted vectors")?
888 {
889 on_neighbors(id, distance)
890 }
891 }
892 }
893 Ok(())
894 }
895}
896
897pub struct QuantAccessor<'a, T>
908where
909 T: VectorRepr,
910{
911 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
912 computer: super::quant::QuantQueryComputer,
914 element: Box<[u8]>,
916}
917
918impl<'a, T> QuantAccessor<'a, T>
919where
920 T: VectorRepr,
921{
922 pub(crate) fn new(
923 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
924 query: &[T],
925 ) -> ANNResult<Self> {
926 let computer = provider.quant_vectors.query_computer(query)?;
927 Ok(Self {
928 provider,
929 computer,
930 element: (0..provider.quant_vectors.quantizer.bytes())
931 .map(|_| u8::default())
932 .collect(),
933 })
934 }
935
936 fn get_distance(&mut self, id: u32) -> Result<f32, AccessError> {
937 match self
938 .provider
939 .quant_vectors
940 .get_vector_into(id.into_usize(), &mut self.element)
941 {
942 Ok(()) => self
943 .computer
944 .evaluate(&self.element)
945 .map_err(RankedError::Error),
946 Err(err) => Err(err),
947 }
948 }
949}
950
951impl<T> HasId for QuantAccessor<'_, T>
952where
953 T: VectorRepr,
954{
955 type Id = u32;
956}
957
958impl<T> glue::SearchAccessor for QuantAccessor<'_, T>
959where
960 T: VectorRepr,
961{
962 fn starting_points(&self) -> impl Future<Output = ANNResult<Vec<u32>>> {
963 std::future::ready(self.provider.starting_points())
964 }
965
966 async fn start_point_distances<F>(&mut self, mut f: F) -> ANNResult<()>
967 where
968 F: FnMut(Self::Id, f32) + Send,
969 {
970 for i in self.provider.starting_points()? {
971 f(
972 i,
973 self.get_distance(i)
974 .escalate("starting point retrieval must succeed")?,
975 )
976 }
977 Ok(())
978 }
979
980 async fn expand_beam<Itr, P, F>(
981 &mut self,
982 ids: Itr,
983 mut pred: P,
984 mut on_neighbors: F,
985 ) -> ANNResult<()>
986 where
987 Itr: Iterator<Item = Self::Id> + Send,
988 P: glue::HybridPredicate<Self::Id> + Send + Sync,
989 F: FnMut(Self::Id, f32) + Send,
990 {
991 let mut neighbors = AdjacencyList::new();
992 for n in ids {
993 self.provider.neighbors().get_neighbors(n, &mut neighbors)?;
994 for &id in neighbors.iter().filter(|i| pred.eval_mut(i)) {
995 if let Some(distance) = self
996 .get_distance(id)
997 .allow_transient("skipping deleted vectors")?
998 {
999 on_neighbors(id, distance)
1000 }
1001 }
1002 }
1003 Ok(())
1004 }
1005}
1006
1007pub struct FullPruneAccessor<'a, T, Q>
1013where
1014 T: VectorRepr,
1015 Q: AsyncFriendly,
1016{
1017 provider: &'a BfTreeProvider<T, Q>,
1018 neighbors: NeighborAccessor<'a, u32>,
1019 set: map::Map<u32, Box<[T]>, map::Ref<[T]>>,
1020 distance: T::Distance,
1021}
1022
1023impl<'a, T, Q> FullPruneAccessor<'a, T, Q>
1024where
1025 T: VectorRepr,
1026 Q: AsyncFriendly,
1027{
1028 fn new(
1029 provider: &'a BfTreeProvider<T, Q>,
1030 set: map::Map<u32, Box<[T]>, map::Ref<[T]>>,
1031 ) -> Self {
1032 Self {
1033 provider,
1034 neighbors: provider.neighbor_provider.scratch(&provider.locks),
1035 set,
1036 distance: T::distance(provider.metric, Some(provider.full_vectors.dim())),
1037 }
1038 }
1039}
1040
1041impl<T, Q> HasId for FullPruneAccessor<'_, T, Q>
1042where
1043 T: VectorRepr,
1044 Q: AsyncFriendly,
1045{
1046 type Id = u32;
1047}
1048
1049impl<'q, T, Q> glue::PruneAccessor for FullPruneAccessor<'q, T, Q>
1050where
1051 T: VectorRepr,
1052 Q: AsyncFriendly,
1053{
1054 type ElementRef<'a> = &'a [T];
1055
1056 type View<'a>
1057 = map::View<'a, u32, Box<[T]>, map::Ref<[T]>>
1058 where
1059 Self: 'a;
1060
1061 type Distance<'a>
1062 = &'a T::Distance
1063 where
1064 Self: 'a;
1065
1066 type Neighbors<'a>
1067 = diskann::provider::Neighbors<'a, NeighborAccessor<'q, u32>>
1068 where
1069 Self: 'a;
1070
1071 fn fill<Itr>(
1072 &mut self,
1073 itr: Itr,
1074 ) -> impl SendFuture<ANNResult<(Self::View<'_>, Self::Distance<'_>)>>
1075 where
1076 Itr: ExactSizeIterator<Item = Self::Id> + Clone + Send + Sync,
1077 {
1078 let mut buf: Option<Box<[T]>> = None;
1079
1080 let view = self.set.fill(itr, |i: u32| -> ANNResult<_> {
1081 let mut b = match buf.take() {
1082 Some(b) => b,
1083 None => std::iter::repeat_n(T::default(), self.provider.dim()).collect(),
1084 };
1085
1086 match self
1087 .provider
1088 .full_vectors
1089 .get_vector_into(i.into_usize(), &mut b)
1090 .allow_transient("transient errors allowed during fill")?
1091 {
1092 Some(()) => Ok(Some(b)),
1093 None => {
1094 buf = Some(b);
1095 Ok(None)
1096 }
1097 }
1098 });
1099
1100 let result = view.map(|v| (v, &self.distance));
1101 std::future::ready(result)
1102 }
1103
1104 fn neighbors(&mut self) -> Self::Neighbors<'_> {
1105 diskann::provider::Neighbors(&mut self.neighbors)
1106 }
1107}
1108
1109pub struct QuantPruneAccessor<'a, T>
1115where
1116 T: VectorRepr,
1117{
1118 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
1119 neighbors: NeighborAccessor<'a, u32>,
1120 set: map::Map<u32, Owned>,
1121 distance: UnwrapErr<spherical_iface::DistanceComputer, spherical_iface::DistanceError>,
1122}
1123
1124impl<'a, T> QuantPruneAccessor<'a, T>
1125where
1126 T: VectorRepr,
1127{
1128 fn new(
1129 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
1130 capacity: usize,
1131 ) -> ANNResult<Self> {
1132 let distance = provider
1133 .quant_vectors
1134 .distance_computer()
1135 .map(UnwrapErr::new)?;
1136 let set = map::Builder::new(map::Capacity::Default).build(capacity);
1137 Ok(Self {
1138 provider,
1139 neighbors: provider.neighbor_provider.scratch(&provider.locks),
1140 set,
1141 distance,
1142 })
1143 }
1144}
1145
1146impl<T> HasId for QuantPruneAccessor<'_, T>
1147where
1148 T: VectorRepr,
1149{
1150 type Id = u32;
1151}
1152
1153impl<'q, T> glue::PruneAccessor for QuantPruneAccessor<'q, T>
1154where
1155 T: VectorRepr,
1156{
1157 type ElementRef<'a> = Opaque<'a>;
1158
1159 type View<'a>
1160 = map::View<'a, u32, Owned>
1161 where
1162 Self: 'a;
1163
1164 type Distance<'a>
1165 = &'a UnwrapErr<spherical_iface::DistanceComputer, spherical_iface::DistanceError>
1166 where
1167 Self: 'a;
1168
1169 type Neighbors<'a>
1170 = diskann::provider::Neighbors<'a, NeighborAccessor<'q, u32>>
1171 where
1172 Self: 'a;
1173
1174 fn fill<Itr>(
1175 &mut self,
1176 itr: Itr,
1177 ) -> impl SendFuture<ANNResult<(Self::View<'_>, Self::Distance<'_>)>>
1178 where
1179 Itr: ExactSizeIterator<Item = Self::Id> + Clone + Send + Sync,
1180 {
1181 let mut buf: Option<Box<[u8]>> = None;
1182 let bytes = self.provider.quant_vectors.quantizer.bytes();
1183
1184 let view = self.set.fill(itr, |i: u32| -> ANNResult<_> {
1185 let mut b = match buf.take() {
1186 Some(b) => b,
1187 None => std::iter::repeat_n(0, bytes).collect(),
1188 };
1189
1190 match self
1191 .provider
1192 .quant_vectors
1193 .get_vector_into(i.into_usize(), &mut b)
1194 .allow_transient("transient errors allowed during fill")?
1195 {
1196 Some(()) => Ok(Some(Owned(b))),
1197 None => {
1198 buf = Some(b);
1199 Ok(None)
1200 }
1201 }
1202 });
1203
1204 let result = view.map(|v| (v, &self.distance));
1205 std::future::ready(result)
1206 }
1207
1208 fn neighbors(&mut self) -> Self::Neighbors<'_> {
1209 diskann::provider::Neighbors(&mut self.neighbors)
1210 }
1211}
1212
1213pub struct Owned(Box<[u8]>);
1221
1222impl<'short> diskann_utils::Reborrow<'short> for Owned {
1223 type Target = Opaque<'short>;
1224 fn reborrow(&'short self) -> Self::Target {
1225 Opaque::new(&self.0)
1226 }
1227}
1228
1229impl<'a, T, Q> SearchStrategy<'a, BfTreeProvider<T, Q>, &'a [T]> for FullPrecision
1237where
1238 T: VectorRepr,
1239 Q: AsyncFriendly,
1240{
1241 type SearchAccessor = FullAccessor<'a, T, Q>;
1242 type SearchAccessorError = Infallible;
1243
1244 fn search_accessor(
1245 &'a self,
1246 provider: &'a BfTreeProvider<T, Q>,
1247 _context: &'a DefaultContext,
1248 query: &'a [T],
1249 ) -> Result<Self::SearchAccessor, Self::SearchAccessorError> {
1250 Ok(FullAccessor::new(provider, query))
1251 }
1252}
1253
1254impl<'a, T, Q> DefaultPostProcessor<'a, BfTreeProvider<T, Q>, &'a [T]> for FullPrecision
1255where
1256 T: VectorRepr,
1257 Q: AsyncFriendly,
1258{
1259 default_post_processor!(glue::Pipeline<glue::FilterStartPoints, CopyIds>);
1260}
1261
1262impl<T, Q> PruneStrategy<BfTreeProvider<T, Q>> for FullPrecision
1264where
1265 T: VectorRepr,
1266 Q: AsyncFriendly,
1267{
1268 type PruneAccessor<'a> = FullPruneAccessor<'a, T, Q>;
1269 type PruneAccessorError = diskann::error::Infallible;
1270
1271 fn prune_accessor<'a>(
1272 &'a self,
1273 provider: &'a BfTreeProvider<T, Q>,
1274 _context: &'a DefaultContext,
1275 capacity: usize,
1276 ) -> Result<Self::PruneAccessor<'a>, Self::PruneAccessorError> {
1277 let set = map::Builder::new(map::Capacity::Default).build(capacity);
1278 Ok(FullPruneAccessor::new(provider, set))
1279 }
1280}
1281
1282impl<'a, T, Q> InsertStrategy<'a, BfTreeProvider<T, Q>, &'a [T]> for FullPrecision
1283where
1284 T: VectorRepr,
1285 Q: AsyncFriendly,
1286{
1287 type PruneStrategy = Self;
1288 fn prune_strategy(&self) -> Self::PruneStrategy {
1289 *self
1290 }
1291}
1292
1293impl<T, Q, B> MultiInsertStrategy<BfTreeProvider<T, Q>, B> for FullPrecision
1294where
1295 T: VectorRepr,
1296 Q: AsyncFriendly,
1297 B: for<'a> Batch<Element<'a> = &'a [T]> + Debug,
1298{
1299 type Seed = map::Builder<u32, map::Ref<[T]>>;
1300 type FinishError = diskann::error::Infallible;
1301 type PruneStrategy = Self;
1302 type InsertStrategy = Self;
1303
1304 fn insert_strategy(&self) -> Self::InsertStrategy {
1305 *self
1306 }
1307
1308 fn finish<Itr>(
1309 &self,
1310 _provider: &BfTreeProvider<T, Q>,
1311 _ctx: &DefaultContext,
1312 batch: &std::sync::Arc<B>,
1313 ids: Itr,
1314 ) -> impl std::future::Future<Output = Result<Self::Seed, Self::FinishError>> + Send
1315 where
1316 Itr: ExactSizeIterator<Item = u32> + Send,
1317 {
1318 let overlay = map::Overlay::from_batch(batch.clone(), ids);
1319 let builder = map::Builder::new(map::Capacity::Default).with_overlay(overlay);
1320 std::future::ready(Ok(builder))
1321 }
1322
1323 fn seeded_prune_accessor<'a>(
1324 &'a self,
1325 provider: &'a BfTreeProvider<T, Q>,
1326 _context: &'a DefaultContext,
1327 seed: &'a Self::Seed,
1328 capacity: usize,
1329 ) -> ANNResult<FullPruneAccessor<'a, T, Q>> {
1330 let set = seed.clone().build(capacity);
1331 Ok(FullPruneAccessor::new(provider, set))
1332 }
1333}
1334
1335impl<T, Q> InplaceDeleteStrategy<BfTreeProvider<T, Q>> for FullPrecision
1344where
1345 T: VectorRepr,
1346 Q: AsyncFriendly,
1347{
1348 type DeleteElementError = ANNError;
1349 type DeleteElement<'a> = &'a [T];
1350 type DeleteElementGuard = Box<[T]>;
1351 type PruneStrategy = Self;
1352 type DeleteSearchAccessor<'a> = FullAccessor<'a, T, Q>;
1353 type SearchPostProcessor = CopyIds;
1354 type SearchStrategy = Self;
1355 fn search_strategy(&self) -> Self::SearchStrategy {
1356 Self
1357 }
1358
1359 fn prune_strategy(&self) -> Self::PruneStrategy {
1360 Self
1361 }
1362
1363 fn search_post_processor(&self) -> Self::SearchPostProcessor {
1364 CopyIds
1365 }
1366
1367 async fn get_delete_element<'a>(
1368 &'a self,
1369 provider: &'a BfTreeProvider<T, Q>,
1370 _context: &'a DefaultContext,
1371 id: u32,
1372 ) -> Result<Self::DeleteElementGuard, Self::DeleteElementError> {
1373 use diskann::error::ErrorExt;
1374 let elt = provider
1375 .full_vectors
1376 .get_vector_sync(id.into_usize())
1377 .escalate("get_delete_element: failed to read vector for inplace delete")?
1378 .into();
1379 Ok(elt)
1380 }
1381}
1382
1383impl<'a, T> SearchStrategy<'a, BfTreeProvider<T, QuantVectorProvider>, &'a [T]> for Quantized
1387where
1388 T: VectorRepr,
1389{
1390 type SearchAccessor = QuantAccessor<'a, T>;
1391 type SearchAccessorError = ANNError;
1392
1393 fn search_accessor(
1394 &'a self,
1395 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
1396 _context: &'a DefaultContext,
1397 query: &'a [T],
1398 ) -> Result<Self::SearchAccessor, Self::SearchAccessorError> {
1399 QuantAccessor::new(provider, query)
1400 }
1401}
1402
1403impl<'a, T> DefaultPostProcessor<'a, BfTreeProvider<T, QuantVectorProvider>, &'a [T]> for Quantized
1404where
1405 T: VectorRepr,
1406{
1407 default_post_processor!(glue::Pipeline<glue::FilterStartPoints, Rerank>);
1408}
1409
1410impl<'a, T> InsertStrategy<'a, BfTreeProvider<T, QuantVectorProvider>, &'a [T]> for Quantized
1411where
1412 T: VectorRepr,
1413{
1414 type PruneStrategy = Self;
1415 fn prune_strategy(&self) -> Self::PruneStrategy {
1416 *self
1417 }
1418}
1419
1420impl<T, B> MultiInsertStrategy<BfTreeProvider<T, QuantVectorProvider>, B> for Quantized
1421where
1422 T: VectorRepr,
1423 B: glue::Batch,
1424 B: for<'a> Batch<Element<'a> = &'a [T]> + Debug,
1425{
1426 type Seed = ();
1427 type FinishError = diskann::error::Infallible;
1428 type PruneStrategy = Self;
1429 type InsertStrategy = Self;
1430
1431 fn insert_strategy(&self) -> Self::InsertStrategy {
1432 *self
1433 }
1434
1435 fn finish<Itr>(
1436 &self,
1437 _provider: &BfTreeProvider<T, QuantVectorProvider>,
1438 _ctx: &DefaultContext,
1439 _batch: &std::sync::Arc<B>,
1440 _ids: Itr,
1441 ) -> impl std::future::Future<Output = Result<Self::Seed, Self::FinishError>> + Send
1442 where
1443 Itr: ExactSizeIterator<Item = u32> + Send,
1444 {
1445 std::future::ready(Ok(()))
1446 }
1447
1448 fn seeded_prune_accessor<'a>(
1449 &'a self,
1450 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
1451 _context: &'a DefaultContext,
1452 _seed: &'a (),
1453 capacity: usize,
1454 ) -> ANNResult<QuantPruneAccessor<'a, T>> {
1455 QuantPruneAccessor::new(provider, capacity)
1456 }
1457}
1458
1459impl<T> InplaceDeleteStrategy<BfTreeProvider<T, QuantVectorProvider>> for Quantized
1466where
1467 T: VectorRepr,
1468{
1469 type DeleteElementError = ANNError;
1470 type DeleteElement<'a> = &'a [T];
1471 type DeleteElementGuard = Box<[T]>;
1472 type PruneStrategy = Self;
1473 type DeleteSearchAccessor<'a> = QuantAccessor<'a, T>;
1474 type SearchPostProcessor = Rerank;
1475 type SearchStrategy = Self;
1476 fn search_strategy(&self) -> Self::SearchStrategy {
1477 *self
1478 }
1479
1480 fn prune_strategy(&self) -> Self::PruneStrategy {
1481 *self
1482 }
1483
1484 fn search_post_processor(&self) -> Self::SearchPostProcessor {
1485 Rerank
1486 }
1487
1488 async fn get_delete_element<'a>(
1489 &'a self,
1490 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
1491 _context: &'a DefaultContext,
1492 id: u32,
1493 ) -> Result<Self::DeleteElementGuard, Self::DeleteElementError> {
1494 use diskann::error::ErrorExt;
1495 provider
1496 .full_vectors
1497 .get_vector_sync(id.into_usize())
1498 .escalate("get_delete_element: failed to read vector for inplace delete")
1499 .map(Into::into)
1500 }
1501}
1502
1503impl<T> PruneStrategy<BfTreeProvider<T, QuantVectorProvider>> for Quantized
1505where
1506 T: VectorRepr,
1507{
1508 type PruneAccessor<'a> = QuantPruneAccessor<'a, T>;
1509 type PruneAccessorError = ANNError;
1510
1511 fn prune_accessor<'a>(
1512 &'a self,
1513 provider: &'a BfTreeProvider<T, QuantVectorProvider>,
1514 _context: &'a DefaultContext,
1515 capacity: usize,
1516 ) -> Result<Self::PruneAccessor<'a>, Self::PruneAccessorError> {
1517 QuantPruneAccessor::new(provider, capacity)
1518 }
1519}
1520
1521#[derive(Debug, Default, Clone, Copy)]
1523pub struct Rerank;
1524
1525impl<'a, T> glue::SearchPostProcess<QuantAccessor<'a, T>, &[T]> for Rerank
1526where
1527 T: VectorRepr,
1528{
1529 type Error = ANNError;
1530
1531 fn post_process<I, B>(
1532 &self,
1533 accessor: &mut QuantAccessor<'a, T>,
1534 query: &[T],
1535 candidates: I,
1536 output: &mut B,
1537 ) -> impl Future<Output = Result<usize, Self::Error>> + Send
1538 where
1539 I: Iterator<Item = Neighbor<u32>> + Send,
1540 B: SearchOutputBuffer<u32> + Send + ?Sized,
1541 {
1542 use diskann::error::ErrorExt;
1543 let provider = accessor.provider;
1544 let f = T::distance(provider.metric, Some(provider.full_vectors.dim()));
1545
1546 let mut reranked: Vec<(u32, f32)> = Vec::new();
1547 for n in candidates {
1548 match provider
1549 .full_vectors
1550 .get_vector_sync(n.id.into_usize())
1551 .allow_transient("stale candidate during rerank")
1552 {
1553 Ok(Some(vec)) => {
1554 reranked.push((n.id, f.evaluate_similarity(query, &vec)));
1555 }
1556 Ok(None) => {
1557 }
1559 Err(e) => return std::future::ready(Err(e)),
1560 }
1561 }
1562
1563 reranked
1564 .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1565 std::future::ready(Ok(output.extend(reranked)))
1566 }
1567}
1568
1569#[derive(Serialize, Deserialize, Clone)]
1570pub struct BfTreeParams {
1571 pub bytes: usize,
1572 pub max_record_size: usize,
1573 pub leaf_page_size: usize,
1574}
1575
1576#[derive(Serialize, Deserialize, Clone)]
1577pub struct QuantParams {
1578 pub params_quant: BfTreeParams,
1579}
1580
1581#[derive(Serialize, Deserialize, Clone)]
1582pub struct SavedParams {
1583 pub max_points: usize,
1584 pub frozen_points: NonZeroUsize,
1585 pub dim: usize,
1586 pub metric: String,
1587 pub max_degree: u32,
1588 pub prefix: String,
1589 pub params_vector: BfTreeParams,
1590 pub params_neighbor: BfTreeParams,
1591 pub quant_params: Option<QuantParams>,
1592 pub graph_params: Option<GraphParams>,
1593 pub is_memory: bool,
1595 #[serde(default)]
1597 pub use_snapshot: bool,
1598}
1599
1600#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
1602#[serde(rename_all = "lowercase")]
1603pub enum VectorDtype {
1604 F32,
1605 F16,
1606 U8,
1607 I8,
1608}
1609
1610pub trait AsVectorDtype {
1613 const DATA_TYPE: VectorDtype;
1614}
1615
1616impl AsVectorDtype for f32 {
1617 const DATA_TYPE: VectorDtype = VectorDtype::F32;
1618}
1619
1620impl AsVectorDtype for half::f16 {
1621 const DATA_TYPE: VectorDtype = VectorDtype::F16;
1622}
1623
1624impl AsVectorDtype for i8 {
1625 const DATA_TYPE: VectorDtype = VectorDtype::I8;
1626}
1627
1628impl AsVectorDtype for u8 {
1629 const DATA_TYPE: VectorDtype = VectorDtype::U8;
1630}
1631
1632#[derive(Serialize, Deserialize, Clone, Debug)]
1635pub struct GraphParams {
1636 pub l_build: usize,
1641 pub alpha: f32,
1646 pub backedge_ratio: f32,
1649 pub vector_dtype: VectorDtype,
1651}
1652
1653pub struct BfTreePaths;
1656
1657impl BfTreePaths {
1658 pub fn params_json(prefix: &str) -> String {
1660 format!("{}_params.json", prefix)
1661 }
1662
1663 pub fn vectors_bftree(prefix: &str) -> std::path::PathBuf {
1665 std::path::PathBuf::from(format!("{}_vectors.bftree", prefix))
1666 }
1667
1668 pub fn neighbors_bftree(prefix: &str) -> std::path::PathBuf {
1670 std::path::PathBuf::from(format!("{}_neighbors.bftree", prefix))
1671 }
1672
1673 pub fn quant_bftree(prefix: &str) -> std::path::PathBuf {
1675 std::path::PathBuf::from(format!("{}_quant.bftree", prefix))
1676 }
1677
1678 pub fn delete_bin(prefix: &str) -> String {
1680 format!("{}_delete.bin", prefix)
1681 }
1682
1683 pub fn quant_data_bin(prefix: &str) -> String {
1685 format!("{}_quant_data.bin", prefix)
1686 }
1687}
1688
1689fn save_bftree(
1696 tree: &BfTree,
1697 target_path: std::path::PathBuf,
1698 use_snapshot: bool,
1699) -> ANNResult<()> {
1700 if !use_snapshot {
1701 return Err(ANNError::log_index_error(
1702 "cannot snapshot a BfTree that was not configured with use_snapshot(true)",
1703 ));
1704 }
1705 tree.cpr_snapshot(&target_path);
1706 Ok(())
1707}
1708
1709fn load_bftree(snapshot_path: std::path::PathBuf, use_snapshot: bool) -> Result<BfTree, ANNError> {
1714 BfTree::new_from_cpr_snapshot(snapshot_path, use_snapshot, None, None, None)
1715 .map_err(|e| ANNError::from(super::ConfigError(e)))
1716}
1717
1718impl<T> SaveWith<String> for BfTreeProvider<T, NoStore>
1723where
1724 T: VectorRepr,
1725{
1726 type Ok = usize;
1727 type Error = ANNError;
1728
1729 async fn save_with<P>(&self, storage: &P, prefix: &String) -> Result<Self::Ok, Self::Error>
1730 where
1731 P: StorageWriteProvider,
1732 {
1733 let saved_params = SavedParams {
1734 max_points: self.max_points(),
1735 frozen_points: NonZeroUsize::new(self.num_start_points())
1736 .ok_or_else(|| ANNError::log_index_error("num_start_points is zero"))?,
1737 dim: self.dim(),
1738 metric: self.metric().as_str().to_string(),
1739 max_degree: self.max_degree(),
1740 prefix: prefix.clone(),
1741 params_vector: BfTreeParams {
1742 bytes: self.full_vectors.config().get_cb_size_byte(),
1743 max_record_size: self.full_vectors.config().get_cb_max_record_size(),
1744 leaf_page_size: self.full_vectors.config().get_leaf_page_size(),
1745 },
1746 params_neighbor: BfTreeParams {
1747 bytes: self.neighbor_provider.config().get_cb_size_byte(),
1748 max_record_size: self.neighbor_provider.config().get_cb_max_record_size(),
1749 leaf_page_size: self.neighbor_provider.config().get_leaf_page_size(),
1750 },
1751 quant_params: None,
1752 graph_params: self.graph_params.clone(),
1753 is_memory: self.full_vectors.config().is_memory_backend(),
1754 use_snapshot: self.use_snapshot,
1755 };
1756
1757 debug_assert_eq!(
1758 self.full_vectors.config().is_memory_backend(),
1759 self.neighbor_provider.config().is_memory_backend(),
1760 "Vector and neighbor stores have mismatched storage backends"
1761 );
1762
1763 {
1764 let params_filename = BfTreePaths::params_json(&saved_params.prefix);
1765 let params_json = serde_json::to_string(&saved_params).map_err(|e| {
1766 ANNError::log_index_error(format!("Failed to serialize params: {}", e))
1767 })?;
1768 let mut params_writer = storage.create_for_write(¶ms_filename)?;
1769 params_writer.write_all(params_json.as_bytes())?;
1770 }
1771
1772 save_bftree(
1773 self.full_vectors.bftree(),
1774 BfTreePaths::vectors_bftree(&saved_params.prefix),
1775 self.use_snapshot,
1776 )?;
1777 save_bftree(
1778 self.neighbor_provider.bftree(),
1779 BfTreePaths::neighbors_bftree(&saved_params.prefix),
1780 self.use_snapshot,
1781 )?;
1782
1783 Ok(0)
1784 }
1785}
1786
1787impl<T> LoadWith<String> for BfTreeProvider<T, NoStore>
1788where
1789 T: VectorRepr,
1790{
1791 type Error = ANNError;
1792
1793 async fn load_with<P>(storage: &P, prefix: &String) -> Result<Self, Self::Error>
1794 where
1795 P: StorageReadProvider,
1796 {
1797 let saved_params: SavedParams = {
1798 let params_filename = BfTreePaths::params_json(prefix);
1799 let mut params_reader = storage.open_reader(¶ms_filename)?;
1800 let mut params_json = String::new();
1801 params_reader.read_to_string(&mut params_json)?;
1802 serde_json::from_str(¶ms_json).map_err(|e| {
1803 ANNError::log_index_error(format!("Failed to deserialize params: {}", e))
1804 })?
1805 };
1806
1807 let metric = Metric::from_str(&saved_params.metric)
1808 .map_err(|e| ANNError::log_index_error(format!("Failed to parse metric: {}", e)))?;
1809
1810 let vector_index = load_bftree(
1811 BfTreePaths::vectors_bftree(&saved_params.prefix),
1812 saved_params.use_snapshot,
1813 )?;
1814 let full_vectors = VectorProvider::<T>::new_from_bftree(
1815 saved_params.max_points,
1816 saved_params.dim,
1817 saved_params.frozen_points.get(),
1818 vector_index,
1819 );
1820
1821 let adjacency_list_index = load_bftree(
1822 BfTreePaths::neighbors_bftree(&saved_params.prefix),
1823 saved_params.use_snapshot,
1824 )?;
1825 let neighbor_provider = NeighborProvider::<u32>::new_from_bftree(
1826 saved_params.max_degree,
1827 adjacency_list_index,
1828 )?;
1829
1830 Ok(Self {
1831 quant_vectors: NoStore,
1832 full_vectors,
1833 neighbor_provider,
1834 metric,
1835 graph_params: saved_params.graph_params,
1836 use_snapshot: saved_params.use_snapshot,
1837 locks: StripedLocks::new(),
1838 })
1839 }
1840}
1841
1842impl<T> SaveWith<String> for BfTreeProvider<T, QuantVectorProvider>
1843where
1844 T: VectorRepr,
1845{
1846 type Ok = usize;
1847 type Error = ANNError;
1848
1849 async fn save_with<P>(&self, storage: &P, prefix: &String) -> Result<Self::Ok, Self::Error>
1850 where
1851 P: StorageWriteProvider,
1852 {
1853 let saved_params = SavedParams {
1854 max_points: self.max_points(),
1855 frozen_points: NonZeroUsize::new(self.num_start_points())
1856 .ok_or_else(|| ANNError::log_index_error("num_start_points is zero"))?,
1857 dim: self.dim(),
1858 metric: self.metric().as_str().to_string(),
1859 max_degree: self.max_degree(),
1860 prefix: prefix.clone(),
1861 params_vector: BfTreeParams {
1862 bytes: self.full_vectors.config().get_cb_size_byte(),
1863 max_record_size: self.full_vectors.config().get_cb_max_record_size(),
1864 leaf_page_size: self.full_vectors.config().get_leaf_page_size(),
1865 },
1866 params_neighbor: BfTreeParams {
1867 bytes: self.neighbor_provider.config().get_cb_size_byte(),
1868 max_record_size: self.neighbor_provider.config().get_cb_max_record_size(),
1869 leaf_page_size: self.neighbor_provider.config().get_leaf_page_size(),
1870 },
1871 quant_params: Some(QuantParams {
1872 params_quant: BfTreeParams {
1873 bytes: self.quant_vectors.config().get_cb_size_byte(),
1874 max_record_size: self.quant_vectors.config().get_cb_max_record_size(),
1875 leaf_page_size: self.quant_vectors.config().get_leaf_page_size(),
1876 },
1877 }),
1878 graph_params: self.graph_params.clone(),
1879 is_memory: self.full_vectors.config().is_memory_backend(),
1880 use_snapshot: self.use_snapshot,
1881 };
1882
1883 debug_assert_eq!(
1884 self.full_vectors.config().is_memory_backend(),
1885 self.neighbor_provider.config().is_memory_backend(),
1886 "Vector and neighbor stores have mismatched storage backends"
1887 );
1888 debug_assert_eq!(
1889 self.full_vectors.config().is_memory_backend(),
1890 self.quant_vectors.config().is_memory_backend(),
1891 "Vector and quant stores have mismatched storage backends"
1892 );
1893
1894 {
1895 let params_filename = BfTreePaths::params_json(&saved_params.prefix);
1896 let params_json = serde_json::to_string(&saved_params).map_err(|e| {
1897 ANNError::log_index_error(format!("Failed to serialize params: {}", e))
1898 })?;
1899 let mut params_writer = storage.create_for_write(¶ms_filename)?;
1900 params_writer.write_all(params_json.as_bytes())?;
1901 }
1902
1903 save_bftree(
1904 self.full_vectors.bftree(),
1905 BfTreePaths::vectors_bftree(&saved_params.prefix),
1906 self.use_snapshot,
1907 )?;
1908 save_bftree(
1909 self.neighbor_provider.bftree(),
1910 BfTreePaths::neighbors_bftree(&saved_params.prefix),
1911 self.use_snapshot,
1912 )?;
1913 save_bftree(
1914 self.quant_vectors.bftree(),
1915 BfTreePaths::quant_bftree(&saved_params.prefix),
1916 self.use_snapshot,
1917 )?;
1918
1919 let filename = BfTreePaths::quant_data_bin(&saved_params.prefix);
1920 let serialized = self
1921 .quant_vectors
1922 .quantizer
1923 .serialize(GlobalAllocator)
1924 .map_err(|e| ANNError::log_index_error(format!("{e}")))?;
1925 let mut writer = storage.create_for_write(&filename)?;
1926 writer.write_all(&serialized)?;
1927
1928 Ok(0)
1929 }
1930}
1931
1932impl<T> LoadWith<String> for BfTreeProvider<T, QuantVectorProvider>
1933where
1934 T: VectorRepr,
1935{
1936 type Error = ANNError;
1937
1938 async fn load_with<P>(storage: &P, prefix: &String) -> Result<Self, Self::Error>
1939 where
1940 P: StorageReadProvider,
1941 {
1942 let saved_params: SavedParams = {
1943 let params_filename = BfTreePaths::params_json(prefix);
1944 let mut params_reader = storage.open_reader(¶ms_filename)?;
1945 let mut params_json = String::new();
1946 params_reader.read_to_string(&mut params_json)?;
1947 serde_json::from_str(¶ms_json).map_err(|e| {
1948 ANNError::log_index_error(format!("Failed to deserialize params: {}", e))
1949 })?
1950 };
1951
1952 let _quant_params = saved_params.quant_params.ok_or_else(|| {
1953 ANNError::log_index_error("Missing quant_params in saved params for quantized provider")
1954 })?;
1955
1956 let metric = Metric::from_str(&saved_params.metric)
1957 .map_err(|e| ANNError::log_index_error(format!("Failed to parse metric: {}", e)))?;
1958
1959 let vector_index = load_bftree(
1960 BfTreePaths::vectors_bftree(&saved_params.prefix),
1961 saved_params.use_snapshot,
1962 )?;
1963 let full_vectors = VectorProvider::<T>::new_from_bftree(
1964 saved_params.max_points,
1965 saved_params.dim,
1966 saved_params.frozen_points.get(),
1967 vector_index,
1968 );
1969
1970 let adjacency_list_index = load_bftree(
1971 BfTreePaths::neighbors_bftree(&saved_params.prefix),
1972 saved_params.use_snapshot,
1973 )?;
1974 let neighbor_provider = NeighborProvider::<u32>::new_from_bftree(
1975 saved_params.max_degree,
1976 adjacency_list_index,
1977 )?;
1978
1979 let filename = BfTreePaths::quant_data_bin(&saved_params.prefix);
1980 let mut reader = storage.open_reader(&filename)?;
1981 let mut bytes = Vec::new();
1982 reader.read_to_end(&mut bytes)?;
1983 let quantizer: Poly<dyn Quantizer> = try_deserialize(&bytes, GlobalAllocator)
1984 .map_err(|e| ANNError::log_index_error(format!("{e}")))?;
1985
1986 let quant_vector_index = load_bftree(
1987 BfTreePaths::quant_bftree(&saved_params.prefix),
1988 saved_params.use_snapshot,
1989 )?;
1990 let quant_vectors = QuantVectorProvider::new_from_bftree(quantizer, quant_vector_index);
1991
1992 Ok(Self {
1993 quant_vectors,
1994 full_vectors,
1995 neighbor_provider,
1996 metric,
1997 graph_params: saved_params.graph_params,
1998 use_snapshot: saved_params.use_snapshot,
1999 locks: StripedLocks::new(),
2000 })
2001 }
2002}
2003
2004#[cfg(test)]
2017mod tests {
2018 use std::sync::Arc;
2019
2020 use super::*;
2021 use crate::neighbors::NeighborProvider;
2022 use crate::quant::create_test_quantizer;
2023 use crate::vectors::VectorProvider;
2024 use diskann::{
2025 graph::DiskANNIndex,
2026 graph::{self, search::Knn},
2027 neighbor::BackInserter,
2028 };
2029 use diskann_providers::storage::FileStorageProvider;
2030 use diskann_utils::views::{Init, Matrix};
2031
2032 fn create_quant_index() -> Arc<DiskANNIndex<BfTreeProvider<f32, QuantVectorProvider>>> {
2033 let start_point = Matrix::new(Init(|| 0.0f32), 1, 5);
2034 let dim = 5;
2035 let logical_max_degree = 6;
2036 let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32;
2037 let metric = Metric::L2;
2038
2039 let provider = BfTreeProvider::new(
2040 BfTreeProviderParameters {
2041 max_points: 20,
2042 num_start_points: NonZeroUsize::new(1).unwrap(),
2043 dim,
2044 metric,
2045 max_degree: physical_max_degree,
2046 vector_provider_config: Config::default(),
2047 quant_vector_provider_config: Config::default(),
2048 neighbor_list_provider_config: Config::default(),
2049 graph_params: None,
2050 use_snapshot: false,
2051 },
2052 start_point.as_view(),
2053 create_test_quantizer(5),
2054 )
2055 .unwrap();
2056
2057 let index_config = graph::config::Builder::new_with(
2058 logical_max_degree as usize,
2059 graph::config::MaxDegree::new(physical_max_degree as usize),
2060 10,
2061 metric.into(),
2062 |_| {},
2063 )
2064 .build()
2065 .unwrap();
2066
2067 Arc::new(DiskANNIndex::new(index_config, provider, None))
2068 }
2069
2070 #[tokio::test]
2071 async fn test_quantized_index_search() {
2072 let index = create_quant_index();
2073 let ctx = &DefaultContext;
2074
2075 for i in 0..15 {
2076 let point = vec![i as f32; 5];
2077 index
2078 .insert(&Quantized, ctx, &i, point.as_slice())
2079 .await
2080 .unwrap();
2081 }
2082
2083 let query = vec![3.0; 5];
2084 let params = Knn::new(5, 10, None).unwrap();
2085
2086 let mut neighbors = vec![Neighbor::<u32>::default(); 5];
2087 let res = index
2088 .search(
2089 params,
2090 &Quantized,
2091 &DefaultContext,
2092 query.as_slice(),
2093 &mut BackInserter::new(neighbors.as_mut_slice()),
2094 )
2095 .await
2096 .unwrap();
2097
2098 assert_eq!(
2099 res.result_count, 5,
2100 "there are 15 points and we're asking for 5, we expect 5"
2101 );
2102 assert_eq!(neighbors[0].id, 3);
2103 }
2104
2105 #[tokio::test]
2106 async fn test_quantized_index_multi_insert_search() {
2107 let index = create_quant_index();
2108 let ctx = &DefaultContext;
2109
2110 let data = Matrix::new(
2111 Init({
2112 let mut row = 0usize;
2113 let mut col = 0usize;
2114 move || {
2115 let val = row as f32;
2116 col += 1;
2117 if col == 5 {
2118 col = 0;
2119 row += 1;
2120 }
2121 val
2122 }
2123 }),
2124 15,
2125 5,
2126 );
2127 let ids: Arc<[u32]> = (0u32..15).collect::<Vec<_>>().into();
2128 let batch: Arc<Matrix<f32>> = Arc::new(data);
2129 index
2130 .multi_insert::<Quantized, Matrix<f32>>(Quantized, ctx, batch, ids)
2131 .await
2132 .unwrap();
2133
2134 let query = vec![3.0; 5];
2135 let params = Knn::new(5, 10, None).unwrap();
2136
2137 let mut neighbors = vec![Neighbor::<u32>::default(); 5];
2138 let res = index
2139 .search(
2140 params,
2141 &Quantized,
2142 &DefaultContext,
2143 query.as_slice(),
2144 &mut BackInserter::new(neighbors.as_mut_slice()),
2145 )
2146 .await
2147 .unwrap();
2148
2149 assert_eq!(
2150 res.result_count, 5,
2151 "there are 15 points and we're asking for 5, we expect 5"
2152 );
2153 let neighbor_ids: Vec<u32> = neighbors.iter().map(|n| n.id).collect();
2154 for expected in 1u32..=5 {
2155 assert!(
2156 neighbor_ids.contains(&expected),
2157 "expected id {expected} in results, got {neighbor_ids:?}"
2158 );
2159 }
2160 }
2161
2162 #[tokio::test]
2163 async fn test_quantized_delete_and_search() {
2164 let index = create_quant_index();
2165 let ctx = &DefaultContext;
2166
2167 for i in 0..15 {
2168 let point = vec![i as f32; 5];
2169 index
2170 .insert(&Quantized, ctx, &i, point.as_slice())
2171 .await
2172 .unwrap();
2173 }
2174
2175 index
2176 .inplace_delete(Quantized, ctx, &2u32, 2, graph::InplaceDeleteMethod::OneHop)
2177 .await
2178 .unwrap();
2179 index
2180 .inplace_delete(Quantized, ctx, &4u32, 2, graph::InplaceDeleteMethod::OneHop)
2181 .await
2182 .unwrap();
2183
2184 let query = vec![3.0; 5];
2185 let params = Knn::new(5, 10, None).unwrap();
2186
2187 let mut neighbors = vec![Neighbor::<u32>::default(); 5];
2188 let res = index
2189 .search(
2190 params,
2191 &Quantized,
2192 &DefaultContext,
2193 query.as_slice(),
2194 &mut BackInserter::new(neighbors.as_mut_slice()),
2195 )
2196 .await
2197 .unwrap();
2198
2199 assert_eq!(res.result_count, 5);
2200 let neighbor_ids: Vec<u32> = neighbors.iter().map(|n| n.id).collect();
2201 assert!(!neighbor_ids.contains(&2u32));
2202 assert!(!neighbor_ids.contains(&4u32));
2203 }
2204
2205 fn create_full_precision_index() -> Arc<DiskANNIndex<BfTreeProvider<f32, NoStore>>> {
2206 let start_point = Matrix::new(Init(|| 0.0f32), 1, 5);
2207 let logical_max_degree = 6;
2208 let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32;
2209 let metric = Metric::L2;
2210
2211 let provider = BfTreeProvider::new(
2212 BfTreeProviderParameters {
2213 max_points: 20,
2214 num_start_points: NonZeroUsize::new(1).unwrap(),
2215 dim: 5,
2216 metric,
2217 max_degree: physical_max_degree,
2218 vector_provider_config: Config::default(),
2219 quant_vector_provider_config: Config::default(),
2220 neighbor_list_provider_config: Config::default(),
2221 graph_params: None,
2222 use_snapshot: false,
2223 },
2224 start_point.as_view(),
2225 NoStore,
2226 )
2227 .unwrap();
2228
2229 let index_config = graph::config::Builder::new_with(
2230 logical_max_degree as usize,
2231 graph::config::MaxDegree::new(physical_max_degree as usize),
2232 10,
2233 metric.into(),
2234 |_| {},
2235 )
2236 .build()
2237 .unwrap();
2238
2239 Arc::new(DiskANNIndex::new(index_config, provider, None))
2240 }
2241
2242 #[tokio::test]
2243 async fn test_full_precision_index_search() {
2244 let index = create_full_precision_index();
2245 let ctx = &DefaultContext;
2246
2247 for i in 0u32..15 {
2248 let point = vec![i as f32; 5];
2249 index
2250 .insert(&FullPrecision, ctx, &i, point.as_slice())
2251 .await
2252 .unwrap();
2253 }
2254
2255 let query = vec![3.0; 5];
2256 let params = Knn::new(5, 10, None).unwrap();
2257
2258 let mut neighbors = vec![Neighbor::<u32>::default(); 5];
2259 let res = index
2260 .search(
2261 params,
2262 &FullPrecision,
2263 &DefaultContext,
2264 query.as_slice(),
2265 &mut BackInserter::new(neighbors.as_mut_slice()),
2266 )
2267 .await
2268 .unwrap();
2269
2270 assert_eq!(
2271 res.result_count, 5,
2272 "there are 15 points and we're asking for 5, we expect 5"
2273 );
2274 assert_eq!(neighbors[0].id, 3);
2275 }
2276
2277 #[tokio::test]
2278 async fn test_full_precision_delete_and_search() {
2279 let index = create_full_precision_index();
2280 let ctx = &DefaultContext;
2281
2282 for i in 0u32..15 {
2283 let point = vec![i as f32; 5];
2284 index
2285 .insert(&FullPrecision, ctx, &i, point.as_slice())
2286 .await
2287 .unwrap();
2288 }
2289
2290 index
2291 .inplace_delete(
2292 FullPrecision,
2293 ctx,
2294 &2u32,
2295 2,
2296 graph::InplaceDeleteMethod::OneHop,
2297 )
2298 .await
2299 .unwrap();
2300 index
2301 .inplace_delete(
2302 FullPrecision,
2303 ctx,
2304 &4u32,
2305 2,
2306 graph::InplaceDeleteMethod::OneHop,
2307 )
2308 .await
2309 .unwrap();
2310
2311 let query = vec![3.0; 5];
2312 let params = Knn::new(5, 10, None).unwrap();
2313
2314 let mut neighbors = vec![Neighbor::<u32>::default(); 5];
2315 let res = index
2316 .search(
2317 params,
2318 &FullPrecision,
2319 &DefaultContext,
2320 query.as_slice(),
2321 &mut BackInserter::new(neighbors.as_mut_slice()),
2322 )
2323 .await
2324 .unwrap();
2325
2326 assert_eq!(res.result_count, 5);
2327 let neighbor_ids: Vec<u32> = neighbors.iter().map(|n| n.id).collect();
2328 assert!(!neighbor_ids.contains(&2u32));
2329 assert!(!neighbor_ids.contains(&4u32));
2330 }
2331
2332 #[tokio::test]
2333 async fn test_data_provider_and_delete_interface() {
2334 let ctx = &DefaultContext;
2335 let num_start_points = 2;
2336 let dim = 5;
2337 let start_points = Matrix::try_from(
2338 vec![0.0f32; dim]
2339 .into_iter()
2340 .chain(vec![0.5f32; dim])
2341 .collect::<Box<[_]>>(),
2342 num_start_points,
2343 dim,
2344 )
2345 .unwrap();
2346
2347 let provider = BfTreeProvider::new(
2348 BfTreeProviderParameters {
2349 max_points: 10,
2350 num_start_points: NonZeroUsize::new(num_start_points).unwrap(),
2351 dim,
2352 metric: Metric::L2,
2353 max_degree: 64,
2354 vector_provider_config: Config::default(),
2355 quant_vector_provider_config: Config::default(),
2356 neighbor_list_provider_config: Config::default(),
2357 graph_params: None,
2358 use_snapshot: false,
2359 },
2360 start_points.as_view(),
2361 NoStore,
2362 )
2363 .unwrap();
2364
2365 assert_eq!((&provider).into_iter(), 0..(10 + 2));
2368
2369 let iter = provider.iter();
2370
2371 for i in iter.clone() {
2373 let vector: Vec<f32> = (0..5).map(|j| (i * 5 + j) as f32).collect();
2374 provider.set_element(ctx, &i, &vector).await.unwrap();
2375 }
2376
2377 for i in iter.clone() {
2378 assert_eq!(provider.to_external_id(ctx, i).unwrap(), i);
2379 assert_eq!(provider.to_internal_id(ctx, &i).unwrap(), i);
2380 assert_eq!(
2381 provider.status_by_internal_id(ctx, i).await.unwrap(),
2382 ElementStatus::Valid
2383 );
2384 assert_eq!(
2385 provider.status_by_external_id(ctx, &i).await.unwrap(),
2386 ElementStatus::Valid
2387 );
2388
2389 provider.delete(ctx, &i).await.unwrap();
2392 assert_eq!(
2393 provider.status_by_internal_id(ctx, i).await.unwrap(),
2394 ElementStatus::Deleted
2395 );
2396 assert_eq!(
2397 provider.status_by_external_id(ctx, &i).await.unwrap(),
2398 ElementStatus::Deleted
2399 );
2400 }
2401
2402 for i in iter.clone() {
2405 provider.release(ctx, i).await.unwrap();
2406 assert_eq!(
2407 provider.status_by_internal_id(ctx, i).await.unwrap(),
2408 ElementStatus::Deleted
2409 );
2410 }
2411
2412 assert!(provider
2415 .set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0])
2416 .await
2417 .is_err());
2418 }
2419
2420 #[tokio::test]
2426 async fn test_empty_neighbor_list() {
2427 let num_points = 100u32;
2428 let ctx = &DefaultContext;
2429
2430 let num_start_points = 2;
2431 let dim = 3;
2432 let start_points = Matrix::new(Init(|| 0.0f32), num_start_points, dim);
2433
2434 let provider = BfTreeProvider::<f32, _>::new(
2435 BfTreeProviderParameters {
2436 max_points: num_points as usize,
2437 num_start_points: NonZeroUsize::new(num_start_points).unwrap(),
2438 dim,
2439 metric: Metric::L2,
2440 max_degree: 64,
2441 vector_provider_config: Config::default(),
2442 quant_vector_provider_config: Config::default(),
2443 neighbor_list_provider_config: Config::default(),
2444 graph_params: None,
2445 use_snapshot: false,
2446 },
2447 start_points.as_view(),
2448 NoStore,
2449 )
2450 .unwrap();
2451
2452 let mut scratch = provider.neighbor_provider.scratch(&provider.locks);
2453
2454 for i in 0..num_points {
2458 let vector = vec![i as f32, (i + 1) as f32, (i + 2) as f32];
2459 provider.set_element(ctx, &i, &vector).await.unwrap();
2460
2461 let mut out = AdjacencyList::new();
2463 provider
2464 .neighbor_provider
2465 .get_neighbors(i, &mut out)
2466 .unwrap();
2467 assert!(out.is_empty());
2468
2469 scratch.write_neighbors(i, &[]).unwrap();
2471 provider
2472 .neighbor_provider
2473 .get_neighbors(i, &mut out)
2474 .unwrap();
2475
2476 assert!(out.is_empty());
2477 }
2478
2479 for i in 0..num_points {
2483 let mut out = AdjacencyList::new();
2484 let neighbors = vec![10, 20, 30, 40, 50, 60, 70, 80, 90, 100];
2485 scratch.write_neighbors(i, &neighbors).unwrap();
2486
2487 provider
2488 .neighbor_provider
2489 .get_neighbors(i, &mut out)
2490 .unwrap();
2491
2492 assert_eq!(&*out, &[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]); scratch.write_neighbors(i, &[]).unwrap();
2495 provider
2496 .neighbor_provider
2497 .get_neighbors(i, &mut out)
2498 .unwrap();
2499
2500 assert!(out.is_empty());
2501 }
2502
2503 let mut out = AdjacencyList::from_iter_untrusted([10, 20, 30, 40, 50, 60, 70, 80, 90, 100]); assert!(provider
2509 .neighbor_provider
2510 .get_neighbors(200, &mut out)
2511 .is_err());
2512 assert!(out.is_empty());
2513 }
2514
2515 use tempfile::tempdir;
2520
2521 #[tokio::test]
2523 async fn test_bf_tree_provider_save_load_no_quant() {
2524 let num_points = 50usize;
2525 let dim = 4usize;
2526 let max_degree = 32u32;
2527 let num_start_points = NonZeroUsize::new(2).unwrap();
2528 let ctx = &DefaultContext;
2529
2530 let temp_dir = tempdir().unwrap();
2532 let temp_path = temp_dir.path();
2533
2534 let prefix = temp_path
2535 .join("test_bf_tree_provider")
2536 .to_string_lossy()
2537 .to_string();
2538 let vector_path = BfTreePaths::vectors_bftree(&prefix);
2540 let neighbor_path = BfTreePaths::neighbors_bftree(&prefix);
2541
2542 let bytes_vector = 1024 * 1024;
2543 let mut vector_config = Config::new(&vector_path, bytes_vector);
2544 vector_config.leaf_page_size(8192);
2545 vector_config.cb_max_record_size(1024);
2546 vector_config.storage_backend(bf_tree::StorageBackend::Std);
2547 vector_config.use_snapshot(true);
2548
2549 let bytes_neighbor = 1024 * 1024;
2550 let mut neighbor_config = Config::new(&neighbor_path, bytes_neighbor);
2551 neighbor_config.storage_backend(bf_tree::StorageBackend::Std);
2552 neighbor_config.use_snapshot(true);
2553
2554 let params = BfTreeProviderParameters {
2556 max_points: num_points,
2557 num_start_points,
2558 dim,
2559 metric: Metric::L2,
2560 max_degree,
2561 vector_provider_config: vector_config.clone(),
2562 quant_vector_provider_config: Config::default(),
2563 neighbor_list_provider_config: neighbor_config.clone(),
2564 graph_params: None,
2565 use_snapshot: true,
2566 };
2567
2568 let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
2569
2570 let provider =
2572 BfTreeProvider::<f32, NoStore>::new(params.clone(), start_points.as_view(), NoStore)
2573 .unwrap();
2574
2575 for i in 0..num_points {
2577 let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
2578 provider
2579 .set_element(ctx, &(i as u32), &vector)
2580 .await
2581 .unwrap();
2582 }
2583
2584 let mut scratch = provider.neighbor_provider.scratch(&provider.locks);
2586 for i in 0..num_points as u32 {
2587 let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
2588 .map(|j| (i + j) % num_points as u32)
2589 .collect();
2590 scratch.write_neighbors(i, &neighbors).unwrap();
2591 }
2592
2593 assert_eq!(vector_config.get_leaf_page_size(), 8192);
2594 assert_eq!(vector_config.get_cb_max_record_size(), 1024);
2595
2596 let storage = FileStorageProvider;
2597
2598 let save_dir = tempdir().unwrap();
2600 let save_prefix = save_dir
2601 .path()
2602 .join("saved_bf_tree_provider")
2603 .to_string_lossy()
2604 .to_string();
2605 provider.save_with(&storage, &save_prefix).await.unwrap();
2606
2607 let loaded_provider = BfTreeProvider::<f32, NoStore>::load_with(&storage, &save_prefix)
2609 .await
2610 .unwrap();
2611
2612 for i in 0..num_points as u32 {
2614 let original = provider.full_vectors.get_vector_sync(i as usize).unwrap();
2615 let loaded = loaded_provider
2616 .full_vectors
2617 .get_vector_sync(i as usize)
2618 .unwrap();
2619 assert_eq!(original, loaded, "Vector mismatch at index {}", i);
2620 }
2621
2622 for i in 0..num_points as u32 {
2624 let mut original_list = AdjacencyList::new();
2625 let mut loaded_list = AdjacencyList::new();
2626
2627 provider
2628 .neighbor_provider
2629 .get_neighbors(i, &mut original_list)
2630 .unwrap();
2631 loaded_provider
2632 .neighbor_provider
2633 .get_neighbors(i, &mut loaded_list)
2634 .unwrap();
2635
2636 assert_eq!(
2637 &*original_list, &*loaded_list,
2638 "Neighbor list mismatch at index {}",
2639 i
2640 );
2641 }
2642
2643 }
2645
2646 #[tokio::test]
2648 async fn test_bf_tree_provider_save_load_quant() {
2649 let num_points = 50usize;
2650 let dim = 8usize;
2651 let max_degree = 32u32;
2652 let num_start_points = NonZeroUsize::new(2).unwrap();
2653 let ctx = &DefaultContext;
2654
2655 let temp_dir = tempdir().unwrap();
2657 let temp_path = temp_dir.path();
2658
2659 let prefix = temp_path
2660 .join("test_bf_tree_provider_quant")
2661 .to_string_lossy()
2662 .to_string();
2663 let vector_path = BfTreePaths::vectors_bftree(&prefix);
2665 let neighbor_path = BfTreePaths::neighbors_bftree(&prefix);
2666 let quant_path = BfTreePaths::quant_bftree(&prefix);
2667
2668 let bytes_vector = 1024 * 1024;
2669 let mut vector_config = Config::new(&vector_path, bytes_vector);
2670 vector_config.storage_backend(bf_tree::StorageBackend::Std);
2671 vector_config.use_snapshot(true);
2672
2673 let bytes_neighbor = 1024 * 1024;
2674 let mut neighbor_config = Config::new(&neighbor_path, bytes_neighbor);
2675 neighbor_config.storage_backend(bf_tree::StorageBackend::Std);
2676 neighbor_config.use_snapshot(true);
2677
2678 let bytes_quant = 1024 * 1024;
2679 let mut quant_config = Config::new(&quant_path, bytes_quant);
2680 quant_config.storage_backend(bf_tree::StorageBackend::Std);
2681 quant_config.use_snapshot(true);
2682
2683 let quantizer = create_test_quantizer(dim);
2685
2686 let params = BfTreeProviderParameters {
2688 max_points: num_points,
2689 num_start_points,
2690 dim,
2691 metric: Metric::L2,
2692 max_degree,
2693 vector_provider_config: vector_config.clone(),
2694 quant_vector_provider_config: quant_config.clone(),
2695 neighbor_list_provider_config: neighbor_config.clone(),
2696 graph_params: None,
2697 use_snapshot: true,
2698 };
2699
2700 let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
2701 let provider = BfTreeProvider::<f32, QuantVectorProvider>::new(
2703 params.clone(),
2704 start_points.as_view(),
2705 quantizer,
2706 )
2707 .unwrap();
2708
2709 for i in 0..num_points {
2711 let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
2712 provider
2713 .set_element(ctx, &(i as u32), &vector)
2714 .await
2715 .unwrap();
2716 }
2717
2718 let mut scratch = provider.neighbor_provider.scratch(&provider.locks);
2720 for i in 0..num_points as u32 {
2721 let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
2722 .map(|j| (i + j) % num_points as u32)
2723 .collect();
2724 scratch.write_neighbors(i, &neighbors).unwrap();
2725 }
2726
2727 let storage = FileStorageProvider;
2728
2729 let save_dir = tempdir().unwrap();
2731 let save_prefix = save_dir
2732 .path()
2733 .join("saved_bf_tree_provider_quant")
2734 .to_string_lossy()
2735 .to_string();
2736 provider.save_with(&storage, &save_prefix).await.unwrap();
2737
2738 let loaded_provider =
2740 BfTreeProvider::<f32, QuantVectorProvider>::load_with(&storage, &save_prefix)
2741 .await
2742 .unwrap();
2743
2744 assert_eq!(
2746 provider.quant_vectors.quantizer.full_dim(),
2747 loaded_provider.quant_vectors.quantizer.full_dim(),
2748 "Quantizer full_dim mismatch"
2749 );
2750 assert_eq!(
2751 provider.quant_vectors.quantizer.bytes(),
2752 loaded_provider.quant_vectors.quantizer.bytes(),
2753 "Quantizer bytes mismatch"
2754 );
2755 assert_eq!(
2756 provider.quant_vectors.quantizer.nbits(),
2757 loaded_provider.quant_vectors.quantizer.nbits(),
2758 "Quantizer nbits mismatch"
2759 );
2760
2761 for i in 0..num_points as u32 {
2763 let original = provider.full_vectors.get_vector_sync(i as usize).unwrap();
2764 let loaded = loaded_provider
2765 .full_vectors
2766 .get_vector_sync(i as usize)
2767 .unwrap();
2768 assert_eq!(original, loaded, "Vector mismatch at index {}", i);
2769 }
2770
2771 for i in 0..num_points as u32 {
2773 let original = provider.quant_vectors.get_vector_sync(i as usize).unwrap();
2774 let loaded = loaded_provider
2775 .quant_vectors
2776 .get_vector_sync(i as usize)
2777 .unwrap();
2778 assert_eq!(original, loaded, "Quant vector mismatch at index {}", i);
2779 }
2780
2781 for i in 0..num_points as u32 {
2783 let mut original_list = AdjacencyList::new();
2784 let mut loaded_list = AdjacencyList::new();
2785
2786 provider
2787 .neighbor_provider
2788 .get_neighbors(i, &mut original_list)
2789 .unwrap();
2790 loaded_provider
2791 .neighbor_provider
2792 .get_neighbors(i, &mut loaded_list)
2793 .unwrap();
2794
2795 assert_eq!(
2796 &*original_list, &*loaded_list,
2797 "Neighbor list mismatch at index {}",
2798 i
2799 );
2800 }
2801
2802 }
2804
2805 #[tokio::test]
2807 async fn test_bf_tree_provider_memory_save_load_no_quant() {
2808 let num_points = 20usize;
2809 let dim = 4usize;
2810 let max_degree = 16u32;
2811 let num_start_points = NonZeroUsize::new(1).unwrap();
2812 let ctx = &DefaultContext;
2813
2814 let mut vector_config = Config::default();
2815 vector_config.use_snapshot(true);
2816 let mut neighbor_config = Config::default();
2817 neighbor_config.use_snapshot(true);
2818
2819 let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
2820 let provider = BfTreeProvider::<f32, NoStore>::new(
2822 BfTreeProviderParameters {
2823 max_points: num_points,
2824 num_start_points,
2825 dim,
2826 metric: Metric::L2,
2827 max_degree,
2828 vector_provider_config: vector_config,
2829 quant_vector_provider_config: Config::default(),
2830 neighbor_list_provider_config: neighbor_config,
2831 graph_params: None,
2832 use_snapshot: true,
2833 },
2834 start_points.as_view(),
2835 NoStore,
2836 )
2837 .unwrap();
2838
2839 for i in 0..num_points {
2841 let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
2842 provider
2843 .set_element(ctx, &(i as u32), &vector)
2844 .await
2845 .unwrap();
2846 }
2847 let mut scratch = provider.neighbor_provider.scratch(&provider.locks);
2848 for i in 0..num_points as u32 {
2849 let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
2850 .map(|j| (i + j) % num_points as u32)
2851 .collect();
2852 scratch.write_neighbors(i, &neighbors).unwrap();
2853 }
2854
2855 provider.delete(ctx, &3u32).await.unwrap();
2857 provider.delete(ctx, &7u32).await.unwrap();
2858
2859 let save_dir = tempdir().unwrap();
2861 let save_prefix = save_dir
2862 .path()
2863 .join("mem_no_quant")
2864 .to_string_lossy()
2865 .to_string();
2866 let storage = FileStorageProvider;
2867 provider.save_with(&storage, &save_prefix).await.unwrap();
2868
2869 let loaded = BfTreeProvider::<f32, NoStore>::load_with(&storage, &save_prefix)
2871 .await
2872 .unwrap();
2873
2874 for i in 0..num_points as u32 {
2876 if i == 3 || i == 7 {
2877 continue;
2878 }
2879 assert_eq!(
2880 provider.full_vectors.get_vector_sync(i as usize).unwrap(),
2881 loaded.full_vectors.get_vector_sync(i as usize).unwrap(),
2882 "Vector mismatch at {}",
2883 i
2884 );
2885 }
2886
2887 for i in 0..num_points as u32 {
2889 let mut orig = AdjacencyList::new();
2890 let mut load = AdjacencyList::new();
2891 provider
2892 .neighbor_provider
2893 .get_neighbors(i, &mut orig)
2894 .unwrap();
2895 loaded
2896 .neighbor_provider
2897 .get_neighbors(i, &mut load)
2898 .unwrap();
2899 assert_eq!(&*orig, &*load, "Neighbor mismatch at {}", i);
2900 }
2901
2902 assert_eq!(
2904 loaded.status_by_internal_id(ctx, 3).await.unwrap(),
2905 ElementStatus::Deleted
2906 );
2907 assert_eq!(
2908 loaded.status_by_internal_id(ctx, 7).await.unwrap(),
2909 ElementStatus::Deleted
2910 );
2911 assert_eq!(
2912 loaded.status_by_internal_id(ctx, 0).await.unwrap(),
2913 ElementStatus::Valid
2914 );
2915 }
2916
2917 #[tokio::test]
2919 async fn test_bf_tree_provider_memory_save_load_quant() {
2920 let num_points = 20usize;
2921 let dim = 8usize;
2922 let max_degree = 16u32;
2923 let num_start_points = NonZeroUsize::new(1).unwrap();
2924 let ctx = &DefaultContext;
2925
2926 let quantizer = create_test_quantizer(dim);
2927 let mut vector_config = Config::default();
2928 vector_config.use_snapshot(true);
2929 let mut neighbor_config = Config::default();
2930 neighbor_config.use_snapshot(true);
2931 let mut quant_config = Config::default();
2932 quant_config.use_snapshot(true);
2933
2934 let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
2935 let provider = BfTreeProvider::<f32, QuantVectorProvider>::new(
2936 BfTreeProviderParameters {
2937 max_points: num_points,
2938 num_start_points,
2939 dim,
2940 metric: Metric::L2,
2941 max_degree,
2942 vector_provider_config: vector_config,
2943 quant_vector_provider_config: quant_config,
2944 neighbor_list_provider_config: neighbor_config,
2945 graph_params: None,
2946 use_snapshot: true,
2947 },
2948 start_points.as_view(),
2949 quantizer,
2950 )
2951 .unwrap();
2952
2953 for i in 0..num_points {
2955 let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
2956 provider
2957 .set_element(ctx, &(i as u32), &vector)
2958 .await
2959 .unwrap();
2960 }
2961 let mut scratch = provider.neighbor_provider.scratch(&provider.locks);
2962 for i in 0..num_points as u32 {
2963 let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
2964 .map(|j| (i + j) % num_points as u32)
2965 .collect();
2966 scratch.write_neighbors(i, &neighbors).unwrap();
2967 }
2968
2969 provider.delete(ctx, &2u32).await.unwrap();
2970
2971 let save_dir = tempdir().unwrap();
2973 let save_prefix = save_dir
2974 .path()
2975 .join("mem_quant")
2976 .to_string_lossy()
2977 .to_string();
2978 let storage = FileStorageProvider;
2979 provider.save_with(&storage, &save_prefix).await.unwrap();
2980
2981 let loaded = BfTreeProvider::<f32, QuantVectorProvider>::load_with(&storage, &save_prefix)
2983 .await
2984 .unwrap();
2985
2986 for i in 0..num_points as u32 {
2988 if i == 2 {
2989 continue;
2990 }
2991 assert_eq!(
2992 provider.full_vectors.get_vector_sync(i as usize).unwrap(),
2993 loaded.full_vectors.get_vector_sync(i as usize).unwrap(),
2994 "Vector mismatch at {}",
2995 i
2996 );
2997 }
2998
2999 for i in 0..num_points as u32 {
3001 if i == 2 {
3002 continue;
3003 }
3004 assert_eq!(
3005 provider.quant_vectors.get_vector_sync(i as usize).unwrap(),
3006 loaded.quant_vectors.get_vector_sync(i as usize).unwrap(),
3007 "Quant vector mismatch at {}",
3008 i
3009 );
3010 }
3011
3012 for i in 0..num_points as u32 {
3014 if i == 2 {
3015 continue;
3016 }
3017 let mut orig = AdjacencyList::new();
3018 let mut load = AdjacencyList::new();
3019 provider
3020 .neighbor_provider
3021 .get_neighbors(i, &mut orig)
3022 .unwrap();
3023 loaded
3024 .neighbor_provider
3025 .get_neighbors(i, &mut load)
3026 .unwrap();
3027 assert_eq!(&*orig, &*load, "Neighbor mismatch at {}", i);
3028 }
3029
3030 assert_eq!(
3032 loaded.status_by_internal_id(ctx, 2).await.unwrap(),
3033 ElementStatus::Deleted
3034 );
3035 assert_eq!(
3036 loaded.status_by_internal_id(ctx, 0).await.unwrap(),
3037 ElementStatus::Valid
3038 );
3039 }
3040
3041 #[test]
3042 fn test_validate_rejects_undersized_vector_config() {
3043 let result = VectorProvider::<f32>::new_with_config(
3045 100,
3046 1536,
3047 1,
3048 Config::default(), );
3050 let err = result.err().expect("should fail").to_string();
3051 assert!(
3052 err.contains("vector_provider"),
3053 "should name the failing config; got: {err}"
3054 );
3055 assert!(
3056 err.contains("6152"),
3057 "should state the required size; got: {err}"
3058 );
3059 }
3060
3061 #[test]
3062 fn test_validate_rejects_undersized_neighbor_config() {
3063 let mut neighbor_config = Config::default();
3065 neighbor_config.cb_max_record_size(1952);
3066
3067 let result = NeighborProvider::<u32>::new_with_config(500, neighbor_config);
3068 let err = result.err().expect("should fail").to_string();
3069 assert!(
3070 err.contains("neighbor_provider"),
3071 "should name the failing config; got: {err}"
3072 );
3073 }
3074
3075 #[test]
3076 fn test_validate_accepts_valid_config() {
3077 if let Err(e) = VectorProvider::<f32>::new_with_config(100, 128, 1, Config::default()) {
3079 panic!("VectorProvider should succeed: {e}");
3080 }
3081 if let Err(e) = NeighborProvider::<u32>::new_with_config(64, Config::default()) {
3083 panic!("NeighborProvider should succeed: {e}");
3084 }
3085 }
3086}