diff --git a/diskann-benchmark/src/flat/search.rs b/diskann-benchmark/src/flat/search.rs index f2f0d422d8..25f93f2687 100644 --- a/diskann-benchmark/src/flat/search.rs +++ b/diskann-benchmark/src/flat/search.rs @@ -5,15 +5,15 @@ //! Backend for flat-index (brute-force kNN) benchmarks. //! -//! This exercises [`diskann::flat::FlatIndex::knn_search`] over an in-memory +//! This exercises [`diskann::flat::knn_search`] over an in-memory //! provider, measuring recall and latency. use std::{io::Write, num::NonZeroUsize, sync::Arc}; use diskann::{ - flat::{DistancesUnordered, FlatIndex, SearchStrategy}, + flat::{knn_search, DistancesUnordered}, graph::{glue::CopyIds, SearchOutputBuffer}, - provider::{DataProvider, DefaultContext, HasId, NoopGuard}, + provider::HasId, utils::VectorRepr, ANNResult, }; @@ -53,29 +53,10 @@ pub(super) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> ///////////////// /// A minimal in-memory provider for flat search benchmarks. -/// -/// Wraps a loaded [`Matrix`] and implements [`DataProvider`] with identity -/// ID mapping. struct InMemProvider { data: Arc>, } -impl DataProvider for InMemProvider { - type Context = DefaultContext; - type InternalId = u32; - type ExternalId = u32; - type Error = diskann::ANNError; - type Guard = NoopGuard; - - fn to_internal_id(&self, _ctx: &DefaultContext, gid: &u32) -> Result { - Ok(*gid) - } - - fn to_external_id(&self, _ctx: &DefaultContext, id: u32) -> Result { - Ok(id) - } -} - struct Flat { _phantom: std::marker::PhantomData, } @@ -132,10 +113,9 @@ where ); writeln!(output, " Loaded {} vectors of dimension {}", nrows, ncols)?; - // Build the provider and wrap in FlatIndex + // Build the provider. let data = Arc::new(data); let provider = InMemProvider { data: data.clone() }; - let index = FlatIndex::new(provider); // Load queries and groundtruth let queries: Matrix = @@ -172,9 +152,9 @@ where let mut results = Vec::new(); let searcher = Arc::new(Searcher { - index, + provider, queries, - strategy: Strategy::new(metric), + metric, }); for &threads in &input.search.num_threads { @@ -202,50 +182,35 @@ where } } -/////////////////////// -// Flat SearchStrategy // -/////////////////////// - -/// A [`SearchStrategy`] implementation for [`InMemProvider`] that drives -/// a full sequential scan over all vectors. -struct Strategy { - metric: Metric, - _phantom: std::marker::PhantomData, +/// The visitor that iterates over all vectors in the provider. +struct Visitor<'a, T: VectorRepr> { + data: &'a Matrix, + computer: T::QueryDistance, } -impl Strategy { - fn new(metric: Metric) -> Self { +impl<'a, T: VectorRepr> Visitor<'a, T> { + fn new(provider: &'a InMemProvider, query: &[T], metric: Metric) -> Self { Self { - metric, - _phantom: std::marker::PhantomData, + data: &provider.data, + computer: T::query_distance(query, metric), } } } -/// The visitor that iterates over all vectors in the provider. -struct Visitor<'a, T> { - data: &'a Matrix, -} - impl HasId for Visitor<'_, T> { type Id = u32; } -impl DistancesUnordered for Visitor<'_, T> { - type ElementRef<'a> = &'a [T]; +impl DistancesUnordered for Visitor<'_, T> { type Error = diskann::error::Infallible; - fn distances_unordered( - &mut self, - computer: &T::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { async move { for (i, vector) in self.data.row_iter().enumerate() { - let dist = computer.evaluate_similarity(vector); + let dist = self.computer.evaluate_similarity(vector); f(i as u32, dist); } Ok(()) @@ -253,44 +218,15 @@ impl DistancesUnordered for Visitor<'_, T> { } } -impl SearchStrategy, &[T]> for Strategy { - type ElementRef<'a> = &'a [T]; - type QueryComputer = T::QueryDistance; - type QueryComputerError = diskann::error::Infallible; - type Visitor<'a> - = Visitor<'a, T> - where - Self: 'a, - InMemProvider: 'a; - type Error = diskann::error::Infallible; - - fn create_visitor<'a>( - &'a self, - provider: &'a InMemProvider, - _context: &'a DefaultContext, - ) -> Result, Self::Error> { - Ok(Visitor { - data: &provider.data, - }) - } - - fn build_query_computer( - &self, - query: &[T], - ) -> Result { - Ok(T::query_distance(query, self.metric)) - } -} - ////////////////////////////////////////// // benchmark_core::search::Search impl // ////////////////////////////////////////// -/// Wraps a [`FlatIndex`] and queries to implement [`search::Search`]. +/// Wraps a flat-search provider and queries to implement [`search::Search`]. struct Searcher { - index: FlatIndex>, + provider: InMemProvider, queries: Matrix, - strategy: Strategy, + metric: Metric, } /// Search parameters for flat-index benchmarks. @@ -331,20 +267,10 @@ where where O: SearchOutputBuffer + Send, { - let context = DefaultContext; let query = self.queries.row(index); + let mut visitor = Visitor::new(&self.provider, query, self.metric); - let stats = self - .index - .knn_search( - parameters.k, - &self.strategy, - CopyIds, - &context, - query, - buffer, - ) - .await?; + let stats = knn_search(&mut visitor, parameters.k, CopyIds, query, buffer).await?; Ok(Metrics { comparisons: stats.cmps, diff --git a/diskann/src/flat/index.rs b/diskann/src/flat/index.rs index 0c9aa6c81c..35f79bdaf0 100644 --- a/diskann/src/flat/index.rs +++ b/diskann/src/flat/index.rs @@ -3,8 +3,7 @@ * Licensed under the MIT license. */ -//! [`FlatIndex`] — the index wrapper for a [`DataProvider`] -//! over which we do flat search. +//! Brute-force k-nearest-neighbor search over a [`DistancesUnordered`] visitor. use std::num::NonZeroUsize; use diskann_utils::future::SendFuture; @@ -12,10 +11,9 @@ use diskann_utils::future::SendFuture; use crate::{ ANNResult, error::{ErrorExt, IntoANNResult}, - flat::{DistancesUnordered, SearchStrategy}, + flat::DistancesUnordered, graph::{SearchOutputBuffer, glue::SearchPostProcess}, neighbor::{Neighbor, NeighborPriorityQueue}, - provider::DataProvider, }; /// Statistics collected during a flat search. @@ -28,74 +26,53 @@ pub struct SearchStats { pub result_count: u32, } -/// A thin wrapper around a [`DataProvider`] used for flat search. -#[derive(Debug)] -pub struct FlatIndex { - /// The backing provider. - provider: P, -} - -impl FlatIndex

{ - /// Construct a new [`FlatIndex`] around `provider`. - pub fn new(provider: P) -> Self { - Self { provider } - } - - /// 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 visitor through the query - /// computer, keeps the best `k` candidates in a [`NeighborPriorityQueue`], then runs - /// `processor` over the survivors to populate `output`. - /// - /// The post-processor [`SearchPostProcess::post_process`] outputs the number - /// of results that survive, which is returned as `SearchStats::result_count`. - pub fn knn_search( - &self, - k: NonZeroUsize, - strategy: &S, - processor: PP, - context: &P::Context, - query: T, - output: &mut OB, - ) -> impl SendFuture> - where - S: SearchStrategy, - T: Copy + Send + Sync, - O: Send, - PP: for<'a> SearchPostProcess, T, O> + Send + Sync, - OB: SearchOutputBuffer + Send + ?Sized, - { - async move { - let mut visitor = strategy - .create_visitor(&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; - - visitor - .distances_unordered(&computer, |id, dist| { - cmps += 1; - queue.insert(Neighbor::new(id, dist)); - }) - .await - .escalate("flat scan must complete to produce correct k-NN results")?; - - let result_count = processor - .post_process(&mut visitor, query, queue.iter().take(k), output) - .await - .into_ann_result()? as u32; - - Ok(SearchStats { cmps, result_count }) - } +/// Brute-force k-nearest-neighbor search over an initialized visitor. +/// +/// Borrows `visitor` for the duration of the search, streams every distance it produces, +/// keeps the best `k` candidates in a [`NeighborPriorityQueue`], then runs `processor` +/// over the same visitor and the surviving candidates to populate `output`. +/// +/// The visitor remains owned by the caller and becomes available again after the returned +/// future completes. Whether it supports another scan is determined by its implementation. +/// +/// # Errors +/// +/// Returns an error if distance scanning or result post-processing fails. Distance-scan +/// errors are escalated because a partial flat scan cannot produce correct k-nearest-neighbor +/// results. +pub fn knn_search( + visitor: &mut V, + k: NonZeroUsize, + processor: PP, + query: T, + output: &mut OB, +) -> impl SendFuture> +where + V: DistancesUnordered, + T: Copy + Send + Sync, + O: Send, + PP: SearchPostProcess + Send + Sync, + OB: SearchOutputBuffer + Send + ?Sized, +{ + async move { + let k = k.get(); + let mut queue = NeighborPriorityQueue::new(k); + let mut cmps: u32 = 0; + + visitor + .distances_unordered(|id, dist| { + cmps += 1; + queue.insert(Neighbor::new(id, dist)); + }) + .await + .escalate("flat scan must complete to produce correct k-NN results")?; + + let result_count = processor + .post_process(visitor, query, queue.iter().take(k), output) + .await + .into_ann_result()? as u32; + + Ok(SearchStats { cmps, result_count }) } } @@ -105,30 +82,27 @@ impl FlatIndex

{ #[cfg(test)] mod tests { - use crate::flat::{ - FlatIndex, - test::{ - harness::{CopyIdsOracle, EvenIdsOnlyOracle, KnnOracleRun, OracleProcessor}, - provider::{self as flat_provider}, - }, + use crate::flat::test::{ + harness::{CopyIdsOracle, EvenIdsOnlyOracle, KnnOracleRun, OracleProcessor}, + provider::{self as flat_provider}, }; use crate::graph::test::synthetic::Grid; - fn fixture(grid: Grid, size: usize) -> (FlatIndex, usize) { + fn fixture(grid: Grid, size: usize) -> (flat_provider::Provider, usize) { let provider = flat_provider::Provider::grid(grid, size).unwrap(); let len = provider.len(); - (FlatIndex::new(provider), len) + (provider, len) } - /// `knn_search` returns a `Send` future, and a shared `&FlatIndex` can serve + /// `knn_search` returns a `Send` future, and a shared provider can serve /// many concurrent searches on a multi-threaded runtime, each producing the /// correct output independently. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn multithreaded_knn_search() { use std::sync::Arc; - let (index, len) = fixture(Grid::Two, 4); - let index = Arc::new(index); + let (provider, len) = fixture(Grid::Two, 4); + let provider = Arc::new(provider); // Mix of corner, axis-aligned, and off-grid queries; k spans 1..=len. let cases: &[(&[f32], usize)] = &[ @@ -145,34 +119,28 @@ mod tests { /// Spawn every `(query, k)` case under `oracle` onto `set`. fn spawn_cases( set: &mut tokio::task::JoinSet<(Vec, usize, KnnOracleRun)>, - index: &Arc>, + provider: &Arc, oracle: O, cases: &[(&[f32], usize)], ) where O: OracleProcessor + Copy + Send + Sync + 'static, { for (query, k) in cases { - let index = Arc::clone(index); + let provider = Arc::clone(provider); let query: Vec = query.to_vec(); let k = *k; set.spawn(async move { - let outcome = KnnOracleRun::run( - &index, - &flat_provider::Strategy::new(index.provider().dim()), - &oracle, - &query, - k, - ) - .await - .expect("knn_search failed"); + let outcome = KnnOracleRun::run(&provider, &oracle, &query, k) + .await + .expect("knn_search failed"); (query, k, outcome) }); } } let mut set = tokio::task::JoinSet::new(); - spawn_cases(&mut set, &index, CopyIdsOracle, cases); - spawn_cases(&mut set, &index, EvenIdsOnlyOracle, cases); + spawn_cases(&mut set, &provider, CopyIdsOracle, cases); + spawn_cases(&mut set, &provider, EvenIdsOnlyOracle, cases); while let Some(joined) = set.join_next().await { let (query, k, outcome) = joined.expect("task panicked"); @@ -197,11 +165,14 @@ mod tests { fn transient_scan_error() { // The flat scan touches every id, so any transient id is guaranteed to be hit. for transient_ids in [&[0u32][..], &[3][..], &[1, 2, 5][..]] { - let strategy = - flat_provider::Strategy::with_transient(2, transient_ids.iter().copied()); - let (index, _) = fixture(Grid::Two, 3); - let err = KnnOracleRun::run_sync(&index, &strategy, &CopyIdsOracle, &[1.0, 0.0], 4) - .expect_err("transient error during full scan must escalate"); + let (provider, _) = fixture(Grid::Two, 3); + let query = &[1.0, 0.0]; + let visitor = + flat_provider::Visitor::flaky(&provider, query, transient_ids.iter().copied()) + .unwrap(); + let err = + KnnOracleRun::run_sync_with_visitor(&provider, visitor, &CopyIdsOracle, query, 4) + .expect_err("transient error during full scan must escalate"); let msg = format!("{err}"); assert!( @@ -216,10 +187,10 @@ mod tests { /// Run `knn_search` via the harness, assert it fails, and check the error /// message contains `expected_msg`. - fn assert_search_error(strategy: &flat_provider::Strategy, query: &[f32], expected_msg: &str) { - let (index, _) = fixture(Grid::Two, 3); - let err = KnnOracleRun::run_sync(&index, strategy, &CopyIdsOracle, query, 4) - .expect_err("expected knn_search to fail"); + fn assert_visitor_error(query: &[f32], expected_msg: &str) { + let (provider, _) = fixture(Grid::Two, 3); + let err = flat_provider::Visitor::new(&provider, query) + .expect_err("expected visitor construction to fail"); let msg = format!("{err}"); assert!( @@ -229,19 +200,7 @@ mod tests { } #[test] - fn strategy_constructor_errors() { - // Strategy/provider expect dim=2, query has dim=3. - assert_search_error( - &flat_provider::Strategy::new(2), - &[0.0, 0.0, 0.0], - "dimension mismatch", - ); - - // Strategy expects dim=5, provider has dim=2. - assert_search_error( - &flat_provider::Strategy::new(5), - &[0.0, 0.0], - "dimension mismatch", - ); + fn visitor_constructor_errors() { + assert_visitor_error(&[0.0, 0.0, 0.0], "dimension mismatch"); } } diff --git a/diskann/src/flat/mod.rs b/diskann/src/flat/mod.rs index f33d97851d..6f3a4ad2cb 100644 --- a/diskann/src/flat/mod.rs +++ b/diskann/src/flat/mod.rs @@ -12,22 +12,19 @@ //! //! # Architecture //! -//! The module mirrors the layering used by graph search: +//! The search algorithm operates directly on an initialized visitor: //! //! | Graph (random access) | Flat (sequential) | Shared? | //! | :------------------------------------ | :----------------------------------------- |:--------- | -//! | [`crate::provider::DataProvider`] | [`crate::provider::DataProvider`] | Yes | -//! | [`crate::graph::DiskANNIndex`] | [`FlatIndex`] | No | //! | [`crate::graph::glue::SearchAccessor`] | [`DistancesUnordered`] | No | -//! | [`crate::graph::glue::SearchStrategy`] | [`SearchStrategy`] | No | -//! | [`crate::graph::Search`] | [`FlatIndex::knn_search`] | No | +//! | [`crate::graph::Search`] | [`knn_search`] | No | //! | [`crate::graph::glue::SearchPostProcess`] | [`crate::graph::glue::SearchPostProcess`] | Yes | //! pub mod index; pub mod strategy; -pub use index::{FlatIndex, SearchStats}; -pub use strategy::{DistancesUnordered, SearchStrategy}; +pub use index::{SearchStats, knn_search}; +pub use strategy::DistancesUnordered; #[cfg(test)] mod test; diff --git a/diskann/src/flat/strategy.rs b/diskann/src/flat/strategy.rs index 1cdb6957fd..f186c06cca 100644 --- a/diskann/src/flat/strategy.rs +++ b/diskann/src/flat/strategy.rs @@ -3,101 +3,40 @@ * Licensed under the MIT license. */ -//! Core flat-search traits: [`DistancesUnordered`] and [`SearchStrategy`]. +//! Core flat-search trait: [`DistancesUnordered`]. use std::fmt::Debug; use diskann_utils::future::SendFuture; -use diskann_vector::PreprocessedDistanceFunction; -use crate::{ - error::{StandardError, ToRanked}, - provider::{DataProvider, HasId}, -}; +use crate::{error::ToRanked, provider::HasId}; -/// Fused iterate-and-score primitive over the elements of a flat index. +/// A complete, unordered scan of a candidate set. /// -/// Implementations drive an entire scan over the underlying data, scoring each element -/// with the supplied computer `C` and invoking `f` with the resulting `(id, distance)` -/// pair. The associated [`Self::ElementRef`] is the reference shape on which `C` must -/// be able to compute distances. -pub trait DistancesUnordered: HasId + Send + Sync -where - C: for<'a> PreprocessedDistanceFunction, f32>, -{ - /// Lifetime is intentionally unconstrained so it can appear under HRTB without - /// inducing a `'static` bound on `Self`. - type ElementRef<'a>; - - /// The error type for [`Self::distances_unordered`]. +/// Implementations drive the scan rather than exposing elements for the caller to fetch. +/// For each candidate in the scan, the implementation computes its distance and invokes +/// the supplied callback with the candidate's `(id, distance)` pair. No ordering of those +/// callbacks is guaranteed. +/// +/// A visitor represents one initialized scan. It may retain any state needed while +/// scanning, and remains available to the search post-processor after the scan completes. +pub trait DistancesUnordered: HasId + Send + Sync { + /// The error type returned when the visitor cannot complete its scan. type Error: ToRanked + Debug + Send + Sync + 'static; - /// 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, - f: F, - ) -> impl SendFuture> + /// Scan all candidates represented by this visitor and invoke `f` for each result. + /// + /// The callback may have been invoked before an error is returned. In that case, + /// those results form an incomplete scan and must not be treated as exhaustive. + fn distances_unordered(&mut self, f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32); } -/// Per-call configuration that knows how to construct a per-query -/// [`DistancesUnordered`] visitor for a provider, and the [`Self::QueryComputer`] used -/// to score each element during the scan. -pub trait SearchStrategy: Send + Sync -where - P: DataProvider, -{ - /// The reference element shape on which [`Self::QueryComputer`] computes - /// distances. - type ElementRef<'a>; - - /// The concrete query-computer type. - type QueryComputer: for<'a> PreprocessedDistanceFunction, f32> - + Send - + Sync - + 'static; - - /// The error type for [`Self::build_query_computer`]. - type QueryComputerError: StandardError; - - /// The visitor type produced by [`Self::create_visitor`]. - type Visitor<'a>: for<'b> DistancesUnordered< - Self::QueryComputer, - ElementRef<'b> = Self::ElementRef<'b>, - Id = P::InternalId, - > - where - Self: 'a, - P: 'a; - - /// The error type for [`Self::create_visitor`]. - type Error: StandardError; - - /// Construct a fresh visitor over `provider` for the given request `context`. - fn create_visitor<'a>( - &'a self, - provider: &'a P, - context: &'a P::Context, - ) -> Result, Self::Error>; - - /// Construct the per-query computer. - fn build_query_computer( - &self, - query: T, - ) -> Result; -} - #[cfg(test)] mod tests { - //! Direct [`DistancesUnordered`] impls over a few in-memory fixtures: a - //! happy-path scanner over `&[f32]` elements, a scanner whose `ElementRef<'a>` - //! is a lifetime-carrying non-reference type, and a scanner that fails - //! mid-stream. - - use std::marker::PhantomData; + //! Direct [`DistancesUnordered`] impls over in-memory fixtures, including a + //! happy-path scanner and one that fails mid-stream. use diskann_utils::future::SendFuture; use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; @@ -121,27 +60,23 @@ mod tests { /// Scans `items` in order, scoring each with the supplied computer. struct Scanner { items: Vec<(u32, Vec)>, + computer: ::QueryDistance, } impl HasId for Scanner { type Id = u32; } - impl DistancesUnordered<::QueryDistance> for Scanner { - type ElementRef<'a> = &'a [f32]; + impl DistancesUnordered for Scanner { type Error = Infallible; - fn distances_unordered( - &mut self, - computer: &::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { async move { for (id, v) in &self.items { - let dist = computer.evaluate_similarity(v.as_slice()); + let dist = self.computer.evaluate_similarity(v.as_slice()); f(*id, dist); } Ok(()) @@ -153,25 +88,69 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn distances_unordered_scanner() { let query = vec![0.5_f32, 0.9]; - let computer = f32::query_distance(&query, Metric::L2); + let expected_computer = f32::query_distance(&query, Metric::L2); let expected: Vec<(u32, f32)> = sample_items() .into_iter() - .map(|(id, v)| (id, computer.evaluate_similarity(v.as_slice()))) + .map(|(id, v)| (id, expected_computer.evaluate_similarity(v.as_slice()))) .collect(); let mut scanner = Scanner { items: sample_items(), + computer: f32::query_distance(&query, Metric::L2), }; let mut seen: Vec<(u32, f32)> = Vec::new(); scanner - .distances_unordered(&computer, |id, d| seen.push((id, d))) + .distances_unordered(|id, d| seen.push((id, d))) .await .unwrap(); assert_eq!(seen, expected); } + struct BorrowingScanner<'a> { + items: &'a [(u32, f32)], + query: &'a f32, + } + + impl HasId for BorrowingScanner<'_> { + type Id = u32; + } + + impl DistancesUnordered for BorrowingScanner<'_> { + type Error = Infallible; + + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> + where + F: Send + FnMut(Self::Id, f32), + { + async move { + for (id, value) in self.items { + f(*id, (*value - *self.query).abs()); + } + Ok(()) + } + } + } + + #[tokio::test] + async fn accessor_can_borrow_query_state() { + let items = [(10, 1.0), (11, 4.0)]; + let query = 2.0; + let mut scanner = BorrowingScanner { + items: &items, + query: &query, + }; + let mut seen = Vec::new(); + + scanner + .distances_unordered(|id, distance| seen.push((id, distance))) + .await + .unwrap(); + + assert_eq!(seen, [(10, 1.0), (11, 2.0)]); + } + /////////////////////////// // Failing scanner // /////////////////////////// @@ -189,21 +168,17 @@ mod tests { struct Failing { items: Vec<(u32, Vec)>, fail_after: usize, + computer: ::QueryDistance, } impl HasId for Failing { type Id = u32; } - impl DistancesUnordered<::QueryDistance> for Failing { - type ElementRef<'a> = &'a [f32]; + impl DistancesUnordered for Failing { type Error = Boom; - fn distances_unordered( - &mut self, - computer: &::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { @@ -212,7 +187,7 @@ mod tests { if i == self.fail_after { return Err(Boom(*id)); } - let dist = computer.evaluate_similarity(v.as_slice()); + let dist = self.computer.evaluate_similarity(v.as_slice()); f(*id, dist); } Ok(()) @@ -227,14 +202,12 @@ mod tests { let mut scanner = Failing { items: sample_items(), fail_after: 1, // Yield item 0 successfully, fail on item 1. + computer: f32::query_distance(&[0.0, 0.0], Metric::L2), }; - let query = vec![0.0_f32, 0.0]; - let computer = f32::query_distance(&query, Metric::L2); - let mut seen: Vec = Vec::new(); let err = scanner - .distances_unordered(&computer, |id, _d| seen.push(id)) + .distances_unordered(|id, _d| seen.push(id)) .await .expect_err("Failing scanner must surface its error"); @@ -245,106 +218,4 @@ mod tests { "the closure must only see items yielded before the failure", ); } - - ///////////////////////////////////////////// - // Lifetime-carrying concrete `ElementRef` // - ///////////////////////////////////////////// - - struct View<'a> { - ptr: *const f32, - len: usize, - _phantom: PhantomData<&'a [f32]>, - } - - // SAFETY: `View<'a>` semantically carries a `&'a [f32]`, which is `Send + Sync`. - unsafe impl Send for View<'_> {} - unsafe impl Sync for View<'_> {} - - /// Computer that reconstructs a `&[f32]` from a [`View`]'s ptr+len and - /// computes inner product against a stored query. - struct ViewComputer { - query: Vec, - } - - impl<'a> PreprocessedDistanceFunction, f32> for ViewComputer { - fn evaluate_similarity(&self, v: View<'a>) -> f32 { - // SAFETY: `v.ptr` / `v.len` were produced from a `&'a [f32]` held by the - // scanner that owns the backing `Vec`; the phantom lifetime ties this view - // to that borrow, so the slice is valid for the duration of this call. - let s = unsafe { std::slice::from_raw_parts(v.ptr, v.len) }; - s.iter().zip(&self.query).map(|(a, b)| a * b).sum() - } - } - - /// Scans `rows`, yielding a [`View`] tied (via its phantom lifetime) to the - /// borrow of the underlying `Vec`. - struct ViewScanner { - rows: Vec<(u32, Vec)>, - } - - impl ViewScanner { - fn iter<'a>(&self) -> impl Iterator)> { - self.rows.iter().map(|(x, y)| { - ( - *x, - View { - ptr: y.as_ptr(), - len: y.len(), - _phantom: PhantomData, - }, - ) - }) - } - } - - impl HasId for ViewScanner { - type Id = u32; - } - - impl DistancesUnordered for ViewScanner { - type ElementRef<'a> = View<'a>; - type Error = Infallible; - - fn distances_unordered( - &mut self, - computer: &ViewComputer, - mut f: F, - ) -> impl SendFuture> - where - F: Send + FnMut(Self::Id, f32), - { - async move { - for (id, v) in self.iter() { - f(id, computer.evaluate_similarity(v)); - } - Ok(()) - } - } - } - - #[tokio::test] - async fn distances_unordered_lifetime_carrying_element_ref() { - let mut scanner = ViewScanner { - rows: vec![ - (10, vec![1.0, 0.0]), - (11, vec![0.5, 0.5]), - (12, vec![0.0, 2.0]), - ], - }; - let computer = ViewComputer { - query: vec![1.0, 3.0], - }; - let expected: Vec<(u32, f32)> = vec![ - (10, 1.0 * 1.0 + 0.0 * 3.0), - (11, 0.5 * 1.0 + 0.5 * 3.0), - (12, 0.0 * 1.0 + 2.0 * 3.0), - ]; - - let mut seen: Vec<(u32, f32)> = Vec::new(); - scanner - .distances_unordered(&computer, |id, d| seen.push((id, d))) - .await - .unwrap(); - assert_eq!(seen, expected); - } } diff --git a/diskann/src/flat/test/cases/flat_knn_search.rs b/diskann/src/flat/test/cases/flat_knn_search.rs index b6beca9801..5bc21ad1f2 100644 --- a/diskann/src/flat/test/cases/flat_knn_search.rs +++ b/diskann/src/flat/test/cases/flat_knn_search.rs @@ -3,20 +3,17 @@ * Licensed under the MIT license. */ -//! Baseline-cached regression sweep for [`crate::flat::FlatIndex::knn_search`]. +//! Baseline-cached regression sweep for [`crate::flat::knn_search`]. //! -//! Bbuilds a fresh index per parameter combination, runs `knn_search` through the +//! Builds a fresh provider per parameter combination, runs `knn_search` through the //! [`crate::flat::test::harness`], snapshots the result + statistics into //! [`FlatKnnBaseline`], and compares the entire batch against the JSON committed under //! `diskann/test/generated/flat/test/cases/flat_knn_search/`. use crate::{ - flat::{ - FlatIndex, - test::{ - harness, - provider::{self as flat_provider, ElementCounter, Strategy}, - }, + flat::test::{ + harness, + provider::{self as flat_provider, ElementCounter}, }, graph::test::synthetic::Grid, test::{ @@ -34,7 +31,7 @@ fn root() -> TestRoot { const KS: [usize; 3] = [1, 4, 10]; /// One row of the baseline JSON: a single `(grid, size, query, k)` execution of -/// `FlatIndex::knn_search` plus the brute-force ground truth, search stats, and +/// `knn_search` plus the brute-force ground truth, search stats, and /// per-row provider metrics. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] struct FlatKnnBaseline { @@ -88,29 +85,23 @@ verbose_eq!(FlatKnnBaseline { metrics, }); -/// Run `knn_search` + brute-force oracle against a *shared* `index`, assert the +/// Run `knn_search` + brute-force oracle against a *shared* provider, assert the /// cross-row invariants, and produce the baseline row. The per-row provider metrics /// captured into the baseline are the *delta* observed during this row, which keeps /// the snapshot independent of how many rows preceded it. fn run_row( - index: &FlatIndex, + provider: &flat_provider::Provider, grid_dim: usize, grid_size: usize, query: &[f32], k: usize, desc: &str, ) -> FlatKnnBaseline { - let len = index.provider().len(); - let metrics_before = index.provider().metrics(); - - let outcome = harness::KnnOracleRun::run_sync( - index, - &Strategy::new(index.provider().dim()), - &harness::CopyIdsOracle, - query, - k, - ) - .unwrap(); + let len = provider.len(); + let metrics_before = provider.metrics(); + + let outcome = + harness::KnnOracleRun::run_sync(provider, &harness::CopyIdsOracle, query, k).unwrap(); let stats = outcome.stats; assert_eq!( @@ -129,7 +120,7 @@ fn run_row( "flat scan top-k distance multiset must agree with brute force", ); - let metrics_after = index.provider().metrics(); + let metrics_after = provider.metrics(); let metrics = ElementCounter { count: metrics_after.count - metrics_before.count, }; @@ -159,8 +150,8 @@ fn run_row( fn _flat_knn_search(grid: Grid, size: usize, mut parent: TestPath<'_>) { let dim: usize = grid.dim().into(); - // Build the provider and index once, mirroring the production pattern where a - // single index serves many queries. + // Build the provider once, mirroring the production pattern where one provider + // serves many queries. let provider = flat_provider::Provider::grid(grid, size).unwrap(); let len = provider.len(); assert_eq!( @@ -168,8 +159,6 @@ fn _flat_knn_search(grid: Grid, size: usize, mut parent: TestPath<'_>) { size.pow(dim as u32), "flat::test::Provider::grid should produce size^dim rows", ); - let index = FlatIndex::new(provider); - let queries: [(Vec, &str); 2] = [ ( vec![-1.0; dim], @@ -181,12 +170,11 @@ fn _flat_knn_search(grid: Grid, size: usize, mut parent: TestPath<'_>) { ), ]; - let index_ref = &index; let results: Vec = queries .iter() .flat_map(|(q, desc)| { KS.iter() - .map(move |&k| run_row(index_ref, dim, size, q, k, desc)) + .map(|&k| run_row(&provider, dim, size, q, k, desc)) }) .collect(); diff --git a/diskann/src/flat/test/harness.rs b/diskann/src/flat/test/harness.rs index 553b9fab82..097760a0e9 100644 --- a/diskann/src/flat/test/harness.rs +++ b/diskann/src/flat/test/harness.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Reusable execution harness for [`crate::flat::FlatIndex`] tests. +//! Reusable execution harness for [`crate::flat::knn_search`] tests. //! //! Use [`KnnOracleRun::run`] to drive `knn_search` under a chosen [`OracleProcessor`] //! and pair the result with the oracle's expected post-processed output. @@ -14,9 +14,10 @@ use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; use crate::{ ANNResult, + error::IntoANNResult, flat::{ - FlatIndex, SearchStats, - test::provider::{Provider, Strategy, Visitor}, + SearchStats, knn_search, + test::provider::{Provider, Visitor}, }, graph::{ SearchOutputBuffer, @@ -28,7 +29,7 @@ use crate::{ utils::VectorRepr, }; -/// Result of running [`FlatIndex::knn_search`] under the harness alongside the +/// Result of running [`knn_search`] under the harness alongside the /// oracle's expected post-processed output. #[derive(Debug, Clone)] pub(crate) struct KnnOracleRun { @@ -46,41 +47,64 @@ pub(crate) struct KnnOracleRun { } impl KnnOracleRun { - /// Run [`FlatIndex::knn_search`] once under `oracle`, blocking on a fresh + /// Run [`knn_search`] once under `oracle`, blocking on a fresh /// single-threaded runtime, and pair the result with the oracle's expected output. pub fn run_sync( - index: &FlatIndex, - strategy: &Strategy, + provider: &Provider, oracle: &O, query: &[f32], k: usize, ) -> ANNResult { - current_thread_runtime().block_on(Self::run(index, strategy, oracle, query, k)) + current_thread_runtime().block_on(Self::run(provider, oracle, query, k)) + } + + /// Run [`knn_search`] with an already initialized visitor. + pub fn run_sync_with_visitor( + provider: &Provider, + mut visitor: Visitor<'_>, + oracle: &O, + query: &[f32], + k: usize, + ) -> ANNResult { + current_thread_runtime().block_on(Self::run_with_visitor( + provider, + &mut visitor, + oracle, + query, + k, + )) } /// Async variant of [`KnnOracleRun::run_sync`]. Use this from tests that already /// have a Tokio runtime (e.g. `#[tokio::test]`) or that need to drive /// `knn_search` concurrently across tasks. pub async fn run( - index: &FlatIndex, - strategy: &Strategy, + provider: &Provider, + oracle: &O, + query: &[f32], + k: usize, + ) -> ANNResult { + let mut visitor = Visitor::new(provider, query).into_ann_result()?; + Self::run_with_visitor(provider, &mut visitor, oracle, query, k).await + } + + async fn run_with_visitor( + provider: &Provider, + visitor: &mut Visitor<'_>, oracle: &O, query: &[f32], k: usize, ) -> ANNResult { - let context = crate::flat::test::provider::Context::new(); let mut buf = vec![Neighbor::::default(); k]; - let stats = index - .knn_search( - NonZeroUsize::new(k).expect("flat::test::harness requires k > 0"), - strategy, - oracle.processor(), - &context, - query, - &mut BackInserter::new(buf.as_mut_slice()), - ) - .await?; + let stats = knn_search( + visitor, + NonZeroUsize::new(k).expect("flat::test::harness requires k > 0"), + oracle.processor(), + query, + &mut BackInserter::new(buf.as_mut_slice()), + ) + .await?; let mut top_k: Vec> = buf .iter() @@ -90,8 +114,7 @@ impl KnnOracleRun { sort_neighbors(&mut top_k); let top_k_distances = top_k.iter().map(|n| *n.distance()).collect(); - let ground_truth = - oracle.expected(brute_force_topk(index.provider(), Metric::L2, query, k)); + let ground_truth = oracle.expected(brute_force_topk(provider, Metric::L2, query, k)); Ok(Self { top_k: top_k.into_iter().map(Neighbor::as_tuple).collect(), @@ -111,7 +134,7 @@ pub(crate) trait OracleProcessor { /// The post-processor exercised by the search. type Processor: for<'a, 'q> SearchPostProcess, &'q [f32], u32> + Send + Sync; - /// Construct the processor instance fed to [`FlatIndex::knn_search`]. + /// Construct the processor instance fed to [`knn_search`]. fn processor(&self) -> Self::Processor; /// Transform the brute-force top-`k` neighbors into the output the processor is diff --git a/diskann/src/flat/test/provider.rs b/diskann/src/flat/test/provider.rs index 3bd0f3d31e..3a8f6b2900 100644 --- a/diskann/src/flat/test/provider.rs +++ b/diskann/src/flat/test/provider.rs @@ -6,11 +6,9 @@ //! Self-contained test provider for the flat-search module. use std::{ - borrow::Cow, collections::HashSet, fmt::{self, Debug}, future::Future, - sync::Arc, }; use diskann_utils::{future::SendFuture, views::Matrix}; @@ -20,7 +18,7 @@ use thiserror::Error; use crate::{ always_escalate, convert_error, error::{RankedError, ToRanked, TransientError}, - flat::{DistancesUnordered, SearchStrategy}, + flat::DistancesUnordered, graph::test::synthetic::Grid, internal::counter::{Counter, LocalCounter}, provider::{self, ExecutionContext, HasId, NoopGuard}, @@ -124,12 +122,6 @@ crate::test::cmp::verbose_eq!(ElementCounter { count }); #[derive(Debug, Clone, Default)] pub struct Context; -impl Context { - pub fn new() -> Self { - Self - } -} - impl ExecutionContext for Context { fn wrap_spawn(&self, f: F) -> impl Future + Send + 'static where @@ -281,28 +273,50 @@ impl provider::DataProvider for Provider { /// set of ids. pub struct Visitor<'a> { provider: &'a Provider, - transient_ids: Option>>, + transient_ids: Option>, get_element: LocalCounter<'a>, + computer: ::QueryDistance, } impl<'a> Visitor<'a> { /// Construct a visitor with no fault injection. - pub fn new(provider: &'a Provider) -> Self { - Self { + pub fn new(provider: &'a Provider, query: &[f32]) -> Result { + let computer = Self::query_computer(provider, query)?; + Ok(Self { provider, transient_ids: None, get_element: provider.get_element.local(), - } + computer, + }) } /// Construct a visitor that returns a [`TransientGetError`] for any id in /// `transient_ids`. Other ids behave normally. - pub fn flaky(provider: &'a Provider, transient_ids: Cow<'a, HashSet>) -> Self { - Self { + pub fn flaky( + provider: &'a Provider, + query: &[f32], + transient_ids: impl IntoIterator, + ) -> Result { + let computer = Self::query_computer(provider, query)?; + Ok(Self { provider, - transient_ids: Some(transient_ids), + transient_ids: Some(transient_ids.into_iter().collect()), get_element: provider.get_element.local(), + computer, + }) + } + + fn query_computer( + provider: &Provider, + query: &[f32], + ) -> Result<::QueryDistance, VisitorError> { + if query.len() != provider.dim() { + return Err(VisitorError { + expected: provider.dim(), + actual: query.len(), + }); } + Ok(f32::query_distance(query, Metric::L2)) } } @@ -319,15 +333,10 @@ impl HasId for Visitor<'_> { type Id = u32; } -impl DistancesUnordered<::QueryDistance> for Visitor<'_> { - type ElementRef<'a> = &'a [f32]; +impl DistancesUnordered for Visitor<'_> { type Error = AccessError; - fn distances_unordered( - &mut self, - computer: &::QueryDistance, - mut f: F, - ) -> impl SendFuture> + fn distances_unordered(&mut self, mut f: F) -> impl SendFuture> where F: Send + FnMut(Self::Id, f32), { @@ -340,7 +349,7 @@ impl DistancesUnordered<::QueryDistance> for Visitor<'_> { return Err(AccessError::Transient(TransientGetError::new(id))); } self.get_element.increment(); - let dist = computer.evaluate_similarity(vector); + let dist = self.computer.evaluate_similarity(vector); f(id, dist); } Ok(()) @@ -348,84 +357,12 @@ impl DistancesUnordered<::QueryDistance> for Visitor<'_> { } } -////////////// -// Strategy // -////////////// - -/// Error from [`Strategy::create_visitor`] or [`Strategy::build_query_computer`] -/// when dimensions don't match. +/// Error from visitor construction when dimensions don't match. #[derive(Debug, Clone, Error)] -#[error("dimension mismatch: strategy expects {expected}, got {actual}")] -pub struct StrategyError { +#[error("dimension mismatch: provider expects {expected}, got {actual}")] +pub struct VisitorError { pub expected: usize, pub actual: usize, } -convert_error!(StrategyError); - -/// Factory of [`Visitor`]s that validates dimensions and optionally injects -/// transient errors into the scan. -#[derive(Clone, Debug)] -pub struct Strategy { - dim: usize, - transient_ids: Option>>, -} - -impl Strategy { - /// Construct a strategy expecting vectors of dimension `dim`. - pub fn new(dim: usize) -> Self { - Self { - dim, - transient_ids: None, - } - } - - /// Construct a strategy whose visitors return a transient error on `get_element` - /// for every id in `transient_ids`. - pub fn with_transient(dim: usize, transient_ids: impl IntoIterator) -> Self { - Self { - dim, - transient_ids: Some(Arc::new(transient_ids.into_iter().collect())), - } - } -} - -impl SearchStrategy for Strategy { - type ElementRef<'a> = &'a [f32]; - type QueryComputer = ::QueryDistance; - type QueryComputerError = StrategyError; - type Visitor<'a> = Visitor<'a>; - type Error = StrategyError; - - fn create_visitor<'a>( - &'a self, - provider: &'a Provider, - _context: &'a Context, - ) -> Result, Self::Error> { - let actual = provider.dim(); - if actual != self.dim { - return Err(StrategyError { - expected: self.dim, - actual, - }); - } - let visitor = match &self.transient_ids { - Some(ids) => Visitor::flaky(provider, Cow::Borrowed(ids)), - None => Visitor::new(provider), - }; - Ok(visitor) - } - - fn build_query_computer( - &self, - from: &[f32], - ) -> Result { - if from.len() != self.dim { - return Err(StrategyError { - expected: self.dim, - actual: from.len(), - }); - } - Ok(f32::query_distance(from, Metric::L2)) - } -} +convert_error!(VisitorError);