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.