diff --git a/diskann-bftree/src/id.rs b/diskann-bftree/src/id.rs new file mode 100644 index 000000000..1220028ba --- /dev/null +++ b/diskann-bftree/src/id.rs @@ -0,0 +1,182 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Vertex-id abstraction for the bf-tree provider. + +use std::marker::PhantomData; + +use diskann::utils::VectorId; +use diskann::{ANNError, ANNResult}; + +/// Identifier type usable as a `BfTreeProvider` vertex id. +/// +/// This bundles the bounds the core algorithm requires of an id ([`VectorId`], +/// mandated by `HasId::Id` and `DataProvider::InternalId`) with the index +/// arithmetic the provider needs: converting an id *to* a `usize` index +/// ([`as_index`](BfTreeId::as_index), used to key the per-vector stores) and +/// *from* a zero-based index ([`from_index`](BfTreeId::from_index)). The +/// provider mints ids densely from `0..total`, so it needs a way to build an +/// `I` from a counter. +/// +/// Implemented for `u32` (the default, capping at ~4.29B vertices) and `u64` +/// (for billion-scale-and-beyond, larger-than-memory datasets). On a 64-bit +/// target `u64` covers every representable `usize`, so its conversions never +/// fail. +pub trait BfTreeId: VectorId { + /// Compile-time proof that this id type's `usize` index conversions are + /// lossless on the target platform. + /// + /// [`as_index`](BfTreeId::as_index) and the bf-tree stores round-trip an id + /// through `usize`, which is only lossless when `usize` is at least as wide + /// as the id type. `u32` fits every supported target; `u64` overrides this to + /// assert a 64-bit `usize`, since on a 32-bit target a `u64` id above + /// `u32::MAX` would silently truncate. Because associated consts are only + /// evaluated when monomorphized, referencing this in a generic id path (see + /// [`validate_id_capacity`]) turns the truncation into a compile error *only* + /// for `u64` providers on 32-bit targets, leaving the default `u32` path + /// (even on 32-bit) unaffected. + const INDEX_CONVERSION_LOSSLESS: () = (); + + /// Build an id from a zero-based index, truncating on overflow. + /// + /// Only call this for indices already known to fit (e.g. ids drawn from + /// `0..total`, which the provider guarantees fit by construction). + fn from_index(index: usize) -> Self; + + /// Build an id from a zero-based index, returning `None` if it does not fit. + fn try_from_index(index: usize) -> Option; + + /// Convert this id to its zero-based `usize` index. + /// + /// The provider uses the identity map, so an id *is* its own index. This is + /// lossless on the 64-bit targets diskann supports. + fn as_index(&self) -> usize; + + /// An iterator over the dense id range `0..total`. + /// + /// The conversion is monomorphized and inlined per id type (no stored + /// function pointer), and the iterator preserves the exact-size and + /// double-ended properties of the underlying index range. + #[inline] + fn id_range(total: usize) -> IdRange { + IdRange::new(total) + } +} + +impl BfTreeId for u32 { + #[inline(always)] + fn from_index(index: usize) -> Self { + index as u32 + } + + #[inline(always)] + fn try_from_index(index: usize) -> Option { + u32::try_from(index).ok() + } + + #[inline(always)] + fn as_index(&self) -> usize { + *self as usize + } +} + +impl BfTreeId for u64 { + const INDEX_CONVERSION_LOSSLESS: () = assert!( + usize::BITS >= u64::BITS, + "u64 bf-tree vertex ids require a 64-bit target: on a 32-bit `usize`, ids above \ + u32::MAX would truncate when converted to a store index" + ); + + #[inline(always)] + fn from_index(index: usize) -> Self { + index as u64 + } + + #[inline(always)] + fn try_from_index(index: usize) -> Option { + u64::try_from(index).ok() + } + + #[inline(always)] + fn as_index(&self) -> usize { + *self as usize + } +} + +/// A dense id iterator yielding `I::from_index(0..total)`. +/// +/// Wrapping a `Range` and mapping inside [`Iterator::next`] keeps the +/// conversion a zero-sized, inlinable function item — unlike `Range::map` with a +/// `fn(usize) -> I` pointer, which forces an indirect call per element — while +/// still delegating length and reverse iteration to the inner range. +#[derive(Debug, Clone)] +pub struct IdRange { + inner: std::ops::Range, + _marker: PhantomData I>, +} + +impl IdRange { + #[inline] + fn new(total: usize) -> Self { + Self { + inner: 0..total, + _marker: PhantomData, + } + } +} + +impl Iterator for IdRange { + type Item = I; + + #[inline] + fn next(&mut self) -> Option { + self.inner.next().map(I::from_index) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} + +impl DoubleEndedIterator for IdRange { + #[inline] + fn next_back(&mut self) -> Option { + self.inner.next_back().map(I::from_index) + } +} + +impl ExactSizeIterator for IdRange { + #[inline] + fn len(&self) -> usize { + self.inner.len() + } +} + +/// Validate that a provider holding `total` ids can represent every id in `0..total`. +/// +/// `BfTreeProvider::iter` mints ids densely via the infallible (truncating) +/// [`BfTreeId::from_index`]; callers must guarantee the range fits in `I`. This check +/// enforces that guarantee up front (at construction and load) so the truncating +/// conversion can never silently wrap a real id. +pub(crate) fn validate_id_capacity(total: usize) -> ANNResult<()> { + // Force evaluation of the id type's compile-time index-conversion guard. This + // is a no-op for `u32`, but fails to compile a `u64` provider on a 32-bit + // target (where `as_index` would truncate). It lives here because this + // generic function is monomorphized for every id type a provider is built + // with, so the check is scoped to actual `u64` usage. + let () = I::INDEX_CONVERSION_LOSSLESS; + + if let Some(last) = total.checked_sub(1) { + if I::try_from_index(last).is_none() { + return Err(ANNError::message(format!( + "provider capacity of {total} ids exceeds the maximum representable by the \ + {}-byte vertex id type", + std::mem::size_of::() + ))); + } + } + Ok(()) +} diff --git a/diskann-bftree/src/lib.rs b/diskann-bftree/src/lib.rs index e0a7f3de0..a9d2f96e7 100644 --- a/diskann-bftree/src/lib.rs +++ b/diskann-bftree/src/lib.rs @@ -9,6 +9,7 @@ //! [`DataProvider`](diskann::provider::DataProvider) trait, enabling indexes that can //! transparently spill to disk for datasets larger than available memory. +pub mod id; pub mod neighbors; pub mod provider; pub mod quant; @@ -16,6 +17,8 @@ pub mod vectors; mod locks; +pub use id::BfTreeId; + // Accessors pub use provider::{ AsVectorDtype, BfTreePaths, BfTreeProvider, BfTreeProviderParameters, CreateQuantProvider, diff --git a/diskann-bftree/src/neighbors.rs b/diskann-bftree/src/neighbors.rs index b3e3c6991..e68a6275f 100644 --- a/diskann-bftree/src/neighbors.rs +++ b/diskann-bftree/src/neighbors.rs @@ -12,15 +12,15 @@ use bytemuck::{cast_slice, cast_slice_mut}; use diskann::{ graph::AdjacencyList, provider::{self, HasId}, - utils::{IntoUsize, VectorId}, ANNError, ANNResult, }; use super::ConfigError; +use crate::id::BfTreeId; use crate::locks::StripedLocks; use crate::{bftree_insert, TestCallCount}; -pub struct NeighborProvider { +pub struct NeighborProvider { adjacency_list_index: BfTree, dim: usize, // Max number of neighbors in a neighbor list + 1 for the neighbor count #[allow(dead_code)] @@ -28,15 +28,18 @@ pub struct NeighborProvider { _phantom: PhantomData, } -impl HasId for NeighborProvider { +impl HasId for NeighborProvider { type Id = I; } -impl NeighborProvider { +impl NeighborProvider { /// Create a new instance based on bf-tree Config directly. pub fn new_with_config(max_degree: u32, config: Config) -> ANNResult { - let key_size = std::mem::size_of::(); - let value_size = (max_degree as usize + 1) * std::mem::size_of::(); + // Records are keyed by an `I`-width id and store a `dim`-cell `I`-width value + // (`dim == max_degree + 1`), so size the validation by the actual id width + // rather than assuming `u32`. + let key_size = std::mem::size_of::(); + let value_size = (max_degree as usize + 1) * std::mem::size_of::(); crate::validate_record_size("neighbor_provider", &config, key_size, value_size)?; let adj_list_index = BfTree::with_config(config, None).map_err(ConfigError)?; @@ -45,7 +48,7 @@ impl NeighborProvider { } fn new(max_degree: u32, adjacency_list_index: BfTree) -> ANNResult { - let dim = 1 + max_degree.into_usize(); + let dim = 1 + max_degree as usize; Ok(Self { adjacency_list_index, @@ -286,7 +289,7 @@ impl NeighborProvider { pub struct NeighborAccessor<'a, I> where - I: VectorId + IntoUsize, + I: BfTreeId, { provider: &'a NeighborProvider, locks: &'a StripedLocks, @@ -295,28 +298,28 @@ where impl<'a, I> NeighborAccessor<'a, I> where - I: VectorId + IntoUsize, + I: BfTreeId, { pub fn write_neighbors(&mut self, id: I, neighbors: &[I]) -> ANNResult<()> { - let _guard = self.locks.lock(id.into_usize()); + let _guard = self.locks.lock(id.as_index()); self.provider.set_neighbors(id, neighbors, &mut self.buf) } pub fn write_append(&mut self, id: I, neighbors: &[I]) -> ANNResult<()> { - let _guard = self.locks.lock(id.into_usize()); + let _guard = self.locks.lock(id.as_index()); self.provider.append_vector(id, neighbors, &mut self.buf) } } impl<'a, I> HasId for NeighborAccessor<'a, I> where - I: VectorId + IntoUsize, + I: BfTreeId, { type Id = I; } impl<'a, I> provider::NeighborAccessor for NeighborAccessor<'a, I> where - I: VectorId + IntoUsize, + I: BfTreeId, { fn get_neighbors( &mut self, @@ -329,14 +332,14 @@ where impl<'a, I> provider::NeighborAccessorMut for NeighborAccessor<'a, I> where - I: VectorId + IntoUsize, + I: BfTreeId, { fn set_neighbors( &mut self, id: Self::Id, neighbors: &[Self::Id], ) -> impl std::future::Future> + Send { - let _guard = self.locks.lock(id.into_usize()); + let _guard = self.locks.lock(id.as_index()); std::future::ready(self.provider.set_neighbors(id, neighbors, &mut self.buf)) } fn append_vector( @@ -344,7 +347,7 @@ where id: Self::Id, neighbors: &[Self::Id], ) -> impl std::future::Future> + Send { - let _guard = self.locks.lock(id.into_usize()); + let _guard = self.locks.lock(id.as_index()); std::future::ready(self.provider.append_vector(id, neighbors, &mut self.buf)) } } @@ -416,6 +419,42 @@ mod tests { } } + /// Exercise the `u64` id path beyond the `u32` range: vertex ids and neighbor + /// values above `u32::MAX` must round-trip, and ids that share their low 32 bits + /// must not collide (proving the bf-tree key uses the full 8-byte width). + #[tokio::test] + async fn test_u64_high_bit_ids() { + let locks = Arc::new(StripedLocks::new()); + let neighbor_provider = + NeighborProvider::::new_with_config(6, Config::default()).unwrap(); + let mut scratch = neighbor_provider.scratch(&locks); + + let high: u64 = (u32::MAX as u64) + 1; + let big_id: u64 = high + 7; + let big_neighbors: Vec = vec![high, high + 1, u64::MAX, 3]; + scratch.write_neighbors(big_id, &big_neighbors).unwrap(); + + let mut result = AdjacencyList::with_capacity(10); + neighbor_provider + .get_neighbors(big_id, &mut result) + .unwrap(); + assert_eq!(&*big_neighbors, &*result); + + // `low` and `low | (1 << 32)` share their low 32 bits; with an 8-byte key they + // are distinct entries. A truncating 4-byte key would alias them. + let low: u64 = 1; + let aliased: u64 = low | (1u64 << 32); + scratch.write_neighbors(low, &[10, 11]).unwrap(); + scratch.write_neighbors(aliased, &[20, 21]).unwrap(); + + neighbor_provider.get_neighbors(low, &mut result).unwrap(); + assert_eq!(&[10u64, 11], &*result); + neighbor_provider + .get_neighbors(aliased, &mut result) + .unwrap(); + assert_eq!(&[20u64, 21], &*result); + } + /// Test corner cases of appending to neighbor list #[tokio::test] async fn test_neighbor_accessors() { diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 3f6bbae2e..7e9e53d63 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -32,7 +32,7 @@ use diskann::{ }, neighbor::Neighbor, provider::{DataProvider, DefaultContext, Delete, ElementStatus, HasId, NoopGuard, SetElement}, - utils::{IntoUsize, VectorRepr}, + utils::VectorRepr, ANNError, ANNResult, }; use diskann_utils::{ @@ -46,7 +46,7 @@ use super::{ neighbors::{NeighborAccessor, NeighborProvider}, quant::QuantVectorProvider, vectors::VectorProvider, - AccessError, NoStore, + AccessError, BfTreeId, NoStore, }; use crate::locks::StripedLocks; use diskann_providers::model::graph::provider::async_::distances::UnwrapErr; @@ -178,9 +178,10 @@ use diskann_providers::storage::{LoadWith, SaveWith, StorageReadProvider, Storag /// quantizer, /// ); /// ``` -pub struct BfTreeProvider +pub struct BfTreeProvider where T: VectorRepr, + I: BfTreeId, { // The quant vector store. If `Q == NoStore`, the quantized operations are disabled. // @@ -192,7 +193,7 @@ where // Provider that holds the graph structure as neighbors of vectors. // - pub(crate) neighbor_provider: NeighborProvider, + pub(crate) neighbor_provider: NeighborProvider, // The metric to use for distances // @@ -266,9 +267,10 @@ pub struct BfTreeProviderParameters { pub use_snapshot: bool, } -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, + I: BfTreeId, { /// Construct a new data provider from empty. Callers of this are required to manually set start /// points before performing search tasks. @@ -300,6 +302,14 @@ where .quant_vector_provider_config .use_snapshot(params.use_snapshot); + let id_capacity = params + .max_points + .checked_add(params.num_start_points.get()) + .ok_or_else(|| { + ANNError::message("max_points + num_start_points overflows usize".to_string()) + })?; + crate::id::validate_id_capacity::(id_capacity)?; + Ok(Self { quant_vectors: quant_precursor.create(params.quant_vector_provider_config)?, full_vectors: VectorProvider::new_with_config( @@ -368,17 +378,17 @@ where // /// // pub(crate) fn is_not_start_point(&self) -> impl Fn(&Neighbor) -> bool { // let range = self.full_vectors.start_point_range(); - // move |neighbor| !range.contains(&neighbor.id.into_usize()) + // move |neighbor| !range.contains(&neighbor.id.as_index()) // } /// Return a vector of starting points. - pub fn starting_points(&self) -> ANNResult> { + pub fn starting_points(&self) -> ANNResult> { self.full_vectors.starting_points() } /// An iterator over all ids including start points (even if they are deleted). - pub fn iter(&self) -> std::ops::Range { - 0..(self.full_vectors.total() as u32) + pub fn iter(&self) -> crate::id::IdRange { + I::id_range(self.full_vectors.total()) } pub fn num_start_points(&self) -> usize { @@ -406,9 +416,10 @@ where } } -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, + I: BfTreeId, { /// Return the number of vector reads for full-precision and quant-vectors respectively /// @@ -420,9 +431,10 @@ where } } -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, + I: BfTreeId, { /// Return the number of vector reads for full-precision and quant-vectors respectively /// @@ -461,10 +473,11 @@ impl DeleteQuant for NoStore { /// [`InplaceDeleteStrategy::get_delete_element`]) *after* the delete has already been committed. /// Use [`InplaceDeleteMethod::OneHop`] or [`InplaceDeleteMethod::TwoHopAndOneHop`] instead, /// as these strategies only require neighbor topology (which remains accessible). -impl Delete for BfTreeProvider +impl Delete for BfTreeProvider where T: VectorRepr, Q: AsyncFriendly + DeleteQuant, + I: BfTreeId, { fn release( &self, @@ -480,12 +493,12 @@ where gid: &Self::ExternalId, ) -> impl std::future::Future> + Send { let id = *gid; - let _guard = self.locks.lock(id as usize); + let _guard = self.locks.lock(id.as_index()); // Only delete vector data here. Neighbor adjacency cleanup (zeroing the // deleted vertex's edge list and patching neighbors-of-neighbors) is // handled by `DiskANNIndex::inplace_delete` → `drop_adj_list`. - self.full_vectors.delete_vector(id as usize); - self.quant_vectors.delete_vector(id as usize); + self.full_vectors.delete_vector(id.as_index()); + self.quant_vectors.delete_vector(id.as_index()); std::future::ready(Ok(())) } @@ -505,7 +518,7 @@ where id: Self::InternalId, ) -> impl std::future::Future> + Send { - let status = match self.full_vectors.get_vector_sync(id.into_usize()) { + let status = match self.full_vectors.get_vector_sync(id.as_index()) { Ok(_) => Ok(ElementStatus::Valid), Err(RankedError::Transient(_)) => Ok(ElementStatus::Deleted), Err(RankedError::Error(e)) => Err(e), @@ -516,12 +529,13 @@ where /// Allow `&BfTreeProvider` to implement `IntoIter` /// -impl IntoIterator for &BfTreeProvider +impl IntoIterator for &BfTreeProvider where T: VectorRepr, + I: BfTreeId, { - type Item = u32; - type IntoIter = std::ops::Range; + type Item = I; + type IntoIter = crate::id::IdRange; fn into_iter(self) -> Self::IntoIter { self.iter() } @@ -559,12 +573,13 @@ impl CreateQuantProvider for Poly { } } -impl BfTreeProvider +impl BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - pub fn neighbors(&self) -> &NeighborProvider { + pub fn neighbors(&self) -> &NeighborProvider { &self.neighbor_provider } } @@ -573,27 +588,28 @@ where // Data Provider // /////////////////// -impl DataProvider for BfTreeProvider +impl DataProvider for BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { type Context = DefaultContext; // The `BfTreeProvider` uses the identity map for IDs. // - type InternalId = u32; + type InternalId = I; // The `BfTreeProvider` uses the identity map for IDs. // - type ExternalId = u32; + type ExternalId = I; // Use a general error type for now. // type Error = ANNError; // No insert-ID recovery. - type Guard = NoopGuard; + type Guard = NoopGuard; // Translate an external id to its corresponding internal id. // @@ -616,12 +632,13 @@ where } } -impl HasId for BfTreeProvider +impl HasId for BfTreeProvider where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - type Id = u32; + type Id = I; } //////////////// @@ -630,9 +647,10 @@ where /// Assign to both the full-precision and quant vector stores /// -impl SetElement<&[T]> for BfTreeProvider +impl SetElement<&[T]> for BfTreeProvider where T: VectorRepr, + I: BfTreeId, { type SetError = ANNError; @@ -645,18 +663,18 @@ where fn set_element( &self, _context: &Self::Context, - id: &u32, + id: &I, element: &[T], ) -> impl Future> + Send { - let _guard = self.locks.lock(id.into_usize()); + let _guard = self.locks.lock(id.as_index()); // First, write to the authoritative full-precision store. - if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) { + if let Err(err) = self.full_vectors.set_vector_sync(id.as_index(), element) { return std::future::ready(Err(err)); } // Then, write the compressed representation to the quant store. - if let Err(err) = self.quant_vectors.set_vector_sync(id.into_usize(), element) { + if let Err(err) = self.quant_vectors.set_vector_sync(id.as_index(), element) { debug_assert!( false, "quant write failed after full-precision success: {err}" @@ -670,9 +688,10 @@ where /// Assign to just the full-precision store /// -impl SetElement<&[T]> for BfTreeProvider +impl SetElement<&[T]> for BfTreeProvider where T: VectorRepr, + I: BfTreeId, { type SetError = ANNError; @@ -681,12 +700,12 @@ where fn set_element( &self, _context: &Self::Context, - id: &u32, + id: &I, element: &[T], ) -> impl Future> + Send { - let _guard = self.locks.lock(id.into_usize()); + let _guard = self.locks.lock(id.as_index()); - if let Err(err) = self.full_vectors.set_vector_sync(id.into_usize(), element) { + if let Err(err) = self.full_vectors.set_vector_sync(id.as_index(), element) { return std::future::ready(Err(err)); } @@ -726,12 +745,13 @@ 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, + I: BfTreeId, { fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> { - let start_point_ids = self.full_vectors.starting_points()?; + let start_point_ids: Vec = self.full_vectors.starting_points()?; if start_points.nrows() != start_point_ids.len() { return Err(ANNError::message(format!( "expected start_points to contain `{}` rows, instead it has {}", @@ -743,8 +763,8 @@ where let mut scratch = self.neighbor_provider.scratch(&self.locks); for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) { // Set the full-precision vector - self.full_vectors.set_vector_sync(id.into_usize(), v)?; - self.quant_vectors.set_vector_sync(id.into_usize(), v)?; + self.full_vectors.set_vector_sync(id.as_index(), v)?; + self.quant_vectors.set_vector_sync(id.as_index(), v)?; // Initialize empty neighbor list scratch.write_neighbors(id, &[])?; } @@ -757,12 +777,13 @@ 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, + I: BfTreeId, { fn set_start_points(&self, _hidden: Hidden, start_points: MatrixView<'_, T>) -> ANNResult<()> { - let start_point_ids = self.full_vectors.starting_points()?; + let start_point_ids: Vec = self.full_vectors.starting_points()?; if start_points.nrows() != start_point_ids.len() { return Err(ANNError::message(format!( "expected start_points to contain `{}` rows, instead it has {}", @@ -774,7 +795,7 @@ where let mut scratch = self.neighbor_provider.scratch(&self.locks); for (id, v) in std::iter::zip(start_point_ids, start_points.row_iter()) { // Set the full-precision vector - self.full_vectors.set_vector_sync(id.into_usize(), v)?; + self.full_vectors.set_vector_sync(id.as_index(), v)?; // Initialize empty neighbor list scratch.write_neighbors(id, &[])?; } @@ -794,25 +815,27 @@ where /// * [`Accessor`] for the [`BfTreeProvider`]. /// * [`BuildQueryComputer`]. /// -pub struct FullAccessor<'a, T, Q> +pub struct FullAccessor<'a, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { /// The host provider. - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, /// The fused query-distance computer. computer: T::QueryDistance, /// A buffer to store retrieved elements. element: Box<[T]>, } -impl<'a, T, Q> FullAccessor<'a, T, Q> +impl<'a, T, Q, I> FullAccessor<'a, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - pub(crate) fn new(provider: &'a BfTreeProvider, query: &[T]) -> Self { + pub(crate) fn new(provider: &'a BfTreeProvider, query: &[T]) -> Self { Self { provider, computer: T::query_distance(query, provider.metric), @@ -822,28 +845,30 @@ where } } - fn get_distance(&mut self, id: u32) -> Result { + fn get_distance(&mut self, id: I) -> Result { self.provider .full_vectors - .get_vector_into(id.into_usize(), &mut self.element) + .get_vector_into(id.as_index(), &mut self.element) .map(|_: ()| self.computer.evaluate_similarity(&self.element)) } } -impl HasId for FullAccessor<'_, T, Q> +impl HasId for FullAccessor<'_, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - type Id = u32; + type Id = I; } -impl glue::SearchAccessor for FullAccessor<'_, T, Q> +impl glue::SearchAccessor for FullAccessor<'_, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - fn starting_points(&self) -> impl Future>> { + fn starting_points(&self) -> impl Future>> { std::future::ready(self.provider.starting_points()) } @@ -898,23 +923,25 @@ where /// /// * [`Accessor`] for the `BfTreeProvider`. /// -pub struct QuantAccessor<'a, T> +pub struct QuantAccessor<'a, T, I> where T: VectorRepr, + I: BfTreeId, { - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, /// The fused query-distance computer. computer: super::quant::QuantQueryComputer, /// A buffer to store retrieved elements. element: Box<[u8]>, } -impl<'a, T> QuantAccessor<'a, T> +impl<'a, T, I> QuantAccessor<'a, T, I> where T: VectorRepr, + I: BfTreeId, { pub(crate) fn new( - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, query: &[T], ) -> ANNResult { let computer = provider.quant_vectors.query_computer(query)?; @@ -927,11 +954,11 @@ where }) } - fn get_distance(&mut self, id: u32) -> Result { + fn get_distance(&mut self, id: I) -> Result { match self .provider .quant_vectors - .get_vector_into(id.into_usize(), &mut self.element) + .get_vector_into(id.as_index(), &mut self.element) { Ok(()) => self .computer @@ -942,18 +969,20 @@ where } } -impl HasId for QuantAccessor<'_, T> +impl HasId for QuantAccessor<'_, T, I> where T: VectorRepr, + I: BfTreeId, { - type Id = u32; + type Id = I; } -impl glue::SearchAccessor for QuantAccessor<'_, T> +impl glue::SearchAccessor for QuantAccessor<'_, T, I> where T: VectorRepr, + I: BfTreeId, { - fn starting_points(&self) -> impl Future>> { + fn starting_points(&self) -> impl Future>> { std::future::ready(self.provider.starting_points()) } @@ -1003,25 +1032,27 @@ where /////////////////////// /// A [`glue::PruneAccessor`] for full-precision vectors in the `BfTreeProvider`. -pub struct FullPruneAccessor<'a, T, Q> +pub struct FullPruneAccessor<'a, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - provider: &'a BfTreeProvider, - neighbors: NeighborAccessor<'a, u32>, - set: map::Map, map::Ref<[T]>>, + provider: &'a BfTreeProvider, + neighbors: NeighborAccessor<'a, I>, + set: map::Map, map::Ref<[T]>>, distance: T::Distance, } -impl<'a, T, Q> FullPruneAccessor<'a, T, Q> +impl<'a, T, Q, I> FullPruneAccessor<'a, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { fn new( - provider: &'a BfTreeProvider, - set: map::Map, map::Ref<[T]>>, + provider: &'a BfTreeProvider, + set: map::Map, map::Ref<[T]>>, ) -> Self { Self { provider, @@ -1032,23 +1063,25 @@ where } } -impl HasId for FullPruneAccessor<'_, T, Q> +impl HasId for FullPruneAccessor<'_, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - type Id = u32; + type Id = I; } -impl<'q, T, Q> glue::PruneAccessor for FullPruneAccessor<'q, T, Q> +impl<'q, T, Q, I> glue::PruneAccessor for FullPruneAccessor<'q, T, Q, I> where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { type ElementRef<'a> = &'a [T]; type View<'a> - = map::View<'a, u32, Box<[T]>, map::Ref<[T]>> + = map::View<'a, I, Box<[T]>, map::Ref<[T]>> where Self: 'a; @@ -1058,7 +1091,7 @@ where Self: 'a; type Neighbors<'a> - = diskann::provider::Neighbors<'a, NeighborAccessor<'q, u32>> + = diskann::provider::Neighbors<'a, NeighborAccessor<'q, I>> where Self: 'a; @@ -1071,7 +1104,7 @@ where { let mut buf: Option> = None; - let view = self.set.fill(itr, |i: u32| -> ANNResult<_> { + let view = self.set.fill(itr, |i: I| -> ANNResult<_> { let mut b = match buf.take() { Some(b) => b, None => std::iter::repeat_n(T::default(), self.provider.dim()).collect(), @@ -1080,7 +1113,7 @@ where match self .provider .full_vectors - .get_vector_into(i.into_usize(), &mut b) + .get_vector_into(i.as_index(), &mut b) .allow_transient("transient errors allowed during fill")? { Some(()) => Ok(Some(b)), @@ -1105,22 +1138,24 @@ where //////////////////////// /// A [`glue::PruneAccessor`] for quantized vectors in the `BfTreeProvider`. -pub struct QuantPruneAccessor<'a, T> +pub struct QuantPruneAccessor<'a, T, I> where T: VectorRepr, + I: BfTreeId, { - provider: &'a BfTreeProvider, - neighbors: NeighborAccessor<'a, u32>, - set: map::Map, + provider: &'a BfTreeProvider, + neighbors: NeighborAccessor<'a, I>, + set: map::Map, distance: UnwrapErr, } -impl<'a, T> QuantPruneAccessor<'a, T> +impl<'a, T, I> QuantPruneAccessor<'a, T, I> where T: VectorRepr, + I: BfTreeId, { fn new( - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, capacity: usize, ) -> ANNResult { let distance = provider @@ -1137,21 +1172,23 @@ where } } -impl HasId for QuantPruneAccessor<'_, T> +impl HasId for QuantPruneAccessor<'_, T, I> where T: VectorRepr, + I: BfTreeId, { - type Id = u32; + type Id = I; } -impl<'q, T> glue::PruneAccessor for QuantPruneAccessor<'q, T> +impl<'q, T, I> glue::PruneAccessor for QuantPruneAccessor<'q, T, I> where T: VectorRepr, + I: BfTreeId, { type ElementRef<'a> = Opaque<'a>; type View<'a> - = map::View<'a, u32, Owned> + = map::View<'a, I, Owned> where Self: 'a; @@ -1161,7 +1198,7 @@ where Self: 'a; type Neighbors<'a> - = diskann::provider::Neighbors<'a, NeighborAccessor<'q, u32>> + = diskann::provider::Neighbors<'a, NeighborAccessor<'q, I>> where Self: 'a; @@ -1175,7 +1212,7 @@ where let mut buf: Option> = None; let bytes = self.provider.quant_vectors.quantizer.bytes(); - let view = self.set.fill(itr, |i: u32| -> ANNResult<_> { + let view = self.set.fill(itr, |i: I| -> ANNResult<_> { let mut b = match buf.take() { Some(b) => b, None => std::iter::repeat_n(0, bytes).collect(), @@ -1184,7 +1221,7 @@ where match self .provider .quant_vectors - .get_vector_into(i.into_usize(), &mut b) + .get_vector_into(i.as_index(), &mut b) .allow_transient("transient errors allowed during fill")? { Some(()) => Ok(Some(Owned(b))), @@ -1227,17 +1264,18 @@ impl<'short> diskann_utils::Reborrow<'short> for Owned { /// Perform a search entirely in the full-precision space. /// /// Starting points are not filtered out of the final results. -impl<'a, T, Q> SearchStrategy<'a, BfTreeProvider, &'a [T]> for FullPrecision +impl<'a, T, Q, I> SearchStrategy<'a, BfTreeProvider, &'a [T]> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - type SearchAccessor = FullAccessor<'a, T, Q>; + type SearchAccessor = FullAccessor<'a, T, Q, I>; type SearchAccessorError = Infallible; fn search_accessor( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, query: &'a [T], ) -> Result { @@ -1245,26 +1283,28 @@ where } } -impl<'a, T, Q> DefaultPostProcessor<'a, BfTreeProvider, &'a [T]> for FullPrecision +impl<'a, T, Q, I> DefaultPostProcessor<'a, BfTreeProvider, &'a [T]> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { default_post_processor!(glue::Pipeline); } // Pruning -impl PruneStrategy> for FullPrecision +impl PruneStrategy> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { - type PruneAccessor<'a> = FullPruneAccessor<'a, T, Q>; + type PruneAccessor<'a> = FullPruneAccessor<'a, T, Q, I>; type PruneAccessorError = diskann::error::Infallible; fn prune_accessor<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, capacity: usize, ) -> Result, Self::PruneAccessorError> { @@ -1273,10 +1313,11 @@ where } } -impl<'a, T, Q> InsertStrategy<'a, BfTreeProvider, &'a [T]> for FullPrecision +impl<'a, T, Q, I> InsertStrategy<'a, BfTreeProvider, &'a [T]> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { type PruneStrategy = Self; fn prune_strategy(&self) -> Self::PruneStrategy { @@ -1284,13 +1325,14 @@ where } } -impl MultiInsertStrategy, B> for FullPrecision +impl MultiInsertStrategy, B> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, B: for<'a> Batch = &'a [T]> + Debug, { - type Seed = map::Builder>; + type Seed = map::Builder>; type FinishError = diskann::error::Infallible; type PruneStrategy = Self; type InsertStrategy = Self; @@ -1301,13 +1343,13 @@ where fn finish( &self, - _provider: &BfTreeProvider, + _provider: &BfTreeProvider, _ctx: &DefaultContext, batch: &std::sync::Arc, ids: Itr, ) -> impl std::future::Future> + Send where - Itr: ExactSizeIterator + Send, + Itr: ExactSizeIterator + Send, { let overlay = map::Overlay::from_batch(batch.clone(), ids); let builder = map::Builder::new(map::Capacity::Default).with_overlay(overlay); @@ -1316,11 +1358,11 @@ where fn seeded_prune_accessor<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, seed: &'a Self::Seed, capacity: usize, - ) -> ANNResult> { + ) -> ANNResult> { let set = seed.clone().build(capacity); Ok(FullPruneAccessor::new(provider, set)) } @@ -1334,16 +1376,17 @@ where /// [`InplaceDeleteMethod::TwoHopAndOneHop`]. It is **not compatible** with /// [`InplaceDeleteMethod::VisitedAndTopK`] because `BfTreeProvider` performs hard deletes — /// the vector data is erased before `get_delete_element` is called, causing it to fail. -impl InplaceDeleteStrategy> for FullPrecision +impl InplaceDeleteStrategy> for FullPrecision where T: VectorRepr, Q: AsyncFriendly, + I: BfTreeId, { type DeleteElementError = ANNError; type DeleteElement<'a> = &'a [T]; type DeleteElementGuard = Box<[T]>; type PruneStrategy = Self; - type DeleteSearchAccessor<'a> = FullAccessor<'a, T, Q>; + type DeleteSearchAccessor<'a> = FullAccessor<'a, T, Q, I>; type SearchPostProcessor = CopyIds; type SearchStrategy = Self; fn search_strategy(&self) -> Self::SearchStrategy { @@ -1360,14 +1403,14 @@ where async fn get_delete_element<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, - id: u32, + id: I, ) -> Result { use diskann::error::ErrorExt; let elt = provider .full_vectors - .get_vector_sync(id.into_usize()) + .get_vector_sync(id.as_index()) .escalate("get_delete_element: failed to read vector for inplace delete")? .into(); Ok(elt) @@ -1377,16 +1420,17 @@ where /// Perform a search entirely in the quantized space. /// /// Starting points are not filtered out of the final results. -impl<'a, T> SearchStrategy<'a, BfTreeProvider, &'a [T]> for Quantized +impl<'a, T, I> SearchStrategy<'a, BfTreeProvider, &'a [T]> for Quantized where T: VectorRepr, + I: BfTreeId, { - type SearchAccessor = QuantAccessor<'a, T>; + type SearchAccessor = QuantAccessor<'a, T, I>; type SearchAccessorError = ANNError; fn search_accessor( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, query: &'a [T], ) -> Result { @@ -1394,16 +1438,19 @@ where } } -impl<'a, T> DefaultPostProcessor<'a, BfTreeProvider, &'a [T]> for Quantized +impl<'a, T, I> DefaultPostProcessor<'a, BfTreeProvider, &'a [T]> + for Quantized where T: VectorRepr, + I: BfTreeId, { default_post_processor!(glue::Pipeline); } -impl<'a, T> InsertStrategy<'a, BfTreeProvider, &'a [T]> for Quantized +impl<'a, T, I> InsertStrategy<'a, BfTreeProvider, &'a [T]> for Quantized where T: VectorRepr, + I: BfTreeId, { type PruneStrategy = Self; fn prune_strategy(&self) -> Self::PruneStrategy { @@ -1411,9 +1458,10 @@ where } } -impl MultiInsertStrategy, B> for Quantized +impl MultiInsertStrategy, B> for Quantized where T: VectorRepr, + I: BfTreeId, B: glue::Batch, B: for<'a> Batch = &'a [T]> + Debug, { @@ -1428,24 +1476,24 @@ where fn finish( &self, - _provider: &BfTreeProvider, + _provider: &BfTreeProvider, _ctx: &DefaultContext, _batch: &std::sync::Arc, _ids: Itr, ) -> impl std::future::Future> + Send where - Itr: ExactSizeIterator + Send, + Itr: ExactSizeIterator + Send, { std::future::ready(Ok(())) } fn seeded_prune_accessor<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, _seed: &'a (), capacity: usize, - ) -> ANNResult> { + ) -> ANNResult> { QuantPruneAccessor::new(provider, capacity) } } @@ -1456,15 +1504,16 @@ where /// /// Same constraint as [`FullPrecision`]'s impl: not compatible with /// [`InplaceDeleteMethod::VisitedAndTopK`] due to hard deletes. -impl InplaceDeleteStrategy> for Quantized +impl InplaceDeleteStrategy> for Quantized where T: VectorRepr, + I: BfTreeId, { type DeleteElementError = ANNError; type DeleteElement<'a> = &'a [T]; type DeleteElementGuard = Box<[T]>; type PruneStrategy = Self; - type DeleteSearchAccessor<'a> = QuantAccessor<'a, T>; + type DeleteSearchAccessor<'a> = QuantAccessor<'a, T, I>; type SearchPostProcessor = Rerank; type SearchStrategy = Self; fn search_strategy(&self) -> Self::SearchStrategy { @@ -1481,30 +1530,31 @@ where async fn get_delete_element<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, - id: u32, + id: I, ) -> Result { use diskann::error::ErrorExt; provider .full_vectors - .get_vector_sync(id.into_usize()) + .get_vector_sync(id.as_index()) .escalate("get_delete_element: failed to read vector for inplace delete") .map(Into::into) } } // Pruning -impl PruneStrategy> for Quantized +impl PruneStrategy> for Quantized where T: VectorRepr, + I: BfTreeId, { - type PruneAccessor<'a> = QuantPruneAccessor<'a, T>; + type PruneAccessor<'a> = QuantPruneAccessor<'a, T, I>; type PruneAccessorError = ANNError; fn prune_accessor<'a>( &'a self, - provider: &'a BfTreeProvider, + provider: &'a BfTreeProvider, _context: &'a DefaultContext, capacity: usize, ) -> Result, Self::PruneAccessorError> { @@ -1516,32 +1566,33 @@ where #[derive(Debug, Default, Clone, Copy)] pub struct Rerank; -impl<'a, T> glue::SearchPostProcess, &[T]> for Rerank +impl<'a, T, I> glue::SearchPostProcess, &[T]> for Rerank where T: VectorRepr, + I: BfTreeId, { type Error = ANNError; - fn post_process( + fn post_process( &self, - accessor: &mut QuantAccessor<'a, T>, + accessor: &mut QuantAccessor<'a, T, I>, query: &[T], - candidates: I, + candidates: Itr, output: &mut B, ) -> impl Future> + Send where - I: Iterator> + Send, - B: SearchOutputBuffer + Send + ?Sized, + Itr: 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(); + let mut reranked: Vec<(I, f32)> = Vec::new(); for n in candidates { match provider .full_vectors - .get_vector_sync(n.id().into_usize()) + .get_vector_sync(n.id().as_index()) .allow_transient("stale candidate during rerank") { Ok(Some(vec)) => { @@ -1589,6 +1640,37 @@ pub struct SavedParams { /// Whether CPR snapshot support was enabled. #[serde(default)] pub use_snapshot: bool, + /// Width in bytes of the vertex id type the index was built with + /// (`size_of::()`). Defaults to 4 (`u32`) for indexes saved before this + /// field existed, all of which used `u32` ids. + #[serde(default = "default_id_width")] + pub id_width: usize, +} + +/// Default [`SavedParams::id_width`] for legacy indexes (all of which were `u32`). +fn default_id_width() -> usize { + std::mem::size_of::() +} + +/// Validate, on load, that the persisted id metadata is compatible with the id type +/// `I` the caller is loading as: the width must match what the index was built with, +/// and the capacity must fit in `I`. +fn validate_loaded_id_params(saved_params: &SavedParams) -> ANNResult<()> { + let expected = std::mem::size_of::(); + if saved_params.id_width != expected { + return Err(ANNError::message(format!( + "index was built with {}-byte vertex ids but is being loaded with a {expected}-byte \ + id type; load it with the matching id type", + saved_params.id_width, + ))); + } + let id_capacity = saved_params + .max_points + .checked_add(saved_params.frozen_points.get()) + .ok_or_else(|| { + ANNError::message("max_points + frozen_points overflows usize".to_string()) + })?; + crate::id::validate_id_capacity::(id_capacity) } /// The element type of the full-precision vectors stored in the index. @@ -1713,9 +1795,10 @@ fn load_bftree(snapshot_path: std::path::PathBuf, use_snapshot: bool) -> Result< // Serialization // ////////////////////// -impl SaveWith for BfTreeProvider +impl SaveWith for BfTreeProvider where T: VectorRepr, + I: BfTreeId, { type Ok = usize; type Error = ANNError; @@ -1746,6 +1829,7 @@ where graph_params: self.graph_params.clone(), is_memory: self.full_vectors.config().is_memory_backend(), use_snapshot: self.use_snapshot, + id_width: std::mem::size_of::(), }; debug_assert_eq!( @@ -1778,9 +1862,10 @@ where } } -impl LoadWith for BfTreeProvider +impl LoadWith for BfTreeProvider where T: VectorRepr, + I: BfTreeId, { type Error = ANNError; @@ -1798,6 +1883,8 @@ where })? }; + validate_loaded_id_params::(&saved_params)?; + let metric = Metric::from_str(&saved_params.metric) .map_err(|e| ANNError::message(lazy_format!(move, "Failed to parse metric: {}", e)))?; @@ -1816,10 +1903,8 @@ where BfTreePaths::neighbors_bftree(&saved_params.prefix), saved_params.use_snapshot, )?; - let neighbor_provider = NeighborProvider::::new_from_bftree( - saved_params.max_degree, - adjacency_list_index, - )?; + let neighbor_provider = + NeighborProvider::::new_from_bftree(saved_params.max_degree, adjacency_list_index)?; Ok(Self { quant_vectors: NoStore, @@ -1833,9 +1918,10 @@ where } } -impl SaveWith for BfTreeProvider +impl SaveWith for BfTreeProvider where T: VectorRepr, + I: BfTreeId, { type Ok = usize; type Error = ANNError; @@ -1872,6 +1958,7 @@ where graph_params: self.graph_params.clone(), is_memory: self.full_vectors.config().is_memory_backend(), use_snapshot: self.use_snapshot, + id_width: std::mem::size_of::(), }; debug_assert_eq!( @@ -1923,9 +2010,10 @@ where } } -impl LoadWith for BfTreeProvider +impl LoadWith for BfTreeProvider where T: VectorRepr, + I: BfTreeId, { type Error = ANNError; @@ -1943,6 +2031,8 @@ where })? }; + validate_loaded_id_params::(&saved_params)?; + let _quant_params = saved_params.quant_params.ok_or_else(|| { ANNError::message("Missing quant_params in saved params for quantized provider") })?; @@ -1965,10 +2055,8 @@ where BfTreePaths::neighbors_bftree(&saved_params.prefix), saved_params.use_snapshot, )?; - let neighbor_provider = NeighborProvider::::new_from_bftree( - saved_params.max_degree, - adjacency_list_index, - )?; + let neighbor_provider = + NeighborProvider::::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)?; @@ -2098,6 +2186,187 @@ mod tests { assert_eq!(*neighbors[0].id(), 3); } + /// Build and search a quantized index whose vertex ids are `u64` rather than the + /// default `u32`. This exercises the full generic pipeline end-to-end (set_element, + /// `u64`-keyed neighbor storage via `bytemuck`, greedy search, and rerank) to prove + /// that the `BfTreeProvider<_, _, u64>` path is functional and not merely compilable. + #[tokio::test] + async fn test_quantized_index_search_u64_ids() { + let start_point = Matrix::new(Init(|| 0.0f32), 1, 5); + let dim = 5; + let logical_max_degree = 6; + let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; + let metric = Metric::L2; + + let provider: BfTreeProvider = BfTreeProvider::new( + BfTreeProviderParameters { + max_points: 20, + num_start_points: NonZeroUsize::new(1).unwrap(), + dim, + metric, + max_degree: physical_max_degree, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + use_snapshot: false, + }, + start_point.as_view(), + create_test_quantizer(5), + ) + .unwrap(); + + let index_config = graph::config::Builder::new_with( + logical_max_degree as usize, + graph::config::MaxDegree::new(physical_max_degree as usize), + 10, + metric.into(), + |_| {}, + ) + .build() + .unwrap(); + + let index = Arc::new(DiskANNIndex::new(index_config, provider, None)); + let ctx = &DefaultContext; + + for i in 0u64..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(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(), 3u64); + } + + /// End-to-end insert + search where the vertex ids themselves live in the + /// `u64` range beyond `u32::MAX`. + /// + /// The other `u64` tests exercise the `u64` id *type* but only ever use id + /// *values* in `0..N`, which still fit in a `u32`; they would pass even if + /// the provider silently truncated keys to 32 bits. Because the provider + /// mints ids densely from `0..max_points + num_start_points` and stores them + /// in sparse bf-trees (construction is `O(num_start_points)`, not + /// `O(max_points)`), we can set `max_points` above `u32::MAX` and insert a + /// handful of points at real u64-range ids without a billion-scale dataset. + /// A truncating key path would alias these ids down into the `u32` range and + /// corrupt the graph, so a correct nearest-neighbor result here proves the + /// full data/quant/neighbor/search stack keys on the complete 8-byte id. + #[tokio::test] + async fn test_quantized_index_search_u64_high_ids() { + let start_point = Matrix::new(Init(|| 0.0f32), 1, 5); + let dim = 5; + let logical_max_degree = 6; + let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; + let metric = Metric::L2; + + // Base id above u32::MAX so every inserted vertex id (and the start + // point at `max_points`) sits in the u64-only range. `max_points` stays + // sparse, so this does not allocate ~4.3B slots. + let base: u64 = (u32::MAX as u64) + 1; + let num_points: u64 = 15; + + let provider: BfTreeProvider = BfTreeProvider::new( + BfTreeProviderParameters { + max_points: base as usize + num_points as usize + 1, + num_start_points: NonZeroUsize::new(1).unwrap(), + dim, + metric, + max_degree: physical_max_degree, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + use_snapshot: false, + }, + start_point.as_view(), + create_test_quantizer(5), + ) + .unwrap(); + + // Every start point id must be beyond the u32 range for this test to + // mean anything. + for id in provider.starting_points().unwrap() { + assert!( + id > u32::MAX as u64, + "start point id {id} should be in the u64-only range" + ); + } + + let index_config = graph::config::Builder::new_with( + logical_max_degree as usize, + graph::config::MaxDegree::new(physical_max_degree as usize), + 10, + metric.into(), + |_| {}, + ) + .build() + .unwrap(); + + let index = Arc::new(DiskANNIndex::new(index_config, provider, None)); + let ctx = &DefaultContext; + + // Insert point k (vector `[k; 5]`) at id `base + k`, so the id carries + // its high bits while the vector value stays comparable to the query. + for k in 0u64..num_points { + let id = base + k; + let point = vec![k as f32; 5]; + index + .insert(&Quantized, ctx, &id, point.as_slice()) + .await + .unwrap(); + } + + let query = vec![3.0; 5]; + let params = Knn::new(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" + ); + // The nearest neighbor to `[3.0; 5]` is the point with value 3, stored at + // id `base + 3`. Getting the full u64 id back (not a truncated variant) + // confirms the id survives the round trip through the whole stack. + assert_eq!(*neighbors[0].id(), base + 3); + assert!( + *neighbors[0].id() > u32::MAX as u64, + "returned id {} must be in the u64-only range", + neighbors[0].id() + ); + } + #[tokio::test] async fn test_quantized_index_multi_insert_search() { let index = create_quant_index(); @@ -2275,6 +2544,211 @@ mod tests { assert_eq!(*neighbors[0].id(), 3); } + /// Recall-parity validation for `u64` ids using a realistically-shaped + /// (but self-contained) dataset. + /// + /// Real ANN datasets ship dense `0..N` ids that all fit in a `u32`, so + /// nothing naturally exercises id values above `u32::MAX`, and building a + /// billion-scale dataset in a unit test is infeasible. Instead we *transpose* + /// a modest random dataset's ids up into the u64-only range by an order- + /// preserving offset (`row -> base + row`, `base = u32::MAX + 1`). Because the + /// offset is a bijection that preserves id ordering, it changes only the key + /// *values*, not the graph the algorithm builds: a correct 8-byte key path + /// must return the exact same neighbors (by row) as the `u32` identity-mapped + /// baseline. A path that truncated keys to 32 bits would alias the offset ids + /// back down, corrupt the graph, and diverge from the baseline. + /// + /// The test therefore asserts three things: + /// 1. every id returned by the `u64` run is above `u32::MAX` (the transpose + /// actually happened), + /// 2. the `u64` and `u32` runs return identical neighbor sets per query + /// (recall parity — the key width is transparent to results), and + /// 3. both runs achieve a sane recall floor against brute-force ground truth + /// (so the parity is between two *working* indexes, not two broken ones). + #[tokio::test] + async fn test_u64_id_recall_parity_on_transposed_dataset() { + const N: usize = 1500; + const DIM: usize = 16; + const NUM_QUERIES: usize = 30; + const K: usize = 10; + const L: usize = 64; + + // Deterministic splitmix64 generator so the dataset (and thus recall) is + // reproducible across runs and platforms without a dependency on `rand`. + struct SplitMix64(u64); + impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn next_f32(&mut self) -> f32 { + // Top 24 bits give a uniform value in [0, 1). + (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32 + } + } + + let mut rng = SplitMix64(0x0DDB_1A5E_5EED_1234); + let gen_vec = |rng: &mut SplitMix64| (0..DIM).map(|_| rng.next_f32()).collect::>(); + + let data: Vec> = (0..N).map(|_| gen_vec(&mut rng)).collect(); + let queries: Vec> = (0..NUM_QUERIES).map(|_| gen_vec(&mut rng)).collect(); + + // Exact brute-force top-K ground truth (row indices), ties broken by row. + let l2 = |a: &[f32], b: &[f32]| -> f32 { + a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::() + }; + let ground_truth: Vec> = queries + .iter() + .map(|q| { + let mut scored: Vec<(f32, usize)> = data + .iter() + .enumerate() + .map(|(i, v)| (l2(q, v), i)) + .collect(); + scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap().then(a.1.cmp(&b.1))); + scored.into_iter().take(K).map(|(_, i)| i).collect() + }) + .collect(); + + // Build a full-precision bf-tree index over id type `I`, insert every row + // at `id_of(row)`, and return the per-query neighbor ids. + async fn build_and_search( + data: &[Vec], + queries: &[Vec], + max_points: usize, + id_of: F, + ) -> Vec> + where + I: BfTreeId, + F: Fn(usize) -> I, + { + let logical_max_degree = 32usize; + let physical_max_degree = (logical_max_degree as f32 * 1.3) as u32; + let metric = Metric::L2; + let start_point = Matrix::new(Init(|| 0.0f32), 1, DIM); + + let provider: BfTreeProvider = BfTreeProvider::new( + BfTreeProviderParameters { + max_points, + num_start_points: NonZeroUsize::new(1).unwrap(), + dim: DIM, + metric, + max_degree: physical_max_degree, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + use_snapshot: false, + }, + start_point.as_view(), + NoStore, + ) + .unwrap(); + + let index_config = graph::config::Builder::new_with( + logical_max_degree, + graph::config::MaxDegree::new(physical_max_degree as usize), + L, + metric.into(), + |_| {}, + ) + .build() + .unwrap(); + + let index = Arc::new(DiskANNIndex::new(index_config, provider, None)); + let ctx = &DefaultContext; + + for (row, vector) in data.iter().enumerate() { + index + .insert(&FullPrecision, ctx, &id_of(row), vector.as_slice()) + .await + .unwrap(); + } + + let mut results = Vec::with_capacity(queries.len()); + for query in queries { + let params = Knn::new(L, None).unwrap(); + let mut neighbors = vec![Neighbor::::default(); K]; + let res = index + .search( + params, + &FullPrecision, + ctx, + query.as_slice(), + &mut BackInserter::new(neighbors.as_mut_slice()), + ) + .await + .unwrap(); + results.push( + neighbors[..res.result_count as usize] + .iter() + .map(|n| *n.id()) + .collect(), + ); + } + results + } + + // Baseline: identity mapping, dense u32 ids in 0..N. + let u32_results = build_and_search::(&data, &queries, N, |row| row as u32).await; + + // Transposed: u64 ids shifted entirely above the u32 range. + let base: u64 = (u32::MAX as u64) + 1; + let u64_results = + build_and_search::(&data, &queries, base as usize + N, |row| base + row as u64) + .await; + + let mut recall_hits_u32 = 0usize; + let mut recall_hits_u64 = 0usize; + let total = NUM_QUERIES * K; + + for qi in 0..NUM_QUERIES { + // Map both id spaces back to row indices. + let mut rows_u32: Vec = u32_results[qi].iter().map(|&id| id as usize).collect(); + let mut rows_u64: Vec = Vec::with_capacity(u64_results[qi].len()); + for &id in &u64_results[qi] { + assert!( + id > u32::MAX as u64, + "query {qi}: transposed id {id} must live in the u64-only range" + ); + rows_u64.push((id - base) as usize); + } + + // Recall parity: the key width must be transparent to the results. + let mut sorted_u32 = rows_u32.clone(); + let mut sorted_u64 = rows_u64.clone(); + sorted_u32.sort_unstable(); + sorted_u64.sort_unstable(); + assert_eq!( + sorted_u32, sorted_u64, + "query {qi}: u64-transposed run returned a different neighbor set than the u32 baseline" + ); + + // Recall against brute-force ground truth for both runs. + let gt: std::collections::HashSet = ground_truth[qi].iter().copied().collect(); + rows_u32.retain(|r| gt.contains(r)); + rows_u64.retain(|r| gt.contains(r)); + recall_hits_u32 += rows_u32.len(); + recall_hits_u64 += rows_u64.len(); + } + + let recall_u32 = recall_hits_u32 as f64 / total as f64; + let recall_u64 = recall_hits_u64 as f64 / total as f64; + + assert_eq!( + recall_hits_u32, recall_hits_u64, + "recall must be identical between the u32 baseline and the u64-transposed run" + ); + assert!( + recall_u64 > 0.8, + "recall floor not met (u32={recall_u32:.4}, u64={recall_u64:.4}); \ + parity is only meaningful between two working indexes" + ); + } + #[tokio::test] async fn test_full_precision_delete_and_search() { let index = create_full_precision_index(); @@ -2366,7 +2840,7 @@ mod tests { // Iterator // - assert_eq!((&provider).into_iter(), 0..(10 + 2)); + assert!((&provider).into_iter().eq(0u32..(10 + 2))); let iter = provider.iter(); @@ -3086,4 +3560,129 @@ mod tests { panic!("NeighborProvider should succeed: {e}"); } } + + /// A `u64`-id provider must round-trip through save/load with the wider on-disk + /// neighbor format, and loading the same bytes with a mismatched id width must fail. + #[tokio::test] + async fn test_bf_tree_provider_save_load_u64_ids() { + let num_points = 16usize; + let dim = 4usize; + let max_degree = 16u32; + let num_start_points = NonZeroUsize::new(2).unwrap(); + let ctx = &DefaultContext; + + let temp_dir = tempdir().unwrap(); + let prefix = temp_dir + .path() + .join("u64_provider") + .to_string_lossy() + .to_string(); + + let mut vector_config = Config::new(BfTreePaths::vectors_bftree(&prefix), 1024 * 1024); + vector_config.storage_backend(bf_tree::StorageBackend::Std); + vector_config.use_snapshot(true); + + let mut neighbor_config = Config::new(BfTreePaths::neighbors_bftree(&prefix), 1024 * 1024); + neighbor_config.storage_backend(bf_tree::StorageBackend::Std); + neighbor_config.use_snapshot(true); + + let params = BfTreeProviderParameters { + max_points: num_points, + num_start_points, + dim, + metric: Metric::L2, + max_degree, + vector_provider_config: vector_config, + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: neighbor_config, + graph_params: None, + use_snapshot: true, + }; + + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim); + let provider = + BfTreeProvider::::new(params, start_points.as_view(), NoStore) + .unwrap(); + + for i in 0..num_points { + let vector: Vec = (0..dim).map(|j| (i * dim + j) as f32 * 0.1).collect(); + provider + .set_element(ctx, &(i as u64), &vector) + .await + .unwrap(); + } + + let mut scratch = provider.neighbor_provider.scratch(&provider.locks); + for i in 0..num_points as u64 { + let neighbors: Vec = (0..std::cmp::min(i, max_degree as u64)) + .map(|j| (i + j) % num_points as u64) + .collect(); + scratch.write_neighbors(i, &neighbors).unwrap(); + } + drop(scratch); + + let storage = FileStorageProvider; + let save_dir = tempdir().unwrap(); + let save_prefix = save_dir + .path() + .join("saved_u64_provider") + .to_string_lossy() + .to_string(); + provider.save_with(&storage, &save_prefix).await.unwrap(); + + let loaded = BfTreeProvider::::load_with(&storage, &save_prefix) + .await + .unwrap(); + + for i in 0..num_points as u64 { + let mut original_list = AdjacencyList::new(); + let mut loaded_list = AdjacencyList::new(); + provider + .neighbor_provider + .get_neighbors(i, &mut original_list) + .unwrap(); + loaded + .neighbor_provider + .get_neighbors(i, &mut loaded_list) + .unwrap(); + assert_eq!(&*original_list, &*loaded_list, "neighbor mismatch at {i}"); + } + + // Loading a u64-saved index as u32 must be rejected, not silently misinterpreted. + let mismatch = BfTreeProvider::::load_with(&storage, &save_prefix).await; + assert!( + mismatch.is_err(), + "loading a u64 index with a u32 id type must fail" + ); + } + + /// Constructing a provider whose vertex count cannot be represented by the id type + /// must fail eagerly rather than silently truncate ids. + #[tokio::test] + async fn test_new_rejects_capacity_exceeding_id_type() { + let dim = 4usize; + let num_start_points = NonZeroUsize::new(1).unwrap(); + let start_points = Matrix::new(Init(|| 0.0f32), num_start_points.into(), dim); + + let params = BfTreeProviderParameters { + // Largest index would be u32::MAX + 1, which a u32 id cannot hold. + max_points: u32::MAX as usize + 2, + num_start_points, + dim, + metric: Metric::L2, + max_degree: 8, + vector_provider_config: Config::default(), + quant_vector_provider_config: Config::default(), + neighbor_list_provider_config: Config::default(), + graph_params: None, + use_snapshot: false, + }; + + let result = + BfTreeProvider::::new(params, start_points.as_view(), NoStore); + assert!( + result.is_err(), + "u32 provider must reject a capacity exceeding u32::MAX" + ); + } } diff --git a/diskann-bftree/src/vectors.rs b/diskann-bftree/src/vectors.rs index ab535401f..96d3cf9a6 100644 --- a/diskann-bftree/src/vectors.rs +++ b/diskann-bftree/src/vectors.rs @@ -96,11 +96,14 @@ impl VectorProvider { /// Return a vector of vector Ids of the starting points /// #[inline(always)] - pub fn starting_points(&self) -> ANNResult> { + pub fn starting_points(&self) -> ANNResult> { (self.max_vectors..self.total()) .map(|i| { - u32::try_from(i).map_err(|_| { - ANNError::message(lazy_format!(move, "start point id {i} exceeds u32::MAX")) + I::try_from_index(i).ok_or_else(|| { + ANNError::message(lazy_format!( + move, + "start point id {i} exceeds the id type's maximum" + )) }) }) .collect()