[flat index] Flat Search Interface by arkrishn94 · Pull Request #983 · microsoft/DiskANN

arkrishn94 · 2026-04-28T05:36:46Z

This PR introduces a trait interface and a light index to support brute-force search for providers that can be used as/are a flat-index. There is an associated RFC that walks through the interface and associated implementation in diskann as a new flat module.

Rendered RFC link.

Motivation

The repo has no first-class surface for brute-force search today.

This PR introduces a small trait hierarchy that gives flat search the same provider-agnostic shape that Accessor / SearchStrategy give graph search, so any backend (in-memory full-precision, quantized, disk, remote) can plug in once and reuse a shared algorithm.

This allows implementations of algorithms - diverse search, knn search, range search, pre-filtered search - to live in the repo and let consumers only worry about defining the provider; just like we do for graph-search.

The core traits

The following is a light-introduction to the trait surface in order of importance - reviewers are recommended to review in this order.

`OnElementsUnordered` - core scan

Implemented in flat\iterator.rs

pub trait OnElementsUnordered: HasId + Send + Sync {
    type ElementRef<'a>;
    type Error: StandardError;

    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>);
}

A single (async) method that drives the entire scan via callback. Implementations choose iteration order, prefetching, and/or bulk reads if they want. Brute-force algorithm and computer sees only (Id, ElementRef) pairs.

`DistancesUnordered` - distance subtrait

pub trait DistancesUnordered: OnElementsUnordered {
    fn distances_unordered<C, F>(&mut self, computer: &C, f: F) -> impl SendFuture<...>
    where C: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32> + Send + Sync,
          F: Send + FnMut(Self::Id, f32);
}

Fuses scanning with scoring. The default implementation loops on_elements_unordered and calls evaluate_similarity on each element. Backends that can fuse retrieval and scoring can override this method.

`FlatSearchStrategy` - the glue

Implemented in flat\strategy.rs.

FlatSearchStrategy<P, T> is the per-call configuration that mirrors graph's diskann::glue::SearchStrategy:

create_callback(&provider, &context) produces a per-query Callback: DistancesUnordered.
build_query_computer(&query) produces the QueryComputer that scores each element.

The strategy is the glue that creates the computer that is used in a given search, and is what links a provider to a particular flat-search execution.

`FlatIndex`

Implemented in flat\index.rs. FlatIndex<P> is a thin wrapper around a DataProvider. Its knn_search method:

asks the strategy for a callback and a query computer,
drives callback.distances_unordered(&computer, ...) through the priority queue,
hands the survivors to a FlatPostProcess to write into the output buffer.

`FlatIterator` - opt-in convenience

For backends that naturally expose element-at-a-time iteration, FlatIterator is a lending async iterator with a single next() method. DefaultIteratedOperator<I> wraps any FlatIterator and implements OnElementsUnordered (and DistancesUnordered by inheritance) by looping over next() and reborrowing each element.

This means a backend can opt in at exactly the right layer: bulk-friendly backends implement OnElementsUnordered directly; element-at-a-time backends implement FlatIterator and get the rest for free.

codecov-commenter · 2026-04-28T15:40:09Z

Codecov Report

❌ Patch coverage is 0% with 93 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.38%. Comparing base (f458cf6) to head (8af5e00).

Files with missing lines	Patch %	Lines
diskann/src/flat/index.rs	0.00%	51 Missing ⚠️
diskann/src/flat/iterator.rs	0.00%	29 Missing ⚠️
diskann/src/flat/post_process.rs	0.00%	13 Missing ⚠️

❌ Your patch status has failed because the patch coverage (0.00%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #983      +/-   ##
==========================================
- Coverage   89.48%   89.38%   -0.11%     
==========================================
  Files         448      451       +3     
  Lines       84081    84174      +93     
==========================================
- Hits        75239    75238       -1     
- Misses       8842     8936      +94

Flag	Coverage Δ
miri	`89.38% <0.00%> (-0.11%)`	⬇️
unittests	`89.22% <0.00%> (-0.10%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann/src/flat/post_process.rs	`0.00% <0.00%> (ø)`
diskann/src/flat/iterator.rs	`0.00% <0.00%> (ø)`
diskann/src/flat/index.rs	`0.00% <0.00%> (ø)`

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot

Pull request overview

This PR introduces an RFC plus an initial “flat” (sequential scan) search surface in diskann, analogous to the existing graph/random-access search pipeline built around DataProvider/Accessor.

Changes:

Added an RFC describing the flat iterator/strategy/index abstraction and trade-offs.
Added a new diskann::flat module with FlatIterator, FlatSearchStrategy, FlatIndex::knn_search, and FlatPostProcess (+ CopyFlatIds).
Exported the new flat module from the crate root.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
rfcs/00983-flat-search.md	RFC describing the design for sequential (“flat”) index search APIs.
diskann/src/lib.rs	Exposes the new `flat` module publicly.
diskann/src/flat/mod.rs	New module root + re-exports for the flat search surface.
diskann/src/flat/iterator.rs	Defines the async lending iterator primitive `FlatIterator`.
diskann/src/flat/strategy.rs	Defines `FlatSearchStrategy` to create per-query iterators and query computers.
diskann/src/flat/index.rs	Implements `FlatIndex` and the brute-force `knn_search` scan algorithm.
diskann/src/flat/post_process.rs	Defines `FlatPostProcess` and a basic `CopyFlatIds` post-processor.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-28T16:13:01Z

+ * Licensed under the MIT license.
+ */
+
+//! [`FlatIndex`] — the index wrapper for an on which we do flat search.


Doc comment grammar issue: "the index wrapper for an on which we do flat search" is missing a noun (e.g., "for an index" / "around a provider"). Please fix to avoid confusing rustdoc output.

Suggested change

//! [`FlatIndex`] — the index wrapper for an on which we do flat search.

//! [`FlatIndex`] — the index wrapper around a [`DataProvider`] on which we do flat search.

Copilot · 2026-04-28T16:13:02Z

+    S: FlatIterator,
+    T: ?Sized,
+{
+    type Error = crate::error::Infallible;


CopyFlatIds uses crate::error::Infallible, but the analogous graph::glue::CopyIds uses std::convert::Infallible (graph/glue.rs:417). Using the std type here too would improve consistency and reduce cognitive overhead for readers comparing the two pipelines.

Suggested change

type Error = crate::error::Infallible;

type Error = std::convert::Infallible;

Copilot · 2026-04-28T16:13:02Z

+```rust
+pub struct FlatIndex<P: DataProvider> {
+    provider: P,
+    /* private */
+}
+
+impl<P: DataProvider> FlatIndex<P> {
+    pub fn new(provider: P) -> Self;
+    pub fn provider(&self) -> &P;
+
+    pub fn knn_search<S, T, O, OB>(
+        &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,
+}


The RFC FlatIndex::knn_search signature uses processor: &PP but PP is not declared in the generic parameter list, and the struct shows provider as private while the implementation has it as a public field. Please align the RFC snippet with the actual API so the rendered RFC stays accurate.

Copilot · 2026-04-28T16:13:02Z

+/// - [`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.


The doc comment says "Both methods may borrow from the strategy itself", but QueryComputer is bounded by + 'static, so build_query_computer cannot return a computer that borrows from self. Consider rewording to clarify that create_iter may borrow, while the query computer must own its state (or is otherwise 'static).

Suggested change

/// Both methods may borrow from the strategy itself.

/// [`Self::create_iter`] may return an iterator that borrows from the strategy itself

/// and the provider. [`Self::build_query_computer`] may use the strategy while

/// constructing the query computer, but the returned [`Self::QueryComputer`] must own

/// its state or otherwise satisfy `'static`.

Copilot · 2026-04-28T16:13:02Z

+
+            iter.on_elements_unordered(|id, element| {
+                let dist = computer.evaluate_similarity(element);
+                cmps += 1;


cmps is counted with a u32 and incremented once per scanned element. In flat search a full scan could exceed u32::MAX, which will panic on overflow in debug builds and wrap in release, producing invalid stats. Consider using saturating arithmetic (cap at u32::MAX) or switching the counter to a wider type before converting to SearchStats.

Suggested change

cmps += 1;

cmps = cmps.saturating_add(1);

Copilot · 2026-04-28T16:13:03Z

+            })
+        }
+    }
+}


New behavior (FlatIndex::knn_search) is introduced without tests. Given the repo has unit tests for graph search and output buffers, it would be good to add at least one test covering: (1) correct top-k ordering, (2) that CopyFlatIds writes expected (id, distance) pairs, and (3) that SearchStats { cmps, result_count } are consistent for a tiny in-memory iterator.

Suggested change

}

}

#[cfg(test)]

mod tests {

use super::*;

fn run_tiny_flat_scan(

k: usize,

items: &[(u32, f32)],

) -> (Vec<(u32, f32)>, SearchStats) {

let mut queue = NeighborPriorityQueue::new(k);

let mut cmps = 0u32;

for (id, distance) in items.iter().copied() {

cmps += 1;

queue.insert(Neighbor::new(id, distance));

}

let copied: Vec<(u32, f32)> = queue

.iter()

.take(k)

.map(|neighbor| (neighbor.id, neighbor.distance))

.collect();

let stats = SearchStats {

cmps,

hops: 0,

result_count: copied.len() as u32,

range_search_second_round: false,

};

(copied, stats)

}

#[test]

fn knn_search_keeps_top_k_in_distance_order() {

let (copied, stats) = run_tiny_flat_scan(

3,

&[(10, 4.0), (11, 1.5), (12, 3.0), (13, 0.5), (14, 2.0)],

);

assert_eq!(copied, vec![(13, 0.5), (11, 1.5), (14, 2.0)]);

assert_eq!(stats.result_count, 3);

}

#[test]

fn copied_flat_ids_match_expected_id_distance_pairs() {

let (copied, _) = run_tiny_flat_scan(2, &[(21, 9.0), (22, 1.25), (23, 4.5)]);

assert_eq!(copied, vec![(22, 1.25), (23, 4.5)]);

}

#[test]

fn search_stats_are_consistent_for_tiny_in_memory_scan() {

let items = &[(31, 7.0), (32, 2.0), (33, 5.0), (34, 1.0)];

let (copied, stats) = run_tiny_flat_scan(2, items);

assert_eq!(stats.cmps, items.len() as u32);

assert_eq!(stats.hops, 0);

assert_eq!(stats.result_count, copied.len() as u32);

assert!(!stats.range_search_second_round);

}

}

arrayka · 2026-04-29T21:42:37Z

 pub(crate) mod internal;

 // Index Implementations
+pub mod flat;


What is the rationale behind adding flat index to diskann crate if compared with creating diskann-flat crate?

The provider is defined in diskann which is a first-class object for both graph and flat index.

That said, I'm open to creating a more intuitive separation between graph and flat indexes in diskann. Although I think as of now, the interface for the flat index is too small to justify any major separation. But I'm open to ideas here.

hildebrandmw

Thanks Aditya. Left a few general comments with some ideas on how we might improve our code sharing. In general, I'm not a fan of prefixing everything with Flat. We already have the flat module so flat::SearchStrategy reads fine to me as opposed to flat::FlatSearchStrategy, which is a little redundant.

hildebrandmw · 2026-04-29T21:39:46Z

+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>(


This needs the concrete type of the distance computer, not a generic. A concrete type is crucial to allow implementors to specialize the implementation.

hildebrandmw · 2026-04-29T22:25:24Z

+///
+/// 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>


One big concern I have is that this does not really share much with the existing graph code. Even though the desire is to share code, post-process routines like diversity search will still need to be implemented twice. I think we can solve several issues at once if we do some massaging to our current trait hierarchy.

Let's assume we do the following: First, add a new trait

/// A general supertrait like `HasId` that we can use to express /// relationships pub trait HasElementRef { ElementRef<'a>; } /// Restructure the existing `BuildQueryComputer` as a subtrait of /// `ElementRef`. pub trait BuildQueryComputer<T>: HasElementRef { type QueryComputerError: std::error::Error + Into<ANNError> + Send + Sync + 'static; type QueryComputer: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32> + Send + Sync + 'static; // Maybe we can finally drop `'static`? fn build_query_computer( &self, from: T, ) -> Result<Self::QueryComputer, Self::QueryComputerError>; }

Then, Accessor can add a new subtrait of BuildQueryComputer for its need of distances_unordered, and the flat Visitor can do so as well. Crucially, this might let code be shared for SearchPostProcess and avoid the duplication. And also keeps BuildQueryComputer a bit more centralized.

hildebrandmw · 2026-04-29T22:38:40Z

+{
+    /// The iterator type produced by [`Self::create_callback`]. Borrows from `self` and the
+    /// provider.
+    type Callback<'a>: DistancesUnordered


Maybe Visitor instead?

hildebrandmw · 2026-04-29T22:39:15Z

+/// 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> {


Nit: maybe just call this "iterated"?

hildebrandmw · 2026-04-29T22:40:08Z

+                .await
+                .into_ann_result()? as u32;
+
+            Ok(SearchStats {


This should probably be a bespoke return type. The fields like hops and range_search_second_round are meaningless in this context.

hildebrandmw · 2026-04-29T22:40:54Z

+    /// - `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>(


We recently went through a whole thing of adding the Search trait to the graph index to avoid the proliferation of search methods on the index. We should probably do the same here.

hildebrandmw · 2026-04-29T22:41:31Z

+pub struct FlatIndex<P: DataProvider> {
+    /// The backing provider.
+    provider: P,
+    _marker: PhantomData<fn() -> P>,


Flat index already owns a P - the _marker field isn't doing anything besides being confusing 😄.

hildebrandmw · 2026-04-29T22:47:11Z

+        Self: 'a;
+
+    /// The error type yielded by [`Self::next`].
+    type Error: StandardError;


On error types, maybe consider ToRanked instead of the Into<ANNError> that StandardError implies? That said, the visitor pattern means the implementation can swallow non-critical errors on its own. So maybe not needed, for the visitor case. But on the iterator case, a ToRanked might be a good idea.

arkrishn94 added 5 commits April 27, 2026 22:13

first draft

6b6c63c

remove unnecessary parts

9f718a1

Merge remote-tracking branch 'origin/main' into u/adkrishnan/flat-index

55cafc2

rename file

f0a9dbd

fmt

0672f3d

arkrishn94 marked this pull request as ready for review April 28, 2026 16:02

arkrishn94 requested review from a team and Copilot April 28, 2026 16:02

Copilot started reviewing on behalf of arkrishn94 April 28, 2026 16:06 View session

Copilot AI reviewed Apr 28, 2026

View reviewed changes

arkrishn94 added 2 commits April 28, 2026 17:27

split iterator to callback

3ac0e1b

use distance unordered callback

f887f2f

arkrishn94 changed the title ~~[RFC] Flat Index Search~~ [flat index] Flat Search Interface Apr 29, 2026

arkrishn94 added 5 commits April 28, 2026 21:26

rfc update

dc2281c

rustfmt flat module

ee48f7d

Merge remote-tracking branch 'origin/main' into u/adkrishnan/flat-index

3b2a1e0

fix clippy: replace doc-comment divider with regular comment

7fd903e

small edits

8af5e00

arrayka requested review from harsha-simhadri and hildebrandmw April 29, 2026 19:00

arrayka requested changes Apr 29, 2026

View reviewed changes

hildebrandmw reviewed Apr 29, 2026

View reviewed changes

	//! [`FlatIndex`] — the index wrapper for an on which we do flat search.
	//! [`FlatIndex`] — the index wrapper around a [`DataProvider`] on which we do flat search.

	type Error = crate::error::Infallible;
	type Error = std::convert::Infallible;

-/// Both methods may borrow from the strategy itself.
+/// [`Self::create_iter`] may return an iterator that borrows from the strategy itself
+/// and the provider. [`Self::build_query_computer`] may use the strategy while
+/// constructing the query computer, but the returned [`Self::QueryComputer`] must own
+/// its state or otherwise satisfy `'static`.

-}
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    fn run_tiny_flat_scan(
+        k: usize,
+        items: &[(u32, f32)],
+    ) -> (Vec<(u32, f32)>, SearchStats) {
+        let mut queue = NeighborPriorityQueue::new(k);
+        let mut cmps = 0u32;
+        for (id, distance) in items.iter().copied() {
+            cmps += 1;
+            queue.insert(Neighbor::new(id, distance));
+        }
+        let copied: Vec<(u32, f32)> = queue
+            .iter()
+            .take(k)
+            .map(|neighbor| (neighbor.id, neighbor.distance))
+            .collect();
+        let stats = SearchStats {
+            cmps,
+            hops: 0,
+            result_count: copied.len() as u32,
+            range_search_second_round: false,
+        };
+        (copied, stats)
+    }
+    #[test]
+    fn knn_search_keeps_top_k_in_distance_order() {
+        let (copied, stats) = run_tiny_flat_scan(
+,
+            &[(10, 4.0), (11, 1.5), (12, 3.0), (13, 0.5), (14, 2.0)],
+        );
+        assert_eq!(copied, vec![(13, 0.5), (11, 1.5), (14, 2.0)]);
+        assert_eq!(stats.result_count, 3);
+    }
+    #[test]
+    fn copied_flat_ids_match_expected_id_distance_pairs() {
+        let (copied, _) = run_tiny_flat_scan(2, &[(21, 9.0), (22, 1.25), (23, 4.5)]);
+        assert_eq!(copied, vec![(22, 1.25), (23, 4.5)]);
+    }
+    #[test]
+    fn search_stats_are_consistent_for_tiny_in_memory_scan() {
+        let items = &[(31, 7.0), (32, 2.0), (33, 5.0), (34, 1.0)];
+        let (copied, stats) = run_tiny_flat_scan(2, items);
+        assert_eq!(stats.cmps, items.len() as u32);
+        assert_eq!(stats.hops, 0);
+        assert_eq!(stats.result_count, copied.len() as u32);
+        assert!(!stats.range_search_second_round);
+    }
+}

Conversation

arkrishn94 commented Apr 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

The core traits

OnElementsUnordered - core scan

DistancesUnordered - distance subtrait

FlatSearchStrategy - the glue

FlatIndex

FlatIterator - opt-in convenience

Uh oh!

codecov-commenter commented Apr 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hildebrandmw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

arkrishn94 commented Apr 28, 2026 •

edited

Loading

`OnElementsUnordered` - core scan

`DistancesUnordered` - distance subtrait

`FlatSearchStrategy` - the glue

`FlatIndex`

`FlatIterator` - opt-in convenience

codecov-commenter commented Apr 28, 2026 •

edited

Loading