[flat index] Flat Search Interface

diskann/src/flat/index.rs

@@ -0,0 +1,109 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! [`FlatIndex`] — the index wrapper for an on which we do 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, FlatSearchStrategy},
+    graph::{SearchOutputBuffer, index::SearchStats},
+    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<P: DataProvider> {
+    /// The backing provider.
+    pub provider: P,
+    _marker: PhantomData<fn() -> P>,
+}
+
+impl<P: DataProvider> FlatIndex<P> {
+    /// 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. See [`FlatSearchStrategy`]
+    /// - `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<S, T, O, OB, PP>(
+        &self,
+        k: NonZeroUsize,
+        strategy: &S,
+        processor: &PP,
+        context: &P::Context,
+        query: &T,
+        output: &mut OB,
+    ) -> impl SendFuture<ANNResult<SearchStats>>
+    where
+        S: FlatSearchStrategy<P, T>,
+        T: ?Sized + Sync,
+        O: Send,
+        OB: SearchOutputBuffer<O> + Send + ?Sized,
+        PP: for<'a> FlatPostProcess<S::Iter<'a>, 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(SearchStats {
+                cmps,
+                hops: 0,
+                result_count,
+                range_search_second_round: false,
+            })
+        }
+    }
+}

diskann/src/flat/iterator.rs

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! [`FlatIterator`] — the sequential access primitive for accessing a flat index.
+
+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.
+///
+/// Algorithms see only `(Id, ElementRef)` pairs and treat the stream as opaque.
+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`.
+    #[allow(clippy::type_complexity)]
+    fn next(
+        &mut self,
+    ) -> impl SendFuture<Result<Option<(Self::Id, Self::Element<'_>)>, Self::Error>>;
+
+    /// Drive the entire scan, invoking `f` for each yielded element.
+    ///
+    /// The default implementation loops over [`Self::next`].
+    fn on_elements_unordered<F>(&mut self, mut f: F) -> impl SendFuture<Result<(), Self::Error>>
+    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(())
+        }
+    }
+}

diskann/src/flat/mod.rs

@@ -0,0 +1,43 @@
+/*
+ * 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 strategy;
+
+pub use index::FlatIndex;
+pub use iterator::FlatIterator;
+pub use post_process::{CopyFlatIds, FlatPostProcess};
+pub use strategy::FlatSearchStrategy;

diskann/src/flat/post_process.rs

@@ -0,0 +1,72 @@
+/*
+ * 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,
+};
+
+/// 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.
+///
+/// The `O` type parameter lets callers pick the output element type (raw `(Id, f32)`
+/// pairs, fully hydrated hits etc.).
+pub trait FlatPostProcess<S, T, O = <S as HasId>::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<I, B>(
+        &self,
+        iter: &mut S,
+        query: &T,
+        candidates: I,
+        output: &mut B,
+    ) -> impl SendFuture<Result<usize, Self::Error>>
+    where
+        I: Iterator<Item = Neighbor<S::Id>> + Send,
+        B: SearchOutputBuffer<O> + 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<S, T> FlatPostProcess<S, T> for CopyFlatIds
+where
+    S: FlatIterator,
+    T: ?Sized,
+{
+    type Error = crate::error::Infallible;
+
+    fn post_process<I, B>(
+        &self,
+        _iter: &mut S,
+        _query: &T,
+        candidates: I,
+        output: &mut B,
+    ) -> impl SendFuture<Result<usize, Self::Error>>
+    where
+        I: Iterator<Item = Neighbor<<S as HasId>::Id>> + Send,
+        B: SearchOutputBuffer<<S as HasId>::Id> + Send + ?Sized,
+    {
+        let count = output.extend(candidates.map(|n| (n.id, n.distance)));
+        std::future::ready(Ok(count))
+    }
+}

diskann/src/flat/strategy.rs

@@ -0,0 +1,73 @@
+/*
+ * 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<P, T>: 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<
+            <Self::Iter<'a> 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::Iter<'a>, 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<Self::QueryComputer, Self::Error>;
+}

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.