Skip to main content

diskann_bftree/
provider.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use 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
54/////////////////////
55// BfTreeProvider //
56/////////////////////
57
58/// An Bf-tree based implementation of a [`DataProvider`] built around the idea of having up to
59/// two layers of vector stores: a full-precision store and an optional quantized store.
60/// This provider uses the identity mapping between external and internal vector IDs.
61///
62/// # Type Parameters:
63///
64/// * `T`: The primitive element type of the full-precision vector. This is typically some
65///   type like `f32` or `half::f16`.
66///
67/// * `Q`: The full type of the quant vector store. This is not constrained by a trait and
68///   rather relies on implementation for several concrete types, including:
69///
70///   - [`QuantVectorProvider`]: A Bf-Tree based spherical quantized vector store.
71///   - [`NoStore`]: Disable quantization altogether. Note that this disables all
72///     methods reached through quantization based [`Accessor`]s at compile-time.
73///
74/// # Indexing Strategies
75///
76/// * [`FullPrecision`]: Only retrieves data from the full-precision portion of the index.
77///   No quantized vectors are used. During search, start points are filtered from the
78///   final results.
79///
80/// * [`Quantized`]: Performs all operations (search, pruning, insert) entirely in the
81///   quantized space using spherical distance functions. Post-processing copies candidate
82///   IDs forward without reranking. Fastest option — full-precision vectors are not
83///   touched at query time.
84///
85/// # Examples
86///
87/// The following code demonstrates how to instantiate and use the `BfTreeProvider` in
88/// a number of different scenarios.
89///
90/// ## Full-Precision Only - No Deletes
91///
92/// This example demonstrates how to create a `BfTreeProvider` that only supports
93/// full-precision vectors.
94/// ```
95/// use diskann_bftree::provider::{
96///     BfTreeProvider, BfTreeProviderParameters
97/// };
98/// use diskann_bftree::NoStore;
99/// use diskann_vector::distance::Metric;
100/// use diskann_utils::views::{Init, Matrix};
101/// use bf_tree::Config;
102/// use std::num::NonZeroUsize;
103///
104/// let parameters = BfTreeProviderParameters {
105///     max_points: 5,
106///     num_start_points: NonZeroUsize::new(1).unwrap(),
107///     dim: 4,
108///     metric: Metric::L2,
109///     max_degree: 32,
110///     vector_provider_config: Config::default(),
111///     quant_vector_provider_config: Config::default(),
112///     neighbor_list_provider_config: Config::default(),
113///     graph_params: None,
114///     use_snapshot: false,
115/// };
116///
117/// // Create a table that supports 5 points and 1 start point.
118/// let start_points = Matrix::new(Init(|| 0.0f32), 1, 4);
119/// let provider = BfTreeProvider::<f32, _>::new(
120///     parameters,
121///     start_points.as_view(),
122///     NoStore,
123/// );
124/// ```
125///
126/// ## Full-Precision and Spherical Quantization
127///
128/// To create a two-level provider with a spherical quantization-based quant vector store,
129/// a `Poly<dyn Quantizer>` can be supplied for the `quant_precursor` argument, as this
130/// implements the [`CreateQuantProvider`] trait.
131/// ```
132/// use diskann_quantization::{
133///     alloc::{GlobalAllocator, Poly, poly},
134///     algorithms::TransformKind,
135///     spherical::{iface, SphericalQuantizer, SupportedMetric, PreScale},
136/// };
137/// use diskann_utils::views::{Init, Matrix};
138/// use diskann_bftree::provider::{
139///     BfTreeProvider, BfTreeProviderParameters
140/// };
141/// use diskann_vector::distance::Metric;
142/// use bf_tree::Config;
143/// use std::num::NonZeroUsize;
144/// use rand::rngs::StdRng;
145/// use rand::SeedableRng;
146///
147/// let dim = 4;
148/// let data = Matrix::new(Init(|| 1.0f32), 4, dim);
149/// let mut rng = StdRng::seed_from_u64(42);
150/// let sq = SphericalQuantizer::train(
151///     data.as_view(), TransformKind::Null,
152///     SupportedMetric::SquaredL2, PreScale::None,
153///     &mut rng, GlobalAllocator,
154/// ).unwrap();
155/// let imp = iface::Impl::<1>::new(sq).unwrap();
156/// let poly = Poly::new(imp, GlobalAllocator).unwrap();
157/// let quantizer: Poly<dyn iface::Quantizer> = poly!(iface::Quantizer, poly);
158///
159/// let parameters = BfTreeProviderParameters {
160///     max_points: 5,
161///     num_start_points: NonZeroUsize::new(1).unwrap(),
162///     dim: 4,
163///     metric: Metric::L2,
164///     max_degree: 32,
165///     vector_provider_config: Config::default(),
166///     quant_vector_provider_config: Config::default(),
167///     neighbor_list_provider_config: Config::default(),
168///     graph_params: None,
169///     use_snapshot: false,
170/// };
171///
172/// // Create a table that supports 5 points and 1 start point.
173/// let start_points = Matrix::new(Init(|| 0.0f32), 1, 4);
174/// let provider = BfTreeProvider::<f32, _>::new(
175///     parameters,
176///     start_points.as_view(),
177///     quantizer,
178/// );
179/// ```
180pub struct BfTreeProvider<T, Q = QuantVectorProvider>
181where
182    T: VectorRepr,
183{
184    // The quant vector store. If `Q == NoStore`, the quantized operations are disabled.
185    //
186    pub(super) quant_vectors: Q,
187
188    // The full vector store.
189    //
190    pub(super) full_vectors: VectorProvider<T>,
191
192    // Provider that holds the graph structure as neighbors of vectors.
193    //
194    pub(crate) neighbor_provider: NeighborProvider<u32>,
195
196    // The metric to use for distances
197    //
198    pub(super) metric: Metric,
199
200    // Graph configuration parameters for persistence
201    //
202    pub(crate) graph_params: Option<GraphParams>,
203
204    // Whether CPR snapshot support is enabled for this provider's trees.
205    pub(crate) use_snapshot: bool,
206    // Striped locks for per-vertex synchronization across all stores.
207    //
208    // # Locking protocol
209    //
210    // Each mutating operation (`set_element`, `delete`, `write_neighbors`,
211    // `set_neighbors`, `append_vector`) acquires exactly **one** stripe lock
212    // for the target vertex ID, performs its work, and drops the guard at
213    // scope exit. No operation acquires a second lock while holding the first,
214    // so **deadlock is structurally impossible** — there is no ordering
215    // concern.
216    //
217    // Reads (`get_neighbors`, `get_vector`) do **not** acquire the stripe
218    // lock. They rely on bf-tree's internal thread safety for reads. A reader
219    // may therefore observe a partially-updated neighbor list (e.g.,
220    // mid-append), but bf-tree guarantees the read will not crash or return
221    // corrupt memory. The DiskANN search algorithm tolerates stale neighbor
222    // reads — it is approximate by design, and momentary inconsistency
223    // affects quality briefly rather than correctness.
224    pub(crate) locks: StripedLocks,
225}
226
227#[derive(Debug, Clone)]
228pub struct BfTreeProviderParameters {
229    // The maximum number of valid points that provider can hold.
230    pub max_points: usize,
231
232    // The number of start points (frozen points) for graph search entry.
233    pub num_start_points: NonZeroUsize,
234
235    // The dimension of the full-precision vectors.
236    pub dim: usize,
237
238    // The metric to use for distance computations
239    pub metric: Metric,
240
241    // The physical maximum degree (maximum neighbor list capacity per vertex).
242    // Callers are responsible for applying any slack factor externally before
243    // passing this value.
244    pub max_degree: u32,
245
246    // bf-tree config for vector provider.
247    // bf-tree requires a minimum circular buffer size relative to leaf_page_size:
248    //   - Cache-only mode: cb_size_byte >= 4 * leaf_page_size
249    //   - Non cache-only mode: cb_size_byte >= 2 * leaf_page_size
250    // The default leaf_page_size is 4096 and default cb_size_byte is 32MB.
251    pub vector_provider_config: Config,
252
253    // bf-tree config for quant vector provider.
254    // Same minimum circular buffer constraint as above.
255    pub quant_vector_provider_config: Config,
256
257    // bf-tree config for neighbor list provider.
258    // Same minimum circular buffer constraint as above.
259    pub neighbor_list_provider_config: Config,
260
261    // Optional graph configuration parameters for persistence
262    pub graph_params: Option<GraphParams>,
263
264    // Whether to enable CPR snapshot support on the underlying bf-trees.
265    pub use_snapshot: bool,
266}
267
268impl<T, Q> BfTreeProvider<T, Q>
269where
270    T: VectorRepr,
271{
272    /// Construct a new data provider from empty. Callers of this are required to manually set start
273    /// points before performing search tasks.
274    ///
275    /// This constructor for `BfTreeProvider` should be used when building and constructing from
276    /// scratch.
277    ///
278    /// # Arguments
279    /// * `params`: An instance of [`BfTreeProviderParameters`] collecting shared
280    ///   configuration information.
281    /// * `quant_precursor`: A precursor type for the quantizer layer.
282    ///
283    /// # Type Constraints
284    /// * `Self: StartPoint<T>` - The provider must implement the `StartPoint` trait.
285    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        // Force all configs to match the provider-level use_snapshot setting so the
291        // underlying bf-trees are always consistent with our snapshot guard.
292        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    /// Construct a new data provider with start points initialized.
322    ///
323    /// This is the primary constructor for `BfTreeProvider` where start points are known from the
324    /// beginning. It creates the provider and sets the start points in one operation.
325    ///
326    /// # Arguments
327    /// * `params`: An instance of [`BfTreeProviderParameters`] collecting shared
328    ///   configuration information.
329    /// * `start_points`: A matrix view containing the start point vectors. The number
330    ///   of rows must match `params.num_start_points.get()`.
331    /// * `quant_precursor`: A precursor type for the quantizer layer.
332    ///
333    /// # Type Constraints
334    /// * `Self: StartPoint<T>` - The provider must implement the `StartPoint` trait.
335    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        // Early validation before allocating resources
345        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            // Initialize all neighborhoods to be empty lists.
357            // This is a temporary solution to the problem of trying to access
358            // an uninitialized neighbor list in functions `consolidate_deletes` and
359            // `consolidate_simple` and getting an error. This is a stop-gap solution
360            // until BF-tree API is improved to handle `exists` queries.
361            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    // /// Return a predicate that can be applied to `Iter::filter` to remove start points
371    // /// from an iterator of neighbors.
372    // ///
373    // /// This is used during post-processing
374    // ///
375    // pub(crate) fn is_not_start_point(&self) -> impl Fn(&Neighbor<u32>) -> bool {
376    //     let range = self.full_vectors.start_point_range();
377    //     move |neighbor| !range.contains(&neighbor.id.into_usize())
378    // }
379
380    /// Return a vector of starting points.
381    pub fn starting_points(&self) -> ANNResult<Vec<u32>> {
382        self.full_vectors.starting_points()
383    }
384
385    /// An iterator over all ids including start points (even if they are deleted).
386    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    /// Return the maximum number of points (excluding frozen/start points)
395    pub fn max_points(&self) -> usize {
396        self.full_vectors.max_vectors
397    }
398
399    /// Return the vector dimension
400    pub fn dim(&self) -> usize {
401        self.full_vectors.dim()
402    }
403
404    /// Return the distance metric
405    pub fn metric(&self) -> Metric {
406        self.metric
407    }
408
409    /// Return the maximum degree from the neighbor provider
410    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    /// Return the number of vector reads for full-precision and quant-vectors respectively
420    ///
421    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    /// Return the number of vector reads for full-precision and quant-vectors respectively
434    ///
435    pub fn counts_for_get_vector(&self) -> (usize, usize) {
436        (self.full_vectors.num_get_calls.get(), 0)
437    }
438}
439
440/// Trait for types that may need cleanup during hard-delete.
441///
442/// `QuantVectorProvider` implements this to delete its quantized embedding.
443/// `NoStore` implements this as a no-op.
444pub(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
458/// Hard-delete implementation for `BfTreeProvider`.
459///
460/// This provider performs **hard deletes**: vector data is immediately and irrecoverably erased
461/// from the underlying bf-tree storage when [`Delete::delete`] is called. When a quantized
462/// store is present, both the full-precision and quantized entries are removed.
463///
464/// This has an important consequence for inplace delete: the
465/// [`InplaceDeleteMethod::VisitedAndTopK`] strategy is **incompatible** with this provider
466/// because it requires reading the deleted vector's data (via
467/// [`InplaceDeleteStrategy::get_delete_element`]) *after* the delete has already been committed.
468/// Use [`InplaceDeleteMethod::OneHop`] or [`InplaceDeleteMethod::TwoHopAndOneHop`] instead,
469/// as these strategies only require neighbor topology (which remains accessible).
470impl<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        // Only delete vector data here. Neighbor adjacency cleanup (zeroing the
491        // deleted vertex's edge list and patching neighbors-of-neighbors) is
492        // handled by `DiskANNIndex::inplace_delete` → `drop_adj_list`.
493        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
523/// Allow `&BfTreeProvider` to implement `IntoIter`
524///
525impl<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
536/// A helper trait to select the quant vector store.
537///
538/// This is also implemented for [`NoStore`], which explicitly disables deletion
539/// related functionality
540///
541pub trait CreateQuantProvider {
542    // The type of the created quant provider.
543    //
544    type Target;
545
546    // Create a quant provider capable of tracking `max_points` with and additional
547    // `frozen_points` at the end.
548    //
549    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
559/// Allow a `FixedChunkPQTable` to be promoted to full quant vector store.
560///
561impl 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
578///////////////////
579// Data Provider //
580///////////////////
581
582impl<T, Q> DataProvider for BfTreeProvider<T, Q>
583where
584    T: VectorRepr,
585    Q: AsyncFriendly,
586{
587    type Context = DefaultContext;
588
589    // The `BfTreeProvider` uses the identity map for IDs.
590    //
591    type InternalId = u32;
592
593    // The `BfTreeProvider` uses the identity map for IDs.
594    //
595    type ExternalId = u32;
596
597    // Use a general error type for now.
598    //
599    type Error = ANNError;
600
601    // No insert-ID recovery.
602    type Guard = NoopGuard<u32>;
603
604    // Translate an external id to its corresponding internal id.
605    //
606    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    // Translate an internal id its corresponding external id.
615    //
616    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
633////////////////
634// SetElement //
635////////////////
636
637/// Assign to both the full-precision and quant vector stores
638///
639impl<T> SetElement<&[T]> for BfTreeProvider<T, QuantVectorProvider>
640where
641    T: VectorRepr,
642{
643    type SetError = ANNError;
644
645    /// Store the provided element in both the full-precision and quant vector stores.
646    ///
647    /// Full-precision is written first as the authoritative store. This ensures that
648    /// if the quant write fails, the worst case is a vector reachable in full-precision
649    /// but not yet in quantized search (benign), rather than a quantized ghost with no
650    /// full-precision backing data.
651    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        // First, write to the authoritative full-precision store.
660        if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) {
661            return std::future::ready(Err(err));
662        }
663
664        // Then, write the compressed representation to the quant store.
665        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
677/// Assign to just the full-precision store
678///
679impl<T> SetElement<&[T]> for BfTreeProvider<T, NoStore>
680where
681    T: VectorRepr,
682{
683    type SetError = ANNError;
684
685    /// Store the provided element in just the full-precision vector stores
686    ///
687    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
703//////////////////////
704// StartPoint Trait //
705//////////////////////
706
707/// A struct with a private member that cannot be constructed outside of this module.
708///
709/// This is used to prevent users from calling internal methods directly.
710pub struct Hidden(());
711
712/// A trait for setting the start points of a BfTreeProvider.
713///
714/// This trait is implemented by `BfTreeProvider` variants that support setting start points.
715/// The `Hidden` parameter ensures that users cannot call `set_start_points` directly;
716/// they must go through the `BfTreeProvider::new` constructor which handles this internally.
717pub trait StartPoint<T> {
718    /// Set the start points of the provider.
719    ///
720    /// # Safety
721    /// This method is internal and should not be called directly by users.
722    /// Use `BfTreeProvider::new` instead.
723    #[doc(hidden)]
724    fn set_start_points(&self, hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()>;
725}
726
727////////////////////
728// SetStartPoints //
729////////////////////
730
731/// Set start points for the BfTreeProvider with quantization.
732///
733/// This implementation sets both the full-precision and quantized vectors for each
734/// start point, as well as initializing empty neighbor lists.
735impl<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            // Set the full-precision vector
752            self.full_vectors.set_vector_sync(id.into_usize(), v)?;
753            self.quant_vectors.set_vector_sync(id.into_usize(), v)?;
754            // Initialize empty neighbor list
755            scratch.write_neighbors(id, &[])?;
756        }
757
758        Ok(())
759    }
760}
761
762/// Set start points for the BfTreeProvider without quantization.
763///
764/// This implementation sets the full-precision vectors for each start point
765/// and initializes empty neighbor lists.
766impl<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            // Set the full-precision vector
783            self.full_vectors.set_vector_sync(id.into_usize(), v)?;
784            // Initialize empty neighbor list
785            scratch.write_neighbors(id, &[])?;
786        }
787
788        Ok(())
789    }
790}
791
792//////////////////
793// FullAccessor //
794//////////////////
795
796/// An accessor for retrieving full-precision vectors from the `BfTreeProvider`.
797///
798/// This type implements the following traits:
799///
800/// * [`Accessor`] for the [`BfTreeProvider`].
801/// * [`BuildQueryComputer`].
802///
803pub struct FullAccessor<'a, T, Q>
804where
805    T: VectorRepr,
806    Q: AsyncFriendly,
807{
808    /// The host provider.
809    provider: &'a BfTreeProvider<T, Q>,
810    /// The fused query-distance computer.
811    computer: T::QueryDistance,
812    /// A buffer to store retrieved elements.
813    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
897///////////////////
898// QuantAccessor //
899///////////////////
900
901/// An accessor that retrieves the quantized portion of the [`BfTreeProvider`].
902///
903/// This type implements the following traits:
904///
905/// * [`Accessor`] for the `BfTreeProvider`.
906///
907pub struct QuantAccessor<'a, T>
908where
909    T: VectorRepr,
910{
911    provider: &'a BfTreeProvider<T, QuantVectorProvider>,
912    /// The fused query-distance computer.
913    computer: super::quant::QuantQueryComputer,
914    /// A buffer to store retrieved elements.
915    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
1007///////////////////////
1008// FullPruneAccessor //
1009///////////////////////
1010
1011/// A [`glue::PruneAccessor`] for full-precision vectors in the `BfTreeProvider`.
1012pub 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
1109////////////////////////
1110// QuantPruneAccessor //
1111////////////////////////
1112
1113/// A [`glue::PruneAccessor`] for quantized vectors in the `BfTreeProvider`.
1114pub 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
1213/// An owned quantized vector that reborrows to [`Opaque`].
1214///
1215/// Unlike inmem providers (which hand back zero-copy references into a contiguous backing
1216/// array), bf_tree copies vector data out of the tree on every access. The
1217/// [`workingset::View`] trait requires `get` to return something that implements
1218/// `Reborrow<'short, Target = Opaque<'short>>`, so we need an owned type that bridges
1219/// bf_tree's copy-out model with the working set's reborrow expectation.
1220pub 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
1229////////////////
1230// Strategies //
1231////////////////
1232
1233/// Perform a search entirely in the full-precision space.
1234///
1235/// Starting points are not filtered out of the final results.
1236impl<'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
1262// Pruning
1263impl<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
1335/// Inplace delete strategy using full-precision vectors.
1336///
1337/// # Compatibility
1338///
1339/// This strategy is used with [`InplaceDeleteMethod::OneHop`] and
1340/// [`InplaceDeleteMethod::TwoHopAndOneHop`]. It is **not compatible** with
1341/// [`InplaceDeleteMethod::VisitedAndTopK`] because `BfTreeProvider` performs hard deletes —
1342/// the vector data is erased before `get_delete_element` is called, causing it to fail.
1343impl<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
1383/// Perform a search entirely in the quantized space.
1384///
1385/// Starting points are not filtered out of the final results.
1386impl<'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
1459/// Inplace delete strategy using quantized vectors.
1460///
1461/// # Compatibility
1462///
1463/// Same constraint as [`FullPrecision`]'s impl: not compatible with
1464/// [`InplaceDeleteMethod::VisitedAndTopK`] due to hard deletes.
1465impl<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
1503// Pruning
1504impl<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/// Post-processor that reranks quantized search results using full-precision distances.
1522#[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                    // Transient (deleted/missing) — skip this candidate.
1558                }
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    /// Whether the original model was in-memory (`true`) or on-disk (`false`).
1594    pub is_memory: bool,
1595    /// Whether CPR snapshot support was enabled.
1596    #[serde(default)]
1597    pub use_snapshot: bool,
1598}
1599
1600/// The element type of the full-precision vectors stored in the index.
1601#[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
1610/// A trait for mapping concrete vector element types to their [`VectorDtype`]
1611/// discriminant at compile time.
1612pub 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/// Graph configuration parameters persisted alongside the index.
1633/// These are needed to reconstruct the `DiskANNIndex` config on load.
1634#[derive(Serialize, Deserialize, Clone, Debug)]
1635pub struct GraphParams {
1636    /// l_build is the search list size used during index construction.
1637    /// When inserting a new vector into the DiskANN graph, the algorithm
1638    /// performs a greedy search to find the best neighbors to connect to.
1639    /// l_build controls how many candidate nodes are tracked during that search.
1640    pub l_build: usize,
1641    /// alpha is the pruning aggressiveness parameter used during graph
1642    /// construction. During pruning, when deciding whether to keep a candidate
1643    /// neighbor k for node i, the algorithm checks if there's already a
1644    /// closer neighbor j that "occludes" k. The occlusion test is is governed by alpha.
1645    pub alpha: f32,
1646    /// backedge_ratio controls how many reverse (back) edges are added after
1647    /// pruning during graph construction.
1648    pub backedge_ratio: f32,
1649    /// vector_dtype indicates the data type of the vectors stored in the index, which is necessary for correctly interpreting the raw bytes of the vectors when loading the index from disk.
1650    pub vector_dtype: VectorDtype,
1651}
1652
1653/// Helper struct for generating consistent file paths for BfTreeProvider persistence.
1654/// Centralizes all path patterns to avoid hardcoded strings throughout the codebase.
1655pub struct BfTreePaths;
1656
1657impl BfTreePaths {
1658    /// Returns the path for the parameters JSON file
1659    pub fn params_json(prefix: &str) -> String {
1660        format!("{}_params.json", prefix)
1661    }
1662
1663    /// Returns the path for the vectors BfTree file
1664    pub fn vectors_bftree(prefix: &str) -> std::path::PathBuf {
1665        std::path::PathBuf::from(format!("{}_vectors.bftree", prefix))
1666    }
1667
1668    /// Returns the path for the neighbors BfTree file
1669    pub fn neighbors_bftree(prefix: &str) -> std::path::PathBuf {
1670        std::path::PathBuf::from(format!("{}_neighbors.bftree", prefix))
1671    }
1672
1673    /// Returns the path for the quantized vectors BfTree file
1674    pub fn quant_bftree(prefix: &str) -> std::path::PathBuf {
1675        std::path::PathBuf::from(format!("{}_quant.bftree", prefix))
1676    }
1677
1678    /// Returns the path for the delete bitmap file
1679    pub fn delete_bin(prefix: &str) -> String {
1680        format!("{}_delete.bin", prefix)
1681    }
1682
1683    /// Returns the path for the spherical quantizer data file
1684    pub fn quant_data_bin(prefix: &str) -> String {
1685        format!("{}_quant_data.bin", prefix)
1686    }
1687}
1688
1689/// Save a BfTree to disk via CPR snapshot.
1690///
1691/// # Errors
1692///
1693/// Returns an error if `use_snapshot` is false, indicating the tree was not
1694/// configured with snapshot support.
1695fn 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
1709/// Load a BfTree from a CPR snapshot file.
1710///
1711/// The 0.5.x loader reconstructs config from the snapshot file header, so
1712/// no external `BfTreeParams` are needed.
1713fn 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
1718//////////////////////
1719// Serialization    //
1720//////////////////////
1721
1722impl<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(&params_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(&params_filename)?;
1800            let mut params_json = String::new();
1801            params_reader.read_to_string(&mut params_json)?;
1802            serde_json::from_str(&params_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(&params_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(&params_filename)?;
1945            let mut params_json = String::new();
1946            params_reader.read_to_string(&mut params_json)?;
1947            serde_json::from_str(&params_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///////////
2005// Tests //
2006///////////
2007
2008/// These unit tests target the correctness and functionality of Bf-Tree data provider
2009/// For functionality tests, Bf-Tree data provider runs against edge cases and is verified with hard-coded ground truth
2010/// For correctness tests, Bf-Tree data provider runs common workflows and is verified with the inmem data provider
2011///
2012/// Note that the test scenarios here should focus on assumptions made at the data provider level such as 'new vector should
2013/// have no neighbors'. Tests at individual index provider level should be added in the correponding provider mod instead while
2014/// tests regarding DiskANN async algorithms should be added to index::diskann_async
2015///
2016#[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        // Iterator
2366        //
2367        assert_eq!((&provider).into_iter(), 0..(10 + 2));
2368
2369        let iter = provider.iter();
2370
2371        // Insert vectors so they exist in the bf_tree (hard-delete checks presence)
2372        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            // Delete by external ID.
2390            //
2391            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        // With hard deletes, `release` is a no-op (data is permanently removed).
2403        // Verify that released IDs remain deleted.
2404        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        // out-of-bound set-element fails.
2413        //
2414        assert!(provider
2415            .set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0])
2416            .await
2417            .is_err());
2418    }
2419
2420    /// This functionality test targets scenarios of empty neighbor lists and ensures:
2421    /// 1. A new vector's neighbor list is empty
2422    /// 2. A vector's neighbor list could be set to empty
2423    /// 3. A non-existant vector's neighbor list is empty
2424    ///
2425    #[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        // Insert new vectors without neighbors and empty neighbor list is
2455        // expected for each newly inserted vector
2456        //
2457        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            // First attempt should return empty
2462            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            // After we set the empty neighbor list, our attempt should succeed
2470            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        // Add a non-empty neighbor list for a vector and then set it to empty
2480        // In the end, an empty neighbor list is expected for the vector
2481        //
2482        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]); // len = 10
2493
2494            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        // Non-existant vectors have empty neighbor lists
2504        //
2505        let mut out = AdjacencyList::from_iter_untrusted([10, 20, 30, 40, 50, 60, 70, 80, 90, 100]); // len = 10
2506
2507        // Attempt to access non-existant vector's neighbor list should fail as NotFound
2508        assert!(provider
2509            .neighbor_provider
2510            .get_neighbors(200, &mut out)
2511            .is_err());
2512        assert!(out.is_empty());
2513    }
2514
2515    ///////////////////////////////////////////////
2516    // SaveWith/LoadWith Tests for BfTree-based //
2517    ///////////////////////////////////////////////
2518
2519    use tempfile::tempdir;
2520
2521    /// Test saving and loading of BfTreeProvider without quantization, including deleted vertices
2522    #[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        // Create a temporary directory for test files
2531        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        // Create configs with storage backends
2539        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        // Create provider parameters
2555        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        // Create provider
2571        let provider =
2572            BfTreeProvider::<f32, NoStore>::new(params.clone(), start_points.as_view(), NoStore)
2573                .unwrap();
2574
2575        // Populate provider with vectors
2576        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        // Populate provider with neighbor lists
2585        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        // Save to a different prefix to exercise the snapshot copy logic
2599        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        // Load using trait method (includes delete bitmap)
2608        let loaded_provider = BfTreeProvider::<f32, NoStore>::load_with(&storage, &save_prefix)
2609            .await
2610            .unwrap();
2611
2612        // Verify vectors
2613        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        // Verify neighbor lists
2623        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        // Cleanup is automatic when temp_dir goes out of scope
2644    }
2645
2646    /// Test saving and loading of BfTreeProvider with quantization, including deleted vertices
2647    #[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        // Create a temporary directory for test files
2656        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        // Create configs with storage backends
2664        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        // Create spherical quantizer
2684        let quantizer = create_test_quantizer(dim);
2685
2686        // Create provider parameters
2687        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        // Create provider with quantization
2702        let provider = BfTreeProvider::<f32, QuantVectorProvider>::new(
2703            params.clone(),
2704            start_points.as_view(),
2705            quantizer,
2706        )
2707        .unwrap();
2708
2709        // Populate provider with vectors
2710        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        // Populate provider with neighbor lists
2719        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        // Save to a different prefix to exercise the snapshot copy logic
2730        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        // Load using trait method (includes delete bitmap and quantization)
2739        let loaded_provider =
2740            BfTreeProvider::<f32, QuantVectorProvider>::load_with(&storage, &save_prefix)
2741                .await
2742                .unwrap();
2743
2744        // Verify quantizer properties match after round-trip
2745        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        // Verify vectors
2762        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        // Verify quantized vectors
2772        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        // Verify neighbor lists
2782        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        // Cleanup is automatic when temp_dir goes out of scope
2803    }
2804
2805    /// Test saving an in-memory (no disk) BfTreeProvider without quantization and loading it back.
2806    #[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        // In-memory config (no file path needed)
2821        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        // Populate vectors and neighbors
2840        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        // Delete a couple of vectors
2856        provider.delete(ctx, &3u32).await.unwrap();
2857        provider.delete(ctx, &7u32).await.unwrap();
2858
2859        // Save to disk from in-memory
2860        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        // Load back
2870        let loaded = BfTreeProvider::<f32, NoStore>::load_with(&storage, &save_prefix)
2871            .await
2872            .unwrap();
2873
2874        // Verify vectors
2875        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        // Verify neighbors
2888        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        // Verify deletes
2903        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    /// Test saving an in-memory BfTreeProvider with PQ quantization and loading it back.
2918    #[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        // Populate vectors and neighbors
2954        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        // Save to disk from in-memory
2972        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        // Load back
2982        let loaded = BfTreeProvider::<f32, QuantVectorProvider>::load_with(&storage, &save_prefix)
2983            .await
2984            .unwrap();
2985
2986        // Verify full vectors (skip deleted id 2)
2987        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        // Verify quant vectors (skip deleted id 2)
3000        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        // Verify neighbors (skip deleted id 2)
3013        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        // Verify delete
3031        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        // 1536 * 4 = 6144 bytes + 8-byte key = 6152 bytes needed
3044        let result = VectorProvider::<f32>::new_with_config(
3045            100,
3046            1536,
3047            1,
3048            Config::default(), // cb_max_record_size = 1952
3049        );
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        // max_degree=500 → value = (500+1)*4 = 2004 bytes + 4-byte key = 2008 bytes
3064        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        // dim=128, f32 → key=8 + value=512 = 520 bytes, fits default 1952
3078        if let Err(e) = VectorProvider::<f32>::new_with_config(100, 128, 1, Config::default()) {
3079            panic!("VectorProvider should succeed: {e}");
3080        }
3081        // max_degree=64 → key=4 + value=(64+1)*4 = 264 bytes, fits default 1952
3082        if let Err(e) = NeighborProvider::<u32>::new_with_config(64, Config::default()) {
3083            panic!("NeighborProvider should succeed: {e}");
3084        }
3085    }
3086}