diff --git a/Cargo.lock b/Cargo.lock index 68dabeac21..831b0ea6d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -438,6 +438,7 @@ dependencies = [ "bytemuck", "dashmap", "diskann-linalg", + "diskann-quantization", "diskann-utils", "diskann-vector", "diskann-wide", diff --git a/diskann-utils/src/views/rowmajor.rs b/diskann-utils/src/views/rowmajor.rs index 74214b8519..ec8db80169 100644 --- a/diskann-utils/src/views/rowmajor.rs +++ b/diskann-utils/src/views/rowmajor.rs @@ -538,6 +538,16 @@ pub unsafe trait MatrixMut: Matrix { }) } + /// Divide the matrix into disjoint mutable sub-matrices with up to `batchsize` + /// rows and the same column count. + /// + /// The final window may have fewer rows. A matrix with zero columns still + /// yields windows with the corresponding row counts. Windows may coexist, + /// and their lifetime is tied to the mutable borrow of the matrix. + fn window_iter_mut(&mut self, batchsize: NonZeroUsize) -> iter::WindowsMut<'_, Self::Element> { + iter::WindowsMut::new(self.as_view_mut(), batchsize) + } + //-------// // Rayon // //-------// @@ -2478,10 +2488,96 @@ mod tests { m.window_iter(NonZeroUsize::new(batchsize).unwrap()), nrows.div_ceil(batchsize), ); + assert_exact_fused( + m.window_iter_mut(NonZeroUsize::new(batchsize).unwrap()), + nrows.div_ceil(batchsize), + ); + } + } + } + + #[test] + fn mutable_windows_are_disjoint_and_match_scalar_indexing() { + for (nrows, ncols) in [ + (0, 0), + (0, 4), + (3, 0), + (1, 1), + (1, 4), + (4, 1), + (5, 3), + (6, 2), + ] { + let original = striped_matrix(nrows, ncols); + for batchsize in [1, 2, 3, usize::MAX] { + let mut matrix = original.clone(); + { + let mut view = matrix.as_view_mut(); + let windows: Vec<_> = view + .window_iter_mut(NonZeroUsize::new(batchsize).unwrap()) + .collect(); + assert_eq!(windows.len(), nrows.div_ceil(batchsize)); + for (index, mut window) in windows.into_iter().enumerate() { + let start = index * batchsize; + assert_eq!(window.nrows(), batchsize.min(nrows - start)); + assert_eq!(window.ncols(), ncols); + for (row, values) in window.rows_mut().enumerate() { + for (col, value) in values.iter_mut().enumerate() { + assert_eq!(*value, *original.element(start + row, col)); + *value += index + 1; + } + } + } + } + for row in 0..nrows { + for col in 0..ncols { + assert_eq!( + *matrix.element(row, col), + original.element(row, col) + row / batchsize + 1, + ); + } + } } } } + #[test] + fn mutable_windows_preserve_zero_sized_shapes() { + let mut matrix = Owned::from_element(5, 3, ()); + let windows: Vec<_> = matrix + .window_iter_mut(NonZeroUsize::new(2).unwrap()) + .collect(); + assert_eq!( + windows + .iter() + .map(|window| (window.nrows(), window.ncols())) + .collect::>(), + [(2, 3), (2, 3), (1, 3)], + ); + assert_eq!( + windows + .iter() + .map(|window| window.as_slice().len()) + .collect::>(), + [6, 6, 3], + ); + + let mut matrix = Owned::from_element(usize::MAX, 0, 0); + let ptr = matrix.as_ptr(); + let mut windows = matrix.window_iter_mut(NonZeroUsize::new(usize::MAX - 1).unwrap()); + assert_eq!(windows.len(), 2); + let first = windows.next().unwrap(); + assert_eq!(windows.len(), 1); + let last = windows.next().unwrap(); + assert_eq!(windows.len(), 0); + assert_eq!((first.nrows(), first.ncols()), (usize::MAX - 1, 0)); + assert_eq!((last.nrows(), last.ncols()), (1, 0)); + assert_eq!(first.as_ptr(), ptr); + assert_eq!(last.as_ptr(), ptr); + assert!(windows.next().is_none()); + assert!(windows.next().is_none()); + } + #[test] fn matrix_transformations_preserve_empty_shapes() { for (nrows, ncols) in [(0, 0), (0, 4), (3, 0)] { diff --git a/diskann-utils/src/views/rowmajor/iter.rs b/diskann-utils/src/views/rowmajor/iter.rs index 3024aa6b4e..a84f0beaf8 100644 --- a/diskann-utils/src/views/rowmajor/iter.rs +++ b/diskann-utils/src/views/rowmajor/iter.rs @@ -5,9 +5,7 @@ use std::{marker::PhantomData, num::NonZeroUsize, ptr::NonNull}; -#[cfg(feature = "rayon")] -use crate::views::rowmajor::MatrixMut; -use crate::views::rowmajor::{Layout, Matrix, Mut, Ref}; +use crate::views::rowmajor::{Layout, Matrix, MatrixMut, Mut, Ref}; //------// // Rows // @@ -190,6 +188,77 @@ impl<'a, T> Iterator for Windows<'a, T> { impl ExactSizeIterator for Windows<'_, T> {} impl std::iter::FusedIterator for Windows<'_, T> {} +//------------// +// WindowsMut // +//------------// + +/// An iterator over disjoint mutable windows. See: [`MatrixMut::window_iter_mut`]. +#[derive(Debug)] +pub struct WindowsMut<'a, T> { + ptr: NonNull, + remaining: usize, + batchsize: NonZeroUsize, + ncols: usize, + _lifetime: PhantomData<&'a mut [T]>, +} + +impl<'a, T> WindowsMut<'a, T> { + pub(super) fn new(mut matrix: Mut<'a, T>, batchsize: NonZeroUsize) -> Self { + let layout = matrix.layout(); + Self { + ptr: matrix.as_nonnull_mut(), + remaining: layout.nrows(), + batchsize, + ncols: layout.ncols(), + _lifetime: PhantomData, + } + } +} + +// SAFETY: `WindowsMut` owns an exclusive slice borrow, so sending it requires `T: Send`. +unsafe impl Send for WindowsMut<'_, T> {} +// SAFETY: Shared access to `WindowsMut` exposes only shared access to `T`. +unsafe impl Sync for WindowsMut<'_, T> {} + +impl<'a, T> Iterator for WindowsMut<'a, T> { + type Item = Mut<'a, T>; + fn next(&mut self) -> Option> { + if self.remaining == 0 { + None + } else { + let next_remaining = self.remaining.saturating_sub(self.batchsize.get()); + let nrows = self.remaining - next_remaining; + + // SAFETY: `self.ptr` starts the remaining suffix of a valid `Mut`, and `nrows` + // does not exceed that suffix. Keeping the parent's column count therefore + // produces a valid subview and a layout no larger than the parent layout. + // Previously yielded windows precede this suffix and overlap no elements. + let window = unsafe { + Mut { + ptr: self.ptr, + layout: Layout::new_unchecked(nrows, self.ncols), + _lifetime: PhantomData, + } + }; + + // SAFETY: Advancing by the yielded window remains within or one past the + // original matrix span. The validated parent layout guarantees that the + // multiplication and pointer offset are representable. + self.ptr = unsafe { self.ptr.add(nrows * self.ncols) }; + self.remaining = next_remaining; + Some(window) + } + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.remaining.div_ceil(self.batchsize.get()); + (remaining, Some(remaining)) + } +} + +impl ExactSizeIterator for WindowsMut<'_, T> {} +impl std::iter::FusedIterator for WindowsMut<'_, T> {} + //--------// // ParMut // //--------// diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 630e0f46d0..771d35e835 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -14,6 +14,7 @@ targets = ["x86_64-unknown-linux-gnu", "aarch64-pc-windows-msvc", "x86_64-pc-win [dependencies] anyhow.workspace = true bytemuck = { workspace = true, features = ["must_cast"]} +diskann-quantization = { workspace = true } diskann-utils = { workspace = true, default-features = false } futures-util = { workspace = true, default-features = false } half = { workspace = true, features = ["bytemuck", "num-traits"] } diff --git a/diskann/src/ivf/build/mod.rs b/diskann/src/ivf/build/mod.rs new file mode 100644 index 0000000000..7a1effdf55 --- /dev/null +++ b/diskann/src/ivf/build/mod.rs @@ -0,0 +1,36 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Complete regional split operations, independent of physical list storage. + +use std::fmt::Debug; + +use crate::{ + ANNError, + ivf::region::{Update, View}, +}; + +pub(crate) mod twomeans; +pub(crate) use twomeans::TwoMeansSplit; + +/// Replace parents and reassign their members together with selected neighbors. +pub(crate) trait Split> { + type Error: Debug + Send + Sync + Into; + + /// Parents and neighbors must be distinct, disjoint filled sources. + /// Other filled sources must not be edited. + /// + /// # Errors + /// + /// A failure may leave partial private updates. The caller must discard the + /// operation rather than finalize it. + fn split>( + &mut self, + region: &R, + update: &mut U, + parents: &[L], + neighbors: &[L], + ) -> Result<(), Self::Error>; +} diff --git a/diskann/src/ivf/build/twomeans.rs b/diskann/src/ivf/build/twomeans.rs new file mode 100644 index 0000000000..550f6a0e17 --- /dev/null +++ b/diskann/src/ivf/build/twomeans.rs @@ -0,0 +1,216 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Two-means splitting and nearest-centroid regional assignment. + +use std::{fmt::Debug, num::NonZeroUsize}; + +use diskann_quantization::algorithms::kmeans::lloyds::lloyds; +use diskann_utils::{ + Reborrow, + views::rowmajor::{self, Matrix, MatrixMut}, +}; +use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; + +use super::Split; +use crate::{ + ANNError, ANNResult, + ivf::region::{AsMatRef, Update, View}, + utils::{TWO, VectorId}, +}; + +//////////// +// Errors // +//////////// + +#[derive(Debug, thiserror::Error)] +pub(crate) enum TwoMeansError { + #[error("IVF two-means could not retrieve list {0:?} from the region")] + UnableToRetrieve(L), + #[error("IVF two-means needs at least two points in list {list:?}, got {actual}")] + NotEnoughPoints { list: L, actual: usize }, + #[error("Centroid count overflows for {0} parents")] + Overflow(usize), + #[error("Parents is empty")] + EmptyParents, +} + +impl From> for ANNError { + #[track_caller] + fn from(error: TwoMeansError) -> Self { + Self::new(error) + } +} + +/// Two-means per parent followed by nearest-centroid assignment over the region. +/// +/// Seeds are the first two rows of each parent. Lloyd iterations reset empty +/// centers to zero. One centroid matrix is allocated up front for all parents. +/// Retained neighbors precede new children when resolving equal-distance assignments. +/// All computation completes before any updates are recorded. +pub struct TwoMeansSplit { + iterations: NonZeroUsize, + dim: NonZeroUsize, +} + +impl TwoMeansSplit { + pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { + Self { iterations, dim } + } + + /// Fit and output the centroids given by the region and the parents. + fn fit(&self, region: &R, parents: &[L]) -> ANNResult> + where + L: VectorId, + R: View, + for<'a> R::VectorsRef<'a>: AsMatRef<'a, f32>, + { + let rows = parents + .len() + .checked_mul(TWO.get()) + .ok_or(TwoMeansError::::Overflow(parents.len()))?; + + let mut centers = rowmajor::Owned::try_from_element(rows, self.dim.get(), 0.0)?; + + for (&list, mut dst) in parents.iter().zip(centers.window_iter_mut(TWO)) { + let vectors = region + .get_vectors(list) + .ok_or(TwoMeansError::UnableToRetrieve(list))?; + + let points = vectors.reborrow().as_view(); + + dst.as_mut_slice().copy_from_slice( + points + .subview(0..TWO.get()) + .ok_or(TwoMeansError::NotEnoughPoints { + list, + actual: points.nrows(), + })? + .as_slice(), + ); + + let _ = lloyds(points, dst, self.iterations.get()); + } + + Ok(centers) + } +} + +///////////// +// Targets // +///////////// + +#[derive(Clone, Copy)] +/// Defines the category of the target centroid. +enum TargetKind { + // Existing centroid, from a neighbor for example + Exists(L), + // New centroid as a result of a split. + New, +} + +#[derive(Clone)] +/// Target object holding the target kind and reference to centroid vector. +struct Target<'a, L> { + target: TargetKind, + centroid: &'a [f32], +} + +/////////// +// Split // +/////////// + +impl Split for TwoMeansSplit +where + I: VectorId, + L: VectorId, + R: View, + for<'a> R::VectorsRef<'a>: AsMatRef<'a, f32>, +{ + type Error = ANNError; + + fn split>( + &mut self, + region: &R, + update: &mut U, + parents: &[L], + neighbors: &[L], + ) -> ANNResult<()> { + if parents.is_empty() { + return Err(TwoMeansError::::EmptyParents.into()); + } + + // Fit children + + let children = self.fit(region, parents)?; + + let mut targets = Vec::new(); + + for &neighbor in neighbors { + let view = region + .get_list(neighbor) + .ok_or(TwoMeansError::UnableToRetrieve(neighbor))?; + + targets.push(Target { + target: TargetKind::Exists(neighbor), + centroid: view.centroid, + }); + } + + targets.extend(children.rows().map(|centroid| Target { + target: TargetKind::New, + centroid, + })); + + // Assign + + let mut memberships = vec![Vec::new(); targets.len()]; + + for &list in parents.iter().chain(neighbors.iter()) { + let view = region + .get_list(list) + .ok_or(TwoMeansError::UnableToRetrieve(list))?; + + let vectors = region + .get_vectors(list) + .ok_or(TwoMeansError::UnableToRetrieve(list))?; + + let points = vectors.reborrow().as_view(); + + for (&id, point) in view.members.iter().zip(points.rows()) { + memberships[nearest(point, &targets)].push(id); + } + } + + // Update + + for (target, members) in targets.iter().zip(&memberships) { + match target.target { + TargetKind::New => update.create(target.centroid, members)?, + TargetKind::Exists(list) => update.set_members(list, members)?, + }; + } + + for &parent in parents { + update.retire(parent)?; + } + Ok(()) + } +} + +/// Destinations are nonempty; the earliest wins ties, including infinite distances. +#[inline(always)] +fn nearest(point: &[f32], targets: &[Target<'_, L>]) -> usize { + let mut best = 0; + let mut distance: f32 = SquaredL2::evaluate(point, targets[0].centroid); + for (index, target) in targets.iter().enumerate().skip(1) { + let candidate = SquaredL2::evaluate(point, target.centroid); + if candidate < distance { + best = index; + distance = candidate; + } + } + best +} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs new file mode 100644 index 0000000000..0b87365fb4 --- /dev/null +++ b/diskann/src/ivf/index.rs @@ -0,0 +1,173 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! IVF policy and orchestration, independent of physical list storage. + +use std::{collections::BTreeSet, marker::PhantomData, num::NonZeroUsize}; + +use diskann_utils::future::SendFuture; +use thiserror::Error; + +use super::traits::{BuildAccessor, NearestList, SearchAccessor, SearchStrategy}; +use crate::{ + ANNResult, + error::IntoANNResult, + graph::SearchOutputBuffer, + ivf::{ + build::{Split, TwoMeansSplit}, + region::AsMatRef, + traits::BuildStrategy, + }, + neighbor::NeighborPriorityQueue, + provider::ExecutionContext, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ReassignCount { + None, + Count(NonZeroUsize), +} + +/// Parameters for incremental two-means splitting. +#[derive(Debug, Clone, Copy)] +pub struct Config { + /// Split a touched list when its pending membership exceeds this size. + pub split_threshold: NonZeroUsize, + /// Nearest-list query count per split parent, before filtering split parents. + /// + /// - [`ReassignCount::None`] corresponds to no neighborhood gathering. + /// - [`ReassignCount::Count`] coresponds to a non-zero neighborhood size to be gathered. + pub reassign_count: ReassignCount, + /// Positive number of Lloyd steps per parent. + pub split_iterations: NonZeroUsize, + pub seed: u64, +} + +#[derive(Debug, Error)] +#[error("split_threshold must be at least 2, got {0}")] +pub struct ConfigError(pub usize); + +/// Work successfully published by one insertion. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct InsertStats { + pub inserted: usize, + pub splits: usize, +} + +/// An IVF index with serialized mutation through the underlying provider. +#[derive(Debug)] +pub struct IVFIndex { + config: Config, + phantom: PhantomData, +} + +impl IVFIndex { + /// # Errors + /// + /// Returns [`ConfigError`] when the split threshold is less than two. + pub fn new(config: Config) -> Result { + Ok(Self { + config, + phantom: PhantomData::default(), + }) + } + + pub fn insert_batch<'a, S, B>( + &'a mut self, + strategy: &'a S, + context: &'a C, + points: B, + ) -> impl SendFuture> + where + S: BuildStrategy<'a, B, C>, + S::Accessor: for<'b> BuildAccessor: AsMatRef<'b, f32>>, + B: Send + Sync, + { + let config = self.config; + + async move { + let mut accessor = strategy.build_accessor(context).into_ann_result()?; + + let dim = accessor.dim(); + + let mut parents = BTreeSet::new(); + + let inserted = accessor + .route(points, |&list, size| { + if size.get() > config.split_threshold.get() { + parents.insert(list); + } + }) + .await?; + + if !parents.is_empty() { + let mut neighbors = BTreeSet::new(); + + if let ReassignCount::Count(n) = config.reassign_count { + let router = accessor.nearest(); + + for &pid in parents.iter() { + router + .find( + pid, + n, + |id| !parents.contains(id), + |id, _| { + neighbors.insert(*id); + }, + ) + .await?; + } + } + + let iter = parents.iter().chain(neighbors.iter()).copied(); + + let (region, mut update) = accessor.fill(iter).await?; + + TwoMeansSplit::new(config.split_iterations, dim).split( + ®ion, + &mut update, + &parents.iter().copied().collect::>(), + &neighbors.iter().copied().collect::>(), + )?; + }; + + accessor.finish().await?; + + Ok(InsertStats { + inserted: inserted.get(), + splits: parents.len(), + }) + } + } + + /// Append the best `k` candidates from the provider-selected lists to `output`. + /// + /// # Errors + /// + /// Accessor construction or scan failures leave `output` unchanged. + pub fn search<'a, S, T, OB>( + &'a self, + strategy: &'a S, + context: &'a C, + query: T, + k: NonZeroUsize, + nprobe: NonZeroUsize, + output: &mut OB, + ) -> impl SendFuture> + where + S: SearchStrategy<'a, T, C>, + T: Copy + Send + Sync, + OB: SearchOutputBuffer + Send + ?Sized, + { + async move { + let accessor = strategy.search_accessor(context, query).into_ann_result()?; + let mut queue = NeighborPriorityQueue::new(k.get()); + accessor.scan(nprobe, &mut queue).await?; + + Ok(output.extend(queue.iter())) + } + } +} diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs new file mode 100644 index 0000000000..0e7f0c6cbb --- /dev/null +++ b/diskann/src/ivf/mod.rs @@ -0,0 +1,51 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Incremental IVF construction and search. +//! +//! The index routes points, splits oversized lists, and assigns their members to +//! replacement centroids. Providers supply routing and operation-local appends +//! through [`traits::BuildAccessor`], borrowed inputs through [`region::View`], +//! and regional writes through [`region::Update`]. +//! +//! Neighbor selection uses [`traits::NearestList`] before filling the unique +//! union of parents and retained neighbors. Children are fitted only from parent +//! points, then all gathered points are reassigned across every child and retained +//! neighbor centroid. This permits movement between retained neighbors as well. +//! Parent centroids are not assignment candidates; retained centroids are unchanged. +//! Regional splitting owns fitting, assignment, and updates as one operation. +//! The built-in two-means implementation fits each parent independently, then assigns +//! all gathered points to their nearest destination. Child centroids are owned; +//! retained centroids and source vectors remain borrowed. +//! +//! Clustering finishes reading its inputs before recording child creation and +//! parent retirement and neighbor membership replacement through the separate +//! update object. The region's input +//! borrows end before the accessor finalizes regional writes and ordinary appends +//! together. Providers may buffer writes or apply them inside a private operation +//! context. +//! Providers choose their own storage, publication, and failure-recovery mechanisms. +//! +//! A query-bound [`traits::SearchAccessor`] selects and scans lists, controlling its +//! own scheduling. It can reduce independent [`traits::Accumulator`] states before +//! returning. The index writes the final top-k only after the scan succeeds. +//! +//! This first implementation supports batch insertion into an already initialized +//! provider, with two-means splitting and optional neighborhood reassignment. +//! Setting [`Config::reassign_count`] to [`ReassignCount::None`] preserves +//! parent-only splitting. +//! Each touched parent is split once; neighbors that become oversized are not +//! recursively split. +//! Initialization, deletion, and colocation are not implemented. + +pub mod build; +pub mod index; +pub mod region; +pub mod traits; + +pub use index::{Config, ConfigError, IVFIndex, InsertStats, ReassignCount}; + +#[cfg(test)] +mod tests; diff --git a/diskann/src/ivf/region.rs b/diskann/src/ivf/region.rs new file mode 100644 index 0000000000..ef70947459 --- /dev/null +++ b/diskann/src/ivf/region.rs @@ -0,0 +1,73 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Regional input views and write operations. + +use std::fmt::Debug; + +use diskann_utils::{Reborrow, views::rowmajor}; + +use crate::ANNResult; + +/// One source list's centroid and complete operation-local membership. +/// +/// Only reference metadata is owned; point and centroid coordinates are borrowed. +/// A region may include both split targets and neighboring lists. +#[derive(Debug)] +pub struct List<'a, I, L> { + pub id: L, + pub centroid: &'a [f32], + pub members: &'a [I], +} + +/// Views for one build operation to perform splits/re-assignments . +pub trait View: Send { + /// The vector representation consumed by numerical algorithms. + type VectorsRef<'a>; + + /// The returned handle supports shorter borrows for numerical algorithms. + type Vectors<'a>: for<'b> Reborrow<'b, Target = Self::VectorsRef<'b>> + where + Self: 'a; + + /// Borrow a filled source's centroid and complete operation-local membership. + fn get_list(&self, i: L) -> Option>; + + /// Borrow canonical rows in the same order as the source's member IDs. + /// + /// There is one row per member, including pending accessor appends. Coordinates + /// have the accessor's canonical dimension and are finite. + fn get_vectors(&self, i: L) -> Option>; +} + +/// Regional writes for one build operation. +pub trait Update: Send { + /// Record a complete membership replacement. + fn set_members(&mut self, i: L, ids: &[I]) -> ANNResult<()>; + + /// Record retirement of a source list and its centroid. + fn retire(&mut self, i: L) -> ANNResult<()>; + + /// Record a new centroid and its complete membership. + fn create(&mut self, centroid: &[f32], members: &[I]) -> ANNResult<()>; +} + +///////////// +/// AsMat /// +///////////// + +pub trait AsMatRef<'a, T> { + fn as_view(self) -> rowmajor::Ref<'a, T>; +} + +impl<'a, T> AsMatRef<'a, T> for rowmajor::Ref<'a, T> { + fn as_view(self) -> rowmajor::Ref<'a, T> { + self + } +} + +pub trait AsMatMut<'a, T>: AsMatRef<'a, T> { + fn as_mut_view(self) -> rowmajor::Mut<'a, T>; +} diff --git a/diskann/src/ivf/tests/mod.rs b/diskann/src/ivf/tests/mod.rs new file mode 100644 index 0000000000..6a6a82d109 --- /dev/null +++ b/diskann/src/ivf/tests/mod.rs @@ -0,0 +1,1775 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + cell::Cell, + collections::{BTreeMap, BTreeSet}, + future::Future, + num::NonZeroUsize, +}; + +use diskann_utils::{ + Reborrow, + future::SendFuture, + views::rowmajor::{self, Matrix, MatrixMut}, +}; +use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; +use tokio::sync::{Mutex, MutexGuard}; + +use self::{ + pending::{Membership, PendingLists}, + support::*, +}; +use super::{ + Config, IVFIndex, InsertStats, ReassignCount, + build::{Split, TwoMeansSplit, twomeans::TwoMeansError}, + region::{AsMatRef, List, Update, View}, + traits::{ + Accumulator, BuildAccessor, BuildStrategy, HasIds, NearestList, SearchAccessor, + SearchStrategy, + }, +}; +use crate::{ + ANNError, ANNResult, + graph::{IdDistance, SearchOutputBuffer}, + neighbor::{BackInserter, Neighbor, NeighborPriorityQueue}, + provider::DefaultContext, + utils::ONE, +}; + +mod pending; +mod support; + +#[tokio::test] +async fn routing_callbacks_emit_ranked_results() { + let expected: [(u32, f32); 3] = [(1, 25.0), (2, 25.0), (3, 8836.0)]; + for store in [seeded(), Store::default()] { + let router = Coarse { + stored: &store.lists, + fail_route: false, + fail_nearest: false, + }; + for limit in [1, 2, usize::MAX] { + let mut emitted = Vec::new(); + require_send(router.route( + &[6u8, 0][..], + NonZeroUsize::new(limit).unwrap(), + |id, distance| emitted.push((*id, distance)), + )) + .await + .unwrap(); + let expected = &expected[..limit.min(store.lists.len())]; + assert_eq!(emitted, expected); + } + } +} + +#[tokio::test] +async fn nearest_callbacks_filter_before_limiting_and_count_emitted_results() { + let store = seeded(); + let nearest = Coarse { + stored: &store.lists, + fail_route: false, + fail_nearest: false, + }; + for (limit, excluded, expected) in [ + (1, vec![1], vec![(2, 100.0_f32)]), + (2, vec![1], vec![(2, 100.0), (3, 9801.0)]), + (usize::MAX, vec![], vec![(1, 0.0), (2, 100.0), (3, 9801.0)]), + (2, vec![1, 2, 3], vec![]), + ] { + let mut emitted = Vec::new(); + let count = require_send(nearest.find( + 1, + NonZeroUsize::new(limit).unwrap(), + |id| !excluded.contains(id), + |id, distance| emitted.push((*id, distance)), + )) + .await + .unwrap(); + assert_eq!(count, expected.len()); + assert_eq!(emitted, expected); + } +} + +#[tokio::test] +async fn nearest_lookup_failures_do_not_emit_results() { + let store = seeded(); + for (fail_nearest, id, expected) in [ + (false, 99, "unknown list 99"), + (true, 1, "injected nearest failure"), + ] { + let nearest = Coarse { + stored: &store.lists, + fail_route: false, + fail_nearest, + }; + let mut emitted = Vec::new(); + let error = require_send(nearest.find( + id, + ONE, + |_| true, + |id, distance| emitted.push((*id, distance)), + )) + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + assert!(emitted.is_empty()); + } +} + +#[tokio::test] +async fn routing_preserves_native_inputs_and_id_order() { + let strategy = Strategy::new(seeded()); + let mut build: Build<'_, u8> = strategy.build_accessor(&DefaultContext).unwrap(); + let points: [(u32, &[u8]); 3] = [(9, &[11, 0]), (7, &[1, 0]), (8, &[0, 0])]; + let mut emitted = Vec::new(); + let inserted = require_send(build.route(&points[..], |&list, size| { + emitted.push((list, size.get())); + })) + .await + .unwrap(); + assert_eq!(inserted.get(), points.len()); + assert_eq!(emitted, [(1, 4), (2, 3)]); + assert_eq!( + build + .pending + .members(1, &build.store.lists[&1].members) + .unwrap() + .collect::>(), + [10_000, 10_001, 10_007, 10_008] + ); + assert_eq!( + build + .pending + .members(2, &build.store.lists[&2].members) + .unwrap() + .collect::>(), + [10_002, 10_003, 10_009] + ); + for point in &points { + let id = point.0 + 10_000; + assert_eq!(build.staged[&id], point.1); + assert!(!build.store.points.contains_key(&id)); + } + let router = build.nearest(); + for (point, expected) in points.iter().zip([2, 1, 1]) { + let mut selected = Vec::new(); + require_send(router.route(point.1, ONE, |&id, _| selected.push(id))) + .await + .unwrap(); + assert_eq!(selected, [expected]); + } + drop(build); + let store = strategy.store(); + assert_eq!(store.points.len(), 5); + assert_eq!(store.finishes, 0); +} + +#[tokio::test] +async fn routing_reports_sizes_including_previous_pending_appends() { + let strategy = Strategy::new(seeded()); + let mut build: Build<'_, u8> = strategy.build_accessor(&DefaultContext).unwrap(); + let points: [(u32, &[u8]); 2] = [(5, &[1, 0]), (6, &[2, 0])]; + let mut emitted = Vec::new(); + for point in points { + let inserted = require_send(build.route(std::slice::from_ref(&point), |&list, size| { + emitted.push((list, size.get())); + })) + .await + .unwrap(); + assert_eq!(inserted.get(), 1); + } + assert_eq!(emitted, [(1, 3), (1, 4)]); + assert_eq!(build.pending.list_size(1, 2).unwrap(), 4); + assert_eq!(build.store.lists[&1].members, [10_000, 10_001]); + assert_eq!(build.store.finishes, 0); + + require_send(build.finish()).await.unwrap(); + let store = strategy.store(); + assert_eq!(store.lists[&1].members, [10_000, 10_001, 10_005, 10_006]); + assert_eq!(store.points[&10_005], [1.0, 0.0]); + assert_eq!(store.points[&10_006], [2.0, 0.0]); + assert_eq!(store.finishes, 1); +} + +#[tokio::test] +async fn failed_routing_leaves_no_staged_points_or_appends() { + for fail_route in [false, true] { + let mut store = seeded(); + store.fail_route = fail_route; + if !fail_route { + store.lists.clear(); + } + let strategy = Strategy::new(store); + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); + let points = [(5, &[1.0, 0.0][..])]; + let mut emitted = Vec::new(); + let error = require_send(build.route(&points[..], |&list, size| { + emitted.push((list, size.get())); + })) + .await + .unwrap_err(); + let expected = if fail_route { + "injected routing failure" + } else { + "IVF routing returned no live centroid" + }; + assert!(error.to_string().contains(expected), "{error}"); + assert!(emitted.is_empty()); + assert!(build.staged.is_empty()); + assert!(build.pending.clone().into_changes().next().is_none()); + assert!(!build.store.points.contains_key(&10_005)); + assert_eq!(build.store.finishes, 0); + } +} + +#[tokio::test] +async fn native_inputs_are_inserted_using_the_assigned_ids() { + let (mut index, strategy) = index(seeded(), 4); + let points: [(u32, &[u8]); 3] = [(9, &[11, 0]), (7, &[1, 0]), (8, &[0, 0])]; + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points[..])) + .await + .unwrap(); + assert_eq!( + stats, + InsertStats { + inserted: 3, + splits: 0 + } + ); + let store = strategy.store(); + assert_eq!(store.lists[&1].members, [10_000, 10_001, 10_007, 10_008]); + assert_eq!(store.lists[&2].members, [10_002, 10_003, 10_009]); + for (external, vector) in points { + let expected: Vec<_> = vector.iter().copied().map(f32::from).collect(); + assert_eq!(store.points[&(external + 10_000)], expected); + } + assert_eq!(store.finishes, 1); + assert_eq!(store.fills, 0); +} + +#[tokio::test] +async fn failed_staging_leaves_no_new_elements() { + let strategy = Strategy::new(seeded()); + for points in [ + vec![(5, &[1.0, 0.0][..]), (5, &[2.0, 0.0][..])], + vec![(5, &[1.0, 0.0][..]), (0, &[2.0, 0.0][..])], + vec![(5, &[1.0, 0.0][..]), (6, &[2.0][..])], + vec![(5, &[1.0, 0.0][..]), (6, &[f32::NAN, 0.0][..])], + vec![(5, &[1.0, 0.0][..]), (u32::MAX, &[2.0, 0.0][..])], + ] { + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); + let mut emitted = Vec::new(); + assert!( + require_send(build.route(&points[..], |&list, size| { + emitted.push((list, size.get())); + })) + .await + .is_err() + ); + assert!(emitted.is_empty()); + assert!(build.staged.is_empty()); + assert!(build.pending.clone().into_changes().next().is_none()); + assert_eq!(build.store.points.len(), 5); + assert_eq!(build.store.finishes, 0); + } +} + +#[tokio::test] +async fn appends_publish_actual_memberships_without_filling() { + let (mut index, strategy) = index( + Store { + dim: 2, + lists: [ + ( + 1, + StoredList { + centroid: vec![0.0, 0.0], + members: Vec::new(), + }, + ), + ( + 2, + StoredList { + centroid: vec![10.0, 0.0], + members: Vec::new(), + }, + ), + ] + .into(), + next_list: 3, + fail_nearest: true, + ..Store::default() + }, + 2, + ); + let points: [(u32, &[f32]); 3] = [ + (7, &[0.0, 1.0][..]), + (8, &[1.0, 0.0][..]), + (9, &[10.0, 1.0][..]), + ]; + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points[..])) + .await + .unwrap(); + + assert_eq!( + stats, + InsertStats { + inserted: 3, + splits: 0 + } + ); + let store = strategy.store(); + assert_eq!(store.lists[&1].members, [10_007, 10_008]); + assert_eq!(store.lists[&2].members, [10_009]); + assert_eq!(store.points[&10_007], [0.0, 1.0]); + assert_eq!(store.to_internal_id(&DefaultContext, &7).unwrap(), 10_007); + assert_eq!(store.to_external_id(&DefaultContext, 10_007).unwrap(), 7); + assert_eq!(store.finishes, 1); + assert_eq!(store.fills, 0); +} + +#[tokio::test] +async fn empty_batches_are_rejected_without_staging_or_publication() { + let (mut index, strategy) = index(seeded(), 2); + let points: &[(u32, &[f32])] = &[]; + let mut emitted = Vec::new(); + { + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); + let error = require_send(build.route(points, |&list, size| { + emitted.push((list, size.get())); + })) + .await + .unwrap_err(); + assert!(error.to_string().contains("cannot insert an empty batch")); + assert!(emitted.is_empty()); + assert!(build.staged.is_empty()); + assert!(build.pending.clone().into_changes().next().is_none()); + } + let error = require_send(index.insert_batch(&strategy, &DefaultContext, points)) + .await + .unwrap_err(); + assert!(error.to_string().contains("cannot insert an empty batch")); + let store = strategy.store(); + assert_eq!(store.points.len(), 5); + assert_eq!(store.lists[&1].members, [10_000, 10_001]); + assert_eq!(store.lists[&2].members, [10_002, 10_003]); + assert_eq!(store.finishes, 0); + assert_eq!(store.fills, 0); +} + +#[tokio::test] +async fn splitting_commits_children_and_ordinary_appends_once() { + let mut store = seeded(); + store.fail_nearest = true; + let (mut index, strategy) = index(store, 2); + let points: [(u32, &[f32]); 3] = [ + (5, &[1.0, 1.0][..]), + (6, &[11.0, 1.0][..]), + (7, &[101.0, 0.0][..]), + ]; + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points[..])) + .await + .unwrap(); + let store = strategy.store(); + + assert_eq!( + stats, + InsertStats { + inserted: 3, + splits: 2 + } + ); + assert_eq!(store.finishes, 1); + assert_eq!(store.fills, 1); + assert!(!store.lists.contains_key(&1)); + assert!(!store.lists.contains_key(&2)); + assert_eq!(store.lists.len(), 5); + assert_eq!(store.lists[&3].members, [10_004, 10_007]); + let mut members: Vec<_> = store + .lists + .values() + .flat_map(|list| list.members.iter().copied()) + .collect(); + members.sort_unstable(); + assert_eq!(members, (10_000..10_008).collect::>()); + + for (&id, list) in &store.lists { + if id == 3 { + continue; + } + for point in &list.members { + let vector = &store.points[point]; + let own: f32 = SquaredL2::evaluate(vector.as_slice(), list.centroid.as_slice()); + for (&other, candidate) in &store.lists { + if other != 3 { + let distance: f32 = + SquaredL2::evaluate(vector.as_slice(), candidate.centroid.as_slice()); + assert!(own <= distance); + } + } + } + } +} + +#[tokio::test] +async fn neighboring_lists_exchange_points_and_pending_appends_once() { + let mut store = seeded(); + store.points.insert(10_001, vec![10.0, 0.0]); + store.points.insert(10_002, vec![5.0, 0.0]); + store.points.insert(10_003, vec![10.1, 0.0]); + store.points.insert(10_005, vec![5.0, 0.0]); + store.points.insert(10_006, vec![5.0, 0.0]); + store + .lists + .get_mut(&1) + .unwrap() + .members + .extend([10_005, 10_006]); + store.lists.get_mut(&2).unwrap().centroid = vec![5.0, 0.0]; + let (mut index, strategy) = index_with_neighbors(store, 3, 2); + let points: [(u32, &[f32]); 2] = [(7, &[1.0, 0.0][..]), (8, &[5.2, 0.0][..])]; + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points[..])) + .await + .unwrap(); + let store = strategy.store(); + + assert_eq!( + stats, + InsertStats { + inserted: 2, + splits: 1 + } + ); + assert_eq!(store.fills, 1); + assert_eq!(store.finishes, 1); + assert_eq!(store.filled_lists, [1, 2, 3]); + assert!(!store.lists.contains_key(&1)); + assert_eq!(store.lists.len(), 4); + assert_eq!(store.lists[&2].centroid, [5.0, 0.0]); + assert_eq!(store.lists[&2].members, [10_005, 10_006, 10_002, 10_008]); + assert!(store.lists[&2].members.len() > 3); + assert!( + store + .lists + .iter() + .any(|(&id, list)| { id >= 4 && list.members.contains(&10_003) }) + ); + assert_eq!(store.lists[&3].centroid, [100.0, 0.0]); + assert_eq!(store.lists[&3].members, [10_004]); + let mut members: Vec<_> = store + .lists + .values() + .flat_map(|list| list.members.iter().copied()) + .collect(); + members.sort_unstable(); + assert_eq!(members, (10_000..10_009).collect::>()); +} + +#[tokio::test] +async fn neighborhood_selection_filters_parents_and_fills_the_union_once() { + for (centers, one_candidate) in [ + (vec![], vec![]), + (vec![100.0, 200.0, 300.0], vec![3]), + (vec![0.0, 12.0, 200.0], vec![3, 4]), + ] { + for count in [0, 1, 2, 3, 4, usize::MAX] { + let mut store = seeded(); + if centers.is_empty() { + store.lists.remove(&3); + store.points.remove(&10_004); + } else { + store.lists.get_mut(&3).unwrap().centroid = vec![centers[0], 0.0]; + for (id, ¢er) in (4..).zip(¢ers[1..]) { + store.lists.insert( + id, + StoredList { + centroid: vec![center, 0.0], + members: Vec::new(), + }, + ); + } + } + store.next_list = 3 + centers.len() as u32; + store.fail_nearest = count == 0; + let (mut index, strategy) = index_with_neighbors(store, 2, count); + let points: [(u32, &[f32]); 2] = [(5, &[1.0, 1.0][..]), (6, &[11.0, 1.0][..])]; + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points[..])) + .await + .unwrap(); + let store = strategy.store(); + let neighbors = match (count, centers.len()) { + (0, _) | (_, 0) => vec![], + (1, _) => one_candidate.clone(), + (2, _) => vec![3, 4], + _ => vec![3, 4, 5], + }; + let sources: Vec<_> = [1, 2].into_iter().chain(neighbors).collect(); + assert_eq!(store.filled_lists, sources); + assert_eq!(store.fills, 1); + assert_eq!(store.finishes, 1); + assert_eq!( + stats, + InsertStats { + inserted: 2, + splits: 2 + } + ); + assert!(!store.lists.contains_key(&1)); + assert!(!store.lists.contains_key(&2)); + assert_eq!(store.lists.len(), 4 + centers.len()); + for (id, ¢er) in (3..).zip(¢ers) { + assert_eq!(store.lists[&id].centroid, [center, 0.0]); + } + let mut members: Vec<_> = store + .lists + .values() + .flat_map(|list| list.members.iter().copied()) + .collect(); + members.sort_unstable(); + assert_eq!(members, store.points.keys().copied().collect::>()); + } + } +} + +#[tokio::test] +async fn region_view_borrows_inputs_while_updates_are_recorded() { + let strategy = Strategy::new(seeded()); + { + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); + let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; + let reversed = [points[1], points[0]]; + let mut emitted = Vec::new(); + build + .route(&reversed[..], |&list, size| { + emitted.push((list, size.get())); + }) + .await + .unwrap(); + assert_eq!(emitted, [(1, 4)]); + { + let (region, mut update) = require_send(build.fill([1].into_iter())).await.unwrap(); + let view = region.get_list(1).unwrap(); + let vectors = region.get_vectors(1).unwrap(); + let source = ®ion.sources[&1]; + assert_eq!(view.members, [10_000, 10_001, 10_006, 10_005]); + assert_eq!(view.centroid.as_ptr(), source.centroid.as_ptr()); + assert_eq!(view.members.as_ptr(), source.members.as_ptr()); + assert_eq!( + vectors.as_slice().as_ptr(), + source.vectors.as_slice().as_ptr() + ); + assert_eq!(vectors.row(0), [0.0, 0.0]); + assert_eq!(vectors.row(2), points[1].1); + assert_eq!(vectors.row(3), points[0].1); + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()) + .split(®ion, &mut update, &[1], &[]) + .unwrap(); + assert_eq!(view.members, [10_000, 10_001, 10_006, 10_005]); + assert_eq!(vectors.row(2), points[1].1); + assert!(update.set_members(1, &[]).is_err()); + } + assert_eq!(build.dim().get(), 2); + assert_eq!(build.created.len(), 2); + let mut members: Vec<_> = build + .created + .iter() + .flat_map(|list| list.members.iter().copied()) + .collect(); + members.sort_unstable(); + assert_eq!(members, [10_000, 10_001, 10_005, 10_006]); + assert!(build.pending.is_retired(1)); + } + let store = strategy.store(); + assert_eq!(store.lists[&1].members, [10_000, 10_001]); + assert!(!store.points.contains_key(&10_005)); + assert_eq!(store.finishes, 0); +} + +#[tokio::test] +async fn borrowed_updates_share_pending_memberships_and_publish_on_finish() { + let strategy = Strategy::new(seeded()); + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); + let points = [ + (5, &[1.0, 0.0][..]), + (6, &[11.0, 0.0][..]), + (7, &[100.0, 0.0][..]), + ]; + build.route(&points[..], |_, _| {}).await.unwrap(); + + { + let (region, mut update) = require_send(build.fill([1, 2].into_iter())).await.unwrap(); + let parent = region.get_list(1).unwrap(); + let neighbor = region.get_list(2).unwrap(); + assert_eq!(parent.members, [10_000, 10_001, 10_005]); + assert_eq!(neighbor.members, [10_002, 10_003, 10_006]); + assert!(update.set_members(3, &[]).is_err()); + assert!(update.retire(3).is_err()); + + update.set_members(2, &[10_002, 10_003, 10_005]).unwrap(); + update + .create(&[1.0, 0.0], &[10_000, 10_001, 10_006]) + .unwrap(); + update.retire(1).unwrap(); + + assert_eq!(parent.members, [10_000, 10_001, 10_005]); + assert_eq!(neighbor.members, [10_002, 10_003, 10_006]); + assert_eq!(region.get_vectors(2).unwrap().row(2), points[1].1); + } + + assert!(build.pending.is_retired(1)); + assert_eq!(build.pending.list_size(2, 2).unwrap(), 3); + assert_eq!(build.created.len(), 1); + assert_eq!(build.store.points.len(), 5); + assert_eq!(build.store.lists[&1].members, [10_000, 10_001]); + assert_eq!(build.store.lists[&2].members, [10_002, 10_003]); + assert_eq!(build.store.lists.len(), 3); + assert_eq!(build.store.next_list, 4); + assert_eq!(build.store.finishes, 0); + + require_send(build.finish()).await.unwrap(); + let store = strategy.store(); + assert!(!store.lists.contains_key(&1)); + assert_eq!(store.lists[&2].members, [10_002, 10_003, 10_005]); + assert_eq!(store.lists[&3].members, [10_004, 10_007]); + assert_eq!(store.lists[&4].members, [10_000, 10_001, 10_006]); + assert_eq!(store.lists[&4].centroid, [1.0, 0.0]); + assert_eq!(store.points.len(), 8); + assert_eq!(store.lists.len(), 3); + assert_eq!(store.next_list, 5); + assert_eq!(store.finishes, 1); +} + +#[tokio::test] +async fn failed_planning_or_publication_is_not_reported_as_success() { + for (fail_nearest, fail_fill, fail_create, fail_set_members, fail_retire, fail_finish) in [ + (true, false, false, false, false, false), + (false, true, false, false, false, false), + (false, false, true, false, false, false), + (false, false, false, true, false, false), + (false, false, false, false, true, false), + (false, false, false, false, false, true), + ] { + let mut store = seeded(); + store.fail_nearest = fail_nearest; + store.fail_fill = fail_fill; + store.fail_create = fail_create; + store.fail_set_members = fail_set_members; + store.fail_retire = fail_retire; + store.fail_finish = fail_finish; + let (mut index, strategy) = index_with_neighbors(store, 2, 2); + let points: [(u32, &[f32]); 1] = [(5, &[1.0, 1.0][..])]; + let error = require_send(index.insert_batch(&strategy, &DefaultContext, &points[..])) + .await + .unwrap_err(); + let expected = if fail_nearest { + "injected nearest failure" + } else if fail_fill { + "injected fill failure" + } else if fail_create { + "injected create failure" + } else if fail_set_members { + "injected set_members failure" + } else if fail_retire { + "injected retire failure" + } else { + "injected finish failure" + }; + assert!(error.to_string().contains(expected), "{error}"); + let store = strategy.store(); + assert_eq!(store.finishes, 0); + assert_eq!(store.lists[&1].members, [10_000, 10_001]); + assert_eq!(store.lists[&2].members, [10_002, 10_003]); + assert_eq!(store.lists.len(), 3); + assert!(store.to_internal_id(&DefaultContext, &5).is_err()); + } +} + +#[tokio::test] +async fn uninitialized_insertion_fails() { + let (mut index, strategy) = index( + Store { + dim: 2, + next_list: 1, + ..Store::default() + }, + 2, + ); + let points: [(u32, &[f32]); 1] = [(7, &[0.0, 1.0][..])]; + let error = index + .insert_batch(&strategy, &DefaultContext, &points[..]) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("IVF routing returned no live centroid"), + "{error}" + ); + assert_eq!(strategy.store().finishes, 0); +} + +#[test] +fn accumulator_empty_like_preserves_search_limit_not_capacity() { + let mut best = NeighborPriorityQueue::auto_resizable_with_search_param_l(2); + for id in 0..3 { + best.push(id, id as f32); + } + assert!(best.capacity() > best.search_l()); + + let empty = best.empty_like(); + assert_eq!(empty.search_l(), 2); + assert_eq!(empty.capacity(), 2); + assert_eq!(empty.size(), 0); + assert_eq!(best.size(), 3); +} + +#[test] +fn accumulator_merges_partitioned_top_k() { + let candidates = [ + (0, 9.0), + (1, 1.0), + (2, 4.0), + (3, 7.0), + (4, 3.0), + (5, 8.0), + (6, 5.0), + (7, 2.0), + ]; + for k in [1, 3, 16] { + let mut reference = NeighborPriorityQueue::new(k); + reference.insert(Neighbor::new(99, -1.0)); + for &(id, distance) in &candidates { + reference.insert(Neighbor::new(id, distance)); + } + let expected: Vec<_> = reference.iter().map(Neighbor::as_tuple).collect(); + + for chunk_size in [1, 3, candidates.len()] { + for reverse in [false, true] { + let mut best = NeighborPriorityQueue::new(k); + best.push(99, -1.0); + let mut partials: Vec<_> = candidates + .chunks(chunk_size) + .map(|chunk| { + let mut partial = best.empty_like(); + assert_eq!(partial.search_l(), k); + assert_eq!(partial.size(), 0); + for &(id, distance) in chunk { + partial.push(id, distance); + } + assert!(partial.size() <= k); + partial + }) + .collect(); + if reverse { + partials.reverse(); + } + for partial in partials { + best.merge(partial); + } + let empty = best.empty_like(); + best.merge(empty); + assert_eq!( + best.iter().map(Neighbor::as_tuple).collect::>(), + expected, + ); + } + } + } +} + +#[test] +fn accumulator_preserves_infinite_distances_and_ignores_nan() { + let mut best = NeighborPriorityQueue::new(3); + let mut left = best.empty_like(); + left.push(1, f32::INFINITY); + left.push(2, f32::NAN); + let mut right = best.empty_like(); + right.push(3, 1.0); + right.push(4, f32::INFINITY); + + best.merge(left); + best.merge(right); + assert_eq!( + best.iter() + .map(|neighbor| *neighbor.distance()) + .collect::>(), + [1.0, f32::INFINITY, f32::INFINITY], + ); + assert_eq!( + best.iter() + .map(|neighbor| *neighbor.id()) + .collect::>(), + BTreeSet::from([1, 3, 4]), + ); +} + +#[test] +#[should_panic(expected = "cannot merge accumulators with different search limits")] +fn accumulator_rejects_incompatible_limits() { + let mut best = NeighborPriorityQueue::::new(2); + best.merge(NeighborPriorityQueue::new(1)); +} + +#[tokio::test] +async fn search_scans_selected_lists_and_returns_the_best_k() { + let mut store = seeded(); + store.points.insert(10_004, vec![7.5, 0.0]); + let (index, strategy) = index(store, 2); + let query = [7.5f32, 0.0]; + for (nprobe, k, expected) in [ + (1, 3, vec![(10_002, 6.25_f32), (10_003, 20.25)]), + (2, 3, vec![(10_002, 6.25), (10_003, 20.25), (10_001, 30.25)]), + (3, 2, vec![(10_004, 0.0), (10_002, 6.25)]), + ( + usize::MAX, + 10, + vec![ + (10_004, 0.0), + (10_002, 6.25), + (10_003, 20.25), + (10_001, 30.25), + (10_000, 56.25), + ], + ), + ] { + let mut output = Vec::new(); + let count = require_send(index.search( + &strategy, + &DefaultContext, + query.as_slice(), + NonZeroUsize::new(k).unwrap(), + NonZeroUsize::new(nprobe).unwrap(), + &mut output, + )) + .await + .unwrap(); + assert_eq!(count, expected.len()); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + expected + ); + } + let store = strategy.store(); + assert_eq!(store.points.len(), 5); + assert_eq!(store.lists.len(), 3); + assert_eq!(store.finishes, 0); +} + +#[tokio::test] +async fn search_respects_output_buffer_capacity_and_existing_results() { + let (index, strategy) = index(seeded(), 2); + let query = [7.5f32, 0.0]; + let mut ids = [u32::MAX; 3]; + let mut distances = [f32::INFINITY; 3]; + { + let mut output = IdDistance::new(&mut ids, &mut distances); + assert!(output.push(Neighbor::new(42, -1.0)).is_available()); + for expected_count in [2, 0] { + let count = require_send(index.search( + &strategy, + &DefaultContext, + query.as_slice(), + NonZeroUsize::new(4).unwrap(), + NonZeroUsize::new(3).unwrap(), + &mut output, + )) + .await + .unwrap(); + assert_eq!(count, expected_count); + assert_eq!(output.current_len(), 3); + } + } + assert_eq!(ids, [42, 10_002, 10_003]); + assert_eq!(distances, [-1.0, 6.25, 20.25]); + + let mut buffer = []; + let mut output = BackInserter::new(&mut buffer); + assert_eq!( + require_send(index.search( + &strategy, + &DefaultContext, + query.as_slice(), + ONE, + ONE, + &mut output, + )) + .await + .unwrap(), + 0 + ); +} + +#[tokio::test] +async fn search_returns_no_results_for_empty_regions() { + for empty_index in [false, true] { + let mut store = seeded(); + if empty_index { + store.lists.clear(); + store.points.clear(); + } else { + store.lists.get_mut(&2).unwrap().members.clear(); + } + let (index, strategy) = index(store, 2); + let mut output = vec![Neighbor::new(42, -1.0)]; + assert_eq!( + require_send(index.search( + &strategy, + &DefaultContext, + &[7.5f32, 0.0][..], + ONE, + ONE, + &mut output, + )) + .await + .unwrap(), + 0 + ); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + [(42, -1.0)] + ); + } +} + +#[tokio::test] +async fn search_failures_do_not_publish_partial_results() { + for (case, expected) in [ + ("dimension", "invalid query vector"), + ("nonfinite", "invalid query vector"), + ("route", "injected routing failure"), + ("scan", "missing point 10003"), + ("later_scan", "missing point 10001"), + ] { + let mut store = seeded(); + if case == "route" { + store.fail_route = true; + } + if case == "scan" { + store.points.remove(&10_003); + } + if case == "later_scan" { + store.points.remove(&10_001); + } + let (index, strategy) = index(store, 2); + let mut query = [7.5f32, 0.0]; + if case == "nonfinite" { + query[0] = f32::NAN; + } + let query = if case == "dimension" { + &query[..1] + } else { + &query[..] + }; + let nprobe = if case == "later_scan" { + NonZeroUsize::new(3).unwrap() + } else { + ONE + }; + let mut output = vec![Neighbor::new(42, -1.0)]; + let error = require_send(index.search( + &strategy, + &DefaultContext, + query, + NonZeroUsize::new(3).unwrap(), + nprobe, + &mut output, + )) + .await + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + [(42, -1.0)] + ); + assert_eq!(strategy.store().finishes, 0); + } +} + +#[tokio::test] +async fn search_supports_native_queries_after_insertion_and_splitting() { + let (mut index, strategy) = index(seeded(), 2); + let points = [(5, &[8u8, 0][..])]; + let stats = index + .insert_batch(&strategy, &DefaultContext, &points[..]) + .await + .unwrap(); + assert_eq!(stats.splits, 1); + + let mut output = Vec::new(); + assert_eq!( + require_send(index.search( + &strategy, + &DefaultContext, + &[8u8, 0][..], + ONE, + NonZeroUsize::new(usize::MAX).unwrap(), + &mut output, + )) + .await + .unwrap(), + 1 + ); + assert_eq!( + output + .into_iter() + .map(Neighbor::as_tuple) + .collect::>(), + [(10_005, 0.0)] + ); +} + +#[test] +fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { + let sources = [ + ( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-10.0], [10.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![5.0], + members: vec![2], + vectors: matrix(&[[50.0]]), + }, + ), + ( + 3, + Source { + centroid: vec![50.0], + members: vec![3], + vectors: matrix(&[[51.0]]), + }, + ), + ( + 4, + Source { + centroid: vec![10.0], + members: Vec::new(), + vectors: matrix::<1>(&[]), + }, + ), + ] + .into(); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), ONE) + .split(®ion, &mut update, &[1], &[2, 3, 4]) + .unwrap(); + + let mut centers: Vec<_> = update.created.iter().map(|list| list.centroid[0]).collect(); + centers.sort_by(f32::total_cmp); + assert_eq!(centers, [-10.0, 10.0]); + let mut child_members: Vec<_> = update + .created + .iter() + .flat_map(|list| list.members.iter().copied()) + .collect(); + child_members.sort_unstable(); + assert_eq!(child_members, [0]); + let mut changes: BTreeMap<_, _> = std::mem::take(update.memberships).into_changes().collect(); + assert!(matches!(changes.remove(&1), Some(Membership::Retire))); + for (id, expected) in [(2, vec![]), (3, vec![2, 3]), (4, vec![1])] { + let Some(Membership::Replace(members)) = changes.remove(&id) else { + panic!("expected replacement for retained list {id}"); + }; + assert_eq!(members, expected); + } + assert!(changes.is_empty()); +} + +#[tokio::test] +async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() { + let mut store = seeded(); + store.points.insert(10_001, vec![10.0, 0.0]); + store.points.insert(10_002, vec![5.0, 0.0]); + store.points.insert(10_003, vec![10.1, 0.0]); + store.lists.get_mut(&1).unwrap().centroid = vec![5.0, 0.0]; + store.lists.get_mut(&2).unwrap().centroid = vec![5.0, 0.0]; + let strategy = Strategy::new(store); + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); + let points = [(5, &[5.0, 0.0][..])]; + let mut emitted = Vec::new(); + build + .route(&points[..], |&list, size| { + emitted.push((list, size.get())); + }) + .await + .unwrap(); + assert_eq!(emitted, [(1, 3)]); + + { + let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); + assert!( + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) + .split(®ion, &mut update, &[1, 99], &[]) + .is_err() + ); + assert!( + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) + .split(®ion, &mut update, &[1], &[99]) + .is_err() + ); + assert!(update.created.is_empty()); + let changes: Vec<_> = update.memberships.clone().into_changes().collect(); + assert!( + matches!(changes.as_slice(), [(1, Membership::Append(members))] if members == &[10_005]) + ); + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()) + .split(®ion, &mut update, &[1], &[]) + .unwrap(); + assert_eq!(update.created.len(), 2); + } + require_send(build.finish()).await.unwrap(); + + let store = strategy.store(); + assert_eq!(store.finishes, 1); + assert!(!store.lists.contains_key(&1)); + assert_eq!(store.lists[&2].centroid, [5.0, 0.0]); + assert_eq!(store.lists[&2].members, [10_002, 10_003]); + assert_eq!(store.lists[&3].members, [10_004]); + assert_eq!(store.lists.len(), 4); + let mut members: Vec<_> = store + .lists + .values() + .flat_map(|list| list.members.iter().copied()) + .collect(); + members.sort_unstable(); + assert_eq!(members, (10_000..10_006).collect::>()); +} + +#[test] +fn assignment_can_cross_parent_boundaries() { + let left = [[-100.0], [-99.0], [-10.0], [-9.0], [-0.1]]; + let right = [[0.1], [0.2], [100.0], [101.0]]; + let sources = [ + ( + 1, + Source { + centroid: vec![-50.0], + members: (0..left.len() as u32).collect(), + vectors: matrix(&left), + }, + ), + ( + 2, + Source { + centroid: vec![50.0], + members: (100..100 + right.len() as u32).collect(), + vectors: matrix(&right), + }, + ), + ] + .into(); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new( + NonZeroUsize::new(10).unwrap(), + NonZeroUsize::new(1).unwrap(), + ) + .split(®ion, &mut update, &[1, 2], &[]) + .unwrap(); + assert!( + update.created[2..] + .iter() + .any(|partition| partition.members.contains(&4)) + ); +} + +#[test] +fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { + for value in [0.0, 7.0] { + let sources = [( + 1, + Source { + centroid: vec![value], + members: vec![0, 1, 2], + vectors: matrix(&[[value]; 3]), + }, + )] + .into(); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[1], &[]) + .unwrap(); + assert_eq!(update.created.len(), 2); + assert_eq!(update.created[0].centroid, [value]); + assert_eq!(update.created[1].centroid, [0.0]); + assert_eq!(update.created[0].members, [0, 1, 2]); + assert!(update.created[1].members.is_empty()); + } +} + +#[test] +fn two_means_prefers_neighbors_in_supplied_order_on_ties() { + let sources = BTreeMap::from([ + ( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![-1.0], + members: vec![2], + vectors: matrix(&[[-1.0]]), + }, + ), + ( + 3, + Source { + centroid: vec![-1.0], + members: vec![3], + vectors: matrix(&[[-1.0]]), + }, + ), + ]); + let region = Region::new(&sources); + for neighbors in [[2, 3], [3, 2]] { + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[1], &neighbors) + .unwrap(); + assert_eq!(created.len(), 2); + assert_eq!(created[0].centroid, [-1.0]); + assert!(created[0].members.is_empty()); + assert_eq!(created[1].centroid, [1.0]); + assert_eq!(created[1].members, [1]); + let changes: BTreeMap<_, _> = memberships.into_changes().collect(); + assert!(matches!(changes[&1], Membership::Retire)); + assert!(matches!( + &changes[&neighbors[0]], + Membership::Replace(members) if members == &[0, neighbors[0], neighbors[1]] + )); + assert!(matches!( + &changes[&neighbors[1]], + Membership::Replace(members) if members.is_empty() + )); + } +} + +#[test] +fn two_means_reuses_splitter_across_regions() { + let sources = [ + ( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-10.0], [10.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![105.0], + members: vec![2, 3], + vectors: matrix(&[[100.0], [110.0]]), + }, + ), + ( + 3, + Source { + centroid: vec![105.0], + members: vec![4, 5, 6], + vectors: matrix(&[[9.9], [100.1], [105.0]]), + }, + ), + ( + 4, + Source { + centroid: vec![1000.0], + members: vec![7], + vectors: matrix(&[[100.2]]), + }, + ), + ] + .into(); + let region = Region::new(&sources); + let parents = [1, 2]; + let neighbors = [3, 4]; + let mut full_memberships = PendingLists::default(); + let mut full_created = Vec::new(); + let mut full = support::Update::new(&sources, 1, &mut full_memberships, &mut full_created); + let mut splitter = TwoMeansSplit::new(ONE, ONE); + splitter + .split(®ion, &mut full, &parents, &neighbors) + .unwrap(); + assert_eq!( + full.created + .iter() + .map(|list| list.centroid[0]) + .collect::>(), + [-10.0, 10.0, 100.0, 110.0] + ); + assert_eq!(full.created[0].members, [0]); + assert_eq!(full.created[1].members, [1, 4]); + assert_eq!(full.created[3].members, [3]); + assert_eq!(full.created[2].members, [2, 5, 7]); + let mut repeated_memberships = PendingLists::default(); + let mut repeated_created = Vec::new(); + let mut repeated = support::Update::new( + &sources, + 1, + &mut repeated_memberships, + &mut repeated_created, + ); + splitter + .split(®ion, &mut repeated, &parents, &neighbors) + .unwrap(); + for (first, repeated) in full.created.iter().zip(repeated.created.iter()) { + assert_eq!(first.centroid, repeated.centroid); + assert_eq!(first.members, repeated.members); + } + let mut smaller_memberships = PendingLists::default(); + let mut smaller_created = Vec::new(); + let mut smaller = + support::Update::new(&sources, 1, &mut smaller_memberships, &mut smaller_created); + splitter.split(®ion, &mut smaller, &[1], &[]).unwrap(); + assert_eq!(smaller.created.len(), 2); + assert_eq!(smaller.created[0].members, [0]); + assert_eq!(smaller.created[1].members, [1]); + let mut changes: BTreeMap<_, _> = std::mem::take(full.memberships).into_changes().collect(); + for parent in [1, 2] { + assert!(matches!(changes.remove(&parent), Some(Membership::Retire))); + } + for (id, expected) in [(3, vec![6]), (4, vec![])] { + let Some(Membership::Replace(members)) = changes.remove(&id) else { + panic!("expected replacement for retained list {id}"); + }; + assert_eq!(members, expected); + } + assert!(changes.is_empty()); +} + +#[test] +fn empty_split_records_no_updates() { + let sources = BTreeMap::new(); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + let error = TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[], &[]) + .unwrap_err(); + assert!(matches!( + error.downcast_ref::>(), + Some(TwoMeansError::EmptyParents) + )); + assert!(created.is_empty()); + assert_eq!(memberships.into_changes().count(), 0); +} + +#[test] +fn region_vector_reborrows_share_source_storage() { + let sources = BTreeMap::from([( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + )]); + let region = Region::new(&sources); + let view = region.get_list(1).unwrap(); + assert!(std::ptr::eq(view.members, sources[&1].members.as_slice())); + let vectors = region.get_vectors(1).unwrap(); + let reborrowed = vectors.reborrow(); + assert!(std::ptr::eq( + reborrowed.as_slice(), + sources[&1].vectors.as_slice() + )); + assert_eq!(reborrowed.row(0), [-1.0]); + assert_eq!(reborrowed.row(1), [1.0]); +} + +#[test] +fn two_means_handles_borrowed_batches() { + struct Batch<'a> { + vectors: rowmajor::Ref<'a, f32>, + reads: &'a Cell, + } + + struct BatchRef<'a> { + vectors: rowmajor::Ref<'a, f32>, + reads: &'a Cell, + } + + impl<'r> Reborrow<'r> for Batch<'_> { + type Target = BatchRef<'r>; + + fn reborrow(&'r self) -> Self::Target { + BatchRef { + vectors: self.vectors.reborrow(), + reads: self.reads, + } + } + } + + impl<'a> AsMatRef<'a, f32> for BatchRef<'a> { + fn as_view(self) -> rowmajor::Ref<'a, f32> { + self.reads.set(self.reads.get() + 1); + self.vectors + } + } + + struct BatchRegion<'a> { + inner: Region<'a>, + reads: Cell, + } + + impl View for BatchRegion<'_> { + type VectorsRef<'a> = BatchRef<'a>; + type Vectors<'a> + = Batch<'a> + where + Self: 'a; + + fn get_list(&self, list: u32) -> Option> { + self.inner.get_list(list) + } + + fn get_vectors(&self, list: u32) -> Option> { + self.inner.get_vectors(list).map(|vectors| Batch { + vectors, + reads: &self.reads, + }) + } + } + + let sources = BTreeMap::from([ + ( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![10.0], + members: vec![2], + vectors: matrix(&[[10.0]]), + }, + ), + ]); + let region = BatchRegion { + inner: Region::new(&sources), + reads: Cell::new(0), + }; + let batch = region.get_vectors(1).unwrap(); + let borrowed = batch.reborrow(); + assert!(std::ptr::eq(borrowed.reads, ®ion.reads)); + assert!(std::ptr::eq( + borrowed.vectors.as_slice(), + sources[&1].vectors.as_slice() + )); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[1], &[2]) + .unwrap(); + assert_eq!(region.reads.get(), 3); + assert_eq!(created.len(), 2); + assert_eq!(created[0].centroid, [-1.0]); + assert_eq!(created[0].members, [0]); + assert_eq!(created[1].centroid, [1.0]); + assert_eq!(created[1].members, [1]); + let changes: BTreeMap<_, _> = memberships.into_changes().collect(); + assert!(matches!(changes[&1], Membership::Retire)); + assert!(matches!(&changes[&2], Membership::Replace(members) if members == &[2])); +} + +#[test] +fn split_accepts_native_batches_without_matrix_access() { + struct NativeBatch<'a>(&'a [[u8; 1]]); + + impl<'r> Reborrow<'r> for NativeBatch<'_> { + type Target = &'r [[u8; 1]]; + + fn reborrow(&'r self) -> Self::Target { + self.0 + } + } + + struct NativeRegion<'a> { + inner: Region<'a>, + vectors: &'a [[u8; 1]], + } + + impl View for NativeRegion<'_> { + type VectorsRef<'a> = &'a [[u8; 1]]; + type Vectors<'a> + = NativeBatch<'a> + where + Self: 'a; + + fn get_list(&self, list: u32) -> Option> { + self.inner.get_list(list) + } + + fn get_vectors(&self, list: u32) -> Option> { + self.inner.get_list(list).map(|_| NativeBatch(self.vectors)) + } + } + + struct NativeSplit; + + impl Split for NativeSplit + where + R: for<'a> View = &'a [[u8; 1]]>, + { + type Error = ANNError; + + fn split>( + &mut self, + region: &R, + update: &mut U, + parents: &[u32], + neighbors: &[u32], + ) -> ANNResult<()> { + assert!(neighbors.is_empty()); + let view = region.get_list(parents[0]).unwrap(); + let vectors = region.get_vectors(parents[0]).unwrap(); + for (point, member) in vectors.reborrow().iter().zip(view.members) { + update.create(&[f32::from(point[0])], std::slice::from_ref(member))?; + } + update.retire(parents[0]) + } + } + + let mut splitter = NativeSplit; + for offset in [0u8, 20] { + let sources = BTreeMap::from([( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[0.0], [10.0]]), + }, + )]); + let vectors = vec![[offset], [offset + 10]]; + let region = NativeRegion { + inner: Region::new(&sources), + vectors: &vectors, + }; + let batch = region.get_vectors(1).unwrap(); + assert!(std::ptr::eq(batch.reborrow(), vectors.as_slice())); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + splitter.split(®ion, &mut update, &[1], &[]).unwrap(); + assert_eq!(created.len(), 2); + assert_eq!(created[0].centroid, [f32::from(offset)]); + assert_eq!(created[0].members, [0]); + assert_eq!(created[1].centroid, [f32::from(offset + 10)]); + assert_eq!(created[1].members, [1]); + assert!(matches!( + memberships.into_changes().collect::>().as_slice(), + [(1, Membership::Retire)] + )); + } +} + +#[test] +fn two_means_matches_quantization_lloyds() { + let sources = BTreeMap::from([ + ( + 1, + Source { + centroid: vec![0.0, 0.0], + members: vec![0, 1, 2, 3], + vectors: matrix(&[[0.0, 2.0], [2.0, 4.0], [6.0, 6.0], [10.0, 12.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![100.0, 100.0], + members: vec![4, 5, 6], + vectors: matrix(&[[100.0, 102.0], [110.0, 112.0], [120.0, 120.0]]), + }, + ), + ]); + let region = Region::new(&sources); + for iterations in [1, 5] { + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 2, &mut memberships, &mut created); + TwoMeansSplit::new( + NonZeroUsize::new(iterations).unwrap(), + NonZeroUsize::new(2).unwrap(), + ) + .split(®ion, &mut update, &[1, 2], &[]) + .unwrap(); + assert_eq!(created.len(), 4); + for (source, actual) in sources.values().zip(created.chunks_exact(2)) { + let points = source.vectors.as_view(); + let mut expected = points.subview(0..2).unwrap().to_rowmajor_owned(); + let _ = diskann_quantization::algorithms::kmeans::lloyds::lloyds( + points, + expected.as_view_mut(), + iterations, + ); + for (actual, expected) in actual.iter().zip(expected.rows()) { + assert_eq!(actual.centroid, expected); + } + } + } +} + +#[test] +fn two_means_reports_insufficient_points_before_writing() { + for points in [matrix::<2>(&[]), matrix(&[[0.0, 2.0]])] { + let expected = points.nrows(); + let sources = BTreeMap::from([( + 1, + Source { + centroid: vec![0.0; points.ncols()], + members: (0..).take(points.nrows()).collect(), + vectors: points, + }, + )]); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 2, &mut memberships, &mut created); + let error = TwoMeansSplit::new(ONE, NonZeroUsize::new(2).unwrap()) + .split(®ion, &mut update, &[1], &[]) + .unwrap_err(); + assert!(matches!( + error.downcast_ref::>().unwrap(), + TwoMeansError::NotEnoughPoints { list: 1, actual } if *actual == expected + )); + assert!(created.is_empty()); + assert_eq!(memberships.into_changes().count(), 0); + } +} + +#[test] +fn two_means_missing_vectors_preserves_typed_error_without_updates() { + let sources = BTreeMap::from([( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + )]); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + let error = TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[1, 99], &[]) + .unwrap_err(); + assert!(matches!( + error.downcast_ref::>(), + Some(TwoMeansError::UnableToRetrieve(99)) + )); + assert!(created.is_empty()); + assert_eq!(memberships.into_changes().count(), 0); +} + +#[test] +fn two_means_reports_missing_neighbors_before_writing() { + let sources = BTreeMap::from([ + ( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[0.0], [1.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![2.0], + members: vec![2], + vectors: matrix(&[[2.0]]), + }, + ), + ]); + for ids in [vec![99], vec![2, 99]] { + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + let region = Region::new(&sources); + let error = TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[1], &ids) + .unwrap_err(); + assert!(matches!( + error.downcast_ref::>(), + Some(TwoMeansError::UnableToRetrieve(99)) + )); + assert!(update.created.is_empty()); + assert_eq!(std::mem::take(update.memberships).into_changes().count(), 0); + } +} + +#[test] +fn regional_assignment_prefers_first_when_all_distances_overflow() { + let sources = BTreeMap::from([ + ( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + ), + ( + 2, + Source { + centroid: vec![f32::MAX], + members: vec![2], + vectors: matrix(&[[-f32::MAX]]), + }, + ), + ]); + let region = Region::new(&sources); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new(ONE, ONE) + .split(®ion, &mut update, &[1], &[2]) + .unwrap(); + assert_eq!(created[0].members, [0]); + assert_eq!(created[1].members, [1]); + let changes: BTreeMap<_, _> = memberships.into_changes().collect(); + assert!(matches!(&changes[&2], Membership::Replace(members) if members == &[2])); +} diff --git a/diskann/src/ivf/tests/pending.rs b/diskann/src/ivf/tests/pending.rs new file mode 100644 index 0000000000..7a107363ed --- /dev/null +++ b/diskann/src/ivf/tests/pending.rs @@ -0,0 +1,106 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Pending memberships for the operation-local test provider. +//! +//! Appends do not materialize stored memberships. Replacement supersedes earlier +//! appends, and retirement supersedes both. A provider can use this helper for +//! read-your-writes behavior and consume it when publishing its build operation. +//! It is not the algorithm's commit protocol; providers may implement that behavior +//! with their own storage or transactions instead. + +use hashbrown::{HashMap, hash_map::Entry}; + +use crate::{ANNError, ANNResult, utils::VectorId}; + +/// A pending change relative to a list's stored membership. +#[derive(Debug, Clone)] +pub enum Membership { + Append(Vec), + Replace(Vec), + Retire, +} + +/// Membership state without vector storage or centroid-ID allocation. +/// +/// Callers supply stored sizes and members, and validate that list IDs exist. +#[derive(Debug, Default, Clone)] +pub struct PendingLists { + changes: HashMap>, +} + +impl PendingLists { + pub fn append(&mut self, list: L, members: Vec) -> ANNResult<()> { + match self.changes.entry(list) { + Entry::Vacant(entry) => { + entry.insert(Membership::Append(members)); + } + Entry::Occupied(mut entry) => match entry.get_mut() { + Membership::Append(current) | Membership::Replace(current) => { + current.extend(members); + } + Membership::Retire => return Err(retired(list)), + }, + } + Ok(()) + } + + pub fn set(&mut self, list: L, members: Vec) -> ANNResult<()> { + if self.is_retired(list) { + return Err(retired(list)); + } + self.changes.insert(list, Membership::Replace(members)); + Ok(()) + } + + pub fn retire(&mut self, list: L) -> ANNResult<()> { + if self.is_retired(list) { + return Err(retired(list)); + } + self.changes.insert(list, Membership::Retire); + Ok(()) + } + + pub fn is_retired(&self, list: L) -> bool { + matches!(self.changes.get(&list), Some(Membership::Retire)) + } + + pub fn list_size(&self, list: L, stored_size: usize) -> ANNResult { + match self.changes.get(&list) { + None => Ok(stored_size), + Some(Membership::Append(members)) => stored_size + .checked_add(members.len()) + .ok_or_else(|| ANNError::message(format!("IVF list {list} size overflows usize"))), + Some(Membership::Replace(members)) => Ok(members.len()), + Some(Membership::Retire) => Err(retired(list)), + } + } + + /// Borrow current member IDs without packing stored and appended IDs together. + pub fn members<'a>( + &'a self, + list: L, + stored: &'a [I], + ) -> ANNResult + Clone + 'a> { + let (head, tail): (&[I], &[I]) = match self.changes.get(&list) { + None => (stored, &[]), + Some(Membership::Append(members)) => (stored, members), + Some(Membership::Replace(members)) => (members, &[]), + Some(Membership::Retire) => return Err(retired(list)), + }; + Ok(head.iter().chain(tail).copied()) + } + + /// Consume pending state. The iteration order is unspecified. + pub fn into_changes(self) -> impl Iterator)> { + self.changes.into_iter() + } +} + +fn retired(list: L) -> ANNError { + ANNError::message(format!( + "IVF list {list} is retired in this build operation" + )) +} diff --git a/diskann/src/ivf/tests/support.rs b/diskann/src/ivf/tests/support.rs new file mode 100644 index 0000000000..ba75e67462 --- /dev/null +++ b/diskann/src/ivf/tests/support.rs @@ -0,0 +1,626 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; + +#[derive(Debug, Clone)] +pub(super) struct StoredList { + pub(super) centroid: Vec, + pub(super) members: Vec, +} + +#[derive(Debug, Default)] +pub(super) struct Store { + pub(super) dim: usize, + pub(super) points: BTreeMap>, + pub(super) lists: BTreeMap, + pub(super) next_list: u32, + pub(super) finishes: usize, + pub(super) fills: usize, + pub(super) filled_lists: Vec, + pub(super) fail_route: bool, + pub(super) fail_nearest: bool, + pub(super) fail_fill: bool, + pub(super) fail_create: bool, + pub(super) fail_set_members: bool, + pub(super) fail_retire: bool, + pub(super) fail_finish: bool, +} + +impl Store { + pub(super) fn to_internal_id(&self, _: &DefaultContext, id: &u32) -> ANNResult { + id.checked_add(10_000) + .filter(|id| self.points.contains_key(id)) + .ok_or_else(|| ANNError::message("point is not visible")) + } + + pub(super) fn to_external_id(&self, _: &DefaultContext, id: u32) -> ANNResult { + if !self.points.contains_key(&id) { + return Err(ANNError::message("point is not visible")); + } + id.checked_sub(10_000) + .ok_or_else(|| ANNError::message("invalid internal point ID")) + } +} + +pub(super) struct Build<'p, T> { + pub(super) store: MutexGuard<'p, Store>, + pub(super) pending: PendingLists, + pub(super) created: Vec, + pub(super) staged: BTreeMap>, + pub(super) sources: BTreeMap, + pub(super) fills: Cell, + pub(super) filled_lists: Vec, +} + +pub(super) struct Coarse<'a> { + pub(super) stored: &'a BTreeMap, + pub(super) fail_route: bool, + pub(super) fail_nearest: bool, +} + +fn squared_distance>(query: &[T], vector: &[f32]) -> f32 { + query + .iter() + .zip(vector) + .map(|(&value, &coordinate)| { + let delta = value.into() - coordinate; + delta * delta + }) + .sum() +} + +impl Coarse<'_> { + fn select(&self, query: &[T], nprobe: NonZeroUsize, pred: P, mut emit: F) -> usize + where + T: Copy + Into, + P: Fn(&u32) -> bool, + F: FnMut(&u32, f32), + { + let mut selected: Vec<_> = self + .stored + .iter() + .filter(|(id, _)| pred(id)) + .map(|(&id, list)| (id, squared_distance(query, &list.centroid))) + .collect(); + selected.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))); + selected.truncate(nprobe.get()); + for (id, distance) in &selected { + emit(id, *distance); + } + selected.len() + } + pub(super) fn route( + &self, + query: &[T], + nprobe: NonZeroUsize, + emit: F, + ) -> impl SendFuture> + where + T: Copy + Into + Sync, + F: FnMut(&u32, f32) + Send, + { + async move { + if self.fail_route { + return Err(ANNError::message("injected routing failure")); + } + self.select(query, nprobe, |_| true, emit); + Ok(()) + } + } +} + +impl NearestList for Coarse<'_> { + type Id = u32; + + fn find( + &self, + id: Self::Id, + nprobe: NonZeroUsize, + pred: P, + emit: F, + ) -> impl SendFuture> + where + P: Fn(&Self::Id) -> bool + Send, + F: FnMut(&Self::Id, f32) + Send, + { + async move { + if self.fail_nearest { + return Err(ANNError::message("injected nearest failure")); + } + let query = &self + .stored + .get(&id) + .ok_or_else(|| ANNError::message(format!("unknown list {id}")))? + .centroid; + Ok(self.select(query, nprobe, pred, emit)) + } + } +} + +pub(super) struct Source { + pub(super) centroid: Vec, + pub(super) members: Vec, + pub(super) vectors: rowmajor::Owned, +} + +pub(super) struct Update<'a> { + pub(super) dim: usize, + pub(super) sources: &'a BTreeMap, + pub(super) memberships: &'a mut PendingLists, + pub(super) created: &'a mut Vec, + pub(super) fail_create: bool, + pub(super) fail_set_members: bool, + pub(super) fail_retire: bool, +} + +impl<'a> Update<'a> { + pub(super) fn new( + sources: &'a BTreeMap, + dim: usize, + memberships: &'a mut PendingLists, + created: &'a mut Vec, + ) -> Self { + Self { + dim, + sources, + memberships, + created, + fail_create: false, + fail_set_members: false, + fail_retire: false, + } + } + + fn check_list(&self, list: u32) -> ANNResult<()> { + if self.sources.contains_key(&list) { + Ok(()) + } else { + Err(ANNError::message(format!("unknown list {list}"))) + } + } +} + +pub(super) struct Region<'a> { + pub(super) sources: &'a BTreeMap, +} + +impl<'a> Region<'a> { + pub(super) fn new(sources: &'a BTreeMap) -> Self { + Self { sources } + } +} + +impl View for Region<'_> { + type VectorsRef<'a> = rowmajor::Ref<'a, f32>; + type Vectors<'a> + = rowmajor::Ref<'a, f32> + where + Self: 'a; + + fn get_list(&self, list: u32) -> Option> { + self.sources.get(&list).map(|source| List { + id: list, + centroid: &source.centroid, + members: &source.members, + }) + } + + fn get_vectors(&self, list: u32) -> Option> { + self.sources + .get(&list) + .map(|source| source.vectors.as_view()) + } +} + +impl crate::ivf::region::Update for Update<'_> { + fn set_members(&mut self, list: u32, members: &[u32]) -> ANNResult<()> { + if self.fail_set_members { + return Err(ANNError::message("injected set_members failure")); + } + self.check_list(list)?; + self.memberships.set(list, members.to_vec()) + } + + fn retire(&mut self, list: u32) -> ANNResult<()> { + if self.fail_retire { + return Err(ANNError::message("injected retire failure")); + } + self.check_list(list)?; + self.memberships.retire(list) + } + + fn create(&mut self, centroid: &[f32], members: &[u32]) -> ANNResult<()> { + if self.fail_create { + return Err(ANNError::message("injected create failure")); + } + if centroid.len() != self.dim || !centroid.iter().all(|value| value.is_finite()) { + return Err(ANNError::message("invalid centroid")); + } + self.created.push(StoredList { + centroid: centroid.to_vec(), + members: members.to_vec(), + }); + Ok(()) + } +} + +impl HasIds for Build<'_, T> { + type Id = u32; + type ListId = u32; +} + +impl<'p, 'b, 'v, T: Copy + Into + Send + Sync> BuildAccessor<&'b [(u32, &'v [T])]> + for Build<'p, T> +{ + type Update<'a> + = Update<'a> + where + Self: 'a; + type Nearest<'a> + = Coarse<'a> + where + Self: 'a; + type VectorsRef<'a> = rowmajor::Ref<'a, f32>; + type Region<'a> + = Region<'a> + where + Self: 'a; + + fn dim(&self) -> NonZeroUsize { + NonZeroUsize::new(self.store.dim).unwrap() + } + + fn nearest(&mut self) -> Self::Nearest<'_> { + Coarse { + stored: &self.store.lists, + fail_route: self.store.fail_route, + fail_nearest: self.store.fail_nearest, + } + } + + fn route( + &mut self, + points: &'b [(u32, &'v [T])], + mut emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::ListId, NonZeroUsize) + Send, + { + async move { + let inserted = NonZeroUsize::new(points.len()) + .ok_or_else(|| ANNError::message("cannot insert an empty batch"))?; + let mut staged = BTreeMap::new(); + let mut routed = BTreeMap::>::new(); + for &(external, vector) in points { + if vector.len() != self.store.dim + || !vector.iter().all(|&value| value.into().is_finite()) + { + return Err(ANNError::message("invalid input vector")); + } + let id = external + .checked_add(10_000) + .ok_or_else(|| ANNError::message("internal point ID overflow"))?; + if self.store.points.contains_key(&id) + || self.staged.contains_key(&id) + || staged.contains_key(&id) + { + return Err(ANNError::message("point already exists")); + } + let mut count = 0usize; + { + let router = self.nearest(); + router + .route(vector, ONE, |&list, _| { + count += 1; + routed.entry(list).or_default().push(id); + }) + .await?; + } + if count != 1 { + return Err(ANNError::message("IVF routing returned no live centroid")); + } + staged.insert(id, vector.to_vec()); + } + self.staged.extend(staged); + for (list, members) in routed { + let size = self + .store + .lists + .get(&list) + .ok_or_else(|| ANNError::message(format!("unknown list {list}")))? + .members + .len(); + self.pending.append(list, members)?; + let size = NonZeroUsize::new(self.pending.list_size(list, size)?) + .ok_or_else(|| ANNError::message("appending must leave a nonempty list"))?; + emit(&list, size); + } + Ok(inserted) + } + } + + fn fill( + &mut self, + lists: Itr, + ) -> impl SendFuture, Self::Update<'_>)>> + where + Itr: Iterator + Clone + Send + Sync, + { + async move { + self.fills.set(self.fills.get() + 1); + if self.store.fail_fill { + return Err(ANNError::message("injected fill failure")); + } + self.filled_lists.extend(lists.clone()); + self.sources.clear(); + for list in lists { + let stored = self + .store + .lists + .get(&list) + .ok_or_else(|| ANNError::message(format!("unknown list {list}")))?; + let members: Vec<_> = self.pending.members(list, &stored.members)?.collect(); + let mut vectors = + rowmajor::Owned::from_element(members.len(), self.dim().get(), 0.0); + for (&id, row) in members.iter().zip(vectors.rows_mut()) { + if let Some(vector) = self.staged.get(&id) { + for (dst, &src) in row.iter_mut().zip(vector) { + *dst = src.into(); + } + } else { + let vector = self + .store + .points + .get(&id) + .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; + row.copy_from_slice(vector); + } + } + self.sources.insert( + list, + Source { + centroid: stored.centroid.clone(), + members, + vectors, + }, + ); + } + let mut update = Update::new( + &self.sources, + self.store.dim, + &mut self.pending, + &mut self.created, + ); + update.fail_create = self.store.fail_create; + update.fail_set_members = self.store.fail_set_members; + update.fail_retire = self.store.fail_retire; + Ok((Region::new(&self.sources), update)) + } + } + + fn finish(mut self) -> impl SendFuture> { + async move { + if self.store.fail_finish { + return Err(ANNError::message("injected finish failure")); + } + let created = u32::try_from(self.created.len()) + .map_err(|_| ANNError::message("list ID overflow"))?; + let next_list = self + .store + .next_list + .checked_add(created) + .ok_or_else(|| ANNError::message("list ID overflow"))?; + for (id, vector) in self.staged { + self.store + .points + .insert(id, vector.into_iter().map(Into::into).collect()); + } + for (id, list) in (self.store.next_list..next_list).zip(self.created) { + self.store.lists.insert(id, list); + } + for (id, change) in self.pending.into_changes() { + match change { + Membership::Append(members) => { + self.store + .lists + .get_mut(&id) + .unwrap() + .members + .extend(members); + } + Membership::Replace(members) => { + self.store.lists.get_mut(&id).unwrap().members = members; + } + Membership::Retire => { + self.store.lists.remove(&id).unwrap(); + } + } + } + self.store.next_list = next_list; + self.store.finishes += 1; + self.store.fills += self.fills.get(); + self.store.filled_lists = self.filled_lists; + Ok(()) + } + } +} + +pub(super) struct Strategy { + pub(super) store: Mutex, +} + +impl Strategy { + pub(super) fn new(store: Store) -> Self { + Self { + store: Mutex::new(store), + } + } + + pub(super) fn store(&self) -> MutexGuard<'_, Store> { + self.store.try_lock().unwrap() + } +} + +impl HasIds for Strategy { + type Id = u32; + type ListId = u32; +} + +impl<'a, T: Copy + Into + Send + Sync> BuildStrategy<'a, &[(u32, &[T])], DefaultContext> + for Strategy +{ + type Accessor = Build<'a, T>; + type Error = ANNError; + + fn build_accessor(&'a self, _: &'a DefaultContext) -> ANNResult { + Ok(Build { + store: self.store.try_lock().map_err(ANNError::new)?, + pending: PendingLists::default(), + created: Vec::new(), + staged: BTreeMap::new(), + sources: BTreeMap::new(), + fills: Cell::new(0), + filled_lists: Vec::new(), + }) + } +} + +pub(super) struct Search<'p, 'q, T> { + pub(super) store: MutexGuard<'p, Store>, + pub(super) query: &'q [T], +} + +impl HasIds for Search<'_, '_, T> { + type Id = u32; + type ListId = u32; +} + +impl<'q, T: Copy + Into + Send + Sync> SearchAccessor<&'q [T]> for Search<'_, 'q, T> { + fn scan(&self, nprobe: NonZeroUsize, results: &mut A) -> impl SendFuture> + where + A: Accumulator, + { + async move { + if self.store.fail_route { + return Err(ANNError::message("injected routing failure")); + } + let coarse = Coarse { + stored: &self.store.lists, + fail_route: false, + fail_nearest: false, + }; + let mut selected = Vec::new(); + coarse.select(self.query, nprobe, |_| true, |id, _| selected.push(*id)); + + for list in selected { + let source = self + .store + .lists + .get(&list) + .ok_or_else(|| ANNError::message(format!("unknown list {list}")))?; + let mut partial = results.empty_like(); + for id in &source.members { + let vector = self + .store + .points + .get(id) + .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; + partial.push(*id, squared_distance(self.query, vector)); + } + results.merge(partial); + } + Ok(()) + } + } +} + +impl<'a, 'q, T: Copy + Into + Send + Sync> SearchStrategy<'a, &'q [T], DefaultContext> + for Strategy +{ + type Accessor = Search<'a, 'q, T>; + type Error = ANNError; + + fn search_accessor( + &'a self, + _: &'a DefaultContext, + query: &'q [T], + ) -> ANNResult { + let store = self.store.try_lock().map_err(ANNError::new)?; + if query.len() != store.dim || !query.iter().all(|&value| value.into().is_finite()) { + return Err(ANNError::message("invalid query vector")); + } + Ok(Search { store, query }) + } +} + +pub(super) fn matrix(rows: &[[f32; D]]) -> rowmajor::Owned { + let data: Box<[f32]> = rows.iter().flatten().copied().collect(); + rowmajor::Owned::try_from_data(data, rows.len(), D).unwrap() +} + +pub(super) fn index(store: Store, threshold: usize) -> (IVFIndex, Strategy) { + index_with_neighbors(store, threshold, 0) +} + +pub(super) fn index_with_neighbors( + store: Store, + threshold: usize, + neighbors: usize, +) -> (IVFIndex, Strategy) { + let index = IVFIndex::new(Config { + split_threshold: NonZeroUsize::new(threshold).unwrap(), + reassign_count: NonZeroUsize::new(neighbors) + .map_or(ReassignCount::None, ReassignCount::Count), + split_iterations: NonZeroUsize::new(5).unwrap(), + seed: 7, + }) + .unwrap(); + (index, Strategy::new(store)) +} + +pub(super) fn seeded() -> Store { + Store { + dim: 2, + points: [ + (10_000, vec![0.0, 0.0]), + (10_001, vec![2.0, 0.0]), + (10_002, vec![10.0, 0.0]), + (10_003, vec![12.0, 0.0]), + (10_004, vec![100.0, 0.0]), + ] + .into(), + lists: [ + ( + 1, + StoredList { + centroid: vec![1.0, 0.0], + members: vec![10_000, 10_001], + }, + ), + ( + 2, + StoredList { + centroid: vec![11.0, 0.0], + members: vec![10_002, 10_003], + }, + ), + ( + 3, + StoredList { + centroid: vec![100.0, 0.0], + members: vec![10_004], + }, + ), + ] + .into(), + next_list: 4, + ..Store::default() + } +} + +pub(super) fn require_send(future: F) -> F { + future +} diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs new file mode 100644 index 0000000000..7a9bafc7ce --- /dev/null +++ b/diskann/src/ivf/traits.rs @@ -0,0 +1,160 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Storage boundaries for IVF construction and search. + +use std::num::NonZeroUsize; + +use diskann_utils::future::SendFuture; + +use super::region::{Update, View}; +use crate::{ + ANNResult, + error::StandardError, + neighbor::{Neighbor, NeighborPriorityQueue, NeighborPriorityQueueIdType}, + provider::ExecutionContext, + utils::VectorId, +}; + +pub trait HasIds { + type Id: VectorId; + type ListId: VectorId; +} + +pub trait NearestList: Send { + type Id; + /// Select `nprobe` distinct live lists nearest to `id`'s centroid, that pass the filter. + /// + /// Results are nearest first and may include `id` itself. + /// + /// # Errors + /// + /// Fails if `id` is not live or the provider cannot complete the lookup. + fn find( + &self, + id: Self::Id, + nprobe: NonZeroUsize, + pred: P, + emit: F, + ) -> impl SendFuture> + where + P: Fn(&Self::Id) -> bool + Send, + F: FnMut(&Self::Id, f32) + Send; +} + +pub trait BuildAccessor: Send + Sized + HasIds { + type Nearest<'a>: NearestList + where + Self: 'a; + + type VectorsRef<'a>; + + type Region<'a>: for<'b> View = Self::VectorsRef<'b>> + where + Self: 'a; + + type Update<'a>: Update + where + Self: 'a; + + /// Positive dimension of all canonical points and centroids. + fn dim(&self) -> NonZeroUsize; + + fn nearest(&mut self) -> Self::Nearest<'_>; + + /// Stage points and append their internal IDs to exactly one routed list each. + /// + /// Emit operation-local list sizes including pending appends. Lists may be emitted + /// more than once, but each touched list must report its final size. + /// Staged points and memberships remain unpublished until [`Self::finish`]. + fn route(&mut self, points: B, emit: F) -> impl SendFuture> + where + F: FnMut(&Self::ListId, NonZeroUsize) + Send; + + /// Make requested list's members, data and centroids available + /// for split and assignment, and register updates to the region. + /// + /// Writes recorded through the update handle remain pending until [`Self::finish`]. + fn fill( + &mut self, + lists: Itr, + ) -> impl SendFuture, Self::Update<'_>)>> + where + Itr: Iterator + Clone + Send + Sync; + + /// Finalize the operation, publishing accessor-side appends and regional writes together. + fn finish(self) -> impl SendFuture>; +} + +////////////// +// Strategy // +////////////// + +pub trait BuildStrategy<'a, B, C: ExecutionContext>: Send + Sync + HasIds { + /// [`BuildAccessor`] used to insert new points, and performs splits if necessary. + type Accessor: BuildAccessor; + + type Error: StandardError; + + fn build_accessor(&'a self, context: &'a C) -> Result; +} + +//////////// +// Search // +//////////// + +pub trait SearchStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { + /// Query-bound accessor responsible for list selection and scanning. + type Accessor: SearchAccessor; + + type Error: StandardError; + + fn search_accessor(&'a self, context: &'a C, query: T) -> Result; +} + +/// Candidate collection with independent partial states for partitioned scans. +pub trait Accumulator: Send + Sized { + /// Create an empty accumulator with the same selection policy and configuration. + fn empty_like(&self) -> Self; + + /// Consider one candidate for inclusion in the results. + fn push(&mut self, id: I, distance: f32); + + /// Combine a compatible partial state, such as one created by [`Self::empty_like`]. + fn merge(&mut self, other: Self); +} + +/// Queue-backed accumulation requires unique candidate IDs across partial states. +/// Equal-distance selections may depend on merge order. +impl Accumulator for NeighborPriorityQueue { + fn empty_like(&self) -> Self { + Self::new(self.search_l()) + } + + fn push(&mut self, id: I, distance: f32) { + self.insert(Neighbor::new(id, distance)); + } + + fn merge(&mut self, other: Self) { + for neighbor in other.iter() { + self.insert(neighbor); + } + } +} + +pub trait SearchAccessor: Send + Sized + HasIds { + /// Scan the `min(nprobe, live_list_count)` lists nearest the accessor's query. + /// + /// Implementations own selection and scheduling, and may reduce independent + /// accumulators from parallel scans into `results`. + /// + /// # Errors + /// + /// A failure may leave partial results. The caller must discard them rather + /// than publish them as a completed search. + fn scan(&self, nprobe: NonZeroUsize, results: &mut A) -> impl SendFuture> + where + A: Accumulator; +} diff --git a/diskann/src/lib.rs b/diskann/src/lib.rs index 0f8a40b931..be0af89e9b 100644 --- a/diskann/src/lib.rs +++ b/diskann/src/lib.rs @@ -17,6 +17,7 @@ pub(crate) mod internal; // Index Implementations pub mod flat; pub mod graph; +pub mod ivf; // Top level exports. pub use error::ann_error::{ANNError, ANNResult}; diff --git a/diskann/src/utils/utils.rs b/diskann/src/utils/utils.rs index 78d3fb33fc..9c74917335 100644 --- a/diskann/src/utils/utils.rs +++ b/diskann/src/utils/utils.rs @@ -6,6 +6,8 @@ use std::num::NonZeroUsize; /// A constant representing the number one as a `NonZeroUsize`. pub const ONE: NonZeroUsize = NonZeroUsize::new(1).unwrap(); +/// A constant representing the number two as a `NonZeroUsize`. +pub const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap(); /// Read exact number of POD type elements from a reader into a vector. pub fn read_exact_into(reader: &mut R, count: usize) -> std::io::Result> diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md new file mode 100644 index 0000000000..57f551a5d4 --- /dev/null +++ b/rfcs/01187-incremental-ivf.md @@ -0,0 +1,1131 @@ +# Incremental IVF Accessor Architecture + +| | | +|---|---| +| **Authors** | Aditya Krishnan | +| **Created** | 2026-08-19 | +| **Updated** | 2026-10-15 | +| **Status** | Draft | + +## Summary + +The implementation in [the IVF module](../diskann/src/ivf/mod.rs) follows the +graph index's coarse accessor boundaries. The index owns routing policy, +overflow selection, two-means fitting, and regional assignment. Providers supply: + +1. Point routing through `BuildAccessor::route`, and nearest-list selection through + `NearestList`. +2. A scoped `BuildAccessor` that stages ordinary appends, returns their resulting + sizes, fills requested sources, and publishes the operation. +3. A read-only `RegionView` for filled source data and an independent + `RegionUpdate` for list creation, complete membership replacement, and + retirement. + +Typed point preparation belongs to the insertion strategy, not to the core +build accessor. Preparation reserves internal IDs and supplies a shared, +contiguous canonical batch. Pending appends participate in operation-local +membership reads and sizes. Region filling resolves both stored and +prepared points, so clustering has no staged-point boundaries or provenance +classification. Each source supplies its centroid, member IDs, and a borrowed +matrix with one canonical vector row per member in the same order. + +Insertion provisionally appends points and identifies oversized parents. With +`ReassignCount::Count(n)`, it queries the nearest `n` lists per parent, excludes +every split parent, then fills the unique union of parents and neighbors once. +Nearest-list queries may return the queried parent or other split parents; +these consume query slots, so fewer than `n` retained neighbors may remain. +Shared neighbors are filled once. Parents and neighbors are ordered sets, and +filling accepts their chained ID iterator without collecting another vector. + +Two children are fitted per parent, using only that parent's complete membership. +All children are fitted before final assignment. Every gathered point is then +assigned across all new children and retained neighbor centroids; parent +centroids are not candidates. Retained lists keep their identities and centroids +while receiving complete replacement memberships. This full regional policy +permits parent-to-neighbor, neighbor-to-child, cross-parent, and neighbor-to-neighbor +movement. It is broader than the superseded per-parent candidate policy below. +The algorithm computes complete centroid and member groups before writing child +creation, retained memberships, and parent retirement through the update object. +Ordinary appends and splitting are published together, never as separate commits. +Each parent is split once per insertion; this is not a recursive size-balancing +operation, and retained neighbors are not split if reassignment makes them +oversized. Setting `reassign_count` to `ReassignCount::None` preserves parent-only +splitting. The count limits nearest-query results per parent, not the point +volume or total union size. Configuration uses nonzero split thresholds and +Lloyd iteration counts; the split threshold must still be at least two. + +Filling returns the source view and its regional writer together. The writer is +the accessor's `Update` type and is independent of the view's temporary input +borrows. The algorithm reads through the view and writes through the update +object; no update extraction step is needed. After the view's borrows end, the +index passes `Some(&update)` to `BuildAccessor::finish`. An insertion without +splitting passes `None`. The algorithm does not inspect the provider's update +representation or require new child IDs during computation. + +Regional writes may be buffered or applied eagerly inside the provider's private +operation context. An update object may own storage access or a transaction +handle; it need not be a passive edit buffer. Such writes must preserve live +source views, and successful finalization publishes one coherent result. +Providers coordinate accessor-side appends and regional writes so complete +replacement or retirement supersedes relevant earlier appends. + +The built-in split reads and edits only explicitly selected parents and neighbors, +even when a region contains additional sources. It covers every source member +exactly once, including pending appends to retained neighbors; other lists are +unchanged. Explicit source IDs avoid requiring enumeration on the region view. + +The split operation receives parent IDs, neighbor IDs, the region, and its writer, +and controls retrieval and conversion of provider-defined vector handles without +per-point reference tables. The built-in strategy borrows vector matrices and member +slices; a provider with scattered storage may pack rows during filling to satisfy +their matrix interface. The built-in two-means fitter seeds the first two input +rows and delegates to quantization's Lloyd routine, which +builds internal transposed scratch and resets empty centers to zero. Its +norm-based distance arithmetic does not preserve support for extreme finite +coordinates whose squared norms overflow. +`Split` is the complete numerical-policy boundary. The built-in `TwoMeansSplitter` +implementation allocates one centroid matrix with two rows per parent up front, +fits each parent into its mutable row window, then assigns +the union of parents and neighbors to the nearest new or retained centroid. +Equal-distance assignments prefer retained neighbors in supplied order, then +children in parent and centroid order. +Retained centroid coordinates are borrowed, not copied. Private helpers perform +fitting and assignment; no intermediate buffer protocol or independent stage +traits are required. After computation succeeds, the operation creates children, +replaces neighbor memberships, and retires the explicit parents. +The index currently selects this implementation directly. Other split strategies +may use different vector representations or combine fitting and assignment +differently. Restricted neighbor policies still require caller-supplied +parent-neighbor relationships, which the index does not yet preserve. +Alternative geometries must keep fitting, assignment, and provider routing +compatible. + +`Split::Error` allows implementation-specific failures. `TwoMeansSplitter` returns +`ANNError`, preserving typed errors for empty parent selections, centroid-count +overflow, missing sources, and insufficient points. Providers supply dimensionally +consistent vectors through infallible matrix-view conversion. Conversion of typed +errors to `ANNError` records the caller location. All numerical failures abort +before regional writes. + +Fitted centroids and assignment/member metadata need new storage. Providers can +reuse the canonical batch and optional pending-membership helper, or implement +their own caches and private transactions. The owned update is provider-defined; +incremental storage providers may derive sparse changes internally. + +Mutations remain serialized through `&mut self`. Pending list writes and point +mappings are not visible outside the operation. Failed computation records no +regional edits. Failed edit recording may leave partial private edits, so the +caller must discard the operation rather than publish it. Successful `finish` +must publish a coherent partition; recovery after failed publication is +provider-defined. Providers with externally shared state must coordinate those +aliases themselves. + +The current implementation supports batched insert/split into a provider that +already has centroids. Initialization, search, neighborhood-selection policies, +other than fixed-count nearest lists, deletion/dissolve, and colocation remain +future work. See +[the current contracts](../diskann/src/ivf/traits.rs) and +[the region interfaces](../diskann/src/ivf/region.rs) for the implemented API. + +The remainder of this RFC records the previous delta-oriented design and its +longer-term goals. Its API sketches, source layout, and `apply`-based algorithms +are superseded by the scoped construction interface described above. + +## Motivation + +### Background + +The initial IVF interface in `diskann/src/ivf` supports a conventional fixed +partition: + +1. A coarse accessor selects `nprobe` list ids. +2. A fine accessor scans those lists and emits `(point_id, distance)` pairs. +3. Insert selects one list and appends one point. + +That interface is appropriate for immutable IVF and fixed-centroid IVF-Flat. +It also follows the graph module's useful strategy/accessor split: algorithms +invoke coarse-grained operations and do not know how vectors or edges are laid +out. + +The experimental online GraphIVF algorithm has a larger mutation unit. One +insert batch may: + +- route many points against one centroid set; +- identify several overflowing lists; +- fit two children per overflowing parent with 2-means; +- create child centroids and retire parent centroids; +- read neighboring lists; and +- reassign all points in the affected regions. +- reassign centroids to other co-location sets, and create and retire co-location sets + +A delete batch removes points, retires underfull centroids, and moves each +victim list's remaining members to nearby survivors. These are set-level +partition changes rather than independent appends. + +### Problem Statement + +The fixed-partition interface does not provide enough behavior for GraphIVF. + +First, coarse and fine accessors are independently constructed. If mutation can +overlap search, centroid selection and list scanning may observe incompatible +states. + +Second, `append(list, id, vector)` cannot represent a split or dissolve. +Adding imperative `remove`, `create_list`, and `replace_centroid` calls would +make the generic algorithm choose a partial-write order for every backend. + +Third, GraphIVF planning needs list sizes, reverse assignments, list members, +canonical vectors, centroid vectors, fresh list ids, and co-location sets. These operations must +be available in batch-friendly forms so disk and blob implementations do not +devolve into one request per point. + +One solution is to standardize version tokens, pinned read snapshots, staged +mutations, compare-and-swap publication, and conflict retries. That provides +strong portable semantics, but it also makes a specific MVCC-like model part of +every provider implementation. It creates many associated types before the +in-memory GraphIVF algorithm has established which guarantees are necessary in +practice. + +This RFC instead makes consistency an accessor contract. The algorithm uses one +accessor for an entire operation; the implementation decides how that accessor +coordinates its provider, centroid index, and inverted-list store. + +### Goals + +1. Keep GraphIVF policy independent from list location, physical layout, + encoding, caching, and I/O scheduling. +2. Keep exact and graph centroid navigation interchangeable. +3. Support batched insert, local split, and regional reassignment. +4. Support batched delete, underfull-list dissolve, and reassignment. +5. Give providers coarse-grained, batch-friendly read operations. +6. Represent the final logical mutation declaratively rather than as ordered + low-level storage calls. +7. Keep canonical maintenance vectors independent from optimized list scan + representations such as PQ codes. +8. Reuse execution contexts, ranked errors, and the graph module's + strategy/accessor style, with one IVF-specific provider trait for identity. +9. Leave locking, reader visibility, durability, rollback, and recovery to each + concrete accessor implementation. + +### Non-Goals + +- A portable snapshot, generation, or compare-and-swap protocol. +- A generic transaction or write-ahead-log abstraction. +- Concurrent mutation through one `DynamicIvfIndex` in the first version. +- Portable concurrent search/mutation semantics. +- Distributed consensus or cross-service transactions. +- Defining a durable list file format. +- Choosing a specific k-means implementation. +- Residual PQ training or codebook lifecycle management. +- Safely composing arbitrary independently implemented centroid and list stores. + +## Terminology And Invariants + +### Terminology + +- **Point id**: the provider's internal vector id. External ids and the + translation between the two are owned by the provider (`Provider`). +- **List id**: the stable logical id shared by a centroid and its inverted list. + Like the point id, it belongs to the provider (`Provider::ListId`), so + every search and maintenance accessor over one provider uses the same list + ids. +- **Unified view**: mutually compatible point data, centroid state, reverse + assignments, and inverted lists presented by one operation accessor. +- **Selected lists**: list ids and their coarse distances, nearest first. A + search accessor scans only lists it selected itself. +- **Canonical vector**: the full-precision, or `f32`-decodable, representation + used for centroid fitting and exact regional reassignment. +- **Scan payload**: the representation stored with a list for query scoring. It + may be full precision, quantized, or otherwise encoded. +- **Partition update**: the complete logical centroid and point-membership + change calculated by the algorithm. +- **Co-Location Set**: a set of list ids that share some spatial similarity. May be the identity mapping. + +### Required Logical Invariants + +After a maintenance accessor returns `Ok(())` from `apply`: + +1. Every visible point is assigned to exactly one live list. +2. `assignment[point] == list` if and only if `point` is a member of `list`. +3. Every live list has exactly one live centroid with the same logical id. +4. Every list member has an available canonical vector and scan payload. +5. Retired list ids are never reused for a different centroid. +6. Point identity mappings agree with point visibility in the partition. +7. Inserts may trigger splits but not dissolves. +8. Deletes may trigger dissolves but not splits. +9. When dissolves are enabled, + `2 * merge_threshold <= split_threshold` provides hysteresis. +10. Each live list may belong to only one co-location set + +The generic API does not specify whether an accessor uses copied state, +exclusive access, or a versioned backend to preserve these invariants. + +### Accessor Consistency Contract + +A correct strategy must construct accessors with these properties: + +1. All methods on one maintenance accessor operate on one compatible logical + view until `apply` consumes it. +2. A search accessor scans only lists it selected itself. +3. Successful `apply` does not expose a mixture of old and new component state + through subsequent index operations. +4. If an operation fails after making irreversible changes, subsequent index + operations must either observe a coherent state or fail closed, for example + through provider poisoning. +5. Any stronger behavior, including rollback, crash atomicity, reader isolation, + durability, and retryability, is documented by the concrete provider. + +These rules are semantic requirements and are not encoded by public version +tokens. Reusable provider contract tests are therefore an important part of the +implementation. + +## Previous Proposal (Superseded) + +### Source Layout + +```text +diskann/src/ivf/ + mod.rs + dynamic.rs contracts proposed by this RFC + update.rs declarative partition updates applied by the maintenance accessor + grouped.rs flat grouped layouts used by planning + index.rs DynamicIvfIndex initialize/insert orchestration + online/ insert planning: staging, split planning, kernels + batch.rs staged and routed insert batch + split.rs the split plan: gather each region, fit children, place points + kernels.rs distance, nearest-candidate, and Lloyd kernels + test/ in-memory provider for index tests +``` + +The initial fixed-partition prototype (`glue.rs`, `index.rs`, and its test +provider) is removed; the contracts in `ivf::dynamic` replace it. + +### Responsibilities + +| Component | Owns | Does not decide | +|---|---|---| +| Dynamic IVF algorithm | split/dissolve policy, affected regions, child fitting, reassignment, co-location sets | physical layout, I/O shape, provider consistency mechanism | +| Aggregate provider/strategy | compatible point, centroid, and list implementations | split/dissolve policy | +| `CentroidIndex` | live centroid catalog and exact or approximate navigation | inverted-list storage | +| `CoLocationSet` | assignment of centroids to co-location groups | how co-location groups dictate ordering within the data provider | +| Search accessor | query-bound selection, list scoring, coherent search view | top-k policy | +| Maintenance accessor | unified planning view, point staging, id reservation, final update application | which centroids and points should move | + +The generic algorithm does not separately open a centroid index and inverted +list store and then attempt to coordinate them. A concrete strategy composes +those components and returns one accessor that already knows how they fit +together. Consequently, arbitrary implementations with matching Rust types are +not automatically safe to combine. + +The canonical vector belongs to the provider rather than to a particular list. +This permits a list backend to store only PQ codes while maintenance still reads +`f32` vectors. Moving a point between lists does not change its identity or +canonical representation, although the destination scan payload may need to be +encoded again. + +### Selected Lists + +Selection returns plain values, nearest first: + +```rust +pub struct SelectedList { + pub id: ListId, + pub distance: f32, +} +``` + +A search accessor may privately retain a lock guard, epoch, immutable root, +manifest, prefetch state, or other read context between selection and scanning. +None of that state appears in the public type, so selected lists are meaningful +only to the accessor that selected them. + +### Search Accessor + +One query-bound accessor owns both coarse and fine operation state: + +```rust +pub trait SearchAccessor: HasId + Send + Sync { + type ListId: VectorId; + type Error: ToRanked + Debug + Send + Sync + 'static; + + fn select_lists( + &mut self, + nprobe: usize, + ) -> impl SendFuture>, Self::Error>>; + + fn scan_lists( + &mut self, + lists: &[SelectedList], + emit: F, + ) -> impl SendFuture> + where + F: FnMut(Self::Id, f32) + Send; +} +``` + +The interface makes no callback-order guarantee. A backend may coalesce disk +ranges, combine blob requests, process lists concurrently, and merge scored +points before invoking `emit`. + +The strategy follows the graph module's accessor factory pattern. The provider +declares the identity types once and translates the ids of visible points; all +reads and writes go through accessors: + +```rust +pub trait Provider: Sized + Send + Sync + 'static { + type Context: ExecutionContext; + type InternalId: VectorId; + type ExternalId: PartialEq + Send + Sync + 'static; + type ListId: VectorId; + type Error: ToRanked + Debug + Send + Sync + 'static; + + fn to_internal_id( + &self, + context: &Self::Context, + gid: &Self::ExternalId, + ) -> Result; + + fn to_external_id( + &self, + context: &Self::Context, + id: Self::InternalId, + ) -> Result; +} + +pub trait SearchStrategy<'a, P: Provider, T>: Send + Sync { + type SearchAccessor: SearchAccessor; + type Error: StandardError; + + fn search_accessor( + &'a self, + provider: &'a P, + context: &'a P::Context, + query: T, + ) -> Result; +} +``` + +`ivf::dynamic::Provider` replaces the crate-wide `DataProvider` for this index. +That trait is shaped around graph insertion, where `SetElement` stores a point +immediately and a `Guard` later confirms or rolls it back; IVF points instead +stay staged in an accessor until `apply`, so the guard type and the separate +list-id trait it would need are not part of the IVF contract. Search returns +internal ids. + +Search orchestration is deliberately unaware of consistency mechanics: + +```text +accessor = strategy.search_accessor(provider, context, query) +lists = accessor.select_lists(nprobe) +accessor.scan_lists(&lists, |id, distance| top_k.insert(id, distance)) +post_process(point statuses and external ids) +``` + +### Centroid Index + +Centroids remain memory resident behind one interface: + +```rust +pub trait CentroidIndex: Send + Sync { + type ListId: VectorId; + type Error: ToRanked + Debug + Send + Sync + 'static; + + fn len(&self) -> usize; + fn is_empty(&self) -> bool; + fn centroid(&self, id: Self::ListId) -> Option<&[f32]>; + + fn select( + &self, + query: &[f32], + nprobe: usize, + ) -> impl SendFuture>, Self::Error>>; +} +``` + +An exact implementation scores a packed matrix of live centroids. A graph +implementation stores the same authoritative centroid catalog plus a DiskANN +graph used for navigation. It filters graph results against the live catalog +and completes an undersized result with exact selection. + +The generic interface does not prescribe how a centroid index is updated or +rebuilt. The maintenance accessor's `apply` implementation coordinates centroid +changes with list and point changes. A provider may mutate its graph in place, +construct a replacement graph, temporarily use exact routing, or poison itself +after an unrecoverable partial graph update. + +### Declarative Partition Updates + +The algorithm describes the desired logical result rather than issuing ordered +storage commands. The insert path's update in `ivf::update` lists primitive +changes: + +```rust +pub enum CentroidDelta { + Install { id: L, centroid: Box<[f32]> }, + Retire { id: L }, +} + +pub enum PointDelta { + // A newly staged point joins `to`. + Append { id: Id, to: L }, + // An existing point leaves `from` and joins `to`. + Move { id: Id, from: L, to: L }, +} + +pub struct InsertionUpdate { + centroids: Vec>, + points: Vec>, // sorted by destination, then id +} +``` + +Bootstrap and insert batches are both mixes of these deltas. Bootstrap installs +centroids. An insert appends its points; when it splits, it also installs two +children per parent, retires the parent, and moves the parent's members and any +neighbor members whose list changed. Each variant has one meaning in every +update, so a provider handles four kinds of change rather than every +combination of operations. The delete path's update is not defined yet; it +will follow the same delta style. + +`InsertionUpdate`'s fields are private and its constructor is private to the +IVF module, so only the index can create updates. Providers receive them +through `MaintenanceAccessor::apply` and can only read them. The delta enums are +necessarily public so providers can match on them, but they do nothing outside +an update. This keeps the provider contract one-directional: the algorithm +decides the update, and providers decide how to store it. + +The planner records each decision as a delta: every staged point is appended to +its final list, and every existing point whose list changed is moved. That +makes "every staged point is placed once" and "every parent member is moved out +once" properties of how the update is built, so the constructor only sorts the +point deltas and performs no validation. + +The enums make a move with neither a source nor a destination unrepresentable. +The planner guarantees the rest by construction, and `InsertionUpdate` +documents them for providers: + +- installed ids are fresh, and their centroids are finite with the accessor's + dimension; +- retired ids are live, and each list appears in at most one centroid delta; +- each point appears in at most one point delta, and every staged point is + appended; +- a move's source holds the point and differs from its destination; +- every destination is live after the update, so no point joins a retired + list; and +- every member of a retired list is moved out. + +Children may end up empty. + +#### Layout + +Point deltas are sorted by destination, then point id, so each list's incoming +points, appended and moved alike, form one contiguous run. A provider can +therefore append each list's incoming points with one write, re-encode residual +payloads against one centroid per run, and process distinct runs in parallel. +Because every member of a retired list is moved out, a provider may drop a +retired list wholesale instead of removing its points one by one. No point +vectors cross the update boundary, and each installed centroid crosses once. + +#### Validating Inputs Instead Of Updates + +Because the planner is trusted, the index validates untrusted data once, where +it enters, and builds updates without further checks: + +- caller input: initial centroids must be non-empty, finite, and of the + provider's dimension; +- vector reads and staging: every vector row must be written with finite + values (rows start as NaN, so unwritten rows are caught by the same check); +- staging: one distinct id must be returned per staged point; and +- reservations: reserved list ids must be distinct and not live. + +A provider that violates the accessor contract therefore fails the operation +before any update is built, rather than receiving an inconsistent update. + +#### Provider Checks + +View-dependent facts cannot be checked by the index. It guarantees them by +building updates only from what the same accessor returned: + +- installed ids were reserved by this accessor; +- retired lists and destinations are live in the accessor's planning view; +- every moved point belonged to its source list in that view; and +- every appended point was staged by this accessor. + +A provider may still verify these facts, for example in debug builds. + +The update remains useful even for an in-place backend. It gives the provider +the entire mutation before it chooses write order, batching, rollback, or +recovery behavior. + +Co-location set updates are expected to follow the same declarative pattern. +They are deferred until the insert path and maintenance accessor stabilize, so +the co-location types are not part of the code yet. + +### Maintenance Accessor + +The maintenance accessor holds planning reads, id reservation, and the final +`apply`; a separate trait stages new points, so only staging depends on the +input element type `T`: + +```rust +pub trait MaintenanceAccessor: Send + Sized { + type Centroids: CentroidIndex; + type Error: ToRanked + Debug + Send + Sync + 'static; + + fn dim(&self) -> usize; + fn centroids(&self) -> &Self::Centroids; + + fn list_size( + &mut self, + list: P::ListId, + ) -> impl SendFuture>; + + fn read_members( + &mut self, + list: P::ListId, + ) -> impl SendFuture>; + + fn read_vectors( + &mut self, + ids: &[P::InternalId], + out: MutMatrixView<'_, f32>, + ) -> impl SendFuture>; + + fn reserve_list_ids( + &mut self, + count: usize, + ) -> impl SendFuture, Self::Error>>; + + fn apply( + self, + update: InsertionUpdate, + ) -> impl SendFuture>; +} + +pub trait StageElements: MaintenanceAccessor

{ + fn stage( + &mut self, + points: &[(P::ExternalId, T)], + out: MutMatrixView<'_, f32>, + ) -> impl SendFuture, Self::Error>>; +} +``` + +The accessor takes its id types from the provider, so they are declared in one +place. + +Reads are shaped by how much data they move. List reads, `list_size` and +`read_members`, take one list and return its result: an insert batch makes at +most a few hundred of them, and a list's members are bounded by the split +threshold. `read_vectors` moves far more data, one vector per region member, so +it takes a batch and writes row `i` of the caller's matrix with the canonical +vector of `ids[i]`, in any order. Providers can therefore sort, coalesce, or +parallelize those reads. `dim` fixes the length of every canonical vector and +centroid, so callers can allocate before reading. A read of several lists at +once, for example a whole co-location group, can be added when co-location sets +are integrated. + +`stage` takes the whole insert batch in one call. The provider allocates +internal ids, rejects duplicate external ids, converts each element to its +stored representation, and writes the canonical `f32` vector of each point into +the caller's matrix. Staged points remain invisible, including to +`to_internal_id`, until the accessor applies an update that places them, and are +discarded if the accessor is dropped first. `read_vectors` therefore only covers +visible points. A provider that accepts several input element types implements +`StageElements` once per type. Reverse-assignment lookups and deletion staging +will be added with the delete path; the latter records point removal while +keeping whatever data the accessor needs for dissolve planning. + +These methods may delegate to existing provider machinery such as +`SetElement`, but the generic GraphIVF algorithm does not independently complete +a `SetElement` guard. The accessor coordinates point lifecycle with its other +components using its private mechanism. + +`apply` consumes the accessor. This prevents further planning reads through the +same operation object and provides one generic handoff containing the full +logical change. Bootstrap and inserts both produce an `InsertionUpdate`, so +`apply` is one accessor method rather than a trait per update type, and the +index needs no extra bound to apply what it planned. `apply` is not a generic +transaction protocol: + +- `Ok(())` guarantees a coherent resulting index. +- `Err` does not guarantee rollback or identify a conflict. +- Dropping an accessor does not have one portable cleanup guarantee. +- Reader visibility and durability are implementation-defined. + +A provider may document a stronger contract. For example, a blob provider can +privately stage immutable objects and conditionally replace a manifest even +though the generic API exposes only `apply`. + +The maintenance accessor is produced by a strategy that borrows the provider +exclusively: + +```rust +pub trait MaintenanceStrategy<'a, P: Provider>: Send + Sync { + type MaintenanceAccessor: MaintenanceAccessor

; + type Error: StandardError; + + fn maintenance_accessor( + &'a self, + provider: &'a mut P, + context: &'a P::Context, + ) -> Result; +} +``` + +Search strategies receive `&'a P`; maintenance strategies receive `&'a mut P`. +See [Index Borrowing Model](#index-borrowing-model). + +Because accessors take their ids from the provider, the index adds a bound only +where it stages: `S::MaintenanceAccessor: StageElements` for inserts, and +none for initialization. + +### Inverted-List Store Boundary + +This RFC does not add a low-level `InvertedListStore` trait. Storage backends +vary substantially in useful operation shape: an in-memory store borrows slices, +a disk store coalesces ranges, and a blob store batches objects. Flattening them +into point-level `get`, `append`, and `remove` methods would discard those +optimization opportunities. + +Search and maintenance accessors are the algorithm-facing inverted-list +abstraction. Concrete crates may define lower-level reusable list-store traits +and compose them inside a strategy. The aggregate accessor remains responsible +for compatibility with the selected centroid index and the provider. + +### Index Borrowing Model + +The initial dynamic index uses exclusive mutation: + +```rust +impl DynamicIvfIndex

{ + pub fn search(&self, /* ... */) -> impl SendFuture>; + + pub fn insert_batch(&mut self, /* ... */) + -> impl SendFuture>; + + pub fn delete_batch(&mut self, /* ... */) + -> impl SendFuture>; +} +``` + +Rust then prevents a search future borrowing this index from overlapping a +mutation future borrowing it mutably. Concurrent searches remain possible. + +Mutations forward that exclusivity to the provider: the index lends `&mut P` to +`MaintenanceStrategy::maintenance_accessor`, and the resulting accessor holds it +until `apply` consumes the accessor or it is dropped. An in-memory provider can +therefore implement `apply` as direct, in-place mutation of plain owned state, +with no locks or interior mutability. If the provider were lent only as `&P`, +every provider would need interior synchronization even though the index +already guarantees exclusivity. + +This rule covers operations through the index API, not aliases or administrative +paths exposed independently by a provider, such as state shared through an +`Arc`. Providers that expose such paths are responsible for coordinating them. + +This choice can later be relaxed to `&self` mutation for providers that support +it, but doing so requires a new concurrency contract and should be deliberate. +Such a relaxation would add a separate shared-borrow maintenance strategy +rather than weaken the exclusive one. + +## Previous Algorithms (Superseded) + +### Bootstrap + +Bootstrap uses the same maintenance interface as later mutations: + +1. Resolve initial full-precision centroids from an explicit matrix or a seed + strategy over caller-provided sample data. +2. Construct a maintenance accessor over an empty provider. +3. Reserve one stable list id per initial centroid. +4. Apply an `InsertionUpdate` installing those centroids, with no point deltas. + +Initial sample points are not implicitly assigned. Callers insert them through +the normal batch path. Reopen skips bootstrap when centroids already exist. + +### Search + +1. Validate `k`, `nprobe`, and query dimensions. +2. Construct one query-bound search accessor. +3. Select `min(nprobe, live_centroids)` lists. +4. Pass the resulting plan back to that accessor for scanning. +5. Accumulate internal ids in the existing neighbor priority queue. +6. Apply point-status filtering and external-id translation as appropriate. +7. Return results and coarse/fine statistics. + +No generic retry or version check occurs between selection and scanning. +Coherence follows from exclusive index mutation and the accessor implementation. + +Note that co-location groups are never accessed during search. + +### Batch Insert + +1. Construct one maintenance accessor for the operation. +2. Stage the batch in one call, obtaining internal ids and the canonical + vectors of the staged points. Duplicate or already-live external ids fail + according to provider semantics. +3. Check that the ids are distinct and every vector is finite. +4. Route all inputs through the accessor's centroid index, and group the batch + by route. +5. Read routed-list sizes and compute projected post-insert sizes. +6. Admit overflow parents subject to the live-cluster cap. +7. Reserve two stable child ids per parent. +8. Select each parent's nearest live lists, excluding every parent, as its + neighbors. Read the centroids of all neighbors, then the members and + canonical vectors of every parent and neighbor. +9. Fit each parent's two children with 2-means over its existing members and + the staged points routed to it, seeded with two of those points. +10. Place each parent's points on the nearest of its neighbors and two + children. +11. Leave each neighbor's points in place unless a child of a region containing + that neighbor is closer than the neighbor's centroid, in which case move + them to the nearest such child. +12. Build one `InsertionUpdate`. Append every staged point to its final list, + which is its route if nothing splits. For each split, install the two + children, retire the parent, and move every parent member and every + neighbor member whose list changed. +13. Consume the accessor with `apply(update)`. + +Steps 1 through 8 perform every accessor read; steps 9 through 12 compute on +owned, densely indexed data. Routes are computed before structural change; +regional placement deliberately revisits points whose routes become stale due +to the split. + +Neighbor points are handled as in SPFresh: they only move toward the new +children, not between existing lists, which bounds each neighbor point's work to +its own centroid and the new children. A neighbor shared by several regions is +compared against the children of all of them, which yields one placement per +point without merging per-region results. + +### Split Semantics + +A split retires one parent and installs two children, for a net increase of one +live cluster. Parent ids are never reused. A parent's candidate centroids are its +selected live neighbors plus its two children; a neighbor's candidates are its +own centroid plus the children of every region containing it. Candidate points +are the parent members, incoming points, and selected-neighbor members. + +The provider may rewrite complete lists, append delta segments, or mutate +in-memory lists. Those choices do not change the semantics of the update. + +### Batch Delete And Dissolve + +1. Construct one maintenance accessor and resolve every requested point's + current assignment. +2. Stage deletion of each point id. +3. Group deletes by list and compute projected post-delete sizes. +4. Admit underfull victims without crossing `min_clusters`. +5. Select survivor candidates for every victim, excluding all victims admitted + in this operation. +6. Read each victim's remaining members and canonical vectors. +7. Assign those members exactly among the saved survivor candidates. +8. Build one update that retires each victim, moves its remaining members to + their survivors, and removes the deleted points. Its type is not defined + yet; it follows the delta style of `InsertionUpdate`, with a point delta for + removal such as `Delete { id, from }`. +9. Consume the accessor with `apply(update)`. + +A dissolve fits no centroid and does not read survivor members. Removing a +centroid cannot make a point already assigned to another surviving centroid +prefer the removed centroid. Deletes never trigger a split in the same +operation. A survivor above the split threshold remains eligible for the next +insert-driven split. + +## Concurrency And Consistency + +The portable baseline is: + +- many concurrent searches through shared index borrows; +- one mutation through an exclusive index borrow, forwarded to the maintenance + accessor as `&mut P`; and +- no search/mutation overlap through the same index value. + +Within an operation, providers may parallelize routing, list reads, decoding, +and scoring through `ExecutionContext::wrap_spawn` and workspace thread-pool +conventions. + +Concrete implementations may offer more: + +| Provider | Possible private unified-view mechanism | +|---|---| +| In memory | exclusive `&mut P` borrow, one lock guard, copy-on-write root, or epoch | +| Local disk | process lock, database transaction, journal, or private manifest swap | +| Blob storage | immutable objects plus conditional manifest update | + +The generic traits neither require nor expose these mechanisms. A provider that +allows concurrent mutation and search must document whether readers see old +state, new state, block, or retry. That behavior is not portable across +strategies. + +## Failure Semantics + +Accessor and provider failures use the existing ranked error machinery. +Transient and critical low-level errors retain their normal meaning, but the +generic algorithm cannot infer the mutation's final state solely from an +`apply` error. + +Each maintenance implementation must choose and document one failure model, +for example: + +| Model | Behavior after an irreversible failure | +|---|---| +| Rollback | Restore the prior coherent state and remain usable | +| Complete-forward | Repair or finish the intended update before reopening | +| Poison | Reject later search and mutation until recovery or rebuild | +| Durable transaction | Rely on backend recovery to expose one coherent state | + +At minimum, an implementation must fail closed rather than return results from +a knowingly incoherent state. Generic GraphIVF returns the error and does not +automatically retry because there is no portable conflict or rollback signal. + +Cancellation has similarly provider-defined behavior. An in-memory accessor may +clean provisional state in `Drop`; an in-place implementation may hold exclusive +access and poison on interruption; a transactional backend may abort. The core +traits do not promise cancellation safety. + +## Configuration + +The initial dynamic algorithm validates at least: + +```rust +pub struct DynamicIvfConfig { + pub split_threshold: usize, + pub merge_threshold: usize, + pub min_clusters: usize, + pub max_clusters: Option, + pub reassign_neighbors: usize, + pub two_means_iterations: usize, +} +``` + +Centroid-index-specific configuration belongs to the strategy. An exact +strategy needs no graph degree or beam; a graph strategy owns its graph build +and search parameters. Durable metadata belongs to concrete provider crates. + +## Implementation Plan + +### Phase 0: Compile-Only Contracts + +- Add `ivf::dynamic` with the traits and data types in this RFC. +- Remove the fixed-partition prototype API and its tests. +- Compile, lint, format, and build rustdoc for the skeleton. + +### Phase 1: Unified In-Memory Provider + +- Build one provider containing canonical vectors, reverse assignments, + in-memory lists, and an exact centroid index. +- Implement maintenance through the exclusive `&mut P` borrow with + lock-free, in-place `apply`. +- Add an invariant checker and provider poisoning for injected partial failures. + +### Phase 2: Dynamic Search + +- Implement search orchestration with one query-bound accessor. +- Compare results to brute force over exactly the selected lists. +- Test that an accessor scans only lists it selected, under its own view. + +### Phase 3: Exact-Centroid GraphIVF + +- Port batch routing, split admission, per-parent child fitting, and + SPFresh-style regional placement from the prototype. +- Start with exact centroid selection to isolate partition correctness. +- Check every successful update against the reference partition model. + +### Phase 4: Graph Centroid Index + +- Implement `CentroidIndex` with `graph::DiskANNIndex` plus an authoritative + centroid catalog. +- Complete undersized graph results with exact selection. +- Define and test its partial-failure behavior under centroid churn. + +### Phase 5: Delete And Dissolve + +- Port delete grouping, victim admission, survivor selection, and reassignment. +- Add hysteresis, minimum-cluster, delete/reinsert, and overlapping-batch tests. + +### Phase 6: Durable Providers + +- Design local-disk consistency and recovery inside its concrete accessor. +- Add crash/reopen tests around each irreversible write. +- Design blob coordination separately rather than forcing disk and blob through + one generic publication protocol. + +## Testing Strategy + +Algorithm tests use a small exact reference model containing a centroid map, +list map, and reverse assignment map. After each successful `apply`, tests +verify every logical invariant and compare provider state to the expected +update. + +Required shared tests include: + +- search equivalence to brute force over selected lists; +- split and dissolve invariant preservation; +- overlapping maintenance regions; +- duplicate insert and repeated delete handling; +- staged canonical-vector decoding; +- stable, never-reused list ids; +- centroid graph exact fallback; and +- provider failure injection with coherent-state or fail-closed verification. + +Durability, cancellation, reader isolation, and recovery tests are +provider-specific because those guarantees are not part of the shared traits. + +## Advantages + +1. Far fewer associated types and generic constraints. +2. In-memory implementations can use normal exclusive mutation without + pretending to implement MVCC. +3. The API resembles the graph module's operation-scoped accessor model. +4. Providers retain their native consistency and durability mechanisms. +5. The GraphIVF algorithm remains readable: it plans one logical update and + hands it to one accessor. +6. The design does not prematurely standardize conflict detection, version + formats, reclamation, or crash recovery. +7. Exact, graph, memory, disk, and blob implementations can evolve without + exposing backend handles in generic types. + +## Disadvantages + +1. Coherence is documented rather than type-enforced. +2. Arbitrary centroid and list implementations cannot safely be mixed merely + because their associated Rust types match. +3. Concurrent search/mutation behavior is not portable. +4. `apply` errors do not communicate whether an operation rolled back, partly + applied, completed forward, or poisoned the provider. +5. Generic code cannot detect write conflicts or automatically replan. +6. Cancellation and crash semantics differ by provider. +7. Disk and blob implementations probably still need versioning or atomic + publication internally; the complexity is hidden, not eliminated. +8. Correctness depends more heavily on provider contract tests and operational + documentation. + +## Alternatives + +### Public Snapshot And Commit Protocol + +The stronger alternative exposes a provider-defined generation, pinned read +snapshot, snapshot-bearing selection plan, staged mutation, commit outcome, and +compare-and-swap conflict. It makes reader coherence, mutation visibility, and +retry behavior portable. + +That design is preferable if concurrent search/mutation, multiple writers, or +uniform crash-safe persistence are immediate requirements. It is not selected +for the first implementation because it imposes substantial machinery on every +provider before the in-memory algorithm exists. + +### Separate Coarse And Fine Accessors + +Independent accessors are smaller individually, but duplicate query setup and +have no common owner for consistency or prefetch state. One query-bound accessor +keeps coarse and fine operations together. + +### Pass The Coarse Accessor Into Fine Scan + +Passing `&mut CoarseAccessor` into a fine accessor shares state but couples two +objects, complicates async lifetimes, and makes selected lists difficult to +inspect. Keeping the selection as a small value and private state in one +accessor is simpler. + +### Imperative List Mutation Methods + +Immediate, independently published list mutations expose a partial-write schedule +to GraphIVF. The current design instead stages ordinary appends in the build +accessor and regional writes through a `RegionUpdate`, independently of the +`RegionView` used for input reads. The writer may buffer changes or apply them +inside a private operation context, and the accessor publishes both kinds of +changes only at `finish`. This keeps get/set-style operations without requiring +an algorithm-facing delta protocol or prescribing the provider's physical write +order. + +### Optional Endpoints Instead Of Delta Variants + +An earlier draft used a single `PartitionUpdate` with a free-form centroid +delta, a flat list of `PointMove { id, from: Option, to: Option }`, and +public fields. Optional endpoints made a move with neither endpoint +representable and left the meaning of each combination implicit. Public fields +let any code build invalid values, such as a retired list that was never +emptied or centroids of different dimensions. Unsorted moves made every +provider regroup them by list. Named delta variants, construction restricted to +the index, and destination-sorted point deltas remove those problems. + +### Per-Operation Update Types + +An intermediate design gave each operation its own update type: `Bootstrap`, +`Appends`, `Splits`, and `Reassignments`, with `SplitInsert` combining the last +three. Each type stored its points pre-grouped, and a provider implemented an +`Apply` trait once per type it supported. The types made some structure +explicit, such as the two children replacing each parent, and let a provider +that never splits leave out `Apply`. However, the insert path alone +needed eight types and three `Apply` implementations, and a destination list's +incoming points were spread across several groups. Nothing consumed the extra +structure, and inserts required every `Apply` bound anyway. Primitive deltas +express every insert-path change in one type. If a graph centroid index or +co-location planning needs a split's lineage, `CentroidDelta::Install` can gain +a parent link. + +### Publicly Constructible Updates + +An intermediate design gave every update type a public, validating constructor +with its own error type, so any code could build updates. Because only the index +ever needs to build them, that validation duplicated guarantees the planner +already provides by construction, re-sorted data the planner had already +grouped, and roughly doubled the size of the update module. Restricting +construction to the index and validating untrusted inputs where they enter +removes that duplication. Provider crates exercise their `apply` +implementations through the index instead of hand-built updates. + +### Use `SetElement` Independently + +Completing a normal `SetElement` guard separately from list and centroid changes +can expose inconsistent point lifecycle. A maintenance accessor may reuse the +same internal machinery, but it coordinates point staging and final application +as one provider-specific operation. + +### Require A Public Low-Level List Store + +A universal low-level store tends toward point-at-a-time methods that are poor +for range I/O and object storage. Accessors preserve coarse operations. A lower +level trait can be introduced later once multiple concrete providers reveal a +useful common shape. + +## Design Questions For Review + +1. **Exclusive mutation**: keep `insert_batch` and `delete_batch` on `&mut self` + for the first implementation, as recommended, or require provider-managed + concurrent access immediately? +2. **Failure health**: should a small common health/poison trait be required, or + should all post-error behavior remain provider-specific? +3. **Point staging**: should `StageElements`, and deletion staging once it is + added, remain on the maintenance accessor, or should a future + transactional provider trait supply them? +4. **Update granularity**: point deltas are sorted by destination list, so bulk + backends can write each touched list once. Should updates also carry a + per-list summary of added and removed counts for exact pre-sizing? +5. **Selected-list ownership**: is the documented same-accessor requirement + sufficient, or should a future private token make misuse dynamically + detectable? + +## Benchmark Results + +This is a design-only RFC. No performance claims are made. Benchmarks begin once +the exact-centroid in-memory implementation provides a correctness baseline. + +## Future Work + +- [ ] Provider-specific concurrent search/mutation implementations. +- [ ] Multiple non-overlapping writers. +- [ ] A common health/poison interface if implementations converge on one. +- [ ] Background list compaction. +- [ ] Residual PQ and codebook lifecycle management. +- [ ] Soft deletes followed by asynchronous dissolve. +- [ ] Filter-aware list selection and scan integration. +- [ ] Durable checkpoint and recovery tooling. +- [ ] A low-level list-store trait informed by concrete providers. +- [ ] Distributed or sharded list providers. + +## References + +1. [PR #1187: WIP IVF index interface](https://github.com/microsoft/DiskANN/pull/1187) +2. [RFC 01067: Refactor Search Accessor](01067-search-accessor.md) +3. [RFC 00983: Flat Search](00983-flat-search.md) \ No newline at end of file