From 64dedefd44d4293d85df7f5e97dfc2899eef41e6 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 16 Jun 2026 10:45:02 -0400 Subject: [PATCH 01/33] init rfc --- rfcs/00000-ivf-index.md | 148 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 rfcs/00000-ivf-index.md diff --git a/rfcs/00000-ivf-index.md b/rfcs/00000-ivf-index.md new file mode 100644 index 0000000000..b0f331f23a --- /dev/null +++ b/rfcs/00000-ivf-index.md @@ -0,0 +1,148 @@ +# Inverted File (IVF) Index + +| | | +|------------------|--------------------------------| +| **Authors** | Aditya Krishnan | +| **Created** | 2026-06-10 | +| **Updated** | 2026-06-10 | +| **Status** | Draft — requirements gathering | + +> **Note**: This is a *requirements* document. It deliberately stays at the +> level of "what we need and why" and defers concrete type/API design to a +> follow-up RFC. Sections are intentionally terse so we can iterate quickly. + +## 1. Summary + +Add an **Inverted File (IVF)** index to DiskANN. An IVF index partitions the +dataset into a fixed number of *inverted lists* (a.k.a. cells / clusters), each +summarized by a **centroid**. Each list stores the representations (full +precision or quantized) of the points assigned to it, laid out contiguously. + +- **Insert**: assign a point to the list whose centroid is closest, then append + the point's representation to that list. +- **Search**: select the `nprobe` lists whose centroids are closest to the + query, score the query against every point in those lists, and return the + top-`k` via the neighbor priority queue. + +## 2. Motivation + +### 2.1 Background + +DiskANN today exposes two index families: + +- **Graph search** — random-access greedy search over a proximity graph + ([`crate::graph::DiskANNIndex`], driven by the [`crate::provider::Accessor`] + trait). +- **Flat search** — sequential one-pass scan-and-score + ([`crate::flat::FlatIndex`], driven by [`crate::flat::DistancesUnordered`]). + +There is currently **no clustering / coarse-quantizer based index**. IVF is the +canonical such structure and a well-understood baseline in the ANN literature +(IVF, IVFPQ, etc.). + +### 2.2 Problem Statement + +We want a partition-based index that: + +- Trades recall for latency via a single tunable knob (`nprobe`). +- Has predictable, scan-friendly memory layout (contiguous lists) that maps + cleanly onto the existing flat-search scan-and-score primitives. +- Composes with the existing quantization stack so lists can store quantized + representations. + +### 2.3 Goals + +1. Define an IVF index that supports **build**, **insert**, and **top-k + search** with an `nprobe` parameter. +2. **Reuse existing primitives** rather than inventing parallel machinery: + - [`crate::provider::DataProvider`] for id mapping / context. + - The flat-search scan-and-score path ([`DistancesUnordered`]) for scoring + points inside a list. + - [`crate::neighbor`] priority queue for top-k accumulation. + - The quantization crate for list representations. +3. Support both **full-precision** and **quantized** list representations behind + a common representation abstraction. +4. Make the **coarse quantizer** (centroid set + assignment) a pluggable + component so we can later swap k-means for alternatives. +5. Make the representation and access of inverted lists to be pluggable + so it works with multiple data-backends (disk, k-v table, caching, blob). + +### 2.4 Non-Goals (for the first cut) + +- IVFPQ residual encoding / per-list re-ranking (future work). +- Deletes and compaction. + +## 3. Requirements + +### 3.1 Functional + +| ID | Requirement | +|----|-------------| +| F1 | **Build** an index from a dataset: derive `n_lists` centroids (initially k-means) and assign every point to its nearest centroid's list. | +| F2 | **Insert** a single point: find nearest inverted list (start by using centroids), append representation to that list. Insert must not require a full rebuild. | +| F3 | **Search**: given a query and parameters `k` and `nprobe`, select the `nprobe` nearest lists by centroid distance, scan-and-score their members, return top-`k`. List scoring should reuse the flat-search scan-and-score primitive [`DistancesUnordered`]; no bespoke distance loop. | | +| F4 | Lists can store any representation of vectors (e.g. **full-precision**, **quantized**), selected at build time. | +| F5 | The **coarse quantizer** (centroid training + assignment) is a distinct, swappable component. | +| F6 | The representation of inverted lists should be swappable, consumers should be able to implement it for different backends - disk, caching, inmemory, blob etc. | + +### 3.2 Non-Functional + +| ID | Requirement | +|----|-------------| +| N2 | Per-list storage is **contiguous** to enable sequential, SIMD-friendly scans. | +| N3 | `nprobe` is a per-query parameter (not fixed at build time). | +| N4 | Errors follow the mid-level [`ANNError`] regime (this is a `diskann`-crate algorithm). | +| N5 | Build and search are parallelizable, but must use the workspace thread-pool conventions (no global rayon pool — see `clippy.toml`). | +| N6 | Index parameters (`n_lists`, representation kind, distance function) are recorded so an index can be reloaded consistently. | + +### 3.3 Open Questions + +- **Crate placement**: new module under `diskann/` (alongside `flat/` and + `graph/`) vs. a dedicated `diskann-ivf` crate? Leaning toward a module first. + + **Answers**: + - Let's add it as a modulke to `diskann` for now. +- **Coarse quantizer ownership**: does the centroid set live inside the IVF + index, or is it a reusable component shared with quantization? + **Answers**: + - It can live within the index for now. We can refactor this later if we want it to be a reusable component. +- **Centroid representation**: full-precision centroids always, or allow + quantized centroids for the coarse search too? + **Answers**: + - Doesn't matter. The component in the index that returns the top inverted lists should be generic enough that we can implement + multiple algorithms/representations to perform this operation. Including full precision centroids, quantized centroids etc. +- **List growth**: per-list `Vec` vs. a single backing arena with list offsets. + Affects insert cost and scan locality. + **Answers**: + - We need to come up with a clean API for inserting a point to an inverted list. Specific implementations will implement this API + depending on their specific data backend. +- **Relationship to `DataProvider`**: is the IVF index a *consumer* of a + provider, or does it own its own storage backend? + **Answers**: + - Critically, the IVF index is a consumer of a provider. It does not own its backend. +- **Serialization**: reuse an existing on-disk format, or define a new one? + **Answers**: + - Let's not worry about serialization for now. + +## 4. Sketch (non-binding) + +A rough mental model to anchor discussion — *not* a committed design: + +```text +IvfIndex +├── coarse: CoarseQuantizer // n_lists centroids + assign(query) -> list ids +├── lists: [InvertedList; n_lists] // each: contiguous representations + local->external id map +└── search(query, k, nprobe): + candidate_lists = coarse.closest(query, nprobe) + for list in candidate_lists: + list.scan_and_score(query, &mut neighbor_queue) // via DistancesUnordered + return neighbor_queue.top_k(k) +``` + +## 5. Future Work + +- [ ] IVFPQ: store PQ residual codes per list; re-rank with full precision. +- [ ] Disk-resident lists / out-of-core search. +- [ ] Deletes, tombstoning, and list compaction. +- [ ] Advanced probing (soft assignment, learned routing). +- [ ] Label / attribute filtering integration (`diskann-label-filter`). From 82e8aeff3378e7cec345bf75c48ced1be2bd1ac8 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Thu, 18 Jun 2026 11:45:49 -0400 Subject: [PATCH 02/33] Add IVF index module with traits, example provider, and tests Introduces diskann::ivf with the IvfIndex wrapper, five core glue traits (ListAccessor, SearchAccessor, SearchStrategy, InsertAccessor, IvfInsertStrategy), an example list-owned IVF-Flat provider, an exact oracle harness, and a test suite covering correctness, concurrency, inserts (including concurrent inserts), and error escalation. --- diskann/src/ivf/glue.rs | 138 ++++ diskann/src/ivf/index.rs | 154 +++++ diskann/src/ivf/mod.rs | 18 + diskann/src/ivf/test/cases/concurrency.rs | 134 ++++ diskann/src/ivf/test/cases/correctness.rs | 117 ++++ diskann/src/ivf/test/cases/errors.rs | 111 +++ diskann/src/ivf/test/cases/insert.rs | 197 ++++++ diskann/src/ivf/test/cases/mod.rs | 76 +++ diskann/src/ivf/test/harness.rs | 182 +++++ diskann/src/ivf/test/mod.rs | 10 + diskann/src/ivf/test/provider.rs | 792 ++++++++++++++++++++++ diskann/src/lib.rs | 1 + 12 files changed, 1930 insertions(+) create mode 100644 diskann/src/ivf/glue.rs create mode 100644 diskann/src/ivf/index.rs create mode 100644 diskann/src/ivf/mod.rs create mode 100644 diskann/src/ivf/test/cases/concurrency.rs create mode 100644 diskann/src/ivf/test/cases/correctness.rs create mode 100644 diskann/src/ivf/test/cases/errors.rs create mode 100644 diskann/src/ivf/test/cases/insert.rs create mode 100644 diskann/src/ivf/test/cases/mod.rs create mode 100644 diskann/src/ivf/test/harness.rs create mode 100644 diskann/src/ivf/test/mod.rs create mode 100644 diskann/src/ivf/test/provider.rs diff --git a/diskann/src/ivf/glue.rs b/diskann/src/ivf/glue.rs new file mode 100644 index 0000000000..f3e846ec6e --- /dev/null +++ b/diskann/src/ivf/glue.rs @@ -0,0 +1,138 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::fmt::Debug; + +use diskann_utils::future::SendFuture; + +use crate::{ + error::{StandardError, ToRanked}, + graph::SearchOutputBuffer, + provider::{DataProvider, HasId}, +}; + +///////////////// +// Search path // +///////////////// + +/// Selects candidate inverted lists for a bound query or insert vector. +pub trait ListAccessor: Send + Sync { + /// Opaque handle identifying an inverted list within the index. + type Id: Copy + Send + Sync; + + /// The error type for [`Self::select_lists`]. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Push the selected lists into `output` in distance order. + fn select_lists( + &mut self, + nprobe: usize, + output: &mut B, + ) -> impl SendFuture> + where + B: SearchOutputBuffer + Send + ?Sized; +} + +/// Scans a set of lists for search. +pub trait SearchAccessor: HasId + Send + Sync { + /// Opaque handle identifying an inverted list within the index. + type ListId: Copy + Send + Sync; + + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Score members of `lists` and invoke `f` for each `(id, distance)` pair. + fn scan_lists(&mut self, lists: Itr, f: F) -> impl SendFuture> + where + Itr: Iterator + Send, + F: Send + FnMut(Self::Id, f32); +} + +/// Per-call factory for IVF search. +pub trait SearchStrategy<'a, Provider, T>: Send + Sync +where + Provider: DataProvider, +{ + /// The inverted-list handle type, shared by both accessors. + type ListId: Copy + Send + Sync; + + /// The fine accessor, keyed to the provider's internal id and the shared list handle. + type SearchAccessor: SearchAccessor; + + /// The coarse accessor, keyed to the shared list handle. + type ListAccessor: ListAccessor; + + /// An error that can occur when constructing either accessor. + type Error: StandardError; + + /// Construct the fine scan accessor. + fn search_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + query: T, + ) -> Result; + + /// Construct the coarse list-selection accessor. + fn list_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + query: T, + ) -> Result; +} + +///////////////// +// Insert path // +///////////////// + +/// Appends a vector to a chosen list during insert. +pub trait InsertAccessor: HasId + Send + Sync { + /// Opaque handle identifying an inverted list within the index. + type ListId: Copy + Send + Sync; + + /// The error type for [`Self::append`]. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Append `vector` to `list` under `id`. + fn append( + &mut self, + list: Self::ListId, + id: Self::Id, + vector: T, + ) -> impl SendFuture>; +} + +/// Per-call factory for IVF insert. +pub trait InsertStrategy<'a, Provider, T>: Send + Sync +where + Provider: DataProvider, +{ + /// The inverted-list handle type, shared by both accessors. + type ListId: Copy + Send + Sync; + + /// The append accessor, keyed to the provider's internal id and the shared list handle. + type InsertAccessor: InsertAccessor; + + /// The coarse accessor, keyed to the shared list handle. + type ListAccessor: ListAccessor; + + /// An error that can occur when constructing either accessor. + type Error: StandardError; + + /// Construct the append accessor. + fn insert_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + ) -> Result; + + /// Construct the coarse list-selection accessor. + fn list_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + vector: T, + ) -> Result; +} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs new file mode 100644 index 0000000000..d7b2d42c50 --- /dev/null +++ b/diskann/src/ivf/index.rs @@ -0,0 +1,154 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! IVF index wrapper. + +use std::num::NonZeroUsize; + +use diskann_utils::future::SendFuture; + +use crate::{ + ANNResult, + error::{ANNError, ANNErrorKind, ErrorExt, IntoANNResult}, + graph::SearchOutputBuffer, + ivf::{InsertAccessor, InsertStrategy, ListAccessor, SearchAccessor, SearchStrategy}, + neighbor::{Neighbor, NeighborPriorityQueue, NeighborPriorityQueueIdType}, + provider::{DataProvider, Guard, SetElement}, +}; + +/// Statistics collected during an IVF search. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct SearchStats { + /// Distance computations performed while scanning lists. + pub cmps: u32, + + /// Results written to the output buffer. + pub result_count: u32, +} + +/// IVF index wrapper over a [`DataProvider`]. +#[derive(Debug)] +pub struct IvfIndex { + provider: P, +} + +impl IvfIndex

{ + /// Construct a new index around `provider`. + pub fn new(provider: P) -> Self { + Self { provider } + } + + /// Borrow the underlying provider. + pub fn provider(&self) -> &P { + &self.provider + } + + /// Run IVF k-nearest-neighbor search. + pub fn knn_search<'a, S, T, OB>( + &'a self, + k: NonZeroUsize, + nprobe: usize, + strategy: &'a S, + context: &'a P::Context, + query: T, + output: &mut OB, + ) -> impl SendFuture> + where + S: SearchStrategy<'a, P, T>, + S::ListId: Eq, + P::InternalId: NeighborPriorityQueueIdType, + T: Copy + Send, + OB: SearchOutputBuffer + Send + ?Sized, + { + async move { + let mut list_accessor = strategy + .list_accessor(&self.provider, context, query) + .into_ann_result()?; + + let mut lists: Vec> = Vec::with_capacity(nprobe); + list_accessor + .select_lists(nprobe, &mut lists) + .await + .escalate("IVF coarse list selection must complete")?; + + let mut search_accessor = strategy + .search_accessor(&self.provider, context, query) + .into_ann_result()?; + + let k = k.get(); + let mut queue = NeighborPriorityQueue::new(k); + let mut cmps: u32 = 0; + + search_accessor + .scan_lists(lists.iter().map(|n| n.id), |id, dist| { + cmps += 1; + queue.insert(Neighbor::new(id, dist)); + }) + .await + .escalate("IVF list scan must complete to produce correct k-NN results")?; + + let result_count = + output.extend(queue.iter().take(k).map(|n| (n.id, n.distance))) as u32; + + Ok(SearchStats { cmps, result_count }) + } + } + + /// Insert a vector under external id `id`. + pub fn insert<'a, S, T>( + &'a self, + strategy: &'a S, + context: &'a P::Context, + id: &P::ExternalId, + vector: T, + ) -> impl SendFuture> + where + S: InsertStrategy<'a, P, T>, + S::ListId: Eq, + P: SetElement, + T: Copy + Send, + { + async move { + let guard = self + .provider + .set_element(context, id, vector) + .await + .escalate("IVF insert requires a successful `set_element`")?; + + let internal_id = guard.id(); + + let mut list_accessor = strategy + .list_accessor(&self.provider, context, vector) + .into_ann_result()?; + + let mut lists: Vec> = Vec::with_capacity(1); + + list_accessor + .select_lists(1, &mut lists) + .await + .escalate("IVF insert must select a target list")?; + + let list = lists.first().map(|n| n.id).ok_or_else(|| { + ANNError::message( + ANNErrorKind::IndexError, + "IVF insert: list selection returned no candidate list", + ) + })?; + + let mut insert_accessor = strategy + .insert_accessor(&self.provider, context) + .into_ann_result()?; + + insert_accessor + .append(list, internal_id, vector) + .await + .escalate("IVF insert must append the vector to its assigned list")?; + + guard.complete().await; + + Ok(()) + } + } +} diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs new file mode 100644 index 0000000000..f93b4c9c80 --- /dev/null +++ b/diskann/src/ivf/mod.rs @@ -0,0 +1,18 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! IVF index traits and wrapper. +//! +//! Search selects candidate lists, scans them, and returns the best `k` results. +//! Insert selects one list and appends the new point to it. + +pub mod glue; +pub mod index; + +pub use glue::{InsertAccessor, InsertStrategy, ListAccessor, SearchAccessor, SearchStrategy}; +pub use index::{IvfIndex, SearchStats}; + +#[cfg(test)] +mod test; diff --git a/diskann/src/ivf/test/cases/concurrency.rs b/diskann/src/ivf/test/cases/concurrency.rs new file mode 100644 index 0000000000..f13e974096 --- /dev/null +++ b/diskann/src/ivf/test/cases/concurrency.rs @@ -0,0 +1,134 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Concurrency behavior of the fine-scan accessor. +//! +//! [`ScanAccessor`](crate::ivf::test::provider::ScanAccessor) exposes two paths -- +//! [`ScanMode::Sequential`] and [`ScanMode::Concurrent`] (one spawned task per probed list). +//! These tests assert the two paths produce identical results, that the concurrent path +//! actually spawned tasks, and that many searches can run concurrently over a shared index. + +use std::sync::Arc; + +use crate::{ + ivf::test::{ + cases::{METRIC, N_LISTS, grid_index, queries}, + harness::{IvfOracleRun, assert_same_distances}, + provider::{Context, ScanMode, Strategy}, + }, + test::tokio::current_thread_runtime, +}; + +/// The sequential and concurrent scan paths must return the same top-`k` for every query. +#[test] +fn sequential_and_concurrent_agree() { + let index = grid_index(); + let dim = index.provider().dim(); + let sequential = Strategy::new(dim, METRIC).with_mode(ScanMode::Sequential); + let concurrent = Strategy::new(dim, METRIC).with_mode(ScanMode::Concurrent); + + for query in queries() { + for k in [1usize, 8] { + let seq = IvfOracleRun::run_sync(&index, &sequential, &query, k, N_LISTS).unwrap(); + let con = IvfOracleRun::run_sync(&index, &concurrent, &query, k, N_LISTS).unwrap(); + + assert_same_distances(&con.top_k, &seq.ground_truth); + assert_same_distances(&con.top_k, &seq.top_k); + assert_eq!(con.stats.cmps, seq.stats.cmps); + } + } +} + +/// The concurrent path must route its list scans through `wrap_spawn`, so the context +/// observes one spawn per probed list. +#[tokio::test] +async fn concurrent_path_spawns_per_list() { + let index = grid_index(); + let dim = index.provider().dim(); + let strategy = Strategy::new(dim, METRIC).with_mode(ScanMode::Concurrent); + + let context = Context::new(); + let mut buf = vec![crate::neighbor::Neighbor::::default(); 4]; + + let stats = index + .knn_search( + std::num::NonZeroUsize::new(4).unwrap(), + N_LISTS, + &strategy, + &context, + [3.5f32, 3.5].as_slice(), + &mut crate::neighbor::BackInserter::new(buf.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!( + context.spawns(), + N_LISTS, + "concurrent scan should spawn one task per probed list", + ); + assert_eq!(stats.cmps as usize, index.provider().len()); +} + +/// The sequential path performs no spawns. +#[tokio::test] +async fn sequential_path_does_not_spawn() { + let index = grid_index(); + let dim = index.provider().dim(); + let strategy = Strategy::new(dim, METRIC).with_mode(ScanMode::Sequential); + + let context = Context::new(); + let mut buf = vec![crate::neighbor::Neighbor::::default(); 4]; + + index + .knn_search( + std::num::NonZeroUsize::new(4).unwrap(), + N_LISTS, + &strategy, + &context, + [3.5f32, 3.5].as_slice(), + &mut crate::neighbor::BackInserter::new(buf.as_mut_slice()), + ) + .await + .unwrap(); + + assert_eq!(context.spawns(), 0); +} + +/// Many concurrent searches over a shared index (on a multi-threaded runtime) must each +/// return the correct restricted top-`k`. +#[test] +fn many_concurrent_searches_share_one_index() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .expect("multi-thread runtime"); + + let index = Arc::new(grid_index()); + let dim = index.provider().dim(); + let strategy = Arc::new(Strategy::new(dim, METRIC).with_mode(ScanMode::Concurrent)); + + runtime.block_on(async move { + let mut set = tokio::task::JoinSet::new(); + for query in queries() { + let index = Arc::clone(&index); + let strategy = Arc::clone(&strategy); + set.spawn(async move { + let run = IvfOracleRun::run(&index, &strategy, &query, 8, N_LISTS - 1) + .await + .unwrap(); + assert_same_distances(&run.top_k, &run.ground_truth); + }); + } + while let Some(joined) = set.join_next().await { + joined.expect("search task panicked"); + } + }); +} + +// Bring the single-thread runtime helper into scope for documentation/intra-doc links. +#[allow(unused_imports)] +use current_thread_runtime as _current_thread_runtime; diff --git a/diskann/src/ivf/test/cases/correctness.rs b/diskann/src/ivf/test/cases/correctness.rs new file mode 100644 index 0000000000..2334908ad4 --- /dev/null +++ b/diskann/src/ivf/test/cases/correctness.rs @@ -0,0 +1,117 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Correctness of [`crate::ivf::IvfIndex::knn_search`]. +//! +//! The central invariant is that an IVF search probing `nprobe` lists returns exactly the +//! brute-force top-`k` over the union of those lists' members. The harness computes that +//! restricted oracle directly, so each case can assert strict equality. + +use crate::ivf::test::{ + cases::{METRIC, N_LISTS, grid_index, queries}, + harness::{IvfOracleRun, assert_same_distances, global_topk}, + provider::Strategy, +}; + +/// For every query and a sweep of `(k, nprobe)`, the search must equal the restricted +/// brute-force oracle. +#[test] +fn matches_restricted_brute_force() { + let index = grid_index(); + let dim = index.provider().dim(); + let strategy = Strategy::new(dim, METRIC); + + for query in queries() { + for nprobe in 1..=N_LISTS { + for k in [1usize, 4, 16] { + let run = IvfOracleRun::run_sync(&index, &strategy, &query, k, nprobe).unwrap(); + + assert_same_distances(&run.top_k, &run.ground_truth); + + // `top_k_distances` is the distance projection of `top_k`. + let projected: Vec = run.top_k.iter().map(|(_, d)| *d).collect(); + assert_eq!(run.top_k_distances, projected); + + // The fine step must compute exactly one distance per probed member. + assert_eq!( + run.stats.cmps as usize, run.probed_members, + "query {query:?} k={k} nprobe={nprobe}: cmps must equal probed members", + ); + + // We never return more than the oracle (which is already truncated to k), + // and never more than k. + assert_eq!(run.top_k.len(), run.ground_truth.len()); + assert!(run.top_k.len() <= k); + } + } + } +} + +/// Probing every list (`nprobe == n_lists`) makes IVF exhaustive: it must equal the global +/// brute-force top-`k`. +#[test] +fn full_probe_equals_global_brute_force() { + let index = grid_index(); + let dim = index.provider().dim(); + let strategy = Strategy::new(dim, METRIC); + + for query in queries() { + for k in [1usize, 4, 16] { + let run = IvfOracleRun::run_sync(&index, &strategy, &query, k, N_LISTS).unwrap(); + + let global: Vec<(u32, f32)> = global_topk(index.provider(), &query, k) + .into_iter() + .map(|n| (n.id, n.distance)) + .collect(); + assert_same_distances(&run.top_k, &global); + + // Probing every list scans the whole dataset exactly once. + assert_eq!(run.stats.cmps as usize, index.provider().len()); + assert_eq!(run.probed_lists.len(), N_LISTS); + } + } +} + +/// `result_count` is `min(k, probed_members)` and matches the buffer contents. +#[test] +fn result_count_is_bounded() { + let index = grid_index(); + let dim = index.provider().dim(); + let strategy = Strategy::new(dim, METRIC); + let query = vec![0.0, 0.0]; + + // nprobe = 1 probes a single 16-member quadrant. + let run = IvfOracleRun::run_sync(&index, &strategy, &query, 100, 1).unwrap(); + assert_eq!(run.stats.result_count as usize, run.probed_members); + assert_eq!(run.top_k.len(), run.probed_members); + assert_eq!(run.stats.result_count as usize, run.top_k.len()); +} + +/// The provider's coarse/fine comparison counters reflect the work a search performs: +/// `n_lists` centroid comparisons (one full coarse pass) and one fine comparison per +/// probed member. +#[test] +fn provider_metrics_track_distance_work() { + let index = grid_index(); + let dim = index.provider().dim(); + let strategy = Strategy::new(dim, METRIC); + let query = vec![3.5f32, 3.5]; + + let before = index.provider().metrics(); + let run = IvfOracleRun::run_sync(&index, &strategy, &query, 4, N_LISTS).unwrap(); + let after = index.provider().metrics(); + + assert_eq!(index.provider().n_lists(), N_LISTS); + assert_eq!( + after.centroid_cmps - before.centroid_cmps, + index.provider().n_lists(), + "coarse step compares the query against every centroid once", + ); + assert_eq!( + after.member_cmps - before.member_cmps, + run.probed_members, + "fine step compares the query against every probed member once", + ); +} diff --git a/diskann/src/ivf/test/cases/errors.rs b/diskann/src/ivf/test/cases/errors.rs new file mode 100644 index 0000000000..4438a0723e --- /dev/null +++ b/diskann/src/ivf/test/cases/errors.rs @@ -0,0 +1,111 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Error propagation through the IVF index API. +//! +//! A failure in any accessor or strategy must surface as an [`ANNError`] out of +//! `knn_search` / `insert` rather than being swallowed or panicking. + +use std::num::NonZeroUsize; + +use crate::{ + ivf::test::{ + cases::{DIM, METRIC, N_LISTS, grid_index}, + provider::{Context, ScanMode, Strategy}, + }, + neighbor::{BackInserter, Neighbor}, + test::tokio::current_thread_runtime, +}; + +/// An injected fine-scan failure must escalate out of `knn_search` (sequential path). +#[test] +fn scan_failure_escalates_sequential() { + let index = grid_index(); + // The center query probes every quadrant; failing on list 0 is guaranteed to be hit. + let strategy = Strategy::new(DIM, METRIC) + .with_mode(ScanMode::Sequential) + .failing_on_list(0); + + let context = Context::new(); + let mut buf = vec![Neighbor::::default(); 4]; + let result = current_thread_runtime().block_on(index.knn_search( + NonZeroUsize::new(4).unwrap(), + N_LISTS, + &strategy, + &context, + [3.5f32, 3.5].as_slice(), + &mut BackInserter::new(buf.as_mut_slice()), + )); + + assert!(result.is_err(), "scan failure must propagate"); +} + +/// The same failure must escalate from the concurrent path, which has to shut down its +/// in-flight tasks before returning the error. +#[test] +fn scan_failure_escalates_concurrent() { + let index = grid_index(); + let strategy = Strategy::new(DIM, METRIC) + .with_mode(ScanMode::Concurrent) + .failing_on_list(0); + + let context = Context::new(); + let mut buf = vec![Neighbor::::default(); 4]; + let result = current_thread_runtime().block_on(index.knn_search( + NonZeroUsize::new(4).unwrap(), + N_LISTS, + &strategy, + &context, + [3.5f32, 3.5].as_slice(), + &mut BackInserter::new(buf.as_mut_slice()), + )); + + assert!( + result.is_err(), + "scan failure must propagate from concurrent path" + ); +} + +/// A query whose dimension disagrees with the strategy must fail at accessor construction. +#[test] +fn dimension_mismatch_errors() { + let index = grid_index(); + let strategy = Strategy::new(DIM, METRIC); + + let context = Context::new(); + let mut buf = vec![Neighbor::::default(); 1]; + let bad_query = vec![1.0f32, 2.0, 3.0]; // DIM is 2. + let result = current_thread_runtime().block_on(index.knn_search( + NonZeroUsize::new(1).unwrap(), + 1, + &strategy, + &context, + bad_query.as_slice(), + &mut BackInserter::new(buf.as_mut_slice()), + )); + + assert!(result.is_err(), "dimension mismatch must be rejected"); +} + +/// A dimension mismatch on insert is likewise rejected. +#[test] +fn insert_dimension_mismatch_errors() { + let index = grid_index(); + let strategy = Strategy::new(DIM, METRIC); + let new_id = index.provider().len() as u32; + let bad_vector = vec![1.0f32]; // DIM is 2. + + let result = current_thread_runtime().block_on(index.insert( + &strategy, + &Context::new(), + &new_id, + bad_vector.as_slice(), + )); + + assert!( + result.is_err(), + "insert dimension mismatch must be rejected" + ); +} diff --git a/diskann/src/ivf/test/cases/insert.rs b/diskann/src/ivf/test/cases/insert.rs new file mode 100644 index 0000000000..2eb9dd496d --- /dev/null +++ b/diskann/src/ivf/test/cases/insert.rs @@ -0,0 +1,197 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Insert behavior of [`crate::ivf::IvfIndex::insert`]. +//! +//! Inserting a vector must (a) store it in the provider, (b) append it to its nearest list, +//! and (c) make it findable by a subsequent search that probes that list. + +use crate::{ + ivf::test::{ + cases::{DIM, METRIC, N_LISTS, grid_index}, + harness::IvfOracleRun, + provider::{Context, Strategy, assign_nearest_for_test}, + }, + neighbor::{BackInserter, Neighbor}, + test::tokio::current_thread_runtime, +}; +use std::{num::NonZeroUsize, sync::Arc}; + +/// After inserting a point, a search probing all lists returns it as the exact nearest +/// neighbor of itself (distance 0). +#[test] +fn inserted_point_is_findable() { + let index = grid_index(); + let strategy = Strategy::new(DIM, METRIC); + let new_point = vec![3.5f32, 3.5]; + let new_id = index.provider().len() as u32; + + current_thread_runtime() + .block_on(index.insert(&strategy, &Context::new(), &new_id, new_point.as_slice())) + .unwrap(); + + assert_eq!(index.provider().len(), (new_id + 1) as usize); + + let run = IvfOracleRun::run_sync(&index, &strategy, &new_point, 1, N_LISTS).unwrap(); + assert_eq!(run.top_k.len(), 1); + let (found_id, dist) = run.top_k[0]; + assert_eq!(found_id, new_id); + assert!(dist.abs() <= f32::EPSILON, "self-distance must be zero"); +} + +/// The inserted point lands in the list whose centroid is nearest to it. +#[test] +fn inserted_point_goes_to_nearest_list() { + let index = grid_index(); + let strategy = Strategy::new(DIM, METRIC); + let new_point = vec![6.0f32, 6.0]; + let new_id = index.provider().len() as u32; + + let expected_list = assign_nearest_for_test(index.provider().centroids(), &new_point, METRIC); + + current_thread_runtime() + .block_on(index.insert(&strategy, &Context::new(), &new_id, new_point.as_slice())) + .unwrap(); + + let members = index.provider().members(expected_list); + assert!( + members.contains(&new_id), + "inserted id {new_id} should be in list {expected_list}, members={members:?}", + ); +} + +/// Inserting then probing only the assigned list still finds the point. +#[test] +fn inserted_point_found_with_single_probe() { + let index = grid_index(); + let strategy = Strategy::new(DIM, METRIC); + let new_point = vec![1.0f32, 1.0]; + let new_id = index.provider().len() as u32; + + current_thread_runtime() + .block_on(index.insert(&strategy, &Context::new(), &new_id, new_point.as_slice())) + .unwrap(); + + let context = Context::new(); + let mut buf = vec![Neighbor::::default(); 1]; + let stats = current_thread_runtime() + .block_on(index.knn_search( + NonZeroUsize::new(1).unwrap(), + 1, + &strategy, + &context, + new_point.as_slice(), + &mut BackInserter::new(buf.as_mut_slice()), + )) + .unwrap(); + + assert_eq!(stats.result_count, 1); + assert_eq!(buf[0].id, new_id); +} + +/// Many `insert` calls racing on a shared `Arc` (multi-threaded runtime) must all +/// take effect: every point gets a unique id, lands in its deterministically-assigned list, +/// and is findable at distance zero afterward. +/// +/// This exercises the insert path's interior mutability under real contention -- the atomic +/// id allocator in `set_element` and the per-list `RwLock::write` in `AppendAccessor::append` +/// (with several points deliberately routed to the same list to force write contention). +#[test] +fn concurrent_inserts_all_land_and_are_findable() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .expect("multi-thread runtime"); + + let index = Arc::new(grid_index()); + let strategy = Arc::new(Strategy::new(DIM, METRIC)); + let base_len = index.provider().len(); + + // Distinct off-grid points spread across all four quadrants. Several share a quadrant so + // their appends contend on the same list's `RwLock`. + let new_points: Vec> = vec![ + vec![0.25, 0.25], // lower-left + vec![1.75, 0.50], // lower-left + vec![0.50, 2.25], // lower-left + vec![1.25, 5.75], // upper-left + vec![0.75, 6.25], // upper-left + vec![6.25, 1.25], // lower-right + vec![5.75, 5.75], // upper-right + vec![6.50, 6.50], // upper-right + vec![5.25, 6.75], // upper-right + vec![6.75, 5.25], // upper-right + ]; + let n = new_points.len(); + + // Deterministic routing: how many of the new points each list should receive. + let mut expected_per_list = [0usize; N_LISTS]; + for point in &new_points { + let list = assign_nearest_for_test(index.provider().centroids(), point, METRIC); + expected_per_list[list as usize] += 1; + } + + runtime.block_on({ + let index = Arc::clone(&index); + let strategy = Arc::clone(&strategy); + let points = new_points.clone(); + async move { + let mut set = tokio::task::JoinSet::new(); + for (i, point) in points.into_iter().enumerate() { + let index = Arc::clone(&index); + let strategy = Arc::clone(&strategy); + set.spawn(async move { + let external_id = (base_len + i) as u32; + index + .insert(&*strategy, &Context::new(), &external_id, point.as_slice()) + .await + .expect("concurrent insert failed"); + }); + } + while let Some(joined) = set.join_next().await { + joined.expect("insert task panicked"); + } + } + }); + + // (1) Every insert allocated a unique, sequential id and committed to some list: the set + // of newly added ids is exactly `base_len .. base_len + n`, with no losses or duplicates. + assert_eq!(index.provider().len(), base_len + n); + let mut new_ids: Vec = (0..N_LISTS as u32) + .flat_map(|list| index.provider().members(list)) + .filter(|id| (*id as usize) >= base_len) + .collect(); + new_ids.sort_unstable(); + let expected_ids: Vec = (base_len as u32..(base_len + n) as u32).collect(); + assert_eq!( + new_ids, expected_ids, + "concurrent inserts must allocate every id exactly once with no loss", + ); + + // (2) Concurrent routing matched the deterministic per-list assignment. + for (list, &expected) in expected_per_list.iter().enumerate() { + let new_in_list = index + .provider() + .members(list as u32) + .into_iter() + .filter(|id| (*id as usize) >= base_len) + .count(); + assert_eq!( + new_in_list, expected, + "list {list} received the wrong number of concurrent inserts", + ); + } + + // (3) Each inserted point is findable at distance zero -- proving its vector bytes were + // written intact (no torn or lost writes under contention). + for point in &new_points { + let run = IvfOracleRun::run_sync(&index, &strategy, point, 1, N_LISTS).unwrap(); + assert_eq!(run.top_k.len(), 1); + assert!( + run.top_k[0].1.abs() <= f32::EPSILON, + "inserted point {point:?} must be findable at distance zero", + ); + } +} diff --git a/diskann/src/ivf/test/cases/mod.rs b/diskann/src/ivf/test/cases/mod.rs new file mode 100644 index 0000000000..3027b4d8ed --- /dev/null +++ b/diskann/src/ivf/test/cases/mod.rs @@ -0,0 +1,76 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Test cases for the IVF module, grouped by concern. +//! +//! * [`correctness`] -- the central invariant (IVF == brute force restricted to probed +//! lists) and its corollaries across `k`/`nprobe`. +//! * [`concurrency`] -- the sequential and concurrent fine-scan paths agree, and the +//! concurrent path actually fans out. +//! * [`insert`] -- inserted points land in the right list and become findable. +//! * [`errors`] -- accessor/strategy failures escalate out of the index API. +//! +//! All cases share the [`fixture`] below: a clean 8x8 integer grid partitioned into four +//! quadrant lists. The geometry is deliberately simple so the expected list assignments are +//! obvious by inspection, but every assertion is still validated against the exact oracle in +//! [`super::harness`] rather than hand-computed answers. + +mod concurrency; +mod correctness; +mod errors; +mod insert; + +use diskann_utils::views::Matrix; +use diskann_vector::distance::Metric; + +use crate::ivf::{IvfIndex, test::provider::Provider}; + +/// Vector/centroid dimension for the shared fixture. +pub(super) const DIM: usize = 2; + +/// Side length of the integer grid used as the dataset (`SIDE * SIDE` points). +pub(super) const SIDE: usize = 8; + +/// Number of inverted lists (quadrant centroids) in the shared fixture. +pub(super) const N_LISTS: usize = 4; + +/// Metric used throughout the IVF tests. +pub(super) const METRIC: Metric = Metric::L2; + +/// The four quadrant centroids of the 8x8 grid, one per inverted list. +fn centroids() -> Matrix { + // Quadrant centers of the [0, 7]^2 grid. + #[rustfmt::skip] + let data = vec![ + 1.5, 1.5, // list 0: lower-left + 1.5, 5.5, // list 1: upper-left + 5.5, 1.5, // list 2: lower-right + 5.5, 5.5, // list 3: upper-right + ]; + Matrix::try_from(data.into_boxed_slice(), N_LISTS, DIM).expect("centroid matrix is well-formed") +} + +/// Build the shared fixture: an [`IvfIndex`] over an 8x8 grid split into four quadrant +/// lists. +pub(super) fn grid_index() -> IvfIndex { + let vectors = crate::graph::test::synthetic::Grid::Two.data(SIDE); + debug_assert_eq!(vectors.ncols(), DIM); + debug_assert_eq!(vectors.nrows(), SIDE * SIDE); + let provider = + Provider::build(vectors, centroids(), METRIC).expect("fixture provider is well-formed"); + IvfIndex::new(provider) +} + +/// A spread of query points: grid corners, the center, and off-grid coordinates. +pub(super) fn queries() -> Vec> { + vec![ + vec![0.0, 0.0], + vec![7.0, 7.0], + vec![3.5, 3.5], + vec![2.0, 6.0], + vec![6.3, 1.1], + vec![4.0, 0.0], + ] +} diff --git a/diskann/src/ivf/test/harness.rs b/diskann/src/ivf/test/harness.rs new file mode 100644 index 0000000000..30ac05b1ce --- /dev/null +++ b/diskann/src/ivf/test/harness.rs @@ -0,0 +1,182 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Reusable execution harness for [`crate::ivf::IvfIndex`] tests. +//! +//! [`IvfOracleRun`] drives [`IvfIndex::knn_search`] and pairs the result with an *exact* +//! oracle. IVF is an approximate algorithm only because it restricts the search to the +//! members of the `nprobe` selected lists -- within that restriction it is exhaustive. The +//! oracle exploits this: it replays the same list selection (via the shared +//! [`nearest_lists`](super::provider::nearest_lists) core) and then brute-forces *only* the +//! union of those lists' members. The IVF result must therefore match the oracle exactly, +//! which lets the tests assert strict equality rather than a recall threshold. + +use std::{cmp::Ordering, num::NonZeroUsize}; + +use diskann_vector::PreprocessedDistanceFunction; + +use crate::{ + ANNResult, + ivf::{ + IvfIndex, SearchStats, + test::provider::{self, Id, ListId, Provider, Strategy}, + }, + neighbor::{BackInserter, Neighbor}, + test::tokio::current_thread_runtime, + utils::VectorRepr, +}; + +/// Result of one [`IvfIndex::knn_search`] run alongside the exact oracle answer. +#[derive(Debug, Clone)] +pub(crate) struct IvfOracleRun { + /// `(id, distance)` pairs returned by the search, re-sorted by `(distance asc, id asc)` + /// so equality checks are deterministic on ties. + pub top_k: Vec<(Id, f32)>, + /// `top_k.iter().map(|(_, d)| d).collect()` -- convenient for distance-multiset checks. + pub top_k_distances: Vec, + /// Statistics returned by `knn_search`. + pub stats: SearchStats, + /// The exact expected top-`k`: brute force restricted to the members of the selected + /// lists, in `(distance asc, id asc)` order. + pub ground_truth: Vec<(Id, f32)>, + /// The lists the oracle (and thus the search) probed for this query. + pub probed_lists: Vec, + /// Total number of members across `probed_lists` -- the exact number of fine-distance + /// computations the search must perform. + pub probed_members: usize, +} + +impl IvfOracleRun { + /// Run [`IvfIndex::knn_search`] once, blocking on a fresh single-threaded runtime. + pub fn run_sync( + index: &IvfIndex, + strategy: &Strategy, + query: &[f32], + k: usize, + nprobe: usize, + ) -> ANNResult { + current_thread_runtime().block_on(Self::run(index, strategy, query, k, nprobe)) + } + + /// Async variant of [`IvfOracleRun::run_sync`], for tests that already own a runtime + /// (e.g. `#[tokio::test]`) or drive searches concurrently across tasks. + pub async fn run( + index: &IvfIndex, + strategy: &Strategy, + query: &[f32], + k: usize, + nprobe: usize, + ) -> ANNResult { + let context = provider::Context::new(); + let mut buf = vec![Neighbor::::default(); k]; + + let stats = index + .knn_search( + NonZeroUsize::new(k).expect("ivf::test::harness requires k > 0"), + nprobe, + strategy, + &context, + query, + &mut BackInserter::new(buf.as_mut_slice()), + ) + .await?; + + let mut top_k: Vec> = buf + .iter() + .copied() + .take(stats.result_count as usize) + .collect(); + sort_neighbors(&mut top_k); + let top_k_distances = top_k.iter().map(|n| n.distance).collect(); + + let (ground_truth, probed_lists, probed_members) = + restricted_topk(index.provider(), query, k, nprobe); + + Ok(Self { + top_k: top_k.into_iter().map(Neighbor::as_tuple).collect(), + top_k_distances, + stats, + ground_truth: ground_truth.into_iter().map(Neighbor::as_tuple).collect(), + probed_lists, + probed_members, + }) + } +} + +/// Exact IVF oracle: select the same `nprobe` lists the search would, then brute-force the +/// union of their members. +/// +/// Returns `(top_k, probed_lists, probed_member_count)`. +pub(crate) fn restricted_topk( + provider: &Provider, + query: &[f32], + k: usize, + nprobe: usize, +) -> (Vec>, Vec, usize) { + let metric = provider.metric(); + let probed_lists = provider::nearest_lists(provider.centroids(), query, metric, nprobe); + + let computer = f32::query_distance(query, metric); + let mut neighbors: Vec> = Vec::new(); + let mut probed_members = 0usize; + for &list in &probed_lists { + for (id, vector) in provider.list_entries(list) { + probed_members += 1; + neighbors.push(Neighbor::new(id, computer.evaluate_similarity(&vector[..]))); + } + } + + sort_neighbors(&mut neighbors); + neighbors.truncate(k); + (neighbors, probed_lists, probed_members) +} + +/// Brute-force top-`k` over the *entire* provider -- used by tests that probe every list +/// (`nprobe == n_lists`) and therefore expect the global answer. +pub(crate) fn global_topk(provider: &Provider, query: &[f32], k: usize) -> Vec> { + let metric = provider.metric(); + let computer = f32::query_distance(query, metric); + let mut neighbors: Vec> = Vec::new(); + for list in 0..provider.n_lists() as ListId { + for (id, vector) in provider.list_entries(list) { + neighbors.push(Neighbor::new(id, computer.evaluate_similarity(&vector[..]))); + } + } + sort_neighbors(&mut neighbors); + neighbors.truncate(k); + neighbors +} + +/// Sort `(distance asc, id asc)` with NaN treated as equal. +fn sort_neighbors(neighbors: &mut [Neighbor]) { + neighbors.sort_by(|a, b| { + a.distance + .partial_cmp(&b.distance) + .unwrap_or(Ordering::Equal) + .then(a.id.cmp(&b.id)) + }); +} + +/// Assert two `(id, distance)` lists agree on the *distance multiset*. +/// +/// The priority queue may break exact-distance ties differently from the oracle's +/// `(distance, id)` ordering, so element-wise id equality is too strict. Comparing the +/// sorted distance sequences (already sorted by the harness) plus the length is the right +/// invariant for a correct top-`k`. +pub(crate) fn assert_same_distances(actual: &[(Id, f32)], expected: &[(Id, f32)]) { + let a: Vec = actual.iter().map(|(_, d)| *d).collect(); + let e: Vec = expected.iter().map(|(_, d)| *d).collect(); + assert_eq!( + a.len(), + e.len(), + "result count mismatch: got {a:?}, expected {e:?}" + ); + for (i, (ad, ed)) in a.iter().zip(e.iter()).enumerate() { + assert!( + (ad - ed).abs() <= f32::EPSILON, + "distance mismatch at rank {i}: got {ad}, expected {ed} (got {a:?}, expected {e:?})" + ); + } +} diff --git a/diskann/src/ivf/test/mod.rs b/diskann/src/ivf/test/mod.rs new file mode 100644 index 0000000000..a19e745212 --- /dev/null +++ b/diskann/src/ivf/test/mod.rs @@ -0,0 +1,10 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Test fixtures and helpers for the IVF module. +pub(crate) mod harness; +pub(crate) mod provider; + +mod cases; diff --git a/diskann/src/ivf/test/provider.rs b/diskann/src/ivf/test/provider.rs new file mode 100644 index 0000000000..983bc5befa --- /dev/null +++ b/diskann/src/ivf/test/provider.rs @@ -0,0 +1,792 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Self-contained, **production-shaped** example backend for the IVF module. +//! +//! This file is intended to double as a worked example of how to implement the three IVF +//! accessor traits ([`ListAccessor`], [`SearchAccessor`], [`InsertAccessor`]) and a +//! [`DataProvider`] for them, with concurrency in mind. +//! +//! # Roles +//! +//! A real deployment may split storage across several collaborators (see +//! `rfcs/00000-ivf-index.md`). For a compact, testable example we fold all of them into one +//! [`Provider`] that plays three roles at once: +//! +//! * **DataProvider** -- owns the external<->internal id mapping and allocates internal ids. +//! * **Coarse quantizer** -- owns the immutable centroids (`list_id -> centroid`). +//! * **Inverted-list store** -- owns the per-list membership. This is a list-owned +//! **IVF-Flat** layout: each list stores its members' full vectors *contiguously* in a +//! single flat buffer, in the same order as the list's ids (see [`ListStore`]). The fine +//! scan therefore streams a list's vectors sequentially out of one cache-friendly block +//! rather than chasing per-id pointers into a global table. +//! +//! Because the vectors live with their list, there is no global `id -> vector` table: a +//! vector is materialized only when [`InsertAccessor::append`] commits it to a list. The +//! provider's [`provider::SetElement::set_element`] therefore only *allocates* an internal +//! id; the append step does the contiguous store. +//! +//! # Concurrency model +//! +//! The centroids are immutable and shared via [`Arc`], so the fine-scan accessor can hand +//! list scans to independent tasks (see [`ScanAccessor`]'s concurrent path). The mutable +//! state -- each list's contiguous store -- is guarded by a per-list [`RwLock`] so inserts +//! can append to one list while searches read others. This mirrors how a production backend +//! would keep the read (search) path lock-light while still supporting concurrent append. + +use std::{ + cmp::Ordering, + fmt::{self, Debug}, + future::Future, + sync::{ + Arc, RwLock, + atomic::{AtomicU32, Ordering as AtomicOrdering}, + }, +}; + +use diskann_utils::{future::SendFuture, views::Matrix}; +use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; +use thiserror::Error; + +use crate::graph::SearchOutputBuffer; +use crate::{ + ANNError, always_escalate, + internal::counter::{Counter, LocalCounter}, + ivf::{InsertAccessor, InsertStrategy, ListAccessor, SearchAccessor, SearchStrategy}, + provider::{self, ExecutionContext, HasId, NoopGuard}, + utils::VectorRepr, +}; + +/// The id type used for inverted lists throughout this example. +pub type ListId = u32; + +/// The internal/external id type used by this provider. +pub type Id = u32; + +////////////////////////// +// Shared selection core // +////////////////////////// + +/// Return the (up to) `nprobe` lists closest to `query` as `(list_id, distance)` pairs, +/// ordered by `(distance asc, list_id asc)`. +/// +/// This is the single source of truth for coarse list selection. Both [`CentroidAccessor`] +/// (the production path, via this function) and the test harness's oracle (via +/// [`nearest_lists`]) build on it, which guarantees they agree on exactly which lists a +/// query probes -- essential for an *exact* approximate-search oracle. +fn scored_lists( + centroids: &Matrix, + query: &[f32], + metric: Metric, + nprobe: usize, +) -> Vec<(ListId, f32)> { + let computer = f32::query_distance(query, metric); + let mut scored: Vec<(ListId, f32)> = centroids + .row_iter() + .enumerate() + .map(|(i, c)| (i as ListId, computer.evaluate_similarity(c))) + .collect(); + scored.sort_by(|a, b| cmp_dist_id(*a, *b)); + scored.truncate(nprobe); + scored +} + +/// Return the (up to) `nprobe` list ids closest to `query`, ordered by +/// `(distance asc, list_id asc)`. Used by the test harness oracle. +pub fn nearest_lists( + centroids: &Matrix, + query: &[f32], + metric: Metric, + nprobe: usize, +) -> Vec { + scored_lists(centroids, query, metric, nprobe) + .into_iter() + .map(|(id, _)| id) + .collect() +} + +/// Order by `(distance asc, id asc)` with NaN treated as equal. +fn cmp_dist_id(a: (ListId, f32), b: (ListId, f32)) -> Ordering { + a.1.partial_cmp(&b.1) + .unwrap_or(Ordering::Equal) + .then(a.0.cmp(&b.0)) +} + +/// Index of the single nearest list for `vector`. +fn assign_nearest(centroids: &Matrix, vector: &[f32], metric: Metric) -> ListId { + let computer = f32::query_distance(vector, metric); + centroids + .row_iter() + .enumerate() + .map(|(i, c)| (i as ListId, computer.evaluate_similarity(c))) + .min_by(|a, b| cmp_dist_id(*a, *b)) + .map(|(id, _)| id) + .expect("provider always has at least one centroid") +} + +/// Public re-export of [`assign_nearest`] for tests that want to predict an insert's target +/// list independently of the index. +pub fn assign_nearest_for_test(centroids: &Matrix, vector: &[f32], metric: Metric) -> ListId { + assign_nearest(centroids, vector, metric) +} + +////////////// +// Provider // +////////////// + +/// Error conditions for [`Provider::build`]. +#[derive(Debug, Error)] +pub enum ProviderError { + #[error("ivf::test::Provider needs at least one vector")] + EmptyVectors, + #[error("ivf::test::Provider needs at least one centroid")] + EmptyCentroids, + #[error("ivf::test::Provider vectors must have non-zero dimension")] + ZeroDimension, + #[error("ivf::test::Provider dimension mismatch: vectors={vectors}, centroids={centroids}")] + DimMismatch { vectors: usize, centroids: usize }, +} + +impl From for ANNError { + #[track_caller] + fn from(err: ProviderError) -> ANNError { + ANNError::opaque(err) + } +} + +/// One inverted list's contiguous storage. +/// +/// `ids[j]` is the internal id of the `j`-th member, and its full vector occupies +/// `data[j * dim .. (j + 1) * dim]`. The two vectors grow together in [`Self::push`], so a +/// member's id and its vector always share the same index -- the property the fine scan +/// relies on to pair a streamed vector back with its id. +#[derive(Debug, Default)] +struct ListStore { + ids: Vec, + data: Vec, +} + +impl ListStore { + /// Append `(id, vector)`, keeping ids and the contiguous `data` buffer in lockstep. + fn push(&mut self, id: Id, vector: &[f32]) { + self.ids.push(id); + self.data.extend_from_slice(vector); + } + + /// Iterate `(id, vector)` pairs in storage order, slicing the flat buffer by `dim`. + fn iter(&self, dim: usize) -> impl Iterator { + self.ids.iter().copied().zip(self.data.chunks_exact(dim)) + } +} + +/// In-memory IVF backend. +#[derive(Debug)] +pub struct Provider { + /// Immutable centroids, one row per inverted list. + centroids: Arc>, + /// Per-list contiguous storage (ids + full vectors). Fixed list count; each list grows + /// on insert. + lists: Arc>>, + /// Allocates the next internal id. Identity external<->internal ids means this is just a + /// monotonically increasing slot counter. + next_id: Arc, + /// Vector dimension. + dim: usize, + /// Distance metric used for both coarse selection and fine scoring (RFC D5). + metric: Metric, + /// Number of query->centroid distance computations performed. + centroid_cmps: Counter, + /// Number of query->member distance computations performed. + member_cmps: Counter, +} + +impl Provider { + /// Build a provider from `vectors` and `centroids`, assigning each vector to its + /// nearest centroid's list (the "bring-your-own-centroids" build of RFC F1/D6). + /// + /// # Errors + /// + /// Returns [`ProviderError`] if either matrix is empty, the dimension is zero, or the + /// vector and centroid dimensions disagree. + pub fn build( + vectors: Matrix, + centroids: Matrix, + metric: Metric, + ) -> Result { + if vectors.nrows() == 0 { + return Err(ProviderError::EmptyVectors); + } + if centroids.nrows() == 0 { + return Err(ProviderError::EmptyCentroids); + } + let dim = vectors.ncols(); + if dim == 0 { + return Err(ProviderError::ZeroDimension); + } + if centroids.ncols() != dim { + return Err(ProviderError::DimMismatch { + vectors: dim, + centroids: centroids.ncols(), + }); + } + + let n_lists = centroids.nrows(); + let lists: Vec> = (0..n_lists) + .map(|_| RwLock::new(ListStore::default())) + .collect(); + for (i, vector) in vectors.row_iter().enumerate() { + let list = assign_nearest(¢roids, vector, metric); + lists[list as usize] + .write() + .expect("list lock poisoned during build") + .push(i as Id, vector); + } + + let next_id = vectors.nrows() as u32; + + Ok(Self { + centroids: Arc::new(centroids), + lists: Arc::new(lists), + next_id: Arc::new(AtomicU32::new(next_id)), + dim, + metric, + centroid_cmps: Counter::new(), + member_cmps: Counter::new(), + }) + } + + /// Vector dimension. + pub fn dim(&self) -> usize { + self.dim + } + + /// Distance metric. + pub fn metric(&self) -> Metric { + self.metric + } + + /// Number of inverted lists (== number of centroids). + pub fn n_lists(&self) -> usize { + self.centroids.nrows() + } + + /// Number of vectors currently stored (== number of allocated internal ids). + pub fn len(&self) -> usize { + self.next_id.load(AtomicOrdering::Relaxed) as usize + } + + /// Borrow the centroids (used by the harness oracle). + pub fn centroids(&self) -> &Matrix { + &self.centroids + } + + /// Snapshot the ids in `list`, in storage order (used by the harness oracle and tests). + pub fn members(&self, list: ListId) -> Vec { + self.lists[list as usize] + .read() + .expect("list lock poisoned") + .ids + .clone() + } + + /// Snapshot the `(id, vector)` entries in `list`, in storage order (used by the harness + /// oracle). The vectors are copied out of the list's contiguous buffer. + pub fn list_entries(&self, list: ListId) -> Vec<(Id, Vec)> { + let store = self.lists[list as usize] + .read() + .expect("list lock poisoned"); + store + .iter(self.dim) + .map(|(id, v)| (id, v.to_vec())) + .collect() + } + + /// Snapshot of the per-provider distance-computation counters. + pub fn metrics(&self) -> ProviderMetrics { + ProviderMetrics { + centroid_cmps: self.centroid_cmps.value(), + member_cmps: self.member_cmps.value(), + } + } +} + +/// Distance-computation counters tracked by [`Provider`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderMetrics { + /// Total query->centroid distance computations (coarse step). + pub centroid_cmps: usize, + /// Total query->member distance computations (fine step). + pub member_cmps: usize, +} + +///////////// +// Context // +///////////// + +/// Per-operation execution context that records how many tasks were spawned through it, +/// so concurrency tests can assert the fine-scan path actually fanned out. +#[derive(Debug, Clone)] +pub struct Context(Arc); + +impl Context { + pub fn new() -> Self { + Self(Arc::new(Counter::new())) + } + + /// Number of spawns routed through [`ExecutionContext::wrap_spawn`]. + pub fn spawns(&self) -> usize { + self.0.value() + } +} + +impl Default for Context { + fn default() -> Self { + Self::new() + } +} + +impl ExecutionContext for Context { + fn wrap_spawn(&self, f: F) -> impl Future + Send + 'static + where + F: Future + Send + 'static, + { + self.0.increment(); + f + } +} + +//////////// +// Errors // +//////////// + +/// Critical id-validation error: a requested id is out of range. +#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] +#[error("ivf::test::Provider has no id {0}")] +pub struct InvalidId(pub u32); + +always_escalate!(InvalidId); + +impl From for ANNError { + #[track_caller] + fn from(err: InvalidId) -> ANNError { + ANNError::opaque(err) + } +} + +/// Critical error injected by [`Strategy::failing_on_list`]: scanning a specific list fails. +/// +/// Used to verify that a fine-scan failure escalates out of +/// [`crate::ivf::IvfIndex::knn_search`]. +#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] +#[error("ivf::test::Provider injected scan failure on list {0}")] +pub struct ScanFailure(pub ListId); + +always_escalate!(ScanFailure); + +impl From for ANNError { + #[track_caller] + fn from(err: ScanFailure) -> ANNError { + ANNError::opaque(err) + } +} + +/// Dimension-mismatch error from strategy construction. +#[derive(Debug, Clone, Error, PartialEq, Eq)] +#[error("dimension mismatch: strategy expects {expected}, got {actual}")] +pub struct StrategyError { + pub expected: usize, + pub actual: usize, +} + +impl From for ANNError { + #[track_caller] + fn from(err: StrategyError) -> ANNError { + ANNError::opaque(err) + } +} + +////////////////// +// DataProvider // +////////////////// + +impl provider::DataProvider for Provider { + type Context = Context; + type InternalId = Id; + type ExternalId = Id; + type Error = InvalidId; + type Guard = NoopGuard; + + fn to_internal_id(&self, _ctx: &Context, gid: &Id) -> Result { + if (*gid as usize) < self.len() { + Ok(*gid) + } else { + Err(InvalidId(*gid)) + } + } + + fn to_external_id(&self, _ctx: &Context, id: Id) -> Result { + if (id as usize) < self.len() { + Ok(id) + } else { + Err(InvalidId(id)) + } + } +} + +impl provider::SetElement<&[f32]> for Provider { + type SetError = InvalidId; + + async fn set_element( + &self, + _context: &Context, + _id: &Id, + _element: &[f32], + ) -> Result { + // List-owned storage: the vector is *not* stored here -- it is committed to its + // list's contiguous buffer by `InsertAccessor::append`, once the target list is + // known. `set_element` therefore only allocates the next internal id. A production + // provider would honor the supplied external id and return a guard that rolls back + // the allocation on drop-without-complete. + let new_id = self.next_id.fetch_add(1, AtomicOrdering::Relaxed); + Ok(NoopGuard::new(new_id)) + } +} + +////////////////////////////// +// ListAccessor (coarse step) // +////////////////////////////// + +/// Coarse accessor: selects the `nprobe` lists nearest to a bound query. +/// +/// Holds an [`Arc`] clone of the immutable centroids and a preprocessed query computer. The +/// `LocalCounter` flushes the coarse comparison count back to the provider on drop. +pub struct CentroidAccessor<'a> { + centroids: Arc>, + metric: Metric, + nprobe_query: Vec, + cmps: LocalCounter<'a>, +} + +impl Debug for CentroidAccessor<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CentroidAccessor") + .field("n_lists", &self.centroids.nrows()) + .finish_non_exhaustive() + } +} + +impl ListAccessor for CentroidAccessor<'_> { + type Id = ListId; + type Error = InvalidId; + + fn select_lists( + &mut self, + nprobe: usize, + output: &mut B, + ) -> impl SendFuture> + where + B: SearchOutputBuffer + Send + ?Sized, + { + async move { + self.cmps.increment_by(self.centroids.nrows()); + for (id, dist) in scored_lists(&self.centroids, &self.nprobe_query, self.metric, nprobe) + { + if output.push(id, dist).is_full() { + break; + } + } + Ok(()) + } + } +} + +///////////////////////////////// +// SearchAccessor (fine step) // +///////////////////////////////// + +/// Whether the fine scan runs sequentially or fans lists out across spawned tasks. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ScanMode { + /// Scan lists one at a time on the current task. + Sequential, + /// Spawn one task per list via [`ExecutionContext::wrap_spawn`] and merge results. + Concurrent, +} + +/// Fine accessor: scans the members of a set of lists, scoring each against a bound query. +/// +/// All data needed for a scan is held behind [`Arc`] so the concurrent path can move clones +/// into `'static` spawned tasks. Each spawned task produces its own `(id, distance)` vector; +/// the driver merges those into the single `f` callback, preserving the trait's +/// single-threaded-callback contract while doing the distance work concurrently. +pub struct ScanAccessor<'a> { + lists: Arc>>, + computer: Arc<::QueryDistance>, + dim: usize, + context: Context, + mode: ScanMode, + fail_on_list: Option, + cmps: LocalCounter<'a>, +} + +impl Debug for ScanAccessor<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ScanAccessor") + .field("mode", &self.mode) + .field("fail_on_list", &self.fail_on_list) + .finish_non_exhaustive() + } +} + +impl HasId for ScanAccessor<'_> { + type Id = Id; +} + +impl ScanAccessor<'_> { + /// Score every member of `list`, returning `(id, distance)` pairs. Shared by both the + /// sequential driver and the spawned-task driver. + /// + /// Streams the list's contiguous vector buffer in storage order, pairing each vector + /// with the id at the matching index. + fn score_list( + lists: &Arc>>, + computer: &::QueryDistance, + dim: usize, + list: ListId, + ) -> Vec<(Id, f32)> { + let store = lists[list as usize].read().expect("list lock poisoned"); + store + .iter(dim) + .map(|(id, v)| (id, computer.evaluate_similarity(v))) + .collect() + } +} + +impl SearchAccessor for ScanAccessor<'_> { + type ListId = ListId; + type Error = ScanFailure; + + fn scan_lists( + &mut self, + lists: Itr, + mut f: F, + ) -> impl SendFuture> + where + Itr: Iterator + Send, + F: Send + FnMut(Self::Id, f32), + { + async move { + match self.mode { + ScanMode::Sequential => { + for list in lists { + if self.fail_on_list == Some(list) { + return Err(ScanFailure(list)); + } + for (id, dist) in + Self::score_list(&self.lists, &self.computer, self.dim, list) + { + self.cmps.increment(); + f(id, dist); + } + } + } + ScanMode::Concurrent => { + let mut set = tokio::task::JoinSet::new(); + for list in lists { + if self.fail_on_list == Some(list) { + set.shutdown().await; + return Err(ScanFailure(list)); + } + let lists = Arc::clone(&self.lists); + let computer = Arc::clone(&self.computer); + let dim = self.dim; + let task = async move { Self::score_list(&lists, &computer, dim, list) }; + set.spawn(self.context.wrap_spawn(task)); + } + while let Some(joined) = set.join_next().await { + let scored = joined.expect("scan task panicked"); + for (id, dist) in scored { + self.cmps.increment(); + f(id, dist); + } + } + } + } + Ok(()) + } + } +} + +///////////////////////////////// +// InsertAccessor (append step) // +///////////////////////////////// + +/// Append accessor: commits `(id, vector)` to a chosen list's contiguous store. +/// +/// This example is **IVF-Flat** with *list-owned* storage: the full vector is appended to +/// the target list's contiguous buffer here (in lockstep with its id), and read back from +/// that same buffer during the fine scan. There is no global `id -> vector` table. +/// +/// An **IVF-PQ** backend would slot in here unchanged in shape: quantize `vector` and push +/// the code instead of the raw floats, e.g. `store.push_code(id, codebook.encode(vector))`. +#[derive(Debug)] +pub struct AppendAccessor { + lists: Arc>>, + dim: usize, +} + +impl HasId for AppendAccessor { + type Id = Id; +} + +impl<'v> InsertAccessor<&'v [f32]> for AppendAccessor { + type ListId = ListId; + type Error = InvalidId; + + fn append( + &mut self, + list: Self::ListId, + id: Self::Id, + vector: &'v [f32], + ) -> impl SendFuture> { + async move { + debug_assert_eq!( + vector.len(), + self.dim, + "append received a vector of the wrong dimension", + ); + self.lists[list as usize] + .write() + .expect("list lock poisoned") + .push(id, vector); + Ok(()) + } + } +} + +////////////// +// Strategy // +////////////// + +/// Per-call factory wiring centroids (coarse) to scanning/appending (fine) for `&[f32]` +/// queries and vectors. Validates query dimension and can inject a scan failure. +#[derive(Debug, Clone)] +pub struct Strategy { + dim: usize, + metric: Metric, + mode: ScanMode, + fail_on_list: Option, +} + +impl Strategy { + /// Construct a strategy expecting vectors of dimension `dim`, scanning sequentially. + pub fn new(dim: usize, metric: Metric) -> Self { + Self { + dim, + metric, + mode: ScanMode::Sequential, + fail_on_list: None, + } + } + + /// Select the fine-scan mode (sequential vs. concurrent). + pub fn with_mode(mut self, mode: ScanMode) -> Self { + self.mode = mode; + self + } + + /// Make the fine scan fail when it reaches `list`. + pub fn failing_on_list(mut self, list: ListId) -> Self { + self.fail_on_list = Some(list); + self + } + + fn check_dim(&self, actual: usize) -> Result<(), StrategyError> { + if actual != self.dim { + Err(StrategyError { + expected: self.dim, + actual, + }) + } else { + Ok(()) + } + } +} + +impl<'a> SearchStrategy<'a, Provider, &'a [f32]> for Strategy { + type ListId = ListId; + type SearchAccessor = ScanAccessor<'a>; + type ListAccessor = CentroidAccessor<'a>; + type Error = StrategyError; + + fn search_accessor( + &'a self, + provider: &'a Provider, + context: &'a Context, + query: &'a [f32], + ) -> Result { + self.check_dim(query.len())?; + Ok(ScanAccessor { + lists: Arc::clone(&provider.lists), + computer: Arc::new(f32::query_distance(query, self.metric)), + dim: provider.dim, + context: context.clone(), + mode: self.mode, + fail_on_list: self.fail_on_list, + cmps: provider.member_cmps.local(), + }) + } + + fn list_accessor( + &'a self, + provider: &'a Provider, + _context: &'a Context, + query: &'a [f32], + ) -> Result { + self.check_dim(query.len())?; + Ok(CentroidAccessor { + centroids: Arc::clone(&provider.centroids), + metric: self.metric, + nprobe_query: query.to_vec(), + cmps: provider.centroid_cmps.local(), + }) + } +} + +impl<'a> InsertStrategy<'a, Provider, &'a [f32]> for Strategy { + type ListId = ListId; + type InsertAccessor = AppendAccessor; + type ListAccessor = CentroidAccessor<'a>; + type Error = StrategyError; + + fn insert_accessor( + &'a self, + provider: &'a Provider, + _context: &'a Context, + ) -> Result { + // The vector is validated by `list_accessor` (the coarse step runs first), and is + // stored contiguously by `append`, so the append accessor only needs the list + // storage handle and the dimension (for its lockstep buffer layout). + Ok(AppendAccessor { + lists: Arc::clone(&provider.lists), + dim: provider.dim, + }) + } + + fn list_accessor( + &'a self, + provider: &'a Provider, + _context: &'a Context, + vector: &'a [f32], + ) -> Result { + self.check_dim(vector.len())?; + Ok(CentroidAccessor { + centroids: Arc::clone(&provider.centroids), + metric: self.metric, + nprobe_query: vector.to_vec(), + cmps: provider.centroid_cmps.local(), + }) + } +} diff --git a/diskann/src/lib.rs b/diskann/src/lib.rs index 9c1f6ac76b..14c20bfdf7 100644 --- a/diskann/src/lib.rs +++ b/diskann/src/lib.rs @@ -15,6 +15,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, ANNErrorKind, ANNResult}; From b70b235bc69ea5a0b000f7bb19bb096cd88f33d7 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Thu, 18 Jun 2026 11:50:46 -0400 Subject: [PATCH 03/33] delete rfc --- rfcs/00000-ivf-index.md | 148 ---------------------------------------- 1 file changed, 148 deletions(-) delete mode 100644 rfcs/00000-ivf-index.md diff --git a/rfcs/00000-ivf-index.md b/rfcs/00000-ivf-index.md deleted file mode 100644 index b0f331f23a..0000000000 --- a/rfcs/00000-ivf-index.md +++ /dev/null @@ -1,148 +0,0 @@ -# Inverted File (IVF) Index - -| | | -|------------------|--------------------------------| -| **Authors** | Aditya Krishnan | -| **Created** | 2026-06-10 | -| **Updated** | 2026-06-10 | -| **Status** | Draft — requirements gathering | - -> **Note**: This is a *requirements* document. It deliberately stays at the -> level of "what we need and why" and defers concrete type/API design to a -> follow-up RFC. Sections are intentionally terse so we can iterate quickly. - -## 1. Summary - -Add an **Inverted File (IVF)** index to DiskANN. An IVF index partitions the -dataset into a fixed number of *inverted lists* (a.k.a. cells / clusters), each -summarized by a **centroid**. Each list stores the representations (full -precision or quantized) of the points assigned to it, laid out contiguously. - -- **Insert**: assign a point to the list whose centroid is closest, then append - the point's representation to that list. -- **Search**: select the `nprobe` lists whose centroids are closest to the - query, score the query against every point in those lists, and return the - top-`k` via the neighbor priority queue. - -## 2. Motivation - -### 2.1 Background - -DiskANN today exposes two index families: - -- **Graph search** — random-access greedy search over a proximity graph - ([`crate::graph::DiskANNIndex`], driven by the [`crate::provider::Accessor`] - trait). -- **Flat search** — sequential one-pass scan-and-score - ([`crate::flat::FlatIndex`], driven by [`crate::flat::DistancesUnordered`]). - -There is currently **no clustering / coarse-quantizer based index**. IVF is the -canonical such structure and a well-understood baseline in the ANN literature -(IVF, IVFPQ, etc.). - -### 2.2 Problem Statement - -We want a partition-based index that: - -- Trades recall for latency via a single tunable knob (`nprobe`). -- Has predictable, scan-friendly memory layout (contiguous lists) that maps - cleanly onto the existing flat-search scan-and-score primitives. -- Composes with the existing quantization stack so lists can store quantized - representations. - -### 2.3 Goals - -1. Define an IVF index that supports **build**, **insert**, and **top-k - search** with an `nprobe` parameter. -2. **Reuse existing primitives** rather than inventing parallel machinery: - - [`crate::provider::DataProvider`] for id mapping / context. - - The flat-search scan-and-score path ([`DistancesUnordered`]) for scoring - points inside a list. - - [`crate::neighbor`] priority queue for top-k accumulation. - - The quantization crate for list representations. -3. Support both **full-precision** and **quantized** list representations behind - a common representation abstraction. -4. Make the **coarse quantizer** (centroid set + assignment) a pluggable - component so we can later swap k-means for alternatives. -5. Make the representation and access of inverted lists to be pluggable - so it works with multiple data-backends (disk, k-v table, caching, blob). - -### 2.4 Non-Goals (for the first cut) - -- IVFPQ residual encoding / per-list re-ranking (future work). -- Deletes and compaction. - -## 3. Requirements - -### 3.1 Functional - -| ID | Requirement | -|----|-------------| -| F1 | **Build** an index from a dataset: derive `n_lists` centroids (initially k-means) and assign every point to its nearest centroid's list. | -| F2 | **Insert** a single point: find nearest inverted list (start by using centroids), append representation to that list. Insert must not require a full rebuild. | -| F3 | **Search**: given a query and parameters `k` and `nprobe`, select the `nprobe` nearest lists by centroid distance, scan-and-score their members, return top-`k`. List scoring should reuse the flat-search scan-and-score primitive [`DistancesUnordered`]; no bespoke distance loop. | | -| F4 | Lists can store any representation of vectors (e.g. **full-precision**, **quantized**), selected at build time. | -| F5 | The **coarse quantizer** (centroid training + assignment) is a distinct, swappable component. | -| F6 | The representation of inverted lists should be swappable, consumers should be able to implement it for different backends - disk, caching, inmemory, blob etc. | - -### 3.2 Non-Functional - -| ID | Requirement | -|----|-------------| -| N2 | Per-list storage is **contiguous** to enable sequential, SIMD-friendly scans. | -| N3 | `nprobe` is a per-query parameter (not fixed at build time). | -| N4 | Errors follow the mid-level [`ANNError`] regime (this is a `diskann`-crate algorithm). | -| N5 | Build and search are parallelizable, but must use the workspace thread-pool conventions (no global rayon pool — see `clippy.toml`). | -| N6 | Index parameters (`n_lists`, representation kind, distance function) are recorded so an index can be reloaded consistently. | - -### 3.3 Open Questions - -- **Crate placement**: new module under `diskann/` (alongside `flat/` and - `graph/`) vs. a dedicated `diskann-ivf` crate? Leaning toward a module first. - - **Answers**: - - Let's add it as a modulke to `diskann` for now. -- **Coarse quantizer ownership**: does the centroid set live inside the IVF - index, or is it a reusable component shared with quantization? - **Answers**: - - It can live within the index for now. We can refactor this later if we want it to be a reusable component. -- **Centroid representation**: full-precision centroids always, or allow - quantized centroids for the coarse search too? - **Answers**: - - Doesn't matter. The component in the index that returns the top inverted lists should be generic enough that we can implement - multiple algorithms/representations to perform this operation. Including full precision centroids, quantized centroids etc. -- **List growth**: per-list `Vec` vs. a single backing arena with list offsets. - Affects insert cost and scan locality. - **Answers**: - - We need to come up with a clean API for inserting a point to an inverted list. Specific implementations will implement this API - depending on their specific data backend. -- **Relationship to `DataProvider`**: is the IVF index a *consumer* of a - provider, or does it own its own storage backend? - **Answers**: - - Critically, the IVF index is a consumer of a provider. It does not own its backend. -- **Serialization**: reuse an existing on-disk format, or define a new one? - **Answers**: - - Let's not worry about serialization for now. - -## 4. Sketch (non-binding) - -A rough mental model to anchor discussion — *not* a committed design: - -```text -IvfIndex -├── coarse: CoarseQuantizer // n_lists centroids + assign(query) -> list ids -├── lists: [InvertedList; n_lists] // each: contiguous representations + local->external id map -└── search(query, k, nprobe): - candidate_lists = coarse.closest(query, nprobe) - for list in candidate_lists: - list.scan_and_score(query, &mut neighbor_queue) // via DistancesUnordered - return neighbor_queue.top_k(k) -``` - -## 5. Future Work - -- [ ] IVFPQ: store PQ residual codes per list; re-rank with full precision. -- [ ] Disk-resident lists / out-of-core search. -- [ ] Deletes, tombstoning, and list compaction. -- [ ] Advanced probing (soft assignment, learned routing). -- [ ] Label / attribute filtering integration (`diskann-label-filter`). From 4a18959757c1990353ad1e22e70af0f118444d56 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Mon, 14 Sep 2026 17:04:56 -0400 Subject: [PATCH 04/33] dynamic index glue/trait initial brainstorm --- diskann/src/ivf/dynamic.rs | 316 +++++++++++++ diskann/src/ivf/mod.rs | 6 +- rfcs/01187-incremental-ivf.md | 820 ++++++++++++++++++++++++++++++++++ 3 files changed, 1140 insertions(+), 2 deletions(-) create mode 100644 diskann/src/ivf/dynamic.rs create mode 100644 rfcs/01187-incremental-ivf.md diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs new file mode 100644 index 0000000000..f01f3d947d --- /dev/null +++ b/diskann/src/ivf/dynamic.rs @@ -0,0 +1,316 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Contracts for an incrementally maintained IVF index. +//! +//! This module is an API skeleton. It defines the boundary between an IVF +//! algorithm and provider-specific centroid navigation, inverted-list access, +//! and partition mutation. Concrete accessors are responsible for presenting a +//! coherent view across those components. The online split/dissolve algorithm +//! is not implemented here yet. + +use std::fmt::Debug; + +use diskann_utils::future::SendFuture; + +use crate::{ + error::{StandardError, ToRanked}, + provider::{DataProvider, HasId}, + utils::VectorId, +}; + +/// One centroid/list selected during coarse search. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct SelectedList { + /// Stable logical list identifier. + pub id: ListId, + /// Distance from the bound query to this list's centroid. + pub distance: f32, +} + +/// The handoff from centroid selection to inverted-list scanning. +/// +/// A plan must be consumed by the same [`SearchAccessor`] that produced it. +/// The accessor is responsible for keeping centroid selection and list reads +/// coherent using provider-specific coordination. +#[derive(Debug, Clone)] +pub struct SelectionPlan { + selected: Vec>, +} + +impl SelectionPlan { + /// Construct a plan from selected logical lists. + pub fn new(selected: Vec>) -> Self { + Self { selected } + } + + /// Lists selected in increasing coarse-distance order. + pub fn selected(&self) -> &[SelectedList] { + &self.selected + } + + /// Consume the plan and return the selected lists. + pub fn into_selected(self) -> Vec> { + self.selected + } +} + +/// Work performed while scanning the selected lists. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ScanStats { + /// Number of point representations scored against the query. + pub comparisons: u64, + /// Number of distinct lists read. + pub lists_scanned: u32, +} + +/// An in-memory index over the live centroids. +/// +/// The centroid catalog is authoritative. Approximate implementations, such as +/// a graph navigator, must retain enough catalog information to reject retired +/// ids and recover with exact selection when navigation cannot produce a usable +/// result. Exact and graph implementations expose the same interface. +pub trait CentroidIndex: Send + Sync { + /// Stable logical centroid id, also used as the inverted-list id. + type ListId: VectorId; + + /// Errors encountered while reading or navigating the centroid index. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Number of live centroids. + fn len(&self) -> usize; + + /// Whether there are no live centroids. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Borrow one authoritative full-precision centroid vector. + fn centroid(&self, id: Self::ListId) -> Option<&[f32]>; + + /// Select exactly `min(nprobe, self.len())` live centroids for `query`. + /// + /// A graph implementation must use exact selection as a recovery path when + /// graph navigation produces too few live ids. + fn select( + &self, + query: &[f32], + nprobe: usize, + ) -> impl SendFuture, Self::Error>>; +} + +/// Query-bound access to centroid selection and inverted-list scanning. +/// +/// This is the dynamic IVF search algorithm's primary extension point, in the +/// same spirit as [`crate::graph::SearchAccessor`]. Implementations are free to +/// batch reads, coalesce blob requests, prefetch, decode quantized payloads, or +/// fan work out across tasks. +pub trait SearchAccessor: HasId + Send + Sync { + /// Stable logical centroid/list id. + type ListId: VectorId; + + /// Errors from list selection or scanning. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Select exactly the available requested number of lists for the bound query. + fn select_lists( + &mut self, + nprobe: usize, + ) -> impl SendFuture, Self::Error>>; + + /// Scan every list in `plan`, scoring members against the bound query. + fn scan_lists( + &mut self, + plan: SelectionPlan, + emit: F, + ) -> impl SendFuture> + where + F: FnMut(Self::Id, f32) + Send; +} + +/// Factory for one dynamic IVF search accessor. +pub trait SearchStrategy<'a, Provider, T>: Send + Sync +where + Provider: DataProvider, +{ + /// Query-bound accessor used for both coarse and fine search. + type SearchAccessor: SearchAccessor; + + /// Error constructing the accessor. + type Error: StandardError; + + /// Construct an accessor around `query`. + fn search_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + query: T, + ) -> Result; +} + +/// Size metadata for one logical inverted list. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ListMetadata { + /// Stable logical list identifier. + pub id: ListId, + /// Number of live point ids in the list. + pub len: usize, +} + +/// A new centroid to install as part of a partition update. +#[derive(Debug, Clone, PartialEq)] +pub struct CentroidRecord { + /// Fresh stable id. Installed ids must never be reused after retirement. + pub id: ListId, + /// Full-precision centroid used by maintenance and exact fallback. + pub vector: Box<[f32]>, +} + +/// Changes to the authoritative live-centroid catalog. +#[derive(Debug, Clone, PartialEq)] +pub struct CentroidDelta { + /// Centroids made live by this update. + pub insert: Vec>, + /// Live centroids retired by this update. + pub retire: Vec, +} + +/// One point membership change in a partition mutation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PointMove { + /// Internal point id being inserted, deleted, or reassigned. + pub id: Id, + /// Previous list, or `None` for a newly inserted point. + pub from: Option, + /// New list, or `None` for a deleted point. + pub to: Option, +} + +/// A complete logical partition change. +/// +/// The maintenance accessor applies this update across point identity, +/// canonical vectors, list membership, reverse assignments, scan payloads, and +/// centroid liveness according to its provider-specific consistency contract. +#[derive(Debug, Clone, PartialEq)] +pub struct PartitionUpdate { + /// Centroids inserted and retired by this mutation. + pub centroids: CentroidDelta, + /// Point membership changes, including inserts and deletes. + pub point_moves: Vec>, +} + +/// Operation-scoped reads and writes needed by split/dissolve maintenance. +/// +/// The accessor is the consistency boundary for one mutation. It must present a +/// unified view across point data, centroids, reverse assignments, and inverted +/// lists for all planning reads. It also owns any provider-specific lock, epoch, +/// transaction, staging area, or poison state needed to apply the final update. +/// +/// Read methods are deliberately coarse grained so disk and blob providers can +/// batch I/O. Callbacks may be invoked in any order. `read_members` may invoke +/// its callback more than once for a list when the backend streams chunks. +pub trait MaintenanceAccessor: HasId + Send + Sized +where + T: Send, +{ + /// External point id accepted by the data provider. + type ExternalId: PartialEq + Send + Sync + 'static; + + /// Stable centroid/list id. + type ListId: VectorId; + + /// In-memory centroid catalog and navigator in this accessor's unified view. + type Centroids: CentroidIndex; + + /// Errors from planning reads, staging, or applying the update. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Borrow the centroid index used for routing and maintenance neighborhoods. + fn centroids(&self) -> &Self::Centroids; + + /// Read list sizes for `lists`. + fn list_metadata( + &mut self, + lists: I, + emit: F, + ) -> impl SendFuture> + where + I: Iterator + Send, + F: FnMut(ListMetadata) + Send; + + /// Resolve current list assignments for point ids. + fn assignments(&mut self, ids: I, emit: F) -> impl SendFuture> + where + I: Iterator + Send, + F: FnMut(Self::Id, Option) + Send; + + /// Stream the live member ids of selected lists. + fn read_members(&mut self, lists: I, emit: F) -> impl SendFuture> + where + I: Iterator + Send, + F: FnMut(Self::ListId, &[Self::Id]) + Send; + + /// Materialize canonical full-precision vectors for maintenance. + fn read_vectors(&mut self, ids: I, emit: F) -> impl SendFuture> + where + I: Iterator + Send, + F: FnMut(Self::Id, &[f32]) + Send; + + /// Stage a canonical point and reserve its internal id. + fn stage_insert( + &mut self, + id: &Self::ExternalId, + element: T, + ) -> impl SendFuture>; + + /// Materialize canonical vectors for points staged by this accessor. + fn read_staged_vectors( + &mut self, + ids: I, + emit: F, + ) -> impl SendFuture> + where + I: Iterator + Send, + F: FnMut(Self::Id, &[f32]) + Send; + + /// Stage deletion of a currently visible internal point id. + fn stage_delete(&mut self, id: Self::Id) -> impl SendFuture>; + + /// Reserve fresh logical list ids that will not alias retired ids. + fn reserve_list_ids( + &mut self, + count: usize, + ) -> impl SendFuture, Self::Error>>; + + /// Apply the complete logical update and finish this mutation operation. + /// + /// Returning `Ok(())` means all components expose one coherent resulting + /// index. Rollback, durability, concurrent-reader visibility, and recovery + /// after `Err` are intentionally provider-defined. + fn apply( + self, + update: PartitionUpdate, + ) -> impl SendFuture>; +} + +/// Factory for an operation-scoped dynamic IVF maintenance accessor. +pub trait MaintenanceStrategy<'a, Provider, T>: Send + Sync +where + Provider: DataProvider, + T: Send, +{ + /// Accessor used to plan and stage split/dissolve operations. + type MaintenanceAccessor: MaintenanceAccessor; + + /// Error constructing the accessor. + type Error: StandardError; + + /// Construct one maintenance accessor. + fn maintenance_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + ) -> Result; +} diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index f93b4c9c80..4806cca308 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -3,11 +3,13 @@ * Licensed under the MIT license. */ -//! IVF index traits and wrapper. +//! IVF index traits and wrappers. //! //! Search selects candidate lists, scans them, and returns the best `k` results. -//! Insert selects one list and appends the new point to it. +//! The fixed-partition API inserts by appending to one selected list, while +//! [`dynamic`] defines accessor and mutation contracts for incremental partitions. +pub mod dynamic; pub mod glue; pub mod index; diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md new file mode 100644 index 0000000000..2805b57227 --- /dev/null +++ b/rfcs/01187-incremental-ivf.md @@ -0,0 +1,820 @@ +# Incremental IVF Accessor Architecture + +| | | +|---|---| +| **Authors** | Aditya Krishnan | +| **Created** | 2026-08-19 | +| **Updated** | 2026-08-24 | +| **Status** | Draft | + +## Summary + +Add an accessor-based architecture for incrementally maintained inverted-file +(IVF) indexes. The architecture separates the GraphIVF algorithm from point +storage, centroid navigation, and inverted-list access without prescribing a +cross-component snapshot or commit protocol. + +The generic API has two principal extension points: + +1. A query-bound search accessor selects centroids and scans the corresponding + lists. +2. An operation-scoped maintenance accessor presents one unified view of point + data, centroids, reverse assignments, and inverted lists while planning and + applying a split or dissolve. + +Each concrete strategy or aggregate provider is responsible for coordinating +its components. It may use exclusive borrowing, locks, epochs, immutable roots, +database transactions, manifests, or provider-specific recovery internally. +Those mechanisms are deliberately absent from the generic IVF traits. + +The first implementation serializes mutation through `&mut self`. This prevents +search and mutation from overlapping through the same index value and gives an +in-memory provider a straightforward implementation path. More concurrent or +durable providers may implement stronger private guarantees without changing +the algorithm-facing interface. + +## 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 one joint k-means; +- create child centroids and retire parent centroids; +- read neighboring lists; and +- reassign all points in the affected regions. + +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, and fresh list ids. 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 `DataProvider` identity types, execution contexts, ranked errors, and + the graph module's strategy/accessor style. +9. Leave locking, reader visibility, durability, rollback, and recovery to each + concrete accessor implementation. +10. Preserve the existing fixed-partition IVF API while the dynamic algorithm + is developed. + +### 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 remain owned by + `DataProvider`. +- **List id**: the stable logical id shared by a centroid and its inverted list. +- **Unified view**: mutually compatible point data, centroid state, reverse + assignments, and inverted lists presented by one operation accessor. +- **Selection plan**: selected list ids and their coarse distances. It is + consumed by the same search accessor that produced it. +- **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. + +### 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. + +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. `SelectionPlan` is passed back to the search accessor that produced it. +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. + +## Proposal + +### Source Layout + +```text +diskann/src/ivf/ + mod.rs + glue.rs existing fixed-partition traits + index.rs existing fixed-partition orchestration + dynamic.rs contracts proposed by this RFC + online/ future GraphIVF algorithm and scratch state +``` + +The dynamic names remain under `ivf::dynamic` while experimental, avoiding +collisions with the current `ivf::SearchAccessor` and `ivf::SearchStrategy`. + +### Responsibilities + +| Component | Owns | Does not decide | +|---|---|---| +| Dynamic IVF algorithm | split/dissolve policy, affected regions, child fitting, reassignment | 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 | +| 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. + +### Selection Plan + +The plan is intentionally small: + +```rust +pub struct SelectedList { + pub id: ListId, + pub distance: f32, +} + +pub struct SelectionPlan { + selected: Vec>, +} +``` + +The accessor may privately retain a lock guard, epoch, immutable root, manifest, +prefetch state, or other read context. None of that state appears in the plan's +public type. A plan is meaningful only to the accessor that created it. + +### 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, + plan: SelectionPlan, + 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: + +```rust +pub trait SearchStrategy<'a, Provider, T>: Send + Sync +where + Provider: DataProvider, +{ + type SearchAccessor: SearchAccessor; + type Error: StandardError; + + fn search_accessor( + &'a self, + provider: &'a Provider, + context: &'a Provider::Context, + query: T, + ) -> Result; +} +``` + +Search orchestration is deliberately unaware of consistency mechanics: + +```text +accessor = strategy.search_accessor(provider, context, query) +plan = accessor.select_lists(nprobe) +accessor.scan_lists(plan, |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 Update + +The algorithm describes the desired logical result rather than issuing ordered +storage commands: + +```rust +pub struct PointMove { + pub id: Id, + pub from: Option, + pub to: Option, +} + +pub struct CentroidDelta { + pub insert: Vec>, + pub retire: Vec, +} + +pub struct PartitionUpdate { + pub centroids: CentroidDelta, + pub point_moves: Vec>, +} +``` + +`from = None` denotes insertion, `to = None` denotes deletion, and two present +values denote reassignment. Both values may not be absent. Each point occurs at +most once in a normalized update, including when planned split regions overlap. + +Before returning success, an implementation validates at least: + +- inserted centroid ids were reserved by this accessor; +- retired centroid ids exist in the accessor's planning view; +- every `from` assignment matches that view; +- every `to` list exists in the resulting live set; +- no point or centroid id occurs in conflicting operations; and +- staged point insertions and deletions agree with point moves. + +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. + +### Maintenance Accessor + +The maintenance accessor combines planning reads and mutation staging: + +```rust +pub trait MaintenanceAccessor: HasId + Send + Sized +where + T: Send, +{ + type ExternalId: PartialEq + Send + Sync + 'static; + type ListId: VectorId; + type Centroids: CentroidIndex< + ListId = Self::ListId, + Error = Self::Error, + >; + type Error: ToRanked + Debug + Send + Sync + 'static; + + fn centroids(&self) -> &Self::Centroids; + + // Batched planning reads: + // list_metadata, assignments, read_members, read_vectors + + fn stage_insert( + &mut self, + id: &Self::ExternalId, + element: T, + ) -> impl SendFuture>; + + fn read_staged_vectors( + &mut self, + ids: I, + emit: F, + ) -> impl SendFuture> + where + I: Iterator + Send, + F: FnMut(Self::Id, &[f32]) + Send; + + fn stage_delete( + &mut self, + id: Self::Id, + ) -> impl SendFuture>; + + fn reserve_list_ids( + &mut self, + count: usize, + ) -> impl SendFuture, Self::Error>>; + + fn apply( + self, + update: PartitionUpdate, + ) -> impl SendFuture>; +} +``` + +The exact read signatures live in `ivf::dynamic`; all use callbacks so providers +can batch or stream results. `read_members` may invoke its callback more than +once per list. + +`stage_insert` owns external/internal id allocation and canonical input storage +for this operation. `read_staged_vectors` exposes the provider's canonical +`f32` decoding without requiring those points to be visible through ordinary +provider reads. `stage_delete` 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. It 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`. + +### 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 `DataProvider`. + +### 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. +This rule covers operations through the index API, not aliases or administrative +paths exposed independently by a provider. 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. + +## Algorithms + +### 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 a partition update containing those centroid records and no points. + +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. + +### Batch Insert + +1. Construct one maintenance accessor for the operation. +2. Stage every external id and input element, obtaining internal ids. + Duplicate or already-live external ids fail according to provider semantics. +3. Batch-materialize canonical vectors for the staged points. +4. Route all inputs through the accessor's centroid index. +5. Read routed-list sizes and compute projected post-insert sizes. +6. Admit overflow parents subject to the live-cluster cap and fresh-id budget. +7. Read each parent's members and nearest live centroid neighbors. +8. Fit `2l` children with one joint k-means over `l` admitted parent regions, + seeded with two members per parent. +9. Reserve stable ids for all child centroids. +10. Reassign affected points exactly among each region's surviving neighbors and + two new children. +11. Normalize overlapping regions into one final move per point. +12. Build one `PartitionUpdate` containing centroid changes and point moves. +13. Consume the accessor with `apply(update)`. + +With one overflow parent, joint `2l`-means reduces to local 2-means. Routes are +computed before structural change; regional reassignment deliberately revisits +points whose routes become stale due to the split. + +### Split Semantics + +A split retires one parent and installs two children, for a net increase of one +live cluster. Parent ids are never reused. Candidate centroids for a region are +the parent's selected live neighbors plus its two children. 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 `PartitionUpdate` semantics. + +### 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 containing point deletes, retired victim centroids, and + reassigned victim members. +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; 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 index 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. +- Keep the fixed-partition API and tests unchanged. +- 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 with exclusive access and 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 selection plans cannot be used meaningfully outside their producing + accessor's configured view. + +### Phase 3: Exact-Centroid GraphIVF + +- Port batch routing, split admission, joint child fitting, and tiled regional + reassignment 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 +`PartitionUpdate`. + +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 selection plans difficult to +inspect. Keeping the plan as a small value and private state in one accessor is +simpler. + +### Imperative List Mutation Methods + +Adding `append`, `remove`, `move`, `create_list`, and `retire_list` exposes a +partial-write schedule to GraphIVF. A complete `PartitionUpdate` lets each +provider choose a suitable physical implementation. + +### 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 `stage_insert` and `stage_delete` remain on the + maintenance accessor, or should a future `TransactionalDataProvider` supply + them? +4. **Update granularity**: keep one normalized move per point, or add optional + whole-list replacement hints for bulk backends? +5. **Selection plan 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 From 092aa7d94b2277aaec1d384a2d6729c824d68055 Mon Sep 17 00:00:00 2001 From: Magdalen Dobson Manohar <58752279+magdalendobson@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:00:38 -0400 Subject: [PATCH 05/33] Add co-location sets to the initial contract (#1438) Co-authored-by: Magdalen Manohar --- diskann/src/ivf/dynamic.rs | 69 +++++++++++++++++++++++++++++++++-- rfcs/01187-incremental-ivf.md | 24 ++++++++---- 2 files changed, 83 insertions(+), 10 deletions(-) diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index f01f3d947d..aa20e3ec8d 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -101,6 +101,33 @@ pub trait CentroidIndex: Send + Sync { ) -> impl SendFuture, Self::Error>>; } +/// An authoritative grouping of live centroids for locality-preserving storage. +/// +/// Every live centroid belongs to exactly one co-location group. Implementations +/// may use an identity grouping, where each group contains one centroid. Group +/// ids remain stable until retired and are independent of centroid/list ids. +pub trait CoLocationSet: Send + Sync { + /// Stable logical centroid/list id. + type ListId: VectorId; + + /// Stable logical co-location group id. + type CoLocationGroupId: VectorId; + + /// Number of live co-location groups. + fn len(&self) -> usize; + + /// Whether there are no live co-location groups. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Borrow the centroid ids in a live co-location group. + fn group(&self, id: Self::CoLocationGroupId) -> Option<&[Self::ListId]>; + + /// Return the live co-location group containing a centroid. + fn group_for(&self, id: Self::ListId) -> Option; +} + /// Query-bound access to centroid selection and inverted-list scanning. /// /// This is the dynamic IVF search algorithm's primary extension point, in the @@ -177,6 +204,24 @@ pub struct CentroidDelta { pub retire: Vec, } +/// A new co-location group to install as part of a partition update. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CoLocationGroup { + /// Fresh stable group id. Installed ids must never be reused after retirement. + pub id: CoLocationGroupId, + /// Live centroids that should be stored together. + pub centroids: Vec, +} + +/// Changes to the authoritative co-location set. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CoLocationDelta { + /// Co-location groups made live by this update. + pub insert: Vec>, + /// Live co-location groups retired by this update. + pub retire: Vec, +} + /// One point membership change in a partition mutation. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct PointMove { @@ -192,11 +237,14 @@ pub struct PointMove { /// /// The maintenance accessor applies this update across point identity, /// canonical vectors, list membership, reverse assignments, scan payloads, and -/// centroid liveness according to its provider-specific consistency contract. +/// centroid liveness, and co-location membership according to its +/// provider-specific consistency contract. #[derive(Debug, Clone, PartialEq)] -pub struct PartitionUpdate { +pub struct PartitionUpdate { /// Centroids inserted and retired by this mutation. pub centroids: CentroidDelta, + /// Co-location groups inserted and retired by this mutation. + pub co_locations: CoLocationDelta, /// Point membership changes, including inserts and deletes. pub point_moves: Vec>, } @@ -221,15 +269,24 @@ where /// Stable centroid/list id. type ListId: VectorId; + /// Stable co-location group id. + type CoLocationGroupId: VectorId; + /// In-memory centroid catalog and navigator in this accessor's unified view. type Centroids: CentroidIndex; + /// Authoritative co-location grouping in this accessor's unified view. + type CoLocations: CoLocationSet; + /// Errors from planning reads, staging, or applying the update. type Error: ToRanked + Debug + Send + Sync + 'static; /// Borrow the centroid index used for routing and maintenance neighborhoods. fn centroids(&self) -> &Self::Centroids; + /// Borrow the co-location set used to plan locality-preserving layout. + fn co_locations(&self) -> &Self::CoLocations; + /// Read list sizes for `lists`. fn list_metadata( &mut self, @@ -284,6 +341,12 @@ where count: usize, ) -> impl SendFuture, Self::Error>>; + /// Reserve fresh co-location group ids that will not alias retired ids. + fn reserve_group_ids( + &mut self, + count: usize, + ) -> impl SendFuture, Self::Error>>; + /// Apply the complete logical update and finish this mutation operation. /// /// Returning `Ok(())` means all components expose one coherent resulting @@ -291,7 +354,7 @@ where /// after `Err` are intentionally provider-defined. fn apply( self, - update: PartitionUpdate, + update: PartitionUpdate, ) -> impl SendFuture>; } diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 2805b57227..c1baa45d9f 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -10,17 +10,19 @@ ## Summary Add an accessor-based architecture for incrementally maintained inverted-file -(IVF) indexes. The architecture separates the GraphIVF algorithm from point -storage, centroid navigation, and inverted-list access without prescribing a -cross-component snapshot or commit protocol. +(IVF) indexes, along with co-location sets of centroids to assist with ordering +spatially similar posting lists on slower storage mediums. The architecture +separates the GraphIVF algorithm from point storage, centroid navigation, and +inverted-list access without prescribing a cross-component snapshot or commit +protocol. The generic API has two principal extension points: 1. A query-bound search accessor selects centroids and scans the corresponding lists. 2. An operation-scoped maintenance accessor presents one unified view of point - data, centroids, reverse assignments, and inverted lists while planning and - applying a split or dissolve. + data, centroids, reverse assignments, inverted lists, and co-location sets + while planning and applying a split or dissolve. Each concrete strategy or aggregate provider is responsible for coordinating its components. It may use exclusive borrowing, locks, epochs, immutable roots, @@ -58,6 +60,7 @@ insert batch may: - 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 @@ -76,7 +79,7 @@ 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, and fresh list ids. These operations must +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. @@ -139,6 +142,7 @@ coordinates its provider, centroid index, and inverted-list store. 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 @@ -154,6 +158,7 @@ After a maintenance accessor returns `Ok(())` from `apply`: 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. @@ -197,9 +202,10 @@ collisions with the current `ivf::SearchAccessor` and `ivf::SearchStrategy`. | Component | Owns | Does not decide | |---|---|---| -| Dynamic IVF algorithm | split/dissolve policy, affected regions, child fitting, reassignment | physical layout, I/O shape, provider consistency mechanism | +| 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 | @@ -362,6 +368,8 @@ 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. +The co-location set updates follow an identical pattern. + ### Maintenance Accessor The maintenance accessor combines planning reads and mutation staging: @@ -510,6 +518,8 @@ the normal batch path. Reopen skips bootstrap when centroids already exist. 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. From d07637723aaac71681a3bfbdfb705790461d1217 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 29 Sep 2026 11:55:47 -0400 Subject: [PATCH 06/33] initial index impl, lots more to go --- diskann/src/ivf/dynamic.rs | 16 +- diskann/src/ivf/glue.rs | 138 ---- diskann/src/ivf/index.rs | 751 +++++++++++++++++--- diskann/src/ivf/mod.rs | 13 +- diskann/src/ivf/test/cases/concurrency.rs | 134 ---- diskann/src/ivf/test/cases/correctness.rs | 117 ---- diskann/src/ivf/test/cases/errors.rs | 111 --- diskann/src/ivf/test/cases/insert.rs | 197 ------ diskann/src/ivf/test/cases/mod.rs | 76 --- diskann/src/ivf/test/harness.rs | 182 ----- diskann/src/ivf/test/mod.rs | 10 - diskann/src/ivf/test/provider.rs | 792 ---------------------- rfcs/01187-incremental-ivf.md | 11 +- 13 files changed, 672 insertions(+), 1876 deletions(-) delete mode 100644 diskann/src/ivf/glue.rs delete mode 100644 diskann/src/ivf/test/cases/concurrency.rs delete mode 100644 diskann/src/ivf/test/cases/correctness.rs delete mode 100644 diskann/src/ivf/test/cases/errors.rs delete mode 100644 diskann/src/ivf/test/cases/insert.rs delete mode 100644 diskann/src/ivf/test/cases/mod.rs delete mode 100644 diskann/src/ivf/test/harness.rs delete mode 100644 diskann/src/ivf/test/mod.rs delete mode 100644 diskann/src/ivf/test/provider.rs diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index aa20e3ec8d..a133ed40bd 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -23,9 +23,9 @@ use crate::{ /// One centroid/list selected during coarse search. #[derive(Debug, Clone, Copy, PartialEq)] -pub struct SelectedList { +pub struct SelectedList { /// Stable logical list identifier. - pub id: ListId, + pub id: L, /// Distance from the bound query to this list's centroid. pub distance: f32, } @@ -36,8 +36,8 @@ pub struct SelectedList { /// The accessor is responsible for keeping centroid selection and list reads /// coherent using provider-specific coordination. #[derive(Debug, Clone)] -pub struct SelectionPlan { - selected: Vec>, +pub struct SelectionPlan { + selected: Vec>, } impl SelectionPlan { @@ -196,7 +196,7 @@ pub struct CentroidRecord { } /// Changes to the authoritative live-centroid catalog. -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct CentroidDelta { /// Centroids made live by this update. pub insert: Vec>, @@ -214,7 +214,7 @@ pub struct CoLocationGroup { } /// Changes to the authoritative co-location set. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct CoLocationDelta { /// Co-location groups made live by this update. pub insert: Vec>, @@ -232,6 +232,7 @@ pub struct PointMove { /// New list, or `None` for a deleted point. pub to: Option, } +/// constrcutors for this too. /// A complete logical partition change. /// @@ -239,7 +240,7 @@ pub struct PointMove { /// canonical vectors, list membership, reverse assignments, scan payloads, and /// centroid liveness, and co-location membership according to its /// provider-specific consistency contract. -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct PartitionUpdate { /// Centroids inserted and retired by this mutation. pub centroids: CentroidDelta, @@ -248,6 +249,7 @@ pub struct PartitionUpdate { /// Point membership changes, including inserts and deletes. pub point_moves: Vec>, } +/// constructor enforces safety + correctness. /// Operation-scoped reads and writes needed by split/dissolve maintenance. /// diff --git a/diskann/src/ivf/glue.rs b/diskann/src/ivf/glue.rs deleted file mode 100644 index f3e846ec6e..0000000000 --- a/diskann/src/ivf/glue.rs +++ /dev/null @@ -1,138 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use std::fmt::Debug; - -use diskann_utils::future::SendFuture; - -use crate::{ - error::{StandardError, ToRanked}, - graph::SearchOutputBuffer, - provider::{DataProvider, HasId}, -}; - -///////////////// -// Search path // -///////////////// - -/// Selects candidate inverted lists for a bound query or insert vector. -pub trait ListAccessor: Send + Sync { - /// Opaque handle identifying an inverted list within the index. - type Id: Copy + Send + Sync; - - /// The error type for [`Self::select_lists`]. - type Error: ToRanked + Debug + Send + Sync + 'static; - - /// Push the selected lists into `output` in distance order. - fn select_lists( - &mut self, - nprobe: usize, - output: &mut B, - ) -> impl SendFuture> - where - B: SearchOutputBuffer + Send + ?Sized; -} - -/// Scans a set of lists for search. -pub trait SearchAccessor: HasId + Send + Sync { - /// Opaque handle identifying an inverted list within the index. - type ListId: Copy + Send + Sync; - - type Error: ToRanked + Debug + Send + Sync + 'static; - - /// Score members of `lists` and invoke `f` for each `(id, distance)` pair. - fn scan_lists(&mut self, lists: Itr, f: F) -> impl SendFuture> - where - Itr: Iterator + Send, - F: Send + FnMut(Self::Id, f32); -} - -/// Per-call factory for IVF search. -pub trait SearchStrategy<'a, Provider, T>: Send + Sync -where - Provider: DataProvider, -{ - /// The inverted-list handle type, shared by both accessors. - type ListId: Copy + Send + Sync; - - /// The fine accessor, keyed to the provider's internal id and the shared list handle. - type SearchAccessor: SearchAccessor; - - /// The coarse accessor, keyed to the shared list handle. - type ListAccessor: ListAccessor; - - /// An error that can occur when constructing either accessor. - type Error: StandardError; - - /// Construct the fine scan accessor. - fn search_accessor( - &'a self, - provider: &'a Provider, - context: &'a Provider::Context, - query: T, - ) -> Result; - - /// Construct the coarse list-selection accessor. - fn list_accessor( - &'a self, - provider: &'a Provider, - context: &'a Provider::Context, - query: T, - ) -> Result; -} - -///////////////// -// Insert path // -///////////////// - -/// Appends a vector to a chosen list during insert. -pub trait InsertAccessor: HasId + Send + Sync { - /// Opaque handle identifying an inverted list within the index. - type ListId: Copy + Send + Sync; - - /// The error type for [`Self::append`]. - type Error: ToRanked + Debug + Send + Sync + 'static; - - /// Append `vector` to `list` under `id`. - fn append( - &mut self, - list: Self::ListId, - id: Self::Id, - vector: T, - ) -> impl SendFuture>; -} - -/// Per-call factory for IVF insert. -pub trait InsertStrategy<'a, Provider, T>: Send + Sync -where - Provider: DataProvider, -{ - /// The inverted-list handle type, shared by both accessors. - type ListId: Copy + Send + Sync; - - /// The append accessor, keyed to the provider's internal id and the shared list handle. - type InsertAccessor: InsertAccessor; - - /// The coarse accessor, keyed to the shared list handle. - type ListAccessor: ListAccessor; - - /// An error that can occur when constructing either accessor. - type Error: StandardError; - - /// Construct the append accessor. - fn insert_accessor( - &'a self, - provider: &'a Provider, - context: &'a Provider::Context, - ) -> Result; - - /// Construct the coarse list-selection accessor. - fn list_accessor( - &'a self, - provider: &'a Provider, - context: &'a Provider::Context, - vector: T, - ) -> Result; -} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index d7b2d42c50..3b4448e853 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -3,41 +3,101 @@ * Licensed under the MIT license. */ -//! IVF index wrapper. +//! Dynamic IVF index orchestration. +//! +//! The index owns split policy and drives one [`MaintenanceAccessor`] per mutation; +//! it never touches point, centroid, or list storage directly. -use std::num::NonZeroUsize; +use std::{ + collections::{BTreeMap, BTreeSet, btree_map::Entry}, + fmt::{Debug, Display}, +}; -use diskann_utils::future::SendFuture; +use diskann_utils::{future::SendFuture, views::MatrixView}; +use rand::{Rng, SeedableRng, rngs::StdRng}; +use thiserror::Error; use crate::{ - ANNResult, - error::{ANNError, ANNErrorKind, ErrorExt, IntoANNResult}, - graph::SearchOutputBuffer, - ivf::{InsertAccessor, InsertStrategy, ListAccessor, SearchAccessor, SearchStrategy}, - neighbor::{Neighbor, NeighborPriorityQueue, NeighborPriorityQueueIdType}, - provider::{DataProvider, Guard, SetElement}, + ANNError, ANNErrorKind, ANNResult, + error::{ErrorExt, IntoANNResult}, + ivf::dynamic::{ + CentroidDelta, CentroidIndex, CentroidRecord, MaintenanceAccessor, MaintenanceStrategy, + PartitionUpdate, PointMove, + }, + provider::DataProvider, + utils::VectorId, }; -/// Statistics collected during an IVF search. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct SearchStats { - /// Distance computations performed while scanning lists. - pub cmps: u32, +/// Parameters for the online split policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DynamicIvfConfig { + /// Split a list once an insert batch would grow it beyond this many points. + pub split_threshold: usize, + /// Hard cap on live lists; `None` allows unbounded growth. + pub max_clusters: Option, + /// Nearby lists reassigned together with each split parent. + pub reassign_neighbors: usize, + /// Lloyd iterations used to fit split children. + pub two_means_iterations: usize, + /// Seed for split-child initialization. + pub seed: u64, +} + +/// Invalid [`DynamicIvfConfig`] values. +#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] +pub enum ConfigError { + #[error("split_threshold must be at least 2, got {0}")] + SplitThreshold(usize), + #[error("reassign_neighbors must be at least 1")] + ReassignNeighbors, +} + +impl From for ANNError { + #[track_caller] + fn from(err: ConfigError) -> Self { + ANNError::new(ANNErrorKind::IndexConfigError, err) + } +} - /// Results written to the output buffer. - pub result_count: u32, +/// Work performed by one [`DynamicIvfIndex::insert_batch`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct InsertStats { + /// Points added to the index. + pub inserted: usize, + /// Lists split into two children. + pub splits: usize, + /// Previously indexed points moved by regional reassignment. + pub reassigned: usize, } -/// IVF index wrapper over a [`DataProvider`]. +/// An incrementally maintained IVF index. +/// +/// Mutations take `&mut self`, so they cannot overlap searches through the same value. #[derive(Debug)] -pub struct IvfIndex { +pub struct DynamicIvfIndex { provider: P, + config: DynamicIvfConfig, + rng: StdRng, } -impl IvfIndex

{ - /// Construct a new index around `provider`. - pub fn new(provider: P) -> Self { - Self { provider } +impl DynamicIvfIndex

{ + /// Construct an index over `provider`. + /// + /// # Errors + /// + /// Returns [`ConfigError`] if `config` is invalid. + pub fn new(provider: P, config: DynamicIvfConfig) -> Result { + if config.split_threshold < 2 { + return Err(ConfigError::SplitThreshold(config.split_threshold)); + } + if config.reassign_neighbors == 0 { + return Err(ConfigError::ReassignNeighbors); + } + Ok(Self { + provider, + rng: StdRng::seed_from_u64(config.seed), + config, + }) } /// Borrow the underlying provider. @@ -45,110 +105,613 @@ impl IvfIndex

{ &self.provider } - /// Run IVF k-nearest-neighbor search. - pub fn knn_search<'a, S, T, OB>( - &'a self, - k: NonZeroUsize, - nprobe: usize, + /// Borrow the split policy. + pub fn config(&self) -> &DynamicIvfConfig { + &self.config + } + + /// Install the initial centroids. + /// + /// # Errors + /// + /// Fails if `centroids` is empty, the index already has centroids, or the + /// maintenance accessor fails. + pub fn initialize<'a, S, T>( + &'a mut self, strategy: &'a S, context: &'a P::Context, - query: T, - output: &mut OB, - ) -> impl SendFuture> + centroids: MatrixView<'a, f32>, + ) -> impl SendFuture> where - S: SearchStrategy<'a, P, T>, - S::ListId: Eq, - P::InternalId: NeighborPriorityQueueIdType, - T: Copy + Send, - OB: SearchOutputBuffer + Send + ?Sized, + S: MaintenanceStrategy<'a, P, T>, + T: Send, { + let provider = &self.provider; async move { - let mut list_accessor = strategy - .list_accessor(&self.provider, context, query) - .into_ann_result()?; - - let mut lists: Vec> = Vec::with_capacity(nprobe); - list_accessor - .select_lists(nprobe, &mut lists) - .await - .escalate("IVF coarse list selection must complete")?; + let count = centroids.nrows(); + if count == 0 { + return Err(index_error("initialize requires at least one centroid")); + } - let mut search_accessor = strategy - .search_accessor(&self.provider, context, query) + let mut accessor = strategy + .maintenance_accessor(provider, context) .into_ann_result()?; - let k = k.get(); - let mut queue = NeighborPriorityQueue::new(k); - let mut cmps: u32 = 0; + if !accessor.centroids().is_empty() { + return Err(index_error("dynamic IVF index is already initialized")); + } - search_accessor - .scan_lists(lists.iter().map(|n| n.id), |id, dist| { - cmps += 1; - queue.insert(Neighbor::new(id, dist)); - }) + let list_ids = accessor + .reserve_list_ids(count) .await - .escalate("IVF list scan must complete to produce correct k-NN results")?; + .escalate("initialize must reserve list ids")?; + check_reserved(count, list_ids.len())?; - let result_count = - output.extend(queue.iter().take(k).map(|n| (n.id, n.distance))) as u32; + let centroids = CentroidDelta { + insert: list_ids + .into_iter() + .zip(centroids.row_iter()) + .map(|(id, vector)| CentroidRecord { + id, + vector: vector.into(), + }) + .collect(), + retire: Vec::new(), + }; - Ok(SearchStats { cmps, result_count }) + accessor + .apply(PartitionUpdate { + centroids, + ..Default::default() + }) + .await + .escalate("initialize must apply the initial centroids") } } - /// Insert a vector under external id `id`. - pub fn insert<'a, S, T>( - &'a self, + /// Insert a batch, splitting every list the batch pushes past `split_threshold`. + /// + /// Routing and split planning run against the accessor's view before any change; + /// the whole batch is then handed to [`MaintenanceAccessor::apply`] as one update. + /// + /// # Errors + /// + /// Fails if the index is not initialized, the provider returns inconsistent + /// data, or any accessor call fails. State after an `apply` failure is + /// provider-defined. + pub fn insert_batch<'a, S, T>( + &'a mut self, strategy: &'a S, context: &'a P::Context, - id: &P::ExternalId, - vector: T, - ) -> impl SendFuture> + points: &'a [(P::ExternalId, T)], + ) -> impl SendFuture> where - S: InsertStrategy<'a, P, T>, - S::ListId: Eq, - P: SetElement, - T: Copy + Send, + S: MaintenanceStrategy<'a, P, T>, + T: Copy + Send + Sync, { + let provider = &self.provider; + let config = self.config; + let rng = &mut self.rng; async move { - let guard = self - .provider - .set_element(context, id, vector) - .await - .escalate("IVF insert requires a successful `set_element`")?; - - let internal_id = guard.id(); + if points.is_empty() { + return Ok(InsertStats::default()); + } - let mut list_accessor = strategy - .list_accessor(&self.provider, context, vector) + let mut accessor = strategy + .maintenance_accessor(provider, context) .into_ann_result()?; - let mut lists: Vec> = Vec::with_capacity(1); + if accessor.centroids().is_empty() { + return Err(index_error( + "insert requires an initialized dynamic IVF index", + )); + } + + let (new_ids, mut vectors) = stage_points(&mut accessor, points).await?; + + let mut routed: BTreeMap<_, Vec<_>> = BTreeMap::new(); + + for &id in &new_ids { + let plan = accessor + .centroids() + .select(vectors.get(id)?, 1) + .await + .escalate("insert must route every point")?; + let list = plan + .selected() + .first() + .map(|selected| selected.id) + .ok_or_else(|| index_error("centroid selection returned no list"))?; + routed.entry(list).or_default().push(id); + } + + let parents = select_split_parents(&mut accessor, &config, &routed).await?; - list_accessor - .select_lists(1, &mut lists) + if parents.is_empty() { + let update = PartitionUpdate { + point_moves: routed_moves(&routed, &BTreeMap::new()), + ..Default::default() + }; + accessor + .apply(update) + .await + .escalate("insert must apply routed points")?; + + return Ok(InsertStats { + inserted: new_ids.len(), + ..Default::default() + }); + } + + let plan = SplitPlan::prepare(&mut accessor, &config, &parents, &mut vectors).await?; + let children = + plan.fit_children(&routed, &vectors, config.two_means_iterations, rng)?; + + let child_ids = accessor + .reserve_list_ids(children.len()) .await - .escalate("IVF insert must select a target list")?; + .escalate("split must reserve child list ids")?; + check_reserved(children.len(), child_ids.len())?; - let list = lists.first().map(|n| n.id).ok_or_else(|| { - ANNError::message( - ANNErrorKind::IndexError, - "IVF insert: list selection returned no candidate list", - ) - })?; + let placements = plan.reassign(&routed, &vectors, &child_ids, &children)?; - let mut insert_accessor = strategy - .insert_accessor(&self.provider, context) - .into_ann_result()?; + let mut point_moves = routed_moves(&routed, &placements); + let mut reassigned = 0; + for (&from, ids) in &plan.members { + for &id in ids { + if let Some(&to) = placements.get(&id) + && to != from + { + reassigned += 1; + point_moves.push(PointMove { + id, + from: Some(from), + to: Some(to), + }); + } + } + } - insert_accessor - .append(list, internal_id, vector) + let splits = parents.len(); + let update = PartitionUpdate { + centroids: CentroidDelta { + insert: child_ids + .iter() + .zip(children) + .map(|(&id, vector)| CentroidRecord { id, vector }) + .collect(), + retire: parents, + }, + point_moves, + ..Default::default() + }; + + accessor + .apply(update) .await - .escalate("IVF insert must append the vector to its assigned list")?; + .escalate("insert must apply the split update")?; + + Ok(InsertStats { + inserted: new_ids.len(), + splits, + reassigned, + }) + } + } +} - guard.complete().await; +/////////////////// +// Insert phases // +/////////////////// + +/// Canonical vectors gathered for one mutation, all of dimension `dim`. +struct Vectors { + dim: usize, + rows: BTreeMap>, +} +impl Vectors { + fn get(&self, id: Id) -> ANNResult<&[f32]> { + self.rows + .get(&id) + .map(|row| &**row) + .ok_or_else(|| index_error(format!("no canonical vector for point {id}"))) + } + + fn check_dim(&self, len: usize) -> ANNResult<()> { + if len == self.dim { Ok(()) + } else { + Err(index_error(format!( + "vector dimension {len} does not match batch dimension {}", + self.dim + ))) + } + } + + fn check_dims(&self) -> ANNResult<()> { + self.rows + .values() + .try_for_each(|row| self.check_dim(row.len())) + } +} + +/// Stage every point and materialize its canonical vector. +async fn stage_points( + accessor: &mut A, + points: &[(A::ExternalId, T)], +) -> ANNResult<(Vec, Vectors)> +where + A: MaintenanceAccessor, + T: Copy + Send + Sync, +{ + let mut ids = Vec::with_capacity(points.len()); + for (external, element) in points { + ids.push( + accessor + .stage_insert(external, *element) + .await + .escalate("insert must stage every point")?, + ); + } + + let mut rows: BTreeMap> = BTreeMap::new(); + accessor + .read_staged_vectors(ids.iter().copied(), |id, vector| { + rows.insert(id, Box::from(vector)); + }) + .await + .escalate("insert must read staged canonical vectors")?; + + let dim = rows.values().next().map_or(0, |row| row.len()); + if dim == 0 { + return Err(index_error("staged canonical vectors must be non-empty")); + } + let vectors = Vectors { dim, rows }; + vectors.check_dims()?; + Ok((ids, vectors)) +} + +/// Route each staged point to its nearest live list. +async fn route( + accessor: &mut A, + ids: &[A::Id], + vectors: &Vectors, +) -> ANNResult>> +where + A: MaintenanceAccessor, + T: Send, +{ + let centroids = accessor.centroids(); + let mut routed: BTreeMap> = BTreeMap::new(); + + for &id in ids { + let plan = centroids + .select(vectors.get(id)?, 1) + .await + .escalate("insert must route every point")?; + let list = plan + .selected() + .first() + .map(|selected| selected.id) + .ok_or_else(|| index_error("centroid selection returned no list"))?; + routed.entry(list).or_default().push(id); + } + + Ok(routed) +} + +/// Return routed lists whose projected size exceeds the threshold, largest first +/// under the cluster cap, in ascending id order. +async fn select_split_parents( + accessor: &mut A, + config: &DynamicIvfConfig, + routed: &BTreeMap>, +) -> ANNResult> +where + A: MaintenanceAccessor, + T: Send, +{ + let mut sizes = BTreeMap::new(); + accessor + .list_metadata(routed.keys().copied(), |metadata| { + sizes.insert(metadata.id, metadata.len); + }) + .await + .escalate("insert must read routed list sizes")?; + + let mut overflowing = Vec::new(); + for (&list, incoming) in routed { + let current = sizes + .get(&list) + .copied() + .ok_or_else(|| index_error(format!("no metadata for list {list}")))?; + let projected = current + incoming.len(); + if projected > config.split_threshold { + overflowing.push((projected, list)); + } + } + + // Each split adds one live list. + let budget = config.max_clusters.map_or(usize::MAX, |max| { + max.saturating_sub(accessor.centroids().len()) + }); + overflowing.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.cmp(&b.1))); + overflowing.truncate(budget); + + let mut parents: Vec<_> = overflowing.into_iter().map(|(_, list)| list).collect(); + parents.sort_unstable(); + Ok(parents) +} + +/// One split parent and the surviving lists reassigned alongside it. +struct Region { + parent: L, + neighbors: Vec, +} + +/// Everything read from the accessor to split a set of parents. +struct SplitPlan { + regions: Vec>, + neighbor_centroids: BTreeMap>, + /// Current members of every parent and neighbor list. + members: BTreeMap>, +} + +impl SplitPlan { + async fn prepare( + accessor: &mut A, + config: &DynamicIvfConfig, + parents: &[L], + vectors: &mut Vectors, + ) -> ANNResult + where + A: MaintenanceAccessor, + T: Send, + { + let parent_set: BTreeSet = parents.iter().copied().collect(); + let centroids = accessor.centroids(); + let mut regions = Vec::with_capacity(parents.len()); + let mut neighbor_centroids = BTreeMap::new(); + + for &parent in parents { + let anchor = centroids + .centroid(parent) + .ok_or_else(|| index_error(format!("split parent {parent} is not live")))?; + let selected = centroids + .select(anchor, config.reassign_neighbors + 1) + .await + .escalate("split must find neighbor lists")?; + + // Other parents are retired by this update, so they are never candidates. + let neighbors: Vec = selected + .into_selected() + .into_iter() + .map(|selected| selected.id) + .filter(|id| !parent_set.contains(id)) + .take(config.reassign_neighbors) + .collect(); + + for &neighbor in &neighbors { + if let Entry::Vacant(entry) = neighbor_centroids.entry(neighbor) { + let vector = centroids.centroid(neighbor).ok_or_else(|| { + index_error(format!("neighbor list {neighbor} is not live")) + })?; + vectors.check_dim(vector.len())?; + entry.insert(Box::from(vector)); + } + } + regions.push(Region { parent, neighbors }); + } + + let lists: BTreeSet = regions + .iter() + .flat_map(|region| { + std::iter::once(region.parent).chain(region.neighbors.iter().copied()) + }) + .collect(); + + let mut members: BTreeMap> = BTreeMap::new(); + accessor + .read_members(lists.iter().copied(), |list, ids| { + members.entry(list).or_default().extend_from_slice(ids) + }) + .await + .escalate("split must read region members")?; + + accessor + .read_vectors(members.values().flatten().copied(), |id, vector| { + vectors.rows.insert(id, Box::from(vector)); + }) + .await + .escalate("split must read region canonical vectors")?; + vectors.check_dims()?; + + Ok(Self { + regions, + neighbor_centroids, + members, + }) + } + + /// Existing members plus incoming points of `list`. + fn points_of<'s>( + &'s self, + list: L, + routed: &'s BTreeMap>, + ) -> impl Iterator + 's { + self.members + .get(&list) + .into_iter() + .chain(routed.get(&list)) + .flatten() + .copied() + } + + /// Fit two children per parent with one joint k-means over every parent's points. + fn fit_children( + &self, + routed: &BTreeMap>, + vectors: &Vectors, + iterations: usize, + rng: &mut StdRng, + ) -> ANNResult>> { + let mut points = Vec::new(); + let mut centers = Vec::with_capacity(2 * self.regions.len()); + for region in &self.regions { + let start = points.len(); + for id in self.points_of(region.parent, routed) { + points.push(vectors.get(id)?); + } + + let span = points.len() - start; + if span < 2 { + return Err(index_error(format!( + "split parent {} has fewer than two points", + region.parent + ))); + } + let a = rng.random_range(0..span); + let b = (a + rng.random_range(1..span)) % span; + centers.push(points[start + a].to_vec()); + centers.push(points[start + b].to_vec()); + } + + lloyd(&points, &mut centers, iterations.max(1)); + Ok(centers.into_iter().map(Vec::into_boxed_slice).collect()) + } + + /// Place every region point on its nearest candidate across all regions containing it. + fn reassign( + &self, + routed: &BTreeMap>, + vectors: &Vectors, + child_ids: &[L], + children: &[Box<[f32]>], + ) -> ANNResult> { + let mut best: BTreeMap = BTreeMap::new(); + for ((region, ids), vecs) in self + .regions + .iter() + .zip(child_ids.chunks_exact(2)) + .zip(children.chunks_exact(2)) + { + let mut candidates: Vec<(L, &[f32])> = Vec::with_capacity(region.neighbors.len() + 2); + for &neighbor in ®ion.neighbors { + let vector = self.neighbor_centroids.get(&neighbor).ok_or_else(|| { + index_error(format!("missing centroid for neighbor list {neighbor}")) + })?; + candidates.push((neighbor, vector)); + } + candidates.extend(ids.iter().copied().zip(vecs.iter().map(|v| &**v))); + + let lists = std::iter::once(region.parent).chain(region.neighbors.iter().copied()); + for list in lists { + for id in self.points_of(list, routed) { + let (to, distance) = nearest(vectors.get(id)?, &candidates) + .ok_or_else(|| index_error("split region has no candidate lists"))?; + match best.entry(id) { + Entry::Vacant(entry) => { + entry.insert((to, distance)); + } + Entry::Occupied(mut entry) => { + if distance < entry.get().1 { + entry.insert((to, distance)); + } + } + } + } + } } + Ok(best.into_iter().map(|(id, (list, _))| (id, list)).collect()) + } +} + +/// Moves for newly inserted points: their split placement, else their route. +fn routed_moves( + routed: &BTreeMap>, + placements: &BTreeMap, +) -> Vec> { + routed + .iter() + .flat_map(|(&route, ids)| { + ids.iter().map(move |&id| PointMove { + id, + from: None, + to: Some(placements.get(&id).copied().unwrap_or(route)), + }) + }) + .collect() +} + +///////////// +// Helpers // +///////////// + +// Clustering and reassignment use squared L2 regardless of the search metric. +fn squared_l2(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() +} + +fn nearest(point: &[f32], candidates: &[(L, &[f32])]) -> Option<(L, f32)> { + candidates + .iter() + .map(|&(list, centroid)| (list, squared_l2(point, centroid))) + .min_by(|a, b| a.1.total_cmp(&b.1)) +} + +fn lloyd(points: &[&[f32]], centers: &mut [Vec], iterations: usize) { + let dim = centers.first().map_or(0, Vec::len); + if dim == 0 { + return; + } + let mut sums = vec![0.0f32; centers.len() * dim]; + let mut counts = vec![0usize; centers.len()]; + for _ in 0..iterations { + sums.fill(0.0); + counts.fill(0); + for &point in points { + let Some(c) = centers + .iter() + .map(|center| squared_l2(point, center)) + .enumerate() + .min_by(|a, b| a.1.total_cmp(&b.1)) + .map(|(c, _)| c) + else { + continue; + }; + counts[c] += 1; + for (sum, x) in sums[c * dim..(c + 1) * dim].iter_mut().zip(point) { + *sum += x; + } + } + // Empty clusters keep their previous center. + for ((center, sum), &count) in centers.iter_mut().zip(sums.chunks_exact(dim)).zip(&counts) { + if count > 0 { + let scale = 1.0 / count as f32; + for (x, s) in center.iter_mut().zip(sum) { + *x = s * scale; + } + } + } + } +} + +#[track_caller] +fn index_error(message: D) -> ANNError +where + D: Display + Debug + Send + Sync + 'static, +{ + ANNError::message(ANNErrorKind::IndexError, message) +} + +#[track_caller] +fn check_reserved(expected: usize, actual: usize) -> ANNResult<()> { + if expected == actual { + Ok(()) + } else { + Err(index_error(format!( + "accessor reserved {actual} ids, expected {expected}" + ))) } } diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 4806cca308..bb417e4de0 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -3,18 +3,9 @@ * Licensed under the MIT license. */ -//! IVF index traits and wrappers. -//! -//! Search selects candidate lists, scans them, and returns the best `k` results. -//! The fixed-partition API inserts by appending to one selected list, while -//! [`dynamic`] defines accessor and mutation contracts for incremental partitions. +//! Incrementally maintained IVF index. pub mod dynamic; -pub mod glue; pub mod index; -pub use glue::{InsertAccessor, InsertStrategy, ListAccessor, SearchAccessor, SearchStrategy}; -pub use index::{IvfIndex, SearchStats}; - -#[cfg(test)] -mod test; +pub use index::{ConfigError, DynamicIvfConfig, DynamicIvfIndex, InsertStats}; diff --git a/diskann/src/ivf/test/cases/concurrency.rs b/diskann/src/ivf/test/cases/concurrency.rs deleted file mode 100644 index f13e974096..0000000000 --- a/diskann/src/ivf/test/cases/concurrency.rs +++ /dev/null @@ -1,134 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Concurrency behavior of the fine-scan accessor. -//! -//! [`ScanAccessor`](crate::ivf::test::provider::ScanAccessor) exposes two paths -- -//! [`ScanMode::Sequential`] and [`ScanMode::Concurrent`] (one spawned task per probed list). -//! These tests assert the two paths produce identical results, that the concurrent path -//! actually spawned tasks, and that many searches can run concurrently over a shared index. - -use std::sync::Arc; - -use crate::{ - ivf::test::{ - cases::{METRIC, N_LISTS, grid_index, queries}, - harness::{IvfOracleRun, assert_same_distances}, - provider::{Context, ScanMode, Strategy}, - }, - test::tokio::current_thread_runtime, -}; - -/// The sequential and concurrent scan paths must return the same top-`k` for every query. -#[test] -fn sequential_and_concurrent_agree() { - let index = grid_index(); - let dim = index.provider().dim(); - let sequential = Strategy::new(dim, METRIC).with_mode(ScanMode::Sequential); - let concurrent = Strategy::new(dim, METRIC).with_mode(ScanMode::Concurrent); - - for query in queries() { - for k in [1usize, 8] { - let seq = IvfOracleRun::run_sync(&index, &sequential, &query, k, N_LISTS).unwrap(); - let con = IvfOracleRun::run_sync(&index, &concurrent, &query, k, N_LISTS).unwrap(); - - assert_same_distances(&con.top_k, &seq.ground_truth); - assert_same_distances(&con.top_k, &seq.top_k); - assert_eq!(con.stats.cmps, seq.stats.cmps); - } - } -} - -/// The concurrent path must route its list scans through `wrap_spawn`, so the context -/// observes one spawn per probed list. -#[tokio::test] -async fn concurrent_path_spawns_per_list() { - let index = grid_index(); - let dim = index.provider().dim(); - let strategy = Strategy::new(dim, METRIC).with_mode(ScanMode::Concurrent); - - let context = Context::new(); - let mut buf = vec![crate::neighbor::Neighbor::::default(); 4]; - - let stats = index - .knn_search( - std::num::NonZeroUsize::new(4).unwrap(), - N_LISTS, - &strategy, - &context, - [3.5f32, 3.5].as_slice(), - &mut crate::neighbor::BackInserter::new(buf.as_mut_slice()), - ) - .await - .unwrap(); - - assert_eq!( - context.spawns(), - N_LISTS, - "concurrent scan should spawn one task per probed list", - ); - assert_eq!(stats.cmps as usize, index.provider().len()); -} - -/// The sequential path performs no spawns. -#[tokio::test] -async fn sequential_path_does_not_spawn() { - let index = grid_index(); - let dim = index.provider().dim(); - let strategy = Strategy::new(dim, METRIC).with_mode(ScanMode::Sequential); - - let context = Context::new(); - let mut buf = vec![crate::neighbor::Neighbor::::default(); 4]; - - index - .knn_search( - std::num::NonZeroUsize::new(4).unwrap(), - N_LISTS, - &strategy, - &context, - [3.5f32, 3.5].as_slice(), - &mut crate::neighbor::BackInserter::new(buf.as_mut_slice()), - ) - .await - .unwrap(); - - assert_eq!(context.spawns(), 0); -} - -/// Many concurrent searches over a shared index (on a multi-threaded runtime) must each -/// return the correct restricted top-`k`. -#[test] -fn many_concurrent_searches_share_one_index() { - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(4) - .enable_all() - .build() - .expect("multi-thread runtime"); - - let index = Arc::new(grid_index()); - let dim = index.provider().dim(); - let strategy = Arc::new(Strategy::new(dim, METRIC).with_mode(ScanMode::Concurrent)); - - runtime.block_on(async move { - let mut set = tokio::task::JoinSet::new(); - for query in queries() { - let index = Arc::clone(&index); - let strategy = Arc::clone(&strategy); - set.spawn(async move { - let run = IvfOracleRun::run(&index, &strategy, &query, 8, N_LISTS - 1) - .await - .unwrap(); - assert_same_distances(&run.top_k, &run.ground_truth); - }); - } - while let Some(joined) = set.join_next().await { - joined.expect("search task panicked"); - } - }); -} - -// Bring the single-thread runtime helper into scope for documentation/intra-doc links. -#[allow(unused_imports)] -use current_thread_runtime as _current_thread_runtime; diff --git a/diskann/src/ivf/test/cases/correctness.rs b/diskann/src/ivf/test/cases/correctness.rs deleted file mode 100644 index 2334908ad4..0000000000 --- a/diskann/src/ivf/test/cases/correctness.rs +++ /dev/null @@ -1,117 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Correctness of [`crate::ivf::IvfIndex::knn_search`]. -//! -//! The central invariant is that an IVF search probing `nprobe` lists returns exactly the -//! brute-force top-`k` over the union of those lists' members. The harness computes that -//! restricted oracle directly, so each case can assert strict equality. - -use crate::ivf::test::{ - cases::{METRIC, N_LISTS, grid_index, queries}, - harness::{IvfOracleRun, assert_same_distances, global_topk}, - provider::Strategy, -}; - -/// For every query and a sweep of `(k, nprobe)`, the search must equal the restricted -/// brute-force oracle. -#[test] -fn matches_restricted_brute_force() { - let index = grid_index(); - let dim = index.provider().dim(); - let strategy = Strategy::new(dim, METRIC); - - for query in queries() { - for nprobe in 1..=N_LISTS { - for k in [1usize, 4, 16] { - let run = IvfOracleRun::run_sync(&index, &strategy, &query, k, nprobe).unwrap(); - - assert_same_distances(&run.top_k, &run.ground_truth); - - // `top_k_distances` is the distance projection of `top_k`. - let projected: Vec = run.top_k.iter().map(|(_, d)| *d).collect(); - assert_eq!(run.top_k_distances, projected); - - // The fine step must compute exactly one distance per probed member. - assert_eq!( - run.stats.cmps as usize, run.probed_members, - "query {query:?} k={k} nprobe={nprobe}: cmps must equal probed members", - ); - - // We never return more than the oracle (which is already truncated to k), - // and never more than k. - assert_eq!(run.top_k.len(), run.ground_truth.len()); - assert!(run.top_k.len() <= k); - } - } - } -} - -/// Probing every list (`nprobe == n_lists`) makes IVF exhaustive: it must equal the global -/// brute-force top-`k`. -#[test] -fn full_probe_equals_global_brute_force() { - let index = grid_index(); - let dim = index.provider().dim(); - let strategy = Strategy::new(dim, METRIC); - - for query in queries() { - for k in [1usize, 4, 16] { - let run = IvfOracleRun::run_sync(&index, &strategy, &query, k, N_LISTS).unwrap(); - - let global: Vec<(u32, f32)> = global_topk(index.provider(), &query, k) - .into_iter() - .map(|n| (n.id, n.distance)) - .collect(); - assert_same_distances(&run.top_k, &global); - - // Probing every list scans the whole dataset exactly once. - assert_eq!(run.stats.cmps as usize, index.provider().len()); - assert_eq!(run.probed_lists.len(), N_LISTS); - } - } -} - -/// `result_count` is `min(k, probed_members)` and matches the buffer contents. -#[test] -fn result_count_is_bounded() { - let index = grid_index(); - let dim = index.provider().dim(); - let strategy = Strategy::new(dim, METRIC); - let query = vec![0.0, 0.0]; - - // nprobe = 1 probes a single 16-member quadrant. - let run = IvfOracleRun::run_sync(&index, &strategy, &query, 100, 1).unwrap(); - assert_eq!(run.stats.result_count as usize, run.probed_members); - assert_eq!(run.top_k.len(), run.probed_members); - assert_eq!(run.stats.result_count as usize, run.top_k.len()); -} - -/// The provider's coarse/fine comparison counters reflect the work a search performs: -/// `n_lists` centroid comparisons (one full coarse pass) and one fine comparison per -/// probed member. -#[test] -fn provider_metrics_track_distance_work() { - let index = grid_index(); - let dim = index.provider().dim(); - let strategy = Strategy::new(dim, METRIC); - let query = vec![3.5f32, 3.5]; - - let before = index.provider().metrics(); - let run = IvfOracleRun::run_sync(&index, &strategy, &query, 4, N_LISTS).unwrap(); - let after = index.provider().metrics(); - - assert_eq!(index.provider().n_lists(), N_LISTS); - assert_eq!( - after.centroid_cmps - before.centroid_cmps, - index.provider().n_lists(), - "coarse step compares the query against every centroid once", - ); - assert_eq!( - after.member_cmps - before.member_cmps, - run.probed_members, - "fine step compares the query against every probed member once", - ); -} diff --git a/diskann/src/ivf/test/cases/errors.rs b/diskann/src/ivf/test/cases/errors.rs deleted file mode 100644 index 4438a0723e..0000000000 --- a/diskann/src/ivf/test/cases/errors.rs +++ /dev/null @@ -1,111 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Error propagation through the IVF index API. -//! -//! A failure in any accessor or strategy must surface as an [`ANNError`] out of -//! `knn_search` / `insert` rather than being swallowed or panicking. - -use std::num::NonZeroUsize; - -use crate::{ - ivf::test::{ - cases::{DIM, METRIC, N_LISTS, grid_index}, - provider::{Context, ScanMode, Strategy}, - }, - neighbor::{BackInserter, Neighbor}, - test::tokio::current_thread_runtime, -}; - -/// An injected fine-scan failure must escalate out of `knn_search` (sequential path). -#[test] -fn scan_failure_escalates_sequential() { - let index = grid_index(); - // The center query probes every quadrant; failing on list 0 is guaranteed to be hit. - let strategy = Strategy::new(DIM, METRIC) - .with_mode(ScanMode::Sequential) - .failing_on_list(0); - - let context = Context::new(); - let mut buf = vec![Neighbor::::default(); 4]; - let result = current_thread_runtime().block_on(index.knn_search( - NonZeroUsize::new(4).unwrap(), - N_LISTS, - &strategy, - &context, - [3.5f32, 3.5].as_slice(), - &mut BackInserter::new(buf.as_mut_slice()), - )); - - assert!(result.is_err(), "scan failure must propagate"); -} - -/// The same failure must escalate from the concurrent path, which has to shut down its -/// in-flight tasks before returning the error. -#[test] -fn scan_failure_escalates_concurrent() { - let index = grid_index(); - let strategy = Strategy::new(DIM, METRIC) - .with_mode(ScanMode::Concurrent) - .failing_on_list(0); - - let context = Context::new(); - let mut buf = vec![Neighbor::::default(); 4]; - let result = current_thread_runtime().block_on(index.knn_search( - NonZeroUsize::new(4).unwrap(), - N_LISTS, - &strategy, - &context, - [3.5f32, 3.5].as_slice(), - &mut BackInserter::new(buf.as_mut_slice()), - )); - - assert!( - result.is_err(), - "scan failure must propagate from concurrent path" - ); -} - -/// A query whose dimension disagrees with the strategy must fail at accessor construction. -#[test] -fn dimension_mismatch_errors() { - let index = grid_index(); - let strategy = Strategy::new(DIM, METRIC); - - let context = Context::new(); - let mut buf = vec![Neighbor::::default(); 1]; - let bad_query = vec![1.0f32, 2.0, 3.0]; // DIM is 2. - let result = current_thread_runtime().block_on(index.knn_search( - NonZeroUsize::new(1).unwrap(), - 1, - &strategy, - &context, - bad_query.as_slice(), - &mut BackInserter::new(buf.as_mut_slice()), - )); - - assert!(result.is_err(), "dimension mismatch must be rejected"); -} - -/// A dimension mismatch on insert is likewise rejected. -#[test] -fn insert_dimension_mismatch_errors() { - let index = grid_index(); - let strategy = Strategy::new(DIM, METRIC); - let new_id = index.provider().len() as u32; - let bad_vector = vec![1.0f32]; // DIM is 2. - - let result = current_thread_runtime().block_on(index.insert( - &strategy, - &Context::new(), - &new_id, - bad_vector.as_slice(), - )); - - assert!( - result.is_err(), - "insert dimension mismatch must be rejected" - ); -} diff --git a/diskann/src/ivf/test/cases/insert.rs b/diskann/src/ivf/test/cases/insert.rs deleted file mode 100644 index 2eb9dd496d..0000000000 --- a/diskann/src/ivf/test/cases/insert.rs +++ /dev/null @@ -1,197 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Insert behavior of [`crate::ivf::IvfIndex::insert`]. -//! -//! Inserting a vector must (a) store it in the provider, (b) append it to its nearest list, -//! and (c) make it findable by a subsequent search that probes that list. - -use crate::{ - ivf::test::{ - cases::{DIM, METRIC, N_LISTS, grid_index}, - harness::IvfOracleRun, - provider::{Context, Strategy, assign_nearest_for_test}, - }, - neighbor::{BackInserter, Neighbor}, - test::tokio::current_thread_runtime, -}; -use std::{num::NonZeroUsize, sync::Arc}; - -/// After inserting a point, a search probing all lists returns it as the exact nearest -/// neighbor of itself (distance 0). -#[test] -fn inserted_point_is_findable() { - let index = grid_index(); - let strategy = Strategy::new(DIM, METRIC); - let new_point = vec![3.5f32, 3.5]; - let new_id = index.provider().len() as u32; - - current_thread_runtime() - .block_on(index.insert(&strategy, &Context::new(), &new_id, new_point.as_slice())) - .unwrap(); - - assert_eq!(index.provider().len(), (new_id + 1) as usize); - - let run = IvfOracleRun::run_sync(&index, &strategy, &new_point, 1, N_LISTS).unwrap(); - assert_eq!(run.top_k.len(), 1); - let (found_id, dist) = run.top_k[0]; - assert_eq!(found_id, new_id); - assert!(dist.abs() <= f32::EPSILON, "self-distance must be zero"); -} - -/// The inserted point lands in the list whose centroid is nearest to it. -#[test] -fn inserted_point_goes_to_nearest_list() { - let index = grid_index(); - let strategy = Strategy::new(DIM, METRIC); - let new_point = vec![6.0f32, 6.0]; - let new_id = index.provider().len() as u32; - - let expected_list = assign_nearest_for_test(index.provider().centroids(), &new_point, METRIC); - - current_thread_runtime() - .block_on(index.insert(&strategy, &Context::new(), &new_id, new_point.as_slice())) - .unwrap(); - - let members = index.provider().members(expected_list); - assert!( - members.contains(&new_id), - "inserted id {new_id} should be in list {expected_list}, members={members:?}", - ); -} - -/// Inserting then probing only the assigned list still finds the point. -#[test] -fn inserted_point_found_with_single_probe() { - let index = grid_index(); - let strategy = Strategy::new(DIM, METRIC); - let new_point = vec![1.0f32, 1.0]; - let new_id = index.provider().len() as u32; - - current_thread_runtime() - .block_on(index.insert(&strategy, &Context::new(), &new_id, new_point.as_slice())) - .unwrap(); - - let context = Context::new(); - let mut buf = vec![Neighbor::::default(); 1]; - let stats = current_thread_runtime() - .block_on(index.knn_search( - NonZeroUsize::new(1).unwrap(), - 1, - &strategy, - &context, - new_point.as_slice(), - &mut BackInserter::new(buf.as_mut_slice()), - )) - .unwrap(); - - assert_eq!(stats.result_count, 1); - assert_eq!(buf[0].id, new_id); -} - -/// Many `insert` calls racing on a shared `Arc` (multi-threaded runtime) must all -/// take effect: every point gets a unique id, lands in its deterministically-assigned list, -/// and is findable at distance zero afterward. -/// -/// This exercises the insert path's interior mutability under real contention -- the atomic -/// id allocator in `set_element` and the per-list `RwLock::write` in `AppendAccessor::append` -/// (with several points deliberately routed to the same list to force write contention). -#[test] -fn concurrent_inserts_all_land_and_are_findable() { - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(4) - .enable_all() - .build() - .expect("multi-thread runtime"); - - let index = Arc::new(grid_index()); - let strategy = Arc::new(Strategy::new(DIM, METRIC)); - let base_len = index.provider().len(); - - // Distinct off-grid points spread across all four quadrants. Several share a quadrant so - // their appends contend on the same list's `RwLock`. - let new_points: Vec> = vec![ - vec![0.25, 0.25], // lower-left - vec![1.75, 0.50], // lower-left - vec![0.50, 2.25], // lower-left - vec![1.25, 5.75], // upper-left - vec![0.75, 6.25], // upper-left - vec![6.25, 1.25], // lower-right - vec![5.75, 5.75], // upper-right - vec![6.50, 6.50], // upper-right - vec![5.25, 6.75], // upper-right - vec![6.75, 5.25], // upper-right - ]; - let n = new_points.len(); - - // Deterministic routing: how many of the new points each list should receive. - let mut expected_per_list = [0usize; N_LISTS]; - for point in &new_points { - let list = assign_nearest_for_test(index.provider().centroids(), point, METRIC); - expected_per_list[list as usize] += 1; - } - - runtime.block_on({ - let index = Arc::clone(&index); - let strategy = Arc::clone(&strategy); - let points = new_points.clone(); - async move { - let mut set = tokio::task::JoinSet::new(); - for (i, point) in points.into_iter().enumerate() { - let index = Arc::clone(&index); - let strategy = Arc::clone(&strategy); - set.spawn(async move { - let external_id = (base_len + i) as u32; - index - .insert(&*strategy, &Context::new(), &external_id, point.as_slice()) - .await - .expect("concurrent insert failed"); - }); - } - while let Some(joined) = set.join_next().await { - joined.expect("insert task panicked"); - } - } - }); - - // (1) Every insert allocated a unique, sequential id and committed to some list: the set - // of newly added ids is exactly `base_len .. base_len + n`, with no losses or duplicates. - assert_eq!(index.provider().len(), base_len + n); - let mut new_ids: Vec = (0..N_LISTS as u32) - .flat_map(|list| index.provider().members(list)) - .filter(|id| (*id as usize) >= base_len) - .collect(); - new_ids.sort_unstable(); - let expected_ids: Vec = (base_len as u32..(base_len + n) as u32).collect(); - assert_eq!( - new_ids, expected_ids, - "concurrent inserts must allocate every id exactly once with no loss", - ); - - // (2) Concurrent routing matched the deterministic per-list assignment. - for (list, &expected) in expected_per_list.iter().enumerate() { - let new_in_list = index - .provider() - .members(list as u32) - .into_iter() - .filter(|id| (*id as usize) >= base_len) - .count(); - assert_eq!( - new_in_list, expected, - "list {list} received the wrong number of concurrent inserts", - ); - } - - // (3) Each inserted point is findable at distance zero -- proving its vector bytes were - // written intact (no torn or lost writes under contention). - for point in &new_points { - let run = IvfOracleRun::run_sync(&index, &strategy, point, 1, N_LISTS).unwrap(); - assert_eq!(run.top_k.len(), 1); - assert!( - run.top_k[0].1.abs() <= f32::EPSILON, - "inserted point {point:?} must be findable at distance zero", - ); - } -} diff --git a/diskann/src/ivf/test/cases/mod.rs b/diskann/src/ivf/test/cases/mod.rs deleted file mode 100644 index 3027b4d8ed..0000000000 --- a/diskann/src/ivf/test/cases/mod.rs +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Test cases for the IVF module, grouped by concern. -//! -//! * [`correctness`] -- the central invariant (IVF == brute force restricted to probed -//! lists) and its corollaries across `k`/`nprobe`. -//! * [`concurrency`] -- the sequential and concurrent fine-scan paths agree, and the -//! concurrent path actually fans out. -//! * [`insert`] -- inserted points land in the right list and become findable. -//! * [`errors`] -- accessor/strategy failures escalate out of the index API. -//! -//! All cases share the [`fixture`] below: a clean 8x8 integer grid partitioned into four -//! quadrant lists. The geometry is deliberately simple so the expected list assignments are -//! obvious by inspection, but every assertion is still validated against the exact oracle in -//! [`super::harness`] rather than hand-computed answers. - -mod concurrency; -mod correctness; -mod errors; -mod insert; - -use diskann_utils::views::Matrix; -use diskann_vector::distance::Metric; - -use crate::ivf::{IvfIndex, test::provider::Provider}; - -/// Vector/centroid dimension for the shared fixture. -pub(super) const DIM: usize = 2; - -/// Side length of the integer grid used as the dataset (`SIDE * SIDE` points). -pub(super) const SIDE: usize = 8; - -/// Number of inverted lists (quadrant centroids) in the shared fixture. -pub(super) const N_LISTS: usize = 4; - -/// Metric used throughout the IVF tests. -pub(super) const METRIC: Metric = Metric::L2; - -/// The four quadrant centroids of the 8x8 grid, one per inverted list. -fn centroids() -> Matrix { - // Quadrant centers of the [0, 7]^2 grid. - #[rustfmt::skip] - let data = vec![ - 1.5, 1.5, // list 0: lower-left - 1.5, 5.5, // list 1: upper-left - 5.5, 1.5, // list 2: lower-right - 5.5, 5.5, // list 3: upper-right - ]; - Matrix::try_from(data.into_boxed_slice(), N_LISTS, DIM).expect("centroid matrix is well-formed") -} - -/// Build the shared fixture: an [`IvfIndex`] over an 8x8 grid split into four quadrant -/// lists. -pub(super) fn grid_index() -> IvfIndex { - let vectors = crate::graph::test::synthetic::Grid::Two.data(SIDE); - debug_assert_eq!(vectors.ncols(), DIM); - debug_assert_eq!(vectors.nrows(), SIDE * SIDE); - let provider = - Provider::build(vectors, centroids(), METRIC).expect("fixture provider is well-formed"); - IvfIndex::new(provider) -} - -/// A spread of query points: grid corners, the center, and off-grid coordinates. -pub(super) fn queries() -> Vec> { - vec![ - vec![0.0, 0.0], - vec![7.0, 7.0], - vec![3.5, 3.5], - vec![2.0, 6.0], - vec![6.3, 1.1], - vec![4.0, 0.0], - ] -} diff --git a/diskann/src/ivf/test/harness.rs b/diskann/src/ivf/test/harness.rs deleted file mode 100644 index 30ac05b1ce..0000000000 --- a/diskann/src/ivf/test/harness.rs +++ /dev/null @@ -1,182 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Reusable execution harness for [`crate::ivf::IvfIndex`] tests. -//! -//! [`IvfOracleRun`] drives [`IvfIndex::knn_search`] and pairs the result with an *exact* -//! oracle. IVF is an approximate algorithm only because it restricts the search to the -//! members of the `nprobe` selected lists -- within that restriction it is exhaustive. The -//! oracle exploits this: it replays the same list selection (via the shared -//! [`nearest_lists`](super::provider::nearest_lists) core) and then brute-forces *only* the -//! union of those lists' members. The IVF result must therefore match the oracle exactly, -//! which lets the tests assert strict equality rather than a recall threshold. - -use std::{cmp::Ordering, num::NonZeroUsize}; - -use diskann_vector::PreprocessedDistanceFunction; - -use crate::{ - ANNResult, - ivf::{ - IvfIndex, SearchStats, - test::provider::{self, Id, ListId, Provider, Strategy}, - }, - neighbor::{BackInserter, Neighbor}, - test::tokio::current_thread_runtime, - utils::VectorRepr, -}; - -/// Result of one [`IvfIndex::knn_search`] run alongside the exact oracle answer. -#[derive(Debug, Clone)] -pub(crate) struct IvfOracleRun { - /// `(id, distance)` pairs returned by the search, re-sorted by `(distance asc, id asc)` - /// so equality checks are deterministic on ties. - pub top_k: Vec<(Id, f32)>, - /// `top_k.iter().map(|(_, d)| d).collect()` -- convenient for distance-multiset checks. - pub top_k_distances: Vec, - /// Statistics returned by `knn_search`. - pub stats: SearchStats, - /// The exact expected top-`k`: brute force restricted to the members of the selected - /// lists, in `(distance asc, id asc)` order. - pub ground_truth: Vec<(Id, f32)>, - /// The lists the oracle (and thus the search) probed for this query. - pub probed_lists: Vec, - /// Total number of members across `probed_lists` -- the exact number of fine-distance - /// computations the search must perform. - pub probed_members: usize, -} - -impl IvfOracleRun { - /// Run [`IvfIndex::knn_search`] once, blocking on a fresh single-threaded runtime. - pub fn run_sync( - index: &IvfIndex, - strategy: &Strategy, - query: &[f32], - k: usize, - nprobe: usize, - ) -> ANNResult { - current_thread_runtime().block_on(Self::run(index, strategy, query, k, nprobe)) - } - - /// Async variant of [`IvfOracleRun::run_sync`], for tests that already own a runtime - /// (e.g. `#[tokio::test]`) or drive searches concurrently across tasks. - pub async fn run( - index: &IvfIndex, - strategy: &Strategy, - query: &[f32], - k: usize, - nprobe: usize, - ) -> ANNResult { - let context = provider::Context::new(); - let mut buf = vec![Neighbor::::default(); k]; - - let stats = index - .knn_search( - NonZeroUsize::new(k).expect("ivf::test::harness requires k > 0"), - nprobe, - strategy, - &context, - query, - &mut BackInserter::new(buf.as_mut_slice()), - ) - .await?; - - let mut top_k: Vec> = buf - .iter() - .copied() - .take(stats.result_count as usize) - .collect(); - sort_neighbors(&mut top_k); - let top_k_distances = top_k.iter().map(|n| n.distance).collect(); - - let (ground_truth, probed_lists, probed_members) = - restricted_topk(index.provider(), query, k, nprobe); - - Ok(Self { - top_k: top_k.into_iter().map(Neighbor::as_tuple).collect(), - top_k_distances, - stats, - ground_truth: ground_truth.into_iter().map(Neighbor::as_tuple).collect(), - probed_lists, - probed_members, - }) - } -} - -/// Exact IVF oracle: select the same `nprobe` lists the search would, then brute-force the -/// union of their members. -/// -/// Returns `(top_k, probed_lists, probed_member_count)`. -pub(crate) fn restricted_topk( - provider: &Provider, - query: &[f32], - k: usize, - nprobe: usize, -) -> (Vec>, Vec, usize) { - let metric = provider.metric(); - let probed_lists = provider::nearest_lists(provider.centroids(), query, metric, nprobe); - - let computer = f32::query_distance(query, metric); - let mut neighbors: Vec> = Vec::new(); - let mut probed_members = 0usize; - for &list in &probed_lists { - for (id, vector) in provider.list_entries(list) { - probed_members += 1; - neighbors.push(Neighbor::new(id, computer.evaluate_similarity(&vector[..]))); - } - } - - sort_neighbors(&mut neighbors); - neighbors.truncate(k); - (neighbors, probed_lists, probed_members) -} - -/// Brute-force top-`k` over the *entire* provider -- used by tests that probe every list -/// (`nprobe == n_lists`) and therefore expect the global answer. -pub(crate) fn global_topk(provider: &Provider, query: &[f32], k: usize) -> Vec> { - let metric = provider.metric(); - let computer = f32::query_distance(query, metric); - let mut neighbors: Vec> = Vec::new(); - for list in 0..provider.n_lists() as ListId { - for (id, vector) in provider.list_entries(list) { - neighbors.push(Neighbor::new(id, computer.evaluate_similarity(&vector[..]))); - } - } - sort_neighbors(&mut neighbors); - neighbors.truncate(k); - neighbors -} - -/// Sort `(distance asc, id asc)` with NaN treated as equal. -fn sort_neighbors(neighbors: &mut [Neighbor]) { - neighbors.sort_by(|a, b| { - a.distance - .partial_cmp(&b.distance) - .unwrap_or(Ordering::Equal) - .then(a.id.cmp(&b.id)) - }); -} - -/// Assert two `(id, distance)` lists agree on the *distance multiset*. -/// -/// The priority queue may break exact-distance ties differently from the oracle's -/// `(distance, id)` ordering, so element-wise id equality is too strict. Comparing the -/// sorted distance sequences (already sorted by the harness) plus the length is the right -/// invariant for a correct top-`k`. -pub(crate) fn assert_same_distances(actual: &[(Id, f32)], expected: &[(Id, f32)]) { - let a: Vec = actual.iter().map(|(_, d)| *d).collect(); - let e: Vec = expected.iter().map(|(_, d)| *d).collect(); - assert_eq!( - a.len(), - e.len(), - "result count mismatch: got {a:?}, expected {e:?}" - ); - for (i, (ad, ed)) in a.iter().zip(e.iter()).enumerate() { - assert!( - (ad - ed).abs() <= f32::EPSILON, - "distance mismatch at rank {i}: got {ad}, expected {ed} (got {a:?}, expected {e:?})" - ); - } -} diff --git a/diskann/src/ivf/test/mod.rs b/diskann/src/ivf/test/mod.rs deleted file mode 100644 index a19e745212..0000000000 --- a/diskann/src/ivf/test/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Test fixtures and helpers for the IVF module. -pub(crate) mod harness; -pub(crate) mod provider; - -mod cases; diff --git a/diskann/src/ivf/test/provider.rs b/diskann/src/ivf/test/provider.rs deleted file mode 100644 index 983bc5befa..0000000000 --- a/diskann/src/ivf/test/provider.rs +++ /dev/null @@ -1,792 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Self-contained, **production-shaped** example backend for the IVF module. -//! -//! This file is intended to double as a worked example of how to implement the three IVF -//! accessor traits ([`ListAccessor`], [`SearchAccessor`], [`InsertAccessor`]) and a -//! [`DataProvider`] for them, with concurrency in mind. -//! -//! # Roles -//! -//! A real deployment may split storage across several collaborators (see -//! `rfcs/00000-ivf-index.md`). For a compact, testable example we fold all of them into one -//! [`Provider`] that plays three roles at once: -//! -//! * **DataProvider** -- owns the external<->internal id mapping and allocates internal ids. -//! * **Coarse quantizer** -- owns the immutable centroids (`list_id -> centroid`). -//! * **Inverted-list store** -- owns the per-list membership. This is a list-owned -//! **IVF-Flat** layout: each list stores its members' full vectors *contiguously* in a -//! single flat buffer, in the same order as the list's ids (see [`ListStore`]). The fine -//! scan therefore streams a list's vectors sequentially out of one cache-friendly block -//! rather than chasing per-id pointers into a global table. -//! -//! Because the vectors live with their list, there is no global `id -> vector` table: a -//! vector is materialized only when [`InsertAccessor::append`] commits it to a list. The -//! provider's [`provider::SetElement::set_element`] therefore only *allocates* an internal -//! id; the append step does the contiguous store. -//! -//! # Concurrency model -//! -//! The centroids are immutable and shared via [`Arc`], so the fine-scan accessor can hand -//! list scans to independent tasks (see [`ScanAccessor`]'s concurrent path). The mutable -//! state -- each list's contiguous store -- is guarded by a per-list [`RwLock`] so inserts -//! can append to one list while searches read others. This mirrors how a production backend -//! would keep the read (search) path lock-light while still supporting concurrent append. - -use std::{ - cmp::Ordering, - fmt::{self, Debug}, - future::Future, - sync::{ - Arc, RwLock, - atomic::{AtomicU32, Ordering as AtomicOrdering}, - }, -}; - -use diskann_utils::{future::SendFuture, views::Matrix}; -use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; -use thiserror::Error; - -use crate::graph::SearchOutputBuffer; -use crate::{ - ANNError, always_escalate, - internal::counter::{Counter, LocalCounter}, - ivf::{InsertAccessor, InsertStrategy, ListAccessor, SearchAccessor, SearchStrategy}, - provider::{self, ExecutionContext, HasId, NoopGuard}, - utils::VectorRepr, -}; - -/// The id type used for inverted lists throughout this example. -pub type ListId = u32; - -/// The internal/external id type used by this provider. -pub type Id = u32; - -////////////////////////// -// Shared selection core // -////////////////////////// - -/// Return the (up to) `nprobe` lists closest to `query` as `(list_id, distance)` pairs, -/// ordered by `(distance asc, list_id asc)`. -/// -/// This is the single source of truth for coarse list selection. Both [`CentroidAccessor`] -/// (the production path, via this function) and the test harness's oracle (via -/// [`nearest_lists`]) build on it, which guarantees they agree on exactly which lists a -/// query probes -- essential for an *exact* approximate-search oracle. -fn scored_lists( - centroids: &Matrix, - query: &[f32], - metric: Metric, - nprobe: usize, -) -> Vec<(ListId, f32)> { - let computer = f32::query_distance(query, metric); - let mut scored: Vec<(ListId, f32)> = centroids - .row_iter() - .enumerate() - .map(|(i, c)| (i as ListId, computer.evaluate_similarity(c))) - .collect(); - scored.sort_by(|a, b| cmp_dist_id(*a, *b)); - scored.truncate(nprobe); - scored -} - -/// Return the (up to) `nprobe` list ids closest to `query`, ordered by -/// `(distance asc, list_id asc)`. Used by the test harness oracle. -pub fn nearest_lists( - centroids: &Matrix, - query: &[f32], - metric: Metric, - nprobe: usize, -) -> Vec { - scored_lists(centroids, query, metric, nprobe) - .into_iter() - .map(|(id, _)| id) - .collect() -} - -/// Order by `(distance asc, id asc)` with NaN treated as equal. -fn cmp_dist_id(a: (ListId, f32), b: (ListId, f32)) -> Ordering { - a.1.partial_cmp(&b.1) - .unwrap_or(Ordering::Equal) - .then(a.0.cmp(&b.0)) -} - -/// Index of the single nearest list for `vector`. -fn assign_nearest(centroids: &Matrix, vector: &[f32], metric: Metric) -> ListId { - let computer = f32::query_distance(vector, metric); - centroids - .row_iter() - .enumerate() - .map(|(i, c)| (i as ListId, computer.evaluate_similarity(c))) - .min_by(|a, b| cmp_dist_id(*a, *b)) - .map(|(id, _)| id) - .expect("provider always has at least one centroid") -} - -/// Public re-export of [`assign_nearest`] for tests that want to predict an insert's target -/// list independently of the index. -pub fn assign_nearest_for_test(centroids: &Matrix, vector: &[f32], metric: Metric) -> ListId { - assign_nearest(centroids, vector, metric) -} - -////////////// -// Provider // -////////////// - -/// Error conditions for [`Provider::build`]. -#[derive(Debug, Error)] -pub enum ProviderError { - #[error("ivf::test::Provider needs at least one vector")] - EmptyVectors, - #[error("ivf::test::Provider needs at least one centroid")] - EmptyCentroids, - #[error("ivf::test::Provider vectors must have non-zero dimension")] - ZeroDimension, - #[error("ivf::test::Provider dimension mismatch: vectors={vectors}, centroids={centroids}")] - DimMismatch { vectors: usize, centroids: usize }, -} - -impl From for ANNError { - #[track_caller] - fn from(err: ProviderError) -> ANNError { - ANNError::opaque(err) - } -} - -/// One inverted list's contiguous storage. -/// -/// `ids[j]` is the internal id of the `j`-th member, and its full vector occupies -/// `data[j * dim .. (j + 1) * dim]`. The two vectors grow together in [`Self::push`], so a -/// member's id and its vector always share the same index -- the property the fine scan -/// relies on to pair a streamed vector back with its id. -#[derive(Debug, Default)] -struct ListStore { - ids: Vec, - data: Vec, -} - -impl ListStore { - /// Append `(id, vector)`, keeping ids and the contiguous `data` buffer in lockstep. - fn push(&mut self, id: Id, vector: &[f32]) { - self.ids.push(id); - self.data.extend_from_slice(vector); - } - - /// Iterate `(id, vector)` pairs in storage order, slicing the flat buffer by `dim`. - fn iter(&self, dim: usize) -> impl Iterator { - self.ids.iter().copied().zip(self.data.chunks_exact(dim)) - } -} - -/// In-memory IVF backend. -#[derive(Debug)] -pub struct Provider { - /// Immutable centroids, one row per inverted list. - centroids: Arc>, - /// Per-list contiguous storage (ids + full vectors). Fixed list count; each list grows - /// on insert. - lists: Arc>>, - /// Allocates the next internal id. Identity external<->internal ids means this is just a - /// monotonically increasing slot counter. - next_id: Arc, - /// Vector dimension. - dim: usize, - /// Distance metric used for both coarse selection and fine scoring (RFC D5). - metric: Metric, - /// Number of query->centroid distance computations performed. - centroid_cmps: Counter, - /// Number of query->member distance computations performed. - member_cmps: Counter, -} - -impl Provider { - /// Build a provider from `vectors` and `centroids`, assigning each vector to its - /// nearest centroid's list (the "bring-your-own-centroids" build of RFC F1/D6). - /// - /// # Errors - /// - /// Returns [`ProviderError`] if either matrix is empty, the dimension is zero, or the - /// vector and centroid dimensions disagree. - pub fn build( - vectors: Matrix, - centroids: Matrix, - metric: Metric, - ) -> Result { - if vectors.nrows() == 0 { - return Err(ProviderError::EmptyVectors); - } - if centroids.nrows() == 0 { - return Err(ProviderError::EmptyCentroids); - } - let dim = vectors.ncols(); - if dim == 0 { - return Err(ProviderError::ZeroDimension); - } - if centroids.ncols() != dim { - return Err(ProviderError::DimMismatch { - vectors: dim, - centroids: centroids.ncols(), - }); - } - - let n_lists = centroids.nrows(); - let lists: Vec> = (0..n_lists) - .map(|_| RwLock::new(ListStore::default())) - .collect(); - for (i, vector) in vectors.row_iter().enumerate() { - let list = assign_nearest(¢roids, vector, metric); - lists[list as usize] - .write() - .expect("list lock poisoned during build") - .push(i as Id, vector); - } - - let next_id = vectors.nrows() as u32; - - Ok(Self { - centroids: Arc::new(centroids), - lists: Arc::new(lists), - next_id: Arc::new(AtomicU32::new(next_id)), - dim, - metric, - centroid_cmps: Counter::new(), - member_cmps: Counter::new(), - }) - } - - /// Vector dimension. - pub fn dim(&self) -> usize { - self.dim - } - - /// Distance metric. - pub fn metric(&self) -> Metric { - self.metric - } - - /// Number of inverted lists (== number of centroids). - pub fn n_lists(&self) -> usize { - self.centroids.nrows() - } - - /// Number of vectors currently stored (== number of allocated internal ids). - pub fn len(&self) -> usize { - self.next_id.load(AtomicOrdering::Relaxed) as usize - } - - /// Borrow the centroids (used by the harness oracle). - pub fn centroids(&self) -> &Matrix { - &self.centroids - } - - /// Snapshot the ids in `list`, in storage order (used by the harness oracle and tests). - pub fn members(&self, list: ListId) -> Vec { - self.lists[list as usize] - .read() - .expect("list lock poisoned") - .ids - .clone() - } - - /// Snapshot the `(id, vector)` entries in `list`, in storage order (used by the harness - /// oracle). The vectors are copied out of the list's contiguous buffer. - pub fn list_entries(&self, list: ListId) -> Vec<(Id, Vec)> { - let store = self.lists[list as usize] - .read() - .expect("list lock poisoned"); - store - .iter(self.dim) - .map(|(id, v)| (id, v.to_vec())) - .collect() - } - - /// Snapshot of the per-provider distance-computation counters. - pub fn metrics(&self) -> ProviderMetrics { - ProviderMetrics { - centroid_cmps: self.centroid_cmps.value(), - member_cmps: self.member_cmps.value(), - } - } -} - -/// Distance-computation counters tracked by [`Provider`]. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderMetrics { - /// Total query->centroid distance computations (coarse step). - pub centroid_cmps: usize, - /// Total query->member distance computations (fine step). - pub member_cmps: usize, -} - -///////////// -// Context // -///////////// - -/// Per-operation execution context that records how many tasks were spawned through it, -/// so concurrency tests can assert the fine-scan path actually fanned out. -#[derive(Debug, Clone)] -pub struct Context(Arc); - -impl Context { - pub fn new() -> Self { - Self(Arc::new(Counter::new())) - } - - /// Number of spawns routed through [`ExecutionContext::wrap_spawn`]. - pub fn spawns(&self) -> usize { - self.0.value() - } -} - -impl Default for Context { - fn default() -> Self { - Self::new() - } -} - -impl ExecutionContext for Context { - fn wrap_spawn(&self, f: F) -> impl Future + Send + 'static - where - F: Future + Send + 'static, - { - self.0.increment(); - f - } -} - -//////////// -// Errors // -//////////// - -/// Critical id-validation error: a requested id is out of range. -#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] -#[error("ivf::test::Provider has no id {0}")] -pub struct InvalidId(pub u32); - -always_escalate!(InvalidId); - -impl From for ANNError { - #[track_caller] - fn from(err: InvalidId) -> ANNError { - ANNError::opaque(err) - } -} - -/// Critical error injected by [`Strategy::failing_on_list`]: scanning a specific list fails. -/// -/// Used to verify that a fine-scan failure escalates out of -/// [`crate::ivf::IvfIndex::knn_search`]. -#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] -#[error("ivf::test::Provider injected scan failure on list {0}")] -pub struct ScanFailure(pub ListId); - -always_escalate!(ScanFailure); - -impl From for ANNError { - #[track_caller] - fn from(err: ScanFailure) -> ANNError { - ANNError::opaque(err) - } -} - -/// Dimension-mismatch error from strategy construction. -#[derive(Debug, Clone, Error, PartialEq, Eq)] -#[error("dimension mismatch: strategy expects {expected}, got {actual}")] -pub struct StrategyError { - pub expected: usize, - pub actual: usize, -} - -impl From for ANNError { - #[track_caller] - fn from(err: StrategyError) -> ANNError { - ANNError::opaque(err) - } -} - -////////////////// -// DataProvider // -////////////////// - -impl provider::DataProvider for Provider { - type Context = Context; - type InternalId = Id; - type ExternalId = Id; - type Error = InvalidId; - type Guard = NoopGuard; - - fn to_internal_id(&self, _ctx: &Context, gid: &Id) -> Result { - if (*gid as usize) < self.len() { - Ok(*gid) - } else { - Err(InvalidId(*gid)) - } - } - - fn to_external_id(&self, _ctx: &Context, id: Id) -> Result { - if (id as usize) < self.len() { - Ok(id) - } else { - Err(InvalidId(id)) - } - } -} - -impl provider::SetElement<&[f32]> for Provider { - type SetError = InvalidId; - - async fn set_element( - &self, - _context: &Context, - _id: &Id, - _element: &[f32], - ) -> Result { - // List-owned storage: the vector is *not* stored here -- it is committed to its - // list's contiguous buffer by `InsertAccessor::append`, once the target list is - // known. `set_element` therefore only allocates the next internal id. A production - // provider would honor the supplied external id and return a guard that rolls back - // the allocation on drop-without-complete. - let new_id = self.next_id.fetch_add(1, AtomicOrdering::Relaxed); - Ok(NoopGuard::new(new_id)) - } -} - -////////////////////////////// -// ListAccessor (coarse step) // -////////////////////////////// - -/// Coarse accessor: selects the `nprobe` lists nearest to a bound query. -/// -/// Holds an [`Arc`] clone of the immutable centroids and a preprocessed query computer. The -/// `LocalCounter` flushes the coarse comparison count back to the provider on drop. -pub struct CentroidAccessor<'a> { - centroids: Arc>, - metric: Metric, - nprobe_query: Vec, - cmps: LocalCounter<'a>, -} - -impl Debug for CentroidAccessor<'_> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CentroidAccessor") - .field("n_lists", &self.centroids.nrows()) - .finish_non_exhaustive() - } -} - -impl ListAccessor for CentroidAccessor<'_> { - type Id = ListId; - type Error = InvalidId; - - fn select_lists( - &mut self, - nprobe: usize, - output: &mut B, - ) -> impl SendFuture> - where - B: SearchOutputBuffer + Send + ?Sized, - { - async move { - self.cmps.increment_by(self.centroids.nrows()); - for (id, dist) in scored_lists(&self.centroids, &self.nprobe_query, self.metric, nprobe) - { - if output.push(id, dist).is_full() { - break; - } - } - Ok(()) - } - } -} - -///////////////////////////////// -// SearchAccessor (fine step) // -///////////////////////////////// - -/// Whether the fine scan runs sequentially or fans lists out across spawned tasks. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ScanMode { - /// Scan lists one at a time on the current task. - Sequential, - /// Spawn one task per list via [`ExecutionContext::wrap_spawn`] and merge results. - Concurrent, -} - -/// Fine accessor: scans the members of a set of lists, scoring each against a bound query. -/// -/// All data needed for a scan is held behind [`Arc`] so the concurrent path can move clones -/// into `'static` spawned tasks. Each spawned task produces its own `(id, distance)` vector; -/// the driver merges those into the single `f` callback, preserving the trait's -/// single-threaded-callback contract while doing the distance work concurrently. -pub struct ScanAccessor<'a> { - lists: Arc>>, - computer: Arc<::QueryDistance>, - dim: usize, - context: Context, - mode: ScanMode, - fail_on_list: Option, - cmps: LocalCounter<'a>, -} - -impl Debug for ScanAccessor<'_> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ScanAccessor") - .field("mode", &self.mode) - .field("fail_on_list", &self.fail_on_list) - .finish_non_exhaustive() - } -} - -impl HasId for ScanAccessor<'_> { - type Id = Id; -} - -impl ScanAccessor<'_> { - /// Score every member of `list`, returning `(id, distance)` pairs. Shared by both the - /// sequential driver and the spawned-task driver. - /// - /// Streams the list's contiguous vector buffer in storage order, pairing each vector - /// with the id at the matching index. - fn score_list( - lists: &Arc>>, - computer: &::QueryDistance, - dim: usize, - list: ListId, - ) -> Vec<(Id, f32)> { - let store = lists[list as usize].read().expect("list lock poisoned"); - store - .iter(dim) - .map(|(id, v)| (id, computer.evaluate_similarity(v))) - .collect() - } -} - -impl SearchAccessor for ScanAccessor<'_> { - type ListId = ListId; - type Error = ScanFailure; - - fn scan_lists( - &mut self, - lists: Itr, - mut f: F, - ) -> impl SendFuture> - where - Itr: Iterator + Send, - F: Send + FnMut(Self::Id, f32), - { - async move { - match self.mode { - ScanMode::Sequential => { - for list in lists { - if self.fail_on_list == Some(list) { - return Err(ScanFailure(list)); - } - for (id, dist) in - Self::score_list(&self.lists, &self.computer, self.dim, list) - { - self.cmps.increment(); - f(id, dist); - } - } - } - ScanMode::Concurrent => { - let mut set = tokio::task::JoinSet::new(); - for list in lists { - if self.fail_on_list == Some(list) { - set.shutdown().await; - return Err(ScanFailure(list)); - } - let lists = Arc::clone(&self.lists); - let computer = Arc::clone(&self.computer); - let dim = self.dim; - let task = async move { Self::score_list(&lists, &computer, dim, list) }; - set.spawn(self.context.wrap_spawn(task)); - } - while let Some(joined) = set.join_next().await { - let scored = joined.expect("scan task panicked"); - for (id, dist) in scored { - self.cmps.increment(); - f(id, dist); - } - } - } - } - Ok(()) - } - } -} - -///////////////////////////////// -// InsertAccessor (append step) // -///////////////////////////////// - -/// Append accessor: commits `(id, vector)` to a chosen list's contiguous store. -/// -/// This example is **IVF-Flat** with *list-owned* storage: the full vector is appended to -/// the target list's contiguous buffer here (in lockstep with its id), and read back from -/// that same buffer during the fine scan. There is no global `id -> vector` table. -/// -/// An **IVF-PQ** backend would slot in here unchanged in shape: quantize `vector` and push -/// the code instead of the raw floats, e.g. `store.push_code(id, codebook.encode(vector))`. -#[derive(Debug)] -pub struct AppendAccessor { - lists: Arc>>, - dim: usize, -} - -impl HasId for AppendAccessor { - type Id = Id; -} - -impl<'v> InsertAccessor<&'v [f32]> for AppendAccessor { - type ListId = ListId; - type Error = InvalidId; - - fn append( - &mut self, - list: Self::ListId, - id: Self::Id, - vector: &'v [f32], - ) -> impl SendFuture> { - async move { - debug_assert_eq!( - vector.len(), - self.dim, - "append received a vector of the wrong dimension", - ); - self.lists[list as usize] - .write() - .expect("list lock poisoned") - .push(id, vector); - Ok(()) - } - } -} - -////////////// -// Strategy // -////////////// - -/// Per-call factory wiring centroids (coarse) to scanning/appending (fine) for `&[f32]` -/// queries and vectors. Validates query dimension and can inject a scan failure. -#[derive(Debug, Clone)] -pub struct Strategy { - dim: usize, - metric: Metric, - mode: ScanMode, - fail_on_list: Option, -} - -impl Strategy { - /// Construct a strategy expecting vectors of dimension `dim`, scanning sequentially. - pub fn new(dim: usize, metric: Metric) -> Self { - Self { - dim, - metric, - mode: ScanMode::Sequential, - fail_on_list: None, - } - } - - /// Select the fine-scan mode (sequential vs. concurrent). - pub fn with_mode(mut self, mode: ScanMode) -> Self { - self.mode = mode; - self - } - - /// Make the fine scan fail when it reaches `list`. - pub fn failing_on_list(mut self, list: ListId) -> Self { - self.fail_on_list = Some(list); - self - } - - fn check_dim(&self, actual: usize) -> Result<(), StrategyError> { - if actual != self.dim { - Err(StrategyError { - expected: self.dim, - actual, - }) - } else { - Ok(()) - } - } -} - -impl<'a> SearchStrategy<'a, Provider, &'a [f32]> for Strategy { - type ListId = ListId; - type SearchAccessor = ScanAccessor<'a>; - type ListAccessor = CentroidAccessor<'a>; - type Error = StrategyError; - - fn search_accessor( - &'a self, - provider: &'a Provider, - context: &'a Context, - query: &'a [f32], - ) -> Result { - self.check_dim(query.len())?; - Ok(ScanAccessor { - lists: Arc::clone(&provider.lists), - computer: Arc::new(f32::query_distance(query, self.metric)), - dim: provider.dim, - context: context.clone(), - mode: self.mode, - fail_on_list: self.fail_on_list, - cmps: provider.member_cmps.local(), - }) - } - - fn list_accessor( - &'a self, - provider: &'a Provider, - _context: &'a Context, - query: &'a [f32], - ) -> Result { - self.check_dim(query.len())?; - Ok(CentroidAccessor { - centroids: Arc::clone(&provider.centroids), - metric: self.metric, - nprobe_query: query.to_vec(), - cmps: provider.centroid_cmps.local(), - }) - } -} - -impl<'a> InsertStrategy<'a, Provider, &'a [f32]> for Strategy { - type ListId = ListId; - type InsertAccessor = AppendAccessor; - type ListAccessor = CentroidAccessor<'a>; - type Error = StrategyError; - - fn insert_accessor( - &'a self, - provider: &'a Provider, - _context: &'a Context, - ) -> Result { - // The vector is validated by `list_accessor` (the coarse step runs first), and is - // stored contiguously by `append`, so the append accessor only needs the list - // storage handle and the dimension (for its lockstep buffer layout). - Ok(AppendAccessor { - lists: Arc::clone(&provider.lists), - dim: provider.dim, - }) - } - - fn list_accessor( - &'a self, - provider: &'a Provider, - _context: &'a Context, - vector: &'a [f32], - ) -> Result { - self.check_dim(vector.len())?; - Ok(CentroidAccessor { - centroids: Arc::clone(&provider.centroids), - metric: self.metric, - nprobe_query: vector.to_vec(), - cmps: provider.centroid_cmps.local(), - }) - } -} diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index c1baa45d9f..b499b137f3 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -110,8 +110,6 @@ coordinates its provider, centroid index, and inverted-list store. the graph module's strategy/accessor style. 9. Leave locking, reader visibility, durability, rollback, and recovery to each concrete accessor implementation. -10. Preserve the existing fixed-partition IVF API while the dynamic algorithm - is developed. ### Non-Goals @@ -189,14 +187,13 @@ implementation. ```text diskann/src/ivf/ mod.rs - glue.rs existing fixed-partition traits - index.rs existing fixed-partition orchestration dynamic.rs contracts proposed by this RFC + index.rs DynamicIvfIndex initialize/insert orchestration online/ future GraphIVF algorithm and scratch state ``` -The dynamic names remain under `ivf::dynamic` while experimental, avoiding -collisions with the current `ivf::SearchAccessor` and `ivf::SearchStrategy`. +The initial fixed-partition prototype (`glue.rs`, `index.rs`, and its test +provider) is removed; the contracts in `ivf::dynamic` replace it. ### Responsibilities @@ -649,7 +646,7 @@ and search parameters. Durable metadata belongs to concrete provider crates. ### Phase 0: Compile-Only Contracts - Add `ivf::dynamic` with the traits and data types in this RFC. -- Keep the fixed-partition API and tests unchanged. +- Remove the fixed-partition prototype API and its tests. - Compile, lint, format, and build rustdoc for the skeleton. ### Phase 1: Unified In-Memory Provider From 83c451ac2c03d6b1118726eeb1890bc5693bb75d Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Wed, 30 Sep 2026 17:16:51 -0400 Subject: [PATCH 07/33] lots to change but we're getting somewhere --- diskann/src/ivf/colocation.rs | 56 +++ diskann/src/ivf/dynamic.rs | 240 +++------ diskann/src/ivf/grouped.rs | 183 +++++++ diskann/src/ivf/index.rs | 807 ++++++++++++------------------ diskann/src/ivf/mod.rs | 7 + diskann/src/ivf/online/batch.rs | 108 ++++ diskann/src/ivf/online/gather.rs | 111 ++++ diskann/src/ivf/online/kernels.rs | 132 +++++ diskann/src/ivf/online/mod.rs | 64 +++ diskann/src/ivf/online/split.rs | 346 +++++++++++++ diskann/src/ivf/test/mod.rs | 8 + diskann/src/ivf/test/provider.rs | 494 ++++++++++++++++++ diskann/src/ivf/update.rs | 533 ++++++++++++++++++++ rfcs/01187-incremental-ivf.md | 377 ++++++++++---- 14 files changed, 2714 insertions(+), 752 deletions(-) create mode 100644 diskann/src/ivf/colocation.rs create mode 100644 diskann/src/ivf/grouped.rs create mode 100644 diskann/src/ivf/online/batch.rs create mode 100644 diskann/src/ivf/online/gather.rs create mode 100644 diskann/src/ivf/online/kernels.rs create mode 100644 diskann/src/ivf/online/mod.rs create mode 100644 diskann/src/ivf/online/split.rs create mode 100644 diskann/src/ivf/test/mod.rs create mode 100644 diskann/src/ivf/test/provider.rs create mode 100644 diskann/src/ivf/update.rs diff --git a/diskann/src/ivf/colocation.rs b/diskann/src/ivf/colocation.rs new file mode 100644 index 0000000000..9f42137f0f --- /dev/null +++ b/diskann/src/ivf/colocation.rs @@ -0,0 +1,56 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Co-location groups of centroids for locality-preserving list storage. +//! +//! Not yet integrated with [`super::dynamic::MaintenanceAccessor`] or the updates in +//! [`super::update`]. + +use crate::utils::VectorId; + +/// An authoritative grouping of live centroids for locality-preserving storage. +/// +/// Every live centroid belongs to exactly one co-location group. Implementations +/// may use an identity grouping, where each group contains one centroid. Group +/// ids remain stable until retired and are independent of centroid/list ids. +pub trait CoLocationSet: Send + Sync { + /// Stable logical centroid/list id. + type ListId: VectorId; + + /// Stable logical co-location group id. + type CoLocationGroupId: VectorId; + + /// Number of live co-location groups. + fn len(&self) -> usize; + + /// Whether there are no live co-location groups. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Borrow the centroid ids in a live co-location group. + fn group(&self, id: Self::CoLocationGroupId) -> Option<&[Self::ListId]>; + + /// Return the live co-location group containing a centroid. + fn group_for(&self, id: Self::ListId) -> Option; +} + +/// A new co-location group to install. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CoLocationGroup { + /// Fresh stable group id. Installed ids must never be reused after retirement. + pub id: CoLocationGroupId, + /// Live centroids that should be stored together. + pub centroids: Vec, +} + +/// Changes to the authoritative co-location set. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CoLocationDelta { + /// Co-location groups made live by this update. + pub insert: Vec>, + /// Live co-location groups retired by this update. + pub retire: Vec, +} diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index a133ed40bd..b84da71b68 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -13,7 +13,7 @@ use std::fmt::Debug; -use diskann_utils::future::SendFuture; +use diskann_utils::{future::SendFuture, views::MutMatrixView}; use crate::{ error::{StandardError, ToRanked}, @@ -66,6 +66,12 @@ pub struct ScanStats { pub lists_scanned: u32, } +/// A data provider whose points are partitioned into stable logical lists. +pub trait ListProvider: DataProvider { + /// Stable id shared by a centroid and its inverted list. + type ListId: VectorId; +} + /// An in-memory index over the live centroids. /// /// The centroid catalog is authoritative. Approximate implementations, such as @@ -101,33 +107,6 @@ pub trait CentroidIndex: Send + Sync { ) -> impl SendFuture, Self::Error>>; } -/// An authoritative grouping of live centroids for locality-preserving storage. -/// -/// Every live centroid belongs to exactly one co-location group. Implementations -/// may use an identity grouping, where each group contains one centroid. Group -/// ids remain stable until retired and are independent of centroid/list ids. -pub trait CoLocationSet: Send + Sync { - /// Stable logical centroid/list id. - type ListId: VectorId; - - /// Stable logical co-location group id. - type CoLocationGroupId: VectorId; - - /// Number of live co-location groups. - fn len(&self) -> usize; - - /// Whether there are no live co-location groups. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Borrow the centroid ids in a live co-location group. - fn group(&self, id: Self::CoLocationGroupId) -> Option<&[Self::ListId]>; - - /// Return the live co-location group containing a centroid. - fn group_for(&self, id: Self::ListId) -> Option; -} - /// Query-bound access to centroid selection and inverted-list scanning. /// /// This is the dynamic IVF search algorithm's primary extension point, in the @@ -135,7 +114,8 @@ pub trait CoLocationSet: Send + Sync { /// batch reads, coalesce blob requests, prefetch, decode quantized payloads, or /// fan work out across tasks. pub trait SearchAccessor: HasId + Send + Sync { - /// Stable logical centroid/list id. + /// Stable logical centroid/list id, fixed to [`ListProvider::ListId`] by the + /// strategy. type ListId: VectorId; /// Errors from list selection or scanning. @@ -160,10 +140,10 @@ pub trait SearchAccessor: HasId + Send + Sync { /// Factory for one dynamic IVF search accessor. pub trait SearchStrategy<'a, Provider, T>: Send + Sync where - Provider: DataProvider, + Provider: ListProvider, { /// Query-bound accessor used for both coarse and fine search. - type SearchAccessor: SearchAccessor; + type SearchAccessor: SearchAccessor; /// Error constructing the accessor. type Error: StandardError; @@ -179,88 +159,22 @@ where /// Size metadata for one logical inverted list. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ListMetadata { - /// Stable logical list identifier. - pub id: ListId, +pub struct ListMetadata { /// Number of live point ids in the list. pub len: usize, } -/// A new centroid to install as part of a partition update. -#[derive(Debug, Clone, PartialEq)] -pub struct CentroidRecord { - /// Fresh stable id. Installed ids must never be reused after retirement. - pub id: ListId, - /// Full-precision centroid used by maintenance and exact fallback. - pub vector: Box<[f32]>, -} - -/// Changes to the authoritative live-centroid catalog. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct CentroidDelta { - /// Centroids made live by this update. - pub insert: Vec>, - /// Live centroids retired by this update. - pub retire: Vec, -} - -/// A new co-location group to install as part of a partition update. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CoLocationGroup { - /// Fresh stable group id. Installed ids must never be reused after retirement. - pub id: CoLocationGroupId, - /// Live centroids that should be stored together. - pub centroids: Vec, -} - -/// Changes to the authoritative co-location set. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct CoLocationDelta { - /// Co-location groups made live by this update. - pub insert: Vec>, - /// Live co-location groups retired by this update. - pub retire: Vec, -} - -/// One point membership change in a partition mutation. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct PointMove { - /// Internal point id being inserted, deleted, or reassigned. - pub id: Id, - /// Previous list, or `None` for a newly inserted point. - pub from: Option, - /// New list, or `None` for a deleted point. - pub to: Option, -} -/// constrcutors for this too. - -/// A complete logical partition change. -/// -/// The maintenance accessor applies this update across point identity, -/// canonical vectors, list membership, reverse assignments, scan payloads, and -/// centroid liveness, and co-location membership according to its -/// provider-specific consistency contract. -#[derive(Debug, Clone, Default, PartialEq)] -pub struct PartitionUpdate { - /// Centroids inserted and retired by this mutation. - pub centroids: CentroidDelta, - /// Co-location groups inserted and retired by this mutation. - pub co_locations: CoLocationDelta, - /// Point membership changes, including inserts and deletes. - pub point_moves: Vec>, -} -/// constructor enforces safety + correctness. - /// Operation-scoped reads and writes needed by split/dissolve maintenance. /// /// The accessor is the consistency boundary for one mutation. It must present a /// unified view across point data, centroids, reverse assignments, and inverted /// lists for all planning reads. It also owns any provider-specific lock, epoch, -/// transaction, staging area, or poison state needed to apply the final update. +/// transaction, staging area, or poison state needed to apply the final update +/// through [`Apply`]. /// -/// Read methods are deliberately coarse grained so disk and blob providers can -/// batch I/O. Callbacks may be invoked in any order. `read_members` may invoke -/// its callback more than once for a list when the backend streams chunks. +/// List reads take one list and return its result. [`Self::read_vectors`], which +/// reads far more data, takes a batch so disk and blob providers can reorder, +/// coalesce, or parallelize I/O. pub trait MaintenanceAccessor: HasId + Send + Sized where T: Send, @@ -268,54 +182,45 @@ where /// External point id accepted by the data provider. type ExternalId: PartialEq + Send + Sync + 'static; - /// Stable centroid/list id. + /// Stable centroid/list id, fixed to [`ListProvider::ListId`] by the strategy. type ListId: VectorId; - /// Stable co-location group id. - type CoLocationGroupId: VectorId; - /// In-memory centroid catalog and navigator in this accessor's unified view. type Centroids: CentroidIndex; - /// Authoritative co-location grouping in this accessor's unified view. - type CoLocations: CoLocationSet; - /// Errors from planning reads, staging, or applying the update. type Error: ToRanked + Debug + Send + Sync + 'static; + /// Dimension of every canonical vector and centroid in this accessor's view. + /// + /// Fixed for the provider's lifetime, including before the index is initialized. + fn dim(&self) -> usize; + /// Borrow the centroid index used for routing and maintenance neighborhoods. fn centroids(&self) -> &Self::Centroids; - /// Borrow the co-location set used to plan locality-preserving layout. - fn co_locations(&self) -> &Self::CoLocations; - - /// Read list sizes for `lists`. - fn list_metadata( + /// Read the size of a live list. + fn list_metadata( &mut self, - lists: I, - emit: F, - ) -> impl SendFuture> - where - I: Iterator + Send, - F: FnMut(ListMetadata) + Send; + list: Self::ListId, + ) -> impl SendFuture>; - /// Resolve current list assignments for point ids. - fn assignments(&mut self, ids: I, emit: F) -> impl SendFuture> - where - I: Iterator + Send, - F: FnMut(Self::Id, Option) + Send; - - /// Stream the live member ids of selected lists. - fn read_members(&mut self, lists: I, emit: F) -> impl SendFuture> - where - I: Iterator + Send, - F: FnMut(Self::ListId, &[Self::Id]) + Send; + /// Read the member ids of a live list. + fn read_members( + &mut self, + list: Self::ListId, + ) -> impl SendFuture>; - /// Materialize canonical full-precision vectors for maintenance. - fn read_vectors(&mut self, ids: I, emit: F) -> impl SendFuture> - where - I: Iterator + Send, - F: FnMut(Self::Id, &[f32]) + Send; + /// Write the canonical vector of `ids[i]` into row `i` of `out`, in any order. + /// + /// `out` has one row per id and [`Self::dim`] columns, and every row must be + /// written. Covers points staged by this accessor as well as visible points, + /// without making staged points visible through ordinary provider reads. + fn read_vectors( + &mut self, + ids: &[Self::Id], + out: MutMatrixView<'_, f32>, + ) -> impl SendFuture>; /// Stage a canonical point and reserve its internal id. fn stage_insert( @@ -324,50 +229,51 @@ where element: T, ) -> impl SendFuture>; - /// Materialize canonical vectors for points staged by this accessor. - fn read_staged_vectors( - &mut self, - ids: I, - emit: F, - ) -> impl SendFuture> - where - I: Iterator + Send, - F: FnMut(Self::Id, &[f32]) + Send; - - /// Stage deletion of a currently visible internal point id. - fn stage_delete(&mut self, id: Self::Id) -> impl SendFuture>; - /// Reserve fresh logical list ids that will not alias retired ids. fn reserve_list_ids( &mut self, count: usize, ) -> impl SendFuture, Self::Error>>; +} - /// Reserve fresh co-location group ids that will not alias retired ids. - fn reserve_group_ids( - &mut self, - count: usize, - ) -> impl SendFuture, Self::Error>>; +/// Applies one kind of partition update and finishes the maintenance operation. +/// +/// Providers implement this once per update they support, from +/// [`crate::ivf::update`]. For example, a provider that never splits implements only +/// `Apply>`. Only the index constructs updates, and it guarantees their +/// documented invariants, so implementations may rely on them without re-checking. +/// +/// Returning `Ok(())` means all components expose one coherent resulting index. +/// Rollback, durability, concurrent-reader visibility, and recovery after `Err` are +/// intentionally provider-defined. +pub trait Apply: Send + Sized { + /// Errors from applying the update. + type Error: ToRanked + Debug + Send + Sync + 'static; - /// Apply the complete logical update and finish this mutation operation. - /// - /// Returning `Ok(())` means all components expose one coherent resulting - /// index. Rollback, durability, concurrent-reader visibility, and recovery - /// after `Err` are intentionally provider-defined. - fn apply( - self, - update: PartitionUpdate, - ) -> impl SendFuture>; + /// Apply `update`, consuming the accessor. + fn apply(self, update: U) -> impl SendFuture>; } /// Factory for an operation-scoped dynamic IVF maintenance accessor. +/// +/// The provider is borrowed exclusively for the lifetime of the accessor, so no +/// search or other mutation can observe it until the accessor is applied or +/// dropped. Providers may therefore mutate plain in-memory state in +/// [`Apply::apply`] without interior synchronization. Providers that share state +/// outside this borrow (for example through an `Arc`) remain responsible for +/// coordinating those aliases. pub trait MaintenanceStrategy<'a, Provider, T>: Send + Sync where - Provider: DataProvider, + Provider: ListProvider, T: Send, { /// Accessor used to plan and stage split/dissolve operations. - type MaintenanceAccessor: MaintenanceAccessor; + type MaintenanceAccessor: MaintenanceAccessor< + T, + Id = Provider::InternalId, + ExternalId = Provider::ExternalId, + ListId = Provider::ListId, + >; /// Error constructing the accessor. type Error: StandardError; @@ -375,7 +281,7 @@ where /// Construct one maintenance accessor. fn maintenance_accessor( &'a self, - provider: &'a Provider, + provider: &'a mut Provider, context: &'a Provider::Context, ) -> Result; } diff --git a/diskann/src/ivf/grouped.rs b/diskann/src/ivf/grouped.rs new file mode 100644 index 0000000000..7fd22cb2c5 --- /dev/null +++ b/diskann/src/ivf/grouped.rs @@ -0,0 +1,183 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Flat grouped layouts shared by the partition updates and the online planner. + +use std::ops::Range; + +/// Values stored in consecutive groups numbered from zero. +/// +/// Group `i` owns `values[ends[i - 1]..ends[i]]` and may be empty. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(super) struct Csr { + ends: Vec, + values: Vec, +} + +impl Csr { + /// Append a group holding `values`. + pub(super) fn push(&mut self, values: impl IntoIterator) { + self.values.extend(values); + self.ends.push(self.values.len()); + } + + /// Number of groups. + pub(super) fn num_groups(&self) -> usize { + self.ends.len() + } + + /// Positions in [`Self::values`] owned by `group`. + pub(super) fn range(&self, group: usize) -> Range { + let start = group.checked_sub(1).map_or(0, |prev| self.ends[prev]); + start..self.ends[group] + } + + /// Values of `group`. + pub(super) fn group(&self, group: usize) -> &[V] { + &self.values[self.range(group)] + } + + /// Every value, group by group. + pub(super) fn values(&self) -> &[V] { + &self.values + } +} + +impl Csr { + /// For each target below `targets`, the groups containing it in ascending order. + /// + /// Every value must be below `targets`. + pub(super) fn transpose(&self, targets: usize) -> Self { + let mut ends = vec![0; targets]; + for &target in &self.values { + ends[target] += 1; + } + let mut next = Vec::with_capacity(targets); + let mut total = 0; + for end in &mut ends { + next.push(total); + total += *end; + *end = total; + } + + let mut values = vec![0; total]; + for group in 0..self.num_groups() { + for &target in self.group(group) { + values[next[target]] = group; + next[target] += 1; + } + } + Self { ends, values } + } +} + +/// Values grouped under ascending unique keys, ascending within each group. +#[derive(Debug, PartialEq, Eq)] +pub(super) struct Grouped { + keys: Box<[K]>, + groups: Csr, +} + +impl Grouped { + /// Group distinct `pairs` by key. + pub(super) fn from_pairs(mut pairs: Vec<(K, V)>) -> Self { + pairs.sort_unstable(); + debug_assert!( + pairs.windows(2).all(|pair| pair[0] != pair[1]), + "grouped pairs must be distinct" + ); + + let same_key = |a: &(K, V), b: &(K, V)| a.0 == b.0; + let keys = pairs.chunk_by(same_key).map(|group| group[0].0).collect(); + let ends = pairs + .chunk_by(same_key) + .scan(0, |end, group| { + *end += group.len(); + Some(*end) + }) + .collect(); + let values = pairs.into_iter().map(|(_, value)| value).collect(); + Self { + keys, + groups: Csr { ends, values }, + } + } + + /// Number of values across all groups. + pub(super) fn len(&self) -> usize { + self.groups.values.len() + } + + /// Group keys, ascending. + pub(super) fn keys(&self) -> &[K] { + &self.keys + } + + /// Every value, group by group. + pub(super) fn values(&self) -> &[V] { + self.groups.values() + } + + /// The values under `key`, or nothing if `key` is absent. + pub(super) fn get(&self, key: K) -> &[V] { + self.keys + .binary_search(&key) + .map_or(&[], |index| self.groups.group(index)) + } + + /// The groups at positions `range`, with their keys. + pub(super) fn groups(&self, range: Range) -> impl ExactSizeIterator { + range.map(|index| (self.keys[index], self.groups.group(index))) + } + + /// Every group with its key, in ascending key order. + pub(super) fn iter(&self) -> impl ExactSizeIterator { + self.groups(0..self.keys.len()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn groups(csr: &Csr) -> Vec> { + (0..csr.num_groups()) + .map(|group| csr.group(group).to_vec()) + .collect() + } + + #[test] + fn csr_push_keeps_empty_groups() { + let mut csr = Csr::default(); + csr.push([1, 2]); + csr.push([]); + csr.push([3]); + assert_eq!(groups(&csr), vec![vec![1, 2], vec![], vec![3]]); + assert_eq!(csr.range(2), 2..3); + assert_eq!(csr.values(), &[1, 2, 3]); + } + + #[test] + fn transpose_lists_containing_groups() { + let mut csr = Csr::default(); + csr.push([1, 2]); + csr.push([]); + csr.push([2]); + let transposed = csr.transpose(4); + assert_eq!( + groups(&transposed), + vec![vec![], vec![0], vec![0, 2], vec![]] + ); + } + + #[test] + fn grouped_get_finds_groups_by_key() { + let grouped = Grouped::from_pairs(vec![(7, 2), (5, 1), (7, 0)]); + assert_eq!(grouped.keys(), &[5, 7]); + assert_eq!(grouped.get(7), &[0, 2]); + assert_eq!(grouped.get(6), &[] as &[i32]); + assert_eq!(grouped.len(), 3); + } +} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 3b4448e853..0c576d19c5 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -5,27 +5,22 @@ //! Dynamic IVF index orchestration. //! -//! The index owns split policy and drives one [`MaintenanceAccessor`] per mutation; -//! it never touches point, centroid, or list storage directly. +//! The index drives one [`MaintenanceAccessor`] per mutation: it plans the update +//! against the accessor's view and hands the result to [`Apply`]. It never touches +//! point, centroid, or list storage directly. -use std::{ - collections::{BTreeMap, BTreeSet, btree_map::Entry}, - fmt::{Debug, Display}, -}; - -use diskann_utils::{future::SendFuture, views::MatrixView}; -use rand::{Rng, SeedableRng, rngs::StdRng}; +use diskann_utils::{future::SendFuture, views::Matrix}; +use rand::{SeedableRng, rngs::StdRng}; use thiserror::Error; use crate::{ ANNError, ANNErrorKind, ANNResult, error::{ErrorExt, IntoANNResult}, - ivf::dynamic::{ - CentroidDelta, CentroidIndex, CentroidRecord, MaintenanceAccessor, MaintenanceStrategy, - PartitionUpdate, PointMove, + ivf::{ + dynamic::{Apply, CentroidIndex, ListProvider, MaintenanceAccessor, MaintenanceStrategy}, + online::{Parent, SplitPlan, StagedBatch, check_reserved, index_error}, + update::{Appends, Bootstrap, CentroidBlock, SplitInsert}, }, - provider::DataProvider, - utils::VectorId, }; /// Parameters for the online split policy. @@ -72,15 +67,16 @@ pub struct InsertStats { /// An incrementally maintained IVF index. /// -/// Mutations take `&mut self`, so they cannot overlap searches through the same value. +/// Mutations take `&mut self` and lend the provider exclusively to one +/// [`MaintenanceAccessor`], so they cannot overlap searches through the same value. #[derive(Debug)] -pub struct DynamicIvfIndex { +pub struct DynamicIvfIndex { provider: P, config: DynamicIvfConfig, rng: StdRng, } -impl DynamicIvfIndex

{ +impl DynamicIvfIndex

{ /// Construct an index over `provider`. /// /// # Errors @@ -114,24 +110,29 @@ impl DynamicIvfIndex

{ /// /// # Errors /// - /// Fails if `centroids` is empty, the index already has centroids, or the - /// maintenance accessor fails. + /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index + /// already has centroids, or the maintenance accessor fails. pub fn initialize<'a, S, T>( &'a mut self, strategy: &'a S, context: &'a P::Context, - centroids: MatrixView<'a, f32>, + centroids: Matrix, ) -> impl SendFuture> where S: MaintenanceStrategy<'a, P, T>, + S::MaintenanceAccessor: Apply>, T: Send, { - let provider = &self.provider; + let provider = &mut self.provider; + async move { let count = centroids.nrows(); if count == 0 { return Err(index_error("initialize requires at least one centroid")); } + if centroids.as_slice().iter().any(|x| !x.is_finite()) { + return Err(index_error("initial centroids must be finite")); + } let mut accessor = strategy .maintenance_accessor(provider, context) @@ -140,30 +141,22 @@ impl DynamicIvfIndex

{ if !accessor.centroids().is_empty() { return Err(index_error("dynamic IVF index is already initialized")); } + let dim = accessor.dim(); + if dim == 0 || centroids.ncols() != dim { + return Err(index_error(format!( + "initial centroids have dimension {}, expected {dim}", + centroids.ncols() + ))); + } let list_ids = accessor .reserve_list_ids(count) .await .escalate("initialize must reserve list ids")?; - check_reserved(count, list_ids.len())?; - - let centroids = CentroidDelta { - insert: list_ids - .into_iter() - .zip(centroids.row_iter()) - .map(|(id, vector)| CentroidRecord { - id, - vector: vector.into(), - }) - .collect(), - retire: Vec::new(), - }; + check_reserved(accessor.centroids(), &list_ids, count)?; accessor - .apply(PartitionUpdate { - centroids, - ..Default::default() - }) + .apply(Bootstrap::new(CentroidBlock::new(list_ids, centroids))) .await .escalate("initialize must apply the initial centroids") } @@ -171,8 +164,9 @@ impl DynamicIvfIndex

{ /// Insert a batch, splitting every list the batch pushes past `split_threshold`. /// - /// Routing and split planning run against the accessor's view before any change; - /// the whole batch is then handed to [`MaintenanceAccessor::apply`] as one update. + /// Routing and split planning run against the accessor's view before any change. + /// The batch is then applied as one [`Appends`], or as one [`SplitInsert`] when + /// lists split. /// /// # Errors /// @@ -187,11 +181,14 @@ impl DynamicIvfIndex

{ ) -> impl SendFuture> where S: MaintenanceStrategy<'a, P, T>, + S::MaintenanceAccessor: + Apply> + Apply>, T: Copy + Send + Sync, { - let provider = &self.provider; + let provider = &mut self.provider; let config = self.config; let rng = &mut self.rng; + async move { if points.is_empty() { return Ok(InsertStats::default()); @@ -207,511 +204,335 @@ impl DynamicIvfIndex

{ )); } - let (new_ids, mut vectors) = stage_points(&mut accessor, points).await?; - - let mut routed: BTreeMap<_, Vec<_>> = BTreeMap::new(); - - for &id in &new_ids { - let plan = accessor - .centroids() - .select(vectors.get(id)?, 1) - .await - .escalate("insert must route every point")?; - let list = plan - .selected() - .first() - .map(|selected| selected.id) - .ok_or_else(|| index_error("centroid selection returned no list"))?; - routed.entry(list).or_default().push(id); - } - - let parents = select_split_parents(&mut accessor, &config, &routed).await?; + let batch = StagedBatch::stage(&mut accessor, points).await?; + let parents = Self::select_parents(&config, &mut accessor, &batch).await?; if parents.is_empty() { - let update = PartitionUpdate { - point_moves: routed_moves(&routed, &BTreeMap::new()), - ..Default::default() - }; + let update = batch.into_appends(); + let inserted = update.len(); accessor .apply(update) .await .escalate("insert must apply routed points")?; - return Ok(InsertStats { - inserted: new_ids.len(), - ..Default::default() + inserted, + splits: 0, + reassigned: 0, }); } - let plan = SplitPlan::prepare(&mut accessor, &config, &parents, &mut vectors).await?; - let children = - plan.fit_children(&routed, &vectors, config.two_means_iterations, rng)?; - - let child_ids = accessor - .reserve_list_ids(children.len()) - .await - .escalate("split must reserve child list ids")?; - check_reserved(children.len(), child_ids.len())?; - - let placements = plan.reassign(&routed, &vectors, &child_ids, &children)?; - - let mut point_moves = routed_moves(&routed, &placements); - let mut reassigned = 0; - for (&from, ids) in &plan.members { - for &id in ids { - if let Some(&to) = placements.get(&id) - && to != from - { - reassigned += 1; - point_moves.push(PointMove { - id, - from: Some(from), - to: Some(to), - }); - } - } - } - - let splits = parents.len(); - let update = PartitionUpdate { - centroids: CentroidDelta { - insert: child_ids - .iter() - .zip(children) - .map(|(&id, vector)| CentroidRecord { id, vector }) - .collect(), - retire: parents, - }, - point_moves, - ..Default::default() + let update = SplitPlan::gather(&mut accessor, &config, batch, parents) + .await? + .solve(config.two_means_iterations, rng)?; + let stats = InsertStats { + inserted: update.appends().len(), + splits: update.splits().len(), + reassigned: update.splits().num_evacuated() + update.reassignments().len(), }; - accessor .apply(update) .await .escalate("insert must apply the split update")?; - - Ok(InsertStats { - inserted: new_ids.len(), - splits, - reassigned, - }) + Ok(stats) } } -} -/////////////////// -// Insert phases // -/////////////////// + /// Choose the routed lists to split. + /// + /// Admits every list the batch pushes past `split_threshold`, largest projected + /// size first, while the live-list count stays within `max_clusters`. Parents are + /// returned in ascending list order. + /// + /// Takes the configuration rather than `&self` because the provider is lent to + /// `accessor` for the whole mutation. + /// + /// # Errors + /// + /// Fails if the accessor fails. + async fn select_parents( + config: &DynamicIvfConfig, + accessor: &mut A, + batch: &StagedBatch, + ) -> ANNResult>> + where + A: MaintenanceAccessor, + T: Send, + { + let mut overflowing = Vec::new(); + for (list, incoming) in batch.routed() { + let len = accessor + .list_metadata(list) + .await + .escalate("insert must read routed list sizes")? + .len; + let projected = len + incoming; + if projected > config.split_threshold { + overflowing.push((projected, Parent { list, len })); + } + } -/// Canonical vectors gathered for one mutation, all of dimension `dim`. -struct Vectors { - dim: usize, - rows: BTreeMap>, -} + // Each split adds one live list. + let budget = config.max_clusters.map_or(usize::MAX, |max| { + max.saturating_sub(accessor.centroids().len()) + }); + overflowing.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.list.cmp(&b.1.list))); + overflowing.truncate(budget); -impl Vectors { - fn get(&self, id: Id) -> ANNResult<&[f32]> { - self.rows - .get(&id) - .map(|row| &**row) - .ok_or_else(|| index_error(format!("no canonical vector for point {id}"))) + let mut parents: Vec<_> = overflowing.into_iter().map(|(_, parent)| parent).collect(); + parents.sort_unstable_by_key(|parent| parent.list); + Ok(parents) } +} - fn check_dim(&self, len: usize) -> ANNResult<()> { - if len == self.dim { - Ok(()) - } else { - Err(index_error(format!( - "vector dimension {len} does not match batch dimension {}", - self.dim - ))) +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + ivf::test::provider::{Faults, Provider, Strategy}, + provider::DefaultContext, + }; + + fn config(split_threshold: usize, reassign_neighbors: usize) -> DynamicIvfConfig { + DynamicIvfConfig { + split_threshold, + max_clusters: None, + reassign_neighbors, + two_means_iterations: 10, + seed: 7, } } - fn check_dims(&self) -> ANNResult<()> { - self.rows - .values() - .try_for_each(|row| self.check_dim(row.len())) + /// Pair each vector with a fresh external id, counting up from `first`. + fn batch(first: u32, vectors: &[[f32; D]]) -> Vec<(u32, &[f32])> { + (first..) + .zip(vectors.iter().map(|vector| &vector[..])) + .collect() } -} -/// Stage every point and materialize its canonical vector. -async fn stage_points( - accessor: &mut A, - points: &[(A::ExternalId, T)], -) -> ANNResult<(Vec, Vectors)> -where - A: MaintenanceAccessor, - T: Copy + Send + Sync, -{ - let mut ids = Vec::with_capacity(points.len()); - for (external, element) in points { - ids.push( - accessor - .stage_insert(external, *element) - .await - .escalate("insert must stage every point")?, - ); + fn matrix(rows: &[[f32; D]]) -> Matrix { + let data: Box<[f32]> = rows.iter().flatten().copied().collect(); + Matrix::try_from(data, rows.len(), D).unwrap() } - let mut rows: BTreeMap> = BTreeMap::new(); - accessor - .read_staged_vectors(ids.iter().copied(), |id, vector| { - rows.insert(id, Box::from(vector)); - }) - .await - .escalate("insert must read staged canonical vectors")?; - - let dim = rows.values().next().map_or(0, |row| row.len()); - if dim == 0 { - return Err(index_error("staged canonical vectors must be non-empty")); - } - let vectors = Vectors { dim, rows }; - vectors.check_dims()?; - Ok((ids, vectors)) -} - -/// Route each staged point to its nearest live list. -async fn route( - accessor: &mut A, - ids: &[A::Id], - vectors: &Vectors, -) -> ANNResult>> -where - A: MaintenanceAccessor, - T: Send, -{ - let centroids = accessor.centroids(); - let mut routed: BTreeMap> = BTreeMap::new(); - - for &id in ids { - let plan = centroids - .select(vectors.get(id)?, 1) + async fn initialized( + provider: Provider, + config: DynamicIvfConfig, + centroids: Matrix, + ) -> DynamicIvfIndex { + let mut index = DynamicIvfIndex::new(provider, config).unwrap(); + index + .initialize::<_, &[f32]>(&Strategy, &DefaultContext, centroids) .await - .escalate("insert must route every point")?; - let list = plan - .selected() - .first() - .map(|selected| selected.id) - .ok_or_else(|| index_error("centroid selection returned no list"))?; - routed.entry(list).or_default().push(id); + .unwrap(); + index } - Ok(routed) -} - -/// Return routed lists whose projected size exceeds the threshold, largest first -/// under the cluster cap, in ascending id order. -async fn select_split_parents( - accessor: &mut A, - config: &DynamicIvfConfig, - routed: &BTreeMap>, -) -> ANNResult> -where - A: MaintenanceAccessor, - T: Send, -{ - let mut sizes = BTreeMap::new(); - accessor - .list_metadata(routed.keys().copied(), |metadata| { - sizes.insert(metadata.id, metadata.len); - }) - .await - .escalate("insert must read routed list sizes")?; - - let mut overflowing = Vec::new(); - for (&list, incoming) in routed { - let current = sizes - .get(&list) - .copied() - .ok_or_else(|| index_error(format!("no metadata for list {list}")))?; - let projected = current + incoming.len(); - if projected > config.split_threshold { - overflowing.push((projected, list)); - } + #[test] + fn new_rejects_invalid_config() { + let new = |config| DynamicIvfIndex::new(Provider::new(1), config).err(); + assert_eq!(new(config(1, 1)), Some(ConfigError::SplitThreshold(1))); + assert_eq!(new(config(4, 0)), Some(ConfigError::ReassignNeighbors)); } - // Each split adds one live list. - let budget = config.max_clusters.map_or(usize::MAX, |max| { - max.saturating_sub(accessor.centroids().len()) - }); - overflowing.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.cmp(&b.1))); - overflowing.truncate(budget); - - let mut parents: Vec<_> = overflowing.into_iter().map(|(_, list)| list).collect(); - parents.sort_unstable(); - Ok(parents) -} - -/// One split parent and the surviving lists reassigned alongside it. -struct Region { - parent: L, - neighbors: Vec, -} - -/// Everything read from the accessor to split a set of parents. -struct SplitPlan { - regions: Vec>, - neighbor_centroids: BTreeMap>, - /// Current members of every parent and neighbor list. - members: BTreeMap>, -} - -impl SplitPlan { - async fn prepare( - accessor: &mut A, - config: &DynamicIvfConfig, - parents: &[L], - vectors: &mut Vectors, - ) -> ANNResult - where - A: MaintenanceAccessor, - T: Send, - { - let parent_set: BTreeSet = parents.iter().copied().collect(); - let centroids = accessor.centroids(); - let mut regions = Vec::with_capacity(parents.len()); - let mut neighbor_centroids = BTreeMap::new(); - - for &parent in parents { - let anchor = centroids - .centroid(parent) - .ok_or_else(|| index_error(format!("split parent {parent} is not live")))?; - let selected = centroids - .select(anchor, config.reassign_neighbors + 1) - .await - .escalate("split must find neighbor lists")?; - - // Other parents are retired by this update, so they are never candidates. - let neighbors: Vec = selected - .into_selected() - .into_iter() - .map(|selected| selected.id) - .filter(|id| !parent_set.contains(id)) - .take(config.reassign_neighbors) - .collect(); - - for &neighbor in &neighbors { - if let Entry::Vacant(entry) = neighbor_centroids.entry(neighbor) { - let vector = centroids.centroid(neighbor).ok_or_else(|| { - index_error(format!("neighbor list {neighbor} is not live")) - })?; - vectors.check_dim(vector.len())?; - entry.insert(Box::from(vector)); - } - } - regions.push(Region { parent, neighbors }); - } - - let lists: BTreeSet = regions - .iter() - .flat_map(|region| { - std::iter::once(region.parent).chain(region.neighbors.iter().copied()) - }) - .collect(); - - let mut members: BTreeMap> = BTreeMap::new(); - accessor - .read_members(lists.iter().copied(), |list, ids| { - members.entry(list).or_default().extend_from_slice(ids) - }) + #[tokio::test] + async fn initialize_rejects_bad_centroids() { + let mut index = DynamicIvfIndex::new(Provider::new(2), config(10, 1)).unwrap(); + let wrong_dim = index + .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0, 0.0]])) .await - .escalate("split must read region members")?; - - accessor - .read_vectors(members.values().flatten().copied(), |id, vector| { - vectors.rows.insert(id, Box::from(vector)); - }) + .unwrap_err(); + assert!(wrong_dim.to_string().contains("dimension"), "{wrong_dim}"); + let non_finite = index + .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[f32::NAN, 0.0]])) .await - .escalate("split must read region canonical vectors")?; - vectors.check_dims()?; + .unwrap_err(); + assert!(non_finite.to_string().contains("finite"), "{non_finite}"); - Ok(Self { - regions, - neighbor_centroids, - members, - }) + index + .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0]])) + .await + .unwrap(); + let again = index + .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[1.0, 1.0]])) + .await + .unwrap_err(); + assert!(again.to_string().contains("already initialized"), "{again}"); } - /// Existing members plus incoming points of `list`. - fn points_of<'s>( - &'s self, - list: L, - routed: &'s BTreeMap>, - ) -> impl Iterator + 's { - self.members - .get(&list) - .into_iter() - .chain(routed.get(&list)) - .flatten() - .copied() + #[tokio::test] + async fn insert_requires_initialization() { + let mut index = DynamicIvfIndex::new(Provider::new(1), config(4, 1)).unwrap(); + let empty = index + .insert_batch(&Strategy, &DefaultContext, &batch::<1>(0, &[])) + .await + .unwrap(); + assert_eq!(empty, InsertStats::default()); + let err = index + .insert_batch(&Strategy, &DefaultContext, &batch(0, &[[1.0]])) + .await + .unwrap_err(); + assert!(err.to_string().contains("initialized"), "{err}"); } - /// Fit two children per parent with one joint k-means over every parent's points. - fn fit_children( - &self, - routed: &BTreeMap>, - vectors: &Vectors, - iterations: usize, - rng: &mut StdRng, - ) -> ANNResult>> { - let mut points = Vec::new(); - let mut centers = Vec::with_capacity(2 * self.regions.len()); - for region in &self.regions { - let start = points.len(); - for id in self.points_of(region.parent, routed) { - points.push(vectors.get(id)?); + #[tokio::test] + async fn insert_without_overflow_appends_to_nearest_lists() { + let centroids = matrix(&[[0.0, 0.0], [10.0, 10.0]]); + let mut index = initialized(Provider::new(2), config(10, 1), centroids).await; + let vectors = [[1.0, 1.0], [9.0, 9.0], [2.0, 0.0]]; + let stats = index + .insert_batch(&Strategy, &DefaultContext, &batch(0, &vectors)) + .await + .unwrap(); + assert_eq!( + stats, + InsertStats { + inserted: 3, + splits: 0, + reassigned: 0 } + ); - let span = points.len() - start; - if span < 2 { - return Err(index_error(format!( - "split parent {} has fewer than two points", - region.parent - ))); - } - let a = rng.random_range(0..span); - let b = (a + rng.random_range(1..span)) % span; - centers.push(points[start + a].to_vec()); - centers.push(points[start + b].to_vec()); - } + let provider = index.provider(); + provider.check(); + assert_eq!( + provider.member_vectors(0), + vec![vec![1.0, 1.0], vec![2.0, 0.0]] + ); + assert_eq!(provider.member_vectors(1), vec![vec![9.0, 9.0]]); + } - lloyd(&points, &mut centers, iterations.max(1)); - Ok(centers.into_iter().map(Vec::into_boxed_slice).collect()) + #[tokio::test] + async fn max_clusters_caps_splits() { + let mut capped = config(2, 1); + capped.max_clusters = Some(2); + let mut index = initialized(Provider::new(1), capped, matrix(&[[0.0], [10.0]])).await; + let stats = index + .insert_batch( + &Strategy, + &DefaultContext, + &batch(0, &[[1.0], [2.0], [3.0]]), + ) + .await + .unwrap(); + assert_eq!(stats.splits, 0); + assert_eq!(index.provider().lists(), vec![0, 1]); + index.provider().check(); } - /// Place every region point on its nearest candidate across all regions containing it. - fn reassign( - &self, - routed: &BTreeMap>, - vectors: &Vectors, - child_ids: &[L], - children: &[Box<[f32]>], - ) -> ANNResult> { - let mut best: BTreeMap = BTreeMap::new(); - for ((region, ids), vecs) in self - .regions - .iter() - .zip(child_ids.chunks_exact(2)) - .zip(children.chunks_exact(2)) - { - let mut candidates: Vec<(L, &[f32])> = Vec::with_capacity(region.neighbors.len() + 2); - for &neighbor in ®ion.neighbors { - let vector = self.neighbor_centroids.get(&neighbor).ok_or_else(|| { - index_error(format!("missing centroid for neighbor list {neighbor}")) - })?; - candidates.push((neighbor, vector)); - } - candidates.extend(ids.iter().copied().zip(vecs.iter().map(|v| &**v))); - - let lists = std::iter::once(region.parent).chain(region.neighbors.iter().copied()); - for list in lists { - for id in self.points_of(list, routed) { - let (to, distance) = nearest(vectors.get(id)?, &candidates) - .ok_or_else(|| index_error("split region has no candidate lists"))?; - match best.entry(id) { - Entry::Vacant(entry) => { - entry.insert((to, distance)); - } - Entry::Occupied(mut entry) => { - if distance < entry.get().1 { - entry.insert((to, distance)); - } - } - } - } + /// Lists A (0), B (20), and C (40) in one dimension. Inserting 0.5, 8, and 9 + /// overflows A, which splits into children near 1/6 and 8.5, with B and C as + /// its neighbors. + #[tokio::test] + async fn split_places_parent_points_and_moves_neighbor_points_only_toward_children() { + let provider = Provider::from_lists( + 1, + &[ + (&[0.0], &[&[-1.0], &[1.0]]), + (&[20.0], &[&[12.0], &[25.0], &[33.0]]), + (&[40.0], &[&[41.0]]), + ], + ); + let mut index = DynamicIvfIndex::new(provider, config(4, 2)).unwrap(); + let stats = index + .insert_batch( + &Strategy, + &DefaultContext, + &batch(100, &[[0.5], [8.0], [9.0]]), + ) + .await + .unwrap(); + // A's members -1 and 1 are evacuated, and 12 moves out of B. + assert_eq!( + stats, + InsertStats { + inserted: 3, + splits: 1, + reassigned: 3 } - } - Ok(best.into_iter().map(|(id, (list, _))| (id, list)).collect()) + ); + + let provider = index.provider(); + provider.check(); + assert_eq!(provider.lists(), vec![1, 2, 3, 4]); + let low = provider.list_of(&[0.5]).unwrap(); + let high = provider.list_of(&[8.0]).unwrap(); + assert!((provider.centroid(low).unwrap()[0] - 1.0 / 6.0).abs() < 1e-6); + assert_eq!(provider.centroid(high).unwrap(), &[8.5]); + assert_eq!( + provider.member_vectors(low), + vec![vec![-1.0], vec![0.5], vec![1.0]] + ); + assert_eq!( + provider.member_vectors(high), + vec![vec![8.0], vec![9.0], vec![12.0]] + ); + // 33 is closer to C than to B, but neighbor points only move toward children. + assert_eq!(provider.member_vectors(1), vec![vec![25.0], vec![33.0]]); + assert_eq!(provider.member_vectors(2), vec![vec![41.0]]); } -} -/// Moves for newly inserted points: their split placement, else their route. -fn routed_moves( - routed: &BTreeMap>, - placements: &BTreeMap, -) -> Vec> { - routed - .iter() - .flat_map(|(&route, ids)| { - ids.iter().map(move |&id| PointMove { - id, - from: None, - to: Some(placements.get(&id).copied().unwrap_or(route)), + #[tokio::test] + async fn repeated_batches_keep_the_partition_valid() { + let centroids = matrix(&[[0.0, 0.0], [10.0, 10.0]]); + let mut index = initialized(Provider::new(2), config(8, 2), centroids).await; + let vectors: Vec<[f32; 2]> = (0..300) + .map(|i| { + let t = i as f32; + [(t * 0.37).sin() * 20.0, (t * 0.11).cos() * 20.0] }) - }) - .collect() -} - -///////////// -// Helpers // -///////////// - -// Clustering and reassignment use squared L2 regardless of the search metric. -fn squared_l2(a: &[f32], b: &[f32]) -> f32 { - a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() -} - -fn nearest(point: &[f32], candidates: &[(L, &[f32])]) -> Option<(L, f32)> { - candidates - .iter() - .map(|&(list, centroid)| (list, squared_l2(point, centroid))) - .min_by(|a, b| a.1.total_cmp(&b.1)) -} + .collect(); -fn lloyd(points: &[&[f32]], centers: &mut [Vec], iterations: usize) { - let dim = centers.first().map_or(0, Vec::len); - if dim == 0 { - return; - } - let mut sums = vec![0.0f32; centers.len() * dim]; - let mut counts = vec![0usize; centers.len()]; - for _ in 0..iterations { - sums.fill(0.0); - counts.fill(0); - for &point in points { - let Some(c) = centers - .iter() - .map(|center| squared_l2(point, center)) - .enumerate() - .min_by(|a, b| a.1.total_cmp(&b.1)) - .map(|(c, _)| c) - else { - continue; - }; - counts[c] += 1; - for (sum, x) in sums[c * dim..(c + 1) * dim].iter_mut().zip(point) { - *sum += x; - } - } - // Empty clusters keep their previous center. - for ((center, sum), &count) in centers.iter_mut().zip(sums.chunks_exact(dim)).zip(&counts) { - if count > 0 { - let scale = 1.0 / count as f32; - for (x, s) in center.iter_mut().zip(sum) { - *x = s * scale; - } - } + // Insert in batches of 25, checking the partition after each batch. + let mut stats = InsertStats::default(); + for (first, chunk) in (0..).step_by(25).zip(vectors.chunks(25)) { + let batch = index + .insert_batch(&Strategy, &DefaultContext, &batch(first, chunk)) + .await + .unwrap(); + index.provider().check(); + stats.inserted += batch.inserted; + stats.splits += batch.splits; + stats.reassigned += batch.reassigned; } - } -} -#[track_caller] -fn index_error(message: D) -> ANNError -where - D: Display + Debug + Send + Sync + 'static, -{ - ANNError::message(ANNErrorKind::IndexError, message) -} + let provider = index.provider(); + assert_eq!(stats.inserted, 300); + assert_eq!(provider.len(), 300); + assert!(stats.splits > 1 && stats.reassigned > 0, "{stats:?}"); + assert_eq!(provider.lists().len(), 2 + stats.splits); + assert_eq!(provider.num_retired(), stats.splits); + } -#[track_caller] -fn check_reserved(expected: usize, actual: usize) -> ANNResult<()> { - if expected == actual { - Ok(()) - } else { - Err(index_error(format!( - "accessor reserved {actual} ids, expected {expected}" - ))) + #[tokio::test] + async fn inconsistent_provider_responses_fail_before_any_change() { + let fault = |set: fn(&mut Faults)| { + let mut faults = Faults::default(); + set(&mut faults); + faults + }; + let cases = [ + ( + fault(|f| f.skip_vector = true), + "no finite canonical vector", + ), + (fault(|f| f.nan_vector = true), "no finite canonical vector"), + (fault(|f| f.repeat_stage = true), "staging returned point"), + (fault(|f| f.reuse_lists = true), "reserved live list id"), + (fault(|f| f.drop_member = true), "but its metadata reports"), + ]; + for (faults, expected) in cases { + // Inserting 0.5 and 8 overflows list 0, so every case reaches the split path. + let provider = + Provider::from_lists(1, &[(&[0.0], &[&[-1.0], &[1.0]]), (&[20.0], &[&[12.0]])]) + .with_faults(faults); + let before = provider.clone(); + let mut index = DynamicIvfIndex::new(provider, config(3, 1)).unwrap(); + let err = index + .insert_batch(&Strategy, &DefaultContext, &batch(100, &[[0.5], [8.0]])) + .await + .unwrap_err(); + assert!(err.to_string().contains(expected), "{faults:?}: {err}"); + assert!(index.provider().same_contents(&before), "{faults:?}"); + } } } diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index bb417e4de0..140caab3cf 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -5,7 +5,14 @@ //! Incrementally maintained IVF index. +pub mod colocation; pub mod dynamic; +mod grouped; pub mod index; +mod online; +pub mod update; + +#[cfg(test)] +mod test; pub use index::{ConfigError, DynamicIvfConfig, DynamicIvfIndex, InsertStats}; diff --git a/diskann/src/ivf/online/batch.rs b/diskann/src/ivf/online/batch.rs new file mode 100644 index 0000000000..9bcf37e61a --- /dev/null +++ b/diskann/src/ivf/online/batch.rs @@ -0,0 +1,108 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! A staged and routed insert batch. + +use diskann_utils::views::Matrix; + +use super::{first_repeat, gather::read_rows, index_error}; +use crate::{ + ANNResult, + error::ErrorExt, + ivf::{ + dynamic::{CentroidIndex, MaintenanceAccessor}, + grouped::Grouped, + update::Appends, + }, + utils::VectorId, +}; + +/// Staged points, their canonical vectors, and their routes. +/// +/// Batch position `i` refers to `ids[i]`, row `i` of `vectors`, and `routes[i]`. +pub(in crate::ivf) struct StagedBatch { + pub(super) ids: Vec, + pub(super) vectors: Matrix, + /// The nearest live list of each point. + pub(super) routes: Vec, + /// Batch positions grouped by route. + pub(super) by_route: Grouped, +} + +impl StagedBatch { + /// Stage every point, read its canonical vector, and route it to its nearest live + /// list. + /// + /// # Errors + /// + /// Fails if any accessor call fails or returns inconsistent data. + pub(in crate::ivf) async fn stage( + accessor: &mut A, + points: &[(A::ExternalId, T)], + ) -> ANNResult + where + A: MaintenanceAccessor, + T: Copy + Send + Sync, + { + let dim = accessor.dim(); + if dim == 0 { + return Err(index_error( + "canonical vectors must have a non-zero dimension", + )); + } + + let mut ids = Vec::with_capacity(points.len()); + for (external, element) in points { + ids.push( + accessor + .stage_insert(external, *element) + .await + .escalate("insert must stage every point")?, + ); + } + if let Some(id) = first_repeat(ids.iter().copied()) { + return Err(index_error(format!( + "staging returned point {id} more than once" + ))); + } + let vectors = read_rows(accessor, &ids, dim).await?; + + let centroids = accessor.centroids(); + let mut routes = Vec::with_capacity(ids.len()); + for vector in vectors.row_iter() { + let plan = centroids + .select(vector, 1) + .await + .escalate("insert must route every point")?; + let list = plan + .selected() + .first() + .map(|selected| selected.id) + .ok_or_else(|| index_error("centroid selection returned no list"))?; + routes.push(list); + } + + let by_route = Grouped::from_pairs(routes.iter().copied().zip(0..).collect()); + Ok(Self { + ids, + vectors, + routes, + by_route, + }) + } + + /// Place every point on its route. + pub(in crate::ivf) fn into_appends(self) -> Appends { + Appends::new(&self.ids, &self.routes) + } + + /// Lists that received points, ascending, with the number of points routed to + /// each. + pub(in crate::ivf) fn routed(&self) -> impl ExactSizeIterator { + self.by_route + .iter() + .map(|(list, positions)| (list, positions.len())) + } +} diff --git a/diskann/src/ivf/online/gather.rs b/diskann/src/ivf/online/gather.rs new file mode 100644 index 0000000000..41b852805e --- /dev/null +++ b/diskann/src/ivf/online/gather.rs @@ -0,0 +1,111 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Validated reads from a maintenance accessor. + +use std::ops::Range; + +use diskann_utils::views::Matrix; + +use super::index_error; +use crate::{ + ANNResult, + error::ErrorExt, + ivf::{dynamic::MaintenanceAccessor, grouped::Csr}, + utils::VectorId, +}; + +/// Read the canonical vector of every point in `ids`; row `i` belongs to `ids[i]`. +/// +/// # Errors +/// +/// Fails if the accessor fails, or leaves a row unwritten or non-finite. +pub(super) async fn read_rows( + accessor: &mut A, + ids: &[A::Id], + dim: usize, +) -> ANNResult> +where + A: MaintenanceAccessor, + T: Send, +{ + // Rows the accessor never writes stay NaN, so the finiteness check catches them. + let mut rows = Matrix::new(f32::NAN, ids.len(), dim); + accessor + .read_vectors(ids, rows.as_mut_view()) + .await + .escalate("maintenance must read canonical vectors")?; + + if let Some((id, _)) = ids + .iter() + .zip(rows.row_iter()) + .find(|(_, row)| row.iter().any(|x| !x.is_finite())) + { + return Err(index_error(format!( + "no finite canonical vector for point {id}" + ))); + } + Ok(rows) +} + +/// The existing members of a table of lists and their canonical vectors. +/// +/// List `i`'s members occupy positions [`Self::range`]`(i)`, which index both the ids +/// and the vector rows. +pub(super) struct Members { + ids: Csr, + vectors: Matrix, +} + +impl Members { + /// Read the members of every list in `lists`, then their canonical vectors. + /// + /// # Errors + /// + /// Fails if the accessor fails or [`read_rows`] fails. + pub(super) async fn read(accessor: &mut A, lists: &[A::ListId]) -> ANNResult + where + A: MaintenanceAccessor, + T: Send, + { + let mut ids = Csr::default(); + for &list in lists { + let members = accessor + .read_members(list) + .await + .escalate("maintenance must read list members")?; + ids.push(members.iter().copied()); + } + + let dim = accessor.dim(); + let vectors = read_rows(accessor, ids.values(), dim).await?; + Ok(Self { ids, vectors }) + } + + /// Total number of members across all lists. + pub(super) fn total(&self) -> usize { + self.ids.values().len() + } + + /// Positions of `list`'s members. + pub(super) fn range(&self, list: usize) -> Range { + self.ids.range(list) + } + + /// Member ids of `list`. + pub(super) fn ids(&self, list: usize) -> &[Id] { + self.ids.group(list) + } + + /// The canonical vector at member position `position`. + pub(super) fn vector(&self, position: usize) -> &[f32] { + self.vectors.row(position) + } + + /// Canonical vectors of `list`'s members, in member order. + pub(super) fn vectors(&self, list: usize) -> impl ExactSizeIterator { + self.range(list).map(|position| self.vectors.row(position)) + } +} diff --git a/diskann/src/ivf/online/kernels.rs b/diskann/src/ivf/online/kernels.rs new file mode 100644 index 0000000000..103f5e7aad --- /dev/null +++ b/diskann/src/ivf/online/kernels.rs @@ -0,0 +1,132 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Distance and clustering kernels for split planning. +//! +//! Clustering and placement use squared Euclidean distance regardless of the search +//! metric. + +use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; + +/// Squared Euclidean distance between two vectors of equal length. +pub(super) fn squared_l2(a: &[f32], b: &[f32]) -> f32 { + SquaredL2::evaluate(a, b) +} + +/// The candidate nearest to `point`, preferring the earliest candidate on ties. +pub(super) fn nearest(point: &[f32], candidates: &[(L, &[f32])]) -> Option { + candidates + .iter() + .map(|&(list, centroid)| (list, squared_l2(point, centroid))) + .min_by(|a, b| a.1.total_cmp(&b.1)) + .map(|(list, _)| list) +} + +/// Scratch buffers for [`lloyd`], reusable across calls. +#[derive(Debug, Default)] +pub(super) struct LloydScratch { + sums: Vec, + counts: Vec, +} + +/// Refine `centers`, stored row-major with `dim` columns, with `iterations` Lloyd +/// steps over the points yielded by `points`. +/// +/// Each point joins its nearest center, preferring the earliest on ties. A center that +/// attracts no points keeps its position. +pub(super) fn lloyd<'p, P, I>( + points: P, + centers: &mut [f32], + dim: usize, + iterations: usize, + scratch: &mut LloydScratch, +) where + P: Fn() -> I, + I: Iterator, +{ + if dim == 0 || centers.is_empty() { + return; + } + let LloydScratch { sums, counts } = scratch; + for _ in 0..iterations { + sums.clear(); + sums.resize(centers.len(), 0.0); + counts.clear(); + counts.resize(centers.len() / dim, 0); + + for point in points() { + let Some(center) = centers + .chunks_exact(dim) + .map(|center| squared_l2(point, center)) + .enumerate() + .min_by(|a, b| a.1.total_cmp(&b.1)) + .map(|(center, _)| center) + else { + continue; + }; + counts[center] += 1; + for (sum, x) in sums[center * dim..(center + 1) * dim].iter_mut().zip(point) { + *sum += x; + } + } + + for ((center, sum), &count) in centers + .chunks_exact_mut(dim) + .zip(sums.chunks_exact(dim)) + .zip(counts.iter()) + { + if count > 0 { + let scale = 1.0 / count as f32; + for (x, s) in center.iter_mut().zip(sum) { + *x = s * scale; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn nearest_prefers_the_earliest_tie() { + let (a, b, c) = ([0.0, 0.0], [2.0, 0.0], [0.0, 2.0]); + let candidates = [(7u32, &a[..]), (8, &b[..]), (9, &c[..])]; + assert_eq!(nearest(&[1.8, 0.1], &candidates), Some(8)); + // Equidistant from 8 and 9. + assert_eq!(nearest(&[1.0, 1.0], &candidates[1..]), Some(8)); + assert_eq!(nearest::(&[1.0, 1.0], &[]), None); + } + + #[test] + fn lloyd_moves_centers_to_cluster_means() { + let points: Vec<[f32; 2]> = vec![[0.0, 0.0], [0.0, 2.0], [10.0, 0.0], [10.0, 2.0]]; + let mut centers = [1.0, 1.0, 9.0, 1.0]; + let mut scratch = LloydScratch::default(); + lloyd( + || points.iter().map(|p| &p[..]), + &mut centers, + 2, + 3, + &mut scratch, + ); + assert_eq!(centers, [0.0, 1.0, 10.0, 1.0]); + } + + #[test] + fn lloyd_keeps_centers_without_points() { + let points: Vec<[f32; 1]> = vec![[0.0], [1.0]]; + let mut centers = [0.5, 100.0]; + lloyd( + || points.iter().map(|p| &p[..]), + &mut centers, + 1, + 2, + &mut LloydScratch::default(), + ); + assert_eq!(centers, [0.5, 100.0]); + } +} diff --git a/diskann/src/ivf/online/mod.rs b/diskann/src/ivf/online/mod.rs new file mode 100644 index 0000000000..d1cc4d9b75 --- /dev/null +++ b/diskann/src/ivf/online/mod.rs @@ -0,0 +1,64 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! The online insert algorithm: stage and route a batch, then split every list the +//! batch overflows. +//! +//! Planning performs all of its accessor reads first and then computes the update +//! from owned, densely indexed data. Everything read from the accessor is validated +//! as it arrives, so the updates built from it need no further checks. Nothing here +//! mutates the index; the caller applies the planned update. + +mod batch; +mod gather; +mod kernels; +mod split; + +use std::fmt::{Debug, Display}; + +pub(super) use batch::StagedBatch; +pub(super) use split::{Parent, SplitPlan}; + +use crate::{ANNError, ANNErrorKind, ANNResult, ivf::dynamic::CentroidIndex}; + +#[track_caller] +pub(super) fn index_error(message: D) -> ANNError +where + D: Display + Debug + Send + Sync + 'static, +{ + ANNError::message(ANNErrorKind::IndexError, message) +} + +/// Check that the accessor reserved `expected` distinct list ids, none of them live. +#[track_caller] +pub(super) fn check_reserved( + centroids: &C, + ids: &[C::ListId], + expected: usize, +) -> ANNResult<()> { + if ids.len() != expected { + return Err(index_error(format!( + "accessor reserved {} list ids, expected {expected}", + ids.len() + ))); + } + if let Some(id) = first_repeat(ids.iter().copied()) { + return Err(index_error(format!("accessor reserved list id {id} twice"))); + } + if let Some(id) = ids.iter().find(|&&id| centroids.centroid(id).is_some()) { + return Err(index_error(format!("accessor reserved live list id {id}"))); + } + Ok(()) +} + +/// The first value that occurs more than once. +fn first_repeat(values: impl Iterator) -> Option { + let mut sorted: Vec = values.collect(); + sorted.sort_unstable(); + sorted + .windows(2) + .find(|pair| pair[0] == pair[1]) + .map(|pair| pair[0]) +} diff --git a/diskann/src/ivf/online/split.rs b/diskann/src/ivf/online/split.rs new file mode 100644 index 0000000000..087a038e9c --- /dev/null +++ b/diskann/src/ivf/online/split.rs @@ -0,0 +1,346 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Split planning: gather each split's region, fit children, and place every affected +//! point. + +use std::ops::Range; + +use diskann_utils::views::Matrix; +use rand::{Rng, rngs::StdRng}; + +use super::{ + batch::StagedBatch, + check_reserved, + gather::Members, + index_error, + kernels::{LloydScratch, lloyd, nearest}, +}; +use crate::{ + ANNResult, + error::ErrorExt, + ivf::{ + DynamicIvfConfig, + dynamic::{CentroidIndex, MaintenanceAccessor}, + grouped::Csr, + update::{Appends, CentroidBlock, Reassignments, SplitInsert, Splits}, + }, + utils::VectorId, +}; + +/// A list admitted for splitting, with its size in the accessor's view. +#[derive(Debug, Clone, Copy)] +pub(in crate::ivf) struct Parent { + pub(in crate::ivf) list: L, + pub(in crate::ivf) len: usize, +} + +/// Everything one split batch needs, gathered before any computation. +/// +/// Lists are indexed densely: `lists[..parents]` are the split parents in ascending +/// order and `lists[parents..]` are their distinct neighbors in ascending order. +/// Region `r` is parent `lists[r]`, which children `2r` and `2r + 1` replace. +/// Neighbor `n` is `lists[parents + n]`. +pub(in crate::ivf) struct SplitPlan { + batch: StagedBatch, + lists: Vec, + parents: usize, + /// Each region's neighbors. + neighbors: Csr, + /// Each neighbor's regions, ascending. + regions: Csr, + /// Row `n` is the centroid of neighbor `n`. + neighbor_centroids: Matrix, + /// Existing members of every list, indexed like `lists`. + members: Members, + /// Reserved ids for the children, two per region. + child_ids: Vec, +} + +impl SplitPlan { + /// Perform every accessor call the split needs. + /// + /// Reserves two child ids per parent, selects each parent's nearest surviving + /// lists as its neighbors, and reads the members and canonical vectors of every + /// parent and neighbor. + /// + /// # Errors + /// + /// Fails if an accessor call fails or the accessor's view is inconsistent. + pub(in crate::ivf) async fn gather( + accessor: &mut A, + config: &DynamicIvfConfig, + batch: StagedBatch, + parents: Vec>, + ) -> ANNResult + where + A: MaintenanceAccessor, + T: Send, + { + let count = 2 * parents.len(); + let child_ids = accessor + .reserve_list_ids(count) + .await + .escalate("split must reserve child list ids")?; + check_reserved(accessor.centroids(), &child_ids, count)?; + + let dim = accessor.dim(); + let mut lists: Vec = parents.iter().map(|parent| parent.list).collect(); + let centroids = accessor.centroids(); + + // Every parent is retired, so no parent is ever a neighbor. + let mut by_region: Csr = Csr::default(); + for &parent in &lists { + let anchor = centroids + .centroid(parent) + .ok_or_else(|| index_error(format!("split parent {parent} is not live")))?; + let selected = centroids + .select(anchor, config.reassign_neighbors.saturating_add(1)) + .await + .escalate("split must find neighbor lists")?; + by_region.push( + selected + .selected() + .iter() + .map(|selected| selected.id) + .filter(|id| lists.binary_search(id).is_err()) + .take(config.reassign_neighbors), + ); + } + + let mut distinct = by_region.values().to_vec(); + distinct.sort_unstable(); + distinct.dedup(); + + let mut neighbor_centroids = Matrix::new(0.0f32, distinct.len(), dim); + for (row, &list) in distinct.iter().enumerate() { + let centroid = centroids + .centroid(list) + .ok_or_else(|| index_error(format!("neighbor list {list} is not live")))?; + if centroid.len() != dim { + return Err(index_error(format!( + "centroid {list} has dimension {}, expected {dim}", + centroid.len() + ))); + } + neighbor_centroids.row_mut(row).copy_from_slice(centroid); + } + + let mut neighbors = Csr::default(); + for region in 0..by_region.num_groups() { + // Every neighbor is in `distinct`, so each search succeeds. + neighbors.push( + by_region + .group(region) + .iter() + .filter_map(|list| distinct.binary_search(list).ok()), + ); + } + let regions = neighbors.transpose(distinct.len()); + lists.extend_from_slice(&distinct); + + let members = Members::read(accessor, &lists).await?; + for (region, parent) in parents.iter().enumerate() { + let len = members.ids(region).len(); + if len != parent.len { + return Err(index_error(format!( + "split parent {} has {len} members, but its metadata reports {}", + parent.list, parent.len + ))); + } + } + + Ok(Self { + batch, + lists, + parents: parents.len(), + neighbors, + regions, + neighbor_centroids, + members, + child_ids, + }) + } + + /// Fit the children, place every affected point, and build the update. + /// + /// Each parent splits into two children fitted by 2-means over its existing + /// members and the batch points routed to it. Those points then move to the + /// nearest of the parent's neighbors and children. A point of a neighbor list moves + /// to a child of a region containing that list only if the child is closer than + /// the list's centroid (as in SPFresh); otherwise it stays. + /// + /// # Errors + /// + /// Fails if a parent has fewer than two points. + pub(in crate::ivf) fn solve( + mut self, + iterations: usize, + rng: &mut StdRng, + ) -> ANNResult> { + let centers = self.fit_children(iterations, rng)?; + let children = CentroidBlock::new(std::mem::take(&mut self.child_ids), centers); + let mut routes = std::mem::take(&mut self.batch.routes); + let destinations = self.place(&children, &mut routes)?; + Ok(self.package(children, &routes, &destinations)) + } + + /// Fit two children per region; rows `2r` and `2r + 1` split region `r`. + fn fit_children(&self, iterations: usize, rng: &mut StdRng) -> ANNResult> { + let dim = self.batch.vectors.ncols(); + let mut centers = Matrix::new(0.0f32, 2 * self.parents, dim); + let mut scratch = LloydScratch::default(); + for region in 0..self.parents { + let span = self.members.ids(region).len() + self.staged(region).len(); + if span < 2 { + return Err(index_error(format!( + "split parent {} has fewer than two points", + self.lists[region] + ))); + } + let a = rng.random_range(0..span); + let b = (a + rng.random_range(1..span)) % span; + + let rows = &mut centers.as_mut_slice()[2 * region * dim..2 * (region + 1) * dim]; + let (first, second) = rows.split_at_mut(dim); + first.copy_from_slice(self.point(region, a)); + second.copy_from_slice(self.point(region, b)); + lloyd( + || self.points(region), + rows, + dim, + iterations.max(1), + &mut scratch, + ); + } + Ok(centers) + } + + /// Choose a list for every point of every region list. + /// + /// Returns the existing members' lists, aligned with their positions, and + /// overwrites the routes of the batch points it places. + fn place(&self, children: &CentroidBlock, routes: &mut [L]) -> ANNResult> { + let mut destinations = vec![L::default(); self.members.total()]; + let child_ids = children.ids(); + let child_vectors = children.vectors(); + let mut candidates: Vec<(L, &[f32])> = Vec::new(); + + for region in 0..self.parents { + candidates.clear(); + for &neighbor in self.neighbors.group(region) { + candidates.push(self.neighbor(neighbor)); + } + for row in [2 * region, 2 * region + 1] { + candidates.push((child_ids[row], child_vectors.row(row))); + } + self.assign(region, &candidates, &mut destinations, routes)?; + } + + // The list's own centroid comes first, so ties keep the point in place. + for neighbor in 0..self.regions.num_groups() { + candidates.clear(); + candidates.push(self.neighbor(neighbor)); + for ®ion in self.regions.group(neighbor) { + for row in [2 * region, 2 * region + 1] { + candidates.push((child_ids[row], child_vectors.row(row))); + } + } + self.assign( + self.parents + neighbor, + &candidates, + &mut destinations, + routes, + )?; + } + Ok(destinations) + } + + /// Place every point of `list` on its nearest candidate. + fn assign( + &self, + list: usize, + candidates: &[(L, &[f32])], + destinations: &mut [L], + routes: &mut [L], + ) -> ANNResult<()> { + let no_candidates = || index_error("split placement has no candidate lists"); + for (destination, vector) in destinations[self.members.range(list)] + .iter_mut() + .zip(self.members.vectors(list)) + { + *destination = nearest(vector, candidates).ok_or_else(no_candidates)?; + } + for &position in self.staged(list) { + routes[position] = + nearest(self.batch.vectors.row(position), candidates).ok_or_else(no_candidates)?; + } + Ok(()) + } + + /// Assemble the update from the chosen lists. + fn package( + &self, + children: CentroidBlock, + routes: &[L], + destinations: &[L], + ) -> SplitInsert { + let parents = self.lists[..self.parents].to_vec(); + SplitInsert::new( + Appends::new(&self.batch.ids, routes), + Splits::new(parents, children, self.moves(0..self.parents, destinations)), + Reassignments::new(self.moves(self.parents..self.lists.len(), destinations)), + ) + } + + /// `(list, members, destinations)` for every list in `lists`. + fn moves<'s>( + &'s self, + lists: Range, + destinations: &'s [L], + ) -> impl Iterator { + lists.map(move |list| { + ( + self.lists[list], + self.members.ids(list), + &destinations[self.members.range(list)], + ) + }) + } + + /// Neighbor `neighbor`'s list id and centroid. + fn neighbor(&self, neighbor: usize) -> (L, &[f32]) { + ( + self.lists[self.parents + neighbor], + self.neighbor_centroids.row(neighbor), + ) + } + + /// Batch positions routed to `list`. + fn staged(&self, list: usize) -> &[usize] { + self.batch.by_route.get(self.lists[list]) + } + + /// Canonical vectors of `list`'s existing members followed by its batch points. + fn points(&self, list: usize) -> impl Iterator { + self.members.vectors(list).chain( + self.staged(list) + .iter() + .map(|&position| self.batch.vectors.row(position)), + ) + } + + /// The `index`-th vector of [`Self::points`]. + fn point(&self, list: usize, index: usize) -> &[f32] { + let members = self.members.range(list); + if index < members.len() { + self.members.vector(members.start + index) + } else { + self.batch + .vectors + .row(self.staged(list)[index - members.len()]) + } + } +} diff --git a/diskann/src/ivf/test/mod.rs b/diskann/src/ivf/test/mod.rs new file mode 100644 index 0000000000..7792688185 --- /dev/null +++ b/diskann/src/ivf/test/mod.rs @@ -0,0 +1,8 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Test infrastructure for the dynamic IVF index. + +pub(crate) mod provider; diff --git a/diskann/src/ivf/test/provider.rs b/diskann/src/ivf/test/provider.rs new file mode 100644 index 0000000000..a1115e2e29 --- /dev/null +++ b/diskann/src/ivf/test/provider.rs @@ -0,0 +1,494 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! In-memory provider for exercising the dynamic IVF insert path. +//! +//! The provider stores plain owned state and mutates it in place through the +//! exclusive borrow lent to its maintenance accessor. It writes vector rows in +//! reverse order to exercise out-of-order reads, and [`Faults`] injects inconsistent +//! responses. + +use std::collections::{BTreeMap, BTreeSet}; + +use diskann_utils::{future::SendFuture, views::MutMatrixView}; + +use crate::{ + ANNError, ANNErrorKind, + ivf::{ + dynamic::{ + Apply, CentroidIndex, ListMetadata, ListProvider, MaintenanceAccessor, + MaintenanceStrategy, SelectedList, SelectionPlan, + }, + update::{Appends, Bootstrap, SplitInsert}, + }, + provider::{DataProvider, DefaultContext, HasId, NoopGuard}, +}; + +fn error(message: impl Into) -> ANNError { + ANNError::message(ANNErrorKind::IndexError, message.into()) +} + +fn squared(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() +} + +/// Inconsistent responses the provider can inject. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) struct Faults { + /// Leave the last row of every `read_vectors` call unwritten. + pub(crate) skip_vector: bool, + /// Write NaN into the last row of every `read_vectors` call. + pub(crate) nan_vector: bool, + /// Omit the first member of every list in `read_members`. + pub(crate) drop_member: bool, + /// Return the first staged id again from every later `stage_insert`. + pub(crate) repeat_stage: bool, + /// Reserve list ids starting from zero, which may be live. + pub(crate) reuse_lists: bool, +} + +/// Exact centroid catalog. +#[derive(Debug, Clone, Default, PartialEq)] +pub(crate) struct Centroids(BTreeMap>); + +impl CentroidIndex for Centroids { + type ListId = u32; + type Error = ANNError; + + fn len(&self) -> usize { + self.0.len() + } + + fn centroid(&self, id: u32) -> Option<&[f32]> { + self.0.get(&id).map(|centroid| &**centroid) + } + + fn select( + &self, + query: &[f32], + nprobe: usize, + ) -> impl SendFuture, ANNError>> { + let mut selected: Vec<_> = self + .0 + .iter() + .map(|(&id, centroid)| SelectedList { + id, + distance: squared(query, centroid), + }) + .collect(); + selected.sort_by(|a, b| a.distance.total_cmp(&b.distance).then(a.id.cmp(&b.id))); + selected.truncate(nprobe); + std::future::ready(Ok(SelectionPlan::new(selected))) + } +} + +/// An in-memory dynamic IVF provider. Point `i` has internal id `i`. +#[derive(Debug, Clone)] +pub(crate) struct Provider { + dim: usize, + /// External id and canonical vector of every committed point. + points: Vec<(u32, Box<[f32]>)>, + /// Current list of every committed point. + assignment: Vec, + lists: BTreeMap>, + centroids: Centroids, + retired: BTreeSet, + next_list: u32, + faults: Faults, +} + +impl Provider { + /// An empty provider for vectors of dimension `dim`. + pub(crate) fn new(dim: usize) -> Self { + Self { + dim, + points: Vec::new(), + assignment: Vec::new(), + lists: BTreeMap::new(), + centroids: Centroids::default(), + retired: BTreeSet::new(), + next_list: 0, + faults: Faults::default(), + } + } + + /// A provider whose list `i` has centroid `lists[i].0` and members `lists[i].1`, + /// which receive consecutive internal and external ids. + pub(crate) fn from_lists(dim: usize, lists: &[(&[f32], &[&[f32]])]) -> Self { + let mut provider = Self::new(dim); + for (list, &(centroid, members)) in (0..).zip(lists) { + provider.centroids.0.insert(list, centroid.into()); + let ids = members + .iter() + .map(|&vector| { + let id = provider.points.len() as u32; + provider.points.push((id, vector.into())); + provider.assignment.push(list); + id + }) + .collect(); + provider.lists.insert(list, ids); + provider.next_list = list + 1; + } + provider + } + + /// Inject `faults` into subsequent reads. + pub(crate) fn with_faults(mut self, faults: Faults) -> Self { + self.faults = faults; + self + } + + /// Whether both providers hold the same points, lists, and centroids. + pub(crate) fn same_contents(&self, other: &Self) -> bool { + self.points == other.points + && self.assignment == other.assignment + && self.lists == other.lists + && self.centroids == other.centroids + && self.retired == other.retired + } + + /// Live list ids, ascending. + pub(crate) fn lists(&self) -> Vec { + self.lists.keys().copied().collect() + } + + /// Number of retired lists. + pub(crate) fn num_retired(&self) -> usize { + self.retired.len() + } + + /// Number of committed points. + pub(crate) fn len(&self) -> usize { + self.points.len() + } + + /// Centroid of a live list. + pub(crate) fn centroid(&self, list: u32) -> Option<&[f32]> { + self.centroids.centroid(list) + } + + /// Vectors of a live list's members, sorted. + pub(crate) fn member_vectors(&self, list: u32) -> Vec> { + let mut vectors: Vec> = self.lists[&list] + .iter() + .map(|&id| self.points[id as usize].1.to_vec()) + .collect(); + vectors.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + vectors + } + + /// The list holding the point whose vector is `vector`. + pub(crate) fn list_of(&self, vector: &[f32]) -> Option { + self.points + .iter() + .position(|(_, point)| &**point == vector) + .map(|id| self.assignment[id]) + } + + /// Assert every structural invariant of the partition. + pub(crate) fn check(&self) { + assert_eq!(self.lists.len(), self.centroids.0.len()); + let mut seen = vec![0usize; self.points.len()]; + for (&list, members) in &self.lists { + assert!( + self.centroids.0.contains_key(&list), + "list {list} has no centroid" + ); + assert!(!self.retired.contains(&list), "list {list} is retired"); + for &id in members { + seen[id as usize] += 1; + assert_eq!(self.assignment[id as usize], list, "point {id}"); + } + } + assert!( + seen.iter().all(|&count| count == 1), + "every point must be in exactly one list" + ); + assert!(self.centroids.0.values().all(|c| c.len() == self.dim)); + assert!(self.points.iter().all(|(_, v)| v.len() == self.dim)); + } +} + +impl DataProvider for Provider { + type Context = DefaultContext; + type InternalId = u32; + type ExternalId = u32; + type Error = ANNError; + type Guard = NoopGuard; + + fn to_internal_id(&self, _: &DefaultContext, gid: &u32) -> Result { + self.points + .iter() + .position(|(external, _)| external == gid) + .map(|id| id as u32) + .ok_or_else(|| error(format!("no point with external id {gid}"))) + } + + fn to_external_id(&self, _: &DefaultContext, id: u32) -> Result { + self.points + .get(id as usize) + .map(|(external, _)| *external) + .ok_or_else(|| error(format!("no point {id}"))) + } +} + +impl ListProvider for Provider { + type ListId = u32; +} + +/// Maintenance accessor over an exclusively borrowed [`Provider`]. +pub(crate) struct Accessor<'a> { + provider: &'a mut Provider, + /// External id and vector of points staged by this accessor, in id order. + staged: Vec<(u32, Box<[f32]>)>, +} + +impl Accessor<'_> { + fn vector(&self, id: u32) -> Result<&[f32], ANNError> { + let id = id as usize; + let committed = self.provider.points.len(); + let point = if id < committed { + self.provider.points.get(id) + } else { + self.staged.get(id - committed) + }; + point + .map(|(_, vector)| &**vector) + .ok_or_else(|| error(format!("no point {id}"))) + } + + fn members(&self, list: u32) -> Result<&[u32], ANNError> { + let members = self + .provider + .lists + .get(&list) + .ok_or_else(|| error(format!("list {list} is not live")))?; + Ok( + match (self.provider.faults.drop_member, members.split_first()) { + (true, Some((_, rest))) => rest, + _ => members, + }, + ) + } + + fn commit_staged(&mut self) { + for point in self.staged.drain(..) { + self.provider.points.push(point); + self.provider.assignment.push(u32::MAX); + } + } + + fn install(&mut self, id: u32, vector: &[f32]) -> Result<(), ANNError> { + if self.provider.retired.contains(&id) || self.provider.lists.contains_key(&id) { + return Err(error(format!("list id {id} is reused"))); + } + self.provider.centroids.0.insert(id, vector.into()); + self.provider.lists.insert(id, Vec::new()); + Ok(()) + } + + fn place(&mut self, list: u32, ids: &[u32]) -> Result<(), ANNError> { + self.provider + .lists + .get_mut(&list) + .ok_or_else(|| error(format!("list {list} is not live")))? + .extend_from_slice(ids); + for &id in ids { + let slot = self + .provider + .assignment + .get_mut(id as usize) + .ok_or_else(|| error(format!("no point {id}")))?; + *slot = list; + } + Ok(()) + } +} + +impl HasId for Accessor<'_> { + type Id = u32; +} + +impl<'b> MaintenanceAccessor<&'b [f32]> for Accessor<'_> { + type ExternalId = u32; + type ListId = u32; + type Centroids = Centroids; + type Error = ANNError; + + fn dim(&self) -> usize { + self.provider.dim + } + + fn centroids(&self) -> &Centroids { + &self.provider.centroids + } + + async fn list_metadata(&mut self, list: u32) -> Result { + let len = self + .provider + .lists + .get(&list) + .ok_or_else(|| error(format!("list {list} is not live")))? + .len(); + Ok(ListMetadata { len }) + } + + async fn read_members(&mut self, list: u32) -> Result<&[u32], ANNError> { + self.members(list) + } + + async fn read_vectors( + &mut self, + ids: &[u32], + mut out: MutMatrixView<'_, f32>, + ) -> Result<(), ANNError> { + let faults = self.provider.faults; + for (row, &id) in ids.iter().enumerate().rev() { + let last = row + 1 == ids.len(); + if last && faults.skip_vector { + continue; + } + let out = out.row_mut(row); + out.copy_from_slice(self.vector(id)?); + if last && faults.nan_vector { + out[0] = f32::NAN; + } + } + Ok(()) + } + + async fn stage_insert(&mut self, id: &u32, element: &'b [f32]) -> Result { + if element.len() != self.provider.dim { + return Err(error(format!("point {id} has the wrong dimension"))); + } + let known = + |points: &[(u32, Box<[f32]>)]| points.iter().any(|(external, _)| external == id); + if known(&self.provider.points) || known(&self.staged) { + return Err(error(format!("external id {id} already exists"))); + } + let committed = self.provider.points.len(); + let internal = if self.provider.faults.repeat_stage && !self.staged.is_empty() { + committed + } else { + committed + self.staged.len() + }; + self.staged.push((*id, element.into())); + Ok(internal as u32) + } + + async fn reserve_list_ids(&mut self, count: usize) -> Result, ANNError> { + if self.provider.faults.reuse_lists { + return Ok((0..count as u32).collect()); + } + let start = self.provider.next_list; + self.provider.next_list += count as u32; + Ok((start..self.provider.next_list).collect()) + } +} + +impl Apply> for Accessor<'_> { + type Error = ANNError; + + async fn apply(mut self, update: Bootstrap) -> Result<(), ANNError> { + for (id, vector) in update.centroids().iter() { + self.install(id, vector)?; + } + Ok(()) + } +} + +impl Apply> for Accessor<'_> { + type Error = ANNError; + + async fn apply(mut self, update: Appends) -> Result<(), ANNError> { + self.commit_staged(); + for (list, ids) in update.iter() { + self.place(list, ids)?; + } + Ok(()) + } +} + +impl Apply> for Accessor<'_> { + type Error = ANNError; + + async fn apply(mut self, update: SplitInsert) -> Result<(), ANNError> { + self.commit_staged(); + let (appends, splits, reassignments) = update.into_parts(); + + // Install every child first, since any part of the update may target one. + for split in splits.iter() { + for (child, centroid) in split.children().into_iter().zip(split.centroids()) { + self.install(child, centroid)?; + } + } + for split in splits.iter() { + let parent = split.parent(); + let mut members = self + .provider + .lists + .remove(&parent) + .ok_or_else(|| error(format!("split parent {parent} is not live")))?; + self.provider.centroids.0.remove(&parent); + self.provider.retired.insert(parent); + + let mut evacuated: Vec = split + .evacuations() + .flat_map(|(_, ids)| ids.iter().copied()) + .collect(); + members.sort_unstable(); + evacuated.sort_unstable(); + if members != evacuated { + return Err(error(format!( + "split of {parent} does not evacuate its members" + ))); + } + for (list, ids) in split.evacuations() { + self.place(list, ids)?; + } + } + for (transfer, ids) in reassignments.iter() { + let source = self + .provider + .lists + .get_mut(&transfer.source()) + .ok_or_else(|| error(format!("list {} is not live", transfer.source())))?; + let before = source.len(); + source.retain(|id| !ids.contains(id)); + if before - source.len() != ids.len() { + return Err(error(format!( + "reassigned points are not all in list {}", + transfer.source() + ))); + } + self.place(transfer.destination(), ids)?; + } + for (list, ids) in appends.iter() { + self.place(list, ids)?; + } + Ok(()) + } +} + +/// Strategy lending a [`Provider`] to one [`Accessor`]. +#[derive(Debug, Default)] +pub(crate) struct Strategy; + +impl<'a> MaintenanceStrategy<'a, Provider, &[f32]> for Strategy { + type MaintenanceAccessor = Accessor<'a>; + type Error = ANNError; + + fn maintenance_accessor( + &'a self, + provider: &'a mut Provider, + _context: &'a DefaultContext, + ) -> Result, ANNError> { + Ok(Accessor { + provider, + staged: Vec::new(), + }) + } +} diff --git a/diskann/src/ivf/update.rs b/diskann/src/ivf/update.rs new file mode 100644 index 0000000000..82d5698dd7 --- /dev/null +++ b/diskann/src/ivf/update.rs @@ -0,0 +1,533 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Declarative partition updates consumed by [`Apply`]. +//! +//! Each distinct operation is its own type: [`Bootstrap`], [`Appends`], [`Splits`], and +//! [`Reassignments`]. [`SplitInsert`] composes the last three for an insert batch that +//! splits lists. +//! +//! Only the index constructs updates. It validates everything it reads from the +//! provider or its caller, and its planner guarantees the invariants documented on +//! each type by construction. Providers receive updates through [`Apply`] and can only +//! read them. +//! +//! Point memberships are stored grouped: group keys are ascending and unique, and each +//! group's point ids are ascending in one contiguous slice. Providers can therefore +//! write each list once and process distinct groups in parallel. +//! +//! [`Apply`]: crate::ivf::dynamic::Apply + +use std::cmp::Ordering; + +use diskann_utils::views::{Matrix, MatrixView}; + +use crate::{ivf::grouped::Grouped, utils::VectorId}; + +/////////////////// +// CentroidBlock // +/////////////////// + +/// Freshly reserved centroids with full-precision vectors, one row per id. +/// +/// A block is never empty, its ids are fresh and unique, its rows share one non-zero +/// dimension, and every value is finite. +#[derive(Debug, PartialEq)] +pub struct CentroidBlock { + ids: Box<[L]>, + vectors: Matrix, +} + +impl CentroidBlock { + /// Pair `ids[i]` with row `i` of `vectors`. + pub(super) fn new(ids: Vec, vectors: Matrix) -> Self { + debug_assert_eq!(ids.len(), vectors.nrows()); + Self { + ids: ids.into(), + vectors, + } + } + + /// Number of centroids. + #[expect(clippy::len_without_is_empty, reason = "a block is never empty")] + pub fn len(&self) -> usize { + self.ids.len() + } + + /// Dimension of every centroid. + pub fn dim(&self) -> usize { + self.vectors.ncols() + } + + /// Centroid ids in row order. + pub fn ids(&self) -> &[L] { + &self.ids + } + + /// Centroid vectors; row `i` belongs to `self.ids()[i]`. + pub fn vectors(&self) -> MatrixView<'_, f32> { + self.vectors.as_view() + } + + /// Iterate `(id, vector)` pairs in row order. + pub fn iter(&self) -> impl ExactSizeIterator { + self.ids.iter().copied().zip(self.vectors.row_iter()) + } + + /// Release the ids and vectors, for example to adopt the matrix without copying. + pub fn into_parts(self) -> (Box<[L]>, Matrix) { + (self.ids, self.vectors) + } +} + +/////////////// +// Bootstrap // +/////////////// + +/// The initial centroids of an empty index. Installs no points. +#[derive(Debug, PartialEq)] +pub struct Bootstrap { + centroids: CentroidBlock, +} + +impl Bootstrap { + /// Install `centroids` as the first live lists. + pub(super) fn new(centroids: CentroidBlock) -> Self { + Self { centroids } + } + + /// The initial centroids. + pub fn centroids(&self) -> &CentroidBlock { + &self.centroids + } + + /// Release the initial centroids. + pub fn into_centroids(self) -> CentroidBlock { + self.centroids + } +} + +///////////// +// Appends // +///////////// + +/// Newly staged points placed into lists, grouped by destination list. +/// +/// Every point staged for the operation appears exactly once. +#[derive(Debug, PartialEq, Eq)] +pub struct Appends { + by_list: Grouped, +} + +impl Appends { + /// Place `ids[i]` into `lists[i]`. + pub(super) fn new(ids: &[Id], lists: &[L]) -> Self { + debug_assert_eq!(ids.len(), lists.len()); + let pairs = lists.iter().copied().zip(ids.iter().copied()).collect(); + Self { + by_list: Grouped::from_pairs(pairs), + } + } + + /// Number of appended points. + pub fn len(&self) -> usize { + self.by_list.len() + } + + /// Whether no points are appended. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Destination lists, ascending. + pub fn lists(&self) -> &[L] { + self.by_list.keys() + } + + /// Every appended id, grouped by destination list. + pub fn ids(&self) -> &[Id] { + self.by_list.values() + } + + /// Iterate `(list, ids)` groups in ascending list order. + pub fn iter(&self) -> impl ExactSizeIterator { + self.by_list.iter() + } +} + +////////////// +// Transfer // +////////////// + +/// The source and destination of reassigned points. The two never match. +/// +/// Ordered by destination, then source. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Transfer { + source: L, + destination: L, +} + +impl Transfer { + /// The list the points leave. + pub fn source(&self) -> L { + self.source + } + + /// The list the points join. + pub fn destination(&self) -> L { + self.destination + } +} + +impl Ord for Transfer { + fn cmp(&self, other: &Self) -> Ordering { + self.destination + .cmp(&other.destination) + .then_with(|| self.source.cmp(&other.source)) + } +} + +impl PartialOrd for Transfer { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +//////////// +// Splits // +//////////// + +/// Parents retired by binary splits, their two children each, and where every +/// existing member of each parent goes. +/// +/// Parents are ascending and unique. Rows `2 * i` and `2 * i + 1` of +/// [`Splits::children`] replace parent `i`; children are never parents and may end up +/// empty. Every member a parent had when the split was planned is evacuated exactly +/// once, never into a parent. Newly staged points routed to a parent are placed by +/// [`Appends`] instead. +#[derive(Debug, PartialEq)] +pub struct Splits { + parents: Box<[L]>, + children: CentroidBlock, + /// Keyed by `(parent, destination)`. + evacuations: Grouped<(L, L), Id>, +} + +impl Splits { + /// Retire `parents`, replacing parent `i` with rows `2i` and `2i + 1` of `children`. + /// + /// Each item of `evacuations` is `(parent, members, lists)`: `members[j]` moves to + /// `lists[j]`. + pub(super) fn new<'a>( + parents: Vec, + children: CentroidBlock, + evacuations: impl IntoIterator, + ) -> Self { + debug_assert_eq!(children.len(), 2 * parents.len()); + let pairs = evacuations + .into_iter() + .flat_map(|(parent, members, lists)| { + debug_assert_eq!(members.len(), lists.len()); + lists + .iter() + .zip(members) + .map(move |(&list, &id)| ((parent, list), id)) + }) + .collect(); + Self { + parents: parents.into(), + children, + evacuations: Grouped::from_pairs(pairs), + } + } + + /// Number of splits. + #[expect(clippy::len_without_is_empty, reason = "there is always a split")] + pub fn len(&self) -> usize { + self.parents.len() + } + + /// Retired parent lists, ascending. + pub fn parents(&self) -> &[L] { + &self.parents + } + + /// Every child centroid; rows `2 * i` and `2 * i + 1` replace `self.parents()[i]`. + pub fn children(&self) -> &CentroidBlock { + &self.children + } + + /// Number of existing points evacuated from all parents. + pub fn num_evacuated(&self) -> usize { + self.evacuations.len() + } + + /// Iterate the individual splits in ascending parent order. + pub fn iter(&self) -> impl ExactSizeIterator> { + (0..self.parents.len()).map(|index| Split { + splits: self, + index, + }) + } +} + +/// One parent of [`Splits`] and the two children replacing it. +#[derive(Debug, Clone, Copy)] +pub struct Split<'a, Id, L> { + splits: &'a Splits, + index: usize, +} + +impl<'a, Id: VectorId, L: VectorId> Split<'a, Id, L> { + /// The retired parent list. + pub fn parent(self) -> L { + self.splits.parents[self.index] + } + + /// The two new child lists. + pub fn children(self) -> [L; 2] { + let ids = self.splits.children.ids(); + [ids[2 * self.index], ids[2 * self.index + 1]] + } + + /// The two child centroids, in the order of [`Self::children`]. + pub fn centroids(self) -> [&'a [f32]; 2] { + let vectors = &self.splits.children.vectors; + [vectors.row(2 * self.index), vectors.row(2 * self.index + 1)] + } + + /// Where the parent's existing members go, as `(destination, ids)` groups in + /// ascending destination order. + pub fn evacuations(self) -> impl ExactSizeIterator { + let evacuations = &self.splits.evacuations; + let parent = self.parent(); + let start = evacuations + .keys() + .partition_point(|&(from, _)| from < parent); + let end = evacuations + .keys() + .partition_point(|&(from, _)| from <= parent); + evacuations + .groups(start..end) + .map(|((_, destination), ids)| (destination, ids)) + } +} + +/////////////////// +// Reassignments // +/////////////////// + +/// Existing points moved between lists that survive the update. +/// +/// Grouped by [`Transfer`] in ascending destination-then-source order; ids ascend +/// within each transfer. +#[derive(Debug, PartialEq, Eq)] +pub struct Reassignments { + by_transfer: Grouped, Id>, +} + +impl Reassignments { + /// Each item of `moves` is `(source, members, lists)`: `members[j]` moves to + /// `lists[j]`, or stays when that is `source`. + pub(super) fn new<'a>(moves: impl IntoIterator) -> Self { + let pairs = moves + .into_iter() + .flat_map(|(source, members, lists)| { + debug_assert_eq!(members.len(), lists.len()); + lists + .iter() + .zip(members) + .filter(move |&(&destination, _)| destination != source) + .map(move |(&destination, &id)| { + ( + Transfer { + source, + destination, + }, + id, + ) + }) + }) + .collect(); + Self { + by_transfer: Grouped::from_pairs(pairs), + } + } + + /// Number of reassigned points. + pub fn len(&self) -> usize { + self.by_transfer.len() + } + + /// Whether no points are reassigned. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Iterate `(transfer, ids)` groups in ascending destination-then-source order. + pub fn iter(&self) -> impl ExactSizeIterator, &[Id])> { + self.by_transfer.iter() + } +} + +///////////////// +// SplitInsert // +///////////////// + +/// An insert batch that splits lists: the new points, the splits they trigger, and +/// the reassignment of existing points around the splits. +/// +/// Every point appears in exactly one part. Nothing is appended or reassigned into a +/// split parent, and no reassignment leaves a parent (its evacuation empties it) or a +/// new child. +#[derive(Debug, PartialEq)] +pub struct SplitInsert { + appends: Appends, + splits: Splits, + reassignments: Reassignments, +} + +impl SplitInsert { + /// Combine the parts of one split insert. + pub(super) fn new( + appends: Appends, + splits: Splits, + reassignments: Reassignments, + ) -> Self { + Self { + appends, + splits, + reassignments, + } + } + + /// Newly staged points and their final lists. + pub fn appends(&self) -> &Appends { + &self.appends + } + + /// The splits and their parents' evacuations. + pub fn splits(&self) -> &Splits { + &self.splits + } + + /// Existing points moved between surviving lists. + pub fn reassignments(&self) -> &Reassignments { + &self.reassignments + } + + /// Release the parts. + pub fn into_parts(self) -> (Appends, Splits, Reassignments) { + (self.appends, self.splits, self.reassignments) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A block whose values are `0, 1, 2, ...` in row-major order. + fn block(ids: &[u32], dim: usize) -> CentroidBlock { + let data = (0..ids.len() * dim).map(|x| x as f32).collect(); + CentroidBlock::new( + ids.to_vec(), + Matrix::try_from(data, ids.len(), dim).unwrap(), + ) + } + + fn owned<'a, K, V: Copy + 'a>(groups: impl Iterator) -> Vec<(K, Vec)> { + groups.map(|(key, values)| (key, values.to_vec())).collect() + } + + #[test] + fn centroid_block_pairs_ids_with_rows() { + let bootstrap = Bootstrap::new(block(&[4, 2], 3)); + let block = bootstrap.centroids(); + assert_eq!((block.len(), block.dim()), (2, 3)); + assert_eq!(block.ids(), &[4, 2]); + assert_eq!(block.vectors().row(1), &[3.0, 4.0, 5.0]); + assert_eq!( + owned(block.iter()), + vec![(4, vec![0.0, 1.0, 2.0]), (2, vec![3.0, 4.0, 5.0])] + ); + + let (ids, vectors) = bootstrap.into_centroids().into_parts(); + assert_eq!(&*ids, &[4, 2]); + assert_eq!(vectors.nrows(), 2); + } + + #[test] + fn appends_group_ids_by_list() { + let appends = Appends::::new(&[5, 1, 3, 2], &[7, 9, 7, 9]); + assert_eq!(appends.len(), 4); + assert_eq!(appends.lists(), &[7, 9]); + assert_eq!(appends.ids(), &[3, 5, 1, 2]); + assert_eq!( + owned(appends.iter()), + vec![(7, vec![3, 5]), (9, vec![1, 2])] + ); + assert!(Appends::::new(&[], &[]).is_empty()); + } + + #[test] + fn splits_expose_each_split() { + let splits = Splits::::new( + vec![10, 20], + block(&[30, 31, 40, 41], 2), + [(10, &[3, 1, 2][..], &[31, 30, 40][..]), (20, &[], &[])], + ); + assert_eq!(splits.len(), 2); + assert_eq!(splits.parents(), &[10, 20]); + assert_eq!(splits.children().ids(), &[30, 31, 40, 41]); + assert_eq!(splits.num_evacuated(), 3); + + let split: Vec<_> = splits.iter().collect(); + assert_eq!(split[0].parent(), 10); + assert_eq!(split[0].children(), [30, 31]); + assert_eq!(split[0].centroids(), [&[0.0f32, 1.0][..], &[2.0, 3.0][..]]); + assert_eq!( + owned(split[0].evacuations()), + vec![(30, vec![1]), (31, vec![3]), (40, vec![2])] + ); + assert_eq!(split[1].children(), [40, 41]); + assert_eq!(split[1].evacuations().len(), 0); + } + + #[test] + fn reassignments_skip_members_that_stay() { + let reassignments = Reassignments::::new([ + (7, &[1, 2, 3][..], &[7, 9, 8][..]), + (8, &[4, 5][..], &[9, 8][..]), + ]); + assert_eq!(reassignments.len(), 3); + let transfers: Vec<_> = reassignments + .iter() + .map(|(t, ids)| (t.source(), t.destination(), ids.to_vec())) + .collect(); + assert_eq!( + transfers, + vec![(7, 8, vec![3]), (7, 9, vec![2]), (8, 9, vec![4])] + ); + assert!(Reassignments::::new([(7, &[1][..], &[7][..])]).is_empty()); + } + + #[test] + fn split_insert_releases_its_parts() { + let update = SplitInsert::new( + Appends::::new(&[100], &[30]), + Splits::new(vec![10], block(&[30, 31], 1), [(10, &[1][..], &[31][..])]), + Reassignments::new([(7, &[2][..], &[30][..])]), + ); + assert_eq!(update.appends().len(), 1); + assert_eq!(update.splits().num_evacuated(), 1); + assert_eq!(update.reassignments().len(), 1); + let (appends, splits, reassignments) = update.into_parts(); + assert_eq!( + (appends.lists(), splits.parents(), reassignments.len()), + (&[30][..], &[10][..], 1) + ); + } +} diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index b499b137f3..2b8fce1353 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -56,7 +56,7 @@ insert batch may: - route many points against one centroid set; - identify several overflowing lists; -- fit two children per overflowing parent with one joint k-means; +- 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. @@ -130,6 +130,9 @@ coordinates its provider, centroid index, and inverted-list store. - **Point id**: the provider's internal vector id. External ids remain owned by `DataProvider`. - **List id**: the stable logical id shared by a centroid and its inverted list. + Like the point id, it belongs to the provider (`ListProvider::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. - **Selection plan**: selected list ids and their coarse distances. It is @@ -188,8 +191,16 @@ implementation. diskann/src/ivf/ mod.rs dynamic.rs contracts proposed by this RFC + update.rs declarative partition updates consumed by Apply + grouped.rs flat grouped layouts shared by updates and planning + colocation.rs co-location set contracts, not yet wired into maintenance index.rs DynamicIvfIndex initialize/insert orchestration - online/ future GraphIVF algorithm and scratch state + online/ insert planning: staging, split planning, kernels + batch.rs staged and routed insert batch + gather.rs validated reads (members and canonical vectors) + split.rs parent admission and the split plan + 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 @@ -265,14 +276,22 @@ 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 strategy follows the graph module's accessor factory pattern. Providers +declare their list id type once, alongside their point id type: ```rust +pub trait ListProvider: DataProvider { + type ListId: VectorId; +} + pub trait SearchStrategy<'a, Provider, T>: Send + Sync where - Provider: DataProvider, + Provider: ListProvider, { - type SearchAccessor: SearchAccessor; + type SearchAccessor: SearchAccessor< + Id = Provider::InternalId, + ListId = Provider::ListId, + >; type Error: StandardError; fn search_accessor( @@ -325,47 +344,121 @@ 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 Update +### Declarative Partition Updates The algorithm describes the desired logical result rather than issuing ordered -storage commands: +storage commands. Each distinct operation is its own type in `ivf::update`: ```rust -pub struct PointMove { - pub id: Id, - pub from: Option, - pub to: Option, -} +// Fresh centroids: one contiguous matrix, row i belongs to ids[i]. +pub struct CentroidBlock { ids: Box<[L]>, vectors: Matrix } + +pub struct Bootstrap { centroids: CentroidBlock } -pub struct CentroidDelta { - pub insert: Vec>, - pub retire: Vec, +// (i) Newly staged points placed into lists, grouped by destination. +pub struct Appends { /* list -> ids */ } + +// (ii) Binary splits: rows 2i and 2i+1 of `children` replace parents[i], and +// every existing member of each parent is evacuated exactly once. +pub struct Splits { + parents: Box<[L]>, + children: CentroidBlock, + /* (parent, destination) -> ids */ } -pub struct PartitionUpdate { - pub centroids: CentroidDelta, - pub point_moves: Vec>, +// (iii) Existing points moved between surviving lists. +pub struct Reassignments { /* Transfer { source, destination } -> ids */ } + +// An insert batch that splits. +pub struct SplitInsert { + appends: Appends, + splits: Splits, + reassignments: Reassignments, } ``` -`from = None` denotes insertion, `to = None` denotes deletion, and two present -values denote reassignment. Both values may not be absent. Each point occurs at -most once in a normalized update, including when planned split regions overlap. +Fields are private and constructors are private to the IVF module, so only the +index can create updates. Providers receive them through `Apply` and can only +read them. This keeps the provider contract one-directional: the algorithm +decides the update, and providers decide how to store it. + +The planner builds each update from the items it read from the accessor, such as +staged ids or a list's members, together with a decision array aligned index for +index. That makes "every staged point is placed once" and "every parent member +is evacuated once" properties of how the update is built, so constructors only +group their input and perform no validation. + +Several invariants are unrepresentable: + +- a move with neither a source nor a destination; +- centroids of different dimensions; +- retiring a list without evacuating it, since only `Splits` retires lists; +- a dissolve inside an insert or a split inside a delete; and +- a split into anything other than two children. + +The planner guarantees the rest by construction, and each update type documents +them for providers: + +- split children are fresh, and no parent receives points, since parents are + never placement candidates; +- a reassignment never has matching source and destination, and never leaves a + parent or a new child; and +- every point appears in exactly one part of an update. + +Children may end up empty. + +#### Layout + +Point memberships are grouped: group keys are ascending and unique, and each +group's point ids are ascending in one contiguous slice. `Appends` groups by +destination list; `Reassignments` groups by transfer in destination-then-source +order; `Splits` groups each parent's evacuation by destination. A provider can +therefore append each list's incoming points with one write, drop a split +parent wholesale instead of removing its points one by one, re-encode residual +payloads against one centroid per group, and process distinct groups in +parallel. No point vectors cross the update boundary; child centroids cross +once, in one matrix that k-means writes in place and that a graph centroid index +can bulk-insert. Owned buffers move into the update in O(1), and `into_parts` +lets a provider adopt them. -Before returning success, an implementation validates at least: +#### Validating Inputs Instead Of Updates -- inserted centroid ids were reserved by this accessor; -- retired centroid ids exist in the accessor's planning view; -- every `from` assignment matches that view; -- every `to` list exists in the resulting live set; -- no point or centroid id occurs in conflicting operations; and -- staged point insertions and deletions agree with point moves. +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; +- accessor reads: every vector row must be written with finite values (rows + start as NaN, so unwritten rows are caught by the same check), and each split + parent's members must match its reported size; +- staging: staged ids must be distinct; 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: + +- child and bootstrap ids were reserved by this accessor; +- split parents and destinations are live in the accessor's planning view; +- every evacuated or reassigned member 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. -The co-location set updates follow an identical pattern. +Co-location set updates are expected to follow the same declarative pattern. +They are deferred until the insert path and maintenance accessor stabilize: +`CoLocationSet`, `CoLocationGroup`, and `CoLocationDelta` live in +`ivf::colocation`, and no update type or `MaintenanceAccessor` references them +yet. ### Maintenance Accessor @@ -384,52 +477,61 @@ where >; type Error: ToRanked + Debug + Send + Sync + 'static; + fn dim(&self) -> usize; fn centroids(&self) -> &Self::Centroids; - // Batched planning reads: - // list_metadata, assignments, read_members, read_vectors - - fn stage_insert( + fn list_metadata( &mut self, - id: &Self::ExternalId, - element: T, - ) -> impl SendFuture>; + list: Self::ListId, + ) -> impl SendFuture>; - fn read_staged_vectors( + fn read_members( &mut self, - ids: I, - emit: F, - ) -> impl SendFuture> - where - I: Iterator + Send, - F: FnMut(Self::Id, &[f32]) + Send; + list: Self::ListId, + ) -> impl SendFuture>; - fn stage_delete( + fn read_vectors( &mut self, - id: Self::Id, + ids: &[Self::Id], + out: MutMatrixView<'_, f32>, ) -> impl SendFuture>; + fn stage_insert( + &mut self, + id: &Self::ExternalId, + element: T, + ) -> impl SendFuture>; + fn reserve_list_ids( &mut self, count: usize, ) -> impl SendFuture, Self::Error>>; +} - fn apply( - self, - update: PartitionUpdate, - ) -> impl SendFuture>; +pub trait Apply: Send + Sized { + type Error: ToRanked + Debug + Send + Sync + 'static; + + fn apply(self, update: U) -> impl SendFuture>; } ``` -The exact read signatures live in `ivf::dynamic`; all use callbacks so providers -can batch or stream results. `read_members` may invoke its callback more than -once per list. +Reads are shaped by how much data they move. List reads, `list_metadata` 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 staged point and +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_insert` owns external/internal id allocation and canonical input storage -for this operation. `read_staged_vectors` exposes the provider's canonical -`f32` decoding without requiring those points to be visible through ordinary -provider reads. `stage_delete` records point removal while keeping whatever data -the accessor needs for dissolve planning. +for this operation. `read_vectors` also covers points staged by this accessor, +exposing the provider's canonical `f32` decoding without making those points +visible through ordinary provider reads. 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 @@ -438,7 +540,11 @@ 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. It is not a generic transaction protocol: +logical change. A provider implements `Apply` once per update type it +supports: `Bootstrap`, `Appends`, and `SplitInsert` for the insert path, or only +`Appends` for a provider that never splits. The index requires exactly the +`Apply` bounds each operation uses. `apply` is not a generic transaction +protocol: - `Ok(())` guarantees a coherent resulting index. - `Err` does not guarantee rollback or identify a conflict. @@ -449,6 +555,38 @@ 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, Provider, T>: Send + Sync +where + Provider: ListProvider, + T: Send, +{ + type MaintenanceAccessor: MaintenanceAccessor< + T, + Id = Provider::InternalId, + ExternalId = Provider::ExternalId, + ListId = Provider::ListId, + >; + type Error: StandardError; + + fn maintenance_accessor( + &'a self, + provider: &'a mut Provider, + context: &'a Provider::Context, + ) -> Result; +} +``` + +Search strategies receive `&'a Provider`; maintenance strategies receive +`&'a mut Provider`. See [Index Borrowing Model](#index-borrowing-model). + +Because both strategies fix their accessors' ids to the provider's, the index +names updates directly in terms of the provider, for example +`S::MaintenanceAccessor: Apply>`. + ### Inverted-List Store Boundary This RFC does not add a low-level `InvertedListStore` trait. Storage backends @@ -467,7 +605,7 @@ for compatibility with the selected centroid index and `DataProvider`. The initial dynamic index uses exclusive mutation: ```rust -impl DynamicIvfIndex

{ +impl DynamicIvfIndex

{ pub fn search(&self, /* ... */) -> impl SendFuture>; pub fn insert_batch(&mut self, /* ... */) @@ -480,12 +618,23 @@ impl DynamicIvfIndex

{ 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 Provider` 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 `&Provider`, 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. Providers that expose such paths are -responsible for coordinating them. +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. ## Algorithms @@ -497,7 +646,8 @@ Bootstrap uses the same maintenance interface as later mutations: 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 a partition update containing those centroid records and no points. +4. Apply a `Bootstrap` holding those centroids as one `CentroidBlock` and no + points. Initial sample points are not implicitly assigned. Callers insert them through the normal batch path. Reopen skips bootstrap when centroids already exist. @@ -522,33 +672,49 @@ Note that co-location groups are never accessed during search. 1. Construct one maintenance accessor for the operation. 2. Stage every external id and input element, obtaining internal ids. Duplicate or already-live external ids fail according to provider semantics. -3. Batch-materialize canonical vectors for the staged points. -4. Route all inputs through the accessor's centroid index. +3. Read the staged points' canonical vectors into one matrix. +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 and fresh-id budget. -7. Read each parent's members and nearest live centroid neighbors. -8. Fit `2l` children with one joint k-means over `l` admitted parent regions, - seeded with two members per parent. -9. Reserve stable ids for all child centroids. -10. Reassign affected points exactly among each region's surviving neighbors and - two new children. -11. Normalize overlapping regions into one final move per point. -12. Build one `PartitionUpdate` containing centroid changes and point moves. +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. If nothing splits, build `Appends` from the routes. Otherwise build + `Appends` for the staged points' final lists, `Splits` evacuating every + parent, and `Reassignments` for moved neighbor members, and combine them + into one `SplitInsert`. 13. Consume the accessor with `apply(update)`. -With one overflow parent, joint `2l`-means reduces to local 2-means. Routes are -computed before structural change; regional reassignment deliberately revisits -points whose routes become stale due to the split. +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. Candidate centroids for a region are -the parent's selected live neighbors plus its two children. Candidate points +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 `PartitionUpdate` semantics. +in-memory lists. Those choices do not change the semantics of the update. ### Batch Delete And Dissolve @@ -561,8 +727,10 @@ in-memory lists. Those choices do not change `PartitionUpdate` semantics. 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 containing point deletes, retired victim centroids, and - reassigned victim members. +8. Build one delete update containing point deletes, retired victim centroids, + and the evacuation of each victim's remaining members. Its types are not + defined yet; they follow the same pattern as the insert updates, with an + evacuation shared with `Splits`. 9. Consume the accessor with `apply(update)`. A dissolve fits no centroid and does not read survivor members. Removing a @@ -576,7 +744,8 @@ insert-driven split. The portable baseline is: - many concurrent searches through shared index borrows; -- one mutation through an exclusive index borrow; and +- one mutation through an exclusive index borrow, forwarded to the maintenance + accessor as `&mut Provider`; and - no search/mutation overlap through the same index value. Within an operation, providers may parallelize routing, list reads, decoding, @@ -587,7 +756,7 @@ Concrete implementations may offer more: | Provider | Possible private unified-view mechanism | |---|---| -| In memory | exclusive index borrow, one lock guard, copy-on-write root, or epoch | +| In memory | exclusive `&mut Provider` 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 | @@ -653,7 +822,8 @@ and search parameters. Durable metadata belongs to concrete provider crates. - Build one provider containing canonical vectors, reverse assignments, in-memory lists, and an exact centroid index. -- Implement maintenance with exclusive access and in-place `apply`. +- Implement maintenance through the exclusive `&mut Provider` borrow with + lock-free, in-place `apply`. - Add an invariant checker and provider poisoning for injected partial failures. ### Phase 2: Dynamic Search @@ -665,8 +835,8 @@ and search parameters. Durable metadata belongs to concrete provider crates. ### Phase 3: Exact-Centroid GraphIVF -- Port batch routing, split admission, joint child fitting, and tiled regional - reassignment from the prototype. +- 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. @@ -694,7 +864,7 @@ and search parameters. Durable metadata belongs to concrete provider crates. 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 -`PartitionUpdate`. +update. Required shared tests include: @@ -769,9 +939,31 @@ simpler. ### Imperative List Mutation Methods Adding `append`, `remove`, `move`, `create_list`, and `retire_list` exposes a -partial-write schedule to GraphIVF. A complete `PartitionUpdate` lets each +partial-write schedule to GraphIVF. A complete declarative update lets each provider choose a suitable physical implementation. +### One Monolithic Update With Per-Point Moves + +An earlier draft used one `PartitionUpdate` holding a free-form centroid delta +and a flat list of `PointMove { id, from: Option, to: Option }`. It was +compact to define but overloaded: one type encoded insert, delete, reassignment, +split, and dissolve, so every provider had to handle every combination and +regroup moves by list. Its public fields also allowed invalid values, such as a +move with neither endpoint, a retired list that was never emptied, or ragged +centroid dimensions. Separate types that only the index can construct remove +those states and hand providers data that is already grouped. + +### 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 @@ -793,11 +985,12 @@ useful common shape. 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 `stage_insert` and `stage_delete` remain on the - maintenance accessor, or should a future `TransactionalDataProvider` supply - them? -4. **Update granularity**: keep one normalized move per point, or add optional - whole-list replacement hints for bulk backends? +3. **Point staging**: should `stage_insert`, and deletion staging once it is + added, remain on the maintenance accessor, or should a future + `TransactionalDataProvider` supply them? +4. **Update granularity**: updates are grouped by destination list, so bulk + backends can rewrite each touched list once. Should they also carry a + per-list summary of added and removed counts for exact pre-sizing? 5. **Selection plan ownership**: is the documented same-accessor requirement sufficient, or should a future private token make misuse dynamically detectable? From a15c4797ca1d0464f873638e629727ced28dbcee Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Thu, 1 Oct 2026 10:01:18 -0400 Subject: [PATCH 08/33] cutting superfluous things --- diskann/src/ivf/colocation.rs | 56 ---------------- diskann/src/ivf/dynamic.rs | 60 +++++------------ diskann/src/ivf/grouped.rs | 5 -- diskann/src/ivf/index.rs | 18 ++--- diskann/src/ivf/mod.rs | 3 +- diskann/src/ivf/online/batch.rs | 8 +-- diskann/src/ivf/online/gather.rs | 111 ------------------------------- diskann/src/ivf/online/mod.rs | 40 ++++++++++- diskann/src/ivf/online/split.rs | 78 +++++++++++----------- diskann/src/ivf/test/provider.rs | 28 +++----- diskann/src/ivf/update.rs | 33 ++------- rfcs/01187-incremental-ivf.md | 55 +++++++-------- 12 files changed, 142 insertions(+), 353 deletions(-) delete mode 100644 diskann/src/ivf/colocation.rs delete mode 100644 diskann/src/ivf/online/gather.rs diff --git a/diskann/src/ivf/colocation.rs b/diskann/src/ivf/colocation.rs deleted file mode 100644 index 9f42137f0f..0000000000 --- a/diskann/src/ivf/colocation.rs +++ /dev/null @@ -1,56 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Co-location groups of centroids for locality-preserving list storage. -//! -//! Not yet integrated with [`super::dynamic::MaintenanceAccessor`] or the updates in -//! [`super::update`]. - -use crate::utils::VectorId; - -/// An authoritative grouping of live centroids for locality-preserving storage. -/// -/// Every live centroid belongs to exactly one co-location group. Implementations -/// may use an identity grouping, where each group contains one centroid. Group -/// ids remain stable until retired and are independent of centroid/list ids. -pub trait CoLocationSet: Send + Sync { - /// Stable logical centroid/list id. - type ListId: VectorId; - - /// Stable logical co-location group id. - type CoLocationGroupId: VectorId; - - /// Number of live co-location groups. - fn len(&self) -> usize; - - /// Whether there are no live co-location groups. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Borrow the centroid ids in a live co-location group. - fn group(&self, id: Self::CoLocationGroupId) -> Option<&[Self::ListId]>; - - /// Return the live co-location group containing a centroid. - fn group_for(&self, id: Self::ListId) -> Option; -} - -/// A new co-location group to install. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CoLocationGroup { - /// Fresh stable group id. Installed ids must never be reused after retirement. - pub id: CoLocationGroupId, - /// Live centroids that should be stored together. - pub centroids: Vec, -} - -/// Changes to the authoritative co-location set. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct CoLocationDelta { - /// Co-location groups made live by this update. - pub insert: Vec>, - /// Live co-location groups retired by this update. - pub retire: Vec, -} diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index b84da71b68..ba539e13ec 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -5,11 +5,10 @@ //! Contracts for an incrementally maintained IVF index. //! -//! This module is an API skeleton. It defines the boundary between an IVF -//! algorithm and provider-specific centroid navigation, inverted-list access, -//! and partition mutation. Concrete accessors are responsible for presenting a -//! coherent view across those components. The online split/dissolve algorithm -//! is not implemented here yet. +//! This module defines the boundary between the IVF algorithm and +//! provider-specific centroid navigation, inverted-list access, and partition +//! mutation. Concrete accessors are responsible for presenting a coherent view +//! across those components. use std::fmt::Debug; @@ -21,42 +20,15 @@ use crate::{ utils::VectorId, }; -/// One centroid/list selected during coarse search. +/// One centroid/list selected for a query. #[derive(Debug, Clone, Copy, PartialEq)] pub struct SelectedList { /// Stable logical list identifier. pub id: L, - /// Distance from the bound query to this list's centroid. + /// Distance from the query to this list's centroid. pub distance: f32, } -/// The handoff from centroid selection to inverted-list scanning. -/// -/// A plan must be consumed by the same [`SearchAccessor`] that produced it. -/// The accessor is responsible for keeping centroid selection and list reads -/// coherent using provider-specific coordination. -#[derive(Debug, Clone)] -pub struct SelectionPlan { - selected: Vec>, -} - -impl SelectionPlan { - /// Construct a plan from selected logical lists. - pub fn new(selected: Vec>) -> Self { - Self { selected } - } - - /// Lists selected in increasing coarse-distance order. - pub fn selected(&self) -> &[SelectedList] { - &self.selected - } - - /// Consume the plan and return the selected lists. - pub fn into_selected(self) -> Vec> { - self.selected - } -} - /// Work performed while scanning the selected lists. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct ScanStats { @@ -96,7 +68,8 @@ pub trait CentroidIndex: Send + Sync { /// Borrow one authoritative full-precision centroid vector. fn centroid(&self, id: Self::ListId) -> Option<&[f32]>; - /// Select exactly `min(nprobe, self.len())` live centroids for `query`. + /// Select exactly `min(nprobe, self.len())` live centroids for `query`, nearest + /// first. /// /// A graph implementation must use exact selection as a recovery path when /// graph navigation produces too few live ids. @@ -104,14 +77,14 @@ pub trait CentroidIndex: Send + Sync { &self, query: &[f32], nprobe: usize, - ) -> impl SendFuture, Self::Error>>; + ) -> impl SendFuture>, Self::Error>>; } /// Query-bound access to centroid selection and inverted-list scanning. /// /// This is the dynamic IVF search algorithm's primary extension point, in the -/// same spirit as [`crate::graph::SearchAccessor`]. Implementations are free to -/// batch reads, coalesce blob requests, prefetch, decode quantized payloads, or +/// same spirit as [`crate::graph::glue::SearchAccessor`]. Implementations are free +/// to batch reads, coalesce blob requests, prefetch, decode quantized payloads, or /// fan work out across tasks. pub trait SearchAccessor: HasId + Send + Sync { /// Stable logical centroid/list id, fixed to [`ListProvider::ListId`] by the @@ -121,16 +94,19 @@ pub trait SearchAccessor: HasId + Send + Sync { /// Errors from list selection or scanning. type Error: ToRanked + Debug + Send + Sync + 'static; - /// Select exactly the available requested number of lists for the bound query. + /// Select up to `nprobe` lists for the bound query, nearest first. fn select_lists( &mut self, nprobe: usize, - ) -> impl SendFuture, Self::Error>>; + ) -> impl SendFuture>, Self::Error>>; - /// Scan every list in `plan`, scoring members against the bound query. + /// Scan `lists`, scoring their members against the bound query. + /// + /// `lists` must come from [`Self::select_lists`] on this accessor, which keeps + /// selection and scanning coherent using provider-specific coordination. fn scan_lists( &mut self, - plan: SelectionPlan, + lists: &[SelectedList], emit: F, ) -> impl SendFuture> where diff --git a/diskann/src/ivf/grouped.rs b/diskann/src/ivf/grouped.rs index 7fd22cb2c5..6ae4828bdb 100644 --- a/diskann/src/ivf/grouped.rs +++ b/diskann/src/ivf/grouped.rs @@ -115,11 +115,6 @@ impl Grouped { &self.keys } - /// Every value, group by group. - pub(super) fn values(&self) -> &[V] { - self.groups.values() - } - /// The values under `key`, or nothing if `key` is absent. pub(super) fn get(&self, key: K) -> &[V] { self.keys diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 0c576d19c5..3eb0e8c774 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -18,7 +18,7 @@ use crate::{ error::{ErrorExt, IntoANNResult}, ivf::{ dynamic::{Apply, CentroidIndex, ListProvider, MaintenanceAccessor, MaintenanceStrategy}, - online::{Parent, SplitPlan, StagedBatch, check_reserved, index_error}, + online::{SplitPlan, StagedBatch, check_reserved, index_error}, update::{Appends, Bootstrap, CentroidBlock, SplitInsert}, }, }; @@ -101,11 +101,6 @@ impl DynamicIvfIndex

{ &self.provider } - /// Borrow the split policy. - pub fn config(&self) -> &DynamicIvfConfig { - &self.config - } - /// Install the initial centroids. /// /// # Errors @@ -253,7 +248,7 @@ impl DynamicIvfIndex

{ config: &DynamicIvfConfig, accessor: &mut A, batch: &StagedBatch, - ) -> ANNResult>> + ) -> ANNResult> where A: MaintenanceAccessor, T: Send, @@ -267,7 +262,7 @@ impl DynamicIvfIndex

{ .len; let projected = len + incoming; if projected > config.split_threshold { - overflowing.push((projected, Parent { list, len })); + overflowing.push((projected, list)); } } @@ -275,11 +270,11 @@ impl DynamicIvfIndex

{ let budget = config.max_clusters.map_or(usize::MAX, |max| { max.saturating_sub(accessor.centroids().len()) }); - overflowing.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.list.cmp(&b.1.list))); + overflowing.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.cmp(&b.1))); overflowing.truncate(budget); - let mut parents: Vec<_> = overflowing.into_iter().map(|(_, parent)| parent).collect(); - parents.sort_unstable_by_key(|parent| parent.list); + let mut parents: Vec<_> = overflowing.into_iter().map(|(_, list)| list).collect(); + parents.sort_unstable(); Ok(parents) } } @@ -518,7 +513,6 @@ mod tests { (fault(|f| f.nan_vector = true), "no finite canonical vector"), (fault(|f| f.repeat_stage = true), "staging returned point"), (fault(|f| f.reuse_lists = true), "reserved live list id"), - (fault(|f| f.drop_member = true), "but its metadata reports"), ]; for (faults, expected) in cases { // Inserting 0.5 and 8 overflows list 0, so every case reaches the split path. diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 140caab3cf..d138b4094a 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -5,10 +5,9 @@ //! Incrementally maintained IVF index. -pub mod colocation; pub mod dynamic; mod grouped; -pub mod index; +mod index; mod online; pub mod update; diff --git a/diskann/src/ivf/online/batch.rs b/diskann/src/ivf/online/batch.rs index 9bcf37e61a..e8b9aebeab 100644 --- a/diskann/src/ivf/online/batch.rs +++ b/diskann/src/ivf/online/batch.rs @@ -7,7 +7,7 @@ use diskann_utils::views::Matrix; -use super::{first_repeat, gather::read_rows, index_error}; +use super::{first_repeat, index_error, read_rows}; use crate::{ ANNResult, error::ErrorExt, @@ -72,12 +72,10 @@ impl StagedBatch { let centroids = accessor.centroids(); let mut routes = Vec::with_capacity(ids.len()); for vector in vectors.row_iter() { - let plan = centroids + let list = centroids .select(vector, 1) .await - .escalate("insert must route every point")?; - let list = plan - .selected() + .escalate("insert must route every point")? .first() .map(|selected| selected.id) .ok_or_else(|| index_error("centroid selection returned no list"))?; diff --git a/diskann/src/ivf/online/gather.rs b/diskann/src/ivf/online/gather.rs deleted file mode 100644 index 41b852805e..0000000000 --- a/diskann/src/ivf/online/gather.rs +++ /dev/null @@ -1,111 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Validated reads from a maintenance accessor. - -use std::ops::Range; - -use diskann_utils::views::Matrix; - -use super::index_error; -use crate::{ - ANNResult, - error::ErrorExt, - ivf::{dynamic::MaintenanceAccessor, grouped::Csr}, - utils::VectorId, -}; - -/// Read the canonical vector of every point in `ids`; row `i` belongs to `ids[i]`. -/// -/// # Errors -/// -/// Fails if the accessor fails, or leaves a row unwritten or non-finite. -pub(super) async fn read_rows( - accessor: &mut A, - ids: &[A::Id], - dim: usize, -) -> ANNResult> -where - A: MaintenanceAccessor, - T: Send, -{ - // Rows the accessor never writes stay NaN, so the finiteness check catches them. - let mut rows = Matrix::new(f32::NAN, ids.len(), dim); - accessor - .read_vectors(ids, rows.as_mut_view()) - .await - .escalate("maintenance must read canonical vectors")?; - - if let Some((id, _)) = ids - .iter() - .zip(rows.row_iter()) - .find(|(_, row)| row.iter().any(|x| !x.is_finite())) - { - return Err(index_error(format!( - "no finite canonical vector for point {id}" - ))); - } - Ok(rows) -} - -/// The existing members of a table of lists and their canonical vectors. -/// -/// List `i`'s members occupy positions [`Self::range`]`(i)`, which index both the ids -/// and the vector rows. -pub(super) struct Members { - ids: Csr, - vectors: Matrix, -} - -impl Members { - /// Read the members of every list in `lists`, then their canonical vectors. - /// - /// # Errors - /// - /// Fails if the accessor fails or [`read_rows`] fails. - pub(super) async fn read(accessor: &mut A, lists: &[A::ListId]) -> ANNResult - where - A: MaintenanceAccessor, - T: Send, - { - let mut ids = Csr::default(); - for &list in lists { - let members = accessor - .read_members(list) - .await - .escalate("maintenance must read list members")?; - ids.push(members.iter().copied()); - } - - let dim = accessor.dim(); - let vectors = read_rows(accessor, ids.values(), dim).await?; - Ok(Self { ids, vectors }) - } - - /// Total number of members across all lists. - pub(super) fn total(&self) -> usize { - self.ids.values().len() - } - - /// Positions of `list`'s members. - pub(super) fn range(&self, list: usize) -> Range { - self.ids.range(list) - } - - /// Member ids of `list`. - pub(super) fn ids(&self, list: usize) -> &[Id] { - self.ids.group(list) - } - - /// The canonical vector at member position `position`. - pub(super) fn vector(&self, position: usize) -> &[f32] { - self.vectors.row(position) - } - - /// Canonical vectors of `list`'s members, in member order. - pub(super) fn vectors(&self, list: usize) -> impl ExactSizeIterator { - self.range(list).map(|position| self.vectors.row(position)) - } -} diff --git a/diskann/src/ivf/online/mod.rs b/diskann/src/ivf/online/mod.rs index d1cc4d9b75..c4831a2326 100644 --- a/diskann/src/ivf/online/mod.rs +++ b/diskann/src/ivf/online/mod.rs @@ -12,16 +12,21 @@ //! mutates the index; the caller applies the planned update. mod batch; -mod gather; mod kernels; mod split; use std::fmt::{Debug, Display}; pub(super) use batch::StagedBatch; -pub(super) use split::{Parent, SplitPlan}; +pub(super) use split::SplitPlan; -use crate::{ANNError, ANNErrorKind, ANNResult, ivf::dynamic::CentroidIndex}; +use diskann_utils::views::Matrix; + +use crate::{ + ANNError, ANNErrorKind, ANNResult, + error::ErrorExt, + ivf::dynamic::{CentroidIndex, MaintenanceAccessor}, +}; #[track_caller] pub(super) fn index_error(message: D) -> ANNError @@ -62,3 +67,32 @@ fn first_repeat(values: impl Iterator) -> Option { .find(|pair| pair[0] == pair[1]) .map(|pair| pair[0]) } + +/// Read the canonical vector of every point in `ids`; row `i` belongs to `ids[i]`. +/// +/// # Errors +/// +/// Fails if the accessor fails, or leaves a row unwritten or non-finite. +async fn read_rows(accessor: &mut A, ids: &[A::Id], dim: usize) -> ANNResult> +where + A: MaintenanceAccessor, + T: Send, +{ + // Rows the accessor never writes stay NaN, so the finiteness check catches them. + let mut rows = Matrix::new(f32::NAN, ids.len(), dim); + accessor + .read_vectors(ids, rows.as_mut_view()) + .await + .escalate("maintenance must read canonical vectors")?; + + if let Some((id, _)) = ids + .iter() + .zip(rows.row_iter()) + .find(|(_, row)| row.iter().any(|x| !x.is_finite())) + { + return Err(index_error(format!( + "no finite canonical vector for point {id}" + ))); + } + Ok(rows) +} diff --git a/diskann/src/ivf/online/split.rs b/diskann/src/ivf/online/split.rs index 087a038e9c..e1adfcd69d 100644 --- a/diskann/src/ivf/online/split.rs +++ b/diskann/src/ivf/online/split.rs @@ -13,10 +13,9 @@ use rand::{Rng, rngs::StdRng}; use super::{ batch::StagedBatch, - check_reserved, - gather::Members, - index_error, + check_reserved, index_error, kernels::{LloydScratch, lloyd, nearest}, + read_rows, }; use crate::{ ANNResult, @@ -30,13 +29,6 @@ use crate::{ utils::VectorId, }; -/// A list admitted for splitting, with its size in the accessor's view. -#[derive(Debug, Clone, Copy)] -pub(in crate::ivf) struct Parent { - pub(in crate::ivf) list: L, - pub(in crate::ivf) len: usize, -} - /// Everything one split batch needs, gathered before any computation. /// /// Lists are indexed densely: `lists[..parents]` are the split parents in ascending @@ -53,8 +45,10 @@ pub(in crate::ivf) struct SplitPlan { regions: Csr, /// Row `n` is the centroid of neighbor `n`. neighbor_centroids: Matrix, - /// Existing members of every list, indexed like `lists`. - members: Members, + /// Existing members of every list, grouped like `lists`. + member_ids: Csr, + /// Row `i` is the canonical vector of `member_ids.values()[i]`. + member_vectors: Matrix, /// Reserved ids for the children, two per region. child_ids: Vec, } @@ -64,7 +58,7 @@ impl SplitPlan { /// /// Reserves two child ids per parent, selects each parent's nearest surviving /// lists as its neighbors, and reads the members and canonical vectors of every - /// parent and neighbor. + /// parent and neighbor. `parents` must be ascending. /// /// # Errors /// @@ -73,7 +67,7 @@ impl SplitPlan { accessor: &mut A, config: &DynamicIvfConfig, batch: StagedBatch, - parents: Vec>, + parents: Vec, ) -> ANNResult where A: MaintenanceAccessor, @@ -87,12 +81,11 @@ impl SplitPlan { check_reserved(accessor.centroids(), &child_ids, count)?; let dim = accessor.dim(); - let mut lists: Vec = parents.iter().map(|parent| parent.list).collect(); let centroids = accessor.centroids(); // Every parent is retired, so no parent is ever a neighbor. let mut by_region: Csr = Csr::default(); - for &parent in &lists { + for &parent in &parents { let anchor = centroids .centroid(parent) .ok_or_else(|| index_error(format!("split parent {parent} is not live")))?; @@ -102,10 +95,9 @@ impl SplitPlan { .escalate("split must find neighbor lists")?; by_region.push( selected - .selected() .iter() .map(|selected| selected.id) - .filter(|id| lists.binary_search(id).is_err()) + .filter(|id| parents.binary_search(id).is_err()) .take(config.reassign_neighbors), ); } @@ -139,27 +131,30 @@ impl SplitPlan { ); } let regions = neighbors.transpose(distinct.len()); + + let num_parents = parents.len(); + let mut lists = parents; lists.extend_from_slice(&distinct); - let members = Members::read(accessor, &lists).await?; - for (region, parent) in parents.iter().enumerate() { - let len = members.ids(region).len(); - if len != parent.len { - return Err(index_error(format!( - "split parent {} has {len} members, but its metadata reports {}", - parent.list, parent.len - ))); - } + let mut member_ids = Csr::default(); + for &list in &lists { + let members = accessor + .read_members(list) + .await + .escalate("split must read list members")?; + member_ids.push(members.iter().copied()); } + let member_vectors = read_rows(accessor, member_ids.values(), dim).await?; Ok(Self { batch, lists, - parents: parents.len(), + parents: num_parents, neighbors, regions, neighbor_centroids, - members, + member_ids, + member_vectors, child_ids, }) } @@ -193,7 +188,7 @@ impl SplitPlan { let mut centers = Matrix::new(0.0f32, 2 * self.parents, dim); let mut scratch = LloydScratch::default(); for region in 0..self.parents { - let span = self.members.ids(region).len() + self.staged(region).len(); + let span = self.member_ids.group(region).len() + self.staged(region).len(); if span < 2 { return Err(index_error(format!( "split parent {} has fewer than two points", @@ -223,7 +218,7 @@ impl SplitPlan { /// Returns the existing members' lists, aligned with their positions, and /// overwrites the routes of the batch points it places. fn place(&self, children: &CentroidBlock, routes: &mut [L]) -> ANNResult> { - let mut destinations = vec![L::default(); self.members.total()]; + let mut destinations = vec![L::default(); self.member_ids.values().len()]; let child_ids = children.ids(); let child_vectors = children.vectors(); let mut candidates: Vec<(L, &[f32])> = Vec::new(); @@ -267,9 +262,9 @@ impl SplitPlan { routes: &mut [L], ) -> ANNResult<()> { let no_candidates = || index_error("split placement has no candidate lists"); - for (destination, vector) in destinations[self.members.range(list)] + for (destination, vector) in destinations[self.member_ids.range(list)] .iter_mut() - .zip(self.members.vectors(list)) + .zip(self.member_rows(list)) { *destination = nearest(vector, candidates).ok_or_else(no_candidates)?; } @@ -304,8 +299,8 @@ impl SplitPlan { lists.map(move |list| { ( self.lists[list], - self.members.ids(list), - &destinations[self.members.range(list)], + self.member_ids.group(list), + &destinations[self.member_ids.range(list)], ) }) } @@ -323,9 +318,16 @@ impl SplitPlan { self.batch.by_route.get(self.lists[list]) } + /// Canonical vectors of `list`'s existing members. + fn member_rows(&self, list: usize) -> impl Iterator { + self.member_ids + .range(list) + .map(|row| self.member_vectors.row(row)) + } + /// Canonical vectors of `list`'s existing members followed by its batch points. fn points(&self, list: usize) -> impl Iterator { - self.members.vectors(list).chain( + self.member_rows(list).chain( self.staged(list) .iter() .map(|&position| self.batch.vectors.row(position)), @@ -334,9 +336,9 @@ impl SplitPlan { /// The `index`-th vector of [`Self::points`]. fn point(&self, list: usize, index: usize) -> &[f32] { - let members = self.members.range(list); + let members = self.member_ids.range(list); if index < members.len() { - self.members.vector(members.start + index) + self.member_vectors.row(members.start + index) } else { self.batch .vectors diff --git a/diskann/src/ivf/test/provider.rs b/diskann/src/ivf/test/provider.rs index a1115e2e29..41c2d4c94f 100644 --- a/diskann/src/ivf/test/provider.rs +++ b/diskann/src/ivf/test/provider.rs @@ -19,7 +19,7 @@ use crate::{ ivf::{ dynamic::{ Apply, CentroidIndex, ListMetadata, ListProvider, MaintenanceAccessor, - MaintenanceStrategy, SelectedList, SelectionPlan, + MaintenanceStrategy, SelectedList, }, update::{Appends, Bootstrap, SplitInsert}, }, @@ -41,8 +41,6 @@ pub(crate) struct Faults { pub(crate) skip_vector: bool, /// Write NaN into the last row of every `read_vectors` call. pub(crate) nan_vector: bool, - /// Omit the first member of every list in `read_members`. - pub(crate) drop_member: bool, /// Return the first staged id again from every later `stage_insert`. pub(crate) repeat_stage: bool, /// Reserve list ids starting from zero, which may be live. @@ -69,7 +67,7 @@ impl CentroidIndex for Centroids { &self, query: &[f32], nprobe: usize, - ) -> impl SendFuture, ANNError>> { + ) -> impl SendFuture>, ANNError>> { let mut selected: Vec<_> = self .0 .iter() @@ -80,7 +78,7 @@ impl CentroidIndex for Centroids { .collect(); selected.sort_by(|a, b| a.distance.total_cmp(&b.distance).then(a.id.cmp(&b.id))); selected.truncate(nprobe); - std::future::ready(Ok(SelectionPlan::new(selected))) + std::future::ready(Ok(selected)) } } @@ -260,20 +258,6 @@ impl Accessor<'_> { .ok_or_else(|| error(format!("no point {id}"))) } - fn members(&self, list: u32) -> Result<&[u32], ANNError> { - let members = self - .provider - .lists - .get(&list) - .ok_or_else(|| error(format!("list {list} is not live")))?; - Ok( - match (self.provider.faults.drop_member, members.split_first()) { - (true, Some((_, rest))) => rest, - _ => members, - }, - ) - } - fn commit_staged(&mut self) { for point in self.staged.drain(..) { self.provider.points.push(point); @@ -337,7 +321,11 @@ impl<'b> MaintenanceAccessor<&'b [f32]> for Accessor<'_> { } async fn read_members(&mut self, list: u32) -> Result<&[u32], ANNError> { - self.members(list) + self.provider + .lists + .get(&list) + .map(Vec::as_slice) + .ok_or_else(|| error(format!("list {list} is not live"))) } async fn read_vectors( diff --git a/diskann/src/ivf/update.rs b/diskann/src/ivf/update.rs index 82d5698dd7..166008f32b 100644 --- a/diskann/src/ivf/update.rs +++ b/diskann/src/ivf/update.rs @@ -56,11 +56,6 @@ impl CentroidBlock { self.ids.len() } - /// Dimension of every centroid. - pub fn dim(&self) -> usize { - self.vectors.ncols() - } - /// Centroid ids in row order. pub fn ids(&self) -> &[L] { &self.ids @@ -141,16 +136,6 @@ impl Appends { self.len() == 0 } - /// Destination lists, ascending. - pub fn lists(&self) -> &[L] { - self.by_list.keys() - } - - /// Every appended id, grouped by destination list. - pub fn ids(&self) -> &[Id] { - self.by_list.values() - } - /// Iterate `(list, ids)` groups in ascending list order. pub fn iter(&self) -> impl ExactSizeIterator { self.by_list.iter() @@ -250,12 +235,8 @@ impl Splits { self.parents.len() } - /// Retired parent lists, ascending. - pub fn parents(&self) -> &[L] { - &self.parents - } - - /// Every child centroid; rows `2 * i` and `2 * i + 1` replace `self.parents()[i]`. + /// Every child centroid; rows `2 * i` and `2 * i + 1` replace the `i`-th split's + /// parent. pub fn children(&self) -> &CentroidBlock { &self.children } @@ -446,7 +427,7 @@ mod tests { fn centroid_block_pairs_ids_with_rows() { let bootstrap = Bootstrap::new(block(&[4, 2], 3)); let block = bootstrap.centroids(); - assert_eq!((block.len(), block.dim()), (2, 3)); + assert_eq!(block.len(), 2); assert_eq!(block.ids(), &[4, 2]); assert_eq!(block.vectors().row(1), &[3.0, 4.0, 5.0]); assert_eq!( @@ -463,8 +444,6 @@ mod tests { fn appends_group_ids_by_list() { let appends = Appends::::new(&[5, 1, 3, 2], &[7, 9, 7, 9]); assert_eq!(appends.len(), 4); - assert_eq!(appends.lists(), &[7, 9]); - assert_eq!(appends.ids(), &[3, 5, 1, 2]); assert_eq!( owned(appends.iter()), vec![(7, vec![3, 5]), (9, vec![1, 2])] @@ -480,7 +459,6 @@ mod tests { [(10, &[3, 1, 2][..], &[31, 30, 40][..]), (20, &[], &[])], ); assert_eq!(splits.len(), 2); - assert_eq!(splits.parents(), &[10, 20]); assert_eq!(splits.children().ids(), &[30, 31, 40, 41]); assert_eq!(splits.num_evacuated(), 3); @@ -492,6 +470,7 @@ mod tests { owned(split[0].evacuations()), vec![(30, vec![1]), (31, vec![3]), (40, vec![2])] ); + assert_eq!(split[1].parent(), 20); assert_eq!(split[1].children(), [40, 41]); assert_eq!(split[1].evacuations().len(), 0); } @@ -526,8 +505,8 @@ mod tests { assert_eq!(update.reassignments().len(), 1); let (appends, splits, reassignments) = update.into_parts(); assert_eq!( - (appends.lists(), splits.parents(), reassignments.len()), - (&[30][..], &[10][..], 1) + (appends.len(), splits.len(), reassignments.len()), + (1, 1, 1) ); } } diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 2b8fce1353..a03d3db2ba 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -135,8 +135,8 @@ coordinates its provider, centroid index, and inverted-list store. ids. - **Unified view**: mutually compatible point data, centroid state, reverse assignments, and inverted lists presented by one operation accessor. -- **Selection plan**: selected list ids and their coarse distances. It is - consumed by the same search accessor that produced it. +- **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 @@ -170,7 +170,7 @@ 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. `SelectionPlan` is passed back to the search accessor that produced 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 @@ -193,12 +193,10 @@ diskann/src/ivf/ dynamic.rs contracts proposed by this RFC update.rs declarative partition updates consumed by Apply grouped.rs flat grouped layouts shared by updates and planning - colocation.rs co-location set contracts, not yet wired into maintenance index.rs DynamicIvfIndex initialize/insert orchestration online/ insert planning: staging, split planning, kernels batch.rs staged and routed insert batch - gather.rs validated reads (members and canonical vectors) - split.rs parent admission and the split plan + 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 ``` @@ -229,24 +227,21 @@ This permits a list backend to store only PQ codes while maintenance still reads canonical representation, although the destination scan payload may need to be encoded again. -### Selection Plan +### Selected Lists -The plan is intentionally small: +Selection returns plain values, nearest first: ```rust pub struct SelectedList { pub id: ListId, pub distance: f32, } - -pub struct SelectionPlan { - selected: Vec>, -} ``` -The accessor may privately retain a lock guard, epoch, immutable root, manifest, -prefetch state, or other read context. None of that state appears in the plan's -public type. A plan is meaningful only to the accessor that created it. +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 @@ -260,11 +255,11 @@ pub trait SearchAccessor: HasId + Send + Sync { fn select_lists( &mut self, nprobe: usize, - ) -> impl SendFuture, Self::Error>>; + ) -> impl SendFuture>, Self::Error>>; fn scan_lists( &mut self, - plan: SelectionPlan, + lists: &[SelectedList], emit: F, ) -> impl SendFuture> where @@ -307,8 +302,8 @@ Search orchestration is deliberately unaware of consistency mechanics: ```text accessor = strategy.search_accessor(provider, context, query) -plan = accessor.select_lists(nprobe) -accessor.scan_lists(plan, |id, distance| top_k.insert(id, distance)) +lists = accessor.select_lists(nprobe) +accessor.scan_lists(&lists, |id, distance| top_k.insert(id, distance)) post_process(point statuses and external ids) ``` @@ -329,7 +324,7 @@ pub trait CentroidIndex: Send + Sync { &self, query: &[f32], nprobe: usize, - ) -> impl SendFuture, Self::Error>>; + ) -> impl SendFuture>, Self::Error>>; } ``` @@ -429,8 +424,7 @@ it enters, and builds updates without further checks: - caller input: initial centroids must be non-empty, finite, and of the provider's dimension; - accessor reads: every vector row must be written with finite values (rows - start as NaN, so unwritten rows are caught by the same check), and each split - parent's members must match its reported size; + start as NaN, so unwritten rows are caught by the same check); - staging: staged ids must be distinct; and - reservations: reserved list ids must be distinct and not live. @@ -455,10 +449,8 @@ 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: -`CoLocationSet`, `CoLocationGroup`, and `CoLocationDelta` live in -`ivf::colocation`, and no update type or `MaintenanceAccessor` references them -yet. +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 @@ -830,8 +822,7 @@ and search parameters. Durable metadata belongs to concrete provider crates. - Implement search orchestration with one query-bound accessor. - Compare results to brute force over exactly the selected lists. -- Test that selection plans cannot be used meaningfully outside their producing - accessor's configured view. +- Test that an accessor scans only lists it selected, under its own view. ### Phase 3: Exact-Centroid GraphIVF @@ -932,9 +923,9 @@ 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 selection plans difficult to -inspect. Keeping the plan as a small value and private state in one accessor is -simpler. +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 @@ -991,7 +982,7 @@ useful common shape. 4. **Update granularity**: updates are grouped by destination list, so bulk backends can rewrite each touched list once. Should they also carry a per-list summary of added and removed counts for exact pre-sizing? -5. **Selection plan ownership**: is the documented same-accessor requirement +5. **Selected-list ownership**: is the documented same-accessor requirement sufficient, or should a future private token make misuse dynamically detectable? From 935b088c3e04b77f8eb489b5a16353ec44a719d0 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Thu, 1 Oct 2026 22:28:08 -0400 Subject: [PATCH 09/33] move to enum based point updates --- diskann/src/ivf/dynamic.rs | 175 +++++----- diskann/src/ivf/grouped.rs | 26 +- diskann/src/ivf/index.rs | 192 +++++------ diskann/src/ivf/online/batch.rs | 51 +-- diskann/src/ivf/online/mod.rs | 33 +- diskann/src/ivf/online/split.rs | 95 +++--- diskann/src/ivf/test/provider.rs | 270 +++++++-------- diskann/src/ivf/update.rs | 547 +++++-------------------------- rfcs/01187-incremental-ivf.md | 401 +++++++++++----------- 9 files changed, 699 insertions(+), 1091 deletions(-) diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index ba539e13ec..84e7d80a84 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -16,7 +16,8 @@ use diskann_utils::{future::SendFuture, views::MutMatrixView}; use crate::{ error::{StandardError, ToRanked}, - provider::{DataProvider, HasId}, + ivf::update::InsertionUpdate, + provider::{ExecutionContext, HasId}, utils::VectorId, }; @@ -38,10 +39,43 @@ pub struct ScanStats { pub lists_scanned: u32, } -/// A data provider whose points are partitioned into stable logical lists. -pub trait ListProvider: DataProvider { - /// Stable id shared by a centroid and its inverted list. +/// The storage behind a dynamic IVF index: point identity and the lists points are +/// partitioned into. +/// +/// Reads and writes go through operation-scoped accessors; the provider itself only +/// names the identity types and translates the ids of visible points. Points staged +/// through [`StageElements`] are not translatable until an update that places them is +/// applied. +pub trait Provider: Sized + Send + Sync + 'static { + /// Per-operation context handed to strategies. + type Context: ExecutionContext; + + /// Internal point id used by lists and updates. + type InternalId: VectorId; + + /// Caller-facing point id. + type ExternalId: PartialEq + Send + Sync + 'static; + + /// Stable id shared by a centroid and its inverted list. Retired ids are never + /// reused. type ListId: VectorId; + + /// Errors from id translation. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Translate an external id to its corresponding internal id. + fn to_internal_id( + &self, + context: &Self::Context, + gid: &Self::ExternalId, + ) -> Result; + + /// Translate an internal id to its corresponding external id. + fn to_external_id( + &self, + context: &Self::Context, + id: Self::InternalId, + ) -> Result; } /// An in-memory index over the live centroids. @@ -87,7 +121,7 @@ pub trait CentroidIndex: Send + Sync { /// to batch reads, coalesce blob requests, prefetch, decode quantized payloads, or /// fan work out across tasks. pub trait SearchAccessor: HasId + Send + Sync { - /// Stable logical centroid/list id, fixed to [`ListProvider::ListId`] by the + /// Stable logical centroid/list id, fixed to [`Provider::ListId`] by the /// strategy. type ListId: VectorId; @@ -114,12 +148,9 @@ pub trait SearchAccessor: HasId + Send + Sync { } /// Factory for one dynamic IVF search accessor. -pub trait SearchStrategy<'a, Provider, T>: Send + Sync -where - Provider: ListProvider, -{ +pub trait SearchStrategy<'a, P: Provider, T>: Send + Sync { /// Query-bound accessor used for both coarse and fine search. - type SearchAccessor: SearchAccessor; + type SearchAccessor: SearchAccessor; /// Error constructing the accessor. type Error: StandardError; @@ -127,42 +158,25 @@ where /// Construct an accessor around `query`. fn search_accessor( &'a self, - provider: &'a Provider, - context: &'a Provider::Context, + provider: &'a P, + context: &'a P::Context, query: T, ) -> Result; } -/// Size metadata for one logical inverted list. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ListMetadata { - /// Number of live point ids in the list. - pub len: usize, -} - /// Operation-scoped reads and writes needed by split/dissolve maintenance. /// /// The accessor is the consistency boundary for one mutation. It must present a -/// unified view across point data, centroids, reverse assignments, and inverted -/// lists for all planning reads. It also owns any provider-specific lock, epoch, -/// transaction, staging area, or poison state needed to apply the final update -/// through [`Apply`]. +/// unified view across point data, centroids, and inverted lists for all planning +/// reads. It also owns any provider-specific lock, epoch, transaction, staging area, +/// or poison state needed to apply the final update through [`Self::apply`]. /// /// List reads take one list and return its result. [`Self::read_vectors`], which /// reads far more data, takes a batch so disk and blob providers can reorder, /// coalesce, or parallelize I/O. -pub trait MaintenanceAccessor: HasId + Send + Sized -where - T: Send, -{ - /// External point id accepted by the data provider. - type ExternalId: PartialEq + Send + Sync + 'static; - - /// Stable centroid/list id, fixed to [`ListProvider::ListId`] by the strategy. - type ListId: VectorId; - +pub trait MaintenanceAccessor: Send + Sized { /// In-memory centroid catalog and navigator in this accessor's unified view. - type Centroids: CentroidIndex; + type Centroids: CentroidIndex; /// Errors from planning reads, staging, or applying the update. type Error: ToRanked + Debug + Send + Sync + 'static; @@ -176,58 +190,64 @@ where fn centroids(&self) -> &Self::Centroids; /// Read the size of a live list. - fn list_metadata( - &mut self, - list: Self::ListId, - ) -> impl SendFuture>; + fn list_size(&mut self, list: P::ListId) -> impl SendFuture>; /// Read the member ids of a live list. fn read_members( &mut self, - list: Self::ListId, - ) -> impl SendFuture>; + list: P::ListId, + ) -> impl SendFuture>; /// Write the canonical vector of `ids[i]` into row `i` of `out`, in any order. /// /// `out` has one row per id and [`Self::dim`] columns, and every row must be - /// written. Covers points staged by this accessor as well as visible points, - /// without making staged points visible through ordinary provider reads. + /// written. Covers visible points only; the canonical vectors of staged points + /// come from [`StageElements::stage`]. fn read_vectors( &mut self, - ids: &[Self::Id], + ids: &[P::InternalId], out: MutMatrixView<'_, f32>, ) -> impl SendFuture>; - /// Stage a canonical point and reserve its internal id. - fn stage_insert( - &mut self, - id: &Self::ExternalId, - element: T, - ) -> impl SendFuture>; - /// Reserve fresh logical list ids that will not alias retired ids. + /// TO DO: Re-evaluate we even need this. fn reserve_list_ids( &mut self, count: usize, - ) -> impl SendFuture, Self::Error>>; + ) -> impl SendFuture, Self::Error>>; + + /// Apply `update` and finish the operation, consuming the accessor. + /// + /// Only the index constructs updates, and it guarantees the invariants documented + /// on [`InsertionUpdate`], so implementations may rely on them without re-checking. + /// + /// Returning `Ok(())` means all components expose one coherent resulting index. + /// Rollback, durability, concurrent-reader visibility, and recovery after `Err` are + /// intentionally provider-defined. + fn apply( + self, + update: InsertionUpdate, + ) -> impl SendFuture>; } -/// Applies one kind of partition update and finishes the maintenance operation. -/// -/// Providers implement this once per update they support, from -/// [`crate::ivf::update`]. For example, a provider that never splits implements only -/// `Apply>`. Only the index constructs updates, and it guarantees their -/// documented invariants, so implementations may rely on them without re-checking. +/// Stages new points of element type `T` for one maintenance operation. /// -/// Returning `Ok(())` means all components expose one coherent resulting index. -/// Rollback, durability, concurrent-reader visibility, and recovery after `Err` are -/// intentionally provider-defined. -pub trait Apply: Send + Sized { - /// Errors from applying the update. - type Error: ToRanked + Debug + Send + Sync + 'static; - - /// Apply `update`, consuming the accessor. - fn apply(self, update: U) -> impl SendFuture>; +/// The provider owns id allocation, duplicate detection, and the conversion of `T` to +/// its stored and canonical representations. Staged points become visible only when +/// the accessor applies an update that places them, and are discarded if the accessor +/// is dropped first. A provider that accepts several input types implements this once +/// per type. +pub trait StageElements: MaintenanceAccessor

{ + /// Stage `points` and return their internal ids in order, writing the canonical + /// vector of `points[i]` into row `i` of `out`. + /// + /// `out` has one row per point and [`MaintenanceAccessor::dim`] columns, and every + /// row must be written. + fn stage( + &mut self, + points: &[(P::ExternalId, T)], + out: MutMatrixView<'_, f32>, + ) -> impl SendFuture, Self::Error>>; } /// Factory for an operation-scoped dynamic IVF maintenance accessor. @@ -235,21 +255,12 @@ pub trait Apply: Send + Sized { /// The provider is borrowed exclusively for the lifetime of the accessor, so no /// search or other mutation can observe it until the accessor is applied or /// dropped. Providers may therefore mutate plain in-memory state in -/// [`Apply::apply`] without interior synchronization. Providers that share state -/// outside this borrow (for example through an `Arc`) remain responsible for -/// coordinating those aliases. -pub trait MaintenanceStrategy<'a, Provider, T>: Send + Sync -where - Provider: ListProvider, - T: Send, -{ +/// [`MaintenanceAccessor::apply`] without interior synchronization. Providers that +/// share state outside this borrow (for example through an `Arc`) remain responsible +/// for coordinating those aliases. +pub trait MaintenanceStrategy<'a, P: Provider>: Send + Sync { /// Accessor used to plan and stage split/dissolve operations. - type MaintenanceAccessor: MaintenanceAccessor< - T, - Id = Provider::InternalId, - ExternalId = Provider::ExternalId, - ListId = Provider::ListId, - >; + type MaintenanceAccessor: MaintenanceAccessor

; /// Error constructing the accessor. type Error: StandardError; @@ -257,7 +268,7 @@ where /// Construct one maintenance accessor. fn maintenance_accessor( &'a self, - provider: &'a mut Provider, - context: &'a Provider::Context, + provider: &'a mut P, + context: &'a P::Context, ) -> Result; } diff --git a/diskann/src/ivf/grouped.rs b/diskann/src/ivf/grouped.rs index 6ae4828bdb..c444a0b588 100644 --- a/diskann/src/ivf/grouped.rs +++ b/diskann/src/ivf/grouped.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Flat grouped layouts shared by the partition updates and the online planner. +//! Flat grouped layouts used by the online planner. use std::ops::Range; @@ -105,16 +105,6 @@ impl Grouped { } } - /// Number of values across all groups. - pub(super) fn len(&self) -> usize { - self.groups.values.len() - } - - /// Group keys, ascending. - pub(super) fn keys(&self) -> &[K] { - &self.keys - } - /// The values under `key`, or nothing if `key` is absent. pub(super) fn get(&self, key: K) -> &[V] { self.keys @@ -122,14 +112,9 @@ impl Grouped { .map_or(&[], |index| self.groups.group(index)) } - /// The groups at positions `range`, with their keys. - pub(super) fn groups(&self, range: Range) -> impl ExactSizeIterator { - range.map(|index| (self.keys[index], self.groups.group(index))) - } - /// Every group with its key, in ascending key order. pub(super) fn iter(&self) -> impl ExactSizeIterator { - self.groups(0..self.keys.len()) + (0..self.keys.len()).map(|index| (self.keys[index], self.groups.group(index))) } } @@ -170,9 +155,12 @@ mod tests { #[test] fn grouped_get_finds_groups_by_key() { let grouped = Grouped::from_pairs(vec![(7, 2), (5, 1), (7, 0)]); - assert_eq!(grouped.keys(), &[5, 7]); + let groups: Vec<_> = grouped + .iter() + .map(|(key, values)| (key, values.to_vec())) + .collect(); + assert_eq!(groups, vec![(5, vec![1]), (7, vec![0, 2])]); assert_eq!(grouped.get(7), &[0, 2]); assert_eq!(grouped.get(6), &[] as &[i32]); - assert_eq!(grouped.len(), 3); } } diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 3eb0e8c774..fef218c7d7 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -5,9 +5,10 @@ //! Dynamic IVF index orchestration. //! -//! The index drives one [`MaintenanceAccessor`] per mutation: it plans the update -//! against the accessor's view and hands the result to [`Apply`]. It never touches -//! point, centroid, or list storage directly. +//! The index drives one [`MaintenanceAccessor`] per mutation: it plans an +//! [`InsertionUpdate`] against the accessor's view and hands it to +//! [`MaintenanceAccessor::apply`]. It never touches point, centroid, or list storage +//! directly. use diskann_utils::{future::SendFuture, views::Matrix}; use rand::{SeedableRng, rngs::StdRng}; @@ -17,10 +18,13 @@ use crate::{ ANNError, ANNErrorKind, ANNResult, error::{ErrorExt, IntoANNResult}, ivf::{ - dynamic::{Apply, CentroidIndex, ListProvider, MaintenanceAccessor, MaintenanceStrategy}, + dynamic::{ + CentroidIndex, MaintenanceAccessor, MaintenanceStrategy, Provider, StageElements, + }, online::{SplitPlan, StagedBatch, check_reserved, index_error}, - update::{Appends, Bootstrap, CentroidBlock, SplitInsert}, + update::{CentroidDelta, InsertionUpdate, PointDelta}, }, + utils::VectorId, }; /// Parameters for the online split policy. @@ -28,8 +32,6 @@ use crate::{ pub struct DynamicIvfConfig { /// Split a list once an insert batch would grow it beyond this many points. pub split_threshold: usize, - /// Hard cap on live lists; `None` allows unbounded growth. - pub max_clusters: Option, /// Nearby lists reassigned together with each split parent. pub reassign_neighbors: usize, /// Lloyd iterations used to fit split children. @@ -65,18 +67,37 @@ pub struct InsertStats { pub reassigned: usize, } +impl InsertStats { + /// Tally an insert's update: every retired list is a split parent. + fn of(update: &InsertionUpdate) -> Self { + let mut stats = Self::default(); + for point in update.points() { + match point { + PointDelta::Append { .. } => stats.inserted += 1, + PointDelta::Move { .. } => stats.reassigned += 1, + } + } + stats.splits = update + .centroids() + .iter() + .filter(|delta| matches!(delta, CentroidDelta::Retire { .. })) + .count(); + stats + } +} + /// An incrementally maintained IVF index. /// /// Mutations take `&mut self` and lend the provider exclusively to one /// [`MaintenanceAccessor`], so they cannot overlap searches through the same value. #[derive(Debug)] -pub struct DynamicIvfIndex { +pub struct DynamicIvfIndex { provider: P, config: DynamicIvfConfig, rng: StdRng, } -impl DynamicIvfIndex

{ +impl DynamicIvfIndex

{ /// Construct an index over `provider`. /// /// # Errors @@ -107,16 +128,14 @@ impl DynamicIvfIndex

{ /// /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index /// already has centroids, or the maintenance accessor fails. - pub fn initialize<'a, S, T>( + pub fn initialize<'a, S>( &'a mut self, strategy: &'a S, context: &'a P::Context, centroids: Matrix, ) -> impl SendFuture> where - S: MaintenanceStrategy<'a, P, T>, - S::MaintenanceAccessor: Apply>, - T: Send, + S: MaintenanceStrategy<'a, P>, { let provider = &mut self.provider; @@ -150,8 +169,17 @@ impl DynamicIvfIndex

{ .escalate("initialize must reserve list ids")?; check_reserved(accessor.centroids(), &list_ids, count)?; + let installs = list_ids + .into_iter() + .zip(centroids.row_iter()) + .map(|(id, centroid)| CentroidDelta::Install { + id, + centroid: centroid.into(), + }) + .collect(); + accessor - .apply(Bootstrap::new(CentroidBlock::new(list_ids, centroids))) + .apply(InsertionUpdate::new(installs, Vec::new())) .await .escalate("initialize must apply the initial centroids") } @@ -160,8 +188,8 @@ impl DynamicIvfIndex

{ /// Insert a batch, splitting every list the batch pushes past `split_threshold`. /// /// Routing and split planning run against the accessor's view before any change. - /// The batch is then applied as one [`Appends`], or as one [`SplitInsert`] when - /// lists split. + /// The batch, and any splits it triggers, are then applied as one + /// [`InsertionUpdate`]. /// /// # Errors /// @@ -175,10 +203,9 @@ impl DynamicIvfIndex

{ points: &'a [(P::ExternalId, T)], ) -> impl SendFuture> where - S: MaintenanceStrategy<'a, P, T>, - S::MaintenanceAccessor: - Apply> + Apply>, - T: Copy + Send + Sync, + S: MaintenanceStrategy<'a, P>, + S::MaintenanceAccessor: StageElements, + T: Sync, { let provider = &mut self.provider; let config = self.config; @@ -199,83 +226,41 @@ impl DynamicIvfIndex

{ )); } - let batch = StagedBatch::stage(&mut accessor, points).await?; - let parents = Self::select_parents(&config, &mut accessor, &batch).await?; + let batch = StagedBatch::stage(&mut accessor, points).await?; // register points with provider -> internal ids + f32 representations of the vectors. - if parents.is_empty() { - let update = batch.into_appends(); - let inserted = update.len(); - accessor - .apply(update) + // Split every routed list the batch pushes past the threshold. `routed` + // visits lists in ascending order, as `SplitPlan::gather` requires. + let mut parents = Vec::new(); + for (list, incoming) in batch.routed() { + let len = accessor + .list_size(list) .await - .escalate("insert must apply routed points")?; - return Ok(InsertStats { - inserted, - splits: 0, - reassigned: 0, - }); + .escalate("insert must read routed list sizes")?; + + let projected = len + incoming; + + if projected > config.split_threshold { + parents.push(list); + } } - let update = SplitPlan::gather(&mut accessor, &config, batch, parents) - .await? - .solve(config.two_means_iterations, rng)?; - let stats = InsertStats { - inserted: update.appends().len(), - splits: update.splits().len(), - reassigned: update.splits().num_evacuated() + update.reassignments().len(), + let update = if parents.is_empty() { + batch.into_update() + } else { + SplitPlan::gather(&mut accessor, &config, batch, parents) + .await? + .solve(config.two_means_iterations, rng)? }; + + let stats = InsertStats::of(&update); + accessor .apply(update) .await - .escalate("insert must apply the split update")?; - Ok(stats) - } - } + .escalate("insert must apply its update")?; - /// Choose the routed lists to split. - /// - /// Admits every list the batch pushes past `split_threshold`, largest projected - /// size first, while the live-list count stays within `max_clusters`. Parents are - /// returned in ascending list order. - /// - /// Takes the configuration rather than `&self` because the provider is lent to - /// `accessor` for the whole mutation. - /// - /// # Errors - /// - /// Fails if the accessor fails. - async fn select_parents( - config: &DynamicIvfConfig, - accessor: &mut A, - batch: &StagedBatch, - ) -> ANNResult> - where - A: MaintenanceAccessor, - T: Send, - { - let mut overflowing = Vec::new(); - for (list, incoming) in batch.routed() { - let len = accessor - .list_metadata(list) - .await - .escalate("insert must read routed list sizes")? - .len; - let projected = len + incoming; - if projected > config.split_threshold { - overflowing.push((projected, list)); - } + Ok(stats) } - - // Each split adds one live list. - let budget = config.max_clusters.map_or(usize::MAX, |max| { - max.saturating_sub(accessor.centroids().len()) - }); - overflowing.sort_unstable_by(|a, b| b.0.cmp(&a.0).then(a.1.cmp(&b.1))); - overflowing.truncate(budget); - - let mut parents: Vec<_> = overflowing.into_iter().map(|(_, list)| list).collect(); - parents.sort_unstable(); - Ok(parents) } } @@ -290,7 +275,6 @@ mod tests { fn config(split_threshold: usize, reassign_neighbors: usize) -> DynamicIvfConfig { DynamicIvfConfig { split_threshold, - max_clusters: None, reassign_neighbors, two_means_iterations: 10, seed: 7, @@ -316,7 +300,7 @@ mod tests { ) -> DynamicIvfIndex { let mut index = DynamicIvfIndex::new(provider, config).unwrap(); index - .initialize::<_, &[f32]>(&Strategy, &DefaultContext, centroids) + .initialize(&Strategy, &DefaultContext, centroids) .await .unwrap(); index @@ -333,22 +317,22 @@ mod tests { async fn initialize_rejects_bad_centroids() { let mut index = DynamicIvfIndex::new(Provider::new(2), config(10, 1)).unwrap(); let wrong_dim = index - .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0, 0.0]])) + .initialize(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0, 0.0]])) .await .unwrap_err(); assert!(wrong_dim.to_string().contains("dimension"), "{wrong_dim}"); let non_finite = index - .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[f32::NAN, 0.0]])) + .initialize(&Strategy, &DefaultContext, matrix(&[[f32::NAN, 0.0]])) .await .unwrap_err(); assert!(non_finite.to_string().contains("finite"), "{non_finite}"); index - .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0]])) + .initialize(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0]])) .await .unwrap(); let again = index - .initialize::<_, &[f32]>(&Strategy, &DefaultContext, matrix(&[[1.0, 1.0]])) + .initialize(&Strategy, &DefaultContext, matrix(&[[1.0, 1.0]])) .await .unwrap_err(); assert!(again.to_string().contains("already initialized"), "{again}"); @@ -396,24 +380,6 @@ mod tests { assert_eq!(provider.member_vectors(1), vec![vec![9.0, 9.0]]); } - #[tokio::test] - async fn max_clusters_caps_splits() { - let mut capped = config(2, 1); - capped.max_clusters = Some(2); - let mut index = initialized(Provider::new(1), capped, matrix(&[[0.0], [10.0]])).await; - let stats = index - .insert_batch( - &Strategy, - &DefaultContext, - &batch(0, &[[1.0], [2.0], [3.0]]), - ) - .await - .unwrap(); - assert_eq!(stats.splits, 0); - assert_eq!(index.provider().lists(), vec![0, 1]); - index.provider().check(); - } - /// Lists A (0), B (20), and C (40) in one dimension. Inserting 0.5, 8, and 9 /// overflows A, which splits into children near 1/6 and 8.5, with B and C as /// its neighbors. @@ -512,6 +478,10 @@ mod tests { ), (fault(|f| f.nan_vector = true), "no finite canonical vector"), (fault(|f| f.repeat_stage = true), "staging returned point"), + ( + fault(|f| f.drop_stage = true), + "staging returned 1 ids for 2 points", + ), (fault(|f| f.reuse_lists = true), "reserved live list id"), ]; for (faults, expected) in cases { diff --git a/diskann/src/ivf/online/batch.rs b/diskann/src/ivf/online/batch.rs index e8b9aebeab..10cf77900f 100644 --- a/diskann/src/ivf/online/batch.rs +++ b/diskann/src/ivf/online/batch.rs @@ -7,14 +7,14 @@ use diskann_utils::views::Matrix; -use super::{first_repeat, index_error, read_rows}; +use super::{check_rows, first_repeat, index_error}; use crate::{ ANNResult, error::ErrorExt, ivf::{ - dynamic::{CentroidIndex, MaintenanceAccessor}, + dynamic::{CentroidIndex, Provider, StageElements}, grouped::Grouped, - update::Appends, + update::{InsertionUpdate, PointDelta}, }, utils::VectorId, }; @@ -32,19 +32,19 @@ pub(in crate::ivf) struct StagedBatch { } impl StagedBatch { - /// Stage every point, read its canonical vector, and route it to its nearest live - /// list. + /// Stage every point with the accessor and route it to its nearest live list. /// /// # Errors /// /// Fails if any accessor call fails or returns inconsistent data. - pub(in crate::ivf) async fn stage( + pub(in crate::ivf) async fn stage( accessor: &mut A, - points: &[(A::ExternalId, T)], + points: &[(P::ExternalId, T)], ) -> ANNResult where - A: MaintenanceAccessor, - T: Copy + Send + Sync, + P: Provider, + A: StageElements, + T: Sync, { let dim = accessor.dim(); if dim == 0 { @@ -53,21 +53,24 @@ impl StagedBatch { )); } - let mut ids = Vec::with_capacity(points.len()); - for (external, element) in points { - ids.push( - accessor - .stage_insert(external, *element) - .await - .escalate("insert must stage every point")?, - ); + let mut vectors = Matrix::new(f32::NAN, points.len(), dim); + let ids = accessor + .stage(points, vectors.as_mut_view()) + .await + .escalate("insert must stage every point")?; + if ids.len() != points.len() { + return Err(index_error(format!( + "staging returned {} ids for {} points", + ids.len(), + points.len() + ))); } if let Some(id) = first_repeat(ids.iter().copied()) { return Err(index_error(format!( "staging returned point {id} more than once" ))); } - let vectors = read_rows(accessor, &ids, dim).await?; + check_rows(&ids, &vectors)?; let centroids = accessor.centroids(); let mut routes = Vec::with_capacity(ids.len()); @@ -91,9 +94,15 @@ impl StagedBatch { }) } - /// Place every point on its route. - pub(in crate::ivf) fn into_appends(self) -> Appends { - Appends::new(&self.ids, &self.routes) + /// Append every point to its route. + pub(in crate::ivf) fn into_update(self) -> InsertionUpdate { + let appends = self + .ids + .iter() + .zip(&self.routes) + .map(|(&id, &to)| PointDelta::Append { id, to }) + .collect(); + InsertionUpdate::new(Vec::new(), appends) } /// Lists that received points, ascending, with the number of points routed to diff --git a/diskann/src/ivf/online/mod.rs b/diskann/src/ivf/online/mod.rs index c4831a2326..10f1a2a4f5 100644 --- a/diskann/src/ivf/online/mod.rs +++ b/diskann/src/ivf/online/mod.rs @@ -25,7 +25,8 @@ use diskann_utils::views::Matrix; use crate::{ ANNError, ANNErrorKind, ANNResult, error::ErrorExt, - ivf::dynamic::{CentroidIndex, MaintenanceAccessor}, + ivf::dynamic::{CentroidIndex, MaintenanceAccessor, Provider}, + utils::VectorId, }; #[track_caller] @@ -68,31 +69,43 @@ fn first_repeat(values: impl Iterator) -> Option { .map(|pair| pair[0]) } -/// Read the canonical vector of every point in `ids`; row `i` belongs to `ids[i]`. +/// Read the canonical vector of every visible point in `ids`; row `i` belongs to +/// `ids[i]`. /// /// # Errors /// /// Fails if the accessor fails, or leaves a row unwritten or non-finite. -async fn read_rows(accessor: &mut A, ids: &[A::Id], dim: usize) -> ANNResult> +async fn read_rows( + accessor: &mut A, + ids: &[P::InternalId], + dim: usize, +) -> ANNResult> where - A: MaintenanceAccessor, - T: Send, + P: Provider, + A: MaintenanceAccessor

, { - // Rows the accessor never writes stay NaN, so the finiteness check catches them. let mut rows = Matrix::new(f32::NAN, ids.len(), dim); accessor .read_vectors(ids, rows.as_mut_view()) .await .escalate("maintenance must read canonical vectors")?; + check_rows(ids, &rows)?; + Ok(rows) +} - if let Some((id, _)) = ids +/// Check that every row the accessor was asked to write is finite. +/// +/// Callers fill rows with NaN before handing them to the accessor, so rows it never +/// writes fail this check too. +fn check_rows(ids: &[Id], rows: &Matrix) -> ANNResult<()> { + match ids .iter() .zip(rows.row_iter()) .find(|(_, row)| row.iter().any(|x| !x.is_finite())) { - return Err(index_error(format!( + Some((id, _)) => Err(index_error(format!( "no finite canonical vector for point {id}" - ))); + ))), + None => Ok(()), } - Ok(rows) } diff --git a/diskann/src/ivf/online/split.rs b/diskann/src/ivf/online/split.rs index e1adfcd69d..ef2fbe3f91 100644 --- a/diskann/src/ivf/online/split.rs +++ b/diskann/src/ivf/online/split.rs @@ -6,8 +6,6 @@ //! Split planning: gather each split's region, fit children, and place every affected //! point. -use std::ops::Range; - use diskann_utils::views::Matrix; use rand::{Rng, rngs::StdRng}; @@ -22,9 +20,9 @@ use crate::{ error::ErrorExt, ivf::{ DynamicIvfConfig, - dynamic::{CentroidIndex, MaintenanceAccessor}, + dynamic::{CentroidIndex, MaintenanceAccessor, Provider}, grouped::Csr, - update::{Appends, CentroidBlock, Reassignments, SplitInsert, Splits}, + update::{CentroidDelta, InsertionUpdate, PointDelta}, }, utils::VectorId, }; @@ -63,15 +61,15 @@ impl SplitPlan { /// # Errors /// /// Fails if an accessor call fails or the accessor's view is inconsistent. - pub(in crate::ivf) async fn gather( + pub(in crate::ivf) async fn gather( accessor: &mut A, config: &DynamicIvfConfig, batch: StagedBatch, parents: Vec, ) -> ANNResult where - A: MaintenanceAccessor, - T: Send, + P: Provider, + A: MaintenanceAccessor

, { let count = 2 * parents.len(); let child_ids = accessor @@ -174,12 +172,11 @@ impl SplitPlan { mut self, iterations: usize, rng: &mut StdRng, - ) -> ANNResult> { - let centers = self.fit_children(iterations, rng)?; - let children = CentroidBlock::new(std::mem::take(&mut self.child_ids), centers); + ) -> ANNResult> { + let children = self.fit_children(iterations, rng)?; let mut routes = std::mem::take(&mut self.batch.routes); let destinations = self.place(&children, &mut routes)?; - Ok(self.package(children, &routes, &destinations)) + Ok(self.package(&children, &routes, &destinations)) } /// Fit two children per region; rows `2r` and `2r + 1` split region `r`. @@ -215,12 +212,12 @@ impl SplitPlan { /// Choose a list for every point of every region list. /// - /// Returns the existing members' lists, aligned with their positions, and - /// overwrites the routes of the batch points it places. - fn place(&self, children: &CentroidBlock, routes: &mut [L]) -> ANNResult> { + /// Row `r` of `children` is the centroid of child `r`. Returns the existing + /// members' lists, aligned with their positions, and overwrites the routes of the + /// batch points it places. + fn place(&self, children: &Matrix, routes: &mut [L]) -> ANNResult> { let mut destinations = vec![L::default(); self.member_ids.values().len()]; - let child_ids = children.ids(); - let child_vectors = children.vectors(); + let child = |row: usize| (self.child_ids[row], children.row(row)); let mut candidates: Vec<(L, &[f32])> = Vec::new(); for region in 0..self.parents { @@ -228,9 +225,7 @@ impl SplitPlan { for &neighbor in self.neighbors.group(region) { candidates.push(self.neighbor(neighbor)); } - for row in [2 * region, 2 * region + 1] { - candidates.push((child_ids[row], child_vectors.row(row))); - } + candidates.extend([child(2 * region), child(2 * region + 1)]); self.assign(region, &candidates, &mut destinations, routes)?; } @@ -239,9 +234,7 @@ impl SplitPlan { candidates.clear(); candidates.push(self.neighbor(neighbor)); for ®ion in self.regions.group(neighbor) { - for row in [2 * region, 2 * region + 1] { - candidates.push((child_ids[row], child_vectors.row(row))); - } + candidates.extend([child(2 * region), child(2 * region + 1)]); } self.assign( self.parents + neighbor, @@ -275,34 +268,46 @@ impl SplitPlan { Ok(()) } - /// Assemble the update from the chosen lists. + /// Assemble the update: install the children, retire the parents, append the + /// batch, and move every existing point whose list changed. fn package( &self, - children: CentroidBlock, + children: &Matrix, routes: &[L], destinations: &[L], - ) -> SplitInsert { - let parents = self.lists[..self.parents].to_vec(); - SplitInsert::new( - Appends::new(&self.batch.ids, routes), - Splits::new(parents, children, self.moves(0..self.parents, destinations)), - Reassignments::new(self.moves(self.parents..self.lists.len(), destinations)), - ) - } + ) -> InsertionUpdate { + let installs = self + .child_ids + .iter() + .zip(children.row_iter()) + .map(|(&id, centroid)| CentroidDelta::Install { + id, + centroid: centroid.into(), + }); + let retires = self.lists[..self.parents] + .iter() + .map(|&id| CentroidDelta::Retire { id }); - /// `(list, members, destinations)` for every list in `lists`. - fn moves<'s>( - &'s self, - lists: Range, - destinations: &'s [L], - ) -> impl Iterator { - lists.map(move |list| { - ( - self.lists[list], - self.member_ids.group(list), - &destinations[self.member_ids.range(list)], - ) - }) + let appends = self + .batch + .ids + .iter() + .zip(routes) + .map(|(&id, &to)| PointDelta::Append { id, to }); + // A parent is never a destination, so all its members move. + let moves = self.lists.iter().enumerate().flat_map(|(list, &from)| { + self.member_ids + .group(list) + .iter() + .zip(&destinations[self.member_ids.range(list)]) + .filter(move |&(_, &to)| to != from) + .map(move |(&id, &to)| PointDelta::Move { id, from, to }) + }); + + InsertionUpdate::new( + installs.chain(retires).collect(), + appends.chain(moves).collect(), + ) } /// Neighbor `neighbor`'s list id and centroid. diff --git a/diskann/src/ivf/test/provider.rs b/diskann/src/ivf/test/provider.rs index 41c2d4c94f..fe505d774e 100644 --- a/diskann/src/ivf/test/provider.rs +++ b/diskann/src/ivf/test/provider.rs @@ -18,12 +18,12 @@ use crate::{ ANNError, ANNErrorKind, ivf::{ dynamic::{ - Apply, CentroidIndex, ListMetadata, ListProvider, MaintenanceAccessor, - MaintenanceStrategy, SelectedList, + self, CentroidIndex, MaintenanceAccessor, MaintenanceStrategy, SelectedList, + StageElements, }, - update::{Appends, Bootstrap, SplitInsert}, + update::{CentroidDelta, InsertionUpdate, PointDelta}, }, - provider::{DataProvider, DefaultContext, HasId, NoopGuard}, + provider::DefaultContext, }; fn error(message: impl Into) -> ANNError { @@ -37,16 +37,34 @@ fn squared(a: &[f32], b: &[f32]) -> f32 { /// Inconsistent responses the provider can inject. #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] pub(crate) struct Faults { - /// Leave the last row of every `read_vectors` call unwritten. + /// Leave the last row of every staged or read batch unwritten. pub(crate) skip_vector: bool, - /// Write NaN into the last row of every `read_vectors` call. + /// Write NaN into the last row of every staged or read batch. pub(crate) nan_vector: bool, - /// Return the first staged id again from every later `stage_insert`. + /// Return the first staged id again for every later point of a `stage` call. pub(crate) repeat_stage: bool, + /// Return one id fewer than the number of points a `stage` call stages. + pub(crate) drop_stage: bool, /// Reserve list ids starting from zero, which may be live. pub(crate) reuse_lists: bool, } +/// Write `rows[i]` into row `i` of `out` in reverse order, injecting vector faults +/// into the last row. +fn write_rows(faults: Faults, mut out: MutMatrixView<'_, f32>, rows: &[&[f32]]) { + for (index, row) in rows.iter().enumerate().rev() { + let last = index + 1 == rows.len(); + if last && faults.skip_vector { + continue; + } + let out = out.row_mut(index); + out.copy_from_slice(row); + if last && faults.nan_vector { + out[0] = f32::NAN; + } + } +} + /// Exact centroid catalog. #[derive(Debug, Clone, Default, PartialEq)] pub(crate) struct Centroids(BTreeMap>); @@ -210,12 +228,12 @@ impl Provider { } } -impl DataProvider for Provider { +impl dynamic::Provider for Provider { type Context = DefaultContext; type InternalId = u32; type ExternalId = u32; + type ListId = u32; type Error = ANNError; - type Guard = NoopGuard; fn to_internal_id(&self, _: &DefaultContext, gid: &u32) -> Result { self.points @@ -233,10 +251,6 @@ impl DataProvider for Provider { } } -impl ListProvider for Provider { - type ListId = u32; -} - /// Maintenance accessor over an exclusively borrowed [`Provider`]. pub(crate) struct Accessor<'a> { provider: &'a mut Provider, @@ -246,14 +260,9 @@ pub(crate) struct Accessor<'a> { impl Accessor<'_> { fn vector(&self, id: u32) -> Result<&[f32], ANNError> { - let id = id as usize; - let committed = self.provider.points.len(); - let point = if id < committed { - self.provider.points.get(id) - } else { - self.staged.get(id - committed) - }; - point + self.provider + .points + .get(id as usize) .map(|(_, vector)| &**vector) .ok_or_else(|| error(format!("no point {id}"))) } @@ -265,40 +274,61 @@ impl Accessor<'_> { } } - fn install(&mut self, id: u32, vector: &[f32]) -> Result<(), ANNError> { + fn install(&mut self, id: u32, centroid: &[f32]) -> Result<(), ANNError> { if self.provider.retired.contains(&id) || self.provider.lists.contains_key(&id) { return Err(error(format!("list id {id} is reused"))); } - self.provider.centroids.0.insert(id, vector.into()); + self.provider.centroids.0.insert(id, centroid.into()); self.provider.lists.insert(id, Vec::new()); Ok(()) } - fn place(&mut self, list: u32, ids: &[u32]) -> Result<(), ANNError> { + /// Retire `list`, dropping its members wholesale. + fn retire(&mut self, list: u32) -> Result<(), ANNError> { self.provider .lists - .get_mut(&list) - .ok_or_else(|| error(format!("list {list} is not live")))? - .extend_from_slice(ids); - for &id in ids { - let slot = self - .provider - .assignment - .get_mut(id as usize) - .ok_or_else(|| error(format!("no point {id}")))?; - *slot = list; + .remove(&list) + .ok_or_else(|| error(format!("retired list {list} is not live")))?; + self.provider.centroids.0.remove(&list); + self.provider.retired.insert(list); + Ok(()) + } + + /// Remove `id` from `list`, which retirement may have dropped already. + fn leave(&mut self, id: u32, list: u32) -> Result<(), ANNError> { + if self.provider.retired.contains(&list) { + return Ok(()); } + let members = self + .provider + .lists + .get_mut(&list) + .ok_or_else(|| error(format!("list {list} is not live")))?; + let position = members + .iter() + .position(|&member| member == id) + .ok_or_else(|| error(format!("point {id} is not in list {list}")))?; + members.remove(position); Ok(()) } -} -impl HasId for Accessor<'_> { - type Id = u32; + fn place(&mut self, id: u32, list: u32) -> Result<(), ANNError> { + self.provider + .lists + .get_mut(&list) + .ok_or_else(|| error(format!("list {list} is not live")))? + .push(id); + let slot = self + .provider + .assignment + .get_mut(id as usize) + .ok_or_else(|| error(format!("no point {id}")))?; + *slot = list; + Ok(()) + } } -impl<'b> MaintenanceAccessor<&'b [f32]> for Accessor<'_> { - type ExternalId = u32; - type ListId = u32; +impl MaintenanceAccessor for Accessor<'_> { type Centroids = Centroids; type Error = ANNError; @@ -310,14 +340,14 @@ impl<'b> MaintenanceAccessor<&'b [f32]> for Accessor<'_> { &self.provider.centroids } - async fn list_metadata(&mut self, list: u32) -> Result { + async fn list_size(&mut self, list: u32) -> Result { let len = self .provider .lists .get(&list) .ok_or_else(|| error(format!("list {list} is not live")))? .len(); - Ok(ListMetadata { len }) + Ok(len) } async fn read_members(&mut self, list: u32) -> Result<&[u32], ANNError> { @@ -331,42 +361,16 @@ impl<'b> MaintenanceAccessor<&'b [f32]> for Accessor<'_> { async fn read_vectors( &mut self, ids: &[u32], - mut out: MutMatrixView<'_, f32>, + out: MutMatrixView<'_, f32>, ) -> Result<(), ANNError> { - let faults = self.provider.faults; - for (row, &id) in ids.iter().enumerate().rev() { - let last = row + 1 == ids.len(); - if last && faults.skip_vector { - continue; - } - let out = out.row_mut(row); - out.copy_from_slice(self.vector(id)?); - if last && faults.nan_vector { - out[0] = f32::NAN; - } - } + let rows = ids + .iter() + .map(|&id| self.vector(id)) + .collect::, _>>()?; + write_rows(self.provider.faults, out, &rows); Ok(()) } - async fn stage_insert(&mut self, id: &u32, element: &'b [f32]) -> Result { - if element.len() != self.provider.dim { - return Err(error(format!("point {id} has the wrong dimension"))); - } - let known = - |points: &[(u32, Box<[f32]>)]| points.iter().any(|(external, _)| external == id); - if known(&self.provider.points) || known(&self.staged) { - return Err(error(format!("external id {id} already exists"))); - } - let committed = self.provider.points.len(); - let internal = if self.provider.faults.repeat_stage && !self.staged.is_empty() { - committed - } else { - committed + self.staged.len() - }; - self.staged.push((*id, element.into())); - Ok(internal as u32) - } - async fn reserve_list_ids(&mut self, count: usize) -> Result, ANNError> { if self.provider.faults.reuse_lists { return Ok((0..count as u32).collect()); @@ -375,89 +379,61 @@ impl<'b> MaintenanceAccessor<&'b [f32]> for Accessor<'_> { self.provider.next_list += count as u32; Ok((start..self.provider.next_list).collect()) } -} - -impl Apply> for Accessor<'_> { - type Error = ANNError; - - async fn apply(mut self, update: Bootstrap) -> Result<(), ANNError> { - for (id, vector) in update.centroids().iter() { - self.install(id, vector)?; - } - Ok(()) - } -} - -impl Apply> for Accessor<'_> { - type Error = ANNError; - async fn apply(mut self, update: Appends) -> Result<(), ANNError> { + async fn apply(mut self, update: InsertionUpdate) -> Result<(), ANNError> { self.commit_staged(); - for (list, ids) in update.iter() { - self.place(list, ids)?; + for delta in update.centroids() { + match delta { + CentroidDelta::Install { id, centroid } => self.install(*id, centroid)?, + CentroidDelta::Retire { id } => self.retire(*id)?, + } + } + for &delta in update.points() { + match delta { + PointDelta::Append { id, to } => self.place(id, to)?, + PointDelta::Move { id, from, to } => { + self.leave(id, from)?; + self.place(id, to)?; + } + } } Ok(()) } } -impl Apply> for Accessor<'_> { - type Error = ANNError; - - async fn apply(mut self, update: SplitInsert) -> Result<(), ANNError> { - self.commit_staged(); - let (appends, splits, reassignments) = update.into_parts(); - - // Install every child first, since any part of the update may target one. - for split in splits.iter() { - for (child, centroid) in split.children().into_iter().zip(split.centroids()) { - self.install(child, centroid)?; - } - } - for split in splits.iter() { - let parent = split.parent(); - let mut members = self - .provider - .lists - .remove(&parent) - .ok_or_else(|| error(format!("split parent {parent} is not live")))?; - self.provider.centroids.0.remove(&parent); - self.provider.retired.insert(parent); - - let mut evacuated: Vec = split - .evacuations() - .flat_map(|(_, ids)| ids.iter().copied()) - .collect(); - members.sort_unstable(); - evacuated.sort_unstable(); - if members != evacuated { - return Err(error(format!( - "split of {parent} does not evacuate its members" - ))); +impl<'b> StageElements for Accessor<'_> { + async fn stage( + &mut self, + points: &[(u32, &'b [f32])], + out: MutMatrixView<'_, f32>, + ) -> Result, ANNError> { + let mut ids = Vec::with_capacity(points.len()); + for &(external, element) in points { + if element.len() != self.provider.dim { + return Err(error(format!("point {external} has the wrong dimension"))); } - for (list, ids) in split.evacuations() { - self.place(list, ids)?; + let known = |points: &[(u32, Box<[f32]>)]| { + points.iter().any(|&(existing, _)| existing == external) + }; + if known(&self.provider.points) || known(&self.staged) { + return Err(error(format!("external id {external} already exists"))); } + let committed = self.provider.points.len(); + let internal = if self.provider.faults.repeat_stage && !self.staged.is_empty() { + committed + } else { + committed + self.staged.len() + }; + self.staged.push((external, element.into())); + ids.push(internal as u32); } - for (transfer, ids) in reassignments.iter() { - let source = self - .provider - .lists - .get_mut(&transfer.source()) - .ok_or_else(|| error(format!("list {} is not live", transfer.source())))?; - let before = source.len(); - source.retain(|id| !ids.contains(id)); - if before - source.len() != ids.len() { - return Err(error(format!( - "reassigned points are not all in list {}", - transfer.source() - ))); - } - self.place(transfer.destination(), ids)?; - } - for (list, ids) in appends.iter() { - self.place(list, ids)?; + + let rows: Vec<&[f32]> = points.iter().map(|&(_, element)| element).collect(); + write_rows(self.provider.faults, out, &rows); + if self.provider.faults.drop_stage { + ids.pop(); } - Ok(()) + Ok(ids) } } @@ -465,7 +441,7 @@ impl Apply> for Accessor<'_> { #[derive(Debug, Default)] pub(crate) struct Strategy; -impl<'a> MaintenanceStrategy<'a, Provider, &[f32]> for Strategy { +impl<'a> MaintenanceStrategy<'a, Provider> for Strategy { type MaintenanceAccessor = Accessor<'a>; type Error = ANNError; diff --git a/diskann/src/ivf/update.rs b/diskann/src/ivf/update.rs index 166008f32b..d2340be3d9 100644 --- a/diskann/src/ivf/update.rs +++ b/diskann/src/ivf/update.rs @@ -3,406 +3,95 @@ * Licensed under the MIT license. */ -//! Declarative partition updates consumed by [`Apply`]. +//! Declarative partition updates applied by [`MaintenanceAccessor::apply`]. //! -//! Each distinct operation is its own type: [`Bootstrap`], [`Appends`], [`Splits`], and -//! [`Reassignments`]. [`SplitInsert`] composes the last three for an insert batch that -//! splits lists. +//! An [`InsertionUpdate`] lists the primitive changes made by bootstrap or one insert +//! batch: lists installed or retired, as [`CentroidDelta`]s, and points appended or +//! moved, as [`PointDelta`]s. //! //! Only the index constructs updates. It validates everything it reads from the //! provider or its caller, and its planner guarantees the invariants documented on -//! each type by construction. Providers receive updates through [`Apply`] and can only -//! read them. +//! [`InsertionUpdate`] by construction. Providers can only read updates. //! -//! Point memberships are stored grouped: group keys are ascending and unique, and each -//! group's point ids are ascending in one contiguous slice. Providers can therefore -//! write each list once and process distinct groups in parallel. -//! -//! [`Apply`]: crate::ivf::dynamic::Apply - -use std::cmp::Ordering; - -use diskann_utils::views::{Matrix, MatrixView}; - -use crate::{ivf::grouped::Grouped, utils::VectorId}; - -/////////////////// -// CentroidBlock // -/////////////////// - -/// Freshly reserved centroids with full-precision vectors, one row per id. -/// -/// A block is never empty, its ids are fresh and unique, its rows share one non-zero -/// dimension, and every value is finite. -#[derive(Debug, PartialEq)] -pub struct CentroidBlock { - ids: Box<[L]>, - vectors: Matrix, -} - -impl CentroidBlock { - /// Pair `ids[i]` with row `i` of `vectors`. - pub(super) fn new(ids: Vec, vectors: Matrix) -> Self { - debug_assert_eq!(ids.len(), vectors.nrows()); - Self { - ids: ids.into(), - vectors, - } - } - - /// Number of centroids. - #[expect(clippy::len_without_is_empty, reason = "a block is never empty")] - pub fn len(&self) -> usize { - self.ids.len() - } - - /// Centroid ids in row order. - pub fn ids(&self) -> &[L] { - &self.ids - } +//! [`MaintenanceAccessor::apply`]: crate::ivf::dynamic::MaintenanceAccessor::apply - /// Centroid vectors; row `i` belongs to `self.ids()[i]`. - pub fn vectors(&self) -> MatrixView<'_, f32> { - self.vectors.as_view() - } - - /// Iterate `(id, vector)` pairs in row order. - pub fn iter(&self) -> impl ExactSizeIterator { - self.ids.iter().copied().zip(self.vectors.row_iter()) - } - - /// Release the ids and vectors, for example to adopt the matrix without copying. - pub fn into_parts(self) -> (Box<[L]>, Matrix) { - (self.ids, self.vectors) - } -} - -/////////////// -// Bootstrap // -/////////////// - -/// The initial centroids of an empty index. Installs no points. -#[derive(Debug, PartialEq)] -pub struct Bootstrap { - centroids: CentroidBlock, -} - -impl Bootstrap { - /// Install `centroids` as the first live lists. - pub(super) fn new(centroids: CentroidBlock) -> Self { - Self { centroids } - } +use crate::utils::VectorId; - /// The initial centroids. - pub fn centroids(&self) -> &CentroidBlock { - &self.centroids - } - - /// Release the initial centroids. - pub fn into_centroids(self) -> CentroidBlock { - self.centroids - } +/// A change to the set of live lists. +#[derive(Debug, Clone, PartialEq)] +pub enum CentroidDelta { + /// Make the fresh list `id` live with the full-precision `centroid`. + Install { id: L, centroid: Box<[f32]> }, + /// Retire the live list `id`. Retired ids are never reused. + Retire { id: L }, } -///////////// -// Appends // -///////////// - -/// Newly staged points placed into lists, grouped by destination list. -/// -/// Every point staged for the operation appears exactly once. -#[derive(Debug, PartialEq, Eq)] -pub struct Appends { - by_list: Grouped, +/// A change to one point's list membership. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PointDelta { + /// The point `id`, staged by this accessor, joins list `to`. + Append { id: Id, to: L }, + /// The point `id` leaves list `from` and joins list `to`. + Move { id: Id, from: L, to: L }, } -impl Appends { - /// Place `ids[i]` into `lists[i]`. - pub(super) fn new(ids: &[Id], lists: &[L]) -> Self { - debug_assert_eq!(ids.len(), lists.len()); - let pairs = lists.iter().copied().zip(ids.iter().copied()).collect(); - Self { - by_list: Grouped::from_pairs(pairs), +impl PointDelta { + /// The point joining a list. + pub fn id(&self) -> Id { + match *self { + Self::Append { id, .. } | Self::Move { id, .. } => id, } } - /// Number of appended points. - pub fn len(&self) -> usize { - self.by_list.len() - } - - /// Whether no points are appended. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Iterate `(list, ids)` groups in ascending list order. - pub fn iter(&self) -> impl ExactSizeIterator { - self.by_list.iter() - } -} - -////////////// -// Transfer // -////////////// - -/// The source and destination of reassigned points. The two never match. -/// -/// Ordered by destination, then source. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct Transfer { - source: L, - destination: L, -} - -impl Transfer { - /// The list the points leave. - pub fn source(&self) -> L { - self.source - } - - /// The list the points join. - pub fn destination(&self) -> L { - self.destination - } -} - -impl Ord for Transfer { - fn cmp(&self, other: &Self) -> Ordering { - self.destination - .cmp(&other.destination) - .then_with(|| self.source.cmp(&other.source)) - } -} - -impl PartialOrd for Transfer { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -//////////// -// Splits // -//////////// - -/// Parents retired by binary splits, their two children each, and where every -/// existing member of each parent goes. -/// -/// Parents are ascending and unique. Rows `2 * i` and `2 * i + 1` of -/// [`Splits::children`] replace parent `i`; children are never parents and may end up -/// empty. Every member a parent had when the split was planned is evacuated exactly -/// once, never into a parent. Newly staged points routed to a parent are placed by -/// [`Appends`] instead. -#[derive(Debug, PartialEq)] -pub struct Splits { - parents: Box<[L]>, - children: CentroidBlock, - /// Keyed by `(parent, destination)`. - evacuations: Grouped<(L, L), Id>, -} - -impl Splits { - /// Retire `parents`, replacing parent `i` with rows `2i` and `2i + 1` of `children`. - /// - /// Each item of `evacuations` is `(parent, members, lists)`: `members[j]` moves to - /// `lists[j]`. - pub(super) fn new<'a>( - parents: Vec, - children: CentroidBlock, - evacuations: impl IntoIterator, - ) -> Self { - debug_assert_eq!(children.len(), 2 * parents.len()); - let pairs = evacuations - .into_iter() - .flat_map(|(parent, members, lists)| { - debug_assert_eq!(members.len(), lists.len()); - lists - .iter() - .zip(members) - .map(move |(&list, &id)| ((parent, list), id)) - }) - .collect(); - Self { - parents: parents.into(), - children, - evacuations: Grouped::from_pairs(pairs), + /// The list the point joins. + pub fn to(&self) -> L { + match *self { + Self::Append { to, .. } | Self::Move { to, .. } => to, } } - - /// Number of splits. - #[expect(clippy::len_without_is_empty, reason = "there is always a split")] - pub fn len(&self) -> usize { - self.parents.len() - } - - /// Every child centroid; rows `2 * i` and `2 * i + 1` replace the `i`-th split's - /// parent. - pub fn children(&self) -> &CentroidBlock { - &self.children - } - - /// Number of existing points evacuated from all parents. - pub fn num_evacuated(&self) -> usize { - self.evacuations.len() - } - - /// Iterate the individual splits in ascending parent order. - pub fn iter(&self) -> impl ExactSizeIterator> { - (0..self.parents.len()).map(|index| Split { - splits: self, - index, - }) - } -} - -/// One parent of [`Splits`] and the two children replacing it. -#[derive(Debug, Clone, Copy)] -pub struct Split<'a, Id, L> { - splits: &'a Splits, - index: usize, -} - -impl<'a, Id: VectorId, L: VectorId> Split<'a, Id, L> { - /// The retired parent list. - pub fn parent(self) -> L { - self.splits.parents[self.index] - } - - /// The two new child lists. - pub fn children(self) -> [L; 2] { - let ids = self.splits.children.ids(); - [ids[2 * self.index], ids[2 * self.index + 1]] - } - - /// The two child centroids, in the order of [`Self::children`]. - pub fn centroids(self) -> [&'a [f32]; 2] { - let vectors = &self.splits.children.vectors; - [vectors.row(2 * self.index), vectors.row(2 * self.index + 1)] - } - - /// Where the parent's existing members go, as `(destination, ids)` groups in - /// ascending destination order. - pub fn evacuations(self) -> impl ExactSizeIterator { - let evacuations = &self.splits.evacuations; - let parent = self.parent(); - let start = evacuations - .keys() - .partition_point(|&(from, _)| from < parent); - let end = evacuations - .keys() - .partition_point(|&(from, _)| from <= parent); - evacuations - .groups(start..end) - .map(|((_, destination), ids)| (destination, ids)) - } } -/////////////////// -// Reassignments // -/////////////////// - -/// Existing points moved between lists that survive the update. +/// The complete partition change made by bootstrap or one insert batch. /// -/// Grouped by [`Transfer`] in ascending destination-then-source order; ids ascend -/// within each transfer. -#[derive(Debug, PartialEq, Eq)] -pub struct Reassignments { - by_transfer: Grouped, Id>, -} - -impl Reassignments { - /// Each item of `moves` is `(source, members, lists)`: `members[j]` moves to - /// `lists[j]`, or stays when that is `source`. - pub(super) fn new<'a>(moves: impl IntoIterator) -> Self { - let pairs = moves - .into_iter() - .flat_map(|(source, members, lists)| { - debug_assert_eq!(members.len(), lists.len()); - lists - .iter() - .zip(members) - .filter(move |&(&destination, _)| destination != source) - .map(move |(&destination, &id)| { - ( - Transfer { - source, - destination, - }, - id, - ) - }) - }) - .collect(); - Self { - by_transfer: Grouped::from_pairs(pairs), - } - } - - /// Number of reassigned points. - pub fn len(&self) -> usize { - self.by_transfer.len() - } - - /// Whether no points are reassigned. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Iterate `(transfer, ids)` groups in ascending destination-then-source order. - pub fn iter(&self) -> impl ExactSizeIterator, &[Id])> { - self.by_transfer.iter() - } -} - -///////////////// -// SplitInsert // -///////////////// - -/// An insert batch that splits lists: the new points, the splits they trigger, and -/// the reassignment of existing points around the splits. +/// The index guarantees, and providers may rely on without re-checking, that: /// -/// Every point appears in exactly one part. Nothing is appended or reassigned into a -/// split parent, and no reassignment leaves a parent (its evacuation empties it) or a -/// new child. +/// - installed ids were reserved by this accessor, and their centroids are finite +/// with the accessor's dimension; +/// - retired ids are live; +/// - each list appears in at most one centroid delta, and each point in at most one +/// point delta; +/// - every point staged by this accessor is appended; +/// - a move's `from` holds the point and differs from `to`; +/// - every destination is live after the update: installed, or live and not retired; +/// - every member of a retired list is moved out, so a provider may drop a retired +/// list wholesale and only place the points moved out of it; and +/// - [`Self::points`] is sorted by destination, then id, so +/// `points().chunk_by(|a, b| a.to() == b.to())` yields each list's incoming points +/// in one run. #[derive(Debug, PartialEq)] -pub struct SplitInsert { - appends: Appends, - splits: Splits, - reassignments: Reassignments, +pub struct InsertionUpdate { + centroids: Vec>, + points: Vec>, } -impl SplitInsert { - /// Combine the parts of one split insert. +impl InsertionUpdate { + /// Combine `centroids` and `points`, sorting `points` by destination, then id. pub(super) fn new( - appends: Appends, - splits: Splits, - reassignments: Reassignments, + centroids: Vec>, + mut points: Vec>, ) -> Self { - Self { - appends, - splits, - reassignments, - } - } - - /// Newly staged points and their final lists. - pub fn appends(&self) -> &Appends { - &self.appends + points.sort_unstable_by_key(|point| (point.to(), point.id())); + Self { centroids, points } } - /// The splits and their parents' evacuations. - pub fn splits(&self) -> &Splits { - &self.splits - } - - /// Existing points moved between surviving lists. - pub fn reassignments(&self) -> &Reassignments { - &self.reassignments + /// The lists installed or retired. + pub fn centroids(&self) -> &[CentroidDelta] { + &self.centroids } - /// Release the parts. - pub fn into_parts(self) -> (Appends, Splits, Reassignments) { - (self.appends, self.splits, self.reassignments) + /// The points appended or moved, sorted by destination, then id. + pub fn points(&self) -> &[PointDelta] { + &self.points } } @@ -410,103 +99,31 @@ impl SplitInsert { mod tests { use super::*; - /// A block whose values are `0, 1, 2, ...` in row-major order. - fn block(ids: &[u32], dim: usize) -> CentroidBlock { - let data = (0..ids.len() * dim).map(|x| x as f32).collect(); - CentroidBlock::new( - ids.to_vec(), - Matrix::try_from(data, ids.len(), dim).unwrap(), - ) - } - - fn owned<'a, K, V: Copy + 'a>(groups: impl Iterator) -> Vec<(K, Vec)> { - groups.map(|(key, values)| (key, values.to_vec())).collect() - } - - #[test] - fn centroid_block_pairs_ids_with_rows() { - let bootstrap = Bootstrap::new(block(&[4, 2], 3)); - let block = bootstrap.centroids(); - assert_eq!(block.len(), 2); - assert_eq!(block.ids(), &[4, 2]); - assert_eq!(block.vectors().row(1), &[3.0, 4.0, 5.0]); - assert_eq!( - owned(block.iter()), - vec![(4, vec![0.0, 1.0, 2.0]), (2, vec![3.0, 4.0, 5.0])] - ); - - let (ids, vectors) = bootstrap.into_centroids().into_parts(); - assert_eq!(&*ids, &[4, 2]); - assert_eq!(vectors.nrows(), 2); - } - #[test] - fn appends_group_ids_by_list() { - let appends = Appends::::new(&[5, 1, 3, 2], &[7, 9, 7, 9]); - assert_eq!(appends.len(), 4); - assert_eq!( - owned(appends.iter()), - vec![(7, vec![3, 5]), (9, vec![1, 2])] + fn points_are_sorted_by_destination_then_id() { + let update = InsertionUpdate::::new( + vec![CentroidDelta::Retire { id: 4 }], + vec![ + PointDelta::Move { + id: 3, + from: 4, + to: 9, + }, + PointDelta::Append { id: 8, to: 7 }, + PointDelta::Append { id: 1, to: 9 }, + PointDelta::Move { + id: 2, + from: 4, + to: 7, + }, + ], ); - assert!(Appends::::new(&[], &[]).is_empty()); - } - - #[test] - fn splits_expose_each_split() { - let splits = Splits::::new( - vec![10, 20], - block(&[30, 31, 40, 41], 2), - [(10, &[3, 1, 2][..], &[31, 30, 40][..]), (20, &[], &[])], - ); - assert_eq!(splits.len(), 2); - assert_eq!(splits.children().ids(), &[30, 31, 40, 41]); - assert_eq!(splits.num_evacuated(), 3); - - let split: Vec<_> = splits.iter().collect(); - assert_eq!(split[0].parent(), 10); - assert_eq!(split[0].children(), [30, 31]); - assert_eq!(split[0].centroids(), [&[0.0f32, 1.0][..], &[2.0, 3.0][..]]); - assert_eq!( - owned(split[0].evacuations()), - vec![(30, vec![1]), (31, vec![3]), (40, vec![2])] - ); - assert_eq!(split[1].parent(), 20); - assert_eq!(split[1].children(), [40, 41]); - assert_eq!(split[1].evacuations().len(), 0); - } - - #[test] - fn reassignments_skip_members_that_stay() { - let reassignments = Reassignments::::new([ - (7, &[1, 2, 3][..], &[7, 9, 8][..]), - (8, &[4, 5][..], &[9, 8][..]), - ]); - assert_eq!(reassignments.len(), 3); - let transfers: Vec<_> = reassignments + let order: Vec<_> = update + .points() .iter() - .map(|(t, ids)| (t.source(), t.destination(), ids.to_vec())) + .map(|point| (point.to(), point.id())) .collect(); - assert_eq!( - transfers, - vec![(7, 8, vec![3]), (7, 9, vec![2]), (8, 9, vec![4])] - ); - assert!(Reassignments::::new([(7, &[1][..], &[7][..])]).is_empty()); - } - - #[test] - fn split_insert_releases_its_parts() { - let update = SplitInsert::new( - Appends::::new(&[100], &[30]), - Splits::new(vec![10], block(&[30, 31], 1), [(10, &[1][..], &[31][..])]), - Reassignments::new([(7, &[2][..], &[30][..])]), - ); - assert_eq!(update.appends().len(), 1); - assert_eq!(update.splits().num_evacuated(), 1); - assert_eq!(update.reassignments().len(), 1); - let (appends, splits, reassignments) = update.into_parts(); - assert_eq!( - (appends.len(), splits.len(), reassignments.len()), - (1, 1, 1) - ); + assert_eq!(order, vec![(7, 2), (7, 8), (9, 1), (9, 3)]); + assert_eq!(update.centroids(), &[CentroidDelta::Retire { id: 4 }]); } } diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index a03d3db2ba..eec9ec55ca 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -106,8 +106,8 @@ coordinates its provider, centroid index, and inverted-list store. low-level storage calls. 7. Keep canonical maintenance vectors independent from optimized list scan representations such as PQ codes. -8. Reuse `DataProvider` identity types, execution contexts, ranked errors, and - the graph module's strategy/accessor style. +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. @@ -127,10 +127,10 @@ coordinates its provider, centroid index, and inverted-list store. ### Terminology -- **Point id**: the provider's internal vector id. External ids remain owned by - `DataProvider`. +- **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 (`ListProvider::ListId`), so + 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 @@ -191,8 +191,8 @@ implementation. diskann/src/ivf/ mod.rs dynamic.rs contracts proposed by this RFC - update.rs declarative partition updates consumed by Apply - grouped.rs flat grouped layouts shared by updates and planning + 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 @@ -271,33 +271,51 @@ 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. Providers -declare their list id type once, alongside their point id type: +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 ListProvider: DataProvider { +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, Provider, T>: Send + Sync -where - Provider: ListProvider, -{ - type SearchAccessor: SearchAccessor< - Id = Provider::InternalId, - ListId = Provider::ListId, - >; +pub trait SearchStrategy<'a, P: Provider, T>: Send + Sync { + type SearchAccessor: SearchAccessor; type Error: StandardError; fn search_accessor( &'a self, - provider: &'a Provider, - context: &'a Provider::Context, + 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 @@ -342,79 +360,74 @@ after an unrecoverable partial graph update. ### Declarative Partition Updates The algorithm describes the desired logical result rather than issuing ordered -storage commands. Each distinct operation is its own type in `ivf::update`: +storage commands. The insert path's update in `ivf::update` lists primitive +changes: ```rust -// Fresh centroids: one contiguous matrix, row i belongs to ids[i]. -pub struct CentroidBlock { ids: Box<[L]>, vectors: Matrix } - -pub struct Bootstrap { centroids: CentroidBlock } - -// (i) Newly staged points placed into lists, grouped by destination. -pub struct Appends { /* list -> ids */ } - -// (ii) Binary splits: rows 2i and 2i+1 of `children` replace parents[i], and -// every existing member of each parent is evacuated exactly once. -pub struct Splits { - parents: Box<[L]>, - children: CentroidBlock, - /* (parent, destination) -> ids */ +pub enum CentroidDelta { + Install { id: L, centroid: Box<[f32]> }, + Retire { id: L }, } -// (iii) Existing points moved between surviving lists. -pub struct Reassignments { /* Transfer { source, destination } -> ids */ } +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 }, +} -// An insert batch that splits. -pub struct SplitInsert { - appends: Appends, - splits: Splits, - reassignments: Reassignments, +pub struct InsertionUpdate { + centroids: Vec>, + points: Vec>, // sorted by destination, then id } ``` -Fields are private and constructors are private to the IVF module, so only the -index can create updates. Providers receive them through `Apply` and can only -read them. This keeps the provider contract one-directional: the algorithm +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 builds each update from the items it read from the accessor, such as -staged ids or a list's members, together with a decision array aligned index for -index. That makes "every staged point is placed once" and "every parent member -is evacuated once" properties of how the update is built, so constructors only -group their input and perform no validation. - -Several invariants are unrepresentable: - -- a move with neither a source nor a destination; -- centroids of different dimensions; -- retiring a list without evacuating it, since only `Splits` retires lists; -- a dissolve inside an insert or a split inside a delete; and -- a split into anything other than two children. - -The planner guarantees the rest by construction, and each update type documents -them for providers: - -- split children are fresh, and no parent receives points, since parents are - never placement candidates; -- a reassignment never has matching source and destination, and never leaves a - parent or a new child; and -- every point appears in exactly one part of an update. +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 memberships are grouped: group keys are ascending and unique, and each -group's point ids are ascending in one contiguous slice. `Appends` groups by -destination list; `Reassignments` groups by transfer in destination-then-source -order; `Splits` groups each parent's evacuation by destination. A provider can -therefore append each list's incoming points with one write, drop a split -parent wholesale instead of removing its points one by one, re-encode residual -payloads against one centroid per group, and process distinct groups in -parallel. No point vectors cross the update boundary; child centroids cross -once, in one matrix that k-means writes in place and that a graph centroid index -can bulk-insert. Owned buffers move into the update in O(1), and `into_parts` -lets a provider adopt them. +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 @@ -423,9 +436,9 @@ it enters, and builds updates without further checks: - caller input: initial centroids must be non-empty, finite, and of the provider's dimension; -- accessor reads: every vector row must be written with finite values (rows - start as NaN, so unwritten rows are caught by the same check); -- staging: staged ids must be distinct; and +- 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 @@ -436,10 +449,9 @@ before any update is built, rather than receiving an inconsistent update. View-dependent facts cannot be checked by the index. It guarantees them by building updates only from what the same accessor returned: -- child and bootstrap ids were reserved by this accessor; -- split parents and destinations are live in the accessor's planning view; -- every evacuated or reassigned member belonged to its source list in that - view; and +- 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. @@ -454,76 +466,78 @@ the co-location types are not part of the code yet. ### Maintenance Accessor -The maintenance accessor combines planning reads and mutation staging: +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: HasId + Send + Sized -where - T: Send, -{ - type ExternalId: PartialEq + Send + Sync + 'static; - type ListId: VectorId; - type Centroids: CentroidIndex< - ListId = Self::ListId, - Error = Self::Error, - >; +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_metadata( + fn list_size( &mut self, - list: Self::ListId, - ) -> impl SendFuture>; + list: P::ListId, + ) -> impl SendFuture>; fn read_members( &mut self, - list: Self::ListId, - ) -> impl SendFuture>; + list: P::ListId, + ) -> impl SendFuture>; fn read_vectors( &mut self, - ids: &[Self::Id], + ids: &[P::InternalId], out: MutMatrixView<'_, f32>, ) -> impl SendFuture>; - fn stage_insert( - &mut self, - id: &Self::ExternalId, - element: T, - ) -> impl SendFuture>; - fn reserve_list_ids( &mut self, count: usize, - ) -> impl SendFuture, Self::Error>>; -} + ) -> impl SendFuture, Self::Error>>; -pub trait Apply: Send + Sized { - type Error: ToRanked + Debug + Send + Sync + 'static; + fn apply( + self, + update: InsertionUpdate, + ) -> impl SendFuture>; +} - fn apply(self, update: U) -> impl SendFuture>; +pub trait StageElements: MaintenanceAccessor

{ + fn stage( + &mut self, + points: &[(P::ExternalId, T)], + out: MutMatrixView<'_, f32>, + ) -> impl SendFuture, Self::Error>>; } ``` -Reads are shaped by how much data they move. List reads, `list_metadata` and +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 staged point and -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_insert` owns external/internal id allocation and canonical input storage -for this operation. `read_vectors` also covers points staged by this accessor, -exposing the provider's canonical `f32` decoding without making those points -visible through ordinary provider reads. 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. +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 @@ -532,11 +546,10 @@ 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. A provider implements `Apply` once per update type it -supports: `Bootstrap`, `Appends`, and `SplitInsert` for the insert path, or only -`Appends` for a provider that never splits. The index requires exactly the -`Apply` bounds each operation uses. `apply` is not a generic transaction -protocol: +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. @@ -551,33 +564,24 @@ The maintenance accessor is produced by a strategy that borrows the provider exclusively: ```rust -pub trait MaintenanceStrategy<'a, Provider, T>: Send + Sync -where - Provider: ListProvider, - T: Send, -{ - type MaintenanceAccessor: MaintenanceAccessor< - T, - Id = Provider::InternalId, - ExternalId = Provider::ExternalId, - ListId = Provider::ListId, - >; +pub trait MaintenanceStrategy<'a, P: Provider>: Send + Sync { + type MaintenanceAccessor: MaintenanceAccessor

; type Error: StandardError; fn maintenance_accessor( &'a self, - provider: &'a mut Provider, - context: &'a Provider::Context, + provider: &'a mut P, + context: &'a P::Context, ) -> Result; } ``` -Search strategies receive `&'a Provider`; maintenance strategies receive -`&'a mut Provider`. See [Index Borrowing Model](#index-borrowing-model). +Search strategies receive `&'a P`; maintenance strategies receive `&'a mut P`. +See [Index Borrowing Model](#index-borrowing-model). -Because both strategies fix their accessors' ids to the provider's, the index -names updates directly in terms of the provider, for example -`S::MaintenanceAccessor: Apply>`. +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 @@ -590,14 +594,14 @@ 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 `DataProvider`. +for compatibility with the selected centroid index and the provider. ### Index Borrowing Model The initial dynamic index uses exclusive mutation: ```rust -impl DynamicIvfIndex

{ +impl DynamicIvfIndex

{ pub fn search(&self, /* ... */) -> impl SendFuture>; pub fn insert_batch(&mut self, /* ... */) @@ -611,13 +615,13 @@ impl DynamicIvfIndex

{ 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 Provider` 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 `&Provider`, every provider would need -interior synchronization even though the index already guarantees exclusivity. +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 @@ -638,8 +642,7 @@ Bootstrap uses the same maintenance interface as later mutations: 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 a `Bootstrap` holding those centroids as one `CentroidBlock` and no - points. +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. @@ -662,9 +665,10 @@ Note that co-location groups are never accessed during search. ### Batch Insert 1. Construct one maintenance accessor for the operation. -2. Stage every external id and input element, obtaining internal ids. - Duplicate or already-live external ids fail according to provider semantics. -3. Read the staged points' canonical vectors into one matrix. +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. @@ -680,10 +684,10 @@ Note that co-location groups are never accessed during search. 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. If nothing splits, build `Appends` from the routes. Otherwise build - `Appends` for the staged points' final lists, `Splits` evacuating every - parent, and `Reassignments` for moved neighbor members, and combine them - into one `SplitInsert`. +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 @@ -719,10 +723,10 @@ in-memory lists. Those choices do not change the semantics of the update. 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 delete update containing point deletes, retired victim centroids, - and the evacuation of each victim's remaining members. Its types are not - defined yet; they follow the same pattern as the insert updates, with an - evacuation shared with `Splits`. +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 @@ -737,7 +741,7 @@ 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 Provider`; and + accessor as `&mut P`; and - no search/mutation overlap through the same index value. Within an operation, providers may parallelize routing, list reads, decoding, @@ -748,7 +752,7 @@ Concrete implementations may offer more: | Provider | Possible private unified-view mechanism | |---|---| -| In memory | exclusive `&mut Provider` borrow, one lock guard, copy-on-write root, or epoch | +| 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 | @@ -814,7 +818,7 @@ and search parameters. Durable metadata belongs to concrete provider crates. - Build one provider containing canonical vectors, reverse assignments, in-memory lists, and an exact centroid index. -- Implement maintenance through the exclusive `&mut Provider` borrow with +- 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. @@ -933,16 +937,31 @@ Adding `append`, `remove`, `move`, `create_list`, and `retire_list` exposes a partial-write schedule to GraphIVF. A complete declarative update lets each provider choose a suitable physical implementation. -### One Monolithic Update With Per-Point Moves - -An earlier draft used one `PartitionUpdate` holding a free-form centroid delta -and a flat list of `PointMove { id, from: Option, to: Option }`. It was -compact to define but overloaded: one type encoded insert, delete, reassignment, -split, and dissolve, so every provider had to handle every combination and -regroup moves by list. Its public fields also allowed invalid values, such as a -move with neither endpoint, a retired list that was never emptied, or ragged -centroid dimensions. Separate types that only the index can construct remove -those states and hand providers data that is already grouped. +### 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 @@ -952,7 +971,7 @@ 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` +removes that duplication. Provider crates exercise their `apply` implementations through the index instead of hand-built updates. ### Use `SetElement` Independently @@ -976,11 +995,11 @@ useful common shape. 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 `stage_insert`, and deletion staging once it is +3. **Point staging**: should `StageElements`, and deletion staging once it is added, remain on the maintenance accessor, or should a future - `TransactionalDataProvider` supply them? -4. **Update granularity**: updates are grouped by destination list, so bulk - backends can rewrite each touched list once. Should they also carry a + 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 From ced8da042298a289bff8e5fe44f8d83cde011e9c Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Tue, 6 Oct 2026 09:56:42 -0400 Subject: [PATCH 10/33] another refactor --- diskann/src/ivf/dynamic.rs | 254 ++++++------- diskann/src/ivf/grouped.rs | 166 --------- diskann/src/ivf/index.rs | 605 +++++++++++++------------------ diskann/src/ivf/mod.rs | 6 +- diskann/src/ivf/online/batch.rs | 115 ------ diskann/src/ivf/online/mod.rs | 86 +---- diskann/src/ivf/online/split.rs | 409 ++++----------------- diskann/src/ivf/test/mod.rs | 8 - diskann/src/ivf/test/provider.rs | 458 ----------------------- diskann/src/ivf/update.rs | 144 +++----- 10 files changed, 530 insertions(+), 1721 deletions(-) delete mode 100644 diskann/src/ivf/grouped.rs delete mode 100644 diskann/src/ivf/online/batch.rs delete mode 100644 diskann/src/ivf/test/mod.rs delete mode 100644 diskann/src/ivf/test/provider.rs diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index 84e7d80a84..e201ab97a8 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -10,14 +10,15 @@ //! mutation. Concrete accessors are responsible for presenting a coherent view //! across those components. -use std::fmt::Debug; +use std::{fmt::Debug, ops::Deref}; use diskann_utils::{future::SendFuture, views::MutMatrixView}; use crate::{ + ANNResult, error::{StandardError, ToRanked}, - ivf::update::InsertionUpdate, - provider::{ExecutionContext, HasId}, + ivf::update::Deltas, + provider::{ExecutionContext, Guard, NoopGuard}, utils::VectorId, }; @@ -78,13 +79,33 @@ pub trait Provider: Sized + Send + Sync + 'static { ) -> Result; } +pub trait StageElements { + /// The kind of error yielded by `set_element`. + type StageError: ToRanked + std::fmt::Debug + Send + Sync + 'static; + + /// Stage a new point; its id is provisional until `update` commits it. + fn stage_point( + &mut self, + id: &P::ExternalId, + element: T, + out: &mut [f32], + ) -> impl SendFuture>; + + /// Stage a new centroid and mint its list id, provisional until `update` commits it. + fn stage_centroid(&mut self, centroid: &[f32]) -> impl SendFuture>; +} + +/////////////// +// Centroids // +/////////////// + /// An in-memory index over the live centroids. /// /// The centroid catalog is authoritative. Approximate implementations, such as /// a graph navigator, must retain enough catalog information to reject retired /// ids and recover with exact selection when navigation cannot produce a usable /// result. Exact and graph implementations expose the same interface. -pub trait CentroidIndex: Send + Sync { +pub trait Centroids: Send + Sync { /// Stable logical centroid id, also used as the inverted-list id. type ListId: VectorId; @@ -94,13 +115,16 @@ pub trait CentroidIndex: Send + Sync { /// Number of live centroids. fn len(&self) -> usize; + /// dimension of full precision centroid vectors. + fn dim(&self) -> usize; + /// Whether there are no live centroids. fn is_empty(&self) -> bool { self.len() == 0 } /// Borrow one authoritative full-precision centroid vector. - fn centroid(&self, id: Self::ListId) -> Option<&[f32]>; + fn centroid(&self, id: Self::ListId) -> Option>; /// Select exactly `min(nprobe, self.len())` live centroids for `query`, nearest /// first. @@ -111,156 +135,81 @@ pub trait CentroidIndex: Send + Sync { &self, query: &[f32], nprobe: usize, - ) -> impl SendFuture>, Self::Error>>; + ) -> Result>, Self::Error>; } -/// Query-bound access to centroid selection and inverted-list scanning. -/// -/// This is the dynamic IVF search algorithm's primary extension point, in the -/// same spirit as [`crate::graph::glue::SearchAccessor`]. Implementations are free -/// to batch reads, coalesce blob requests, prefetch, decode quantized payloads, or -/// fan work out across tasks. -pub trait SearchAccessor: HasId + Send + Sync { - /// Stable logical centroid/list id, fixed to [`Provider::ListId`] by the - /// strategy. - type ListId: VectorId; - - /// Errors from list selection or scanning. - type Error: ToRanked + Debug + Send + Sync + 'static; - - /// Select up to `nprobe` lists for the bound query, nearest first. - fn select_lists( - &mut self, - nprobe: usize, - ) -> impl SendFuture>, Self::Error>>; +////////////// +// Accessor // +////////////// - /// Scan `lists`, scoring their members against the bound query. - /// - /// `lists` must come from [`Self::select_lists`] on this accessor, which keeps - /// selection and scanning coherent using provider-specific coordination. - fn scan_lists( - &mut self, - lists: &[SelectedList], - emit: F, - ) -> impl SendFuture> +pub trait InsertAccessor: Send + Sized + StageElements { + /// In-memory centroid catalog and navigator in this accessor's unified view. + type Centroids<'a>: Centroids where - F: FnMut(Self::Id, f32) + Send; -} - -/// Factory for one dynamic IVF search accessor. -pub trait SearchStrategy<'a, P: Provider, T>: Send + Sync { - /// Query-bound accessor used for both coarse and fine search. - type SearchAccessor: SearchAccessor; - - /// Error constructing the accessor. - type Error: StandardError; - - /// Construct an accessor around `query`. - fn search_accessor( - &'a self, - provider: &'a P, - context: &'a P::Context, - query: T, - ) -> Result; -} + Self: 'a; -/// Operation-scoped reads and writes needed by split/dissolve maintenance. -/// -/// The accessor is the consistency boundary for one mutation. It must present a -/// unified view across point data, centroids, and inverted lists for all planning -/// reads. It also owns any provider-specific lock, epoch, transaction, staging area, -/// or poison state needed to apply the final update through [`Self::apply`]. -/// -/// List reads take one list and return its result. [`Self::read_vectors`], which -/// reads far more data, takes a batch so disk and blob providers can reorder, -/// coalesce, or parallelize I/O. -pub trait MaintenanceAccessor: Send + Sized { - /// In-memory centroid catalog and navigator in this accessor's unified view. - type Centroids: CentroidIndex; + // Scoped reader to read vectors mapped to list ids. + type Reader<'a>: Reader + where + Self: 'a; /// Errors from planning reads, staging, or applying the update. type Error: ToRanked + Debug + Send + Sync + 'static; - /// Dimension of every canonical vector and centroid in this accessor's view. - /// - /// Fixed for the provider's lifetime, including before the index is initialized. - fn dim(&self) -> usize; + /// Borrow the list reader to read vectors from posting lists. + fn reader(&self) -> Self::Reader<'_>; /// Borrow the centroid index used for routing and maintenance neighborhoods. - fn centroids(&self) -> &Self::Centroids; + fn centroids(&self) -> Self::Centroids<'_>; + + fn dim(&self) -> usize; /// Read the size of a live list. - fn list_size(&mut self, list: P::ListId) -> impl SendFuture>; + fn list_size(&self, list: P::ListId) -> Result; /// Read the member ids of a live list. - fn read_members( - &mut self, - list: P::ListId, - ) -> impl SendFuture>; + fn get_members(&self, list: P::ListId) -> Result<&[P::InternalId], Self::Error>; - /// Write the canonical vector of `ids[i]` into row `i` of `out`, in any order. - /// - /// `out` has one row per id and [`Self::dim`] columns, and every row must be - /// written. Covers visible points only; the canonical vectors of staged points - /// come from [`StageElements::stage`]. - fn read_vectors( - &mut self, - ids: &[P::InternalId], - out: MutMatrixView<'_, f32>, - ) -> impl SendFuture>; + /// Consume the accessor and write the update + fn update(self, value: Deltas) -> impl SendFuture>; +} - /// Reserve fresh logical list ids that will not alias retired ids. - /// TO DO: Re-evaluate we even need this. - fn reserve_list_ids( - &mut self, - count: usize, - ) -> impl SendFuture, Self::Error>>; +pub trait Reader: Send + Sync { + /// Error type in case a read fails. + type Error: ToRanked + Debug + Send + Sync + 'static; - /// Apply `update` and finish the operation, consuming the accessor. - /// - /// Only the index constructs updates, and it guarantees the invariants documented - /// on [`InsertionUpdate`], so implementations may rely on them without re-checking. + /// Id type + type Id; + + /// dimension of the output of each Id. + fn dim(&self) -> usize; + + /// Write the canonical vector of `ids[i]` into row `i` of `out`, in any order. /// - /// Returning `Ok(())` means all components expose one coherent resulting index. - /// Rollback, durability, concurrent-reader visibility, and recovery after `Err` are - /// intentionally provider-defined. - fn apply( - self, - update: InsertionUpdate, + /// `out` has one row per id and [`Self::dim`] columns, and every row must be + /// written. + fn read_into( + &self, + id: Self::Id, + out: MutMatrixView<'_, T>, ) -> impl SendFuture>; } -/// Stages new points of element type `T` for one maintenance operation. -/// -/// The provider owns id allocation, duplicate detection, and the conversion of `T` to -/// its stored and canonical representations. Staged points become visible only when -/// the accessor applies an update that places them, and are discarded if the accessor -/// is dropped first. A provider that accepts several input types implements this once -/// per type. -pub trait StageElements: MaintenanceAccessor

{ - /// Stage `points` and return their internal ids in order, writing the canonical - /// vector of `points[i]` into row `i` of `out`. - /// - /// `out` has one row per point and [`MaintenanceAccessor::dim`] columns, and every - /// row must be written. - fn stage( - &mut self, - points: &[(P::ExternalId, T)], - out: MutMatrixView<'_, f32>, - ) -> impl SendFuture, Self::Error>>; -} +////////////// +// Strategy // +////////////// /// Factory for an operation-scoped dynamic IVF maintenance accessor. /// /// The provider is borrowed exclusively for the lifetime of the accessor, so no /// search or other mutation can observe it until the accessor is applied or /// dropped. Providers may therefore mutate plain in-memory state in -/// [`MaintenanceAccessor::apply`] without interior synchronization. Providers that +/// [`InsertAccessor::apply`] without interior synchronization. Providers that /// share state outside this borrow (for example through an `Arc`) remain responsible /// for coordinating those aliases. -pub trait MaintenanceStrategy<'a, P: Provider>: Send + Sync { +pub trait MaintenanceStrategy<'a, P: Provider, T: Sync>: Send + Sync { /// Accessor used to plan and stage split/dissolve operations. - type MaintenanceAccessor: MaintenanceAccessor

; + type MaintenanceAccessor: InsertAccessor; /// Error constructing the accessor. type Error: StandardError; @@ -272,3 +221,60 @@ pub trait MaintenanceStrategy<'a, P: Provider>: Send + Sync { context: &'a P::Context, ) -> Result; } + +/// Factory for one dynamic IVF search accessor. +pub trait SearchStrategy<'a, P: Provider, T>: Send + Sync { + /// Query-bound accessor used for both coarse and fine search. + type SearchAccessor: SearchAccessor; + + /// Error constructing the accessor. + type Error: StandardError; + + /// Construct an accessor around `query`. + fn search_accessor( + &'a self, + provider: &'a P, + context: &'a P::Context, + query: T, + ) -> Result; +} + +//////////// +// Search // +//////////// + +/// Query-bound access to centroid selection and inverted-list scanning. +/// +/// This is the dynamic IVF search algorithm's primary extension point, in the +/// same spirit as [`crate::graph::glue::SearchAccessor`]. Implementations are free +/// to batch reads, coalesce blob requests, prefetch, decode quantized payloads, or +/// fan work out across tasks. +pub trait SearchAccessor: Send + Sync { + /// Stable logical centroid/list id, fixed by the strategy. + type ListId; + + type InternalId; + + /// In-memory centroid view. + type Centroids<'a>: Centroids + where + Self: 'a; + + /// Errors from list selection or scanning. + type Error: ToRanked + Debug + Send + Sync + 'static; + + /// Borrow the centroid index used for routing. + fn centroids(&self) -> Self::Centroids<'_>; + + /// Scan `lists`, scoring their members against the bound query. + /// + /// `lists` must come from [`Self::select_lists`] on this accessor, which keeps + /// selection and scanning coherent using provider-specific coordination. + fn scan_lists( + &mut self, + lists: &[SelectedList], + emit: F, + ) -> impl SendFuture> + where + F: FnMut(Self::InternalId, f32) + Send; +} diff --git a/diskann/src/ivf/grouped.rs b/diskann/src/ivf/grouped.rs deleted file mode 100644 index c444a0b588..0000000000 --- a/diskann/src/ivf/grouped.rs +++ /dev/null @@ -1,166 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Flat grouped layouts used by the online planner. - -use std::ops::Range; - -/// Values stored in consecutive groups numbered from zero. -/// -/// Group `i` owns `values[ends[i - 1]..ends[i]]` and may be empty. -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct Csr { - ends: Vec, - values: Vec, -} - -impl Csr { - /// Append a group holding `values`. - pub(super) fn push(&mut self, values: impl IntoIterator) { - self.values.extend(values); - self.ends.push(self.values.len()); - } - - /// Number of groups. - pub(super) fn num_groups(&self) -> usize { - self.ends.len() - } - - /// Positions in [`Self::values`] owned by `group`. - pub(super) fn range(&self, group: usize) -> Range { - let start = group.checked_sub(1).map_or(0, |prev| self.ends[prev]); - start..self.ends[group] - } - - /// Values of `group`. - pub(super) fn group(&self, group: usize) -> &[V] { - &self.values[self.range(group)] - } - - /// Every value, group by group. - pub(super) fn values(&self) -> &[V] { - &self.values - } -} - -impl Csr { - /// For each target below `targets`, the groups containing it in ascending order. - /// - /// Every value must be below `targets`. - pub(super) fn transpose(&self, targets: usize) -> Self { - let mut ends = vec![0; targets]; - for &target in &self.values { - ends[target] += 1; - } - let mut next = Vec::with_capacity(targets); - let mut total = 0; - for end in &mut ends { - next.push(total); - total += *end; - *end = total; - } - - let mut values = vec![0; total]; - for group in 0..self.num_groups() { - for &target in self.group(group) { - values[next[target]] = group; - next[target] += 1; - } - } - Self { ends, values } - } -} - -/// Values grouped under ascending unique keys, ascending within each group. -#[derive(Debug, PartialEq, Eq)] -pub(super) struct Grouped { - keys: Box<[K]>, - groups: Csr, -} - -impl Grouped { - /// Group distinct `pairs` by key. - pub(super) fn from_pairs(mut pairs: Vec<(K, V)>) -> Self { - pairs.sort_unstable(); - debug_assert!( - pairs.windows(2).all(|pair| pair[0] != pair[1]), - "grouped pairs must be distinct" - ); - - let same_key = |a: &(K, V), b: &(K, V)| a.0 == b.0; - let keys = pairs.chunk_by(same_key).map(|group| group[0].0).collect(); - let ends = pairs - .chunk_by(same_key) - .scan(0, |end, group| { - *end += group.len(); - Some(*end) - }) - .collect(); - let values = pairs.into_iter().map(|(_, value)| value).collect(); - Self { - keys, - groups: Csr { ends, values }, - } - } - - /// The values under `key`, or nothing if `key` is absent. - pub(super) fn get(&self, key: K) -> &[V] { - self.keys - .binary_search(&key) - .map_or(&[], |index| self.groups.group(index)) - } - - /// Every group with its key, in ascending key order. - pub(super) fn iter(&self) -> impl ExactSizeIterator { - (0..self.keys.len()).map(|index| (self.keys[index], self.groups.group(index))) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn groups(csr: &Csr) -> Vec> { - (0..csr.num_groups()) - .map(|group| csr.group(group).to_vec()) - .collect() - } - - #[test] - fn csr_push_keeps_empty_groups() { - let mut csr = Csr::default(); - csr.push([1, 2]); - csr.push([]); - csr.push([3]); - assert_eq!(groups(&csr), vec![vec![1, 2], vec![], vec![3]]); - assert_eq!(csr.range(2), 2..3); - assert_eq!(csr.values(), &[1, 2, 3]); - } - - #[test] - fn transpose_lists_containing_groups() { - let mut csr = Csr::default(); - csr.push([1, 2]); - csr.push([]); - csr.push([2]); - let transposed = csr.transpose(4); - assert_eq!( - groups(&transposed), - vec![vec![], vec![0], vec![0, 2], vec![]] - ); - } - - #[test] - fn grouped_get_finds_groups_by_key() { - let grouped = Grouped::from_pairs(vec![(7, 2), (5, 1), (7, 0)]); - let groups: Vec<_> = grouped - .iter() - .map(|(key, values)| (key, values.to_vec())) - .collect(); - assert_eq!(groups, vec![(5, vec![1]), (7, vec![0, 2])]); - assert_eq!(grouped.get(7), &[0, 2]); - assert_eq!(grouped.get(6), &[] as &[i32]); - } -} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index fef218c7d7..a314688a25 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -5,12 +5,14 @@ //! Dynamic IVF index orchestration. //! -//! The index drives one [`MaintenanceAccessor`] per mutation: it plans an -//! [`InsertionUpdate`] against the accessor's view and hands it to -//! [`MaintenanceAccessor::apply`]. It never touches point, centroid, or list storage -//! directly. +//! The index drives one [`InsertAccessor`] per mutation: it plans a [`Deltas`] against +//! the accessor's view and hands it to [`InsertAccessor::update`]. It never touches +//! point, centroid, or list storage directly. -use diskann_utils::{future::SendFuture, views::Matrix}; +use diskann_utils::{ + future::SendFuture, + views::{Matrix, MatrixView}, +}; use rand::{SeedableRng, rngs::StdRng}; use thiserror::Error; @@ -19,17 +21,16 @@ use crate::{ error::{ErrorExt, IntoANNResult}, ivf::{ dynamic::{ - CentroidIndex, MaintenanceAccessor, MaintenanceStrategy, Provider, StageElements, + Centroids, InsertAccessor, MaintenanceStrategy, Provider, Reader, StageElements, }, - online::{SplitPlan, StagedBatch, check_reserved, index_error}, - update::{CentroidDelta, InsertionUpdate, PointDelta}, + online::{TwoMeans, index_error, two_means}, + update::{CentroidDelta, Delta, Deltas, MoveTo}, }, - utils::VectorId, }; /// Parameters for the online split policy. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct DynamicIvfConfig { +pub struct Config { /// Split a list once an insert batch would grow it beyond this many points. pub split_threshold: usize, /// Nearby lists reassigned together with each split parent. @@ -67,43 +68,24 @@ pub struct InsertStats { pub reassigned: usize, } -impl InsertStats { - /// Tally an insert's update: every retired list is a split parent. - fn of(update: &InsertionUpdate) -> Self { - let mut stats = Self::default(); - for point in update.points() { - match point { - PointDelta::Append { .. } => stats.inserted += 1, - PointDelta::Move { .. } => stats.reassigned += 1, - } - } - stats.splits = update - .centroids() - .iter() - .filter(|delta| matches!(delta, CentroidDelta::Retire { .. })) - .count(); - stats - } -} - /// An incrementally maintained IVF index. /// /// Mutations take `&mut self` and lend the provider exclusively to one /// [`MaintenanceAccessor`], so they cannot overlap searches through the same value. #[derive(Debug)] -pub struct DynamicIvfIndex { +pub struct IVFIndex { provider: P, - config: DynamicIvfConfig, + config: Config, rng: StdRng, } -impl DynamicIvfIndex

{ +impl IVFIndex

{ /// Construct an index over `provider`. /// /// # Errors /// /// Returns [`ConfigError`] if `config` is invalid. - pub fn new(provider: P, config: DynamicIvfConfig) -> Result { + pub fn new(provider: P, config: Config) -> Result { if config.split_threshold < 2 { return Err(ConfigError::SplitThreshold(config.split_threshold)); } @@ -122,68 +104,68 @@ impl DynamicIvfIndex

{ &self.provider } - /// Install the initial centroids. - /// - /// # Errors - /// - /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index - /// already has centroids, or the maintenance accessor fails. - pub fn initialize<'a, S>( - &'a mut self, - strategy: &'a S, - context: &'a P::Context, - centroids: Matrix, - ) -> impl SendFuture> - where - S: MaintenanceStrategy<'a, P>, - { - let provider = &mut self.provider; - - async move { - let count = centroids.nrows(); - if count == 0 { - return Err(index_error("initialize requires at least one centroid")); - } - if centroids.as_slice().iter().any(|x| !x.is_finite()) { - return Err(index_error("initial centroids must be finite")); - } - - let mut accessor = strategy - .maintenance_accessor(provider, context) - .into_ann_result()?; - - if !accessor.centroids().is_empty() { - return Err(index_error("dynamic IVF index is already initialized")); - } - let dim = accessor.dim(); - if dim == 0 || centroids.ncols() != dim { - return Err(index_error(format!( - "initial centroids have dimension {}, expected {dim}", - centroids.ncols() - ))); - } - - let list_ids = accessor - .reserve_list_ids(count) - .await - .escalate("initialize must reserve list ids")?; - check_reserved(accessor.centroids(), &list_ids, count)?; - - let installs = list_ids - .into_iter() - .zip(centroids.row_iter()) - .map(|(id, centroid)| CentroidDelta::Install { - id, - centroid: centroid.into(), - }) - .collect(); - - accessor - .apply(InsertionUpdate::new(installs, Vec::new())) - .await - .escalate("initialize must apply the initial centroids") - } - } + // /// Install the initial centroids. + // /// + // /// # Errors + // /// + // /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index + // /// already has centroids, or the maintenance accessor fails. + // pub fn initialize<'a, S>( + // &'a mut self, + // strategy: &'a S, + // context: &'a P::Context, + // centroids: Matrix, + // ) -> impl SendFuture> + // where + // S: MaintenanceStrategy<'a, P>, + // { + // let provider = &mut self.provider; + + // async move { + // let count = centroids.nrows(); + // if count == 0 { + // return Err(index_error("initialize requires at least one centroid")); + // } + // if centroids.as_slice().iter().any(|x| !x.is_finite()) { + // return Err(index_error("initial centroids must be finite")); + // } + + // let mut accessor = strategy + // .maintenance_accessor(provider, context) + // .into_ann_result()?; + + // if !accessor.centroids().is_empty() { + // return Err(index_error("dynamic IVF index is already initialized")); + // } + // let dim = accessor.dim(); + // if dim == 0 || centroids.ncols() != dim { + // return Err(index_error(format!( + // "initial centroids have dimension {}, expected {dim}", + // centroids.ncols() + // ))); + // } + + // let list_ids = accessor + // .reserve_list_ids(count) + // .await + // .escalate("initialize must reserve list ids")?; + // check_reserved(&accessor.centroids(), &list_ids, count)?; + + // let installs = list_ids + // .into_iter() + // .zip(centroids.row_iter()) + // .map(|(id, centroid)| CentroidDelta::Install { + // id, + // centroid: centroid.into(), + // }) + // .collect(); + + // accessor + // .apply(InsertionUpdate::new(installs, Vec::new())) + // .await + // .escalate("initialize must apply the initial centroids") + // } + // } /// Insert a batch, splitting every list the batch pushes past `split_threshold`. /// @@ -203,9 +185,8 @@ impl DynamicIvfIndex

{ points: &'a [(P::ExternalId, T)], ) -> impl SendFuture> where - S: MaintenanceStrategy<'a, P>, - S::MaintenanceAccessor: StageElements, - T: Sync, + S: MaintenanceStrategy<'a, P, T>, + T: Sync + Copy, { let provider = &mut self.provider; let config = self.config; @@ -220,283 +201,219 @@ impl DynamicIvfIndex

{ .maintenance_accessor(provider, context) .into_ann_result()?; - if accessor.centroids().is_empty() { - return Err(index_error( - "insert requires an initialized dynamic IVF index", - )); - } - - let batch = StagedBatch::stage(&mut accessor, points).await?; // register points with provider -> internal ids + f32 representations of the vectors. + // stage the input points + let mut vectors = Matrix::::new(f32::NAN, points.len(), accessor.dim()); + let mut ids = Vec::with_capacity(points.len()); + let mut routed = hashbrown::HashMap::>::new(); - // Split every routed list the batch pushes past the threshold. `routed` - // visits lists in ascending order, as `SplitPlan::gather` requires. - let mut parents = Vec::new(); - for (list, incoming) in batch.routed() { - let len = accessor - .list_size(list) + for (position, ((id, v), out)) in points.iter().zip(vectors.row_iter_mut()).enumerate() + { + //stage the input points first. + let id = accessor + .stage_point(id, *v, out) .await - .escalate("insert must read routed list sizes")?; + .escalate("Unable to stage point")?; + ids.push(id); + + // route the points to their nearest centroid + let parent = accessor + .centroids() + .select(out, 1) + .escalate("Unable to select")? + .first() + .map(|s| s.id) + .ok_or_else(|| { + ANNError::message(ANNErrorKind::IndexError, "Didn't return list") + })?; + + routed.entry(parent).or_default().push(position); + } - let projected = len + incoming; + // filter the routed lists for ones that need splitting. + let mut parents = Vec::new(); + for (&list, positions) in &routed { + let len = accessor.list_size(list).escalate("Unable to get len")?; - if projected > config.split_threshold { + if len + positions.len() > config.split_threshold { parents.push(list); } } - - let update = if parents.is_empty() { - batch.into_update() - } else { - SplitPlan::gather(&mut accessor, &config, batch, parents) - .await? - .solve(config.two_means_iterations, rng)? - }; - - let stats = InsertStats::of(&update); + parents.sort_unstable(); + + // for the ones that need splitting - split and re-assign necessary points in neighbors + let update = Self::split::<_, T>( + &mut accessor, + &config, + rng, + &ids, + vectors.as_view(), + &routed, + &parents, + ) + .await?; accessor - .apply(update) + .update(update) .await .escalate("insert must apply its update")?; - Ok(stats) + Ok(InsertStats::default()) } } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::{ - ivf::test::provider::{Faults, Provider, Strategy}, - provider::DefaultContext, - }; - - fn config(split_threshold: usize, reassign_neighbors: usize) -> DynamicIvfConfig { - DynamicIvfConfig { - split_threshold, - reassign_neighbors, - two_means_iterations: 10, - seed: 7, + /// Build the update for one insert batch: split every list in `parents` and append + /// every staged point to the list that ends up holding it. + /// + /// Batch position `i` is the staged point `ids[i]`, whose full-precision vector is + /// row `i` of `vectors`. `routed` maps every routed list to the batch positions + /// routed to it, and `parents` names the subset of those lists that overflow. + /// + /// # Errors + /// + /// Fails if an accessor call fails or a split list holds fewer than two points. + async fn split( + accessor: &mut A, + config: &Config, + rng: &mut StdRng, + ids: &[P::InternalId], + vectors: MatrixView<'_, f32>, + routed: &hashbrown::HashMap>, + parents: &[P::ListId], + ) -> ANNResult> + where + A: InsertAccessor, + T: Sync, + { + let mut deltas = Vec::new(); + + // Lists that keep their centroid take their staged points as they are. + for (&list, positions) in routed { + if parents.contains(&list) { + continue; + } + deltas.push(Delta::PointAppends { + to: list, + ids: positions.iter().map(|&position| ids[position]).collect(), + }); } - } - /// Pair each vector with a fresh external id, counting up from `first`. - fn batch(first: u32, vectors: &[[f32; D]]) -> Vec<(u32, &[f32])> { - (first..) - .zip(vectors.iter().map(|vector| &vector[..])) - .collect() - } + // Overflowing lists are replaced by two children, which take their staged + // points instead. + for &parent in parents { + let staged: &[usize] = routed.get(&parent).map_or(&[], Vec::as_slice); + deltas.extend( + Self::split_list::<_, T>(accessor, config, rng, ids, vectors, parent, staged) + .await?, + ); + } - fn matrix(rows: &[[f32; D]]) -> Matrix { - let data: Box<[f32]> = rows.iter().flatten().copied().collect(); - Matrix::try_from(data, rows.len(), D).unwrap() + Ok(Deltas::new(deltas)) } - async fn initialized( - provider: Provider, - config: DynamicIvfConfig, - centroids: Matrix, - ) -> DynamicIvfIndex { - let mut index = DynamicIvfIndex::new(provider, config).unwrap(); - index - .initialize(&Strategy, &DefaultContext, centroids) - .await - .unwrap(); - index - } + /// Plan one list's split: fit two children over every point the list would hold, + /// then place each of those points on the nearer child and retire the parent. + /// + /// `staged` holds the batch positions routed to `parent`. + /// + /// # Errors + /// + /// Fails if an accessor call fails or the list holds fewer than two points. + async fn split_list( + accessor: &mut A, + config: &Config, + rng: &mut StdRng, + ids: &[P::InternalId], + vectors: MatrixView<'_, f32>, + parent: P::ListId, + staged: &[usize], + ) -> ANNResult>> + where + A: InsertAccessor, + T: Sync, + { + // let reader = accessor.reader(); - #[test] - fn new_rejects_invalid_config() { - let new = |config| DynamicIvfIndex::new(Provider::new(1), config).err(); - assert_eq!(new(config(1, 1)), Some(ConfigError::SplitThreshold(1))); - assert_eq!(new(config(4, 0)), Some(ConfigError::ReassignNeighbors)); - } + // let num_members = accessor.list_size(parent)?; - #[tokio::test] - async fn initialize_rejects_bad_centroids() { - let mut index = DynamicIvfIndex::new(Provider::new(2), config(10, 1)).unwrap(); - let wrong_dim = index - .initialize(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0, 0.0]])) - .await - .unwrap_err(); - assert!(wrong_dim.to_string().contains("dimension"), "{wrong_dim}"); - let non_finite = index - .initialize(&Strategy, &DefaultContext, matrix(&[[f32::NAN, 0.0]])) - .await - .unwrap_err(); - assert!(non_finite.to_string().contains("finite"), "{non_finite}"); - - index - .initialize(&Strategy, &DefaultContext, matrix(&[[0.0, 0.0]])) - .await - .unwrap(); - let again = index - .initialize(&Strategy, &DefaultContext, matrix(&[[1.0, 1.0]])) - .await - .unwrap_err(); - assert!(again.to_string().contains("already initialized"), "{again}"); - } + // let mut points = Matrix::new(f32::NAN, num_members + staged, reader.dim()); - #[tokio::test] - async fn insert_requires_initialization() { - let mut index = DynamicIvfIndex::new(Provider::new(1), config(4, 1)).unwrap(); - let empty = index - .insert_batch(&Strategy, &DefaultContext, &batch::<1>(0, &[])) - .await - .unwrap(); - assert_eq!(empty, InsertStats::default()); - let err = index - .insert_batch(&Strategy, &DefaultContext, &batch(0, &[[1.0]])) - .await - .unwrap_err(); - assert!(err.to_string().contains("initialized"), "{err}"); - } + // reader.read_into(parent, points.as_mut_view()).await?; + let dim = accessor.dim(); + let members = accessor + .get_members(parent) + .escalate("split must read the parent's members")? + .to_vec(); - #[tokio::test] - async fn insert_without_overflow_appends_to_nearest_lists() { - let centroids = matrix(&[[0.0, 0.0], [10.0, 10.0]]); - let mut index = initialized(Provider::new(2), config(10, 1), centroids).await; - let vectors = [[1.0, 1.0], [9.0, 9.0], [2.0, 0.0]]; - let stats = index - .insert_batch(&Strategy, &DefaultContext, &batch(0, &vectors)) - .await - .unwrap(); - assert_eq!( - stats, - InsertStats { - inserted: 3, - splits: 0, - reassigned: 0 - } - ); - - let provider = index.provider(); - provider.check(); - assert_eq!( - provider.member_vectors(0), - vec![vec![1.0, 1.0], vec![2.0, 0.0]] - ); - assert_eq!(provider.member_vectors(1), vec![vec![9.0, 9.0]]); - } + // Everything the parent would hold: its current members, then its staged points. + let mut points = Matrix::new(f32::NAN, members.len() + staged.len(), dim); + { + let mut member_vectors = Matrix::new(f32::NAN, members.len(), dim); + accessor + .reader() + .read_into(parent, member_vectors.as_mut_view()) + .await + .escalate("split must read the parent's vectors")?; - /// Lists A (0), B (20), and C (40) in one dimension. Inserting 0.5, 8, and 9 - /// overflows A, which splits into children near 1/6 and 8.5, with B and C as - /// its neighbors. - #[tokio::test] - async fn split_places_parent_points_and_moves_neighbor_points_only_toward_children() { - let provider = Provider::from_lists( - 1, - &[ - (&[0.0], &[&[-1.0], &[1.0]]), - (&[20.0], &[&[12.0], &[25.0], &[33.0]]), - (&[40.0], &[&[41.0]]), - ], - ); - let mut index = DynamicIvfIndex::new(provider, config(4, 2)).unwrap(); - let stats = index - .insert_batch( - &Strategy, - &DefaultContext, - &batch(100, &[[0.5], [8.0], [9.0]]), - ) - .await - .unwrap(); - // A's members -1 and 1 are evacuated, and 12 moves out of B. - assert_eq!( - stats, - InsertStats { - inserted: 3, - splits: 1, - reassigned: 3 + let sources = member_vectors + .row_iter() + .chain(staged.iter().map(|&position| vectors.row(position))); + for (row, source) in points.row_iter_mut().zip(sources) { + row.copy_from_slice(source); } - ); - - let provider = index.provider(); - provider.check(); - assert_eq!(provider.lists(), vec![1, 2, 3, 4]); - let low = provider.list_of(&[0.5]).unwrap(); - let high = provider.list_of(&[8.0]).unwrap(); - assert!((provider.centroid(low).unwrap()[0] - 1.0 / 6.0).abs() < 1e-6); - assert_eq!(provider.centroid(high).unwrap(), &[8.5]); - assert_eq!( - provider.member_vectors(low), - vec![vec![-1.0], vec![0.5], vec![1.0]] - ); - assert_eq!( - provider.member_vectors(high), - vec![vec![8.0], vec![9.0], vec![12.0]] - ); - // 33 is closer to C than to B, but neighbor points only move toward children. - assert_eq!(provider.member_vectors(1), vec![vec![25.0], vec![33.0]]); - assert_eq!(provider.member_vectors(2), vec![vec![41.0]]); - } + } - #[tokio::test] - async fn repeated_batches_keep_the_partition_valid() { - let centroids = matrix(&[[0.0, 0.0], [10.0, 10.0]]); - let mut index = initialized(Provider::new(2), config(8, 2), centroids).await; - let vectors: Vec<[f32; 2]> = (0..300) - .map(|i| { - let t = i as f32; - [(t * 0.37).sin() * 20.0, (t * 0.11).cos() * 20.0] - }) - .collect(); + let TwoMeans { + centroids, + children, + } = two_means(points.as_view(), config.two_means_iterations, rng)?; - // Insert in batches of 25, checking the partition after each batch. - let mut stats = InsertStats::default(); - for (first, chunk) in (0..).step_by(25).zip(vectors.chunks(25)) { - let batch = index - .insert_batch(&Strategy, &DefaultContext, &batch(first, chunk)) - .await - .unwrap(); - index.provider().check(); - stats.inserted += batch.inserted; - stats.splits += batch.splits; - stats.reassigned += batch.reassigned; + let mut deltas = Vec::new(); + let mut child_ids = Vec::new(); + + // Install both children before anything is placed into them. + for centroid in centroids.row_iter() { + let id = accessor.stage_centroid(centroid).await?; + child_ids.push(id); + + deltas.push(Delta::CentroidDelta { + id, + delta: CentroidDelta::Install { + centroid: centroid.into(), + }, + }); } - let provider = index.provider(); - assert_eq!(stats.inserted, 300); - assert_eq!(provider.len(), 300); - assert!(stats.splits > 1 && stats.reassigned > 0, "{stats:?}"); - assert_eq!(provider.lists().len(), 2 + stats.splits); - assert_eq!(provider.num_retired(), stats.splits); - } + // Retire the parent, moving every member it held onto that member's child. + let moves = members + .iter() + .zip(&children) + .map(|(&id, &child)| MoveTo::new(id, child_ids[child])) + .collect(); - #[tokio::test] - async fn inconsistent_provider_responses_fail_before_any_change() { - let fault = |set: fn(&mut Faults)| { - let mut faults = Faults::default(); - set(&mut faults); - faults - }; - let cases = [ - ( - fault(|f| f.skip_vector = true), - "no finite canonical vector", - ), - (fault(|f| f.nan_vector = true), "no finite canonical vector"), - (fault(|f| f.repeat_stage = true), "staging returned point"), - ( - fault(|f| f.drop_stage = true), - "staging returned 1 ids for 2 points", - ), - (fault(|f| f.reuse_lists = true), "reserved live list id"), - ]; - for (faults, expected) in cases { - // Inserting 0.5 and 8 overflows list 0, so every case reaches the split path. - let provider = - Provider::from_lists(1, &[(&[0.0], &[&[-1.0], &[1.0]]), (&[20.0], &[&[12.0]])]) - .with_faults(faults); - let before = provider.clone(); - let mut index = DynamicIvfIndex::new(provider, config(3, 1)).unwrap(); - let err = index - .insert_batch(&Strategy, &DefaultContext, &batch(100, &[[0.5], [8.0]])) - .await - .unwrap_err(); - assert!(err.to_string().contains(expected), "{faults:?}: {err}"); - assert!(index.provider().same_contents(&before), "{faults:?}"); + deltas.push(Delta::CentroidDelta { + id: parent, + delta: CentroidDelta::Retire { moves }, + }); + + // Append the staged points routed here onto their child. + for (child, &list) in child_ids.iter().enumerate() { + let appended: Box<[_]> = staged + .iter() + .zip(&children[members.len()..]) + .filter(|&(_, &group)| group == child) + .map(|(&position, _)| ids[position]) + .collect(); + if !appended.is_empty() { + deltas.push(Delta::PointAppends { + to: list, + ids: appended, + }); + } } + + Ok(deltas) } } + +#[cfg(test)] +mod tests {} diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index d138b4094a..228957c6a8 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -6,12 +6,8 @@ //! Incrementally maintained IVF index. pub mod dynamic; -mod grouped; mod index; mod online; pub mod update; -#[cfg(test)] -mod test; - -pub use index::{ConfigError, DynamicIvfConfig, DynamicIvfIndex, InsertStats}; +pub use index::{Config, ConfigError, IVFIndex, InsertStats}; diff --git a/diskann/src/ivf/online/batch.rs b/diskann/src/ivf/online/batch.rs deleted file mode 100644 index 10cf77900f..0000000000 --- a/diskann/src/ivf/online/batch.rs +++ /dev/null @@ -1,115 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! A staged and routed insert batch. - -use diskann_utils::views::Matrix; - -use super::{check_rows, first_repeat, index_error}; -use crate::{ - ANNResult, - error::ErrorExt, - ivf::{ - dynamic::{CentroidIndex, Provider, StageElements}, - grouped::Grouped, - update::{InsertionUpdate, PointDelta}, - }, - utils::VectorId, -}; - -/// Staged points, their canonical vectors, and their routes. -/// -/// Batch position `i` refers to `ids[i]`, row `i` of `vectors`, and `routes[i]`. -pub(in crate::ivf) struct StagedBatch { - pub(super) ids: Vec, - pub(super) vectors: Matrix, - /// The nearest live list of each point. - pub(super) routes: Vec, - /// Batch positions grouped by route. - pub(super) by_route: Grouped, -} - -impl StagedBatch { - /// Stage every point with the accessor and route it to its nearest live list. - /// - /// # Errors - /// - /// Fails if any accessor call fails or returns inconsistent data. - pub(in crate::ivf) async fn stage( - accessor: &mut A, - points: &[(P::ExternalId, T)], - ) -> ANNResult - where - P: Provider, - A: StageElements, - T: Sync, - { - let dim = accessor.dim(); - if dim == 0 { - return Err(index_error( - "canonical vectors must have a non-zero dimension", - )); - } - - let mut vectors = Matrix::new(f32::NAN, points.len(), dim); - let ids = accessor - .stage(points, vectors.as_mut_view()) - .await - .escalate("insert must stage every point")?; - if ids.len() != points.len() { - return Err(index_error(format!( - "staging returned {} ids for {} points", - ids.len(), - points.len() - ))); - } - if let Some(id) = first_repeat(ids.iter().copied()) { - return Err(index_error(format!( - "staging returned point {id} more than once" - ))); - } - check_rows(&ids, &vectors)?; - - let centroids = accessor.centroids(); - let mut routes = Vec::with_capacity(ids.len()); - for vector in vectors.row_iter() { - let list = centroids - .select(vector, 1) - .await - .escalate("insert must route every point")? - .first() - .map(|selected| selected.id) - .ok_or_else(|| index_error("centroid selection returned no list"))?; - routes.push(list); - } - - let by_route = Grouped::from_pairs(routes.iter().copied().zip(0..).collect()); - Ok(Self { - ids, - vectors, - routes, - by_route, - }) - } - - /// Append every point to its route. - pub(in crate::ivf) fn into_update(self) -> InsertionUpdate { - let appends = self - .ids - .iter() - .zip(&self.routes) - .map(|(&id, &to)| PointDelta::Append { id, to }) - .collect(); - InsertionUpdate::new(Vec::new(), appends) - } - - /// Lists that received points, ascending, with the number of points routed to - /// each. - pub(in crate::ivf) fn routed(&self) -> impl ExactSizeIterator { - self.by_route - .iter() - .map(|(list, positions)| (list, positions.len())) - } -} diff --git a/diskann/src/ivf/online/mod.rs b/diskann/src/ivf/online/mod.rs index 10f1a2a4f5..00f74c26e1 100644 --- a/diskann/src/ivf/online/mod.rs +++ b/diskann/src/ivf/online/mod.rs @@ -11,23 +11,14 @@ //! as it arrives, so the updates built from it need no further checks. Nothing here //! mutates the index; the caller applies the planned update. -mod batch; mod kernels; mod split; use std::fmt::{Debug, Display}; -pub(super) use batch::StagedBatch; -pub(super) use split::SplitPlan; +pub(super) use split::{TwoMeans, two_means}; -use diskann_utils::views::Matrix; - -use crate::{ - ANNError, ANNErrorKind, ANNResult, - error::ErrorExt, - ivf::dynamic::{CentroidIndex, MaintenanceAccessor, Provider}, - utils::VectorId, -}; +use crate::{ANNError, ANNErrorKind}; #[track_caller] pub(super) fn index_error(message: D) -> ANNError @@ -36,76 +27,3 @@ where { ANNError::message(ANNErrorKind::IndexError, message) } - -/// Check that the accessor reserved `expected` distinct list ids, none of them live. -#[track_caller] -pub(super) fn check_reserved( - centroids: &C, - ids: &[C::ListId], - expected: usize, -) -> ANNResult<()> { - if ids.len() != expected { - return Err(index_error(format!( - "accessor reserved {} list ids, expected {expected}", - ids.len() - ))); - } - if let Some(id) = first_repeat(ids.iter().copied()) { - return Err(index_error(format!("accessor reserved list id {id} twice"))); - } - if let Some(id) = ids.iter().find(|&&id| centroids.centroid(id).is_some()) { - return Err(index_error(format!("accessor reserved live list id {id}"))); - } - Ok(()) -} - -/// The first value that occurs more than once. -fn first_repeat(values: impl Iterator) -> Option { - let mut sorted: Vec = values.collect(); - sorted.sort_unstable(); - sorted - .windows(2) - .find(|pair| pair[0] == pair[1]) - .map(|pair| pair[0]) -} - -/// Read the canonical vector of every visible point in `ids`; row `i` belongs to -/// `ids[i]`. -/// -/// # Errors -/// -/// Fails if the accessor fails, or leaves a row unwritten or non-finite. -async fn read_rows( - accessor: &mut A, - ids: &[P::InternalId], - dim: usize, -) -> ANNResult> -where - P: Provider, - A: MaintenanceAccessor

, -{ - let mut rows = Matrix::new(f32::NAN, ids.len(), dim); - accessor - .read_vectors(ids, rows.as_mut_view()) - .await - .escalate("maintenance must read canonical vectors")?; - check_rows(ids, &rows)?; - Ok(rows) -} - -/// Check that every row the accessor was asked to write is finite. -/// -/// Callers fill rows with NaN before handing them to the accessor, so rows it never -/// writes fail this check too. -fn check_rows(ids: &[Id], rows: &Matrix) -> ANNResult<()> { - match ids - .iter() - .zip(rows.row_iter()) - .find(|(_, row)| row.iter().any(|x| !x.is_finite())) - { - Some((id, _)) => Err(index_error(format!( - "no finite canonical vector for point {id}" - ))), - None => Ok(()), - } -} diff --git a/diskann/src/ivf/online/split.rs b/diskann/src/ivf/online/split.rs index ef2fbe3f91..434665d362 100644 --- a/diskann/src/ivf/online/split.rs +++ b/diskann/src/ivf/online/split.rs @@ -3,351 +3,102 @@ * Licensed under the MIT license. */ -//! Split planning: gather each split's region, fit children, and place every affected -//! point. +//! Fitting the two centroids that replace a split list. -use diskann_utils::views::Matrix; +use diskann_utils::views::{Matrix, MatrixView}; use rand::{Rng, rngs::StdRng}; use super::{ - batch::StagedBatch, - check_reserved, index_error, + index_error, kernels::{LloydScratch, lloyd, nearest}, - read_rows, }; -use crate::{ - ANNResult, - error::ErrorExt, - ivf::{ - DynamicIvfConfig, - dynamic::{CentroidIndex, MaintenanceAccessor, Provider}, - grouped::Csr, - update::{CentroidDelta, InsertionUpdate, PointDelta}, - }, - utils::VectorId, -}; - -/// Everything one split batch needs, gathered before any computation. -/// -/// Lists are indexed densely: `lists[..parents]` are the split parents in ascending -/// order and `lists[parents..]` are their distinct neighbors in ascending order. -/// Region `r` is parent `lists[r]`, which children `2r` and `2r + 1` replace. -/// Neighbor `n` is `lists[parents + n]`. -pub(in crate::ivf) struct SplitPlan { - batch: StagedBatch, - lists: Vec, - parents: usize, - /// Each region's neighbors. - neighbors: Csr, - /// Each neighbor's regions, ascending. - regions: Csr, - /// Row `n` is the centroid of neighbor `n`. - neighbor_centroids: Matrix, - /// Existing members of every list, grouped like `lists`. - member_ids: Csr, - /// Row `i` is the canonical vector of `member_ids.values()[i]`. - member_vectors: Matrix, - /// Reserved ids for the children, two per region. - child_ids: Vec, +use crate::ANNResult; + +/// A list's two replacement centroids and the child each of its points joins. +#[derive(Debug)] +pub(in crate::ivf) struct TwoMeans { + /// Row `c` is the centroid of child `c`. + pub(in crate::ivf) centroids: Matrix, + /// `children[i]` is the child row `i` of the clustered points joins. + pub(in crate::ivf) children: Vec, } -impl SplitPlan { - /// Perform every accessor call the split needs. - /// - /// Reserves two child ids per parent, selects each parent's nearest surviving - /// lists as its neighbors, and reads the members and canonical vectors of every - /// parent and neighbor. `parents` must be ascending. - /// - /// # Errors - /// - /// Fails if an accessor call fails or the accessor's view is inconsistent. - pub(in crate::ivf) async fn gather( - accessor: &mut A, - config: &DynamicIvfConfig, - batch: StagedBatch, - parents: Vec, - ) -> ANNResult - where - P: Provider, - A: MaintenanceAccessor

, - { - let count = 2 * parents.len(); - let child_ids = accessor - .reserve_list_ids(count) - .await - .escalate("split must reserve child list ids")?; - check_reserved(accessor.centroids(), &child_ids, count)?; - - let dim = accessor.dim(); - let centroids = accessor.centroids(); - - // Every parent is retired, so no parent is ever a neighbor. - let mut by_region: Csr = Csr::default(); - for &parent in &parents { - let anchor = centroids - .centroid(parent) - .ok_or_else(|| index_error(format!("split parent {parent} is not live")))?; - let selected = centroids - .select(anchor, config.reassign_neighbors.saturating_add(1)) - .await - .escalate("split must find neighbor lists")?; - by_region.push( - selected - .iter() - .map(|selected| selected.id) - .filter(|id| parents.binary_search(id).is_err()) - .take(config.reassign_neighbors), - ); - } - - let mut distinct = by_region.values().to_vec(); - distinct.sort_unstable(); - distinct.dedup(); - - let mut neighbor_centroids = Matrix::new(0.0f32, distinct.len(), dim); - for (row, &list) in distinct.iter().enumerate() { - let centroid = centroids - .centroid(list) - .ok_or_else(|| index_error(format!("neighbor list {list} is not live")))?; - if centroid.len() != dim { - return Err(index_error(format!( - "centroid {list} has dimension {}, expected {dim}", - centroid.len() - ))); - } - neighbor_centroids.row_mut(row).copy_from_slice(centroid); - } - - let mut neighbors = Csr::default(); - for region in 0..by_region.num_groups() { - // Every neighbor is in `distinct`, so each search succeeds. - neighbors.push( - by_region - .group(region) - .iter() - .filter_map(|list| distinct.binary_search(list).ok()), - ); - } - let regions = neighbors.transpose(distinct.len()); - - let num_parents = parents.len(); - let mut lists = parents; - lists.extend_from_slice(&distinct); - - let mut member_ids = Csr::default(); - for &list in &lists { - let members = accessor - .read_members(list) - .await - .escalate("split must read list members")?; - member_ids.push(members.iter().copied()); - } - let member_vectors = read_rows(accessor, member_ids.values(), dim).await?; - - Ok(Self { - batch, - lists, - parents: num_parents, - neighbors, - regions, - neighbor_centroids, - member_ids, - member_vectors, - child_ids, - }) - } - - /// Fit the children, place every affected point, and build the update. - /// - /// Each parent splits into two children fitted by 2-means over its existing - /// members and the batch points routed to it. Those points then move to the - /// nearest of the parent's neighbors and children. A point of a neighbor list moves - /// to a child of a region containing that list only if the child is closer than - /// the list's centroid (as in SPFresh); otherwise it stays. - /// - /// # Errors - /// - /// Fails if a parent has fewer than two points. - pub(in crate::ivf) fn solve( - mut self, - iterations: usize, - rng: &mut StdRng, - ) -> ANNResult> { - let children = self.fit_children(iterations, rng)?; - let mut routes = std::mem::take(&mut self.batch.routes); - let destinations = self.place(&children, &mut routes)?; - Ok(self.package(&children, &routes, &destinations)) - } - - /// Fit two children per region; rows `2r` and `2r + 1` split region `r`. - fn fit_children(&self, iterations: usize, rng: &mut StdRng) -> ANNResult> { - let dim = self.batch.vectors.ncols(); - let mut centers = Matrix::new(0.0f32, 2 * self.parents, dim); - let mut scratch = LloydScratch::default(); - for region in 0..self.parents { - let span = self.member_ids.group(region).len() + self.staged(region).len(); - if span < 2 { - return Err(index_error(format!( - "split parent {} has fewer than two points", - self.lists[region] - ))); - } - let a = rng.random_range(0..span); - let b = (a + rng.random_range(1..span)) % span; - - let rows = &mut centers.as_mut_slice()[2 * region * dim..2 * (region + 1) * dim]; - let (first, second) = rows.split_at_mut(dim); - first.copy_from_slice(self.point(region, a)); - second.copy_from_slice(self.point(region, b)); - lloyd( - || self.points(region), - rows, - dim, - iterations.max(1), - &mut scratch, - ); - } - Ok(centers) - } - - /// Choose a list for every point of every region list. - /// - /// Row `r` of `children` is the centroid of child `r`. Returns the existing - /// members' lists, aligned with their positions, and overwrites the routes of the - /// batch points it places. - fn place(&self, children: &Matrix, routes: &mut [L]) -> ANNResult> { - let mut destinations = vec![L::default(); self.member_ids.values().len()]; - let child = |row: usize| (self.child_ids[row], children.row(row)); - let mut candidates: Vec<(L, &[f32])> = Vec::new(); - - for region in 0..self.parents { - candidates.clear(); - for &neighbor in self.neighbors.group(region) { - candidates.push(self.neighbor(neighbor)); - } - candidates.extend([child(2 * region), child(2 * region + 1)]); - self.assign(region, &candidates, &mut destinations, routes)?; - } - - // The list's own centroid comes first, so ties keep the point in place. - for neighbor in 0..self.regions.num_groups() { - candidates.clear(); - candidates.push(self.neighbor(neighbor)); - for ®ion in self.regions.group(neighbor) { - candidates.extend([child(2 * region), child(2 * region + 1)]); - } - self.assign( - self.parents + neighbor, - &candidates, - &mut destinations, - routes, - )?; - } - Ok(destinations) - } - - /// Place every point of `list` on its nearest candidate. - fn assign( - &self, - list: usize, - candidates: &[(L, &[f32])], - destinations: &mut [L], - routes: &mut [L], - ) -> ANNResult<()> { - let no_candidates = || index_error("split placement has no candidate lists"); - for (destination, vector) in destinations[self.member_ids.range(list)] - .iter_mut() - .zip(self.member_rows(list)) - { - *destination = nearest(vector, candidates).ok_or_else(no_candidates)?; - } - for &position in self.staged(list) { - routes[position] = - nearest(self.batch.vectors.row(position), candidates).ok_or_else(no_candidates)?; - } - Ok(()) +/// Cluster `points` into two groups, seeding Lloyd's algorithm with two distinct +/// points. +/// +/// # Errors +/// +/// Fails if `points` holds fewer than two rows. +pub(in crate::ivf) fn two_means( + points: MatrixView<'_, f32>, + iterations: usize, + rng: &mut StdRng, +) -> ANNResult { + let (count, dim) = (points.nrows(), points.ncols()); + if count < 2 { + return Err(index_error(format!( + "a split needs at least two points, got {count}" + ))); } - /// Assemble the update: install the children, retire the parents, append the - /// batch, and move every existing point whose list changed. - fn package( - &self, - children: &Matrix, - routes: &[L], - destinations: &[L], - ) -> InsertionUpdate { - let installs = self - .child_ids - .iter() - .zip(children.row_iter()) - .map(|(&id, centroid)| CentroidDelta::Install { - id, - centroid: centroid.into(), - }); - let retires = self.lists[..self.parents] - .iter() - .map(|&id| CentroidDelta::Retire { id }); - - let appends = self - .batch - .ids - .iter() - .zip(routes) - .map(|(&id, &to)| PointDelta::Append { id, to }); - // A parent is never a destination, so all its members move. - let moves = self.lists.iter().enumerate().flat_map(|(list, &from)| { - self.member_ids - .group(list) - .iter() - .zip(&destinations[self.member_ids.range(list)]) - .filter(move |&(_, &to)| to != from) - .map(move |(&id, &to)| PointDelta::Move { id, from, to }) - }); - - InsertionUpdate::new( - installs.chain(retires).collect(), - appends.chain(moves).collect(), - ) - } + let first = rng.random_range(0..count); + let second = (first + rng.random_range(1..count)) % count; + let mut centroids = Matrix::new(0.0f32, 2, dim); + centroids.row_mut(0).copy_from_slice(points.row(first)); + centroids.row_mut(1).copy_from_slice(points.row(second)); + + lloyd( + || points.row_iter(), + centroids.as_mut_slice(), + dim, + iterations.max(1), + &mut LloydScratch::default(), + ); + + let children = { + let candidates = [(0, centroids.row(0)), (1, centroids.row(1))]; + points + .row_iter() + // There are always two candidates, so every point joins one of them. + .map(|point| nearest(point, &candidates).unwrap_or(0)) + .collect() + }; + + Ok(TwoMeans { + centroids, + children, + }) +} - /// Neighbor `neighbor`'s list id and centroid. - fn neighbor(&self, neighbor: usize) -> (L, &[f32]) { - ( - self.lists[self.parents + neighbor], - self.neighbor_centroids.row(neighbor), - ) - } +#[cfg(test)] +mod tests { + use rand::SeedableRng; - /// Batch positions routed to `list`. - fn staged(&self, list: usize) -> &[usize] { - self.batch.by_route.get(self.lists[list]) - } + use super::*; - /// Canonical vectors of `list`'s existing members. - fn member_rows(&self, list: usize) -> impl Iterator { - self.member_ids - .range(list) - .map(|row| self.member_vectors.row(row)) + fn matrix(rows: &[[f32; D]]) -> Matrix { + let data: Box<[f32]> = rows.iter().flatten().copied().collect(); + Matrix::try_from(data, rows.len(), D).unwrap() } - /// Canonical vectors of `list`'s existing members followed by its batch points. - fn points(&self, list: usize) -> impl Iterator { - self.member_rows(list).chain( - self.staged(list) - .iter() - .map(|&position| self.batch.vectors.row(position)), - ) + #[test] + fn two_means_separates_distant_clusters() { + let points = matrix(&[[0.0, 0.0], [10.0, 0.0], [0.0, 1.0], [10.0, 1.0], [1.0, 0.0]]); + let fit = two_means(points.as_view(), 10, &mut StdRng::seed_from_u64(7)).unwrap(); + + assert_eq!(fit.centroids.nrows(), 2); + assert_eq!(fit.children.len(), 5); + // Rows 0, 2 and 4 are near the origin; rows 1 and 3 are near x = 10. + let low = fit.children[0]; + assert_eq!(fit.children, vec![low, 1 - low, low, 1 - low, low]); + assert_eq!(fit.centroids.row(low), &[1.0 / 3.0, 1.0 / 3.0]); + assert_eq!(fit.centroids.row(1 - low), &[10.0, 0.5]); } - /// The `index`-th vector of [`Self::points`]. - fn point(&self, list: usize, index: usize) -> &[f32] { - let members = self.member_ids.range(list); - if index < members.len() { - self.member_vectors.row(members.start + index) - } else { - self.batch - .vectors - .row(self.staged(list)[index - members.len()]) - } + #[test] + fn two_means_needs_two_points() { + let points = matrix(&[[1.0]]); + let err = two_means(points.as_view(), 1, &mut StdRng::seed_from_u64(7)).unwrap_err(); + assert!(err.to_string().contains("at least two points"), "{err}"); } } diff --git a/diskann/src/ivf/test/mod.rs b/diskann/src/ivf/test/mod.rs deleted file mode 100644 index 7792688185..0000000000 --- a/diskann/src/ivf/test/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Test infrastructure for the dynamic IVF index. - -pub(crate) mod provider; diff --git a/diskann/src/ivf/test/provider.rs b/diskann/src/ivf/test/provider.rs deleted file mode 100644 index fe505d774e..0000000000 --- a/diskann/src/ivf/test/provider.rs +++ /dev/null @@ -1,458 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! In-memory provider for exercising the dynamic IVF insert path. -//! -//! The provider stores plain owned state and mutates it in place through the -//! exclusive borrow lent to its maintenance accessor. It writes vector rows in -//! reverse order to exercise out-of-order reads, and [`Faults`] injects inconsistent -//! responses. - -use std::collections::{BTreeMap, BTreeSet}; - -use diskann_utils::{future::SendFuture, views::MutMatrixView}; - -use crate::{ - ANNError, ANNErrorKind, - ivf::{ - dynamic::{ - self, CentroidIndex, MaintenanceAccessor, MaintenanceStrategy, SelectedList, - StageElements, - }, - update::{CentroidDelta, InsertionUpdate, PointDelta}, - }, - provider::DefaultContext, -}; - -fn error(message: impl Into) -> ANNError { - ANNError::message(ANNErrorKind::IndexError, message.into()) -} - -fn squared(a: &[f32], b: &[f32]) -> f32 { - a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() -} - -/// Inconsistent responses the provider can inject. -#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] -pub(crate) struct Faults { - /// Leave the last row of every staged or read batch unwritten. - pub(crate) skip_vector: bool, - /// Write NaN into the last row of every staged or read batch. - pub(crate) nan_vector: bool, - /// Return the first staged id again for every later point of a `stage` call. - pub(crate) repeat_stage: bool, - /// Return one id fewer than the number of points a `stage` call stages. - pub(crate) drop_stage: bool, - /// Reserve list ids starting from zero, which may be live. - pub(crate) reuse_lists: bool, -} - -/// Write `rows[i]` into row `i` of `out` in reverse order, injecting vector faults -/// into the last row. -fn write_rows(faults: Faults, mut out: MutMatrixView<'_, f32>, rows: &[&[f32]]) { - for (index, row) in rows.iter().enumerate().rev() { - let last = index + 1 == rows.len(); - if last && faults.skip_vector { - continue; - } - let out = out.row_mut(index); - out.copy_from_slice(row); - if last && faults.nan_vector { - out[0] = f32::NAN; - } - } -} - -/// Exact centroid catalog. -#[derive(Debug, Clone, Default, PartialEq)] -pub(crate) struct Centroids(BTreeMap>); - -impl CentroidIndex for Centroids { - type ListId = u32; - type Error = ANNError; - - fn len(&self) -> usize { - self.0.len() - } - - fn centroid(&self, id: u32) -> Option<&[f32]> { - self.0.get(&id).map(|centroid| &**centroid) - } - - fn select( - &self, - query: &[f32], - nprobe: usize, - ) -> impl SendFuture>, ANNError>> { - let mut selected: Vec<_> = self - .0 - .iter() - .map(|(&id, centroid)| SelectedList { - id, - distance: squared(query, centroid), - }) - .collect(); - selected.sort_by(|a, b| a.distance.total_cmp(&b.distance).then(a.id.cmp(&b.id))); - selected.truncate(nprobe); - std::future::ready(Ok(selected)) - } -} - -/// An in-memory dynamic IVF provider. Point `i` has internal id `i`. -#[derive(Debug, Clone)] -pub(crate) struct Provider { - dim: usize, - /// External id and canonical vector of every committed point. - points: Vec<(u32, Box<[f32]>)>, - /// Current list of every committed point. - assignment: Vec, - lists: BTreeMap>, - centroids: Centroids, - retired: BTreeSet, - next_list: u32, - faults: Faults, -} - -impl Provider { - /// An empty provider for vectors of dimension `dim`. - pub(crate) fn new(dim: usize) -> Self { - Self { - dim, - points: Vec::new(), - assignment: Vec::new(), - lists: BTreeMap::new(), - centroids: Centroids::default(), - retired: BTreeSet::new(), - next_list: 0, - faults: Faults::default(), - } - } - - /// A provider whose list `i` has centroid `lists[i].0` and members `lists[i].1`, - /// which receive consecutive internal and external ids. - pub(crate) fn from_lists(dim: usize, lists: &[(&[f32], &[&[f32]])]) -> Self { - let mut provider = Self::new(dim); - for (list, &(centroid, members)) in (0..).zip(lists) { - provider.centroids.0.insert(list, centroid.into()); - let ids = members - .iter() - .map(|&vector| { - let id = provider.points.len() as u32; - provider.points.push((id, vector.into())); - provider.assignment.push(list); - id - }) - .collect(); - provider.lists.insert(list, ids); - provider.next_list = list + 1; - } - provider - } - - /// Inject `faults` into subsequent reads. - pub(crate) fn with_faults(mut self, faults: Faults) -> Self { - self.faults = faults; - self - } - - /// Whether both providers hold the same points, lists, and centroids. - pub(crate) fn same_contents(&self, other: &Self) -> bool { - self.points == other.points - && self.assignment == other.assignment - && self.lists == other.lists - && self.centroids == other.centroids - && self.retired == other.retired - } - - /// Live list ids, ascending. - pub(crate) fn lists(&self) -> Vec { - self.lists.keys().copied().collect() - } - - /// Number of retired lists. - pub(crate) fn num_retired(&self) -> usize { - self.retired.len() - } - - /// Number of committed points. - pub(crate) fn len(&self) -> usize { - self.points.len() - } - - /// Centroid of a live list. - pub(crate) fn centroid(&self, list: u32) -> Option<&[f32]> { - self.centroids.centroid(list) - } - - /// Vectors of a live list's members, sorted. - pub(crate) fn member_vectors(&self, list: u32) -> Vec> { - let mut vectors: Vec> = self.lists[&list] - .iter() - .map(|&id| self.points[id as usize].1.to_vec()) - .collect(); - vectors.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - vectors - } - - /// The list holding the point whose vector is `vector`. - pub(crate) fn list_of(&self, vector: &[f32]) -> Option { - self.points - .iter() - .position(|(_, point)| &**point == vector) - .map(|id| self.assignment[id]) - } - - /// Assert every structural invariant of the partition. - pub(crate) fn check(&self) { - assert_eq!(self.lists.len(), self.centroids.0.len()); - let mut seen = vec![0usize; self.points.len()]; - for (&list, members) in &self.lists { - assert!( - self.centroids.0.contains_key(&list), - "list {list} has no centroid" - ); - assert!(!self.retired.contains(&list), "list {list} is retired"); - for &id in members { - seen[id as usize] += 1; - assert_eq!(self.assignment[id as usize], list, "point {id}"); - } - } - assert!( - seen.iter().all(|&count| count == 1), - "every point must be in exactly one list" - ); - assert!(self.centroids.0.values().all(|c| c.len() == self.dim)); - assert!(self.points.iter().all(|(_, v)| v.len() == self.dim)); - } -} - -impl dynamic::Provider for Provider { - type Context = DefaultContext; - type InternalId = u32; - type ExternalId = u32; - type ListId = u32; - type Error = ANNError; - - fn to_internal_id(&self, _: &DefaultContext, gid: &u32) -> Result { - self.points - .iter() - .position(|(external, _)| external == gid) - .map(|id| id as u32) - .ok_or_else(|| error(format!("no point with external id {gid}"))) - } - - fn to_external_id(&self, _: &DefaultContext, id: u32) -> Result { - self.points - .get(id as usize) - .map(|(external, _)| *external) - .ok_or_else(|| error(format!("no point {id}"))) - } -} - -/// Maintenance accessor over an exclusively borrowed [`Provider`]. -pub(crate) struct Accessor<'a> { - provider: &'a mut Provider, - /// External id and vector of points staged by this accessor, in id order. - staged: Vec<(u32, Box<[f32]>)>, -} - -impl Accessor<'_> { - fn vector(&self, id: u32) -> Result<&[f32], ANNError> { - self.provider - .points - .get(id as usize) - .map(|(_, vector)| &**vector) - .ok_or_else(|| error(format!("no point {id}"))) - } - - fn commit_staged(&mut self) { - for point in self.staged.drain(..) { - self.provider.points.push(point); - self.provider.assignment.push(u32::MAX); - } - } - - fn install(&mut self, id: u32, centroid: &[f32]) -> Result<(), ANNError> { - if self.provider.retired.contains(&id) || self.provider.lists.contains_key(&id) { - return Err(error(format!("list id {id} is reused"))); - } - self.provider.centroids.0.insert(id, centroid.into()); - self.provider.lists.insert(id, Vec::new()); - Ok(()) - } - - /// Retire `list`, dropping its members wholesale. - fn retire(&mut self, list: u32) -> Result<(), ANNError> { - self.provider - .lists - .remove(&list) - .ok_or_else(|| error(format!("retired list {list} is not live")))?; - self.provider.centroids.0.remove(&list); - self.provider.retired.insert(list); - Ok(()) - } - - /// Remove `id` from `list`, which retirement may have dropped already. - fn leave(&mut self, id: u32, list: u32) -> Result<(), ANNError> { - if self.provider.retired.contains(&list) { - return Ok(()); - } - let members = self - .provider - .lists - .get_mut(&list) - .ok_or_else(|| error(format!("list {list} is not live")))?; - let position = members - .iter() - .position(|&member| member == id) - .ok_or_else(|| error(format!("point {id} is not in list {list}")))?; - members.remove(position); - Ok(()) - } - - fn place(&mut self, id: u32, list: u32) -> Result<(), ANNError> { - self.provider - .lists - .get_mut(&list) - .ok_or_else(|| error(format!("list {list} is not live")))? - .push(id); - let slot = self - .provider - .assignment - .get_mut(id as usize) - .ok_or_else(|| error(format!("no point {id}")))?; - *slot = list; - Ok(()) - } -} - -impl MaintenanceAccessor for Accessor<'_> { - type Centroids = Centroids; - type Error = ANNError; - - fn dim(&self) -> usize { - self.provider.dim - } - - fn centroids(&self) -> &Centroids { - &self.provider.centroids - } - - async fn list_size(&mut self, list: u32) -> Result { - let len = self - .provider - .lists - .get(&list) - .ok_or_else(|| error(format!("list {list} is not live")))? - .len(); - Ok(len) - } - - async fn read_members(&mut self, list: u32) -> Result<&[u32], ANNError> { - self.provider - .lists - .get(&list) - .map(Vec::as_slice) - .ok_or_else(|| error(format!("list {list} is not live"))) - } - - async fn read_vectors( - &mut self, - ids: &[u32], - out: MutMatrixView<'_, f32>, - ) -> Result<(), ANNError> { - let rows = ids - .iter() - .map(|&id| self.vector(id)) - .collect::, _>>()?; - write_rows(self.provider.faults, out, &rows); - Ok(()) - } - - async fn reserve_list_ids(&mut self, count: usize) -> Result, ANNError> { - if self.provider.faults.reuse_lists { - return Ok((0..count as u32).collect()); - } - let start = self.provider.next_list; - self.provider.next_list += count as u32; - Ok((start..self.provider.next_list).collect()) - } - - async fn apply(mut self, update: InsertionUpdate) -> Result<(), ANNError> { - self.commit_staged(); - for delta in update.centroids() { - match delta { - CentroidDelta::Install { id, centroid } => self.install(*id, centroid)?, - CentroidDelta::Retire { id } => self.retire(*id)?, - } - } - for &delta in update.points() { - match delta { - PointDelta::Append { id, to } => self.place(id, to)?, - PointDelta::Move { id, from, to } => { - self.leave(id, from)?; - self.place(id, to)?; - } - } - } - Ok(()) - } -} - -impl<'b> StageElements for Accessor<'_> { - async fn stage( - &mut self, - points: &[(u32, &'b [f32])], - out: MutMatrixView<'_, f32>, - ) -> Result, ANNError> { - let mut ids = Vec::with_capacity(points.len()); - for &(external, element) in points { - if element.len() != self.provider.dim { - return Err(error(format!("point {external} has the wrong dimension"))); - } - let known = |points: &[(u32, Box<[f32]>)]| { - points.iter().any(|&(existing, _)| existing == external) - }; - if known(&self.provider.points) || known(&self.staged) { - return Err(error(format!("external id {external} already exists"))); - } - let committed = self.provider.points.len(); - let internal = if self.provider.faults.repeat_stage && !self.staged.is_empty() { - committed - } else { - committed + self.staged.len() - }; - self.staged.push((external, element.into())); - ids.push(internal as u32); - } - - let rows: Vec<&[f32]> = points.iter().map(|&(_, element)| element).collect(); - write_rows(self.provider.faults, out, &rows); - if self.provider.faults.drop_stage { - ids.pop(); - } - Ok(ids) - } -} - -/// Strategy lending a [`Provider`] to one [`Accessor`]. -#[derive(Debug, Default)] -pub(crate) struct Strategy; - -impl<'a> MaintenanceStrategy<'a, Provider> for Strategy { - type MaintenanceAccessor = Accessor<'a>; - type Error = ANNError; - - fn maintenance_accessor( - &'a self, - provider: &'a mut Provider, - _context: &'a DefaultContext, - ) -> Result, ANNError> { - Ok(Accessor { - provider, - staged: Vec::new(), - }) - } -} diff --git a/diskann/src/ivf/update.rs b/diskann/src/ivf/update.rs index d2340be3d9..674bd5b1d7 100644 --- a/diskann/src/ivf/update.rs +++ b/diskann/src/ivf/update.rs @@ -3,54 +3,59 @@ * Licensed under the MIT license. */ -//! Declarative partition updates applied by [`MaintenanceAccessor::apply`]. -//! -//! An [`InsertionUpdate`] lists the primitive changes made by bootstrap or one insert -//! batch: lists installed or retired, as [`CentroidDelta`]s, and points appended or -//! moved, as [`PointDelta`]s. -//! -//! Only the index constructs updates. It validates everything it reads from the -//! provider or its caller, and its planner guarantees the invariants documented on -//! [`InsertionUpdate`] by construction. Providers can only read updates. -//! -//! [`MaintenanceAccessor::apply`]: crate::ivf::dynamic::MaintenanceAccessor::apply - use crate::utils::VectorId; /// A change to the set of live lists. #[derive(Debug, Clone, PartialEq)] -pub enum CentroidDelta { - /// Make the fresh list `id` live with the full-precision `centroid`. - Install { id: L, centroid: Box<[f32]> }, - /// Retire the live list `id`. Retired ids are never reused. - Retire { id: L }, +pub enum CentroidDelta { + /// Install a new centroid for this Id. + Install { centroid: Box<[f32]> }, + /// Retire the Id and move the points in this list to other list Ids. + Retire { moves: Box<[MoveTo]> }, } -/// A change to one point's list membership. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PointDelta { - /// The point `id`, staged by this accessor, joins list `to`. - Append { id: Id, to: L }, - /// The point `id` leaves list `from` and joins list `to`. - Move { id: Id, from: L, to: L }, +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MoveTo { + id: I, + to: L, } -impl PointDelta { - /// The point joining a list. - pub fn id(&self) -> Id { - match *self { - Self::Append { id, .. } | Self::Move { id, .. } => id, - } +impl MoveTo { + /// Move the point `id` into the list `to`. + pub(super) fn new(id: I, to: L) -> Self { + Self { id, to } + } + + /// The point that moves. + pub fn id(&self) -> I { + self.id } /// The list the point joins. pub fn to(&self) -> L { - match *self { - Self::Append { to, .. } | Self::Move { to, .. } => to, - } + self.to } } +#[derive(Debug, Clone, PartialEq)] +pub enum Delta { + // Centroid delta keyed by `id` + CentroidDelta { + id: L, + delta: CentroidDelta, + }, + // Point moves due to re-assignments, keyed by source centroid Id. + PointMoves { + from: L, + moves: Box<[MoveTo]>, + }, + // Point insertions from staged points, keyed by destination centroid Id. + PointAppends { + to: L, + ids: Box<[Id]>, + }, +} + /// The complete partition change made by bootstrap or one insert batch. /// /// The index guarantees, and providers may rely on without re-checking, that: @@ -61,69 +66,32 @@ impl PointDelta { /// - each list appears in at most one centroid delta, and each point in at most one /// point delta; /// - every point staged by this accessor is appended; -/// - a move's `from` holds the point and differs from `to`; +/// - a move's source holds the point and differs from its destination; /// - every destination is live after the update: installed, or live and not retired; +/// and /// - every member of a retired list is moved out, so a provider may drop a retired -/// list wholesale and only place the points moved out of it; and -/// - [`Self::points`] is sorted by destination, then id, so -/// `points().chunk_by(|a, b| a.to() == b.to())` yields each list's incoming points -/// in one run. +/// list wholesale and only place the points moved out of it. +/// +/// Deltas are in no particular order, so a provider that cannot apply them as one +/// atomic change should install centroids before placing points into them. #[derive(Debug, PartialEq)] -pub struct InsertionUpdate { - centroids: Vec>, - points: Vec>, +pub struct Deltas { + inner: Box<[Delta]>, } -impl InsertionUpdate { - /// Combine `centroids` and `points`, sorting `points` by destination, then id. - pub(super) fn new( - centroids: Vec>, - mut points: Vec>, - ) -> Self { - points.sort_unstable_by_key(|point| (point.to(), point.id())); - Self { centroids, points } - } - - /// The lists installed or retired. - pub fn centroids(&self) -> &[CentroidDelta] { - &self.centroids +impl Deltas { + /// Collect `deltas` into one partition change. + pub(super) fn new(deltas: Vec>) -> Self { + Self { + inner: deltas.into(), + } } - /// The points appended or moved, sorted by destination, then id. - pub fn points(&self) -> &[PointDelta] { - &self.points + /// The individual changes this update makes. + pub fn deltas(&self) -> &[Delta] { + &self.inner } } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn points_are_sorted_by_destination_then_id() { - let update = InsertionUpdate::::new( - vec![CentroidDelta::Retire { id: 4 }], - vec![ - PointDelta::Move { - id: 3, - from: 4, - to: 9, - }, - PointDelta::Append { id: 8, to: 7 }, - PointDelta::Append { id: 1, to: 9 }, - PointDelta::Move { - id: 2, - from: 4, - to: 7, - }, - ], - ); - let order: Vec<_> = update - .points() - .iter() - .map(|point| (point.to(), point.id())) - .collect(); - assert_eq!(order, vec![(7, 2), (7, 8), (9, 1), (9, 3)]); - assert_eq!(update.centroids(), &[CentroidDelta::Retire { id: 4 }]); - } -} +mod tests {} From f2a41db259407e4c204b7e0c575e3c0d2db5f249 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Tue, 6 Oct 2026 15:41:30 -0400 Subject: [PATCH 11/33] small edits --- diskann/src/ivf/dynamic.rs | 11 +++--- diskann/src/ivf/index.rs | 69 ++++++++++++++++++-------------------- diskann/src/ivf/update.rs | 10 +++--- 3 files changed, 42 insertions(+), 48 deletions(-) diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index ab834efb08..ce00cf0a91 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -80,9 +80,6 @@ pub trait Provider: Sized + Send + Sync + 'static { } pub trait StageElements { - /// The kind of error yielded by `set_element`. - type StageError: ToRanked + std::fmt::Debug + Send + Sync + 'static; - /// Stage a new point; its id is provisional until `update` commits it. fn stage_point( &mut self, @@ -144,7 +141,7 @@ pub trait Centroids: Send + Sync { pub trait InsertAccessor: Send + Sized + StageElements { /// In-memory centroid catalog and navigator in this accessor's unified view. - type Centroids<'a>: Centroids + type Centroids<'a>: Centroids where Self: 'a; @@ -154,7 +151,7 @@ pub trait InsertAccessor: Send + Sized + StageElements Self::Reader<'_>; @@ -256,12 +253,12 @@ pub trait SearchAccessor: Send + Sync { type InternalId; /// In-memory centroid view. - type Centroids<'a>: Centroids + type Centroids<'a>: Centroids where Self: 'a; /// Errors from list selection or scanning. - type Error: ToRanked + Debug + Send + Sync + 'static; + type Error: StandardError; /// Borrow the centroid index used for routing. fn centroids(&self) -> Self::Centroids<'_>; diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index ad6dce6c34..3dab7eff61 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -208,17 +208,15 @@ impl IVFIndex

{ for (position, ((id, v), out)) in points.iter().zip(vectors.rows_mut()).enumerate() { //stage the input points first. - let id = accessor - .stage_point(id, *v, out) - .await - .escalate("Unable to stage point")?; + let id = accessor.stage_point(id, *v, out).await.into_ann_result()?; ids.push(id); - // route the points to their nearest centroid + // route the points to their nearest centroid. Every staged point must be + // placed, so a routing failure cannot be skipped. let parent = accessor .centroids() .select(out, 1) - .escalate("Unable to select")? + .escalate("insert must route every staged point")? .first() .map(|s| s.id) .ok_or_else(|| ANNError::message("Didn't return list"))?; @@ -229,13 +227,12 @@ impl IVFIndex

{ // filter the routed lists for ones that need splitting. let mut parents = Vec::new(); for (&list, positions) in &routed { - let len = accessor.list_size(list).escalate("Unable to get len")?; + let len = accessor.list_size(list).into_ann_result()?; if len + positions.len() > config.split_threshold { parents.push(list); } } - parents.sort_unstable(); // for the ones that need splitting - split and re-assign necessary points in neighbors let update = Self::split::<_, T>( @@ -281,6 +278,16 @@ impl IVFIndex

{ A: InsertAccessor, T: Sync, { + // Kidns of deltas - + // 1. Staged points that are not part of splitting. -> [`Delta::PointAppends`] -> available before split. + // 2. Staged points that are part of splitting. -> [`Delta::PointAppends`] -> needs new centroid, centroid id and assignment. + // 3. New centroids -> [`Delta::CentroidDelta::Install`] -> needs clustering resulting from spolit. + // 4. Retire centroids -> [`Delta::CentroidDelta::Retire`] -> needs the assignments to new centroids from old one, result from clustering + // 5. (Optional) Reassignments + + // Split - + // - Get: staged vectors, + let mut deltas = Vec::new(); // Lists that keep their centroid take their staged points as they are. @@ -328,35 +335,25 @@ impl IVFIndex

{ A: InsertAccessor, T: Sync, { - // let reader = accessor.reader(); + let dim = accessor.dim(); - // let num_members = accessor.list_size(parent)?; + let num_members = accessor.list_size(parent).into_ann_result()?; - // let mut points = Matrix::new(f32::NAN, num_members + staged, reader.dim()); + let mut points = rowmajor::Owned::from_element(num_members + staged.len(), dim, f32::NAN); - // reader.read_into(parent, points.as_mut_view()).await?; - let dim = accessor.dim(); - let members = accessor - .get_members(parent) - .escalate("split must read the parent's members")? - .to_vec(); - - // Everything the parent would hold: its current members, then its staged points. - let mut points = rowmajor::Owned::from_element(members.len() + staged.len(), dim, f32::NAN); - { - let mut member_vectors = rowmajor::Owned::from_element(members.len(), dim, f32::NAN); - accessor - .reader() - .read_into(parent, member_vectors.as_view_mut()) - .await - .escalate("split must read the parent's vectors")?; + let (head, tail) = points.as_mut_slice().split_at_mut(num_members * dim); - let sources = member_vectors - .rows() - .chain(staged.iter().map(|&position| vectors.row(position))); - for (row, source) in points.rows_mut().zip(sources) { - row.copy_from_slice(source); - } + let head_view = rowmajor::Mut::try_from_data(head, num_members, dim)?; + + accessor + .reader() + .read_into(parent, head_view) + .await + .escalate("split must read the parent's vectors")?; + + let mut tail_view = rowmajor::Mut::try_from_data(tail, staged.len(), dim)?; + for (row, &position) in tail_view.rows_mut().zip(staged) { + row.copy_from_slice(vectors.row(position)); } let TwoMeans { @@ -374,12 +371,12 @@ impl IVFIndex

{ deltas.push(Delta::CentroidDelta { id, - delta: CentroidDelta::Install { - centroid: centroid.into(), - }, + delta: CentroidDelta::Install, }); } + let members = accessor.get_members(parent).into_ann_result()?; + // Retire the parent, moving every member it held onto that member's child. let moves = members .iter() diff --git a/diskann/src/ivf/update.rs b/diskann/src/ivf/update.rs index 674bd5b1d7..f7ab7c74f1 100644 --- a/diskann/src/ivf/update.rs +++ b/diskann/src/ivf/update.rs @@ -8,13 +8,13 @@ use crate::utils::VectorId; /// A change to the set of live lists. #[derive(Debug, Clone, PartialEq)] pub enum CentroidDelta { - /// Install a new centroid for this Id. - Install { centroid: Box<[f32]> }, + /// Install a the staged centroid with this Id. + Install, /// Retire the Id and move the points in this list to other list Ids. Retire { moves: Box<[MoveTo]> }, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Copy, Eq)] pub struct MoveTo { id: I, to: L, @@ -44,12 +44,12 @@ pub enum Delta { id: L, delta: CentroidDelta, }, - // Point moves due to re-assignments, keyed by source centroid Id. + // Point moves due to re-assignments, grouped by source centroid Id. PointMoves { from: L, moves: Box<[MoveTo]>, }, - // Point insertions from staged points, keyed by destination centroid Id. + // Point insertions from staged points, grouped by destination centroid Id. PointAppends { to: L, ids: Box<[Id]>, From de717fcc17d94f8eec2cbc31453130d16cec512c Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Thu, 8 Oct 2026 09:23:06 -0400 Subject: [PATCH 12/33] ok looks like the split framework is coming together --- diskann/src/ivf/dynamic.rs | 27 +- diskann/src/ivf/index.rs | 281 ++----- diskann/src/ivf/mod.rs | 1 + diskann/src/ivf/online/clustering.rs | 226 ++++++ diskann/src/ivf/online/kernels.rs | 2 +- diskann/src/ivf/online/mod.rs | 13 +- diskann/src/ivf/online/split.rs | 104 --- diskann/src/ivf/split.rs | 1119 ++++++++++++++++++++++++++ 8 files changed, 1420 insertions(+), 353 deletions(-) create mode 100644 diskann/src/ivf/online/clustering.rs delete mode 100644 diskann/src/ivf/online/split.rs create mode 100644 diskann/src/ivf/split.rs diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/dynamic.rs index ce00cf0a91..dba1c4f59d 100644 --- a/diskann/src/ivf/dynamic.rs +++ b/diskann/src/ivf/dynamic.rs @@ -18,7 +18,7 @@ use crate::{ ANNResult, error::{StandardError, ToRanked}, ivf::update::Deltas, - provider::{ExecutionContext, Guard, NoopGuard}, + provider::ExecutionContext, utils::VectorId, }; @@ -79,7 +79,7 @@ pub trait Provider: Sized + Send + Sync + 'static { ) -> Result; } -pub trait StageElements { +pub trait Stage { /// Stage a new point; its id is provisional until `update` commits it. fn stage_point( &mut self, @@ -104,7 +104,7 @@ pub trait StageElements { /// result. Exact and graph implementations expose the same interface. pub trait Centroids: Send + Sync { /// Stable logical centroid id, also used as the inverted-list id. - type ListId: VectorId; + type Id: VectorId; /// Errors encountered while reading or navigating the centroid index. type Error: ToRanked + Debug + Send + Sync + 'static; @@ -121,7 +121,7 @@ pub trait Centroids: Send + Sync { } /// Borrow one authoritative full-precision centroid vector. - fn centroid(&self, id: Self::ListId) -> Option>; + fn centroid(&self, id: Self::Id) -> Option>; /// Select exactly `min(nprobe, self.len())` live centroids for `query`, nearest /// first. @@ -132,16 +132,16 @@ pub trait Centroids: Send + Sync { &self, query: &[f32], nprobe: usize, - ) -> Result>, Self::Error>; + ) -> Result>, Self::Error>; } ////////////// // Accessor // ////////////// -pub trait InsertAccessor: Send + Sized + StageElements { +pub trait InsertAccessor: Send + Sized + Stage { /// In-memory centroid catalog and navigator in this accessor's unified view. - type Centroids<'a>: Centroids + type Centroids<'a>: Centroids where Self: 'a; @@ -164,7 +164,7 @@ pub trait InsertAccessor: Send + Sized + StageElements Result; - /// Read the member ids of a live list. + /// Read the member ids of a live list, in the reader's vector-row order. fn get_members(&self, list: P::ListId) -> Result<&[P::InternalId], Self::Error>; /// Consume the accessor and write the update @@ -181,10 +181,11 @@ pub trait Reader: Send + Sync { /// dimension of the output of each Id. fn dim(&self) -> usize; - /// Write the canonical vector of `ids[i]` into row `i` of `out`, in any order. + /// Read every canonical vector for `id`, preserving its member order in `out`. /// - /// `out` has one row per id and [`Self::dim`] columns, and every row must be - /// written. + /// For an insert accessor's list reader, row `i` corresponds to member `i` + /// returned by [`InsertAccessor::get_members`] in the same accessor view. + /// `out` has one row per member and [`Self::dim`] columns; every row must be written. fn read_into( &self, id: Self::Id, @@ -201,7 +202,7 @@ pub trait Reader: Send + Sync { /// The provider is borrowed exclusively for the lifetime of the accessor, so no /// search or other mutation can observe it until the accessor is applied or /// dropped. Providers may therefore mutate plain in-memory state in -/// [`InsertAccessor::apply`] without interior synchronization. Providers that +/// [`InsertAccessor::update`] without interior synchronization. Providers that /// share state outside this borrow (for example through an `Arc`) remain responsible /// for coordinating those aliases. pub trait MaintenanceStrategy<'a, P: Provider, T: Sync>: Send + Sync { @@ -253,7 +254,7 @@ pub trait SearchAccessor: Send + Sync { type InternalId; /// In-memory centroid view. - type Centroids<'a>: Centroids + type Centroids<'a>: Centroids where Self: 'a; diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 3dab7eff61..37f29ff231 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -5,9 +5,9 @@ //! Dynamic IVF index orchestration. //! -//! The index drives one [`InsertAccessor`] per mutation: it plans a [`Deltas`] against -//! the accessor's view and hands it to [`InsertAccessor::update`]. It never touches -//! point, centroid, or list storage directly. +//! The index drives one [`InsertAccessor`] per mutation: it plans a +//! [`Deltas`](super::update::Deltas) against the accessor's view and hands it to +//! [`InsertAccessor::update`]. It never touches point, centroid, or list storage directly. use diskann_utils::{ future::SendFuture, @@ -20,11 +20,9 @@ use crate::{ ANNError, ANNResult, error::{ErrorExt, IntoANNResult}, ivf::{ - dynamic::{ - Centroids, InsertAccessor, MaintenanceStrategy, Provider, Reader, StageElements, - }, - online::{TwoMeans, index_error, two_means}, - update::{CentroidDelta, Delta, Deltas, MoveTo}, + dynamic::{Centroids, InsertAccessor, MaintenanceStrategy, Provider, Stage}, + split::plan_split_update, + update::{self, Deltas}, }, }; @@ -104,79 +102,64 @@ impl IVFIndex

{ &self.provider } - // /// Install the initial centroids. - // /// - // /// # Errors - // /// - // /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index - // /// already has centroids, or the maintenance accessor fails. - // pub fn initialize<'a, S>( - // &'a mut self, - // strategy: &'a S, - // context: &'a P::Context, - // centroids: Matrix, - // ) -> impl SendFuture> - // where - // S: MaintenanceStrategy<'a, P>, - // { - // let provider = &mut self.provider; + /// Install the initial centroids. + /// + /// # Errors + /// + /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index + /// already has centroids, or the maintenance accessor fails. + pub fn initialize<'a, S, T>( + &'a mut self, + strategy: &'a S, + context: &'a P::Context, + centroids: rowmajor::Ref<'_, f32>, + ) -> impl SendFuture> + where + S: MaintenanceStrategy<'a, P, T>, + T: Sync, + { + let provider = &mut self.provider; - // async move { - // let count = centroids.nrows(); - // if count == 0 { - // return Err(index_error("initialize requires at least one centroid")); - // } - // if centroids.as_slice().iter().any(|x| !x.is_finite()) { - // return Err(index_error("initial centroids must be finite")); - // } + async move { + let count = centroids.nrows(); + if count == 0 { + return Err(ANNError::message("Cannot initialize with zero centroids.")); + } - // let mut accessor = strategy - // .maintenance_accessor(provider, context) - // .into_ann_result()?; + let mut accessor = strategy + .maintenance_accessor(provider, context) + .into_ann_result()?; - // if !accessor.centroids().is_empty() { - // return Err(index_error("dynamic IVF index is already initialized")); - // } - // let dim = accessor.dim(); - // if dim == 0 || centroids.ncols() != dim { - // return Err(index_error(format!( - // "initial centroids have dimension {}, expected {dim}", - // centroids.ncols() - // ))); - // } + let mut installs = Vec::with_capacity(centroids.nrows()); - // let list_ids = accessor - // .reserve_list_ids(count) - // .await - // .escalate("initialize must reserve list ids")?; - // check_reserved(&accessor.centroids(), &list_ids, count)?; + for c in centroids.rows() { + let id = accessor.stage_centroid(c).await?; + installs.push(update::Delta::CentroidDelta { + id, + delta: update::CentroidDelta::Install, + }); + } - // let installs = list_ids - // .into_iter() - // .zip(centroids.row_iter()) - // .map(|(id, centroid)| CentroidDelta::Install { - // id, - // centroid: centroid.into(), - // }) - // .collect(); + accessor + .update(Deltas::::new(installs)) + .await + .escalate("initialize must apply the initial centroids")?; - // accessor - // .apply(InsertionUpdate::new(installs, Vec::new())) - // .await - // .escalate("initialize must apply the initial centroids") - // } - // } + Ok(()) + } + } /// Insert a batch, splitting every list the batch pushes past `split_threshold`. /// - /// Routing and split planning run against the accessor's view before any change. - /// The batch, and any splits it triggers, are then applied as one - /// [`InsertionUpdate`]. + /// All split parents form one region. Their points are assigned across all fitted + /// children, never back to retiring parents. Other staged points keep their routes. + /// The batch and its splits are then applied as one + /// [`Deltas`](super::update::Deltas). /// /// # Errors /// /// Fails if the index is not initialized, the provider returns inconsistent - /// data, or any accessor call fails. State after an `apply` failure is + /// data, or any accessor call fails. State after an `update` failure is /// provider-defined. pub fn insert_batch<'a, S, T>( &'a mut self, @@ -211,8 +194,7 @@ impl IVFIndex

{ let id = accessor.stage_point(id, *v, out).await.into_ann_result()?; ids.push(id); - // route the points to their nearest centroid. Every staged point must be - // placed, so a routing failure cannot be skipped. + // route the points to their nearest centroid. let parent = accessor .centroids() .select(out, 1) @@ -234,10 +216,10 @@ impl IVFIndex

{ } } - // for the ones that need splitting - split and re-assign necessary points in neighbors - let update = Self::split::<_, T>( + // Fit every split's children before assigning points and building the update. + let update = plan_split_update::( &mut accessor, - &config, + config.two_means_iterations, rng, &ids, vectors.as_view(), @@ -254,159 +236,6 @@ impl IVFIndex

{ Ok(InsertStats::default()) } } - - /// Build the update for one insert batch: split every list in `parents` and append - /// every staged point to the list that ends up holding it. - /// - /// Batch position `i` is the staged point `ids[i]`, whose full-precision vector is - /// row `i` of `vectors`. `routed` maps every routed list to the batch positions - /// routed to it, and `parents` names the subset of those lists that overflow. - /// - /// # Errors - /// - /// Fails if an accessor call fails or a split list holds fewer than two points. - async fn split( - accessor: &mut A, - config: &Config, - rng: &mut StdRng, - ids: &[P::InternalId], - vectors: rowmajor::Ref<'_, f32>, - routed: &hashbrown::HashMap>, - parents: &[P::ListId], - ) -> ANNResult> - where - A: InsertAccessor, - T: Sync, - { - // Kidns of deltas - - // 1. Staged points that are not part of splitting. -> [`Delta::PointAppends`] -> available before split. - // 2. Staged points that are part of splitting. -> [`Delta::PointAppends`] -> needs new centroid, centroid id and assignment. - // 3. New centroids -> [`Delta::CentroidDelta::Install`] -> needs clustering resulting from spolit. - // 4. Retire centroids -> [`Delta::CentroidDelta::Retire`] -> needs the assignments to new centroids from old one, result from clustering - // 5. (Optional) Reassignments - - // Split - - // - Get: staged vectors, - - let mut deltas = Vec::new(); - - // Lists that keep their centroid take their staged points as they are. - for (&list, positions) in routed { - if parents.contains(&list) { - continue; - } - deltas.push(Delta::PointAppends { - to: list, - ids: positions.iter().map(|&position| ids[position]).collect(), - }); - } - - // Overflowing lists are replaced by two children, which take their staged - // points instead. - for &parent in parents { - let staged: &[usize] = routed.get(&parent).map_or(&[], Vec::as_slice); - deltas.extend( - Self::split_list::<_, T>(accessor, config, rng, ids, vectors, parent, staged) - .await?, - ); - } - - Ok(Deltas::new(deltas)) - } - - /// Plan one list's split: fit two children over every point the list would hold, - /// then place each of those points on the nearer child and retire the parent. - /// - /// `staged` holds the batch positions routed to `parent`. - /// - /// # Errors - /// - /// Fails if an accessor call fails or the list holds fewer than two points. - async fn split_list( - accessor: &mut A, - config: &Config, - rng: &mut StdRng, - ids: &[P::InternalId], - vectors: rowmajor::Ref<'_, f32>, - parent: P::ListId, - staged: &[usize], - ) -> ANNResult>> - where - A: InsertAccessor, - T: Sync, - { - let dim = accessor.dim(); - - let num_members = accessor.list_size(parent).into_ann_result()?; - - let mut points = rowmajor::Owned::from_element(num_members + staged.len(), dim, f32::NAN); - - let (head, tail) = points.as_mut_slice().split_at_mut(num_members * dim); - - let head_view = rowmajor::Mut::try_from_data(head, num_members, dim)?; - - accessor - .reader() - .read_into(parent, head_view) - .await - .escalate("split must read the parent's vectors")?; - - let mut tail_view = rowmajor::Mut::try_from_data(tail, staged.len(), dim)?; - for (row, &position) in tail_view.rows_mut().zip(staged) { - row.copy_from_slice(vectors.row(position)); - } - - let TwoMeans { - centroids, - children, - } = two_means(points.as_view(), config.two_means_iterations, rng)?; - - let mut deltas = Vec::new(); - let mut child_ids = Vec::new(); - - // Install both children before anything is placed into them. - for centroid in centroids.rows() { - let id = accessor.stage_centroid(centroid).await?; - child_ids.push(id); - - deltas.push(Delta::CentroidDelta { - id, - delta: CentroidDelta::Install, - }); - } - - let members = accessor.get_members(parent).into_ann_result()?; - - // Retire the parent, moving every member it held onto that member's child. - let moves = members - .iter() - .zip(&children) - .map(|(&id, &child)| MoveTo::new(id, child_ids[child])) - .collect(); - - deltas.push(Delta::CentroidDelta { - id: parent, - delta: CentroidDelta::Retire { moves }, - }); - - // Append the staged points routed here onto their child. - for (child, &list) in child_ids.iter().enumerate() { - let appended: Box<[_]> = staged - .iter() - .zip(&children[members.len()..]) - .filter(|&(_, &group)| group == child) - .map(|(&position, _)| ids[position]) - .collect(); - if !appended.is_empty() { - deltas.push(Delta::PointAppends { - to: list, - ids: appended, - }); - } - } - - Ok(deltas) - } } #[cfg(test)] diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 228957c6a8..7a271700f7 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -8,6 +8,7 @@ pub mod dynamic; mod index; mod online; +mod split; pub mod update; pub use index::{Config, ConfigError, IVFIndex, InsertStats}; diff --git a/diskann/src/ivf/online/clustering.rs b/diskann/src/ivf/online/clustering.rs new file mode 100644 index 0000000000..4b1269a2dd --- /dev/null +++ b/diskann/src/ivf/online/clustering.rs @@ -0,0 +1,226 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Numerical fitting and assignment over borrowed matrices. + +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; +use rand::{Rng, rngs::StdRng}; + +use super::{ + index_error, + kernels::{LloydScratch, lloyd, nearest}, +}; +use crate::ANNResult; + +/// Fit two children directly into `centroids`, using distinct random seeds. +/// +/// `centroids` must have two rows and the same dimension as `points`. +/// Zero iterations still performs one Lloyd step. +/// +/// # Errors +/// +/// Returns an error when fewer than two training points are supplied. +pub(in crate::ivf) fn fit_two_means( + points: rowmajor::Ref<'_, f32>, + mut centroids: rowmajor::Mut<'_, f32>, + iterations: usize, + rng: &mut StdRng, + scratch: &mut LloydScratch, +) -> ANNResult<()> { + let count = points.nrows(); + if count < 2 { + return Err(index_error(format!( + "a split needs at least two points, got {count}" + ))); + } + + let first = rng.random_range(0..count); + let second = (first + rng.random_range(1..count)) % count; + centroids.row_mut(0).copy_from_slice(points.row(first)); + centroids.row_mut(1).copy_from_slice(points.row(second)); + lloyd( + || points.rows(), + centroids.as_mut_slice(), + points.ncols(), + iterations.max(1), + scratch, + ); + Ok(()) +} + +/// Assign by squared Euclidean distance, preferring the earliest centroid on ties. +/// +/// # Errors +/// +/// Returns an error when there are no centroids or the dimensions differ. +pub(in crate::ivf) fn assign_nearest( + points: rowmajor::Ref<'_, f32>, + centroids: rowmajor::Ref<'_, f32>, +) -> ANNResult> { + if centroids.nrows() == 0 { + return Err(index_error("assignment needs at least one centroid")); + } + if points.ncols() != centroids.ncols() { + return Err(index_error(format!( + "assignment points have dimension {}, but centroids have dimension {}", + points.ncols(), + centroids.ncols() + ))); + } + + let candidates: Vec<_> = centroids.rows().enumerate().collect(); + let mut assigned = Vec::with_capacity(points.nrows()); + for point in points.rows() { + assigned.push( + nearest(point, &candidates) + .ok_or_else(|| index_error("assignment needs at least one centroid"))?, + ); + } + Ok(assigned.into_boxed_slice()) +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng; + + use 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() + } + + fn fit( + points: rowmajor::Ref<'_, f32>, + iterations: usize, + rng: &mut StdRng, + ) -> ANNResult> { + let mut centroids = rowmajor::Owned::from_element(2, points.ncols(), 0.0); + fit_two_means( + points, + centroids.as_view_mut(), + iterations, + rng, + &mut LloydScratch::default(), + )?; + Ok(centroids) + } + + #[test] + fn two_means_separates_distant_clusters() { + let points = matrix(&[[0.0, 0.0], [10.0, 0.0], [0.0, 1.0], [10.0, 1.0], [1.0, 0.0]]); + let centroids = fit(points.as_view(), 10, &mut StdRng::seed_from_u64(7)).unwrap(); + let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); + + let low = children[0]; + assert_eq!(children.as_ref(), &[low, 1 - low, low, 1 - low, low]); + assert_eq!(centroids.row(low), &[1.0 / 3.0, 1.0 / 3.0]); + assert_eq!(centroids.row(1 - low), &[10.0, 0.5]); + } + + #[test] + fn two_means_needs_two_points() { + for count in [0, 1] { + let points = rowmajor::Owned::from_element(count, 1, 1.0f32); + let err = fit(points.as_view(), 1, &mut StdRng::seed_from_u64(7)).unwrap_err(); + assert!(err.to_string().contains("at least two points"), "{err}"); + } + } + + #[test] + fn two_means_preserves_seed_order_and_rng_consumption() { + let points = matrix(&[[0.0], [10.0]]); + let mut rng = StdRng::seed_from_u64(7); + let mut expected_rng = rng.clone(); + let first = expected_rng.random_range(0..points.nrows()); + let second = (first + expected_rng.random_range(1..points.nrows())) % points.nrows(); + let centroids = fit(points.as_view(), 10, &mut rng).unwrap(); + + assert_eq!(centroids.row(0), points.row(first)); + assert_eq!(centroids.row(1), points.row(second)); + assert_eq!(rng.random::(), expected_rng.random::()); + } + + #[test] + fn two_means_zero_iterations_still_refines_centroids() { + let points = matrix(&[[0.0], [2.0], [10.0], [12.0]]); + let zero = fit(points.as_view(), 0, &mut StdRng::seed_from_u64(7)).unwrap(); + let one = fit(points.as_view(), 1, &mut StdRng::seed_from_u64(7)).unwrap(); + assert_eq!(zero.as_slice(), one.as_slice()); + } + + #[test] + fn two_means_writes_only_the_output_subview_and_reuses_scratch() { + let points = matrix(&[[0.0], [10.0]]); + let mut output = rowmajor::Owned::from_element(6, 1, -1.0); + let mut scratch = LloydScratch::default(); + let mut rng = StdRng::seed_from_u64(7); + for rows in [1..3, 3..5] { + let out = rowmajor::Mut::try_from_data(&mut output.as_mut_slice()[rows], 2, 1).unwrap(); + fit_two_means(points.as_view(), out, 10, &mut rng, &mut scratch).unwrap(); + } + assert_eq!(output.row(0), &[-1.0]); + assert_eq!(output.row(5), &[-1.0]); + for rows in [1..3, 3..5] { + let assigned = assign_nearest(points.as_view(), output.subview(rows).unwrap()).unwrap(); + assert_ne!(assigned[0], assigned[1]); + } + } + + #[test] + fn nearest_centroid_assigns_points_to_an_independent_candidate_set() { + let points = matrix(&[[-1.0], [12.0], [19.0]]); + let centroids = matrix(&[[0.0], [10.0], [20.0]]); + let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); + assert_eq!(children.as_ref(), &[0, 1, 2]); + } + + #[test] + fn nearest_centroid_prefers_the_earliest_row_on_ties() { + let points = matrix(&[[1.8, 0.1], [1.0, 1.0]]); + let centroids = matrix(&[[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]]); + let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); + assert_eq!(children.as_ref(), &[1, 0]); + } + + #[test] + fn nearest_centroid_supports_one_candidate() { + let points = matrix(&[[-1.0], [12.0], [19.0]]); + let centroids = matrix(&[[10.0]]); + let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); + assert_eq!(children.as_ref(), &[0, 0, 0]); + } + + #[test] + fn nearest_centroid_supports_an_empty_point_batch() { + let points = matrix::<1>(&[]); + let centroids = matrix(&[[10.0]]); + assert!( + assign_nearest(points.as_view(), centroids.as_view()) + .unwrap() + .is_empty() + ); + } + + #[test] + fn nearest_centroid_needs_a_candidate() { + let points = matrix(&[[1.0]]); + let centroids = matrix::<1>(&[]); + let err = assign_nearest(points.as_view(), centroids.as_view()).unwrap_err(); + assert!(err.to_string().contains("at least one centroid"), "{err}"); + } + + #[test] + fn nearest_centroid_requires_matching_dimensions() { + let points = matrix(&[[1.0]]); + let centroids = matrix(&[[0.0, 0.0]]); + let err = assign_nearest(points.as_view(), centroids.as_view()).unwrap_err(); + assert!( + err.to_string() + .contains("points have dimension 1, but centroids have dimension 2"), + "{err}" + ); + } +} diff --git a/diskann/src/ivf/online/kernels.rs b/diskann/src/ivf/online/kernels.rs index 103f5e7aad..54809f6aeb 100644 --- a/diskann/src/ivf/online/kernels.rs +++ b/diskann/src/ivf/online/kernels.rs @@ -26,7 +26,7 @@ pub(super) fn nearest(point: &[f32], candidates: &[(L, &[f32])]) -> Opt /// Scratch buffers for [`lloyd`], reusable across calls. #[derive(Debug, Default)] -pub(super) struct LloydScratch { +pub(in crate::ivf) struct LloydScratch { sums: Vec, counts: Vec, } diff --git a/diskann/src/ivf/online/mod.rs b/diskann/src/ivf/online/mod.rs index 962a665dfd..fb30c527b8 100644 --- a/diskann/src/ivf/online/mod.rs +++ b/diskann/src/ivf/online/mod.rs @@ -3,20 +3,15 @@ * Licensed under the MIT license. */ -//! The online insert algorithm: stage and route a batch, then split every list the -//! batch overflows. -//! -//! Planning performs all of its accessor reads first and then computes the update -//! from owned, densely indexed data. Everything read from the accessor is validated -//! as it arrives, so the updates built from it need no further checks. Nothing here -//! mutates the index; the caller applies the planned update. +//! Borrowed-matrix fitting and assignment helpers and squared-Euclidean kernels. +mod clustering; mod kernels; -mod split; use std::fmt::{Debug, Display}; -pub(super) use split::{TwoMeans, two_means}; +pub(super) use clustering::{assign_nearest, fit_two_means}; +pub(super) use kernels::LloydScratch; use crate::ANNError; diff --git a/diskann/src/ivf/online/split.rs b/diskann/src/ivf/online/split.rs deleted file mode 100644 index ba1f2f0bc4..0000000000 --- a/diskann/src/ivf/online/split.rs +++ /dev/null @@ -1,104 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Fitting the two centroids that replace a split list. - -use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; -use rand::{Rng, rngs::StdRng}; - -use super::{ - index_error, - kernels::{LloydScratch, lloyd, nearest}, -}; -use crate::ANNResult; - -/// A list's two replacement centroids and the child each of its points joins. -#[derive(Debug)] -pub(in crate::ivf) struct TwoMeans { - /// Row `c` is the centroid of child `c`. - pub(in crate::ivf) centroids: rowmajor::Owned, - /// `children[i]` is the child row `i` of the clustered points joins. - pub(in crate::ivf) children: Vec, -} - -/// Cluster `points` into two groups, seeding Lloyd's algorithm with two distinct -/// points. -/// -/// # Errors -/// -/// Fails if `points` holds fewer than two rows. -pub(in crate::ivf) fn two_means( - points: rowmajor::Ref<'_, f32>, - iterations: usize, - rng: &mut StdRng, -) -> ANNResult { - let (count, dim) = (points.nrows(), points.ncols()); - if count < 2 { - return Err(index_error(format!( - "a split needs at least two points, got {count}" - ))); - } - - let first = rng.random_range(0..count); - let second = (first + rng.random_range(1..count)) % count; - let mut centroids = rowmajor::Owned::from_element(2, dim, 0.0f32); - centroids.row_mut(0).copy_from_slice(points.row(first)); - centroids.row_mut(1).copy_from_slice(points.row(second)); - - lloyd( - || points.rows(), - centroids.as_mut_slice(), - dim, - iterations.max(1), - &mut LloydScratch::default(), - ); - - let children = { - let candidates = [(0, centroids.row(0)), (1, centroids.row(1))]; - points - .rows() - // There are always two candidates, so every point joins one of them. - .map(|point| nearest(point, &candidates).unwrap_or(0)) - .collect() - }; - - Ok(TwoMeans { - centroids, - children, - }) -} - -#[cfg(test)] -mod tests { - use rand::SeedableRng; - - use 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() - } - - #[test] - fn two_means_separates_distant_clusters() { - let points = matrix(&[[0.0, 0.0], [10.0, 0.0], [0.0, 1.0], [10.0, 1.0], [1.0, 0.0]]); - let fit = two_means(points.as_view(), 10, &mut StdRng::seed_from_u64(7)).unwrap(); - - assert_eq!(fit.centroids.nrows(), 2); - assert_eq!(fit.children.len(), 5); - // Rows 0, 2 and 4 are near the origin; rows 1 and 3 are near x = 10. - let low = fit.children[0]; - assert_eq!(fit.children, vec![low, 1 - low, low, 1 - low, low]); - assert_eq!(fit.centroids.row(low), &[1.0 / 3.0, 1.0 / 3.0]); - assert_eq!(fit.centroids.row(1 - low), &[10.0, 0.5]); - } - - #[test] - fn two_means_needs_two_points() { - let points = matrix(&[[1.0]]); - let err = two_means(points.as_view(), 1, &mut StdRng::seed_from_u64(7)).unwrap_err(); - assert!(err.to_string().contains("at least two points"), "{err}"); - } -} diff --git a/diskann/src/ivf/split.rs b/diskann/src/ivf/split.rs new file mode 100644 index 0000000000..67eb6100b9 --- /dev/null +++ b/diskann/src/ivf/split.rs @@ -0,0 +1,1119 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Gather one region, replace parent centroids, assign points, and build an update. +//! +//! Source membership stays fixed while splitting changes the candidate centroids. +//! Strategies and providers are trusted to preserve row alignment and valid destinations. + +use std::ops::Range; + +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; +use hashbrown::HashMap; +use rand::rngs::StdRng; + +use crate::{ + ANNResult, + error::{ErrorExt, IntoANNResult}, + ivf::{ + dynamic::{Centroids, InsertAccessor, Provider, Reader}, + online::{LloydScratch, assign_nearest, fit_two_means, index_error}, + update::{CentroidDelta, Delta, Deltas, MoveTo}, + }, + utils::VectorId, +}; + +/// Original membership, independent of the region's eventual destinations. +struct SourceList { + id: L, + rows: Range, + /// Committed members precede staged points within `rows`. + member_count: usize, +} + +struct Region { + lists: Vec>, + point_ids: Vec, + points: rowmajor::Owned, + /// `None` denotes a proposed centroid whose list ID has not been staged. + centroid_ids: Vec>, + centroids: rowmajor::Owned, +} + +impl Region { + fn list_points(&self, index: usize) -> ANNResult> { + self.points + .subview(self.lists[index].rows.clone()) + .ok_or_else(|| index_error("source rows lie outside the region's point matrix")) + } + + fn member_ids(&self, index: usize) -> &[I] { + let list = &self.lists[index]; + &self.point_ids[list.rows.start..list.rows.start + list.member_count] + } + + fn staged_ids(&self, index: usize) -> &[I] { + let list = &self.lists[index]; + &self.point_ids[list.rows.start + list.member_count..list.rows.end] + } +} + +/// Replace parent centroids without changing source point storage or row order. +trait Split: Send + Sync { + fn split(&self, region: &mut Region, parents: &[L], rng: &mut StdRng) -> ANNResult<()>; +} + +/// Return one valid candidate-centroid row per point, in the region's row order. +trait Assign: Send + Sync { + fn assign(&self, region: &Region) -> ANNResult>; +} + +struct TwoMeansSplit { + iterations: usize, +} + +impl Split for TwoMeansSplit { + fn split(&self, region: &mut Region, parents: &[L], rng: &mut StdRng) -> ANNResult<()> { + let is_parent = |id: &Option| id.is_some_and(|id| parents.contains(&id)); + let survivors = region + .centroid_ids + .iter() + .filter(|id| !is_parent(id)) + .count(); + let count = 2 * parents.len() + survivors; + let dim = region.points.ncols(); + let mut centroids = rowmajor::Owned::try_from_element(count, dim, 0.0)?; + let mut centroid_ids = Vec::with_capacity(count); + let mut scratch = LloydScratch::default(); + + for (index, parent) in parents.iter().enumerate() { + let source = region + .lists + .iter() + .position(|list| list.id == *parent) + .ok_or_else(|| { + index_error(format!("split parent {parent} is not in the region")) + })?; + let start = 2 * index * dim; + let out = rowmajor::Mut::try_from_data( + &mut centroids.as_mut_slice()[start..start + 2 * dim], + 2, + dim, + )?; + fit_two_means( + region.list_points(source)?, + out, + self.iterations, + rng, + &mut scratch, + )?; + centroid_ids.extend([None, None]); + } + for (id, centroid) in region.centroid_ids.iter().zip(region.centroids.rows()) { + if !is_parent(id) { + centroids + .row_mut(centroid_ids.len()) + .copy_from_slice(centroid); + centroid_ids.push(*id); + } + } + region.centroids = centroids; + region.centroid_ids = centroid_ids; + Ok(()) + } +} + +struct NearestCentroid; + +impl Assign for NearestCentroid { + fn assign(&self, region: &Region) -> ANNResult> { + assign_nearest(region.points.as_view(), region.centroids.as_view()) + } +} + +async fn gather_region( + accessor: &mut A, + list_ids: &[P::ListId], + staged_ids: &[P::InternalId], + staged_vectors: rowmajor::Ref<'_, f32>, + routed: &HashMap>, +) -> ANNResult> +where + P: Provider, + A: InsertAccessor, + T: Sync, +{ + let mut lists = Vec::with_capacity(list_ids.len()); + let mut rows = 0; + for &id in list_ids { + let member_count = accessor.get_members(id).into_ann_result()?.len(); + let staged_count = routed.get(&id).map_or(0, Vec::len); + let end = rows + member_count + staged_count; + lists.push(SourceList { + id, + rows: rows..end, + member_count, + }); + rows = end; + } + + let dim = staged_vectors.ncols(); + let mut point_ids = Vec::with_capacity(rows); + let mut points = rowmajor::Owned::try_from_element(rows, dim, 0.0)?; + let mut centroids = rowmajor::Owned::try_from_element(list_ids.len(), dim, 0.0)?; + let mut centroid_ids = Vec::with_capacity(list_ids.len()); + for (index, list) in lists.iter().enumerate() { + point_ids.extend_from_slice(accessor.get_members(list.id).into_ann_result()?); + let staged = routed.get(&list.id).map_or(&[][..], Vec::as_slice); + point_ids.extend(staged.iter().map(|&position| staged_ids[position])); + + let values = &mut points.as_mut_slice()[list.rows.start * dim..list.rows.end * dim]; + let (members, tail) = values.split_at_mut(list.member_count * dim); + accessor + .reader() + .read_into( + list.id, + rowmajor::Mut::try_from_data(members, list.member_count, dim)?, + ) + .await + .escalate("split must read the requested list's vectors")?; + let mut tail = rowmajor::Mut::try_from_data(tail, staged.len(), dim)?; + for (&position, out) in staged.iter().zip(tail.rows_mut()) { + out.copy_from_slice(staged_vectors.row(position)); + } + + let catalog = accessor.centroids(); + let centroid = catalog + .centroid(list.id) + .ok_or_else(|| index_error(format!("centroid {} is unavailable", list.id)))?; + centroids.row_mut(index).copy_from_slice(¢roid); + centroid_ids.push(Some(list.id)); + } + Ok(Region { + lists, + point_ids, + points, + centroid_ids, + centroids, + }) +} + +async fn build_deltas( + accessor: &mut A, + region: &Region, + assignments: &[usize], +) -> ANNResult>> +where + P: Provider, + A: InsertAccessor, + T: Sync, +{ + let mut deltas = Vec::with_capacity(2 * region.centroid_ids.len() + region.lists.len()); + let mut destinations = Vec::with_capacity(region.centroid_ids.len()); + for (id, centroid) in region.centroid_ids.iter().zip(region.centroids.rows()) { + let id = match id { + Some(id) => *id, + None => { + let id = accessor.stage_centroid(centroid).await?; + deltas.push(Delta::CentroidDelta { + id, + delta: CentroidDelta::Install, + }); + id + } + }; + destinations.push(id); + } + + let mut append_counts = vec![0; destinations.len()]; + for (index, list) in region.lists.iter().enumerate() { + let retired = !region.centroid_ids.contains(&Some(list.id)); + let mut moves = Vec::with_capacity(list.member_count); + for (&id, &to) in region + .member_ids(index) + .iter() + .zip(&assignments[list.rows.start..list.rows.start + list.member_count]) + { + let to = destinations[to]; + if to != list.id { + moves.push(MoveTo::new(id, to)); + } + } + if retired { + deltas.push(Delta::CentroidDelta { + id: list.id, + delta: CentroidDelta::Retire { + moves: moves.into_boxed_slice(), + }, + }); + } else if !moves.is_empty() { + deltas.push(Delta::PointMoves { + from: list.id, + moves: moves.into_boxed_slice(), + }); + } + for &to in &assignments[list.rows.start + list.member_count..list.rows.end] { + append_counts[to] += 1; + } + } + + let mut appends: Vec<_> = append_counts.into_iter().map(Vec::with_capacity).collect(); + for (index, list) in region.lists.iter().enumerate() { + for (&id, &to) in region + .staged_ids(index) + .iter() + .zip(&assignments[list.rows.start + list.member_count..list.rows.end]) + { + appends[to].push(id); + } + } + for (to, ids) in destinations.into_iter().zip(appends) { + if !ids.is_empty() { + deltas.push(Delta::PointAppends { + to, + ids: ids.into_boxed_slice(), + }); + } + } + Ok(deltas) +} + +/// Split all parents in one region; append points routed outside it unchanged. +/// +/// Every regional point can choose any child, but no retiring parent. +pub(super) async fn plan_split_update( + accessor: &mut A, + two_means_iterations: usize, + rng: &mut StdRng, + ids: &[P::InternalId], + vectors: rowmajor::Ref<'_, f32>, + routed: &HashMap>, + parents: &[P::ListId], +) -> ANNResult> +where + P: Provider, + A: InsertAccessor, + T: Sync, +{ + let mut deltas = if parents.is_empty() { + Vec::with_capacity(routed.len()) + } else { + let mut region = gather_region::(accessor, parents, ids, vectors, routed).await?; + TwoMeansSplit { + iterations: two_means_iterations, + } + .split(&mut region, parents, rng)?; + let assignments = NearestCentroid.assign(®ion)?; + build_deltas::(accessor, ®ion, &assignments).await? + }; + for (&list, positions) in routed { + if !parents.contains(&list) && !positions.is_empty() { + deltas.push(Delta::PointAppends { + to: list, + ids: positions.iter().map(|&position| ids[position]).collect(), + }); + } + } + Ok(Deltas::new(deltas)) +} + +#[cfg(test)] +mod tests { + use std::{ + cell::Cell, + sync::{Arc, Mutex}, + }; + + use rand::{Rng, SeedableRng}; + + use super::*; + use crate::{ + ANNError, + ivf::{ + Config, IVFIndex, + dynamic::{MaintenanceStrategy, SelectedList, Stage}, + }, + provider::DefaultContext, + }; + + struct TestProvider; + + impl Provider for TestProvider { + type Context = DefaultContext; + type InternalId = u32; + type ExternalId = u32; + type ListId = u32; + type Error = ANNError; + + fn to_internal_id(&self, _context: &DefaultContext, id: &u32) -> ANNResult { + Ok(*id) + } + + fn to_external_id(&self, _context: &DefaultContext, id: u32) -> ANNResult { + Ok(id) + } + } + + struct TestList { + members: Box<[u32]>, + vectors: rowmajor::Owned, + centroid: Box<[f32]>, + } + + struct TestAccessor { + lists: HashMap, + reads: Arc>>, + stage_calls: Cell, + staged: Vec<(u32, Box<[f32]>)>, + applied: Arc>>>, + fail_read: Option, + fail_stage: Option, + missing_centroid: Option, + } + + fn matrix(values: &[f32]) -> rowmajor::Owned { + rowmajor::Owned::try_from_data(values.into(), values.len(), 1).unwrap() + } + + fn routes() -> HashMap> { + [(10, vec![3, 0]), (20, vec![1]), (30, vec![2])] + .into_iter() + .collect() + } + + impl TestAccessor { + fn new() -> Self { + Self { + lists: [ + ( + 10, + TestList { + members: Box::new([100, 101]), + vectors: matrix(&[0.0, 2.0]), + centroid: Box::new([1.0]), + }, + ), + ( + 20, + TestList { + members: Box::new([200, 201]), + vectors: matrix(&[10.0, 12.0]), + centroid: Box::new([11.0]), + }, + ), + ( + 30, + TestList { + members: Box::new([300]), + vectors: matrix(&[100.0]), + centroid: Box::new([100.0]), + }, + ), + ] + .into_iter() + .collect(), + reads: Arc::default(), + stage_calls: Cell::new(0), + staged: Vec::new(), + applied: Arc::default(), + fail_read: None, + fail_stage: None, + missing_centroid: None, + } + } + } + + struct TestCentroids<'a> { + lists: &'a HashMap, + missing: Option, + } + + impl Centroids for TestCentroids<'_> { + type Id = u32; + type Error = ANNError; + + fn len(&self) -> usize { + self.lists.len() + } + + fn dim(&self) -> usize { + self.lists.values().next().unwrap().centroid.len() + } + + fn centroid(&self, id: u32) -> Option> { + self.lists + .get(&id) + .filter(|_| self.missing != Some(id)) + .map(|list| list.centroid.as_ref()) + } + + fn select(&self, query: &[f32], nprobe: usize) -> ANNResult>> { + let mut selected: Vec<_> = self + .lists + .iter() + .map(|(&id, list)| SelectedList { + id, + distance: query + .iter() + .zip(list.centroid.iter()) + .map(|(a, b)| (a - b).powi(2)) + .sum(), + }) + .collect(); + selected.sort_by(|a, b| a.distance.total_cmp(&b.distance).then(a.id.cmp(&b.id))); + selected.truncate(nprobe); + Ok(selected) + } + } + + struct TestReader<'a> { + lists: &'a HashMap, + reads: &'a Mutex>, + fail: Option, + } + + impl Reader for TestReader<'_> { + type Error = ANNError; + type Id = u32; + + fn dim(&self) -> usize { + self.lists.values().next().unwrap().centroid.len() + } + + async fn read_into(&self, id: u32, mut out: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + self.reads.lock().unwrap().push(id); + if self.fail == Some(id) { + return Err(index_error("fixture read failed")); + } + out.as_mut_slice() + .copy_from_slice(self.lists[&id].vectors.as_slice()); + Ok(()) + } + } + + impl Stage for TestAccessor { + async fn stage_point(&mut self, id: &u32, element: f32, out: &mut [f32]) -> ANNResult { + out[0] = element; + Ok(*id) + } + + async fn stage_centroid(&mut self, centroid: &[f32]) -> ANNResult { + let call = self.stage_calls.get(); + self.stage_calls.set(call + 1); + if self.fail_stage == Some(call) { + return Err(index_error("fixture staging failed")); + } + let id = 1000 + call; + self.staged.push((id, centroid.into())); + Ok(id) + } + } + + impl InsertAccessor for TestAccessor { + type Centroids<'a> = TestCentroids<'a>; + type Reader<'a> = TestReader<'a>; + type Error = ANNError; + + fn reader(&self) -> TestReader<'_> { + TestReader { + lists: &self.lists, + reads: &self.reads, + fail: self.fail_read, + } + } + + fn centroids(&self) -> TestCentroids<'_> { + TestCentroids { + lists: &self.lists, + missing: self.missing_centroid, + } + } + + fn dim(&self) -> usize { + self.lists.values().next().unwrap().centroid.len() + } + + fn list_size(&self, list: u32) -> ANNResult { + Ok(self.get_members(list)?.len()) + } + + fn get_members(&self, list: u32) -> ANNResult<&[u32]> { + self.lists + .get(&list) + .map(|list| list.members.as_ref()) + .ok_or_else(|| index_error(format!("fixture list {list} is missing"))) + } + + async fn update(self, update: Deltas) -> ANNResult<()> { + self.applied.lock().unwrap().push(update); + Ok(()) + } + } + + struct TestStrategy(Mutex>); + + impl<'a> MaintenanceStrategy<'a, TestProvider, f32> for TestStrategy { + type MaintenanceAccessor = TestAccessor; + type Error = ANNError; + + fn maintenance_accessor( + &'a self, + _provider: &'a mut TestProvider, + _context: &'a DefaultContext, + ) -> ANNResult { + self.0 + .lock() + .unwrap() + .take() + .ok_or_else(|| index_error("fixture accessor was consumed")) + } + } + + fn require_send(value: T) -> T { + value + } + + fn appends(update: &Deltas) -> HashMap> { + update + .deltas() + .iter() + .filter_map(|delta| match delta { + Delta::PointAppends { to, ids } => Some((*to, ids.to_vec())), + _ => None, + }) + .collect() + } + + fn retired(update: &Deltas, parent: u32) -> Vec<(u32, u32)> { + update + .deltas() + .iter() + .find_map(|delta| match delta { + Delta::CentroidDelta { + id, + delta: CentroidDelta::Retire { moves }, + } if *id == parent => { + Some(moves.iter().map(|point| (point.id(), point.to())).collect()) + } + _ => None, + }) + .unwrap() + } + + #[tokio::test] + async fn gather_preserves_list_and_staged_order_in_one_buffer() { + let mut accessor = TestAccessor::new(); + let region = require_send(gather_region::( + &mut accessor, + &[10, 20, 30], + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + )) + .await + .unwrap(); + + assert_eq!( + region.point_ids, + [100, 101, 503, 500, 200, 201, 501, 300, 502] + ); + assert_eq!( + region.points.as_slice(), + &[0.0, 2.0, 3.0, 1.0, 10.0, 12.0, 11.0, 100.0, 100.0] + ); + assert_eq!(region.centroid_ids, [Some(10), Some(20), Some(30)]); + assert_eq!(region.centroids.as_slice(), &[1.0, 11.0, 100.0]); + for (index, (rows, members, staged)) in [ + (0..4, &[100, 101][..], &[503, 500][..]), + (4..7, &[200, 201][..], &[501][..]), + (7..9, &[300][..], &[502][..]), + ] + .into_iter() + .enumerate() + { + assert_eq!(region.lists[index].rows, rows); + assert_eq!(region.member_ids(index), members); + assert_eq!(region.staged_ids(index), staged); + assert_eq!( + region.list_points(index).unwrap().as_ptr(), + region.points.row(rows.start).as_ptr() + ); + } + assert_eq!(*accessor.reads.lock().unwrap(), [10, 20, 30]); + } + + #[tokio::test] + async fn gather_handles_multiple_dimensions_and_an_empty_member_prefix() { + let mut accessor = TestAccessor::new(); + accessor.lists = [( + 10, + TestList { + members: Box::new([]), + vectors: rowmajor::Owned::from_element(0, 2, 0.0), + centroid: Box::new([0.0, 0.0]), + }, + )] + .into_iter() + .collect(); + let staged = rowmajor::Owned::try_from_data(Box::new([1.0, 2.0, 3.0, 4.0]), 2, 2).unwrap(); + let region = gather_region::( + &mut accessor, + &[10], + &[500, 501], + staged.as_view(), + &[(10, vec![1, 0])].into_iter().collect(), + ) + .await + .unwrap(); + assert_eq!(region.points.as_slice(), &[3.0, 4.0, 1.0, 2.0]); + assert_eq!(region.point_ids, [501, 500]); + assert!(region.member_ids(0).is_empty()); + assert_eq!(*accessor.reads.lock().unwrap(), [10]); + } + + #[tokio::test] + async fn split_can_retire_an_empty_list_and_assign_only_staged_points() { + let mut accessor = TestAccessor::new(); + accessor.lists.get_mut(&10).unwrap().members = Box::new([]); + accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[]); + let update = plan_split_update::( + &mut accessor, + 10, + &mut StdRng::seed_from_u64(7), + &[500, 501], + matrix(&[0.0, 10.0]).as_view(), + &[(10, vec![1, 0])].into_iter().collect(), + &[10], + ) + .await + .unwrap(); + assert!(retired(&update, 10).is_empty()); + assert_eq!(accessor.staged.len(), 2); + let appended = appends(&update); + assert_eq!(appended.len(), 2); + for (to, ids) in appended { + assert!(to >= 1000); + assert_eq!(ids.len(), 1); + assert!([500, 501].contains(&ids[0])); + } + } + + #[tokio::test] + async fn split_replaces_parents_preserves_points_and_numerical_behavior() { + let mut accessor = TestAccessor::new(); + let mut region = gather_region::( + &mut accessor, + &[10, 20, 30], + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + ) + .await + .unwrap(); + let points_ptr = region.points.as_ptr(); + let ids_ptr = region.point_ids.as_ptr(); + let point_values = region.points.as_slice().to_vec(); + let mut rng = StdRng::seed_from_u64(7); + let mut expected_rng = rng.clone(); + let mut expected = rowmajor::Owned::from_element(4, 1, 0.0); + for (out, source) in expected.as_mut_slice().chunks_exact_mut(2).zip([1, 0]) { + fit_two_means( + region.list_points(source).unwrap(), + rowmajor::Mut::try_from_data(out, 2, 1).unwrap(), + 10, + &mut expected_rng, + &mut LloydScratch::default(), + ) + .unwrap(); + } + + TwoMeansSplit { iterations: 10 } + .split(&mut region, &[20, 10], &mut rng) + .unwrap(); + assert_eq!(region.centroid_ids, [None, None, None, None, Some(30)]); + assert_eq!(®ion.centroids.as_slice()[..4], expected.as_slice()); + assert_eq!(region.centroids.row(4), &[100.0]); + assert_eq!(region.points.as_ptr(), points_ptr); + assert_eq!(region.point_ids.as_ptr(), ids_ptr); + assert_eq!(region.points.as_slice(), point_values); + assert_eq!(rng.random::(), expected_rng.random::()); + assert_eq!(accessor.stage_calls.get(), 0); + } + + #[tokio::test] + async fn regional_assignment_can_choose_another_parents_children() { + let mut accessor = TestAccessor::new(); + accessor.lists.get_mut(&10).unwrap().members = Box::new([100, 101, 102]); + accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[0.0, 10.0, 100.0]); + accessor.lists.get_mut(&20).unwrap().vectors = matrix(&[1.0, 2.0]); + let mut region = gather_region::( + &mut accessor, + &[10, 20], + &[], + matrix(&[]).as_view(), + &HashMap::new(), + ) + .await + .unwrap(); + TwoMeansSplit { iterations: 10 } + .split(&mut region, &[10, 20], &mut StdRng::seed_from_u64(7)) + .unwrap(); + let assigned = NearestCentroid.assign(®ion).unwrap(); + assert_eq!(region.centroid_ids, [None, None, None, None]); + assert!( + assigned[0] >= 2, + "A's point at zero should prefer B's child at one" + ); + assert_eq!(region.centroids.row(assigned[0]), &[1.0]); + let update = Deltas::new( + build_deltas::(&mut accessor, ®ion, &assigned) + .await + .unwrap(), + ); + assert_eq!(retired(&update, 10)[0], (100, 1000 + assigned[0] as u32)); + for parent in [10, 20] { + assert!(retired(&update, parent).iter().all(|(_, to)| *to >= 1000)); + } + } + + #[tokio::test] + async fn planner_moves_all_parent_members_and_appends_every_staged_point() { + let mut accessor = TestAccessor::new(); + let update = require_send(plan_split_update::( + &mut accessor, + 10, + &mut StdRng::seed_from_u64(7), + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + &[10, 20], + )) + .await + .unwrap(); + assert_eq!(accessor.staged.len(), 4); + assert_eq!(*accessor.reads.lock().unwrap(), [10, 20]); + for (parent, ids) in [(10, [100, 101]), (20, [200, 201])] { + let moves = retired(&update, parent); + assert_eq!(moves.iter().map(|(id, _)| *id).collect::>(), ids); + assert!(moves.iter().all(|(_, to)| *to >= 1000)); + } + let appended = appends(&update); + assert_eq!(appended[&30], [502]); + let mut ids: Vec<_> = appended.values().flatten().copied().collect(); + ids.sort_unstable(); + assert_eq!(ids, [500, 501, 502, 503]); + let destination = appended.values().find(|ids| ids.contains(&503)).unwrap(); + assert_eq!(destination, &[503]); + } + + struct ThreeChildren; + + impl Split for ThreeChildren { + fn split( + &self, + region: &mut Region, + _parents: &[u32], + _rng: &mut StdRng, + ) -> ANNResult<()> { + region.centroids = matrix(&[0.0, 2.0, 100.0, 100.0]); + region.centroid_ids = vec![None, None, None, Some(30)]; + Ok(()) + } + } + + #[tokio::test] + async fn custom_split_supports_three_children_and_neighbor_reassignment() { + let mut accessor = TestAccessor::new(); + let mut region = gather_region::( + &mut accessor, + &[10, 30], + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + ) + .await + .unwrap(); + ThreeChildren + .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) + .unwrap(); + let assigned = NearestCentroid.assign(®ion).unwrap(); + let update = Deltas::new( + build_deltas::(&mut accessor, ®ion, &assigned) + .await + .unwrap(), + ); + assert_eq!(accessor.staged.len(), 3); + let neighbor_moves: Vec<_> = update + .deltas() + .iter() + .filter_map(|delta| match delta { + Delta::PointMoves { from: 30, moves } => Some(moves.as_ref()), + _ => None, + }) + .flatten() + .map(|point| (point.id(), point.to())) + .collect(); + assert_eq!(neighbor_moves, [(300, 1002)]); + assert_eq!(appends(&update)[&1002], [502]); + assert!( + update + .deltas() + .iter() + .all(|delta| !matches!(delta, Delta::CentroidDelta { id: 30, .. })) + ); + } + + #[tokio::test] + async fn surviving_neighbor_members_and_staged_points_can_stay_put() { + let mut accessor = TestAccessor::new(); + let mut region = gather_region::( + &mut accessor, + &[10, 30], + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + ) + .await + .unwrap(); + TwoMeansSplit { iterations: 10 } + .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) + .unwrap(); + let assigned = NearestCentroid.assign(®ion).unwrap(); + let update = Deltas::new( + build_deltas::(&mut accessor, ®ion, &assigned) + .await + .unwrap(), + ); + assert_eq!(accessor.staged.len(), 2); + assert!( + update + .deltas() + .iter() + .all(|delta| !matches!(delta, Delta::PointMoves { .. })) + ); + assert_eq!(appends(&update)[&30], [502]); + } + + #[tokio::test] + async fn custom_assignment_uses_the_same_flat_row_contract() { + struct LastCentroid; + impl Assign for LastCentroid { + fn assign(&self, region: &Region) -> ANNResult> { + Ok(vec![region.centroids.nrows() - 1; region.points.nrows()].into_boxed_slice()) + } + } + let mut accessor = TestAccessor::new(); + let mut region = gather_region::( + &mut accessor, + &[10, 30], + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + ) + .await + .unwrap(); + TwoMeansSplit { iterations: 10 } + .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) + .unwrap(); + let update = Deltas::new( + build_deltas::( + &mut accessor, + ®ion, + &LastCentroid.assign(®ion).unwrap(), + ) + .await + .unwrap(), + ); + assert_eq!(retired(&update, 10), [(100, 30), (101, 30)]); + assert_eq!(appends(&update)[&30], [503, 500, 502]); + } + + #[tokio::test] + async fn no_splits_and_empty_batches_skip_reads_and_centroid_staging() { + for (ids, vectors, routed) in [ + ( + vec![500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]), + routes(), + ), + (Vec::new(), matrix(&[]), HashMap::new()), + ] { + let mut accessor = TestAccessor::new(); + let update = plan_split_update::( + &mut accessor, + 10, + &mut StdRng::seed_from_u64(7), + &ids, + vectors.as_view(), + &routed, + &[], + ) + .await + .unwrap(); + assert!(accessor.reads.lock().unwrap().is_empty()); + assert_eq!(accessor.stage_calls.get(), 0); + assert_eq!(update.deltas().len(), routed.len()); + if !ids.is_empty() { + assert_eq!( + appends(&update), + [(10, vec![503, 500]), (20, vec![501]), (30, vec![502]),] + .into_iter() + .collect() + ); + } + } + } + + #[tokio::test] + async fn gathering_errors_propagate_before_centroid_staging() { + for failure in ["members", "reader", "centroid"] { + let mut accessor = TestAccessor::new(); + match failure { + "members" => { + accessor.lists.remove(&20); + } + "reader" => { + accessor.fail_read = Some(20); + } + "centroid" => { + accessor.missing_centroid = Some(20); + } + _ => unreachable!(), + } + let err = plan_split_update::( + &mut accessor, + 10, + &mut StdRng::seed_from_u64(7), + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + &[10, 20], + ) + .await + .unwrap_err(); + let expected = match failure { + "members" => "fixture list 20 is missing", + "reader" => "fixture read failed", + "centroid" => "centroid 20 is unavailable", + _ => unreachable!(), + }; + assert!(err.to_string().contains(expected), "{err}"); + assert_eq!(accessor.stage_calls.get(), 0); + } + } + + #[tokio::test] + async fn numerical_and_staging_errors_propagate() { + let mut accessor = TestAccessor::new(); + accessor.lists.get_mut(&10).unwrap().members = Box::new([100]); + accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[0.0]); + let err = plan_split_update::( + &mut accessor, + 10, + &mut StdRng::seed_from_u64(7), + &[], + matrix(&[]).as_view(), + &HashMap::new(), + &[10], + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("at least two points"), "{err}"); + assert_eq!(accessor.stage_calls.get(), 0); + + let mut region = gather_region::( + &mut accessor, + &[10], + &[], + matrix(&[]).as_view(), + &HashMap::new(), + ) + .await + .unwrap(); + region.centroids = matrix(&[]); + region.centroid_ids.clear(); + let err = NearestCentroid.assign(®ion).unwrap_err(); + assert!(err.to_string().contains("at least one centroid"), "{err}"); + + accessor.fail_stage = Some(1); + let err = plan_split_update::( + &mut accessor, + 10, + &mut StdRng::seed_from_u64(7), + &[500, 501, 502, 503], + matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), + &routes(), + &[10, 20], + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("fixture staging failed"), "{err}"); + assert_eq!(accessor.staged.len(), 1); + } + + #[tokio::test] + async fn insert_batch_applies_splits_and_unsplit_appends_together() { + for threshold in [2, 100] { + let accessor = TestAccessor::new(); + let applied = Arc::clone(&accessor.applied); + let reads = Arc::clone(&accessor.reads); + let strategy = TestStrategy(Mutex::new(Some(accessor))); + let mut index = IVFIndex::new( + TestProvider, + Config { + split_threshold: threshold, + reassign_neighbors: 1, + two_means_iterations: 10, + seed: 7, + }, + ) + .unwrap(); + require_send(index.insert_batch( + &strategy, + &DefaultContext, + &[(500, 1.0), (501, 11.0), (502, 100.0), (503, 3.0)], + )) + .await + .unwrap(); + let updates = applied.lock().unwrap(); + assert_eq!(updates.len(), 1); + let appended = appends(&updates[0]); + let mut ids: Vec<_> = appended.values().flatten().copied().collect(); + ids.sort_unstable(); + assert_eq!(ids, [500, 501, 502, 503]); + assert_eq!(appended[&30], [502]); + if threshold == 2 { + assert_eq!(retired(&updates[0], 10).len(), 2); + assert_eq!(retired(&updates[0], 20).len(), 2); + let mut read_ids = reads.lock().unwrap().clone(); + read_ids.sort_unstable(); + assert_eq!(read_ids, [10, 20]); + } else { + assert_eq!(appended[&10], [500, 503]); + assert_eq!(appended[&20], [501]); + assert!(reads.lock().unwrap().is_empty()); + } + } + } + + #[tokio::test] + async fn empty_insert_does_not_construct_an_accessor_or_apply_an_update() { + let strategy = TestStrategy(Mutex::new(None)); + let mut index = IVFIndex::new( + TestProvider, + Config { + split_threshold: 2, + reassign_neighbors: 1, + two_means_iterations: 10, + seed: 7, + }, + ) + .unwrap(); + let result = require_send(index.insert_batch(&strategy, &DefaultContext, &[])) + .await + .unwrap(); + assert_eq!(result.inserted, 0); + } +} From 74bd9b0f1ff691edc14deb6e54893b21b91a6079 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Thu, 8 Oct 2026 17:41:22 -0400 Subject: [PATCH 13/33] small --- diskann/src/ivf/index.rs | 228 +++++++- diskann/src/ivf/mod.rs | 2 +- diskann/src/ivf/split.rs | 649 ++++++++++------------ diskann/src/ivf/{dynamic.rs => traits.rs} | 0 4 files changed, 486 insertions(+), 393 deletions(-) rename diskann/src/ivf/{dynamic.rs => traits.rs} (100%) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 37f29ff231..d7c9ac03e7 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -7,7 +7,8 @@ //! //! The index drives one [`InsertAccessor`] per mutation: it plans a //! [`Deltas`](super::update::Deltas) against the accessor's view and hands it to -//! [`InsertAccessor::update`]. It never touches point, centroid, or list storage directly. +//! [`InsertAccessor::update`]. Gathering and update construction use only the accessor; +//! splitting and assignment operate on an in-memory region. use diskann_utils::{ future::SendFuture, @@ -20,9 +21,10 @@ use crate::{ ANNError, ANNResult, error::{ErrorExt, IntoANNResult}, ivf::{ - dynamic::{Centroids, InsertAccessor, MaintenanceStrategy, Provider, Stage}, - split::plan_split_update, - update::{self, Deltas}, + traits::{Centroids, InsertAccessor, MaintenanceStrategy, Provider, Reader, Stage}, + online::index_error, + split::{Assign, NearestCentroid, Region, SourceList, Split, TwoMeansSplit}, + update::{self, CentroidDelta, Delta, Deltas, MoveTo}, }, }; @@ -66,6 +68,18 @@ pub struct InsertStats { pub reassigned: usize, } +/// A staged point and its borrowed canonical vector, ready for regional gathering. +pub(super) struct StagedPoint<'a, I> { + pub(super) id: I, + pub(super) vector: &'a [f32], +} + +/// Staged points routed to one existing list, in insertion order. +pub(super) struct ListInsertion<'a, I, L> { + pub(super) list: L, + pub(super) staged: Vec>, +} + /// An incrementally maintained IVF index. /// /// Mutations take `&mut self` and lend the provider exclusively to one @@ -184,17 +198,13 @@ impl IVFIndex

{ .maintenance_accessor(provider, context) .into_ann_result()?; - // stage the input points let mut vectors = rowmajor::Owned::from_element(points.len(), accessor.dim(), f32::NAN); - let mut ids = Vec::with_capacity(points.len()); - let mut routed = hashbrown::HashMap::>::new(); + let mut routed = + hashbrown::HashMap::>>::new(); - for (position, ((id, v), out)) in points.iter().zip(vectors.rows_mut()).enumerate() { - //stage the input points first. + for ((id, v), out) in points.iter().zip(vectors.rows_mut()) { let id = accessor.stage_point(id, *v, out).await.into_ann_result()?; - ids.push(id); - // route the points to their nearest centroid. let parent = accessor .centroids() .select(out, 1) @@ -203,39 +213,199 @@ impl IVFIndex

{ .map(|s| s.id) .ok_or_else(|| ANNError::message("Didn't return list"))?; - routed.entry(parent).or_default().push(position); + routed + .entry(parent) + .or_default() + .push(StagedPoint { id, vector: out }); } - // filter the routed lists for ones that need splitting. - let mut parents = Vec::new(); - for (&list, positions) in &routed { + let mut split_inputs = Vec::with_capacity(routed.len()); + let mut deltas = Vec::with_capacity(routed.len()); + for (list, staged) in routed { let len = accessor.list_size(list).into_ann_result()?; - if len + positions.len() > config.split_threshold { - parents.push(list); + if len + staged.len() > config.split_threshold { + split_inputs.push(ListInsertion { list, staged }); + } else { + deltas.push(update::Delta::PointAppends { + to: list, + ids: staged.into_iter().map(|point| point.id).collect(), + }); } } - // Fit every split's children before assigning points and building the update. - let update = plan_split_update::( - &mut accessor, - config.two_means_iterations, - rng, - &ids, - vectors.as_view(), - &routed, - &parents, - ) - .await?; + if !split_inputs.is_empty() { + let mut region = Self::gather_region::<_, T>(&mut accessor, &split_inputs).await?; + let parents: Vec<_> = split_inputs.iter().map(|input| input.list).collect(); + TwoMeansSplit { + iterations: config.two_means_iterations, + } + .split(&mut region, &parents, rng)?; + let assignments = NearestCentroid.assign(®ion)?; + deltas.extend( + Self::build_deltas::<_, T>(&mut accessor, ®ion, &assignments).await?, + ); + } accessor - .update(update) + .update(Deltas::new(deltas)) .await .escalate("insert must apply its update")?; Ok(InsertStats::default()) } } + + /// Gather committed members and staged points into their final region rows. + pub(super) async fn gather_region( + accessor: &mut A, + inputs: &[ListInsertion<'_, P::InternalId, P::ListId>], + ) -> ANNResult> + where + A: InsertAccessor, + T: Sync, + { + let mut lists = Vec::with_capacity(inputs.len()); + let mut rows = 0; + for input in inputs { + let id = input.list; + let member_count = accessor.get_members(id).into_ann_result()?.len(); + let end = rows + member_count + input.staged.len(); + lists.push(SourceList { + id, + rows: rows..end, + member_count, + }); + rows = end; + } + + let dim = accessor.dim(); + let mut point_ids = Vec::with_capacity(rows); + let mut points = rowmajor::Owned::try_from_element(rows, dim, 0.0)?; + let mut centroids = rowmajor::Owned::try_from_element(inputs.len(), dim, 0.0)?; + let mut centroid_ids = Vec::with_capacity(inputs.len()); + + for (index, (list, input)) in lists.iter().zip(inputs).enumerate() { + point_ids.extend_from_slice(accessor.get_members(list.id).into_ann_result()?); + + let member_end = list.rows.start + list.member_count; + accessor + .reader() + .read_into( + list.id, + points + .subview_mut(list.rows.start..member_end) + .ok_or_else(|| { + index_error("member rows lie outside the region's point matrix") + })?, + ) + .await + .escalate("split must read the requested list's vectors")?; + let mut tail = points + .subview_mut(member_end..list.rows.end) + .ok_or_else(|| index_error("staged rows lie outside the region's point matrix"))?; + for (point, out) in input.staged.iter().zip(tail.rows_mut()) { + point_ids.push(point.id); + out.copy_from_slice(point.vector); + } + + let catalog = accessor.centroids(); + let centroid = catalog + .centroid(list.id) + .ok_or_else(|| index_error(format!("centroid {} is unavailable", list.id)))?; + centroids.row_mut(index).copy_from_slice(¢roid); + centroid_ids.push(Some(list.id)); + } + + Ok(Region { + lists, + point_ids, + points, + centroid_ids, + centroids, + }) + } + + /// Stage proposed centroids and translate aligned assignments into partition deltas. + pub(super) async fn build_deltas( + accessor: &mut A, + region: &Region, + assignments: &[usize], + ) -> ANNResult>> + where + A: InsertAccessor, + T: Sync, + { + let mut deltas = Vec::with_capacity(2 * region.centroid_ids.len() + region.lists.len()); + let mut destinations = Vec::with_capacity(region.centroid_ids.len()); + for (id, centroid) in region.centroid_ids.iter().zip(region.centroids.rows()) { + let id = match id { + Some(id) => *id, + None => { + let id = accessor.stage_centroid(centroid).await?; + deltas.push(Delta::CentroidDelta { + id, + delta: CentroidDelta::Install, + }); + id + } + }; + destinations.push(id); + } + + let mut append_counts = vec![0; destinations.len()]; + + for (index, list) in region.lists.iter().enumerate() { + let retired = !region.centroid_ids.contains(&Some(list.id)); + let mut moves = Vec::with_capacity(list.member_count); + for (&id, &to) in region + .member_ids(index) + .iter() + .zip(&assignments[list.rows.start..list.rows.start + list.member_count]) + { + let to = destinations[to]; + if to != list.id { + moves.push(MoveTo::new(id, to)); + } + } + if retired { + deltas.push(Delta::CentroidDelta { + id: list.id, + delta: CentroidDelta::Retire { + moves: moves.into_boxed_slice(), + }, + }); + } else if !moves.is_empty() { + deltas.push(Delta::PointMoves { + from: list.id, + moves: moves.into_boxed_slice(), + }); + } + for &to in &assignments[list.rows.start + list.member_count..list.rows.end] { + append_counts[to] += 1; + } + } + + let mut appends: Vec<_> = append_counts.into_iter().map(Vec::with_capacity).collect(); + for (index, list) in region.lists.iter().enumerate() { + for (&id, &to) in region + .staged_ids(index) + .iter() + .zip(&assignments[list.rows.start + list.member_count..list.rows.end]) + { + appends[to].push(id); + } + } + for (to, ids) in destinations.into_iter().zip(appends) { + if !ids.is_empty() { + deltas.push(Delta::PointAppends { + to, + ids: ids.into_boxed_slice(), + }); + } + } + Ok(deltas) + } } #[cfg(test)] diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 7a271700f7..0316bb3258 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -5,7 +5,7 @@ //! Incrementally maintained IVF index. -pub mod dynamic; +pub mod traits; mod index; mod online; mod split; diff --git a/diskann/src/ivf/split.rs b/diskann/src/ivf/split.rs index 67eb6100b9..15c183c5c7 100644 --- a/diskann/src/ivf/split.rs +++ b/diskann/src/ivf/split.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Gather one region, replace parent centroids, assign points, and build an update. +//! In-memory regions and strategies for centroid splitting and point assignment. //! //! Source membership stays fixed while splitting changes the candidate centroids. //! Strategies and providers are trusted to preserve row alignment and valid destinations. @@ -11,35 +11,29 @@ use std::ops::Range; use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; -use hashbrown::HashMap; use rand::rngs::StdRng; use crate::{ ANNResult, - error::{ErrorExt, IntoANNResult}, - ivf::{ - dynamic::{Centroids, InsertAccessor, Provider, Reader}, - online::{LloydScratch, assign_nearest, fit_two_means, index_error}, - update::{CentroidDelta, Delta, Deltas, MoveTo}, - }, + ivf::online::{LloydScratch, assign_nearest, fit_two_means, index_error}, utils::VectorId, }; /// Original membership, independent of the region's eventual destinations. -struct SourceList { - id: L, - rows: Range, +pub(super) struct SourceList { + pub(super) id: L, + pub(super) rows: Range, /// Committed members precede staged points within `rows`. - member_count: usize, + pub(super) member_count: usize, } -struct Region { - lists: Vec>, - point_ids: Vec, - points: rowmajor::Owned, +pub(super) struct Region { + pub(super) lists: Vec>, + pub(super) point_ids: Vec, + pub(super) points: rowmajor::Owned, /// `None` denotes a proposed centroid whose list ID has not been staged. - centroid_ids: Vec>, - centroids: rowmajor::Owned, + pub(super) centroid_ids: Vec>, + pub(super) centroids: rowmajor::Owned, } impl Region { @@ -49,29 +43,29 @@ impl Region { .ok_or_else(|| index_error("source rows lie outside the region's point matrix")) } - fn member_ids(&self, index: usize) -> &[I] { + pub(super) fn member_ids(&self, index: usize) -> &[I] { let list = &self.lists[index]; &self.point_ids[list.rows.start..list.rows.start + list.member_count] } - fn staged_ids(&self, index: usize) -> &[I] { + pub(super) fn staged_ids(&self, index: usize) -> &[I] { let list = &self.lists[index]; &self.point_ids[list.rows.start + list.member_count..list.rows.end] } } /// Replace parent centroids without changing source point storage or row order. -trait Split: Send + Sync { +pub(super) trait Split: Send + Sync { fn split(&self, region: &mut Region, parents: &[L], rng: &mut StdRng) -> ANNResult<()>; } /// Return one valid candidate-centroid row per point, in the region's row order. -trait Assign: Send + Sync { +pub(super) trait Assign: Send + Sync { fn assign(&self, region: &Region) -> ANNResult>; } -struct TwoMeansSplit { - iterations: usize, +pub(super) struct TwoMeansSplit { + pub(super) iterations: usize, } impl Split for TwoMeansSplit { @@ -125,7 +119,7 @@ impl Split for TwoMeansSplit { } } -struct NearestCentroid; +pub(super) struct NearestCentroid; impl Assign for NearestCentroid { fn assign(&self, region: &Region) -> ANNResult> { @@ -133,192 +127,6 @@ impl Assign for NearestCentroid { } } -async fn gather_region( - accessor: &mut A, - list_ids: &[P::ListId], - staged_ids: &[P::InternalId], - staged_vectors: rowmajor::Ref<'_, f32>, - routed: &HashMap>, -) -> ANNResult> -where - P: Provider, - A: InsertAccessor, - T: Sync, -{ - let mut lists = Vec::with_capacity(list_ids.len()); - let mut rows = 0; - for &id in list_ids { - let member_count = accessor.get_members(id).into_ann_result()?.len(); - let staged_count = routed.get(&id).map_or(0, Vec::len); - let end = rows + member_count + staged_count; - lists.push(SourceList { - id, - rows: rows..end, - member_count, - }); - rows = end; - } - - let dim = staged_vectors.ncols(); - let mut point_ids = Vec::with_capacity(rows); - let mut points = rowmajor::Owned::try_from_element(rows, dim, 0.0)?; - let mut centroids = rowmajor::Owned::try_from_element(list_ids.len(), dim, 0.0)?; - let mut centroid_ids = Vec::with_capacity(list_ids.len()); - for (index, list) in lists.iter().enumerate() { - point_ids.extend_from_slice(accessor.get_members(list.id).into_ann_result()?); - let staged = routed.get(&list.id).map_or(&[][..], Vec::as_slice); - point_ids.extend(staged.iter().map(|&position| staged_ids[position])); - - let values = &mut points.as_mut_slice()[list.rows.start * dim..list.rows.end * dim]; - let (members, tail) = values.split_at_mut(list.member_count * dim); - accessor - .reader() - .read_into( - list.id, - rowmajor::Mut::try_from_data(members, list.member_count, dim)?, - ) - .await - .escalate("split must read the requested list's vectors")?; - let mut tail = rowmajor::Mut::try_from_data(tail, staged.len(), dim)?; - for (&position, out) in staged.iter().zip(tail.rows_mut()) { - out.copy_from_slice(staged_vectors.row(position)); - } - - let catalog = accessor.centroids(); - let centroid = catalog - .centroid(list.id) - .ok_or_else(|| index_error(format!("centroid {} is unavailable", list.id)))?; - centroids.row_mut(index).copy_from_slice(¢roid); - centroid_ids.push(Some(list.id)); - } - Ok(Region { - lists, - point_ids, - points, - centroid_ids, - centroids, - }) -} - -async fn build_deltas( - accessor: &mut A, - region: &Region, - assignments: &[usize], -) -> ANNResult>> -where - P: Provider, - A: InsertAccessor, - T: Sync, -{ - let mut deltas = Vec::with_capacity(2 * region.centroid_ids.len() + region.lists.len()); - let mut destinations = Vec::with_capacity(region.centroid_ids.len()); - for (id, centroid) in region.centroid_ids.iter().zip(region.centroids.rows()) { - let id = match id { - Some(id) => *id, - None => { - let id = accessor.stage_centroid(centroid).await?; - deltas.push(Delta::CentroidDelta { - id, - delta: CentroidDelta::Install, - }); - id - } - }; - destinations.push(id); - } - - let mut append_counts = vec![0; destinations.len()]; - for (index, list) in region.lists.iter().enumerate() { - let retired = !region.centroid_ids.contains(&Some(list.id)); - let mut moves = Vec::with_capacity(list.member_count); - for (&id, &to) in region - .member_ids(index) - .iter() - .zip(&assignments[list.rows.start..list.rows.start + list.member_count]) - { - let to = destinations[to]; - if to != list.id { - moves.push(MoveTo::new(id, to)); - } - } - if retired { - deltas.push(Delta::CentroidDelta { - id: list.id, - delta: CentroidDelta::Retire { - moves: moves.into_boxed_slice(), - }, - }); - } else if !moves.is_empty() { - deltas.push(Delta::PointMoves { - from: list.id, - moves: moves.into_boxed_slice(), - }); - } - for &to in &assignments[list.rows.start + list.member_count..list.rows.end] { - append_counts[to] += 1; - } - } - - let mut appends: Vec<_> = append_counts.into_iter().map(Vec::with_capacity).collect(); - for (index, list) in region.lists.iter().enumerate() { - for (&id, &to) in region - .staged_ids(index) - .iter() - .zip(&assignments[list.rows.start + list.member_count..list.rows.end]) - { - appends[to].push(id); - } - } - for (to, ids) in destinations.into_iter().zip(appends) { - if !ids.is_empty() { - deltas.push(Delta::PointAppends { - to, - ids: ids.into_boxed_slice(), - }); - } - } - Ok(deltas) -} - -/// Split all parents in one region; append points routed outside it unchanged. -/// -/// Every regional point can choose any child, but no retiring parent. -pub(super) async fn plan_split_update( - accessor: &mut A, - two_means_iterations: usize, - rng: &mut StdRng, - ids: &[P::InternalId], - vectors: rowmajor::Ref<'_, f32>, - routed: &HashMap>, - parents: &[P::ListId], -) -> ANNResult> -where - P: Provider, - A: InsertAccessor, - T: Sync, -{ - let mut deltas = if parents.is_empty() { - Vec::with_capacity(routed.len()) - } else { - let mut region = gather_region::(accessor, parents, ids, vectors, routed).await?; - TwoMeansSplit { - iterations: two_means_iterations, - } - .split(&mut region, parents, rng)?; - let assignments = NearestCentroid.assign(®ion)?; - build_deltas::(accessor, ®ion, &assignments).await? - }; - for (&list, positions) in routed { - if !parents.contains(&list) && !positions.is_empty() { - deltas.push(Delta::PointAppends { - to: list, - ids: positions.iter().map(|&position| ids[position]).collect(), - }); - } - } - Ok(Deltas::new(deltas)) -} - #[cfg(test)] mod tests { use std::{ @@ -326,6 +134,7 @@ mod tests { sync::{Arc, Mutex}, }; + use hashbrown::HashMap; use rand::{Rng, SeedableRng}; use super::*; @@ -333,7 +142,12 @@ mod tests { ANNError, ivf::{ Config, IVFIndex, - dynamic::{MaintenanceStrategy, SelectedList, Stage}, + traits::{ + Centroids, InsertAccessor, MaintenanceStrategy, Provider, Reader, SelectedList, + Stage, + }, + index::{ListInsertion, StagedPoint}, + update::{CentroidDelta, Delta, Deltas}, }, provider::DefaultContext, }; @@ -377,9 +191,30 @@ mod tests { rowmajor::Owned::try_from_data(values.into(), values.len(), 1).unwrap() } - fn routes() -> HashMap> { - [(10, vec![3, 0]), (20, vec![1]), (30, vec![2])] - .into_iter() + fn inputs<'a>( + vectors: &'a rowmajor::Owned, + lists: &[u32], + ) -> Vec> { + lists + .iter() + .map(|&list| { + let points: &[(u32, usize)] = match list { + 10 => &[(503, 3), (500, 0)], + 20 => &[(501, 1)], + 30 => &[(502, 2)], + _ => unreachable!(), + }; + ListInsertion { + list, + staged: points + .iter() + .map(|&(id, row)| StagedPoint { + id, + vector: vectors.row(row), + }) + .collect(), + } + }) .collect() } @@ -496,7 +331,7 @@ mod tests { impl Stage for TestAccessor { async fn stage_point(&mut self, id: &u32, element: f32, out: &mut [f32]) -> ANNResult { out[0] = element; - Ok(*id) + Ok(*id + 10_000) } async fn stage_centroid(&mut self, centroid: &[f32]) -> ANNResult { @@ -575,15 +410,30 @@ mod tests { value } + async fn split_update( + accessor: &mut TestAccessor, + inputs: &[ListInsertion<'_, u32, u32>], + ) -> ANNResult>> { + let mut region = + IVFIndex::::gather_region::<_, f32>(accessor, inputs).await?; + let parents: Vec<_> = inputs.iter().map(|input| input.list).collect(); + TwoMeansSplit { iterations: 10 }.split( + &mut region, + &parents, + &mut StdRng::seed_from_u64(7), + )?; + let assignments = NearestCentroid.assign(®ion)?; + IVFIndex::::build_deltas::<_, f32>(accessor, ®ion, &assignments).await + } + fn appends(update: &Deltas) -> HashMap> { - update - .deltas() - .iter() - .filter_map(|delta| match delta { - Delta::PointAppends { to, ids } => Some((*to, ids.to_vec())), - _ => None, - }) - .collect() + let mut appends = HashMap::new(); + for delta in update.deltas() { + if let Delta::PointAppends { to, ids } = delta { + assert!(appends.insert(*to, ids.to_vec()).is_none()); + } + } + appends } fn retired(update: &Deltas, parent: u32) -> Vec<(u32, u32)> { @@ -605,12 +455,13 @@ mod tests { #[tokio::test] async fn gather_preserves_list_and_staged_order_in_one_buffer() { let mut accessor = TestAccessor::new(); - let region = require_send(gather_region::( + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let inputs = inputs(&vectors, &[10, 20, 30]); + assert_eq!(inputs[0].staged[0].vector.as_ptr(), vectors.row(3).as_ptr()); + assert_eq!(inputs[0].staged[1].vector.as_ptr(), vectors.row(0).as_ptr()); + let region = require_send(IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10, 20, 30], - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), + &inputs, )) .await .unwrap(); @@ -658,12 +509,21 @@ mod tests { .into_iter() .collect(); let staged = rowmajor::Owned::try_from_data(Box::new([1.0, 2.0, 3.0, 4.0]), 2, 2).unwrap(); - let region = gather_region::( + let region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10], - &[500, 501], - staged.as_view(), - &[(10, vec![1, 0])].into_iter().collect(), + &[ListInsertion { + list: 10, + staged: vec![ + StagedPoint { + id: 501, + vector: staged.row(1), + }, + StagedPoint { + id: 500, + vector: staged.row(0), + }, + ], + }], ) .await .unwrap(); @@ -674,21 +534,93 @@ mod tests { } #[tokio::test] - async fn split_can_retire_an_empty_list_and_assign_only_staged_points() { + async fn gather_fills_disjoint_multidimensional_subviews_in_list_order() { let mut accessor = TestAccessor::new(); - accessor.lists.get_mut(&10).unwrap().members = Box::new([]); - accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[]); - let update = plan_split_update::( + accessor.lists = [ + ( + 10, + TestList { + members: Box::new([100, 101]), + vectors: rowmajor::Owned::try_from_data(Box::new([1.0, 2.0, 3.0, 4.0]), 2, 2) + .unwrap(), + centroid: Box::new([2.0, 3.0]), + }, + ), + ( + 20, + TestList { + members: Box::new([200]), + vectors: rowmajor::Owned::try_from_data(Box::new([5.0, 6.0]), 1, 2).unwrap(), + centroid: Box::new([5.0, 6.0]), + }, + ), + ] + .into_iter() + .collect(); + let staged = rowmajor::Owned::try_from_data(Box::new([7.0, 8.0, 9.0, 10.0]), 2, 2).unwrap(); + let region = require_send(IVFIndex::::gather_region::<_, f32>( &mut accessor, - 10, - &mut StdRng::seed_from_u64(7), - &[500, 501], - matrix(&[0.0, 10.0]).as_view(), - &[(10, vec![1, 0])].into_iter().collect(), - &[10], - ) + &[ + ListInsertion { + list: 10, + staged: vec![ + StagedPoint { + id: 501, + vector: staged.row(1), + }, + StagedPoint { + id: 500, + vector: staged.row(0), + }, + ], + }, + ListInsertion { + list: 20, + staged: Vec::new(), + }, + ], + )) .await .unwrap(); + + assert_eq!(region.lists[0].rows, 0..4); + assert_eq!(region.lists[1].rows, 4..5); + assert_eq!(region.point_ids, [100, 101, 501, 500, 200]); + assert_eq!( + region.points.as_slice(), + &[1.0, 2.0, 3.0, 4.0, 9.0, 10.0, 7.0, 8.0, 5.0, 6.0], + ); + assert_eq!(region.centroids.as_slice(), &[2.0, 3.0, 5.0, 6.0]); + assert!(region.staged_ids(1).is_empty()); + assert_eq!(*accessor.reads.lock().unwrap(), [10, 20]); + } + + #[tokio::test] + async fn split_can_retire_an_empty_list_and_assign_only_staged_points() { + let mut accessor = TestAccessor::new(); + accessor.lists.get_mut(&10).unwrap().members = Box::new([]); + accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[]); + let vectors = matrix(&[0.0, 10.0]); + let update = Deltas::new( + split_update( + &mut accessor, + &[ListInsertion { + list: 10, + staged: vec![ + StagedPoint { + id: 501, + vector: vectors.row(1), + }, + StagedPoint { + id: 500, + vector: vectors.row(0), + }, + ], + }], + ) + .await + .unwrap(), + ); assert!(retired(&update, 10).is_empty()); assert_eq!(accessor.staged.len(), 2); let appended = appends(&update); @@ -703,12 +635,10 @@ mod tests { #[tokio::test] async fn split_replaces_parents_preserves_points_and_numerical_behavior() { let mut accessor = TestAccessor::new(); - let mut region = gather_region::( + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let mut region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10, 20, 30], - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), + &inputs(&vectors, &[10, 20, 30]), ) .await .unwrap(); @@ -748,12 +678,18 @@ mod tests { accessor.lists.get_mut(&10).unwrap().members = Box::new([100, 101, 102]); accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[0.0, 10.0, 100.0]); accessor.lists.get_mut(&20).unwrap().vectors = matrix(&[1.0, 2.0]); - let mut region = gather_region::( + let mut region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10, 20], - &[], - matrix(&[]).as_view(), - &HashMap::new(), + &[ + ListInsertion { + list: 10, + staged: Vec::new(), + }, + ListInsertion { + list: 20, + staged: Vec::new(), + }, + ], ) .await .unwrap(); @@ -768,7 +704,7 @@ mod tests { ); assert_eq!(region.centroids.row(assigned[0]), &[1.0]); let update = Deltas::new( - build_deltas::(&mut accessor, ®ion, &assigned) + IVFIndex::::build_deltas::<_, f32>(&mut accessor, ®ion, &assigned) .await .unwrap(), ); @@ -779,19 +715,14 @@ mod tests { } #[tokio::test] - async fn planner_moves_all_parent_members_and_appends_every_staged_point() { + async fn regional_deltas_move_parent_members_and_append_only_regional_staged_points() { let mut accessor = TestAccessor::new(); - let update = require_send(plan_split_update::( - &mut accessor, - 10, - &mut StdRng::seed_from_u64(7), - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), - &[10, 20], - )) - .await - .unwrap(); + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let update = Deltas::new( + require_send(split_update(&mut accessor, &inputs(&vectors, &[10, 20]))) + .await + .unwrap(), + ); assert_eq!(accessor.staged.len(), 4); assert_eq!(*accessor.reads.lock().unwrap(), [10, 20]); for (parent, ids) in [(10, [100, 101]), (20, [200, 201])] { @@ -800,10 +731,10 @@ mod tests { assert!(moves.iter().all(|(_, to)| *to >= 1000)); } let appended = appends(&update); - assert_eq!(appended[&30], [502]); + assert!(!appended.contains_key(&30)); let mut ids: Vec<_> = appended.values().flatten().copied().collect(); ids.sort_unstable(); - assert_eq!(ids, [500, 501, 502, 503]); + assert_eq!(ids, [500, 501, 503]); let destination = appended.values().find(|ids| ids.contains(&503)).unwrap(); assert_eq!(destination, &[503]); } @@ -826,12 +757,10 @@ mod tests { #[tokio::test] async fn custom_split_supports_three_children_and_neighbor_reassignment() { let mut accessor = TestAccessor::new(); - let mut region = gather_region::( + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let mut region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10, 30], - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), + &inputs(&vectors, &[10, 30]), ) .await .unwrap(); @@ -840,7 +769,7 @@ mod tests { .unwrap(); let assigned = NearestCentroid.assign(®ion).unwrap(); let update = Deltas::new( - build_deltas::(&mut accessor, ®ion, &assigned) + IVFIndex::::build_deltas::<_, f32>(&mut accessor, ®ion, &assigned) .await .unwrap(), ); @@ -868,12 +797,10 @@ mod tests { #[tokio::test] async fn surviving_neighbor_members_and_staged_points_can_stay_put() { let mut accessor = TestAccessor::new(); - let mut region = gather_region::( + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let mut region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10, 30], - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), + &inputs(&vectors, &[10, 30]), ) .await .unwrap(); @@ -882,7 +809,7 @@ mod tests { .unwrap(); let assigned = NearestCentroid.assign(®ion).unwrap(); let update = Deltas::new( - build_deltas::(&mut accessor, ®ion, &assigned) + IVFIndex::::build_deltas::<_, f32>(&mut accessor, ®ion, &assigned) .await .unwrap(), ); @@ -905,12 +832,10 @@ mod tests { } } let mut accessor = TestAccessor::new(); - let mut region = gather_region::( + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let mut region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10, 30], - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), + &inputs(&vectors, &[10, 30]), ) .await .unwrap(); @@ -918,7 +843,7 @@ mod tests { .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) .unwrap(); let update = Deltas::new( - build_deltas::( + IVFIndex::::build_deltas::<_, f32>( &mut accessor, ®ion, &LastCentroid.assign(®ion).unwrap(), @@ -931,43 +856,25 @@ mod tests { } #[tokio::test] - async fn no_splits_and_empty_batches_skip_reads_and_centroid_staging() { - for (ids, vectors, routed) in [ - ( - vec![500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]), - routes(), - ), - (Vec::new(), matrix(&[]), HashMap::new()), - ] { - let mut accessor = TestAccessor::new(); - let update = plan_split_update::( - &mut accessor, - 10, - &mut StdRng::seed_from_u64(7), - &ids, - vectors.as_view(), - &routed, - &[], - ) - .await - .unwrap(); - assert!(accessor.reads.lock().unwrap().is_empty()); - assert_eq!(accessor.stage_calls.get(), 0); - assert_eq!(update.deltas().len(), routed.len()); - if !ids.is_empty() { - assert_eq!( - appends(&update), - [(10, vec![503, 500]), (20, vec![501]), (30, vec![502]),] - .into_iter() - .collect() - ); - } - } + async fn empty_region_inputs_skip_reads_and_centroid_staging() { + let mut accessor = TestAccessor::new(); + let region = require_send(IVFIndex::::gather_region::<_, f32>( + &mut accessor, + &[], + )) + .await + .unwrap(); + assert!(region.lists.is_empty()); + assert!(region.point_ids.is_empty()); + assert_eq!(region.points.nrows(), 0); + assert_eq!(region.centroids.nrows(), 0); + assert!(accessor.reads.lock().unwrap().is_empty()); + assert_eq!(accessor.stage_calls.get(), 0); } #[tokio::test] async fn gathering_errors_propagate_before_centroid_staging() { + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); for failure in ["members", "reader", "centroid"] { let mut accessor = TestAccessor::new(); match failure { @@ -982,17 +889,9 @@ mod tests { } _ => unreachable!(), } - let err = plan_split_update::( - &mut accessor, - 10, - &mut StdRng::seed_from_u64(7), - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), - &[10, 20], - ) - .await - .unwrap_err(); + let err = split_update(&mut accessor, &inputs(&vectors, &[10, 20])) + .await + .unwrap_err(); let expected = match failure { "members" => "fixture list 20 is missing", "reader" => "fixture read failed", @@ -1009,26 +908,24 @@ mod tests { let mut accessor = TestAccessor::new(); accessor.lists.get_mut(&10).unwrap().members = Box::new([100]); accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[0.0]); - let err = plan_split_update::( + let err = split_update( &mut accessor, - 10, - &mut StdRng::seed_from_u64(7), - &[], - matrix(&[]).as_view(), - &HashMap::new(), - &[10], + &[ListInsertion { + list: 10, + staged: Vec::new(), + }], ) .await .unwrap_err(); assert!(err.to_string().contains("at least two points"), "{err}"); assert_eq!(accessor.stage_calls.get(), 0); - let mut region = gather_region::( + let mut region = IVFIndex::::gather_region::<_, f32>( &mut accessor, - &[10], - &[], - matrix(&[]).as_view(), - &HashMap::new(), + &[ListInsertion { + list: 10, + staged: Vec::new(), + }], ) .await .unwrap(); @@ -1038,24 +935,17 @@ mod tests { assert!(err.to_string().contains("at least one centroid"), "{err}"); accessor.fail_stage = Some(1); - let err = plan_split_update::( - &mut accessor, - 10, - &mut StdRng::seed_from_u64(7), - &[500, 501, 502, 503], - matrix(&[1.0, 11.0, 100.0, 3.0]).as_view(), - &routes(), - &[10, 20], - ) - .await - .unwrap_err(); + let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); + let err = split_update(&mut accessor, &inputs(&vectors, &[10, 20])) + .await + .unwrap_err(); assert!(err.to_string().contains("fixture staging failed"), "{err}"); assert_eq!(accessor.staged.len(), 1); } #[tokio::test] async fn insert_batch_applies_splits_and_unsplit_appends_together() { - for threshold in [2, 100] { + for threshold in [2, 3, 100] { let accessor = TestAccessor::new(); let applied = Arc::clone(&accessor.applied); let reads = Arc::clone(&accessor.reads); @@ -1082,22 +972,55 @@ mod tests { let appended = appends(&updates[0]); let mut ids: Vec<_> = appended.values().flatten().copied().collect(); ids.sort_unstable(); - assert_eq!(ids, [500, 501, 502, 503]); - assert_eq!(appended[&30], [502]); - if threshold == 2 { + assert_eq!(ids, [10500, 10501, 10502, 10503]); + assert_eq!(appended[&30], [10502]); + if threshold < 100 { assert_eq!(retired(&updates[0], 10).len(), 2); - assert_eq!(retired(&updates[0], 20).len(), 2); let mut read_ids = reads.lock().unwrap().clone(); read_ids.sort_unstable(); - assert_eq!(read_ids, [10, 20]); + if threshold == 2 { + assert_eq!(retired(&updates[0], 20).len(), 2); + assert_eq!(read_ids, [10, 20]); + } else { + assert_eq!(appended[&20], [10501]); + assert_eq!(read_ids, [10]); + } } else { - assert_eq!(appended[&10], [500, 503]); - assert_eq!(appended[&20], [501]); + assert_eq!(appended[&10], [10500, 10503]); + assert_eq!(appended[&20], [10501]); assert!(reads.lock().unwrap().is_empty()); + assert_eq!(updates[0].deltas().len(), 3); } } } + #[tokio::test] + async fn regional_failure_does_not_commit_already_planned_ordinary_appends() { + let mut accessor = TestAccessor::new(); + accessor.fail_read = Some(10); + let applied = Arc::clone(&accessor.applied); + let strategy = TestStrategy(Mutex::new(Some(accessor))); + let mut index = IVFIndex::new( + TestProvider, + Config { + split_threshold: 3, + reassign_neighbors: 1, + two_means_iterations: 10, + seed: 7, + }, + ) + .unwrap(); + let err = require_send(index.insert_batch( + &strategy, + &DefaultContext, + &[(500, 1.0), (501, 11.0), (502, 100.0), (503, 3.0)], + )) + .await + .unwrap_err(); + assert!(err.to_string().contains("fixture read failed"), "{err}"); + assert!(applied.lock().unwrap().is_empty()); + } + #[tokio::test] async fn empty_insert_does_not_construct_an_accessor_or_apply_an_update() { let strategy = TestStrategy(Mutex::new(None)); diff --git a/diskann/src/ivf/dynamic.rs b/diskann/src/ivf/traits.rs similarity index 100% rename from diskann/src/ivf/dynamic.rs rename to diskann/src/ivf/traits.rs From 83c01093a3a41705791a8b12803f3340967dd89b Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Fri, 9 Oct 2026 10:24:51 -0400 Subject: [PATCH 14/33] another approach --- diskann/src/ivf/clustering.rs | 110 +++ diskann/src/ivf/index.rs | 388 ++-------- diskann/src/ivf/mod.rs | 25 +- diskann/src/ivf/online/clustering.rs | 226 ------ diskann/src/ivf/online/kernels.rs | 132 ---- diskann/src/ivf/online/mod.rs | 24 - diskann/src/ivf/pending.rs | 106 +++ diskann/src/ivf/split.rs | 1042 -------------------------- diskann/src/ivf/tests.rs | 713 ++++++++++++++++++ diskann/src/ivf/traits.rs | 284 ++----- diskann/src/ivf/update.rs | 97 --- diskann/src/ivf/workingset.rs | 77 ++ rfcs/01187-incremental-ivf.md | 83 +- 13 files changed, 1239 insertions(+), 2068 deletions(-) create mode 100644 diskann/src/ivf/clustering.rs delete mode 100644 diskann/src/ivf/online/clustering.rs delete mode 100644 diskann/src/ivf/online/kernels.rs delete mode 100644 diskann/src/ivf/online/mod.rs create mode 100644 diskann/src/ivf/pending.rs delete mode 100644 diskann/src/ivf/split.rs create mode 100644 diskann/src/ivf/tests.rs delete mode 100644 diskann/src/ivf/update.rs create mode 100644 diskann/src/ivf/workingset.rs diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs new file mode 100644 index 0000000000..fee3ce0731 --- /dev/null +++ b/diskann/src/ivf/clustering.rs @@ -0,0 +1,110 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Two-means fitting and region-wide assignment over borrowed point rows. + +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; +use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; +use rand::{Rng, rngs::StdRng}; + +use super::workingset::{ListView, PointRef}; +use crate::{ANNError, ANNResult}; + +pub(super) struct Children { + pub centroids: rowmajor::Owned, + pub members: Vec>, +} + +/// Fit two children per source, then assign every point across all children. +pub(super) fn split( + lists: &[ListView<'_, I, L>], + dim: usize, + iterations: usize, + rng: &mut StdRng, +) -> ANNResult> { + if lists.is_empty() || dim == 0 { + return Err(ANNError::message( + "IVF splitting needs a nonempty working set and positive dimension", + )); + } + let mut centroids = rowmajor::Owned::try_from_element(2 * lists.len(), dim, 0.0)?; + let mut sums = vec![0.0; 2 * dim]; + for (list, centers) in lists + .iter() + .zip(centroids.as_mut_slice().chunks_exact_mut(2 * dim)) + { + fit_two_means(&list.points, centers, dim, iterations, rng, &mut sums)?; + } + + let count = lists.iter().map(|list| list.points.len()).sum(); + let mut assignments = Vec::with_capacity(count); + let mut sizes = vec![0; centroids.nrows()]; + for point in lists.iter().flat_map(|list| &list.points) { + let to = nearest(point.vector, centroids.as_slice(), dim); + assignments.push(to); + sizes[to] += 1; + } + let mut members: Vec> = sizes.into_iter().map(Vec::with_capacity).collect(); + for (point, to) in lists.iter().flat_map(|list| &list.points).zip(assignments) { + members[to].push(point.id); + } + Ok(Children { centroids, members }) +} + +fn fit_two_means( + points: &[PointRef<'_, I>], + centers: &mut [f32], + dim: usize, + iterations: usize, + rng: &mut StdRng, + sums: &mut [f64], +) -> ANNResult<()> { + if points.len() < 2 { + return Err(ANNError::message("an IVF split needs at least two points")); + } + let first = rng.random_range(0..points.len()); + let second = (first + rng.random_range(1..points.len())) % points.len(); + centers[..dim].copy_from_slice(points[first].vector); + centers[dim..].copy_from_slice(points[second].vector); + + for _ in 0..iterations.max(1) { + sums.fill(0.0); + let mut counts = [0; 2]; + for point in points { + let to = nearest(point.vector, centers, dim); + counts[to] += 1; + for (sum, &value) in sums[to * dim..(to + 1) * dim].iter_mut().zip(point.vector) { + // A finite mean must not overflow while accumulating its inputs. + *sum += f64::from(value); + } + } + for ((center, sum), count) in centers + .chunks_exact_mut(dim) + .zip(sums.chunks_exact(dim)) + .zip(counts) + { + if count != 0 { + for (value, &sum) in center.iter_mut().zip(sum) { + *value = (sum / count as f64) as f32; + } + } + } + } + Ok(()) +} + +/// Prefer the first centroid on ties. `centers` contains at least one full row. +fn nearest(point: &[f32], centers: &[f32], dim: usize) -> usize { + let mut best = 0; + let mut best_distance: f32 = SquaredL2::evaluate(point, ¢ers[..dim]); + for (index, center) in centers.chunks_exact(dim).enumerate().skip(1) { + let distance: f32 = SquaredL2::evaluate(point, center); + if distance.total_cmp(&best_distance).is_lt() { + best = index; + best_distance = distance; + } + } + best +} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index d7c9ac03e7..b096c02903 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -3,87 +3,47 @@ * Licensed under the MIT license. */ -//! Dynamic IVF index orchestration. -//! -//! The index drives one [`InsertAccessor`] per mutation: it plans a -//! [`Deltas`](super::update::Deltas) against the accessor's view and hands it to -//! [`InsertAccessor::update`]. Gathering and update construction use only the accessor; -//! splitting and assignment operate on an in-memory region. +//! IVF policy and orchestration, independent of physical list storage. + +use std::collections::BTreeMap; use diskann_utils::{ future::SendFuture, - views::rowmajor::{self, Matrix, MatrixMut}, + views::rowmajor::{self, Matrix}, }; use rand::{SeedableRng, rngs::StdRng}; use thiserror::Error; -use crate::{ - ANNError, ANNResult, - error::{ErrorExt, IntoANNResult}, - ivf::{ - traits::{Centroids, InsertAccessor, MaintenanceStrategy, Provider, Reader, Stage}, - online::index_error, - split::{Assign, NearestCentroid, Region, SourceList, Split, TwoMeansSplit}, - update::{self, CentroidDelta, Delta, Deltas, MoveTo}, +use super::{ + clustering, + traits::{ + BuildAccessor, BuildStrategy, CoarseAccessor, InsertStrategy, ListAccessor, Provider, }, }; +use crate::{ANNError, ANNResult, error::IntoANNResult}; -/// Parameters for the online split policy. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Parameters for incremental two-means splitting. +#[derive(Debug, Clone, Copy)] pub struct Config { - /// Split a list once an insert batch would grow it beyond this many points. + /// Split a touched list when its pending membership exceeds this size. pub split_threshold: usize, - /// Nearby lists reassigned together with each split parent. - pub reassign_neighbors: usize, - /// Lloyd iterations used to fit split children. + /// Lloyd steps per parent. Zero still performs one step. pub two_means_iterations: usize, - /// Seed for split-child initialization. pub seed: u64, } -/// Invalid [`DynamicIvfConfig`] values. -#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)] -pub enum ConfigError { - #[error("split_threshold must be at least 2, got {0}")] - SplitThreshold(usize), - #[error("reassign_neighbors must be at least 1")] - ReassignNeighbors, -} - -impl From for ANNError { - #[track_caller] - fn from(err: ConfigError) -> Self { - ANNError::new(err) - } -} +#[derive(Debug, Error)] +#[error("split_threshold must be at least 2, got {0}")] +pub struct ConfigError(pub usize); -/// Work performed by one [`DynamicIvfIndex::insert_batch`]. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +/// Work successfully published by one insertion. +#[derive(Debug, Default, PartialEq, Eq)] pub struct InsertStats { - /// Points added to the index. pub inserted: usize, - /// Lists split into two children. pub splits: usize, - /// Previously indexed points moved by regional reassignment. - pub reassigned: usize, -} - -/// A staged point and its borrowed canonical vector, ready for regional gathering. -pub(super) struct StagedPoint<'a, I> { - pub(super) id: I, - pub(super) vector: &'a [f32], } -/// Staged points routed to one existing list, in insertion order. -pub(super) struct ListInsertion<'a, I, L> { - pub(super) list: L, - pub(super) staged: Vec>, -} - -/// An incrementally maintained IVF index. -/// -/// Mutations take `&mut self` and lend the provider exclusively to one -/// [`MaintenanceAccessor`], so they cannot overlap searches through the same value. +/// An IVF index with serialized mutation through the underlying provider. #[derive(Debug)] pub struct IVFIndex { provider: P, @@ -92,17 +52,12 @@ pub struct IVFIndex { } impl IVFIndex

{ - /// Construct an index over `provider`. - /// /// # Errors /// - /// Returns [`ConfigError`] if `config` is invalid. + /// Returns [`ConfigError`] when the split threshold is less than two. pub fn new(provider: P, config: Config) -> Result { if config.split_threshold < 2 { - return Err(ConfigError::SplitThreshold(config.split_threshold)); - } - if config.reassign_neighbors == 0 { - return Err(ConfigError::ReassignNeighbors); + return Err(ConfigError(config.split_threshold)); } Ok(Self { provider, @@ -111,70 +66,58 @@ impl IVFIndex

{ }) } - /// Borrow the underlying provider. pub fn provider(&self) -> &P { &self.provider } - /// Install the initial centroids. + /// Install initial empty lists and their centroids in one build operation. /// /// # Errors /// - /// Fails if `centroids` is empty, non-finite, or of the wrong dimension, the index - /// already has centroids, or the maintenance accessor fails. - pub fn initialize<'a, S, T>( + /// Fails for an already initialized index, or a provider error. + /// Failed publication is provider-defined. + pub fn initialize<'a, S>( &'a mut self, strategy: &'a S, context: &'a P::Context, centroids: rowmajor::Ref<'_, f32>, ) -> impl SendFuture> where - S: MaintenanceStrategy<'a, P, T>, - T: Sync, + S: BuildStrategy<'a, P>, { - let provider = &mut self.provider; - async move { - let count = centroids.nrows(); - if count == 0 { - return Err(ANNError::message("Cannot initialize with zero centroids.")); + if centroids.nrows() == 0 || centroids.ncols() == 0 { + return Err(ANNError::message( + "IVF initialization needs nonempty centroids", + )); } - let mut accessor = strategy - .maintenance_accessor(provider, context) + .build_accessor(&mut self.provider, context) .into_ann_result()?; - let mut installs = Vec::with_capacity(centroids.nrows()); - - for c in centroids.rows() { - let id = accessor.stage_centroid(c).await?; - installs.push(update::Delta::CentroidDelta { - id, - delta: update::CentroidDelta::Install, - }); + if !accessor.coarse().is_empty() { + return Err(ANNError::message("the IVF index is already initialized")); } - accessor - .update(Deltas::::new(installs)) - .await - .escalate("initialize must apply the initial centroids")?; - - Ok(()) + for centroid in centroids.rows() { + // TO DO: Need to probably introduce an initialization strategy + accessor.lists().create_list(centroid, Vec::new()).await?; + } + accessor.finish().await } } - /// Insert a batch, splitting every list the batch pushes past `split_threshold`. + /// Route and provisionally append a batch, split oversized lists, then publish. /// - /// All split parents form one region. Their points are assigned across all fitted - /// children, never back to retiring parents. Other staged points keep their routes. - /// The batch and its splits are then applied as one - /// [`Deltas`](super::update::Deltas). + /// All split parents form one working set. Their members choose among every + /// fitted child, never among retiring parent centroids. Other appends keep their + /// routes. Point coordinates are borrowed throughout fitting and assignment. /// /// # Errors /// - /// Fails if the index is not initialized, the provider returns inconsistent - /// data, or any accessor call fails. State after an `update` failure is - /// provider-defined. + /// Fails before publication if the index is uninitialized or preparation, + /// routing, filling, fitting, or list mutation fails. State after a failed + /// `finish` is provider-defined. An empty batch is a no-op. pub fn insert_batch<'a, S, T>( &'a mut self, strategy: &'a S, @@ -182,8 +125,8 @@ impl IVFIndex

{ points: &'a [(P::ExternalId, T)], ) -> impl SendFuture> where - S: MaintenanceStrategy<'a, P, T>, - T: Sync + Copy, + S: InsertStrategy<'a, P, T>, + T: Sync, { let provider = &mut self.provider; let config = self.config; @@ -193,220 +136,55 @@ impl IVFIndex

{ if points.is_empty() { return Ok(InsertStats::default()); } - let mut accessor = strategy - .maintenance_accessor(provider, context) + .build_accessor(provider, context) .into_ann_result()?; - - let mut vectors = rowmajor::Owned::from_element(points.len(), accessor.dim(), f32::NAN); - let mut routed = - hashbrown::HashMap::>>::new(); - - for ((id, v), out) in points.iter().zip(vectors.rows_mut()) { - let id = accessor.stage_point(id, *v, out).await.into_ann_result()?; - - let parent = accessor - .centroids() - .select(out, 1) - .escalate("insert must route every staged point")? - .first() - .map(|s| s.id) - .ok_or_else(|| ANNError::message("Didn't return list"))?; - - routed - .entry(parent) - .or_default() - .push(StagedPoint { id, vector: out }); + if accessor.coarse().is_empty() { + return Err(ANNError::message("the IVF index is not initialized")); } + let batch = strategy.prepare(&mut accessor, points).await?; - let mut split_inputs = Vec::with_capacity(routed.len()); - let mut deltas = Vec::with_capacity(routed.len()); - for (list, staged) in routed { - let len = accessor.list_size(list).into_ann_result()?; + let mut routed = BTreeMap::>::new(); - if len + staged.len() > config.split_threshold { - split_inputs.push(ListInsertion { list, staged }); - } else { - deltas.push(update::Delta::PointAppends { - to: list, - ids: staged.into_iter().map(|point| point.id).collect(), - }); - } + for point in batch.points() { + let selected = accessor.coarse().select(point.vector, 1).await?; + let list = selected + .first() + .ok_or_else(|| ANNError::message("IVF routing returned no live centroid"))? + .id; + routed.entry(list).or_default().push(point.id); } - if !split_inputs.is_empty() { - let mut region = Self::gather_region::<_, T>(&mut accessor, &split_inputs).await?; - let parents: Vec<_> = split_inputs.iter().map(|input| input.list).collect(); - TwoMeansSplit { - iterations: config.two_means_iterations, + // Stable list order also makes seeded fitting independent of hash order. + let mut parents = Vec::with_capacity(routed.len()); + for (list, members) in routed { + let mut lists = accessor.lists(); + lists.append_members(list, members).await?; + if lists.list_size(list).await? > config.split_threshold { + parents.push(list); } - .split(&mut region, &parents, rng)?; - let assignments = NearestCentroid.assign(®ion)?; - deltas.extend( - Self::build_deltas::<_, T>(&mut accessor, ®ion, &assignments).await?, - ); - } - - accessor - .update(Deltas::new(deltas)) - .await - .escalate("insert must apply its update")?; - - Ok(InsertStats::default()) - } - } - - /// Gather committed members and staged points into their final region rows. - pub(super) async fn gather_region( - accessor: &mut A, - inputs: &[ListInsertion<'_, P::InternalId, P::ListId>], - ) -> ANNResult> - where - A: InsertAccessor, - T: Sync, - { - let mut lists = Vec::with_capacity(inputs.len()); - let mut rows = 0; - for input in inputs { - let id = input.list; - let member_count = accessor.get_members(id).into_ann_result()?.len(); - let end = rows + member_count + input.staged.len(); - lists.push(SourceList { - id, - rows: rows..end, - member_count, - }); - rows = end; - } - - let dim = accessor.dim(); - let mut point_ids = Vec::with_capacity(rows); - let mut points = rowmajor::Owned::try_from_element(rows, dim, 0.0)?; - let mut centroids = rowmajor::Owned::try_from_element(inputs.len(), dim, 0.0)?; - let mut centroid_ids = Vec::with_capacity(inputs.len()); - - for (index, (list, input)) in lists.iter().zip(inputs).enumerate() { - point_ids.extend_from_slice(accessor.get_members(list.id).into_ann_result()?); - - let member_end = list.rows.start + list.member_count; - accessor - .reader() - .read_into( - list.id, - points - .subview_mut(list.rows.start..member_end) - .ok_or_else(|| { - index_error("member rows lie outside the region's point matrix") - })?, - ) - .await - .escalate("split must read the requested list's vectors")?; - let mut tail = points - .subview_mut(member_end..list.rows.end) - .ok_or_else(|| index_error("staged rows lie outside the region's point matrix"))?; - for (point, out) in input.staged.iter().zip(tail.rows_mut()) { - point_ids.push(point.id); - out.copy_from_slice(point.vector); } - - let catalog = accessor.centroids(); - let centroid = catalog - .centroid(list.id) - .ok_or_else(|| index_error(format!("centroid {} is unavailable", list.id)))?; - centroids.row_mut(index).copy_from_slice(¢roid); - centroid_ids.push(Some(list.id)); - } - - Ok(Region { - lists, - point_ids, - points, - centroid_ids, - centroids, - }) - } - - /// Stage proposed centroids and translate aligned assignments into partition deltas. - pub(super) async fn build_deltas( - accessor: &mut A, - region: &Region, - assignments: &[usize], - ) -> ANNResult>> - where - A: InsertAccessor, - T: Sync, - { - let mut deltas = Vec::with_capacity(2 * region.centroid_ids.len() + region.lists.len()); - let mut destinations = Vec::with_capacity(region.centroid_ids.len()); - for (id, centroid) in region.centroid_ids.iter().zip(region.centroids.rows()) { - let id = match id { - Some(id) => *id, - None => { - let id = accessor.stage_centroid(centroid).await?; - deltas.push(Delta::CentroidDelta { - id, - delta: CentroidDelta::Install, - }); - id + let splits = parents.len(); + if !parents.is_empty() { + let children = { + let dim = accessor.dim(); + let view = accessor.fill(&parents).await?; + clustering::split(&view, dim, config.two_means_iterations, rng)? + }; + for (centroid, members) in children.centroids.rows().zip(children.members) { + accessor.lists().create_list(centroid, members).await?; } - }; - destinations.push(id); - } - - let mut append_counts = vec![0; destinations.len()]; - - for (index, list) in region.lists.iter().enumerate() { - let retired = !region.centroid_ids.contains(&Some(list.id)); - let mut moves = Vec::with_capacity(list.member_count); - for (&id, &to) in region - .member_ids(index) - .iter() - .zip(&assignments[list.rows.start..list.rows.start + list.member_count]) - { - let to = destinations[to]; - if to != list.id { - moves.push(MoveTo::new(id, to)); + for parent in parents { + accessor.lists().retire_list(parent).await?; } } - if retired { - deltas.push(Delta::CentroidDelta { - id: list.id, - delta: CentroidDelta::Retire { - moves: moves.into_boxed_slice(), - }, - }); - } else if !moves.is_empty() { - deltas.push(Delta::PointMoves { - from: list.id, - moves: moves.into_boxed_slice(), - }); - } - for &to in &assignments[list.rows.start + list.member_count..list.rows.end] { - append_counts[to] += 1; - } - } - let mut appends: Vec<_> = append_counts.into_iter().map(Vec::with_capacity).collect(); - for (index, list) in region.lists.iter().enumerate() { - for (&id, &to) in region - .staged_ids(index) - .iter() - .zip(&assignments[list.rows.start + list.member_count..list.rows.end]) - { - appends[to].push(id); - } - } - for (to, ids) in destinations.into_iter().zip(appends) { - if !ids.is_empty() { - deltas.push(Delta::PointAppends { - to, - ids: ids.into_boxed_slice(), - }); - } + accessor.finish().await?; + + Ok(InsertStats { + inserted: points.len(), + splits, + }) } - Ok(deltas) } } - -#[cfg(test)] -mod tests {} diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 0316bb3258..7918b5af53 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -3,12 +3,27 @@ * Licensed under the MIT license. */ -//! Incrementally maintained IVF index. +//! Incremental IVF construction. +//! +//! The index routes points, splits oversized lists, and assigns their members to +//! replacement centroids. Providers supply centroid selection, logical list +//! operations, and borrowed working sets through [`traits::BuildAccessor`]. +//! +//! Appends and replacements remain private to one build operation until it +//! finishes. This lets clustering treat prepared and stored points identically. +//! Providers choose their own storage, publication, and failure-recovery mechanisms. +//! +//! This first implementation supports initialization and batch insertion with +//! two-means splitting. Search, neighborhood reassignment, deletion, and colocation +//! are not implemented. -pub mod traits; +mod clustering; mod index; -mod online; -mod split; -pub mod update; +pub mod pending; +pub mod traits; +pub mod workingset; pub use index::{Config, ConfigError, IVFIndex, InsertStats}; + +#[cfg(test)] +mod tests; diff --git a/diskann/src/ivf/online/clustering.rs b/diskann/src/ivf/online/clustering.rs deleted file mode 100644 index 4b1269a2dd..0000000000 --- a/diskann/src/ivf/online/clustering.rs +++ /dev/null @@ -1,226 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Numerical fitting and assignment over borrowed matrices. - -use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; -use rand::{Rng, rngs::StdRng}; - -use super::{ - index_error, - kernels::{LloydScratch, lloyd, nearest}, -}; -use crate::ANNResult; - -/// Fit two children directly into `centroids`, using distinct random seeds. -/// -/// `centroids` must have two rows and the same dimension as `points`. -/// Zero iterations still performs one Lloyd step. -/// -/// # Errors -/// -/// Returns an error when fewer than two training points are supplied. -pub(in crate::ivf) fn fit_two_means( - points: rowmajor::Ref<'_, f32>, - mut centroids: rowmajor::Mut<'_, f32>, - iterations: usize, - rng: &mut StdRng, - scratch: &mut LloydScratch, -) -> ANNResult<()> { - let count = points.nrows(); - if count < 2 { - return Err(index_error(format!( - "a split needs at least two points, got {count}" - ))); - } - - let first = rng.random_range(0..count); - let second = (first + rng.random_range(1..count)) % count; - centroids.row_mut(0).copy_from_slice(points.row(first)); - centroids.row_mut(1).copy_from_slice(points.row(second)); - lloyd( - || points.rows(), - centroids.as_mut_slice(), - points.ncols(), - iterations.max(1), - scratch, - ); - Ok(()) -} - -/// Assign by squared Euclidean distance, preferring the earliest centroid on ties. -/// -/// # Errors -/// -/// Returns an error when there are no centroids or the dimensions differ. -pub(in crate::ivf) fn assign_nearest( - points: rowmajor::Ref<'_, f32>, - centroids: rowmajor::Ref<'_, f32>, -) -> ANNResult> { - if centroids.nrows() == 0 { - return Err(index_error("assignment needs at least one centroid")); - } - if points.ncols() != centroids.ncols() { - return Err(index_error(format!( - "assignment points have dimension {}, but centroids have dimension {}", - points.ncols(), - centroids.ncols() - ))); - } - - let candidates: Vec<_> = centroids.rows().enumerate().collect(); - let mut assigned = Vec::with_capacity(points.nrows()); - for point in points.rows() { - assigned.push( - nearest(point, &candidates) - .ok_or_else(|| index_error("assignment needs at least one centroid"))?, - ); - } - Ok(assigned.into_boxed_slice()) -} - -#[cfg(test)] -mod tests { - use rand::SeedableRng; - - use 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() - } - - fn fit( - points: rowmajor::Ref<'_, f32>, - iterations: usize, - rng: &mut StdRng, - ) -> ANNResult> { - let mut centroids = rowmajor::Owned::from_element(2, points.ncols(), 0.0); - fit_two_means( - points, - centroids.as_view_mut(), - iterations, - rng, - &mut LloydScratch::default(), - )?; - Ok(centroids) - } - - #[test] - fn two_means_separates_distant_clusters() { - let points = matrix(&[[0.0, 0.0], [10.0, 0.0], [0.0, 1.0], [10.0, 1.0], [1.0, 0.0]]); - let centroids = fit(points.as_view(), 10, &mut StdRng::seed_from_u64(7)).unwrap(); - let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); - - let low = children[0]; - assert_eq!(children.as_ref(), &[low, 1 - low, low, 1 - low, low]); - assert_eq!(centroids.row(low), &[1.0 / 3.0, 1.0 / 3.0]); - assert_eq!(centroids.row(1 - low), &[10.0, 0.5]); - } - - #[test] - fn two_means_needs_two_points() { - for count in [0, 1] { - let points = rowmajor::Owned::from_element(count, 1, 1.0f32); - let err = fit(points.as_view(), 1, &mut StdRng::seed_from_u64(7)).unwrap_err(); - assert!(err.to_string().contains("at least two points"), "{err}"); - } - } - - #[test] - fn two_means_preserves_seed_order_and_rng_consumption() { - let points = matrix(&[[0.0], [10.0]]); - let mut rng = StdRng::seed_from_u64(7); - let mut expected_rng = rng.clone(); - let first = expected_rng.random_range(0..points.nrows()); - let second = (first + expected_rng.random_range(1..points.nrows())) % points.nrows(); - let centroids = fit(points.as_view(), 10, &mut rng).unwrap(); - - assert_eq!(centroids.row(0), points.row(first)); - assert_eq!(centroids.row(1), points.row(second)); - assert_eq!(rng.random::(), expected_rng.random::()); - } - - #[test] - fn two_means_zero_iterations_still_refines_centroids() { - let points = matrix(&[[0.0], [2.0], [10.0], [12.0]]); - let zero = fit(points.as_view(), 0, &mut StdRng::seed_from_u64(7)).unwrap(); - let one = fit(points.as_view(), 1, &mut StdRng::seed_from_u64(7)).unwrap(); - assert_eq!(zero.as_slice(), one.as_slice()); - } - - #[test] - fn two_means_writes_only_the_output_subview_and_reuses_scratch() { - let points = matrix(&[[0.0], [10.0]]); - let mut output = rowmajor::Owned::from_element(6, 1, -1.0); - let mut scratch = LloydScratch::default(); - let mut rng = StdRng::seed_from_u64(7); - for rows in [1..3, 3..5] { - let out = rowmajor::Mut::try_from_data(&mut output.as_mut_slice()[rows], 2, 1).unwrap(); - fit_two_means(points.as_view(), out, 10, &mut rng, &mut scratch).unwrap(); - } - assert_eq!(output.row(0), &[-1.0]); - assert_eq!(output.row(5), &[-1.0]); - for rows in [1..3, 3..5] { - let assigned = assign_nearest(points.as_view(), output.subview(rows).unwrap()).unwrap(); - assert_ne!(assigned[0], assigned[1]); - } - } - - #[test] - fn nearest_centroid_assigns_points_to_an_independent_candidate_set() { - let points = matrix(&[[-1.0], [12.0], [19.0]]); - let centroids = matrix(&[[0.0], [10.0], [20.0]]); - let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); - assert_eq!(children.as_ref(), &[0, 1, 2]); - } - - #[test] - fn nearest_centroid_prefers_the_earliest_row_on_ties() { - let points = matrix(&[[1.8, 0.1], [1.0, 1.0]]); - let centroids = matrix(&[[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]]); - let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); - assert_eq!(children.as_ref(), &[1, 0]); - } - - #[test] - fn nearest_centroid_supports_one_candidate() { - let points = matrix(&[[-1.0], [12.0], [19.0]]); - let centroids = matrix(&[[10.0]]); - let children = assign_nearest(points.as_view(), centroids.as_view()).unwrap(); - assert_eq!(children.as_ref(), &[0, 0, 0]); - } - - #[test] - fn nearest_centroid_supports_an_empty_point_batch() { - let points = matrix::<1>(&[]); - let centroids = matrix(&[[10.0]]); - assert!( - assign_nearest(points.as_view(), centroids.as_view()) - .unwrap() - .is_empty() - ); - } - - #[test] - fn nearest_centroid_needs_a_candidate() { - let points = matrix(&[[1.0]]); - let centroids = matrix::<1>(&[]); - let err = assign_nearest(points.as_view(), centroids.as_view()).unwrap_err(); - assert!(err.to_string().contains("at least one centroid"), "{err}"); - } - - #[test] - fn nearest_centroid_requires_matching_dimensions() { - let points = matrix(&[[1.0]]); - let centroids = matrix(&[[0.0, 0.0]]); - let err = assign_nearest(points.as_view(), centroids.as_view()).unwrap_err(); - assert!( - err.to_string() - .contains("points have dimension 1, but centroids have dimension 2"), - "{err}" - ); - } -} diff --git a/diskann/src/ivf/online/kernels.rs b/diskann/src/ivf/online/kernels.rs deleted file mode 100644 index 54809f6aeb..0000000000 --- a/diskann/src/ivf/online/kernels.rs +++ /dev/null @@ -1,132 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Distance and clustering kernels for split planning. -//! -//! Clustering and placement use squared Euclidean distance regardless of the search -//! metric. - -use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; - -/// Squared Euclidean distance between two vectors of equal length. -pub(super) fn squared_l2(a: &[f32], b: &[f32]) -> f32 { - SquaredL2::evaluate(a, b) -} - -/// The candidate nearest to `point`, preferring the earliest candidate on ties. -pub(super) fn nearest(point: &[f32], candidates: &[(L, &[f32])]) -> Option { - candidates - .iter() - .map(|&(list, centroid)| (list, squared_l2(point, centroid))) - .min_by(|a, b| a.1.total_cmp(&b.1)) - .map(|(list, _)| list) -} - -/// Scratch buffers for [`lloyd`], reusable across calls. -#[derive(Debug, Default)] -pub(in crate::ivf) struct LloydScratch { - sums: Vec, - counts: Vec, -} - -/// Refine `centers`, stored row-major with `dim` columns, with `iterations` Lloyd -/// steps over the points yielded by `points`. -/// -/// Each point joins its nearest center, preferring the earliest on ties. A center that -/// attracts no points keeps its position. -pub(super) fn lloyd<'p, P, I>( - points: P, - centers: &mut [f32], - dim: usize, - iterations: usize, - scratch: &mut LloydScratch, -) where - P: Fn() -> I, - I: Iterator, -{ - if dim == 0 || centers.is_empty() { - return; - } - let LloydScratch { sums, counts } = scratch; - for _ in 0..iterations { - sums.clear(); - sums.resize(centers.len(), 0.0); - counts.clear(); - counts.resize(centers.len() / dim, 0); - - for point in points() { - let Some(center) = centers - .chunks_exact(dim) - .map(|center| squared_l2(point, center)) - .enumerate() - .min_by(|a, b| a.1.total_cmp(&b.1)) - .map(|(center, _)| center) - else { - continue; - }; - counts[center] += 1; - for (sum, x) in sums[center * dim..(center + 1) * dim].iter_mut().zip(point) { - *sum += x; - } - } - - for ((center, sum), &count) in centers - .chunks_exact_mut(dim) - .zip(sums.chunks_exact(dim)) - .zip(counts.iter()) - { - if count > 0 { - let scale = 1.0 / count as f32; - for (x, s) in center.iter_mut().zip(sum) { - *x = s * scale; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn nearest_prefers_the_earliest_tie() { - let (a, b, c) = ([0.0, 0.0], [2.0, 0.0], [0.0, 2.0]); - let candidates = [(7u32, &a[..]), (8, &b[..]), (9, &c[..])]; - assert_eq!(nearest(&[1.8, 0.1], &candidates), Some(8)); - // Equidistant from 8 and 9. - assert_eq!(nearest(&[1.0, 1.0], &candidates[1..]), Some(8)); - assert_eq!(nearest::(&[1.0, 1.0], &[]), None); - } - - #[test] - fn lloyd_moves_centers_to_cluster_means() { - let points: Vec<[f32; 2]> = vec![[0.0, 0.0], [0.0, 2.0], [10.0, 0.0], [10.0, 2.0]]; - let mut centers = [1.0, 1.0, 9.0, 1.0]; - let mut scratch = LloydScratch::default(); - lloyd( - || points.iter().map(|p| &p[..]), - &mut centers, - 2, - 3, - &mut scratch, - ); - assert_eq!(centers, [0.0, 1.0, 10.0, 1.0]); - } - - #[test] - fn lloyd_keeps_centers_without_points() { - let points: Vec<[f32; 1]> = vec![[0.0], [1.0]]; - let mut centers = [0.5, 100.0]; - lloyd( - || points.iter().map(|p| &p[..]), - &mut centers, - 1, - 2, - &mut LloydScratch::default(), - ); - assert_eq!(centers, [0.5, 100.0]); - } -} diff --git a/diskann/src/ivf/online/mod.rs b/diskann/src/ivf/online/mod.rs deleted file mode 100644 index fb30c527b8..0000000000 --- a/diskann/src/ivf/online/mod.rs +++ /dev/null @@ -1,24 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Borrowed-matrix fitting and assignment helpers and squared-Euclidean kernels. - -mod clustering; -mod kernels; - -use std::fmt::{Debug, Display}; - -pub(super) use clustering::{assign_nearest, fit_two_means}; -pub(super) use kernels::LloydScratch; - -use crate::ANNError; - -#[track_caller] -pub(super) fn index_error(message: D) -> ANNError -where - D: Display + Debug + Send + Sync + 'static, -{ - ANNError::message(message) -} diff --git a/diskann/src/ivf/pending.rs b/diskann/src/ivf/pending.rs new file mode 100644 index 0000000000..6c7991821b --- /dev/null +++ b/diskann/src/ivf/pending.rs @@ -0,0 +1,106 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Optional buffering for operation-local list memberships. +//! +//! 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)] +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)] +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/split.rs b/diskann/src/ivf/split.rs deleted file mode 100644 index 15c183c5c7..0000000000 --- a/diskann/src/ivf/split.rs +++ /dev/null @@ -1,1042 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! In-memory regions and strategies for centroid splitting and point assignment. -//! -//! Source membership stays fixed while splitting changes the candidate centroids. -//! Strategies and providers are trusted to preserve row alignment and valid destinations. - -use std::ops::Range; - -use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; -use rand::rngs::StdRng; - -use crate::{ - ANNResult, - ivf::online::{LloydScratch, assign_nearest, fit_two_means, index_error}, - utils::VectorId, -}; - -/// Original membership, independent of the region's eventual destinations. -pub(super) struct SourceList { - pub(super) id: L, - pub(super) rows: Range, - /// Committed members precede staged points within `rows`. - pub(super) member_count: usize, -} - -pub(super) struct Region { - pub(super) lists: Vec>, - pub(super) point_ids: Vec, - pub(super) points: rowmajor::Owned, - /// `None` denotes a proposed centroid whose list ID has not been staged. - pub(super) centroid_ids: Vec>, - pub(super) centroids: rowmajor::Owned, -} - -impl Region { - fn list_points(&self, index: usize) -> ANNResult> { - self.points - .subview(self.lists[index].rows.clone()) - .ok_or_else(|| index_error("source rows lie outside the region's point matrix")) - } - - pub(super) fn member_ids(&self, index: usize) -> &[I] { - let list = &self.lists[index]; - &self.point_ids[list.rows.start..list.rows.start + list.member_count] - } - - pub(super) fn staged_ids(&self, index: usize) -> &[I] { - let list = &self.lists[index]; - &self.point_ids[list.rows.start + list.member_count..list.rows.end] - } -} - -/// Replace parent centroids without changing source point storage or row order. -pub(super) trait Split: Send + Sync { - fn split(&self, region: &mut Region, parents: &[L], rng: &mut StdRng) -> ANNResult<()>; -} - -/// Return one valid candidate-centroid row per point, in the region's row order. -pub(super) trait Assign: Send + Sync { - fn assign(&self, region: &Region) -> ANNResult>; -} - -pub(super) struct TwoMeansSplit { - pub(super) iterations: usize, -} - -impl Split for TwoMeansSplit { - fn split(&self, region: &mut Region, parents: &[L], rng: &mut StdRng) -> ANNResult<()> { - let is_parent = |id: &Option| id.is_some_and(|id| parents.contains(&id)); - let survivors = region - .centroid_ids - .iter() - .filter(|id| !is_parent(id)) - .count(); - let count = 2 * parents.len() + survivors; - let dim = region.points.ncols(); - let mut centroids = rowmajor::Owned::try_from_element(count, dim, 0.0)?; - let mut centroid_ids = Vec::with_capacity(count); - let mut scratch = LloydScratch::default(); - - for (index, parent) in parents.iter().enumerate() { - let source = region - .lists - .iter() - .position(|list| list.id == *parent) - .ok_or_else(|| { - index_error(format!("split parent {parent} is not in the region")) - })?; - let start = 2 * index * dim; - let out = rowmajor::Mut::try_from_data( - &mut centroids.as_mut_slice()[start..start + 2 * dim], - 2, - dim, - )?; - fit_two_means( - region.list_points(source)?, - out, - self.iterations, - rng, - &mut scratch, - )?; - centroid_ids.extend([None, None]); - } - for (id, centroid) in region.centroid_ids.iter().zip(region.centroids.rows()) { - if !is_parent(id) { - centroids - .row_mut(centroid_ids.len()) - .copy_from_slice(centroid); - centroid_ids.push(*id); - } - } - region.centroids = centroids; - region.centroid_ids = centroid_ids; - Ok(()) - } -} - -pub(super) struct NearestCentroid; - -impl Assign for NearestCentroid { - fn assign(&self, region: &Region) -> ANNResult> { - assign_nearest(region.points.as_view(), region.centroids.as_view()) - } -} - -#[cfg(test)] -mod tests { - use std::{ - cell::Cell, - sync::{Arc, Mutex}, - }; - - use hashbrown::HashMap; - use rand::{Rng, SeedableRng}; - - use super::*; - use crate::{ - ANNError, - ivf::{ - Config, IVFIndex, - traits::{ - Centroids, InsertAccessor, MaintenanceStrategy, Provider, Reader, SelectedList, - Stage, - }, - index::{ListInsertion, StagedPoint}, - update::{CentroidDelta, Delta, Deltas}, - }, - provider::DefaultContext, - }; - - struct TestProvider; - - impl Provider for TestProvider { - type Context = DefaultContext; - type InternalId = u32; - type ExternalId = u32; - type ListId = u32; - type Error = ANNError; - - fn to_internal_id(&self, _context: &DefaultContext, id: &u32) -> ANNResult { - Ok(*id) - } - - fn to_external_id(&self, _context: &DefaultContext, id: u32) -> ANNResult { - Ok(id) - } - } - - struct TestList { - members: Box<[u32]>, - vectors: rowmajor::Owned, - centroid: Box<[f32]>, - } - - struct TestAccessor { - lists: HashMap, - reads: Arc>>, - stage_calls: Cell, - staged: Vec<(u32, Box<[f32]>)>, - applied: Arc>>>, - fail_read: Option, - fail_stage: Option, - missing_centroid: Option, - } - - fn matrix(values: &[f32]) -> rowmajor::Owned { - rowmajor::Owned::try_from_data(values.into(), values.len(), 1).unwrap() - } - - fn inputs<'a>( - vectors: &'a rowmajor::Owned, - lists: &[u32], - ) -> Vec> { - lists - .iter() - .map(|&list| { - let points: &[(u32, usize)] = match list { - 10 => &[(503, 3), (500, 0)], - 20 => &[(501, 1)], - 30 => &[(502, 2)], - _ => unreachable!(), - }; - ListInsertion { - list, - staged: points - .iter() - .map(|&(id, row)| StagedPoint { - id, - vector: vectors.row(row), - }) - .collect(), - } - }) - .collect() - } - - impl TestAccessor { - fn new() -> Self { - Self { - lists: [ - ( - 10, - TestList { - members: Box::new([100, 101]), - vectors: matrix(&[0.0, 2.0]), - centroid: Box::new([1.0]), - }, - ), - ( - 20, - TestList { - members: Box::new([200, 201]), - vectors: matrix(&[10.0, 12.0]), - centroid: Box::new([11.0]), - }, - ), - ( - 30, - TestList { - members: Box::new([300]), - vectors: matrix(&[100.0]), - centroid: Box::new([100.0]), - }, - ), - ] - .into_iter() - .collect(), - reads: Arc::default(), - stage_calls: Cell::new(0), - staged: Vec::new(), - applied: Arc::default(), - fail_read: None, - fail_stage: None, - missing_centroid: None, - } - } - } - - struct TestCentroids<'a> { - lists: &'a HashMap, - missing: Option, - } - - impl Centroids for TestCentroids<'_> { - type Id = u32; - type Error = ANNError; - - fn len(&self) -> usize { - self.lists.len() - } - - fn dim(&self) -> usize { - self.lists.values().next().unwrap().centroid.len() - } - - fn centroid(&self, id: u32) -> Option> { - self.lists - .get(&id) - .filter(|_| self.missing != Some(id)) - .map(|list| list.centroid.as_ref()) - } - - fn select(&self, query: &[f32], nprobe: usize) -> ANNResult>> { - let mut selected: Vec<_> = self - .lists - .iter() - .map(|(&id, list)| SelectedList { - id, - distance: query - .iter() - .zip(list.centroid.iter()) - .map(|(a, b)| (a - b).powi(2)) - .sum(), - }) - .collect(); - selected.sort_by(|a, b| a.distance.total_cmp(&b.distance).then(a.id.cmp(&b.id))); - selected.truncate(nprobe); - Ok(selected) - } - } - - struct TestReader<'a> { - lists: &'a HashMap, - reads: &'a Mutex>, - fail: Option, - } - - impl Reader for TestReader<'_> { - type Error = ANNError; - type Id = u32; - - fn dim(&self) -> usize { - self.lists.values().next().unwrap().centroid.len() - } - - async fn read_into(&self, id: u32, mut out: rowmajor::Mut<'_, f32>) -> ANNResult<()> { - self.reads.lock().unwrap().push(id); - if self.fail == Some(id) { - return Err(index_error("fixture read failed")); - } - out.as_mut_slice() - .copy_from_slice(self.lists[&id].vectors.as_slice()); - Ok(()) - } - } - - impl Stage for TestAccessor { - async fn stage_point(&mut self, id: &u32, element: f32, out: &mut [f32]) -> ANNResult { - out[0] = element; - Ok(*id + 10_000) - } - - async fn stage_centroid(&mut self, centroid: &[f32]) -> ANNResult { - let call = self.stage_calls.get(); - self.stage_calls.set(call + 1); - if self.fail_stage == Some(call) { - return Err(index_error("fixture staging failed")); - } - let id = 1000 + call; - self.staged.push((id, centroid.into())); - Ok(id) - } - } - - impl InsertAccessor for TestAccessor { - type Centroids<'a> = TestCentroids<'a>; - type Reader<'a> = TestReader<'a>; - type Error = ANNError; - - fn reader(&self) -> TestReader<'_> { - TestReader { - lists: &self.lists, - reads: &self.reads, - fail: self.fail_read, - } - } - - fn centroids(&self) -> TestCentroids<'_> { - TestCentroids { - lists: &self.lists, - missing: self.missing_centroid, - } - } - - fn dim(&self) -> usize { - self.lists.values().next().unwrap().centroid.len() - } - - fn list_size(&self, list: u32) -> ANNResult { - Ok(self.get_members(list)?.len()) - } - - fn get_members(&self, list: u32) -> ANNResult<&[u32]> { - self.lists - .get(&list) - .map(|list| list.members.as_ref()) - .ok_or_else(|| index_error(format!("fixture list {list} is missing"))) - } - - async fn update(self, update: Deltas) -> ANNResult<()> { - self.applied.lock().unwrap().push(update); - Ok(()) - } - } - - struct TestStrategy(Mutex>); - - impl<'a> MaintenanceStrategy<'a, TestProvider, f32> for TestStrategy { - type MaintenanceAccessor = TestAccessor; - type Error = ANNError; - - fn maintenance_accessor( - &'a self, - _provider: &'a mut TestProvider, - _context: &'a DefaultContext, - ) -> ANNResult { - self.0 - .lock() - .unwrap() - .take() - .ok_or_else(|| index_error("fixture accessor was consumed")) - } - } - - fn require_send(value: T) -> T { - value - } - - async fn split_update( - accessor: &mut TestAccessor, - inputs: &[ListInsertion<'_, u32, u32>], - ) -> ANNResult>> { - let mut region = - IVFIndex::::gather_region::<_, f32>(accessor, inputs).await?; - let parents: Vec<_> = inputs.iter().map(|input| input.list).collect(); - TwoMeansSplit { iterations: 10 }.split( - &mut region, - &parents, - &mut StdRng::seed_from_u64(7), - )?; - let assignments = NearestCentroid.assign(®ion)?; - IVFIndex::::build_deltas::<_, f32>(accessor, ®ion, &assignments).await - } - - fn appends(update: &Deltas) -> HashMap> { - let mut appends = HashMap::new(); - for delta in update.deltas() { - if let Delta::PointAppends { to, ids } = delta { - assert!(appends.insert(*to, ids.to_vec()).is_none()); - } - } - appends - } - - fn retired(update: &Deltas, parent: u32) -> Vec<(u32, u32)> { - update - .deltas() - .iter() - .find_map(|delta| match delta { - Delta::CentroidDelta { - id, - delta: CentroidDelta::Retire { moves }, - } if *id == parent => { - Some(moves.iter().map(|point| (point.id(), point.to())).collect()) - } - _ => None, - }) - .unwrap() - } - - #[tokio::test] - async fn gather_preserves_list_and_staged_order_in_one_buffer() { - let mut accessor = TestAccessor::new(); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let inputs = inputs(&vectors, &[10, 20, 30]); - assert_eq!(inputs[0].staged[0].vector.as_ptr(), vectors.row(3).as_ptr()); - assert_eq!(inputs[0].staged[1].vector.as_ptr(), vectors.row(0).as_ptr()); - let region = require_send(IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &inputs, - )) - .await - .unwrap(); - - assert_eq!( - region.point_ids, - [100, 101, 503, 500, 200, 201, 501, 300, 502] - ); - assert_eq!( - region.points.as_slice(), - &[0.0, 2.0, 3.0, 1.0, 10.0, 12.0, 11.0, 100.0, 100.0] - ); - assert_eq!(region.centroid_ids, [Some(10), Some(20), Some(30)]); - assert_eq!(region.centroids.as_slice(), &[1.0, 11.0, 100.0]); - for (index, (rows, members, staged)) in [ - (0..4, &[100, 101][..], &[503, 500][..]), - (4..7, &[200, 201][..], &[501][..]), - (7..9, &[300][..], &[502][..]), - ] - .into_iter() - .enumerate() - { - assert_eq!(region.lists[index].rows, rows); - assert_eq!(region.member_ids(index), members); - assert_eq!(region.staged_ids(index), staged); - assert_eq!( - region.list_points(index).unwrap().as_ptr(), - region.points.row(rows.start).as_ptr() - ); - } - assert_eq!(*accessor.reads.lock().unwrap(), [10, 20, 30]); - } - - #[tokio::test] - async fn gather_handles_multiple_dimensions_and_an_empty_member_prefix() { - let mut accessor = TestAccessor::new(); - accessor.lists = [( - 10, - TestList { - members: Box::new([]), - vectors: rowmajor::Owned::from_element(0, 2, 0.0), - centroid: Box::new([0.0, 0.0]), - }, - )] - .into_iter() - .collect(); - let staged = rowmajor::Owned::try_from_data(Box::new([1.0, 2.0, 3.0, 4.0]), 2, 2).unwrap(); - let region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &[ListInsertion { - list: 10, - staged: vec![ - StagedPoint { - id: 501, - vector: staged.row(1), - }, - StagedPoint { - id: 500, - vector: staged.row(0), - }, - ], - }], - ) - .await - .unwrap(); - assert_eq!(region.points.as_slice(), &[3.0, 4.0, 1.0, 2.0]); - assert_eq!(region.point_ids, [501, 500]); - assert!(region.member_ids(0).is_empty()); - assert_eq!(*accessor.reads.lock().unwrap(), [10]); - } - - #[tokio::test] - async fn gather_fills_disjoint_multidimensional_subviews_in_list_order() { - let mut accessor = TestAccessor::new(); - accessor.lists = [ - ( - 10, - TestList { - members: Box::new([100, 101]), - vectors: rowmajor::Owned::try_from_data(Box::new([1.0, 2.0, 3.0, 4.0]), 2, 2) - .unwrap(), - centroid: Box::new([2.0, 3.0]), - }, - ), - ( - 20, - TestList { - members: Box::new([200]), - vectors: rowmajor::Owned::try_from_data(Box::new([5.0, 6.0]), 1, 2).unwrap(), - centroid: Box::new([5.0, 6.0]), - }, - ), - ] - .into_iter() - .collect(); - let staged = rowmajor::Owned::try_from_data(Box::new([7.0, 8.0, 9.0, 10.0]), 2, 2).unwrap(); - let region = require_send(IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &[ - ListInsertion { - list: 10, - staged: vec![ - StagedPoint { - id: 501, - vector: staged.row(1), - }, - StagedPoint { - id: 500, - vector: staged.row(0), - }, - ], - }, - ListInsertion { - list: 20, - staged: Vec::new(), - }, - ], - )) - .await - .unwrap(); - - assert_eq!(region.lists[0].rows, 0..4); - assert_eq!(region.lists[1].rows, 4..5); - assert_eq!(region.point_ids, [100, 101, 501, 500, 200]); - assert_eq!( - region.points.as_slice(), - &[1.0, 2.0, 3.0, 4.0, 9.0, 10.0, 7.0, 8.0, 5.0, 6.0], - ); - assert_eq!(region.centroids.as_slice(), &[2.0, 3.0, 5.0, 6.0]); - assert!(region.staged_ids(1).is_empty()); - assert_eq!(*accessor.reads.lock().unwrap(), [10, 20]); - } - - #[tokio::test] - async fn split_can_retire_an_empty_list_and_assign_only_staged_points() { - let mut accessor = TestAccessor::new(); - accessor.lists.get_mut(&10).unwrap().members = Box::new([]); - accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[]); - let vectors = matrix(&[0.0, 10.0]); - let update = Deltas::new( - split_update( - &mut accessor, - &[ListInsertion { - list: 10, - staged: vec![ - StagedPoint { - id: 501, - vector: vectors.row(1), - }, - StagedPoint { - id: 500, - vector: vectors.row(0), - }, - ], - }], - ) - .await - .unwrap(), - ); - assert!(retired(&update, 10).is_empty()); - assert_eq!(accessor.staged.len(), 2); - let appended = appends(&update); - assert_eq!(appended.len(), 2); - for (to, ids) in appended { - assert!(to >= 1000); - assert_eq!(ids.len(), 1); - assert!([500, 501].contains(&ids[0])); - } - } - - #[tokio::test] - async fn split_replaces_parents_preserves_points_and_numerical_behavior() { - let mut accessor = TestAccessor::new(); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let mut region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &inputs(&vectors, &[10, 20, 30]), - ) - .await - .unwrap(); - let points_ptr = region.points.as_ptr(); - let ids_ptr = region.point_ids.as_ptr(); - let point_values = region.points.as_slice().to_vec(); - let mut rng = StdRng::seed_from_u64(7); - let mut expected_rng = rng.clone(); - let mut expected = rowmajor::Owned::from_element(4, 1, 0.0); - for (out, source) in expected.as_mut_slice().chunks_exact_mut(2).zip([1, 0]) { - fit_two_means( - region.list_points(source).unwrap(), - rowmajor::Mut::try_from_data(out, 2, 1).unwrap(), - 10, - &mut expected_rng, - &mut LloydScratch::default(), - ) - .unwrap(); - } - - TwoMeansSplit { iterations: 10 } - .split(&mut region, &[20, 10], &mut rng) - .unwrap(); - assert_eq!(region.centroid_ids, [None, None, None, None, Some(30)]); - assert_eq!(®ion.centroids.as_slice()[..4], expected.as_slice()); - assert_eq!(region.centroids.row(4), &[100.0]); - assert_eq!(region.points.as_ptr(), points_ptr); - assert_eq!(region.point_ids.as_ptr(), ids_ptr); - assert_eq!(region.points.as_slice(), point_values); - assert_eq!(rng.random::(), expected_rng.random::()); - assert_eq!(accessor.stage_calls.get(), 0); - } - - #[tokio::test] - async fn regional_assignment_can_choose_another_parents_children() { - let mut accessor = TestAccessor::new(); - accessor.lists.get_mut(&10).unwrap().members = Box::new([100, 101, 102]); - accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[0.0, 10.0, 100.0]); - accessor.lists.get_mut(&20).unwrap().vectors = matrix(&[1.0, 2.0]); - let mut region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &[ - ListInsertion { - list: 10, - staged: Vec::new(), - }, - ListInsertion { - list: 20, - staged: Vec::new(), - }, - ], - ) - .await - .unwrap(); - TwoMeansSplit { iterations: 10 } - .split(&mut region, &[10, 20], &mut StdRng::seed_from_u64(7)) - .unwrap(); - let assigned = NearestCentroid.assign(®ion).unwrap(); - assert_eq!(region.centroid_ids, [None, None, None, None]); - assert!( - assigned[0] >= 2, - "A's point at zero should prefer B's child at one" - ); - assert_eq!(region.centroids.row(assigned[0]), &[1.0]); - let update = Deltas::new( - IVFIndex::::build_deltas::<_, f32>(&mut accessor, ®ion, &assigned) - .await - .unwrap(), - ); - assert_eq!(retired(&update, 10)[0], (100, 1000 + assigned[0] as u32)); - for parent in [10, 20] { - assert!(retired(&update, parent).iter().all(|(_, to)| *to >= 1000)); - } - } - - #[tokio::test] - async fn regional_deltas_move_parent_members_and_append_only_regional_staged_points() { - let mut accessor = TestAccessor::new(); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let update = Deltas::new( - require_send(split_update(&mut accessor, &inputs(&vectors, &[10, 20]))) - .await - .unwrap(), - ); - assert_eq!(accessor.staged.len(), 4); - assert_eq!(*accessor.reads.lock().unwrap(), [10, 20]); - for (parent, ids) in [(10, [100, 101]), (20, [200, 201])] { - let moves = retired(&update, parent); - assert_eq!(moves.iter().map(|(id, _)| *id).collect::>(), ids); - assert!(moves.iter().all(|(_, to)| *to >= 1000)); - } - let appended = appends(&update); - assert!(!appended.contains_key(&30)); - let mut ids: Vec<_> = appended.values().flatten().copied().collect(); - ids.sort_unstable(); - assert_eq!(ids, [500, 501, 503]); - let destination = appended.values().find(|ids| ids.contains(&503)).unwrap(); - assert_eq!(destination, &[503]); - } - - struct ThreeChildren; - - impl Split for ThreeChildren { - fn split( - &self, - region: &mut Region, - _parents: &[u32], - _rng: &mut StdRng, - ) -> ANNResult<()> { - region.centroids = matrix(&[0.0, 2.0, 100.0, 100.0]); - region.centroid_ids = vec![None, None, None, Some(30)]; - Ok(()) - } - } - - #[tokio::test] - async fn custom_split_supports_three_children_and_neighbor_reassignment() { - let mut accessor = TestAccessor::new(); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let mut region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &inputs(&vectors, &[10, 30]), - ) - .await - .unwrap(); - ThreeChildren - .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) - .unwrap(); - let assigned = NearestCentroid.assign(®ion).unwrap(); - let update = Deltas::new( - IVFIndex::::build_deltas::<_, f32>(&mut accessor, ®ion, &assigned) - .await - .unwrap(), - ); - assert_eq!(accessor.staged.len(), 3); - let neighbor_moves: Vec<_> = update - .deltas() - .iter() - .filter_map(|delta| match delta { - Delta::PointMoves { from: 30, moves } => Some(moves.as_ref()), - _ => None, - }) - .flatten() - .map(|point| (point.id(), point.to())) - .collect(); - assert_eq!(neighbor_moves, [(300, 1002)]); - assert_eq!(appends(&update)[&1002], [502]); - assert!( - update - .deltas() - .iter() - .all(|delta| !matches!(delta, Delta::CentroidDelta { id: 30, .. })) - ); - } - - #[tokio::test] - async fn surviving_neighbor_members_and_staged_points_can_stay_put() { - let mut accessor = TestAccessor::new(); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let mut region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &inputs(&vectors, &[10, 30]), - ) - .await - .unwrap(); - TwoMeansSplit { iterations: 10 } - .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) - .unwrap(); - let assigned = NearestCentroid.assign(®ion).unwrap(); - let update = Deltas::new( - IVFIndex::::build_deltas::<_, f32>(&mut accessor, ®ion, &assigned) - .await - .unwrap(), - ); - assert_eq!(accessor.staged.len(), 2); - assert!( - update - .deltas() - .iter() - .all(|delta| !matches!(delta, Delta::PointMoves { .. })) - ); - assert_eq!(appends(&update)[&30], [502]); - } - - #[tokio::test] - async fn custom_assignment_uses_the_same_flat_row_contract() { - struct LastCentroid; - impl Assign for LastCentroid { - fn assign(&self, region: &Region) -> ANNResult> { - Ok(vec![region.centroids.nrows() - 1; region.points.nrows()].into_boxed_slice()) - } - } - let mut accessor = TestAccessor::new(); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let mut region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &inputs(&vectors, &[10, 30]), - ) - .await - .unwrap(); - TwoMeansSplit { iterations: 10 } - .split(&mut region, &[10], &mut StdRng::seed_from_u64(7)) - .unwrap(); - let update = Deltas::new( - IVFIndex::::build_deltas::<_, f32>( - &mut accessor, - ®ion, - &LastCentroid.assign(®ion).unwrap(), - ) - .await - .unwrap(), - ); - assert_eq!(retired(&update, 10), [(100, 30), (101, 30)]); - assert_eq!(appends(&update)[&30], [503, 500, 502]); - } - - #[tokio::test] - async fn empty_region_inputs_skip_reads_and_centroid_staging() { - let mut accessor = TestAccessor::new(); - let region = require_send(IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &[], - )) - .await - .unwrap(); - assert!(region.lists.is_empty()); - assert!(region.point_ids.is_empty()); - assert_eq!(region.points.nrows(), 0); - assert_eq!(region.centroids.nrows(), 0); - assert!(accessor.reads.lock().unwrap().is_empty()); - assert_eq!(accessor.stage_calls.get(), 0); - } - - #[tokio::test] - async fn gathering_errors_propagate_before_centroid_staging() { - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - for failure in ["members", "reader", "centroid"] { - let mut accessor = TestAccessor::new(); - match failure { - "members" => { - accessor.lists.remove(&20); - } - "reader" => { - accessor.fail_read = Some(20); - } - "centroid" => { - accessor.missing_centroid = Some(20); - } - _ => unreachable!(), - } - let err = split_update(&mut accessor, &inputs(&vectors, &[10, 20])) - .await - .unwrap_err(); - let expected = match failure { - "members" => "fixture list 20 is missing", - "reader" => "fixture read failed", - "centroid" => "centroid 20 is unavailable", - _ => unreachable!(), - }; - assert!(err.to_string().contains(expected), "{err}"); - assert_eq!(accessor.stage_calls.get(), 0); - } - } - - #[tokio::test] - async fn numerical_and_staging_errors_propagate() { - let mut accessor = TestAccessor::new(); - accessor.lists.get_mut(&10).unwrap().members = Box::new([100]); - accessor.lists.get_mut(&10).unwrap().vectors = matrix(&[0.0]); - let err = split_update( - &mut accessor, - &[ListInsertion { - list: 10, - staged: Vec::new(), - }], - ) - .await - .unwrap_err(); - assert!(err.to_string().contains("at least two points"), "{err}"); - assert_eq!(accessor.stage_calls.get(), 0); - - let mut region = IVFIndex::::gather_region::<_, f32>( - &mut accessor, - &[ListInsertion { - list: 10, - staged: Vec::new(), - }], - ) - .await - .unwrap(); - region.centroids = matrix(&[]); - region.centroid_ids.clear(); - let err = NearestCentroid.assign(®ion).unwrap_err(); - assert!(err.to_string().contains("at least one centroid"), "{err}"); - - accessor.fail_stage = Some(1); - let vectors = matrix(&[1.0, 11.0, 100.0, 3.0]); - let err = split_update(&mut accessor, &inputs(&vectors, &[10, 20])) - .await - .unwrap_err(); - assert!(err.to_string().contains("fixture staging failed"), "{err}"); - assert_eq!(accessor.staged.len(), 1); - } - - #[tokio::test] - async fn insert_batch_applies_splits_and_unsplit_appends_together() { - for threshold in [2, 3, 100] { - let accessor = TestAccessor::new(); - let applied = Arc::clone(&accessor.applied); - let reads = Arc::clone(&accessor.reads); - let strategy = TestStrategy(Mutex::new(Some(accessor))); - let mut index = IVFIndex::new( - TestProvider, - Config { - split_threshold: threshold, - reassign_neighbors: 1, - two_means_iterations: 10, - seed: 7, - }, - ) - .unwrap(); - require_send(index.insert_batch( - &strategy, - &DefaultContext, - &[(500, 1.0), (501, 11.0), (502, 100.0), (503, 3.0)], - )) - .await - .unwrap(); - let updates = applied.lock().unwrap(); - assert_eq!(updates.len(), 1); - let appended = appends(&updates[0]); - let mut ids: Vec<_> = appended.values().flatten().copied().collect(); - ids.sort_unstable(); - assert_eq!(ids, [10500, 10501, 10502, 10503]); - assert_eq!(appended[&30], [10502]); - if threshold < 100 { - assert_eq!(retired(&updates[0], 10).len(), 2); - let mut read_ids = reads.lock().unwrap().clone(); - read_ids.sort_unstable(); - if threshold == 2 { - assert_eq!(retired(&updates[0], 20).len(), 2); - assert_eq!(read_ids, [10, 20]); - } else { - assert_eq!(appended[&20], [10501]); - assert_eq!(read_ids, [10]); - } - } else { - assert_eq!(appended[&10], [10500, 10503]); - assert_eq!(appended[&20], [10501]); - assert!(reads.lock().unwrap().is_empty()); - assert_eq!(updates[0].deltas().len(), 3); - } - } - } - - #[tokio::test] - async fn regional_failure_does_not_commit_already_planned_ordinary_appends() { - let mut accessor = TestAccessor::new(); - accessor.fail_read = Some(10); - let applied = Arc::clone(&accessor.applied); - let strategy = TestStrategy(Mutex::new(Some(accessor))); - let mut index = IVFIndex::new( - TestProvider, - Config { - split_threshold: 3, - reassign_neighbors: 1, - two_means_iterations: 10, - seed: 7, - }, - ) - .unwrap(); - let err = require_send(index.insert_batch( - &strategy, - &DefaultContext, - &[(500, 1.0), (501, 11.0), (502, 100.0), (503, 3.0)], - )) - .await - .unwrap_err(); - assert!(err.to_string().contains("fixture read failed"), "{err}"); - assert!(applied.lock().unwrap().is_empty()); - } - - #[tokio::test] - async fn empty_insert_does_not_construct_an_accessor_or_apply_an_update() { - let strategy = TestStrategy(Mutex::new(None)); - let mut index = IVFIndex::new( - TestProvider, - Config { - split_threshold: 2, - reassign_neighbors: 1, - two_means_iterations: 10, - seed: 7, - }, - ) - .unwrap(); - let result = require_send(index.insert_batch(&strategy, &DefaultContext, &[])) - .await - .unwrap(); - assert_eq!(result.inserted, 0); - } -} diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs new file mode 100644 index 0000000000..f9c2ffde27 --- /dev/null +++ b/diskann/src/ivf/tests.rs @@ -0,0 +1,713 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{cell::Cell, collections::BTreeMap, future::Future, sync::Arc}; + +use diskann_utils::{ + future::SendFuture, + views::rowmajor::{self, Matrix, MatrixMut}, +}; +use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; +use rand::{SeedableRng, rngs::StdRng}; + +use super::{ + Config, IVFIndex, InsertStats, clustering, + pending::{Membership, PendingLists}, + traits::{ + BuildAccessor, BuildStrategy, CoarseAccessor, InsertStrategy, ListAccessor, Provider, + SelectedList, + }, + workingset::{Batch, ListView, PointRef, WorkingSet}, +}; +use crate::{ANNError, ANNResult, provider::DefaultContext}; + +#[derive(Debug)] +struct StoredList { + centroid: Vec, + members: Vec, +} + +#[derive(Debug, Default)] +struct Store { + dim: usize, + points: BTreeMap>, + lists: BTreeMap, + next_list: u32, + finishes: usize, + fills: usize, + fail_fill: bool, + fail_finish: bool, +} + +impl Provider for Store { + type Context = DefaultContext; + type InternalId = u32; + type ExternalId = u32; + type ListId = u32; + type Error = ANNError; + + 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")) + } + + 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")) + } +} + +struct Build<'p> { + store: &'p mut Store, + pending: PendingLists, + created: BTreeMap>, + batch: Option>>, + next_list: u32, + fills: Cell, +} + +impl Build<'_> { + fn stored_members(&self, list: u32) -> ANNResult<&[u32]> { + if let Some(stored) = self.store.lists.get(&list) { + Ok(&stored.members) + } else if self.created.contains_key(&list) { + Ok(&[]) + } else { + Err(ANNError::message(format!("unknown list {list}"))) + } + } +} + +struct Coarse<'a> { + stored: &'a BTreeMap, + created: &'a BTreeMap>, + pending: &'a PendingLists, +} + +impl CoarseAccessor for Coarse<'_> { + type ListId = u32; + + fn len(&self) -> usize { + self.stored + .keys() + .chain(self.created.keys()) + .filter(|&&id| !self.pending.is_retired(id)) + .count() + } + + fn select( + &mut self, + query: &[f32], + nprobe: usize, + ) -> impl SendFuture>>> { + async move { + let mut selected: Vec<_> = self + .stored + .iter() + .map(|(&id, list)| (id, list.centroid.as_slice())) + .chain( + self.created + .iter() + .map(|(&id, vector)| (id, vector.as_slice())), + ) + .filter(|(id, _)| !self.pending.is_retired(*id)) + .map(|(id, vector)| SelectedList { + id, + distance: SquaredL2::evaluate(query, vector), + }) + .collect(); + selected.sort_by(|a, b| { + a.distance + .total_cmp(&b.distance) + .then_with(|| a.id.cmp(&b.id)) + }); + selected.truncate(nprobe); + Ok(selected) + } + } +} + +struct Lists<'a, 'p>(&'a mut Build<'p>); + +impl ListAccessor for Lists<'_, '_> { + type InternalId = u32; + type ListId = u32; + + fn list_size(&mut self, list: u32) -> impl SendFuture> { + async move { + self.0 + .pending + .list_size(list, self.0.stored_members(list)?.len()) + } + } + + fn append_members(&mut self, list: u32, members: Vec) -> impl SendFuture> { + async move { + self.0.stored_members(list)?; + self.0.pending.append(list, members) + } + } + + fn set_members(&mut self, list: u32, members: Vec) -> impl SendFuture> { + async move { + self.0.stored_members(list)?; + self.0.pending.set(list, members) + } + } + + fn create_list( + &mut self, + centroid: &[f32], + members: Vec, + ) -> impl SendFuture> { + async move { + if centroid.len() != self.0.store.dim { + return Err(ANNError::message("wrong centroid dimension")); + } + let id = self.0.next_list; + self.0.next_list += 1; + self.0.created.insert(id, centroid.to_vec()); + self.0.pending.set(id, members)?; + Ok(id) + } + } + + fn retire_list(&mut self, list: u32) -> impl SendFuture> { + async move { + self.0.stored_members(list)?; + self.0.pending.retire(list) + } + } +} + +impl<'p> BuildAccessor for Build<'p> { + type InternalId = u32; + type ListId = u32; + type Coarse<'a> + = Coarse<'a> + where + Self: 'a; + type Lists<'a> + = Lists<'a, 'p> + where + Self: 'a; + + fn dim(&self) -> usize { + self.store.dim + } + + fn coarse(&mut self) -> Self::Coarse<'_> { + Coarse { + stored: &self.store.lists, + created: &self.created, + pending: &self.pending, + } + } + + fn lists(&mut self) -> Self::Lists<'_> { + Lists(self) + } + + fn fill(&mut self, lists: &[u32]) -> impl SendFuture>> { + async move { + self.fills.set(self.fills.get() + 1); + if self.store.fail_fill { + return Err(ANNError::message("injected fill failure")); + } + let mut view = Vec::with_capacity(lists.len()); + for &list in lists { + let size = self + .pending + .list_size(list, self.stored_members(list)?.len())?; + let mut points = Vec::with_capacity(size); + for id in self.pending.members(list, self.stored_members(list)?)? { + let vector = self + .batch + .as_ref() + .and_then(|batch| batch.get(id)) + .or_else(|| self.store.points.get(&id).map(Vec::as_slice)) + .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; + points.push(PointRef { id, vector }); + } + view.push(ListView { id: list, points }); + } + Ok(view) + } + } + + fn finish(self) -> impl SendFuture> { + async move { + if self.store.fail_finish { + return Err(ANNError::message("injected finish failure")); + } + if let Some(batch) = self.batch { + for point in batch.points() { + self.store.points.insert(point.id, point.vector.to_vec()); + } + } + for (id, centroid) in self.created { + self.store.lists.insert( + id, + StoredList { + centroid, + members: Vec::new(), + }, + ); + } + 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 = self.next_list; + self.store.finishes += 1; + self.store.fills += self.fills.get(); + Ok(()) + } + } +} + +struct Strategy; + +impl<'a> BuildStrategy<'a, Store> for Strategy { + type Accessor = Build<'a>; + type Error = ANNError; + + fn build_accessor( + &'a self, + store: &'a mut Store, + _: &'a DefaultContext, + ) -> ANNResult> { + let next_list = store.next_list; + Ok(Build { + store, + pending: PendingLists::default(), + created: BTreeMap::new(), + batch: None, + next_list, + fills: Cell::new(0), + }) + } +} + +impl<'a, 'p> InsertStrategy<'a, Store, &'p [f32]> for Strategy { + fn prepare( + &self, + accessor: &mut Self::Accessor, + points: &[(u32, &'p [f32])], + ) -> impl SendFuture>>> { + async move { + let mut vectors = rowmajor::Owned::from_element(points.len(), accessor.dim(), 0.0); + let mut ids = Vec::with_capacity(points.len()); + for ((external, vector), out) in points.iter().zip(vectors.rows_mut()) { + if vector.len() != out.len() || !vector.iter().all(|value| value.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 accessor.store.points.contains_key(&id) || ids.contains(&id) { + return Err(ANNError::message("point already exists")); + } + out.copy_from_slice(vector); + ids.push(id); + } + let batch = Arc::new(Batch::new(ids, vectors)?); + accessor.batch = Some(Arc::clone(&batch)); + Ok(batch) + } + } +} + +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() +} + +fn index(store: Store, threshold: usize) -> IVFIndex { + IVFIndex::new( + store, + Config { + split_threshold: threshold, + two_means_iterations: 5, + seed: 7, + }, + ) + .unwrap() +} + +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() + } +} + +fn require_send(future: F) -> F { + future +} + +#[tokio::test] +async fn initialize_and_append_publish_actual_memberships() { + let mut index = index( + Store { + dim: 2, + next_list: 1, + ..Store::default() + }, + 2, + ); + let centers = matrix(&[[0.0, 0.0], [10.0, 0.0]]); + require_send(index.initialize(&Strategy, &DefaultContext, centers.as_view())) + .await + .unwrap(); + let points = [ + (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 + } + ); + assert_eq!(index.provider().lists[&1].members, [10_007, 10_008]); + assert_eq!(index.provider().lists[&2].members, [10_009]); + assert_eq!(index.provider().points[&10_007], [0.0, 1.0]); + assert_eq!( + index + .provider() + .to_internal_id(&DefaultContext, &7) + .unwrap(), + 10_007 + ); + assert_eq!( + index + .provider() + .to_external_id(&DefaultContext, 10_007) + .unwrap(), + 7 + ); + assert_eq!(index.provider().finishes, 2); + assert_eq!(index.provider().fills, 0); + + let empty: [(u32, &[f32]); 0] = []; + assert_eq!( + index + .insert_batch(&Strategy, &DefaultContext, &empty) + .await + .unwrap(), + InsertStats::default() + ); + assert_eq!(index.provider().finishes, 2); +} + +#[tokio::test] +async fn splitting_commits_children_and_ordinary_appends_once() { + let mut index = index(seeded(), 2); + let points = [ + (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 = index.provider(); + + 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 working_set_borrows_stored_and_prepared_rows() { + let mut store = seeded(); + { + let mut build = Strategy + .build_accessor(&mut store, &DefaultContext) + .unwrap(); + let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; + let batch = Strategy.prepare(&mut build, &points).await.unwrap(); + let old_pointer = build.store.points[&10_000].as_ptr(); + build + .lists() + .append_members(1, vec![10_006, 10_005]) + .await + .unwrap(); + assert_eq!(build.lists().list_size(1).await.unwrap(), 4); + { + let view = require_send(build.fill(&[1])).await.unwrap(); + assert_eq!( + view[0] + .points + .iter() + .map(|point| point.id) + .collect::>(), + [10_000, 10_001, 10_006, 10_005] + ); + assert_eq!(view[0].points[0].vector.as_ptr(), old_pointer); + assert!(std::ptr::eq( + view[0].points[2].vector, + batch.get(10_006).unwrap() + )); + assert!(std::ptr::eq( + view[0].points[3].vector, + batch.get(10_005).unwrap() + )); + } + build + .lists() + .set_members(1, vec![10_000, 10_005]) + .await + .unwrap(); + assert_eq!(build.lists().list_size(1).await.unwrap(), 2); + assert_eq!(build.fill(&[1]).await.unwrap()[0].points.len(), 2); + build.lists().retire_list(1).await.unwrap(); + assert!(build.lists().list_size(1).await.is_err()); + assert!(build.fill(&[1]).await.is_err()); + assert!(build.lists().append_members(1, vec![10_001]).await.is_err()); + } + 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 failed_planning_or_publication_is_not_reported_as_success() { + for fail_fill in [true, false] { + let mut store = seeded(); + store.fail_fill = fail_fill; + store.fail_finish = !fail_fill; + let mut index = index(store, 2); + let points = [(5, &[1.0, 1.0][..])]; + let error = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) + .await + .unwrap_err(); + let expected = if fail_fill { + "injected fill failure" + } else { + "injected finish failure" + }; + assert!(error.to_string().contains(expected), "{error}"); + assert_eq!(index.provider().finishes, 0); + assert_eq!(index.provider().lists[&1].members, [10_000, 10_001]); + assert_eq!(index.provider().lists.len(), 3); + assert!( + index + .provider() + .to_internal_id(&DefaultContext, &5) + .is_err() + ); + } +} + +#[tokio::test] +async fn invalid_initialization_and_uninitialized_insertion_fail() { + assert!( + IVFIndex::new( + Store::default(), + Config { + split_threshold: 1, + two_means_iterations: 5, + seed: 7, + }, + ) + .is_err() + ); + let mut index = index( + Store { + dim: 2, + next_list: 1, + ..Store::default() + }, + 2, + ); + for centers in [ + matrix::<2>(&[]), + matrix(&[[0.0]]), + matrix(&[[f32::NAN, 0.0]]), + ] { + assert!( + index + .initialize(&Strategy, &DefaultContext, centers.as_view()) + .await + .is_err() + ); + } + let points = [(7, &[0.0, 1.0][..])]; + assert!( + index + .insert_batch(&Strategy, &DefaultContext, &points) + .await + .is_err() + ); + assert_eq!(index.provider().finishes, 0); + let centers = matrix(&[[0.0, 0.0]]); + index + .initialize(&Strategy, &DefaultContext, centers.as_view()) + .await + .unwrap(); + assert!( + index + .initialize(&Strategy, &DefaultContext, centers.as_view()) + .await + .is_err() + ); + assert_eq!(index.provider().finishes, 1); +} + +#[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 view = vec![ + ListView { + id: 1u32, + points: left + .iter() + .enumerate() + .map(|(id, vector)| PointRef { + id: id as u32, + vector, + }) + .collect(), + }, + ListView { + id: 2, + points: right + .iter() + .enumerate() + .map(|(id, vector)| PointRef { + id: 100 + id as u32, + vector, + }) + .collect(), + }, + ]; + let children = clustering::split(&view, 1, 10, &mut StdRng::seed_from_u64(7)).unwrap(); + assert!( + children.members[2..] + .iter() + .any(|members| members.contains(&4)) + ); +} + +#[test] +fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { + let rows = [[f32::MAX], [f32::MAX], [f32::MAX]]; + let view = vec![ListView { + id: 1u32, + points: rows + .iter() + .enumerate() + .map(|(id, row)| PointRef { + id: id as u32, + vector: row, + }) + .collect(), + }]; + let zero = clustering::split(&view, 1, 0, &mut StdRng::seed_from_u64(7)).unwrap(); + let one = clustering::split(&view, 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); + assert_eq!(zero.centroids.as_slice(), one.centroids.as_slice()); + assert!( + zero.centroids + .as_slice() + .iter() + .all(|value| value.is_finite()) + ); + assert_eq!(zero.members[0], [0, 1, 2]); + assert!(zero.members[1].is_empty()); +} diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index dba1c4f59d..75953e654d 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -3,75 +3,37 @@ * Licensed under the MIT license. */ -//! Contracts for an incrementally maintained IVF index. -//! -//! This module defines the boundary between the IVF algorithm and -//! provider-specific centroid navigation, inverted-list access, and partition -//! mutation. Concrete accessors are responsible for presenting a coherent view -//! across those components. +//! Storage boundaries for IVF construction. -use std::{fmt::Debug, ops::Deref}; +use std::{fmt::Debug, sync::Arc}; -use diskann_utils::{future::SendFuture, views::rowmajor}; +use diskann_utils::future::SendFuture; +use super::workingset::{Batch, WorkingSet}; use crate::{ ANNResult, error::{StandardError, ToRanked}, - ivf::update::Deltas, provider::ExecutionContext, utils::VectorId, }; -/// One centroid/list selected for a query. -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct SelectedList { - /// Stable logical list identifier. - pub id: L, - /// Distance from the query to this list's centroid. - pub distance: f32, -} - -/// Work performed while scanning the selected lists. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct ScanStats { - /// Number of point representations scored against the query. - pub comparisons: u64, - /// Number of distinct lists read. - pub lists_scanned: u32, -} - -/// The storage behind a dynamic IVF index: point identity and the lists points are -/// partitioned into. +/// Point identity and execution context shared by an IVF provider's accessors. /// -/// Reads and writes go through operation-scoped accessors; the provider itself only -/// names the identity types and translates the ids of visible points. Points staged -/// through [`StageElements`] are not translatable until an update that places them is -/// applied. +/// Prepared points become translatable only after the build operation finishes. pub trait Provider: Sized + Send + Sync + 'static { - /// Per-operation context handed to strategies. type Context: ExecutionContext; - - /// Internal point id used by lists and updates. type InternalId: VectorId; - - /// Caller-facing point id. type ExternalId: PartialEq + Send + Sync + 'static; - - /// Stable id shared by a centroid and its inverted list. Retired ids are never - /// reused. + /// Stable identity shared by a centroid and its inverted list. Never reused. type ListId: VectorId; - - /// Errors from id translation. type Error: ToRanked + Debug + Send + Sync + 'static; - /// Translate an external id to its corresponding internal id. fn to_internal_id( &self, context: &Self::Context, - gid: &Self::ExternalId, + id: &Self::ExternalId, ) -> Result; - /// Translate an internal id to its corresponding external id. fn to_external_id( &self, context: &Self::Context, @@ -79,200 +41,110 @@ pub trait Provider: Sized + Send + Sync + 'static { ) -> Result; } -pub trait Stage { - /// Stage a new point; its id is provisional until `update` commits it. - fn stage_point( - &mut self, - id: &P::ExternalId, - element: T, - out: &mut [f32], - ) -> impl SendFuture>; - - /// Stage a new centroid and mint its list id, provisional until `update` commits it. - fn stage_centroid(&mut self, centroid: &[f32]) -> impl SendFuture>; +/// A live list selected by centroid distance. +#[derive(Debug, Clone, Copy)] +pub struct SelectedList { + pub id: L, + pub distance: f32, } -/////////////// -// Centroids // -/////////////// - -/// An in-memory index over the live centroids. -/// -/// The centroid catalog is authoritative. Approximate implementations, such as -/// a graph navigator, must retain enough catalog information to reject retired -/// ids and recover with exact selection when navigation cannot produce a usable -/// result. Exact and graph implementations expose the same interface. -pub trait Centroids: Send + Sync { - /// Stable logical centroid id, also used as the inverted-list id. - type Id: VectorId; - - /// Errors encountered while reading or navigating the centroid index. - type Error: ToRanked + Debug + Send + Sync + 'static; +/// Centroid selection within one build operation. +pub trait CoarseAccessor: Send { + type ListId: VectorId; - /// Number of live centroids. fn len(&self) -> usize; - /// dimension of full precision centroid vectors. - fn dim(&self) -> usize; - - /// Whether there are no live centroids. fn is_empty(&self) -> bool { self.len() == 0 } - /// Borrow one authoritative full-precision centroid vector. - fn centroid(&self, id: Self::Id) -> Option>; - - /// Select exactly `min(nprobe, self.len())` live centroids for `query`, nearest - /// first. - /// - /// A graph implementation must use exact selection as a recovery path when - /// graph navigation produces too few live ids. + /// Select `min(nprobe, len())` live lists, nearest first. fn select( - &self, + &mut self, query: &[f32], nprobe: usize, - ) -> Result>, Self::Error>; + ) -> impl SendFuture>>>; } -////////////// -// Accessor // -////////////// - -pub trait InsertAccessor: Send + Sized + Stage { - /// In-memory centroid catalog and navigator in this accessor's unified view. - type Centroids<'a>: Centroids - where - Self: 'a; - - // Scoped reader to read vectors mapped to list ids. - type Reader<'a>: Reader - where - Self: 'a; - - /// Errors from planning reads, staging, or applying the update. - type Error: StandardError; - - /// Borrow the list reader to read vectors from posting lists. - fn reader(&self) -> Self::Reader<'_>; +/// Logical list mutations +pub trait ListAccessor: Send { + type InternalId: VectorId; + type ListId: VectorId; - /// Borrow the centroid index used for routing and maintenance neighborhoods. - fn centroids(&self) -> Self::Centroids<'_>; + fn list_size(&mut self, list: Self::ListId) -> impl SendFuture>; - fn dim(&self) -> usize; + /// Append members to a list. + fn append_members( + &mut self, + list: Self::ListId, + members: Vec, + ) -> impl SendFuture>; - /// Read the size of a live list. - fn list_size(&self, list: P::ListId) -> Result; + /// Replace the complete membership of a list. + fn set_members( + &mut self, + list: Self::ListId, + members: Vec, + ) -> impl SendFuture>; - /// Read the member ids of a live list, in the reader's vector-row order. - fn get_members(&self, list: P::ListId) -> Result<&[P::InternalId], Self::Error>; + /// Reserve a fresh list ID and stage its centroid and complete membership. + fn create_list( + &mut self, + centroid: &[f32], + members: Vec, + ) -> impl SendFuture>; - /// Consume the accessor and write the update - fn update(self, value: Deltas) -> impl SendFuture>; + /// Stage retirement. All mutations to `list` should be applied first before retiring. + fn retire_list(&mut self, list: Self::ListId) -> impl SendFuture>; } -pub trait Reader: Send + Sync { - /// Error type in case a read fails. - type Error: ToRanked + Debug + Send + Sync + 'static; +pub trait BuildAccessor: Send + Sized { + type InternalId: VectorId; + type ListId: VectorId; - /// Id type - type Id; + type Coarse<'a>: CoarseAccessor + where + Self: 'a; + type Lists<'a>: ListAccessor + where + Self: 'a; - /// dimension of the output of each Id. + /// Positive dimension of all canonical points and centroids. fn dim(&self) -> usize; + fn coarse(&mut self) -> Self::Coarse<'_>; + fn lists(&mut self) -> Self::Lists<'_>; - /// Read every canonical vector for `id`, preserving its member order in `out`. + /// Make every requested list's complete current membership available. /// - /// For an insert accessor's list reader, row `i` corresponds to member `i` - /// returned by [`InsertAccessor::get_members`] in the same accessor view. - /// `out` has one row per member and [`Self::dim`] columns; every row must be written. - fn read_into( - &self, - id: Self::Id, - out: rowmajor::Mut<'_, T>, - ) -> impl SendFuture>; -} - -////////////// -// Strategy // -////////////// - -/// Factory for an operation-scoped dynamic IVF maintenance accessor. -/// -/// The provider is borrowed exclusively for the lifetime of the accessor, so no -/// search or other mutation can observe it until the accessor is applied or -/// dropped. Providers may therefore mutate plain in-memory state in -/// [`InsertAccessor::update`] without interior synchronization. Providers that -/// share state outside this borrow (for example through an `Arc`) remain responsible -/// for coordinating those aliases. -pub trait MaintenanceStrategy<'a, P: Provider, T: Sync>: Send + Sync { - /// Accessor used to plan and stage split/dissolve operations. - type MaintenanceAccessor: InsertAccessor; - - /// Error constructing the accessor. - type Error: StandardError; + /// Lists appear in request order. + /// TO DO: Make [`WorkingSet`] an associated type here. + fn fill( + &mut self, + lists: &[Self::ListId], + ) -> impl SendFuture>>; - /// Construct one maintenance accessor. - fn maintenance_accessor( - &'a self, - provider: &'a mut P, - context: &'a P::Context, - ) -> Result; + /// Publish the operation's point mappings, centroids, and memberships together. + fn finish(self) -> impl SendFuture>; } -/// Factory for one dynamic IVF search accessor. -pub trait SearchStrategy<'a, P: Provider, T>: Send + Sync { - /// Query-bound accessor used for both coarse and fine search. - type SearchAccessor: SearchAccessor; - - /// Error constructing the accessor. +pub trait BuildStrategy<'a, P: Provider>: Send + Sync { + type Accessor: BuildAccessor; type Error: StandardError; - /// Construct an accessor around `query`. - fn search_accessor( + fn build_accessor( &'a self, - provider: &'a P, + provider: &'a mut P, context: &'a P::Context, - query: T, - ) -> Result; + ) -> Result; } -//////////// -// Search // -//////////// - -/// Query-bound access to centroid selection and inverted-list scanning. +/// Prepare typed inputs and seed a build accessor with their canonical data. /// -/// This is the dynamic IVF search algorithm's primary extension point, in the -/// same spirit as [`crate::graph::glue::SearchAccessor`]. Implementations are free -/// to batch reads, coalesce blob requests, prefetch, decode quantized payloads, or -/// fan work out across tasks. -pub trait SearchAccessor: Send + Sync { - /// Stable logical centroid/list id, fixed by the strategy. - type ListId; - - type InternalId; - - /// In-memory centroid view. - type Centroids<'a>: Centroids - where - Self: 'a; - - /// Errors from list selection or scanning. - type Error: StandardError; - - /// Borrow the centroid index used for routing. - fn centroids(&self) -> Self::Centroids<'_>; - - /// Scan `lists`, scoring their members against the bound query. - /// - /// `lists` must come from [`Self::select_lists`] on this accessor, which keeps - /// selection and scanning coherent using provider-specific coordination. - fn scan_lists( - &mut self, - lists: &[SelectedList], - emit: F, - ) -> impl SendFuture> - where - F: FnMut(Self::InternalId, f32) + Send; +/// Only this boundary depends on `T`; list operations and clustering do not. +pub trait InsertStrategy<'a, P: Provider, T: Sync>: BuildStrategy<'a, P> { + fn prepare( + &self, + accessor: &mut Self::Accessor, + points: &[(P::ExternalId, T)], + ) -> impl SendFuture>>>; } diff --git a/diskann/src/ivf/update.rs b/diskann/src/ivf/update.rs deleted file mode 100644 index f7ab7c74f1..0000000000 --- a/diskann/src/ivf/update.rs +++ /dev/null @@ -1,97 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use crate::utils::VectorId; - -/// A change to the set of live lists. -#[derive(Debug, Clone, PartialEq)] -pub enum CentroidDelta { - /// Install a the staged centroid with this Id. - Install, - /// Retire the Id and move the points in this list to other list Ids. - Retire { moves: Box<[MoveTo]> }, -} - -#[derive(Debug, Clone, PartialEq, Copy, Eq)] -pub struct MoveTo { - id: I, - to: L, -} - -impl MoveTo { - /// Move the point `id` into the list `to`. - pub(super) fn new(id: I, to: L) -> Self { - Self { id, to } - } - - /// The point that moves. - pub fn id(&self) -> I { - self.id - } - - /// The list the point joins. - pub fn to(&self) -> L { - self.to - } -} - -#[derive(Debug, Clone, PartialEq)] -pub enum Delta { - // Centroid delta keyed by `id` - CentroidDelta { - id: L, - delta: CentroidDelta, - }, - // Point moves due to re-assignments, grouped by source centroid Id. - PointMoves { - from: L, - moves: Box<[MoveTo]>, - }, - // Point insertions from staged points, grouped by destination centroid Id. - PointAppends { - to: L, - ids: Box<[Id]>, - }, -} - -/// The complete partition change made by bootstrap or one insert batch. -/// -/// The index guarantees, and providers may rely on without re-checking, that: -/// -/// - installed ids were reserved by this accessor, and their centroids are finite -/// with the accessor's dimension; -/// - retired ids are live; -/// - each list appears in at most one centroid delta, and each point in at most one -/// point delta; -/// - every point staged by this accessor is appended; -/// - a move's source holds the point and differs from its destination; -/// - every destination is live after the update: installed, or live and not retired; -/// and -/// - every member of a retired list is moved out, so a provider may drop a retired -/// list wholesale and only place the points moved out of it. -/// -/// Deltas are in no particular order, so a provider that cannot apply them as one -/// atomic change should install centroids before placing points into them. -#[derive(Debug, PartialEq)] -pub struct Deltas { - inner: Box<[Delta]>, -} - -impl Deltas { - /// Collect `deltas` into one partition change. - pub(super) fn new(deltas: Vec>) -> Self { - Self { - inner: deltas.into(), - } - } - - /// The individual changes this update makes. - pub fn deltas(&self) -> &[Delta] { - &self.inner - } -} - -#[cfg(test)] -mod tests {} diff --git a/diskann/src/ivf/workingset.rs b/diskann/src/ivf/workingset.rs new file mode 100644 index 0000000000..88299779f2 --- /dev/null +++ b/diskann/src/ivf/workingset.rs @@ -0,0 +1,77 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Borrowed clustering views and a reusable canonical insertion batch. + +use diskann_utils::views::rowmajor::{self, Matrix}; +use hashbrown::HashMap; + +use crate::{ANNError, ANNResult, utils::VectorId}; + +/// A point's identity and borrowed canonical vector. +#[derive(Debug, Clone, Copy)] +pub struct PointRef<'a, I> { + pub id: I, + pub vector: &'a [f32], +} + +/// The complete operation-local membership of one list. +/// +/// Only reference metadata is owned; point coordinates are borrowed. +#[derive(Debug)] +pub struct ListView<'a, I, L> { + pub id: L, + pub points: Vec>, +} + +/// Lists supplied by a build accessor, in the order requested by `fill`. +pub type WorkingSet<'a, I, L> = Vec>; + +/// Canonical rows shared between routing and a build accessor's working-set seed. +/// +/// IDs must be unique. The contiguous matrix is owned once; iteration and lookup +/// borrow its rows without copying coordinates. +#[derive(Debug)] +pub struct Batch { + ids: Vec, + vectors: rowmajor::Owned, + rows: HashMap, +} + +impl Batch { + /// # Errors + /// + /// Fails if the matrix has zero columns or its row count differs from `ids`. + pub fn new(ids: Vec, vectors: rowmajor::Owned) -> ANNResult { + if vectors.ncols() == 0 || vectors.nrows() != ids.len() { + return Err(ANNError::message( + "an IVF batch needs one nonempty canonical row per point ID", + )); + } + let rows = ids + .iter() + .copied() + .enumerate() + .map(|(row, id)| (id, row)) + .collect(); + Ok(Self { ids, vectors, rows }) + } + + pub fn vectors(&self) -> rowmajor::Ref<'_, f32> { + self.vectors.as_view() + } + + pub fn points(&self) -> impl Iterator> { + self.ids + .iter() + .copied() + .zip(self.vectors.rows()) + .map(|(id, vector)| PointRef { id, vector }) + } + + pub fn get(&self, id: I) -> Option<&[f32]> { + self.rows.get(&id).map(|&row| self.vectors.row(row)) + } +} diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index eec9ec55ca..bc16753a30 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -4,36 +4,55 @@ |---|---| | **Authors** | Aditya Krishnan | | **Created** | 2026-08-19 | -| **Updated** | 2026-08-24 | +| **Updated** | 2026-10-08 | | **Status** | Draft | ## Summary -Add an accessor-based architecture for incrementally maintained inverted-file -(IVF) indexes, along with co-location sets of centroids to assist with ordering -spatially similar posting lists on slower storage mediums. The architecture -separates the GraphIVF algorithm from point storage, centroid navigation, and -inverted-list access without prescribing a cross-component snapshot or commit -protocol. - -The generic API has two principal extension points: - -1. A query-bound search accessor selects centroids and scans the corresponding - lists. -2. An operation-scoped maintenance accessor presents one unified view of point - data, centroids, reverse assignments, inverted lists, and co-location sets - while planning and applying a split or dissolve. - -Each concrete strategy or aggregate provider is responsible for coordinating -its components. It may use exclusive borrowing, locks, epochs, immutable roots, -database transactions, manifests, or provider-specific recovery internally. -Those mechanisms are deliberately absent from the generic IVF traits. - -The first implementation serializes mutation through `&mut self`. This prevents -search and mutation from overlapping through the same index value and gives an -in-memory provider a straightforward implementation path. More concurrent or -durable providers may implement stronger private guarantees without changing -the algorithm-facing interface. +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. Centroid selection through a coarse accessor. +2. Logical list operations: sizes, appends, complete membership replacement, + creation with a centroid and members, and retirement. +3. One scoped build accessor that supplies borrowed working sets and publishes + the operation with `finish`. + +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. Working-set filling resolves both stored and +prepared points, so clustering has no staged-point boundaries or provenance +classification. + +Insertion provisionally appends points, gathers all oversized parents into one +working set, fits two children per parent, and assigns all regional members +against all child centroids. Parent centroids are not candidates. The algorithm +creates the resulting child lists, retires the parents, and finishes once. +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. + +Working sets borrow vector coordinates. Only fitted centroid coordinates and +assignment/member metadata need new storage. Providers can reuse the canonical +batch and the optional pending-membership helper, or implement their own caches +and private transactions. No delta transport is required by the algorithm; +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. 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 includes initialization and batched insert/split. +Search, neighborhood reassignment policies, deletion/dissolve, and colocation +remain future work. See [the current contracts](../diskann/src/ivf/traits.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 @@ -183,7 +202,7 @@ 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. -## Proposal +## Previous Proposal (Superseded) ### Source Layout @@ -632,7 +651,7 @@ 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. -## Algorithms +## Previous Algorithms (Superseded) ### Bootstrap @@ -933,9 +952,11 @@ accessor is simpler. ### Imperative List Mutation Methods -Adding `append`, `remove`, `move`, `create_list`, and `retire_list` exposes a -partial-write schedule to GraphIVF. A complete declarative update lets each -provider choose a suitable physical implementation. +Immediate, independently published list mutations expose a partial-write schedule +to GraphIVF. The current design instead uses logical list operations inside one +build accessor: they change private operation state and publish 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 From 3f06365ca3898b1f71eca35d3631794dfe8f303c Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 12 Oct 2026 10:29:08 -0400 Subject: [PATCH 15/33] some rejigging --- diskann/src/ivf/clustering.rs | 56 ++++++--- diskann/src/ivf/index.rs | 66 ++++++----- diskann/src/ivf/mod.rs | 5 +- diskann/src/ivf/partition.rs | 31 +++++ diskann/src/ivf/tests.rs | 206 ++++++++++++++++++++++++---------- diskann/src/ivf/traits.rs | 42 ++++--- diskann/src/ivf/workingset.rs | 6 +- rfcs/01187-incremental-ivf.md | 29 +++-- 8 files changed, 306 insertions(+), 135 deletions(-) create mode 100644 diskann/src/ivf/partition.rs diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index fee3ce0731..6f52525ed4 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -9,33 +9,48 @@ use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use rand::{Rng, rngs::StdRng}; -use super::workingset::{ListView, PointRef}; +use super::{ + partition::{Partition, Partitioning}, + workingset::{ListView, PointRef}, +}; use crate::{ANNError, ANNResult}; -pub(super) struct Children { - pub centroids: rowmajor::Owned, - pub members: Vec>, -} - -/// Fit two children per source, then assign every point across all children. -pub(super) fn split( +/// Replace each parent with two children, retaining the other source centroids. +/// +/// `parents` is a distinct subset of `lists`. All regional points choose among +/// the children and retained centroids, never among retiring parent centroids. +pub(super) fn split( lists: &[ListView<'_, I, L>], + parents: &[L], dim: usize, iterations: usize, rng: &mut StdRng, -) -> ANNResult> { +) -> ANNResult> { if lists.is_empty() || dim == 0 { return Err(ANNError::message( "IVF splitting needs a nonempty working set and positive dimension", )); } - let mut centroids = rowmajor::Owned::try_from_element(2 * lists.len(), dim, 0.0)?; + let destinations = lists.len() + parents.len(); + let mut centroids = rowmajor::Owned::try_from_element(destinations, dim, 0.0)?; + let mut partitions = Vec::with_capacity(destinations); let mut sums = vec![0.0; 2 * dim]; - for (list, centers) in lists - .iter() - .zip(centroids.as_mut_slice().chunks_exact_mut(2 * dim)) - { - fit_two_means(&list.points, centers, dim, iterations, rng, &mut sums)?; + for list in lists { + let row = partitions.len(); + if parents.contains(&list.id) { + let centers = &mut centroids.as_mut_slice()[row * dim..(row + 2) * dim]; + fit_two_means(&list.points, centers, dim, iterations, rng, &mut sums)?; + partitions.extend((0..2).map(|_| Partition { + id: None, + members: Vec::new(), + })); + } else { + centroids.row_mut(row).copy_from_slice(list.centroid); + partitions.push(Partition { + id: Some(list.id), + members: Vec::new(), + }); + } } let count = lists.iter().map(|list| list.points.len()).sum(); @@ -46,11 +61,16 @@ pub(super) fn split( assignments.push(to); sizes[to] += 1; } - let mut members: Vec> = sizes.into_iter().map(Vec::with_capacity).collect(); + for (partition, size) in partitions.iter_mut().zip(sizes) { + partition.members = Vec::with_capacity(size); + } for (point, to) in lists.iter().flat_map(|list| &list.points).zip(assignments) { - members[to].push(point.id); + partitions[to].members.push(point.id); } - Ok(Children { centroids, members }) + Ok(Partitioning { + centroids, + partitions, + }) } fn fit_two_means( diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index b096c02903..20c6702069 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -16,6 +16,7 @@ use thiserror::Error; use super::{ clustering, + partition::Partitioning, traits::{ BuildAccessor, BuildStrategy, CoarseAccessor, InsertStrategy, ListAccessor, Provider, }, @@ -72,10 +73,7 @@ impl IVFIndex

{ /// Install initial empty lists and their centroids in one build operation. /// - /// # Errors - /// - /// Fails for an already initialized index, or a provider error. - /// Failed publication is provider-defined. + /// TO DO: needs pub fn initialize<'a, S>( &'a mut self, strategy: &'a S, @@ -100,24 +98,12 @@ impl IVFIndex

{ } for centroid in centroids.rows() { - // TO DO: Need to probably introduce an initialization strategy accessor.lists().create_list(centroid, Vec::new()).await?; } accessor.finish().await } } - /// Route and provisionally append a batch, split oversized lists, then publish. - /// - /// All split parents form one working set. Their members choose among every - /// fitted child, never among retiring parent centroids. Other appends keep their - /// routes. Point coordinates are borrowed throughout fitting and assignment. - /// - /// # Errors - /// - /// Fails before publication if the index is uninitialized or preparation, - /// routing, filling, fitting, or list mutation fails. State after a failed - /// `finish` is provider-defined. An empty batch is a no-op. pub fn insert_batch<'a, S, T>( &'a mut self, strategy: &'a S, @@ -151,32 +137,28 @@ impl IVFIndex

{ let list = selected .first() .ok_or_else(|| ANNError::message("IVF routing returned no live centroid"))? - .id; + .0; + routed.entry(list).or_default().push(point.id); } - // Stable list order also makes seeded fitting independent of hash order. let mut parents = Vec::with_capacity(routed.len()); + for (list, members) in routed { - let mut lists = accessor.lists(); - lists.append_members(list, members).await?; - if lists.list_size(list).await? > config.split_threshold { + let size = accessor.lists().append_members(list, members).await?; + if size > config.split_threshold { parents.push(list); } } + let splits = parents.len(); if !parents.is_empty() { - let children = { + let replacement = { let dim = accessor.dim(); let view = accessor.fill(&parents).await?; - clustering::split(&view, dim, config.two_means_iterations, rng)? + clustering::split(&view, &parents, dim, config.two_means_iterations, rng)? }; - for (centroid, members) in children.centroids.rows().zip(children.members) { - accessor.lists().create_list(centroid, members).await?; - } - for parent in parents { - accessor.lists().retire_list(parent).await?; - } + Self::apply_partitioning(&mut accessor.lists(), &parents, replacement).await?; } accessor.finish().await?; @@ -187,4 +169,30 @@ impl IVFIndex

{ }) } } + + /// Apply destination memberships before retiring the replaced lists. + pub(super) async fn apply_partitioning( + accessor: &mut A, + retired: &[A::ListId], + partitioning: Partitioning, +) -> ANNResult<()> +where + A: ListAccessor, +{ + for (centroid, partition) in partitioning.centroids.rows().zip(partitioning.partitions) { + match partition.id { + Some(id) => accessor.set_members(id, partition.members).await?, + None => { + accessor.create_list(centroid, partition.members).await?; + } + } + } + for &id in retired { + accessor.retire_list(id).await?; + } + Ok(()) +} } + + + diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 7918b5af53..c4e1325633 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -12,13 +12,16 @@ //! Appends and replacements remain private to one build operation until it //! finishes. This lets clustering treat prepared and stored points identically. //! Providers choose their own storage, publication, and failure-recovery mechanisms. +//! A [`partition::Partitioning`] describes new and retained destinations. The index +//! creates new lists, sets retained memberships, and retires split parents. //! //! This first implementation supports initialization and batch insertion with -//! two-means splitting. Search, neighborhood reassignment, deletion, and colocation +//! two-means splitting. Neighborhood selection, search, deletion, and colocation //! are not implemented. mod clustering; mod index; +pub mod partition; pub mod pending; pub mod traits; pub mod workingset; diff --git a/diskann/src/ivf/partition.rs b/diskann/src/ivf/partition.rs new file mode 100644 index 0000000000..a1be388f1a --- /dev/null +++ b/diskann/src/ivf/partition.rs @@ -0,0 +1,31 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Complete destination memberships produced by regional clustering. + +use diskann_utils::views::rowmajor; + +/// One destination list in a [`Partitioning`]. +#[derive(Debug)] +pub struct Partition { + /// `Some` retains a source list's identity and centroid; `None` creates a new list. + pub id: Option, + pub members: Vec, +} + +/// A region's resulting centroids and complete memberships. +/// +/// There is one centroid row per partition; row `i` corresponds to `partitions[i]`. +/// Retained centroid rows must match their source centroids. The member groups +/// together contain every source member exactly once, including pending appends. +/// +/// Coordinates of points are not included: providers already own or have prepared +/// those vectors. Membership buffers can be transferred without copying their +/// contents; centroid rows are available for creating new lists. +#[derive(Debug)] +pub struct Partitioning { + pub centroids: rowmajor::Owned, + pub partitions: Vec>, +} diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index f9c2ffde27..3ff275623f 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -14,10 +14,10 @@ use rand::{SeedableRng, rngs::StdRng}; use super::{ Config, IVFIndex, InsertStats, clustering, + index::apply_partitioning, pending::{Membership, PendingLists}, traits::{ BuildAccessor, BuildStrategy, CoarseAccessor, InsertStrategy, ListAccessor, Provider, - SelectedList, }, workingset::{Batch, ListView, PointRef, WorkingSet}, }; @@ -38,6 +38,8 @@ struct Store { finishes: usize, fills: usize, fail_fill: bool, + fail_create: bool, + fail_retire: bool, fail_finish: bool, } @@ -101,13 +103,9 @@ impl CoarseAccessor for Coarse<'_> { .count() } - fn select( - &mut self, - query: &[f32], - nprobe: usize, - ) -> impl SendFuture>>> { + fn select(&self, query: &[f32], nprobe: usize) -> impl SendFuture>> { async move { - let mut selected: Vec<_> = self + let mut selected: Vec<(u32, f32)> = self .stored .iter() .map(|(&id, list)| (id, list.centroid.as_slice())) @@ -117,16 +115,9 @@ impl CoarseAccessor for Coarse<'_> { .map(|(&id, vector)| (id, vector.as_slice())), ) .filter(|(id, _)| !self.pending.is_retired(*id)) - .map(|(id, vector)| SelectedList { - id, - distance: SquaredL2::evaluate(query, vector), - }) + .map(|(id, vector)| (id, SquaredL2::evaluate(query, vector))) .collect(); - selected.sort_by(|a, b| { - a.distance - .total_cmp(&b.distance) - .then_with(|| a.id.cmp(&b.id)) - }); + selected.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))); selected.truncate(nprobe); Ok(selected) } @@ -139,25 +130,15 @@ impl ListAccessor for Lists<'_, '_> { type InternalId = u32; type ListId = u32; - fn list_size(&mut self, list: u32) -> impl SendFuture> { - async move { - self.0 - .pending - .list_size(list, self.0.stored_members(list)?.len()) - } - } - - fn append_members(&mut self, list: u32, members: Vec) -> impl SendFuture> { + fn append_members( + &mut self, + list: u32, + members: Vec, + ) -> impl SendFuture> { async move { - self.0.stored_members(list)?; - self.0.pending.append(list, members) - } - } - - fn set_members(&mut self, list: u32, members: Vec) -> impl SendFuture> { - async move { - self.0.stored_members(list)?; - self.0.pending.set(list, members) + let stored_size = self.0.stored_members(list)?.len(); + self.0.pending.append(list, members)?; + self.0.pending.list_size(list, stored_size) } } @@ -167,19 +148,35 @@ impl ListAccessor for Lists<'_, '_> { members: Vec, ) -> impl SendFuture> { async move { - if centroid.len() != self.0.store.dim { - return Err(ANNError::message("wrong centroid dimension")); + if self.0.store.fail_create { + return Err(ANNError::message("injected create_list failure")); + } + if centroid.len() != self.0.store.dim || !centroid.iter().all(|value| value.is_finite()) + { + return Err(ANNError::message("invalid centroid")); } let id = self.0.next_list; - self.0.next_list += 1; + self.0.next_list = id + .checked_add(1) + .ok_or_else(|| ANNError::message("list ID overflow"))?; self.0.created.insert(id, centroid.to_vec()); self.0.pending.set(id, members)?; Ok(id) } } + fn set_members(&mut self, list: u32, members: Vec) -> impl SendFuture> { + async move { + self.0.stored_members(list)?; + self.0.pending.set(list, members) + } + } + fn retire_list(&mut self, list: u32) -> impl SendFuture> { async move { + if self.0.store.fail_retire { + return Err(ANNError::message("injected retire_list failure")); + } self.0.stored_members(list)?; self.0.pending.retire(list) } @@ -235,7 +232,22 @@ impl<'p> BuildAccessor for Build<'p> { .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; points.push(PointRef { id, vector }); } - view.push(ListView { id: list, points }); + let centroid = self + .created + .get(&list) + .map(Vec::as_slice) + .or_else(|| { + self.store + .lists + .get(&list) + .map(|stored| stored.centroid.as_slice()) + }) + .ok_or_else(|| ANNError::message(format!("missing centroid {list}")))?; + view.push(ListView { + id: list, + centroid, + points, + }); } Ok(view) } @@ -522,13 +534,16 @@ async fn working_set_borrows_stored_and_prepared_rows() { let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; let batch = Strategy.prepare(&mut build, &points).await.unwrap(); let old_pointer = build.store.points[&10_000].as_ptr(); - build - .lists() - .append_members(1, vec![10_006, 10_005]) - .await - .unwrap(); - assert_eq!(build.lists().list_size(1).await.unwrap(), 4); - { + let centroid_pointer = build.store.lists[&1].centroid.as_ptr(); + assert_eq!( + build + .lists() + .append_members(1, vec![10_006, 10_005]) + .await + .unwrap(), + 4 + ); + let replacement = { let view = require_send(build.fill(&[1])).await.unwrap(); assert_eq!( view[0] @@ -538,6 +553,7 @@ async fn working_set_borrows_stored_and_prepared_rows() { .collect::>(), [10_000, 10_001, 10_006, 10_005] ); + assert_eq!(view[0].centroid.as_ptr(), centroid_pointer); assert_eq!(view[0].points[0].vector.as_ptr(), old_pointer); assert!(std::ptr::eq( view[0].points[2].vector, @@ -547,18 +563,14 @@ async fn working_set_borrows_stored_and_prepared_rows() { view[0].points[3].vector, batch.get(10_005).unwrap() )); - } - build - .lists() - .set_members(1, vec![10_000, 10_005]) + clustering::split(&view, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap() + }; + require_send(apply_partitioning(&mut build.lists(), &[1], replacement)) .await .unwrap(); - assert_eq!(build.lists().list_size(1).await.unwrap(), 2); - assert_eq!(build.fill(&[1]).await.unwrap()[0].points.len(), 2); - build.lists().retire_list(1).await.unwrap(); - assert!(build.lists().list_size(1).await.is_err()); assert!(build.fill(&[1]).await.is_err()); assert!(build.lists().append_members(1, vec![10_001]).await.is_err()); + assert!(build.lists().set_members(1, Vec::new()).await.is_err()); } assert_eq!(store.lists[&1].members, [10_000, 10_001]); assert!(!store.points.contains_key(&10_005)); @@ -567,10 +579,17 @@ async fn working_set_borrows_stored_and_prepared_rows() { #[tokio::test] async fn failed_planning_or_publication_is_not_reported_as_success() { - for fail_fill in [true, false] { + for (fail_fill, fail_create, fail_retire, fail_finish) in [ + (true, false, false, false), + (false, true, false, false), + (false, false, true, false), + (false, false, false, true), + ] { let mut store = seeded(); store.fail_fill = fail_fill; - store.fail_finish = !fail_fill; + store.fail_create = fail_create; + store.fail_retire = fail_retire; + store.fail_finish = fail_finish; let mut index = index(store, 2); let points = [(5, &[1.0, 1.0][..])]; let error = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) @@ -578,6 +597,10 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { .unwrap_err(); let expected = if fail_fill { "injected fill failure" + } else if fail_create { + "injected create_list failure" + } else if fail_retire { + "injected retire_list failure" } else { "injected finish failure" }; @@ -649,6 +672,61 @@ async fn invalid_initialization_and_uninitialized_insertion_fail() { assert_eq!(index.provider().finishes, 1); } +#[tokio::test] +async fn list_mutations_retain_neighbors_and_reassign_both_directions() { + 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(&2).unwrap().centroid = vec![5.0, 0.0]; + let mut build = Strategy + .build_accessor(&mut store, &DefaultContext) + .unwrap(); + let points = [(5, &[5.0, 0.0][..])]; + Strategy.prepare(&mut build, &points).await.unwrap(); + assert_eq!( + build.lists().append_members(1, vec![10_005]).await.unwrap(), + 3 + ); + + let sources = [1, 2]; + let replacement = { + let view = build.fill(&sources).await.unwrap(); + clustering::split(&view, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap() + }; + assert_eq!(replacement.partitions.len(), 3); + let retained = replacement + .partitions + .iter() + .find(|partition| partition.id == Some(2)) + .unwrap(); + assert_eq!(retained.members, [10_005, 10_002]); + assert!( + replacement + .partitions + .iter() + .any(|partition| { partition.id.is_none() && partition.members.contains(&10_003) }) + ); + require_send(apply_partitioning(&mut build.lists(), &[1], replacement)) + .await + .unwrap(); + require_send(build.finish()).await.unwrap(); + + 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_005, 10_002]); + 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]]; @@ -656,6 +734,7 @@ fn assignment_can_cross_parent_boundaries() { let view = vec![ ListView { id: 1u32, + centroid: &[-50.0], points: left .iter() .enumerate() @@ -667,6 +746,7 @@ fn assignment_can_cross_parent_boundaries() { }, ListView { id: 2, + centroid: &[50.0], points: right .iter() .enumerate() @@ -677,11 +757,12 @@ fn assignment_can_cross_parent_boundaries() { .collect(), }, ]; - let children = clustering::split(&view, 1, 10, &mut StdRng::seed_from_u64(7)).unwrap(); + let replacement = + clustering::split(&view, &[1, 2], 1, 10, &mut StdRng::seed_from_u64(7)).unwrap(); assert!( - children.members[2..] + replacement.partitions[2..] .iter() - .any(|members| members.contains(&4)) + .any(|partition| partition.members.contains(&4)) ); } @@ -690,6 +771,7 @@ fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { let rows = [[f32::MAX], [f32::MAX], [f32::MAX]]; let view = vec![ListView { id: 1u32, + centroid: &[f32::MAX], points: rows .iter() .enumerate() @@ -699,8 +781,8 @@ fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { }) .collect(), }]; - let zero = clustering::split(&view, 1, 0, &mut StdRng::seed_from_u64(7)).unwrap(); - let one = clustering::split(&view, 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); + let zero = clustering::split(&view, &[1], 1, 0, &mut StdRng::seed_from_u64(7)).unwrap(); + let one = clustering::split(&view, &[1], 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); assert_eq!(zero.centroids.as_slice(), one.centroids.as_slice()); assert!( zero.centroids @@ -708,6 +790,6 @@ fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { .iter() .all(|value| value.is_finite()) ); - assert_eq!(zero.members[0], [0, 1, 2]); - assert!(zero.members[1].is_empty()); + assert_eq!(zero.partitions[0].members, [0, 1, 2]); + assert!(zero.partitions[1].members.is_empty()); } diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 75953e654d..4ff7cb5584 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -60,41 +60,49 @@ pub trait CoarseAccessor: Send { /// Select `min(nprobe, len())` live lists, nearest first. fn select( - &mut self, + &self, query: &[f32], nprobe: usize, - ) -> impl SendFuture>>>; + ) -> impl SendFuture>>; } -/// Logical list mutations +/// Logical list mutations, private to the enclosing build operation. +/// +/// Reads and working-set fills include pending writes. Publication occurs only +/// when the build accessor finishes. pub trait ListAccessor: Send { type InternalId: VectorId; type ListId: VectorId; - fn list_size(&mut self, list: Self::ListId) -> impl SendFuture>; - - /// Append members to a list. + /// Append members and return the resulting operation-local list size. fn append_members( &mut self, list: Self::ListId, members: Vec, - ) -> impl SendFuture>; - - /// Replace the complete membership of a list. - fn set_members( - &mut self, - list: Self::ListId, - members: Vec, - ) -> impl SendFuture>; + ) -> impl SendFuture>; /// Reserve a fresh list ID and stage its centroid and complete membership. + /// + /// The centroid must be finite and have the canonical dimension. Empty + /// memberships are valid, including during initialization. fn create_list( &mut self, centroid: &[f32], members: Vec, ) -> impl SendFuture>; - /// Stage retirement. All mutations to `list` should be applied first before retiring. + /// Replace a live list's complete membership, leaving its centroid unchanged. + /// + /// This supersedes earlier pending membership changes to the list. + fn set_members( + &mut self, + list: Self::ListId, + members: Vec, + ) -> impl SendFuture>; + + /// Stage retirement of a live list and its centroid. + /// + /// All its members must be placed elsewhere before the build operation finishes. fn retire_list(&mut self, list: Self::ListId) -> impl SendFuture>; } @@ -116,7 +124,9 @@ pub trait BuildAccessor: Send + Sized { /// Make every requested list's complete current membership available. /// - /// Lists appear in request order. + /// Lists appear in request order, with their canonical centroids and every + /// member, including pending appends. Coordinates have the canonical dimension + /// and are finite. Missing required data is an error, not a partial working set. /// TO DO: Make [`WorkingSet`] an associated type here. fn fill( &mut self, diff --git a/diskann/src/ivf/workingset.rs b/diskann/src/ivf/workingset.rs index 88299779f2..ef6e4afc7f 100644 --- a/diskann/src/ivf/workingset.rs +++ b/diskann/src/ivf/workingset.rs @@ -17,12 +17,14 @@ pub struct PointRef<'a, I> { pub vector: &'a [f32], } -/// The complete operation-local membership of one list. +/// One source list's centroid and complete operation-local membership. /// -/// Only reference metadata is owned; point coordinates are borrowed. +/// Only reference metadata is owned; point and centroid coordinates are borrowed. +/// A working set may include both split targets and retained neighboring lists. #[derive(Debug)] pub struct ListView<'a, I, L> { pub id: L, + pub centroid: &'a [f32], pub points: Vec>, } diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index bc16753a30..bf4619fe0f 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -4,7 +4,7 @@ |---|---| | **Authors** | Aditya Krishnan | | **Created** | 2026-08-19 | -| **Updated** | 2026-10-08 | +| **Updated** | 2026-10-12 | | **Status** | Draft | ## Summary @@ -14,8 +14,8 @@ graph index's coarse accessor boundaries. The index owns routing policy, overflow selection, two-means fitting, and regional assignment. Providers supply: 1. Centroid selection through a coarse accessor. -2. Logical list operations: sizes, appends, complete membership replacement, - creation with a centroid and members, and retirement. +2. Logical list operations: append members and return the resulting size, create + a list, set a list's complete membership, and retire a list. 3. One scoped build accessor that supplies borrowed working sets and publishes the operation with `finish`. @@ -24,15 +24,30 @@ build accessor. Preparation reserves internal IDs and supplies a shared, contiguous canonical batch. Pending appends participate in operation-local membership reads and sizes. Working-set filling resolves both stored and prepared points, so clustering has no staged-point boundaries or provenance -classification. +classification. Working sets also borrow each source centroid, allowing a +gathered region to include retained neighbors as well as split targets. Insertion provisionally appends points, gathers all oversized parents into one working set, fits two children per parent, and assigns all regional members against all child centroids. Parent centroids are not candidates. The algorithm -creates the resulting child lists, retires the parents, and finishes once. +applies the resulting partitioning by creating new lists, setting retained-list +memberships, and retiring split parents, then finishes once. 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. +operation. Initialization creates empty lists with the supplied centroids. + +A partitioning contains a contiguous centroid matrix and a complete member group +per destination. Each destination optionally names an existing source list to +retain; unnamed destinations receive fresh list IDs through `create_list`. +The algorithm explicitly retires split parents after applying destination +memberships. Retained lists keep their identities and centroids while +receiving their complete new memberships. The result covers every regional member +exactly once, including pending appends; lists outside the region are unchanged. + +Split targets are a subset of the gathered sources. Numerical fitting replaces +only their centroids, and assignment considers the fitted children and any +retained source centroids. The default insertion path still gathers only split +parents; choosing neighboring lists is not implemented. Working sets borrow vector coordinates. Only fitted centroid coordinates and assignment/member metadata need new storage. Providers can reuse the canonical @@ -46,7 +61,7 @@ a coherent partition; recovery after failed publication is provider-defined. Providers with externally shared state must coordinate those aliases themselves. The current implementation includes initialization and batched insert/split. -Search, neighborhood reassignment policies, deletion/dissolve, and colocation +Search, neighborhood-selection policies, deletion/dissolve, and colocation remain future work. See [the current contracts](../diskann/src/ivf/traits.rs) for the implemented API. From 6c7908dda4e85887e988f7e5ee53f07c8a9f641b Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 12 Oct 2026 10:29:38 -0400 Subject: [PATCH 16/33] missed a bit --- diskann/src/ivf/index.rs | 39 ++++++++++++++++++--------------------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 20c6702069..13cb57eaf7 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -73,7 +73,7 @@ impl IVFIndex

{ /// Install initial empty lists and their centroids in one build operation. /// - /// TO DO: needs + /// TO DO: needs pub fn initialize<'a, S>( &'a mut self, strategy: &'a S, @@ -172,27 +172,24 @@ impl IVFIndex

{ /// Apply destination memberships before retiring the replaced lists. pub(super) async fn apply_partitioning( - accessor: &mut A, - retired: &[A::ListId], - partitioning: Partitioning, -) -> ANNResult<()> -where - A: ListAccessor, -{ - for (centroid, partition) in partitioning.centroids.rows().zip(partitioning.partitions) { - match partition.id { - Some(id) => accessor.set_members(id, partition.members).await?, - None => { - accessor.create_list(centroid, partition.members).await?; + accessor: &mut A, + retired: &[A::ListId], + partitioning: Partitioning, + ) -> ANNResult<()> + where + A: ListAccessor, + { + for (centroid, partition) in partitioning.centroids.rows().zip(partitioning.partitions) { + match partition.id { + Some(id) => accessor.set_members(id, partition.members).await?, + None => { + accessor.create_list(centroid, partition.members).await?; + } } } + for &id in retired { + accessor.retire_list(id).await?; + } + Ok(()) } - for &id in retired { - accessor.retire_list(id).await?; - } - Ok(()) } -} - - - From 93dc333bd910a63ff79eb47a19f35e5ec1a29f30 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 12 Oct 2026 18:29:54 -0400 Subject: [PATCH 17/33] something sensible --- diskann/src/ivf/clustering.rs | 123 +++++--- diskann/src/ivf/index.rs | 108 ++----- diskann/src/ivf/mod.rs | 19 +- diskann/src/ivf/tests.rs | 549 ++++++++++++++++------------------ diskann/src/ivf/traits.rs | 92 ++---- diskann/src/ivf/workingset.rs | 37 ++- rfcs/01187-incremental-ivf.md | 96 +++--- 7 files changed, 487 insertions(+), 537 deletions(-) diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index 6f52525ed4..643c96ea7a 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -11,61 +11,90 @@ use rand::{Rng, rngs::StdRng}; use super::{ partition::{Partition, Partitioning}, - workingset::{ListView, PointRef}, + workingset::RegionSet, }; use crate::{ANNError, ANNResult}; -/// Replace each parent with two children, retaining the other source centroids. +/// Fit two children per parent, assign their members across all children, and record edits. /// -/// `parents` is a distinct subset of `lists`. All regional points choose among -/// the children and retained centroids, never among retiring parent centroids. -pub(super) fn split( - lists: &[ListView<'_, I, L>], +/// Parents are distinct filled sources. Other filled lists are not read or edited. +/// Computation completes before any edits are recorded. A recording error may +/// leave partial pending edits, which the caller must discard rather than publish. +pub(super) fn split>( + region: &mut R, + parents: &[L], + dim: usize, + iterations: usize, + rng: &mut StdRng, +) -> ANNResult<()> { + let result = compute_split(&*region, parents, dim, iterations, rng)?; + for (centroid, partition) in result.centroids.rows().zip(result.partitions) { + match partition.id { + Some(id) => region.set_members(id, &partition.members)?, + None => region.create(centroid, &partition.members)?, + } + } + for &parent in parents { + region.retire(parent)?; + } + Ok(()) +} + +fn compute_split>( + region: &R, parents: &[L], dim: usize, iterations: usize, rng: &mut StdRng, ) -> ANNResult> { - if lists.is_empty() || dim == 0 { + if parents.is_empty() || dim == 0 { return Err(ANNError::message( - "IVF splitting needs a nonempty working set and positive dimension", + "IVF splitting needs parents and a positive dimension", )); } - let destinations = lists.len() + parents.len(); - let mut centroids = rowmajor::Owned::try_from_element(destinations, dim, 0.0)?; - let mut partitions = Vec::with_capacity(destinations); + let mut inputs = Vec::with_capacity(parents.len()); + for &parent in parents { + let view = region + .get_view(parent) + .ok_or_else(|| ANNError::message("IVF split parent is missing from the region"))?; + let vectors = region.get_vectors(parent).ok_or_else(|| { + ANNError::message("IVF split parent vectors are missing from the region") + })?; + inputs.push((view.members, vectors)); + } + + let mut centroids = rowmajor::Owned::try_from_element(2 * parents.len(), dim, 0.0)?; let mut sums = vec![0.0; 2 * dim]; - for list in lists { - let row = partitions.len(); - if parents.contains(&list.id) { - let centers = &mut centroids.as_mut_slice()[row * dim..(row + 2) * dim]; - fit_two_means(&list.points, centers, dim, iterations, rng, &mut sums)?; - partitions.extend((0..2).map(|_| Partition { - id: None, - members: Vec::new(), - })); - } else { - centroids.row_mut(row).copy_from_slice(list.centroid); - partitions.push(Partition { - id: Some(list.id), - members: Vec::new(), - }); - } + for ((_, vectors), centers) in inputs + .iter() + .zip(centroids.as_mut_slice().chunks_exact_mut(2 * dim)) + { + fit_two_means(vectors.as_view(), centers, dim, iterations, rng, &mut sums)?; } - let count = lists.iter().map(|list| list.points.len()).sum(); + let count = inputs.iter().map(|(members, _)| members.len()).sum(); let mut assignments = Vec::with_capacity(count); let mut sizes = vec![0; centroids.nrows()]; - for point in lists.iter().flat_map(|list| &list.points) { - let to = nearest(point.vector, centroids.as_slice(), dim); - assignments.push(to); - sizes[to] += 1; - } - for (partition, size) in partitions.iter_mut().zip(sizes) { - partition.members = Vec::with_capacity(size); + for (_, vectors) in &inputs { + for vector in vectors.rows() { + let to = nearest(vector, centroids.as_slice(), dim); + assignments.push(to); + sizes[to] += 1; + } } - for (point, to) in lists.iter().flat_map(|list| &list.points).zip(assignments) { - partitions[to].members.push(point.id); + let mut partitions: Vec> = sizes + .into_iter() + .map(|size| Partition { + id: None, + members: Vec::with_capacity(size), + }) + .collect(); + for (id, to) in inputs + .iter() + .flat_map(|(members, _)| members.iter().copied()) + .zip(assignments) + { + partitions[to].members.push(id); } Ok(Partitioning { centroids, @@ -73,29 +102,29 @@ pub(super) fn split( }) } -fn fit_two_means( - points: &[PointRef<'_, I>], +fn fit_two_means( + points: rowmajor::Ref<'_, f32>, centers: &mut [f32], dim: usize, iterations: usize, rng: &mut StdRng, sums: &mut [f64], ) -> ANNResult<()> { - if points.len() < 2 { + if points.nrows() < 2 { return Err(ANNError::message("an IVF split needs at least two points")); } - let first = rng.random_range(0..points.len()); - let second = (first + rng.random_range(1..points.len())) % points.len(); - centers[..dim].copy_from_slice(points[first].vector); - centers[dim..].copy_from_slice(points[second].vector); + let first = rng.random_range(0..points.nrows()); + let second = (first + rng.random_range(1..points.nrows())) % points.nrows(); + centers[..dim].copy_from_slice(points.row(first)); + centers[dim..].copy_from_slice(points.row(second)); for _ in 0..iterations.max(1) { sums.fill(0.0); let mut counts = [0; 2]; - for point in points { - let to = nearest(point.vector, centers, dim); + for point in points.rows() { + let to = nearest(point, centers, dim); counts[to] += 1; - for (sum, &value) in sums[to * dim..(to + 1) * dim].iter_mut().zip(point.vector) { + for (sum, &value) in sums[to * dim..(to + 1) * dim].iter_mut().zip(point) { // A finite mean must not overflow while accumulating its inputs. *sum += f64::from(value); } diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 13cb57eaf7..7efa883331 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -7,19 +7,14 @@ use std::collections::BTreeMap; -use diskann_utils::{ - future::SendFuture, - views::rowmajor::{self, Matrix}, -}; +use diskann_utils::future::SendFuture; use rand::{SeedableRng, rngs::StdRng}; use thiserror::Error; use super::{ clustering, - partition::Partitioning, - traits::{ - BuildAccessor, BuildStrategy, CoarseAccessor, InsertStrategy, ListAccessor, Provider, - }, + traits::{BuildAccessor, BuildStrategy, InsertStrategy, Provider, Router}, + workingset::RegionSet, }; use crate::{ANNError, ANNResult, error::IntoANNResult}; @@ -71,39 +66,6 @@ impl IVFIndex

{ &self.provider } - /// Install initial empty lists and their centroids in one build operation. - /// - /// TO DO: needs - pub fn initialize<'a, S>( - &'a mut self, - strategy: &'a S, - context: &'a P::Context, - centroids: rowmajor::Ref<'_, f32>, - ) -> impl SendFuture> - where - S: BuildStrategy<'a, P>, - { - async move { - if centroids.nrows() == 0 || centroids.ncols() == 0 { - return Err(ANNError::message( - "IVF initialization needs nonempty centroids", - )); - } - let mut accessor = strategy - .build_accessor(&mut self.provider, context) - .into_ann_result()?; - - if !accessor.coarse().is_empty() { - return Err(ANNError::message("the IVF index is already initialized")); - } - - for centroid in centroids.rows() { - accessor.lists().create_list(centroid, Vec::new()).await?; - } - accessor.finish().await - } - } - pub fn insert_batch<'a, S, T>( &'a mut self, strategy: &'a S, @@ -125,43 +87,48 @@ impl IVFIndex

{ let mut accessor = strategy .build_accessor(provider, context) .into_ann_result()?; - if accessor.coarse().is_empty() { + if accessor.router().is_empty() { return Err(ANNError::message("the IVF index is not initialized")); } let batch = strategy.prepare(&mut accessor, points).await?; let mut routed = BTreeMap::>::new(); - for point in batch.points() { - let selected = accessor.coarse().select(point.vector, 1).await?; - let list = selected - .first() - .ok_or_else(|| ANNError::message("IVF routing returned no live centroid"))? - .0; + { + let router = accessor.router(); + + for point in batch.points() { + let selected = router.select(point.vector, 1).await?; + let list = selected + .first() + .ok_or_else(|| ANNError::message("IVF routing returned no live centroid"))? + .0; - routed.entry(list).or_default().push(point.id); + routed.entry(list).or_default().push(point.id); + } } let mut parents = Vec::with_capacity(routed.len()); + // append the points to their respective list ids and determine which need to be split. for (list, members) in routed { - let size = accessor.lists().append_members(list, members).await?; + let size = accessor.append_to(list, &members).await?; if size > config.split_threshold { parents.push(list); } } let splits = parents.len(); - if !parents.is_empty() { - let replacement = { - let dim = accessor.dim(); - let view = accessor.fill(&parents).await?; - clustering::split(&view, &parents, dim, config.two_means_iterations, rng)? - }; - Self::apply_partitioning(&mut accessor.lists(), &parents, replacement).await?; - } + let update = if !parents.is_empty() { + let dim = accessor.dim(); + let mut view = accessor.fill(&parents).await?; + clustering::split(&mut view, &parents, dim, config.two_means_iterations, rng)?; + Some(view.into_update()) + } else { + None + }; - accessor.finish().await?; + accessor.finish(update).await?; Ok(InsertStats { inserted: points.len(), @@ -169,27 +136,4 @@ impl IVFIndex

{ }) } } - - /// Apply destination memberships before retiring the replaced lists. - pub(super) async fn apply_partitioning( - accessor: &mut A, - retired: &[A::ListId], - partitioning: Partitioning, - ) -> ANNResult<()> - where - A: ListAccessor, - { - for (centroid, partition) in partitioning.centroids.rows().zip(partitioning.partitions) { - match partition.id { - Some(id) => accessor.set_members(id, partition.members).await?, - None => { - accessor.create_list(centroid, partition.members).await?; - } - } - } - for &id in retired { - accessor.retire_list(id).await?; - } - Ok(()) - } } diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index c4e1325633..281f05d989 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -6,18 +6,19 @@ //! Incremental IVF construction. //! //! The index routes points, splits oversized lists, and assigns their members to -//! replacement centroids. Providers supply centroid selection, logical list -//! operations, and borrowed working sets through [`traits::BuildAccessor`]. +//! replacement centroids. Providers supply routing and operation-local appends +//! through [`traits::BuildAccessor`], and borrowed inputs and private regional +//! edits through [`workingset::RegionSet`]. //! -//! Appends and replacements remain private to one build operation until it -//! finishes. This lets clustering treat prepared and stored points identically. +//! Clustering finishes reading its inputs before recording child creation and +//! parent retirement. Consuming the region releases its input borrows and yields +//! an owned update; the accessor publishes it together with ordinary appends once. //! Providers choose their own storage, publication, and failure-recovery mechanisms. -//! A [`partition::Partitioning`] describes new and retained destinations. The index -//! creates new lists, sets retained memberships, and retires split parents. //! -//! This first implementation supports initialization and batch insertion with -//! two-means splitting. Neighborhood selection, search, deletion, and colocation -//! are not implemented. +//! This first implementation supports batch insertion into an already initialized +//! provider, with two-means splitting of the explicitly requested parents. +//! Initialization, neighborhood reassignment, search, deletion, and colocation are +//! not implemented. mod clustering; mod index; diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index 3ff275623f..e0ddc3617e 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -14,12 +14,9 @@ use rand::{SeedableRng, rngs::StdRng}; use super::{ Config, IVFIndex, InsertStats, clustering, - index::apply_partitioning, pending::{Membership, PendingLists}, - traits::{ - BuildAccessor, BuildStrategy, CoarseAccessor, InsertStrategy, ListAccessor, Provider, - }, - workingset::{Batch, ListView, PointRef, WorkingSet}, + traits::{BuildAccessor, BuildStrategy, HasIds, InsertStrategy, Provider, Router}, + workingset::{Batch, ListView, RegionSet}, }; use crate::{ANNError, ANNResult, provider::DefaultContext}; @@ -68,39 +65,20 @@ impl Provider for Store { struct Build<'p> { store: &'p mut Store, pending: PendingLists, - created: BTreeMap>, batch: Option>>, - next_list: u32, + sources: BTreeMap, fills: Cell, } -impl Build<'_> { - fn stored_members(&self, list: u32) -> ANNResult<&[u32]> { - if let Some(stored) = self.store.lists.get(&list) { - Ok(&stored.members) - } else if self.created.contains_key(&list) { - Ok(&[]) - } else { - Err(ANNError::message(format!("unknown list {list}"))) - } - } -} - struct Coarse<'a> { stored: &'a BTreeMap, - created: &'a BTreeMap>, - pending: &'a PendingLists, } -impl CoarseAccessor for Coarse<'_> { +impl Router for Coarse<'_> { type ListId = u32; fn len(&self) -> usize { - self.stored - .keys() - .chain(self.created.keys()) - .filter(|&&id| !self.pending.is_retired(id)) - .count() + self.stored.len() } fn select(&self, query: &[f32], nprobe: usize) -> impl SendFuture>> { @@ -108,14 +86,7 @@ impl CoarseAccessor for Coarse<'_> { let mut selected: Vec<(u32, f32)> = self .stored .iter() - .map(|(&id, list)| (id, list.centroid.as_slice())) - .chain( - self.created - .iter() - .map(|(&id, vector)| (id, vector.as_slice())), - ) - .filter(|(id, _)| !self.pending.is_retired(*id)) - .map(|(id, vector)| (id, SquaredL2::evaluate(query, vector))) + .map(|(&id, list)| (id, SquaredL2::evaluate(query, list.centroid.as_slice()))) .collect(); selected.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))); selected.truncate(nprobe); @@ -124,74 +95,112 @@ impl CoarseAccessor for Coarse<'_> { } } -struct Lists<'a, 'p>(&'a mut Build<'p>); +struct Source { + centroid: Vec, + members: Vec, + vectors: rowmajor::Owned, +} -impl ListAccessor for Lists<'_, '_> { - type InternalId = u32; - type ListId = u32; +#[derive(Debug, Default)] +struct Update { + memberships: PendingLists, + created: Vec, +} - fn append_members( - &mut self, - list: u32, - members: Vec, - ) -> impl SendFuture> { - async move { - let stored_size = self.0.stored_members(list)?.len(); - self.0.pending.append(list, members)?; - self.0.pending.list_size(list, stored_size) +struct Region<'a> { + sources: &'a BTreeMap, + dim: usize, + update: Update, + fail_create: bool, + fail_retire: bool, +} + +impl<'a> Region<'a> { + fn new(sources: &'a BTreeMap, dim: usize) -> Self { + Self { + sources, + dim, + update: Update::default(), + fail_create: false, + fail_retire: false, } } - fn create_list( - &mut self, - centroid: &[f32], - members: Vec, - ) -> impl SendFuture> { - async move { - if self.0.store.fail_create { - return Err(ANNError::message("injected create_list failure")); - } - if centroid.len() != self.0.store.dim || !centroid.iter().all(|value| value.is_finite()) - { - return Err(ANNError::message("invalid centroid")); - } - let id = self.0.next_list; - self.0.next_list = id - .checked_add(1) - .ok_or_else(|| ANNError::message("list ID overflow"))?; - self.0.created.insert(id, centroid.to_vec()); - self.0.pending.set(id, members)?; - Ok(id) - } + fn source(&self, list: u32) -> ANNResult<&Source> { + self.sources + .get(&list) + .ok_or_else(|| ANNError::message(format!("unknown list {list}"))) } +} - fn set_members(&mut self, list: u32, members: Vec) -> impl SendFuture> { - async move { - self.0.stored_members(list)?; - self.0.pending.set(list, members) +impl RegionSet for Region<'_> { + type Update = Update; + + fn get_view(&self, list: u32) -> Option> { + self.sources.get(&list).map(|source| ListView { + 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()) + } + + fn set_members(&mut self, list: u32, members: &[u32]) -> ANNResult<()> { + self.source(list)?; + self.update.memberships.set(list, members.to_vec()) + } + + fn append_members(&mut self, list: u32, members: &[u32]) -> ANNResult { + let size = self.source(list)?.members.len(); + self.update.memberships.append(list, members.to_vec())?; + self.update.memberships.list_size(list, size) + } + + fn retire(&mut self, list: u32) -> ANNResult<()> { + if self.fail_retire { + return Err(ANNError::message("injected retire failure")); } + self.source(list)?; + self.update.memberships.retire(list) } - fn retire_list(&mut self, list: u32) -> impl SendFuture> { - async move { - if self.0.store.fail_retire { - return Err(ANNError::message("injected retire_list failure")); - } - self.0.stored_members(list)?; - self.0.pending.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.update.created.push(StoredList { + centroid: centroid.to_vec(), + members: members.to_vec(), + }); + Ok(()) + } + + fn into_update(self) -> Self::Update { + self.update } } -impl<'p> BuildAccessor for Build<'p> { +impl HasIds for Build<'_> { type InternalId = u32; type ListId = u32; - type Coarse<'a> +} + +impl<'p> BuildAccessor for Build<'p> { + type Update = Update; + type Router<'a> = Coarse<'a> where Self: 'a; - type Lists<'a> - = Lists<'a, 'p> + type Region<'a> + = Region<'a> where Self: 'a; @@ -199,80 +208,90 @@ impl<'p> BuildAccessor for Build<'p> { self.store.dim } - fn coarse(&mut self) -> Self::Coarse<'_> { + fn router(&mut self) -> Self::Router<'_> { Coarse { stored: &self.store.lists, - created: &self.created, - pending: &self.pending, } } - fn lists(&mut self) -> Self::Lists<'_> { - Lists(self) + fn append_to(&mut self, list: u32, members: &[u32]) -> impl SendFuture> { + async move { + let size = self + .store + .lists + .get(&list) + .ok_or_else(|| ANNError::message(format!("unknown list {list}")))? + .members + .len(); + self.pending.append(list, members.to_vec())?; + self.pending.list_size(list, size) + } } - fn fill(&mut self, lists: &[u32]) -> impl SendFuture>> { + fn fill(&mut self, lists: &[u32]) -> impl SendFuture>> { async move { self.fills.set(self.fills.get() + 1); if self.store.fail_fill { return Err(ANNError::message("injected fill failure")); } - let mut view = Vec::with_capacity(lists.len()); + self.sources.clear(); for &list in lists { - let size = self - .pending - .list_size(list, self.stored_members(list)?.len())?; - let mut points = Vec::with_capacity(size); - for id in self.pending.members(list, self.stored_members(list)?)? { + 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(), 0.0); + for (&id, row) in members.iter().zip(vectors.rows_mut()) { let vector = self .batch .as_ref() .and_then(|batch| batch.get(id)) .or_else(|| self.store.points.get(&id).map(Vec::as_slice)) .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; - points.push(PointRef { id, vector }); + row.copy_from_slice(vector); } - let centroid = self - .created - .get(&list) - .map(Vec::as_slice) - .or_else(|| { - self.store - .lists - .get(&list) - .map(|stored| stored.centroid.as_slice()) - }) - .ok_or_else(|| ANNError::message(format!("missing centroid {list}")))?; - view.push(ListView { - id: list, - centroid, - points, - }); + self.sources.insert( + list, + Source { + centroid: stored.centroid.clone(), + members, + vectors, + }, + ); } - Ok(view) + let mut region = Region::new(&self.sources, self.dim()); + region.fail_create = self.store.fail_create; + region.fail_retire = self.store.fail_retire; + Ok(region) } } - fn finish(self) -> impl SendFuture> { + fn finish(&mut self, update: Self::Update) -> impl SendFuture> { async move { if self.store.fail_finish { return Err(ANNError::message("injected finish failure")); } - if let Some(batch) = self.batch { + let created = u32::try_from(update.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"))?; + if let Some(batch) = self.batch.take() { for point in batch.points() { self.store.points.insert(point.id, point.vector.to_vec()); } } - for (id, centroid) in self.created { - self.store.lists.insert( - id, - StoredList { - centroid, - members: Vec::new(), - }, - ); + for (id, list) in (self.store.next_list..next_list).zip(update.created) { + self.store.lists.insert(id, list); } - for (id, change) in self.pending.into_changes() { + for (id, change) in std::mem::take(&mut self.pending) + .into_changes() + .chain(update.memberships.into_changes()) + { match change { Membership::Append(members) => { self.store @@ -290,7 +309,7 @@ impl<'p> BuildAccessor for Build<'p> { } } } - self.store.next_list = self.next_list; + self.store.next_list = next_list; self.store.finishes += 1; self.store.fills += self.fills.get(); Ok(()) @@ -309,13 +328,11 @@ impl<'a> BuildStrategy<'a, Store> for Strategy { store: &'a mut Store, _: &'a DefaultContext, ) -> ANNResult> { - let next_list = store.next_list; Ok(Build { store, pending: PendingLists::default(), - created: BTreeMap::new(), batch: None, - next_list, + sources: BTreeMap::new(), fills: Cell::new(0), }) } @@ -412,19 +429,32 @@ fn require_send(future: F) -> F { } #[tokio::test] -async fn initialize_and_append_publish_actual_memberships() { +async fn appends_publish_actual_memberships_without_filling() { let mut index = index( Store { dim: 2, - next_list: 1, + 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, ..Store::default() }, 2, ); - let centers = matrix(&[[0.0, 0.0], [10.0, 0.0]]); - require_send(index.initialize(&Strategy, &DefaultContext, centers.as_view())) - .await - .unwrap(); let points = [ (7, &[0.0, 1.0][..]), (8, &[1.0, 0.0][..]), @@ -458,7 +488,7 @@ async fn initialize_and_append_publish_actual_memberships() { .unwrap(), 7 ); - assert_eq!(index.provider().finishes, 2); + assert_eq!(index.provider().finishes, 1); assert_eq!(index.provider().fills, 0); let empty: [(u32, &[f32]); 0] = []; @@ -469,7 +499,7 @@ async fn initialize_and_append_publish_actual_memberships() { .unwrap(), InsertStats::default() ); - assert_eq!(index.provider().finishes, 2); + assert_eq!(index.provider().finishes, 1); } #[tokio::test] @@ -525,53 +555,50 @@ async fn splitting_commits_children_and_ordinary_appends_once() { } #[tokio::test] -async fn working_set_borrows_stored_and_prepared_rows() { +async fn region_borrows_filled_rows_and_extracts_owned_edits() { let mut store = seeded(); - { + let update = { let mut build = Strategy .build_accessor(&mut store, &DefaultContext) .unwrap(); let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; let batch = Strategy.prepare(&mut build, &points).await.unwrap(); - let old_pointer = build.store.points[&10_000].as_ptr(); - let centroid_pointer = build.store.lists[&1].centroid.as_ptr(); + assert_eq!(build.append_to(1, &[10_006, 10_005]).await.unwrap(), 4); + let mut region = require_send(build.fill(&[1])).await.unwrap(); + let view = region.get_view(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!( - build - .lists() - .append_members(1, vec![10_006, 10_005]) - .await - .unwrap(), - 4 + vectors.as_slice().as_ptr(), + source.vectors.as_slice().as_ptr() ); - let replacement = { - let view = require_send(build.fill(&[1])).await.unwrap(); - assert_eq!( - view[0] - .points - .iter() - .map(|point| point.id) - .collect::>(), - [10_000, 10_001, 10_006, 10_005] - ); - assert_eq!(view[0].centroid.as_ptr(), centroid_pointer); - assert_eq!(view[0].points[0].vector.as_ptr(), old_pointer); - assert!(std::ptr::eq( - view[0].points[2].vector, - batch.get(10_006).unwrap() - )); - assert!(std::ptr::eq( - view[0].points[3].vector, - batch.get(10_005).unwrap() - )); - clustering::split(&view, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap() - }; - require_send(apply_partitioning(&mut build.lists(), &[1], replacement)) - .await - .unwrap(); - assert!(build.fill(&[1]).await.is_err()); - assert!(build.lists().append_members(1, vec![10_001]).await.is_err()); - assert!(build.lists().set_members(1, Vec::new()).await.is_err()); - } + assert_eq!(vectors.row(0), [0.0, 0.0]); + assert_eq!(vectors.row(2), batch.get(10_006).unwrap()); + assert_eq!(vectors.row(3), batch.get(10_005).unwrap()); + clustering::split(&mut region, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap(); + assert!(region.append_members(1, &[10_001]).is_err()); + assert!(region.set_members(1, &[]).is_err()); + let update = region.into_update(); + assert_eq!(build.dim(), 2); + update + }; + assert_eq!(update.created.len(), 2); + let mut members: Vec<_> = update + .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!( + update + .memberships + .into_changes() + .any(|(id, change)| id == 1 && matches!(change, Membership::Retire)) + ); assert_eq!(store.lists[&1].members, [10_000, 10_001]); assert!(!store.points.contains_key(&10_005)); assert_eq!(store.finishes, 0); @@ -598,9 +625,9 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { let expected = if fail_fill { "injected fill failure" } else if fail_create { - "injected create_list failure" + "injected create failure" } else if fail_retire { - "injected retire_list failure" + "injected retire failure" } else { "injected finish failure" }; @@ -618,7 +645,7 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { } #[tokio::test] -async fn invalid_initialization_and_uninitialized_insertion_fail() { +async fn invalid_config_and_uninitialized_insertion_fail() { assert!( IVFIndex::new( Store::default(), @@ -638,18 +665,6 @@ async fn invalid_initialization_and_uninitialized_insertion_fail() { }, 2, ); - for centers in [ - matrix::<2>(&[]), - matrix(&[[0.0]]), - matrix(&[[f32::NAN, 0.0]]), - ] { - assert!( - index - .initialize(&Strategy, &DefaultContext, centers.as_view()) - .await - .is_err() - ); - } let points = [(7, &[0.0, 1.0][..])]; assert!( index @@ -658,22 +673,10 @@ async fn invalid_initialization_and_uninitialized_insertion_fail() { .is_err() ); assert_eq!(index.provider().finishes, 0); - let centers = matrix(&[[0.0, 0.0]]); - index - .initialize(&Strategy, &DefaultContext, centers.as_view()) - .await - .unwrap(); - assert!( - index - .initialize(&Strategy, &DefaultContext, centers.as_view()) - .await - .is_err() - ); - assert_eq!(index.provider().finishes, 1); } #[tokio::test] -async fn list_mutations_retain_neighbors_and_reassign_both_directions() { +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]); @@ -684,38 +687,22 @@ async fn list_mutations_retain_neighbors_and_reassign_both_directions() { .unwrap(); let points = [(5, &[5.0, 0.0][..])]; Strategy.prepare(&mut build, &points).await.unwrap(); - assert_eq!( - build.lists().append_members(1, vec![10_005]).await.unwrap(), - 3 - ); - - let sources = [1, 2]; - let replacement = { - let view = build.fill(&sources).await.unwrap(); - clustering::split(&view, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap() - }; - assert_eq!(replacement.partitions.len(), 3); - let retained = replacement - .partitions - .iter() - .find(|partition| partition.id == Some(2)) - .unwrap(); - assert_eq!(retained.members, [10_005, 10_002]); - assert!( - replacement - .partitions - .iter() - .any(|partition| { partition.id.is_none() && partition.members.contains(&10_003) }) - ); - require_send(apply_partitioning(&mut build.lists(), &[1], replacement)) - .await - .unwrap(); - require_send(build.finish()).await.unwrap(); + assert_eq!(build.append_to(1, &[10_005]).await.unwrap(), 3); + + let mut region = build.fill(&[1, 2]).await.unwrap(); + assert!(clustering::split(&mut region, &[1, 99], 2, 5, &mut StdRng::seed_from_u64(7)).is_err()); + let failed = std::mem::take(&mut region.update); + assert!(failed.created.is_empty()); + assert_eq!(failed.memberships.into_changes().count(), 0); + clustering::split(&mut region, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap(); + let update = region.into_update(); + assert_eq!(update.created.len(), 2); + require_send(build.finish(update)).await.unwrap(); 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_005, 10_002]); + 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 @@ -731,36 +718,30 @@ async fn list_mutations_retain_neighbors_and_reassign_both_directions() { 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 view = vec![ - ListView { - id: 1u32, - centroid: &[-50.0], - points: left - .iter() - .enumerate() - .map(|(id, vector)| PointRef { - id: id as u32, - vector, - }) - .collect(), - }, - ListView { - id: 2, - centroid: &[50.0], - points: right - .iter() - .enumerate() - .map(|(id, vector)| PointRef { - id: 100 + id as u32, - vector, - }) - .collect(), - }, - ]; - let replacement = - clustering::split(&view, &[1, 2], 1, 10, &mut StdRng::seed_from_u64(7)).unwrap(); + 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 mut region = Region::new(&sources, 1); + clustering::split(&mut region, &[1, 2], 1, 10, &mut StdRng::seed_from_u64(7)).unwrap(); + let update = region.into_update(); assert!( - replacement.partitions[2..] + update.created[2..] .iter() .any(|partition| partition.members.contains(&4)) ); @@ -769,27 +750,27 @@ fn assignment_can_cross_parent_boundaries() { #[test] fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { let rows = [[f32::MAX], [f32::MAX], [f32::MAX]]; - let view = vec![ListView { - id: 1u32, - centroid: &[f32::MAX], - points: rows - .iter() - .enumerate() - .map(|(id, row)| PointRef { - id: id as u32, - vector: row, - }) - .collect(), - }]; - let zero = clustering::split(&view, &[1], 1, 0, &mut StdRng::seed_from_u64(7)).unwrap(); - let one = clustering::split(&view, &[1], 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); - assert_eq!(zero.centroids.as_slice(), one.centroids.as_slice()); - assert!( - zero.centroids - .as_slice() - .iter() - .all(|value| value.is_finite()) - ); - assert_eq!(zero.partitions[0].members, [0, 1, 2]); - assert!(zero.partitions[1].members.is_empty()); + let sources = [( + 1, + Source { + centroid: vec![f32::MAX], + members: vec![0, 1, 2], + vectors: matrix(&rows), + }, + )] + .into(); + let mut region = Region::new(&sources, 1); + clustering::split(&mut region, &[1], 1, 0, &mut StdRng::seed_from_u64(7)).unwrap(); + let zero = region.into_update(); + let mut region = Region::new(&sources, 1); + clustering::split(&mut region, &[1], 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); + let one = region.into_update(); + assert_eq!(zero.created.len(), 2); + assert_eq!(one.created.len(), 2); + for (zero, one) in zero.created.iter().zip(&one.created) { + assert_eq!(zero.centroid, one.centroid); + assert!(zero.centroid.iter().all(|value| value.is_finite())); + } + assert_eq!(zero.created[0].members, [0, 1, 2]); + assert!(zero.created[1].members.is_empty()); } diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 4ff7cb5584..943b58e80e 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -9,10 +9,11 @@ use std::{fmt::Debug, sync::Arc}; use diskann_utils::future::SendFuture; -use super::workingset::{Batch, WorkingSet}; +use super::workingset::Batch; use crate::{ ANNResult, error::{StandardError, ToRanked}, + ivf::workingset::RegionSet, provider::ExecutionContext, utils::VectorId, }; @@ -41,15 +42,13 @@ pub trait Provider: Sized + Send + Sync + 'static { ) -> Result; } -/// A live list selected by centroid distance. -#[derive(Debug, Clone, Copy)] -pub struct SelectedList { - pub id: L, - pub distance: f32, +pub trait HasIds { + type InternalId: VectorId; + type ListId: VectorId; } /// Centroid selection within one build operation. -pub trait CoarseAccessor: Send { +pub trait Router: Send { type ListId: VectorId; fn len(&self) -> usize; @@ -66,75 +65,36 @@ pub trait CoarseAccessor: Send { ) -> impl SendFuture>>; } -/// Logical list mutations, private to the enclosing build operation. -/// -/// Reads and working-set fills include pending writes. Publication occurs only -/// when the build accessor finishes. -pub trait ListAccessor: Send { - type InternalId: VectorId; - type ListId: VectorId; - - /// Append members and return the resulting operation-local list size. - fn append_members( - &mut self, - list: Self::ListId, - members: Vec, - ) -> impl SendFuture>; - - /// Reserve a fresh list ID and stage its centroid and complete membership. - /// - /// The centroid must be finite and have the canonical dimension. Empty - /// memberships are valid, including during initialization. - fn create_list( - &mut self, - centroid: &[f32], - members: Vec, - ) -> impl SendFuture>; - - /// Replace a live list's complete membership, leaving its centroid unchanged. - /// - /// This supersedes earlier pending membership changes to the list. - fn set_members( - &mut self, - list: Self::ListId, - members: Vec, - ) -> impl SendFuture>; - - /// Stage retirement of a live list and its centroid. - /// - /// All its members must be placed elsewhere before the build operation finishes. - fn retire_list(&mut self, list: Self::ListId) -> impl SendFuture>; -} - -pub trait BuildAccessor: Send + Sized { - type InternalId: VectorId; - type ListId: VectorId; +pub trait BuildAccessor: Send + Sized + HasIds { + /// Owned regional edits. The default records no regional changes. + type Update: Send; - type Coarse<'a>: CoarseAccessor + type Router<'a>: Router where Self: 'a; - type Lists<'a>: ListAccessor + + type Region<'a>: RegionSet where Self: 'a; /// Positive dimension of all canonical points and centroids. fn dim(&self) -> usize; - fn coarse(&mut self) -> Self::Coarse<'_>; - fn lists(&mut self) -> Self::Lists<'_>; - - /// Make every requested list's complete current membership available. - /// - /// Lists appear in request order, with their canonical centroids and every - /// member, including pending appends. Coordinates have the canonical dimension - /// and are finite. Missing required data is an error, not a partial working set. - /// TO DO: Make [`WorkingSet`] an associated type here. - fn fill( + + fn router(&mut self) -> Self::Router<'_>; + + /// Append members and return the resulting operation-local list size. + fn append_to( &mut self, - lists: &[Self::ListId], - ) -> impl SendFuture>>; + list: Self::ListId, + members: &[Self::InternalId], + ) -> impl SendFuture>; + + /// Make every requested list's complete operation-local membership, vectors + /// and centroids available for computation. + fn fill(&mut self, lists: &[Self::ListId]) -> impl SendFuture>>; - /// Publish the operation's point mappings, centroids, and memberships together. - fn finish(self) -> impl SendFuture>; + /// Publish accessor-side appends and owned regional edits together. + fn finish(&mut self, update: Option) -> impl SendFuture>; } pub trait BuildStrategy<'a, P: Provider>: Send + Sync { diff --git a/diskann/src/ivf/workingset.rs b/diskann/src/ivf/workingset.rs index ef6e4afc7f..ec91a65d2e 100644 --- a/diskann/src/ivf/workingset.rs +++ b/diskann/src/ivf/workingset.rs @@ -20,18 +20,45 @@ pub struct PointRef<'a, I> { /// One source list's centroid and complete operation-local membership. /// /// Only reference metadata is owned; point and centroid coordinates are borrowed. -/// A working set may include both split targets and retained neighboring lists. +/// A region may include both split targets and neighboring lists. #[derive(Debug)] pub struct ListView<'a, I, L> { pub id: L, pub centroid: &'a [f32], - pub points: Vec>, + pub members: &'a [I], } -/// Lists supplied by a build accessor, in the order requested by `fill`. -pub type WorkingSet<'a, I, L> = Vec>; +/// Provider-defined source data and private edits for one build operation. +/// +/// Recorded edits remain private until their extracted update is published. +/// Discard the region if computation or edit recording fails. +pub trait RegionSet: Send { + /// Owned edits, independent of the region's temporary input borrows. + type Update: Send; + + /// Borrow a filled source's centroid and complete operation-local membership. + fn get_view(&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>; + + /// 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<()>; + + /// Release input borrows and transfer the recorded edits for publication. + fn into_update(self) -> Self::Update; +} -/// Canonical rows shared between routing and a build accessor's working-set seed. +/// Canonical rows shared between routing and an operation's build accessor. /// /// IDs must be unique. The contiguous matrix is owned once; iteration and lookup /// borrow its rows without copying coordinates. diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index bf4619fe0f..89e61a37dd 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -13,57 +13,63 @@ 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. Centroid selection through a coarse accessor. -2. Logical list operations: append members and return the resulting size, create - a list, set a list's complete membership, and retire a list. -3. One scoped build accessor that supplies borrowed working sets and publishes - the operation with `finish`. +1. Centroid selection through a `Router`. +2. A scoped `BuildAccessor` that stages ordinary appends, returns their resulting + sizes, fills requested sources, and publishes the operation. +3. An editable `RegionSet` that borrows source data and records list creation, + membership replacement or append, 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. Working-set filling resolves both stored and +membership reads and sizes. Region filling resolves both stored and prepared points, so clustering has no staged-point boundaries or provenance -classification. Working sets also borrow each source centroid, allowing a -gathered region to include retained neighbors as well as split targets. +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, gathers all oversized parents into one -working set, fits two children per parent, and assigns all regional members -against all child centroids. Parent centroids are not candidates. The algorithm -applies the resulting partitioning by creating new lists, setting retained-list -memberships, and retiring split parents, then finishes once. -Ordinary appends and splitting are published together, never as separate commits. +region, fits two children per parent, and assigns the requested parents' members +against all child centroids. All children are fitted before final assignment, and +parent centroids are not candidates. The algorithm computes complete centroid +and member groups before recording child creation and parent retirement in the +region. 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. Initialization creates empty lists with the supplied centroids. - -A partitioning contains a contiguous centroid matrix and a complete member group -per destination. Each destination optionally names an existing source list to -retain; unnamed destinations receive fresh list IDs through `create_list`. -The algorithm explicitly retires split parents after applying destination -memberships. Retained lists keep their identities and centroids while -receiving their complete new memberships. The result covers every regional member -exactly once, including pending appends; lists outside the region are unchanged. - -Split targets are a subset of the gathered sources. Numerical fitting replaces -only their centroids, and assignment considers the fitted children and any -retained source centroids. The default insertion path still gathers only split -parents; choosing neighboring lists is not implemented. - -Working sets borrow vector coordinates. Only fitted centroid coordinates and -assignment/member metadata need new storage. Providers can reuse the canonical -batch and the optional pending-membership helper, or implement their own caches -and private transactions. No delta transport is required by the algorithm; +operation. + +Consuming `RegionSet::into_update` releases its temporary input borrows and moves +its private edits into the accessor's owned `Update`. The index passes that value +to `BuildAccessor::finish`, which publishes it together with accessor-side +appends. An insertion without splitting passes the default, empty regional update. +The algorithm does not inspect the provider's update representation or require +new child IDs during computation. + +The built-in split reads and edits only explicitly requested parents, even when +a region contains additional sources. It covers every parent member exactly once, +including pending appends; other lists are unchanged. Source enumeration and +neighborhood reassignment are deferred. + +Numerical computation borrows vector matrices and member slices, without building +per-point reference tables or copying vector coordinates. A provider with +scattered storage may pack rows during filling to satisfy the matrix interface. +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. 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 includes initialization and batched insert/split. -Search, neighborhood-selection policies, deletion/dissolve, and colocation -remain future work. See [the current contracts](../diskann/src/ivf/traits.rs) for -the implemented API. +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, +deletion/dissolve, and colocation remain future work. See +[the current contracts](../diskann/src/ivf/traits.rs) and +[the region interface](../diskann/src/ivf/workingset.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 @@ -968,10 +974,12 @@ accessor is simpler. ### Imperative List Mutation Methods Immediate, independently published list mutations expose a partial-write schedule -to GraphIVF. The current design instead uses logical list operations inside one -build accessor: they change private operation state and publish only at `finish`. -This keeps get/set-style operations without requiring an algorithm-facing delta -protocol or prescribing the provider's physical write order. +to GraphIVF. The current design instead stages ordinary appends in the build +accessor and regional list edits in a `RegionSet`. Consuming the region produces +an owned update, 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 From cf044f9f47a8a744bc052f858284da5319b35f68 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Wed, 14 Oct 2026 09:17:49 -0400 Subject: [PATCH 18/33] route with nearest --- diskann/src/ivf/clustering.rs | 17 +- diskann/src/ivf/index.rs | 24 +- diskann/src/ivf/mod.rs | 12 +- diskann/src/ivf/pending.rs | 4 +- diskann/src/ivf/{workingset.rs => region.rs} | 36 ++- diskann/src/ivf/tests.rs | 259 ++++++++++++------- diskann/src/ivf/traits.rs | 34 ++- rfcs/01187-incremental-ivf.md | 42 +-- 8 files changed, 265 insertions(+), 163 deletions(-) rename diskann/src/ivf/{workingset.rs => region.rs} (81%) diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index 643c96ea7a..6e0a563b45 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -11,7 +11,7 @@ use rand::{Rng, rngs::StdRng}; use super::{ partition::{Partition, Partitioning}, - workingset::RegionSet, + region::{RegionUpdate, RegionView}, }; use crate::{ANNError, ANNResult}; @@ -20,27 +20,28 @@ use crate::{ANNError, ANNResult}; /// Parents are distinct filled sources. Other filled lists are not read or edited. /// Computation completes before any edits are recorded. A recording error may /// leave partial pending edits, which the caller must discard rather than publish. -pub(super) fn split>( - region: &mut R, +pub(super) fn split, U: RegionUpdate>( + region: &R, + update: &mut U, parents: &[L], dim: usize, iterations: usize, rng: &mut StdRng, ) -> ANNResult<()> { - let result = compute_split(&*region, parents, dim, iterations, rng)?; + let result = compute_split(region, parents, dim, iterations, rng)?; for (centroid, partition) in result.centroids.rows().zip(result.partitions) { match partition.id { - Some(id) => region.set_members(id, &partition.members)?, - None => region.create(centroid, &partition.members)?, + Some(id) => update.set_members(id, &partition.members)?, + None => update.create(centroid, &partition.members)?, } } for &parent in parents { - region.retire(parent)?; + update.retire(parent)?; } Ok(()) } -fn compute_split>( +fn compute_split>( region: &R, parents: &[L], dim: usize, diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 7efa883331..f660ac2e44 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -13,8 +13,7 @@ use thiserror::Error; use super::{ clustering, - traits::{BuildAccessor, BuildStrategy, InsertStrategy, Provider, Router}, - workingset::RegionSet, + traits::{BuildAccessor, InsertStrategy, Provider, Router}, }; use crate::{ANNError, ANNResult, error::IntoANNResult}; @@ -87,11 +86,11 @@ impl IVFIndex

{ let mut accessor = strategy .build_accessor(provider, context) .into_ann_result()?; - if accessor.router().is_empty() { - return Err(ANNError::message("the IVF index is not initialized")); - } + + // Step 1: Compute the batch let batch = strategy.prepare(&mut accessor, points).await?; + // Step 2: Route the points to their nearest centroid. let mut routed = BTreeMap::>::new(); { @@ -121,14 +120,21 @@ impl IVFIndex

{ let splits = parents.len(); let update = if !parents.is_empty() { let dim = accessor.dim(); - let mut view = accessor.fill(&parents).await?; - clustering::split(&mut view, &parents, dim, config.two_means_iterations, rng)?; - Some(view.into_update()) + let (region, mut update) = accessor.fill(&parents).await?; + clustering::split( + ®ion, + &mut update, + &parents, + dim, + config.two_means_iterations, + rng, + )?; + Some(update) } else { None }; - accessor.finish(update).await?; + accessor.finish(update.as_ref()).await?; Ok(InsertStats { inserted: points.len(), diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 281f05d989..e7172a43f4 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -7,12 +7,14 @@ //! //! The index routes points, splits oversized lists, and assigns their members to //! replacement centroids. Providers supply routing and operation-local appends -//! through [`traits::BuildAccessor`], and borrowed inputs and private regional -//! edits through [`workingset::RegionSet`]. +//! through [`traits::BuildAccessor`], borrowed inputs through [`regionset::RegionView`], +//! and regional writes through [`regionset::RegionUpdate`]. //! //! Clustering finishes reading its inputs before recording child creation and -//! parent retirement. Consuming the region releases its input borrows and yields -//! an owned update; the accessor publishes it together with ordinary appends once. +//! parent retirement 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. //! //! This first implementation supports batch insertion into an already initialized @@ -24,8 +26,8 @@ mod clustering; mod index; pub mod partition; pub mod pending; +pub mod region; pub mod traits; -pub mod workingset; pub use index::{Config, ConfigError, IVFIndex, InsertStats}; diff --git a/diskann/src/ivf/pending.rs b/diskann/src/ivf/pending.rs index 6c7991821b..1a0f33a242 100644 --- a/diskann/src/ivf/pending.rs +++ b/diskann/src/ivf/pending.rs @@ -16,7 +16,7 @@ use hashbrown::{HashMap, hash_map::Entry}; use crate::{ANNError, ANNResult, utils::VectorId}; /// A pending change relative to a list's stored membership. -#[derive(Debug)] +#[derive(Debug, Clone)] pub enum Membership { Append(Vec), Replace(Vec), @@ -26,7 +26,7 @@ pub enum Membership { /// Membership state without vector storage or centroid-ID allocation. /// /// Callers supply stored sizes and members, and validate that list IDs exist. -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] pub struct PendingLists { changes: HashMap>, } diff --git a/diskann/src/ivf/workingset.rs b/diskann/src/ivf/region.rs similarity index 81% rename from diskann/src/ivf/workingset.rs rename to diskann/src/ivf/region.rs index ec91a65d2e..0054b9423f 100644 --- a/diskann/src/ivf/workingset.rs +++ b/diskann/src/ivf/region.rs @@ -3,59 +3,53 @@ * Licensed under the MIT license. */ -//! Borrowed clustering views and a reusable canonical insertion batch. +//! Regional input views, write operations, and a reusable canonical insertion batch. use diskann_utils::views::rowmajor::{self, Matrix}; use hashbrown::HashMap; use crate::{ANNError, ANNResult, utils::VectorId}; -/// A point's identity and borrowed canonical vector. -#[derive(Debug, Clone, Copy)] -pub struct PointRef<'a, I> { - pub id: I, - pub vector: &'a [f32], -} - /// 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 ListView<'a, I, L> { +pub struct View<'a, I, L> { pub id: L, pub centroid: &'a [f32], pub members: &'a [I], } -/// Provider-defined source data and private edits for one build operation. -/// -/// Recorded edits remain private until their extracted update is published. -/// Discard the region if computation or edit recording fails. -pub trait RegionSet: Send { - /// Owned edits, independent of the region's temporary input borrows. - type Update: Send; - +/// Views for one build operation to perform splits/re-assignments . +pub trait RegionView: Send { /// Borrow a filled source's centroid and complete operation-local membership. - fn get_view(&self, i: L) -> Option>; + fn get_view(&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 RegionUpdate: 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<()>; +} - /// Release input borrows and transfer the recorded edits for publication. - fn into_update(self) -> Self::Update; +/// A point's identity and borrowed canonical vector. +#[derive(Debug, Clone, Copy)] +pub struct PointRef<'a, I> { + pub id: I, + pub vector: &'a [f32], } /// Canonical rows shared between routing and an operation's build accessor. diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index e0ddc3617e..75d1131574 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -3,7 +3,12 @@ * Licensed under the MIT license. */ -use std::{cell::Cell, collections::BTreeMap, future::Future, sync::Arc}; +use std::{ + cell::Cell, + collections::{BTreeMap, BTreeSet}, + future::Future, + sync::Arc, +}; use diskann_utils::{ future::SendFuture, @@ -15,12 +20,12 @@ use rand::{SeedableRng, rngs::StdRng}; use super::{ Config, IVFIndex, InsertStats, clustering, pending::{Membership, PendingLists}, + region::{Batch, RegionUpdate, RegionView, View}, traits::{BuildAccessor, BuildStrategy, HasIds, InsertStrategy, Provider, Router}, - workingset::{Batch, ListView, RegionSet}, }; -use crate::{ANNError, ANNResult, provider::DefaultContext}; +use crate::{ANNError, ANNResult, ivf::traits::RouterAndNearest, provider::DefaultContext}; -#[derive(Debug)] +#[derive(Debug, Clone)] struct StoredList { centroid: Vec, members: Vec, @@ -95,49 +100,69 @@ impl Router for Coarse<'_> { } } +impl RouterAndNearest for Coarse<'_> { + fn nearest( + &self, + id: Self::ListId, + nprobe: usize, + ) -> impl SendFuture>> { + async move { + let query = &self.stored.get(&id).unwrap().centroid; + self.select(&query, nprobe).await + } + } +} + struct Source { centroid: Vec, members: Vec, vectors: rowmajor::Owned, } -#[derive(Debug, Default)] +#[derive(Debug)] struct Update { + dim: usize, + lists: BTreeSet, memberships: PendingLists, created: Vec, -} - -struct Region<'a> { - sources: &'a BTreeMap, - dim: usize, - update: Update, fail_create: bool, fail_retire: bool, } -impl<'a> Region<'a> { - fn new(sources: &'a BTreeMap, dim: usize) -> Self { +impl Update { + fn new(sources: &BTreeMap, dim: usize) -> Self { Self { - sources, dim, - update: Update::default(), + lists: sources.keys().copied().collect(), + memberships: PendingLists::default(), + created: Vec::new(), fail_create: false, fail_retire: false, } } - fn source(&self, list: u32) -> ANNResult<&Source> { - self.sources - .get(&list) - .ok_or_else(|| ANNError::message(format!("unknown list {list}"))) + fn check_list(&self, list: u32) -> ANNResult<()> { + if self.lists.contains(&list) { + Ok(()) + } else { + Err(ANNError::message(format!("unknown list {list}"))) + } } } -impl RegionSet for Region<'_> { - type Update = Update; +struct Region<'a> { + sources: &'a BTreeMap, +} - fn get_view(&self, list: u32) -> Option> { - self.sources.get(&list).map(|source| ListView { +impl<'a> Region<'a> { + fn new(sources: &'a BTreeMap) -> Self { + Self { sources } + } +} + +impl RegionView for Region<'_> { + fn get_view(&self, list: u32) -> Option> { + self.sources.get(&list).map(|source| View { id: list, centroid: &source.centroid, members: &source.members, @@ -149,24 +174,20 @@ impl RegionSet for Region<'_> { .get(&list) .map(|source| source.vectors.as_view()) } +} +impl RegionUpdate for Update { fn set_members(&mut self, list: u32, members: &[u32]) -> ANNResult<()> { - self.source(list)?; - self.update.memberships.set(list, members.to_vec()) - } - - fn append_members(&mut self, list: u32, members: &[u32]) -> ANNResult { - let size = self.source(list)?.members.len(); - self.update.memberships.append(list, members.to_vec())?; - self.update.memberships.list_size(list, size) + 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.source(list)?; - self.update.memberships.retire(list) + self.check_list(list)?; + self.memberships.retire(list) } fn create(&mut self, centroid: &[f32], members: &[u32]) -> ANNResult<()> { @@ -176,16 +197,12 @@ impl RegionSet for Region<'_> { if centroid.len() != self.dim || !centroid.iter().all(|value| value.is_finite()) { return Err(ANNError::message("invalid centroid")); } - self.update.created.push(StoredList { + self.created.push(StoredList { centroid: centroid.to_vec(), members: members.to_vec(), }); Ok(()) } - - fn into_update(self) -> Self::Update { - self.update - } } impl HasIds for Build<'_> { @@ -228,7 +245,10 @@ impl<'p> BuildAccessor for Build<'p> { } } - fn fill(&mut self, lists: &[u32]) -> impl SendFuture>> { + fn fill( + &mut self, + lists: &[u32], + ) -> impl SendFuture, Self::Update)>> { async move { self.fills.set(self.fills.get() + 1); if self.store.fail_fill { @@ -261,37 +281,44 @@ impl<'p> BuildAccessor for Build<'p> { }, ); } - let mut region = Region::new(&self.sources, self.dim()); - region.fail_create = self.store.fail_create; - region.fail_retire = self.store.fail_retire; - Ok(region) + let mut update = Update::new(&self.sources, self.dim()); + update.fail_create = self.store.fail_create; + update.fail_retire = self.store.fail_retire; + Ok((Region::new(&self.sources), update)) } } - fn finish(&mut self, update: Self::Update) -> impl SendFuture> { + fn finish(self, update: Option<&Self::Update>) -> impl SendFuture> { async move { if self.store.fail_finish { return Err(ANNError::message("injected finish failure")); } - let created = u32::try_from(update.created.len()) - .map_err(|_| ANNError::message("list ID overflow"))?; + let created = update.as_ref().map_or(0, |update| update.created.len()); + let created = + u32::try_from(created).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"))?; - if let Some(batch) = self.batch.take() { + if let Some(batch) = self.batch { for point in batch.points() { self.store.points.insert(point.id, point.vector.to_vec()); } } - for (id, list) in (self.store.next_list..next_list).zip(update.created) { - self.store.lists.insert(id, list); - } - for (id, change) in std::mem::take(&mut self.pending) - .into_changes() - .chain(update.memberships.into_changes()) - { + let regional_memberships = if let Some(update) = update { + for (id, list) in (self.store.next_list..next_list).zip(update.created.clone()) { + self.store.lists.insert(id, list); + } + Some(update.memberships.clone()) + } else { + None + }; + for (id, change) in self.pending.into_changes().chain( + regional_memberships + .into_iter() + .flat_map(PendingLists::into_changes), + ) { match change { Membership::Append(members) => { self.store @@ -555,7 +582,7 @@ async fn splitting_commits_children_and_ordinary_appends_once() { } #[tokio::test] -async fn region_borrows_filled_rows_and_extracts_owned_edits() { +async fn region_view_borrows_inputs_while_updates_are_recorded() { let mut store = seeded(); let update = { let mut build = Strategy @@ -564,24 +591,35 @@ async fn region_borrows_filled_rows_and_extracts_owned_edits() { let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; let batch = Strategy.prepare(&mut build, &points).await.unwrap(); assert_eq!(build.append_to(1, &[10_006, 10_005]).await.unwrap(), 4); - let mut region = require_send(build.fill(&[1])).await.unwrap(); - let view = region.get_view(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), batch.get(10_006).unwrap()); - assert_eq!(vectors.row(3), batch.get(10_005).unwrap()); - clustering::split(&mut region, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap(); - assert!(region.append_members(1, &[10_001]).is_err()); - assert!(region.set_members(1, &[]).is_err()); - let update = region.into_update(); + let update = { + let (region, mut update) = require_send(build.fill(&[1])).await.unwrap(); + let view = region.get_view(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), batch.get(10_006).unwrap()); + assert_eq!(vectors.row(3), batch.get(10_005).unwrap()); + clustering::split( + ®ion, + &mut update, + &[1], + 2, + 5, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); + assert_eq!(view.members, [10_000, 10_001, 10_006, 10_005]); + assert_eq!(vectors.row(2), batch.get(10_006).unwrap()); + assert!(update.set_members(1, &[]).is_err()); + update + }; assert_eq!(build.dim(), 2); update }; @@ -689,15 +727,39 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() Strategy.prepare(&mut build, &points).await.unwrap(); assert_eq!(build.append_to(1, &[10_005]).await.unwrap(), 3); - let mut region = build.fill(&[1, 2]).await.unwrap(); - assert!(clustering::split(&mut region, &[1, 99], 2, 5, &mut StdRng::seed_from_u64(7)).is_err()); - let failed = std::mem::take(&mut region.update); - assert!(failed.created.is_empty()); - assert_eq!(failed.memberships.into_changes().count(), 0); - clustering::split(&mut region, &[1], 2, 5, &mut StdRng::seed_from_u64(7)).unwrap(); - let update = region.into_update(); - assert_eq!(update.created.len(), 2); - require_send(build.finish(update)).await.unwrap(); + let update = { + let (region, mut update) = build.fill(&[1, 2]).await.unwrap(); + assert!( + clustering::split( + ®ion, + &mut update, + &[1, 99], + 2, + 5, + &mut StdRng::seed_from_u64(7) + ) + .is_err() + ); + assert!(update.created.is_empty()); + assert_eq!( + std::mem::take(&mut update.memberships) + .into_changes() + .count(), + 0 + ); + clustering::split( + ®ion, + &mut update, + &[1], + 2, + 5, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); + assert_eq!(update.created.len(), 2); + update + }; + require_send(build.finish(Some(&update))).await.unwrap(); assert_eq!(store.finishes, 1); assert!(!store.lists.contains_key(&1)); @@ -737,9 +799,17 @@ fn assignment_can_cross_parent_boundaries() { ), ] .into(); - let mut region = Region::new(&sources, 1); - clustering::split(&mut region, &[1, 2], 1, 10, &mut StdRng::seed_from_u64(7)).unwrap(); - let update = region.into_update(); + let region = Region::new(&sources); + let mut update = Update::new(&sources, 1); + clustering::split( + ®ion, + &mut update, + &[1, 2], + 1, + 10, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); assert!( update.created[2..] .iter() @@ -759,12 +829,19 @@ fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { }, )] .into(); - let mut region = Region::new(&sources, 1); - clustering::split(&mut region, &[1], 1, 0, &mut StdRng::seed_from_u64(7)).unwrap(); - let zero = region.into_update(); - let mut region = Region::new(&sources, 1); - clustering::split(&mut region, &[1], 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); - let one = region.into_update(); + let region = Region::new(&sources); + let mut zero = Update::new(&sources, 1); + clustering::split( + ®ion, + &mut zero, + &[1], + 1, + 0, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); + let mut one = Update::new(&sources, 1); + clustering::split(®ion, &mut one, &[1], 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); assert_eq!(zero.created.len(), 2); assert_eq!(one.created.len(), 2); for (zero, one) in zero.created.iter().zip(&one.created) { diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 943b58e80e..1ec1365fe9 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -9,11 +9,10 @@ use std::{fmt::Debug, sync::Arc}; use diskann_utils::future::SendFuture; -use super::workingset::Batch; +use super::region::{Batch, RegionUpdate, RegionView}; use crate::{ ANNResult, error::{StandardError, ToRanked}, - ivf::workingset::RegionSet, provider::ExecutionContext, utils::VectorId, }; @@ -65,18 +64,26 @@ pub trait Router: Send { ) -> impl SendFuture>>; } -pub trait BuildAccessor: Send + Sized + HasIds { - /// Owned regional edits. The default records no regional changes. - type Update: Send; +pub trait RouterAndNearest: Router { + /// Find `nprobe` centroids nearest to list `id`. + fn nearest( + &self, + id: Self::ListId, + nprobe: usize, + ) -> impl SendFuture>>; +} - type Router<'a>: Router +pub trait BuildAccessor: Send + Sized + HasIds { + type Router<'a>: RouterAndNearest where Self: 'a; - type Region<'a>: RegionSet + type Region<'a>: RegionView where Self: 'a; + type Update: RegionUpdate; + /// Positive dimension of all canonical points and centroids. fn dim(&self) -> usize; @@ -90,11 +97,16 @@ pub trait BuildAccessor: Send + Sized + HasIds { ) -> impl SendFuture>; /// Make every requested list's complete operation-local membership, vectors - /// and centroids available for computation. - fn fill(&mut self, lists: &[Self::ListId]) -> impl SendFuture>>; + /// and centroids available for computation, together with its regional writer. + fn fill( + &mut self, + lists: &[Self::ListId], + ) -> impl SendFuture, Self::Update)>>; - /// Publish accessor-side appends and owned regional edits together. - fn finish(&mut self, update: Option) -> impl SendFuture>; + /// Finalize the operation, publishing accessor-side appends and regional writes together. + /// + /// `None` means there are no regional re-writes. + fn finish(self, update: Option<&Self::Update>) -> impl SendFuture>; } pub trait BuildStrategy<'a, P: Provider>: Send + Sync { diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 89e61a37dd..683a100a8c 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -4,7 +4,7 @@ |---|---| | **Authors** | Aditya Krishnan | | **Created** | 2026-08-19 | -| **Updated** | 2026-10-12 | +| **Updated** | 2026-10-13 | | **Status** | Draft | ## Summary @@ -16,8 +16,9 @@ overflow selection, two-means fitting, and regional assignment. Providers supply 1. Centroid selection through a `Router`. 2. A scoped `BuildAccessor` that stages ordinary appends, returns their resulting sizes, fills requested sources, and publishes the operation. -3. An editable `RegionSet` that borrows source data and records list creation, - membership replacement or append, and retirement. +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, @@ -31,18 +32,26 @@ Insertion provisionally appends points, gathers all oversized parents into one region, fits two children per parent, and assigns the requested parents' members against all child centroids. All children are fitted before final assignment, and parent centroids are not candidates. The algorithm computes complete centroid -and member groups before recording child creation and parent retirement in the -region. Ordinary appends and splitting are published together, never as separate +and member groups before writing child creation 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. -Consuming `RegionSet::into_update` releases its temporary input borrows and moves -its private edits into the accessor's owned `Update`. The index passes that value -to `BuildAccessor::finish`, which publishes it together with accessor-side -appends. An insertion without splitting passes the default, empty regional update. -The algorithm does not inspect the provider's update representation or require -new child IDs during computation. +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 requested parents, even when a region contains additional sources. It covers every parent member exactly once, @@ -69,7 +78,7 @@ The current implementation supports batched insert/split into a provider that already has centroids. Initialization, search, neighborhood-selection policies, deletion/dissolve, and colocation remain future work. See [the current contracts](../diskann/src/ivf/traits.rs) and -[the region interface](../diskann/src/ivf/workingset.rs) for the implemented API. +[the region interfaces](../diskann/src/ivf/regionset.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 @@ -975,10 +984,11 @@ accessor is simpler. 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 list edits in a `RegionSet`. Consuming the region produces -an owned update, 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 +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 From 2bd795d04ca2b315498ce97e5eaedf4647765ef0 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Thu, 15 Oct 2026 14:42:36 -0400 Subject: [PATCH 19/33] cleanup --- diskann/src/ivf/clustering.rs | 49 +++-- diskann/src/ivf/index.rs | 81 ++++++-- diskann/src/ivf/mod.rs | 24 ++- diskann/src/ivf/tests.rs | 367 ++++++++++++++++++++++++++++++---- diskann/src/ivf/traits.rs | 53 +++-- rfcs/01187-incremental-ivf.md | 51 +++-- 6 files changed, 499 insertions(+), 126 deletions(-) diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index 6e0a563b45..b102251d62 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -5,6 +5,8 @@ //! Two-means fitting and region-wide assignment over borrowed point rows. +use std::collections::BTreeSet; + use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use rand::{Rng, rngs::StdRng}; @@ -15,20 +17,22 @@ use super::{ }; use crate::{ANNError, ANNResult}; -/// Fit two children per parent, assign their members across all children, and record edits. +/// Fit two children per parent and reassign sources across children and retained neighbors. /// -/// Parents are distinct filled sources. Other filled lists are not read or edited. +/// Parents and neighbors are distinct, disjoint filled sources. Neighbor centroids +/// are retained unchanged. Other filled lists are not read or edited. /// Computation completes before any edits are recorded. A recording error may /// leave partial pending edits, which the caller must discard rather than publish. pub(super) fn split, U: RegionUpdate>( region: &R, update: &mut U, - parents: &[L], + parents: &BTreeSet, + neighbors: &BTreeSet, dim: usize, iterations: usize, rng: &mut StdRng, ) -> ANNResult<()> { - let result = compute_split(region, parents, dim, iterations, rng)?; + let result = compute_split(region, parents, neighbors, dim, iterations, rng)?; for (centroid, partition) in result.centroids.rows().zip(result.partitions) { match partition.id { Some(id) => update.set_members(id, &partition.members)?, @@ -43,7 +47,8 @@ pub(super) fn split, U: RegionUpdate fn compute_split>( region: &R, - parents: &[L], + parents: &BTreeSet, + neighbors: &BTreeSet, dim: usize, iterations: usize, rng: &mut StdRng, @@ -53,22 +58,29 @@ fn compute_split>( "IVF splitting needs parents and a positive dimension", )); } - let mut inputs = Vec::with_capacity(parents.len()); - for &parent in parents { + + let child_count = 2 * parents.len(); + let mut centroids = rowmajor::Owned::try_from_element(child_count + neighbors.len(), dim, 0.0)?; + let mut inputs = Vec::with_capacity(parents.len() + neighbors.len()); + for (index, &source) in parents.iter().chain(neighbors).enumerate() { let view = region - .get_view(parent) - .ok_or_else(|| ANNError::message("IVF split parent is missing from the region"))?; - let vectors = region.get_vectors(parent).ok_or_else(|| { - ANNError::message("IVF split parent vectors are missing from the region") + .get_view(source) + .ok_or_else(|| ANNError::message("IVF split source is missing from the region"))?; + let vectors = region.get_vectors(source).ok_or_else(|| { + ANNError::message("IVF split source vectors are missing from the region") })?; + if index >= parents.len() { + centroids + .row_mut(child_count + index - parents.len()) + .copy_from_slice(view.centroid); + } inputs.push((view.members, vectors)); } - let mut centroids = rowmajor::Owned::try_from_element(2 * parents.len(), dim, 0.0)?; let mut sums = vec![0.0; 2 * dim]; - for ((_, vectors), centers) in inputs + for ((_, vectors), centers) in inputs[..parents.len()] .iter() - .zip(centroids.as_mut_slice().chunks_exact_mut(2 * dim)) + .zip(centroids.as_mut_slice()[..child_count * dim].chunks_exact_mut(2 * dim)) { fit_two_means(vectors.as_view(), centers, dim, iterations, rng, &mut sums)?; } @@ -83,10 +95,11 @@ fn compute_split>( sizes[to] += 1; } } - let mut partitions: Vec> = sizes - .into_iter() - .map(|size| Partition { - id: None, + let mut partitions: Vec> = std::iter::repeat_n(None, child_count) + .chain(neighbors.iter().copied().map(Some)) + .zip(sizes) + .map(|(id, size)| Partition { + id, members: Vec::with_capacity(size), }) .collect(); diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index f660ac2e44..fe03f76b62 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -5,7 +5,10 @@ //! IVF policy and orchestration, independent of physical list storage. -use std::collections::BTreeMap; +use std::{ + collections::{BTreeMap, BTreeSet}, + num::NonZeroUsize, +}; use diskann_utils::future::SendFuture; use rand::{SeedableRng, rngs::StdRng}; @@ -13,17 +16,28 @@ use thiserror::Error; use super::{ clustering, - traits::{BuildAccessor, InsertStrategy, Provider, Router}, + traits::{BuildAccessor, InsertStrategy, Provider, Router, RouterAndNearest}, }; use crate::{ANNError, ANNResult, error::IntoANNResult}; +#[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: usize, - /// Lloyd steps per parent. Zero still performs one step. - pub two_means_iterations: usize, + pub split_threshold: NonZeroUsize, + /// Nearest-list query count per split parent, before filtering split parents. + /// + /// [`ReassignCount::None`] disables neighborhood gathering. + /// All gathered points are reassigned across all children and retained neighbors. + pub reassign_count: ReassignCount, + /// Positive number of Lloyd steps per parent. + pub two_means_iterations: NonZeroUsize, pub seed: u64, } @@ -51,8 +65,8 @@ impl IVFIndex

{ /// /// Returns [`ConfigError`] when the split threshold is less than two. pub fn new(provider: P, config: Config) -> Result { - if config.split_threshold < 2 { - return Err(ConfigError(config.split_threshold)); + if config.split_threshold.get() < 2 { + return Err(ConfigError(config.split_threshold.get())); } Ok(Self { provider, @@ -97,38 +111,63 @@ impl IVFIndex

{ let router = accessor.router(); for point in batch.points() { - let selected = router.select(point.vector, 1).await?; - let list = selected - .first() - .ok_or_else(|| ANNError::message("IVF routing returned no live centroid"))? - .0; + let selected = router.route(point.vector, 1).await?; + + let (list, _) = selected.first().ok_or_else(|| { + ANNError::message("IVF routing returned no live centroid") + })?; - routed.entry(list).or_default().push(point.id); + routed.entry(*list).or_default().push(point.id); } } - let mut parents = Vec::with_capacity(routed.len()); + let mut parents = BTreeSet::new(); + + // Step 3: append the points to their respective list ids and determine which need to be split. + for (list, members) in routed.iter() { + let size = accessor.append_to(*list, members.iter().copied()).await?; - // append the points to their respective list ids and determine which need to be split. - for (list, members) in routed { - let size = accessor.append_to(list, &members).await?; - if size > config.split_threshold { - parents.push(list); + if size.get() > config.split_threshold.get() { + parents.insert(*list); } } let splits = parents.len(); let update = if !parents.is_empty() { + let mut neighbors = BTreeSet::new(); + + if let ReassignCount::Count(n) = config.reassign_count { + let router = accessor.router(); + + for &pid in parents.iter() { + let nids = router + .nearest(pid, n.get()) + .await? + .into_iter() + .map(|(id, _)| id) + .filter(|id| !parents.contains(id)); + + for id in nids { + neighbors.insert(id); + } + } + } + let dim = accessor.dim(); - let (region, mut update) = accessor.fill(&parents).await?; + let iter = parents.iter().chain(neighbors.iter()).copied(); + + let (region, mut update) = accessor.fill(iter).await?; + clustering::split( ®ion, &mut update, &parents, + &neighbors, dim, - config.two_means_iterations, + config.two_means_iterations.get(), rng, )?; + Some(update) } else { None diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index e7172a43f4..1df037eac6 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -7,20 +7,30 @@ //! //! 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 [`regionset::RegionView`], -//! and regional writes through [`regionset::RegionUpdate`]. +//! through [`traits::BuildAccessor`], borrowed inputs through [`region::RegionView`], +//! and regional writes through [`region::RegionUpdate`]. +//! +//! Neighbor selection uses [`traits::RouterAndNearest`] 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. //! //! Clustering finishes reading its inputs before recording child creation and -//! parent retirement through the separate update object. The region's input +//! 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. //! //! This first implementation supports batch insertion into an already initialized -//! provider, with two-means splitting of the explicitly requested parents. -//! Initialization, neighborhood reassignment, search, deletion, and colocation are -//! not implemented. +//! 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, search, deletion, and colocation are not +//! implemented. mod clustering; mod index; @@ -29,7 +39,7 @@ pub mod pending; pub mod region; pub mod traits; -pub use index::{Config, ConfigError, IVFIndex, InsertStats}; +pub use index::{Config, ConfigError, IVFIndex, InsertStats, ReassignCount}; #[cfg(test)] mod tests; diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index 75d1131574..72c4f7086a 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -7,6 +7,7 @@ use std::{ cell::Cell, collections::{BTreeMap, BTreeSet}, future::Future, + num::NonZeroUsize, sync::Arc, }; @@ -18,12 +19,14 @@ use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use rand::{SeedableRng, rngs::StdRng}; use super::{ - Config, IVFIndex, InsertStats, clustering, + Config, IVFIndex, InsertStats, ReassignCount, clustering, pending::{Membership, PendingLists}, region::{Batch, RegionUpdate, RegionView, View}, - traits::{BuildAccessor, BuildStrategy, HasIds, InsertStrategy, Provider, Router}, + traits::{ + BuildAccessor, BuildStrategy, HasIds, InsertStrategy, Provider, Router, RouterAndNearest, + }, }; -use crate::{ANNError, ANNResult, ivf::traits::RouterAndNearest, provider::DefaultContext}; +use crate::{ANNError, ANNResult, provider::DefaultContext}; #[derive(Debug, Clone)] struct StoredList { @@ -39,8 +42,11 @@ struct Store { next_list: u32, finishes: usize, fills: usize, + filled_lists: Vec, + fail_nearest: bool, fail_fill: bool, fail_create: bool, + fail_set_members: bool, fail_retire: bool, fail_finish: bool, } @@ -73,20 +79,22 @@ struct Build<'p> { batch: Option>>, sources: BTreeMap, fills: Cell, + filled_lists: Vec, } struct Coarse<'a> { stored: &'a BTreeMap, + fail_nearest: bool, } impl Router for Coarse<'_> { - type ListId = u32; + type Id = u32; fn len(&self) -> usize { self.stored.len() } - fn select(&self, query: &[f32], nprobe: usize) -> impl SendFuture>> { + fn route(&self, query: &[f32], nprobe: usize) -> impl SendFuture>> { async move { let mut selected: Vec<(u32, f32)> = self .stored @@ -101,14 +109,17 @@ impl Router for Coarse<'_> { } impl RouterAndNearest for Coarse<'_> { - fn nearest( - &self, - id: Self::ListId, - nprobe: usize, - ) -> impl SendFuture>> { + fn nearest(&self, id: Self::Id, nprobe: usize) -> impl SendFuture>> { async move { - let query = &self.stored.get(&id).unwrap().centroid; - self.select(&query, nprobe).await + 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; + self.route(query, nprobe).await } } } @@ -126,6 +137,7 @@ struct Update { memberships: PendingLists, created: Vec, fail_create: bool, + fail_set_members: bool, fail_retire: bool, } @@ -137,6 +149,7 @@ impl Update { memberships: PendingLists::default(), created: Vec::new(), fail_create: false, + fail_set_members: false, fail_retire: false, } } @@ -178,6 +191,9 @@ impl RegionView for Region<'_> { impl RegionUpdate 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()) } @@ -206,7 +222,7 @@ impl RegionUpdate for Update { } impl HasIds for Build<'_> { - type InternalId = u32; + type Id = u32; type ListId = u32; } @@ -228,10 +244,18 @@ impl<'p> BuildAccessor for Build<'p> { fn router(&mut self) -> Self::Router<'_> { Coarse { stored: &self.store.lists, + fail_nearest: self.store.fail_nearest, } } - fn append_to(&mut self, list: u32, members: &[u32]) -> impl SendFuture> { + fn append_to( + &mut self, + list: u32, + members: Itr, + ) -> impl SendFuture> + where + Itr: ExactSizeIterator + Clone + Send + Sync, + { async move { let size = self .store @@ -240,22 +264,27 @@ impl<'p> BuildAccessor for Build<'p> { .ok_or_else(|| ANNError::message(format!("unknown list {list}")))? .members .len(); - self.pending.append(list, members.to_vec())?; - self.pending.list_size(list, size) + self.pending.append(list, members.collect())?; + NonZeroUsize::new(self.pending.list_size(list, size)?) + .ok_or_else(|| ANNError::message("appending must leave a nonempty list")) } } - fn fill( + fn fill( &mut self, - lists: &[u32], - ) -> impl SendFuture, Self::Update)>> { + 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 { + for list in lists { let stored = self .store .lists @@ -283,6 +312,7 @@ impl<'p> BuildAccessor for Build<'p> { } let mut update = Update::new(&self.sources, self.dim()); 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)) } @@ -339,6 +369,7 @@ impl<'p> BuildAccessor for Build<'p> { self.store.next_list = next_list; self.store.finishes += 1; self.store.fills += self.fills.get(); + self.store.filled_lists = self.filled_lists; Ok(()) } } @@ -361,6 +392,7 @@ impl<'a> BuildStrategy<'a, Store> for Strategy { batch: None, sources: BTreeMap::new(), fills: Cell::new(0), + filled_lists: Vec::new(), }) } } @@ -400,11 +432,17 @@ fn matrix(rows: &[[f32; D]]) -> rowmajor::Owned { } fn index(store: Store, threshold: usize) -> IVFIndex { + index_with_neighbors(store, threshold, 0) +} + +fn index_with_neighbors(store: Store, threshold: usize, neighbors: usize) -> IVFIndex { IVFIndex::new( store, Config { - split_threshold: threshold, - two_means_iterations: 5, + split_threshold: NonZeroUsize::new(threshold).unwrap(), + reassign_count: NonZeroUsize::new(neighbors) + .map_or(ReassignCount::None, ReassignCount::Count), + two_means_iterations: NonZeroUsize::new(5).unwrap(), seed: 7, }, ) @@ -478,6 +516,7 @@ async fn appends_publish_actual_memberships_without_filling() { ] .into(), next_list: 3, + fail_nearest: true, ..Store::default() }, 2, @@ -531,7 +570,9 @@ async fn appends_publish_actual_memberships_without_filling() { #[tokio::test] async fn splitting_commits_children_and_ordinary_appends_once() { - let mut index = index(seeded(), 2); + let mut store = seeded(); + store.fail_nearest = true; + let mut index = index(store, 2); let points = [ (5, &[1.0, 1.0][..]), (6, &[11.0, 1.0][..]), @@ -581,6 +622,198 @@ async fn splitting_commits_children_and_ordinary_appends_once() { } } +#[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 = index_with_neighbors(store, 3, 2); + let points = [(7, &[1.0, 0.0][..]), (8, &[5.2, 0.0][..])]; + let stats = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) + .await + .unwrap(); + let store = index.provider(); + + 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]); + 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, two_candidates, three_candidates) in [ + (vec![], vec![], vec![]), + (vec![100.0, 200.0, 300.0], vec![], vec![3]), + (vec![0.0, 12.0, 200.0], vec![3, 4], 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 = index_with_neighbors(store, 2, count); + let points = [(5, &[1.0, 1.0][..]), (6, &[11.0, 1.0][..])]; + let stats = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) + .await + .unwrap(); + let store = index.provider(); + let neighbors = match (count, centers.len()) { + (0 | 1, _) | (_, 0) => vec![], + (2, _) => two_candidates.clone(), + (3, _) => three_candidates.clone(), + (4, _) => 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::>()); + } + } +} + +#[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 update = Update::new(&sources, 1); + clustering::split( + ®ion, + &mut update, + &BTreeSet::from([1]), + &BTreeSet::from([2, 3, 4]), + 1, + 5, + &mut StdRng::seed_from_u64(7), + ) + .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, 1]); + let mut changes: BTreeMap<_, _> = 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![])] { + 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 region_view_borrows_inputs_while_updates_are_recorded() { let mut store = seeded(); @@ -590,9 +823,16 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { .unwrap(); let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; let batch = Strategy.prepare(&mut build, &points).await.unwrap(); - assert_eq!(build.append_to(1, &[10_006, 10_005]).await.unwrap(), 4); + assert_eq!( + build + .append_to(1, [10_006, 10_005].into_iter()) + .await + .unwrap() + .get(), + 4 + ); let update = { - let (region, mut update) = require_send(build.fill(&[1])).await.unwrap(); + let (region, mut update) = require_send(build.fill([1].into_iter())).await.unwrap(); let view = region.get_view(1).unwrap(); let vectors = region.get_vectors(1).unwrap(); let source = ®ion.sources[&1]; @@ -609,7 +849,8 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { clustering::split( ®ion, &mut update, - &[1], + &BTreeSet::from([1]), + &BTreeSet::new(), 2, 5, &mut StdRng::seed_from_u64(7), @@ -644,26 +885,34 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { #[tokio::test] async fn failed_planning_or_publication_is_not_reported_as_success() { - for (fail_fill, fail_create, fail_retire, fail_finish) in [ - (true, false, false, false), - (false, true, false, false), - (false, false, true, false), - (false, false, false, true), + 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 = index(store, 2); + let mut index = index_with_neighbors(store, 2, 2); let points = [(5, &[1.0, 1.0][..])]; let error = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) .await .unwrap_err(); - let expected = if fail_fill { + 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 { @@ -672,6 +921,7 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { assert!(error.to_string().contains(expected), "{error}"); assert_eq!(index.provider().finishes, 0); assert_eq!(index.provider().lists[&1].members, [10_000, 10_001]); + assert_eq!(index.provider().lists[&2].members, [10_002, 10_003]); assert_eq!(index.provider().lists.len(), 3); assert!( index @@ -688,8 +938,9 @@ async fn invalid_config_and_uninitialized_insertion_fail() { IVFIndex::new( Store::default(), Config { - split_threshold: 1, - two_means_iterations: 5, + split_threshold: NonZeroUsize::new(1).unwrap(), + reassign_count: ReassignCount::None, + two_means_iterations: NonZeroUsize::new(5).unwrap(), seed: 7, }, ) @@ -725,15 +976,35 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() .unwrap(); let points = [(5, &[5.0, 0.0][..])]; Strategy.prepare(&mut build, &points).await.unwrap(); - assert_eq!(build.append_to(1, &[10_005]).await.unwrap(), 3); + assert_eq!( + build + .append_to(1, [10_005].into_iter()) + .await + .unwrap() + .get(), + 3 + ); let update = { - let (region, mut update) = build.fill(&[1, 2]).await.unwrap(); + let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); assert!( clustering::split( ®ion, &mut update, - &[1, 99], + &BTreeSet::from([1, 99]), + &BTreeSet::new(), + 2, + 5, + &mut StdRng::seed_from_u64(7) + ) + .is_err() + ); + assert!( + clustering::split( + ®ion, + &mut update, + &BTreeSet::from([1]), + &BTreeSet::from([99]), 2, 5, &mut StdRng::seed_from_u64(7) @@ -750,7 +1021,8 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() clustering::split( ®ion, &mut update, - &[1], + &BTreeSet::from([1]), + &BTreeSet::new(), 2, 5, &mut StdRng::seed_from_u64(7), @@ -804,7 +1076,8 @@ fn assignment_can_cross_parent_boundaries() { clustering::split( ®ion, &mut update, - &[1, 2], + &BTreeSet::from([1, 2]), + &BTreeSet::new(), 1, 10, &mut StdRng::seed_from_u64(7), @@ -834,14 +1107,24 @@ fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { clustering::split( ®ion, &mut zero, - &[1], + &BTreeSet::from([1]), + &BTreeSet::new(), 1, 0, &mut StdRng::seed_from_u64(7), ) .unwrap(); let mut one = Update::new(&sources, 1); - clustering::split(®ion, &mut one, &[1], 1, 1, &mut StdRng::seed_from_u64(7)).unwrap(); + clustering::split( + ®ion, + &mut one, + &BTreeSet::from([1]), + &BTreeSet::new(), + 1, + 1, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); assert_eq!(zero.created.len(), 2); assert_eq!(one.created.len(), 2); for (zero, one) in zero.created.iter().zip(&one.created) { diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 1ec1365fe9..da4ce3f6ce 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -5,7 +5,7 @@ //! Storage boundaries for IVF construction. -use std::{fmt::Debug, sync::Arc}; +use std::{fmt::Debug, num::NonZeroUsize, sync::Arc}; use diskann_utils::future::SendFuture; @@ -42,13 +42,13 @@ pub trait Provider: Sized + Send + Sync + 'static { } pub trait HasIds { - type InternalId: VectorId; + type Id: VectorId; type ListId: VectorId; } /// Centroid selection within one build operation. pub trait Router: Send { - type ListId: VectorId; + type Id: VectorId; fn len(&self) -> usize; @@ -56,33 +56,40 @@ pub trait Router: Send { self.len() == 0 } - /// Select `min(nprobe, len())` live lists, nearest first. - fn select( + /// Select `min(nprobe, len())` distinct live lists, nearest first. + fn route( &self, query: &[f32], nprobe: usize, - ) -> impl SendFuture>>; + ) -> impl SendFuture>>; } pub trait RouterAndNearest: Router { - /// Find `nprobe` centroids nearest to list `id`. + /// Select `min(nprobe, len())` distinct live lists nearest to `id`'s centroid. + /// + /// Results are nearest first and may include `id` itself. Callers exclude + /// lists being replaced when selecting retained neighbors. + /// + /// # Errors + /// + /// Fails if `id` is not live or the provider cannot complete the lookup. fn nearest( &self, - id: Self::ListId, + id: Self::Id, nprobe: usize, - ) -> impl SendFuture>>; + ) -> impl SendFuture>>; } pub trait BuildAccessor: Send + Sized + HasIds { - type Router<'a>: RouterAndNearest + type Router<'a>: RouterAndNearest where Self: 'a; - type Region<'a>: RegionView + type Region<'a>: RegionView where Self: 'a; - type Update: RegionUpdate; + type Update: RegionUpdate; /// Positive dimension of all canonical points and centroids. fn dim(&self) -> usize; @@ -90,18 +97,22 @@ pub trait BuildAccessor: Send + Sized + HasIds { fn router(&mut self) -> Self::Router<'_>; /// Append members and return the resulting operation-local list size. - fn append_to( + fn append_to( &mut self, list: Self::ListId, - members: &[Self::InternalId], - ) -> impl SendFuture>; + members: Itr, + ) -> impl SendFuture> + where + Itr: ExactSizeIterator + Clone + Send + Sync; - /// Make every requested list's complete operation-local membership, vectors - /// and centroids available for computation, together with its regional writer. - fn fill( + /// Make every requested list's members, data and centroids available + /// for split and assignment. Register updates to the region. + fn fill( &mut self, - lists: &[Self::ListId], - ) -> impl SendFuture, Self::Update)>>; + lists: Itr, + ) -> impl SendFuture, Self::Update)>> + where + Itr: Iterator + Clone + Send + Sync; /// Finalize the operation, publishing accessor-side appends and regional writes together. /// @@ -110,7 +121,7 @@ pub trait BuildAccessor: Send + Sized + HasIds { } pub trait BuildStrategy<'a, P: Provider>: Send + Sync { - type Accessor: BuildAccessor; + type Accessor: BuildAccessor; type Error: StandardError; fn build_accessor( diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 683a100a8c..06269c4e05 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -4,7 +4,7 @@ |---|---| | **Authors** | Aditya Krishnan | | **Created** | 2026-08-19 | -| **Updated** | 2026-10-13 | +| **Updated** | 2026-10-15 | | **Status** | Draft | ## Summary @@ -13,7 +13,8 @@ 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. Centroid selection through a `Router`. +1. Point routing through a `Router`, and nearest-list selection through + `RouterAndNearest`. 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 @@ -28,21 +29,36 @@ 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, gathers all oversized parents into one -region, fits two children per parent, and assigns the requested parents' members -against all child centroids. All children are fitted before final assignment, and -parent centroids are not candidates. The algorithm computes complete centroid -and member groups before writing child creation and parent retirement through -the update object. Ordinary appends and splitting are published together, never as separate -commits. +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. +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 +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. @@ -53,10 +69,10 @@ 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 requested parents, even when -a region contains additional sources. It covers every parent member exactly once, -including pending appends; other lists are unchanged. Source enumeration and -neighborhood reassignment are deferred. +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. Numerical computation borrows vector matrices and member slices, without building per-point reference tables or copying vector coordinates. A provider with @@ -76,9 +92,10 @@ aliases themselves. The current implementation supports batched insert/split into a provider that already has centroids. Initialization, search, neighborhood-selection policies, -deletion/dissolve, and colocation remain future work. See +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/regionset.rs) for the implemented API. +[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 From 855fb98704170fa29159df3e28fe7f40147d7f44 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 19 Oct 2026 14:59:34 -0400 Subject: [PATCH 20/33] clusterer and assigner --- diskann/src/ivf/clustering.rs | 408 +++++++++++++++++-------- diskann/src/ivf/index.rs | 27 +- diskann/src/ivf/mod.rs | 5 + diskann/src/ivf/partition.rs | 19 +- diskann/src/ivf/tests.rs | 555 +++++++++++++++++++++++++++++++--- diskann/src/ivf/traits.rs | 2 +- rfcs/01187-incremental-ivf.md | 14 + 7 files changed, 838 insertions(+), 192 deletions(-) diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index b102251d62..11d9323261 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -3,171 +3,335 @@ * Licensed under the MIT license. */ -//! Two-means fitting and region-wide assignment over borrowed point rows. +//! Replaceable child fitting and regional assignment over borrowed point rows. +//! +//! Fitting writes new centroids into reusable scratch. Assignment borrows the +//! destinations, groups source members, and records regional writes directly. -use std::collections::BTreeSet; +use std::{collections::BTreeSet, num::NonZeroUsize}; use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use rand::{Rng, rngs::StdRng}; -use super::{ - partition::{Partition, Partitioning}, - region::{RegionUpdate, RegionView}, -}; +use super::region::{RegionUpdate, RegionView}; use crate::{ANNError, ANNResult}; +/// Fit a parent's replacement centroids without assigning regional memberships. +pub(super) trait Clusterer { + /// Fill `dst` with finite centroids of the canonical dimension. + /// + /// Implementations own initialization, convergence, empty-cluster handling, + /// and reusable fitting scratch. `dst` has `num_centers()` rows and `dim()` + /// columns. The number of centers must be at least two. + fn fit( + &mut self, + points: rowmajor::Ref<'_, T>, + dst: rowmajor::Mut<'_, T>, + rng: &mut StdRng, + ) -> ANNResult<()>; + + fn num_centers(&self) -> NonZeroUsize; + + fn dim(&self) -> NonZeroUsize; +} + +/// Assign the whole region, allowing policies that coordinate multiple sources. +pub(super) trait Assigner { + /// Assign each source member to a destination and record complete memberships. + /// + /// Create new destinations, replace retained memberships, and retire sources + /// not retained as destinations. Validate and compute assignments before + /// recording writes. A write error may leave partial private updates. + fn assign( + &mut self, + region: &R, + update: &mut U, + sources: &[L], + destinations: &[Destination<'_, L>], + ) -> ANNResult<()> + where + R: RegionView, + U: RegionUpdate; +} + +pub(super) enum Type { + /// Existing centroid. + Exists(L), + /// New centroid as a result of the split. + New, +} + +/// A borrowed centroid; only retained destinations have an existing list ID. +pub(super) struct Destination<'a, L> { + pub list: Type, + pub centroid: &'a [f32], +} + +pub(super) struct SplitAndAssign { + pub clusterer: C, + pub assigner: A, + scratch: Option>, // scratch for the centroids. +} + +impl SplitAndAssign { + pub(super) fn new(clusterer: C, assigner: A) -> Self { + Self { + clusterer, + assigner, + scratch: None, + } + } + + /// Initialize the scratch for a particular number of output cluster centroids. + fn initialize(&mut self, num_jobs: NonZeroUsize) -> ANNResult<()> { + let c = self.clusterer.num_centers(); + + let rows = num_jobs + .checked_mul(c) + .ok_or_else(|| ANNError::message("IVF centroid count overflow"))? + .get(); + + let dim = self.clusterer.dim().get(); + + self.scratch = Some(rowmajor::Owned::try_from_element(rows, dim, f32::NAN)?); + + Ok(()) + } +} + /// Fit two children per parent and reassign sources across children and retained neighbors. /// /// Parents and neighbors are distinct, disjoint filled sources. Neighbor centroids /// are retained unchanged. Other filled lists are not read or edited. /// Computation completes before any edits are recorded. A recording error may /// leave partial pending edits, which the caller must discard rather than publish. -pub(super) fn split, U: RegionUpdate>( +/// An empty set of parents is a no-op. +pub(super) fn split, U: RegionUpdate>( region: &R, update: &mut U, parents: &BTreeSet, neighbors: &BTreeSet, - dim: usize, - iterations: usize, + iterations: NonZeroUsize, + dim: NonZeroUsize, rng: &mut StdRng, ) -> ANNResult<()> { - let result = compute_split(region, parents, neighbors, dim, iterations, rng)?; - for (centroid, partition) in result.centroids.rows().zip(result.partitions) { - match partition.id { - Some(id) => update.set_members(id, &partition.members)?, - None => update.create(centroid, &partition.members)?, + SplitAndAssign::new(TwoMeans::new(iterations, dim), RegionalNearest) + .run(region, update, parents, neighbors, rng) +} + +impl SplitAndAssign { + /// Compute the complete result before registering any provider writes. + pub(super) fn run( + &mut self, + region: &R, + update: &mut U, + parents: &BTreeSet, + neighbors: &BTreeSet, + rng: &mut StdRng, + ) -> ANNResult<()> + where + A: Assigner, + R: RegionView, + U: RegionUpdate, + { + self.initialize(NonZeroUsize::new(parents.len()).ok_or(ANNError::message("err"))?)?; + + let n_centroids = self.clusterer.num_centers().get(); + let mut s = 0; + + let mut clusters = self + .scratch + .as_mut() + .ok_or_else(|| ANNError::message("IVF centroid scratch was not initialized"))? + .as_view_mut(); + + for &p in parents { + let vecs = region.get_vectors(p).ok_or_else(|| { + ANNError::message("IVF split source vectors are missing from the region") + })?; + + let end = s + n_centroids; + + self.clusterer.fit( + vecs, + clusters.subview_mut(s..end).ok_or_else(|| { + ANNError::message("IVF centroid scratch range is out of bounds") + })?, + rng, + )?; + + s = end; } + + let mut destinations: Vec<_> = clusters + .rows() + .map(|centroid| Destination { + list: Type::New, + centroid, + }) + .collect(); + + for &n in neighbors { + let view = region + .get_view(n) + .ok_or_else(|| ANNError::message("IVF neighbor is missing from the region"))?; + + destinations.push(Destination { + list: Type::Exists(n), + centroid: view.centroid, + }); + } + + let sources = parents.iter().chain(neighbors).copied().collect::>(); + + self.assigner + .assign(region, update, &sources, &destinations) } - for &parent in parents { - update.retire(parent)?; - } - Ok(()) } -fn compute_split>( - region: &R, - parents: &BTreeSet, - neighbors: &BTreeSet, - dim: usize, - iterations: usize, - rng: &mut StdRng, -) -> ANNResult> { - if parents.is_empty() || dim == 0 { - return Err(ANNError::message( - "IVF splitting needs parents and a positive dimension", - )); - } +pub(super) struct TwoMeans { + iterations: NonZeroUsize, + sums: Vec, + dim: NonZeroUsize, +} - let child_count = 2 * parents.len(); - let mut centroids = rowmajor::Owned::try_from_element(child_count + neighbors.len(), dim, 0.0)?; - let mut inputs = Vec::with_capacity(parents.len() + neighbors.len()); - for (index, &source) in parents.iter().chain(neighbors).enumerate() { - let view = region - .get_view(source) - .ok_or_else(|| ANNError::message("IVF split source is missing from the region"))?; - let vectors = region.get_vectors(source).ok_or_else(|| { - ANNError::message("IVF split source vectors are missing from the region") - })?; - if index >= parents.len() { - centroids - .row_mut(child_count + index - parents.len()) - .copy_from_slice(view.centroid); +impl TwoMeans { + pub(super) fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { + Self { + iterations, + sums: Vec::new(), + dim, } - inputs.push((view.members, vectors)); } +} - let mut sums = vec![0.0; 2 * dim]; - for ((_, vectors), centers) in inputs[..parents.len()] - .iter() - .zip(centroids.as_mut_slice()[..child_count * dim].chunks_exact_mut(2 * dim)) - { - fit_two_means(vectors.as_view(), centers, dim, iterations, rng, &mut sums)?; - } +impl Clusterer for TwoMeans { + fn fit( + &mut self, + points: rowmajor::Ref<'_, f32>, + mut dst: rowmajor::Mut<'_, f32>, + rng: &mut StdRng, + ) -> ANNResult<()> { + let dim = self.dim.get(); + if points.nrows() < 2 { + return Err(ANNError::message("an IVF split needs at least two points")); + } + if points.ncols() != dim || dst.ncols() != dim || dst.nrows() != self.num_centers().get() { + return Err(ANNError::message( + "two-means needs canonical input rows and a two-row destination of the canonical dimension", + )); + } + if !points.as_slice().iter().all(|value| value.is_finite()) { + return Err(ANNError::message( + "two-means input coordinates must be finite", + )); + } + let centers = dst.as_mut_slice(); + self.sums.resize(centers.len(), 0.0); + let first = rng.random_range(0..points.nrows()); + let second = (first + rng.random_range(1..points.nrows())) % points.nrows(); + centers[..dim].copy_from_slice(points.row(first)); + centers[dim..].copy_from_slice(points.row(second)); - let count = inputs.iter().map(|(members, _)| members.len()).sum(); - let mut assignments = Vec::with_capacity(count); - let mut sizes = vec![0; centroids.nrows()]; - for (_, vectors) in &inputs { - for vector in vectors.rows() { - let to = nearest(vector, centroids.as_slice(), dim); - assignments.push(to); - sizes[to] += 1; + for _ in 0..self.iterations.get() { + self.sums.fill(0.0); + let mut counts = [0usize; 2]; + for point in points.rows() { + let to = nearest(point, centers.chunks_exact(dim)); + counts[to] += 1; + for (sum, &value) in self.sums[to * dim..(to + 1) * dim].iter_mut().zip(point) { + // A finite mean must not overflow while accumulating its inputs. + *sum += f64::from(value); + } + } + for ((center, sum), count) in centers + .chunks_exact_mut(dim) + .zip(self.sums.chunks_exact(dim)) + .zip(counts) + { + if count != 0 { + for (value, &sum) in center.iter_mut().zip(sum) { + *value = (sum / count as f64) as f32; + } + } + } } + Ok(()) } - let mut partitions: Vec> = std::iter::repeat_n(None, child_count) - .chain(neighbors.iter().copied().map(Some)) - .zip(sizes) - .map(|(id, size)| Partition { - id, - members: Vec::with_capacity(size), - }) - .collect(); - for (id, to) in inputs - .iter() - .flat_map(|(members, _)| members.iter().copied()) - .zip(assignments) - { - partitions[to].members.push(id); + + fn num_centers(&self) -> NonZeroUsize { + const { NonZeroUsize::new(2).unwrap() } } - Ok(Partitioning { - centroids, - partitions, - }) -} -fn fit_two_means( - points: rowmajor::Ref<'_, f32>, - centers: &mut [f32], - dim: usize, - iterations: usize, - rng: &mut StdRng, - sums: &mut [f64], -) -> ANNResult<()> { - if points.nrows() < 2 { - return Err(ANNError::message("an IVF split needs at least two points")); + fn dim(&self) -> NonZeroUsize { + self.dim } - let first = rng.random_range(0..points.nrows()); - let second = (first + rng.random_range(1..points.nrows())) % points.nrows(); - centers[..dim].copy_from_slice(points.row(first)); - centers[dim..].copy_from_slice(points.row(second)); - - for _ in 0..iterations.max(1) { - sums.fill(0.0); - let mut counts = [0; 2]; - for point in points.rows() { - let to = nearest(point, centers, dim); - counts[to] += 1; - for (sum, &value) in sums[to * dim..(to + 1) * dim].iter_mut().zip(point) { - // A finite mean must not overflow while accumulating its inputs. - *sum += f64::from(value); +} + +pub(super) struct RegionalNearest; + +impl Assigner for RegionalNearest { + fn assign( + &mut self, + region: &R, + update: &mut U, + sources: &[L], + destinations: &[Destination<'_, L>], + ) -> ANNResult<()> + where + R: RegionView, + U: RegionUpdate, + { + let mut memberships: Vec> = (0..destinations.len()).map(|_| Vec::new()).collect(); + + for &src in sources { + let view = region.get_view(src).ok_or_else(|| { + ANNError::message("IVF assignment source is missing from the region") + })?; + let vectors = region.get_vectors(src).ok_or_else(|| { + ANNError::message("IVF assignment source vectors are missing from the region") + })?; + + for (&id, point) in view.members.iter().zip(vectors.rows()) { + let dst = nearest(point, destinations.iter().map(|d| d.centroid)); + memberships[dst].push(id); } } - for ((center, sum), count) in centers - .chunks_exact_mut(dim) - .zip(sums.chunks_exact(dim)) - .zip(counts) - { - if count != 0 { - for (value, &sum) in center.iter_mut().zip(sum) { - *value = (sum / count as f64) as f32; + + let mut retained = BTreeSet::new(); + + for (destination, members) in destinations.iter().zip(memberships) { + match &destination.list { + Type::Exists(id) => { + update.set_members(*id, &members)?; + retained.insert(*id); } + Type::New => update.create(destination.centroid, &members)?, } } + + for &source in sources { + if !retained.contains(&source) { + update.retire(source)?; + } + } + + Ok(()) } - Ok(()) } -/// Prefer the first centroid on ties. `centers` contains at least one full row. -fn nearest(point: &[f32], centers: &[f32], dim: usize) -> usize { - let mut best = 0; - let mut best_distance: f32 = SquaredL2::evaluate(point, ¢ers[..dim]); - for (index, center) in centers.chunks_exact(dim).enumerate().skip(1) { - let distance: f32 = SquaredL2::evaluate(point, center); - if distance.total_cmp(&best_distance).is_lt() { - best = index; - best_distance = distance; +/// Prefer the first centroid on ties. +fn nearest<'a>(point: &[f32], targets: impl Iterator) -> usize { + let mut min = f32::MAX; + let mut best = usize::MAX; + for (i, center) in targets.enumerate() { + let dist = SquaredL2::evaluate(point, center); + if min > dist { + best = i; + min = dist; } } + best } diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index fe03f76b62..5d5c5dc443 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -33,11 +33,11 @@ pub struct Config { pub split_threshold: NonZeroUsize, /// Nearest-list query count per split parent, before filtering split parents. /// - /// [`ReassignCount::None`] disables neighborhood gathering. - /// All gathered points are reassigned across all children and retained neighbors. + /// - [`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 two_means_iterations: NonZeroUsize, + pub split_iterations: NonZeroUsize, pub seed: u64, } @@ -97,10 +97,13 @@ impl IVFIndex

{ if points.is_empty() { return Ok(InsertStats::default()); } + let mut accessor = strategy .build_accessor(provider, context) .into_ann_result()?; + let dim = accessor.dim(); + // Step 1: Compute the batch let batch = strategy.prepare(&mut accessor, points).await?; @@ -133,7 +136,9 @@ impl IVFIndex

{ } let splits = parents.len(); - let update = if !parents.is_empty() { + let mut update_final = None; + + if !parents.is_empty() { let mut neighbors = BTreeSet::new(); if let ReassignCount::Count(n) = config.reassign_count { @@ -144,16 +149,14 @@ impl IVFIndex

{ .nearest(pid, n.get()) .await? .into_iter() - .map(|(id, _)| id) - .filter(|id| !parents.contains(id)); + .filter(|(id, _)| !parents.contains(id)); - for id in nids { + for (id, _) in nids { neighbors.insert(id); } } } - let dim = accessor.dim(); let iter = parents.iter().chain(neighbors.iter()).copied(); let (region, mut update) = accessor.fill(iter).await?; @@ -163,17 +166,15 @@ impl IVFIndex

{ &mut update, &parents, &neighbors, + config.split_iterations, dim, - config.two_means_iterations.get(), rng, )?; - Some(update) - } else { - None + update_final = Some(update) }; - accessor.finish(update.as_ref()).await?; + accessor.finish(update_final.as_ref()).await?; Ok(InsertStats { inserted: points.len(), diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 1df037eac6..2929671ef1 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -15,6 +15,11 @@ //! 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. +//! The numerical driver separates per-parent fitting from whole-region assignment. +//! Fitting writes a fixed number of children per parent into reusable centroid +//! scratch. Assignment borrows those rows and retained centroids, groups members, +//! and records writes directly through the regional update object. +//! The index currently selects two-means and regional-nearest assignment. //! //! Clustering finishes reading its inputs before recording child creation and //! parent retirement and neighbor membership replacement through the separate diff --git a/diskann/src/ivf/partition.rs b/diskann/src/ivf/partition.rs index a1be388f1a..e4ee137568 100644 --- a/diskann/src/ivf/partition.rs +++ b/diskann/src/ivf/partition.rs @@ -5,27 +5,10 @@ //! Complete destination memberships produced by regional clustering. -use diskann_utils::views::rowmajor; - -/// One destination list in a [`Partitioning`]. +/// Complete membership of a new or retained destination list. #[derive(Debug)] pub struct Partition { /// `Some` retains a source list's identity and centroid; `None` creates a new list. pub id: Option, pub members: Vec, } - -/// A region's resulting centroids and complete memberships. -/// -/// There is one centroid row per partition; row `i` corresponds to `partitions[i]`. -/// Retained centroid rows must match their source centroids. The member groups -/// together contain every source member exactly once, including pending appends. -/// -/// Coordinates of points are not included: providers already own or have prepared -/// those vectors. Membership buffers can be transferred without copying their -/// contents; centroid rows are available for creating new lists. -#[derive(Debug)] -pub struct Partitioning { - pub centroids: rowmajor::Owned, - pub partitions: Vec>, -} diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index 72c4f7086a..c3f8cdf236 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -237,8 +237,8 @@ impl<'p> BuildAccessor for Build<'p> { where Self: 'a; - fn dim(&self) -> usize { - self.store.dim + fn dim(&self) -> NonZeroUsize { + NonZeroUsize::new(self.store.dim).unwrap() } fn router(&mut self) -> Self::Router<'_> { @@ -291,7 +291,8 @@ impl<'p> BuildAccessor for Build<'p> { .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(), 0.0); + let mut vectors = + rowmajor::Owned::from_element(members.len(), self.dim().get(), 0.0); for (&id, row) in members.iter().zip(vectors.rows_mut()) { let vector = self .batch @@ -310,7 +311,7 @@ impl<'p> BuildAccessor for Build<'p> { }, ); } - let mut update = Update::new(&self.sources, self.dim()); + let mut update = Update::new(&self.sources, self.dim().get()); update.fail_create = self.store.fail_create; update.fail_set_members = self.store.fail_set_members; update.fail_retire = self.store.fail_retire; @@ -404,7 +405,8 @@ impl<'a, 'p> InsertStrategy<'a, Store, &'p [f32]> for Strategy { points: &[(u32, &'p [f32])], ) -> impl SendFuture>>> { async move { - let mut vectors = rowmajor::Owned::from_element(points.len(), accessor.dim(), 0.0); + let mut vectors = + rowmajor::Owned::from_element(points.len(), accessor.dim().get(), 0.0); let mut ids = Vec::with_capacity(points.len()); for ((external, vector), out) in points.iter().zip(vectors.rows_mut()) { if vector.len() != out.len() || !vector.iter().all(|value| value.is_finite()) { @@ -442,7 +444,7 @@ fn index_with_neighbors(store: Store, threshold: usize, neighbors: usize) -> IVF split_threshold: NonZeroUsize::new(threshold).unwrap(), reassign_count: NonZeroUsize::new(neighbors) .map_or(ReassignCount::None, ReassignCount::Count), - two_means_iterations: NonZeroUsize::new(5).unwrap(), + split_iterations: NonZeroUsize::new(5).unwrap(), seed: 7, }, ) @@ -787,8 +789,8 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { &mut update, &BTreeSet::from([1]), &BTreeSet::from([2, 3, 4]), - 1, - 5, + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(1).unwrap(), &mut StdRng::seed_from_u64(7), ) .unwrap(); @@ -851,8 +853,8 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { &mut update, &BTreeSet::from([1]), &BTreeSet::new(), - 2, - 5, + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), &mut StdRng::seed_from_u64(7), ) .unwrap(); @@ -861,7 +863,7 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { assert!(update.set_members(1, &[]).is_err()); update }; - assert_eq!(build.dim(), 2); + assert_eq!(build.dim().get(), 2); update }; assert_eq!(update.created.len(), 2); @@ -940,7 +942,7 @@ async fn invalid_config_and_uninitialized_insertion_fail() { Config { split_threshold: NonZeroUsize::new(1).unwrap(), reassign_count: ReassignCount::None, - two_means_iterations: NonZeroUsize::new(5).unwrap(), + split_iterations: NonZeroUsize::new(5).unwrap(), seed: 7, }, ) @@ -993,8 +995,8 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() &mut update, &BTreeSet::from([1, 99]), &BTreeSet::new(), - 2, - 5, + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), &mut StdRng::seed_from_u64(7) ) .is_err() @@ -1005,12 +1007,34 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() &mut update, &BTreeSet::from([1]), &BTreeSet::from([99]), - 2, - 5, + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), + &mut StdRng::seed_from_u64(7) + ) + .is_err() + ); + assert!( + clustering::split( + ®ion, + &mut update, + &BTreeSet::from([1]), + &BTreeSet::from([1]), + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), &mut StdRng::seed_from_u64(7) ) .is_err() ); + clustering::split( + ®ion, + &mut update, + &BTreeSet::new(), + &BTreeSet::new(), + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); assert!(update.created.is_empty()); assert_eq!( std::mem::take(&mut update.memberships) @@ -1023,8 +1047,8 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() &mut update, &BTreeSet::from([1]), &BTreeSet::new(), - 2, - 5, + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), &mut StdRng::seed_from_u64(7), ) .unwrap(); @@ -1078,8 +1102,8 @@ fn assignment_can_cross_parent_boundaries() { &mut update, &BTreeSet::from([1, 2]), &BTreeSet::new(), - 1, - 10, + NonZeroUsize::new(10).unwrap(), + NonZeroUsize::new(1).unwrap(), &mut StdRng::seed_from_u64(7), ) .unwrap(); @@ -1103,34 +1127,489 @@ fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { )] .into(); let region = Region::new(&sources); - let mut zero = Update::new(&sources, 1); + let mut one = Update::new(&sources, 1); clustering::split( ®ion, - &mut zero, + &mut one, &BTreeSet::from([1]), &BTreeSet::new(), - 1, - 0, + NonZeroUsize::new(1).unwrap(), + NonZeroUsize::new(1).unwrap(), &mut StdRng::seed_from_u64(7), ) .unwrap(); - let mut one = Update::new(&sources, 1); - clustering::split( + assert_eq!(one.created.len(), 2); + for child in &one.created { + assert_eq!(child.centroid, [f32::MAX]); + } + assert_eq!(one.created[0].members, [0, 1, 2]); + assert!(one.created[1].members.is_empty()); +} + +#[test] +fn split_strategies_support_in_place_fitting_and_source_local_candidates() { + struct RowsAsCenters; + + impl clustering::Clusterer for RowsAsCenters { + fn fit( + &mut self, + points: rowmajor::Ref<'_, f32>, + mut dst: rowmajor::Mut<'_, f32>, + _: &mut StdRng, + ) -> ANNResult<()> { + assert_eq!(points.nrows(), dst.nrows()); + dst.as_mut_slice().copy_from_slice(points.as_slice()); + Ok(()) + } + + fn num_centers(&self) -> NonZeroUsize { + NonZeroUsize::new(2).unwrap() + } + + fn dim(&self) -> NonZeroUsize { + NonZeroUsize::new(1).unwrap() + } + } + + struct ParentLocal { + owners: BTreeMap>, + } + + impl clustering::Assigner for ParentLocal { + fn assign( + &mut self, + region: &R, + update: &mut U, + sources: &[u32], + destinations: &[clustering::Destination<'_, u32>], + ) -> ANNResult<()> + where + R: RegionView, + U: RegionUpdate, + { + assert_eq!(sources, [1, 2, 3, 4]); + assert_eq!(destinations.len(), 6); + assert!(std::ptr::eq( + destinations[4].centroid, + region.get_view(3).unwrap().centroid + )); + let children = [(1, 0..2), (2, 2..4)]; + let mut memberships = vec![Vec::new(); destinations.len()]; + for &source in sources { + let view = region.get_view(source).unwrap(); + let vectors = region.get_vectors(source).unwrap(); + let mut candidates: Vec<_> = destinations + .iter() + .enumerate() + .filter_map(|(index, destination)| { + matches!(&destination.list, clustering::Type::Exists(id) if *id == source) + .then_some(index) + }) + .collect(); + for (parent, range) in &children { + if *parent == source + || self + .owners + .get(&source) + .is_some_and(|owners| owners.contains(parent)) + { + candidates.extend(range.clone()); + } + } + for (&id, point) in view.members.iter().zip(vectors.rows()) { + let destination = candidates + .iter() + .copied() + .min_by(|&a, &b| { + let left: f32 = SquaredL2::evaluate(point, destinations[a].centroid); + let right: f32 = SquaredL2::evaluate(point, destinations[b].centroid); + left.total_cmp(&right) + }) + .unwrap(); + memberships[destination].push(id); + } + } + for (destination, members) in destinations.iter().zip(memberships) { + match &destination.list { + clustering::Type::New => update.create(destination.centroid, &members)?, + clustering::Type::Exists(id) => update.set_members(*id, &members)?, + } + } + for (parent, _) in children { + update.retire(parent)?; + } + Ok(()) + } + } + + 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 = BTreeSet::from([1, 2]); + let neighbors = BTreeSet::from([3, 4]); + let mut full = Update::new(&sources, 1); + let mut splitter = clustering::SplitAndAssign::new(RowsAsCenters, clustering::RegionalNearest); + splitter + .run( + ®ion, + &mut full, + &parents, + &neighbors, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); + let mut restricted = Update::new(&sources, 1); + clustering::SplitAndAssign::new( + RowsAsCenters, + ParentLocal { + owners: [(3, vec![1]), (4, vec![1, 2])].into(), + }, + ) + .run( ®ion, - &mut one, - &BTreeSet::from([1]), - &BTreeSet::new(), - 1, - 1, + &mut restricted, + &parents, + &neighbors, &mut StdRng::seed_from_u64(7), ) .unwrap(); - assert_eq!(zero.created.len(), 2); - assert_eq!(one.created.len(), 2); - for (zero, one) in zero.created.iter().zip(&one.created) { - assert_eq!(zero.centroid, one.centroid); - assert!(zero.centroid.iter().all(|value| value.is_finite())); + + for update in [&full, &restricted] { + assert_eq!( + update + .created + .iter() + .map(|list| list.centroid[0]) + .collect::>(), + [-10.0, 10.0, 100.0, 110.0] + ); + assert_eq!(update.created[0].members, [0]); + assert_eq!(update.created[1].members, [1, 4]); + assert_eq!(update.created[3].members, [3]); + } + assert_eq!(full.created[2].members, [2, 5, 7]); + assert_eq!(restricted.created[2].members, [2, 7]); + let mut repeated = Update::new(&sources, 1); + splitter + .run( + ®ion, + &mut repeated, + &parents, + &neighbors, + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); + for (first, repeated) in full.created.iter().zip(&repeated.created) { + assert_eq!(first.centroid, repeated.centroid); + assert_eq!(first.members, repeated.members); + } + let mut smaller = Update::new(&sources, 1); + splitter + .run( + ®ion, + &mut smaller, + &BTreeSet::from([1]), + &BTreeSet::new(), + &mut StdRng::seed_from_u64(7), + ) + .unwrap(); + assert_eq!(smaller.created.len(), 2); + assert_eq!(smaller.created[0].members, [0]); + assert_eq!(smaller.created[1].members, [1]); + for (update, expected) in [(full, vec![6]), (restricted, vec![5, 6])] { + let mut changes: BTreeMap<_, _> = update.memberships.into_changes().collect(); + for parent in [1, 2] { + assert!(matches!(changes.remove(&parent), Some(Membership::Retire))); + } + for (id, expected) in [(3, expected), (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 numerical_strategy_failures_record_no_updates() { + enum InvalidFit { + Error, + TooFew, + WrongDimension, + NonFinite, + Incomplete, + CountOverflow, + } + + impl clustering::Clusterer for InvalidFit { + fn fit( + &mut self, + _: rowmajor::Ref<'_, f32>, + mut dst: rowmajor::Mut<'_, f32>, + _: &mut StdRng, + ) -> ANNResult<()> { + match self { + Self::Error => Err(ANNError::message("injected fitting failure")), + Self::NonFinite => { + dst.as_mut_slice().fill(f32::INFINITY); + Ok(()) + } + Self::Incomplete => { + dst.row_mut(0).fill(0.0); + Ok(()) + } + _ => unreachable!("invalid centroid layout should be rejected before fitting"), + } + } + + fn num_centers(&self) -> NonZeroUsize { + NonZeroUsize::new(match self { + Self::TooFew => 1, + Self::CountOverflow => usize::MAX, + _ => 2, + }) + .unwrap() + } + + fn dim(&self) -> NonZeroUsize { + NonZeroUsize::new(if matches!(self, Self::WrongDimension) { + 2 + } else { + 1 + }) + .unwrap() + } + } + + struct InvalidAssignment; + + impl clustering::Assigner for InvalidAssignment { + fn assign( + &mut self, + _: &R, + _: &mut U, + _: &[u32], + _: &[clustering::Destination<'_, u32>], + ) -> ANNResult<()> + where + R: RegionView, + U: RegionUpdate, + { + Err(ANNError::message("injected assignment failure")) + } + } + + fn assert_unwritten(update: Update) { + assert!(update.created.is_empty()); + assert_eq!(update.memberships.into_changes().count(), 0); + } + + let sources = [( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + )] + .into(); + let region = Region::new(&sources); + let parents = BTreeSet::from([1]); + let neighbors = BTreeSet::new(); + for (clusterer, expected) in [ + (InvalidFit::Error, "injected fitting failure"), + (InvalidFit::TooFew, "at least two centroids"), + (InvalidFit::WrongDimension, "canonical dimension"), + (InvalidFit::NonFinite, "finite value"), + (InvalidFit::Incomplete, "every centroid coordinate"), + (InvalidFit::CountOverflow, "centroid count overflow"), + ] { + let mut update = Update::new(&sources, 1); + let selected = if matches!(clusterer, InvalidFit::CountOverflow) { + BTreeSet::from([1, 2]) + } else { + parents.clone() + }; + let error = clustering::SplitAndAssign::new(clusterer, clustering::RegionalNearest) + .run( + ®ion, + &mut update, + &selected, + &neighbors, + &mut StdRng::seed_from_u64(7), + ) + .unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + assert_unwritten(update); + } + let mut update = Update::new(&sources, 1); + let error = clustering::SplitAndAssign::new( + clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), + InvalidAssignment, + ) + .run( + ®ion, + &mut update, + &parents, + &neighbors, + &mut StdRng::seed_from_u64(7), + ) + .unwrap_err(); + assert!( + error.to_string().contains("injected assignment failure"), + "{error}" + ); + assert_unwritten(update); +} + +#[test] +fn two_means_fits_only_the_destination_window() { + use clustering::{Clusterer, TwoMeans}; + + let points = matrix(&[[0.0, 2.0], [2.0, 4.0], [10.0, 12.0], [12.0, 14.0]]); + let mut scratch = rowmajor::Owned::from_element(4, 2, -1.0); + let mut clusterer = TwoMeans::new( + NonZeroUsize::new(10).unwrap(), + NonZeroUsize::new(2).unwrap(), + ); + let mut rng = StdRng::seed_from_u64(7); + clusterer + .fit( + points.as_view(), + scratch.subview_mut(1..3).unwrap(), + &mut rng, + ) + .unwrap(); + assert_eq!(scratch.row(0), [-1.0, -1.0]); + assert_eq!(scratch.row(3), [-1.0, -1.0]); + let window = scratch.subview(1..3).unwrap(); + let mut centers: Vec<_> = window.rows().collect(); + centers.sort_by(|a, b| a[0].total_cmp(&b[0])); + assert_eq!(centers, [&[1.0, 3.0][..], &[11.0, 13.0][..]]); + + for (points, rows, dim) in [ + (matrix(&[[0.0, 2.0]]), 2, 2), + (matrix(&[[0.0], [1.0]]), 2, 2), + (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 1, 2), + (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 2, 1), + (matrix(&[[f32::NAN, 2.0], [1.0, 3.0]]), 2, 2), + ] { + let mut dst = rowmajor::Owned::from_element(rows, dim, -1.0); + assert!( + clusterer + .fit(points.as_view(), dst.as_view_mut(), &mut rng) + .is_err() + ); + assert!(dst.as_slice().iter().all(|&value| value == -1.0)); + } +} + +#[test] +fn regional_assignment_rejects_invalid_inputs_before_writing() { + use clustering::{Assigner, Destination, RegionalNearest, Type}; + + let mut 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]]), + }, + ), + ]); + let destinations = [Destination { + list: Type::New, + centroid: &[0.0], + }]; + for invalid in [ + vec![], + vec![Destination { + list: Type::New, + centroid: &[], + }], + vec![Destination { + list: Type::New, + centroid: &[f32::NAN], + }], + vec![ + Destination { + list: Type::New, + centroid: &[0.0], + }, + Destination { + list: Type::New, + centroid: &[1.0, 2.0], + }, + ], + ] { + let mut update = Update::new(&sources, 1); + assert!( + RegionalNearest + .assign(&Region::new(&sources), &mut update, &[1, 2], &invalid) + .is_err() + ); + assert!(update.created.is_empty()); + assert_eq!(update.memberships.into_changes().count(), 0); + } + for (ids, members, vectors) in [ + (vec![1, 99], vec![2], matrix(&[[2.0]])), + (vec![1, 2], vec![], matrix(&[[2.0]])), + (vec![1, 2], vec![2], matrix(&[[0.0, 1.0]])), + (vec![1, 2], vec![2], matrix(&[[f32::INFINITY]])), + ] { + let mut update = Update::new(&sources, 1); + sources.get_mut(&2).unwrap().members = members; + sources.get_mut(&2).unwrap().vectors = vectors; + let region = Region::new(&sources); + assert!( + RegionalNearest + .assign(®ion, &mut update, &ids, &destinations) + .is_err() + ); + assert!(update.created.is_empty()); + assert_eq!(update.memberships.into_changes().count(), 0); } - assert_eq!(zero.created[0].members, [0, 1, 2]); - assert!(zero.created[1].members.is_empty()); } diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index da4ce3f6ce..2a0fbccd3b 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -92,7 +92,7 @@ pub trait BuildAccessor: Send + Sized + HasIds { type Update: RegionUpdate; /// Positive dimension of all canonical points and centroids. - fn dim(&self) -> usize; + fn dim(&self) -> NonZeroUsize; fn router(&mut self) -> Self::Router<'_>; diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 06269c4e05..37b5bf9ae9 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -77,6 +77,20 @@ unchanged. Explicit source IDs avoid requiring enumeration on the region view. Numerical computation borrows vector matrices and member slices, without building per-point reference tables or copying vector coordinates. A provider with scattered storage may pack rows during filling to satisfy the matrix interface. +The numerical driver separates per-parent fitting (`Clusterer`) from whole-region +placement (`Assigner`). A clusterer declares its dimension and fixed number of +centers per parent, then fills mutable row windows of reusable centroid scratch. +A borrowed destination catalog combines the completed rows with unchanged retained +centroids; retained coordinates are not copied. The assigner receives the region, +source IDs, destinations, and regional writer directly. It computes complete +destination memberships before creating new lists, replacing retained memberships, +and retiring replaced sources. The index still selects two-means and full regional +nearest-centroid assignment. Restricted assignment can keep caller-supplied +candidate relationships in its own state, but the index does not yet preserve those +relationships or expose numerical strategy selection. +Alternative geometries must keep fitting, assignment, and provider routing +compatible. + 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; From e3beea5837f9927768f177e8b90cec08bf779517 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 19 Oct 2026 16:07:38 -0400 Subject: [PATCH 21/33] move to lloyd impl in diskann-quantization --- Cargo.lock | 1 + diskann/Cargo.toml | 1 + diskann/src/ivf/clustering.rs | 76 ++++++++--------------------- diskann/src/ivf/index.rs | 10 +--- diskann/src/ivf/tests.rs | 90 ++++++++++++++++++++++++----------- rfcs/01187-incremental-ivf.md | 10 ++-- 6 files changed, 92 insertions(+), 96 deletions(-) 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/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/clustering.rs b/diskann/src/ivf/clustering.rs index 11d9323261..a5d26a48cd 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -10,9 +10,9 @@ use std::{collections::BTreeSet, num::NonZeroUsize}; +use diskann_quantization::algorithms::kmeans::lloyds::lloyds; use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; -use rand::{Rng, rngs::StdRng}; use super::region::{RegionUpdate, RegionView}; use crate::{ANNError, ANNResult}; @@ -24,12 +24,7 @@ pub(super) trait Clusterer { /// Implementations own initialization, convergence, empty-cluster handling, /// and reusable fitting scratch. `dst` has `num_centers()` rows and `dim()` /// columns. The number of centers must be at least two. - fn fit( - &mut self, - points: rowmajor::Ref<'_, T>, - dst: rowmajor::Mut<'_, T>, - rng: &mut StdRng, - ) -> ANNResult<()>; + fn fit(&mut self, points: rowmajor::Ref<'_, T>, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; fn num_centers(&self) -> NonZeroUsize; @@ -114,10 +109,9 @@ pub(super) fn split, U: RegionUpdate neighbors: &BTreeSet, iterations: NonZeroUsize, dim: NonZeroUsize, - rng: &mut StdRng, ) -> ANNResult<()> { SplitAndAssign::new(TwoMeans::new(iterations, dim), RegionalNearest) - .run(region, update, parents, neighbors, rng) + .run(region, update, parents, neighbors) } impl SplitAndAssign { @@ -128,7 +122,6 @@ impl SplitAndAssign { update: &mut U, parents: &BTreeSet, neighbors: &BTreeSet, - rng: &mut StdRng, ) -> ANNResult<()> where A: Assigner, @@ -158,7 +151,6 @@ impl SplitAndAssign { clusters.subview_mut(s..end).ok_or_else(|| { ANNError::message("IVF centroid scratch range is out of bounds") })?, - rng, )?; s = end; @@ -190,19 +182,15 @@ impl SplitAndAssign { } } +/// Random input seeds followed by Lloyd iterations, which reset empty centers to zero. pub(super) struct TwoMeans { iterations: NonZeroUsize, - sums: Vec, dim: NonZeroUsize, } impl TwoMeans { pub(super) fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { - Self { - iterations, - sums: Vec::new(), - dim, - } + Self { iterations, dim } } } @@ -211,52 +199,28 @@ impl Clusterer for TwoMeans { &mut self, points: rowmajor::Ref<'_, f32>, mut dst: rowmajor::Mut<'_, f32>, - rng: &mut StdRng, ) -> ANNResult<()> { - let dim = self.dim.get(); - if points.nrows() < 2 { + let dim = self.dim().get(); + let num_centers = self.num_centers().get(); + + if points.nrows() < num_centers { return Err(ANNError::message("an IVF split needs at least two points")); } - if points.ncols() != dim || dst.ncols() != dim || dst.nrows() != self.num_centers().get() { + if points.ncols() != dim || dst.ncols() != dim || dst.nrows() != num_centers { return Err(ANNError::message( "two-means needs canonical input rows and a two-row destination of the canonical dimension", )); } - if !points.as_slice().iter().all(|value| value.is_finite()) { - return Err(ANNError::message( - "two-means input coordinates must be finite", - )); - } - let centers = dst.as_mut_slice(); - self.sums.resize(centers.len(), 0.0); - let first = rng.random_range(0..points.nrows()); - let second = (first + rng.random_range(1..points.nrows())) % points.nrows(); - centers[..dim].copy_from_slice(points.row(first)); - centers[dim..].copy_from_slice(points.row(second)); - - for _ in 0..self.iterations.get() { - self.sums.fill(0.0); - let mut counts = [0usize; 2]; - for point in points.rows() { - let to = nearest(point, centers.chunks_exact(dim)); - counts[to] += 1; - for (sum, &value) in self.sums[to * dim..(to + 1) * dim].iter_mut().zip(point) { - // A finite mean must not overflow while accumulating its inputs. - *sum += f64::from(value); - } - } - for ((center, sum), count) in centers - .chunks_exact_mut(dim) - .zip(self.sums.chunks_exact(dim)) - .zip(counts) - { - if count != 0 { - for (value, &sum) in center.iter_mut().zip(sum) { - *value = (sum / count as f64) as f32; - } - } - } - } + + dst.as_mut_slice().copy_from_slice( + points + .subview(0..num_centers) + .ok_or(ANNError::message("oh no!"))? + .as_slice(), + ); + + let _ = lloyds(points, dst, self.iterations.get()); + Ok(()) } diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 5d5c5dc443..5c9c80d2ab 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -11,7 +11,6 @@ use std::{ }; use diskann_utils::future::SendFuture; -use rand::{SeedableRng, rngs::StdRng}; use thiserror::Error; use super::{ @@ -57,7 +56,6 @@ pub struct InsertStats { pub struct IVFIndex { provider: P, config: Config, - rng: StdRng, } impl IVFIndex

{ @@ -68,11 +66,7 @@ impl IVFIndex

{ if config.split_threshold.get() < 2 { return Err(ConfigError(config.split_threshold.get())); } - Ok(Self { - provider, - rng: StdRng::seed_from_u64(config.seed), - config, - }) + Ok(Self { provider, config }) } pub fn provider(&self) -> &P { @@ -91,7 +85,6 @@ impl IVFIndex

{ { let provider = &mut self.provider; let config = self.config; - let rng = &mut self.rng; async move { if points.is_empty() { @@ -168,7 +161,6 @@ impl IVFIndex

{ &neighbors, config.split_iterations, dim, - rng, )?; update_final = Some(update) diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index c3f8cdf236..66cc106bde 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -1115,35 +1115,35 @@ fn assignment_can_cross_parent_boundaries() { } #[test] -fn two_means_keeps_empty_children_finite_and_prefers_first_ties() { - let rows = [[f32::MAX], [f32::MAX], [f32::MAX]]; - let sources = [( - 1, - Source { - centroid: vec![f32::MAX], - members: vec![0, 1, 2], - vectors: matrix(&rows), - }, - )] - .into(); - let region = Region::new(&sources); - let mut one = Update::new(&sources, 1); - clustering::split( - ®ion, - &mut one, - &BTreeSet::from([1]), - &BTreeSet::new(), - NonZeroUsize::new(1).unwrap(), - NonZeroUsize::new(1).unwrap(), - &mut StdRng::seed_from_u64(7), - ) - .unwrap(); - assert_eq!(one.created.len(), 2); - for child in &one.created { - assert_eq!(child.centroid, [f32::MAX]); +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 update = Update::new(&sources, 1); + clustering::split( + ®ion, + &mut update, + &BTreeSet::from([1]), + &BTreeSet::new(), + NonZeroUsize::new(1).unwrap(), + NonZeroUsize::new(1).unwrap(), + &mut StdRng::seed_from_u64(7), + ) + .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()); } - assert_eq!(one.created[0].members, [0, 1, 2]); - assert!(one.created[1].members.is_empty()); } #[test] @@ -1496,6 +1496,40 @@ fn numerical_strategy_failures_record_no_updates() { assert_unwritten(update); } +#[test] +fn two_means_matches_quantization_lloyds() { + use clustering::{Clusterer, TwoMeans}; + use rand::Rng; + + let points = matrix(&[[0.0, 2.0], [2.0, 4.0], [6.0, 6.0], [10.0, 12.0]]); + for seed in [0, 7, 42] { + for iterations in [1, 5] { + let mut expected = rowmajor::Owned::from_element(2, 2, 0.0); + let mut reference_rng = StdRng::seed_from_u64(seed); + let first = reference_rng.random_range(0..points.nrows()); + let second = (first + reference_rng.random_range(1..points.nrows())) % points.nrows(); + expected.row_mut(0).copy_from_slice(points.row(first)); + expected.row_mut(1).copy_from_slice(points.row(second)); + let _ = diskann_quantization::algorithms::kmeans::lloyds::lloyds( + points.as_view(), + expected.as_view_mut(), + iterations, + ); + + let mut actual = rowmajor::Owned::from_element(2, 2, 0.0); + let mut rng = StdRng::seed_from_u64(seed); + TwoMeans::new( + NonZeroUsize::new(iterations).unwrap(), + NonZeroUsize::new(2).unwrap(), + ) + .fit(points.as_view(), actual.as_view_mut(), &mut rng) + .unwrap(); + assert_eq!(actual.as_slice(), expected.as_slice()); + assert_eq!(rng.random::(), reference_rng.random::()); + } + } +} + #[test] fn two_means_fits_only_the_destination_window() { use clustering::{Clusterer, TwoMeans}; diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 37b5bf9ae9..3217e93ea0 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -74,9 +74,13 @@ 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. -Numerical computation borrows vector matrices and member slices, without building -per-point reference tables or copying vector coordinates. A provider with -scattered storage may pack rows during filling to satisfy the matrix interface. +The numerical driver borrows vector matrices and member slices, without building +per-point reference tables. A provider with scattered storage may pack rows during +filling to satisfy the matrix interface. The built-in two-means fitter seeds two +distinct random 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. The numerical driver separates per-parent fitting (`Clusterer`) from whole-region placement (`Assigner`). A clusterer declares its dimension and fixed number of centers per parent, then fills mutable row windows of reusable centroid scratch. From a4cee326a3e1e3d6b5d7fc9bc6d9ba9204cd0928 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 19 Oct 2026 16:17:06 -0400 Subject: [PATCH 22/33] some more small edits to clustering --- diskann/src/ivf/clustering.rs | 19 ------------------- diskann/src/ivf/index.rs | 14 +++++++------- 2 files changed, 7 insertions(+), 26 deletions(-) diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index a5d26a48cd..abd8f8604e 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -95,25 +95,6 @@ impl SplitAndAssign { } } -/// Fit two children per parent and reassign sources across children and retained neighbors. -/// -/// Parents and neighbors are distinct, disjoint filled sources. Neighbor centroids -/// are retained unchanged. Other filled lists are not read or edited. -/// Computation completes before any edits are recorded. A recording error may -/// leave partial pending edits, which the caller must discard rather than publish. -/// An empty set of parents is a no-op. -pub(super) fn split, U: RegionUpdate>( - region: &R, - update: &mut U, - parents: &BTreeSet, - neighbors: &BTreeSet, - iterations: NonZeroUsize, - dim: NonZeroUsize, -) -> ANNResult<()> { - SplitAndAssign::new(TwoMeans::new(iterations, dim), RegionalNearest) - .run(region, update, parents, neighbors) -} - impl SplitAndAssign { /// Compute the complete result before registering any provider writes. pub(super) fn run( diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 5c9c80d2ab..796f0ef7d4 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -13,11 +13,12 @@ use std::{ use diskann_utils::future::SendFuture; use thiserror::Error; -use super::{ - clustering, - traits::{BuildAccessor, InsertStrategy, Provider, Router, RouterAndNearest}, +use super::traits::{BuildAccessor, InsertStrategy, Provider, Router, RouterAndNearest}; +use crate::{ + ANNError, ANNResult, + error::IntoANNResult, + ivf::clustering::{RegionalNearest, SplitAndAssign, TwoMeans}, }; -use crate::{ANNError, ANNResult, error::IntoANNResult}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ReassignCount { @@ -154,13 +155,12 @@ impl IVFIndex

{ let (region, mut update) = accessor.fill(iter).await?; - clustering::split( + let clusterer = TwoMeans::new(config.split_iterations, dim); + SplitAndAssign::new(clusterer, RegionalNearest).run( ®ion, &mut update, &parents, &neighbors, - config.split_iterations, - dim, )?; update_final = Some(update) From 41c19b87d1d4219f646ac3048c6329a765e61f22 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 19 Oct 2026 17:46:50 -0400 Subject: [PATCH 23/33] clean up staging and remove provider --- diskann/src/ivf/index.rs | 50 ++++++++++++------------- diskann/src/ivf/region.rs | 54 -------------------------- diskann/src/ivf/traits.rs | 79 +++++++++++---------------------------- 3 files changed, 45 insertions(+), 138 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 796f0ef7d4..203c14ae69 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -7,17 +7,22 @@ use std::{ collections::{BTreeMap, BTreeSet}, + marker::PhantomData, num::NonZeroUsize, }; use diskann_utils::future::SendFuture; use thiserror::Error; -use super::traits::{BuildAccessor, InsertStrategy, Provider, Router, RouterAndNearest}; +use super::traits::{BuildAccessor, Router, RouterAndNearest, SetElements}; use crate::{ ANNError, ANNResult, error::IntoANNResult, - ivf::clustering::{RegionalNearest, SplitAndAssign, TwoMeans}, + ivf::{ + clustering::{RegionalNearest, SplitAndAssign, TwoMeans}, + traits::BuildStrategy, + }, + provider::ExecutionContext, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -54,37 +59,32 @@ pub struct InsertStats { /// An IVF index with serialized mutation through the underlying provider. #[derive(Debug)] -pub struct IVFIndex { - provider: P, +pub struct IVFIndex { config: Config, + phantom: PhantomData, } -impl IVFIndex

{ +impl IVFIndex { /// # Errors /// /// Returns [`ConfigError`] when the split threshold is less than two. - pub fn new(provider: P, config: Config) -> Result { - if config.split_threshold.get() < 2 { - return Err(ConfigError(config.split_threshold.get())); - } - Ok(Self { provider, config }) - } - - pub fn provider(&self) -> &P { - &self.provider + pub fn new(config: Config) -> Result { + Ok(Self { + config, + phantom: PhantomData::default(), + }) } pub fn insert_batch<'a, S, T>( &'a mut self, strategy: &'a S, - context: &'a P::Context, - points: &'a [(P::ExternalId, T)], + context: &'a C, + points: &'a [T], ) -> impl SendFuture> where - S: InsertStrategy<'a, P, T>, + S: BuildStrategy<'a, T, C>, T: Sync, { - let provider = &mut self.provider; let config = self.config; async move { @@ -92,29 +92,27 @@ impl IVFIndex

{ return Ok(InsertStats::default()); } - let mut accessor = strategy - .build_accessor(provider, context) - .into_ann_result()?; + let mut accessor = strategy.build_accessor(context).into_ann_result()?; let dim = accessor.dim(); // Step 1: Compute the batch - let batch = strategy.prepare(&mut accessor, points).await?; + let ids = accessor.set_elements(context, points).await?; // Step 2: Route the points to their nearest centroid. - let mut routed = BTreeMap::>::new(); + let mut routed = BTreeMap::>::new(); { let router = accessor.router(); - for point in batch.points() { - let selected = router.route(point.vector, 1).await?; + for (id, point) in ids.iter().zip(points.iter()) { + let selected = router.route(point, 1).await?; let (list, _) = selected.first().ok_or_else(|| { ANNError::message("IVF routing returned no live centroid") })?; - routed.entry(*list).or_default().push(point.id); + routed.entry(*list).or_default().push(*id); } } diff --git a/diskann/src/ivf/region.rs b/diskann/src/ivf/region.rs index 0054b9423f..132dda7d38 100644 --- a/diskann/src/ivf/region.rs +++ b/diskann/src/ivf/region.rs @@ -44,57 +44,3 @@ pub trait RegionUpdate: Send { /// Record a new centroid and its complete membership. fn create(&mut self, centroid: &[f32], members: &[I]) -> ANNResult<()>; } - -/// A point's identity and borrowed canonical vector. -#[derive(Debug, Clone, Copy)] -pub struct PointRef<'a, I> { - pub id: I, - pub vector: &'a [f32], -} - -/// Canonical rows shared between routing and an operation's build accessor. -/// -/// IDs must be unique. The contiguous matrix is owned once; iteration and lookup -/// borrow its rows without copying coordinates. -#[derive(Debug)] -pub struct Batch { - ids: Vec, - vectors: rowmajor::Owned, - rows: HashMap, -} - -impl Batch { - /// # Errors - /// - /// Fails if the matrix has zero columns or its row count differs from `ids`. - pub fn new(ids: Vec, vectors: rowmajor::Owned) -> ANNResult { - if vectors.ncols() == 0 || vectors.nrows() != ids.len() { - return Err(ANNError::message( - "an IVF batch needs one nonempty canonical row per point ID", - )); - } - let rows = ids - .iter() - .copied() - .enumerate() - .map(|(row, id)| (id, row)) - .collect(); - Ok(Self { ids, vectors, rows }) - } - - pub fn vectors(&self) -> rowmajor::Ref<'_, f32> { - self.vectors.as_view() - } - - pub fn points(&self) -> impl Iterator> { - self.ids - .iter() - .copied() - .zip(self.vectors.rows()) - .map(|(id, vector)| PointRef { id, vector }) - } - - pub fn get(&self, id: I) -> Option<&[f32]> { - self.rows.get(&id).map(|&row| self.vectors.row(row)) - } -} diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 2a0fbccd3b..399c1f62e0 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -5,41 +5,12 @@ //! Storage boundaries for IVF construction. -use std::{fmt::Debug, num::NonZeroUsize, sync::Arc}; +use std::num::NonZeroUsize; use diskann_utils::future::SendFuture; -use super::region::{Batch, RegionUpdate, RegionView}; -use crate::{ - ANNResult, - error::{StandardError, ToRanked}, - provider::ExecutionContext, - utils::VectorId, -}; - -/// Point identity and execution context shared by an IVF provider's accessors. -/// -/// Prepared points become translatable only after the build operation finishes. -pub trait Provider: Sized + Send + Sync + 'static { - type Context: ExecutionContext; - type InternalId: VectorId; - type ExternalId: PartialEq + Send + Sync + 'static; - /// Stable identity shared by a centroid and its inverted list. Never reused. - type ListId: VectorId; - type Error: ToRanked + Debug + Send + Sync + 'static; - - fn to_internal_id( - &self, - context: &Self::Context, - id: &Self::ExternalId, - ) -> Result; - - fn to_external_id( - &self, - context: &Self::Context, - id: Self::InternalId, - ) -> Result; -} +use super::region::{RegionUpdate, RegionView}; +use crate::{ANNResult, error::StandardError, provider::ExecutionContext, utils::VectorId}; pub trait HasIds { type Id: VectorId; @@ -47,7 +18,7 @@ pub trait HasIds { } /// Centroid selection within one build operation. -pub trait Router: Send { +pub trait Router: Send { type Id: VectorId; fn len(&self) -> usize; @@ -57,14 +28,10 @@ pub trait Router: Send { } /// Select `min(nprobe, len())` distinct live lists, nearest first. - fn route( - &self, - query: &[f32], - nprobe: usize, - ) -> impl SendFuture>>; + fn route(&self, query: &T, nprobe: usize) -> impl SendFuture>>; } -pub trait RouterAndNearest: Router { +pub trait RouterAndNearest: Router { /// Select `min(nprobe, len())` distinct live lists nearest to `id`'s centroid. /// /// Results are nearest first and may include `id` itself. Callers exclude @@ -80,8 +47,8 @@ pub trait RouterAndNearest: Router { ) -> impl SendFuture>>; } -pub trait BuildAccessor: Send + Sized + HasIds { - type Router<'a>: RouterAndNearest +pub trait BuildAccessor: Send + Sized + HasIds { + type Router<'a>: RouterAndNearest where Self: 'a; @@ -100,7 +67,7 @@ pub trait BuildAccessor: Send + Sized + HasIds { fn append_to( &mut self, list: Self::ListId, - members: Itr, + ids: Itr, ) -> impl SendFuture> where Itr: ExactSizeIterator + Clone + Send + Sync; @@ -120,24 +87,20 @@ pub trait BuildAccessor: Send + Sized + HasIds { fn finish(self, update: Option<&Self::Update>) -> impl SendFuture>; } -pub trait BuildStrategy<'a, P: Provider>: Send + Sync { - type Accessor: BuildAccessor; +pub trait BuildStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { + type Accessor: BuildAccessor + + SetElements; type Error: StandardError; - fn build_accessor( - &'a self, - provider: &'a mut P, - context: &'a P::Context, - ) -> Result; + fn build_accessor(&'a self, context: &'a C) -> Result; } -/// Prepare typed inputs and seed a build accessor with their canonical data. -/// -/// Only this boundary depends on `T`; list operations and clustering do not. -pub trait InsertStrategy<'a, P: Provider, T: Sync>: BuildStrategy<'a, P> { - fn prepare( - &self, - accessor: &mut Self::Accessor, - points: &[(P::ExternalId, T)], - ) -> impl SendFuture>>>; +pub trait SetElements { + type Id: VectorId; + + fn set_elements( + &mut self, + context: &C, + elements: &[T], + ) -> impl std::future::Future>> + Send; } From b9f0d93e79148cf7f5362418a31e14c307fea815 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Mon, 19 Oct 2026 19:38:45 -0400 Subject: [PATCH 24/33] propagate to tests --- diskann/src/ivf/index.rs | 2 +- diskann/src/ivf/region.rs | 7 +- diskann/src/ivf/tests.rs | 721 ++++++++++++++++++-------------------- diskann/src/ivf/traits.rs | 2 +- 4 files changed, 341 insertions(+), 391 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 203c14ae69..e515104856 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -164,7 +164,7 @@ impl IVFIndex { update_final = Some(update) }; - accessor.finish(update_final.as_ref()).await?; + accessor.finish(update_final).await?; Ok(InsertStats { inserted: points.len(), diff --git a/diskann/src/ivf/region.rs b/diskann/src/ivf/region.rs index 132dda7d38..c0550eb158 100644 --- a/diskann/src/ivf/region.rs +++ b/diskann/src/ivf/region.rs @@ -3,12 +3,11 @@ * Licensed under the MIT license. */ -//! Regional input views, write operations, and a reusable canonical insertion batch. +//! Regional input views and write operations. -use diskann_utils::views::rowmajor::{self, Matrix}; -use hashbrown::HashMap; +use diskann_utils::views::rowmajor; -use crate::{ANNError, ANNResult, utils::VectorId}; +use crate::ANNResult; /// One source list's centroid and complete operation-local membership. /// diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index 66cc106bde..cb87dd729a 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -8,7 +8,6 @@ use std::{ collections::{BTreeMap, BTreeSet}, future::Future, num::NonZeroUsize, - sync::Arc, }; use diskann_utils::{ @@ -16,15 +15,13 @@ use diskann_utils::{ views::rowmajor::{self, Matrix, MatrixMut}, }; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; -use rand::{SeedableRng, rngs::StdRng}; +use tokio::sync::{Mutex, MutexGuard}; use super::{ Config, IVFIndex, InsertStats, ReassignCount, clustering, pending::{Membership, PendingLists}, - region::{Batch, RegionUpdate, RegionView, View}, - traits::{ - BuildAccessor, BuildStrategy, HasIds, InsertStrategy, Provider, Router, RouterAndNearest, - }, + region::{RegionUpdate, RegionView, View}, + traits::{BuildAccessor, BuildStrategy, HasIds, Router, RouterAndNearest, SetElements}, }; use crate::{ANNError, ANNResult, provider::DefaultContext}; @@ -51,13 +48,7 @@ struct Store { fail_finish: bool, } -impl Provider for Store { - type Context = DefaultContext; - type InternalId = u32; - type ExternalId = u32; - type ListId = u32; - type Error = ANNError; - +impl Store { fn to_internal_id(&self, _: &DefaultContext, id: &u32) -> ANNResult { id.checked_add(10_000) .filter(|id| self.points.contains_key(id)) @@ -73,10 +64,10 @@ impl Provider for Store { } } -struct Build<'p> { - store: &'p mut Store, +struct Build<'p, T> { + store: MutexGuard<'p, Store>, pending: PendingLists, - batch: Option>>, + staged: BTreeMap>, sources: BTreeMap, fills: Cell, filled_lists: Vec, @@ -87,28 +78,46 @@ struct Coarse<'a> { fail_nearest: bool, } -impl Router for Coarse<'_> { +impl Coarse<'_> { + fn select>(&self, query: &[T], nprobe: usize) -> Vec<(u32, f32)> { + let mut selected: Vec<_> = self + .stored + .iter() + .map(|(&id, list)| { + let distance: f32 = query + .iter() + .zip(&list.centroid) + .map(|(&value, ¢er)| { + let delta = value.into() - center; + delta * delta + }) + .sum(); + (id, distance) + }) + .collect(); + selected.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))); + selected.truncate(nprobe); + selected + } +} + +impl + Sync> Router<(u32, &[T])> for Coarse<'_> { type Id = u32; fn len(&self) -> usize { self.stored.len() } - fn route(&self, query: &[f32], nprobe: usize) -> impl SendFuture>> { - async move { - let mut selected: Vec<(u32, f32)> = self - .stored - .iter() - .map(|(&id, list)| (id, SquaredL2::evaluate(query, list.centroid.as_slice()))) - .collect(); - selected.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))); - selected.truncate(nprobe); - Ok(selected) - } + fn route( + &self, + query: &(u32, &[T]), + nprobe: usize, + ) -> impl SendFuture>> { + async move { Ok(self.select(query.1, nprobe)) } } } -impl RouterAndNearest for Coarse<'_> { +impl + Sync> RouterAndNearest<(u32, &[T])> for Coarse<'_> { fn nearest(&self, id: Self::Id, nprobe: usize) -> impl SendFuture>> { async move { if self.fail_nearest { @@ -119,7 +128,7 @@ impl RouterAndNearest for Coarse<'_> { .get(&id) .ok_or_else(|| ANNError::message(format!("unknown list {id}")))? .centroid; - self.route(query, nprobe).await + Ok(self.select(query, nprobe)) } } } @@ -221,12 +230,12 @@ impl RegionUpdate for Update { } } -impl HasIds for Build<'_> { +impl HasIds for Build<'_, T> { type Id = u32; type ListId = u32; } -impl<'p> BuildAccessor for Build<'p> { +impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build<'p, T> { type Update = Update; type Router<'a> = Coarse<'a> @@ -294,13 +303,18 @@ impl<'p> BuildAccessor for Build<'p> { let mut vectors = rowmajor::Owned::from_element(members.len(), self.dim().get(), 0.0); for (&id, row) in members.iter().zip(vectors.rows_mut()) { - let vector = self - .batch - .as_ref() - .and_then(|batch| batch.get(id)) - .or_else(|| self.store.points.get(&id).map(Vec::as_slice)) - .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; - row.copy_from_slice(vector); + 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, @@ -319,7 +333,7 @@ impl<'p> BuildAccessor for Build<'p> { } } - fn finish(self, update: Option<&Self::Update>) -> impl SendFuture> { + fn finish(mut self, update: Option) -> impl SendFuture> { async move { if self.store.fail_finish { return Err(ANNError::message("injected finish failure")); @@ -332,16 +346,16 @@ impl<'p> BuildAccessor for Build<'p> { .next_list .checked_add(created) .ok_or_else(|| ANNError::message("list ID overflow"))?; - if let Some(batch) = self.batch { - for point in batch.points() { - self.store.points.insert(point.id, point.vector.to_vec()); - } + for (id, vector) in self.staged { + self.store + .points + .insert(id, vector.into_iter().map(Into::into).collect()); } let regional_memberships = if let Some(update) = update { - for (id, list) in (self.store.next_list..next_list).zip(update.created.clone()) { + for (id, list) in (self.store.next_list..next_list).zip(update.created) { self.store.lists.insert(id, list); } - Some(update.memberships.clone()) + Some(update.memberships) } else { None }; @@ -376,21 +390,38 @@ impl<'p> BuildAccessor for Build<'p> { } } -struct Strategy; +struct Strategy { + store: Mutex, +} + +impl Strategy { + fn new(store: Store) -> Self { + Self { + store: Mutex::new(store), + } + } + + fn store(&self) -> MutexGuard<'_, Store> { + self.store.try_lock().unwrap() + } +} + +impl HasIds for Strategy { + type Id = u32; + type ListId = u32; +} -impl<'a> BuildStrategy<'a, Store> for Strategy { - type Accessor = Build<'a>; +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, - store: &'a mut Store, - _: &'a DefaultContext, - ) -> ANNResult> { + fn build_accessor(&'a self, _: &'a DefaultContext) -> ANNResult { Ok(Build { - store, + store: self.store.try_lock().map_err(ANNError::new)?, pending: PendingLists::default(), - batch: None, + staged: BTreeMap::new(), sources: BTreeMap::new(), fills: Cell::new(0), filled_lists: Vec::new(), @@ -398,33 +429,38 @@ impl<'a> BuildStrategy<'a, Store> for Strategy { } } -impl<'a, 'p> InsertStrategy<'a, Store, &'p [f32]> for Strategy { - fn prepare( - &self, - accessor: &mut Self::Accessor, - points: &[(u32, &'p [f32])], - ) -> impl SendFuture>>> { - async move { - let mut vectors = - rowmajor::Owned::from_element(points.len(), accessor.dim().get(), 0.0); - let mut ids = Vec::with_capacity(points.len()); - for ((external, vector), out) in points.iter().zip(vectors.rows_mut()) { - if vector.len() != out.len() || !vector.iter().all(|value| value.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 accessor.store.points.contains_key(&id) || ids.contains(&id) { - return Err(ANNError::message("point already exists")); - } - out.copy_from_slice(vector); - ids.push(id); +impl<'p, T: Copy + Into + Send + Sync> SetElements<(u32, &'p [T]), DefaultContext> + for Build<'_, T> +{ + type Id = u32; + + async fn set_elements( + &mut self, + _: &DefaultContext, + points: &[(u32, &'p [T])], + ) -> ANNResult> { + let mut staged = BTreeMap::new(); + let mut ids = Vec::with_capacity(points.len()); + 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 batch = Arc::new(Batch::new(ids, vectors)?); - accessor.batch = Some(Arc::clone(&batch)); - Ok(batch) + staged.insert(id, vector.to_vec()); + ids.push(id); } + self.staged.extend(staged); + Ok(ids) } } @@ -433,22 +469,24 @@ fn matrix(rows: &[[f32; D]]) -> rowmajor::Owned { rowmajor::Owned::try_from_data(data, rows.len(), D).unwrap() } -fn index(store: Store, threshold: usize) -> IVFIndex { +fn index(store: Store, threshold: usize) -> (IVFIndex, Strategy) { index_with_neighbors(store, threshold, 0) } -fn index_with_neighbors(store: Store, threshold: usize, neighbors: usize) -> IVFIndex { - IVFIndex::new( - store, - 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() +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)) } fn seeded() -> Store { @@ -495,9 +533,82 @@ fn require_send(future: F) -> F { future } +#[tokio::test] +async fn staging_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 ids = require_send(build.set_elements(&DefaultContext, &points)) + .await + .unwrap(); + assert_eq!(ids, [10_009, 10_007, 10_008]); + for (&id, point) in ids.iter().zip(&points) { + assert_eq!(build.staged[&id], point.1); + assert!(!build.store.points.contains_key(&id)); + } + let router = build.router(); + for (point, expected) in points.iter().zip([2, 1, 1]) { + assert_eq!( + require_send(router.route(point, 1)).await.unwrap()[0].0, + expected + ); + } + drop(build); + let store = strategy.store(); + assert_eq!(store.points.len(), 5); + assert_eq!(store.finishes, 0); +} + +#[tokio::test] +async fn native_inputs_are_inserted_using_the_returned_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(); + assert!( + require_send(build.set_elements(&DefaultContext, &points)) + .await + .is_err() + ); + assert!(build.staged.is_empty()); + 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 = index( + let (mut index, strategy) = index( Store { dim: 2, lists: [ @@ -523,12 +634,12 @@ async fn appends_publish_actual_memberships_without_filling() { }, 2, ); - let points = [ + 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)) + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points)) .await .unwrap(); @@ -539,51 +650,40 @@ async fn appends_publish_actual_memberships_without_filling() { splits: 0 } ); - assert_eq!(index.provider().lists[&1].members, [10_007, 10_008]); - assert_eq!(index.provider().lists[&2].members, [10_009]); - assert_eq!(index.provider().points[&10_007], [0.0, 1.0]); - assert_eq!( - index - .provider() - .to_internal_id(&DefaultContext, &7) - .unwrap(), - 10_007 - ); - assert_eq!( - index - .provider() - .to_external_id(&DefaultContext, 10_007) - .unwrap(), - 7 - ); - assert_eq!(index.provider().finishes, 1); - assert_eq!(index.provider().fills, 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); let empty: [(u32, &[f32]); 0] = []; assert_eq!( index - .insert_batch(&Strategy, &DefaultContext, &empty) + .insert_batch(&strategy, &DefaultContext, &empty) .await .unwrap(), InsertStats::default() ); - assert_eq!(index.provider().finishes, 1); + assert_eq!(store.finishes, 1); } #[tokio::test] async fn splitting_commits_children_and_ordinary_appends_once() { let mut store = seeded(); store.fail_nearest = true; - let mut index = index(store, 2); - let points = [ + 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)) + let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points)) .await .unwrap(); - let store = index.provider(); + let store = strategy.store(); assert_eq!( stats, @@ -639,12 +739,12 @@ async fn neighboring_lists_exchange_points_and_pending_appends_once() { .members .extend([10_005, 10_006]); store.lists.get_mut(&2).unwrap().centroid = vec![5.0, 0.0]; - let mut index = index_with_neighbors(store, 3, 2); - let points = [(7, &[1.0, 0.0][..]), (8, &[5.2, 0.0][..])]; - let stats = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) + 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 = index.provider(); + let store = strategy.store(); assert_eq!( stats, @@ -704,12 +804,12 @@ async fn neighborhood_selection_filters_parents_and_fills_the_union_once() { } store.next_list = 3 + centers.len() as u32; store.fail_nearest = count == 0; - let mut index = index_with_neighbors(store, 2, count); - let points = [(5, &[1.0, 1.0][..]), (6, &[11.0, 1.0][..])]; - let stats = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) + 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 = index.provider(); + let store = strategy.store(); let neighbors = match (count, centers.len()) { (0 | 1, _) | (_, 0) => vec![], (2, _) => two_candidates.clone(), @@ -784,14 +884,15 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { .into(); let region = Region::new(&sources); let mut update = Update::new(&sources, 1); - clustering::split( + clustering::SplitAndAssign::new( + clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), + clustering::RegionalNearest, + ) + .run( ®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::from([2, 3, 4]), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(1).unwrap(), - &mut StdRng::seed_from_u64(7), ) .unwrap(); @@ -818,13 +919,12 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { #[tokio::test] async fn region_view_borrows_inputs_while_updates_are_recorded() { - let mut store = seeded(); + let strategy = Strategy::new(seeded()); let update = { - let mut build = Strategy - .build_accessor(&mut store, &DefaultContext) - .unwrap(); + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; - let batch = Strategy.prepare(&mut build, &points).await.unwrap(); + let ids = build.set_elements(&DefaultContext, &points).await.unwrap(); + assert_eq!(ids, [10_005, 10_006]); assert_eq!( build .append_to(1, [10_006, 10_005].into_iter()) @@ -846,20 +946,19 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { source.vectors.as_slice().as_ptr() ); assert_eq!(vectors.row(0), [0.0, 0.0]); - assert_eq!(vectors.row(2), batch.get(10_006).unwrap()); - assert_eq!(vectors.row(3), batch.get(10_005).unwrap()); - clustering::split( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::new(), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - &mut StdRng::seed_from_u64(7), + assert_eq!(vectors.row(2), points[1].1); + assert_eq!(vectors.row(3), points[0].1); + clustering::SplitAndAssign::new( + clustering::TwoMeans::new( + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), + ), + clustering::RegionalNearest, ) + .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) .unwrap(); assert_eq!(view.members, [10_000, 10_001, 10_006, 10_005]); - assert_eq!(vectors.row(2), batch.get(10_006).unwrap()); + assert_eq!(vectors.row(2), points[1].1); assert!(update.set_members(1, &[]).is_err()); update }; @@ -880,6 +979,7 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { .into_changes() .any(|(id, change)| id == 1 && matches!(change, Membership::Retire)) ); + 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); @@ -902,9 +1002,9 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { store.fail_set_members = fail_set_members; store.fail_retire = fail_retire; store.fail_finish = fail_finish; - let mut index = index_with_neighbors(store, 2, 2); - let points = [(5, &[1.0, 1.0][..])]; - let error = require_send(index.insert_batch(&Strategy, &DefaultContext, &points)) + 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 { @@ -921,34 +1021,18 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { "injected finish failure" }; assert!(error.to_string().contains(expected), "{error}"); - assert_eq!(index.provider().finishes, 0); - assert_eq!(index.provider().lists[&1].members, [10_000, 10_001]); - assert_eq!(index.provider().lists[&2].members, [10_002, 10_003]); - assert_eq!(index.provider().lists.len(), 3); - assert!( - index - .provider() - .to_internal_id(&DefaultContext, &5) - .is_err() - ); + 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 invalid_config_and_uninitialized_insertion_fail() { - assert!( - IVFIndex::new( - Store::default(), - Config { - split_threshold: NonZeroUsize::new(1).unwrap(), - reassign_count: ReassignCount::None, - split_iterations: NonZeroUsize::new(5).unwrap(), - seed: 7, - }, - ) - .is_err() - ); - let mut index = index( +async fn uninitialized_insertion_fails() { + let (mut index, strategy) = index( Store { dim: 2, next_list: 1, @@ -956,14 +1040,14 @@ async fn invalid_config_and_uninitialized_insertion_fail() { }, 2, ); - let points = [(7, &[0.0, 1.0][..])]; + let points: [(u32, &[f32]); 1] = [(7, &[0.0, 1.0][..])]; assert!( index - .insert_batch(&Strategy, &DefaultContext, &points) + .insert_batch(&strategy, &DefaultContext, &points) .await .is_err() ); - assert_eq!(index.provider().finishes, 0); + assert_eq!(strategy.store().finishes, 0); } #[tokio::test] @@ -973,11 +1057,10 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() store.points.insert(10_002, vec![5.0, 0.0]); store.points.insert(10_003, vec![10.1, 0.0]); store.lists.get_mut(&2).unwrap().centroid = vec![5.0, 0.0]; - let mut build = Strategy - .build_accessor(&mut store, &DefaultContext) - .unwrap(); + let strategy = Strategy::new(store); + let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); let points = [(5, &[5.0, 0.0][..])]; - Strategy.prepare(&mut build, &points).await.unwrap(); + build.set_elements(&DefaultContext, &points).await.unwrap(); assert_eq!( build .append_to(1, [10_005].into_iter()) @@ -990,51 +1073,37 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() let update = { let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); assert!( - clustering::split( + clustering::SplitAndAssign::new( + clustering::TwoMeans::new( + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), + ), + clustering::RegionalNearest, + ) + .run( ®ion, &mut update, &BTreeSet::from([1, 99]), - &BTreeSet::new(), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - &mut StdRng::seed_from_u64(7) + &BTreeSet::new() ) .is_err() ); assert!( - clustering::split( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::from([99]), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - &mut StdRng::seed_from_u64(7) + clustering::SplitAndAssign::new( + clustering::TwoMeans::new( + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), + ), + clustering::RegionalNearest, ) - .is_err() - ); - assert!( - clustering::split( + .run( ®ion, &mut update, &BTreeSet::from([1]), - &BTreeSet::from([1]), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - &mut StdRng::seed_from_u64(7) + &BTreeSet::from([99]) ) .is_err() ); - clustering::split( - ®ion, - &mut update, - &BTreeSet::new(), - &BTreeSet::new(), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - &mut StdRng::seed_from_u64(7), - ) - .unwrap(); assert!(update.created.is_empty()); assert_eq!( std::mem::take(&mut update.memberships) @@ -1042,21 +1111,18 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() .count(), 0 ); - clustering::split( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::new(), - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - &mut StdRng::seed_from_u64(7), + clustering::SplitAndAssign::new( + clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()), + clustering::RegionalNearest, ) + .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) .unwrap(); assert_eq!(update.created.len(), 2); update }; - require_send(build.finish(Some(&update))).await.unwrap(); + require_send(build.finish(Some(update))).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]); @@ -1097,14 +1163,18 @@ fn assignment_can_cross_parent_boundaries() { .into(); let region = Region::new(&sources); let mut update = Update::new(&sources, 1); - clustering::split( + clustering::SplitAndAssign::new( + clustering::TwoMeans::new( + NonZeroUsize::new(10).unwrap(), + NonZeroUsize::new(1).unwrap(), + ), + clustering::RegionalNearest, + ) + .run( ®ion, &mut update, &BTreeSet::from([1, 2]), &BTreeSet::new(), - NonZeroUsize::new(10).unwrap(), - NonZeroUsize::new(1).unwrap(), - &mut StdRng::seed_from_u64(7), ) .unwrap(); assert!( @@ -1128,15 +1198,11 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { .into(); let region = Region::new(&sources); let mut update = Update::new(&sources, 1); - clustering::split( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::new(), - NonZeroUsize::new(1).unwrap(), - NonZeroUsize::new(1).unwrap(), - &mut StdRng::seed_from_u64(7), + clustering::SplitAndAssign::new( + clustering::TwoMeans::new(NonZeroUsize::new(1).unwrap(), NonZeroUsize::new(1).unwrap()), + clustering::RegionalNearest, ) + .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) .unwrap(); assert_eq!(update.created.len(), 2); assert_eq!(update.created[0].centroid, [value]); @@ -1155,7 +1221,6 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { &mut self, points: rowmajor::Ref<'_, f32>, mut dst: rowmajor::Mut<'_, f32>, - _: &mut StdRng, ) -> ANNResult<()> { assert_eq!(points.nrows(), dst.nrows()); dst.as_mut_slice().copy_from_slice(points.as_slice()); @@ -1283,13 +1348,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut full = Update::new(&sources, 1); let mut splitter = clustering::SplitAndAssign::new(RowsAsCenters, clustering::RegionalNearest); splitter - .run( - ®ion, - &mut full, - &parents, - &neighbors, - &mut StdRng::seed_from_u64(7), - ) + .run(®ion, &mut full, &parents, &neighbors) .unwrap(); let mut restricted = Update::new(&sources, 1); clustering::SplitAndAssign::new( @@ -1298,13 +1357,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { owners: [(3, vec![1]), (4, vec![1, 2])].into(), }, ) - .run( - ®ion, - &mut restricted, - &parents, - &neighbors, - &mut StdRng::seed_from_u64(7), - ) + .run(®ion, &mut restricted, &parents, &neighbors) .unwrap(); for update in [&full, &restricted] { @@ -1324,13 +1377,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { assert_eq!(restricted.created[2].members, [2, 7]); let mut repeated = Update::new(&sources, 1); splitter - .run( - ®ion, - &mut repeated, - &parents, - &neighbors, - &mut StdRng::seed_from_u64(7), - ) + .run(®ion, &mut repeated, &parents, &neighbors) .unwrap(); for (first, repeated) in full.created.iter().zip(&repeated.created) { assert_eq!(first.centroid, repeated.centroid); @@ -1343,7 +1390,6 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { &mut smaller, &BTreeSet::from([1]), &BTreeSet::new(), - &mut StdRng::seed_from_u64(7), ) .unwrap(); assert_eq!(smaller.created.len(), 2); @@ -1368,50 +1414,29 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { fn numerical_strategy_failures_record_no_updates() { enum InvalidFit { Error, - TooFew, - WrongDimension, - NonFinite, - Incomplete, CountOverflow, } impl clustering::Clusterer for InvalidFit { - fn fit( - &mut self, - _: rowmajor::Ref<'_, f32>, - mut dst: rowmajor::Mut<'_, f32>, - _: &mut StdRng, - ) -> ANNResult<()> { + fn fit(&mut self, _: rowmajor::Ref<'_, f32>, _: rowmajor::Mut<'_, f32>) -> ANNResult<()> { match self { Self::Error => Err(ANNError::message("injected fitting failure")), - Self::NonFinite => { - dst.as_mut_slice().fill(f32::INFINITY); - Ok(()) - } - Self::Incomplete => { - dst.row_mut(0).fill(0.0); - Ok(()) + Self::CountOverflow => { + unreachable!("centroid count overflow should be rejected before fitting") } - _ => unreachable!("invalid centroid layout should be rejected before fitting"), } } fn num_centers(&self) -> NonZeroUsize { NonZeroUsize::new(match self { - Self::TooFew => 1, Self::CountOverflow => usize::MAX, - _ => 2, + Self::Error => 2, }) .unwrap() } fn dim(&self) -> NonZeroUsize { - NonZeroUsize::new(if matches!(self, Self::WrongDimension) { - 2 - } else { - 1 - }) - .unwrap() + NonZeroUsize::new(1).unwrap() } } @@ -1452,10 +1477,6 @@ fn numerical_strategy_failures_record_no_updates() { let neighbors = BTreeSet::new(); for (clusterer, expected) in [ (InvalidFit::Error, "injected fitting failure"), - (InvalidFit::TooFew, "at least two centroids"), - (InvalidFit::WrongDimension, "canonical dimension"), - (InvalidFit::NonFinite, "finite value"), - (InvalidFit::Incomplete, "every centroid coordinate"), (InvalidFit::CountOverflow, "centroid count overflow"), ] { let mut update = Update::new(&sources, 1); @@ -1465,13 +1486,7 @@ fn numerical_strategy_failures_record_no_updates() { parents.clone() }; let error = clustering::SplitAndAssign::new(clusterer, clustering::RegionalNearest) - .run( - ®ion, - &mut update, - &selected, - &neighbors, - &mut StdRng::seed_from_u64(7), - ) + .run(®ion, &mut update, &selected, &neighbors) .unwrap_err(); assert!(error.to_string().contains(expected), "{error}"); assert_unwritten(update); @@ -1481,13 +1496,7 @@ fn numerical_strategy_failures_record_no_updates() { clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), InvalidAssignment, ) - .run( - ®ion, - &mut update, - &parents, - &neighbors, - &mut StdRng::seed_from_u64(7), - ) + .run(®ion, &mut update, &parents, &neighbors) .unwrap_err(); assert!( error.to_string().contains("injected assignment failure"), @@ -1499,34 +1508,24 @@ fn numerical_strategy_failures_record_no_updates() { #[test] fn two_means_matches_quantization_lloyds() { use clustering::{Clusterer, TwoMeans}; - use rand::Rng; let points = matrix(&[[0.0, 2.0], [2.0, 4.0], [6.0, 6.0], [10.0, 12.0]]); - for seed in [0, 7, 42] { - for iterations in [1, 5] { - let mut expected = rowmajor::Owned::from_element(2, 2, 0.0); - let mut reference_rng = StdRng::seed_from_u64(seed); - let first = reference_rng.random_range(0..points.nrows()); - let second = (first + reference_rng.random_range(1..points.nrows())) % points.nrows(); - expected.row_mut(0).copy_from_slice(points.row(first)); - expected.row_mut(1).copy_from_slice(points.row(second)); - let _ = diskann_quantization::algorithms::kmeans::lloyds::lloyds( - points.as_view(), - expected.as_view_mut(), - iterations, - ); + for iterations in [1, 5] { + let mut expected = points.as_view().subview(0..2).unwrap().to_rowmajor_owned(); + let _ = diskann_quantization::algorithms::kmeans::lloyds::lloyds( + points.as_view(), + expected.as_view_mut(), + iterations, + ); - let mut actual = rowmajor::Owned::from_element(2, 2, 0.0); - let mut rng = StdRng::seed_from_u64(seed); - TwoMeans::new( - NonZeroUsize::new(iterations).unwrap(), - NonZeroUsize::new(2).unwrap(), - ) - .fit(points.as_view(), actual.as_view_mut(), &mut rng) - .unwrap(); - assert_eq!(actual.as_slice(), expected.as_slice()); - assert_eq!(rng.random::(), reference_rng.random::()); - } + let mut actual = rowmajor::Owned::from_element(2, 2, 0.0); + TwoMeans::new( + NonZeroUsize::new(iterations).unwrap(), + NonZeroUsize::new(2).unwrap(), + ) + .fit(points.as_view(), actual.as_view_mut()) + .unwrap(); + assert_eq!(actual.as_slice(), expected.as_slice()); } } @@ -1540,13 +1539,8 @@ fn two_means_fits_only_the_destination_window() { NonZeroUsize::new(10).unwrap(), NonZeroUsize::new(2).unwrap(), ); - let mut rng = StdRng::seed_from_u64(7); clusterer - .fit( - points.as_view(), - scratch.subview_mut(1..3).unwrap(), - &mut rng, - ) + .fit(points.as_view(), scratch.subview_mut(1..3).unwrap()) .unwrap(); assert_eq!(scratch.row(0), [-1.0, -1.0]); assert_eq!(scratch.row(3), [-1.0, -1.0]); @@ -1560,23 +1554,18 @@ fn two_means_fits_only_the_destination_window() { (matrix(&[[0.0], [1.0]]), 2, 2), (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 1, 2), (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 2, 1), - (matrix(&[[f32::NAN, 2.0], [1.0, 3.0]]), 2, 2), ] { let mut dst = rowmajor::Owned::from_element(rows, dim, -1.0); - assert!( - clusterer - .fit(points.as_view(), dst.as_view_mut(), &mut rng) - .is_err() - ); + assert!(clusterer.fit(points.as_view(), dst.as_view_mut()).is_err()); assert!(dst.as_slice().iter().all(|&value| value == -1.0)); } } #[test] -fn regional_assignment_rejects_invalid_inputs_before_writing() { +fn regional_assignment_reports_missing_sources_before_writing() { use clustering::{Assigner, Destination, RegionalNearest, Type}; - let mut sources = BTreeMap::from([ + let sources = BTreeMap::from([ ( 1, Source { @@ -1598,51 +1587,13 @@ fn regional_assignment_rejects_invalid_inputs_before_writing() { list: Type::New, centroid: &[0.0], }]; - for invalid in [ - vec![], - vec![Destination { - list: Type::New, - centroid: &[], - }], - vec![Destination { - list: Type::New, - centroid: &[f32::NAN], - }], - vec![ - Destination { - list: Type::New, - centroid: &[0.0], - }, - Destination { - list: Type::New, - centroid: &[1.0, 2.0], - }, - ], - ] { + for ids in [vec![99], vec![1, 99], vec![1, 2, 99]] { let mut update = Update::new(&sources, 1); - assert!( - RegionalNearest - .assign(&Region::new(&sources), &mut update, &[1, 2], &invalid) - .is_err() - ); - assert!(update.created.is_empty()); - assert_eq!(update.memberships.into_changes().count(), 0); - } - for (ids, members, vectors) in [ - (vec![1, 99], vec![2], matrix(&[[2.0]])), - (vec![1, 2], vec![], matrix(&[[2.0]])), - (vec![1, 2], vec![2], matrix(&[[0.0, 1.0]])), - (vec![1, 2], vec![2], matrix(&[[f32::INFINITY]])), - ] { - let mut update = Update::new(&sources, 1); - sources.get_mut(&2).unwrap().members = members; - sources.get_mut(&2).unwrap().vectors = vectors; let region = Region::new(&sources); - assert!( - RegionalNearest - .assign(®ion, &mut update, &ids, &destinations) - .is_err() - ); + let error = RegionalNearest + .assign(®ion, &mut update, &ids, &destinations) + .unwrap_err(); + assert!(error.to_string().contains("source is missing"), "{error}"); assert!(update.created.is_empty()); assert_eq!(update.memberships.into_changes().count(), 0); } diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 399c1f62e0..239077e0fb 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -84,7 +84,7 @@ pub trait BuildAccessor: Send + Sized + HasIds { /// Finalize the operation, publishing accessor-side appends and regional writes together. /// /// `None` means there are no regional re-writes. - fn finish(self, update: Option<&Self::Update>) -> impl SendFuture>; + fn finish(self, update: Option) -> impl SendFuture>; } pub trait BuildStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { From deffb3c85e89a96586d04c1dc798a940d864cb77 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Tue, 20 Oct 2026 09:42:37 -0400 Subject: [PATCH 25/33] update gets a lifetime --- diskann/src/ivf/index.rs | 5 +- diskann/src/ivf/tests.rs | 236 +++++++++++++++++++++++++------------- diskann/src/ivf/traits.rs | 14 ++- 3 files changed, 168 insertions(+), 87 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index e515104856..c8f8471648 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -128,7 +128,6 @@ impl IVFIndex { } let splits = parents.len(); - let mut update_final = None; if !parents.is_empty() { let mut neighbors = BTreeSet::new(); @@ -160,11 +159,9 @@ impl IVFIndex { &parents, &neighbors, )?; - - update_final = Some(update) }; - accessor.finish(update_final).await?; + accessor.finish().await?; Ok(InsertStats { inserted: points.len(), diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index cb87dd729a..f7079bbb56 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -67,6 +67,7 @@ impl Store { struct Build<'p, T> { store: MutexGuard<'p, Store>, pending: PendingLists, + created: Vec, staged: BTreeMap>, sources: BTreeMap, fills: Cell, @@ -139,24 +140,28 @@ struct Source { vectors: rowmajor::Owned, } -#[derive(Debug)] -struct Update { +struct Update<'a> { dim: usize, - lists: BTreeSet, - memberships: PendingLists, - created: Vec, + sources: &'a BTreeMap, + memberships: &'a mut PendingLists, + created: &'a mut Vec, fail_create: bool, fail_set_members: bool, fail_retire: bool, } -impl Update { - fn new(sources: &BTreeMap, dim: usize) -> Self { +impl<'a> Update<'a> { + fn new( + sources: &'a BTreeMap, + dim: usize, + memberships: &'a mut PendingLists, + created: &'a mut Vec, + ) -> Self { Self { dim, - lists: sources.keys().copied().collect(), - memberships: PendingLists::default(), - created: Vec::new(), + sources, + memberships, + created, fail_create: false, fail_set_members: false, fail_retire: false, @@ -164,7 +169,7 @@ impl Update { } fn check_list(&self, list: u32) -> ANNResult<()> { - if self.lists.contains(&list) { + if self.sources.contains_key(&list) { Ok(()) } else { Err(ANNError::message(format!("unknown list {list}"))) @@ -198,7 +203,7 @@ impl RegionView for Region<'_> { } } -impl RegionUpdate for Update { +impl RegionUpdate 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")); @@ -236,7 +241,10 @@ impl HasIds for Build<'_, T> { } impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build<'p, T> { - type Update = Update; + type Update<'a> + = Update<'a> + where + Self: 'a; type Router<'a> = Coarse<'a> where @@ -282,7 +290,7 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build fn fill( &mut self, lists: Itr, - ) -> impl SendFuture, Self::Update)>> + ) -> impl SendFuture, Self::Update<'_>)>> where Itr: Iterator + Clone + Send + Sync, { @@ -325,7 +333,12 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build }, ); } - let mut update = Update::new(&self.sources, self.dim().get()); + 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; @@ -333,14 +346,13 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build } } - fn finish(mut self, update: Option) -> impl SendFuture> { + fn finish(mut self) -> impl SendFuture> { async move { if self.store.fail_finish { return Err(ANNError::message("injected finish failure")); } - let created = update.as_ref().map_or(0, |update| update.created.len()); - let created = - u32::try_from(created).map_err(|_| ANNError::message("list ID overflow"))?; + let created = u32::try_from(self.created.len()) + .map_err(|_| ANNError::message("list ID overflow"))?; let next_list = self .store .next_list @@ -351,19 +363,10 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build .points .insert(id, vector.into_iter().map(Into::into).collect()); } - let regional_memberships = if let Some(update) = update { - for (id, list) in (self.store.next_list..next_list).zip(update.created) { - self.store.lists.insert(id, list); - } - Some(update.memberships) - } else { - None - }; - for (id, change) in self.pending.into_changes().chain( - regional_memberships - .into_iter() - .flat_map(PendingLists::into_changes), - ) { + 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 @@ -421,6 +424,7 @@ impl<'a, T: Copy + Into + Send + Sync> BuildStrategy<'a, (u32, &[T]), Defau 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), @@ -883,7 +887,9 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { ] .into(); let region = Region::new(&sources); - let mut update = Update::new(&sources, 1); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); clustering::SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), clustering::RegionalNearest, @@ -906,7 +912,7 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { .collect(); child_members.sort_unstable(); assert_eq!(child_members, [0, 1]); - let mut changes: BTreeMap<_, _> = update.memberships.into_changes().collect(); + 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![])] { let Some(Membership::Replace(members)) = changes.remove(&id) else { @@ -920,7 +926,7 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { #[tokio::test] async fn region_view_borrows_inputs_while_updates_are_recorded() { let strategy = Strategy::new(seeded()); - let update = { + { let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; let ids = build.set_elements(&DefaultContext, &points).await.unwrap(); @@ -933,7 +939,7 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { .get(), 4 ); - let update = { + { let (region, mut update) = require_send(build.fill([1].into_iter())).await.unwrap(); let view = region.get_view(1).unwrap(); let vectors = region.get_vectors(1).unwrap(); @@ -960,31 +966,81 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { 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()); - update - }; + } assert_eq!(build.dim().get(), 2); - update - }; - assert_eq!(update.created.len(), 2); - let mut members: Vec<_> = update - .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!( - update - .memberships - .into_changes() - .any(|(id, change)| id == 1 && matches!(change, Membership::Retire)) - ); + 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.set_elements(&DefaultContext, &points).await.unwrap(); + for (list, id) in [(1, 10_005), (2, 10_006), (3, 10_007)] { + build.append_to(list, [id].into_iter()).await.unwrap(); + } + + { + let (region, mut update) = require_send(build.fill([1, 2].into_iter())).await.unwrap(); + let parent = region.get_view(1).unwrap(); + let neighbor = region.get_view(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 [ @@ -1070,7 +1126,7 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() 3 ); - let update = { + { let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); assert!( clustering::SplitAndAssign::new( @@ -1105,11 +1161,9 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() .is_err() ); assert!(update.created.is_empty()); - assert_eq!( - std::mem::take(&mut update.memberships) - .into_changes() - .count(), - 0 + let changes: Vec<_> = update.memberships.clone().into_changes().collect(); + assert!( + matches!(changes.as_slice(), [(1, Membership::Append(members))] if members == &[10_005]) ); clustering::SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()), @@ -1118,9 +1172,8 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) .unwrap(); assert_eq!(update.created.len(), 2); - update - }; - require_send(build.finish(Some(update))).await.unwrap(); + } + require_send(build.finish()).await.unwrap(); let store = strategy.store(); assert_eq!(store.finishes, 1); @@ -1162,7 +1215,9 @@ fn assignment_can_cross_parent_boundaries() { ] .into(); let region = Region::new(&sources); - let mut update = Update::new(&sources, 1); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); clustering::SplitAndAssign::new( clustering::TwoMeans::new( NonZeroUsize::new(10).unwrap(), @@ -1197,7 +1252,9 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { )] .into(); let region = Region::new(&sources); - let mut update = Update::new(&sources, 1); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); clustering::SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(1).unwrap(), NonZeroUsize::new(1).unwrap()), clustering::RegionalNearest, @@ -1345,12 +1402,21 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let region = Region::new(&sources); let parents = BTreeSet::from([1, 2]); let neighbors = BTreeSet::from([3, 4]); - let mut full = Update::new(&sources, 1); + let mut full_memberships = PendingLists::default(); + let mut full_created = Vec::new(); + let mut full = Update::new(&sources, 1, &mut full_memberships, &mut full_created); let mut splitter = clustering::SplitAndAssign::new(RowsAsCenters, clustering::RegionalNearest); splitter .run(®ion, &mut full, &parents, &neighbors) .unwrap(); - let mut restricted = Update::new(&sources, 1); + let mut restricted_memberships = PendingLists::default(); + let mut restricted_created = Vec::new(); + let mut restricted = Update::new( + &sources, + 1, + &mut restricted_memberships, + &mut restricted_created, + ); clustering::SplitAndAssign::new( RowsAsCenters, ParentLocal { @@ -1375,15 +1441,24 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { } assert_eq!(full.created[2].members, [2, 5, 7]); assert_eq!(restricted.created[2].members, [2, 7]); - let mut repeated = Update::new(&sources, 1); + let mut repeated_memberships = PendingLists::default(); + let mut repeated_created = Vec::new(); + let mut repeated = Update::new( + &sources, + 1, + &mut repeated_memberships, + &mut repeated_created, + ); splitter .run(®ion, &mut repeated, &parents, &neighbors) .unwrap(); - for (first, repeated) in full.created.iter().zip(&repeated.created) { + 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 = Update::new(&sources, 1); + let mut smaller_memberships = PendingLists::default(); + let mut smaller_created = Vec::new(); + let mut smaller = Update::new(&sources, 1, &mut smaller_memberships, &mut smaller_created); splitter .run( ®ion, @@ -1396,7 +1471,8 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { assert_eq!(smaller.created[0].members, [0]); assert_eq!(smaller.created[1].members, [1]); for (update, expected) in [(full, vec![6]), (restricted, vec![5, 6])] { - let mut changes: BTreeMap<_, _> = update.memberships.into_changes().collect(); + let mut changes: BTreeMap<_, _> = + std::mem::take(update.memberships).into_changes().collect(); for parent in [1, 2] { assert!(matches!(changes.remove(&parent), Some(Membership::Retire))); } @@ -1458,9 +1534,9 @@ fn numerical_strategy_failures_record_no_updates() { } } - fn assert_unwritten(update: Update) { + fn assert_unwritten(update: Update<'_>) { assert!(update.created.is_empty()); - assert_eq!(update.memberships.into_changes().count(), 0); + assert_eq!(std::mem::take(update.memberships).into_changes().count(), 0); } let sources = [( @@ -1479,7 +1555,9 @@ fn numerical_strategy_failures_record_no_updates() { (InvalidFit::Error, "injected fitting failure"), (InvalidFit::CountOverflow, "centroid count overflow"), ] { - let mut update = Update::new(&sources, 1); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); let selected = if matches!(clusterer, InvalidFit::CountOverflow) { BTreeSet::from([1, 2]) } else { @@ -1491,7 +1569,9 @@ fn numerical_strategy_failures_record_no_updates() { assert!(error.to_string().contains(expected), "{error}"); assert_unwritten(update); } - let mut update = Update::new(&sources, 1); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); let error = clustering::SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), InvalidAssignment, @@ -1588,13 +1668,15 @@ fn regional_assignment_reports_missing_sources_before_writing() { centroid: &[0.0], }]; for ids in [vec![99], vec![1, 99], vec![1, 2, 99]] { - let mut update = Update::new(&sources, 1); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); let region = Region::new(&sources); let error = RegionalNearest .assign(®ion, &mut update, &ids, &destinations) .unwrap_err(); assert!(error.to_string().contains("source is missing"), "{error}"); assert!(update.created.is_empty()); - assert_eq!(update.memberships.into_changes().count(), 0); + assert_eq!(std::mem::take(update.memberships).into_changes().count(), 0); } } diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 239077e0fb..52630681b3 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -56,7 +56,9 @@ pub trait BuildAccessor: Send + Sized + HasIds { where Self: 'a; - type Update: RegionUpdate; + type Update<'a>: RegionUpdate + where + Self: 'a; /// Positive dimension of all canonical points and centroids. fn dim(&self) -> NonZeroUsize; @@ -72,19 +74,19 @@ pub trait BuildAccessor: Send + Sized + HasIds { where Itr: ExactSizeIterator + Clone + Send + Sync; - /// Make every requested list's members, data and centroids available + /// Make requested list's members, data and centroids available /// for split and assignment. 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)>> + ) -> impl SendFuture, Self::Update<'_>)>> where Itr: Iterator + Clone + Send + Sync; /// Finalize the operation, publishing accessor-side appends and regional writes together. - /// - /// `None` means there are no regional re-writes. - fn finish(self, update: Option) -> impl SendFuture>; + fn finish(self) -> impl SendFuture>; } pub trait BuildStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { From 4e9e3a3a8f35cbacd5990aedfcf37209679ddd07 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Tue, 20 Oct 2026 11:29:41 -0400 Subject: [PATCH 26/33] route and nearest gets a callback --- diskann/src/ivf/index.rs | 36 +++++---- diskann/src/ivf/tests.rs | 160 +++++++++++++++++++++++++++++++------- diskann/src/ivf/traits.rs | 39 +++++++--- 3 files changed, 181 insertions(+), 54 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index c8f8471648..efb11ea78a 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -23,6 +23,7 @@ use crate::{ traits::BuildStrategy, }, provider::ExecutionContext, + utils::ONE, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -106,13 +107,15 @@ impl IVFIndex { let router = accessor.router(); for (id, point) in ids.iter().zip(points.iter()) { - let selected = router.route(point, 1).await?; - - let (list, _) = selected.first().ok_or_else(|| { - ANNError::message("IVF routing returned no live centroid") - })?; - - routed.entry(*list).or_default().push(*id); + let count = router + .route(point, ONE, |list, _| { + routed.entry(*list).or_default().push(*id); + }) + .await?; + + if count != 1 { + return Err(ANNError::message("IVF routing returned no live centroid")); + } } } @@ -136,15 +139,16 @@ impl IVFIndex { let router = accessor.router(); for &pid in parents.iter() { - let nids = router - .nearest(pid, n.get()) - .await? - .into_iter() - .filter(|(id, _)| !parents.contains(id)); - - for (id, _) in nids { - neighbors.insert(id); - } + router + .nearest( + pid, + n, + |id| !parents.contains(id), + |id, _| { + neighbors.insert(*id); + }, + ) + .await?; } } diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index f7079bbb56..65de9d1c1d 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -23,7 +23,7 @@ use super::{ region::{RegionUpdate, RegionView, View}, traits::{BuildAccessor, BuildStrategy, HasIds, Router, RouterAndNearest, SetElements}, }; -use crate::{ANNError, ANNResult, provider::DefaultContext}; +use crate::{ANNError, ANNResult, provider::DefaultContext, utils::ONE}; #[derive(Debug, Clone)] struct StoredList { @@ -80,10 +80,16 @@ struct Coarse<'a> { } impl Coarse<'_> { - fn select>(&self, query: &[T], nprobe: usize) -> Vec<(u32, f32)> { + 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)| { let distance: f32 = query .iter() @@ -97,8 +103,11 @@ impl Coarse<'_> { }) .collect(); selected.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))); - selected.truncate(nprobe); - selected + selected.truncate(nprobe.get()); + for (id, distance) in &selected { + emit(id, *distance); + } + selected.len() } } @@ -109,17 +118,31 @@ impl + Sync> Router<(u32, &[T])> for Coarse<'_> { self.stored.len() } - fn route( + fn route( &self, query: &(u32, &[T]), - nprobe: usize, - ) -> impl SendFuture>> { - async move { Ok(self.select(query.1, nprobe)) } + nprobe: NonZeroUsize, + emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::Id, f32) + Send, + { + async move { Ok(self.select(query.1, nprobe, |_| true, emit)) } } } impl + Sync> RouterAndNearest<(u32, &[T])> for Coarse<'_> { - fn nearest(&self, id: Self::Id, nprobe: usize) -> impl SendFuture>> { + fn nearest( + &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")); @@ -129,7 +152,7 @@ impl + Sync> RouterAndNearest<(u32, &[T])> for Coarse<'_> { .get(&id) .ok_or_else(|| ANNError::message(format!("unknown list {id}")))? .centroid; - Ok(self.select(query, nprobe)) + Ok(self.select(query, nprobe, pred, emit)) } } } @@ -537,6 +560,84 @@ fn require_send(future: F) -> F { future } +#[tokio::test] +async fn routing_callbacks_emit_ranked_results_and_counts() { + 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_nearest: false, + }; + for limit in [1, 2, usize::MAX] { + let mut emitted = Vec::new(); + let count = require_send(router.route( + &(9, &[6u8, 0][..]), + NonZeroUsize::new(limit).unwrap(), + |id, distance| emitted.push((*id, distance)), + )) + .await + .unwrap(); + let expected = &expected[..limit.min(store.lists.len())]; + assert_eq!(count, expected.len()); + assert_eq!(emitted, expected); + } + } +} + +#[tokio::test] +async fn nearest_callbacks_filter_before_limiting_and_count_emitted_results() { + let store = seeded(); + let router = Coarse { + stored: &store.lists, + 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(RouterAndNearest::<(u32, &[f32])>::nearest( + &router, + 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 router = Coarse { + stored: &store.lists, + fail_nearest, + }; + let mut emitted = Vec::new(); + let error = require_send(RouterAndNearest::<(u32, &[f32])>::nearest( + &router, + 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 staging_preserves_native_inputs_and_id_order() { let strategy = Strategy::new(seeded()); @@ -552,10 +653,12 @@ async fn staging_preserves_native_inputs_and_id_order() { } let router = build.router(); for (point, expected) in points.iter().zip([2, 1, 1]) { - assert_eq!( - require_send(router.route(point, 1)).await.unwrap()[0].0, - expected - ); + let mut selected = None; + let count = require_send(router.route(point, ONE, |id, _| selected = Some(*id))) + .await + .unwrap(); + assert_eq!(count, 1); + assert_eq!(selected, Some(expected)); } drop(build); let store = strategy.store(); @@ -759,7 +862,7 @@ async fn neighboring_lists_exchange_points_and_pending_appends_once() { ); assert_eq!(store.fills, 1); assert_eq!(store.finishes, 1); - assert_eq!(store.filled_lists, [1, 2]); + 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]); @@ -784,10 +887,10 @@ async fn neighboring_lists_exchange_points_and_pending_appends_once() { #[tokio::test] async fn neighborhood_selection_filters_parents_and_fills_the_union_once() { - for (centers, two_candidates, three_candidates) in [ - (vec![], vec![], vec![]), - (vec![100.0, 200.0, 300.0], vec![], vec![3]), - (vec![0.0, 12.0, 200.0], vec![3, 4], vec![3, 4]), + 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(); @@ -815,10 +918,9 @@ async fn neighborhood_selection_filters_parents_and_fills_the_union_once() { .unwrap(); let store = strategy.store(); let neighbors = match (count, centers.len()) { - (0 | 1, _) | (_, 0) => vec![], - (2, _) => two_candidates.clone(), - (3, _) => three_candidates.clone(), - (4, _) => vec![3, 4], + (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(); @@ -1097,11 +1199,15 @@ async fn uninitialized_insertion_fails() { 2, ); let points: [(u32, &[f32]); 1] = [(7, &[0.0, 1.0][..])]; + let error = index + .insert_batch(&strategy, &DefaultContext, &points) + .await + .unwrap_err(); assert!( - index - .insert_batch(&strategy, &DefaultContext, &points) - .await - .is_err() + error + .to_string() + .contains("IVF routing returned no live centroid"), + "{error}" ); assert_eq!(strategy.store().finishes, 0); } diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 52630681b3..55ff5f1427 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -19,7 +19,7 @@ pub trait HasIds { /// Centroid selection within one build operation. pub trait Router: Send { - type Id: VectorId; + type Id; fn len(&self) -> usize; @@ -28,23 +28,34 @@ pub trait Router: Send { } /// Select `min(nprobe, len())` distinct live lists, nearest first. - fn route(&self, query: &T, nprobe: usize) -> impl SendFuture>>; + fn route( + &self, + query: &T, + nprobe: NonZeroUsize, + emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::Id, f32) + Send; } pub trait RouterAndNearest: Router { - /// Select `min(nprobe, len())` distinct live lists nearest to `id`'s centroid. + /// Select `nprobe` distinct live lists nearest to `id`'s centroid, that pass the filter. /// - /// Results are nearest first and may include `id` itself. Callers exclude - /// lists being replaced when selecting retained neighbors. + /// Results are nearest first and may include `id` itself. /// /// # Errors /// /// Fails if `id` is not live or the provider cannot complete the lookup. - fn nearest( + fn nearest( &self, id: Self::Id, - nprobe: usize, - ) -> impl SendFuture>>; + 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 { @@ -75,7 +86,7 @@ pub trait BuildAccessor: Send + Sized + HasIds { Itr: ExactSizeIterator + Clone + Send + Sync; /// Make requested list's members, data and centroids available - /// for split and assignment. Register updates to the region. + /// for split and assignment, and register updates to the region. /// /// Writes recorded through the update handle remain pending until [`Self::finish`]. fn fill( @@ -89,9 +100,15 @@ pub trait BuildAccessor: Send + Sized + HasIds { fn finish(self) -> impl SendFuture>; } +/////////////// +// Strategy // +////////////// + pub trait BuildStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { - type Accessor: BuildAccessor - + SetElements; + /// [`BuildAccessor`] used to insert new points, and performs splits if necessary. + type Accessor: SetElements + + BuildAccessor; + type Error: StandardError; fn build_accessor(&'a self, context: &'a C) -> Result; From d4b1cd92a661c6235237e3cde4aa8ebcd7024c06 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Wed, 21 Oct 2026 10:33:13 -0400 Subject: [PATCH 27/33] add search and generalize vector representation --- diskann/src/ivf/clustering.rs | 103 ++--- diskann/src/ivf/index.rs | 47 ++- diskann/src/ivf/mod.rs | 2 +- diskann/src/ivf/region.rs | 34 +- diskann/src/ivf/tests.rs | 722 ++++++++++++++++++++++++++++++++-- diskann/src/ivf/traits.rs | 39 +- diskann/src/utils/utils.rs | 2 + 7 files changed, 861 insertions(+), 88 deletions(-) diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs index abd8f8604e..c6c0219f9a 100644 --- a/diskann/src/ivf/clustering.rs +++ b/diskann/src/ivf/clustering.rs @@ -11,20 +11,23 @@ use std::{collections::BTreeSet, num::NonZeroUsize}; use diskann_quantization::algorithms::kmeans::lloyds::lloyds; -use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; +use diskann_utils::{ + Reborrow, + views::rowmajor::{self, Matrix, MatrixMut}, +}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use super::region::{RegionUpdate, RegionView}; -use crate::{ANNError, ANNResult}; +use crate::{ANNError, ANNResult, error::IntoANNResult, ivf::region::AsMatRef}; /// Fit a parent's replacement centroids without assigning regional memberships. -pub(super) trait Clusterer { +pub trait Clusterer { /// Fill `dst` with finite centroids of the canonical dimension. /// /// Implementations own initialization, convergence, empty-cluster handling, /// and reusable fitting scratch. `dst` has `num_centers()` rows and `dim()` /// columns. The number of centers must be at least two. - fn fit(&mut self, points: rowmajor::Ref<'_, T>, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; + fn fit(&mut self, points: V, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; fn num_centers(&self) -> NonZeroUsize; @@ -32,25 +35,24 @@ pub(super) trait Clusterer { } /// Assign the whole region, allowing policies that coordinate multiple sources. -pub(super) trait Assigner { +pub trait Assigner> { /// Assign each source member to a destination and record complete memberships. /// /// Create new destinations, replace retained memberships, and retire sources /// not retained as destinations. Validate and compute assignments before /// recording writes. A write error may leave partial private updates. - fn assign( + fn assign( &mut self, region: &R, update: &mut U, sources: &[L], - destinations: &[Destination<'_, L>], + destinations: &[Target<'_, L>], ) -> ANNResult<()> where - R: RegionView, U: RegionUpdate; } -pub(super) enum Type { +pub enum Type { /// Existing centroid. Exists(L), /// New centroid as a result of the split. @@ -58,19 +60,19 @@ pub(super) enum Type { } /// A borrowed centroid; only retained destinations have an existing list ID. -pub(super) struct Destination<'a, L> { +pub struct Target<'a, L> { pub list: Type, pub centroid: &'a [f32], } -pub(super) struct SplitAndAssign { +pub struct SplitAndAssign { pub clusterer: C, pub assigner: A, scratch: Option>, // scratch for the centroids. } -impl SplitAndAssign { - pub(super) fn new(clusterer: C, assigner: A) -> Self { +impl SplitAndAssign { + pub fn new(clusterer: C, assigner: A) -> Self { Self { clusterer, assigner, @@ -79,7 +81,10 @@ impl SplitAndAssign { } /// Initialize the scratch for a particular number of output cluster centroids. - fn initialize(&mut self, num_jobs: NonZeroUsize) -> ANNResult<()> { + fn initialize(&mut self, num_jobs: NonZeroUsize) -> ANNResult + where + C: Clusterer, + { let c = self.clusterer.num_centers(); let rows = num_jobs @@ -91,13 +96,14 @@ impl SplitAndAssign { self.scratch = Some(rowmajor::Owned::try_from_element(rows, dim, f32::NAN)?); - Ok(()) + Ok(c) } -} -impl SplitAndAssign { /// Compute the complete result before registering any provider writes. - pub(super) fn run( + /// + /// Concrete-reference clusterers may need an explicit `R` type argument because + /// of a compiler limitation when inferring higher-ranked associated types. + pub fn run( &mut self, region: &R, update: &mut U, @@ -105,13 +111,17 @@ impl SplitAndAssign { neighbors: &BTreeSet, ) -> ANNResult<()> where - A: Assigner, + A: Assigner, R: RegionView, + C: for<'a> Clusterer>, U: RegionUpdate, { - self.initialize(NonZeroUsize::new(parents.len()).ok_or(ANNError::message("err"))?)?; - - let n_centroids = self.clusterer.num_centers().get(); + let n_centroids = self + .initialize::>( + NonZeroUsize::new(parents.len()) + .ok_or_else(|| ANNError::message("IVF splitting needs at least one parent"))?, + )? + .get(); let mut s = 0; let mut clusters = self @@ -128,7 +138,7 @@ impl SplitAndAssign { let end = s + n_centroids; self.clusterer.fit( - vecs, + vecs.reborrow(), clusters.subview_mut(s..end).ok_or_else(|| { ANNError::message("IVF centroid scratch range is out of bounds") })?, @@ -139,7 +149,7 @@ impl SplitAndAssign { let mut destinations: Vec<_> = clusters .rows() - .map(|centroid| Destination { + .map(|centroid| Target { list: Type::New, centroid, }) @@ -150,7 +160,7 @@ impl SplitAndAssign { .get_view(n) .ok_or_else(|| ANNError::message("IVF neighbor is missing from the region"))?; - destinations.push(Destination { + destinations.push(Target { list: Type::Exists(n), centroid: view.centroid, }); @@ -164,49 +174,47 @@ impl SplitAndAssign { } /// Random input seeds followed by Lloyd iterations, which reset empty centers to zero. -pub(super) struct TwoMeans { +pub struct TwoMeans { iterations: NonZeroUsize, dim: NonZeroUsize, } impl TwoMeans { - pub(super) fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { + pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { Self { iterations, dim } } } -impl Clusterer for TwoMeans { - fn fit( - &mut self, - points: rowmajor::Ref<'_, f32>, - mut dst: rowmajor::Mut<'_, f32>, - ) -> ANNResult<()> { - let dim = self.dim().get(); - let num_centers = self.num_centers().get(); +impl> Clusterer for TwoMeans { + fn fit(&mut self, points: B, mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + let dim = self.dim.get(); + let num_centers = >::num_centers(self).get(); - if points.nrows() < num_centers { + let vectors = points.as_view().into_ann_result()?; + + if vectors.nrows() < num_centers { return Err(ANNError::message("an IVF split needs at least two points")); } - if points.ncols() != dim || dst.ncols() != dim || dst.nrows() != num_centers { + if vectors.ncols() != dim || dst.ncols() != dim || dst.nrows() != num_centers { return Err(ANNError::message( "two-means needs canonical input rows and a two-row destination of the canonical dimension", )); } dst.as_mut_slice().copy_from_slice( - points + vectors .subview(0..num_centers) .ok_or(ANNError::message("oh no!"))? .as_slice(), ); - let _ = lloyds(points, dst, self.iterations.get()); + let _ = lloyds(vectors, dst, self.iterations.get()); Ok(()) } fn num_centers(&self) -> NonZeroUsize { - const { NonZeroUsize::new(2).unwrap() } + const { crate::utils::TWO } } fn dim(&self) -> NonZeroUsize { @@ -214,18 +222,21 @@ impl Clusterer for TwoMeans { } } -pub(super) struct RegionalNearest; +pub struct RegionalNearest; -impl Assigner for RegionalNearest { - fn assign( +impl Assigner for RegionalNearest +where + R: RegionView, + for<'a> R::VectorsRef<'a>: AsMatRef, +{ + fn assign( &mut self, region: &R, update: &mut U, sources: &[L], - destinations: &[Destination<'_, L>], + destinations: &[Target<'_, L>], ) -> ANNResult<()> where - R: RegionView, U: RegionUpdate, { let mut memberships: Vec> = (0..destinations.len()).map(|_| Vec::new()).collect(); @@ -237,6 +248,8 @@ impl Assigner for RegionalNearest { let vectors = region.get_vectors(src).ok_or_else(|| { ANNError::message("IVF assignment source vectors are missing from the region") })?; + let vectors = vectors.reborrow(); + let vectors = vectors.as_view().into_ann_result()?; for (&id, point) in view.members.iter().zip(vectors.rows()) { let dst = nearest(point, destinations.iter().map(|d| d.centroid)); diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index efb11ea78a..cfc0984308 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -14,14 +14,20 @@ use std::{ use diskann_utils::future::SendFuture; use thiserror::Error; -use super::traits::{BuildAccessor, Router, RouterAndNearest, SetElements}; +use super::traits::{ + BuildAccessor, Router, RouterAndNearest, SearchAccessor, SearchStrategy, SetElements, +}; use crate::{ ANNError, ANNResult, - error::IntoANNResult, + error::{ErrorExt, IntoANNResult}, + flat::DistancesUnordered, + graph::SearchOutputBuffer, ivf::{ clustering::{RegionalNearest, SplitAndAssign, TwoMeans}, + region::AsMatRef, traits::BuildStrategy, }, + neighbor::{Neighbor, NeighborPriorityQueue}, provider::ExecutionContext, utils::ONE, }; @@ -84,6 +90,7 @@ impl IVFIndex { ) -> impl SendFuture> where S: BuildStrategy<'a, T, C>, + S::Accessor: for<'b> BuildAccessor: AsMatRef>, T: Sync, { let config = self.config; @@ -173,4 +180,40 @@ impl IVFIndex { }) } } + + 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 lists = Vec::new(); + { + let router = accessor.router(); + router + .route(&query, nprobe, |id, scan| lists.push((*id, scan))) + .await?; + } + + let mut best = NeighborPriorityQueue::new(k.get()); + + for (_, scan) in lists.iter_mut() { + scan.distances_unordered(|id, dist| best.insert(Neighbor::new(id, dist))) + .await + .escalate("why")?; + } + + Ok(output.extend(best.iter())) + } + } } diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 2929671ef1..6919fce5bd 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -37,7 +37,7 @@ //! recursively split. Initialization, search, deletion, and colocation are not //! implemented. -mod clustering; +pub mod clustering; mod index; pub mod partition; pub mod pending; diff --git a/diskann/src/ivf/region.rs b/diskann/src/ivf/region.rs index c0550eb158..54290ac6e2 100644 --- a/diskann/src/ivf/region.rs +++ b/diskann/src/ivf/region.rs @@ -5,9 +5,11 @@ //! Regional input views and write operations. -use diskann_utils::views::rowmajor; +use std::fmt::Debug; -use crate::ANNResult; +use diskann_utils::{Reborrow, views::rowmajor}; + +use crate::{ANNError, ANNResult}; /// One source list's centroid and complete operation-local membership. /// @@ -22,6 +24,14 @@ pub struct View<'a, I, L> { /// Views for one build operation to perform splits/re-assignments . pub trait RegionView: 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_view(&self, i: L) -> Option>; @@ -29,7 +39,7 @@ pub trait RegionView: Send { /// /// 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>; + fn get_vectors(&self, i: L) -> Option>; } /// Regional writes for one build operation. @@ -43,3 +53,21 @@ pub trait RegionUpdate: Send { /// Record a new centroid and its complete membership. fn create(&mut self, centroid: &[f32], members: &[I]) -> ANNResult<()>; } + +pub trait AsMatRef { + type Error: Send + Sync + Debug + Into; + + fn as_view(&self) -> Result, Self::Error>; +} + +impl AsMatRef for rowmajor::Ref<'_, T> { + type Error = crate::error::Infallible; + + fn as_view(&self) -> Result, Self::Error> { + Ok(*self) + } +} + +pub trait AsMatMut: AsMatRef { + fn as_mut_view(&self) -> Result, Self::Error>; +} diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index 65de9d1c1d..bd05c99922 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -11,6 +11,7 @@ use std::{ }; use diskann_utils::{ + Reborrow, future::SendFuture, views::rowmajor::{self, Matrix, MatrixMut}, }; @@ -20,10 +21,20 @@ use tokio::sync::{Mutex, MutexGuard}; use super::{ Config, IVFIndex, InsertStats, ReassignCount, clustering, pending::{Membership, PendingLists}, - region::{RegionUpdate, RegionView, View}, - traits::{BuildAccessor, BuildStrategy, HasIds, Router, RouterAndNearest, SetElements}, + region::{AsMatRef, RegionUpdate, RegionView, View}, + traits::{ + BuildAccessor, BuildStrategy, HasIds, Router, RouterAndNearest, SearchAccessor, + SearchStrategy, SetElements, + }, +}; +use crate::{ + ANNError, ANNResult, + flat::DistancesUnordered, + graph::{IdDistance, SearchOutputBuffer}, + neighbor::{BackInserter, Neighbor}, + provider::{DefaultContext, HasId}, + utils::ONE, }; -use crate::{ANNError, ANNResult, provider::DefaultContext, utils::ONE}; #[derive(Debug, Clone)] struct StoredList { @@ -40,6 +51,7 @@ struct Store { finishes: usize, fills: usize, filled_lists: Vec, + fail_route: bool, fail_nearest: bool, fail_fill: bool, fail_create: bool, @@ -76,9 +88,21 @@ struct Build<'p, T> { struct Coarse<'a> { stored: &'a BTreeMap, + fail_route: bool, 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 @@ -90,17 +114,7 @@ impl Coarse<'_> { .stored .iter() .filter(|(id, _)| pred(id)) - .map(|(&id, list)| { - let distance: f32 = query - .iter() - .zip(&list.centroid) - .map(|(&value, ¢er)| { - let delta = value.into() - center; - delta * delta - }) - .sum(); - (id, distance) - }) + .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()); @@ -113,6 +127,7 @@ impl Coarse<'_> { impl + Sync> Router<(u32, &[T])> for Coarse<'_> { type Id = u32; + type Ret = f32; fn len(&self) -> usize { self.stored.len() @@ -127,7 +142,38 @@ impl + Sync> Router<(u32, &[T])> for Coarse<'_> { where F: FnMut(&Self::Id, f32) + Send, { - async move { Ok(self.select(query.1, nprobe, |_| true, emit)) } + async move { + if self.fail_route { + return Err(ANNError::message("injected routing failure")); + } + Ok(self.select(query.1, nprobe, |_| true, emit)) + } + } +} + +impl + Sync> Router<&[T]> for Coarse<'_> { + type Id = u32; + type Ret = f32; + + fn len(&self) -> usize { + self.stored.len() + } + + fn route( + &self, + query: &&[T], + nprobe: NonZeroUsize, + emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::Id, f32) + Send, + { + async move { + if self.fail_route { + return Err(ANNError::message("injected routing failure")); + } + Ok(self.select(query, nprobe, |_| true, emit)) + } } } @@ -211,6 +257,12 @@ impl<'a> Region<'a> { } impl RegionView for Region<'_> { + type VectorsRef<'a> = rowmajor::Ref<'a, f32>; + type Vectors<'a> + = rowmajor::Ref<'a, f32> + where + Self: 'a; + fn get_view(&self, list: u32) -> Option> { self.sources.get(&list).map(|source| View { id: list, @@ -272,6 +324,7 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build = Coarse<'a> where Self: 'a; + type VectorsRef<'a> = rowmajor::Ref<'a, f32>; type Region<'a> = Region<'a> where @@ -284,6 +337,7 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build fn router(&mut self) -> Self::Router<'_> { Coarse { stored: &self.store.lists, + fail_route: self.store.fail_route, fail_nearest: self.store.fail_nearest, } } @@ -456,6 +510,139 @@ impl<'a, T: Copy + Into + Send + Sync> BuildStrategy<'a, (u32, &[T]), Defau } } +struct Search<'p, 'q, T> { + store: MutexGuard<'p, Store>, + query: &'q [T], +} + +struct SearchRouter<'a, T> { + store: &'a Store, + query: &'a [T], +} + +struct ListScan<'a, T> { + store: &'a Store, + list: u32, + query: &'a [T], +} + +impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { + type Id = u32; + type Ret = ListScan<'a, T>; + + fn len(&self) -> usize { + self.store.lists.len() + } + + fn route( + &self, + query: &&[T], + nprobe: NonZeroUsize, + mut emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::Id, Self::Ret) + Send, + { + 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, + }; + Ok(coarse.select( + query, + nprobe, + |_| true, + |id, _| { + emit( + id, + ListScan { + store: self.store, + list: *id, + query: self.query, + }, + ); + }, + )) + } + } +} + +impl HasId for ListScan<'_, T> { + type Id = u32; +} + +impl + Sync> DistancesUnordered for ListScan<'_, T> { + type Error = ANNError; + + fn distances_unordered(&mut self, mut emit: F) -> impl SendFuture> + where + F: FnMut(Self::Id, f32) + Send, + { + async move { + let list = self + .store + .lists + .get(&self.list) + .ok_or_else(|| ANNError::message(format!("unknown list {}", self.list)))?; + for id in &list.members { + let vector = self + .store + .points + .get(id) + .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; + emit(*id, squared_distance(self.query, vector)); + } + Ok(()) + } + } +} + +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> { + type Router<'a> + = SearchRouter<'a, T> + where + Self: 'a; + type Dist<'a> + = ListScan<'a, T> + where + Self: 'a; + + fn router(&self) -> Self::Router<'_> { + SearchRouter { + store: &self.store, + query: self.query, + } + } +} + +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 }) + } +} + impl<'p, T: Copy + Into + Send + Sync> SetElements<(u32, &'p [T]), DefaultContext> for Build<'_, T> { @@ -560,12 +747,221 @@ fn require_send(future: F) -> F { future } +#[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"), + ] { + let mut store = seeded(); + if case == "route" { + store.fail_route = true; + } + if case == "scan" { + store.points.remove(&10_003); + } + 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 mut output = vec![Neighbor::new(42, -1.0)]; + let error = require_send(index.search( + &strategy, + &DefaultContext, + query, + NonZeroUsize::new(3).unwrap(), + ONE, + &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)] + ); +} + #[tokio::test] async fn routing_callbacks_emit_ranked_results_and_counts() { 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] { @@ -589,6 +985,7 @@ async fn nearest_callbacks_filter_before_limiting_and_count_emitted_results() { let store = seeded(); let router = Coarse { stored: &store.lists, + fail_route: false, fail_nearest: false, }; for (limit, excluded, expected) in [ @@ -621,6 +1018,7 @@ async fn nearest_lookup_failures_do_not_emit_results() { ] { let router = Coarse { stored: &store.lists, + fail_route: false, fail_nearest, }; let mut emitted = Vec::new(); @@ -1379,10 +1777,10 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { fn split_strategies_support_in_place_fitting_and_source_local_candidates() { struct RowsAsCenters; - impl clustering::Clusterer for RowsAsCenters { + impl<'v> clustering::Clusterer> for RowsAsCenters { fn fit( &mut self, - points: rowmajor::Ref<'_, f32>, + points: rowmajor::Ref<'v, f32>, mut dst: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { assert_eq!(points.nrows(), dst.nrows()); @@ -1403,16 +1801,18 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { owners: BTreeMap>, } - impl clustering::Assigner for ParentLocal { - fn assign( + impl clustering::Assigner for ParentLocal + where + R: for<'a> RegionView = rowmajor::Ref<'a, f32>>, + { + fn assign( &mut self, region: &R, update: &mut U, sources: &[u32], - destinations: &[clustering::Destination<'_, u32>], + destinations: &[clustering::Target<'_, u32>], ) -> ANNResult<()> where - R: RegionView, U: RegionUpdate, { assert_eq!(sources, [1, 2, 3, 4]); @@ -1426,6 +1826,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { for &source in sources { let view = region.get_view(source).unwrap(); let vectors = region.get_vectors(source).unwrap(); + let vectors = vectors.reborrow(); let mut candidates: Vec<_> = destinations .iter() .enumerate() @@ -1513,7 +1914,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut full = Update::new(&sources, 1, &mut full_memberships, &mut full_created); let mut splitter = clustering::SplitAndAssign::new(RowsAsCenters, clustering::RegionalNearest); splitter - .run(®ion, &mut full, &parents, &neighbors) + .run::<_, _, Region<'_>, _>(®ion, &mut full, &parents, &neighbors) .unwrap(); let mut restricted_memberships = PendingLists::default(); let mut restricted_created = Vec::new(); @@ -1529,7 +1930,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { owners: [(3, vec![1]), (4, vec![1, 2])].into(), }, ) - .run(®ion, &mut restricted, &parents, &neighbors) + .run::<_, _, Region<'_>, _>(®ion, &mut restricted, &parents, &neighbors) .unwrap(); for update in [&full, &restricted] { @@ -1556,7 +1957,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { &mut repeated_created, ); splitter - .run(®ion, &mut repeated, &parents, &neighbors) + .run::<_, _, Region<'_>, _>(®ion, &mut repeated, &parents, &neighbors) .unwrap(); for (first, repeated) in full.created.iter().zip(repeated.created.iter()) { assert_eq!(first.centroid, repeated.centroid); @@ -1566,7 +1967,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut smaller_created = Vec::new(); let mut smaller = Update::new(&sources, 1, &mut smaller_memberships, &mut smaller_created); splitter - .run( + .run::<_, _, Region<'_>, _>( ®ion, &mut smaller, &BTreeSet::from([1]), @@ -1599,8 +2000,8 @@ fn numerical_strategy_failures_record_no_updates() { CountOverflow, } - impl clustering::Clusterer for InvalidFit { - fn fit(&mut self, _: rowmajor::Ref<'_, f32>, _: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + impl clustering::Clusterer for InvalidFit { + fn fit(&mut self, _: V, _: rowmajor::Mut<'_, f32>) -> ANNResult<()> { match self { Self::Error => Err(ANNError::message("injected fitting failure")), Self::CountOverflow => { @@ -1624,16 +2025,15 @@ fn numerical_strategy_failures_record_no_updates() { struct InvalidAssignment; - impl clustering::Assigner for InvalidAssignment { - fn assign( + impl> clustering::Assigner for InvalidAssignment { + fn assign( &mut self, _: &R, _: &mut U, _: &[u32], - _: &[clustering::Destination<'_, u32>], + _: &[clustering::Target<'_, u32>], ) -> ANNResult<()> where - R: RegionView, U: RegionUpdate, { Err(ANNError::message("injected assignment failure")) @@ -1691,6 +2091,262 @@ fn numerical_strategy_failures_record_no_updates() { assert_unwritten(update); } +#[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_view(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 split_driver_reuses_clusterer_with_fallible_borrowed_batches() { + struct Batch<'a> { + vectors: rowmajor::Ref<'a, f32>, + reads: &'a Cell, + fail_on: Option, + } + + struct BatchRef<'a> { + vectors: rowmajor::Ref<'a, f32>, + reads: &'a Cell, + fail_on: Option, + } + + impl<'r> Reborrow<'r> for Batch<'_> { + type Target = BatchRef<'r>; + + fn reborrow(&'r self) -> Self::Target { + BatchRef { + vectors: self.vectors.reborrow(), + reads: self.reads, + fail_on: self.fail_on, + } + } + } + + impl AsMatRef for BatchRef<'_> { + type Error = ANNError; + + fn as_view(&self) -> ANNResult> { + let count = self.reads.get() + 1; + self.reads.set(count); + if self.fail_on == Some(count) { + return Err(ANNError::message("injected batch conversion failure")); + } + Ok(self.vectors) + } + } + + struct BatchRegion<'a> { + inner: Region<'a>, + reads: Cell, + fail_on: Option, + } + + impl RegionView for BatchRegion<'_> { + type VectorsRef<'a> = BatchRef<'a>; + type Vectors<'a> + = Batch<'a> + where + Self: 'a; + + fn get_view(&self, list: u32) -> Option> { + self.inner.get_view(list) + } + + fn get_vectors(&self, list: u32) -> Option> { + self.inner.get_vectors(list).map(|vectors| Batch { + vectors, + reads: &self.reads, + fail_on: self.fail_on, + }) + } + } + + let mut driver = clustering::SplitAndAssign::new( + clustering::TwoMeans::new(ONE, ONE), + clustering::RegionalNearest, + ); + for fail_on in [Some(1), Some(2), None] { + let sources = BTreeMap::from([( + 1, + Source { + centroid: vec![0.0], + members: vec![0, 1], + vectors: matrix(&[[-1.0], [1.0]]), + }, + )]); + let region = BatchRegion { + inner: Region::new(&sources), + reads: Cell::new(0), + fail_on, + }; + 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 = Update::new(&sources, 1, &mut memberships, &mut created); + let result = driver.run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()); + if let Some(count) = fail_on { + let error = result.unwrap_err(); + assert!( + error + .to_string() + .contains("injected batch conversion failure") + ); + assert_eq!(region.reads.get(), count); + assert!(created.is_empty()); + assert_eq!(memberships.into_changes().count(), 0); + } else { + result.unwrap(); + assert_eq!(region.reads.get(), 2); + 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]); + assert!(matches!( + memberships.into_changes().collect::>().as_slice(), + [(1, Membership::Retire)] + )); + } + } +} + +#[test] +fn split_driver_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 RegionView for NativeRegion<'_> { + type VectorsRef<'a> = &'a [[u8; 1]]; + type Vectors<'a> + = NativeBatch<'a> + where + Self: 'a; + + fn get_view(&self, list: u32) -> Option> { + self.inner.get_view(list) + } + + fn get_vectors(&self, list: u32) -> Option> { + self.inner.get_view(list).map(|_| NativeBatch(self.vectors)) + } + } + + struct NativeCenters; + + impl<'v> clustering::Clusterer<&'v [[u8; 1]]> for NativeCenters { + fn fit(&mut self, points: &'v [[u8; 1]], mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + for (point, center) in points.iter().zip(dst.rows_mut()) { + center[0] = point[0].into(); + } + Ok(()) + } + + fn num_centers(&self) -> NonZeroUsize { + NonZeroUsize::new(2).unwrap() + } + + fn dim(&self) -> NonZeroUsize { + ONE + } + } + + struct OrderedAssignment; + + impl> clustering::Assigner for OrderedAssignment { + fn assign( + &mut self, + region: &R, + update: &mut U, + sources: &[u32], + destinations: &[clustering::Target<'_, u32>], + ) -> ANNResult<()> + where + U: RegionUpdate, + { + let view = region.get_view(sources[0]).unwrap(); + for (destination, member) in destinations.iter().zip(view.members) { + update.create(destination.centroid, std::slice::from_ref(member))?; + } + update.retire(sources[0]) + } + } + + let mut driver = clustering::SplitAndAssign::new(NativeCenters, OrderedAssignment); + 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 = Update::new(&sources, 1, &mut memberships, &mut created); + driver + .run::<_, _, NativeRegion<'_>, _>( + ®ion, + &mut update, + &BTreeSet::from([1]), + &BTreeSet::new(), + ) + .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() { use clustering::{Clusterer, TwoMeans}; @@ -1749,7 +2405,7 @@ fn two_means_fits_only_the_destination_window() { #[test] fn regional_assignment_reports_missing_sources_before_writing() { - use clustering::{Assigner, Destination, RegionalNearest, Type}; + use clustering::{Assigner, RegionalNearest, Target, Type}; let sources = BTreeMap::from([ ( @@ -1769,7 +2425,7 @@ fn regional_assignment_reports_missing_sources_before_writing() { }, ), ]); - let destinations = [Destination { + let destinations = [Target { list: Type::New, centroid: &[0.0], }]; diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 55ff5f1427..96fd309149 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -10,7 +10,10 @@ use std::num::NonZeroUsize; use diskann_utils::future::SendFuture; use super::region::{RegionUpdate, RegionView}; -use crate::{ANNResult, error::StandardError, provider::ExecutionContext, utils::VectorId}; +use crate::{ + ANNResult, error::StandardError, flat::DistancesUnordered, provider::ExecutionContext, + utils::VectorId, +}; pub trait HasIds { type Id: VectorId; @@ -20,6 +23,7 @@ pub trait HasIds { /// Centroid selection within one build operation. pub trait Router: Send { type Id; + type Ret; fn len(&self) -> usize; @@ -35,7 +39,7 @@ pub trait Router: Send { emit: F, ) -> impl SendFuture> where - F: FnMut(&Self::Id, f32) + Send; + F: FnMut(&Self::Id, Self::Ret) + Send; } pub trait RouterAndNearest: Router { @@ -63,7 +67,9 @@ pub trait BuildAccessor: Send + Sized + HasIds { where Self: 'a; - type Region<'a>: RegionView + type VectorsRef<'a>; + + type Region<'a>: for<'b> RegionView = Self::VectorsRef<'b>> where Self: 'a; @@ -100,7 +106,7 @@ pub trait BuildAccessor: Send + Sized + HasIds { fn finish(self) -> impl SendFuture>; } -/////////////// +////////////// // Strategy // ////////////// @@ -123,3 +129,28 @@ pub trait SetElements { elements: &[T], ) -> impl std::future::Future>> + Send; } + +//////////// +// Search // +//////////// + +pub trait SearchStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { + /// [`BuildAccessor`] used to insert new points, and performs splits if necessary. + type Accessor: SearchAccessor; + + type Error: StandardError; + + fn search_accessor(&'a self, context: &'a C, query: T) -> Result; +} + +pub trait SearchAccessor: Send + Sized + HasIds { + type Router<'a>: Router> + where + Self: 'a; + + type Dist<'a>: DistancesUnordered + where + Self: 'a; + + fn router(&self) -> Self::Router<'_>; +} 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> From 2aae82650afa1209248b4cda5ac2c1467ea74225 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Wed, 21 Oct 2026 13:35:14 -0400 Subject: [PATCH 28/33] consolidate setelements, routing and appending --- diskann/src/ivf/index.rs | 67 ++----- diskann/src/ivf/tests.rs | 355 ++++++++++++++++++++++++-------------- diskann/src/ivf/traits.rs | 44 ++--- 3 files changed, 255 insertions(+), 211 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index cfc0984308..3975fccdc8 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -5,20 +5,14 @@ //! IVF policy and orchestration, independent of physical list storage. -use std::{ - collections::{BTreeMap, BTreeSet}, - marker::PhantomData, - num::NonZeroUsize, -}; +use std::{collections::BTreeSet, marker::PhantomData, num::NonZeroUsize}; use diskann_utils::future::SendFuture; use thiserror::Error; -use super::traits::{ - BuildAccessor, Router, RouterAndNearest, SearchAccessor, SearchStrategy, SetElements, -}; +use super::traits::{BuildAccessor, NearestList, Router, SearchAccessor, SearchStrategy}; use crate::{ - ANNError, ANNResult, + ANNResult, error::{ErrorExt, IntoANNResult}, flat::DistancesUnordered, graph::SearchOutputBuffer, @@ -29,7 +23,6 @@ use crate::{ }, neighbor::{Neighbor, NeighborPriorityQueue}, provider::ExecutionContext, - utils::ONE, }; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -82,13 +75,14 @@ impl IVFIndex { }) } - pub fn insert_batch<'a, S, T>( + pub fn insert_batch<'a, S, T, Itr>( &'a mut self, strategy: &'a S, context: &'a C, - points: &'a [T], + points: Itr, ) -> impl SendFuture> where + Itr: ExactSizeIterator + Clone + Send + Sync, S: BuildStrategy<'a, T, C>, S::Accessor: for<'b> BuildAccessor: AsMatRef>, T: Sync, @@ -96,7 +90,8 @@ impl IVFIndex { let config = self.config; async move { - if points.is_empty() { + let inserted = points.len(); + if inserted == 0 { return Ok(InsertStats::default()); } @@ -104,38 +99,15 @@ impl IVFIndex { let dim = accessor.dim(); - // Step 1: Compute the batch - let ids = accessor.set_elements(context, points).await?; - - // Step 2: Route the points to their nearest centroid. - let mut routed = BTreeMap::>::new(); - - { - let router = accessor.router(); - - for (id, point) in ids.iter().zip(points.iter()) { - let count = router - .route(point, ONE, |list, _| { - routed.entry(*list).or_default().push(*id); - }) - .await?; - - if count != 1 { - return Err(ANNError::message("IVF routing returned no live centroid")); - } - } - } - let mut parents = BTreeSet::new(); - // Step 3: append the points to their respective list ids and determine which need to be split. - for (list, members) in routed.iter() { - let size = accessor.append_to(*list, members.iter().copied()).await?; - - if size.get() > config.split_threshold.get() { - parents.insert(*list); - } - } + accessor + .route(points, |&list, size| { + if size.get() > config.split_threshold.get() { + parents.insert(list); + } + }) + .await?; let splits = parents.len(); @@ -143,11 +115,11 @@ impl IVFIndex { let mut neighbors = BTreeSet::new(); if let ReassignCount::Count(n) = config.reassign_count { - let router = accessor.router(); + let router = accessor.nearest(); for &pid in parents.iter() { router - .nearest( + .find( pid, n, |id| !parents.contains(id), @@ -174,10 +146,7 @@ impl IVFIndex { accessor.finish().await?; - Ok(InsertStats { - inserted: points.len(), - splits, - }) + Ok(InsertStats { inserted, splits }) } } diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests.rs index bd05c99922..b0706dab3a 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests.rs @@ -23,8 +23,7 @@ use super::{ pending::{Membership, PendingLists}, region::{AsMatRef, RegionUpdate, RegionView, View}, traits::{ - BuildAccessor, BuildStrategy, HasIds, Router, RouterAndNearest, SearchAccessor, - SearchStrategy, SetElements, + BuildAccessor, BuildStrategy, HasIds, NearestList, Router, SearchAccessor, SearchStrategy, }, }; use crate::{ @@ -138,7 +137,7 @@ impl + Sync> Router<(u32, &[T])> for Coarse<'_> { query: &(u32, &[T]), nprobe: NonZeroUsize, emit: F, - ) -> impl SendFuture> + ) -> impl SendFuture> where F: FnMut(&Self::Id, f32) + Send, { @@ -146,7 +145,8 @@ impl + Sync> Router<(u32, &[T])> for Coarse<'_> { if self.fail_route { return Err(ANNError::message("injected routing failure")); } - Ok(self.select(query.1, nprobe, |_| true, emit)) + self.select(query.1, nprobe, |_| true, emit); + Ok(()) } } } @@ -164,7 +164,7 @@ impl + Sync> Router<&[T]> for Coarse<'_> { query: &&[T], nprobe: NonZeroUsize, emit: F, - ) -> impl SendFuture> + ) -> impl SendFuture> where F: FnMut(&Self::Id, f32) + Send, { @@ -172,13 +172,16 @@ impl + Sync> Router<&[T]> for Coarse<'_> { if self.fail_route { return Err(ANNError::message("injected routing failure")); } - Ok(self.select(query, nprobe, |_| true, emit)) + self.select(query, nprobe, |_| true, emit); + Ok(()) } } } -impl + Sync> RouterAndNearest<(u32, &[T])> for Coarse<'_> { - fn nearest( +impl NearestList for Coarse<'_> { + type Id = u32; + + fn find( &self, id: Self::Id, nprobe: NonZeroUsize, @@ -315,12 +318,12 @@ impl HasIds for Build<'_, T> { type ListId = u32; } -impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build<'p, T> { +impl<'p, 'v, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &'v [T])> for Build<'p, T> { type Update<'a> = Update<'a> where Self: 'a; - type Router<'a> + type Nearest<'a> = Coarse<'a> where Self: 'a; @@ -334,7 +337,7 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build NonZeroUsize::new(self.store.dim).unwrap() } - fn router(&mut self) -> Self::Router<'_> { + fn nearest(&mut self) -> Self::Nearest<'_> { Coarse { stored: &self.store.lists, fail_route: self.store.fail_route, @@ -342,25 +345,60 @@ impl<'p, T: Copy + Into + Send + Sync> BuildAccessor<(u32, &[T])> for Build } } - fn append_to( - &mut self, - list: u32, - members: Itr, - ) -> impl SendFuture> + fn route(&mut self, points: Itr, mut emit: F) -> impl SendFuture> where - Itr: ExactSizeIterator + Clone + Send + Sync, + Itr: ExactSizeIterator + Clone + Send + Sync, + F: FnMut(&Self::ListId, NonZeroUsize) + Send, { async move { - let size = self - .store - .lists - .get(&list) - .ok_or_else(|| ANNError::message(format!("unknown list {list}")))? - .members - .len(); - self.pending.append(list, members.collect())?; - NonZeroUsize::new(self.pending.list_size(list, size)?) - .ok_or_else(|| ANNError::message("appending must leave a nonempty list")) + let mut staged = BTreeMap::new(); + let mut routed = BTreeMap::>::new(); + for point in points { + let (external, vector) = point; + 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(&point, 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(()) } } @@ -539,7 +577,7 @@ impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { query: &&[T], nprobe: NonZeroUsize, mut emit: F, - ) -> impl SendFuture> + ) -> impl SendFuture> where F: FnMut(&Self::Id, Self::Ret) + Send, { @@ -552,7 +590,7 @@ impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { fail_route: false, fail_nearest: false, }; - Ok(coarse.select( + coarse.select( query, nprobe, |_| true, @@ -566,7 +604,8 @@ impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { }, ); }, - )) + ); + Ok(()) } } } @@ -643,41 +682,6 @@ impl<'a, 'q, T: Copy + Into + Send + Sync> SearchStrategy<'a, &'q [T], Defa } } -impl<'p, T: Copy + Into + Send + Sync> SetElements<(u32, &'p [T]), DefaultContext> - for Build<'_, T> -{ - type Id = u32; - - async fn set_elements( - &mut self, - _: &DefaultContext, - points: &[(u32, &'p [T])], - ) -> ANNResult> { - let mut staged = BTreeMap::new(); - let mut ids = Vec::with_capacity(points.len()); - 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")); - } - staged.insert(id, vector.to_vec()); - ids.push(id); - } - self.staged.extend(staged); - Ok(ids) - } -} - 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() @@ -927,7 +931,7 @@ 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) + .insert_batch(&strategy, &DefaultContext, points.iter().copied()) .await .unwrap(); assert_eq!(stats.splits, 1); @@ -956,7 +960,7 @@ async fn search_supports_native_queries_after_insertion_and_splitting() { } #[tokio::test] -async fn routing_callbacks_emit_ranked_results_and_counts() { +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 { @@ -966,7 +970,7 @@ async fn routing_callbacks_emit_ranked_results_and_counts() { }; for limit in [1, 2, usize::MAX] { let mut emitted = Vec::new(); - let count = require_send(router.route( + require_send(router.route( &(9, &[6u8, 0][..]), NonZeroUsize::new(limit).unwrap(), |id, distance| emitted.push((*id, distance)), @@ -974,7 +978,6 @@ async fn routing_callbacks_emit_ranked_results_and_counts() { .await .unwrap(); let expected = &expected[..limit.min(store.lists.len())]; - assert_eq!(count, expected.len()); assert_eq!(emitted, expected); } } @@ -983,7 +986,7 @@ async fn routing_callbacks_emit_ranked_results_and_counts() { #[tokio::test] async fn nearest_callbacks_filter_before_limiting_and_count_emitted_results() { let store = seeded(); - let router = Coarse { + let nearest = Coarse { stored: &store.lists, fail_route: false, fail_nearest: false, @@ -995,8 +998,7 @@ async fn nearest_callbacks_filter_before_limiting_and_count_emitted_results() { (2, vec![1, 2, 3], vec![]), ] { let mut emitted = Vec::new(); - let count = require_send(RouterAndNearest::<(u32, &[f32])>::nearest( - &router, + let count = require_send(nearest.find( 1, NonZeroUsize::new(limit).unwrap(), |id| !excluded.contains(id), @@ -1016,14 +1018,13 @@ async fn nearest_lookup_failures_do_not_emit_results() { (false, 99, "unknown list 99"), (true, 1, "injected nearest failure"), ] { - let router = Coarse { + let nearest = Coarse { stored: &store.lists, fail_route: false, fail_nearest, }; let mut emitted = Vec::new(); - let error = require_send(RouterAndNearest::<(u32, &[f32])>::nearest( - &router, + let error = require_send(nearest.find( id, ONE, |_| true, @@ -1037,26 +1038,45 @@ async fn nearest_lookup_failures_do_not_emit_results() { } #[tokio::test] -async fn staging_preserves_native_inputs_and_id_order() { +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 ids = require_send(build.set_elements(&DefaultContext, &points)) - .await - .unwrap(); - assert_eq!(ids, [10_009, 10_007, 10_008]); - for (&id, point) in ids.iter().zip(&points) { + let mut emitted = Vec::new(); + require_send(build.route(points.iter().copied(), |&list, size| { + emitted.push((list, size.get())); + })) + .await + .unwrap(); + 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.router(); + let router = build.nearest(); for (point, expected) in points.iter().zip([2, 1, 1]) { - let mut selected = None; - let count = require_send(router.route(point, ONE, |id, _| selected = Some(*id))) + let mut selected = Vec::new(); + require_send(router.route(point, ONE, |&id, _| selected.push(id))) .await .unwrap(); - assert_eq!(count, 1); - assert_eq!(selected, Some(expected)); + assert_eq!(selected, [expected]); } drop(build); let store = strategy.store(); @@ -1065,12 +1085,70 @@ async fn staging_preserves_native_inputs_and_id_order() { } #[tokio::test] -async fn native_inputs_are_inserted_using_the_returned_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)) +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 { + require_send(build.route(std::iter::once(point), |&list, size| { + emitted.push((list, size.get())); + })) .await .unwrap(); + } + 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.into_iter(), |&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.iter().copied())) + .await + .unwrap(); assert_eq!( stats, InsertStats { @@ -1100,12 +1178,17 @@ async fn failed_staging_leaves_no_new_elements() { 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.set_elements(&DefaultContext, &points)) - .await - .is_err() + require_send(build.route(points.iter().copied(), |&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); } @@ -1144,9 +1227,10 @@ async fn appends_publish_actual_memberships_without_filling() { (8, &[1.0, 0.0][..]), (9, &[10.0, 1.0][..]), ]; - let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points)) - .await - .unwrap(); + let stats = + require_send(index.insert_batch(&strategy, &DefaultContext, points.iter().copied())) + .await + .unwrap(); assert_eq!( stats, @@ -1167,7 +1251,7 @@ async fn appends_publish_actual_memberships_without_filling() { let empty: [(u32, &[f32]); 0] = []; assert_eq!( index - .insert_batch(&strategy, &DefaultContext, &empty) + .insert_batch(&strategy, &DefaultContext, empty.into_iter()) .await .unwrap(), InsertStats::default() @@ -1185,9 +1269,10 @@ async fn splitting_commits_children_and_ordinary_appends_once() { (6, &[11.0, 1.0][..]), (7, &[101.0, 0.0][..]), ]; - let stats = require_send(index.insert_batch(&strategy, &DefaultContext, &points)) - .await - .unwrap(); + let stats = + require_send(index.insert_batch(&strategy, &DefaultContext, points.iter().copied())) + .await + .unwrap(); let store = strategy.store(); assert_eq!( @@ -1246,9 +1331,10 @@ async fn neighboring_lists_exchange_points_and_pending_appends_once() { 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 stats = + require_send(index.insert_batch(&strategy, &DefaultContext, points.iter().copied())) + .await + .unwrap(); let store = strategy.store(); assert_eq!( @@ -1311,9 +1397,13 @@ async fn neighborhood_selection_filters_parents_and_fills_the_union_once() { 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 stats = require_send(index.insert_batch( + &strategy, + &DefaultContext, + points.iter().copied(), + )) + .await + .unwrap(); let store = strategy.store(); let neighbors = match (count, centers.len()) { (0, _) | (_, 0) => vec![], @@ -1429,16 +1519,14 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { { let mut build: Build<'_, f32> = strategy.build_accessor(&DefaultContext).unwrap(); let points = [(5, &[1.0, 1.0][..]), (6, &[3.0, 1.0][..])]; - let ids = build.set_elements(&DefaultContext, &points).await.unwrap(); - assert_eq!(ids, [10_005, 10_006]); - assert_eq!( - build - .append_to(1, [10_006, 10_005].into_iter()) - .await - .unwrap() - .get(), - 4 - ); + let mut emitted = Vec::new(); + build + .route(points.iter().rev().copied(), |&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_view(1).unwrap(); @@ -1493,10 +1581,10 @@ async fn borrowed_updates_share_pending_memberships_and_publish_on_finish() { (6, &[11.0, 0.0][..]), (7, &[100.0, 0.0][..]), ]; - build.set_elements(&DefaultContext, &points).await.unwrap(); - for (list, id) in [(1, 10_005), (2, 10_006), (3, 10_007)] { - build.append_to(list, [id].into_iter()).await.unwrap(); - } + build + .route(points.iter().copied(), |_, _| {}) + .await + .unwrap(); { let (region, mut update) = require_send(build.fill([1, 2].into_iter())).await.unwrap(); @@ -1560,9 +1648,10 @@ async fn failed_planning_or_publication_is_not_reported_as_success() { 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 error = + require_send(index.insert_batch(&strategy, &DefaultContext, points.iter().copied())) + .await + .unwrap_err(); let expected = if fail_nearest { "injected nearest failure" } else if fail_fill { @@ -1598,7 +1687,7 @@ async fn uninitialized_insertion_fails() { ); let points: [(u32, &[f32]); 1] = [(7, &[0.0, 1.0][..])]; let error = index - .insert_batch(&strategy, &DefaultContext, &points) + .insert_batch(&strategy, &DefaultContext, points.iter().copied()) .await .unwrap_err(); assert!( @@ -1616,19 +1705,19 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() 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][..])]; - build.set_elements(&DefaultContext, &points).await.unwrap(); - assert_eq!( - build - .append_to(1, [10_005].into_iter()) - .await - .unwrap() - .get(), - 3 - ); + let mut emitted = Vec::new(); + build + .route(points.iter().copied(), |&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(); diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index 96fd309149..48902783a9 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -32,17 +32,13 @@ pub trait Router: Send { } /// Select `min(nprobe, len())` distinct live lists, nearest first. - fn route( - &self, - query: &T, - nprobe: NonZeroUsize, - emit: F, - ) -> impl SendFuture> + fn route(&self, query: &T, nprobe: NonZeroUsize, emit: F) -> impl SendFuture> where F: FnMut(&Self::Id, Self::Ret) + Send; } -pub trait RouterAndNearest: Router { +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. @@ -50,7 +46,7 @@ pub trait RouterAndNearest: Router { /// # Errors /// /// Fails if `id` is not live or the provider cannot complete the lookup. - fn nearest( + fn find( &self, id: Self::Id, nprobe: NonZeroUsize, @@ -63,7 +59,7 @@ pub trait RouterAndNearest: Router { } pub trait BuildAccessor: Send + Sized + HasIds { - type Router<'a>: RouterAndNearest + type Nearest<'a>: NearestList where Self: 'a; @@ -80,16 +76,17 @@ pub trait BuildAccessor: Send + Sized + HasIds { /// Positive dimension of all canonical points and centroids. fn dim(&self) -> NonZeroUsize; - fn router(&mut self) -> Self::Router<'_>; + fn nearest(&mut self) -> Self::Nearest<'_>; - /// Append members and return the resulting operation-local list size. - fn append_to( - &mut self, - list: Self::ListId, - ids: Itr, - ) -> impl SendFuture> + /// 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: Itr, emit: F) -> impl SendFuture> where - Itr: ExactSizeIterator + Clone + Send + Sync; + Itr: ExactSizeIterator + Clone + Send + Sync, + 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. @@ -112,24 +109,13 @@ pub trait BuildAccessor: Send + Sized + HasIds { pub trait BuildStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { /// [`BuildAccessor`] used to insert new points, and performs splits if necessary. - type Accessor: SetElements - + BuildAccessor; + type Accessor: BuildAccessor; type Error: StandardError; fn build_accessor(&'a self, context: &'a C) -> Result; } -pub trait SetElements { - type Id: VectorId; - - fn set_elements( - &mut self, - context: &C, - elements: &[T], - ) -> impl std::future::Future>> + Send; -} - //////////// // Search // //////////// From a6e53c927f48f8c8171b13c1a5e0c557f45fa547 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Thu, 22 Oct 2026 17:42:46 -0400 Subject: [PATCH 29/33] refactor modules a bit --- diskann/src/ivf/build/assignment.rs | 120 ++ diskann/src/ivf/build/clustering.rs | 76 + diskann/src/ivf/build/mod.rs | 12 + diskann/src/ivf/build/split.rs | 130 ++ diskann/src/ivf/clustering.rs | 295 --- diskann/src/ivf/index.rs | 29 +- diskann/src/ivf/mod.rs | 19 +- diskann/src/ivf/partition.rs | 14 - diskann/src/ivf/{tests.rs => tests/mod.rs} | 2143 +++++++------------- diskann/src/ivf/{ => tests}/pending.rs | 2 +- diskann/src/ivf/tests/support.rs | 724 +++++++ diskann/src/ivf/traits.rs | 11 +- 12 files changed, 1811 insertions(+), 1764 deletions(-) create mode 100644 diskann/src/ivf/build/assignment.rs create mode 100644 diskann/src/ivf/build/clustering.rs create mode 100644 diskann/src/ivf/build/mod.rs create mode 100644 diskann/src/ivf/build/split.rs delete mode 100644 diskann/src/ivf/clustering.rs delete mode 100644 diskann/src/ivf/partition.rs rename diskann/src/ivf/{tests.rs => tests/mod.rs} (70%) rename diskann/src/ivf/{ => tests}/pending.rs (98%) create mode 100644 diskann/src/ivf/tests/support.rs diff --git a/diskann/src/ivf/build/assignment.rs b/diskann/src/ivf/build/assignment.rs new file mode 100644 index 0000000000..17c2e966f2 --- /dev/null +++ b/diskann/src/ivf/build/assignment.rs @@ -0,0 +1,120 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Assign regional memberships and record the corresponding updates. + +use std::collections::BTreeSet; + +use diskann_utils::{Reborrow, views::rowmajor::Matrix}; +use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; + +use crate::{ + ANNError, ANNResult, + error::IntoANNResult, + ivf::region::{AsMatRef, RegionUpdate, RegionView}, +}; + +/// Assign the whole region, allowing policies that coordinate multiple sources. +pub trait Assigner> { + /// Assign each source member to a destination and record complete memberships. + /// + /// Create new destinations, replace retained memberships, and retire sources + /// not retained as destinations. Validate and compute assignments before + /// recording writes. A write error may leave partial private updates. + fn assign( + &mut self, + region: &R, + update: &mut U, + sources: &[L], + destinations: &[Target<'_, L>], + ) -> ANNResult<()> + where + U: RegionUpdate; +} + +pub enum Type { + /// Existing centroid. + Exists(L), + /// New centroid as a result of the split. + New, +} + +/// A borrowed centroid; only retained destinations have an existing list ID. +pub struct Target<'a, L> { + pub list: Type, + pub centroid: &'a [f32], +} + +pub struct RegionalNearest; + +impl Assigner for RegionalNearest +where + R: RegionView, + for<'a> R::VectorsRef<'a>: AsMatRef, +{ + fn assign( + &mut self, + region: &R, + update: &mut U, + sources: &[L], + destinations: &[Target<'_, L>], + ) -> ANNResult<()> + where + U: RegionUpdate, + { + let mut memberships: Vec> = (0..destinations.len()).map(|_| Vec::new()).collect(); + + for &src in sources { + let view = region.get_view(src).ok_or_else(|| { + ANNError::message("IVF assignment source is missing from the region") + })?; + let vectors = region.get_vectors(src).ok_or_else(|| { + ANNError::message("IVF assignment source vectors are missing from the region") + })?; + let vectors = vectors.reborrow(); + let vectors = vectors.as_view().into_ann_result()?; + + for (&id, point) in view.members.iter().zip(vectors.rows()) { + let dst = nearest(point, destinations.iter().map(|d| d.centroid)); + memberships[dst].push(id); + } + } + + let mut retained = BTreeSet::new(); + + for (destination, members) in destinations.iter().zip(memberships) { + match &destination.list { + Type::Exists(id) => { + update.set_members(*id, &members)?; + retained.insert(*id); + } + Type::New => update.create(destination.centroid, &members)?, + } + } + + for &source in sources { + if !retained.contains(&source) { + update.retire(source)?; + } + } + + Ok(()) + } +} + +/// Prefer the first centroid on ties. +fn nearest<'a>(point: &[f32], targets: impl Iterator) -> usize { + let mut min = f32::MAX; + let mut best = usize::MAX; + for (i, center) in targets.enumerate() { + let dist = SquaredL2::evaluate(point, center); + if min > dist { + best = i; + min = dist; + } + } + + best +} diff --git a/diskann/src/ivf/build/clustering.rs b/diskann/src/ivf/build/clustering.rs new file mode 100644 index 0000000000..1a379c0ef7 --- /dev/null +++ b/diskann/src/ivf/build/clustering.rs @@ -0,0 +1,76 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Fit replacement centroids over provider vector views. + +use std::num::NonZeroUsize; + +use diskann_quantization::algorithms::kmeans::lloyds::lloyds; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; + +use crate::{ANNError, ANNResult, error::IntoANNResult, ivf::region::AsMatRef}; + +/// Fit a parent's replacement centroids without assigning regional memberships. +pub trait Clusterer { + /// Fill `dst` with finite centroids of the canonical dimension. + /// + /// Implementations own initialization, convergence, empty-cluster handling, + /// and reusable fitting scratch. `dst` has `num_centers()` rows and `dim()` + /// columns. The number of centers must be at least two. + fn fit(&mut self, points: V, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; + + fn num_centers(&self) -> NonZeroUsize; + + fn dim(&self) -> NonZeroUsize; +} + +/// Random input seeds followed by Lloyd iterations, which reset empty centers to zero. +pub struct TwoMeans { + iterations: NonZeroUsize, + dim: NonZeroUsize, +} + +impl TwoMeans { + pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { + Self { iterations, dim } + } +} + +impl> Clusterer for TwoMeans { + fn fit(&mut self, points: B, mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + let dim = self.dim.get(); + let num_centers = >::num_centers(self).get(); + + let vectors = points.as_view().into_ann_result()?; + + if vectors.nrows() < num_centers { + return Err(ANNError::message("an IVF split needs at least two points")); + } + if vectors.ncols() != dim || dst.ncols() != dim || dst.nrows() != num_centers { + return Err(ANNError::message( + "two-means needs canonical input rows and a two-row destination of the canonical dimension", + )); + } + + dst.as_mut_slice().copy_from_slice( + vectors + .subview(0..num_centers) + .ok_or(ANNError::message("oh no!"))? + .as_slice(), + ); + + let _ = lloyds(vectors, dst, self.iterations.get()); + + Ok(()) + } + + fn num_centers(&self) -> NonZeroUsize { + const { crate::utils::TWO } + } + + fn dim(&self) -> NonZeroUsize { + self.dim + } +} diff --git a/diskann/src/ivf/build/mod.rs b/diskann/src/ivf/build/mod.rs new file mode 100644 index 0000000000..80a22b46f3 --- /dev/null +++ b/diskann/src/ivf/build/mod.rs @@ -0,0 +1,12 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Numerical fitting and regional assignment for IVF construction. + +pub mod assignment; +pub mod clustering; +mod split; + +pub use split::SplitAndAssign; diff --git a/diskann/src/ivf/build/split.rs b/diskann/src/ivf/build/split.rs new file mode 100644 index 0000000000..a01dc61c9e --- /dev/null +++ b/diskann/src/ivf/build/split.rs @@ -0,0 +1,130 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Coordinate per-parent fitting and whole-region assignment. + +use std::{collections::BTreeSet, num::NonZeroUsize}; + +use diskann_utils::{ + Reborrow, + views::rowmajor::{self, Matrix, MatrixMut}, +}; + +use super::{ + assignment::{Assigner, Target, Type}, + clustering::Clusterer, +}; +use crate::{ + ANNError, ANNResult, + ivf::region::{RegionUpdate, RegionView}, +}; + +pub struct SplitAndAssign { + pub clusterer: C, + pub assigner: A, + scratch: Option>, // scratch for the centroids. +} + +impl SplitAndAssign { + pub fn new(clusterer: C, assigner: A) -> Self { + Self { + clusterer, + assigner, + scratch: None, + } + } + + /// Initialize the scratch for a particular number of output cluster centroids. + fn initialize(&mut self, num_jobs: NonZeroUsize) -> ANNResult + where + C: Clusterer, + { + let c = self.clusterer.num_centers(); + + let rows = num_jobs + .checked_mul(c) + .ok_or_else(|| ANNError::message("IVF centroid count overflow"))? + .get(); + + let dim = self.clusterer.dim().get(); + + self.scratch = Some(rowmajor::Owned::try_from_element(rows, dim, f32::NAN)?); + + Ok(c) + } + + /// Compute the complete result before registering any provider writes. + /// + /// Concrete-reference clusterers may need an explicit `R` type argument because + /// of a compiler limitation when inferring higher-ranked associated types. + pub fn run( + &mut self, + region: &R, + update: &mut U, + parents: &BTreeSet, + neighbors: &BTreeSet, + ) -> ANNResult<()> + where + A: Assigner, + R: RegionView, + C: for<'a> Clusterer>, + U: RegionUpdate, + { + let n_centroids = self + .initialize::>( + NonZeroUsize::new(parents.len()) + .ok_or_else(|| ANNError::message("IVF splitting needs at least one parent"))?, + )? + .get(); + let mut s = 0; + + let mut clusters = self + .scratch + .as_mut() + .ok_or_else(|| ANNError::message("IVF centroid scratch was not initialized"))? + .as_view_mut(); + + for &p in parents { + let vecs = region.get_vectors(p).ok_or_else(|| { + ANNError::message("IVF split source vectors are missing from the region") + })?; + + let end = s + n_centroids; + + self.clusterer.fit( + vecs.reborrow(), + clusters.subview_mut(s..end).ok_or_else(|| { + ANNError::message("IVF centroid scratch range is out of bounds") + })?, + )?; + + s = end; + } + + let mut destinations: Vec<_> = clusters + .rows() + .map(|centroid| Target { + list: Type::New, + centroid, + }) + .collect(); + + for &n in neighbors { + let view = region + .get_view(n) + .ok_or_else(|| ANNError::message("IVF neighbor is missing from the region"))?; + + destinations.push(Target { + list: Type::Exists(n), + centroid: view.centroid, + }); + } + + let sources = parents.iter().chain(neighbors).copied().collect::>(); + + self.assigner + .assign(region, update, &sources, &destinations) + } +} diff --git a/diskann/src/ivf/clustering.rs b/diskann/src/ivf/clustering.rs deleted file mode 100644 index c6c0219f9a..0000000000 --- a/diskann/src/ivf/clustering.rs +++ /dev/null @@ -1,295 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Replaceable child fitting and regional assignment over borrowed point rows. -//! -//! Fitting writes new centroids into reusable scratch. Assignment borrows the -//! destinations, groups source members, and records regional writes directly. - -use std::{collections::BTreeSet, 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::region::{RegionUpdate, RegionView}; -use crate::{ANNError, ANNResult, error::IntoANNResult, ivf::region::AsMatRef}; - -/// Fit a parent's replacement centroids without assigning regional memberships. -pub trait Clusterer { - /// Fill `dst` with finite centroids of the canonical dimension. - /// - /// Implementations own initialization, convergence, empty-cluster handling, - /// and reusable fitting scratch. `dst` has `num_centers()` rows and `dim()` - /// columns. The number of centers must be at least two. - fn fit(&mut self, points: V, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; - - fn num_centers(&self) -> NonZeroUsize; - - fn dim(&self) -> NonZeroUsize; -} - -/// Assign the whole region, allowing policies that coordinate multiple sources. -pub trait Assigner> { - /// Assign each source member to a destination and record complete memberships. - /// - /// Create new destinations, replace retained memberships, and retire sources - /// not retained as destinations. Validate and compute assignments before - /// recording writes. A write error may leave partial private updates. - fn assign( - &mut self, - region: &R, - update: &mut U, - sources: &[L], - destinations: &[Target<'_, L>], - ) -> ANNResult<()> - where - U: RegionUpdate; -} - -pub enum Type { - /// Existing centroid. - Exists(L), - /// New centroid as a result of the split. - New, -} - -/// A borrowed centroid; only retained destinations have an existing list ID. -pub struct Target<'a, L> { - pub list: Type, - pub centroid: &'a [f32], -} - -pub struct SplitAndAssign { - pub clusterer: C, - pub assigner: A, - scratch: Option>, // scratch for the centroids. -} - -impl SplitAndAssign { - pub fn new(clusterer: C, assigner: A) -> Self { - Self { - clusterer, - assigner, - scratch: None, - } - } - - /// Initialize the scratch for a particular number of output cluster centroids. - fn initialize(&mut self, num_jobs: NonZeroUsize) -> ANNResult - where - C: Clusterer, - { - let c = self.clusterer.num_centers(); - - let rows = num_jobs - .checked_mul(c) - .ok_or_else(|| ANNError::message("IVF centroid count overflow"))? - .get(); - - let dim = self.clusterer.dim().get(); - - self.scratch = Some(rowmajor::Owned::try_from_element(rows, dim, f32::NAN)?); - - Ok(c) - } - - /// Compute the complete result before registering any provider writes. - /// - /// Concrete-reference clusterers may need an explicit `R` type argument because - /// of a compiler limitation when inferring higher-ranked associated types. - pub fn run( - &mut self, - region: &R, - update: &mut U, - parents: &BTreeSet, - neighbors: &BTreeSet, - ) -> ANNResult<()> - where - A: Assigner, - R: RegionView, - C: for<'a> Clusterer>, - U: RegionUpdate, - { - let n_centroids = self - .initialize::>( - NonZeroUsize::new(parents.len()) - .ok_or_else(|| ANNError::message("IVF splitting needs at least one parent"))?, - )? - .get(); - let mut s = 0; - - let mut clusters = self - .scratch - .as_mut() - .ok_or_else(|| ANNError::message("IVF centroid scratch was not initialized"))? - .as_view_mut(); - - for &p in parents { - let vecs = region.get_vectors(p).ok_or_else(|| { - ANNError::message("IVF split source vectors are missing from the region") - })?; - - let end = s + n_centroids; - - self.clusterer.fit( - vecs.reborrow(), - clusters.subview_mut(s..end).ok_or_else(|| { - ANNError::message("IVF centroid scratch range is out of bounds") - })?, - )?; - - s = end; - } - - let mut destinations: Vec<_> = clusters - .rows() - .map(|centroid| Target { - list: Type::New, - centroid, - }) - .collect(); - - for &n in neighbors { - let view = region - .get_view(n) - .ok_or_else(|| ANNError::message("IVF neighbor is missing from the region"))?; - - destinations.push(Target { - list: Type::Exists(n), - centroid: view.centroid, - }); - } - - let sources = parents.iter().chain(neighbors).copied().collect::>(); - - self.assigner - .assign(region, update, &sources, &destinations) - } -} - -/// Random input seeds followed by Lloyd iterations, which reset empty centers to zero. -pub struct TwoMeans { - iterations: NonZeroUsize, - dim: NonZeroUsize, -} - -impl TwoMeans { - pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { - Self { iterations, dim } - } -} - -impl> Clusterer for TwoMeans { - fn fit(&mut self, points: B, mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { - let dim = self.dim.get(); - let num_centers = >::num_centers(self).get(); - - let vectors = points.as_view().into_ann_result()?; - - if vectors.nrows() < num_centers { - return Err(ANNError::message("an IVF split needs at least two points")); - } - if vectors.ncols() != dim || dst.ncols() != dim || dst.nrows() != num_centers { - return Err(ANNError::message( - "two-means needs canonical input rows and a two-row destination of the canonical dimension", - )); - } - - dst.as_mut_slice().copy_from_slice( - vectors - .subview(0..num_centers) - .ok_or(ANNError::message("oh no!"))? - .as_slice(), - ); - - let _ = lloyds(vectors, dst, self.iterations.get()); - - Ok(()) - } - - fn num_centers(&self) -> NonZeroUsize { - const { crate::utils::TWO } - } - - fn dim(&self) -> NonZeroUsize { - self.dim - } -} - -pub struct RegionalNearest; - -impl Assigner for RegionalNearest -where - R: RegionView, - for<'a> R::VectorsRef<'a>: AsMatRef, -{ - fn assign( - &mut self, - region: &R, - update: &mut U, - sources: &[L], - destinations: &[Target<'_, L>], - ) -> ANNResult<()> - where - U: RegionUpdate, - { - let mut memberships: Vec> = (0..destinations.len()).map(|_| Vec::new()).collect(); - - for &src in sources { - let view = region.get_view(src).ok_or_else(|| { - ANNError::message("IVF assignment source is missing from the region") - })?; - let vectors = region.get_vectors(src).ok_or_else(|| { - ANNError::message("IVF assignment source vectors are missing from the region") - })?; - let vectors = vectors.reborrow(); - let vectors = vectors.as_view().into_ann_result()?; - - for (&id, point) in view.members.iter().zip(vectors.rows()) { - let dst = nearest(point, destinations.iter().map(|d| d.centroid)); - memberships[dst].push(id); - } - } - - let mut retained = BTreeSet::new(); - - for (destination, members) in destinations.iter().zip(memberships) { - match &destination.list { - Type::Exists(id) => { - update.set_members(*id, &members)?; - retained.insert(*id); - } - Type::New => update.create(destination.centroid, &members)?, - } - } - - for &source in sources { - if !retained.contains(&source) { - update.retire(source)?; - } - } - - Ok(()) - } -} - -/// Prefer the first centroid on ties. -fn nearest<'a>(point: &[f32], targets: impl Iterator) -> usize { - let mut min = f32::MAX; - let mut best = usize::MAX; - for (i, center) in targets.enumerate() { - let dist = SquaredL2::evaluate(point, center); - if min > dist { - best = i; - min = dist; - } - } - - best -} diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 3975fccdc8..aaa2b3dfc9 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -17,7 +17,7 @@ use crate::{ flat::DistancesUnordered, graph::SearchOutputBuffer, ivf::{ - clustering::{RegionalNearest, SplitAndAssign, TwoMeans}, + build::{SplitAndAssign, assignment::RegionalNearest, clustering::TwoMeans}, region::AsMatRef, traits::BuildStrategy, }, @@ -75,33 +75,27 @@ impl IVFIndex { }) } - pub fn insert_batch<'a, S, T, Itr>( + pub fn insert_batch<'a, S, B>( &'a mut self, strategy: &'a S, context: &'a C, - points: Itr, + points: B, ) -> impl SendFuture> where - Itr: ExactSizeIterator + Clone + Send + Sync, - S: BuildStrategy<'a, T, C>, - S::Accessor: for<'b> BuildAccessor: AsMatRef>, - T: Sync, + S: BuildStrategy<'a, B, C>, + S::Accessor: for<'b> BuildAccessor: AsMatRef>, + B: Send + Sync, { let config = self.config; async move { - let inserted = points.len(); - if inserted == 0 { - return Ok(InsertStats::default()); - } - let mut accessor = strategy.build_accessor(context).into_ann_result()?; let dim = accessor.dim(); let mut parents = BTreeSet::new(); - accessor + let inserted = accessor .route(points, |&list, size| { if size.get() > config.split_threshold.get() { parents.insert(list); @@ -109,8 +103,6 @@ impl IVFIndex { }) .await?; - let splits = parents.len(); - if !parents.is_empty() { let mut neighbors = BTreeSet::new(); @@ -146,7 +138,10 @@ impl IVFIndex { accessor.finish().await?; - Ok(InsertStats { inserted, splits }) + Ok(InsertStats { + inserted: inserted.get(), + splits: parents.len(), + }) } } @@ -170,7 +165,7 @@ impl IVFIndex { { let router = accessor.router(); router - .route(&query, nprobe, |id, scan| lists.push((*id, scan))) + .route(query, nprobe, |id, scan| lists.push((*id, scan))) .await?; } diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 6919fce5bd..0e1506cb2f 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -3,23 +3,24 @@ * Licensed under the MIT license. */ -//! Incremental IVF construction. +//! 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::RegionView`], //! and regional writes through [`region::RegionUpdate`]. //! -//! Neighbor selection uses [`traits::RouterAndNearest`] before filling the unique +//! 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. -//! The numerical driver separates per-parent fitting from whole-region assignment. +//! [`build::SplitAndAssign`] separates per-parent fitting from whole-region assignment. //! Fitting writes a fixed number of children per parent into reusable centroid //! scratch. Assignment borrows those rows and retained centroids, groups members, //! and records writes directly through the regional update object. -//! The index currently selects two-means and regional-nearest assignment. +//! The index currently selects [`build::clustering::TwoMeans`] and +//! [`build::assignment::RegionalNearest`]. //! //! Clustering finishes reading its inputs before recording child creation and //! parent retirement and neighbor membership replacement through the separate @@ -34,13 +35,11 @@ //! 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, search, deletion, and colocation are not -//! implemented. +//! recursively split. Search scans routed lists and collects their nearest points. +//! Initialization, deletion, and colocation are not implemented. -pub mod clustering; -mod index; -pub mod partition; -pub mod pending; +pub mod build; +pub mod index; pub mod region; pub mod traits; diff --git a/diskann/src/ivf/partition.rs b/diskann/src/ivf/partition.rs deleted file mode 100644 index e4ee137568..0000000000 --- a/diskann/src/ivf/partition.rs +++ /dev/null @@ -1,14 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Complete destination memberships produced by regional clustering. - -/// Complete membership of a new or retained destination list. -#[derive(Debug)] -pub struct Partition { - /// `Some` retains a source list's identity and centroid; `None` creates a new list. - pub id: Option, - pub members: Vec, -} diff --git a/diskann/src/ivf/tests.rs b/diskann/src/ivf/tests/mod.rs similarity index 70% rename from diskann/src/ivf/tests.rs rename to diskann/src/ivf/tests/mod.rs index b0706dab3a..25c52525c4 100644 --- a/diskann/src/ivf/tests.rs +++ b/diskann/src/ivf/tests/mod.rs @@ -18,9 +18,13 @@ use diskann_utils::{ use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use tokio::sync::{Mutex, MutexGuard}; -use super::{ - Config, IVFIndex, InsertStats, ReassignCount, clustering, +use self::{ pending::{Membership, PendingLists}, + support::*, +}; +use super::{ + Config, IVFIndex, InsertStats, ReassignCount, + build::{SplitAndAssign, assignment, clustering}, region::{AsMatRef, RegionUpdate, RegionView, View}, traits::{ BuildAccessor, BuildStrategy, HasIds, NearestList, Router, SearchAccessor, SearchStrategy, @@ -35,720 +39,683 @@ use crate::{ utils::ONE, }; -#[derive(Debug, Clone)] -struct StoredList { - centroid: Vec, - members: Vec, -} - -#[derive(Debug, Default)] -struct Store { - dim: usize, - points: BTreeMap>, - lists: BTreeMap, - next_list: u32, - finishes: usize, - fills: usize, - filled_lists: Vec, - fail_route: bool, - fail_nearest: bool, - fail_fill: bool, - fail_create: bool, - fail_set_members: bool, - fail_retire: bool, - fail_finish: bool, -} - -impl Store { - 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")) - } +mod pending; +mod support; - fn to_external_id(&self, _: &DefaultContext, id: u32) -> ANNResult { - if !self.points.contains_key(&id) { - return Err(ANNError::message("point is not visible")); +#[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( + (9, &[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); } - id.checked_sub(10_000) - .ok_or_else(|| ANNError::message("invalid internal point ID")) } } -struct Build<'p, T> { - store: MutexGuard<'p, Store>, - pending: PendingLists, - created: Vec, - staged: BTreeMap>, - sources: BTreeMap, - fills: Cell, - filled_lists: Vec, -} - -struct Coarse<'a> { - stored: &'a BTreeMap, - fail_route: bool, - 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() +#[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); } } -impl + Sync> Router<(u32, &[T])> for Coarse<'_> { - type Id = u32; - type Ret = f32; - - fn len(&self) -> usize { - self.stored.len() - } - - fn route( - &self, - query: &(u32, &[T]), - nprobe: NonZeroUsize, - emit: F, - ) -> impl SendFuture> - where - F: FnMut(&Self::Id, f32) + Send, - { - async move { - if self.fail_route { - return Err(ANNError::message("injected routing failure")); - } - self.select(query.1, nprobe, |_| true, emit); - Ok(()) - } +#[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()); } } -impl + Sync> Router<&[T]> for Coarse<'_> { - type Id = u32; - type Ret = f32; - - fn len(&self) -> usize { - self.stored.len() +#[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)); } - - fn route( - &self, - query: &&[T], - nprobe: NonZeroUsize, - emit: F, - ) -> impl SendFuture> - where - F: FnMut(&Self::Id, f32) + Send, - { - async move { - if self.fail_route { - return Err(ANNError::message("injected routing failure")); - } - self.select(query, nprobe, |_| true, emit); - Ok(()) - } + let router = build.nearest(); + for (point, expected) in points.iter().zip([2, 1, 1]) { + let mut selected = Vec::new(); + require_send(router.route(*point, 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); } -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)) - } +#[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); } -} - -struct Source { - centroid: Vec, - members: Vec, - vectors: rowmajor::Owned, -} + 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); -struct Update<'a> { - dim: usize, - sources: &'a BTreeMap, - memberships: &'a mut PendingLists, - created: &'a mut Vec, - fail_create: bool, - fail_set_members: bool, - fail_retire: bool, + 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); } -impl<'a> Update<'a> { - 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, +#[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(); } - } - - fn check_list(&self, list: u32) -> ANNResult<()> { - if self.sources.contains_key(&list) { - Ok(()) + 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 { - Err(ANNError::message(format!("unknown list {list}"))) - } + "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); } } -struct Region<'a> { - sources: &'a BTreeMap, -} - -impl<'a> Region<'a> { - fn new(sources: &'a BTreeMap) -> Self { - Self { sources } +#[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); } -impl RegionView for Region<'_> { - type VectorsRef<'a> = rowmajor::Ref<'a, f32>; - type Vectors<'a> - = rowmajor::Ref<'a, f32> - where - Self: 'a; - - fn get_view(&self, list: u32) -> Option> { - self.sources.get(&list).map(|source| View { - 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()) +#[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); } } -impl RegionUpdate 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) - } +#[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(); - 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")); + assert_eq!( + stats, + InsertStats { + inserted: 3, + splits: 0 } - self.created.push(StoredList { - centroid: centroid.to_vec(), - members: members.to_vec(), - }); - Ok(()) - } -} - -impl HasIds for Build<'_, T> { - type Id = u32; - type ListId = u32; + ); + 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); } -impl<'p, 'v, T: Copy + Into + Send + Sync> BuildAccessor<(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() +#[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); +} - fn nearest(&mut self) -> Self::Nearest<'_> { - Coarse { - stored: &self.store.lists, - fail_route: self.store.fail_route, - fail_nearest: self.store.fail_nearest, - } - } +#[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(); - fn route(&mut self, points: Itr, mut emit: F) -> impl SendFuture> - where - Itr: ExactSizeIterator + Clone + Send + Sync, - F: FnMut(&Self::ListId, NonZeroUsize) + Send, - { - async move { - let mut staged = BTreeMap::new(); - let mut routed = BTreeMap::>::new(); - for point in points { - let (external, vector) = point; - 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(&point, 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(()) + 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::>()); - 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)) + for (&id, list) in &store.lists { + if id == 3 { + continue; } - } - - 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(); - } + 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); } } - self.store.next_list = next_list; - self.store.finishes += 1; - self.store.fills += self.fills.get(); - self.store.filled_lists = self.filled_lists; - Ok(()) } } } -struct Strategy { - store: Mutex, -} +#[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(); -impl Strategy { - fn new(store: Store) -> Self { - Self { - store: Mutex::new(store), + assert_eq!( + stats, + InsertStats { + inserted: 2, + splits: 1 } - } - - 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(), - }) - } -} - -struct Search<'p, 'q, T> { - store: MutexGuard<'p, Store>, - query: &'q [T], -} - -struct SearchRouter<'a, T> { - store: &'a Store, - query: &'a [T], -} - -struct ListScan<'a, T> { - store: &'a Store, - list: u32, - query: &'a [T], + ); + 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::>()); } -impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { - type Id = u32; - type Ret = ListScan<'a, T>; - - fn len(&self) -> usize { - self.store.lists.len() - } - - fn route( - &self, - query: &&[T], - nprobe: NonZeroUsize, - mut emit: F, - ) -> impl SendFuture> - where - F: FnMut(&Self::Id, Self::Ret) + Send, - { - 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, - }; - coarse.select( - query, - nprobe, - |_| true, - |id, _| { - emit( +#[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, - ListScan { - store: self.store, - list: *id, - query: self.query, + StoredList { + centroid: vec![center, 0.0], + members: Vec::new(), }, ); - }, - ); - Ok(()) - } - } -} - -impl HasId for ListScan<'_, T> { - type Id = u32; -} - -impl + Sync> DistancesUnordered for ListScan<'_, T> { - type Error = ANNError; - - fn distances_unordered(&mut self, mut emit: F) -> impl SendFuture> - where - F: FnMut(Self::Id, f32) + Send, - { - async move { - let list = self - .store - .lists - .get(&self.list) - .ok_or_else(|| ANNError::message(format!("unknown list {}", self.list)))?; - for id in &list.members { - let vector = self - .store - .points - .get(id) - .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; - emit(*id, squared_distance(self.query, vector)); + } } - Ok(()) + 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::>()); } } } -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> { - type Router<'a> - = SearchRouter<'a, T> - where - Self: 'a; - type Dist<'a> - = ListScan<'a, T> - where - Self: 'a; - - fn router(&self) -> Self::Router<'_> { - SearchRouter { - store: &self.store, - query: self.query, +#[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_view(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); + SplitAndAssign::new( + clustering::TwoMeans::new( + NonZeroUsize::new(5).unwrap(), + NonZeroUsize::new(2).unwrap(), + ), + assignment::RegionalNearest, + ) + .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) + .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); } -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 }) - } -} +#[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(); -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() -} + { + let (region, mut update) = require_send(build.fill([1, 2].into_iter())).await.unwrap(); + let parent = region.get_view(1).unwrap(); + let neighbor = region.get_view(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()); -fn index(store: Store, threshold: usize) -> (IVFIndex, Strategy) { - index_with_neighbors(store, threshold, 0) -} + 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(); -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)) + 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); } -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() +#[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()); } } -fn require_send(future: F) -> F { - future +#[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); } #[tokio::test] @@ -876,569 +843,89 @@ async fn search_returns_no_results_for_empty_regions() { [(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"), - ] { - let mut store = seeded(); - if case == "route" { - store.fail_route = true; - } - if case == "scan" { - store.points.remove(&10_003); - } - 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 mut output = vec![Neighbor::new(42, -1.0)]; - let error = require_send(index.search( - &strategy, - &DefaultContext, - query, - NonZeroUsize::new(3).unwrap(), - ONE, - &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.iter().copied()) - .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)] - ); -} - -#[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( - &(9, &[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(); - require_send(build.route(points.iter().copied(), |&list, size| { - emitted.push((list, size.get())); - })) - .await - .unwrap(); - 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, 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 { - require_send(build.route(std::iter::once(point), |&list, size| { - emitted.push((list, size.get())); - })) - .await - .unwrap(); - } - 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.into_iter(), |&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.iter().copied())) - .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.iter().copied(), |&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.iter().copied())) - .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); - - let empty: [(u32, &[f32]); 0] = []; - assert_eq!( - index - .insert_batch(&strategy, &DefaultContext, empty.into_iter()) - .await - .unwrap(), - InsertStats::default() - ); - assert_eq!(store.finishes, 1); -} - -#[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.iter().copied())) - .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.iter().copied())) - .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.iter().copied(), - )) - .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 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"), + ] { + let mut store = seeded(); + if case == "route" { + store.fail_route = true; + } + if case == "scan" { + store.points.remove(&10_003); } + 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 mut output = vec![Neighbor::new(42, -1.0)]; + let error = require_send(index.search( + &strategy, + &DefaultContext, + query, + NonZeroUsize::new(3).unwrap(), + ONE, + &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 = [ @@ -1480,9 +967,9 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - clustering::SplitAndAssign::new( + SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), - clustering::RegionalNearest, + assignment::RegionalNearest, ) .run( ®ion, @@ -1513,192 +1000,6 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { assert!(changes.is_empty()); } -#[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 mut emitted = Vec::new(); - build - .route(points.iter().rev().copied(), |&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_view(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); - clustering::SplitAndAssign::new( - clustering::TwoMeans::new( - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - ), - clustering::RegionalNearest, - ) - .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) - .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.iter().copied(), |_, _| {}) - .await - .unwrap(); - - { - let (region, mut update) = require_send(build.fill([1, 2].into_iter())).await.unwrap(); - let parent = region.get_view(1).unwrap(); - let neighbor = region.get_view(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.iter().copied())) - .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.iter().copied()) - .await - .unwrap_err(); - assert!( - error - .to_string() - .contains("IVF routing returned no live centroid"), - "{error}" - ); - assert_eq!(strategy.store().finishes, 0); -} - #[tokio::test] async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() { let mut store = seeded(); @@ -1712,7 +1013,7 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() let points = [(5, &[5.0, 0.0][..])]; let mut emitted = Vec::new(); build - .route(points.iter().copied(), |&list, size| { + .route(&points[..], |&list, size| { emitted.push((list, size.get())); }) .await @@ -1722,12 +1023,12 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() { let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); assert!( - clustering::SplitAndAssign::new( + SplitAndAssign::new( clustering::TwoMeans::new( NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(), ), - clustering::RegionalNearest, + assignment::RegionalNearest, ) .run( ®ion, @@ -1738,12 +1039,12 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() .is_err() ); assert!( - clustering::SplitAndAssign::new( + SplitAndAssign::new( clustering::TwoMeans::new( NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(), ), - clustering::RegionalNearest, + assignment::RegionalNearest, ) .run( ®ion, @@ -1758,9 +1059,9 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() assert!( matches!(changes.as_slice(), [(1, Membership::Append(members))] if members == &[10_005]) ); - clustering::SplitAndAssign::new( + SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()), - clustering::RegionalNearest, + assignment::RegionalNearest, ) .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) .unwrap(); @@ -1811,12 +1112,12 @@ fn assignment_can_cross_parent_boundaries() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - clustering::SplitAndAssign::new( + SplitAndAssign::new( clustering::TwoMeans::new( NonZeroUsize::new(10).unwrap(), NonZeroUsize::new(1).unwrap(), ), - clustering::RegionalNearest, + assignment::RegionalNearest, ) .run( ®ion, @@ -1848,9 +1149,9 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - clustering::SplitAndAssign::new( + SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(1).unwrap(), NonZeroUsize::new(1).unwrap()), - clustering::RegionalNearest, + assignment::RegionalNearest, ) .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) .unwrap(); @@ -1890,7 +1191,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { owners: BTreeMap>, } - impl clustering::Assigner for ParentLocal + impl assignment::Assigner for ParentLocal where R: for<'a> RegionView = rowmajor::Ref<'a, f32>>, { @@ -1899,7 +1200,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { region: &R, update: &mut U, sources: &[u32], - destinations: &[clustering::Target<'_, u32>], + destinations: &[assignment::Target<'_, u32>], ) -> ANNResult<()> where U: RegionUpdate, @@ -1920,7 +1221,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { .iter() .enumerate() .filter_map(|(index, destination)| { - matches!(&destination.list, clustering::Type::Exists(id) if *id == source) + matches!(&destination.list, assignment::Type::Exists(id) if *id == source) .then_some(index) }) .collect(); @@ -1949,8 +1250,8 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { } for (destination, members) in destinations.iter().zip(memberships) { match &destination.list { - clustering::Type::New => update.create(destination.centroid, &members)?, - clustering::Type::Exists(id) => update.set_members(*id, &members)?, + assignment::Type::New => update.create(destination.centroid, &members)?, + assignment::Type::Exists(id) => update.set_members(*id, &members)?, } } for (parent, _) in children { @@ -2001,7 +1302,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut full_memberships = PendingLists::default(); let mut full_created = Vec::new(); let mut full = Update::new(&sources, 1, &mut full_memberships, &mut full_created); - let mut splitter = clustering::SplitAndAssign::new(RowsAsCenters, clustering::RegionalNearest); + let mut splitter = SplitAndAssign::new(RowsAsCenters, assignment::RegionalNearest); splitter .run::<_, _, Region<'_>, _>(®ion, &mut full, &parents, &neighbors) .unwrap(); @@ -2013,7 +1314,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { &mut restricted_memberships, &mut restricted_created, ); - clustering::SplitAndAssign::new( + SplitAndAssign::new( RowsAsCenters, ParentLocal { owners: [(3, vec![1]), (4, vec![1, 2])].into(), @@ -2114,13 +1415,13 @@ fn numerical_strategy_failures_record_no_updates() { struct InvalidAssignment; - impl> clustering::Assigner for InvalidAssignment { + impl> assignment::Assigner for InvalidAssignment { fn assign( &mut self, _: &R, _: &mut U, _: &[u32], - _: &[clustering::Target<'_, u32>], + _: &[assignment::Target<'_, u32>], ) -> ANNResult<()> where U: RegionUpdate, @@ -2158,7 +1459,7 @@ fn numerical_strategy_failures_record_no_updates() { } else { parents.clone() }; - let error = clustering::SplitAndAssign::new(clusterer, clustering::RegionalNearest) + let error = SplitAndAssign::new(clusterer, assignment::RegionalNearest) .run(®ion, &mut update, &selected, &neighbors) .unwrap_err(); assert!(error.to_string().contains(expected), "{error}"); @@ -2167,7 +1468,7 @@ fn numerical_strategy_failures_record_no_updates() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let error = clustering::SplitAndAssign::new( + let error = SplitAndAssign::new( clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), InvalidAssignment, ) @@ -2268,9 +1569,9 @@ fn split_driver_reuses_clusterer_with_fallible_borrowed_batches() { } } - let mut driver = clustering::SplitAndAssign::new( + let mut driver = SplitAndAssign::new( clustering::TwoMeans::new(ONE, ONE), - clustering::RegionalNearest, + assignment::RegionalNearest, ); for fail_on in [Some(1), Some(2), None] { let sources = BTreeMap::from([( @@ -2377,13 +1678,13 @@ fn split_driver_accepts_native_batches_without_matrix_access() { struct OrderedAssignment; - impl> clustering::Assigner for OrderedAssignment { + impl> assignment::Assigner for OrderedAssignment { fn assign( &mut self, region: &R, update: &mut U, sources: &[u32], - destinations: &[clustering::Target<'_, u32>], + destinations: &[assignment::Target<'_, u32>], ) -> ANNResult<()> where U: RegionUpdate, @@ -2396,7 +1697,7 @@ fn split_driver_accepts_native_batches_without_matrix_access() { } } - let mut driver = clustering::SplitAndAssign::new(NativeCenters, OrderedAssignment); + let mut driver = SplitAndAssign::new(NativeCenters, OrderedAssignment); for offset in [0u8, 20] { let sources = BTreeMap::from([( 1, @@ -2494,7 +1795,7 @@ fn two_means_fits_only_the_destination_window() { #[test] fn regional_assignment_reports_missing_sources_before_writing() { - use clustering::{Assigner, RegionalNearest, Target, Type}; + use assignment::{Assigner, RegionalNearest, Target, Type}; let sources = BTreeMap::from([ ( diff --git a/diskann/src/ivf/pending.rs b/diskann/src/ivf/tests/pending.rs similarity index 98% rename from diskann/src/ivf/pending.rs rename to diskann/src/ivf/tests/pending.rs index 1a0f33a242..7a107363ed 100644 --- a/diskann/src/ivf/pending.rs +++ b/diskann/src/ivf/tests/pending.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Optional buffering for operation-local list memberships. +//! 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 diff --git a/diskann/src/ivf/tests/support.rs b/diskann/src/ivf/tests/support.rs new file mode 100644 index 0000000000..273c73bfb1 --- /dev/null +++ b/diskann/src/ivf/tests/support.rs @@ -0,0 +1,724 @@ +/* + * 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() + } +} + +impl + Sync> Router<(u32, &[T])> for Coarse<'_> { + type Id = u32; + type Ret = f32; + + fn len(&self) -> usize { + self.stored.len() + } + + fn route( + &self, + query: (u32, &[T]), + nprobe: NonZeroUsize, + emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::Id, f32) + Send, + { + async move { + if self.fail_route { + return Err(ANNError::message("injected routing failure")); + } + self.select(query.1, nprobe, |_| true, emit); + Ok(()) + } + } +} + +impl + Sync> Router<&[T]> for Coarse<'_> { + type Id = u32; + type Ret = f32; + + fn len(&self) -> usize { + self.stored.len() + } + + fn route(&self, query: &[T], nprobe: NonZeroUsize, emit: F) -> impl SendFuture> + where + F: FnMut(&Self::Id, 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 RegionView for Region<'_> { + type VectorsRef<'a> = rowmajor::Ref<'a, f32>; + type Vectors<'a> + = rowmajor::Ref<'a, f32> + where + Self: 'a; + + fn get_view(&self, list: u32) -> Option> { + self.sources.get(&list).map(|source| View { + 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 RegionUpdate 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 &point in points { + let (external, vector) = point; + 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(point, 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], +} + +pub(super) struct SearchRouter<'a, T> { + pub(super) store: &'a Store, + pub(super) query: &'a [T], +} + +pub(super) struct ListScan<'a, T> { + pub(super) store: &'a Store, + pub(super) list: u32, + pub(super) query: &'a [T], +} + +impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { + type Id = u32; + type Ret = ListScan<'a, T>; + + fn len(&self) -> usize { + self.store.lists.len() + } + + fn route( + &self, + query: &[T], + nprobe: NonZeroUsize, + mut emit: F, + ) -> impl SendFuture> + where + F: FnMut(&Self::Id, Self::Ret) + Send, + { + 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, + }; + coarse.select( + query, + nprobe, + |_| true, + |id, _| { + emit( + id, + ListScan { + store: self.store, + list: *id, + query: self.query, + }, + ); + }, + ); + Ok(()) + } + } +} + +impl HasId for ListScan<'_, T> { + type Id = u32; +} + +impl + Sync> DistancesUnordered for ListScan<'_, T> { + type Error = ANNError; + + fn distances_unordered(&mut self, mut emit: F) -> impl SendFuture> + where + F: FnMut(Self::Id, f32) + Send, + { + async move { + let list = self + .store + .lists + .get(&self.list) + .ok_or_else(|| ANNError::message(format!("unknown list {}", self.list)))?; + for id in &list.members { + let vector = self + .store + .points + .get(id) + .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; + emit(*id, squared_distance(self.query, vector)); + } + Ok(()) + } + } +} + +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> { + type Router<'a> + = SearchRouter<'a, T> + where + Self: 'a; + type Dist<'a> + = ListScan<'a, T> + where + Self: 'a; + + fn router(&self) -> Self::Router<'_> { + SearchRouter { + store: &self.store, + query: self.query, + } + } +} + +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 index 48902783a9..be113e6209 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -32,7 +32,7 @@ pub trait Router: Send { } /// Select `min(nprobe, len())` distinct live lists, nearest first. - fn route(&self, query: &T, nprobe: NonZeroUsize, emit: F) -> impl SendFuture> + fn route(&self, query: T, nprobe: NonZeroUsize, emit: F) -> impl SendFuture> where F: FnMut(&Self::Id, Self::Ret) + Send; } @@ -58,7 +58,7 @@ pub trait NearestList: Send { F: FnMut(&Self::Id, f32) + Send; } -pub trait BuildAccessor: Send + Sized + HasIds { +pub trait BuildAccessor: Send + Sized + HasIds { type Nearest<'a>: NearestList where Self: 'a; @@ -83,9 +83,8 @@ pub trait BuildAccessor: Send + Sized + HasIds { /// 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: Itr, emit: F) -> impl SendFuture> + fn route(&mut self, points: B, emit: F) -> impl SendFuture> where - Itr: ExactSizeIterator + Clone + Send + Sync, F: FnMut(&Self::ListId, NonZeroUsize) + Send; /// Make requested list's members, data and centroids available @@ -107,9 +106,9 @@ pub trait BuildAccessor: Send + Sized + HasIds { // Strategy // ////////////// -pub trait BuildStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { +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 Accessor: BuildAccessor; type Error: StandardError; From 291bc5b1033dd336c1b255d7365b287d5c0e9a4f Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Fri, 23 Oct 2026 11:01:20 -0400 Subject: [PATCH 30/33] fold clustering into --- diskann/src/ivf/build/clustering.rs | 107 ++++++++++++----- diskann/src/ivf/build/split.rs | 72 +++-------- diskann/src/ivf/index.rs | 4 +- diskann/src/ivf/mod.rs | 5 +- diskann/src/ivf/tests/mod.rs | 180 ++++++++++++++++------------ rfcs/01187-incremental-ivf.md | 17 +-- 6 files changed, 210 insertions(+), 175 deletions(-) diff --git a/diskann/src/ivf/build/clustering.rs b/diskann/src/ivf/build/clustering.rs index 1a379c0ef7..b21ee5ef6d 100644 --- a/diskann/src/ivf/build/clustering.rs +++ b/diskann/src/ivf/build/clustering.rs @@ -8,25 +8,36 @@ use std::num::NonZeroUsize; use diskann_quantization::algorithms::kmeans::lloyds::lloyds; -use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; - -use crate::{ANNError, ANNResult, error::IntoANNResult, ivf::region::AsMatRef}; - -/// Fit a parent's replacement centroids without assigning regional memberships. -pub trait Clusterer { +use diskann_utils::{ + Reborrow, + views::rowmajor::{self, Matrix, MatrixMut}, +}; + +use crate::{ + ANNError, ANNResult, + error::IntoANNResult, + ivf::region::{AsMatRef, RegionView}, +}; + +/// Fit replacement centroids for selected sources without assigning regional memberships. +pub trait Clusterer +where + R: RegionView, +{ /// Fill `dst` with finite centroids of the canonical dimension. /// /// Implementations own initialization, convergence, empty-cluster handling, - /// and reusable fitting scratch. `dst` has `num_centers()` rows and `dim()` - /// columns. The number of centers must be at least two. - fn fit(&mut self, points: V, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; + /// and reusable fitting scratch. `dst` has `num_centers(sources.len())` rows + /// and `dim()` columns. The total number of centers must be at least two. + fn fit(&mut self, sources: &[L], region: &R, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; - fn num_centers(&self) -> NonZeroUsize; + /// Total output row count for this source count, or `None` on overflow. + fn num_centers(&self, sources: usize) -> Option; fn dim(&self) -> NonZeroUsize; } -/// Random input seeds followed by Lloyd iterations, which reset empty centers to zero. +/// First-two-row seeds followed by Lloyd iterations, which reset empty centers to zero. pub struct TwoMeans { iterations: NonZeroUsize, dim: NonZeroUsize, @@ -36,38 +47,76 @@ impl TwoMeans { pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { Self { iterations, dim } } -} -impl> Clusterer for TwoMeans { - fn fit(&mut self, points: B, mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { - let dim = self.dim.get(); - let num_centers = >::num_centers(self).get(); - - let vectors = points.as_view().into_ann_result()?; + const fn inner_centers() -> usize { + 2 + } - if vectors.nrows() < num_centers { - return Err(ANNError::message("an IVF split needs at least two points")); - } - if vectors.ncols() != dim || dst.ncols() != dim || dst.nrows() != num_centers { + fn fit_one( + &self, + points: rowmajor::Ref<'_, f32>, + mut dst: rowmajor::Mut<'_, f32>, + ) -> ANNResult<()> { + if points.nrows() < Self::inner_centers() || points.ncols() != self.dim.get() { return Err(ANNError::message( - "two-means needs canonical input rows and a two-row destination of the canonical dimension", + "IVF two-means needs at least two points with the configured dimension", )); } dst.as_mut_slice().copy_from_slice( - vectors - .subview(0..num_centers) - .ok_or(ANNError::message("oh no!"))? + points + .subview(0..Self::inner_centers()) + .ok_or_else(|| ANNError::message("IVF two-means seed rows are missing"))? .as_slice(), ); - let _ = lloyds(vectors, dst, self.iterations.get()); + let _ = lloyds(points, dst, self.iterations.get()); + + Ok(()) + } +} + +impl Clusterer for TwoMeans +where + R: RegionView, + for<'a> R::VectorsRef<'a>: AsMatRef, +{ + fn fit(&mut self, sources: &[L], region: &R, mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + if dst.ncols() != self.dim.get() + || sources.len().checked_mul(Self::inner_centers()) != Some(dst.nrows()) + { + return Err(ANNError::message( + "IVF two-means scratch needs two rows per source and the configured dimension", + )); + } + + let mut s = 0; + for &p in sources { + let vecs = region.get_vectors(p).ok_or_else(|| { + ANNError::message("IVF split source vectors are missing from the region") + })?; + + let vecs = vecs.reborrow(); + + let points = vecs.as_view().into_ann_result()?; + + let end = s + Self::inner_centers(); + + self.fit_one( + points, + dst.subview_mut(s..end).ok_or_else(|| { + ANNError::message("IVF centroid scratch range is out of bounds") + })?, + )?; + + s = end; + } Ok(()) } - fn num_centers(&self) -> NonZeroUsize { - const { crate::utils::TWO } + fn num_centers(&self, sources: usize) -> Option { + sources.checked_mul(Self::inner_centers()) } fn dim(&self) -> NonZeroUsize { diff --git a/diskann/src/ivf/build/split.rs b/diskann/src/ivf/build/split.rs index a01dc61c9e..ae9c6aed95 100644 --- a/diskann/src/ivf/build/split.rs +++ b/diskann/src/ivf/build/split.rs @@ -5,12 +5,7 @@ //! Coordinate per-parent fitting and whole-region assignment. -use std::{collections::BTreeSet, num::NonZeroUsize}; - -use diskann_utils::{ - Reborrow, - views::rowmajor::{self, Matrix, MatrixMut}, -}; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use super::{ assignment::{Assigner, Target, Type}, @@ -24,7 +19,6 @@ use crate::{ pub struct SplitAndAssign { pub clusterer: C, pub assigner: A, - scratch: Option>, // scratch for the centroids. } impl SplitAndAssign { @@ -32,29 +26,9 @@ impl SplitAndAssign { Self { clusterer, assigner, - scratch: None, } } - /// Initialize the scratch for a particular number of output cluster centroids. - fn initialize(&mut self, num_jobs: NonZeroUsize) -> ANNResult - where - C: Clusterer, - { - let c = self.clusterer.num_centers(); - - let rows = num_jobs - .checked_mul(c) - .ok_or_else(|| ANNError::message("IVF centroid count overflow"))? - .get(); - - let dim = self.clusterer.dim().get(); - - self.scratch = Some(rowmajor::Owned::try_from_element(rows, dim, f32::NAN)?); - - Ok(c) - } - /// Compute the complete result before registering any provider writes. /// /// Concrete-reference clusterers may need an explicit `R` type argument because @@ -63,45 +37,29 @@ impl SplitAndAssign { &mut self, region: &R, update: &mut U, - parents: &BTreeSet, - neighbors: &BTreeSet, + parents: &[L], + neighbors: &[L], ) -> ANNResult<()> where A: Assigner, + C: Clusterer, R: RegionView, - C: for<'a> Clusterer>, U: RegionUpdate, { - let n_centroids = self - .initialize::>( - NonZeroUsize::new(parents.len()) - .ok_or_else(|| ANNError::message("IVF splitting needs at least one parent"))?, - )? - .get(); - let mut s = 0; - - let mut clusters = self - .scratch - .as_mut() - .ok_or_else(|| ANNError::message("IVF centroid scratch was not initialized"))? - .as_view_mut(); - - for &p in parents { - let vecs = region.get_vectors(p).ok_or_else(|| { - ANNError::message("IVF split source vectors are missing from the region") - })?; + if parents.is_empty() { + return Err(ANNError::message("IVF splitting needs at least one parent")); + } - let end = s + n_centroids; + let rows = self + .clusterer + .num_centers(parents.len()) + .ok_or_else(|| ANNError::message("IVF centroid count overflow"))?; - self.clusterer.fit( - vecs.reborrow(), - clusters.subview_mut(s..end).ok_or_else(|| { - ANNError::message("IVF centroid scratch range is out of bounds") - })?, - )?; + let mut clusters = + rowmajor::Owned::try_from_element(rows, self.clusterer.dim().get(), f32::NAN)?; - s = end; - } + self.clusterer + .fit(parents, region, clusters.as_view_mut())?; let mut destinations: Vec<_> = clusters .rows() diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index aaa2b3dfc9..21eb6438cd 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -131,8 +131,8 @@ impl IVFIndex { SplitAndAssign::new(clusterer, RegionalNearest).run( ®ion, &mut update, - &parents, - &neighbors, + &parents.iter().copied().collect::>(), + &neighbors.iter().copied().collect::>(), )?; }; diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 0e1506cb2f..661fdf7f3e 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -15,8 +15,9 @@ //! 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. -//! [`build::SplitAndAssign`] separates per-parent fitting from whole-region assignment. -//! Fitting writes a fixed number of children per parent into reusable centroid +//! [`build::SplitAndAssign`] separates regional fitting from assignment. Each +//! numerical stage receives source IDs and the region and controls vector access. +//! Fitting declares its total output size and fills coordinator-owned centroid //! scratch. Assignment borrows those rows and retained centroids, groups members, //! and records writes directly through the regional update object. //! The index currently selects [`build::clustering::TwoMeans`] and diff --git a/diskann/src/ivf/tests/mod.rs b/diskann/src/ivf/tests/mod.rs index 25c52525c4..56e810020d 100644 --- a/diskann/src/ivf/tests/mod.rs +++ b/diskann/src/ivf/tests/mod.rs @@ -3,12 +3,7 @@ * Licensed under the MIT license. */ -use std::{ - cell::Cell, - collections::{BTreeMap, BTreeSet}, - future::Future, - num::NonZeroUsize, -}; +use std::{cell::Cell, collections::BTreeMap, future::Future, num::NonZeroUsize}; use diskann_utils::{ Reborrow, @@ -572,7 +567,7 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { ), assignment::RegionalNearest, ) - .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) + .run(®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); @@ -971,12 +966,7 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), assignment::RegionalNearest, ) - .run( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::from([2, 3, 4]), - ) + .run(®ion, &mut update, &[1], &[2, 3, 4]) .unwrap(); let mut centers: Vec<_> = update.created.iter().map(|list| list.centroid[0]).collect(); @@ -1030,12 +1020,7 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() ), assignment::RegionalNearest, ) - .run( - ®ion, - &mut update, - &BTreeSet::from([1, 99]), - &BTreeSet::new() - ) + .run(®ion, &mut update, &[1, 99], &[]) .is_err() ); assert!( @@ -1046,12 +1031,7 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() ), assignment::RegionalNearest, ) - .run( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::from([99]) - ) + .run(®ion, &mut update, &[1], &[99]) .is_err() ); assert!(update.created.is_empty()); @@ -1063,7 +1043,7 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()), assignment::RegionalNearest, ) - .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) + .run(®ion, &mut update, &[1], &[]) .unwrap(); assert_eq!(update.created.len(), 2); } @@ -1119,12 +1099,7 @@ fn assignment_can_cross_parent_boundaries() { ), assignment::RegionalNearest, ) - .run( - ®ion, - &mut update, - &BTreeSet::from([1, 2]), - &BTreeSet::new(), - ) + .run(®ion, &mut update, &[1, 2], &[]) .unwrap(); assert!( update.created[2..] @@ -1153,7 +1128,7 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { clustering::TwoMeans::new(NonZeroUsize::new(1).unwrap(), NonZeroUsize::new(1).unwrap()), assignment::RegionalNearest, ) - .run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()) + .run(®ion, &mut update, &[1], &[]) .unwrap(); assert_eq!(update.created.len(), 2); assert_eq!(update.created[0].centroid, [value]); @@ -1167,19 +1142,29 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { fn split_strategies_support_in_place_fitting_and_source_local_candidates() { struct RowsAsCenters; - impl<'v> clustering::Clusterer> for RowsAsCenters { + impl clustering::Clusterer for RowsAsCenters + where + R: for<'a> RegionView = rowmajor::Ref<'a, f32>>, + { fn fit( &mut self, - points: rowmajor::Ref<'v, f32>, + sources: &[u32], + region: &R, mut dst: rowmajor::Mut<'_, f32>, ) -> ANNResult<()> { - assert_eq!(points.nrows(), dst.nrows()); - dst.as_mut_slice().copy_from_slice(points.as_slice()); + assert_eq!(dst.nrows(), sources.len() * 2); + assert_eq!(dst.ncols(), 1); + for (&source, centers) in sources.iter().zip(dst.as_mut_slice().chunks_exact_mut(2)) { + let vectors = region.get_vectors(source).unwrap(); + let points = vectors.reborrow(); + assert_eq!(points.nrows(), 2); + centers.copy_from_slice(points.as_slice()); + } Ok(()) } - fn num_centers(&self) -> NonZeroUsize { - NonZeroUsize::new(2).unwrap() + fn num_centers(&self, sources: usize) -> Option { + sources.checked_mul(2) } fn dim(&self) -> NonZeroUsize { @@ -1297,8 +1282,8 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { ] .into(); let region = Region::new(&sources); - let parents = BTreeSet::from([1, 2]); - let neighbors = BTreeSet::from([3, 4]); + let parents = [1, 2]; + let neighbors = [3, 4]; let mut full_memberships = PendingLists::default(); let mut full_created = Vec::new(); let mut full = Update::new(&sources, 1, &mut full_memberships, &mut full_created); @@ -1357,12 +1342,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut smaller_created = Vec::new(); let mut smaller = Update::new(&sources, 1, &mut smaller_memberships, &mut smaller_created); splitter - .run::<_, _, Region<'_>, _>( - ®ion, - &mut smaller, - &BTreeSet::from([1]), - &BTreeSet::new(), - ) + .run::<_, _, Region<'_>, _>(®ion, &mut smaller, &[1], &[]) .unwrap(); assert_eq!(smaller.created.len(), 2); assert_eq!(smaller.created[0].members, [0]); @@ -1390,8 +1370,8 @@ fn numerical_strategy_failures_record_no_updates() { CountOverflow, } - impl clustering::Clusterer for InvalidFit { - fn fit(&mut self, _: V, _: rowmajor::Mut<'_, f32>) -> ANNResult<()> { + impl> clustering::Clusterer for InvalidFit { + fn fit(&mut self, _: &[L], _: &R, _: rowmajor::Mut<'_, f32>) -> ANNResult<()> { match self { Self::Error => Err(ANNError::message("injected fitting failure")), Self::CountOverflow => { @@ -1400,12 +1380,11 @@ fn numerical_strategy_failures_record_no_updates() { } } - fn num_centers(&self) -> NonZeroUsize { - NonZeroUsize::new(match self { + fn num_centers(&self, sources: usize) -> Option { + sources.checked_mul(match self { Self::CountOverflow => usize::MAX, Self::Error => 2, }) - .unwrap() } fn dim(&self) -> NonZeroUsize { @@ -1445,8 +1424,8 @@ fn numerical_strategy_failures_record_no_updates() { )] .into(); let region = Region::new(&sources); - let parents = BTreeSet::from([1]); - let neighbors = BTreeSet::new(); + let parents = [1]; + let neighbors = []; for (clusterer, expected) in [ (InvalidFit::Error, "injected fitting failure"), (InvalidFit::CountOverflow, "centroid count overflow"), @@ -1454,17 +1433,29 @@ fn numerical_strategy_failures_record_no_updates() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let selected = if matches!(clusterer, InvalidFit::CountOverflow) { - BTreeSet::from([1, 2]) + let selected: &[u32] = if matches!(clusterer, InvalidFit::CountOverflow) { + &[1, 2] } else { - parents.clone() + &parents }; let error = SplitAndAssign::new(clusterer, assignment::RegionalNearest) - .run(®ion, &mut update, &selected, &neighbors) + .run(®ion, &mut update, selected, &neighbors) .unwrap_err(); assert!(error.to_string().contains(expected), "{error}"); assert_unwritten(update); } + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 1, &mut memberships, &mut created); + let error = SplitAndAssign::new( + clustering::TwoMeans::new(ONE, ONE), + assignment::RegionalNearest, + ) + .run(®ion, &mut update, &[], &neighbors) + .unwrap_err(); + assert!(error.to_string().contains("at least one parent"), "{error}"); + assert_unwritten(update); + let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); @@ -1597,7 +1588,7 @@ fn split_driver_reuses_clusterer_with_fallible_borrowed_batches() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let result = driver.run(®ion, &mut update, &BTreeSet::from([1]), &BTreeSet::new()); + let result = driver.run(®ion, &mut update, &[1], &[]); if let Some(count) = fail_on { let error = result.unwrap_err(); assert!( @@ -1659,16 +1650,28 @@ fn split_driver_accepts_native_batches_without_matrix_access() { struct NativeCenters; - impl<'v> clustering::Clusterer<&'v [[u8; 1]]> for NativeCenters { - fn fit(&mut self, points: &'v [[u8; 1]], mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { - for (point, center) in points.iter().zip(dst.rows_mut()) { - center[0] = point[0].into(); + impl clustering::Clusterer for NativeCenters + where + R: for<'a> RegionView = &'a [[u8; 1]]>, + { + fn fit( + &mut self, + sources: &[u32], + region: &R, + mut dst: rowmajor::Mut<'_, f32>, + ) -> ANNResult<()> { + for (&source, centers) in sources.iter().zip(dst.as_mut_slice().chunks_exact_mut(2)) { + let vectors = region.get_vectors(source).unwrap(); + let points = vectors.reborrow(); + for (point, center) in points.iter().zip(centers) { + *center = point[0].into(); + } } Ok(()) } - fn num_centers(&self) -> NonZeroUsize { - NonZeroUsize::new(2).unwrap() + fn num_centers(&self, sources: usize) -> Option { + sources.checked_mul(2) } fn dim(&self) -> NonZeroUsize { @@ -1718,12 +1721,7 @@ fn split_driver_accepts_native_batches_without_matrix_access() { let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); driver - .run::<_, _, NativeRegion<'_>, _>( - ®ion, - &mut update, - &BTreeSet::from([1]), - &BTreeSet::new(), - ) + .run::<_, _, NativeRegion<'_>, _>(®ion, &mut update, &[1], &[]) .unwrap(); assert_eq!(created.len(), 2); assert_eq!(created[0].centroid, [f32::from(offset)]); @@ -1741,11 +1739,20 @@ fn split_driver_accepts_native_batches_without_matrix_access() { fn two_means_matches_quantization_lloyds() { use clustering::{Clusterer, TwoMeans}; - let points = matrix(&[[0.0, 2.0], [2.0, 4.0], [6.0, 6.0], [10.0, 12.0]]); + 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]]), + }, + )]); + let region = Region::new(&sources); + let points = sources[&1].vectors.as_view(); for iterations in [1, 5] { - let mut expected = points.as_view().subview(0..2).unwrap().to_rowmajor_owned(); + let mut expected = points.subview(0..2).unwrap().to_rowmajor_owned(); let _ = diskann_quantization::algorithms::kmeans::lloyds::lloyds( - points.as_view(), + points, expected.as_view_mut(), iterations, ); @@ -1755,7 +1762,7 @@ fn two_means_matches_quantization_lloyds() { NonZeroUsize::new(iterations).unwrap(), NonZeroUsize::new(2).unwrap(), ) - .fit(points.as_view(), actual.as_view_mut()) + .fit(&[1], ®ion, actual.as_view_mut()) .unwrap(); assert_eq!(actual.as_slice(), expected.as_slice()); } @@ -1765,14 +1772,22 @@ fn two_means_matches_quantization_lloyds() { fn two_means_fits_only_the_destination_window() { use clustering::{Clusterer, TwoMeans}; - let points = matrix(&[[0.0, 2.0], [2.0, 4.0], [10.0, 12.0], [12.0, 14.0]]); + 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], [10.0, 12.0], [12.0, 14.0]]), + }, + )]); + let region = Region::new(&sources); let mut scratch = rowmajor::Owned::from_element(4, 2, -1.0); let mut clusterer = TwoMeans::new( NonZeroUsize::new(10).unwrap(), NonZeroUsize::new(2).unwrap(), ); clusterer - .fit(points.as_view(), scratch.subview_mut(1..3).unwrap()) + .fit(&[1], ®ion, scratch.subview_mut(1..3).unwrap()) .unwrap(); assert_eq!(scratch.row(0), [-1.0, -1.0]); assert_eq!(scratch.row(3), [-1.0, -1.0]); @@ -1787,8 +1802,17 @@ fn two_means_fits_only_the_destination_window() { (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 1, 2), (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 2, 1), ] { + 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 dst = rowmajor::Owned::from_element(rows, dim, -1.0); - assert!(clusterer.fit(points.as_view(), dst.as_view_mut()).is_err()); + assert!(clusterer.fit(&[1], ®ion, dst.as_view_mut()).is_err()); assert!(dst.as_slice().iter().all(|&value| value == -1.0)); } } diff --git a/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 3217e93ea0..28f0a6bd4e 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -74,16 +74,19 @@ 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 numerical driver borrows vector matrices and member slices, without building -per-point reference tables. A provider with scattered storage may pack rows during -filling to satisfy the matrix interface. The built-in two-means fitter seeds two -distinct random input rows and delegates to quantization's Lloyd routine, which +The numerical stages receive source IDs and the region, and control retrieval and +conversion of provider-defined vector handles without coordinator-built per-point +reference tables. The built-in strategies borrow 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. -The numerical driver separates per-parent fitting (`Clusterer`) from whole-region -placement (`Assigner`). A clusterer declares its dimension and fixed number of -centers per parent, then fills mutable row windows of reusable centroid scratch. +The numerical driver separates regional fitting (`Clusterer`) from placement +(`Assigner`). A clusterer declares its dimension and total number of centers for +the source count, then fills coordinator-owned centroid scratch. Fitting receives +only split parents; assignment receives the union of parents and retained neighbors. A borrowed destination catalog combines the completed rows with unchanged retained centroids; retained coordinates are not copied. The assigner receives the region, source IDs, destinations, and regional writer directly. It computes complete From 794731ddd26a715aff96ae052eeeac98fdb3d402 Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Fri, 23 Oct 2026 16:08:49 -0400 Subject: [PATCH 31/33] conslidate clustering and assignment --- diskann-utils/src/views/rowmajor.rs | 96 +++ diskann-utils/src/views/rowmajor/iter.rs | 75 ++- diskann/src/ivf/build/assignment.rs | 120 ---- diskann/src/ivf/build/clustering.rs | 125 ---- diskann/src/ivf/build/mod.rs | 34 +- diskann/src/ivf/build/split.rs | 88 --- diskann/src/ivf/build/twomeans.rs | 216 +++++++ diskann/src/ivf/index.rs | 7 +- diskann/src/ivf/mod.rs | 11 +- diskann/src/ivf/region.rs | 20 +- diskann/src/ivf/tests/mod.rs | 774 +++++++++-------------- rfcs/01187-incremental-ivf.md | 39 +- 12 files changed, 739 insertions(+), 866 deletions(-) delete mode 100644 diskann/src/ivf/build/assignment.rs delete mode 100644 diskann/src/ivf/build/clustering.rs delete mode 100644 diskann/src/ivf/build/split.rs create mode 100644 diskann/src/ivf/build/twomeans.rs 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/src/ivf/build/assignment.rs b/diskann/src/ivf/build/assignment.rs deleted file mode 100644 index 17c2e966f2..0000000000 --- a/diskann/src/ivf/build/assignment.rs +++ /dev/null @@ -1,120 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Assign regional memberships and record the corresponding updates. - -use std::collections::BTreeSet; - -use diskann_utils::{Reborrow, views::rowmajor::Matrix}; -use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; - -use crate::{ - ANNError, ANNResult, - error::IntoANNResult, - ivf::region::{AsMatRef, RegionUpdate, RegionView}, -}; - -/// Assign the whole region, allowing policies that coordinate multiple sources. -pub trait Assigner> { - /// Assign each source member to a destination and record complete memberships. - /// - /// Create new destinations, replace retained memberships, and retire sources - /// not retained as destinations. Validate and compute assignments before - /// recording writes. A write error may leave partial private updates. - fn assign( - &mut self, - region: &R, - update: &mut U, - sources: &[L], - destinations: &[Target<'_, L>], - ) -> ANNResult<()> - where - U: RegionUpdate; -} - -pub enum Type { - /// Existing centroid. - Exists(L), - /// New centroid as a result of the split. - New, -} - -/// A borrowed centroid; only retained destinations have an existing list ID. -pub struct Target<'a, L> { - pub list: Type, - pub centroid: &'a [f32], -} - -pub struct RegionalNearest; - -impl Assigner for RegionalNearest -where - R: RegionView, - for<'a> R::VectorsRef<'a>: AsMatRef, -{ - fn assign( - &mut self, - region: &R, - update: &mut U, - sources: &[L], - destinations: &[Target<'_, L>], - ) -> ANNResult<()> - where - U: RegionUpdate, - { - let mut memberships: Vec> = (0..destinations.len()).map(|_| Vec::new()).collect(); - - for &src in sources { - let view = region.get_view(src).ok_or_else(|| { - ANNError::message("IVF assignment source is missing from the region") - })?; - let vectors = region.get_vectors(src).ok_or_else(|| { - ANNError::message("IVF assignment source vectors are missing from the region") - })?; - let vectors = vectors.reborrow(); - let vectors = vectors.as_view().into_ann_result()?; - - for (&id, point) in view.members.iter().zip(vectors.rows()) { - let dst = nearest(point, destinations.iter().map(|d| d.centroid)); - memberships[dst].push(id); - } - } - - let mut retained = BTreeSet::new(); - - for (destination, members) in destinations.iter().zip(memberships) { - match &destination.list { - Type::Exists(id) => { - update.set_members(*id, &members)?; - retained.insert(*id); - } - Type::New => update.create(destination.centroid, &members)?, - } - } - - for &source in sources { - if !retained.contains(&source) { - update.retire(source)?; - } - } - - Ok(()) - } -} - -/// Prefer the first centroid on ties. -fn nearest<'a>(point: &[f32], targets: impl Iterator) -> usize { - let mut min = f32::MAX; - let mut best = usize::MAX; - for (i, center) in targets.enumerate() { - let dist = SquaredL2::evaluate(point, center); - if min > dist { - best = i; - min = dist; - } - } - - best -} diff --git a/diskann/src/ivf/build/clustering.rs b/diskann/src/ivf/build/clustering.rs deleted file mode 100644 index b21ee5ef6d..0000000000 --- a/diskann/src/ivf/build/clustering.rs +++ /dev/null @@ -1,125 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Fit replacement centroids over provider vector views. - -use std::num::NonZeroUsize; - -use diskann_quantization::algorithms::kmeans::lloyds::lloyds; -use diskann_utils::{ - Reborrow, - views::rowmajor::{self, Matrix, MatrixMut}, -}; - -use crate::{ - ANNError, ANNResult, - error::IntoANNResult, - ivf::region::{AsMatRef, RegionView}, -}; - -/// Fit replacement centroids for selected sources without assigning regional memberships. -pub trait Clusterer -where - R: RegionView, -{ - /// Fill `dst` with finite centroids of the canonical dimension. - /// - /// Implementations own initialization, convergence, empty-cluster handling, - /// and reusable fitting scratch. `dst` has `num_centers(sources.len())` rows - /// and `dim()` columns. The total number of centers must be at least two. - fn fit(&mut self, sources: &[L], region: &R, dst: rowmajor::Mut<'_, T>) -> ANNResult<()>; - - /// Total output row count for this source count, or `None` on overflow. - fn num_centers(&self, sources: usize) -> Option; - - fn dim(&self) -> NonZeroUsize; -} - -/// First-two-row seeds followed by Lloyd iterations, which reset empty centers to zero. -pub struct TwoMeans { - iterations: NonZeroUsize, - dim: NonZeroUsize, -} - -impl TwoMeans { - pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { - Self { iterations, dim } - } - - const fn inner_centers() -> usize { - 2 - } - - fn fit_one( - &self, - points: rowmajor::Ref<'_, f32>, - mut dst: rowmajor::Mut<'_, f32>, - ) -> ANNResult<()> { - if points.nrows() < Self::inner_centers() || points.ncols() != self.dim.get() { - return Err(ANNError::message( - "IVF two-means needs at least two points with the configured dimension", - )); - } - - dst.as_mut_slice().copy_from_slice( - points - .subview(0..Self::inner_centers()) - .ok_or_else(|| ANNError::message("IVF two-means seed rows are missing"))? - .as_slice(), - ); - - let _ = lloyds(points, dst, self.iterations.get()); - - Ok(()) - } -} - -impl Clusterer for TwoMeans -where - R: RegionView, - for<'a> R::VectorsRef<'a>: AsMatRef, -{ - fn fit(&mut self, sources: &[L], region: &R, mut dst: rowmajor::Mut<'_, f32>) -> ANNResult<()> { - if dst.ncols() != self.dim.get() - || sources.len().checked_mul(Self::inner_centers()) != Some(dst.nrows()) - { - return Err(ANNError::message( - "IVF two-means scratch needs two rows per source and the configured dimension", - )); - } - - let mut s = 0; - for &p in sources { - let vecs = region.get_vectors(p).ok_or_else(|| { - ANNError::message("IVF split source vectors are missing from the region") - })?; - - let vecs = vecs.reborrow(); - - let points = vecs.as_view().into_ann_result()?; - - let end = s + Self::inner_centers(); - - self.fit_one( - points, - dst.subview_mut(s..end).ok_or_else(|| { - ANNError::message("IVF centroid scratch range is out of bounds") - })?, - )?; - - s = end; - } - - Ok(()) - } - - fn num_centers(&self, sources: usize) -> Option { - sources.checked_mul(Self::inner_centers()) - } - - fn dim(&self) -> NonZeroUsize { - self.dim - } -} diff --git a/diskann/src/ivf/build/mod.rs b/diskann/src/ivf/build/mod.rs index 80a22b46f3..2a1862f16c 100644 --- a/diskann/src/ivf/build/mod.rs +++ b/diskann/src/ivf/build/mod.rs @@ -3,10 +3,34 @@ * Licensed under the MIT license. */ -//! Numerical fitting and regional assignment for IVF construction. +//! Complete regional split operations, independent of physical list storage. -pub mod assignment; -pub mod clustering; -mod split; +use std::fmt::Debug; -pub use split::SplitAndAssign; +use crate::{ + ANNError, + ivf::region::{RegionUpdate, RegionView}, +}; + +pub(crate) mod twomeans; +pub(crate) use twomeans::TwoMeansSplitter; + +/// 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/split.rs b/diskann/src/ivf/build/split.rs deleted file mode 100644 index ae9c6aed95..0000000000 --- a/diskann/src/ivf/build/split.rs +++ /dev/null @@ -1,88 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Coordinate per-parent fitting and whole-region assignment. - -use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; - -use super::{ - assignment::{Assigner, Target, Type}, - clustering::Clusterer, -}; -use crate::{ - ANNError, ANNResult, - ivf::region::{RegionUpdate, RegionView}, -}; - -pub struct SplitAndAssign { - pub clusterer: C, - pub assigner: A, -} - -impl SplitAndAssign { - pub fn new(clusterer: C, assigner: A) -> Self { - Self { - clusterer, - assigner, - } - } - - /// Compute the complete result before registering any provider writes. - /// - /// Concrete-reference clusterers may need an explicit `R` type argument because - /// of a compiler limitation when inferring higher-ranked associated types. - pub fn run( - &mut self, - region: &R, - update: &mut U, - parents: &[L], - neighbors: &[L], - ) -> ANNResult<()> - where - A: Assigner, - C: Clusterer, - R: RegionView, - U: RegionUpdate, - { - if parents.is_empty() { - return Err(ANNError::message("IVF splitting needs at least one parent")); - } - - let rows = self - .clusterer - .num_centers(parents.len()) - .ok_or_else(|| ANNError::message("IVF centroid count overflow"))?; - - let mut clusters = - rowmajor::Owned::try_from_element(rows, self.clusterer.dim().get(), f32::NAN)?; - - self.clusterer - .fit(parents, region, clusters.as_view_mut())?; - - let mut destinations: Vec<_> = clusters - .rows() - .map(|centroid| Target { - list: Type::New, - centroid, - }) - .collect(); - - for &n in neighbors { - let view = region - .get_view(n) - .ok_or_else(|| ANNError::message("IVF neighbor is missing from the region"))?; - - destinations.push(Target { - list: Type::Exists(n), - centroid: view.centroid, - }); - } - - let sources = parents.iter().chain(neighbors).copied().collect::>(); - - self.assigner - .assign(region, update, &sources, &destinations) - } -} diff --git a/diskann/src/ivf/build/twomeans.rs b/diskann/src/ivf/build/twomeans.rs new file mode 100644 index 0000000000..0c10629739 --- /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, RegionUpdate, RegionView}, + 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 TwoMeansSplitter { + iterations: NonZeroUsize, + dim: NonZeroUsize, +} + +impl TwoMeansSplitter { + 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: RegionView, + 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 TwoMeansSplitter +where + I: VectorId, + L: VectorId, + R: RegionView, + 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_view(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_view(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 index 21eb6438cd..4a04f62e4f 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -17,7 +17,7 @@ use crate::{ flat::DistancesUnordered, graph::SearchOutputBuffer, ivf::{ - build::{SplitAndAssign, assignment::RegionalNearest, clustering::TwoMeans}, + build::{Split, TwoMeansSplitter}, region::AsMatRef, traits::BuildStrategy, }, @@ -83,7 +83,7 @@ impl IVFIndex { ) -> impl SendFuture> where S: BuildStrategy<'a, B, C>, - S::Accessor: for<'b> BuildAccessor: AsMatRef>, + S::Accessor: for<'b> BuildAccessor: AsMatRef<'b, f32>>, B: Send + Sync, { let config = self.config; @@ -127,8 +127,7 @@ impl IVFIndex { let (region, mut update) = accessor.fill(iter).await?; - let clusterer = TwoMeans::new(config.split_iterations, dim); - SplitAndAssign::new(clusterer, RegionalNearest).run( + TwoMeansSplitter::new(config.split_iterations, dim).split( ®ion, &mut update, &parents.iter().copied().collect::>(), diff --git a/diskann/src/ivf/mod.rs b/diskann/src/ivf/mod.rs index 661fdf7f3e..c18a780d1e 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -15,13 +15,10 @@ //! 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. -//! [`build::SplitAndAssign`] separates regional fitting from assignment. Each -//! numerical stage receives source IDs and the region and controls vector access. -//! Fitting declares its total output size and fills coordinator-owned centroid -//! scratch. Assignment borrows those rows and retained centroids, groups members, -//! and records writes directly through the regional update object. -//! The index currently selects [`build::clustering::TwoMeans`] and -//! [`build::assignment::RegionalNearest`]. +//! [`build::Split`] owns fitting, assignment, and regional updates as one operation. +//! The built-in [`build::TwoMeansSplitter`] 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 diff --git a/diskann/src/ivf/region.rs b/diskann/src/ivf/region.rs index 54290ac6e2..d5bb9a87a9 100644 --- a/diskann/src/ivf/region.rs +++ b/diskann/src/ivf/region.rs @@ -9,7 +9,7 @@ use std::fmt::Debug; use diskann_utils::{Reborrow, views::rowmajor}; -use crate::{ANNError, ANNResult}; +use crate::ANNResult; /// One source list's centroid and complete operation-local membership. /// @@ -54,20 +54,16 @@ pub trait RegionUpdate: Send { fn create(&mut self, centroid: &[f32], members: &[I]) -> ANNResult<()>; } -pub trait AsMatRef { - type Error: Send + Sync + Debug + Into; - - fn as_view(&self) -> Result, Self::Error>; +pub trait AsMatRef<'a, T> { + fn as_view(self) -> rowmajor::Ref<'a, T>; } -impl AsMatRef for rowmajor::Ref<'_, T> { - type Error = crate::error::Infallible; - - fn as_view(&self) -> Result, Self::Error> { - Ok(*self) +impl<'a, T> AsMatRef<'a, T> for rowmajor::Ref<'a, T> { + fn as_view(self) -> rowmajor::Ref<'a, T> { + self } } -pub trait AsMatMut: AsMatRef { - fn as_mut_view(&self) -> Result, Self::Error>; +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 index 56e810020d..0e2c187a8f 100644 --- a/diskann/src/ivf/tests/mod.rs +++ b/diskann/src/ivf/tests/mod.rs @@ -19,7 +19,7 @@ use self::{ }; use super::{ Config, IVFIndex, InsertStats, ReassignCount, - build::{SplitAndAssign, assignment, clustering}, + build::{Split, TwoMeansSplitter, twomeans::TwoMeansError}, region::{AsMatRef, RegionUpdate, RegionView, View}, traits::{ BuildAccessor, BuildStrategy, HasIds, NearestList, Router, SearchAccessor, SearchStrategy, @@ -560,15 +560,9 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { 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); - SplitAndAssign::new( - clustering::TwoMeans::new( - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - ), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[1], &[]) - .unwrap(); + TwoMeansSplitter::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()); @@ -962,12 +956,9 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - SplitAndAssign::new( - clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[1], &[2, 3, 4]) - .unwrap(); + TwoMeansSplitter::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); @@ -978,10 +969,10 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { .flat_map(|list| list.members.iter().copied()) .collect(); child_members.sort_unstable(); - assert_eq!(child_members, [0, 1]); + 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![])] { + 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}"); }; @@ -1013,38 +1004,23 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() { let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); assert!( - SplitAndAssign::new( - clustering::TwoMeans::new( - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - ), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[1, 99], &[]) - .is_err() + TwoMeansSplitter::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) + .split(®ion, &mut update, &[1, 99], &[]) + .is_err() ); assert!( - SplitAndAssign::new( - clustering::TwoMeans::new( - NonZeroUsize::new(5).unwrap(), - NonZeroUsize::new(2).unwrap(), - ), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[1], &[99]) - .is_err() + TwoMeansSplitter::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]) ); - SplitAndAssign::new( - clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[1], &[]) - .unwrap(); + TwoMeansSplitter::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(); @@ -1092,14 +1068,11 @@ fn assignment_can_cross_parent_boundaries() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - SplitAndAssign::new( - clustering::TwoMeans::new( - NonZeroUsize::new(10).unwrap(), - NonZeroUsize::new(1).unwrap(), - ), - assignment::RegionalNearest, + TwoMeansSplitter::new( + NonZeroUsize::new(10).unwrap(), + NonZeroUsize::new(1).unwrap(), ) - .run(®ion, &mut update, &[1, 2], &[]) + .split(®ion, &mut update, &[1, 2], &[]) .unwrap(); assert!( update.created[2..] @@ -1124,12 +1097,9 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - SplitAndAssign::new( - clustering::TwoMeans::new(NonZeroUsize::new(1).unwrap(), NonZeroUsize::new(1).unwrap()), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[1], &[]) - .unwrap(); + TwoMeansSplitter::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]); @@ -1139,113 +1109,61 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { } #[test] -fn split_strategies_support_in_place_fitting_and_source_local_candidates() { - struct RowsAsCenters; - - impl clustering::Clusterer for RowsAsCenters - where - R: for<'a> RegionView = rowmajor::Ref<'a, f32>>, - { - fn fit( - &mut self, - sources: &[u32], - region: &R, - mut dst: rowmajor::Mut<'_, f32>, - ) -> ANNResult<()> { - assert_eq!(dst.nrows(), sources.len() * 2); - assert_eq!(dst.ncols(), 1); - for (&source, centers) in sources.iter().zip(dst.as_mut_slice().chunks_exact_mut(2)) { - let vectors = region.get_vectors(source).unwrap(); - let points = vectors.reborrow(); - assert_eq!(points.nrows(), 2); - centers.copy_from_slice(points.as_slice()); - } - Ok(()) - } - - fn num_centers(&self, sources: usize) -> Option { - sources.checked_mul(2) - } - - fn dim(&self) -> NonZeroUsize { - NonZeroUsize::new(1).unwrap() - } - } - - struct ParentLocal { - owners: BTreeMap>, - } - - impl assignment::Assigner for ParentLocal - where - R: for<'a> RegionView = rowmajor::Ref<'a, f32>>, - { - fn assign( - &mut self, - region: &R, - update: &mut U, - sources: &[u32], - destinations: &[assignment::Target<'_, u32>], - ) -> ANNResult<()> - where - U: RegionUpdate, - { - assert_eq!(sources, [1, 2, 3, 4]); - assert_eq!(destinations.len(), 6); - assert!(std::ptr::eq( - destinations[4].centroid, - region.get_view(3).unwrap().centroid - )); - let children = [(1, 0..2), (2, 2..4)]; - let mut memberships = vec![Vec::new(); destinations.len()]; - for &source in sources { - let view = region.get_view(source).unwrap(); - let vectors = region.get_vectors(source).unwrap(); - let vectors = vectors.reborrow(); - let mut candidates: Vec<_> = destinations - .iter() - .enumerate() - .filter_map(|(index, destination)| { - matches!(&destination.list, assignment::Type::Exists(id) if *id == source) - .then_some(index) - }) - .collect(); - for (parent, range) in &children { - if *parent == source - || self - .owners - .get(&source) - .is_some_and(|owners| owners.contains(parent)) - { - candidates.extend(range.clone()); - } - } - for (&id, point) in view.members.iter().zip(vectors.rows()) { - let destination = candidates - .iter() - .copied() - .min_by(|&a, &b| { - let left: f32 = SquaredL2::evaluate(point, destinations[a].centroid); - let right: f32 = SquaredL2::evaluate(point, destinations[b].centroid); - left.total_cmp(&right) - }) - .unwrap(); - memberships[destination].push(id); - } - } - for (destination, members) in destinations.iter().zip(memberships) { - match &destination.list { - assignment::Type::New => update.create(destination.centroid, &members)?, - assignment::Type::Exists(id) => update.set_members(*id, &members)?, - } - } - for (parent, _) in children { - update.retire(parent)?; - } - Ok(()) - } +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 = Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplitter::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, @@ -1287,42 +1205,21 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut full_memberships = PendingLists::default(); let mut full_created = Vec::new(); let mut full = Update::new(&sources, 1, &mut full_memberships, &mut full_created); - let mut splitter = SplitAndAssign::new(RowsAsCenters, assignment::RegionalNearest); + let mut splitter = TwoMeansSplitter::new(ONE, ONE); splitter - .run::<_, _, Region<'_>, _>(®ion, &mut full, &parents, &neighbors) + .split(®ion, &mut full, &parents, &neighbors) .unwrap(); - let mut restricted_memberships = PendingLists::default(); - let mut restricted_created = Vec::new(); - let mut restricted = Update::new( - &sources, - 1, - &mut restricted_memberships, - &mut restricted_created, + assert_eq!( + full.created + .iter() + .map(|list| list.centroid[0]) + .collect::>(), + [-10.0, 10.0, 100.0, 110.0] ); - SplitAndAssign::new( - RowsAsCenters, - ParentLocal { - owners: [(3, vec![1]), (4, vec![1, 2])].into(), - }, - ) - .run::<_, _, Region<'_>, _>(®ion, &mut restricted, &parents, &neighbors) - .unwrap(); - - for update in [&full, &restricted] { - assert_eq!( - update - .created - .iter() - .map(|list| list.centroid[0]) - .collect::>(), - [-10.0, 10.0, 100.0, 110.0] - ); - assert_eq!(update.created[0].members, [0]); - assert_eq!(update.created[1].members, [1, 4]); - assert_eq!(update.created[3].members, [3]); - } + 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]); - assert_eq!(restricted.created[2].members, [2, 7]); let mut repeated_memberships = PendingLists::default(); let mut repeated_created = Vec::new(); let mut repeated = Update::new( @@ -1332,7 +1229,7 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { &mut repeated_created, ); splitter - .run::<_, _, Region<'_>, _>(®ion, &mut repeated, &parents, &neighbors) + .split(®ion, &mut repeated, &parents, &neighbors) .unwrap(); for (first, repeated) in full.created.iter().zip(repeated.created.iter()) { assert_eq!(first.centroid, repeated.centroid); @@ -1341,135 +1238,39 @@ fn split_strategies_support_in_place_fitting_and_source_local_candidates() { let mut smaller_memberships = PendingLists::default(); let mut smaller_created = Vec::new(); let mut smaller = Update::new(&sources, 1, &mut smaller_memberships, &mut smaller_created); - splitter - .run::<_, _, Region<'_>, _>(®ion, &mut smaller, &[1], &[]) - .unwrap(); + 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]); - for (update, expected) in [(full, vec![6]), (restricted, vec![5, 6])] { - let mut changes: BTreeMap<_, _> = - std::mem::take(update.memberships).into_changes().collect(); - for parent in [1, 2] { - assert!(matches!(changes.remove(&parent), Some(Membership::Retire))); - } - for (id, expected) in [(3, expected), (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()); + 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 numerical_strategy_failures_record_no_updates() { - enum InvalidFit { - Error, - CountOverflow, - } - - impl> clustering::Clusterer for InvalidFit { - fn fit(&mut self, _: &[L], _: &R, _: rowmajor::Mut<'_, f32>) -> ANNResult<()> { - match self { - Self::Error => Err(ANNError::message("injected fitting failure")), - Self::CountOverflow => { - unreachable!("centroid count overflow should be rejected before fitting") - } - } - } - - fn num_centers(&self, sources: usize) -> Option { - sources.checked_mul(match self { - Self::CountOverflow => usize::MAX, - Self::Error => 2, - }) - } - - fn dim(&self) -> NonZeroUsize { - NonZeroUsize::new(1).unwrap() - } - } - - struct InvalidAssignment; - - impl> assignment::Assigner for InvalidAssignment { - fn assign( - &mut self, - _: &R, - _: &mut U, - _: &[u32], - _: &[assignment::Target<'_, u32>], - ) -> ANNResult<()> - where - U: RegionUpdate, - { - Err(ANNError::message("injected assignment failure")) - } - } - - fn assert_unwritten(update: Update<'_>) { - assert!(update.created.is_empty()); - assert_eq!(std::mem::take(update.memberships).into_changes().count(), 0); - } - - let sources = [( - 1, - Source { - centroid: vec![0.0], - members: vec![0, 1], - vectors: matrix(&[[-1.0], [1.0]]), - }, - )] - .into(); +fn empty_split_records_no_updates() { + let sources = BTreeMap::new(); let region = Region::new(&sources); - let parents = [1]; - let neighbors = []; - for (clusterer, expected) in [ - (InvalidFit::Error, "injected fitting failure"), - (InvalidFit::CountOverflow, "centroid count overflow"), - ] { - let mut memberships = PendingLists::default(); - let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let selected: &[u32] = if matches!(clusterer, InvalidFit::CountOverflow) { - &[1, 2] - } else { - &parents - }; - let error = SplitAndAssign::new(clusterer, assignment::RegionalNearest) - .run(®ion, &mut update, selected, &neighbors) - .unwrap_err(); - assert!(error.to_string().contains(expected), "{error}"); - assert_unwritten(update); - } - let mut memberships = PendingLists::default(); - let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let error = SplitAndAssign::new( - clustering::TwoMeans::new(ONE, ONE), - assignment::RegionalNearest, - ) - .run(®ion, &mut update, &[], &neighbors) - .unwrap_err(); - assert!(error.to_string().contains("at least one parent"), "{error}"); - assert_unwritten(update); - let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let error = SplitAndAssign::new( - clustering::TwoMeans::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(1).unwrap()), - InvalidAssignment, - ) - .run(®ion, &mut update, &parents, &neighbors) - .unwrap_err(); - assert!( - error.to_string().contains("injected assignment failure"), - "{error}" - ); - assert_unwritten(update); + let error = TwoMeansSplitter::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] @@ -1496,17 +1297,15 @@ fn region_vector_reborrows_share_source_storage() { } #[test] -fn split_driver_reuses_clusterer_with_fallible_borrowed_batches() { +fn two_means_handles_borrowed_batches() { struct Batch<'a> { vectors: rowmajor::Ref<'a, f32>, reads: &'a Cell, - fail_on: Option, } struct BatchRef<'a> { vectors: rowmajor::Ref<'a, f32>, reads: &'a Cell, - fail_on: Option, } impl<'r> Reborrow<'r> for Batch<'_> { @@ -1516,28 +1315,20 @@ fn split_driver_reuses_clusterer_with_fallible_borrowed_batches() { BatchRef { vectors: self.vectors.reborrow(), reads: self.reads, - fail_on: self.fail_on, } } } - impl AsMatRef for BatchRef<'_> { - type Error = ANNError; - - fn as_view(&self) -> ANNResult> { - let count = self.reads.get() + 1; - self.reads.set(count); - if self.fail_on == Some(count) { - return Err(ANNError::message("injected batch conversion failure")); - } - Ok(self.vectors) + 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, - fail_on: Option, } impl RegionView for BatchRegion<'_> { @@ -1555,68 +1346,58 @@ fn split_driver_reuses_clusterer_with_fallible_borrowed_batches() { self.inner.get_vectors(list).map(|vectors| Batch { vectors, reads: &self.reads, - fail_on: self.fail_on, }) } } - let mut driver = SplitAndAssign::new( - clustering::TwoMeans::new(ONE, ONE), - assignment::RegionalNearest, - ); - for fail_on in [Some(1), Some(2), None] { - let sources = BTreeMap::from([( + let sources = BTreeMap::from([ + ( 1, Source { centroid: vec![0.0], members: vec![0, 1], vectors: matrix(&[[-1.0], [1.0]]), }, - )]); - let region = BatchRegion { - inner: Region::new(&sources), - reads: Cell::new(0), - fail_on, - }; - 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 = Update::new(&sources, 1, &mut memberships, &mut created); - let result = driver.run(®ion, &mut update, &[1], &[]); - if let Some(count) = fail_on { - let error = result.unwrap_err(); - assert!( - error - .to_string() - .contains("injected batch conversion failure") - ); - assert_eq!(region.reads.get(), count); - assert!(created.is_empty()); - assert_eq!(memberships.into_changes().count(), 0); - } else { - result.unwrap(); - assert_eq!(region.reads.get(), 2); - 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]); - assert!(matches!( - memberships.into_changes().collect::>().as_slice(), - [(1, Membership::Retire)] - )); - } - } + ), + ( + 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 = Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplitter::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_driver_accepts_native_batches_without_matrix_access() { +fn split_accepts_native_batches_without_matrix_access() { struct NativeBatch<'a>(&'a [[u8; 1]]); impl<'r> Reborrow<'r> for NativeBatch<'_> { @@ -1648,59 +1429,32 @@ fn split_driver_accepts_native_batches_without_matrix_access() { } } - struct NativeCenters; + struct NativeSplit; - impl clustering::Clusterer for NativeCenters + impl Split for NativeSplit where R: for<'a> RegionView = &'a [[u8; 1]]>, { - fn fit( - &mut self, - sources: &[u32], - region: &R, - mut dst: rowmajor::Mut<'_, f32>, - ) -> ANNResult<()> { - for (&source, centers) in sources.iter().zip(dst.as_mut_slice().chunks_exact_mut(2)) { - let vectors = region.get_vectors(source).unwrap(); - let points = vectors.reborrow(); - for (point, center) in points.iter().zip(centers) { - *center = point[0].into(); - } - } - Ok(()) - } - - fn num_centers(&self, sources: usize) -> Option { - sources.checked_mul(2) - } - - fn dim(&self) -> NonZeroUsize { - ONE - } - } - - struct OrderedAssignment; + type Error = ANNError; - impl> assignment::Assigner for OrderedAssignment { - fn assign( + fn split>( &mut self, region: &R, update: &mut U, - sources: &[u32], - destinations: &[assignment::Target<'_, u32>], - ) -> ANNResult<()> - where - U: RegionUpdate, - { - let view = region.get_view(sources[0]).unwrap(); - for (destination, member) in destinations.iter().zip(view.members) { - update.create(destination.centroid, std::slice::from_ref(member))?; + parents: &[u32], + neighbors: &[u32], + ) -> ANNResult<()> { + assert!(neighbors.is_empty()); + let view = region.get_view(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(sources[0]) + update.retire(parents[0]) } } - let mut driver = SplitAndAssign::new(NativeCenters, OrderedAssignment); + let mut splitter = NativeSplit; for offset in [0u8, 20] { let sources = BTreeMap::from([( 1, @@ -1720,9 +1474,7 @@ fn split_driver_accepts_native_batches_without_matrix_access() { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - driver - .run::<_, _, NativeRegion<'_>, _>(®ion, &mut update, &[1], &[]) - .unwrap(); + 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]); @@ -1737,71 +1489,55 @@ fn split_driver_accepts_native_batches_without_matrix_access() { #[test] fn two_means_matches_quantization_lloyds() { - use clustering::{Clusterer, TwoMeans}; - - 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]]), - }, - )]); + 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); - let points = sources[&1].vectors.as_view(); for iterations in [1, 5] { - let mut expected = points.subview(0..2).unwrap().to_rowmajor_owned(); - let _ = diskann_quantization::algorithms::kmeans::lloyds::lloyds( - points, - expected.as_view_mut(), - iterations, - ); - - let mut actual = rowmajor::Owned::from_element(2, 2, 0.0); - TwoMeans::new( + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 2, &mut memberships, &mut created); + TwoMeansSplitter::new( NonZeroUsize::new(iterations).unwrap(), NonZeroUsize::new(2).unwrap(), ) - .fit(&[1], ®ion, actual.as_view_mut()) + .split(®ion, &mut update, &[1, 2], &[]) .unwrap(); - assert_eq!(actual.as_slice(), expected.as_slice()); + 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_fits_only_the_destination_window() { - use clustering::{Clusterer, TwoMeans}; - - 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], [10.0, 12.0], [12.0, 14.0]]), - }, - )]); - let region = Region::new(&sources); - let mut scratch = rowmajor::Owned::from_element(4, 2, -1.0); - let mut clusterer = TwoMeans::new( - NonZeroUsize::new(10).unwrap(), - NonZeroUsize::new(2).unwrap(), - ); - clusterer - .fit(&[1], ®ion, scratch.subview_mut(1..3).unwrap()) - .unwrap(); - assert_eq!(scratch.row(0), [-1.0, -1.0]); - assert_eq!(scratch.row(3), [-1.0, -1.0]); - let window = scratch.subview(1..3).unwrap(); - let mut centers: Vec<_> = window.rows().collect(); - centers.sort_by(|a, b| a[0].total_cmp(&b[0])); - assert_eq!(centers, [&[1.0, 3.0][..], &[11.0, 13.0][..]]); - - for (points, rows, dim) in [ - (matrix(&[[0.0, 2.0]]), 2, 2), - (matrix(&[[0.0], [1.0]]), 2, 2), - (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 1, 2), - (matrix(&[[0.0, 2.0], [1.0, 3.0]]), 2, 1), - ] { +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 { @@ -1811,16 +1547,48 @@ fn two_means_fits_only_the_destination_window() { }, )]); let region = Region::new(&sources); - let mut dst = rowmajor::Owned::from_element(rows, dim, -1.0); - assert!(clusterer.fit(&[1], ®ion, dst.as_view_mut()).is_err()); - assert!(dst.as_slice().iter().all(|&value| value == -1.0)); + let mut memberships = PendingLists::default(); + let mut created = Vec::new(); + let mut update = Update::new(&sources, 2, &mut memberships, &mut created); + let error = TwoMeansSplitter::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 regional_assignment_reports_missing_sources_before_writing() { - use assignment::{Assigner, RegionalNearest, Target, Type}; +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 = Update::new(&sources, 1, &mut memberships, &mut created); + let error = TwoMeansSplitter::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, @@ -1839,20 +1607,52 @@ fn regional_assignment_reports_missing_sources_before_writing() { }, ), ]); - let destinations = [Target { - list: Type::New, - centroid: &[0.0], - }]; - for ids in [vec![99], vec![1, 99], vec![1, 2, 99]] { + for ids in [vec![99], vec![2, 99]] { let mut memberships = PendingLists::default(); let mut created = Vec::new(); let mut update = Update::new(&sources, 1, &mut memberships, &mut created); let region = Region::new(&sources); - let error = RegionalNearest - .assign(®ion, &mut update, &ids, &destinations) + let error = TwoMeansSplitter::new(ONE, ONE) + .split(®ion, &mut update, &[1], &ids) .unwrap_err(); - assert!(error.to_string().contains("source is missing"), "{error}"); + 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 = Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplitter::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/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index 28f0a6bd4e..aa6781d048 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -74,30 +74,39 @@ 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 numerical stages receive source IDs and the region, and control retrieval and -conversion of provider-defined vector handles without coordinator-built per-point -reference tables. The built-in strategies borrow vector matrices and member +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. -The numerical driver separates regional fitting (`Clusterer`) from placement -(`Assigner`). A clusterer declares its dimension and total number of centers for -the source count, then fills coordinator-owned centroid scratch. Fitting receives -only split parents; assignment receives the union of parents and retained neighbors. -A borrowed destination catalog combines the completed rows with unchanged retained -centroids; retained coordinates are not copied. The assigner receives the region, -source IDs, destinations, and regional writer directly. It computes complete -destination memberships before creating new lists, replacing retained memberships, -and retiring replaced sources. The index still selects two-means and full regional -nearest-centroid assignment. Restricted assignment can keep caller-supplied -candidate relationships in its own state, but the index does not yet preserve those -relationships or expose numerical strategy selection. +`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; From dbe1dd7f0a99434c40842b322a9d1ad4ad223cee Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Fri, 23 Oct 2026 16:22:43 -0400 Subject: [PATCH 32/33] some renaming --- diskann/src/ivf/build/mod.rs | 8 +-- diskann/src/ivf/build/twomeans.rs | 18 +++--- diskann/src/ivf/index.rs | 4 +- diskann/src/ivf/region.rs | 12 ++-- diskann/src/ivf/tests/mod.rs | 95 ++++++++++++++++--------------- diskann/src/ivf/tests/support.rs | 8 +-- diskann/src/ivf/traits.rs | 6 +- 7 files changed, 78 insertions(+), 73 deletions(-) diff --git a/diskann/src/ivf/build/mod.rs b/diskann/src/ivf/build/mod.rs index 2a1862f16c..7a1effdf55 100644 --- a/diskann/src/ivf/build/mod.rs +++ b/diskann/src/ivf/build/mod.rs @@ -9,14 +9,14 @@ use std::fmt::Debug; use crate::{ ANNError, - ivf::region::{RegionUpdate, RegionView}, + ivf::region::{Update, View}, }; pub(crate) mod twomeans; -pub(crate) use twomeans::TwoMeansSplitter; +pub(crate) use twomeans::TwoMeansSplit; /// Replace parents and reassign their members together with selected neighbors. -pub(crate) trait Split> { +pub(crate) trait Split> { type Error: Debug + Send + Sync + Into; /// Parents and neighbors must be distinct, disjoint filled sources. @@ -26,7 +26,7 @@ pub(crate) trait Split> { /// /// A failure may leave partial private updates. The caller must discard the /// operation rather than finalize it. - fn split>( + fn split>( &mut self, region: &R, update: &mut U, diff --git a/diskann/src/ivf/build/twomeans.rs b/diskann/src/ivf/build/twomeans.rs index 0c10629739..550f6a0e17 100644 --- a/diskann/src/ivf/build/twomeans.rs +++ b/diskann/src/ivf/build/twomeans.rs @@ -17,7 +17,7 @@ use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use super::Split; use crate::{ ANNError, ANNResult, - ivf::region::{AsMatRef, RegionUpdate, RegionView}, + ivf::region::{AsMatRef, Update, View}, utils::{TWO, VectorId}, }; @@ -50,12 +50,12 @@ impl From> for ANNError { /// 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 TwoMeansSplitter { +pub struct TwoMeansSplit { iterations: NonZeroUsize, dim: NonZeroUsize, } -impl TwoMeansSplitter { +impl TwoMeansSplit { pub fn new(iterations: NonZeroUsize, dim: NonZeroUsize) -> Self { Self { iterations, dim } } @@ -64,7 +64,7 @@ impl TwoMeansSplitter { fn fit(&self, region: &R, parents: &[L]) -> ANNResult> where L: VectorId, - R: RegionView, + R: View, for<'a> R::VectorsRef<'a>: AsMatRef<'a, f32>, { let rows = parents @@ -122,16 +122,16 @@ struct Target<'a, L> { // Split // /////////// -impl Split for TwoMeansSplitter +impl Split for TwoMeansSplit where I: VectorId, L: VectorId, - R: RegionView, + R: View, for<'a> R::VectorsRef<'a>: AsMatRef<'a, f32>, { type Error = ANNError; - fn split>( + fn split>( &mut self, region: &R, update: &mut U, @@ -150,7 +150,7 @@ where for &neighbor in neighbors { let view = region - .get_view(neighbor) + .get_list(neighbor) .ok_or(TwoMeansError::UnableToRetrieve(neighbor))?; targets.push(Target { @@ -170,7 +170,7 @@ where for &list in parents.iter().chain(neighbors.iter()) { let view = region - .get_view(list) + .get_list(list) .ok_or(TwoMeansError::UnableToRetrieve(list))?; let vectors = region diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index 4a04f62e4f..c748c64431 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -17,7 +17,7 @@ use crate::{ flat::DistancesUnordered, graph::SearchOutputBuffer, ivf::{ - build::{Split, TwoMeansSplitter}, + build::{Split, TwoMeansSplit}, region::AsMatRef, traits::BuildStrategy, }, @@ -127,7 +127,7 @@ impl IVFIndex { let (region, mut update) = accessor.fill(iter).await?; - TwoMeansSplitter::new(config.split_iterations, dim).split( + TwoMeansSplit::new(config.split_iterations, dim).split( ®ion, &mut update, &parents.iter().copied().collect::>(), diff --git a/diskann/src/ivf/region.rs b/diskann/src/ivf/region.rs index d5bb9a87a9..ef70947459 100644 --- a/diskann/src/ivf/region.rs +++ b/diskann/src/ivf/region.rs @@ -16,14 +16,14 @@ use crate::ANNResult; /// 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 View<'a, I, L> { +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 RegionView: Send { +pub trait View: Send { /// The vector representation consumed by numerical algorithms. type VectorsRef<'a>; @@ -33,7 +33,7 @@ pub trait RegionView: Send { Self: 'a; /// Borrow a filled source's centroid and complete operation-local membership. - fn get_view(&self, i: L) -> Option>; + fn get_list(&self, i: L) -> Option>; /// Borrow canonical rows in the same order as the source's member IDs. /// @@ -43,7 +43,7 @@ pub trait RegionView: Send { } /// Regional writes for one build operation. -pub trait RegionUpdate: Send { +pub trait Update: Send { /// Record a complete membership replacement. fn set_members(&mut self, i: L, ids: &[I]) -> ANNResult<()>; @@ -54,6 +54,10 @@ pub trait RegionUpdate: Send { fn create(&mut self, centroid: &[f32], members: &[I]) -> ANNResult<()>; } +///////////// +/// AsMat /// +///////////// + pub trait AsMatRef<'a, T> { fn as_view(self) -> rowmajor::Ref<'a, T>; } diff --git a/diskann/src/ivf/tests/mod.rs b/diskann/src/ivf/tests/mod.rs index 0e2c187a8f..f0420909f6 100644 --- a/diskann/src/ivf/tests/mod.rs +++ b/diskann/src/ivf/tests/mod.rs @@ -19,8 +19,8 @@ use self::{ }; use super::{ Config, IVFIndex, InsertStats, ReassignCount, - build::{Split, TwoMeansSplitter, twomeans::TwoMeansError}, - region::{AsMatRef, RegionUpdate, RegionView, View}, + build::{Split, TwoMeansSplit, twomeans::TwoMeansError}, + region::{AsMatRef, List, Update, View}, traits::{ BuildAccessor, BuildStrategy, HasIds, NearestList, Router, SearchAccessor, SearchStrategy, }, @@ -547,7 +547,7 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { assert_eq!(emitted, [(1, 4)]); { let (region, mut update) = require_send(build.fill([1].into_iter())).await.unwrap(); - let view = region.get_view(1).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]); @@ -560,7 +560,7 @@ async fn region_view_borrows_inputs_while_updates_are_recorded() { 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); - TwoMeansSplitter::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()) + 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]); @@ -597,8 +597,8 @@ async fn borrowed_updates_share_pending_memberships_and_publish_on_finish() { { let (region, mut update) = require_send(build.fill([1, 2].into_iter())).await.unwrap(); - let parent = region.get_view(1).unwrap(); - let neighbor = region.get_view(2).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()); @@ -955,8 +955,8 @@ fn regional_assignment_moves_between_neighbors_and_retains_empty_lists() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - TwoMeansSplitter::new(NonZeroUsize::new(5).unwrap(), ONE) + 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(); @@ -1004,12 +1004,12 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() { let (region, mut update) = build.fill([1, 2].into_iter()).await.unwrap(); assert!( - TwoMeansSplitter::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) .split(®ion, &mut update, &[1, 99], &[]) .is_err() ); assert!( - TwoMeansSplitter::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap(),) .split(®ion, &mut update, &[1], &[99]) .is_err() ); @@ -1018,7 +1018,7 @@ async fn split_only_edits_explicit_parents_and_missing_sources_leave_no_edits() assert!( matches!(changes.as_slice(), [(1, Membership::Append(members))] if members == &[10_005]) ); - TwoMeansSplitter::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()) + TwoMeansSplit::new(NonZeroUsize::new(5).unwrap(), NonZeroUsize::new(2).unwrap()) .split(®ion, &mut update, &[1], &[]) .unwrap(); assert_eq!(update.created.len(), 2); @@ -1067,8 +1067,8 @@ fn assignment_can_cross_parent_boundaries() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - TwoMeansSplitter::new( + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); + TwoMeansSplit::new( NonZeroUsize::new(10).unwrap(), NonZeroUsize::new(1).unwrap(), ) @@ -1096,8 +1096,8 @@ fn two_means_resets_empty_children_to_zero_and_prefers_first_ties() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - TwoMeansSplitter::new(ONE, ONE) + 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); @@ -1140,8 +1140,8 @@ fn two_means_prefers_neighbors_in_supplied_order_on_ties() { for neighbors in [[2, 3], [3, 2]] { let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - TwoMeansSplitter::new(ONE, ONE) + 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); @@ -1204,8 +1204,8 @@ fn two_means_reuses_splitter_across_regions() { let neighbors = [3, 4]; let mut full_memberships = PendingLists::default(); let mut full_created = Vec::new(); - let mut full = Update::new(&sources, 1, &mut full_memberships, &mut full_created); - let mut splitter = TwoMeansSplitter::new(ONE, ONE); + 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(); @@ -1222,7 +1222,7 @@ fn two_means_reuses_splitter_across_regions() { assert_eq!(full.created[2].members, [2, 5, 7]); let mut repeated_memberships = PendingLists::default(); let mut repeated_created = Vec::new(); - let mut repeated = Update::new( + let mut repeated = support::Update::new( &sources, 1, &mut repeated_memberships, @@ -1237,7 +1237,8 @@ fn two_means_reuses_splitter_across_regions() { } let mut smaller_memberships = PendingLists::default(); let mut smaller_created = Vec::new(); - let mut smaller = Update::new(&sources, 1, &mut smaller_memberships, &mut smaller_created); + 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]); @@ -1261,8 +1262,8 @@ fn empty_split_records_no_updates() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let error = TwoMeansSplitter::new(ONE, ONE) + 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!( @@ -1284,7 +1285,7 @@ fn region_vector_reborrows_share_source_storage() { }, )]); let region = Region::new(&sources); - let view = region.get_view(1).unwrap(); + 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(); @@ -1331,15 +1332,15 @@ fn two_means_handles_borrowed_batches() { reads: Cell, } - impl RegionView for BatchRegion<'_> { + impl View for BatchRegion<'_> { type VectorsRef<'a> = BatchRef<'a>; type Vectors<'a> = Batch<'a> where Self: 'a; - fn get_view(&self, list: u32) -> Option> { - self.inner.get_view(list) + fn get_list(&self, list: u32) -> Option> { + self.inner.get_list(list) } fn get_vectors(&self, list: u32) -> Option> { @@ -1381,8 +1382,8 @@ fn two_means_handles_borrowed_batches() { )); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - TwoMeansSplitter::new(ONE, ONE) + 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); @@ -1413,19 +1414,19 @@ fn split_accepts_native_batches_without_matrix_access() { vectors: &'a [[u8; 1]], } - impl RegionView for NativeRegion<'_> { + impl View for NativeRegion<'_> { type VectorsRef<'a> = &'a [[u8; 1]]; type Vectors<'a> = NativeBatch<'a> where Self: 'a; - fn get_view(&self, list: u32) -> Option> { - self.inner.get_view(list) + fn get_list(&self, list: u32) -> Option> { + self.inner.get_list(list) } fn get_vectors(&self, list: u32) -> Option> { - self.inner.get_view(list).map(|_| NativeBatch(self.vectors)) + self.inner.get_list(list).map(|_| NativeBatch(self.vectors)) } } @@ -1433,11 +1434,11 @@ fn split_accepts_native_batches_without_matrix_access() { impl Split for NativeSplit where - R: for<'a> RegionView = &'a [[u8; 1]]>, + R: for<'a> View = &'a [[u8; 1]]>, { type Error = ANNError; - fn split>( + fn split>( &mut self, region: &R, update: &mut U, @@ -1445,7 +1446,7 @@ fn split_accepts_native_batches_without_matrix_access() { neighbors: &[u32], ) -> ANNResult<()> { assert!(neighbors.is_empty()); - let view = region.get_view(parents[0]).unwrap(); + 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))?; @@ -1473,7 +1474,7 @@ fn split_accepts_native_batches_without_matrix_access() { assert!(std::ptr::eq(batch.reborrow(), vectors.as_slice())); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); + 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)]); @@ -1511,8 +1512,8 @@ fn two_means_matches_quantization_lloyds() { for iterations in [1, 5] { let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 2, &mut memberships, &mut created); - TwoMeansSplitter::new( + let mut update = support::Update::new(&sources, 2, &mut memberships, &mut created); + TwoMeansSplit::new( NonZeroUsize::new(iterations).unwrap(), NonZeroUsize::new(2).unwrap(), ) @@ -1549,8 +1550,8 @@ fn two_means_reports_insufficient_points_before_writing() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 2, &mut memberships, &mut created); - let error = TwoMeansSplitter::new(ONE, NonZeroUsize::new(2).unwrap()) + 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!( @@ -1575,8 +1576,8 @@ fn two_means_missing_vectors_preserves_typed_error_without_updates() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - let error = TwoMeansSplitter::new(ONE, ONE) + 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!( @@ -1610,9 +1611,9 @@ fn two_means_reports_missing_neighbors_before_writing() { for ids in [vec![99], vec![2, 99]] { let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); + let mut update = support::Update::new(&sources, 1, &mut memberships, &mut created); let region = Region::new(&sources); - let error = TwoMeansSplitter::new(ONE, ONE) + let error = TwoMeansSplit::new(ONE, ONE) .split(®ion, &mut update, &[1], &ids) .unwrap_err(); assert!(matches!( @@ -1647,8 +1648,8 @@ fn regional_assignment_prefers_first_when_all_distances_overflow() { let region = Region::new(&sources); let mut memberships = PendingLists::default(); let mut created = Vec::new(); - let mut update = Update::new(&sources, 1, &mut memberships, &mut created); - TwoMeansSplitter::new(ONE, ONE) + 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]); diff --git a/diskann/src/ivf/tests/support.rs b/diskann/src/ivf/tests/support.rs index 273c73bfb1..6da384a631 100644 --- a/diskann/src/ivf/tests/support.rs +++ b/diskann/src/ivf/tests/support.rs @@ -224,15 +224,15 @@ impl<'a> Region<'a> { } } -impl RegionView for Region<'_> { +impl View for Region<'_> { type VectorsRef<'a> = rowmajor::Ref<'a, f32>; type Vectors<'a> = rowmajor::Ref<'a, f32> where Self: 'a; - fn get_view(&self, list: u32) -> Option> { - self.sources.get(&list).map(|source| View { + fn get_list(&self, list: u32) -> Option> { + self.sources.get(&list).map(|source| List { id: list, centroid: &source.centroid, members: &source.members, @@ -246,7 +246,7 @@ impl RegionView for Region<'_> { } } -impl RegionUpdate for Update<'_> { +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")); diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index be113e6209..c79d515ef1 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -9,7 +9,7 @@ use std::num::NonZeroUsize; use diskann_utils::future::SendFuture; -use super::region::{RegionUpdate, RegionView}; +use super::region::{Update, View}; use crate::{ ANNResult, error::StandardError, flat::DistancesUnordered, provider::ExecutionContext, utils::VectorId, @@ -65,11 +65,11 @@ pub trait BuildAccessor: Send + Sized + HasIds { type VectorsRef<'a>; - type Region<'a>: for<'b> RegionView = Self::VectorsRef<'b>> + type Region<'a>: for<'b> View = Self::VectorsRef<'b>> where Self: 'a; - type Update<'a>: RegionUpdate + type Update<'a>: Update where Self: 'a; From 06ce7fce8ac0231ec6d4cf42ad36f61d399d32df Mon Sep 17 00:00:00 2001 From: Aditya Krishnan Date: Fri, 23 Oct 2026 17:52:23 -0400 Subject: [PATCH 33/33] search accessor, accumulator --- diskann/src/ivf/index.rs | 33 +++---- diskann/src/ivf/mod.rs | 14 ++- diskann/src/ivf/tests/mod.rs | 132 ++++++++++++++++++++++++-- diskann/src/ivf/tests/support.rs | 156 ++++++------------------------- diskann/src/ivf/traits.rs | 75 +++++++++------ rfcs/01187-incremental-ivf.md | 4 +- 6 files changed, 223 insertions(+), 191 deletions(-) diff --git a/diskann/src/ivf/index.rs b/diskann/src/ivf/index.rs index c748c64431..0b87365fb4 100644 --- a/diskann/src/ivf/index.rs +++ b/diskann/src/ivf/index.rs @@ -10,18 +10,17 @@ use std::{collections::BTreeSet, marker::PhantomData, num::NonZeroUsize}; use diskann_utils::future::SendFuture; use thiserror::Error; -use super::traits::{BuildAccessor, NearestList, Router, SearchAccessor, SearchStrategy}; +use super::traits::{BuildAccessor, NearestList, SearchAccessor, SearchStrategy}; use crate::{ ANNResult, - error::{ErrorExt, IntoANNResult}, - flat::DistancesUnordered, + error::IntoANNResult, graph::SearchOutputBuffer, ivf::{ build::{Split, TwoMeansSplit}, region::AsMatRef, traits::BuildStrategy, }, - neighbor::{Neighbor, NeighborPriorityQueue}, + neighbor::NeighborPriorityQueue, provider::ExecutionContext, }; @@ -144,6 +143,11 @@ impl IVFIndex { } } + /// 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, @@ -160,23 +164,10 @@ impl IVFIndex { { async move { let accessor = strategy.search_accessor(context, query).into_ann_result()?; - let mut lists = Vec::new(); - { - let router = accessor.router(); - router - .route(query, nprobe, |id, scan| lists.push((*id, scan))) - .await?; - } - - let mut best = NeighborPriorityQueue::new(k.get()); - - for (_, scan) in lists.iter_mut() { - scan.distances_unordered(|id, dist| best.insert(Neighbor::new(id, dist))) - .await - .escalate("why")?; - } - - Ok(output.extend(best.iter())) + 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 index c18a780d1e..0e7f0c6cbb 100644 --- a/diskann/src/ivf/mod.rs +++ b/diskann/src/ivf/mod.rs @@ -7,16 +7,16 @@ //! //! 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::RegionView`], -//! and regional writes through [`region::RegionUpdate`]. +//! 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. -//! [`build::Split`] owns fitting, assignment, and regional updates as one operation. -//! The built-in [`build::TwoMeansSplitter`] fits each parent independently, then assigns +//! 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. //! @@ -28,12 +28,16 @@ //! 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. Search scans routed lists and collects their nearest points. +//! recursively split. //! Initialization, deletion, and colocation are not implemented. pub mod build; diff --git a/diskann/src/ivf/tests/mod.rs b/diskann/src/ivf/tests/mod.rs index f0420909f6..6a6a82d109 100644 --- a/diskann/src/ivf/tests/mod.rs +++ b/diskann/src/ivf/tests/mod.rs @@ -3,7 +3,12 @@ * Licensed under the MIT license. */ -use std::{cell::Cell, collections::BTreeMap, future::Future, num::NonZeroUsize}; +use std::{ + cell::Cell, + collections::{BTreeMap, BTreeSet}, + future::Future, + num::NonZeroUsize, +}; use diskann_utils::{ Reborrow, @@ -22,15 +27,15 @@ use super::{ build::{Split, TwoMeansSplit, twomeans::TwoMeansError}, region::{AsMatRef, List, Update, View}, traits::{ - BuildAccessor, BuildStrategy, HasIds, NearestList, Router, SearchAccessor, SearchStrategy, + Accumulator, BuildAccessor, BuildStrategy, HasIds, NearestList, SearchAccessor, + SearchStrategy, }, }; use crate::{ ANNError, ANNResult, - flat::DistancesUnordered, graph::{IdDistance, SearchOutputBuffer}, - neighbor::{BackInserter, Neighbor}, - provider::{DefaultContext, HasId}, + neighbor::{BackInserter, Neighbor, NeighborPriorityQueue}, + provider::DefaultContext, utils::ONE, }; @@ -49,7 +54,7 @@ async fn routing_callbacks_emit_ranked_results() { for limit in [1, 2, usize::MAX] { let mut emitted = Vec::new(); require_send(router.route( - (9, &[6u8, 0][..]), + &[6u8, 0][..], NonZeroUsize::new(limit).unwrap(), |id, distance| emitted.push((*id, distance)), )) @@ -152,7 +157,7 @@ async fn routing_preserves_native_inputs_and_id_order() { let router = build.nearest(); for (point, expected) in points.iter().zip([2, 1, 1]) { let mut selected = Vec::new(); - require_send(router.route(*point, ONE, |&id, _| selected.push(id))) + require_send(router.route(point.1, ONE, |&id, _| selected.push(id))) .await .unwrap(); assert_eq!(selected, [expected]); @@ -707,6 +712,108 @@ async fn uninitialized_insertion_fails() { 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(); @@ -841,6 +948,7 @@ async fn search_failures_do_not_publish_partial_results() { ("nonfinite", "invalid query vector"), ("route", "injected routing failure"), ("scan", "missing point 10003"), + ("later_scan", "missing point 10001"), ] { let mut store = seeded(); if case == "route" { @@ -849,6 +957,9 @@ async fn search_failures_do_not_publish_partial_results() { 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" { @@ -859,13 +970,18 @@ async fn search_failures_do_not_publish_partial_results() { } 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(), - ONE, + nprobe, &mut output, )) .await diff --git a/diskann/src/ivf/tests/support.rs b/diskann/src/ivf/tests/support.rs index 6da384a631..ba75e67462 100644 --- a/diskann/src/ivf/tests/support.rs +++ b/diskann/src/ivf/tests/support.rs @@ -92,46 +92,15 @@ impl Coarse<'_> { } selected.len() } -} - -impl + Sync> Router<(u32, &[T])> for Coarse<'_> { - type Id = u32; - type Ret = f32; - - fn len(&self) -> usize { - self.stored.len() - } - - fn route( + pub(super) fn route( &self, - query: (u32, &[T]), + query: &[T], nprobe: NonZeroUsize, emit: F, ) -> impl SendFuture> where - F: FnMut(&Self::Id, f32) + Send, - { - async move { - if self.fail_route { - return Err(ANNError::message("injected routing failure")); - } - self.select(query.1, nprobe, |_| true, emit); - Ok(()) - } - } -} - -impl + Sync> Router<&[T]> for Coarse<'_> { - type Id = u32; - type Ret = f32; - - fn len(&self) -> usize { - self.stored.len() - } - - fn route(&self, query: &[T], nprobe: NonZeroUsize, emit: F) -> impl SendFuture> - where - F: FnMut(&Self::Id, f32) + Send, + T: Copy + Into + Sync, + F: FnMut(&u32, f32) + Send, { async move { if self.fail_route { @@ -325,8 +294,7 @@ impl<'p, 'b, 'v, T: Copy + Into + Send + Sync> BuildAccessor<&'b [(u32, &'v .ok_or_else(|| ANNError::message("cannot insert an empty batch"))?; let mut staged = BTreeMap::new(); let mut routed = BTreeMap::>::new(); - for &point in points { - let (external, vector) = point; + for &(external, vector) in points { if vector.len() != self.store.dim || !vector.iter().all(|&value| value.into().is_finite()) { @@ -345,7 +313,7 @@ impl<'p, 'b, 'v, T: Copy + Into + Send + Sync> BuildAccessor<&'b [(u32, &'v { let router = self.nearest(); router - .route(point, ONE, |&list, _| { + .route(vector, ONE, |&list, _| { count += 1; routed.entry(list).or_default().push(id); }) @@ -525,33 +493,15 @@ pub(super) struct Search<'p, 'q, T> { pub(super) query: &'q [T], } -pub(super) struct SearchRouter<'a, T> { - pub(super) store: &'a Store, - pub(super) query: &'a [T], -} - -pub(super) struct ListScan<'a, T> { - pub(super) store: &'a Store, - pub(super) list: u32, - pub(super) query: &'a [T], -} - -impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { +impl HasIds for Search<'_, '_, T> { type Id = u32; - type Ret = ListScan<'a, T>; - - fn len(&self) -> usize { - self.store.lists.len() - } + type ListId = u32; +} - fn route( - &self, - query: &[T], - nprobe: NonZeroUsize, - mut emit: F, - ) -> impl SendFuture> +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 - F: FnMut(&Self::Id, Self::Ret) + Send, + A: Accumulator, { async move { if self.store.fail_route { @@ -562,79 +512,31 @@ impl<'a, T: Copy + Into + Sync> Router<&[T]> for SearchRouter<'a, T> { fail_route: false, fail_nearest: false, }; - coarse.select( - query, - nprobe, - |_| true, - |id, _| { - emit( - id, - ListScan { - store: self.store, - list: *id, - query: self.query, - }, - ); - }, - ); - Ok(()) - } - } -} + let mut selected = Vec::new(); + coarse.select(self.query, nprobe, |_| true, |id, _| selected.push(*id)); -impl HasId for ListScan<'_, T> { - type Id = u32; -} - -impl + Sync> DistancesUnordered for ListScan<'_, T> { - type Error = ANNError; - - fn distances_unordered(&mut self, mut emit: F) -> impl SendFuture> - where - F: FnMut(Self::Id, f32) + Send, - { - async move { - let list = self - .store - .lists - .get(&self.list) - .ok_or_else(|| ANNError::message(format!("unknown list {}", self.list)))?; - for id in &list.members { - let vector = self + for list in selected { + let source = self .store - .points - .get(id) - .ok_or_else(|| ANNError::message(format!("missing point {id}")))?; - emit(*id, squared_distance(self.query, vector)); + .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 HasIds for Search<'_, '_, T> { - type Id = u32; - type ListId = u32; -} - -impl<'q, T: Copy + Into + Send + Sync> SearchAccessor<&'q [T]> for Search<'_, 'q, T> { - type Router<'a> - = SearchRouter<'a, T> - where - Self: 'a; - type Dist<'a> - = ListScan<'a, T> - where - Self: 'a; - - fn router(&self) -> Self::Router<'_> { - SearchRouter { - store: &self.store, - query: self.query, - } - } -} - impl<'a, 'q, T: Copy + Into + Send + Sync> SearchStrategy<'a, &'q [T], DefaultContext> for Strategy { diff --git a/diskann/src/ivf/traits.rs b/diskann/src/ivf/traits.rs index c79d515ef1..7a9bafc7ce 100644 --- a/diskann/src/ivf/traits.rs +++ b/diskann/src/ivf/traits.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Storage boundaries for IVF construction. +//! Storage boundaries for IVF construction and search. use std::num::NonZeroUsize; @@ -11,7 +11,10 @@ use diskann_utils::future::SendFuture; use super::region::{Update, View}; use crate::{ - ANNResult, error::StandardError, flat::DistancesUnordered, provider::ExecutionContext, + ANNResult, + error::StandardError, + neighbor::{Neighbor, NeighborPriorityQueue, NeighborPriorityQueueIdType}, + provider::ExecutionContext, utils::VectorId, }; @@ -20,23 +23,6 @@ pub trait HasIds { type ListId: VectorId; } -/// Centroid selection within one build operation. -pub trait Router: Send { - type Id; - type Ret; - - fn len(&self) -> usize; - - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Select `min(nprobe, len())` distinct live lists, nearest first. - fn route(&self, query: T, nprobe: NonZeroUsize, emit: F) -> impl SendFuture> - where - F: FnMut(&Self::Id, Self::Ret) + Send; -} - pub trait NearestList: Send { type Id; /// Select `nprobe` distinct live lists nearest to `id`'s centroid, that pass the filter. @@ -120,7 +106,7 @@ pub trait BuildStrategy<'a, B, C: ExecutionContext>: Send + Sync + HasIds { //////////// pub trait SearchStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds { - /// [`BuildAccessor`] used to insert new points, and performs splits if necessary. + /// Query-bound accessor responsible for list selection and scanning. type Accessor: SearchAccessor; type Error: StandardError; @@ -128,14 +114,47 @@ pub trait SearchStrategy<'a, T: Sync, C: ExecutionContext>: Send + Sync + HasIds fn search_accessor(&'a self, context: &'a C, query: T) -> Result; } -pub trait SearchAccessor: Send + Sized + HasIds { - type Router<'a>: Router> - where - Self: 'a; +/// 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; - type Dist<'a>: DistancesUnordered - where - Self: 'a; + /// Consider one candidate for inclusion in the results. + fn push(&mut self, id: I, distance: f32); - fn router(&self) -> Self::Router<'_>; + /// 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/rfcs/01187-incremental-ivf.md b/rfcs/01187-incremental-ivf.md index aa6781d048..57f551a5d4 100644 --- a/rfcs/01187-incremental-ivf.md +++ b/rfcs/01187-incremental-ivf.md @@ -13,8 +13,8 @@ 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 a `Router`, and nearest-list selection through - `RouterAndNearest`. +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