[flat index] Flat Search Interface

diskann/src/flat/index.rs

@@ -0,0 +1,223 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! [`FlatIndex`] — the index wrapper for a [`DataProvider`](crate::provider::DataProvider)
+//! over which we do flat search.
+use std::num::NonZeroUsize;
+
+use diskann_utils::future::SendFuture;
+
+use crate::{
+    ANNResult,
+    error::{ErrorExt, IntoANNResult},
+    flat::{DistancesUnordered, SearchStrategy},
+    graph::SearchOutputBuffer,
+    neighbor::{Neighbor, NeighborPriorityQueue, NeighborPriorityQueueIdType},
+    provider::DataProvider,
+};
+
+/// Statistics collected during a flat search.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub struct SearchStats {
+    /// The total number of distance computations performed during the scan.
+    pub cmps: u32,
+
+    /// The total number of results written to the output buffer.
+    pub result_count: u32,
+}
+
+/// A thin wrapper around a [`DataProvider`] used for flat search.
+#[derive(Debug)]
+pub struct FlatIndex<P: DataProvider> {
+    /// The backing provider.
+    provider: P,
+}
+
+impl<P: DataProvider> FlatIndex<P> {
+    /// 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`], and
+    /// writes the `(id, distance)` survivors into `output` in best-first order.
+    pub fn knn_search<S, T, OB>(
+        &self,
+        k: NonZeroUsize,
+        strategy: &S,
+        context: &P::Context,
+        query: T,
+        output: &mut OB,
+    ) -> impl SendFuture<ANNResult<SearchStats>>
+    where
+        S: SearchStrategy<P, T>,
+        S::Id: NeighborPriorityQueueIdType,
+        T: Send + Sync,
+        OB: SearchOutputBuffer<S::Id> + 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 =
+                output.extend(queue.iter().take(k).map(|n| (n.id, n.distance))) as u32;
+
+            Ok(SearchStats { cmps, result_count })
+        }
+    }
+}
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use crate::flat::{
+        FlatIndex,
+        test::{
+            harness::KnnOracleRun,
+            provider::{self as flat_provider},
+        },
+    };
+    use crate::graph::test::synthetic::Grid;
+
+    fn fixture(grid: Grid, size: usize) -> (FlatIndex<flat_provider::Provider>, usize) {
+        let provider = flat_provider::Provider::grid(grid, size).unwrap();
+        let len = provider.len();
+        (FlatIndex::new(provider), len)
+    }
+
+    /// `knn_search` returns a `Send` future, and a shared `&FlatIndex` can serve
+    /// many concurrent searches on a multi-threaded runtime, each producing the
+    /// correct top-k 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);
+
+        // Mix of corner, axis-aligned, and off-grid queries; k spans 1..=len.
+        let cases: &[(&[f32], usize)] = &[
+            (&[-1.0, -1.0], 1),
+            (&[1.0, 1.0], len),
+            (&[-1.0, 1.0], len / 2),
+            (&[1.0, -1.0], len - 1),
+            (&[0.0, 0.0], 3),
+            (&[3.0, 3.0], len),
+            (&[-2.0, 0.5], 2),
+            (&[0.5, -0.5], len),
+        ];
+
+        let mut set = tokio::task::JoinSet::new();
+        for (query, k) in cases {
+            let index = Arc::clone(&index);
+            let query: Vec<f32> = query.to_vec();
+            let k = *k;
+            set.spawn(async move {
+                let outcome = KnnOracleRun::run(
+                    &index,
+                    &flat_provider::Strategy::new(index.provider().dim()),
+                    &query,
+                    k,
+                )
+                .await
+                .expect("knn_search failed");
+                (query, k, outcome)
+            });
+        }
+
+        while let Some(joined) = set.join_next().await {
+            let (query, k, outcome) = joined.expect("task panicked");
+            assert_eq!(
+                outcome.top_k, outcome.ground_truth,
+                "query = {query:?}, k = {k}: top-k must match brute force",
+            );
+            assert_eq!(outcome.stats.cmps as usize, len);
+            assert_eq!(outcome.stats.result_count as usize, k.min(len));
+        }
+    }
+
+    ////////////
+    // Errors //
+    ////////////
+
+    /// A transient error from the visitor's scan must escalate up through `knn_search`.
+    #[test]
+    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, &[1.0, 0.0], 4)
+                .expect_err("transient error during full scan must escalate");
+
+            let msg = format!("{err}");
+            assert!(
+                transient_ids
+                    .iter()
+                    .any(|id| msg.contains(&format!("id {id}"))),
+                "transients = {transient_ids:?}: expected error to name one of the \
+                 transient ids, got: {msg}",
+            );
+        }
+    }
+
+    /// 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, query, 4)
+            .expect_err("expected knn_search to fail");
+
+        let msg = format!("{err}");
+        assert!(
+            msg.contains(expected_msg),
+            "expected error containing {expected_msg:?}, got: {msg}",
+        );
+    }
+
+    #[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",
+        );
+    }
+}

diskann/src/flat/mod.rs

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! Sequential ("flat") search.
+//!
+//! 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 or
+//! on-disk shards streamed via I/O.
+//!
+//! # Architecture
+//!
+//! The module mirrors the layering used by graph search:
+//!
+//! | Graph (random access)                       | Flat (sequential)                          |  Shared?  |
+//! | :------------------------------------       | :----------------------------------------- |:--------- |
+//! | [`crate::provider::DataProvider`]           | [`crate::provider::DataProvider`]          | Yes       |
+//! | [`crate::graph::DiskANNIndex`]              | [`FlatIndex`]                              | No        |
+//! | [`crate::graph::glue::ExpandBeam`]          | [`DistancesUnordered`]                     | No        |
+//! | [`crate::graph::glue::SearchStrategy`]      | [`SearchStrategy`]                         | No        |
+//! | [`crate::graph::Search`]                    | [`FlatIndex::knn_search`]                  | No        |
+//!
+pub mod index;
+pub mod strategy;
+
+pub use index::{FlatIndex, SearchStats};
+pub use strategy::{DistancesUnordered, SearchStrategy};
+
+#[cfg(test)]
+mod test;