[flat index] Flat Search Interface

[arkrishn94]

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. This PR adds a small trait hierarchy that gives flat search the same provider-agnostic shape that graph search has, so any backend (in-memory, quantized, disk, remote) can plug in once and reuse a shared algorithm.

What's in this PR

Traits (flat/strategy.rs + flat/iterator.rs)

DistancesUnordered<C> — the single trait a backend must implement. Fuses iteration and scoring into one method: the implementation drives a full scan, scoring each element with a precomputed query computer C, and invokes a callback with (id, distance) pairs. Key associated types:

• ElementRef<'a> -- the reference shape C scores against.
• C : for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32> -- the precomputer query computer.

pub trait DistancesUnordered<C>: HasId + Send + Sync
where
    C: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32>,
{
    type ElementRef<'a>;
    type Error: ToRanked + Debug + Send + Sync + 'static;

    fn distances_unordered<F>(
        &mut self,
        computer: &C,
        f: F,
    ) -> impl SendFuture<Result<(), Self::Error>>
    where
        F: Send + FnMut(<Self as HasId>::Id, f32);
}

SearchStrategy<P, T> — factory that creates a DistancesUnordered visitor from a provider + context, and builds the per-query computer. Mirrors the graph-side strategy pattern. Two fallible methods:

• create_visitor — borrows provider + context, returns a Visitor
• build_query_computer — preprocesses the query T into a QueryComputer

pub trait SearchStrategy<P, T>: Send + Sync
where
    P: DataProvider,
{
    type ElementRef<'a>;
    type QueryComputer: for<'a> PreprocessedDistanceFunction<Self::ElementRef<'a>, f32>;
    type QueryComputerError: StandardError;
    type Visitor<'a>: DistancesUnordered<Self::QueryComputer, Id = P::InternalId>;
    type Error: StandardError;

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

    fn build_query_computer(&self, query: T)
        -> Result<Self::QueryComputer, Self::QueryComputerError>;
}

FlatIterator + Iterated<I> — opt-in lending iterator for element-at-a-time backends; Iterated blanket-impls DistancesUnordered by looping over next() and scoring each element.

Index (flat/index.rs)

FlatIndex<P> — thin 'static wrapper around a DataProvider. Currently we have implemented the naive kNN search algorithm for the flat index. knn_search asks the strategy for a visitor, builds the query computer, drives distances_unordered through a priority queue, and writes results via SearchPostProcess.

Test infrastructure (flat/test/)

A self-contained test provider with dimension-validated Strategy, transient-error injection, and a KnnOracleRun harness that compares knn_search results against a brute-force reference with baseline caching for regression detection.

Future work

• Post-processing — knn_search currently uses the graph-side SearchPostProcess trait to write results into the output buffer. Simplify the DataProvider contract for graph search #1067 will introduce a flat-specific post-processing step that decouples flat search from the graph module's output machinery.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.69401% with 59 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.46%. Comparing base (d00cac0) to head (8f3b383).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann/src/flat/test/provider.rs	73.14%	47 Missing ⚠️
diskann/src/flat/test/cases/flat_knn_search.rs	94.11%	5 Missing ⚠️
diskann/src/flat/iterator.rs	97.87%	4 Missing ⚠️
diskann/src/flat/index.rs	97.45%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #983      +/-   ##
==========================================
- Coverage   89.51%   89.46%   -0.05%     
==========================================
  Files         461      464       +3     
  Lines       85920    86116     +196     
==========================================
+ Hits        76914    77047     +133     
- Misses       9006     9069      +63     

Flag	Coverage Δ
miri	89.46% <90.69%> (-0.05%)	⬇️
unittests	89.09% <90.69%> (-0.05%)	⬇️

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

Files with missing lines	Coverage Δ
diskann/src/flat/test/harness.rs	100.00% <100.00%> (ø)
diskann/src/flat/index.rs	97.45% <97.45%> (ø)
diskann/src/flat/iterator.rs	97.87% <97.87%> (ø)
diskann/src/flat/test/cases/flat_knn_search.rs	94.11% <94.11%> (ø)
diskann/src/flat/test/provider.rs	73.14% <73.14%> (ø)

... and 55 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]

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.

[arkrishn94]

Ack, makes sense.

[harsha-simhadri]

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]

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]

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

[arrayka]

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

[arrayka]

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.

[arkrishn94]

• 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]

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

[arkrishn94]

Couple reasons -

1. visitor implements DistancesUnordered so it is not unreasonable it should know how to construct the correct type to apply the streaming distance computation.
2. 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]

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

[arkrishn94]

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]

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

[arrayka]

I would name the test something like this: flat_index_supports_multi_threading

[arrayka]

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

[hildebrandmw]

Thanks Aditya,

Here's my take: The new HasElementRef and DistancesUnordered (the one in provider, not flat) add a hefty boilerplate burden to users of the graph index. Especially in light of something like #1067 which pushes in the other direction and while that PR is still experimental, I'm getting more convinced is the right API for the graph index. This puts us in kind of an awkward place. If we get this in, users consuming a new version will have churn for these little misc traits that will then need to be undone by #1067 (or some spiritual successor) which is not great.

If you want to start programming against this API, what if you introduce your own temporary TemporaryBuildQueryComputer (I believe this is the biggest reason why this PR needs to add HasElementRef and DistancesUnrodered) and we hold off on post processing for a short period of time to bottom out on #1067. If the latter doesn't pan out, then we can continue with these changes and accept it. Otherwise, the flat index implementations (which presumably will be less numerous than the current graph implementations) can be brought inline with the new trait organization and we'll cause less overall churn for our users.

I also remain unconvinced that the Iterator is any simpler than just implementing the new DistancesUnordered trait directly.

[hildebrandmw]

I'm still not convinced that this is justified. If someone can write a coherent next implementation and all of the state-tracking shenanigans that that will require, why can they not write a handful more lines of code and simply implement all of DistancesUnordered?

[arkrishn94]

I don't expect all (or even most) consumers to implement the flat index using the FlatIterator API. This is to provide an easy entry point for a naive implementation of the flat index for providers. Fwiw it also helps a bit with testing providers the flat index search path for providers that can implement an iteration pattern easily.

[hildebrandmw]

I'm not convinced that this is any easier an entry point than DistancesUnordered. You still need all the BuildQueryComputer implementations, there are more associated types, the return type is more complicated, users need to implement state tracking, and users have to know to use the Iterated proxy type. If one can manage to do all of that, then it's trivial to implement DistancesUnordered directly and be done with it.

[hildebrandmw]

I suppose the error bound on this super trait can be Into<ANNError> instead of ToRanked. Since it is coarse grained, implementations can decide internally to accept transient errors or now. I'm not sure I see a compelling reason to return a ToRanked error here.

[hildebrandmw]

That this somewhat general trait shares the same name as the much more specific trait in the flat index is quite confusing. I support traits having the same name if they serve the same purpose in different modules. For example, graph::SearchExt and flat::SearchExt - these (potentially hypothetical) traits serve the same purpose for different index types.

[arkrishn94]

If you want to start programming against this API, what if you introduce your own temporary TemporaryBuildQueryComputer (I believe this is the biggest reason why this PR needs to add HasElementRef and DistancesUnrodered) and we hold off on post processing for a short period of time to bottom out on #1067. If the latter doesn't pan out, then we can continue with these changes and accept it. Otherwise, the flat index implementations (which presumably will be less numerous than the current graph implementations) can be brought inline with the new trait organization and we'll cause less overall churn for our users.

Thanks for the feedback and thorough review @hildebrandmw. To make sure I understand, are you suggesting we temporarily create a disjoint trait structure for the flat index (avoiding supporting post-processing) and then once #1067 settles down, we can re-evaluate how much shared surface we can have for these two indexes (both on the building/construction side and the post-processing)?

[hildebrandmw]

If you want to start programming against this API, what if you introduce your own temporary TemporaryBuildQueryComputer (I believe this is the biggest reason why this PR needs to add HasElementRef and DistancesUnrodered) and we hold off on post processing for a short period of time to bottom out on #1067. If the latter doesn't pan out, then we can continue with these changes and accept it. Otherwise, the flat index implementations (which presumably will be less numerous than the current graph implementations) can be brought inline with the new trait organization and we'll cause less overall churn for our users.

Thanks for the feedback and thorough review @hildebrandmw. To make sure I understand, are you suggesting we temporarily create a disjoint trait structure for the flat index (avoiding supporting post-processing) and then once #1067 settles down, we can re-evaluate how much shared surface we can have for these two indexes (both on the building/construction side and the post-processing)?

Yeah, that's what I'm proposing. Something to minimize code churn, or at least defer it temporarily until we know that we aren't going to introduce these intermediate traits (HasElementRef + provider::DistancesUnordered) just to remove them again immediately.