Make flat search visitors query-aware

[juchen-ms (partychen)]

• Does this PR have a descriptive title that could go in our release notes?
• Does this PR add any new dependencies?
• Does this PR modify any existing APIs?
• Is the change to the API backwards compatible?
• Should this result in any changes to our documentation, either updating existing docs or adding new ones?

Reference Issues/PRs

Prerequisite API refactor requested during review of #1341.

What does this implement/fix? Briefly explain your changes.

Makes flat search visitors query-aware, moves the search algorithm to the free flat::knn_search entry point, and removes the unnecessary FlatIndex wrapper. It also updates the generic flat tests, test providers, benchmark integration, and API rustdoc for the redesigned public API.

The redesigned interface has several benefits:

• A visitor is constructed for a specific query, so it can own or borrow query preprocessing results and combine them with backend-specific I/O, batching, filtering, and distance-computation state. This is particularly useful for streamed and quantized backends such as the disk PQ scan in Migrate disk PQ flat scan to flat API #1341.
• DistancesUnordered now emits (id, distance) pairs directly. Backends can fuse scanning and distance computation instead of exposing every stored element through a common ElementRef and external QueryComputer abstraction.
• Implementations have a smaller and less brittle type surface. The redesign removes the ElementRef, QueryComputer, and QueryComputerError associated types, the visitor GAT, and their HRTB/lifetime constraints.
• The free flat::knn_search(&provider, ...) function borrows the provider directly, removing a stateless ownership wrapper and making shared providers and concurrent searches more natural.
• Responsibilities are clearer: the generic algorithm manages top-k selection and post-processing, while the query-aware visitor owns the backend-specific complete scan.

Any other comments?

This intentionally changes the existing public flat-search API and is not backwards compatible. The trade-off is a breaking migration for current callers in exchange for an interface that can naturally represent query-aware, streaming, and quantized backends. #1341 will remain open and be rebased onto main after this prerequisite merges.

[Copilot]

Pull request overview

This PR redesigns DiskANN’s flat (sequential) k-NN search API so flat-scan visitors become query-aware, shifting query preprocessing/distance computation into the visitor and exposing the search algorithm as a free flat::knn_search entry point (removing the thin FlatIndex wrapper).

Changes:

• Reworks the flat-search traits so DistancesUnordered visitors own per-query state, and SearchStrategy constructs a visitor using the query.
• Moves the brute-force flat k-NN algorithm to flat::knn_search(provider, ...) and updates tests accordingly.
• Updates benchmark integration to call flat::knn_search directly and to use query-aware visitors.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
diskann/src/flat/test/provider.rs	Updates the test visitor/strategy to be query-aware by storing a query distance computer in the visitor.
diskann/src/flat/test/harness.rs	Switches the reusable harness from FlatIndex::knn_search to the free flat::knn_search function.
diskann/src/flat/test/cases/flat_knn_search.rs	Updates the baseline-cached regression sweep to use a shared provider (no FlatIndex).
diskann/src/flat/strategy.rs	Refactors core flat-search traits to make visitors query-aware and simplifies the DistancesUnordered contract.
diskann/src/flat/mod.rs	Updates module exports/docs to expose knn_search as the flat-search entry point.
diskann/src/flat/index.rs	Removes FlatIndex and implements knn_search as a borrowed-provider free function; updates flat-search tests.
diskann-benchmark/src/flat/search.rs	Migrates benchmark backend from FlatIndex::knn_search to flat::knn_search and updates the benchmark visitor/strategy accordingly.

Suppressed comments (1)

diskann/src/flat/test/cases/flat_knn_search.rs:90

• This doc comment still refers to a shared index, but the test now shares only a provider (the FlatIndex wrapper was removed). Update the wording so it matches the new API.

/// Run `knn_search` + brute-force oracle against a *shared* `index`, assert the
/// cross-row invariants, and produce the baseline row. The per-row provider metrics
/// captured into the baseline are the *delta* observed during this row, which keeps

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.53%. Comparing base (07af709) to head (dae3d62).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1359      +/-   ##
==========================================
+ Coverage   91.55%   92.53%   +0.97%     
==========================================
  Files         521      521              
  Lines      100371   100313      -58     
==========================================
+ Hits        91899    92822     +923     
+ Misses       8472     7491     -981     

Flag	Coverage Δ
miri	92.53% <100.00%> (+0.97%)	⬆️
unittests	92.49% <100.00%> (+1.25%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-benchmark/src/flat/search.rs	95.60% <100.00%> (-0.15%)	⬇️
diskann/src/flat/index.rs	100.00% <100.00%> (ø)
diskann/src/flat/strategy.rs	98.88% <100.00%> (-0.32%)	⬇️
diskann/src/flat/test/cases/flat_knn_search.rs	100.00% <100.00%> (ø)
diskann/src/flat/test/harness.rs	100.00% <100.00%> (ø)
diskann/src/flat/test/provider.rs	73.45% <100.00%> (-0.49%)	⬇️

... and 43 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[Aditya Krishnan (arkrishn94)]

Thanks Junkui, I like the simplification here. Only left some minor comments, but apart from that happy to approve.