From 6b6c63c2c97773ae7634c3306ed724a41a3eb7b2 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Mon, 27 Apr 2026 22:13:05 -0700 Subject: [PATCH 01/10] first draft --- diskann/src/flat/index.rs | 106 +++++++++++ diskann/src/flat/iterator.rs | 91 ++++++++++ diskann/src/flat/mod.rs | 50 ++++++ diskann/src/flat/post_process.rs | 73 ++++++++ diskann/src/flat/search.rs | 35 ++++ diskann/src/flat/stats.rs | 16 ++ diskann/src/flat/strategy.rs | 77 ++++++++ diskann/src/flat/test/mod.rs | 293 +++++++++++++++++++++++++++++++ diskann/src/lib.rs | 1 + rfcs/00000-flat-search.md | 217 +++++++++++++++++++++++ 10 files changed, 959 insertions(+) create mode 100644 diskann/src/flat/index.rs create mode 100644 diskann/src/flat/iterator.rs create mode 100644 diskann/src/flat/mod.rs create mode 100644 diskann/src/flat/post_process.rs create mode 100644 diskann/src/flat/search.rs create mode 100644 diskann/src/flat/stats.rs create mode 100644 diskann/src/flat/strategy.rs create mode 100644 diskann/src/flat/test/mod.rs create mode 100644 rfcs/00000-flat-search.md diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs new file mode 100644 index 000000000..bebc1ea6a --- /dev/null +++ b/diskann/src/flat/index.rs @@ -0,0 +1,106 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! [`FlatIndex`] — the index wrapper for flat search. + +use std::marker::PhantomData; +use std::num::NonZeroUsize; + +use diskann_utils::future::SendFuture; +use diskann_vector::PreprocessedDistanceFunction; + +use crate::{ + ANNResult, + error::IntoANNResult, + flat::{ + FlatIterator, FlatPostProcess, FlatSearchStats, FlatSearchStrategy, + }, + graph::SearchOutputBuffer, + neighbor::{Neighbor, NeighborPriorityQueue}, + provider::DataProvider, +}; + +/// A `'static` thin wrapper around a [`DataProvider`] used for flat search. +/// +/// The provider is owned by the index. The index is constructed once at process startup and +/// shared across requests; per-query state lives in the [`crate::flat::FlatIterator`] that +/// the [`crate::flat::FlatSearchStrategy`] produces. +#[derive(Debug)] +pub struct FlatIndex { + /// The backing provider. + pub provider: P, + _marker: PhantomData P>, +} + +impl FlatIndex

{ + /// Construct a new [`FlatIndex`] around `provider`. + pub fn new(provider: P) -> Self { + Self { + provider, + _marker: PhantomData, + } + } + + /// Borrow the underlying provider. + pub fn provider(&self) -> &P { + &self.provider + } + + /// Brute-force k-nearest-neighbor flat search. + /// + /// Streams every element produced by the strategy's iterator through the query + /// computer, keeps the best `k` candidates in a [`NeighborPriorityQueue`], and hands + /// the survivors to the post-processor. + /// + /// # Arguments + /// - `k`: number of nearest neighbors to return. + /// - `strategy`: produces the per-query iterator and the query computer. + /// - `processor`: post-processes the survivor candidates into the output type. + /// - `context`: per-request context threaded through to the provider. + /// - `query`: the query. + /// - `output`: caller-owned output buffer. + pub fn knn_search( + &self, + k: NonZeroUsize, + strategy: &S, + processor: &PP, + context: &P::Context, + query: &T, + output: &mut OB, + ) -> impl SendFuture> + where + S: FlatSearchStrategy, + T: ?Sized + Sync, + O: Send, + OB: SearchOutputBuffer + Send + ?Sized, + PP: for<'a> FlatPostProcess, T, O> + Send + Sync, + { + async move { + let mut iter = strategy + .create_iter(&self.provider, context) + .into_ann_result()?; + let computer = strategy.build_query_computer(query).into_ann_result()?; + + let k = k.get(); + let mut queue = NeighborPriorityQueue::new(k); + let mut cmps: u32 = 0; + + iter.on_elements_unordered(|id, element| { + let dist = computer.evaluate_similarity(element); + cmps += 1; + queue.insert(Neighbor::new(id, dist)); + }) + .await + .into_ann_result()?; + + let result_count = processor + .post_process(&mut iter, query, queue.iter().take(k), output) + .await + .into_ann_result()? as u32; + + Ok(FlatSearchStats { cmps, result_count }) + } + } +} diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs new file mode 100644 index 000000000..02ef66f99 --- /dev/null +++ b/diskann/src/flat/iterator.rs @@ -0,0 +1,91 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! [`FlatIterator`] — the sequential access primitive for flat search. + +use diskann_utils::{Reborrow, future::SendFuture}; + +use crate::{error::StandardError, provider::HasId}; + +/// A lending, asynchronous iterator over the elements of a flat index. +/// +/// `FlatIterator` is the streaming counterpart to [`crate::provider::Accessor`]. Where an +/// accessor exposes random retrieval by id, a flat iterator exposes a *sequential* walk — +/// each call to [`Self::next`] advances an internal cursor and yields the next element. +/// +/// Like [`crate::provider::Accessor::get_element`], advancing the cursor is **async**: it +/// may need to await an I/O fetch (e.g., reading the next disk page, awaiting a network +/// response, etc.). Iterators backed by purely in-memory data should return a ready +/// future. +/// +/// The iterator is responsible for: +/// - Choosing the iteration order (buffer-sequential, hash-walked, partitioned, …). +/// - Skipping items that should not be visible to the algorithm (deleted, obsolete, …). +/// - Holding any borrows / locks needed to keep the underlying storage alive. +/// +/// Algorithms see only `(Id, ElementRef)` pairs and treat the stream as opaque. +/// +/// # `Element` vs `ElementRef` +/// +/// Same pattern as [`crate::provider::Accessor`]: +/// +/// - `Element<'a>` is the type returned by `next`. Its lifetime is bound to the iterator +/// borrow at the call site, so only one element is live at a time. +/// - `ElementRef<'a>` is an unconstrained-lifetime reborrow used in distance-function +/// bounds. Required to keep [HRTB](https://doc.rust-lang.org/nomicon/hrtb.html) bounds +/// on query computers from forcing `Self: 'static`. +/// +/// # Hot path +/// +/// Algorithms drive the scan via [`Self::on_elements_unordered`]. The provided +/// implementation simply loops over [`Self::next`]; iterators that can amortize +/// per-element cost (prefetching the next chunk, batching distance computation, +/// performing SIMD-friendly bulk reads) should override it. +pub trait FlatIterator: HasId + Send + Sync { + /// A reference to a yielded element with an unconstrained lifetime, suitable for + /// distance-function HRTB bounds. + type ElementRef<'a>; + + /// The concrete element returned by [`Self::next`]. Reborrows to [`Self::ElementRef`]. + type Element<'a>: for<'b> Reborrow<'b, Target = Self::ElementRef<'b>> + Send + Sync + where + Self: 'a; + + /// The error type yielded by [`Self::next`] and [`Self::on_elements_unordered`]. + type Error: StandardError; + + /// Advance the iterator and asynchronously yield the next `(id, element)` pair. + /// + /// Returns `Ok(None)` when the scan is exhausted. The yielded element borrows from + /// the iterator and is invalidated by the next call to `next`. + fn next( + &mut self, + ) -> impl SendFuture)>, Self::Error>>; + + /// Drive the entire scan, invoking `f` for each yielded element. + /// + /// The default implementation loops over [`Self::next`]. Implementations that benefit + /// from bulk dispatch (prefetching, batched SIMD distance computation, etc.) should + /// override this method. + /// + /// The order of invocation is unspecified and may differ between calls. The closure + /// `f` is **synchronous**; if you need to await inside the per-element handler, drive + /// the iterator manually with [`Self::next`]. + fn on_elements_unordered( + &mut self, + mut f: F, + ) -> impl SendFuture> + where + F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>), + { + async move { + while let Some((id, element)) = self.next().await? { + f(id, element.reborrow()); + } + Ok(()) + } + } +} + diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs new file mode 100644 index 000000000..0cd9dcc33 --- /dev/null +++ b/diskann/src/flat/mod.rs @@ -0,0 +1,50 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Sequential ("flat") search infrastructure. +//! +//! This module is the streaming counterpart to the random-access [`crate::provider::Accessor`] +//! family. It is designed for backends whose natural access pattern is a one-pass scan over +//! their data — for example append-only buffered stores, on-disk shards streamed via I/O, +//! or any provider where random access is significantly more expensive than sequential. +//! +//! # Architecture +//! +//! The module mirrors the layering used by graph search: +//! +//! | Graph (random access) | Flat (sequential) | +//! | :------------------------------------ | :-------------------------------- | +//! | [`crate::provider::DataProvider`] | [`crate::provider::DataProvider`] | +//! | [`crate::graph::DiskANNIndex`] | [`FlatIndex`] | +//! | [`crate::provider::Accessor`] | [`FlatIterator`] | +//! | [`crate::graph::glue::SearchStrategy`] | [`FlatSearchStrategy`] | +//! | [`crate::graph::glue::SearchPostProcess`] | [`FlatPostProcess`] | +//! | [`crate::graph::Search`] | [`FlatIndex::knn_search`] | +//! +//! # Hot loop +//! +//! Algorithms drive the scan via [`FlatIterator::next`] (lending iterator) or override +//! [`FlatIterator::on_elements_unordered`] when batching/prefetching wins. The default +//! implementation of `on_elements_unordered` simply loops over `next`. +//! +//! See [`FlatIndex::knn_search`] for the canonical brute-force k-NN algorithm built on these +//! primitives. + +pub mod index; +pub mod iterator; +pub mod post_process; +pub mod search; +pub mod stats; +pub mod strategy; + +#[cfg(any(test, feature = "testing"))] +pub mod test; + +pub use index::FlatIndex; +pub use iterator::FlatIterator; +pub use post_process::{CopyFlatIds, FlatPostProcess}; +pub use search::{KnnFlatError, validate_k}; +pub use stats::FlatSearchStats; +pub use strategy::FlatSearchStrategy; diff --git a/diskann/src/flat/post_process.rs b/diskann/src/flat/post_process.rs new file mode 100644 index 000000000..38267fef6 --- /dev/null +++ b/diskann/src/flat/post_process.rs @@ -0,0 +1,73 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! [`FlatPostProcess`] — terminal stage of the flat search pipeline. + +use diskann_utils::future::SendFuture; + +use crate::{ + error::StandardError, flat::FlatIterator, graph::SearchOutputBuffer, neighbor::Neighbor, provider::HasId, +}; + +/// Hydrate / filter / transform the survivor candidates produced by a flat search and +/// write them into an output buffer. +/// +/// This is the flat counterpart to [`crate::graph::glue::SearchPostProcess`]. Processors +/// receive `&mut S` so they can consult any iterator-owned lookup state (e.g., an +/// `Id -> rich-record` table built up during the scan) when assembling outputs. As with +/// the graph counterpart, [`Self::post_process`] is **async** so that processors can +/// hydrate via I/O without blocking. +/// +/// The `O` type parameter lets callers pick the output element type (raw `(Id, f32)` +/// pairs, fully hydrated hits, etc.). +pub trait FlatPostProcess::Id> +where + S: FlatIterator, + T: ?Sized, +{ + /// Errors yielded by [`Self::post_process`]. + type Error: StandardError; + + /// Consume `candidates` (in distance order) and write at most `k` results into + /// `output`. Returns the number of results written. + fn post_process( + &self, + iter: &mut S, + query: &T, + candidates: I, + output: &mut B, + ) -> impl SendFuture> + where + I: Iterator> + Send, + B: SearchOutputBuffer + Send + ?Sized; +} + +/// A trivial [`FlatPostProcess`] that copies each `(Id, distance)` pair straight into the +/// output buffer. +#[derive(Debug, Default, Clone, Copy)] +pub struct CopyFlatIds; + +impl FlatPostProcess for CopyFlatIds +where + S: FlatIterator, + T: ?Sized, +{ + type Error = crate::error::Infallible; + + fn post_process( + &self, + _iter: &mut S, + _query: &T, + candidates: I, + output: &mut B, + ) -> impl SendFuture> + where + I: Iterator::Id>> + Send, + B: SearchOutputBuffer<::Id> + Send + ?Sized, + { + let count = output.extend(candidates.map(|n| (n.id, n.distance))); + std::future::ready(Ok(count)) + } +} diff --git a/diskann/src/flat/search.rs b/diskann/src/flat/search.rs new file mode 100644 index 000000000..b51b561f1 --- /dev/null +++ b/diskann/src/flat/search.rs @@ -0,0 +1,35 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Error types for flat search parameter validation. + +use std::num::NonZeroUsize; + +use thiserror::Error; + +use crate::{ANNError, ANNErrorKind}; + +/// Errors raised when validating flat search parameters. +#[derive(Debug, Error)] +pub enum KnnFlatError { + /// `k` was zero. + #[error("k cannot be zero")] + KZero, +} + +impl From for ANNError { + #[track_caller] + fn from(err: KnnFlatError) -> Self { + Self::new(ANNErrorKind::IndexError, err) + } +} + +/// Validate and wrap a `k` value as [`NonZeroUsize`]. +/// +/// This is a convenience for callers that want to validate `k` before passing it to +/// [`FlatIndex::knn_search`](crate::flat::FlatIndex::knn_search). +pub fn validate_k(k: usize) -> Result { + NonZeroUsize::new(k).ok_or(KnnFlatError::KZero) +} diff --git a/diskann/src/flat/stats.rs b/diskann/src/flat/stats.rs new file mode 100644 index 000000000..faf14f888 --- /dev/null +++ b/diskann/src/flat/stats.rs @@ -0,0 +1,16 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Statistics returned by a flat search. + +/// Statistics collected during a single flat search invocation. +#[derive(Debug, Clone, Copy, Default)] +pub struct FlatSearchStats { + /// Number of distance computations performed (i.e., elements visited by the scanner). + pub cmps: u32, + + /// Number of results written into the caller-provided output buffer. + pub result_count: u32, +} diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs new file mode 100644 index 000000000..8b0f5e445 --- /dev/null +++ b/diskann/src/flat/strategy.rs @@ -0,0 +1,77 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! [`FlatSearchStrategy`] — glue between [`DataProvider`] and per-query [`FlatIterator`]s. + +use diskann_vector::PreprocessedDistanceFunction; + +use crate::{ + error::StandardError, + flat::FlatIterator, + provider::DataProvider, +}; + +/// Per-call configuration that knows how to construct a [`FlatIterator`] for a provider +/// and how to pre-process queries of type `T` into a distance computer. +/// +/// `FlatSearchStrategy` is the flat counterpart to [`crate::graph::glue::SearchStrategy`]. +/// A strategy instance is stateless config — typically constructed at the call site, used +/// for one search, and dropped. +/// +/// # Why two methods? +/// +/// - [`Self::create_iter`] is query-independent and may be called multiple times per +/// request (e.g., once per parallel query in a batched search). +/// - [`Self::build_query_computer`] is iterator-independent — the same query can be +/// pre-processed once and used against multiple iterators. +/// +/// Both methods may borrow from the strategy itself. +/// +/// # Type parameters +/// +/// - `Provider`: the [`DataProvider`] that backs the index. +/// - `T`: the query type. Often `[E]` for vector queries; can be any `?Sized` type. +pub trait FlatSearchStrategy: Send + Sync +where + P: DataProvider, + T: ?Sized, +{ + /// The iterator type produced by [`Self::create_iter`]. Borrows from `self` and the + /// provider. + type Iter<'a>: FlatIterator + where + Self: 'a, + P: 'a; + + /// The query computer produced by [`Self::build_query_computer`]. + /// + /// The HRTB on `ElementRef` ensures the same computer can score every element yielded + /// by every lifetime of `Iter`. Two lifetimes are needed: `'a` for the iterator + /// instance and `'b` for the reborrowed element. + type QueryComputer: for<'a, 'b> PreprocessedDistanceFunction< + as FlatIterator>::ElementRef<'b>, + f32, + > + Send + + Sync + + 'static; + + /// The error type for both factory methods. + type Error: StandardError; + + /// Construct a fresh iterator over `provider` for the given request `context`. + /// + /// This is where lock acquisition, snapshot pinning, and any other per-query setup + /// should happen. The returned iterator owns whatever borrows / guards it needs to + /// remain valid until it is dropped. + fn create_iter<'a>( + &'a self, + provider: &'a P, + context: &'a P::Context, + ) -> Result, Self::Error>; + + /// Pre-process a query into a [`Self::QueryComputer`] usable for distance computation + /// against any iterator produced by [`Self::create_iter`]. + fn build_query_computer(&self, query: &T) -> Result; +} diff --git a/diskann/src/flat/test/mod.rs b/diskann/src/flat/test/mod.rs new file mode 100644 index 000000000..61307ab61 --- /dev/null +++ b/diskann/src/flat/test/mod.rs @@ -0,0 +1,293 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Trivial in-memory provider, iterator, and strategy used for unit-testing the flat +//! search infrastructure. +//! +//! This is intentionally simple: vectors live in a `Vec>`, ids are `u32`, and +//! distance is squared Euclidean. It exists so the trait shapes in [`crate::flat`] can be +//! exercised end-to-end without dragging in any provider-side machinery. + +use diskann_utils::future::SendFuture; +use diskann_vector::PreprocessedDistanceFunction; +use thiserror::Error; + +use crate::{ + always_escalate, + ANNError, ANNErrorKind, + flat::{FlatIterator, FlatPostProcess, FlatSearchStrategy}, + graph::SearchOutputBuffer, + neighbor::Neighbor, + provider::{DataProvider, DefaultContext, HasId, NoopGuard}, +}; + +/// Trivial flat provider holding a list of fixed-dimension `f32` vectors. +#[derive(Debug)] +pub struct InMemoryFlatProvider { + pub dim: usize, + pub vectors: Vec>, +} + +impl InMemoryFlatProvider { + pub fn new(dim: usize, vectors: Vec>) -> Self { + Self { dim, vectors } + } +} + +#[derive(Debug, Error)] +#[error("invalid vector id {0}")] +pub struct InMemoryProviderError(u32); + +impl From for ANNError { + #[track_caller] + fn from(err: InMemoryProviderError) -> Self { + ANNError::new(ANNErrorKind::IndexError, err) + } +} + +always_escalate!(InMemoryProviderError); + +impl DataProvider for InMemoryFlatProvider { + type Context = DefaultContext; + type InternalId = u32; + type ExternalId = u32; + type Error = InMemoryProviderError; + type Guard = NoopGuard; + + fn to_internal_id( + &self, + _context: &Self::Context, + gid: &u32, + ) -> Result { + if (*gid as usize) < self.vectors.len() { + Ok(*gid) + } else { + Err(InMemoryProviderError(*gid)) + } + } + + fn to_external_id( + &self, + _context: &Self::Context, + id: u32, + ) -> Result { + if (id as usize) < self.vectors.len() { + Ok(id) + } else { + Err(InMemoryProviderError(id)) + } + } +} + +/// Sequential iterator over [`InMemoryFlatProvider`]. +pub struct InMemoryIterator<'a> { + vectors: &'a [Vec], + cursor: u32, +} + +#[derive(Debug, Error)] +#[error("in-memory iterator does not error")] +pub struct InMemoryIteratorError; + +impl From for ANNError { + #[track_caller] + fn from(err: InMemoryIteratorError) -> Self { + ANNError::new(ANNErrorKind::IndexError, err) + } +} + +always_escalate!(InMemoryIteratorError); + +impl<'a> HasId for InMemoryIterator<'a> { + type Id = u32; +} + +impl<'a> FlatIterator for InMemoryIterator<'a> { + type ElementRef<'b> = &'b [f32]; + type Element<'b> + = &'b [f32] + where + Self: 'b; + type Error = InMemoryIteratorError; + + fn next( + &mut self, + ) -> impl SendFuture)>, Self::Error>> { + let idx = self.cursor as usize; + let result = self.vectors.get(idx).map(|v| { + self.cursor += 1; + (idx as u32, v.as_slice()) + }); + std::future::ready(Ok(result)) + } +} + +/// Squared Euclidean computer: holds a copy of the query and scores against `&[f32]`. +#[derive(Debug, Clone)] +pub struct L2QueryComputer { + query: Vec, +} + +impl<'a> PreprocessedDistanceFunction<&'a [f32], f32> for L2QueryComputer { + fn evaluate_similarity(&self, changing: &'a [f32]) -> f32 { + debug_assert_eq!(self.query.len(), changing.len()); + self.query + .iter() + .zip(changing.iter()) + .map(|(a, b)| { + let d = a - b; + d * d + }) + .sum() + } +} + +/// Strategy: produces an [`InMemoryIterator`] and an [`L2QueryComputer`]. +#[derive(Debug, Default, Clone, Copy)] +pub struct InMemoryStrategy; + +#[derive(Debug, Error)] +pub enum InMemoryStrategyError { + #[error("query length {query} does not match provider dimension {dim}")] + DimMismatch { query: usize, dim: usize }, +} + +impl From for ANNError { + #[track_caller] + fn from(err: InMemoryStrategyError) -> Self { + ANNError::new(ANNErrorKind::IndexError, err) + } +} + +impl FlatSearchStrategy for InMemoryStrategy { + type Iter<'a> = InMemoryIterator<'a>; + type QueryComputer = L2QueryComputer; + type Error = InMemoryStrategyError; + + fn create_iter<'a>( + &'a self, + provider: &'a InMemoryFlatProvider, + _context: &'a DefaultContext, + ) -> Result, Self::Error> { + Ok(InMemoryIterator { + vectors: &provider.vectors, + cursor: 0, + }) + } + + fn build_query_computer(&self, query: &[f32]) -> Result { + Ok(L2QueryComputer { + query: query.to_vec(), + }) + } +} + +/// Post-processor that copies the surviving `(id, distance)` pairs straight to output. +/// +/// Identical in behavior to [`crate::flat::CopyFlatIds`] but typed concretely against the +/// in-memory iterator, useful in tests where we want to assert against the exact output +/// shape. +#[derive(Debug, Default, Clone, Copy)] +pub struct CopyInMemoryHits; + +impl<'a> FlatPostProcess, [f32]> for CopyInMemoryHits { + type Error = crate::error::Infallible; + + fn post_process( + &self, + _iter: &mut InMemoryIterator<'a>, + _query: &[f32], + candidates: I, + output: &mut B, + ) -> impl SendFuture> + where + I: Iterator> + Send, + B: SearchOutputBuffer + Send + ?Sized, + { + let count = output.extend(candidates.map(|n| (n.id, n.distance))); + std::future::ready(Ok(count)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + flat::{CopyFlatIds, FlatIndex, validate_k}, + neighbor::Neighbor, + }; + + fn build_provider() -> InMemoryFlatProvider { + // 5 two-dimensional points; the closest to (0.0, 0.0) is index 0. + InMemoryFlatProvider::new( + 2, + vec![ + vec![0.1, 0.0], // d^2 = 0.01 + vec![1.0, 0.0], // d^2 = 1.00 + vec![0.0, 0.5], // d^2 = 0.25 + vec![5.0, 5.0], // d^2 = 50.0 + vec![-0.2, 0.1], // d^2 = 0.05 + ], + ) + } + + #[tokio::test] + async fn knn_flat_returns_top_k_in_distance_order() { + let provider = build_provider(); + let index = FlatIndex::new(provider); + let strategy = InMemoryStrategy; + let processor = CopyInMemoryHits; + let query = vec![0.0_f32, 0.0]; + + let mut output: Vec> = Vec::new(); + let stats = index + .knn_search( + validate_k(3).unwrap(), + &strategy, + &processor, + &DefaultContext, + query.as_slice(), + &mut output, + ) + .await + .expect("search succeeds"); + + assert_eq!(stats.cmps, 5); + assert_eq!(stats.result_count, 3); + + let ids: Vec = output.iter().map(|n| n.id).collect(); + assert_eq!(ids, vec![0, 4, 2]); + } + + #[tokio::test] + async fn knn_flat_with_k_larger_than_n_returns_all() { + let provider = build_provider(); + let index = FlatIndex::new(provider); + let strategy = InMemoryStrategy; + let processor = CopyFlatIds; + let query = vec![0.0_f32, 0.0]; + + let mut output: Vec> = Vec::new(); + let stats = index + .knn_search( + validate_k(100).unwrap(), + &strategy, + &processor, + &DefaultContext, + query.as_slice(), + &mut output, + ) + .await + .expect("search succeeds"); + + assert_eq!(stats.cmps, 5); + assert_eq!(stats.result_count, 5); + } + + #[test] + fn knn_flat_rejects_zero_k() { + assert!(validate_k(0).is_err()); + } +} diff --git a/diskann/src/lib.rs b/diskann/src/lib.rs index 71cb3ed41..9c1f6ac76 100644 --- a/diskann/src/lib.rs +++ b/diskann/src/lib.rs @@ -13,6 +13,7 @@ pub mod utils; pub(crate) mod internal; // Index Implementations +pub mod flat; pub mod graph; // Top level exports. diff --git a/rfcs/00000-flat-search.md b/rfcs/00000-flat-search.md new file mode 100644 index 000000000..78606ddb1 --- /dev/null +++ b/rfcs/00000-flat-search.md @@ -0,0 +1,217 @@ +# Flat Search + +| | | +|------------------|--------------------------------| +| **Authors** | Aditya Krishnan, Alex Razumov, Dongliang Wu | +| **Created** | 2026-04-24 | +| **Updated** | 2026-04-27 | + +## Motivation + +### Background + +DiskANN today exposes a single abstraction family centered on the +[`crate::provider::Accessor`] trait. Accessors are random access by design since the graph greedy search algorithm needs to decide which ids to fetch and the accessor materializes the corresponding elements (vectors, quantized vectors and neighbor lists) on demand. This is the right contract for graph search, where neighborhood expansion is inherently random-access against the [`crate::provider::DataProvider`]. + +A growing class of consumers diverge from our current pattern of use by accesssing their index **sequentially**. Some consumers build their index in an "append-only" fashion and require that they walk the index in a sequential, fixed order, relying on iteration position to enforce versioning / deduplication invariants. + +### Problem Statement + +The problem-statement here is simple: provide first-class support for sequential, one-pass scans over a data backend without +stuffing the algorithm or the backend through the `Accessor` trait surface. + +### Goals + +1. Define a streaming access primitive — `FlatIterator` — that mirrors the role + `Accessor` plays for graph search but exposes a lending-iterator interface instead of + a random-access one. +2. Provide flat-search algorithm implementations (with `knn_search` as default and filtered and diverse variants to opt-into) built on the new + primitives, so consumers can use this against their own providers / backends. +3. Expose support for features and implementations native to the repo like quantized distance computers out-of-the-box. + +## Proposal + +Let's start with the main analog to the `Accessor` trait for the `FlatIndex` - `FlatIterator`. + + +### `FlatIterator` + +```rust +pub trait FlatIterator: HasId + Send + Sync { // Has Id support + // Element yielded by iterator + type ElementRef<'a>; + + // Mostly machinery to play nice with HRTB + type Element<'a>: for<'b> Reborrow<'b, Target = Self::ElementRef<'b>> + Send + Sync + where + Self: 'a; + + type Error: StandardError; + + fn next( + &mut self, + ) -> impl SendFuture)>, Self::Error>>; + + // Default implementation for driving a closure on the items in the index. + fn on_elements_unordered( + &mut self, + mut f: F, + ) -> impl SendFuture> + where F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>), + { + async move { + while let Some((id, element)) = self.next().await? { + f(id, element.reborrow()); + } + + Ok(()) + } + } +} +``` + +The trait combines two access patterns: + +- A required lending-iterator `next()`. +- A defaulted bulk method `on_elements_unordered` that consumes the entire scan via a + callback. The default impl loops over `next`; iterators that benefit from prefetching, + SIMD batching, or amortized per-element cost could override it. + +Both methods are **async** (returning `impl SendFuture<...>`), matching +[`crate::provider::Accessor::get_element`]. Iterators backed by I/O — disk pages, +remote shards — return a real future; in-memory iterators wrap their result in +`std::future::ready`. + +The `Element` / `ElementRef` split is identical to `Accessor` and exists for the same +reason: to keep HRTB bounds on query computers from inducing `'static` requirements on +the iterator type. + + +### The glue: `FlatSearchStrategy` + +While the `FlatIterator` is the primary object that provides access to the elements in the index for the algorithm, it is scoped to each query. We intorduce a constructor - `FlatSearchStrategy` - similar to `SearchStrategy` for `Accessor` to instantiate this object. A strategy is per-call configuration: stateless, cheap to construct, scoped to one +search. It produces both a per-query iterator and a query computer. + +```rust +pub trait FlatSearchStrategy: Send + Sync +where + P: DataProvider, + T: ?Sized, +{ + /// The iterator type produced by [`Self::create_iter`]. Borrows from `self` and the + /// provider. + type Iter<'a>: FlatIterator + where + Self: 'a, + + /// The query computer produced by [`Self::build_query_computer`]. + type QueryComputer: for<'a, 'b> PreprocessedDistanceFunction< + as FlatIterator>::ElementRef<'b>, + f32, + > + Send + + Sync + + 'static; + + /// The error type for both factory methods. + type Error: StandardError; + + /// Construct a fresh iterator over `provider` for the given request `context`. + fn create_iter<'a>( + &'a self, + provider: &'a P, + context: &'a P::Context, + ) -> Result, Self::Error>; + + /// Pre-process a query into a [`Self::QueryComputer`] usable for distance computation + /// against any iterator produced by [`Self::create_iter`]. + fn build_query_computer(&self, query: &T) -> Result; +} +``` + +The `ElementRef<'b>` that the distance function `QueryComputer` acts on is tied to the (reborrowed) element yielded by the `FlatIterator::next()`. + +### `FlatIndex` + +`FlatIndex` is a thin `'static` wrapper around a `DataProvider`. The same `DataProvider` +trait used by graph search is reused here — flat and graph subsystems share a single +provider surface and the same `Context` / id-mapping / error machinery. + +```rust +pub struct FlatIndex { + provider: P, + /* private */ +} + +impl FlatIndex

{ + pub fn new(provider: P) -> Self; + pub fn provider(&self) -> &P; + + pub fn knn_search( + &self, + k: NonZeroUsize, + strategy: &S, + processor: &PP, + context: &P::Context, + query: &T, + output: &mut OB, + ) -> impl SendFuture> + where + S: FlatSearchStrategy, + T: ?Sized + Sync, + O: Send, + OB: SearchOutputBuffer + Send + ?Sized, +} +``` + +The `knn_search` method is the canonical brute-force search algorithm: + +1. Construct the iterator via `strategy.create_iter` to obtain a scoped iterator over the elements. +2. Build the query computer via `strategy.build_query_computer`. +3. Drive the scan via `iter.on_elements_unordered`, scoring each element and + inserting `Neighbor`s into a `NeighborPriorityQueue` of capacity `k`. +4. Hand the survivors (in distance order) to `processor.post_process`. +5. Return search stats. + +Other algorithms (filtered, range, diverse) can be added later as additional methods on +`FlatIndex`. + +## Trade-offs + +### Reusing `DataProvider` + +This design leans into using the `DataProvider` trait which requires implementations to implement `InternalId` and `ExternalId` conversions (via the context). Arguably, this requirement is too restrictive for some consumers of a flat-index. Reasons for sticking with `DataProvider`: + +- Every concrete provider already implements `DataProvider`, so a separate trait adds + an abstraction that existing consumers will have to implement if they want to opt-in to the flat-index path. +- Sharing `DataProvider` means the `Context`, id-mapping (`to_internal_id` / + `to_external_id`), and error machinery are identical across graph and flat search, + reducing the learning surface for new contributors. + +### Async vs sync API for `FlatIterator` + +`next()` and `on_elements_unordered` return a future, making the trait +async. This is the right default for disk-backed and network-backed iterators +where advancing the cursor involves real I/O. It also matches the `Accessor` surface, +keeping the two subsystems shaped the same way. + +The cost is paid by in-memory consumers: every call to `next()` goes through the future +machinery even when the result is immediately available via `std::future::ready`. In a +tight brute-force loop this overhead — poll scaffolding, pinning etc — could be measurable. + +We chose async because the wider audience of consumers (disk, network, mixed) benefits +more than in-memory consumers lose. + +### Expand `Element` to support batched distance computation? + +The current design yields one element per `next()` call, and the query computer scores +elements one at a time via `PreprocessedDistanceFunction::evaluate_similarity`. This could leave some optimization and performance on the table; especially with the upcoming effort around batched distance kernels. + +An alternative is to make `next()` yield a *batch* instead of a single vector representation like `Element<'_>`. Some work will need to be done to define the right interaction between the batch type, the element type in the batch, the interaction with `QueryComputer`'s types and way IDs and distances are collected in the queue. + +We opted for the scalar-per-element design for now because it is simpler to implement and +reason about. The hope is that batched distance computation can be layered on later as an opt-in sub-trait without breaking +existing iterators. + +## Future Work +- Support for other flat-search algorithms like - filtered, range and diverse flat algorithms as additional methods on `FlatIndex`. + From 9f718a13387745befc908688ded5f064066ded6f Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Mon, 27 Apr 2026 22:31:41 -0700 Subject: [PATCH 02/10] remove unnecessary parts --- diskann/src/flat/index.rs | 17 +- diskann/src/flat/iterator.rs | 37 +--- diskann/src/flat/mod.rs | 7 - diskann/src/flat/post_process.rs | 8 +- diskann/src/flat/search.rs | 35 ---- diskann/src/flat/stats.rs | 16 -- diskann/src/flat/test/mod.rs | 293 ------------------------------- 7 files changed, 16 insertions(+), 397 deletions(-) delete mode 100644 diskann/src/flat/search.rs delete mode 100644 diskann/src/flat/stats.rs delete mode 100644 diskann/src/flat/test/mod.rs diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index bebc1ea6a..04227b1da 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! [`FlatIndex`] — the index wrapper for flat search. +//! [`FlatIndex`] — the index wrapper for an on which we do flat search. use std::marker::PhantomData; use std::num::NonZeroUsize; @@ -15,9 +15,9 @@ use crate::{ ANNResult, error::IntoANNResult, flat::{ - FlatIterator, FlatPostProcess, FlatSearchStats, FlatSearchStrategy, + FlatIterator, FlatPostProcess, FlatSearchStrategy, }, - graph::SearchOutputBuffer, + graph::{SearchOutputBuffer, index::SearchStats}, neighbor::{Neighbor, NeighborPriorityQueue}, provider::DataProvider, }; @@ -56,7 +56,7 @@ impl FlatIndex

{ /// /// # Arguments /// - `k`: number of nearest neighbors to return. - /// - `strategy`: produces the per-query iterator and the query computer. + /// - `strategy`: produces the per-query iterator and the query computer. See [`FlatSearchStrategy`] /// - `processor`: post-processes the survivor candidates into the output type. /// - `context`: per-request context threaded through to the provider. /// - `query`: the query. @@ -69,7 +69,7 @@ impl FlatIndex

{ context: &P::Context, query: &T, output: &mut OB, - ) -> impl SendFuture> + ) -> impl SendFuture> where S: FlatSearchStrategy, T: ?Sized + Sync, @@ -100,7 +100,12 @@ impl FlatIndex

{ .await .into_ann_result()? as u32; - Ok(FlatSearchStats { cmps, result_count }) + Ok(SearchStats { + cmps, + hops: 0, + result_count, + range_search_second_round: false, + }) } } } diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index 02ef66f99..895c909ad 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! [`FlatIterator`] — the sequential access primitive for flat search. +//! [`FlatIterator`] — the sequential access primitive for accessing a flat index. use diskann_utils::{Reborrow, future::SendFuture}; @@ -15,34 +15,7 @@ use crate::{error::StandardError, provider::HasId}; /// accessor exposes random retrieval by id, a flat iterator exposes a *sequential* walk — /// each call to [`Self::next`] advances an internal cursor and yields the next element. /// -/// Like [`crate::provider::Accessor::get_element`], advancing the cursor is **async**: it -/// may need to await an I/O fetch (e.g., reading the next disk page, awaiting a network -/// response, etc.). Iterators backed by purely in-memory data should return a ready -/// future. -/// -/// The iterator is responsible for: -/// - Choosing the iteration order (buffer-sequential, hash-walked, partitioned, …). -/// - Skipping items that should not be visible to the algorithm (deleted, obsolete, …). -/// - Holding any borrows / locks needed to keep the underlying storage alive. -/// /// Algorithms see only `(Id, ElementRef)` pairs and treat the stream as opaque. -/// -/// # `Element` vs `ElementRef` -/// -/// Same pattern as [`crate::provider::Accessor`]: -/// -/// - `Element<'a>` is the type returned by `next`. Its lifetime is bound to the iterator -/// borrow at the call site, so only one element is live at a time. -/// - `ElementRef<'a>` is an unconstrained-lifetime reborrow used in distance-function -/// bounds. Required to keep [HRTB](https://doc.rust-lang.org/nomicon/hrtb.html) bounds -/// on query computers from forcing `Self: 'static`. -/// -/// # Hot path -/// -/// Algorithms drive the scan via [`Self::on_elements_unordered`]. The provided -/// implementation simply loops over [`Self::next`]; iterators that can amortize -/// per-element cost (prefetching the next chunk, batching distance computation, -/// performing SIMD-friendly bulk reads) should override it. pub trait FlatIterator: HasId + Send + Sync { /// A reference to a yielded element with an unconstrained lifetime, suitable for /// distance-function HRTB bounds. @@ -66,13 +39,7 @@ pub trait FlatIterator: HasId + Send + Sync { /// Drive the entire scan, invoking `f` for each yielded element. /// - /// The default implementation loops over [`Self::next`]. Implementations that benefit - /// from bulk dispatch (prefetching, batched SIMD distance computation, etc.) should - /// override this method. - /// - /// The order of invocation is unspecified and may differ between calls. The closure - /// `f` is **synchronous**; if you need to await inside the per-element handler, drive - /// the iterator manually with [`Self::next`]. + /// The default implementation loops over [`Self::next`]. fn on_elements_unordered( &mut self, mut f: F, diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs index 0cd9dcc33..34fe62ac8 100644 --- a/diskann/src/flat/mod.rs +++ b/diskann/src/flat/mod.rs @@ -35,16 +35,9 @@ pub mod index; pub mod iterator; pub mod post_process; -pub mod search; -pub mod stats; pub mod strategy; -#[cfg(any(test, feature = "testing"))] -pub mod test; - pub use index::FlatIndex; pub use iterator::FlatIterator; pub use post_process::{CopyFlatIds, FlatPostProcess}; -pub use search::{KnnFlatError, validate_k}; -pub use stats::FlatSearchStats; pub use strategy::FlatSearchStrategy; diff --git a/diskann/src/flat/post_process.rs b/diskann/src/flat/post_process.rs index 38267fef6..2f95c2932 100644 --- a/diskann/src/flat/post_process.rs +++ b/diskann/src/flat/post_process.rs @@ -11,17 +11,15 @@ use crate::{ error::StandardError, flat::FlatIterator, graph::SearchOutputBuffer, neighbor::Neighbor, provider::HasId, }; -/// Hydrate / filter / transform the survivor candidates produced by a flat search and +/// Post-process the survivor candidates produced by a flat search and /// write them into an output buffer. /// /// This is the flat counterpart to [`crate::graph::glue::SearchPostProcess`]. Processors /// receive `&mut S` so they can consult any iterator-owned lookup state (e.g., an -/// `Id -> rich-record` table built up during the scan) when assembling outputs. As with -/// the graph counterpart, [`Self::post_process`] is **async** so that processors can -/// hydrate via I/O without blocking. +/// `Id -> rich-record` table built up during the scan) when assembling outputs. /// /// The `O` type parameter lets callers pick the output element type (raw `(Id, f32)` -/// pairs, fully hydrated hits, etc.). +/// pairs, fully hydrated hits etc.). pub trait FlatPostProcess::Id> where S: FlatIterator, diff --git a/diskann/src/flat/search.rs b/diskann/src/flat/search.rs deleted file mode 100644 index b51b561f1..000000000 --- a/diskann/src/flat/search.rs +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Error types for flat search parameter validation. - -use std::num::NonZeroUsize; - -use thiserror::Error; - -use crate::{ANNError, ANNErrorKind}; - -/// Errors raised when validating flat search parameters. -#[derive(Debug, Error)] -pub enum KnnFlatError { - /// `k` was zero. - #[error("k cannot be zero")] - KZero, -} - -impl From for ANNError { - #[track_caller] - fn from(err: KnnFlatError) -> Self { - Self::new(ANNErrorKind::IndexError, err) - } -} - -/// Validate and wrap a `k` value as [`NonZeroUsize`]. -/// -/// This is a convenience for callers that want to validate `k` before passing it to -/// [`FlatIndex::knn_search`](crate::flat::FlatIndex::knn_search). -pub fn validate_k(k: usize) -> Result { - NonZeroUsize::new(k).ok_or(KnnFlatError::KZero) -} diff --git a/diskann/src/flat/stats.rs b/diskann/src/flat/stats.rs deleted file mode 100644 index faf14f888..000000000 --- a/diskann/src/flat/stats.rs +++ /dev/null @@ -1,16 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Statistics returned by a flat search. - -/// Statistics collected during a single flat search invocation. -#[derive(Debug, Clone, Copy, Default)] -pub struct FlatSearchStats { - /// Number of distance computations performed (i.e., elements visited by the scanner). - pub cmps: u32, - - /// Number of results written into the caller-provided output buffer. - pub result_count: u32, -} diff --git a/diskann/src/flat/test/mod.rs b/diskann/src/flat/test/mod.rs deleted file mode 100644 index 61307ab61..000000000 --- a/diskann/src/flat/test/mod.rs +++ /dev/null @@ -1,293 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Trivial in-memory provider, iterator, and strategy used for unit-testing the flat -//! search infrastructure. -//! -//! This is intentionally simple: vectors live in a `Vec>`, ids are `u32`, and -//! distance is squared Euclidean. It exists so the trait shapes in [`crate::flat`] can be -//! exercised end-to-end without dragging in any provider-side machinery. - -use diskann_utils::future::SendFuture; -use diskann_vector::PreprocessedDistanceFunction; -use thiserror::Error; - -use crate::{ - always_escalate, - ANNError, ANNErrorKind, - flat::{FlatIterator, FlatPostProcess, FlatSearchStrategy}, - graph::SearchOutputBuffer, - neighbor::Neighbor, - provider::{DataProvider, DefaultContext, HasId, NoopGuard}, -}; - -/// Trivial flat provider holding a list of fixed-dimension `f32` vectors. -#[derive(Debug)] -pub struct InMemoryFlatProvider { - pub dim: usize, - pub vectors: Vec>, -} - -impl InMemoryFlatProvider { - pub fn new(dim: usize, vectors: Vec>) -> Self { - Self { dim, vectors } - } -} - -#[derive(Debug, Error)] -#[error("invalid vector id {0}")] -pub struct InMemoryProviderError(u32); - -impl From for ANNError { - #[track_caller] - fn from(err: InMemoryProviderError) -> Self { - ANNError::new(ANNErrorKind::IndexError, err) - } -} - -always_escalate!(InMemoryProviderError); - -impl DataProvider for InMemoryFlatProvider { - type Context = DefaultContext; - type InternalId = u32; - type ExternalId = u32; - type Error = InMemoryProviderError; - type Guard = NoopGuard; - - fn to_internal_id( - &self, - _context: &Self::Context, - gid: &u32, - ) -> Result { - if (*gid as usize) < self.vectors.len() { - Ok(*gid) - } else { - Err(InMemoryProviderError(*gid)) - } - } - - fn to_external_id( - &self, - _context: &Self::Context, - id: u32, - ) -> Result { - if (id as usize) < self.vectors.len() { - Ok(id) - } else { - Err(InMemoryProviderError(id)) - } - } -} - -/// Sequential iterator over [`InMemoryFlatProvider`]. -pub struct InMemoryIterator<'a> { - vectors: &'a [Vec], - cursor: u32, -} - -#[derive(Debug, Error)] -#[error("in-memory iterator does not error")] -pub struct InMemoryIteratorError; - -impl From for ANNError { - #[track_caller] - fn from(err: InMemoryIteratorError) -> Self { - ANNError::new(ANNErrorKind::IndexError, err) - } -} - -always_escalate!(InMemoryIteratorError); - -impl<'a> HasId for InMemoryIterator<'a> { - type Id = u32; -} - -impl<'a> FlatIterator for InMemoryIterator<'a> { - type ElementRef<'b> = &'b [f32]; - type Element<'b> - = &'b [f32] - where - Self: 'b; - type Error = InMemoryIteratorError; - - fn next( - &mut self, - ) -> impl SendFuture)>, Self::Error>> { - let idx = self.cursor as usize; - let result = self.vectors.get(idx).map(|v| { - self.cursor += 1; - (idx as u32, v.as_slice()) - }); - std::future::ready(Ok(result)) - } -} - -/// Squared Euclidean computer: holds a copy of the query and scores against `&[f32]`. -#[derive(Debug, Clone)] -pub struct L2QueryComputer { - query: Vec, -} - -impl<'a> PreprocessedDistanceFunction<&'a [f32], f32> for L2QueryComputer { - fn evaluate_similarity(&self, changing: &'a [f32]) -> f32 { - debug_assert_eq!(self.query.len(), changing.len()); - self.query - .iter() - .zip(changing.iter()) - .map(|(a, b)| { - let d = a - b; - d * d - }) - .sum() - } -} - -/// Strategy: produces an [`InMemoryIterator`] and an [`L2QueryComputer`]. -#[derive(Debug, Default, Clone, Copy)] -pub struct InMemoryStrategy; - -#[derive(Debug, Error)] -pub enum InMemoryStrategyError { - #[error("query length {query} does not match provider dimension {dim}")] - DimMismatch { query: usize, dim: usize }, -} - -impl From for ANNError { - #[track_caller] - fn from(err: InMemoryStrategyError) -> Self { - ANNError::new(ANNErrorKind::IndexError, err) - } -} - -impl FlatSearchStrategy for InMemoryStrategy { - type Iter<'a> = InMemoryIterator<'a>; - type QueryComputer = L2QueryComputer; - type Error = InMemoryStrategyError; - - fn create_iter<'a>( - &'a self, - provider: &'a InMemoryFlatProvider, - _context: &'a DefaultContext, - ) -> Result, Self::Error> { - Ok(InMemoryIterator { - vectors: &provider.vectors, - cursor: 0, - }) - } - - fn build_query_computer(&self, query: &[f32]) -> Result { - Ok(L2QueryComputer { - query: query.to_vec(), - }) - } -} - -/// Post-processor that copies the surviving `(id, distance)` pairs straight to output. -/// -/// Identical in behavior to [`crate::flat::CopyFlatIds`] but typed concretely against the -/// in-memory iterator, useful in tests where we want to assert against the exact output -/// shape. -#[derive(Debug, Default, Clone, Copy)] -pub struct CopyInMemoryHits; - -impl<'a> FlatPostProcess, [f32]> for CopyInMemoryHits { - type Error = crate::error::Infallible; - - fn post_process( - &self, - _iter: &mut InMemoryIterator<'a>, - _query: &[f32], - candidates: I, - output: &mut B, - ) -> impl SendFuture> - where - I: Iterator> + Send, - B: SearchOutputBuffer + Send + ?Sized, - { - let count = output.extend(candidates.map(|n| (n.id, n.distance))); - std::future::ready(Ok(count)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{ - flat::{CopyFlatIds, FlatIndex, validate_k}, - neighbor::Neighbor, - }; - - fn build_provider() -> InMemoryFlatProvider { - // 5 two-dimensional points; the closest to (0.0, 0.0) is index 0. - InMemoryFlatProvider::new( - 2, - vec![ - vec![0.1, 0.0], // d^2 = 0.01 - vec![1.0, 0.0], // d^2 = 1.00 - vec![0.0, 0.5], // d^2 = 0.25 - vec![5.0, 5.0], // d^2 = 50.0 - vec![-0.2, 0.1], // d^2 = 0.05 - ], - ) - } - - #[tokio::test] - async fn knn_flat_returns_top_k_in_distance_order() { - let provider = build_provider(); - let index = FlatIndex::new(provider); - let strategy = InMemoryStrategy; - let processor = CopyInMemoryHits; - let query = vec![0.0_f32, 0.0]; - - let mut output: Vec> = Vec::new(); - let stats = index - .knn_search( - validate_k(3).unwrap(), - &strategy, - &processor, - &DefaultContext, - query.as_slice(), - &mut output, - ) - .await - .expect("search succeeds"); - - assert_eq!(stats.cmps, 5); - assert_eq!(stats.result_count, 3); - - let ids: Vec = output.iter().map(|n| n.id).collect(); - assert_eq!(ids, vec![0, 4, 2]); - } - - #[tokio::test] - async fn knn_flat_with_k_larger_than_n_returns_all() { - let provider = build_provider(); - let index = FlatIndex::new(provider); - let strategy = InMemoryStrategy; - let processor = CopyFlatIds; - let query = vec![0.0_f32, 0.0]; - - let mut output: Vec> = Vec::new(); - let stats = index - .knn_search( - validate_k(100).unwrap(), - &strategy, - &processor, - &DefaultContext, - query.as_slice(), - &mut output, - ) - .await - .expect("search succeeds"); - - assert_eq!(stats.cmps, 5); - assert_eq!(stats.result_count, 5); - } - - #[test] - fn knn_flat_rejects_zero_k() { - assert!(validate_k(0).is_err()); - } -} From f0a9dbd2ebb68ef89060ece0b0317edf1b3a247e Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Mon, 27 Apr 2026 22:37:48 -0700 Subject: [PATCH 03/10] rename file --- rfcs/{00000-flat-search.md => 00983-flat-search.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename rfcs/{00000-flat-search.md => 00983-flat-search.md} (100%) diff --git a/rfcs/00000-flat-search.md b/rfcs/00983-flat-search.md similarity index 100% rename from rfcs/00000-flat-search.md rename to rfcs/00983-flat-search.md From 0672f3d5ab35cc82a42309e48e1c66dbcddb5d9b Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 28 Apr 2026 08:19:33 -0700 Subject: [PATCH 04/10] fmt --- diskann/src/flat/index.rs | 4 +--- diskann/src/flat/iterator.rs | 9 +++------ diskann/src/flat/post_process.rs | 5 +++-- diskann/src/flat/strategy.rs | 6 +----- 4 files changed, 8 insertions(+), 16 deletions(-) diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 04227b1da..6cf4c87c2 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -14,9 +14,7 @@ use diskann_vector::PreprocessedDistanceFunction; use crate::{ ANNResult, error::IntoANNResult, - flat::{ - FlatIterator, FlatPostProcess, FlatSearchStrategy, - }, + flat::{FlatIterator, FlatPostProcess, FlatSearchStrategy}, graph::{SearchOutputBuffer, index::SearchStats}, neighbor::{Neighbor, NeighborPriorityQueue}, provider::DataProvider, diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index 895c909ad..c822e61d8 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -33,17 +33,15 @@ pub trait FlatIterator: HasId + Send + Sync { /// /// Returns `Ok(None)` when the scan is exhausted. The yielded element borrows from /// the iterator and is invalidated by the next call to `next`. + #[allow(clippy::type_complexity)] fn next( &mut self, ) -> impl SendFuture)>, Self::Error>>; /// Drive the entire scan, invoking `f` for each yielded element. /// - /// The default implementation loops over [`Self::next`]. - fn on_elements_unordered( - &mut self, - mut f: F, - ) -> impl SendFuture> + /// The default implementation loops over [`Self::next`]. + fn on_elements_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>), { @@ -55,4 +53,3 @@ pub trait FlatIterator: HasId + Send + Sync { } } } - diff --git a/diskann/src/flat/post_process.rs b/diskann/src/flat/post_process.rs index 2f95c2932..3e688e5bd 100644 --- a/diskann/src/flat/post_process.rs +++ b/diskann/src/flat/post_process.rs @@ -8,7 +8,8 @@ use diskann_utils::future::SendFuture; use crate::{ - error::StandardError, flat::FlatIterator, graph::SearchOutputBuffer, neighbor::Neighbor, provider::HasId, + error::StandardError, flat::FlatIterator, graph::SearchOutputBuffer, neighbor::Neighbor, + provider::HasId, }; /// Post-process the survivor candidates produced by a flat search and @@ -16,7 +17,7 @@ use crate::{ /// /// This is the flat counterpart to [`crate::graph::glue::SearchPostProcess`]. Processors /// receive `&mut S` so they can consult any iterator-owned lookup state (e.g., an -/// `Id -> rich-record` table built up during the scan) when assembling outputs. +/// `Id -> rich-record` table built up during the scan) when assembling outputs. /// /// The `O` type parameter lets callers pick the output element type (raw `(Id, f32)` /// pairs, fully hydrated hits etc.). diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 8b0f5e445..423b2817c 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -7,11 +7,7 @@ use diskann_vector::PreprocessedDistanceFunction; -use crate::{ - error::StandardError, - flat::FlatIterator, - provider::DataProvider, -}; +use crate::{error::StandardError, flat::FlatIterator, provider::DataProvider}; /// Per-call configuration that knows how to construct a [`FlatIterator`] for a provider /// and how to pre-process queries of type `T` into a distance computer. From 3ac0e1b2f4eac2fa11545f04a5908d016be400bb Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 28 Apr 2026 17:27:11 -0700 Subject: [PATCH 05/10] split iterator to callback --- diskann/src/flat/index.rs | 15 +++--- diskann/src/flat/iterator.rs | 81 +++++++++++++++++++++++++++----- diskann/src/flat/mod.rs | 2 +- diskann/src/flat/post_process.rs | 6 +-- diskann/src/flat/strategy.rs | 18 +++---- 5 files changed, 91 insertions(+), 31 deletions(-) diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 6cf4c87c2..9087e4b4b 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -14,7 +14,7 @@ use diskann_vector::PreprocessedDistanceFunction; use crate::{ ANNResult, error::IntoANNResult, - flat::{FlatIterator, FlatPostProcess, FlatSearchStrategy}, + flat::{OnElementsUnordered, FlatPostProcess, FlatSearchStrategy}, graph::{SearchOutputBuffer, index::SearchStats}, neighbor::{Neighbor, NeighborPriorityQueue}, provider::DataProvider, @@ -28,7 +28,7 @@ use crate::{ #[derive(Debug)] pub struct FlatIndex { /// The backing provider. - pub provider: P, + provider: P, _marker: PhantomData P>, } @@ -73,19 +73,20 @@ impl FlatIndex

{ T: ?Sized + Sync, O: Send, OB: SearchOutputBuffer + Send + ?Sized, - PP: for<'a> FlatPostProcess, T, O> + Send + Sync, + PP: for<'a> FlatPostProcess, T, O> + Send + Sync, { async move { - let mut iter = strategy - .create_iter(&self.provider, context) + let mut callback = strategy + .create_callback(&self.provider, context) .into_ann_result()?; + let computer = strategy.build_query_computer(query).into_ann_result()?; let k = k.get(); let mut queue = NeighborPriorityQueue::new(k); let mut cmps: u32 = 0; - iter.on_elements_unordered(|id, element| { + callback.on_elements_unordered(|id, element| { let dist = computer.evaluate_similarity(element); cmps += 1; queue.insert(Neighbor::new(id, dist)); @@ -94,7 +95,7 @@ impl FlatIndex

{ .into_ann_result()?; let result_count = processor - .post_process(&mut iter, query, queue.iter().take(k), output) + .post_process(&mut callback, query, queue.iter().take(k), output) .await .into_ann_result()? as u32; diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index c822e61d8..01a5d7cea 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -3,19 +3,41 @@ * Licensed under the MIT license. */ -//! [`FlatIterator`] — the sequential access primitive for accessing a flat index. +//! [`OnElementsUnordered`] — the sequential access primitive for accessing a flat index. +//! +//! [`FlatIterator`] — a lending async iterator that can be bridged into +//! [`OnElementsUnordered`] via [`DefaultIteratedOperator`]. use diskann_utils::{Reborrow, future::SendFuture}; use crate::{error::StandardError, provider::HasId}; -/// A lending, asynchronous iterator over the elements of a flat index. +/// Callback-driven sequential scan over the elements of a flat index. /// -/// `FlatIterator` is the streaming counterpart to [`crate::provider::Accessor`]. Where an -/// accessor exposes random retrieval by id, a flat iterator exposes a *sequential* walk — -/// each call to [`Self::next`] advances an internal cursor and yields the next element. +/// `OnElementsUnordered` is the streaming counterpart to [`crate::provider::Accessor`]. +/// Where an accessor exposes random retrieval by id, this trait exposes a *sequential* +/// walk that invokes a caller-supplied closure for every element. /// /// Algorithms see only `(Id, ElementRef)` pairs and treat the stream as opaque. +pub trait OnElementsUnordered: HasId + Send + Sync { + /// A reference to a yielded element with an unconstrained lifetime, suitable for + /// distance-function HRTB bounds. + type ElementRef<'a>; + + /// The error type yielded by [`Self::on_elements_unordered`]. + type Error: StandardError; + + /// Drive the entire scan, invoking `f` for each yielded element. + fn on_elements_unordered(&mut self, f: F) -> impl SendFuture> + where + F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>); +} + +/// A lending, asynchronous iterator over the elements of a flat index. +/// +/// Implementations provide element-at-a-time access via [`Self::next`]. Providers that +/// only implement `FlatIterator` can be wrapped in [`DefaultIteratedOperator`] to obtain +/// an [`OnElementsUnordered`] implementation automatically. pub trait FlatIterator: HasId + Send + Sync { /// A reference to a yielded element with an unconstrained lifetime, suitable for /// distance-function HRTB bounds. @@ -26,7 +48,7 @@ pub trait FlatIterator: HasId + Send + Sync { where Self: 'a; - /// The error type yielded by [`Self::next`] and [`Self::on_elements_unordered`]. + /// The error type yielded by [`Self::next`]. type Error: StandardError; /// Advance the iterator and asynchronously yield the next `(id, element)` pair. @@ -37,19 +59,56 @@ pub trait FlatIterator: HasId + Send + Sync { fn next( &mut self, ) -> impl SendFuture)>, Self::Error>>; +} + + +/////////////// +/// Default /// +/////////////// + + +/// Bridges a [`FlatIterator`] into an [`OnElementsUnordered`] by looping over +/// [`FlatIterator::next`] and reborrowing each element into the closure. +/// +/// This is the default adapter for providers that implement element-at-a-time iteration. +/// Providers that can do better (prefetching, SIMD batching, bulk I/O) should implement +/// [`OnElementsUnordered`] directly. +pub struct DefaultIteratedOperator { + inner: I, +} + +impl DefaultIteratedOperator { + /// Wrap an iterator to produce an [`OnElementsUnordered`] implementation. + pub fn new(inner: I) -> Self { + Self { inner } + } + + /// Unwrap, returning the inner iterator. + pub fn into_inner(self) -> I { + self.inner + } +} + +impl HasId for DefaultIteratedOperator { + type Id = I::Id; +} + +impl OnElementsUnordered for DefaultIteratedOperator +where + I: FlatIterator + HasId + Send + Sync, +{ + type ElementRef<'a> = I::ElementRef<'a>; + type Error = I::Error; - /// Drive the entire scan, invoking `f` for each yielded element. - /// - /// The default implementation loops over [`Self::next`]. fn on_elements_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>), { async move { - while let Some((id, element)) = self.next().await? { + while let Some((id, element)) = self.inner.next().await? { f(id, element.reborrow()); } Ok(()) } } -} +} \ No newline at end of file diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs index 34fe62ac8..8754ae9d4 100644 --- a/diskann/src/flat/mod.rs +++ b/diskann/src/flat/mod.rs @@ -38,6 +38,6 @@ pub mod post_process; pub mod strategy; pub use index::FlatIndex; -pub use iterator::FlatIterator; +pub use iterator::{DefaultIteratedOperator, FlatIterator, OnElementsUnordered}; pub use post_process::{CopyFlatIds, FlatPostProcess}; pub use strategy::FlatSearchStrategy; diff --git a/diskann/src/flat/post_process.rs b/diskann/src/flat/post_process.rs index 3e688e5bd..71ffabf3c 100644 --- a/diskann/src/flat/post_process.rs +++ b/diskann/src/flat/post_process.rs @@ -8,7 +8,7 @@ use diskann_utils::future::SendFuture; use crate::{ - error::StandardError, flat::FlatIterator, graph::SearchOutputBuffer, neighbor::Neighbor, + error::StandardError, flat::OnElementsUnordered, graph::SearchOutputBuffer, neighbor::Neighbor, provider::HasId, }; @@ -23,7 +23,7 @@ use crate::{ /// pairs, fully hydrated hits etc.). pub trait FlatPostProcess::Id> where - S: FlatIterator, + S: OnElementsUnordered, T: ?Sized, { /// Errors yielded by [`Self::post_process`]. @@ -50,7 +50,7 @@ pub struct CopyFlatIds; impl FlatPostProcess for CopyFlatIds where - S: FlatIterator, + S: OnElementsUnordered, T: ?Sized, { type Error = crate::error::Infallible; diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 423b2817c..0e77df36b 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -7,7 +7,7 @@ use diskann_vector::PreprocessedDistanceFunction; -use crate::{error::StandardError, flat::FlatIterator, provider::DataProvider}; +use crate::{error::StandardError, flat::OnElementsUnordered, provider::DataProvider}; /// Per-call configuration that knows how to construct a [`FlatIterator`] for a provider /// and how to pre-process queries of type `T` into a distance computer. @@ -18,7 +18,7 @@ use crate::{error::StandardError, flat::FlatIterator, provider::DataProvider}; /// /// # Why two methods? /// -/// - [`Self::create_iter`] is query-independent and may be called multiple times per +/// - [`Self::create_callback`] is query-independent and may be called multiple times per /// request (e.g., once per parallel query in a batched search). /// - [`Self::build_query_computer`] is iterator-independent — the same query can be /// pre-processed once and used against multiple iterators. @@ -34,9 +34,9 @@ where P: DataProvider, T: ?Sized, { - /// The iterator type produced by [`Self::create_iter`]. Borrows from `self` and the + /// The iterator type produced by [`Self::create_callback`]. Borrows from `self` and the /// provider. - type Iter<'a>: FlatIterator + type Callback<'a>: OnElementsUnordered where Self: 'a, P: 'a; @@ -47,7 +47,7 @@ where /// by every lifetime of `Iter`. Two lifetimes are needed: `'a` for the iterator /// instance and `'b` for the reborrowed element. type QueryComputer: for<'a, 'b> PreprocessedDistanceFunction< - as FlatIterator>::ElementRef<'b>, + as OnElementsUnordered>::ElementRef<'b>, f32, > + Send + Sync @@ -59,15 +59,15 @@ where /// Construct a fresh iterator over `provider` for the given request `context`. /// /// This is where lock acquisition, snapshot pinning, and any other per-query setup - /// should happen. The returned iterator owns whatever borrows / guards it needs to + /// should happen. The returned callback object owns whatever borrows / guards it needs to /// remain valid until it is dropped. - fn create_iter<'a>( + fn create_callback<'a>( &'a self, provider: &'a P, context: &'a P::Context, - ) -> Result, Self::Error>; + ) -> Result, Self::Error>; /// Pre-process a query into a [`Self::QueryComputer`] usable for distance computation - /// against any iterator produced by [`Self::create_iter`]. + /// against any iterator produced by [`Self::create_callback`]. fn build_query_computer(&self, query: &T) -> Result; } From f887f2fad9112303bf58f128ce43765d47ddea45 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 28 Apr 2026 18:42:16 -0700 Subject: [PATCH 06/10] use distance unordered callback --- diskann/src/flat/index.rs | 10 ++++------ diskann/src/flat/iterator.rs | 33 +++++++++++++++++++++++++++++++- diskann/src/flat/mod.rs | 4 ++-- diskann/src/flat/post_process.rs | 4 ++-- diskann/src/flat/strategy.rs | 4 ++-- 5 files changed, 42 insertions(+), 13 deletions(-) diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 9087e4b4b..825567f8a 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -9,12 +9,11 @@ use std::marker::PhantomData; use std::num::NonZeroUsize; use diskann_utils::future::SendFuture; -use diskann_vector::PreprocessedDistanceFunction; use crate::{ ANNResult, error::IntoANNResult, - flat::{OnElementsUnordered, FlatPostProcess, FlatSearchStrategy}, + flat::{DistancesUnordered, FlatPostProcess, FlatSearchStrategy}, graph::{SearchOutputBuffer, index::SearchStats}, neighbor::{Neighbor, NeighborPriorityQueue}, provider::DataProvider, @@ -23,8 +22,8 @@ use crate::{ /// A `'static` thin wrapper around a [`DataProvider`] used for flat search. /// /// The provider is owned by the index. The index is constructed once at process startup and -/// shared across requests; per-query state lives in the [`crate::flat::FlatIterator`] that -/// the [`crate::flat::FlatSearchStrategy`] produces. +/// shared across requests; per-query state lives in the [`crate::flat::OnElementsUnordered`] +/// implementation that the [`crate::flat::FlatSearchStrategy`] produces. #[derive(Debug)] pub struct FlatIndex { /// The backing provider. @@ -86,8 +85,7 @@ impl FlatIndex

{ let mut queue = NeighborPriorityQueue::new(k); let mut cmps: u32 = 0; - callback.on_elements_unordered(|id, element| { - let dist = computer.evaluate_similarity(element); + callback.distances_unordered(&computer, |id, dist| { cmps += 1; queue.insert(Neighbor::new(id, dist)); }) diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index 01a5d7cea..fd69b312a 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -9,6 +9,7 @@ //! [`OnElementsUnordered`] via [`DefaultIteratedOperator`]. use diskann_utils::{Reborrow, future::SendFuture}; +use diskann_vector::PreprocessedDistanceFunction; use crate::{error::StandardError, provider::HasId}; @@ -33,6 +34,34 @@ pub trait OnElementsUnordered: HasId + Send + Sync { F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>); } + +/// Extension of [`OnElementsUnordered`] that drives the scan with a pre-built query +/// computer, invoking a callback with `(id, distance)` pairs instead of raw elements. +/// +/// The concrete computer is insantiated and supplied externally +/// by the [`FlatSearchStrategy`](crate::flat::FlatSearchStrategy). +/// +/// The default implementation delegates to [`OnElementsUnordered::on_elements_unordered`], +/// calling `computer.evaluate_similarity` on each element. +pub trait DistancesUnordered: OnElementsUnordered { + /// Drive the entire scan, scoring each element with `computer` and invoking `f` with + /// the resulting `(id, distance)` pair. + fn distances_unordered( + &mut self, + computer: &C, + mut f: F, + ) -> impl SendFuture> + where + C: for<'a> PreprocessedDistanceFunction, f32> + Send + Sync, + F: Send + FnMut(Self::Id, f32), + { + self.on_elements_unordered(move |id, element| { + let dist = computer.evaluate_similarity(element); + f(id, dist); + }) + } +} + /// A lending, asynchronous iterator over the elements of a flat index. /// /// Implementations provide element-at-a-time access via [`Self::next`]. Providers that @@ -111,4 +140,6 @@ where Ok(()) } } -} \ No newline at end of file +} + +impl DistancesUnordered for DefaultIteratedOperator where I: FlatIterator + HasId + Send + Sync {} \ No newline at end of file diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs index 8754ae9d4..2516509da 100644 --- a/diskann/src/flat/mod.rs +++ b/diskann/src/flat/mod.rs @@ -38,6 +38,6 @@ pub mod post_process; pub mod strategy; pub use index::FlatIndex; -pub use iterator::{DefaultIteratedOperator, FlatIterator, OnElementsUnordered}; -pub use post_process::{CopyFlatIds, FlatPostProcess}; +pub use iterator::{DefaultIteratedOperator, DistancesUnordered, FlatIterator, OnElementsUnordered}; +pub use post_process::{CopyIds, FlatPostProcess}; pub use strategy::FlatSearchStrategy; diff --git a/diskann/src/flat/post_process.rs b/diskann/src/flat/post_process.rs index 71ffabf3c..2cab763dd 100644 --- a/diskann/src/flat/post_process.rs +++ b/diskann/src/flat/post_process.rs @@ -46,9 +46,9 @@ where /// A trivial [`FlatPostProcess`] that copies each `(Id, distance)` pair straight into the /// output buffer. #[derive(Debug, Default, Clone, Copy)] -pub struct CopyFlatIds; +pub struct CopyIds; -impl FlatPostProcess for CopyFlatIds +impl FlatPostProcess for CopyIds where S: OnElementsUnordered, T: ?Sized, diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 0e77df36b..9fa2c6e00 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -7,7 +7,7 @@ use diskann_vector::PreprocessedDistanceFunction; -use crate::{error::StandardError, flat::OnElementsUnordered, provider::DataProvider}; +use crate::{error::StandardError, flat::{DistancesUnordered, OnElementsUnordered}, provider::DataProvider}; /// Per-call configuration that knows how to construct a [`FlatIterator`] for a provider /// and how to pre-process queries of type `T` into a distance computer. @@ -36,7 +36,7 @@ where { /// The iterator type produced by [`Self::create_callback`]. Borrows from `self` and the /// provider. - type Callback<'a>: OnElementsUnordered + type Callback<'a>: DistancesUnordered where Self: 'a, P: 'a; From dc2281cfb26607eab6a47f6f97eab446c6badd38 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 28 Apr 2026 21:26:24 -0700 Subject: [PATCH 07/10] rfc update --- rfcs/00983-flat-search.md | 155 ++++++++++++++++++++------------------ 1 file changed, 83 insertions(+), 72 deletions(-) diff --git a/rfcs/00983-flat-search.md b/rfcs/00983-flat-search.md index 78606ddb1..dfef8f516 100644 --- a/rfcs/00983-flat-search.md +++ b/rfcs/00983-flat-search.md @@ -22,75 +22,88 @@ stuffing the algorithm or the backend through the `Accessor` trait surface. ### Goals -1. Define a streaming access primitive — `FlatIterator` — that mirrors the role - `Accessor` plays for graph search but exposes a lending-iterator interface instead of - a random-access one. +1. Define a streaming access primitive — `OnElementsUnordered` — that mirrors the role + `Accessor` plays for graph search but exposes a callback-driven scan instead of + random access. 2. Provide flat-search algorithm implementations (with `knn_search` as default and filtered and diverse variants to opt-into) built on the new primitives, so consumers can use this against their own providers / backends. 3. Expose support for features and implementations native to the repo like quantized distance computers out-of-the-box. ## Proposal -Let's start with the main analog to the `Accessor` trait for the `FlatIndex` - `FlatIterator`. +The flat-search infrastructure is built on a small sequence of traits. The only required traits for the algorithm is `OnElementsUnordered` and its subtrait `DistancesUnordered`. A strategy - `FlatSearchStrategy` - instantiates these implementations for specific providers. An opt-in iterator trait `FlatIterator` and default implementations of the core traits - `DefaultIteratedOperator` - exist for convenience for backends that naturally expose element-at-a-time iteration. - -### `FlatIterator` +### `OnElementsUnordered` — the core scan ```rust -pub trait FlatIterator: HasId + Send + Sync { // Has Id support - // Element yielded by iterator - type ElementRef<'a>; +pub trait OnElementsUnordered: HasId + Send + Sync { + type ElementRef<'a>; + type Error: StandardError; - // Mostly machinery to play nice with HRTB - type Element<'a>: for<'b> Reborrow<'b, Target = Self::ElementRef<'b>> + Send + Sync - where - Self: 'a; + fn on_elements_unordered(&mut self, f: F) -> impl SendFuture> + where + F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>); +} +``` - type Error: StandardError; +A single required method: drive the entire scan via a callback. Async to match +[`crate::provider::Accessor`]. Implementations choose iteration order, prefetching, and +any SIMD-friendly bulk reads if they want; algorithms see only `(Id, ElementRef)` pairs. - fn next( - &mut self, - ) -> impl SendFuture)>, Self::Error>>; +### `DistancesUnordered` — the distance subtrait - // Default implementation for driving a closure on the items in the index. - fn on_elements_unordered( - &mut self, - mut f: F, +```rust +pub trait DistancesUnordered: OnElementsUnordered { + fn distances_unordered( + &mut self, computer: &C, mut f: F, ) -> impl SendFuture> - where F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>), + where + C: for<'a> PreprocessedDistanceFunction, f32> + Send + Sync, + F: Send + FnMut(Self::Id, f32), { - async move { - while let Some((id, element)) = self.next().await? { - f(id, element.reborrow()); - } - - Ok(()) - } + // default delegates to on_elements_unordered + evaluate_similarity } } ``` -The trait combines two access patterns: +A subtrait that fuses scanning with scoring. The default implementation loops +`on_elements_unordered` and calls `computer.evaluate_similarity` on each element. -- A required lending-iterator `next()`. -- A defaulted bulk method `on_elements_unordered` that consumes the entire scan via a - callback. The default impl loops over `next`; iterators that benefit from prefetching, - SIMD batching, or amortized per-element cost could override it. +The query computer is a generic parameter rather than an associated type, so the same +callback type can be driven by different computers. The `FlatSearchStrategy` is the +source of truth for which computer is used in any given search. -Both methods are **async** (returning `impl SendFuture<...>`), matching -[`crate::provider::Accessor::get_element`]. Iterators backed by I/O — disk pages, -remote shards — return a real future; in-memory iterators wrap their result in -`std::future::ready`. +### `FlatIterator` and `DefaultIteratedOperator` — convenience for element-at-a-time backends -The `Element` / `ElementRef` split is identical to `Accessor` and exists for the same -reason: to keep HRTB bounds on query computers from inducing `'static` requirements on -the iterator type. +For backends that naturally expose element-at-a-time iteration, `FlatIterator` is a +lending async iterator: + +```rust +pub trait FlatIterator: HasId + Send + Sync { + type ElementRef<'a>; + // lifetime gymnastics to make lifetime of `Element<'_>` to play nice with HRTB + type Element<'a>: for<'b> Reborrow<'b, Target = Self::ElementRef<'b>> + Send + Sync + where Self: 'a; + type Error: StandardError; + + fn next( + &mut self, + ) -> impl SendFuture)>, Self::Error>>; +} +``` + +`DefaultIteratedOperator` wraps any `FlatIterator` and implements `OnElementsUnordered` +(and `DistancesUnordered` by inheritance) by looping over `next()` and reborrowing each +element. ### The glue: `FlatSearchStrategy` -While the `FlatIterator` is the primary object that provides access to the elements in the index for the algorithm, it is scoped to each query. We intorduce a constructor - `FlatSearchStrategy` - similar to `SearchStrategy` for `Accessor` to instantiate this object. A strategy is per-call configuration: stateless, cheap to construct, scoped to one -search. It produces both a per-query iterator and a query computer. +While `OnElementsUnordered` is the primary handle the algorithm uses to walk the index, +it is scoped to each query. We introduce a constructor — `FlatSearchStrategy` — similar +to `SearchStrategy` for `Accessor`, to instantiate the per-query callback object. +A strategy is per-call configuration that is stateless, cheap to construct and scoped to one +search. It produces both a per-query callback and a query computer. ```rust pub trait FlatSearchStrategy: Send + Sync @@ -98,55 +111,55 @@ where P: DataProvider, T: ?Sized, { - /// The iterator type produced by [`Self::create_iter`]. Borrows from `self` and the - /// provider. - type Iter<'a>: FlatIterator + /// The per-query callback type produced by [`Self::create_callback`]. Borrows from + /// `self` and the provider. + type Callback<'a>: DistancesUnordered where Self: 'a, /// The query computer produced by [`Self::build_query_computer`]. type QueryComputer: for<'a, 'b> PreprocessedDistanceFunction< - as FlatIterator>::ElementRef<'b>, + as OnElementsUnordered>::ElementRef<'b>, f32, > + Send + Sync + 'static; - /// The error type for both factory methods. + /// The error type type Error: StandardError; - /// Construct a fresh iterator over `provider` for the given request `context`. - fn create_iter<'a>( + /// Construct a fresh callback over `provider` for the given request `context`. + fn create_callback<'a>( &'a self, provider: &'a P, context: &'a P::Context, - ) -> Result, Self::Error>; + ) -> Result, Self::Error>; /// Pre-process a query into a [`Self::QueryComputer`] usable for distance computation - /// against any iterator produced by [`Self::create_iter`]. + /// against any callback produced by [`Self::create_callback`]. fn build_query_computer(&self, query: &T) -> Result; } ``` -The `ElementRef<'b>` that the distance function `QueryComputer` acts on is tied to the (reborrowed) element yielded by the `FlatIterator::next()`. +The `ElementRef<'b>` that the `QueryComputer` acts on is tied to the +`OnElementsUnordered::ElementRef` of the callback produced by `create_callback`. ### `FlatIndex` `FlatIndex` is a thin `'static` wrapper around a `DataProvider`. The same `DataProvider` -trait used by graph search is reused here — flat and graph subsystems share a single +trait used by graph search is reused here - flat and graph subsystems share a single provider surface and the same `Context` / id-mapping / error machinery. ```rust pub struct FlatIndex { provider: P, - /* private */ } impl FlatIndex

{ pub fn new(provider: P) -> Self; pub fn provider(&self) -> &P; - pub fn knn_search( + pub fn knn_search( &self, k: NonZeroUsize, strategy: &S, @@ -160,15 +173,16 @@ impl FlatIndex

{ T: ?Sized + Sync, O: Send, OB: SearchOutputBuffer + Send + ?Sized, + PP: for<'a> FlatPostProcess, T, O> + Send + Sync, } ``` The `knn_search` method is the canonical brute-force search algorithm: -1. Construct the iterator via `strategy.create_iter` to obtain a scoped iterator over the elements. +1. Construct the per-query callback via `strategy.create_callback`. 2. Build the query computer via `strategy.build_query_computer`. -3. Drive the scan via `iter.on_elements_unordered`, scoring each element and - inserting `Neighbor`s into a `NeighborPriorityQueue` of capacity `k`. +3. Drive the scan via `callback.distances_unordered(&computer, ...)`, inserting each + `(id, distance)` pair into a `NeighborPriorityQueue` of capacity `k`. 4. Hand the survivors (in distance order) to `processor.post_process`. 5. Return search stats. @@ -187,31 +201,28 @@ This design leans into using the `DataProvider` trait which requires implementat `to_external_id`), and error machinery are identical across graph and flat search, reducing the learning surface for new contributors. -### Async vs sync API for `FlatIterator` +### Async vs sync scan API -`next()` and `on_elements_unordered` return a future, making the trait -async. This is the right default for disk-backed and network-backed iterators -where advancing the cursor involves real I/O. It also matches the `Accessor` surface, +`on_elements_unordered` and `distances_unordered` return a future, making the scan +surface async. This is the right default for disk-backed and network-backed backends +where advancing the scan involves real I/O. It also matches the `Accessor` surface, keeping the two subsystems shaped the same way. -The cost is paid by in-memory consumers: every call to `next()` goes through the future -machinery even when the result is immediately available via `std::future::ready`. In a -tight brute-force loop this overhead — poll scaffolding, pinning etc — could be measurable. +The cost is paid by in-memory consumers: the scan goes through the future machinery +even when results are immediately available. In a tight brute-force loop this overhead — +poll scaffolding, pinning etc — could be measurable. We chose async because the wider audience of consumers (disk, network, mixed) benefits -more than in-memory consumers lose. +more than in-memory consumers lose. ### Expand `Element` to support batched distance computation? The current design yields one element per `next()` call, and the query computer scores -elements one at a time via `PreprocessedDistanceFunction::evaluate_similarity`. This could leave some optimization and performance on the table; especially with the upcoming effort around batched distance kernels. +elements one at a time via `PreprocessedDistanceFunction::evaluate_similarity`. This could leave some optimization and performance on the table; especially with the upcoming effort around batched distance kernels. Of course, a consumer can choose to implement their own optimized implementation of `distances_unordered` that uses batching. An alternative is to make `next()` yield a *batch* instead of a single vector representation like `Element<'_>`. Some work will need to be done to define the right interaction between the batch type, the element type in the batch, the interaction with `QueryComputer`'s types and way IDs and distances are collected in the queue. -We opted for the scalar-per-element design for now because it is simpler to implement and -reason about. The hope is that batched distance computation can be layered on later as an opt-in sub-trait without breaking -existing iterators. - ## Future Work - Support for other flat-search algorithms like - filtered, range and diverse flat algorithms as additional methods on `FlatIndex`. +- Index build -- this is just one part of the picture; more work needs to be done around how this fits in with any traits / interface we need for index build. From ee48f7dc752f6500b379c448ca42d29430ae49ff Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 28 Apr 2026 21:27:13 -0700 Subject: [PATCH 08/10] rustfmt flat module --- diskann/src/flat/index.rs | 15 ++++++++------- diskann/src/flat/iterator.rs | 8 +++----- diskann/src/flat/mod.rs | 4 +++- diskann/src/flat/strategy.rs | 6 +++++- 4 files changed, 19 insertions(+), 14 deletions(-) diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 825567f8a..f91a9ccac 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -22,7 +22,7 @@ use crate::{ /// A `'static` thin wrapper around a [`DataProvider`] used for flat search. /// /// The provider is owned by the index. The index is constructed once at process startup and -/// shared across requests; per-query state lives in the [`crate::flat::OnElementsUnordered`] +/// shared across requests; per-query state lives in the [`crate::flat::OnElementsUnordered`] /// implementation that the [`crate::flat::FlatSearchStrategy`] produces. #[derive(Debug)] pub struct FlatIndex { @@ -85,12 +85,13 @@ impl FlatIndex

{ let mut queue = NeighborPriorityQueue::new(k); let mut cmps: u32 = 0; - callback.distances_unordered(&computer, |id, dist| { - cmps += 1; - queue.insert(Neighbor::new(id, dist)); - }) - .await - .into_ann_result()?; + callback + .distances_unordered(&computer, |id, dist| { + cmps += 1; + queue.insert(Neighbor::new(id, dist)); + }) + .await + .into_ann_result()?; let result_count = processor .post_process(&mut callback, query, queue.iter().take(k), output) diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index fd69b312a..6ccebc214 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -34,7 +34,6 @@ pub trait OnElementsUnordered: HasId + Send + Sync { F: Send + for<'a> FnMut(Self::Id, Self::ElementRef<'a>); } - /// Extension of [`OnElementsUnordered`] that drives the scan with a pre-built query /// computer, invoking a callback with `(id, distance)` pairs instead of raw elements. /// @@ -42,7 +41,7 @@ pub trait OnElementsUnordered: HasId + Send + Sync { /// by the [`FlatSearchStrategy`](crate::flat::FlatSearchStrategy). /// /// The default implementation delegates to [`OnElementsUnordered::on_elements_unordered`], -/// calling `computer.evaluate_similarity` on each element. +/// calling `computer.evaluate_similarity` on each element. pub trait DistancesUnordered: OnElementsUnordered { /// Drive the entire scan, scoring each element with `computer` and invoking `f` with /// the resulting `(id, distance)` pair. @@ -90,12 +89,10 @@ pub trait FlatIterator: HasId + Send + Sync { ) -> impl SendFuture)>, Self::Error>>; } - /////////////// /// Default /// /////////////// - /// Bridges a [`FlatIterator`] into an [`OnElementsUnordered`] by looping over /// [`FlatIterator::next`] and reborrowing each element into the closure. /// @@ -142,4 +139,5 @@ where } } -impl DistancesUnordered for DefaultIteratedOperator where I: FlatIterator + HasId + Send + Sync {} \ No newline at end of file +impl DistancesUnordered for DefaultIteratedOperator where I: FlatIterator + HasId + Send + Sync +{} diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs index 2516509da..bf1290ed8 100644 --- a/diskann/src/flat/mod.rs +++ b/diskann/src/flat/mod.rs @@ -38,6 +38,8 @@ pub mod post_process; pub mod strategy; pub use index::FlatIndex; -pub use iterator::{DefaultIteratedOperator, DistancesUnordered, FlatIterator, OnElementsUnordered}; +pub use iterator::{ + DefaultIteratedOperator, DistancesUnordered, FlatIterator, OnElementsUnordered, +}; pub use post_process::{CopyIds, FlatPostProcess}; pub use strategy::FlatSearchStrategy; diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 9fa2c6e00..4b8ad8fe1 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -7,7 +7,11 @@ use diskann_vector::PreprocessedDistanceFunction; -use crate::{error::StandardError, flat::{DistancesUnordered, OnElementsUnordered}, provider::DataProvider}; +use crate::{ + error::StandardError, + flat::{DistancesUnordered, OnElementsUnordered}, + provider::DataProvider, +}; /// Per-call configuration that knows how to construct a [`FlatIterator`] for a provider /// and how to pre-process queries of type `T` into a distance computer. From 7fd903ecb31c63992340411b41bf919014383065 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Tue, 28 Apr 2026 21:31:55 -0700 Subject: [PATCH 09/10] fix clippy: replace doc-comment divider with regular comment --- diskann/src/flat/iterator.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index 6ccebc214..b725eeae9 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -89,9 +89,7 @@ pub trait FlatIterator: HasId + Send + Sync { ) -> impl SendFuture)>, Self::Error>>; } -/////////////// -/// Default /// -/////////////// +// ─── Default adapter ──────────────────────────────────────────────────────── /// Bridges a [`FlatIterator`] into an [`OnElementsUnordered`] by looping over /// [`FlatIterator::next`] and reborrowing each element into the closure. From 8af5e004255872b8306e8d78c7527fc7a6aa3820 Mon Sep 17 00:00:00 2001 From: "Aditya Krishnan (from Dev Box)" Date: Wed, 29 Apr 2026 10:31:05 -0700 Subject: [PATCH 10/10] small edits --- diskann/src/flat/iterator.rs | 8 +++++++- diskann/src/flat/strategy.rs | 2 +- rfcs/00983-flat-search.md | 2 +- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/diskann/src/flat/iterator.rs b/diskann/src/flat/iterator.rs index b725eeae9..fa07e92b4 100644 --- a/diskann/src/flat/iterator.rs +++ b/diskann/src/flat/iterator.rs @@ -61,6 +61,10 @@ pub trait DistancesUnordered: OnElementsUnordered { } } +////////////// +// Iterator // +////////////// + /// A lending, asynchronous iterator over the elements of a flat index. /// /// Implementations provide element-at-a-time access via [`Self::next`]. Providers that @@ -89,7 +93,9 @@ pub trait FlatIterator: HasId + Send + Sync { ) -> impl SendFuture)>, Self::Error>>; } -// ─── Default adapter ──────────────────────────────────────────────────────── +///////////// +// Default // +///////////// /// Bridges a [`FlatIterator`] into an [`OnElementsUnordered`] by looping over /// [`FlatIterator::next`] and reborrowing each element into the closure. diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 4b8ad8fe1..5da3349ed 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -13,7 +13,7 @@ use crate::{ provider::DataProvider, }; -/// Per-call configuration that knows how to construct a [`FlatIterator`] for a provider +/// Per-call configuration that knows how to construct a [`DistancesUnordered`] for a provider /// and how to pre-process queries of type `T` into a distance computer. /// /// `FlatSearchStrategy` is the flat counterpart to [`crate::graph::glue::SearchStrategy`]. diff --git a/rfcs/00983-flat-search.md b/rfcs/00983-flat-search.md index dfef8f516..b79018a5d 100644 --- a/rfcs/00983-flat-search.md +++ b/rfcs/00983-flat-search.md @@ -4,7 +4,7 @@ |------------------|--------------------------------| | **Authors** | Aditya Krishnan, Alex Razumov, Dongliang Wu | | **Created** | 2026-04-24 | -| **Updated** | 2026-04-27 | +| **Updated** | 2026-04-28 | ## Motivation