[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 the data is accessed / provided; just like we do for graph search.

Important refactor

We start with an important refactor of BuildQueryComputer and its associated ElementRef<'a> that it acts on in diskann::providers.

`HasElementRef` - a new minimal shape trait

This zero-method traits was extracted so the streaming flat traits and the random-access Accessor can share their associated types without one inheriting from the other:

pub trait HasElementRef { type ElementRef<'a>; }

Accessor: HasId + HasElementRef, and the flat traits below depend on the same two pieces independently.

`BuildQueryComputer<T>` - shared query preprocessing

We split the trait in diskann::provider to into:

a core trait that has the build_query_computer method and the QueryComputer associated type, and,
a sub-trait titled DistancesUnordered that has the distances_unordered method.

This split allows both graph and flat indexes to require building of a computer without dragging in random access.

pub trait BuildQueryComputer<T>: HasElementRef {
    type QueryComputer: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32> + ...;
    type Error: StandardError;
    fn build_query_computer(&self, query: T) -> Result<Self::QueryComputer, Self::Error>;
}

Flat search - core components

The following trait and its subtrait are the core traits that define how a brute-force scan over the index is implemented.

`OnElementsUnordered` — the core scan

Implemented in flat/iterator.rs, this is a single (async) method that drives the entire scan via callback. Implementations choose iteration order, prefetching, and bulk reads. Algorithms only see (Id, ElementRef) pairs.

pub trait OnElementsUnordered: HasId + HasElementRef + Send + Sync {
    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>);
}

`flat::DistancesUnordered<T>` — fused scan + score

This subtrait of OnElementsUnordered fuses scanning with scoring - it takes in a QueryComputer in addition to the closure. The default body loops on_elements_unordered and calls evaluate_similarity on each element. Backends that can fuse retrieval with scoring can override it.

The computer type comes from the implementor's own BuildQueryComputer<T> impl, so the visitor produces the computer that drives its own scan.

`flat::SearchStrategy<P, T>` — the glue

Implemented in flat/strategy.rs. Mirrors graph::glue::SearchStrategy (disambiguated by module path) and simply creates a concrete type that implements the scan:

pub trait SearchStrategy<P, T>: Send + Sync where P: DataProvider {
    type Visitor<'a>: DistancesUnordered<T> where Self: 'a, P: 'a;
    type Error: StandardError;

    fn create_visitor<'a>(
        &'a self,
        provider: &'a P,
        context: &'a P::Context,
    ) -> Result<Self::Visitor<'a>, Self::Error>;
}

Strategies are stateless per-call config — constructed at the call site, used for one search and then dropped.

`FlatIndex<P>` — the top-level handle

Implemented in flat/index.rs, this is a thin 'static wrapper around a DataProvider. The search is implemented as:

asks the strategy for a per-query visitor,
builds the query computer from the visitor (BuildQueryComputer::build_query_computer),
drives visitor.distances_unordered(&computer, ...) through a NeighborPriorityQueue,
hands the survivors to a graph::glue::SearchPostProcess to write into the output buffer.

Note the SearchPostProcess bound: the same trait used by graph search!

`FlatIterator` + `Iterated<I>` — opt-in convenience

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

A backend opts in at exactly the right layer: bulk-friendly backends implement OnElementsUnordered directly; element-at-a-time backends implement FlatIterator and use Iterated for the rest.

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

Codecov Report

❌ Patch coverage is 89.00763% with 72 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.47%. Comparing base (95b6593) to head (f338a83).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann/src/flat/test/provider.rs	67.06%	55 Missing ⚠️
diskann/src/flat/iterator.rs	95.63%	10 Missing ⚠️
diskann/src/flat/test/cases/flat_knn_search.rs	93.97%	5 Missing ⚠️
diskann/src/flat/index.rs	98.00%	2 Missing ⚠️

❌ Your patch status has failed because the patch coverage (89.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.51%   89.47%   -0.05%     
==========================================
  Files         460      467       +7     
  Lines       85424    86145     +721     
==========================================
+ Hits        76469    77074     +605     
- Misses       8955     9071     +116

Flag	Coverage Δ
miri	`89.47% <89.00%> (-0.05%)`	⬇️
unittests	`89.31% <89.00%> (-0.05%)`	⬇️

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

Files with missing lines	Coverage Δ
diskann-disk/src/search/provider/disk_provider.rs	`90.89% <ø> (ø)`
diskann-garnet/src/provider.rs	`83.36% <ø> (ø)`
...rc/inline_beta_search/encoded_document_accessor.rs	`0.00% <ø> (ø)`
...ilter/src/inline_beta_search/inline_beta_filter.rs	`0.00% <ø> (ø)`
...odel/graph/provider/async_/inmem/full_precision.rs	`98.39% <ø> (ø)`
...s/src/model/graph/provider/async_/inmem/product.rs	`90.17% <ø> (ø)`
...rs/src/model/graph/provider/async_/inmem/scalar.rs	`92.78% <ø> (ø)`
...src/model/graph/provider/async_/inmem/spherical.rs	`88.13% <ø> (ø)`
...ders/src/model/graph/provider/async_/inmem/test.rs	`69.34% <ø> (ø)`
...ers/src/model/graph/provider/async_/postprocess.rs	`100.00% <ø> (ø)`
... and 10 more

... and 18 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.

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-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.

Ack, makes sense.

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.

harsha-simhadri · 2026-05-05T16:42:24Z

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


This is a nice description comparing traits that enable algorithms based on random access vs sequential scans. Could this be in a higher level directory, either in providers.rs file or in diskann/src/agents.md

harsha-simhadri · 2026-05-05T20:37:10Z

+
+
+### The glue: `FlatSearchStrategy`
+


Since we are introducing substantial new machinery, can we think about whether we can reuse some of this for IVF/SPANN type of indices that rely on clustering and then scanning entire data in specific clusters to support queries.

I can imagine that the OnElementsUnordered and DistancesUnordered could be adapted to the scope of a cluster.

And SearchStrategies for IVF would need to be added.

Even if we dont have a fully fleshed our proposal for clustering based indices, it would be ideal to document how the abstractions can be reused or adapted in the near future and avoid another set of abstractions for clustering based indices

harsha-simhadri · 2026-05-05T20:37:40Z

+5. Return search stats.
+
+Other algorithms (filtered, range, diverse) can be added later as additional methods on
+`FlatIndex`.


How would we support predicates on flat index?

it's a good point - type-1 bitmap filters can be supported by extending the scan trait (DistancesUnordered) with a typed predicate applied before computing distances. This can then be wired through at the top-level search API depending on whether we're doing brute-force or filtered search.

harsha-simhadri · 2026-05-05T20:53:31Z

+//! | [`crate::provider::Accessor`]             | [`FlatIterator`]                  |
+//! | [`crate::graph::glue::SearchStrategy`]    | [`FlatSearchStrategy`]            |
+//! | [`crate::graph::glue::SearchPostProcess`] | [`FlatPostProcess`]               |
+//! | [`crate::graph::Search`]                  | [`FlatIndex::knn_search`]         |


This table is very useful.
Could this description be upleveled to diskann/src/agents.md or readme.

harsha-simhadri · 2026-05-05T20:54:00Z

+//! # Hot loop
+//!
+//! Algorithms drive the scan via [`FlatIterator::next`] (lending iterator) or override
+//! [`FlatIterator::on_elements_unordered`] when batching/prefetching wins. The default


Is this paragraph supposed to be here. Seems a bit out of place.

Agreed, I need to clean up.

harsha-simhadri · 2026-05-05T20:55:16Z

+    provider::HasId,
+};
+
+/// Post-process the survivor candidates produced by a flat search and


Is the intention to support filters with post_process? If so where is the clause?

Post-processing for flat is now common with diskann::graph::glue::SearchPostProcess. So filter will have to live in the object implementing this trait like in the graph case - e.g. inline_beta_filter. Or it can live in the Visitor<'_> itself.

# Conflicts: # diskann/src/graph/glue.rs

arrayka · 2026-05-05T21:59:05Z

+//!
+//! The module mirrors the layering used by graph search:
+//!
+//! | Graph (random access)                     | Flat (sequential)                 |  Shared?  |


Consider adding Responsibility column and provide one-sentence description for each layer.

arrayka · 2026-05-05T22:14:11Z

+pub trait DistancesUnordered<T>: OnElementsUnordered + BuildQueryComputer<T> {
+    /// Drive the entire scan, scoring each element with `computer` and invoking `f` with
+    /// the resulting `(id, distance)` pair.
+    fn distances_unordered<F>(


The algorithm only needs DistancesUnordered. Making it a supertrait of OnElementsUnordered forces all distance-capable providers to also expose raw element streaming, which leaks a lower-level capability and may block implementations that can compute distances without exposing element refs.

Could we decouple the traits and provide a blanket impl of DistancesUnordered for OnElementsUnordered + BuildQueryComputer instead:

Make DistancesUnordered independent, and provide a separate adapter/blanket impl when OnElementsUnordered is available. Conceptually:

trait DistancesUnordered<T> { fn distances_unordered(...) ... } (no supertrait)

impl<T, S> DistancesUnordered<T> for S where S: OnElementsUnordered + BuildQueryComputer<T> { ... }

That keeps the algorithm-facing surface minimal, preserves encapsulation, and still keeps the convenience default behavior for types that do implement OnElementsUnordered.

I'm a little confused; how would we define DistancedUnordered without the BuildQueryComputer trait bound?Specifically, the computer signature in the definition of distances_unordered. If possible, can you flesh this out a bit?

On your suggestion, if we implement the default implementation on any S : OnElementsUnordered + BuildQueryComputer then we won't allow a consumer to specialize the implementation for their S no?

I'm not sure I'm following why we might want to have implementations of DistancesUnordered that don't implement OnElementsUnordered :)

arrayka · 2026-05-05T22:31:36Z

+                .into_ann_result()?;
+
+            let computer =
+                BuildQueryComputer::build_query_computer(&visitor, query).into_ann_result()?;


Could you please provide an example of why visitor is needed to build a query computer?

Couple reasons -

visitor implements DistancesUnordered so it is not unreasonable it should know how to construct the correct type to apply the streaming distance computation.

Making the visitor implement DistancesUnordered has the advantage of being symmetric to the graph case: where the search_accessor implements ExpandBeam which must implement BuildQueryComputer<T>. This symmetry helps a bit with sharing the post-process trait SearchPostProcess for both graph and flat search - since trait bounds are identical for both paths.

arrayka · 2026-05-05T22:34:15Z

    type Id = u32;
 }

+impl provider::HasElementRef for Accessor<'_> {


As an option, consider extracting this prerequisite change into a separate PR to keep this one smaller and easier to review.

Yeah, I'm open to separating this change to a pre-cursor PR.

My only worry is that, since this refactor is specific to enabling the new flat search trait structure I'm proposing here, I don't want to rush the refactor in an earlier PR only to change the flat search trait architecture.

arrayka · 2026-05-08T19:17:00Z

+/// brute-force ground-truth oracle.
+#[derive(Debug, Clone)]
+pub(crate) struct KnnOracleRun {
+    /// Top-`k` `(id, distance)` pairs in canonical `(distance asc, id asc)` order.


Minor: shouldn't id be in the first position: (distance asc, id asc).

arrayka · 2026-05-08T19:21:07Z

+    /// 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 knn_search() {


I would name the test something like this: flat_index_supports_multi_threading

arrayka · 2026-05-08T19:31:33Z

+        self.items.is_empty()
+    }
+
+    /// Distance metric the provider was constructed with.


Why does data provider (which responsibility is to provide data vectors) care about the distance metric?
Why is data provider defining the allowed distance metric?

arrayka · 2026-05-08T19:32:51Z

+    }
+
+    /// Snapshot of the per-provider counters.
+    pub fn metrics(&self) -> Metrics {


Minor: Metrics is a little bit confusing with Metric. Consider renaming Metrics to Counters to avoid confusion.

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

Comment thread diskann/src/flat/index.rs

Comment thread diskann/src/flat/iterator.rs Outdated

Comment thread diskann/src/lib.rs

hildebrandmw reviewed Apr 29, 2026

View reviewed changes

arkrishn94 added 2 commits May 4, 2026 18:09

renames and uplevel query computer

1dd1c72

split buildquerycomputer, haselementref and distancesunordered

3d24ef7

harsha-simhadri reviewed May 5, 2026

View reviewed changes

delete flatpostprocess, cleanup and docs

0a63759

arrayka linked an issue May 5, 2026 that may be closed by this pull request

Create RFC and initial implementation for flat-index search support #999

Open

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

1f48945

# Conflicts: # diskann/src/graph/glue.rs

arrayka reviewed May 5, 2026

View reviewed changes

arkrishn94 added 6 commits May 5, 2026 23:21

update rfc

57d86e9

error types

715150a

small doc fixes

0627cc0

iterator docs

1d4e90b

strong pass on testing

b4d9df0

minor cleanup docs

f338a83

arrayka approved these changes May 8, 2026

View reviewed changes

Conversation

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

Uh oh!

Motivation

Important refactor

HasElementRef - a new minimal shape trait

BuildQueryComputer<T> - shared query preprocessing

Flat search - core components

OnElementsUnordered — the core scan

flat::DistancesUnordered<T> — fused scan + score

flat::SearchStrategy<P, T> — the glue

FlatIndex<P> — the top-level handle

FlatIterator + Iterated<I> — 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hildebrandmw left a comment

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!

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!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

arkrishn94 commented Apr 28, 2026 •

edited

Loading

`HasElementRef` - a new minimal shape trait

`BuildQueryComputer<T>` - shared query preprocessing

`OnElementsUnordered` — the core scan

`flat::DistancesUnordered<T>` — fused scan + score

`flat::SearchStrategy<P, T>` — the glue

`FlatIndex<P>` — the top-level handle

`FlatIterator` + `Iterated<I>` — opt-in convenience

codecov-commenter commented Apr 28, 2026 •

edited

Loading

arkrishn94 May 6, 2026 •

edited

Loading