use std::{
fmt::Debug,
future::Future,
io::{Read, Write},
num::NonZeroUsize,
str::FromStr,
};
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::{
default_post_processor,
error::{ErrorExt, Infallible, RankedError},
graph::{
glue::{
self, Batch, CopyIds, DefaultPostProcessor, InplaceDeleteStrategy, InsertStrategy,
MultiInsertStrategy, PruneStrategy, SearchStrategy,
},
strategy::{FullPrecision, Quantized},
workingset::map,
AdjacencyList, SearchOutputBuffer,
},
neighbor::Neighbor,
provider::{DataProvider, DefaultContext, Delete, ElementStatus, HasId, NoopGuard, SetElement},
utils::{IntoUsize, VectorRepr},
ANNError, ANNResult,
};
use diskann_utils::{
future::{AsyncFriendly, SendFuture},
views::MatrixView,
};
use diskann_vector::{distance::Metric, DistanceFunction, PreprocessedDistanceFunction};
use super::{
neighbors::{NeighborAccessor, NeighborProvider},
quant::QuantVectorProvider,
vectors::VectorProvider,
AccessError, NoStore,
};
use diskann_providers::model::graph::provider::async_::distances::UnwrapErr;
use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, StorageWriteProvider};
pub struct BfTreeProvider<T, Q = QuantVectorProvider>
where
T: VectorRepr,
{
pub(super) quant_vectors: Q,
pub(super) full_vectors: VectorProvider<T>,
pub(crate) neighbor_provider: NeighborProvider<u32>,
pub(super) metric: Metric,
pub(crate) graph_params: Option<GraphParams>,
}
#[derive(Debug, Clone)]
pub struct BfTreeProviderParameters {
pub max_points: usize,
pub num_start_points: NonZeroUsize,
pub dim: usize,
pub metric: Metric,
pub max_degree: u32,
pub vector_provider_config: Config,
pub quant_vector_provider_config: Config,
pub neighbor_list_provider_config: Config,
pub graph_params: Option<GraphParams>,
}
impl<T, Q> BfTreeProvider<T, Q>
where
T: VectorRepr,
{
fn new_empty<TQ>(params: BfTreeProviderParameters, quant_precursor: TQ) -> ANNResult<Self>
where
Self: StartPoint<T>,
TQ: CreateQuantProvider<Target = Q>,
{
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,
})
}
pub fn new<TQ>(
params: BfTreeProviderParameters,
start_points: MatrixView<'_, T>,
quant_precursor: TQ,
) -> ANNResult<Self>
where
Self: StartPoint<T>,
TQ: CreateQuantProvider<Target = Q>,
{
if start_points.nrows() != params.num_start_points.get() {
return Err(ANNError::log_async_index_error(format!(
"start_points matrix has {} rows, but params.num_start_points is {}",
start_points.nrows(),
params.num_start_points.get(),
)));
}
let provider = Self::new_empty(params.clone(), quant_precursor)?;
provider.set_start_points(Hidden(()), start_points)?;
{
let mut scratch = provider.neighbor_provider.scratch();
for i in 0..params.max_points {
let vector_id = i as u32;
scratch.write_neighbors(vector_id, &[])?;
}
}
Ok(provider)
}
pub fn starting_points(&self) -> ANNResult<Vec<u32>> {
self.full_vectors.starting_points()
}
pub fn iter(&self) -> std::ops::Range<u32> {
0..(self.full_vectors.total() as u32)
}
pub fn num_start_points(&self) -> usize {
self.full_vectors.num_start_points
}
pub fn max_points(&self) -> usize {
self.full_vectors.max_vectors
}
pub fn dim(&self) -> usize {
self.full_vectors.dim()
}
pub fn metric(&self) -> Metric {
self.metric
}
pub fn max_degree(&self) -> u32 {
self.neighbor_provider.max_degree()
}
}
impl<T> BfTreeProvider<T, QuantVectorProvider>
where
T: VectorRepr,
{
pub fn counts_for_get_vector(&self) -> (usize, usize) {
(
self.full_vectors.num_get_calls.get(),
self.quant_vectors.num_get_calls.get(),
)
}
}
impl<T> BfTreeProvider<T, NoStore>
where
T: VectorRepr,
{
pub fn counts_for_get_vector(&self) -> (usize, usize) {
(self.full_vectors.num_get_calls.get(), 0)
}
}
impl<T, Q> Delete for BfTreeProvider<T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
fn release(
&self,
_context: &Self::Context,
_id: Self::InternalId,
) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send {
std::future::ready(Ok(()))
}
fn delete(
&self,
_context: &Self::Context,
gid: &Self::ExternalId,
) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send {
let id = *gid;
self.full_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<Output = Result<diskann::provider::ElementStatus, Self::Error>> + Send
{
self.status_by_internal_id(context, *gid)
}
fn status_by_internal_id(
&self,
_context: &Self::Context,
id: Self::InternalId,
) -> impl std::future::Future<Output = Result<diskann::provider::ElementStatus, Self::Error>> + 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)
}
}
impl<T, Q> IntoIterator for &BfTreeProvider<T, Q>
where
T: VectorRepr,
{
type Item = u32;
type IntoIter = std::ops::Range<u32>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub trait CreateQuantProvider {
type Target;
fn create(self, bf_tree_config: Config) -> ANNResult<Self::Target>;
}
impl CreateQuantProvider for NoStore {
type Target = NoStore;
fn create(self, _bf_tree_config: Config) -> ANNResult<Self::Target> {
Ok(self)
}
}
impl CreateQuantProvider for Poly<dyn Quantizer> {
type Target = QuantVectorProvider;
fn create(self, bf_tree_config: Config) -> ANNResult<Self::Target> {
QuantVectorProvider::new_with_config(self, bf_tree_config)
}
}
impl<T, Q> BfTreeProvider<T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
pub fn neighbors(&self) -> &NeighborProvider<u32> {
&self.neighbor_provider
}
}
impl<T, Q> DataProvider for BfTreeProvider<T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
type Context = DefaultContext;
type InternalId = u32;
type ExternalId = u32;
type Error = ANNError;
type Guard = NoopGuard<u32>;
fn to_internal_id(
&self,
_context: &DefaultContext,
gid: &Self::ExternalId,
) -> Result<Self::InternalId, Self::Error> {
Ok(*gid)
}
fn to_external_id(
&self,
_context: &DefaultContext,
id: Self::InternalId,
) -> Result<Self::ExternalId, Self::Error> {
Ok(id)
}
}
impl<T, Q> HasId for BfTreeProvider<T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
type Id = u32;
}
impl<T> SetElement<&[T]> for BfTreeProvider<T, QuantVectorProvider>
where
T: VectorRepr,
{
type SetError = ANNError;
fn set_element(
&self,
_context: &Self::Context,
id: &u32,
element: &[T],
) -> impl Future<Output = Result<Self::Guard, Self::SetError>> + Send {
if let Err(err) = self.quant_vectors.set_vector_sync(id.into_usize(), element) {
return std::future::ready(Err(err));
}
if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) {
return std::future::ready(Err(err));
}
std::future::ready(Ok(NoopGuard::new(*id)))
}
}
impl<T> SetElement<&[T]> for BfTreeProvider<T, NoStore>
where
T: VectorRepr,
{
type SetError = ANNError;
fn set_element(
&self,
_context: &Self::Context,
id: &u32,
element: &[T],
) -> impl Future<Output = Result<Self::Guard, Self::SetError>> + Send {
if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) {
return std::future::ready(Err(err));
}
std::future::ready(Ok(NoopGuard::new(*id)))
}
}
pub struct Hidden(());
pub trait StartPoint<T> {
#[doc(hidden)]
fn set_start_points(&self, hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()>;
}
impl<T> StartPoint<T> for BfTreeProvider<T, QuantVectorProvider>
where
T: VectorRepr,
{
fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> {
let start_point_ids = self.full_vectors.starting_points()?;
if start_points.nrows() != start_point_ids.len() {
return Err(ANNError::log_async_index_error(format!(
"expected start_points to contain `{}` rows, instead it has {}",
start_point_ids.len(),
start_points.nrows(),
)));
}
let mut scratch = self.neighbor_provider.scratch();
for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) {
self.full_vectors.set_vector_sync(id.into_usize(), v)?;
self.quant_vectors.set_vector_sync(id.into_usize(), v)?;
scratch.write_neighbors(id, &[])?;
}
Ok(())
}
}
impl<T> StartPoint<T> for BfTreeProvider<T, NoStore>
where
T: VectorRepr,
{
fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> {
let start_point_ids = self.full_vectors.starting_points()?;
if start_points.nrows() != start_point_ids.len() {
return Err(ANNError::log_async_index_error(format!(
"expected start_points to contain `{}` rows, instead it has {}",
start_point_ids.len(),
start_points.nrows(),
)));
}
let mut scratch = self.neighbor_provider.scratch();
for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) {
self.full_vectors.set_vector_sync(id.into_usize(), v)?;
scratch.write_neighbors(id, &[])?;
}
Ok(())
}
}
pub struct FullAccessor<'a, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
provider: &'a BfTreeProvider<T, Q>,
computer: T::QueryDistance,
element: Box<[T]>,
}
impl<'a, T, Q> FullAccessor<'a, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
pub(crate) fn new(provider: &'a BfTreeProvider<T, Q>, query: &[T]) -> Self {
Self {
provider,
computer: T::query_distance(query, provider.metric),
element: (0..provider.full_vectors.dim())
.map(|_| T::default())
.collect(),
}
}
fn get_distance(&mut self, id: u32) -> Result<f32, AccessError> {
self.provider
.full_vectors
.get_vector_into(id.into_usize(), &mut self.element)
.map(|_: ()| self.computer.evaluate_similarity(&self.element))
}
}
impl<T, Q> HasId for FullAccessor<'_, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
type Id = u32;
}
impl<T, Q> glue::SearchAccessor for FullAccessor<'_, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
fn starting_points(&self) -> impl Future<Output = ANNResult<Vec<u32>>> {
std::future::ready(self.provider.starting_points())
}
async fn start_point_distances<F>(&mut self, mut f: F) -> ANNResult<()>
where
F: FnMut(Self::Id, f32) + Send,
{
for i in self.provider.starting_points()? {
f(
i,
self.get_distance(i)
.escalate("starting point retrieval must succeed")?,
)
}
Ok(())
}
async fn expand_beam<Itr, P, F>(
&mut self,
ids: Itr,
mut pred: P,
mut on_neighbors: F,
) -> ANNResult<()>
where
Itr: Iterator<Item = Self::Id> + Send,
P: glue::HybridPredicate<Self::Id> + Send + Sync,
F: FnMut(Self::Id, f32) + Send,
{
let mut neighbors = AdjacencyList::new();
for n in ids {
self.provider.neighbors().get_neighbors(n, &mut neighbors)?;
for &id in neighbors.iter().filter(|i| pred.eval_mut(i)) {
if let Some(distance) = self
.get_distance(id)
.allow_transient("skipping deleted vectors")?
{
on_neighbors(id, distance)
}
}
}
Ok(())
}
}
pub struct QuantAccessor<'a, T>
where
T: VectorRepr,
{
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
computer: super::quant::QuantQueryComputer,
element: Box<[u8]>,
}
impl<'a, T> QuantAccessor<'a, T>
where
T: VectorRepr,
{
pub(crate) fn new(
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
query: &[T],
) -> ANNResult<Self> {
let computer = provider.quant_vectors.query_computer(query)?;
Ok(Self {
provider,
computer,
element: (0..provider.quant_vectors.quantizer.bytes())
.map(|_| u8::default())
.collect(),
})
}
fn get_distance(&mut self, id: u32) -> Result<f32, AccessError> {
self.provider
.quant_vectors
.get_vector_into(id.into_usize(), &mut self.element)
.map(|_: ()| self.computer.evaluate_similarity(&self.element))
}
}
impl<T> HasId for QuantAccessor<'_, T>
where
T: VectorRepr,
{
type Id = u32;
}
impl<T> glue::SearchAccessor for QuantAccessor<'_, T>
where
T: VectorRepr,
{
fn starting_points(&self) -> impl Future<Output = ANNResult<Vec<u32>>> {
std::future::ready(self.provider.starting_points())
}
async fn start_point_distances<F>(&mut self, mut f: F) -> ANNResult<()>
where
F: FnMut(Self::Id, f32) + Send,
{
for i in self.provider.starting_points()? {
f(
i,
self.get_distance(i)
.escalate("starting point retrieval must succeed")?,
)
}
Ok(())
}
async fn expand_beam<Itr, P, F>(
&mut self,
ids: Itr,
mut pred: P,
mut on_neighbors: F,
) -> ANNResult<()>
where
Itr: Iterator<Item = Self::Id> + Send,
P: glue::HybridPredicate<Self::Id> + Send + Sync,
F: FnMut(Self::Id, f32) + Send,
{
let mut neighbors = AdjacencyList::new();
for n in ids {
self.provider.neighbors().get_neighbors(n, &mut neighbors)?;
for &id in neighbors.iter().filter(|i| pred.eval_mut(i)) {
if let Some(distance) = self
.get_distance(id)
.allow_transient("skipping deleted vectors")?
{
on_neighbors(id, distance)
}
}
}
Ok(())
}
}
pub struct FullPruneAccessor<'a, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
provider: &'a BfTreeProvider<T, Q>,
neighbors: NeighborAccessor<'a, u32>,
set: map::Map<u32, Box<[T]>, map::Ref<[T]>>,
distance: T::Distance,
}
impl<'a, T, Q> FullPruneAccessor<'a, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
fn new(
provider: &'a BfTreeProvider<T, Q>,
set: map::Map<u32, Box<[T]>, map::Ref<[T]>>,
) -> Self {
Self {
provider,
neighbors: provider.neighbor_provider.scratch(),
set,
distance: T::distance(provider.metric, Some(provider.full_vectors.dim())),
}
}
}
impl<T, Q> HasId for FullPruneAccessor<'_, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
type Id = u32;
}
impl<'q, T, Q> glue::PruneAccessor for FullPruneAccessor<'q, T, Q>
where
T: VectorRepr,
Q: AsyncFriendly,
{
type ElementRef<'a> = &'a [T];
type View<'a>
= map::View<'a, u32, Box<[T]>, map::Ref<[T]>>
where
Self: 'a;
type Distance<'a>
= &'a T::Distance
where
Self: 'a;
type Neighbors<'a>
= diskann::provider::Neighbors<'a, NeighborAccessor<'q, u32>>
where
Self: 'a;
fn fill<Itr>(
&mut self,
itr: Itr,
) -> impl SendFuture<ANNResult<(Self::View<'_>, Self::Distance<'_>)>>
where
Itr: ExactSizeIterator<Item = Self::Id> + Clone + Send + Sync,
{
let mut buf: Option<Box<[T]>> = None;
let view = self.set.fill(itr, |i: u32| -> ANNResult<_> {
let mut b = match buf.take() {
Some(b) => b,
None => std::iter::repeat_n(T::default(), self.provider.dim()).collect(),
};
match self
.provider
.full_vectors
.get_vector_into(i.into_usize(), &mut b)
.allow_transient("transient errors allowed during fill")?
{
Some(()) => Ok(Some(b)),
None => {
buf = Some(b);
Ok(None)
}
}
});
let result = view.map(|v| (v, &self.distance));
std::future::ready(result)
}
fn neighbors(&mut self) -> Self::Neighbors<'_> {
diskann::provider::Neighbors(&mut self.neighbors)
}
}
pub struct QuantPruneAccessor<'a, T>
where
T: VectorRepr,
{
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
neighbors: NeighborAccessor<'a, u32>,
set: map::Map<u32, Owned>,
distance: UnwrapErr<spherical_iface::DistanceComputer, spherical_iface::DistanceError>,
}
impl<'a, T> QuantPruneAccessor<'a, T>
where
T: VectorRepr,
{
fn new(
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
capacity: usize,
) -> ANNResult<Self> {
let distance = provider
.quant_vectors
.distance_computer()
.map(UnwrapErr::new)?;
let set = map::Builder::new(map::Capacity::Default).build(capacity);
Ok(Self {
provider,
neighbors: provider.neighbor_provider.scratch(),
set,
distance,
})
}
}
impl<T> HasId for QuantPruneAccessor<'_, T>
where
T: VectorRepr,
{
type Id = u32;
}
impl<'q, T> glue::PruneAccessor for QuantPruneAccessor<'q, T>
where
T: VectorRepr,
{
type ElementRef<'a> = Opaque<'a>;
type View<'a>
= map::View<'a, u32, Owned>
where
Self: 'a;
type Distance<'a>
= &'a UnwrapErr<spherical_iface::DistanceComputer, spherical_iface::DistanceError>
where
Self: 'a;
type Neighbors<'a>
= diskann::provider::Neighbors<'a, NeighborAccessor<'q, u32>>
where
Self: 'a;
fn fill<Itr>(
&mut self,
itr: Itr,
) -> impl SendFuture<ANNResult<(Self::View<'_>, Self::Distance<'_>)>>
where
Itr: ExactSizeIterator<Item = Self::Id> + Clone + Send + Sync,
{
let mut buf: Option<Box<[u8]>> = None;
let bytes = self.provider.quant_vectors.quantizer.bytes();
let view = self.set.fill(itr, |i: u32| -> ANNResult<_> {
let mut b = match buf.take() {
Some(b) => b,
None => std::iter::repeat_n(0, bytes).collect(),
};
match self
.provider
.quant_vectors
.get_vector_into(i.into_usize(), &mut b)
.allow_transient("transient errors allowed during fill")?
{
Some(()) => Ok(Some(Owned(b))),
None => {
buf = Some(b);
Ok(None)
}
}
});
let result = view.map(|v| (v, &self.distance));
std::future::ready(result)
}
fn neighbors(&mut self) -> Self::Neighbors<'_> {
diskann::provider::Neighbors(&mut self.neighbors)
}
}
pub struct Owned(Box<[u8]>);
impl<'short> diskann_utils::Reborrow<'short> for Owned {
type Target = Opaque<'short>;
fn reborrow(&'short self) -> Self::Target {
Opaque::new(&self.0)
}
}
impl<'a, T, Q> SearchStrategy<'a, BfTreeProvider<T, Q>, &'a [T]> for FullPrecision
where
T: VectorRepr,
Q: AsyncFriendly,
{
type SearchAccessor = FullAccessor<'a, T, Q>;
type SearchAccessorError = Infallible;
fn search_accessor(
&'a self,
provider: &'a BfTreeProvider<T, Q>,
_context: &'a DefaultContext,
query: &'a [T],
) -> Result<Self::SearchAccessor, Self::SearchAccessorError> {
Ok(FullAccessor::new(provider, query))
}
}
impl<'a, T, Q> DefaultPostProcessor<'a, BfTreeProvider<T, Q>, &'a [T]> for FullPrecision
where
T: VectorRepr,
Q: AsyncFriendly,
{
default_post_processor!(glue::Pipeline<glue::FilterStartPoints, CopyIds>);
}
impl<T, Q> PruneStrategy<BfTreeProvider<T, Q>> for FullPrecision
where
T: VectorRepr,
Q: AsyncFriendly,
{
type PruneAccessor<'a> = FullPruneAccessor<'a, T, Q>;
type PruneAccessorError = diskann::error::Infallible;
fn prune_accessor<'a>(
&'a self,
provider: &'a BfTreeProvider<T, Q>,
_context: &'a DefaultContext,
capacity: usize,
) -> Result<Self::PruneAccessor<'a>, Self::PruneAccessorError> {
let set = map::Builder::new(map::Capacity::Default).build(capacity);
Ok(FullPruneAccessor::new(provider, set))
}
}
impl<'a, T, Q> InsertStrategy<'a, BfTreeProvider<T, Q>, &'a [T]> for FullPrecision
where
T: VectorRepr,
Q: AsyncFriendly,
{
type PruneStrategy = Self;
fn prune_strategy(&self) -> Self::PruneStrategy {
*self
}
}
impl<T, Q, B> MultiInsertStrategy<BfTreeProvider<T, Q>, B> for FullPrecision
where
T: VectorRepr,
Q: AsyncFriendly,
B: for<'a> Batch<Element<'a> = &'a [T]> + Debug,
{
type Seed = map::Builder<u32, map::Ref<[T]>>;
type FinishError = diskann::error::Infallible;
type PruneStrategy = Self;
type InsertStrategy = Self;
fn insert_strategy(&self) -> Self::InsertStrategy {
*self
}
fn finish<Itr>(
&self,
_provider: &BfTreeProvider<T, Q>,
_ctx: &DefaultContext,
batch: &std::sync::Arc<B>,
ids: Itr,
) -> impl std::future::Future<Output = Result<Self::Seed, Self::FinishError>> + Send
where
Itr: ExactSizeIterator<Item = u32> + Send,
{
let overlay = map::Overlay::from_batch(batch.clone(), ids);
let builder = map::Builder::new(map::Capacity::Default).with_overlay(overlay);
std::future::ready(Ok(builder))
}
fn seeded_prune_accessor<'a>(
&'a self,
provider: &'a BfTreeProvider<T, Q>,
_context: &'a DefaultContext,
seed: &'a Self::Seed,
capacity: usize,
) -> ANNResult<FullPruneAccessor<'a, T, Q>> {
let set = seed.clone().build(capacity);
Ok(FullPruneAccessor::new(provider, set))
}
}
impl<T, Q> InplaceDeleteStrategy<BfTreeProvider<T, Q>> for FullPrecision
where
T: VectorRepr,
Q: AsyncFriendly,
{
type DeleteElementError = ANNError;
type DeleteElement<'a> = &'a [T];
type DeleteElementGuard = Box<[T]>;
type PruneStrategy = Self;
type DeleteSearchAccessor<'a> = FullAccessor<'a, T, Q>;
type SearchPostProcessor = CopyIds;
type SearchStrategy = Self;
fn search_strategy(&self) -> Self::SearchStrategy {
Self
}
fn prune_strategy(&self) -> Self::PruneStrategy {
Self
}
fn search_post_processor(&self) -> Self::SearchPostProcessor {
CopyIds
}
async fn get_delete_element<'a>(
&'a self,
provider: &'a BfTreeProvider<T, Q>,
_context: &'a DefaultContext,
id: u32,
) -> Result<Self::DeleteElementGuard, Self::DeleteElementError> {
use diskann::error::ErrorExt;
let elt = provider
.full_vectors
.get_vector_sync(id.into_usize())
.escalate("get_delete_element: failed to read vector for inplace delete")?
.into();
Ok(elt)
}
}
impl<'a, T> SearchStrategy<'a, BfTreeProvider<T, QuantVectorProvider>, &'a [T]> for Quantized
where
T: VectorRepr,
{
type SearchAccessor = QuantAccessor<'a, T>;
type SearchAccessorError = ANNError;
fn search_accessor(
&'a self,
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
_context: &'a DefaultContext,
query: &'a [T],
) -> Result<Self::SearchAccessor, Self::SearchAccessorError> {
QuantAccessor::new(provider, query)
}
}
impl<'a, T> DefaultPostProcessor<'a, BfTreeProvider<T, QuantVectorProvider>, &'a [T]> for Quantized
where
T: VectorRepr,
{
default_post_processor!(glue::Pipeline<glue::FilterStartPoints, Rerank>);
}
impl<'a, T> InsertStrategy<'a, BfTreeProvider<T, QuantVectorProvider>, &'a [T]> for Quantized
where
T: VectorRepr,
{
type PruneStrategy = Self;
fn prune_strategy(&self) -> Self::PruneStrategy {
*self
}
}
impl<T, B> MultiInsertStrategy<BfTreeProvider<T, QuantVectorProvider>, B> for Quantized
where
T: VectorRepr,
B: glue::Batch,
B: for<'a> Batch<Element<'a> = &'a [T]> + Debug,
{
type Seed = ();
type FinishError = diskann::error::Infallible;
type PruneStrategy = Self;
type InsertStrategy = Self;
fn insert_strategy(&self) -> Self::InsertStrategy {
*self
}
fn finish<Itr>(
&self,
_provider: &BfTreeProvider<T, QuantVectorProvider>,
_ctx: &DefaultContext,
_batch: &std::sync::Arc<B>,
_ids: Itr,
) -> impl std::future::Future<Output = Result<Self::Seed, Self::FinishError>> + Send
where
Itr: ExactSizeIterator<Item = u32> + Send,
{
std::future::ready(Ok(()))
}
fn seeded_prune_accessor<'a>(
&'a self,
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
_context: &'a DefaultContext,
_seed: &'a (),
capacity: usize,
) -> ANNResult<QuantPruneAccessor<'a, T>> {
QuantPruneAccessor::new(provider, capacity)
}
}
impl<T> InplaceDeleteStrategy<BfTreeProvider<T, QuantVectorProvider>> 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>;
type SearchPostProcessor = Rerank;
type SearchStrategy = Self;
fn search_strategy(&self) -> Self::SearchStrategy {
*self
}
fn prune_strategy(&self) -> Self::PruneStrategy {
*self
}
fn search_post_processor(&self) -> Self::SearchPostProcessor {
Rerank
}
async fn get_delete_element<'a>(
&'a self,
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
_context: &'a DefaultContext,
id: u32,
) -> Result<Self::DeleteElementGuard, Self::DeleteElementError> {
use diskann::error::ErrorExt;
provider
.full_vectors
.get_vector_sync(id.into_usize())
.escalate("get_delete_element: failed to read vector for inplace delete")
.map(Into::into)
}
}
impl<T> PruneStrategy<BfTreeProvider<T, QuantVectorProvider>> for Quantized
where
T: VectorRepr,
{
type PruneAccessor<'a> = QuantPruneAccessor<'a, T>;
type PruneAccessorError = ANNError;
fn prune_accessor<'a>(
&'a self,
provider: &'a BfTreeProvider<T, QuantVectorProvider>,
_context: &'a DefaultContext,
capacity: usize,
) -> Result<Self::PruneAccessor<'a>, Self::PruneAccessorError> {
QuantPruneAccessor::new(provider, capacity)
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct Rerank;
impl<'a, T> glue::SearchPostProcess<QuantAccessor<'a, T>, &[T]> for Rerank
where
T: VectorRepr,
{
type Error = ANNError;
fn post_process<I, B>(
&self,
accessor: &mut QuantAccessor<'a, T>,
query: &[T],
candidates: I,
output: &mut B,
) -> impl Future<Output = Result<usize, Self::Error>> + Send
where
I: Iterator<Item = Neighbor<u32>> + Send,
B: SearchOutputBuffer<u32> + 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) => {
}
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)))
}
}
#[derive(Serialize, Deserialize, Clone)]
pub struct BfTreeParams {
pub bytes: usize,
pub max_record_size: usize,
pub leaf_page_size: usize,
}
impl BfTreeParams {
pub fn to_config(&self, path: &std::path::Path, is_memory: bool) -> Config {
let mut config = Config::new(path, self.bytes);
config.cb_max_record_size(self.max_record_size);
config.leaf_page_size(self.leaf_page_size);
if is_memory {
config.storage_backend(bf_tree::StorageBackend::Memory);
} else {
config.storage_backend(bf_tree::StorageBackend::Std);
}
config
}
}
#[derive(Serialize, Deserialize, Clone)]
pub struct QuantParams {
pub params_quant: BfTreeParams,
}
#[derive(Serialize, Deserialize, Clone)]
pub struct SavedParams {
pub max_points: usize,
pub frozen_points: NonZeroUsize,
pub dim: usize,
pub metric: String,
pub max_degree: u32,
pub prefix: String,
pub params_vector: BfTreeParams,
pub params_neighbor: BfTreeParams,
pub quant_params: Option<QuantParams>,
pub graph_params: Option<GraphParams>,
pub is_memory: bool,
}
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum VectorDtype {
F32,
F16,
U8,
I8,
}
pub trait AsVectorDtype {
const DATA_TYPE: VectorDtype;
}
impl AsVectorDtype for f32 {
const DATA_TYPE: VectorDtype = VectorDtype::F32;
}
impl AsVectorDtype for half::f16 {
const DATA_TYPE: VectorDtype = VectorDtype::F16;
}
impl AsVectorDtype for i8 {
const DATA_TYPE: VectorDtype = VectorDtype::I8;
}
impl AsVectorDtype for u8 {
const DATA_TYPE: VectorDtype = VectorDtype::U8;
}
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct GraphParams {
pub l_build: usize,
pub alpha: f32,
pub backedge_ratio: f32,
pub vector_dtype: VectorDtype,
}
pub struct BfTreePaths;
impl BfTreePaths {
pub fn params_json(prefix: &str) -> String {
format!("{}_params.json", prefix)
}
pub fn vectors_bftree(prefix: &str) -> std::path::PathBuf {
std::path::PathBuf::from(format!("{}_vectors.bftree", prefix))
}
pub fn neighbors_bftree(prefix: &str) -> std::path::PathBuf {
std::path::PathBuf::from(format!("{}_neighbors.bftree", prefix))
}
pub fn quant_bftree(prefix: &str) -> std::path::PathBuf {
std::path::PathBuf::from(format!("{}_quant.bftree", prefix))
}
pub fn delete_bin(prefix: &str) -> String {
format!("{}_delete.bin", prefix)
}
pub fn quant_data_bin(prefix: &str) -> String {
format!("{}_quant_data.bin", prefix)
}
}
async fn copy_snapshot_if_needed(
snapshot_path: std::path::PathBuf,
target_path: std::path::PathBuf,
) -> ANNResult<()> {
if snapshot_path != target_path {
tokio::task::spawn_blocking(move || {
std::fs::copy(&snapshot_path, &target_path).map_err(|e| {
ANNError::log_index_error(format!(
"Failed to copy snapshot from {:?} to {:?}: {}",
snapshot_path, target_path, e
))
})
})
.await
.map_err(|e| ANNError::log_index_error(format!("Blocking copy task failed: {}", e)))??;
}
Ok(())
}
async fn save_bftree(tree: &BfTree, target_path: std::path::PathBuf) -> ANNResult<()> {
if tree.config().is_memory_backend() {
tree.snapshot_memory_to_disk(&target_path);
} else {
let snapshot_path = tree.snapshot();
copy_snapshot_if_needed(snapshot_path, target_path).await?;
}
Ok(())
}
fn load_bftree(
params: &BfTreeParams,
snapshot_path: std::path::PathBuf,
is_memory: bool,
) -> Result<BfTree, ANNError> {
let config = params.to_config(&snapshot_path, is_memory);
if is_memory {
BfTree::new_from_snapshot_disk_to_memory(snapshot_path, config)
.map_err(|e| ANNError::from(super::ConfigError(e)))
} else {
BfTree::new_from_snapshot(config, None).map_err(|e| ANNError::from(super::ConfigError(e)))
}
}
impl<T> SaveWith<String> for BfTreeProvider<T, NoStore>
where
T: VectorRepr,
{
type Ok = usize;
type Error = ANNError;
async fn save_with<P>(&self, storage: &P, prefix: &String) -> Result<Self::Ok, Self::Error>
where
P: StorageWriteProvider,
{
let saved_params = SavedParams {
max_points: self.max_points(),
frozen_points: NonZeroUsize::new(self.num_start_points())
.ok_or_else(|| ANNError::log_index_error("num_start_points is zero"))?,
dim: self.dim(),
metric: self.metric().as_str().to_string(),
max_degree: self.max_degree(),
prefix: prefix.clone(),
params_vector: BfTreeParams {
bytes: self.full_vectors.config().get_cb_size_byte(),
max_record_size: self.full_vectors.config().get_cb_max_record_size(),
leaf_page_size: self.full_vectors.config().get_leaf_page_size(),
},
params_neighbor: BfTreeParams {
bytes: self.neighbor_provider.config().get_cb_size_byte(),
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,
graph_params: self.graph_params.clone(),
is_memory: self.full_vectors.config().is_memory_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"
);
{
let params_filename = BfTreePaths::params_json(&saved_params.prefix);
let params_json = serde_json::to_string(&saved_params).map_err(|e| {
ANNError::log_index_error(format!("Failed to serialize params: {}", e))
})?;
let mut params_writer = storage.create_for_write(¶ms_filename)?;
params_writer.write_all(params_json.as_bytes())?;
}
save_bftree(
self.full_vectors.bftree(),
BfTreePaths::vectors_bftree(&saved_params.prefix),
)
.await?;
save_bftree(
self.neighbor_provider.bftree(),
BfTreePaths::neighbors_bftree(&saved_params.prefix),
)
.await?;
Ok(0)
}
}
impl<T> LoadWith<String> for BfTreeProvider<T, NoStore>
where
T: VectorRepr,
{
type Error = ANNError;
async fn load_with<P>(storage: &P, prefix: &String) -> Result<Self, Self::Error>
where
P: StorageReadProvider,
{
let saved_params: SavedParams = {
let params_filename = BfTreePaths::params_json(prefix);
let mut params_reader = storage.open_reader(¶ms_filename)?;
let mut params_json = String::new();
params_reader.read_to_string(&mut params_json)?;
serde_json::from_str(¶ms_json).map_err(|e| {
ANNError::log_index_error(format!("Failed to deserialize params: {}", e))
})?
};
let metric = Metric::from_str(&saved_params.metric)
.map_err(|e| ANNError::log_index_error(format!("Failed to parse metric: {}", e)))?;
let vector_index = load_bftree(
&saved_params.params_vector,
BfTreePaths::vectors_bftree(&saved_params.prefix),
saved_params.is_memory,
)?;
let full_vectors = VectorProvider::<T>::new_from_bftree(
saved_params.max_points,
saved_params.dim,
saved_params.frozen_points.get(),
vector_index,
);
let adjacency_list_index = load_bftree(
&saved_params.params_neighbor,
BfTreePaths::neighbors_bftree(&saved_params.prefix),
saved_params.is_memory,
)?;
let neighbor_provider = NeighborProvider::<u32>::new_from_bftree(
saved_params.max_degree,
adjacency_list_index,
)?;
Ok(Self {
quant_vectors: NoStore,
full_vectors,
neighbor_provider,
metric,
graph_params: saved_params.graph_params,
})
}
}
impl<T> SaveWith<String> for BfTreeProvider<T, QuantVectorProvider>
where
T: VectorRepr,
{
type Ok = usize;
type Error = ANNError;
async fn save_with<P>(&self, storage: &P, prefix: &String) -> Result<Self::Ok, Self::Error>
where
P: StorageWriteProvider,
{
let saved_params = SavedParams {
max_points: self.max_points(),
frozen_points: NonZeroUsize::new(self.num_start_points())
.ok_or_else(|| ANNError::log_index_error("num_start_points is zero"))?,
dim: self.dim(),
metric: self.metric().as_str().to_string(),
max_degree: self.max_degree(),
prefix: prefix.clone(),
params_vector: BfTreeParams {
bytes: self.full_vectors.config().get_cb_size_byte(),
max_record_size: self.full_vectors.config().get_cb_max_record_size(),
leaf_page_size: self.full_vectors.config().get_leaf_page_size(),
},
params_neighbor: BfTreeParams {
bytes: self.neighbor_provider.config().get_cb_size_byte(),
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: Some(QuantParams {
params_quant: BfTreeParams {
bytes: self.quant_vectors.config().get_cb_size_byte(),
max_record_size: self.quant_vectors.config().get_cb_max_record_size(),
leaf_page_size: self.quant_vectors.config().get_leaf_page_size(),
},
}),
graph_params: self.graph_params.clone(),
is_memory: self.full_vectors.config().is_memory_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"
);
debug_assert_eq!(
self.full_vectors.config().is_memory_backend(),
self.quant_vectors.config().is_memory_backend(),
"Vector and quant stores have mismatched storage backends"
);
{
let params_filename = BfTreePaths::params_json(&saved_params.prefix);
let params_json = serde_json::to_string(&saved_params).map_err(|e| {
ANNError::log_index_error(format!("Failed to serialize params: {}", e))
})?;
let mut params_writer = storage.create_for_write(¶ms_filename)?;
params_writer.write_all(params_json.as_bytes())?;
}
save_bftree(
self.full_vectors.bftree(),
BfTreePaths::vectors_bftree(&saved_params.prefix),
)
.await?;
save_bftree(
self.neighbor_provider.bftree(),
BfTreePaths::neighbors_bftree(&saved_params.prefix),
)
.await?;
save_bftree(
self.quant_vectors.bftree(),
BfTreePaths::quant_bftree(&saved_params.prefix),
)
.await?;
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<T> LoadWith<String> for BfTreeProvider<T, QuantVectorProvider>
where
T: VectorRepr,
{
type Error = ANNError;
async fn load_with<P>(storage: &P, prefix: &String) -> Result<Self, Self::Error>
where
P: StorageReadProvider,
{
let saved_params: SavedParams = {
let params_filename = BfTreePaths::params_json(prefix);
let mut params_reader = storage.open_reader(¶ms_filename)?;
let mut params_json = String::new();
params_reader.read_to_string(&mut params_json)?;
serde_json::from_str(¶ms_json).map_err(|e| {
ANNError::log_index_error(format!("Failed to deserialize params: {}", e))
})?
};
let quant_params = saved_params.quant_params.ok_or_else(|| {
ANNError::log_index_error("Missing quant_params in saved params for quantized provider")
})?;
let metric = Metric::from_str(&saved_params.metric)
.map_err(|e| ANNError::log_index_error(format!("Failed to parse metric: {}", e)))?;
let vector_index = load_bftree(
&saved_params.params_vector,
BfTreePaths::vectors_bftree(&saved_params.prefix),
saved_params.is_memory,
)?;
let full_vectors = VectorProvider::<T>::new_from_bftree(
saved_params.max_points,
saved_params.dim,
saved_params.frozen_points.get(),
vector_index,
);
let adjacency_list_index = load_bftree(
&saved_params.params_neighbor,
BfTreePaths::neighbors_bftree(&saved_params.prefix),
saved_params.is_memory,
)?;
let neighbor_provider = NeighborProvider::<u32>::new_from_bftree(
saved_params.max_degree,
adjacency_list_index,
)?;
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<dyn Quantizer> = 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(quantizer, quant_vector_index);
Ok(Self {
quant_vectors,
full_vectors,
neighbor_provider,
metric,
graph_params: saved_params.graph_params,
})
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::*;
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<DiskANNIndex<BfTreeProvider<f32, QuantVectorProvider>>> {
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::<u32>::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 data = Matrix::new(
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,
);
let ids: Arc<[u32]> = (0u32..15).collect::<Vec<_>>().into();
let batch: Arc<Matrix<f32>> = Arc::new(data);
index
.multi_insert::<Quantized, Matrix<f32>>(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::<u32>::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"
);
let neighbor_ids: Vec<u32> = 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]
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::<u32>::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<u32> = neighbors.iter().map(|n| n.id).collect();
assert!(!neighbor_ids.contains(&2u32));
assert!(!neighbor_ids.contains(&4u32));
}
fn create_full_precision_index() -> Arc<DiskANNIndex<BfTreeProvider<f32, NoStore>>> {
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::<u32>::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::<u32>::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<u32> = 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 num_start_points = 2;
let dim = 5;
let start_points = Matrix::try_from(
vec![0.0f32; dim]
.into_iter()
.chain(vec![0.5f32; dim])
.collect::<Box<[_]>>(),
num_start_points,
dim,
)
.unwrap();
let provider = BfTreeProvider::new(
BfTreeProviderParameters {
max_points: 10,
num_start_points: NonZeroUsize::new(num_start_points).unwrap(),
dim,
metric: Metric::L2,
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,
)
.unwrap();
assert_eq!((&provider).into_iter(), 0..(10 + 2));
let iter = provider.iter();
for i in iter.clone() {
let vector: Vec<f32> = (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);
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
);
provider.delete(ctx, &i).await.unwrap();
assert_eq!(
provider.status_by_internal_id(ctx, i).await.unwrap(),
ElementStatus::Deleted
);
assert_eq!(
provider.status_by_external_id(ctx, &i).await.unwrap(),
ElementStatus::Deleted
);
}
for i in iter.clone() {
provider.release(ctx, i).await.unwrap();
assert_eq!(
provider.status_by_internal_id(ctx, i).await.unwrap(),
ElementStatus::Deleted
);
}
assert!(provider
.set_element(ctx, &100, &[1.0, 2.0, 3.0, 4.0])
.await
.is_err());
}
#[tokio::test]
async fn test_empty_neighbor_list() {
let num_points = 100u32;
let ctx = &DefaultContext;
let num_start_points = 2;
let dim = 3;
let start_points = Matrix::new(Init(|| 0.0f32), num_start_points, dim);
let provider = BfTreeProvider::<f32, _>::new(
BfTreeProviderParameters {
max_points: num_points as usize,
num_start_points: NonZeroUsize::new(num_start_points).unwrap(),
dim,
metric: Metric::L2,
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,
)
.unwrap();
let mut scratch = provider.neighbor_provider.scratch();
for i in 0..num_points {
let vector = vec![i as f32, (i + 1) as f32, (i + 2) as f32];
provider.set_element(ctx, &i, &vector).await.unwrap();
let mut out = AdjacencyList::new();
provider
.neighbor_provider
.get_neighbors(i, &mut out)
.unwrap();
assert!(out.is_empty());
scratch.write_neighbors(i, &[]).unwrap();
provider
.neighbor_provider
.get_neighbors(i, &mut out)
.unwrap();
assert!(out.is_empty());
}
for i in 0..num_points {
let mut out = AdjacencyList::new();
let neighbors = vec![10, 20, 30, 40, 50, 60, 70, 80, 90, 100];
scratch.write_neighbors(i, &neighbors).unwrap();
provider
.neighbor_provider
.get_neighbors(i, &mut out)
.unwrap();
assert_eq!(&*out, &[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]);
scratch.write_neighbors(i, &[]).unwrap();
provider
.neighbor_provider
.get_neighbors(i, &mut out)
.unwrap();
assert!(out.is_empty());
}
let mut out = AdjacencyList::from_iter_untrusted([10, 20, 30, 40, 50, 60, 70, 80, 90, 100]);
assert!(provider
.neighbor_provider
.get_neighbors(200, &mut out)
.is_err());
assert!(out.is_empty());
}
use tempfile::tempdir;
#[tokio::test]
async fn test_bf_tree_provider_save_load_no_quant() {
let num_points = 50usize;
let dim = 4usize;
let max_degree = 32u32;
let num_start_points = NonZeroUsize::new(2).unwrap();
let ctx = &DefaultContext;
let temp_dir = tempdir().unwrap();
let temp_path = temp_dir.path();
let prefix = temp_path
.join("test_bf_tree_provider")
.to_string_lossy()
.to_string();
let vector_path = BfTreePaths::vectors_bftree(&prefix);
let neighbor_path = BfTreePaths::neighbors_bftree(&prefix);
let bytes_vector = 1024 * 1024;
let mut vector_config = Config::new(&vector_path, bytes_vector);
vector_config.leaf_page_size(8192);
vector_config.cb_max_record_size(1024);
vector_config.storage_backend(bf_tree::StorageBackend::Std);
let bytes_neighbor = 1024 * 1024;
let mut neighbor_config = Config::new(&neighbor_path, bytes_neighbor);
neighbor_config.storage_backend(bf_tree::StorageBackend::Std);
let params = BfTreeProviderParameters {
max_points: num_points,
num_start_points,
dim,
metric: Metric::L2,
max_degree,
vector_provider_config: vector_config.clone(),
quant_vector_provider_config: Config::default(),
neighbor_list_provider_config: neighbor_config.clone(),
graph_params: None,
};
let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
let provider =
BfTreeProvider::<f32, NoStore>::new(params.clone(), start_points.as_view(), NoStore)
.unwrap();
for i in 0..num_points {
let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
provider
.set_element(ctx, &(i as u32), &vector)
.await
.unwrap();
}
let mut scratch = provider.neighbor_provider.scratch();
for i in 0..num_points as u32 {
let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
.map(|j| (i + j) % num_points as u32)
.collect();
scratch.write_neighbors(i, &neighbors).unwrap();
}
assert_eq!(vector_config.get_leaf_page_size(), 8192);
assert_eq!(vector_config.get_cb_max_record_size(), 1024);
let storage = FileStorageProvider;
let save_dir = tempdir().unwrap();
let save_prefix = save_dir
.path()
.join("saved_bf_tree_provider")
.to_string_lossy()
.to_string();
provider.save_with(&storage, &save_prefix).await.unwrap();
let loaded_provider = BfTreeProvider::<f32, NoStore>::load_with(&storage, &save_prefix)
.await
.unwrap();
for i in 0..num_points as u32 {
let original = provider.full_vectors.get_vector_sync(i as usize).unwrap();
let loaded = loaded_provider
.full_vectors
.get_vector_sync(i as usize)
.unwrap();
assert_eq!(original, loaded, "Vector mismatch at index {}", i);
}
for i in 0..num_points as u32 {
let mut original_list = AdjacencyList::new();
let mut loaded_list = AdjacencyList::new();
provider
.neighbor_provider
.get_neighbors(i, &mut original_list)
.unwrap();
loaded_provider
.neighbor_provider
.get_neighbors(i, &mut loaded_list)
.unwrap();
assert_eq!(
&*original_list, &*loaded_list,
"Neighbor list mismatch at index {}",
i
);
}
}
#[tokio::test]
async fn test_bf_tree_provider_save_load_quant() {
let num_points = 50usize;
let dim = 8usize;
let max_degree = 32u32;
let num_start_points = NonZeroUsize::new(2).unwrap();
let ctx = &DefaultContext;
let temp_dir = tempdir().unwrap();
let temp_path = temp_dir.path();
let prefix = temp_path
.join("test_bf_tree_provider_quant")
.to_string_lossy()
.to_string();
let vector_path = BfTreePaths::vectors_bftree(&prefix);
let neighbor_path = BfTreePaths::neighbors_bftree(&prefix);
let quant_path = BfTreePaths::quant_bftree(&prefix);
let bytes_vector = 1024 * 1024;
let mut vector_config = Config::new(&vector_path, bytes_vector);
vector_config.storage_backend(bf_tree::StorageBackend::Std);
let bytes_neighbor = 1024 * 1024;
let mut neighbor_config = Config::new(&neighbor_path, bytes_neighbor);
neighbor_config.storage_backend(bf_tree::StorageBackend::Std);
let bytes_quant = 1024 * 1024;
let mut quant_config = Config::new(&quant_path, bytes_quant);
quant_config.storage_backend(bf_tree::StorageBackend::Std);
let quantizer = create_test_quantizer(dim);
let params = BfTreeProviderParameters {
max_points: num_points,
num_start_points,
dim,
metric: Metric::L2,
max_degree,
vector_provider_config: vector_config.clone(),
quant_vector_provider_config: quant_config.clone(),
neighbor_list_provider_config: neighbor_config.clone(),
graph_params: None,
};
let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
let provider = BfTreeProvider::<f32, QuantVectorProvider>::new(
params.clone(),
start_points.as_view(),
quantizer,
)
.unwrap();
for i in 0..num_points {
let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
provider
.set_element(ctx, &(i as u32), &vector)
.await
.unwrap();
}
let mut scratch = provider.neighbor_provider.scratch();
for i in 0..num_points as u32 {
let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
.map(|j| (i + j) % num_points as u32)
.collect();
scratch.write_neighbors(i, &neighbors).unwrap();
}
let storage = FileStorageProvider;
let save_dir = tempdir().unwrap();
let save_prefix = save_dir
.path()
.join("saved_bf_tree_provider_quant")
.to_string_lossy()
.to_string();
provider.save_with(&storage, &save_prefix).await.unwrap();
let loaded_provider =
BfTreeProvider::<f32, QuantVectorProvider>::load_with(&storage, &save_prefix)
.await
.unwrap();
assert_eq!(
provider.quant_vectors.quantizer.full_dim(),
loaded_provider.quant_vectors.quantizer.full_dim(),
"Quantizer full_dim mismatch"
);
assert_eq!(
provider.quant_vectors.quantizer.bytes(),
loaded_provider.quant_vectors.quantizer.bytes(),
"Quantizer bytes mismatch"
);
assert_eq!(
provider.quant_vectors.quantizer.nbits(),
loaded_provider.quant_vectors.quantizer.nbits(),
"Quantizer nbits mismatch"
);
for i in 0..num_points as u32 {
let original = provider.full_vectors.get_vector_sync(i as usize).unwrap();
let loaded = loaded_provider
.full_vectors
.get_vector_sync(i as usize)
.unwrap();
assert_eq!(original, loaded, "Vector mismatch at index {}", i);
}
for i in 0..num_points as u32 {
let original = provider.quant_vectors.get_vector_sync(i as usize).unwrap();
let loaded = loaded_provider
.quant_vectors
.get_vector_sync(i as usize)
.unwrap();
assert_eq!(original, loaded, "Quant vector mismatch at index {}", i);
}
for i in 0..num_points as u32 {
let mut original_list = AdjacencyList::new();
let mut loaded_list = AdjacencyList::new();
provider
.neighbor_provider
.get_neighbors(i, &mut original_list)
.unwrap();
loaded_provider
.neighbor_provider
.get_neighbors(i, &mut loaded_list)
.unwrap();
assert_eq!(
&*original_list, &*loaded_list,
"Neighbor list mismatch at index {}",
i
);
}
}
#[tokio::test]
async fn test_bf_tree_provider_memory_save_load_no_quant() {
let num_points = 20usize;
let dim = 4usize;
let max_degree = 16u32;
let num_start_points = NonZeroUsize::new(1).unwrap();
let ctx = &DefaultContext;
let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
let provider = BfTreeProvider::<f32, NoStore>::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(),
NoStore,
)
.unwrap();
for i in 0..num_points {
let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
provider
.set_element(ctx, &(i as u32), &vector)
.await
.unwrap();
}
let mut scratch = provider.neighbor_provider.scratch();
for i in 0..num_points as u32 {
let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
.map(|j| (i + j) % num_points as u32)
.collect();
scratch.write_neighbors(i, &neighbors).unwrap();
}
provider.delete(ctx, &3u32).await.unwrap();
provider.delete(ctx, &7u32).await.unwrap();
let save_dir = tempdir().unwrap();
let save_prefix = save_dir
.path()
.join("mem_no_quant")
.to_string_lossy()
.to_string();
let storage = FileStorageProvider;
provider.save_with(&storage, &save_prefix).await.unwrap();
let loaded = BfTreeProvider::<f32, NoStore>::load_with(&storage, &save_prefix)
.await
.unwrap();
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(),
"Vector mismatch at {}",
i
);
}
for i in 0..num_points as u32 {
let mut orig = AdjacencyList::new();
let mut load = AdjacencyList::new();
provider
.neighbor_provider
.get_neighbors(i, &mut orig)
.unwrap();
loaded
.neighbor_provider
.get_neighbors(i, &mut load)
.unwrap();
assert_eq!(&*orig, &*load, "Neighbor mismatch at {}", i);
}
assert_eq!(
loaded.status_by_internal_id(ctx, 3).await.unwrap(),
ElementStatus::Deleted
);
assert_eq!(
loaded.status_by_internal_id(ctx, 7).await.unwrap(),
ElementStatus::Deleted
);
assert_eq!(
loaded.status_by_internal_id(ctx, 0).await.unwrap(),
ElementStatus::Valid
);
}
#[tokio::test]
async fn test_bf_tree_provider_memory_save_load_quant() {
let num_points = 20usize;
let dim = 8usize;
let max_degree = 16u32;
let num_start_points = NonZeroUsize::new(1).unwrap();
let ctx = &DefaultContext;
let quantizer = create_test_quantizer(dim);
let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim);
let provider = BfTreeProvider::<f32, QuantVectorProvider>::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();
for i in 0..num_points {
let vector: Vec<f32> = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect();
provider
.set_element(ctx, &(i as u32), &vector)
.await
.unwrap();
}
let mut scratch = provider.neighbor_provider.scratch();
for i in 0..num_points as u32 {
let neighbors: Vec<u32> = (0..std::cmp::min(i, max_degree))
.map(|j| (i + j) % num_points as u32)
.collect();
scratch.write_neighbors(i, &neighbors).unwrap();
}
provider.delete(ctx, &2u32).await.unwrap();
let save_dir = tempdir().unwrap();
let save_prefix = save_dir
.path()
.join("mem_quant")
.to_string_lossy()
.to_string();
let storage = FileStorageProvider;
provider.save_with(&storage, &save_prefix).await.unwrap();
let loaded = BfTreeProvider::<f32, QuantVectorProvider>::load_with(&storage, &save_prefix)
.await
.unwrap();
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(),
"Vector mismatch at {}",
i
);
}
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(),
"Quant vector mismatch at {}",
i
);
}
for i in 0..num_points as u32 {
if i == 2 {
continue;
}
let mut orig = AdjacencyList::new();
let mut load = AdjacencyList::new();
provider
.neighbor_provider
.get_neighbors(i, &mut orig)
.unwrap();
loaded
.neighbor_provider
.get_neighbors(i, &mut load)
.unwrap();
assert_eq!(&*orig, &*load, "Neighbor mismatch at {}", i);
}
assert_eq!(
loaded.status_by_internal_id(ctx, 2).await.unwrap(),
ElementStatus::Deleted
);
assert_eq!(
loaded.status_by_internal_id(ctx, 0).await.unwrap(),
ElementStatus::Valid
);
}
}