[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 crate::{
+    ANNResult,
+    error::IntoANNResult,
+    flat::{DistancesUnordered, 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::OnElementsUnordered`]
+/// implementation that the [`crate::flat::FlatSearchStrategy`] produces.
+#[derive(Debug)]
+pub struct FlatIndex<P: DataProvider> {
+    /// The backing provider.
+    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::Callback<'a>, T, O> + Send + Sync,
+    {
+        async move {
+            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;
+
+            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)
+                .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,147 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! [`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 diskann_vector::PreprocessedDistanceFunction;
+
+use crate::{error::StandardError, provider::HasId};
+
+/// Callback-driven sequential scan over the elements of a flat index.
+///
+/// `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<F>(&mut self, f: F) -> impl SendFuture<Result<(), Self::Error>>
+    where
+        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<C, F>(
+        &mut self,
+        computer: &C,
+        mut f: F,
+    ) -> impl SendFuture<Result<(), Self::Error>>
+    where
+        C: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, 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);
+        })
+    }
+}
+
+//////////////
+// Iterator //
+//////////////
+
+/// 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.
+    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`].
+    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>>;
+}
+
+/////////////
+// 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<I> {
+    inner: I,
+}
+
+impl<I> DefaultIteratedOperator<I> {
+    /// 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<I: HasId> HasId for DefaultIteratedOperator<I> {
+    type Id = I::Id;
+}
+
+impl<I> OnElementsUnordered for DefaultIteratedOperator<I>
+where
+    I: FlatIterator + HasId + Send + Sync,
+{
+    type ElementRef<'a> = I::ElementRef<'a>;
+    type Error = I::Error;
+
+    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.inner.next().await? {
+                f(id, element.reborrow());
+            }
+            Ok(())
+        }
+    }
+}
+
+impl<I> DistancesUnordered for DefaultIteratedOperator<I> where I: FlatIterator + HasId + Send + Sync
+{}

diskann/src/flat/mod.rs

@@ -0,0 +1,45 @@
+/*
+ * 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::{
+    DefaultIteratedOperator, DistancesUnordered, FlatIterator, OnElementsUnordered,
+};
+pub use post_process::{CopyIds, 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::OnElementsUnordered, 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: OnElementsUnordered,
+    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 CopyIds;
+
+impl<S, T> FlatPostProcess<S, T> for CopyIds
+where
+    S: OnElementsUnordered,
+    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))
+    }
+}