From a4c41f362855304475bdc9dc03a8e8ce712f964f Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Mon, 18 May 2026 08:38:49 -0700 Subject: [PATCH 1/7] Migrate bf_tree provider from PQ to spherical quantization Squashed 30 commits for rebase. Key changes: - Replace PQ with spherical quantization in bf_tree-provider - Remove D generic parameter from BfTreeProvider - Add QuantAccessor and quantized search support - Implement DeletionCheck traits - Remove hybrid computer (quant distances work alone) - Remove benches/ directory - Add streaming benchmark support - Various fixes and cleanups Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/workflows/ci.yml | 2 +- Cargo.lock | 24 +- Cargo.toml | 1 + diskann-bftree/Cargo.toml | 41 + diskann-bftree/src/lib.rs | 172 ++ .../src/neighbors.rs | 16 +- .../src}/provider.rs | 1965 ++++++++--------- diskann-bftree/src/quant.rs | 411 ++++ .../src/vectors.rs | 53 +- diskann-providers/Cargo.toml | 3 - .../graph/provider/async_/bf_tree/mod.rs | 39 - .../async_/bf_tree/quant_vector_provider.rs | 422 ---- .../src/model/graph/provider/async_/common.rs | 15 +- .../src/model/graph/provider/async_/mod.rs | 4 - .../provider/async_/table_delete_provider.rs | 63 - diskann/src/graph/mod.rs | 2 + diskann/src/graph/strategy.rs | 17 + 17 files changed, 1607 insertions(+), 1643 deletions(-) create mode 100644 diskann-bftree/Cargo.toml create mode 100644 diskann-bftree/src/lib.rs rename diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs => diskann-bftree/src/neighbors.rs (98%) rename {diskann-providers/src/model/graph/provider/async_/bf_tree => diskann-bftree/src}/provider.rs (65%) create mode 100644 diskann-bftree/src/quant.rs rename diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs => diskann-bftree/src/vectors.rs (88%) delete mode 100644 diskann-providers/src/model/graph/provider/async_/bf_tree/mod.rs delete mode 100644 diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs create mode 100644 diskann/src/graph/strategy.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8ccce495f..dae60b368 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ env: CARGO_TERM_COLOR: always # The features we want to explicitly test. For example, the `flatbuffers-build` feature # of `diskann-quantization` requires additional setup and so must not be included by default. - DISKANN_FEATURES: "virtual_storage,bf_tree,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen" + DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen" # Intel SDE version used for baseline and AVX-512 emulation jobs. SDE_VERSION: "sde-external-10.7.0-2026-02-18-lin" diff --git a/Cargo.lock b/Cargo.lock index 94c12d370..db6280ed3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -728,6 +728,28 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "diskann-bftree" +version = "0.52.0" +dependencies = [ + "bf-tree", + "bytemuck", + "diskann", + "diskann-providers", + "diskann-quantization", + "diskann-utils", + "diskann-vector", + "futures-util", + "half", + "rand 0.9.4", + "rstest", + "serde", + "serde_json", + "tempfile", + "thiserror 2.0.17", + "tokio", +] + [[package]] name = "diskann-disk" version = "0.52.0" @@ -835,7 +857,6 @@ version = "0.52.0" dependencies = [ "approx", "arc-swap", - "bf-tree", "bincode", "bytemuck", "byteorder", @@ -861,7 +882,6 @@ dependencies = [ "rayon", "rstest", "serde", - "serde_json", "tempfile", "thiserror 2.0.17", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 42c41a7ab..c9f1fb8b2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ members = [ "diskann-benchmark", "diskann-tools", "vectorset", + "diskann-bftree", ] default-members = [ diff --git a/diskann-bftree/Cargo.toml b/diskann-bftree/Cargo.toml new file mode 100644 index 000000000..ce7033271 --- /dev/null +++ b/diskann-bftree/Cargo.toml @@ -0,0 +1,41 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +[package] +name = "diskann-bftree" +version.workspace = true +description.workspace = true +authors.workspace = true +documentation.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +bf-tree.workspace = true +bytemuck = { workspace = true, features = ["must_cast"] } +diskann.workspace = true +diskann-providers.workspace = true +diskann-quantization = { workspace = true, features = ["flatbuffers"] } +diskann-utils.workspace = true +diskann-vector.workspace = true +half = { workspace = true, features = ["bytemuck", "num-traits"] } +futures-util.workspace = true +rand.workspace = true +serde = { workspace = true, features = ["derive"] } +serde_json.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } + +[dev-dependencies] +diskann = { workspace = true, features = ["testing"] } +diskann-providers = { workspace = true, features = ["testing"] } +diskann-utils = { workspace = true, features = ["testing"] } +rstest.workspace = true +tempfile.workspace = true +tokio = { workspace = true, features = ["full"] } + +[features] +default = [] +experimental_diversity_search = ["diskann/experimental_diversity_search"] + +[lints] +workspace = true diff --git a/diskann-bftree/src/lib.rs b/diskann-bftree/src/lib.rs new file mode 100644 index 000000000..111f2085d --- /dev/null +++ b/diskann-bftree/src/lib.rs @@ -0,0 +1,172 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! BfTree-based data provider for DiskANN async indexes. +//! +//! This crate provides a [`BfTree`](bf_tree::BfTree)-backed implementation of the DiskANN +//! [`DataProvider`](diskann::provider::DataProvider) trait, enabling indexes that can +//! transparently spill to disk for datasets larger than available memory. + +pub mod neighbors; +pub mod provider; +pub mod quant; +pub mod vectors; + +// Accessors +pub use provider::{ + AsVectorDtype, BfTreePaths, BfTreeProvider, BfTreeProviderParameters, CreateQuantProvider, + FullAccessor, GraphParams, Hidden, QuantAccessor, StartPoint, VectorDtype, +}; + +pub use bf_tree::Config; + +use diskann::{ + error::{RankedError, TransientError}, + ANNError, +}; + +#[derive(Debug, Clone, Copy)] +pub struct NoStore; + +/// Wrapper around [`bf_tree::ConfigError`] that implements [`std::error::Error`]. +#[derive(Debug, Clone)] +pub struct ConfigError(pub bf_tree::ConfigError); + +impl std::fmt::Display for ConfigError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "BfTree configuration error: {:?}", self.0) + } +} + +impl std::error::Error for ConfigError {} + +impl From for ANNError { + #[track_caller] + #[inline(never)] + fn from(error: ConfigError) -> ANNError { + ANNError::new(diskann::ANNErrorKind::IndexError, error) + } +} + +trait AsKey { + fn as_key(&self) -> &[u8]; +} + +impl AsKey for usize { + fn as_key(&self) -> &[u8] { + bytemuck::bytes_of(self) + } +} + +//////////// +// Errors // +//////////// +#[derive(Debug)] +pub enum VectorError { + /// the vector has been explicitly deleted + Deleted, + /// the key was not found + NotFound, +} + +#[derive(Debug)] +pub struct VectorUnavailable { + pub id: usize, + pub err: VectorError, +} + +impl std::fmt::Display for VectorUnavailable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.err { + VectorError::Deleted => write!(f, "vector {} was deleted", self.id), + VectorError::NotFound => write!(f, "vector {} not found", self.id), + } + } +} + +impl TransientError for VectorUnavailable { + fn acknowledge(self, _why: D) + where + D: std::fmt::Display, + { + // no-op: we are expecting transient deletion errors during traversal + } + + fn escalate(self, why: D) -> ANNError + where + D: std::fmt::Display, + { + ANNError::log_index_error(format!("{self}, escalated: {why}")) + } +} + +pub type AccessError = RankedError; + +/// Metrics recorded by [`DefaultContext`](diskann::provider::DefaultContext). +#[derive(Debug, Clone)] +#[cfg_attr(test, derive(serde::Serialize, serde::Deserialize))] +pub struct ContextMetrics { + pub spawns: usize, + pub clones: usize, +} + +/// An atomic call counter used for test instrumentation. +/// +/// Under `#[cfg(test)]`, this is a real atomic counter. In production builds, +/// all methods are no-ops that the compiler can eliminate entirely. +#[cfg(test)] +pub(crate) struct TestCallCount { + count: std::sync::atomic::AtomicUsize, +} + +#[cfg(test)] +impl TestCallCount { + pub fn new() -> Self { + Self { + count: std::sync::atomic::AtomicUsize::new(0), + } + } + + pub fn enabled() -> bool { + true + } + + pub fn increment(&self) { + self.count + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + + pub fn get(&self) -> usize { + self.count.load(std::sync::atomic::Ordering::Relaxed) + } +} + +#[cfg(not(test))] +#[allow(dead_code)] +pub(crate) struct TestCallCount {} + +#[cfg(not(test))] +#[allow(dead_code)] +impl TestCallCount { + pub fn new() -> Self { + Self {} + } + + pub fn enabled() -> bool { + false + } + + pub fn increment(&self) {} + + pub fn get(&self) -> usize { + 0 + } +} + +impl Default for TestCallCount { + fn default() -> Self { + Self::new() + } +} diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs b/diskann-bftree/src/neighbors.rs similarity index 98% rename from diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs rename to diskann-bftree/src/neighbors.rs index 16299b622..bf6ace7b4 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/neighbor_provider.rs +++ b/diskann-bftree/src/neighbors.rs @@ -7,22 +7,24 @@ use std::marker::PhantomData; +use crate::AsKey; use bf_tree::{BfTree, Config}; use bytemuck::{bytes_of, cast_slice, cast_slice_mut}; use diskann::{ - ANNError, ANNResult, graph::AdjacencyList, provider::HasId, utils::{IntoUsize, TryIntoVectorId, VectorId}, + ANNError, ANNResult, }; -use super::super::common::TestCallCount; use super::ConfigError; +use crate::TestCallCount; pub struct NeighborProvider { adjacency_list_index: BfTree, dim: usize, // Max number of neighbors in a neighbor list + 1 for the neighbor count - pub num_get_calls: TestCallCount, + #[allow(dead_code)] + pub(crate) num_get_calls: TestCallCount, _phantom: PhantomData, } @@ -81,7 +83,7 @@ impl NeighborProvider { // Serialize the key, vector_id, into a byte string, &[u8] let i = vector_id.into_usize(); - let key = bytes_of::(&i); + let key = i.as_key(); // Search and retrieve the corresponding neighbor list data as a byte string, &[u8], in the format of // |VectorId|VectorId|...|Invalid|Invalid|VectorId (list length)| @@ -166,7 +168,7 @@ impl NeighborProvider { // Serialize the key, vector_id, into a byte string, &[u8] let i = vector_id.into_usize(); - let key = bytes_of::(&i); + let key = i.as_key(); // Serialize the value, neighbor list, into a byte string, &u[8] let neighbor_list_edges_in_byte = cast_slice::(neighbors); @@ -218,7 +220,7 @@ impl NeighborProvider { // We avoid one data copy by directly writing to bf-tree instead of invoking set_neighbor() // Also avoid a bunch of unnecssary checks let i = vector_id.into_usize(); - let key = bytes_of::(&i); + let key = i.as_key(); let value = cast_slice::(&neighbor_list); self.adjacency_list_index.insert(key, value); } @@ -229,7 +231,7 @@ impl NeighborProvider { pub fn delete_vector(&self, vector_id: I) -> ANNResult<()> { // Serialize the key, vector_id, into a byte string, &[u8] let i = vector_id.into_usize(); - let key = bytes_of::(&i); + let key = i.as_key(); self.adjacency_list_index.delete(key); Ok(()) diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs b/diskann-bftree/src/provider.rs similarity index 65% rename from diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs rename to diskann-bftree/src/provider.rs index 7838d403d..2ba804db4 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -9,21 +9,26 @@ use std::{ io::{Read, Write}, num::NonZeroUsize, str::FromStr, - sync::Arc, }; +use diskann_quantization::{ + alloc::{GlobalAllocator, Poly}, + spherical::iface::{self as spherical_iface, try_deserialize, Opaque, Quantizer}, +}; use serde::{Deserialize, Serialize}; use bf_tree::{BfTree, Config}; use diskann::{ - ANNError, ANNResult, default_post_processor, + default_post_processor, + error::{Infallible, RankedError}, graph::{ - AdjacencyList, DiskANNIndex, SearchOutputBuffer, glue::{ - self, Batch, DefaultPostProcessor, ExpandBeam, InplaceDeleteStrategy, InsertStrategy, - MultiInsertStrategy, PruneStrategy, SearchExt, SearchStrategy, + self, Batch, CopyIds, DefaultPostProcessor, ExpandBeam, InplaceDeleteStrategy, + InsertStrategy, MultiInsertStrategy, PruneStrategy, SearchExt, SearchStrategy, }, - workingset::{self, map}, + strategy::{FullPrecision, Quantized}, + workingset::{self, map, Map}, + AdjacencyList, SearchOutputBuffer, }, neighbor::Neighbor, provider::{ @@ -32,27 +37,17 @@ use diskann::{ NoopGuard, SetElement, }, utils::{IntoUsize, VectorRepr}, + ANNError, ANNResult, }; use diskann_utils::{future::AsyncFriendly, views::MatrixView}; -use diskann_vector::{DistanceFunction, distance::Metric}; - -use crate::model::{ - graph::provider::async_::{ - TableDeleteProviderAsync, - bf_tree::{ - neighbor_provider::NeighborProvider, quant_vector_provider::QuantVectorProvider, - vector_provider::VectorProvider, - }, - common::{CreateDeleteProvider, FullPrecision, Hybrid, NoDeletes, NoStore, Panics}, - distances, - postprocess::{AsDeletionCheck, DeletionCheck, RemoveDeletedIdsAndCopy}, - }, - pq::{self, FixedChunkPQTable, NUM_PQ_CENTROIDS}, -}; +use diskann_vector::{distance::Metric, DistanceFunction}; -use crate::storage::{LoadWith, PQStorage, SaveWith}; - -use crate::storage::{StorageReadProvider, StorageWriteProvider}; +use super::{ + neighbors::NeighborProvider, quant::QuantVectorProvider, vectors::VectorProvider, AccessError, + NoStore, +}; +use diskann_providers::model::graph::provider::async_::distances::UnwrapErr; +use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, StorageWriteProvider}; ///////////////////// // BfTreeProvider // @@ -70,40 +65,20 @@ use crate::storage::{StorageReadProvider, StorageWriteProvider}; /// * `Q`: The full type of the quant vector store. This is not constrained by a trait and /// rather relies on implementation for several concrete types, including: /// -/// - [`BfTreeQuantVectorProviderAsync`]: A Bf-Tree based PQ-based quantized vector store. +/// - [`QuantVectorProvider`]: A Bf-Tree based spherical quantized vector store. /// - [`NoStore`]: Disable quantization altogether. Note that this disables all /// methods reached through quantization based [`Accessor`]s at compile-time. /// -/// * `D`: The type of the deleted vector store. Like the quantized store, this is also -/// not constrained by a trait and rather relies on implementation for concrete types. -/// These are: -/// -/// - [`NoDeletes`]: Do not support deletion at all (this disables implementation of -/// the [`Delete`] trait. -/// - [`TableDeleteProviderAsync`]: A bitmap storing deletion information. -/// -/// * `Ctx`: A parameter controlling the [`ExecutionContext`] to be associated with this -/// provider. For the majority of cases, this is [`DefaultContext`], but is left as -/// a parameter to allow extension. -/// /// # Indexing Strategies /// -/// * [`FullPrecision`]: The strategies implemented by [`FullPrecision`] only retrieve data -/// from the full-precision portion of the index. No quantized vectors are used. -/// -/// During search, start points are filtered from the final results. -/// -/// * [`Hybrid`]: The strategies implemented by [`Hybrid`] can use a mix of quantized -/// and full-precision vectors. -/// -/// - Search: During search, quantized vectors are used with reranking applied to the -/// results before returning. +/// * [`FullPrecision`]: Only retrieves data from the full-precision portion of the index. +/// No quantized vectors are used. During search, start points are filtered from the +/// final results. /// -/// - Insertion: Quantized vectors are used during the search phase. During the pruning -/// phase, a hybrid of quantized and full-precision vectors are used. -/// -/// The ratio of full-precision and quantized vectors is controlled by the -/// `max_fp_vecs_per_prune` parameter, which adjusts the implementation of [`Fill`]. +/// * [`Quantized`]: Performs all operations (search, pruning, insert) entirely in the +/// quantized space using spherical distance functions. Post-processing copies candidate +/// IDs forward without reranking. Fastest option — full-precision vectors are not +/// touched at query time. /// /// # Examples /// @@ -115,13 +90,12 @@ use crate::storage::{StorageReadProvider, StorageWriteProvider}; /// This example demonstrates how to create a `BfTreeProvider` that only supports /// full-precision vectors. /// ``` -/// use diskann_providers::model::graph::provider::async_::{ -/// bf_tree::{ -/// BfTreeProvider, BfTreeProviderParameters -/// }, -/// common::{NoStore, NoDeletes}, +/// use diskann_bftree::provider::{ +/// BfTreeProvider, BfTreeProviderParameters /// }; +/// use diskann_bftree::NoStore; /// use diskann_vector::distance::Metric; +/// use diskann_utils::views::{Init, Matrix}; /// use bf_tree::Config; /// use std::num::NonZeroUsize; /// @@ -130,7 +104,6 @@ use crate::storage::{StorageReadProvider, StorageWriteProvider}; /// num_start_points: NonZeroUsize::new(1).unwrap(), /// dim: 4, /// metric: Metric::L2, -/// max_fp_vecs_per_fill: None, /// max_degree: 32, /// vector_provider_config: Config::default(), /// quant_vector_provider_config: Config::default(), @@ -139,95 +112,52 @@ use crate::storage::{StorageReadProvider, StorageWriteProvider}; /// }; /// /// // Create a table that supports 5 points and 1 start point. -/// let provider = BfTreeProvider::::new_empty( +/// let start_points = Matrix::new(Init(|| 0.0f32), 1, 4); +/// let provider = BfTreeProvider::::new( /// parameters, +/// start_points.as_view(), /// NoStore, -/// NoDeletes, /// ); /// ``` /// -/// ## Full-Precision and PQ - No Deletes +/// ## Full-Precision and Spherical Quantization /// -/// To create a two-level provider with a PQ-based quant vector store, a -/// [`FixedChunkPQTable`] can be supplied for the `quant_precursor` argument, as this +/// To create a two-level provider with a spherical quantization-based quant vector store, +/// a `Poly` can be supplied for the `quant_precursor` argument, as this /// implements the [`CreateQuantProvider`] trait. /// ``` -/// use diskann_providers::model::{ -/// pq::FixedChunkPQTable, -/// graph::provider::async_::{ -/// bf_tree::{ -/// BfTreeProvider, BfTreeProviderParameters -/// }, -/// common::NoDeletes, -/// }, -/// }; -/// use diskann_vector::distance::Metric; -/// use bf_tree::Config; -/// use std::num::NonZeroUsize; -/// -/// // An example PQ table. -/// let dim = 4; -/// let table = FixedChunkPQTable::new( -/// dim, -/// Box::new([0.0, 0.0, 0.0, 0.0]), -/// Box::new([0.0, 0.0, 0.0, 0.0]), -/// Box::new([0, dim]), -/// ).unwrap(); -/// -/// let parameters = BfTreeProviderParameters { -/// max_points: 5, -/// num_start_points: NonZeroUsize::new(1).unwrap(), -/// dim: 4, -/// metric: Metric::L2, -/// max_fp_vecs_per_fill: None, -/// max_degree: 32, -/// vector_provider_config: Config::default(), -/// quant_vector_provider_config: Config::default(), -/// neighbor_list_provider_config: Config::default(), -/// graph_params: None, +/// use diskann_quantization::{ +/// alloc::{GlobalAllocator, Poly, poly}, +/// algorithms::TransformKind, +/// spherical::{iface, SphericalQuantizer, SupportedMetric, PreScale}, /// }; -/// -/// // Create a table that supports 5 points and 1 start point. -/// let provider = BfTreeProvider::::new_empty( -/// parameters, -/// table, -/// NoDeletes, -/// ); -/// ``` -/// -/// ## Full-Precision and PQ - With Deletes. -/// -/// If deletes are desired, than the type [`TableBasedDeletes`] can be passed to the -/// constructor. -/// ``` -/// use diskann_providers::model::{ -/// pq::FixedChunkPQTable, -/// graph::provider::async_::{ -/// bf_tree::{ -/// BfTreeProvider, BfTreeProviderParameters -/// }, -/// common::TableBasedDeletes, -/// }, +/// use diskann_utils::views::{Init, Matrix}; +/// use diskann_bftree::provider::{ +/// BfTreeProvider, BfTreeProviderParameters /// }; /// use diskann_vector::distance::Metric; /// use bf_tree::Config; /// use std::num::NonZeroUsize; +/// use rand::rngs::StdRng; +/// use rand::SeedableRng; /// -/// // An example PQ table. /// let dim = 4; -/// let table = FixedChunkPQTable::new( -/// dim, -/// Box::new([0.0, 0.0, 0.0, 0.0]), -/// Box::new([0.0, 0.0, 0.0, 0.0]), -/// Box::new([0, dim]), +/// let data = Matrix::new(Init(|| 1.0f32), 4, dim); +/// let mut rng = StdRng::seed_from_u64(42); +/// let sq = SphericalQuantizer::train( +/// data.as_view(), TransformKind::Null, +/// SupportedMetric::SquaredL2, PreScale::None, +/// &mut rng, GlobalAllocator, /// ).unwrap(); +/// let imp = iface::Impl::<1>::new(sq).unwrap(); +/// let poly = Poly::new(imp, GlobalAllocator).unwrap(); +/// let quantizer: Poly = poly!(iface::Quantizer, poly); /// /// let parameters = BfTreeProviderParameters { /// max_points: 5, /// num_start_points: NonZeroUsize::new(1).unwrap(), /// dim: 4, /// metric: Metric::L2, -/// max_fp_vecs_per_fill: None, /// max_degree: 32, /// vector_provider_config: Config::default(), /// quant_vector_provider_config: Config::default(), @@ -236,13 +166,14 @@ use crate::storage::{StorageReadProvider, StorageWriteProvider}; /// }; /// /// // Create a table that supports 5 points and 1 start point. -/// let provider = BfTreeProvider::::new_empty( +/// let start_points = Matrix::new(Init(|| 0.0f32), 1, 4); +/// let provider = BfTreeProvider::::new( /// parameters, -/// table, -/// TableBasedDeletes, +/// start_points.as_view(), +/// quantizer, /// ); /// ``` -pub struct BfTreeProvider +pub struct BfTreeProvider where T: VectorRepr, { @@ -258,15 +189,6 @@ where // pub(crate) neighbor_provider: NeighborProvider, - // The delete provider. If `D == NoDeletes`, then delete related operations are disabled. - // - pub(super) deleted: D, - - // A parameter controlling hybrid pruning, where some set of full-precision vectors are - // fetched and the rest are quantized vectors - // - pub(super) max_fp_vecs_per_fill: usize, - // The metric to use for distances // pub(super) metric: Metric, @@ -290,10 +212,6 @@ pub struct BfTreeProviderParameters { // The metric to use for distance computations pub metric: Metric, - // If quantization is used, this parameter controls how many full-precision - // vectors are retrieved for each fill operation - pub max_fp_vecs_per_fill: Option, - // The maximum number of neighbors to store for each vector pub max_degree: u32, @@ -310,57 +228,10 @@ pub struct BfTreeProviderParameters { pub graph_params: Option, } -pub type Index = Arc>>; -pub type QuantIndex = Arc>>; - -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, { - /// Construct a new, unpopulated data provider. - /// - /// # Arguments - /// * `params`: An instance of [`BfTreeProviderParameters`] collecting shared - /// configuration information. - /// * `quant_precursor`: A precursor type for the quantizer layer. - /// * `delete_precursor`: A precursor type for the delete layer. - /// * `neighbor_precursor`: A precursor type for the neighbor layer. - /// or the neighbor layer - pub fn new_empty( - params: BfTreeProviderParameters, - quant_precursor: TQ, - delete_precursor: TD, - ) -> ANNResult - where - TQ: CreateQuantProvider, - TD: CreateDeleteProvider, - { - let num_start_points = params.num_start_points.get(); - - Ok(Self { - quant_vectors: quant_precursor.create( - params.max_points, - num_start_points, - params.metric, - params.quant_vector_provider_config, - )?, - full_vectors: VectorProvider::new_with_config( - params.max_points, - params.dim, - num_start_points, - params.vector_provider_config, - )?, - neighbor_provider: NeighborProvider::new_with_config( - params.max_degree, - params.neighbor_list_provider_config, - )?, - deleted: delete_precursor.create(params.max_points + num_start_points), - max_fp_vecs_per_fill: params.max_fp_vecs_per_fill.unwrap_or(usize::MAX), - metric: params.metric, - graph_params: params.graph_params, - }) - } - /// Construct a new data provider with start points initialized. /// /// This is the primary constructor for `BfTreeProvider`. It creates the provider @@ -376,16 +247,14 @@ where /// /// # Type Constraints /// * `Self: StartPoint` - The provider must implement the `StartPoint` trait. - pub fn new( + pub fn new( params: BfTreeProviderParameters, start_points: MatrixView<'_, T>, quant_precursor: TQ, - delete_precursor: TD, ) -> ANNResult where Self: StartPoint, TQ: CreateQuantProvider, - TD: CreateDeleteProvider, { // Early validation before allocating resources if start_points.nrows() != params.num_start_points.get() { @@ -396,7 +265,21 @@ where ))); } - let provider = Self::new_empty(params.clone(), quant_precursor, delete_precursor)?; + let provider = Self { + quant_vectors: quant_precursor.create(params.quant_vector_provider_config)?, + full_vectors: VectorProvider::new_with_config( + params.max_points, + params.dim, + params.num_start_points.get(), + params.vector_provider_config, + )?, + neighbor_provider: NeighborProvider::new_with_config( + params.max_degree, + params.neighbor_list_provider_config, + )?, + metric: params.metric, + graph_params: params.graph_params, + }; provider.set_start_points(Hidden(()), start_points)?; { // Initialize all neighborhoods to be empty lists. @@ -457,18 +340,7 @@ where } } -impl BfTreeProvider -where - T: VectorRepr, -{ - /// A temporary method while development of deletion is in progress - /// - pub fn clear_delete_set(&self) { - self.deleted.clear(); - } -} - -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, { @@ -482,7 +354,7 @@ where } } -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, { @@ -493,9 +365,63 @@ where } } +impl Delete for BfTreeProvider +where + T: VectorRepr, + Q: AsyncFriendly + QuantDelete, +{ + fn release( + &self, + _context: &Self::Context, + _id: Self::InternalId, + ) -> impl std::future::Future> + Send { + // no-op: hard delete removes the data + std::future::ready(Ok(())) + } + + fn delete( + &self, + _context: &Self::Context, + gid: &Self::ExternalId, + ) -> impl std::future::Future> + Send { + let id = *gid; + + if let Err(e) = self.neighbor_provider.set_neighbors(id, &[]) { + return std::future::ready(Err(e)); + } + self.full_vectors.delete_vector(id as usize); + self.quant_vectors.delete_vector(id as usize); + + std::future::ready(Ok(())) + } + + fn status_by_external_id( + &self, + context: &Self::Context, + gid: &Self::ExternalId, + ) -> impl std::future::Future> + Send + { + self.status_by_internal_id(context, *gid) + } + + fn status_by_internal_id( + &self, + _context: &Self::Context, + id: Self::InternalId, + ) -> impl std::future::Future> + Send + { + let status = match self.full_vectors.get_vector_sync(id.into_usize()) { + Ok(_) => Ok(ElementStatus::Valid), + Err(RankedError::Transient(_)) => Ok(ElementStatus::Deleted), + Err(RankedError::Error(e)) => Err(e), + }; + std::future::ready(status) + } +} + /// Allow `&BfTreeProvider` to implement `IntoIter` /// -impl IntoIterator for &BfTreeProvider +impl IntoIterator for &BfTreeProvider where T: VectorRepr, { @@ -519,54 +445,45 @@ pub trait CreateQuantProvider { // Create a quant provider capable of tracking `max_points` with and additional // `frozen_points` at the end. // - fn create( - self, - max_points: usize, - frozen_points: usize, - metric: Metric, - bf_tree_config: Config, - ) -> ANNResult; + fn create(self, bf_tree_config: Config) -> ANNResult; } impl CreateQuantProvider for NoStore { type Target = NoStore; - fn create( - self, - _max_points: usize, - _frozen_points: usize, - _metric: Metric, - _bf_tree_config: Config, - ) -> ANNResult { + fn create(self, _bf_tree_config: Config) -> ANNResult { Ok(self) } } /// Allow a `FixedChunkPQTable` to be promoted to full quant vector store. /// -impl CreateQuantProvider for FixedChunkPQTable { +impl CreateQuantProvider for Poly { type Target = QuantVectorProvider; - fn create( - self, - max_points: usize, - frozen_points: usize, - metric: Metric, - bf_tree_config: Config, - ) -> ANNResult { - QuantVectorProvider::new_with_config( - metric, - max_points, - frozen_points, - self, - bf_tree_config, - ) + fn create(self, bf_tree_config: Config) -> ANNResult { + QuantVectorProvider::new_with_config(self, bf_tree_config) + } +} + +pub(crate) trait QuantDelete { + fn delete_vector(&self, id: usize); +} + +impl QuantDelete for NoStore { + fn delete_vector(&self, _id: usize) { + //no-op + } +} + +impl QuantDelete for QuantVectorProvider { + fn delete_vector(&self, id: usize) { + self.delete_vector(id); } } -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { pub fn neighbors(&self) -> &NeighborProvider { &self.neighbor_provider @@ -577,11 +494,10 @@ where // Data Provider // /////////////////// -impl DataProvider for BfTreeProvider +impl DataProvider for BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { type Context = DefaultContext; @@ -621,20 +537,18 @@ where } } -impl HasId for BfTreeProvider +impl HasId for BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { type Id = u32; } -impl<'a, T, Q, D> DelegateNeighbor<'a> for BfTreeProvider +impl<'a, T, Q> DelegateNeighbor<'a> for BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { type Delegate = &'a NeighborProvider; @@ -643,73 +557,6 @@ where } } -/// Support deletes when we have a valid delete provider. -/// -impl Delete for BfTreeProvider -where - Q: AsyncFriendly, - T: VectorRepr, -{ - fn release( - &self, - _: &DefaultContext, - id: Self::InternalId, - ) -> impl Future> + Send { - // delete the vector from bf-tree - if let Err(e) = self.neighbor_provider.delete_vector(id) { - return std::future::ready(Err(e)); - } - self.deleted.undelete(id.into_usize()); - // set its neighbors to an empty list in the neighbor provider - // self.neighbor_provider.set_neighbors(id, &[]); - let res = self - .neighbor_provider - .set_neighbors(id, &[]) - .map_err(|err| err.context(format!("resetting neighbors for undeleted id {}", id))); - std::future::ready(res) - } - - /// Delete an item by external ID - /// - #[inline] - fn delete( - &self, - _context: &DefaultContext, - gid: &Self::ExternalId, - ) -> impl Future> + Send { - self.deleted.delete(gid.into_usize()); - std::future::ready(Ok(())) - } - - /// Check the status via external ID - /// - #[inline] - fn status_by_external_id( - &self, - context: &DefaultContext, - gid: &Self::ExternalId, - ) -> impl Future> + Send { - // NOTE: ID translation is the identity, so we can refer to `status_by_internal_id`. - self.status_by_internal_id(context, *gid) - } - - /// Check the status via internal ID - /// - #[inline] - fn status_by_internal_id( - &self, - _context: &DefaultContext, - id: Self::InternalId, - ) -> impl Future> + Send { - let status = if self.deleted.is_deleted(id.into_usize()) { - ElementStatus::Deleted - } else { - ElementStatus::Valid - }; - std::future::ready(Ok(status)) - } -} - impl NeighborAccessor for &NeighborProvider { fn get_neighbors( self, @@ -747,10 +594,9 @@ impl NeighborAccessorMut for &NeighborProvider { /// Assign to both the full-precision and quant vector stores /// -impl SetElement<&[T]> for BfTreeProvider +impl SetElement<&[T]> for BfTreeProvider where T: VectorRepr, - D: AsyncFriendly, { type SetError = ANNError; @@ -784,10 +630,9 @@ where /// Assign to just the full-precision store /// -impl SetElement<&[T]> for BfTreeProvider +impl SetElement<&[T]> for BfTreeProvider where T: VectorRepr, - D: AsyncFriendly, { type SetError = ANNError; @@ -843,10 +688,9 @@ pub trait StartPoint { /// /// This implementation sets both the full-precision and quantized vectors for each /// start point, as well as initializing empty neighbor lists. -impl StartPoint for BfTreeProvider +impl StartPoint for BfTreeProvider where T: VectorRepr, - D: AsyncFriendly, { fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> { let start_point_ids = self.full_vectors.starting_points()?; @@ -875,10 +719,9 @@ where /// /// This implementation sets the full-precision vectors for each start point /// and initializes empty neighbor lists. -impl StartPoint for BfTreeProvider +impl StartPoint for BfTreeProvider where T: VectorRepr, - D: AsyncFriendly, { fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> { let start_point_ids = self.full_vectors.starting_points()?; @@ -910,48 +753,43 @@ where /// This type implements the following traits: /// /// * [`Accessor`] for the [`BfTreeProvider`]. -/// * [`ComputerAccessor`] for comparing full-precision distances. /// * [`BuildQueryComputer`]. /// -pub struct FullAccessor<'a, T, Q, D> +pub struct FullAccessor<'a, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { /// The host provider. - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, /// A buffer to store retrieved elements. element: Box<[T]>, } -impl HasId for FullAccessor<'_, T, Q, D> +impl HasId for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { type Id = u32; } -impl SearchExt for FullAccessor<'_, T, Q, D> +impl SearchExt for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { fn starting_points(&self) -> impl Future>> { std::future::ready(self.provider.starting_points()) } } -impl<'a, T, Q, D> FullAccessor<'a, T, Q, D> +impl<'a, T, Q> FullAccessor<'a, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { - pub(crate) fn new(provider: &'a BfTreeProvider) -> Self { + pub(crate) fn new(provider: &'a BfTreeProvider) -> Self { Self { provider, element: (0..provider.full_vectors.dim()) @@ -961,11 +799,10 @@ where } } -impl<'a, T, Q, D> DelegateNeighbor<'a> for FullAccessor<'_, T, Q, D> +impl<'a, T, Q> DelegateNeighbor<'a> for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { type Delegate = &'a NeighborProvider; @@ -974,11 +811,10 @@ where } } -impl Accessor for FullAccessor<'_, T, Q, D> +impl Accessor for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { /// This accessor returns a reference to a local copy of the vector. type Element<'a> @@ -989,9 +825,9 @@ where /// The reference version of `Element` is the same as `Element`. type ElementRef<'a> = &'a [T]; - // Choose to panic on an out-of-bounds access rather than propagate an error. + // Hard-deleted entries may be encountered via stale graph edges. // - type GetError = Panics; + type GetError = AccessError; /// Return the full-precision vector stored at index `i`. /// @@ -1002,26 +838,55 @@ where &mut self, id: Self::Id, ) -> impl Future, Self::GetError>> + Send { - // SAFETY: We've decided to live with UB (undefined behavior) that can result from - // potentially mixing unsynchronized reads and writes on the underlying memory - // - #[allow(clippy::expect_used)] - self.provider + let v = self + .provider .full_vectors .get_vector_into(id.into_usize(), &mut self.element) - .expect("Full vector provider failed to retrieve element"); + .map(|_: ()| &*self.element); + + std::future::ready(v) + } - std::future::ready(Ok(&*self.element)) + /// Perform a bulk operation, silently skipping entries that cannot be read + /// (e.g., hard-deleted vectors whose graph edges have not yet been cleaned up). + /// + fn on_elements_unordered( + &mut self, + itr: Itr, + mut f: F, + ) -> impl Future> + Send + where + Self: Sync, + Itr: Iterator + Send, + F: Send + FnMut(Self::ElementRef<'_>, Self::Id), + { + for i in itr { + match self + .provider + .full_vectors + .get_vector_into(i.into_usize(), &mut self.element) + { + Ok(()) => { + f(&self.element, i); + } + Err(RankedError::Transient(_)) => { + // Deleted or missing vector — expected during graph traversal, skip. + } + Err(e @ RankedError::Error(_)) => { + return std::future::ready(Err(e)); + } + } + } + std::future::ready(Ok(())) } } -impl BuildDistanceComputer for FullAccessor<'_, T, Q, D> +impl BuildDistanceComputer for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { - type DistanceComputerError = Panics; + type DistanceComputerError = Infallible; type DistanceComputer = T::Distance; fn build_distance_computer( @@ -1034,13 +899,12 @@ where } } -impl BuildQueryComputer<&[T]> for FullAccessor<'_, T, Q, D> +impl BuildQueryComputer<&[T]> for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { - type QueryComputerError = Panics; + type QueryComputerError = Infallible; type QueryComputer = T::QueryDistance; fn build_query_computer( @@ -1050,24 +914,11 @@ where Ok(T::query_distance(from, self.provider.metric)) } } -impl ExpandBeam<&[T]> for FullAccessor<'_, T, Q, D> -where - T: VectorRepr, - Q: AsyncFriendly, - D: AsyncFriendly, -{ -} - -impl<'a, T, Q, D> AsDeletionCheck for FullAccessor<'a, T, Q, D> +impl ExpandBeam<&[T]> for FullAccessor<'_, T, Q> where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly + DeletionCheck, { - type Checker = D; - fn as_deletion_check(&self) -> &D { - &self.provider.deleted - } } /////////////////// @@ -1081,52 +932,47 @@ where /// * [`Accessor`] for the `BfTreeProvider`. /// * [`BuildQueryComputer`]. /// -pub struct QuantAccessor<'a, T, D> +pub struct QuantAccessor<'a, T> where T: VectorRepr, - D: AsyncFriendly, { - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, element: Box<[u8]>, } -impl HasId for QuantAccessor<'_, T, D> +impl HasId for QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly, { type Id = u32; } -impl SearchExt for QuantAccessor<'_, T, D> +impl SearchExt for QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly, { fn starting_points(&self) -> impl Future>> { std::future::ready(self.provider.starting_points()) } } -impl<'a, T, D> QuantAccessor<'a, T, D> +impl<'a, T> QuantAccessor<'a, T> where T: VectorRepr, - D: AsyncFriendly, { - pub(crate) fn new(provider: &'a BfTreeProvider) -> Self { + pub(crate) fn new(provider: &'a BfTreeProvider) -> Self { Self { provider, - element: (0..provider.quant_vectors.pq_chunks()) + element: (0..provider.quant_vectors.quantizer.bytes()) .map(|_| u8::default()) .collect(), } } } -impl<'a, T, D> DelegateNeighbor<'a> for QuantAccessor<'_, T, D> +impl<'a, T> DelegateNeighbor<'a> for QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly, { type Delegate = &'a NeighborProvider; fn delegate_neighbor(&'a mut self) -> Self::Delegate { @@ -1134,23 +980,22 @@ where } } -impl Accessor for QuantAccessor<'_, T, D> +impl Accessor for QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly, { /// This accessor returns a reference to a local copy of the element. type Element<'a> - = &'a [u8] + = Opaque<'a> where Self: 'a; /// The reference version of `Element` is simply `Element`. - type ElementRef<'a> = &'a [u8]; + type ElementRef<'a> = Opaque<'a>; // ANNError on access failures in bf-tree // - type GetError = ANNError; + type GetError = AccessError; /// Return the quantized vector stored at index `i`. /// @@ -1164,12 +1009,13 @@ where .provider .quant_vectors .get_vector_into(id.into_usize(), &mut self.element) - .map(|_: ()| &*self.element); + .map(|_: ()| Opaque::new(&self.element)); std::future::ready(v) } - /// Perform a bulk operation + /// Perform a bulk operation, silently skipping entries that cannot be read + /// (e.g., hard-deleted vectors whose graph edges have not yet been cleaned up). /// fn on_elements_unordered( &mut self, @@ -1187,8 +1033,13 @@ where .quant_vectors .get_vector_into(i.into_usize(), &mut self.element) { - Ok(()) => f(&self.element, i), - Err(e) => { + Ok(()) => { + f(Opaque::new(&self.element), i); + } + Err(RankedError::Transient(_)) => { + // Deleted or missing vector — expected during graph traversal, skip. + } + Err(e @ RankedError::Error(_)) => { return std::future::ready(Err(e)); } } @@ -1197,208 +1048,89 @@ where } } -impl BuildQueryComputer<&[T]> for QuantAccessor<'_, T, D> +impl BuildQueryComputer<&[T]> for QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly, { type QueryComputerError = ANNError; - type QueryComputer = pq::distance::QueryComputer>; + type QueryComputer = UnwrapErr< + spherical_iface::QueryComputer, + spherical_iface::QueryDistanceError, + >; fn build_query_computer( &self, from: &[T], ) -> Result { - self.provider.quant_vectors.query_computer(from) + self.provider + .quant_vectors + .query_computer(from) + .map(|qc| UnwrapErr::new(qc.into_inner())) } } -impl ExpandBeam<&[T]> for QuantAccessor<'_, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ -} +impl ExpandBeam<&[T]> for QuantAccessor<'_, T> where T: VectorRepr {} -impl<'a, T, D> AsDeletionCheck for QuantAccessor<'a, T, D> +impl BuildDistanceComputer for QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly + DeletionCheck, { - type Checker = D; - fn as_deletion_check(&self) -> &D { - &self.provider.deleted + type DistanceComputerError = ANNError; + type DistanceComputer = + UnwrapErr; + + fn build_distance_computer( + &self, + ) -> Result { + self.provider + .quant_vectors + .distance_computer() + .map(UnwrapErr::new) } } -///////////////////// -// Hybrid Accessor // -///////////////////// - -/// A hybrid accessor that fetches a mixture of full-precision and quantized vectors during -/// pruning. This allows the application to trade full-precision fetches for accuracy. -/// -/// This type implements the following traits: -/// -/// * [`Accessor`] for the [`BfTreeProvider`]. -/// * [`BuildDistanceComputer`] for computing distances among [`distances::pq::Hybrid`] -/// element types. -/// * [`Fill`] for populating a mixture of full-precision and quant vectors. +/// An owned quantized vector that reborrows to [`Opaque`]. /// -pub struct HybridAccessor<'a, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ - provider: &'a BfTreeProvider, +/// Unlike inmem providers (which hand back zero-copy references into a contiguous backing +/// array), bf_tree copies vector data out of the tree on every access. The +/// [`workingset::View`] trait requires `get` to return something that implements +/// `Reborrow<'short, Target = Opaque<'short>>`, so we need an owned type that bridges +/// bf_tree's copy-out model with the working set's reborrow expectation. +pub struct OwnedOpaque(Vec); + +impl<'short> diskann_utils::Reborrow<'short> for OwnedOpaque { + type Target = Opaque<'short>; + fn reborrow(&'short self) -> Self::Target { + Opaque::new(&self.0) + } } -impl<'a, T, D> HybridAccessor<'a, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ - fn new(provider: &'a BfTreeProvider) -> Self { - Self { provider } +impl<'a> From> for OwnedOpaque { + fn from(value: Opaque<'a>) -> Self { + OwnedOpaque(value.to_vec()) } } -impl HasId for HybridAccessor<'_, T, D> +// Pass-through view — reads quantized vectors directly from the provider. +impl workingset::View for &QuantAccessor<'_, T> where T: VectorRepr, - D: AsyncFriendly, { - type Id = u32; -} - -impl<'a, T, D> DelegateNeighbor<'a> for HybridAccessor<'_, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ - type Delegate = &'a NeighborProvider; - fn delegate_neighbor(&'a mut self) -> Self::Delegate { - self.provider.neighbors() - } -} - -impl Accessor for HybridAccessor<'_, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ - /// The [`distances::pq::Hybrid`] is an enum consisting of either a full-precision - /// vector or a quantized vector. - /// - /// This accessor can return either. + type ElementRef<'a> = Opaque<'a>; type Element<'a> - = distances::pq::Hybrid, Vec> - where - Self: 'a; - - /// The generalized reference form of `Element`. - type ElementRef<'a> = distances::pq::Hybrid<&'a [T], &'a [u8]>; - - // Choose to panic on an out-of-bounds access rather than propagate an error. - type GetError = Panics; - - /// The default behavior of `get_element` returns a full-precision vector. The - /// implementation of [`Fill`] is how the `max_fp_vecs_per_fill` is used - /// - fn get_element( - &mut self, - id: Self::Id, - ) -> impl Future, Self::GetError>> + Send { - // SAFETY: We've decided to live with UB that can result from potentially mixing - // unsynchronized reads and writes on the underlying memory. - #[allow(clippy::expect_used)] - std::future::ready(Ok(distances::pq::Hybrid::Full( - self.provider - .full_vectors - .get_vector_sync(id.into_usize()) - .expect("Full vector provider failed to retrieve element"), - ))) - } -} - -impl BuildDistanceComputer for HybridAccessor<'_, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ - type DistanceComputerError = ANNError; - type DistanceComputer = distances::pq::HybridComputer; - - fn build_distance_computer( - &self, - ) -> Result { - let metric = self.provider.quant_vectors.metric(); - Ok(distances::pq::HybridComputer::new( - self.provider.quant_vectors.distance_computer(), - T::distance(metric, Some(self.provider.full_vectors.dim())), - )) - } -} - -impl workingset::Fill> for HybridAccessor<'_, T, D> -where - T: VectorRepr, - D: AsyncFriendly, -{ - type Error = ANNError; - type View<'a> - = distances::pq::View<'a, T, u8> + = OwnedOpaque where Self: 'a; - async fn fill<'a, Itr>( - &'a mut self, - state: &'a mut distances::pq::HybridMap, - itr: Itr, - ) -> Result, Self::Error> - where - Itr: ExactSizeIterator + Clone + Send + Sync, - Self: 'a, - { - let map = state.get_mut(); - map.prepare(itr.clone()); - let threshold = self.provider.max_fp_vecs_per_fill; - itr.enumerate().try_for_each(|(i, id)| -> ANNResult<()> { - match map.entry(id) { - workingset::map::Entry::Seeded(_) => {} - workingset::map::Entry::Occupied(occupied) => { - if i < threshold && !occupied.get().is_full() { - *occupied.into_mut() = distances::pq::Hybrid::Full( - self.provider - .full_vectors - .get_vector_sync(id.into_usize())?, - ); - } - } - workingset::map::Entry::Vacant(vacant) => { - let element = if i < threshold { - let vec = self - .provider - .full_vectors - .get_vector_sync(id.into_usize())?; - - distances::pq::Hybrid::Full(vec) - } else { - let vec = self - .provider - .quant_vectors - .get_vector_sync(id.into_usize())?; - - distances::pq::Hybrid::Quant(vec) - }; - - vacant.insert(element); - } + fn get(&self, id: u32) -> Option> { + match self.provider.quant_vectors.get_vector_sync(id.into_usize()) { + Ok(v) => Some(OwnedOpaque(v)), + Err(RankedError::Transient(_)) => None, + Err(RankedError::Error(_)) => { + // View::get returns Option — can't propagate; treat as missing. + None } - Ok(()) - })?; - - Ok(map.view()) + } } } @@ -1409,129 +1141,46 @@ where /// Perform a search entirely in the full-precision space. /// /// Starting points are not filtered out of the final results. -impl SearchStrategy, &[T]> for FullPrecision +impl SearchStrategy, &[T]> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly + DeletionCheck, { type QueryComputer = T::QueryDistance; - type SearchAccessor<'a> = FullAccessor<'a, T, Q, D>; - type SearchAccessorError = Panics; + type SearchAccessor<'a> = FullAccessor<'a, T, Q>; + type SearchAccessorError = Infallible; fn search_accessor<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, ) -> Result, Self::SearchAccessorError> { Ok(FullAccessor::new(provider)) } } -impl DefaultPostProcessor, &[T]> for FullPrecision +impl DefaultPostProcessor, &[T]> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly + DeletionCheck, -{ - default_post_processor!(glue::Pipeline); -} - -/// An [`glue::SearchPostProcess`] implementation that reranks PQ vectors. -#[derive(Debug, Default, Clone, Copy)] -pub struct Rerank; - -impl<'a, T, D> glue::SearchPostProcess, &[T]> for Rerank -where - T: VectorRepr, - D: AsyncFriendly + DeletionCheck, -{ - type Error = Panics; - - fn post_process( - &self, - accessor: &mut QuantAccessor<'a, T, D>, - query: &[T], - _computer: &pq::distance::QueryComputer>, - candidates: I, - output: &mut B, - ) -> impl Future> + Send - where - I: Iterator>, - B: SearchOutputBuffer + ?Sized, - { - let provider = &accessor.provider; - let checker = accessor.as_deletion_check(); - let f = T::distance(provider.metric, Some(provider.full_vectors.dim())); - - // Filter before computing the full precision distances. - let mut reranked: Vec<(u32, f32)> = candidates - .filter_map(|n| { - if checker.deletion_check(n.id) { - None - } else { - #[allow(clippy::expect_used)] - let vec = provider - .full_vectors - .get_vector_sync(n.id.into_usize()) - .expect("Full vector provider failed to retrieve element"); - Some((n.id, f.evaluate_similarity(query, &vec))) - } - }) - .collect(); - - // Sort the full precision distances. - reranked - .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); - // Store the reranked results. - std::future::ready(Ok(output.extend(reranked))) - } -} - -/// Perform a search entirely in the quantized space. -impl SearchStrategy, &[T]> for Hybrid -where - T: VectorRepr, - D: AsyncFriendly + DeletionCheck, { - type QueryComputer = pq::distance::QueryComputer>; - type SearchAccessor<'a> = QuantAccessor<'a, T, D>; - type SearchAccessorError = Panics; - - fn search_accessor<'a>( - &'a self, - provider: &'a BfTreeProvider, - _context: &'a DefaultContext, - ) -> Result, Self::SearchAccessorError> { - Ok(QuantAccessor::new(provider)) - } -} - -/// Starting points are filtered out of the final results and results are reranked using -/// the full-precision data. -impl DefaultPostProcessor, &[T]> for Hybrid -where - T: VectorRepr, - D: AsyncFriendly + DeletionCheck, -{ - default_post_processor!(glue::Pipeline); + default_post_processor!(glue::Pipeline); } // Pruning -impl PruneStrategy> for FullPrecision +impl PruneStrategy> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly, { type WorkingSet = map::Map, map::Ref<[T]>>; type DistanceComputer<'a> = T::Distance; - type PruneAccessor<'a> = FullAccessor<'a, T, Q, D>; + type PruneAccessor<'a> = FullAccessor<'a, T, Q>; type PruneAccessorError = diskann::error::Infallible; fn prune_accessor<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, ) -> Result, Self::PruneAccessorError> { Ok(FullAccessor::new(provider)) @@ -1542,45 +1191,10 @@ where } } -impl PruneStrategy> for Hybrid -where - T: VectorRepr, - D: AsyncFriendly, -{ - type WorkingSet = distances::pq::HybridMap; - type DistanceComputer<'a> = distances::pq::HybridComputer; - type PruneAccessor<'a> = HybridAccessor<'a, T, D>; - type PruneAccessorError = diskann::error::Infallible; - - fn prune_accessor<'a>( - &'a self, - provider: &'a BfTreeProvider, - _context: &'a DefaultContext, - ) -> Result, Self::PruneAccessorError> { - Ok(HybridAccessor::new(provider)) - } - - fn create_working_set(&self, capacity: usize) -> Self::WorkingSet { - distances::pq::HybridMap::with_capacity(capacity) - } -} - -impl InsertStrategy, &[T]> for FullPrecision +impl InsertStrategy, &[T]> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly + DeletionCheck, -{ - type PruneStrategy = Self; - fn prune_strategy(&self) -> Self::PruneStrategy { - *self - } -} - -impl InsertStrategy, &[T]> for Hybrid -where - T: VectorRepr, - D: AsyncFriendly + DeletionCheck, { type PruneStrategy = Self; fn prune_strategy(&self) -> Self::PruneStrategy { @@ -1588,11 +1202,10 @@ where } } -impl MultiInsertStrategy, B> for FullPrecision +impl MultiInsertStrategy, B> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly + DeletionCheck, B: for<'a> Batch = &'a [T]> + Debug, { type Seed = map::Builder>; @@ -1606,7 +1219,7 @@ where fn finish( &self, - _provider: &BfTreeProvider, + _provider: &BfTreeProvider, _ctx: &DefaultContext, batch: &std::sync::Arc, ids: Itr, @@ -1620,50 +1233,19 @@ where } } -impl MultiInsertStrategy, B> for Hybrid -where - T: VectorRepr, - D: AsyncFriendly + DeletionCheck, - B: for<'a> Batch = &'a [T]> + Debug, -{ - type Seed = distances::pq::Overlay; - type WorkingSet = distances::pq::HybridMap; - type FinishError = diskann::error::Infallible; - type InsertStrategy = Self; - - fn insert_strategy(&self) -> Self::InsertStrategy { - *self - } - - fn finish( - &self, - _provider: &BfTreeProvider, - _ctx: &DefaultContext, - batch: &std::sync::Arc, - ids: Itr, - ) -> impl std::future::Future> + Send - where - Itr: ExactSizeIterator + Send, - { - let overlay = Self::Seed::from_batch(batch.clone(), ids); - std::future::ready(Ok(overlay)) - } -} - /// Inplace Delete /// -impl InplaceDeleteStrategy> for FullPrecision +impl InplaceDeleteStrategy> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, - D: AsyncFriendly + DeletionCheck, { - type DeleteElementError = Panics; + type DeleteElementError = ANNError; type DeleteElement<'a> = &'a [T]; type DeleteElementGuard = Box<[T]>; type PruneStrategy = Self; - type DeleteSearchAccessor<'a> = FullAccessor<'a, T, Q, D>; - type SearchPostProcessor = RemoveDeletedIdsAndCopy; + type DeleteSearchAccessor<'a> = FullAccessor<'a, T, Q>; + type SearchPostProcessor = CopyIds; type SearchStrategy = Self; fn search_strategy(&self) -> Self::SearchStrategy { Self @@ -1674,35 +1256,110 @@ where } fn search_post_processor(&self) -> Self::SearchPostProcessor { - RemoveDeletedIdsAndCopy + CopyIds } async fn get_delete_element<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, id: u32, ) -> Result { - #[allow(clippy::expect_used)] + use diskann::error::ErrorExt; let elt = provider .full_vectors .get_vector_sync(id.into_usize()) - .expect("Failed to get delete element") + .escalate("delete target must exist")? .into(); Ok(elt) } } -impl InplaceDeleteStrategy> for Hybrid +/// Perform a search entirely in the quantized space. +/// +/// Starting points are not filtered out of the final results. +impl SearchStrategy, &[T]> for Quantized +where + T: VectorRepr, +{ + type QueryComputer = UnwrapErr< + spherical_iface::QueryComputer, + spherical_iface::QueryDistanceError, + >; + type SearchAccessor<'a> = QuantAccessor<'a, T>; + type SearchAccessorError = ANNError; + + fn search_accessor<'a>( + &'a self, + provider: &'a BfTreeProvider, + _context: &'a DefaultContext, + ) -> Result, Self::SearchAccessorError> { + Ok(QuantAccessor::new(provider)) + } +} + +impl DefaultPostProcessor, &[T]> for Quantized +where + T: VectorRepr, +{ + default_post_processor!(glue::Pipeline); +} + +impl InsertStrategy, &[T]> for Quantized +where + T: VectorRepr, +{ + type PruneStrategy = Self; + fn prune_strategy(&self) -> Self::PruneStrategy { + *self + } +} + +impl MultiInsertStrategy, B> for Quantized where T: VectorRepr, - D: AsyncFriendly + DeletionCheck, + B: glue::Batch, + Self: for<'a> InsertStrategy< + BfTreeProvider, + B::Element<'a>, + PruneStrategy = Self, + >, { - type DeleteElementError = Panics; + type Seed = map::Builder>; + type WorkingSet = Map; + type FinishError = diskann::error::Infallible; + type InsertStrategy = Self; + + fn insert_strategy(&self) -> Self::InsertStrategy { + *self + } + + fn finish( + &self, + _provider: &BfTreeProvider, + _ctx: &DefaultContext, + _batch: &std::sync::Arc, + _ids: Itr, + ) -> impl std::future::Future> + Send + where + Itr: ExactSizeIterator + Send, + { + let builder = map::Builder::new(map::Capacity::Default); + std::future::ready(Ok(builder)) + } +} + +/// Inplace Delete +/// +impl InplaceDeleteStrategy> for Quantized +where + T: VectorRepr, +{ + type DeleteElementError = ANNError; type DeleteElement<'a> = &'a [T]; type DeleteElementGuard = Box<[T]>; type PruneStrategy = Self; - type DeleteSearchAccessor<'a> = QuantAccessor<'a, T, D>; + type DeleteSearchAccessor<'a> = QuantAccessor<'a, T>; type SearchPostProcessor = Rerank; type SearchStrategy = Self; fn search_strategy(&self) -> Self::SearchStrategy { @@ -1719,17 +1376,92 @@ where async fn get_delete_element<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, id: u32, ) -> Result { - #[allow(clippy::expect_used)] - let elt = provider + use diskann::error::ErrorExt; + provider .full_vectors .get_vector_sync(id.into_usize()) - .expect("Failed to get delete element") - .into(); - Ok(elt) + .escalate("delete target must exist") + .map(Into::into) + } +} + +// Pruning +impl PruneStrategy> for Quantized +where + T: VectorRepr, +{ + type WorkingSet = Map; + type DistanceComputer<'a> = + UnwrapErr; + type PruneAccessor<'a> = QuantAccessor<'a, T>; + type PruneAccessorError = diskann::error::Infallible; + + fn prune_accessor<'a>( + &'a self, + provider: &'a BfTreeProvider, + _context: &'a DefaultContext, + ) -> Result, Self::PruneAccessorError> { + Ok(QuantAccessor::new(provider)) + } + + fn create_working_set(&self, capacity: usize) -> Self::WorkingSet { + map::Builder::new(map::Capacity::Default).build(capacity) + } +} + +/// Post-processor that reranks quantized search results using full-precision distances. +#[derive(Debug, Default, Clone, Copy)] +pub struct Rerank; + +impl<'a, T> glue::SearchPostProcess, &[T]> for Rerank +where + T: VectorRepr, +{ + type Error = ANNError; + + fn post_process( + &self, + accessor: &mut QuantAccessor<'a, T>, + query: &[T], + _computer: &UnwrapErr< + spherical_iface::QueryComputer, + spherical_iface::QueryDistanceError, + >, + candidates: I, + output: &mut B, + ) -> impl Future> + Send + where + I: Iterator> + Send, + B: SearchOutputBuffer + Send + ?Sized, + { + use diskann::error::ErrorExt; + let provider = accessor.provider; + let f = T::distance(provider.metric, Some(provider.full_vectors.dim())); + + let mut reranked: Vec<(u32, f32)> = Vec::new(); + for n in candidates { + match provider + .full_vectors + .get_vector_sync(n.id.into_usize()) + .allow_transient("stale candidate during rerank") + { + Ok(Some(vec)) => { + reranked.push((n.id, f.evaluate_similarity(query, &vec))); + } + Ok(None) => { + // Transient (deleted/missing) — skip this candidate. + } + Err(e) => return std::future::ready(Err(e)), + } + } + + reranked + .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); + std::future::ready(Ok(output.extend(reranked))) } } @@ -1759,8 +1491,6 @@ impl BfTreeParams { #[derive(Serialize, Deserialize, Clone)] pub struct QuantParams { - pub num_pq_bytes: usize, - pub max_fp_vecs_per_fill: usize, pub params_quant: BfTreeParams, } @@ -1863,9 +1593,9 @@ impl BfTreePaths { format!("{}_delete.bin", prefix) } - /// Returns the path for the PQ pivots file - pub fn pq_pivots_bin(prefix: &str) -> String { - format!("{}_pq_pivots.bin", prefix) + /// Returns the path for the spherical quantizer data file + pub fn quant_data_bin(prefix: &str) -> String { + format!("{}_quant_data.bin", prefix) } } @@ -1920,9 +1650,11 @@ fn load_bftree( } } -// SaveWith/LoadWith for BfTreeProvider with TableDeleteProviderAsync +////////////////////// +// Serialization // +////////////////////// -impl SaveWith for BfTreeProvider +impl SaveWith for BfTreeProvider where T: VectorRepr, { @@ -1951,19 +1683,17 @@ where max_record_size: self.neighbor_provider.config().get_cb_max_record_size(), leaf_page_size: self.neighbor_provider.config().get_leaf_page_size(), }, - quant_params: None, // No quantization parameters + quant_params: None, graph_params: self.graph_params.clone(), is_memory: self.full_vectors.config().is_memory_backend(), }; - // All stores must use the same storage backend. debug_assert_eq!( self.full_vectors.config().is_memory_backend(), self.neighbor_provider.config().is_memory_backend(), "Vector and neighbor stores have mismatched storage backends" ); - // Save only essential parameters as JSON { let params_filename = BfTreePaths::params_json(&saved_params.prefix); let params_json = serde_json::to_string(&saved_params).map_err(|e| { @@ -1973,7 +1703,6 @@ where params_writer.write_all(params_json.as_bytes())?; } - // Save vectors and neighbors save_bftree( self.full_vectors.bftree(), BfTreePaths::vectors_bftree(&saved_params.prefix), @@ -1985,19 +1714,11 @@ where ) .await?; - // Save delete bitmap - { - let filename = BfTreePaths::delete_bin(&saved_params.prefix); - let bitmap_bytes = self.deleted.to_bytes(); - let mut writer = storage.create_for_write(&filename)?; - writer.write_all(&bitmap_bytes)?; - } - Ok(0) } } -impl LoadWith for BfTreeProvider +impl LoadWith for BfTreeProvider where T: VectorRepr, { @@ -2007,7 +1728,6 @@ where where P: StorageReadProvider, { - // Read SavedParams from JSON file let saved_params: SavedParams = { let params_filename = BfTreePaths::params_json(prefix); let mut params_reader = storage.open_reader(¶ms_filename)?; @@ -2016,9 +1736,8 @@ where serde_json::from_str(¶ms_json).map_err(|e| { ANNError::log_index_error(format!("Failed to deserialize params: {}", e)) })? - }; // params_reader is dropped here + }; - // Convert metric string back to Metric enum let metric = Metric::from_str(&saved_params.metric) .map_err(|e| ANNError::log_index_error(format!("Failed to parse metric: {}", e)))?; @@ -2042,34 +1761,17 @@ where let neighbor_provider = NeighborProvider::::new_from_bftree(saved_params.max_degree, adjacency_list_index); - // Load delete bitmap - let total_points = saved_params.max_points + saved_params.frozen_points.get(); - let filename = BfTreePaths::delete_bin(&saved_params.prefix); - - let deleted = if storage.exists(&filename) { - let mut reader = storage.open_reader(&filename)?; - let mut bitmap_bytes = Vec::new(); - reader.read_to_end(&mut bitmap_bytes)?; - TableDeleteProviderAsync::from_bytes(&bitmap_bytes, total_points) - .map_err(|e| ANNError::log_index_error(e))? - } else { - // If file doesn't exist, create a new empty delete provider - TableDeleteProviderAsync::new(total_points) - }; - Ok(Self { quant_vectors: NoStore, full_vectors, neighbor_provider, - deleted, - max_fp_vecs_per_fill: 0, metric, graph_params: saved_params.graph_params, }) } } -impl SaveWith for BfTreeProvider +impl SaveWith for BfTreeProvider where T: VectorRepr, { @@ -2099,8 +1801,6 @@ where leaf_page_size: self.neighbor_provider.config().get_leaf_page_size(), }, quant_params: Some(QuantParams { - num_pq_bytes: self.quant_vectors.pq_chunks(), - max_fp_vecs_per_fill: self.max_fp_vecs_per_fill, params_quant: BfTreeParams { bytes: self.quant_vectors.config().get_cb_size_byte(), max_record_size: self.quant_vectors.config().get_cb_max_record_size(), @@ -2111,7 +1811,6 @@ where is_memory: self.full_vectors.config().is_memory_backend(), }; - // All stores must use the same storage backend. debug_assert_eq!( self.full_vectors.config().is_memory_backend(), self.neighbor_provider.config().is_memory_backend(), @@ -2123,7 +1822,6 @@ where "Vector and quant stores have mismatched storage backends" ); - // Save only essential parameters as JSON { let params_filename = BfTreePaths::params_json(&saved_params.prefix); let params_json = serde_json::to_string(&saved_params).map_err(|e| { @@ -2133,7 +1831,6 @@ where params_writer.write_all(params_json.as_bytes())?; } - // Save vectors, neighbors, and quant vectors save_bftree( self.full_vectors.bftree(), BfTreePaths::vectors_bftree(&saved_params.prefix), @@ -2150,32 +1847,20 @@ where ) .await?; - // Save PQ table metadata and data using PQStorage format - let filename = BfTreePaths::pq_pivots_bin(&saved_params.prefix); - let pq_storage = PQStorage::new(&filename, "", None); - let pq_table = &self.quant_vectors.pq_chunk_table; - pq_storage.write_pivot_data( - pq_table.get_pq_table(), - None, - pq_table.get_chunk_offsets(), - NUM_PQ_CENTROIDS, - pq_table.get_dim(), - storage, - )?; - - // Save delete bitmap - { - let filename = BfTreePaths::delete_bin(&saved_params.prefix); - let bitmap_bytes = self.deleted.to_bytes(); - let mut writer = storage.create_for_write(&filename)?; - writer.write_all(&bitmap_bytes)?; - } + let filename = BfTreePaths::quant_data_bin(&saved_params.prefix); + let serialized = self + .quant_vectors + .quantizer + .serialize(GlobalAllocator) + .map_err(|e| ANNError::log_index_error(format!("{e}")))?; + let mut writer = storage.create_for_write(&filename)?; + writer.write_all(&serialized)?; Ok(0) } } -impl LoadWith for BfTreeProvider +impl LoadWith for BfTreeProvider where T: VectorRepr, { @@ -2185,7 +1870,6 @@ where where P: StorageReadProvider, { - // Read SavedParams from JSON file let saved_params: SavedParams = { let params_filename = BfTreePaths::params_json(prefix); let mut params_reader = storage.open_reader(¶ms_filename)?; @@ -2194,14 +1878,12 @@ where serde_json::from_str(¶ms_json).map_err(|e| { ANNError::log_index_error(format!("Failed to deserialize params: {}", e)) })? - }; // params_reader is dropped here + }; - // Extract quant_params - required for quantized provider let quant_params = saved_params.quant_params.ok_or_else(|| { ANNError::log_index_error("Missing quant_params in saved params for quantized provider") })?; - // Convert metric string back to Metric enum let metric = Metric::from_str(&saved_params.metric) .map_err(|e| ANNError::log_index_error(format!("Failed to parse metric: {}", e)))?; @@ -2225,46 +1907,24 @@ where let neighbor_provider = NeighborProvider::::new_from_bftree(saved_params.max_degree, adjacency_list_index); - // Read PQ table from file using PQStorage format - let filename = BfTreePaths::pq_pivots_bin(&saved_params.prefix); - let pq_storage = PQStorage::new(&filename, "", None); - let pq_table = - pq_storage.load_pq_pivots_bin(&filename, quant_params.num_pq_bytes, storage)?; + let filename = BfTreePaths::quant_data_bin(&saved_params.prefix); + let mut reader = storage.open_reader(&filename)?; + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes)?; + let quantizer: Poly = try_deserialize(&bytes, GlobalAllocator) + .map_err(|e| ANNError::log_index_error(format!("{e}")))?; - let quant_vector_index = load_bftree( - &quant_params.params_quant, - BfTreePaths::quant_bftree(&saved_params.prefix), - saved_params.is_memory, - )?; - let quant_vectors = QuantVectorProvider::new_from_bftree( - metric, - saved_params.max_points, - saved_params.frozen_points.get(), - pq_table.clone(), - quant_vector_index, - ); - - // Load delete bitmap - let total_points = saved_params.max_points + saved_params.frozen_points.get(); - let filename = BfTreePaths::delete_bin(&saved_params.prefix); - - let deleted = if storage.exists(&filename) { - let mut reader = storage.open_reader(&filename)?; - let mut bitmap_bytes = Vec::new(); - reader.read_to_end(&mut bitmap_bytes)?; - TableDeleteProviderAsync::from_bytes(&bitmap_bytes, total_points) - .map_err(|e| ANNError::log_index_error(e))? - } else { - // If file doesn't exist, create a new empty delete provider - TableDeleteProviderAsync::new(total_points) - }; + let quant_vector_index = load_bftree( + &quant_params.params_quant, + BfTreePaths::quant_bftree(&saved_params.prefix), + saved_params.is_memory, + )?; + let quant_vectors = QuantVectorProvider::new_from_bftree(quantizer, quant_vector_index); Ok(Self { quant_vectors, full_vectors, neighbor_provider, - deleted, - max_fp_vecs_per_fill: quant_params.max_fp_vecs_per_fill, metric, graph_params: saved_params.graph_params, }) @@ -2285,28 +1945,320 @@ where /// #[cfg(test)] mod tests { + use std::sync::Arc; + use super::*; - use crate::model::graph::provider::async_::common::TableBasedDeletes; - use crate::storage::file_storage_provider::FileStorageProvider; + use crate::quant::create_test_quantizer; + use diskann::{ + graph::DiskANNIndex, + graph::{self, search::Knn}, + neighbor::BackInserter, + }; + use diskann_providers::storage::FileStorageProvider; + use diskann_utils::views::{Init, Matrix}; + + fn create_quant_index() -> Arc>> { + let start_point = Matrix::new(Init(|| 0.0f32), 1, 5); + let dim = 5; + let max_degree = 8; + let metric = Metric::L2; + + let provider = BfTreeProvider::new( + BfTreeProviderParameters { + max_points: 20, + num_start_points: NonZeroUsize::new(1).unwrap(), + dim, + metric, + max_degree, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + }, + start_point.as_view(), + create_test_quantizer(5), + ) + .unwrap(); + + let index_config = graph::config::Builder::new_with( + 4, + graph::config::MaxDegree::new(max_degree as usize), + 10, + metric.into(), + |_| {}, + ) + .build() + .unwrap(); + + Arc::new(DiskANNIndex::new(index_config, provider, None)) + } + + #[tokio::test] + async fn test_quantized_index_search() { + let index = create_quant_index(); + let ctx = &DefaultContext; + + for i in 0..15 { + let point = vec![i as f32; 5]; + index + .insert(Quantized, ctx, &i, point.as_slice()) + .await + .unwrap(); + } + + let query = vec![3.0; 5]; + let params = Knn::new(5, 10, None).unwrap(); + + let mut neighbors = vec![Neighbor::::default(); 5]; + let res = index + .search( + params, + &Quantized, + &DefaultContext, + query.as_slice(), + &mut BackInserter::new(neighbors.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!( + res.result_count, 5, + "there are 15 points and we're asking for 5, we expect 5" + ); + assert_eq!(neighbors[0].id, 3); + } + + #[tokio::test] + async fn test_quantized_index_multi_insert_search() { + let index = create_quant_index(); + let ctx = &DefaultContext; + + let mut counter = 0.0f32; + let data = Matrix::new( + Init(move || { + counter += 1.0; + counter + }), + 15, + 5, + ); + let ids: Arc<[u32]> = (0u32..15).collect::>().into(); + let batch: Arc> = Arc::new(data); + index + .multi_insert::>(Quantized, ctx, batch, ids) + .await + .unwrap(); + + let query = vec![3.0; 5]; + let params = Knn::new(5, 10, None).unwrap(); + + let mut neighbors = vec![Neighbor::::default(); 5]; + let res = index + .search( + params, + &Quantized, + &DefaultContext, + query.as_slice(), + &mut BackInserter::new(neighbors.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!( + res.result_count, 5, + "there are 15 points and we're asking for 5, we expect 5" + ); + } + + #[tokio::test] + async fn test_quantized_delete_and_search() { + let index = create_quant_index(); + let ctx = &DefaultContext; + + for i in 0..15 { + let point = vec![i as f32; 5]; + index + .insert(Quantized, ctx, &i, point.as_slice()) + .await + .unwrap(); + } + + index + .inplace_delete(Quantized, ctx, &2u32, 2, graph::InplaceDeleteMethod::OneHop) + .await + .unwrap(); + index + .inplace_delete(Quantized, ctx, &4u32, 2, graph::InplaceDeleteMethod::OneHop) + .await + .unwrap(); + + let query = vec![3.0; 5]; + let params = Knn::new(5, 10, None).unwrap(); + + let mut neighbors = vec![Neighbor::::default(); 5]; + let res = index + .search( + params, + &Quantized, + &DefaultContext, + query.as_slice(), + &mut BackInserter::new(neighbors.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!(res.result_count, 5); + let neighbor_ids: Vec = neighbors.iter().map(|n| n.id).collect(); + assert!(!neighbor_ids.contains(&2u32)); + assert!(!neighbor_ids.contains(&4u32)); + } + + fn create_full_precision_index() -> Arc>> { + let start_point = Matrix::new(Init(|| 0.0f32), 1, 5); + let max_degree = 8; + let metric = Metric::L2; + + let provider = BfTreeProvider::new( + BfTreeProviderParameters { + max_points: 20, + num_start_points: NonZeroUsize::new(1).unwrap(), + dim: 5, + metric, + max_degree, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + }, + start_point.as_view(), + NoStore, + ) + .unwrap(); + + let index_config = graph::config::Builder::new_with( + 4, + graph::config::MaxDegree::new(max_degree as usize), + 10, + metric.into(), + |_| {}, + ) + .build() + .unwrap(); + + Arc::new(DiskANNIndex::new(index_config, provider, None)) + } + + #[tokio::test] + async fn test_full_precision_index_search() { + let index = create_full_precision_index(); + let ctx = &DefaultContext; + + for i in 0u32..15 { + let point = vec![i as f32; 5]; + index + .insert(FullPrecision, ctx, &i, point.as_slice()) + .await + .unwrap(); + } + + let query = vec![3.0; 5]; + let params = Knn::new(5, 10, None).unwrap(); + + let mut neighbors = vec![Neighbor::::default(); 5]; + let res = index + .search( + params, + &FullPrecision, + &DefaultContext, + query.as_slice(), + &mut BackInserter::new(neighbors.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!( + res.result_count, 5, + "there are 15 points and we're asking for 5, we expect 5" + ); + assert_eq!(neighbors[0].id, 3); + } + + #[tokio::test] + async fn test_full_precision_delete_and_search() { + let index = create_full_precision_index(); + let ctx = &DefaultContext; + + for i in 0u32..15 { + let point = vec![i as f32; 5]; + index + .insert(FullPrecision, ctx, &i, point.as_slice()) + .await + .unwrap(); + } + + index + .inplace_delete( + FullPrecision, + ctx, + &2u32, + 2, + graph::InplaceDeleteMethod::OneHop, + ) + .await + .unwrap(); + index + .inplace_delete( + FullPrecision, + ctx, + &4u32, + 2, + graph::InplaceDeleteMethod::OneHop, + ) + .await + .unwrap(); + + let query = vec![3.0; 5]; + let params = Knn::new(5, 10, None).unwrap(); + + let mut neighbors = vec![Neighbor::::default(); 5]; + let res = index + .search( + params, + &FullPrecision, + &DefaultContext, + query.as_slice(), + &mut BackInserter::new(neighbors.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!(res.result_count, 5); + let neighbor_ids: Vec = neighbors.iter().map(|n| n.id).collect(); + assert!(!neighbor_ids.contains(&2u32)); + assert!(!neighbor_ids.contains(&4u32)); + } #[tokio::test] async fn test_data_provider_and_delete_interface() { let ctx = &DefaultContext; - let provider = BfTreeProvider::new_empty( + let num_start_points = 2; + let dim = 5; + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points, dim); + + let provider = BfTreeProvider::new( BfTreeProviderParameters { max_points: 10, - num_start_points: NonZeroUsize::new(2).unwrap(), - dim: 5, + num_start_points: NonZeroUsize::new(num_start_points).unwrap(), + dim, metric: Metric::L2, - max_fp_vecs_per_fill: None, max_degree: 64, vector_provider_config: Config::default(), quant_vector_provider_config: Config::default(), neighbor_list_provider_config: Config::default(), graph_params: None, }, + start_points.as_view(), NoStore, - TableBasedDeletes, ) .unwrap(); @@ -2315,6 +2267,13 @@ mod tests { assert_eq!((&provider).into_iter(), 0..(10 + 2)); let iter = provider.iter(); + + // Insert vectors so they exist in the bf_tree (hard-delete checks presence) + for i in iter.clone() { + let vector: Vec = (0..5).map(|j| (i * 5 + j) as f32).collect(); + provider.set_element(ctx, &i, &vector).await.unwrap(); + } + for i in iter.clone() { assert_eq!(provider.to_external_id(ctx, i).unwrap(), i); assert_eq!(provider.to_internal_id(ctx, &i).unwrap(), i); @@ -2340,56 +2299,22 @@ mod tests { ); } - // Call `release` to "undelete" it ID. - // + // With hard deletes, `release` is a no-op (data is permanently removed). + // Verify that released IDs remain deleted. for i in iter.clone() { - // set adjacency list to non-empty before release - provider - .neighbor_provider - .set_neighbors(i, &[1, 2]) - .unwrap(); provider.release(ctx, i).await.unwrap(); assert_eq!( provider.status_by_internal_id(ctx, i).await.unwrap(), - ElementStatus::Valid - ); - assert_eq!( - provider.status_by_external_id(ctx, &i).await.unwrap(), - ElementStatus::Valid - ); - // check that adjacency list was reset after release - let mut neighbors = AdjacencyList::new(); - provider - .neighbor_provider - .get_neighbors(i, &mut neighbors) - .unwrap(); - assert!(neighbors.to_vec().is_empty()); - - // Put it back to "deleted" to test `clear`. - // - provider.delete(ctx, &i).await.unwrap(); - } - - provider.clear_delete_set(); - for i in iter.clone() { - assert_eq!( - provider.status_by_internal_id(ctx, i).await.unwrap(), - ElementStatus::Valid - ); - assert_eq!( - provider.status_by_external_id(ctx, &i).await.unwrap(), - ElementStatus::Valid + ElementStatus::Deleted ); } // out-of-bound set-element fails. // - assert!( - provider - .set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0]) - .await - .is_err() - ); + assert!(provider + .set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0]) + .await + .is_err()); } /// This functionality test targets scenarios of empty neighbor lists and ensures: @@ -2401,21 +2326,25 @@ mod tests { async fn test_empty_neighbor_list() { let num_points = 100u32; let ctx = &DefaultContext; - let provider = BfTreeProvider::::new_empty( + + let num_start_points = 2; + let dim = 3; + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points, dim); + + let provider = BfTreeProvider::::new( BfTreeProviderParameters { max_points: num_points as usize, - num_start_points: NonZeroUsize::new(2).unwrap(), - dim: 3, + num_start_points: NonZeroUsize::new(num_start_points).unwrap(), + dim, metric: Metric::L2, - max_fp_vecs_per_fill: None, max_degree: 64, vector_provider_config: Config::default(), quant_vector_provider_config: Config::default(), neighbor_list_provider_config: Config::default(), graph_params: None, }, + start_points.as_view(), NoStore, - TableBasedDeletes, ) .unwrap(); @@ -2428,9 +2357,10 @@ mod tests { let vector = vec![i as f32, (i + 1) as f32, (i + 2) as f32]; provider.set_element(ctx, &i, &vector).await.unwrap(); - // First attempt should fail as NotFound + // First attempt should return empty let mut out = AdjacencyList::new(); - assert!(neighbor_accessor.get_neighbors(i, &mut out).await.is_err()); + neighbor_accessor.get_neighbors(i, &mut out).await.unwrap(); + assert!(out.is_empty()); // After we set the empty neighbor list, our attempt should succeed neighbor_accessor.set_neighbors(i, &[]).await.unwrap(); @@ -2465,12 +2395,10 @@ mod tests { let mut out = AdjacencyList::from_iter_untrusted([10, 20, 30, 40, 50, 60, 70, 80, 90, 100]); // len = 10 // Attempt to access non-existant vector's neighbor list should fail as NotFound - assert!( - neighbor_accessor - .get_neighbors(200, &mut out) - .await - .is_err() - ); + assert!(neighbor_accessor + .get_neighbors(200, &mut out) + .await + .is_err()); assert!(out.is_empty()); } @@ -2517,7 +2445,6 @@ mod tests { num_start_points, dim, metric: Metric::L2, - max_fp_vecs_per_fill: None, max_degree, vector_provider_config: vector_config.clone(), quant_vector_provider_config: Config::default(), @@ -2525,13 +2452,12 @@ mod tests { graph_params: None, }; + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim); + // Create provider - let provider = BfTreeProvider::::new_empty( - params.clone(), - NoStore, - TableBasedDeletes, - ) - .unwrap(); + let provider = + BfTreeProvider::::new(params.clone(), start_points.as_view(), NoStore) + .unwrap(); // Populate provider with vectors for i in 0..num_points { @@ -2554,16 +2480,6 @@ mod tests { .unwrap(); } - // Delete some vectors to test deletion persistence - let deleted_ids = vec![5u32, 10u32, 15u32, 20u32, 25u32]; - for id in &deleted_ids { - provider.delete(ctx, id).await.unwrap(); - assert_eq!( - provider.status_by_internal_id(ctx, *id).await.unwrap(), - ElementStatus::Deleted - ); - } - assert_eq!(vector_config.get_leaf_page_size(), 8192); assert_eq!(vector_config.get_cb_max_record_size(), 1024); @@ -2579,12 +2495,9 @@ mod tests { provider.save_with(&storage, &save_prefix).await.unwrap(); // Load using trait method (includes delete bitmap) - let loaded_provider = BfTreeProvider::::load_with( - &storage, - &save_prefix, - ) - .await - .unwrap(); + let loaded_provider = BfTreeProvider::::load_with(&storage, &save_prefix) + .await + .unwrap(); // Verify vectors for i in 0..num_points as u32 { @@ -2617,31 +2530,6 @@ mod tests { ); } - // Verify deleted status persists across save/load - for id in &deleted_ids { - assert_eq!( - loaded_provider - .status_by_internal_id(ctx, *id) - .await - .unwrap(), - ElementStatus::Deleted, - "Deletion status not preserved for id {}", - id - ); - } - - // Verify non-deleted vectors remain valid - for i in 0..num_points as u32 { - if !deleted_ids.contains(&i) { - assert_eq!( - loaded_provider.status_by_internal_id(ctx, i).await.unwrap(), - ElementStatus::Valid, - "Non-deleted vector {} incorrectly marked as deleted", - i - ); - } - } - // Cleanup is automatic when temp_dir goes out of scope } @@ -2679,13 +2567,8 @@ mod tests { let mut quant_config = Config::new(&quant_path, bytes_quant); quant_config.storage_backend(bf_tree::StorageBackend::Std); - // Create PQ table - let pq_table = FixedChunkPQTable::new( - dim, - vec![0.0; dim * 256].into_boxed_slice(), - Box::new([0, 4, dim]), - ) - .unwrap(); + // Create spherical quantizer + let quantizer = create_test_quantizer(dim); // Create provider parameters let params = BfTreeProviderParameters { @@ -2693,7 +2576,6 @@ mod tests { num_start_points, dim, metric: Metric::L2, - max_fp_vecs_per_fill: Some(10), max_degree, vector_provider_config: vector_config.clone(), quant_vector_provider_config: quant_config.clone(), @@ -2701,14 +2583,14 @@ mod tests { graph_params: None, }; + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim); // Create provider with quantization - let provider = - BfTreeProvider::::new_empty( - params.clone(), - pq_table.clone(), - TableBasedDeletes, - ) - .unwrap(); + let provider = BfTreeProvider::::new( + params.clone(), + start_points.as_view(), + quantizer, + ) + .unwrap(); // Populate provider with vectors for i in 0..num_points { @@ -2731,16 +2613,6 @@ mod tests { .unwrap(); } - // Delete some vectors to test deletion persistence - let deleted_ids = vec![3u32, 8u32, 15u32, 22u32, 30u32]; - for id in &deleted_ids { - provider.delete(ctx, id).await.unwrap(); - assert_eq!( - provider.status_by_internal_id(ctx, *id).await.unwrap(), - ElementStatus::Deleted - ); - } - let storage = FileStorageProvider; // Save to a different prefix to exercise the snapshot copy logic @@ -2754,40 +2626,25 @@ mod tests { // Load using trait method (includes delete bitmap and quantization) let loaded_provider = - BfTreeProvider::::load_with( - &storage, - &save_prefix, - ) - .await - .unwrap(); + BfTreeProvider::::load_with(&storage, &save_prefix) + .await + .unwrap(); - // Verify PQ table - let original_pq = &provider.quant_vectors.pq_chunk_table; - let loaded_pq = &loaded_provider.quant_vectors.pq_chunk_table; - assert_eq!( - original_pq.get_dim(), - loaded_pq.get_dim(), - "PQ table dim mismatch" - ); - assert_eq!( - original_pq.get_num_chunks(), - loaded_pq.get_num_chunks(), - "PQ table num_chunks mismatch" - ); + // Verify quantizer properties match after round-trip assert_eq!( - original_pq.get_num_centers(), - loaded_pq.get_num_centers(), - "PQ table num_centers mismatch" + provider.quant_vectors.quantizer.full_dim(), + loaded_provider.quant_vectors.quantizer.full_dim(), + "Quantizer full_dim mismatch" ); assert_eq!( - original_pq.get_pq_table(), - loaded_pq.get_pq_table(), - "PQ table data mismatch" + provider.quant_vectors.quantizer.bytes(), + loaded_provider.quant_vectors.quantizer.bytes(), + "Quantizer bytes mismatch" ); assert_eq!( - original_pq.get_chunk_offsets(), - loaded_pq.get_chunk_offsets(), - "PQ table chunk_offsets mismatch" + provider.quant_vectors.quantizer.nbits(), + loaded_provider.quant_vectors.quantizer.nbits(), + "Quantizer nbits mismatch" ); // Verify vectors @@ -2831,31 +2688,6 @@ mod tests { ); } - // Verify deleted status persists across save/load - for id in &deleted_ids { - assert_eq!( - loaded_provider - .status_by_internal_id(ctx, *id) - .await - .unwrap(), - ElementStatus::Deleted, - "Deletion status not preserved for id {}", - id - ); - } - - // Verify non-deleted vectors remain valid - for i in 0..num_points as u32 { - if !deleted_ids.contains(&i) { - assert_eq!( - loaded_provider.status_by_internal_id(ctx, i).await.unwrap(), - ElementStatus::Valid, - "Non-deleted vector {} incorrectly marked as deleted", - i - ); - } - } - // Cleanup is automatic when temp_dir goes out of scope } @@ -2868,22 +2700,22 @@ mod tests { let num_start_points = NonZeroUsize::new(1).unwrap(); let ctx = &DefaultContext; + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim); // In-memory config (no file path needed) - let provider = BfTreeProvider::::new_empty( + let provider = BfTreeProvider::::new( BfTreeProviderParameters { max_points: num_points, num_start_points, dim, metric: Metric::L2, - max_fp_vecs_per_fill: None, max_degree, vector_provider_config: Config::default(), quant_vector_provider_config: Config::default(), neighbor_list_provider_config: Config::default(), graph_params: None, }, + start_points.as_view(), NoStore, - TableBasedDeletes, ) .unwrap(); @@ -2921,15 +2753,15 @@ mod tests { provider.save_with(&storage, &save_prefix).await.unwrap(); // Load back - let loaded = BfTreeProvider::::load_with( - &storage, - &save_prefix, - ) - .await - .unwrap(); + let loaded = BfTreeProvider::::load_with(&storage, &save_prefix) + .await + .unwrap(); // Verify vectors for i in 0..num_points as u32 { + if i == 3 || i == 7 { + continue; + } assert_eq!( provider.full_vectors.get_vector_sync(i as usize).unwrap(), loaded.full_vectors.get_vector_sync(i as usize).unwrap(), @@ -2977,32 +2809,26 @@ mod tests { let num_start_points = NonZeroUsize::new(1).unwrap(); let ctx = &DefaultContext; - let pq_table = FixedChunkPQTable::new( - dim, - vec![0.0; dim * 256].into_boxed_slice(), - Box::new([0, 4, dim]), + let quantizer = create_test_quantizer(dim); + + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim); + let provider = BfTreeProvider::::new( + BfTreeProviderParameters { + max_points: num_points, + num_start_points, + dim, + metric: Metric::L2, + max_degree, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + }, + start_points.as_view(), + quantizer, ) .unwrap(); - let provider = - BfTreeProvider::::new_empty( - BfTreeProviderParameters { - max_points: num_points, - num_start_points, - dim, - metric: Metric::L2, - max_fp_vecs_per_fill: Some(5), - max_degree, - vector_provider_config: Config::default(), - quant_vector_provider_config: Config::default(), - neighbor_list_provider_config: Config::default(), - graph_params: None, - }, - pq_table, - TableBasedDeletes, - ) - .unwrap(); - // Populate vectors and neighbors for i in 0..num_points { let vector: Vec = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect(); @@ -3035,16 +2861,15 @@ mod tests { provider.save_with(&storage, &save_prefix).await.unwrap(); // Load back - let loaded = - BfTreeProvider::::load_with( - &storage, - &save_prefix, - ) + let loaded = BfTreeProvider::::load_with(&storage, &save_prefix) .await .unwrap(); - // Verify full vectors + // Verify full vectors (skip deleted id 2) for i in 0..num_points as u32 { + if i == 2 { + continue; + } assert_eq!( provider.full_vectors.get_vector_sync(i as usize).unwrap(), loaded.full_vectors.get_vector_sync(i as usize).unwrap(), @@ -3053,8 +2878,11 @@ mod tests { ); } - // Verify quant vectors + // Verify quant vectors (skip deleted id 2) for i in 0..num_points as u32 { + if i == 2 { + continue; + } assert_eq!( provider.quant_vectors.get_vector_sync(i as usize).unwrap(), loaded.quant_vectors.get_vector_sync(i as usize).unwrap(), @@ -3063,8 +2891,11 @@ mod tests { ); } - // Verify neighbors + // Verify neighbors (skip deleted id 2) for i in 0..num_points as u32 { + if i == 2 { + continue; + } let mut orig = AdjacencyList::new(); let mut load = AdjacencyList::new(); provider diff --git a/diskann-bftree/src/quant.rs b/diskann-bftree/src/quant.rs new file mode 100644 index 000000000..fd36f1fae --- /dev/null +++ b/diskann-bftree/src/quant.rs @@ -0,0 +1,411 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Bf-Tree quant vector provider. + +use crate::{AccessError, AsKey, VectorError, VectorUnavailable}; +use bf_tree::{BfTree, Config}; +use diskann::{error::IntoANNResult, utils::VectorRepr, ANNError, ANNResult}; +use diskann_quantization::{ + alloc::{GlobalAllocator, Poly, ScopedAllocator}, + spherical::iface::{ + DistanceComputer, Opaque, OpaqueMut, Quantizer, QueryComputer, QueryLayout, + }, +}; +use diskann_vector::PreprocessedDistanceFunction; + +use super::ConfigError; +use crate::TestCallCount; + +pub struct QuantQueryComputer(QueryComputer); + +impl QuantQueryComputer { + pub(crate) fn into_inner(self) -> QueryComputer { + self.0 + } +} + +impl PreprocessedDistanceFunction<&[u8], f32> for QuantQueryComputer { + fn evaluate_similarity(&self, x: &[u8]) -> f32 { + self.0 + .evaluate_similarity(Opaque::new(x)) + .expect("spherical query distance failed") + } +} + +pub struct QuantVectorProvider { + quant_vector_index: BfTree, + pub(crate) quantizer: Poly, + pub(super) num_get_calls: TestCallCount, +} + +impl QuantVectorProvider { + pub fn new_with_config(quantizer: Poly, config: Config) -> ANNResult { + let quant_vector_index = BfTree::with_config(config, None).map_err(ConfigError)?; + + Ok(Self { + quant_vector_index, + quantizer, + num_get_calls: TestCallCount::default(), + }) + } + + /// Access the BfTree config + pub(crate) fn config(&self) -> &Config { + self.quant_vector_index.config() + } + + /// Access the underlying BfTree + pub(crate) fn bftree(&self) -> &BfTree { + &self.quant_vector_index + } + + /// Create a new instance from an existing BfTree (for loading from snapshot) + /// + pub(crate) fn new_from_bftree( + quantizer: Poly, + quant_vector_index: BfTree, + ) -> Self { + Self { + quant_vector_index, + quantizer, + num_get_calls: TestCallCount::default(), + } + } + + /// Return the dimension of the full-precision data associated with this provider + pub fn full_dim(&self) -> usize { + self.quantizer.full_dim() + } + + /// Create a query computer for the provided query vector + pub fn query_computer(&self, query: &[T]) -> ANNResult + where + T: VectorRepr, + { + let query_f32 = T::as_f32(query).into_ann_result()?; + let inner = self + .quantizer + .fused_query_computer( + &query_f32, + QueryLayout::FullPrecision, + true, + GlobalAllocator, + ScopedAllocator::global(), + ) + .map_err(|e| ANNError::log_sq_error(e))?; + Ok(QuantQueryComputer(inner)) + } + + /// Create a distance computer for the underlying schema + pub fn distance_computer(&self) -> ANNResult { + self.quantizer + .distance_computer(GlobalAllocator) + .map_err(|e| ANNError::log_sq_error(e)) + } + + pub(crate) fn get_vector_into(&self, i: usize, buffer: &mut [u8]) -> Result<(), AccessError> { + use diskann::ANNErrorKind; + use thiserror::Error; + + let expected = self.quantizer.bytes(); + if buffer.len() != expected { + #[derive(Debug, Error)] + #[error("expected a buffer with dim {0}, instead got {1}")] + struct WrongDim(usize, usize); + + return Err(AccessError::Error(ANNError::new( + ANNErrorKind::IndexError, + WrongDim(expected, buffer.len()), + ))); + } + + self.num_get_calls.increment(); + match self.quant_vector_index.read(i.as_key(), buffer) { + bf_tree::LeafReadResult::Found(read_size) => { + if read_size as usize != expected { + return Err(AccessError::Error(ANNError::log_index_error(format!( + "The bf-tree entry for vector id {} is marked as found but has size {} instead of the expected size {}", + i, read_size, expected, + )))); + } + } + bf_tree::LeafReadResult::Deleted => { + return Err(AccessError::Transient(VectorUnavailable { + id: i, + err: VectorError::Deleted, + })); + } + bf_tree::LeafReadResult::InvalidKey => { + return Err(AccessError::Error(ANNError::log_index_error(format!( + "The bf-tree entry for vector id {} is marked as invalid", + i, + )))); + } + bf_tree::LeafReadResult::NotFound => { + return Err(AccessError::Transient(VectorUnavailable { + id: i, + err: VectorError::NotFound, + })); + } + }; + + Ok(()) + } + + /// Return the quant vector at index `i`. + pub(crate) fn get_vector_sync(&self, i: usize) -> Result, AccessError> { + let mut value = vec![0u8; self.quantizer.bytes()]; + self.get_vector_into(i, &mut value)?; + Ok(value) + } + + /// Compress the vector, `v`, and set the compressed quant vector with Id, `i`, to it + /// + /// Errors if: + /// + /// * `v.dim() != self.full_dim()`: The slice must have the proper length. + /// * PQ compression encounters an error (such as the presence of `NaN`s). + pub(crate) fn set_vector_sync(&self, i: usize, v: &[T]) -> ANNResult<()> + where + T: Copy + VectorRepr, + { + let vf32: &[f32] = &T::as_f32(v).into_ann_result()?; + + if vf32.len() != self.full_dim() { + return Err(ANNError::log_dimension_mismatch_error( + "Vector f32 dimension is not equal to the expected dimension.".to_string(), + )); + } + + // Serialize the key into a byte string, &[u8] + let key = i.as_key(); + + let dim = self.quantizer.bytes(); + let quant_vector = &mut vec![0u8; dim]; + self.quantizer + .compress( + vf32, + OpaqueMut::new(quant_vector), + ScopedAllocator::global(), + ) + .map_err(|e| ANNError::log_sq_error(e))?; + + self.quant_vector_index.insert(key, quant_vector); + + Ok(()) + } + + /// Set the quant vector with Id, `i`, to `v` + /// + /// Errors if: + /// + /// * `v.len() != self.pq_chunks()`: `v` must have the right length. + #[cfg(test)] + pub(crate) fn set_quant_vector(&self, i: usize, v: &[u8]) -> ANNResult<()> { + if v.len() != self.quantizer.bytes() { + return Err(ANNError::log_index_error( + "Vector dimension is not equal to the expected dimension.", + )); + } + + // Update pq vector with id = i to v + let key = i.as_key(); + + self.quant_vector_index.insert(key, v); + + Ok(()) + } + + pub(crate) fn delete_vector(&self, i: usize) { + let key = i.as_key(); + self.quant_vector_index.delete(key); + } +} + +/// Train a spherical quantizer on simple data and return it as a `Poly`. +#[cfg(test)] +pub(crate) fn create_test_quantizer(dim: usize) -> Poly { + use diskann_quantization::{ + algorithms::TransformKind, + alloc::poly, + spherical::{iface, PreScale, SphericalQuantizer, SupportedMetric}, + }; + use diskann_utils::views::Init; + use diskann_utils::views::Matrix; + use rand::{rngs::StdRng, SeedableRng}; + + // Create training data with spread-out values. + let nrows = 8; + let mut counter = 0.0f32; + let data = Matrix::new( + Init(move || { + counter += 0.5; + counter + }), + nrows, + dim, + ); + + let mut rng = StdRng::seed_from_u64(42); + let quantizer = SphericalQuantizer::train( + data.as_view(), + TransformKind::Null, + SupportedMetric::SquaredL2, + PreScale::None, + &mut rng, + GlobalAllocator, + ) + .unwrap(); + + let imp = iface::Impl::<1>::new(quantizer).unwrap(); + poly!(Quantizer, imp, GlobalAllocator).unwrap() +} + +/////////// +// Tests // +/////////// +/// These unit tests target the functionality of Bf-Tree quant vector provider alone +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use diskann::ANNErrorKind; + use diskann_quantization::spherical::iface::Opaque; + use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction}; + use tokio::task::JoinSet; + + use super::*; + + /// Test edge cases of the Bf-Tree quant vector provider + #[tokio::test] + async fn common_errors() { + let dim = 5; + let quantizer = create_test_quantizer(dim); + let quant_bytes = quantizer.bytes(); + + let bf_tree_config = Config::default(); + let provider = QuantVectorProvider::new_with_config(quantizer, bf_tree_config).unwrap(); + + // try to set an out of bounds vector + let result = provider.set_quant_vector(20, &[]).unwrap_err(); + assert_eq!(result.kind(), ANNErrorKind::IndexError); + + // try to set an out of bounds vector via set_vector_sync + let result = provider.set_vector_sync::(20, &[]).unwrap_err(); + assert_eq!(result.kind(), ANNErrorKind::DimensionMismatchError); + + // try to set a quant vector with the wrong dimension + let result = provider.set_quant_vector(0, &[]).unwrap_err(); + assert_eq!(result.kind(), ANNErrorKind::IndexError); + + // verify expected quant vector byte count + assert_eq!(quant_bytes, provider.quantizer.bytes()); + } + + fn create_test_provider() -> QuantVectorProvider { + let dim = 2; + + let quantizer = create_test_quantizer(dim); + + let bf_tree_config = Config::default(); + let provider = QuantVectorProvider::new_with_config(quantizer, bf_tree_config).unwrap(); + + assert_eq!(provider.full_dim(), dim); + + // Set vectors. + provider.set_vector_sync(0, &[-1.5, -1.5]).unwrap(); + provider.set_vector_sync(1, &[-0.5, -0.5]).unwrap(); + provider.set_vector_sync(2, &[0.5, 0.5]).unwrap(); + provider.set_vector_sync(3, &[1.5, 1.5]).unwrap(); + provider.set_vector_sync(4, &[2.5, 2.5]).unwrap(); + provider + } + + /// Test the distance computation functions of the provider + #[tokio::test] + async fn test_similarity_function() { + let provider = create_test_provider(); + let quant_bytes = provider.quantizer.bytes(); + + // Verify compressed vectors are the expected size. + for i in 0..5 { + let v = provider.get_vector_sync(i).unwrap(); + assert_eq!(v.len(), quant_bytes); + } + + // Error checking. + assert!(provider.set_vector_sync(2, &[0.0]).is_err()); + + // Query Computer — verify it returns finite distances. + let c = provider.query_computer(&[-0.5f32, -0.5]).unwrap(); + let dist = c.evaluate_similarity(&provider.get_vector_sync(3).unwrap()); + assert!(dist.is_finite(), "query distance should be finite"); + + // Distance Computer — verify distances between compressed vectors are finite + // and that identical vectors produce zero distance. + let d = provider.distance_computer().unwrap(); + let v0 = provider.get_vector_sync(0).unwrap(); + let v3 = provider.get_vector_sync(3).unwrap(); + let dist = d + .evaluate_similarity(Opaque::new(&v0), Opaque::new(&v3)) + .unwrap(); + assert!(dist.is_finite(), "distance should be finite"); + + // Same vector should have small self-distance (may not be exactly zero + // due to quantization loss, especially at low bit-widths). + let self_dist = d + .evaluate_similarity(Opaque::new(&v0), Opaque::new(&v0)) + .unwrap(); + assert!( + self_dist.abs() < 1.0, + "self-distance should be small, got {}", + self_dist + ); + } + + /// Test the interleaved and parallel traversal of the Bf-Tree + /// by invoking the async accessors of the quant vector provider + #[tokio::test(flavor = "multi_thread", worker_threads = 5)] + async fn test_parallel_tree_traversal() { + let dim = 2; + let quantizer = create_test_quantizer(dim); + + let bf_tree_config = Config::default(); + let provider = + Arc::new(QuantVectorProvider::new_with_config(quantizer, bf_tree_config).unwrap()); + let mut set = JoinSet::new(); + for i in 0..11 { + let vector = vec![i as f32, (i + 1) as f32]; + let provider_clone = Arc::clone(&provider); + set.spawn(async move { provider_clone.set_vector_sync(i as usize, &vector).unwrap() }); + } + + while let Some(res) = set.join_next().await { + res.unwrap(); + } + + // Verify that each vector was stored and can be retrieved with the correct size. + let quant_bytes = provider.quantizer.bytes(); + let mut expected_buf = vec![0u8; quant_bytes]; + + for i in 0..11 { + let stored = provider.get_vector_sync(i).unwrap(); + assert_eq!(stored.len(), quant_bytes); + + // Compress the same input again and verify we get the same output + // (spherical compression is deterministic). + provider + .quantizer + .compress( + &[i as f32, (i + 1) as f32], + OpaqueMut::new(&mut expected_buf), + ScopedAllocator::global(), + ) + .unwrap(); + assert_eq!(stored, expected_buf); + } + } +} diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs b/diskann-bftree/src/vectors.rs similarity index 88% rename from diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs rename to diskann-bftree/src/vectors.rs index 676df50ac..4de978ff2 100644 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/vector_provider.rs +++ b/diskann-bftree/src/vectors.rs @@ -7,16 +7,18 @@ use std::marker::PhantomData; +use crate::{AccessError, AsKey, VectorError, VectorUnavailable}; use bf_tree::{BfTree, Config}; -use bytemuck::{bytes_of, cast_slice}; +use bytemuck::cast_slice; use diskann::{ - ANNError, ANNErrorKind, ANNResult, + error::RankedError, utils::{ErrorToVectorId, TryIntoVectorId, VectorId, VectorRepr}, + ANNError, ANNErrorKind, ANNResult, }; use thiserror::Error; -use super::super::common::TestCallCount; use super::ConfigError; +use crate::TestCallCount; pub struct VectorProvider { dim: usize, @@ -126,7 +128,7 @@ impl VectorProvider { } // Serialize the key, vector_id, into a byte string, &[u8] - let key = bytes_of::(&i); + let key = i.as_key(); let value = cast_slice::(v); self.vector_index.insert(key, value); @@ -134,49 +136,49 @@ impl VectorProvider { Ok(()) } - pub(crate) fn get_vector_into(&self, i: usize, buffer: &mut [T]) -> ANNResult<()> { + pub(crate) fn get_vector_into(&self, i: usize, buffer: &mut [T]) -> Result<(), AccessError> { if buffer.len() != self.dim { #[derive(Debug, Error)] #[error("expected a buffer with dim {0}, instead got {1}")] struct WrongDim(usize, usize); - return Err(ANNError::new( + return Err(RankedError::Error(ANNError::new( ANNErrorKind::IndexError, WrongDim(self.dim(), buffer.len()), - )); + ))); } self.num_get_calls.increment(); match self .vector_index - .read(bytes_of(&i), bytemuck::must_cast_slice_mut::<_, u8>(buffer)) + .read(i.as_key(), bytemuck::must_cast_slice_mut::<_, u8>(buffer)) { bf_tree::LeafReadResult::Found(read_size) => { let vector_size = std::mem::size_of::() * self.dim; if read_size as usize != vector_size { - return Err(ANNError::log_index_error(format!( + return Err(RankedError::Error(ANNError::log_index_error(format!( "The bf-tree entry for vector id {} is marked as found but has size {} instead of the expected size {}", i, read_size, vector_size, - ))); + )))); } } bf_tree::LeafReadResult::Deleted => { - return Err(ANNError::log_index_error(format!( - "The bf-tree entry for vector id {} is marked as deleted", - i - ))); + return Err(RankedError::Transient(VectorUnavailable { + id: i, + err: VectorError::Deleted, + })); } bf_tree::LeafReadResult::InvalidKey => { - return Err(ANNError::log_index_error(format!( + return Err(RankedError::Error(ANNError::log_index_error(format!( "The bf-tree entry for vector id {} is marked as invalid", i - ))); + )))); } bf_tree::LeafReadResult::NotFound => { - return Err(ANNError::log_index_error(format!( - "The bf-tree entry for vector id {} is marked as not found", - i - ))); + return Err(RankedError::Transient(VectorUnavailable { + id: i, + err: VectorError::NotFound, + })); } }; @@ -185,12 +187,17 @@ impl VectorProvider { /// Return the vector at index `i` #[inline(always)] - pub(crate) fn get_vector_sync(&self, i: usize) -> ANNResult> { + pub(crate) fn get_vector_sync(&self, i: usize) -> Result, AccessError> { // Search for the corresponding vector let mut vector = vec![T::default(); self.dim]; self.get_vector_into(i, &mut vector)?; Ok(vector) } + + pub(crate) fn delete_vector(&self, i: usize) { + let key = i.as_key(); + self.vector_index.delete(key); + } } /////////// @@ -241,7 +248,9 @@ mod tests { .unwrap(); assert_eq!(&vector, &vec![(i as f32), (i + 1) as f32, (i + 2) as f32]); } - assert_eq!(vector_provider.num_get_calls.get(), num_points); + if TestCallCount::enabled() { + assert_eq!(vector_provider.num_get_calls.get(), num_points); + } } /// Test other methods and edge cases of the vector provider and sycrhnoization mechanism of Bf-Tree diff --git a/diskann-providers/Cargo.toml b/diskann-providers/Cargo.toml index c5447dca2..4dae263aa 100644 --- a/diskann-providers/Cargo.toml +++ b/diskann-providers/Cargo.toml @@ -34,10 +34,8 @@ diskann-utils = { workspace = true } diskann-quantization = { workspace = true, features = ["rayon"] } tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } tempfile = { workspace = true, optional = true } -bf-tree = { workspace = true, optional = true } prost = "0.14.1" futures-util.workspace = true -serde_json = { workspace = true, optional = true } vfs = { workspace = true, optional = true } [dev-dependencies] @@ -75,7 +73,6 @@ targets = [ default = [] perf_test = ["dep:opentelemetry"] testing = ["dep:tempfile"] -bf_tree = ["dep:bf-tree", "dep:serde_json"] experimental_diversity_search = ["diskann/experimental_diversity_search"] virtual_storage = ["dep:vfs"] diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/mod.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/mod.rs deleted file mode 100644 index 23647de20..000000000 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/mod.rs +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -mod neighbor_provider; -mod provider; -mod quant_vector_provider; -mod vector_provider; - -// Accessors -pub use provider::{ - AsVectorDtype, BfTreePaths, BfTreeProvider, BfTreeProviderParameters, CreateQuantProvider, - FullAccessor, GraphParams, Hidden, Index, QuantAccessor, QuantIndex, StartPoint, VectorDtype, -}; - -pub use bf_tree::Config; - -use diskann::ANNError; - -/// Wrapper around [`bf_tree::ConfigError`] that implements [`std::error::Error`]. -#[derive(Debug, Clone)] -pub struct ConfigError(pub bf_tree::ConfigError); - -impl std::fmt::Display for ConfigError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "BfTree configuration error: {:?}", self.0) - } -} - -impl std::error::Error for ConfigError {} - -impl From for ANNError { - #[track_caller] - #[inline(never)] - fn from(error: ConfigError) -> ANNError { - ANNError::new(diskann::ANNErrorKind::IndexError, error) - } -} diff --git a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs b/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs deleted file mode 100644 index 2a1a3992a..000000000 --- a/diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs +++ /dev/null @@ -1,422 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Bf-Tree quant vector provider. - -use std::sync::Arc; - -use bf_tree::{BfTree, Config}; -use bytemuck::bytes_of; -use diskann::{ANNError, ANNErrorKind, ANNResult, error::IntoANNResult, utils::VectorRepr}; -use diskann_quantization::CompressInto; -use diskann_utils::object_pool::ObjectPool; -use diskann_vector::distance::Metric; -use thiserror::Error; - -use super::super::common::TestCallCount; -use super::ConfigError; -use crate::{ - model::{ - distance::common::distance_table_pool, - pq::{self, FixedChunkPQTable}, - }, - utils::BridgeErr, -}; - -pub struct QuantVectorProvider { - quant_vector_index: BfTree, - max_vectors: usize, - num_start_points: usize, - pub pq_chunk_table: Arc, - metric: Metric, - pub(super) num_get_calls: TestCallCount, - - vec_pool: Arc>>, -} - -type DistanceComputer = pq::distance::DistanceComputer>; -type QueryComputer = pq::distance::QueryComputer>; - -impl QuantVectorProvider { - pub fn new_with_config( - dist_metric: Metric, - max_vectors: usize, - num_start_points: usize, - pq_chunk_table: FixedChunkPQTable, - config: Config, - ) -> ANNResult { - let quant_vector_index = BfTree::with_config(config, None).map_err(ConfigError)?; - let vec_pool = Arc::new(distance_table_pool(&pq_chunk_table)); - - Ok(Self { - max_vectors, - num_start_points, - quant_vector_index, - pq_chunk_table: Arc::new(pq_chunk_table), - metric: dist_metric, - num_get_calls: TestCallCount::default(), - vec_pool, - }) - } - - /// Return the metric associated with this provider - pub(crate) fn metric(&self) -> Metric { - self.metric - } - - /// Access the BfTree config - pub(crate) fn config(&self) -> &Config { - self.quant_vector_index.config() - } - - /// Access the underlying BfTree - pub(crate) fn bftree(&self) -> &BfTree { - &self.quant_vector_index - } - - /// Create a new instance from an existing BfTree (for loading from snapshot) - /// - pub(crate) fn new_from_bftree( - dist_metric: Metric, - max_vectors: usize, - num_start_points: usize, - pq_chunk_table: FixedChunkPQTable, - quant_vector_index: BfTree, - ) -> Self { - let vec_pool = Arc::new(distance_table_pool(&pq_chunk_table)); - Self { - max_vectors, - num_start_points, - quant_vector_index, - pq_chunk_table: Arc::new(pq_chunk_table), - metric: dist_metric, - num_get_calls: TestCallCount::default(), - vec_pool, - } - } - - /// Return the total number of points including starting points - #[inline(always)] - pub fn total(&self) -> usize { - self.max_vectors + self.num_start_points - } - - /// Return the dimension of the full-precision data associated with this provider - pub fn full_dim(&self) -> usize { - self.pq_chunk_table.get_dim() - } - - /// Return the number of PQ chunks in the underlying PQ schema - pub fn pq_chunks(&self) -> usize { - self.pq_chunk_table.get_num_chunks() - } - - /// Create a query computer for the provided query vector - pub fn query_computer(&self, query: &[T]) -> ANNResult - where - T: VectorRepr, - { - QueryComputer::new( - self.pq_chunk_table.clone(), - self.metric, - &T::as_f32(query).into_ann_result()?, - Some(self.vec_pool.clone()), - ) - } - - /// Create a distance computer for the underlying schema - pub fn distance_computer(&self) -> DistanceComputer { - DistanceComputer::new(self.pq_chunk_table.clone(), self.metric) - } - - pub(crate) fn get_vector_into(&self, i: usize, buffer: &mut [u8]) -> ANNResult<()> { - let expected = buffer.len(); - if buffer.len() != expected { - #[derive(Debug, Error)] - #[error("expected a buffer with dim {0}, instead got {1}")] - struct WrongDim(usize, usize); - - return Err(ANNError::new( - ANNErrorKind::IndexError, - WrongDim(expected, buffer.len()), - )); - } - - self.num_get_calls.increment(); - match self.quant_vector_index.read(bytes_of(&i), buffer) { - bf_tree::LeafReadResult::Found(read_size) => { - if read_size as usize != expected { - return ANNResult::Err(ANNError::log_index_error(format!( - "The bf-tree entry for vector id {} is marked as found but has size {} instead of the expected size {}", - i, read_size, expected, - ))); - } - } - bf_tree::LeafReadResult::Deleted => { - return ANNResult::Err(ANNError::log_index_error(format!( - "The bf-tree entry for vector id {} is marked as deleted", - i, - ))); - } - bf_tree::LeafReadResult::InvalidKey => { - return ANNResult::Err(ANNError::log_index_error(format!( - "The bf-tree entry for vector id {} is marked as invalid", - i, - ))); - } - bf_tree::LeafReadResult::NotFound => { - return ANNResult::Err(ANNError::log_index_error(format!( - "The bf-tree entry for vector id {} is marked as not found", - i, - ))); - } - }; - - Ok(()) - } - - /// Return the quant vector at index `i`. - pub(crate) fn get_vector_sync(&self, i: usize) -> ANNResult> { - let mut value = vec![0u8; self.pq_chunks()]; - self.get_vector_into(i, &mut value)?; - Ok(value) - } - - /// Compress the vector, `v`, and set the compressed quant vector with Id, `i`, to it - /// - /// Errors if: - /// - /// * `i > self.total()`: `i` must be in bounds. - /// * `v.dim() != self.full_dim()`: The slice must have the proper length. - /// * PQ compression encounters an error (such as the presence of `NaN`s). - pub(crate) fn set_vector_sync(&self, i: usize, v: &[T]) -> ANNResult<()> - where - T: Copy + VectorRepr, - { - if i >= self.total() { - return Err(ANNError::log_index_error( - "Vector id is out of boundary in the dataset.", - )); - } - - let vf32: &[f32] = &T::as_f32(v).into_ann_result()?; - - if vf32.len() != self.full_dim() { - return Err(ANNError::log_index_error( - "Vector f32 dimension is not equal to the expected dimension.", - )); - } - - // Serialize the key into a byte string, &[u8] - let key = bytes_of::(&i); - - // Quantize the full vector and de-serialize it as byte string - let dim = self.pq_chunk_table.get_num_chunks(); - let quant_vector = &mut vec![0u8; dim]; - - self.pq_chunk_table - .compress_into(vf32, quant_vector) - .bridge_err()?; - - self.quant_vector_index.insert(key, quant_vector); - - Ok(()) - } - - /// Set the quant vecotr with Id, `i``, to `v` - /// - /// Errors if: - /// - /// * `i >= self.total()`: `i` must be in bounds. - /// * `v.len() != self.pq_chunks()`: `v` must have the right length. - #[cfg(test)] - pub(crate) fn set_quant_vector(&self, i: usize, v: &[u8]) -> ANNResult<()> { - if i >= self.total() { - return Err(ANNError::log_index_error( - "Vector id is out of boundary in the dataset.", - )); - } - if v.len() != self.pq_chunks() { - return Err(ANNError::log_index_error( - "Vector dimension is not equal to the expected dimension.", - )); - } - - // Update pq vector with id = i to v - let key = bytes_of::(&i); - - self.quant_vector_index.insert(key, v); - - Ok(()) - } -} - -/////////// -// Tests // -/////////// - -/// These unit tests target the functionality of Bf-Tree quant vector provider alone -#[cfg(test)] -mod tests { - use diskann::ANNErrorKind; - use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; - use tokio::task::JoinSet; - - use super::*; - - /// Test edges cases of the Bf-Tree quant vector provider - #[tokio::test] - async fn common_errors() { - let dim = 5; - let offsets = vec![0, dim]; - let full_pivot_data = vec![0.0; 256 * dim]; - - let pq_chunk_table = - FixedChunkPQTable::new(dim, full_pivot_data.into(), offsets.into()).unwrap(); - - let bf_tree_config = Config::default(); - let provider = - QuantVectorProvider::new_with_config(Metric::L2, 10, 1, pq_chunk_table, bf_tree_config) - .unwrap(); - - // try to set an out of bounds vector - let result = provider.set_quant_vector(20, &[]).unwrap_err(); - assert_eq!(result.kind(), ANNErrorKind::IndexError); - - // SAFETY: We have exclusive ownership of `provider` - let result = provider.set_vector_sync::(20, &[]).unwrap_err(); - assert_eq!(result.kind(), ANNErrorKind::IndexError); - - // try to set a vector with the wrong dimension - let result = provider.set_quant_vector(0, &[]).unwrap_err(); - assert_eq!(result.kind(), ANNErrorKind::IndexError); - } - - fn create_test_provider() -> QuantVectorProvider { - let num_points = 3; - let frozen_points = 2; - let dim = 2; - - // We can create a really simple, 1 chunk PQ table with known entries to allow - // us to easily verify results. - let table = FixedChunkPQTable::new( - dim, - Box::new([0.0, 0.0, 1.0, 1.0, 2.0, 2.0]), - Box::new([0, dim]), - ) - .unwrap(); - - let bf_tree_config = Config::default(); - let provider = QuantVectorProvider::new_with_config( - Metric::L2, - num_points, - frozen_points, - table, - bf_tree_config, - ) - .unwrap(); - - assert_eq!(provider.total(), num_points + frozen_points); - assert_eq!(provider.full_dim(), dim); - - // Set Vector. - provider.set_vector_sync(0, &[-1.5, -1.5]).unwrap(); - provider.set_vector_sync(1, &[-0.5, -0.5]).unwrap(); - provider.set_vector_sync(2, &[0.5, 0.5]).unwrap(); - provider.set_vector_sync(3, &[1.5, 1.5]).unwrap(); - provider.set_vector_sync(4, &[2.5, 2.5]).unwrap(); - provider - } - - /// Test the similarity functions of the provider - #[tokio::test] - async fn test_similarity_function() { - let provider = create_test_provider(); - - // Get Vector. - assert_eq!(provider.get_vector_sync(0).unwrap(), &[0]); - assert_eq!(provider.get_vector_sync(1).unwrap(), &[0]); - assert_eq!(provider.get_vector_sync(2).unwrap(), &[0]); - assert_eq!(provider.get_vector_sync(3).unwrap(), &[1]); - assert_eq!(provider.get_vector_sync(4).unwrap(), &[2]); - - // Error checking. - assert!(provider.set_vector_sync(5, &[0.0, 0.0]).is_err()); - assert!(provider.set_vector_sync(2, &[0.0]).is_err()); - - // Query Computer. - let c = provider.query_computer(&[-0.5, -0.5]).unwrap(); - let expected: f32 = 1.5 * 1.5 * 2.0; - assert_eq!( - c.evaluate_similarity(provider.get_vector_sync(3).unwrap().as_slice()), - expected - ); - - // Distance Computer. - let d = provider.distance_computer(); - assert_eq!( - d.evaluate_similarity( - provider.get_vector_sync(0).unwrap().as_slice(), - provider.get_vector_sync(3).unwrap().as_slice() - ), - 2.0 - ); - - let slice: &[f32] = &[-0.5, -0.5]; - assert_eq!( - d.evaluate_similarity(slice, provider.get_vector_sync(3).unwrap().as_slice()), - expected, - ); - } - - /// Test the interleaved and parallell traversal of the Bf-Tree - /// by invoking the async accessors of the quant vector provider - #[tokio::test(flavor = "multi_thread", worker_threads = 5)] - async fn test_parallel_tree_traversal() { - let dim = 2; - let offsets = vec![0, dim]; - let full_pivot_data = vec![0.0; 256 * dim]; - let pq_chunk_table = - FixedChunkPQTable::new(dim, full_pivot_data.into(), offsets.into()).unwrap(); - - let bf_tree_config = Config::default(); - let provider = Arc::new( - QuantVectorProvider::new_with_config(Metric::L2, 10, 1, pq_chunk_table, bf_tree_config) - .unwrap(), - ); - let mut set = JoinSet::new(); - for i in 0..11 { - let vector = vec![i as f32, (i + 1) as f32]; - let provider_clone = Arc::clone(&provider); - set.spawn(async move { - // One tokio task per vector insertion - provider_clone.set_vector_sync(i as usize, &vector).unwrap() - }); - } - - while let Some(res) = set.join_next().await { - res.unwrap(); - } - - let dim = provider.pq_chunk_table.get_num_chunks(); - let mut quant_vector: Vec = vec![0; dim]; - let quant_vector_ref: &mut [u8] = &mut quant_vector; - - for i in 0..11 { - // SAFETY: We're only accessing one at a time. - let quant_vector = provider.get_vector_sync(i as usize).unwrap(); - match provider - .pq_chunk_table - .compress_into(&[(i as f32), (i + 1) as f32], quant_vector_ref) - { - Ok(_) => {} - Err(e) => { - panic!("{}", e) - } - }; - assert_eq!(&quant_vector_ref, &quant_vector); - } - } -} diff --git a/diskann-providers/src/model/graph/provider/async_/common.rs b/diskann-providers/src/model/graph/provider/async_/common.rs index d24900404..993888c35 100644 --- a/diskann-providers/src/model/graph/provider/async_/common.rs +++ b/diskann-providers/src/model/graph/provider/async_/common.rs @@ -16,6 +16,8 @@ use crate::{ storage::{AsyncIndexMetadata, AsyncQuantLoadContext, LoadWith, SaveWith}, }; +pub use diskann::graph::strategy::{FullPrecision, Quantized}; + /// Represents a range of start points for an index. /// The range includes `start` and excludes `end`. /// `start` is the first valid point, and `end - 1` is the last valid point. @@ -379,19 +381,6 @@ impl CreateDeleteProvider for TableBasedDeletes { } } -/// Operates entirely in full precision. -/// -/// All indexing and search operations use the uncompressed full-precision vectors. -#[derive(Debug, Clone, Copy)] -pub struct FullPrecision; - -/// Operates entirely in the quantized space. -/// -/// All indexing and search operations use quantized vectors. -/// If full-precision vectors are available, they are only used for the final reranking step. -#[derive(Debug, Clone, Copy)] -pub struct Quantized; - /// Operates primarily in the quantized space with selective use of full precision. /// /// # Search diff --git a/diskann-providers/src/model/graph/provider/async_/mod.rs b/diskann-providers/src/model/graph/provider/async_/mod.rs index 46a7199ca..2cd108974 100644 --- a/diskann-providers/src/model/graph/provider/async_/mod.rs +++ b/diskann-providers/src/model/graph/provider/async_/mod.rs @@ -31,7 +31,3 @@ pub use fast_memory_quant_vector_provider::FastMemoryQuantVectorProviderAsync; // The default `inmem` data provider for the async index. pub mod inmem; - -// Bf-tree based data provider for the async index -#[cfg(feature = "bf_tree")] -pub mod bf_tree; diff --git a/diskann-providers/src/model/graph/provider/async_/table_delete_provider.rs b/diskann-providers/src/model/graph/provider/async_/table_delete_provider.rs index 337d1465b..d6903f01a 100644 --- a/diskann-providers/src/model/graph/provider/async_/table_delete_provider.rs +++ b/diskann-providers/src/model/graph/provider/async_/table_delete_provider.rs @@ -53,7 +53,6 @@ impl TableDeleteProviderAsync { self.delete_table[slot].fetch_or(mask, Ordering::AcqRel); } - // private method for now, but may need to become public pub(crate) fn undelete(&self, vector_id: usize) { assert!(vector_id < self.max_size); let slot = vector_id / 32; @@ -68,45 +67,6 @@ impl TableDeleteProviderAsync { } } - /// Serialize the delete bitmap to bytes (little-endian u32 values) - #[cfg(feature = "bf_tree")] - pub(crate) fn to_bytes(&self) -> Vec { - let mut bytes = Vec::with_capacity(std::mem::size_of_val(self.delete_table.as_slice())); - for atomic_val in &self.delete_table { - let val = atomic_val.load(Ordering::Relaxed); - bytes.extend_from_slice(&val.to_le_bytes()); - } - bytes - } - - /// Create a TableDeleteProviderAsync from serialized bytes - #[cfg(feature = "bf_tree")] - pub(crate) fn from_bytes(bytes: &[u8], max_size: usize) -> Result { - let expected_len = max_size.div_ceil(32); - - let (chunks, remainder) = bytes.as_chunks::<{ std::mem::size_of::() }>(); - if chunks.len() != expected_len { - return Err(format!( - "Delete bitmap size mismatch: expected {} u32 values, got {}", - expected_len, - chunks.len() - )); - } - if !remainder.is_empty() { - return Err("Length of bytes is not a multiple of 4".to_string()); - } - - let delete_table = chunks - .iter() - .map(|chunk| AtomicU32::new(u32::from_le_bytes(*chunk))) - .collect(); - - Ok(TableDeleteProviderAsync { - delete_table, - max_size, - }) - } - #[cfg(test)] pub(crate) fn count(&self) -> usize { let mut count = 0; @@ -234,27 +194,4 @@ mod tests { } delete_provider } - - #[cfg(feature = "bf_tree")] - #[test] - fn test_save_load_roundtrip() { - let original = get_test_delete_table_provider(50, &[0, 5, 20, 34, 48]); - let bytes = original.to_bytes(); - let loaded = TableDeleteProviderAsync::from_bytes(&bytes, 50).unwrap(); - - assert_eq!(original.count(), loaded.count()); - for i in 0..50 { - assert_eq!(original.is_deleted(i), loaded.is_deleted(i)); - } - } - - #[cfg(feature = "bf_tree")] - #[test] - fn test_from_bytes_size_mismatch() { - // max_size=50 requires ceil(50/32) = 2 u32 values = 8 bytes - // Provide wrong number of bytes (e.g., 4 bytes = 1 u32) - let bytes = vec![0u8; 4]; // Only 1 u32 worth of bytes - let result = TableDeleteProviderAsync::from_bytes(&bytes, 50); - assert!(result.is_err()); - } } diff --git a/diskann/src/graph/mod.rs b/diskann/src/graph/mod.rs index 40f07857f..1ff868c1f 100644 --- a/diskann/src/graph/mod.rs +++ b/diskann/src/graph/mod.rs @@ -36,6 +36,8 @@ pub use search::{KnnSearchError, RangeSearchError, Search}; mod internal; +pub mod strategy; + // Integration tests and test providers. #[cfg(any(test, feature = "testing"))] pub mod test; diff --git a/diskann/src/graph/strategy.rs b/diskann/src/graph/strategy.rs new file mode 100644 index 000000000..aa89a69ea --- /dev/null +++ b/diskann/src/graph/strategy.rs @@ -0,0 +1,17 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +/// Operates entirely with full precision +/// +/// All indexing and search operations use the uncompressed full-precision vectors. +#[derive(Debug, Clone, Copy)] +pub struct FullPrecision; + +/// Operates entirely in the quantized space +/// +/// All indexing and search operations use quantized vectors. +/// If full-precision vectors are available, they are only used for the final reranking step. +#[derive(Debug, Clone, Copy)] +pub struct Quantized; From 21e81a7eaed56f611b8a7af86b11262016447372 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Tue, 19 May 2026 12:23:43 -0700 Subject: [PATCH 2/7] new_empty and cargo remove publish step --- diskann-bftree/Cargo.toml | 2 ++ diskann-bftree/src/provider.rs | 57 +++++++++++++++++++++++----------- 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/diskann-bftree/Cargo.toml b/diskann-bftree/Cargo.toml index ce7033271..c46848631 100644 --- a/diskann-bftree/Cargo.toml +++ b/diskann-bftree/Cargo.toml @@ -8,6 +8,7 @@ authors.workspace = true documentation.workspace = true license.workspace = true edition.workspace = true +publish = false [dependencies] bf-tree.workspace = true @@ -39,3 +40,4 @@ experimental_diversity_search = ["diskann/experimental_diversity_search"] [lints] workspace = true + diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 2ba804db4..5d784926d 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -232,10 +232,45 @@ impl BfTreeProvider where T: VectorRepr, { + /// Construct a new data provider from empty. Callers of this are required to manually set start + /// points before performing search tasks. + /// + /// This constructor for `BfTreeProvider` should be used when building and constructing from + /// scratch. + /// + /// # Arguments + /// * `params`: An instance of [`BfTreeProviderParameters`] collecting shared + /// configuration information. + /// * `quant_precursor`: A precursor type for the quantizer layer. + /// + /// # Type Constraints + /// * `Self: StartPoint` - The provider must implement the `StartPoint` trait. + pub fn new_empty(params: BfTreeProviderParameters, quant_precursor: TQ) -> ANNResult + where + Self: StartPoint, + TQ: CreateQuantProvider, + { + Ok(Self { + quant_vectors: quant_precursor.create(params.quant_vector_provider_config)?, + full_vectors: VectorProvider::new_with_config( + params.max_points, + params.dim, + params.num_start_points.get(), + params.vector_provider_config, + )?, + neighbor_provider: NeighborProvider::new_with_config( + params.max_degree, + params.neighbor_list_provider_config, + )?, + metric: params.metric, + graph_params: params.graph_params, + }) + } + /// Construct a new data provider with start points initialized. /// - /// This is the primary constructor for `BfTreeProvider`. It creates the provider - /// and sets the start points in one operation. + /// This is the primary constructor for `BfTreeProvider` where start points are known from the + /// beginning. It creates the provider and sets the start points in one operation. /// /// # Arguments /// * `params`: An instance of [`BfTreeProviderParameters`] collecting shared @@ -243,7 +278,6 @@ where /// * `start_points`: A matrix view containing the start point vectors. The number /// of rows must match `params.num_start_points.get()`. /// * `quant_precursor`: A precursor type for the quantizer layer. - /// * `delete_precursor`: A precursor type for the delete layer. /// /// # Type Constraints /// * `Self: StartPoint` - The provider must implement the `StartPoint` trait. @@ -265,21 +299,7 @@ where ))); } - let provider = Self { - quant_vectors: quant_precursor.create(params.quant_vector_provider_config)?, - full_vectors: VectorProvider::new_with_config( - params.max_points, - params.dim, - params.num_start_points.get(), - params.vector_provider_config, - )?, - neighbor_provider: NeighborProvider::new_with_config( - params.max_degree, - params.neighbor_list_provider_config, - )?, - metric: params.metric, - graph_params: params.graph_params, - }; + let provider = Self::new_empty(params.clone(), quant_precursor)?; provider.set_start_points(Hidden(()), start_points)?; { // Initialize all neighborhoods to be empty lists. @@ -914,6 +934,7 @@ where Ok(T::query_distance(from, self.provider.metric)) } } + impl ExpandBeam<&[T]> for FullAccessor<'_, T, Q> where T: VectorRepr, From eb2098eb9df794c789c918da24131340153b7cbb Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Tue, 19 May 2026 14:28:25 -0700 Subject: [PATCH 3/7] fix publish workspace in cargo --- diskann-bftree/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-bftree/Cargo.toml b/diskann-bftree/Cargo.toml index c46848631..16877d1e3 100644 --- a/diskann-bftree/Cargo.toml +++ b/diskann-bftree/Cargo.toml @@ -8,7 +8,7 @@ authors.workspace = true documentation.workspace = true license.workspace = true edition.workspace = true -publish = false +publish.workspace = false [dependencies] bf-tree.workspace = true From da86885175693906b1ab686c2c125633575304c1 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Tue, 19 May 2026 15:00:13 -0700 Subject: [PATCH 4/7] remove the workspace specification from the cargo --- diskann-bftree/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-bftree/Cargo.toml b/diskann-bftree/Cargo.toml index 16877d1e3..c46848631 100644 --- a/diskann-bftree/Cargo.toml +++ b/diskann-bftree/Cargo.toml @@ -8,7 +8,7 @@ authors.workspace = true documentation.workspace = true license.workspace = true edition.workspace = true -publish.workspace = false +publish = false [dependencies] bf-tree.workspace = true From 6a51d78feccb08859a11a5077adf426e36520e6f Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Wed, 20 May 2026 08:04:58 -0700 Subject: [PATCH 5/7] comments from pr --- diskann-bftree/src/provider.rs | 27 ++------------------------- diskann-bftree/src/quant.rs | 1 + 2 files changed, 3 insertions(+), 25 deletions(-) diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 5d784926d..9de6f3ed9 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -27,7 +27,7 @@ use diskann::{ InsertStrategy, MultiInsertStrategy, PruneStrategy, SearchExt, SearchStrategy, }, strategy::{FullPrecision, Quantized}, - workingset::{self, map, Map}, + workingset::{map, Map}, AdjacencyList, SearchOutputBuffer, }, neighbor::Neighbor, @@ -245,7 +245,7 @@ where /// /// # Type Constraints /// * `Self: StartPoint` - The provider must implement the `StartPoint` trait. - pub fn new_empty(params: BfTreeProviderParameters, quant_precursor: TQ) -> ANNResult + fn new_empty(params: BfTreeProviderParameters, quant_precursor: TQ) -> ANNResult where Self: StartPoint, TQ: CreateQuantProvider, @@ -1132,29 +1132,6 @@ impl<'a> From> for OwnedOpaque { } } -// Pass-through view — reads quantized vectors directly from the provider. -impl workingset::View for &QuantAccessor<'_, T> -where - T: VectorRepr, -{ - type ElementRef<'a> = Opaque<'a>; - type Element<'a> - = OwnedOpaque - where - Self: 'a; - - fn get(&self, id: u32) -> Option> { - match self.provider.quant_vectors.get_vector_sync(id.into_usize()) { - Ok(v) => Some(OwnedOpaque(v)), - Err(RankedError::Transient(_)) => None, - Err(RankedError::Error(_)) => { - // View::get returns Option — can't propagate; treat as missing. - None - } - } - } -} - //////////////// // Strategies // //////////////// diff --git a/diskann-bftree/src/quant.rs b/diskann-bftree/src/quant.rs index fd36f1fae..fd196229e 100644 --- a/diskann-bftree/src/quant.rs +++ b/diskann-bftree/src/quant.rs @@ -156,6 +156,7 @@ impl QuantVectorProvider { } /// Return the quant vector at index `i`. + #[cfg(test)] pub(crate) fn get_vector_sync(&self, i: usize) -> Result, AccessError> { let mut value = vec![0u8; self.quantizer.bytes()]; self.get_vector_into(i, &mut value)?; From 1bd362bb2e33ad67da02964e324bb0df19e62253 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Wed, 20 May 2026 14:36:13 -0700 Subject: [PATCH 6/7] change test to uniform initialization of the vectors in the matrix and verify that 1..5 are returned --- diskann-bftree/src/provider.rs | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 9de6f3ed9..6e482d060 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -2031,11 +2031,19 @@ mod tests { let index = create_quant_index(); let ctx = &DefaultContext; - let mut counter = 0.0f32; let data = Matrix::new( - Init(move || { - counter += 1.0; - counter + Init({ + let mut row = 0usize; + let mut col = 0usize; + move || { + let val = row as f32; + col += 1; + if col == 5 { + col = 0; + row += 1; + } + val + } }), 15, 5, @@ -2066,6 +2074,13 @@ mod tests { res.result_count, 5, "there are 15 points and we're asking for 5, we expect 5" ); + let neighbor_ids: Vec = neighbors.iter().map(|n| n.id).collect(); + for expected in 1u32..=5 { + assert!( + neighbor_ids.contains(&expected), + "expected id {expected} in results, got {neighbor_ids:?}" + ); + } } #[tokio::test] From 70da9e84e533aae1984e6ace7b1e5964a9e8cca6 Mon Sep 17 00:00:00 2001 From: Jordan Maples Date: Thu, 21 May 2026 08:25:28 -0700 Subject: [PATCH 7/7] making start points distinct --- diskann-bftree/src/provider.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 6e482d060..1e46af3fa 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -2256,7 +2256,15 @@ mod tests { let ctx = &DefaultContext; let num_start_points = 2; let dim = 5; - let start_points = Matrix::new(Init(|| 0.0f32), num_start_points, dim); + let start_points = Matrix::try_from( + vec![0.0f32; dim] + .into_iter() + .chain(vec![0.5f32; dim]) + .collect::>(), + num_start_points, + dim, + ) + .unwrap(); let provider = BfTreeProvider::new( BfTreeProviderParameters {