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. The query-aware flat-search API change is isolated in prerequisite PR Make flat search visitors query-aware #1359.
• 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, as described and reviewed separately in Make flat search visitors query-aware #1359.
• Should this result in any changes to our documentation, either updating existing docs or adding new ones? Yes. The affected API rustdoc is updated.

Reference Issues/PRs

Depends on #1359. After #1359 merges, this PR will be rebased onto main so its diff contains only the disk PQ migration.

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

• Migrates disk PQ flat search to the shared flat k-NN API.
• Adds batched unordered PQ-distance scanning.
• Reuses DiskAccessor and one pooled DiskSearchScratch across graph and flat PQ search.
• Shares per-query PQ preprocessing between graph and flat search.
• Preserves scan-time filtering and full-precision reranking.

The generic query-aware flat API redesign is reviewed independently in #1359. This PR focuses on adapting the disk backend to that API without maintaining a separate flat visitor or duplicate scratch initialization path.

Any other comments?

Until #1359 merges, this diff temporarily includes its prerequisite API changes. Those changes will disappear from this PR after it is rebased onto the updated main branch.

Architecture simplification

Before this change, graph search and flat search used separate query-state owners. DiskAccessor and FlatVisitor duplicated the provider reference, pooled scratch handle, cache configuration, constructor logic, query initialization, and post-processing integration.

flowchart TB
    subgraph Before["Before: parallel query-state owners"]
        direction LR
        G1["Graph search"] --> DA1["DiskAccessor"] --> S1["DiskSearchScratch"]
        F1["Flat search"] --> FV["FlatVisitor"] --> S2["DiskSearchScratch"]
    end

    DUP["Duplicated ownership and lifecycle:<br/>provider · pooled scratch · cache flag<br/>constructor · query initialization · post-process adapter"]
    DA1 -.-> DUP
    FV -.-> DUP

    subgraph After["After: one query-aware accessor"]
        direction LR
        G2["Graph search"] --> DA2["DiskAccessor"]
        F2["Flat search"] --> DA2
        DA2 --> API["SearchAccessor<br/>+ DistancesUnordered"]
        API --> S3["DiskSearchScratch<br/>one pooled-query lifecycle"]
    end

Loading

FlatVisitor is removed. DiskAccessor now supports both graph traversal and unordered flat scanning, so both paths share one query-state model and one scratch initialization path while preserving their different filtering stages.

[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 96.98276% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.55%. Comparing base (07af709) to head (fb22e35).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-disk/src/search/provider/disk_provider.rs	95.53%	13 Missing ⚠️
diskann-disk/src/search/pq/quantizer_preprocess.rs	96.96%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1341      +/-   ##
==========================================
- Coverage   91.55%   91.55%   -0.01%     
==========================================
  Files         521      521              
  Lines      100371   100516     +145     
==========================================
+ Hits        91899    92027     +128     
- Misses       8472     8489      +17     

Flag	Coverage Δ
miri	91.55% <96.98%> (-0.01%)	⬇️
unittests	91.23% <96.98%> (-0.01%)	⬇️

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%)	⬇️
diskann-disk/src/search/pq/quantizer_preprocess.rs	65.30% <96.96%> (-31.92%)	⬇️
diskann-disk/src/search/provider/disk_provider.rs	95.81% <95.53%> (-0.01%)	⬇️

... and 3 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)

[Mark Hildebrand (hildebrandmw)]

Thanks - reworking the flat API to resemble the graph API and merging the compute into DiskSearchScratch is much cleaner.

However, my larger concern still applies. I do not see what the flat API is enabling in diskann-disk to justify the increased complexity. Aditya Krishnan (@arkrishn94) - can you weigh in?

[Aditya Krishnan (arkrishn94)]

Thanks Junkui. Firstly, I apologize for the severely delayed review for this PR.

I'm largely on board with the direction of these changes. I like the simplification of the flat search API and might even suggest getting rid of the FlatIndex entirely. I guess if we go down this path, it might make sense to open a pre-cursor PR to this with just the changes to the flat API.

The simplification to DiskSearchScratch looks good, although as Mark said it would be nice to consolidate the initialization for it.

The one comment I had about the complexity of introducing the flat API here is- can we get rid of the FlatVisitor struct entirely and just work over the DiskAccessor?