Migrate disk PQ flat scan to flat API

[juchen-ms (partychen)]

Summary

• implement a disk PQ FlatVisitor with batched DistancesUnordered scanning
• share PQQueryComputer between graph and flat PQ search preprocessing
• route disk flat search through the borrowed-provider flat k-NN entrypoint
• preserve scan-time filtering and full-precision reranking

Validation

• targeted flat framework, PQ scratch, and disk filtered-search tests
• clippy with warnings denied for affected crates

Closes #1104

[Copilot]

Pull request overview

This PR migrates the disk PQ “flat scan” path onto the shared diskann::flat API, introducing a dedicated disk PQ FlatSearchStrategy + visitor that scans PQ-compressed rows and then reuses the existing full-precision reranking + filtering pipeline. It also factors PQ query preprocessing into a reusable owned query-computer (TransposedQueryComputer) so both graph and flat PQ search can share the same preprocessing approach.

Changes:

• Update FlatIndex::knn_search to return a lifetime-bound SendFuture so it can borrow strategy/context/output across .await.
• Add TransposedQueryComputer (+ error type) to build per-query PQ lookup tables for transposed PQ tables.
• Route disk flat scan through FlatIndex using a new disk-specific flat strategy/visitor, and remove now-unused PQ scratch batching API.

Reviewed changes

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

Show a summary per file

File	Description
diskann/src/flat/index.rs	Adjusts knn_search signature/lifetimes to support borrowed-provider flat search entrypoints.
diskann-quantization/src/product/tables/transposed/query.rs	Introduces an owned PQ query computer for transposed tables (L2/IP), with unit tests.
diskann-quantization/src/product/tables/transposed/mod.rs	Wires the new transposed query module into the transposed table submodule exports.
diskann-quantization/src/product/tables/mod.rs	Re-exports the new transposed query computer + error at the tables module boundary.
diskann-quantization/src/product/mod.rs	Re-exports the new transposed query types at the product module boundary.
diskann-disk/src/search/provider/disk_provider.rs	Implements disk PQ flat scan via diskann::flat (DiskFlatProvider/DiskFlatSearchStrategy/DiskFlatVisitor) while preserving scan-time filtering and rerank behavior.
diskann-disk/src/search/pq/pq_scratch.rs	Removes PQScratch::max_vectors and updates tests accordingly (no longer needed after migration).

---

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

[juchen-ms (partychen)]

Aditya Krishnan (@arkrishn94) I ended up making a few design changes beyond the Visitor implementation, and I’d appreciate a sanity check on whether these are the right tradeoffs:

1. Borrowed-provider flat search
   DiskProvider is already owned by DiskANNIndex, so I added a borrowed-provider flat::knn_search entry point instead of creating another FlatIndex, cloning/wrapping the provider, or introducing a BorrowedFlatIndex type. The existing FlatIndex::knn_search delegates to it. Does this seem like the right API shape?

2. Shared rerank implementation
   RerankAndFilter needs to work with both DiskAccessor and FlatVisitor. I extracted the common implementation into rerank_and_filter, leaving two thin SearchPostProcess adapters. The alternative would be a shared accessor trait exposing provider/scratch through associated types. I felt that trait would be more abstraction than the two callers justify, but I’d like your opinion.

3. PQ query-computer construction and pooling
   I separated the query-to-centroid lookup state from PQScratch. Both graph and flat search now obtain a PQQueryComputer from a dedicated object pool, while DiskSearchScratch only retains the batch distance and coordinate buffers. This also preserves the existing SearchStrategy::build_query_computer(query) API: DiskSearchStrategy borrows the PQ schema, metric, and query-computer pool from DiskIndexSearcher instead of passing the provider into the trait method. Does this separation and pooling boundary seem appropriate?

One related detail: filtering happens in FlatVisitor before candidates enter the top-k queue, so reranking uses AcceptAll to avoid evaluating the predicate twice.

These were the main areas where the migration required broader architectural choices, so feedback on them would be helpful before finalizing the approach.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 93.16239% with 32 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.54%. Comparing base (45819de) to head (f270f6d).

Files with missing lines	Patch %	Lines
diskann-disk/src/search/provider/disk_provider.rs	95.32%	13 Missing ⚠️
diskann-disk/src/search/pq/quantizer_preprocess.rs	60.00%	12 Missing ⚠️
diskann-disk/src/search/pq/pq_scratch.rs	93.75%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1341      +/-   ##
==========================================
- Coverage   91.56%   91.54%   -0.02%     
==========================================
  Files         521      521              
  Lines       99527    99872     +345     
==========================================
+ Hits        91132    91431     +299     
- Misses       8395     8441      +46     

Flag	Coverage Δ
miri	91.54% <93.16%> (-0.02%)	⬇️
unittests	91.22% <93.16%> (-0.02%)	⬇️

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

Files with missing lines	Coverage Δ
diskann/src/flat/index.rs	100.00% <100.00%> (ø)
diskann-disk/src/search/pq/pq_scratch.rs	95.75% <93.75%> (-4.25%)	⬇️
diskann-disk/src/search/pq/quantizer_preprocess.rs	59.25% <60.00%> (-37.97%)	⬇️
diskann-disk/src/search/provider/disk_provider.rs	95.76% <95.32%> (+<0.01%)	⬆️

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[Mark Hildebrand (hildebrandmw)]

As usual, I will defer to the maintainers of diskann-disk to make the judgement calls here, but what immediately stands out to me is that trying to fit the flat scan into the diskann flat-scan API is essentially recreating the custom flat-scan implementation but with significantly more code. That is, this appears to be working hard to fit the API (and indeed changing the API in diskann) without materially benefiting from doing so.

To me, this indicates two things:

1. There is an ergonomic gap in the flat API that needs to be fixed. For example - it requires a QueryComputer which is causing some of the churn in this PR [1]. I don't think that's a good direction since it separates the compute engine from the internal of the FlatAccessor, when closer coupling (e.g. how SearchAccessor works now for the graph index) allows for safer optimization.
2. We're missing even lower-level infrastructure (e.g. generic batch PQ computation independent of diskann-disk) that would help with reusability. Think: a more generally reuseable version of compute_pq_distance.

There are parts that look good. Extracting rerank_and_filter to a synchronous function (instead of the current unfortunate bounce through async) is a good improvement. Simplifying PQ scratch initialization is good - though I might suggest keeping it in DiskSearchScratch fusing it with the DiskSearchScratch's pooled API to avoid the multi-stage initialization that is currently done.

[1] The graph portion of diskann used to work this way and it turns out to be way better for a huge number of reasons to not.

[juchen-ms (partychen)]

As usual, I will defer to the maintainers of diskann-disk to make the judgement calls here, but what immediately stands out to me is that trying to fit the flat scan into the diskann flat-scan API is essentially recreating the custom flat-scan implementation but with significantly more code. That is, this appears to be working hard to fit the API (and indeed changing the API in diskann) without materially benefiting from doing so.

To me, this indicates two things:

1. There is an ergonomic gap in the flat API that needs to be fixed. For example - it requires a QueryComputer which is causing some of the churn in this PR [1]. I don't think that's a good direction since it separates the compute engine from the internal of the FlatAccessor, when closer coupling (e.g. how SearchAccessor works now for the graph index) allows for safer optimization.
2. We're missing even lower-level infrastructure (e.g. generic batch PQ computation independent of diskann-disk) that would help with reusability. Think: a more generally reuseable version of compute_pq_distance.

There are parts that look good. Extracting rerank_and_filter to a synchronous function (instead of the current unfortunate bounce through async) is a good improvement. Simplifying PQ scratch initialization is good - though I might suggest keeping it in DiskSearchScratch fusing it with the DiskSearchScratch's pooled API to avoid the multi-stage initialization that is currently done.

[1] The graph portion of diskann used to work this way and it turns out to be way better for a huge number of reasons to not.

I agree that the current migration is adapting around the flat API rather than benefiting from it.

The disk PQ implementation exposed a limitation in the current abstraction: separating the visitor from QueryComputer works for simple element-wise scans, but forces backends with batched computation to split one query across multiple objects and initialization stages.

Following the principle established in PR #1067, I propose passing the query when constructing the flat visitor and making the visitor responsible for producing distances. The generic flat layer would still own top-k selection and post-processing; only backend-specific query execution moves behind the visitor.

For disk PQ, this allows graph and flat search to share one pooled DiskSearchScratch, removing the separate PQQueryComputer and its pool. Filtering, query preprocessing, batching, and PQ distance computation can remain together in the backend.

I also checked the visible consumers of this API. Outside DiskANN’s own benchmarks and tests, I could not find an independent public implementation; the other matches are forks or vendored copies. That suggests the migration cost is currently limited, while changing it now avoids establishing the separated-computer model as a long-term public contract.

The main trade-off is a public trait change, but we can minimize it by keeping the existing trait and method names:

let mut visitor = strategy.create_visitor(provider, context, query)?;
visitor.distances_unordered(callback).await?;

I have a local prototype of this design. It preserves the existing flat-search behavior and removes the additional query-computer machinery rather than replacing it with another abstraction.

Would this direction better match what you had in mind when referencing PR #1067? Mark Hildebrand (@hildebrandmw)
Aditya Krishnan (@arkrishn94)