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 1/3] 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 000000000..b0f331f23 --- /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 2/3] 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 000000000..f3e846ec6 --- /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 000000000..d7b2d42c5 --- /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 000000000..f93b4c9c8 --- /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 000000000..f13e97409 --- /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 000000000..2334908ad --- /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 000000000..4438a0723 --- /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 000000000..2eb9dd496 --- /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 000000000..3027b4d8e --- /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 000000000..30ac05b1c --- /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 000000000..a19e74521 --- /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 000000000..983bc5bef --- /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 9c1f6ac76..14c20bfdf 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 3/3] 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 b0f331f23..000000000 --- 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`).