Migrate disk PQ flat scan to flat API

[juchen-ms (partychen)]

• Does this PR have a descriptive title that could go in our release notes? Yes.
• Does this PR add any new dependencies? No.
• Does this PR modify any existing APIs? Yes. It makes the public flat-search visitor API query-aware and adds a borrowed-provider flat::knn_search entry point.
• Is the change to the API backwards compatible? No. Implementations of the public flat-search traits must move query-specific distance computation into their visitor. The RFC and in-repository implementations are updated in this PR.
• Should this result in any changes to our documentation, either updating existing docs or adding new ones? Yes. RFC 00983 and the affected rustdoc are updated.

Reference Issues/PRs

Closes #1104. The query-aware ownership direction follows the graph-search accessor model discussed in #1067.

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

• Migrates disk PQ flat search to the shared flat k-NN API using query-aware visitors.
• Adds a borrowed-provider flat::knn_search entry point while retaining FlatIndex::knn_search as a convenience wrapper.
• Shares one pooled DiskSearchScratch between graph and flat PQ search.
• Preserves scan-time filtering, top-k selection, and full-precision reranking.
• Updates RFC 00983 to describe the query-aware visitor design.

The previous flat API separated visitors from query-specific distance computation. For disk PQ, that split required query state to be spread across multiple objects, pools, and initialization stages. Making the visitor query-aware lets each backend combine data access and distance computation while the generic flat layer continues to own top-k selection, comparison accounting, error escalation, and post-processing.

Any other comments?

Validation includes targeted flat framework, PQ scratch, and disk filtered-search tests; graph and flat PQ search, filtering, pooled-query isolation, and indexed-vector coverage; and workspace clippy with warnings denied.

[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 94.47853% with 27 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.54%. Comparing base (07af709) to head (75aa7c3).

Files with missing lines	Patch %	Lines
diskann-disk/src/search/provider/disk_provider.rs	95.53%	15 Missing ⚠️
diskann-disk/src/search/pq/quantizer_preprocess.rs	63.63%	12 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1341      +/-   ##
==========================================
- Coverage   91.55%   91.54%   -0.02%     
==========================================
  Files         521      521              
  Lines      100371   100609     +238     
==========================================
+ Hits        91899    92098     +199     
- Misses       8472     8511      +39     

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

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.63% <100.00%> (-0.12%)	⬇️
diskann-disk/src/search/pq/pq_scratch.rs	100.00% <100.00%> (ø)
diskann/src/flat/index.rs	100.00% <100.00%> (ø)
diskann/src/flat/strategy.rs	98.88% <100.00%> (-0.32%)	⬇️
diskann/src/flat/test/provider.rs	73.45% <100.00%> (-0.49%)	⬇️
diskann-disk/src/search/pq/quantizer_preprocess.rs	61.40% <63.63%> (-35.82%)	⬇️
diskann-disk/src/search/provider/disk_provider.rs	95.81% <95.53%> (-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)]

I agree with your assessment. This migration exposed a limitation in the current flat API: separating the visitor from QueryComputer works for simple element-wise scans, but forces a backend such as disk PQ to split one query across multiple objects, pools, and initialization stages. Continuing to adapt around that API would recreate the original custom flat scan without gaining much from the shared abstraction.

I also agree that the better direction is to improve the flat API itself. Following the principle established in PR #1067, I propose making flat visitors query-aware and responsible for producing distances.

Proposed API

pub trait DistancesUnordered: HasId + Send + Sync {
    type Error: ToRanked + Debug + Send + Sync + 'static;

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

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

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

The generic flat-search flow becomes:

let mut visitor = strategy.create_visitor(provider, context, query)?;
visitor.distances_unordered(callback).await?;
processor.post_process(&mut visitor, query, candidates, output).await?;

The responsibility boundary would be:

• The generic flat layer owns top-k selection, comparison accounting, scan-error escalation, and post-processing.
• The backend visitor owns query preprocessing, filtering, batching, data access, and distance computation.

For disk PQ, graph and flat search can then use one pooled DiskSearchScratch containing the query buffer, PQ lookup table, batch buffers, vertex provider, and reranking cache. This removes the separate PQQueryComputer, its object pool, and the additional initialization stage.

The main advantages are:

• It follows the accessor-ownership model from PR Simplify the DataProvider contract for graph search #1067.
• It gives backends a coarse boundary for fusing access and computation.
• It supports both simple element-wise visitors and optimized batch implementations.
• It simplifies disk PQ query state and resource management.

The main trade-off is a public flat-trait change. To limit migration cost, the existing trait and method names remain. I also searched for visible consumers and did not find an independent public implementation outside DiskANN itself, forks, and vendored copies.

I have tried this proposal in the latest revision of the PR so that the design can be reviewed through a concrete implementation:

• create_visitor now receives the query.
• The visitor produces distances without a separate QueryComputer.
• Graph and flat disk PQ search share one pooled DiskSearchScratch.
• Flat filtering still happens before candidates enter top-k, so reranking uses AcceptAll and does not evaluate the predicate twice.
• RFC 00983 has been updated to describe the proposed design.

I also agree that a generic batch PQ primitive independent of diskann-disk would improve reuse. I have kept that as separate follow-up work so this proposal stays focused on the flat API boundary.

I would appreciate your review of both the proposed API direction and the implementation in this revision. Does this align with what you had in mind? Mark Hildebrand (@hildebrandmw) Aditya Krishnan (@arkrishn94)