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? No. The query-aware flat-search API was introduced separately in Make flat search visitors query-aware #1359; this PR adopts it in diskann-disk.
• Is the change to the API backwards compatible? Yes. Existing disk search modes and result semantics are preserved.
• Should this result in any changes to our documentation, either updating existing docs or adding new ones? Yes. The affected implementation rustdoc is updated.

Reference Issues/PRs

Built on the query-aware flat-search API merged in #1359.

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

• Replaces the disk-specific manual PQ flat-scan pipeline with the shared flat k-NN API.
• Implements DistancesUnordered on DiskAccessor to expose complete, batched PQ-distance scanning.
• Initializes query-dependent PQ state after every pooled scratch checkout.
• Preserves scan-time filtering before approximate top-k selection and full-precision reranking afterward.
• Preserves the existing pooled scratch and indexed-vector result behavior across graph and flat search modes.

The disk backend constructs a query-aware DiskAccessor and passes it directly to flat::knn_search. The generic flat layer now owns top-k selection, comparison accounting, error escalation, and post-processing. DiskAccessor continues to own disk-specific PQ preprocessing, batching, filtering, data access, and distance computation.

Any other comments?

This PR has been rebased onto main after #1359 merged. Its diff is limited to the two diskann-disk implementation files.

Architecture simplification

Before this change, disk flat search manually coordinated filtering, batching, PQ-distance collection, top-k selection, comparison accounting, and post-processing inside DiskANNIndex::flat_search. Graph and flat search already used the same DiskAccessor and scratch pool, but the flat algorithm duplicated orchestration now provided by the shared flat API.

flowchart TB
    subgraph Before["Before: disk-specific flat orchestration"]
        direction LR
        F1["FlatScan"] --> M["DiskANNIndex::flat_search"]
        M --> FI["filter IDs"]
        FI --> B["manual batch loop"]
        B --> PQ1["DiskAccessor::pq_distances"]
        PQ1 --> K1["local NeighborPriorityQueue"]
        K1 --> PP1["disk post-processor"]
    end

    subgraph After["After: shared flat orchestration"]
        direction LR
        F2["FlatScan"] --> K2["flat::knn_search"]
        K2 --> DU["DiskAccessor<br/>DistancesUnordered"]
        DU --> PQ2["filtered, batched PQ scan"]
        K2 --> TK["shared top-k · stats · errors"]
        TK --> PP2["RerankAndFilter"]
    end

    G["Graph search"] --> SA["DiskAccessor<br/>SearchAccessor"]
    DU --> S["pooled DiskSearchScratch<br/>per-query PQ preparation"]
    SA --> S

Loading

DiskAccessor now exposes the disk scan through DistancesUnordered, allowing flat::knn_search to drive the common k-NN workflow while graph traversal continues to use the existing SearchAccessor implementation. Both paths preserve their distinct filtering stages and share the same pooled query-state lifecycle.

[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.61883% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.56%. Comparing base (1632651) to head (cb06394).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-disk/src/search/provider/disk_provider.rs	93.75%	12 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1341      +/-   ##
==========================================
+ Coverage   91.55%   92.56%   +1.01%     
==========================================
  Files         521      521              
  Lines      100302   100850     +548     
==========================================
+ Hits        91827    93353    +1526     
+ Misses       8475     7497     -978     

Flag	Coverage Δ
miri	92.56% <94.61%> (+1.01%)	⬆️
unittests	92.50% <94.61%> (+1.27%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-disk/src/search/pq/quantizer_preprocess.rs	100.00% <100.00%> (+2.77%)	⬆️
diskann-disk/src/search/provider/disk_provider.rs	95.65% <93.75%> (-0.17%)	⬇️

... and 72 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?

[juchen-ms (partychen)]

Mark Hildebrand (@hildebrandmw) Aditya Krishnan (@arkrishn94) Thanks for the detailed feedback. I reorganized the implementation around the following design.

Disk PQ flat-scan design

Disk PQ flat search constructs a query-initialized DiskAccessor and passes it directly to flat::knn_search.

DiskAccessor implements both SearchAccessor for graph traversal and DistancesUnordered for exhaustive flat scanning. There is no separate FlatVisitor. Both search modes use the same accessor, pooled DiskSearchScratch, PQ preprocessing, vertex loading, reranking cache, and indexed-vector handling.

For flat search, DistancesUnordered enumerates every eligible vector ID, applies the optional filter before top-k selection, computes approximate PQ distances in scratch-sized batches, and emits each (id, distance) pair to flat::knn_search.

flat::knn_search owns the backend-independent workflow: bounded top-k selection, comparison accounting, scan-error handling, post-processing, and output statistics. The surviving candidates are then reranked using their full-precision indexed vectors.

Why use the flat API here?

The flat API removes algorithm orchestration from diskann-disk. Without it, the disk backend must maintain its own priority queue, comparison counter, batch loop, error handling, and post-processing invocation.

The abstraction does not attempt to hide disk PQ details. DiskAccessor retains control over filtering, batching, PQ computation, disk access, and scratch state. The shared layer only owns behavior common to exhaustive k-NN search. This provides reuse without preventing backend-specific optimizations.

Scratch initialization

DiskAccessor::new checks out a DiskSearchScratch through DiskSearchScratch::pooled_for_query. After every checkout, PQScratch::prepare_query copies and preprocesses the current query. This must happen per checkout rather than only when a pool entry is first allocated, because the same scratch object is reused by subsequent queries.

Graph and flat search therefore share one initialization path while keeping query state isolated between pool uses.

Filtering and reranking

Flat filtering occurs during the scan, before candidates enter the approximate top-k queue. Consequently, reranking uses AcceptAll; applying the predicate again would be redundant.

Graph search retains its existing filtering stages because traversal and flat scan have different candidate-selection semantics.

[Wei Wu (wuw92)]

Could we pass the flat-scan filter explicitly instead of routing it through postprocess_filter? The current flow converts postprocess_filter into a scan-time filter and then hard-codes AcceptAll for post-processing, which makes it difficult to tell where filtering actually belongs. Passing AcceptAll to the strategy and a separate filter argument to flat_search would make the stage ownership explicit.

[juchen-ms (partychen)]

Agreed. Routing the filter through postprocess_filter was another leftover from the earlier design with separate accessors, and I missed restoring the explicit boundary after consolidating on DiskAccessor. In 0b0e333, the flat-scan filter is passed directly to flat_search while the strategy uses AcceptAll, so the scan-time and post-processing responsibilities are explicit.