Migrate disk PQ flat scan to flat API (#1341)

- [x] 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 #1359; this PR adopts it in
`diskann-disk`.
- [x] Is the change to the API backwards compatible? **Yes.** Existing
disk search modes and result semantics are preserved.
- [x] 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.

```mermaid
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
```

`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.

---------

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: partychen

diskann-disk/src/search/pq/mod.rs

@@ -11,4 +11,3 @@ pub use pq_scratch::PQScratch;
 pub(crate) use crate::storage::quant::pq::PQData;
 
 mod quantizer_preprocess;
-pub use quantizer_preprocess::quantizer_preprocess;

diskann-disk/src/search/pq/quantizer_preprocess.rs

@@ -6,59 +6,51 @@
 use diskann::ANNResult;
 use diskann_vector::distance::Metric;
 
-use diskann_providers::model::compute_pq_distance;
 use diskann_providers::utils::BridgeErr;
 
 use super::{PQData, PQScratch};
 
-/// Preprocesses the query vector for PQ distance calculations.
-/// This function rotates the query vector and prepares the PQ table distances
-/// for efficient computation during search operations.
-pub fn quantizer_preprocess(
-    pq_scratch: &mut PQScratch,
-    pq_data: &PQData,
-    metric: Metric,
-    id_to_calculate_pq_distance: &[u32],
-) -> ANNResult<()> {
-    let table = pq_data.pq_table();
-    let expected_len = table.ncenters() * table.nchunks();
-    let dst = diskann_utils::views::MutMatrixView::try_from(
-        &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len],
-        table.nchunks(),
-        table.ncenters(),
-    )
-    .bridge_err()?;
-
-    match metric {
-        // Prior to the introduction of the `quantizer_preprocess` method, the
-        // disk index was hard-coded to use L2 distance for comparisons.
-        //
-        // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
-        // as L2 until a more thorough evaluation can be made.
-        Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
-            table.process_into::<diskann_quantization::distances::SquaredL2>(
-                &pq_scratch.query_scratch,
-                dst,
-            );
-        }
-        Metric::InnerProduct => {
-            table.process_into::<diskann_quantization::distances::InnerProduct>(
-                &pq_scratch.query_scratch,
-                dst,
-            );
-        }
+impl PQScratch {
+    pub(crate) fn prepare_query(
+        &mut self,
+        pq_data: &PQData,
+        metric: Metric,
+        query: &[f32],
+    ) -> ANNResult<()> {
+        self.set(query)?;
+        self.preprocess_query(pq_data, metric)
     }
 
-    // Compute the pq distance between query vector to all the vertex in the pq
-    // calculation id scratch.
-    compute_pq_distance(
-        id_to_calculate_pq_distance,
-        pq_data.get_num_chunks(),
-        &pq_scratch.aligned_pqtable_dist_scratch,
-        pq_data.pq_compressed_data().as_slice(),
-        &mut pq_scratch.aligned_pq_coord_scratch,
-        &mut pq_scratch.aligned_dist_scratch,
-    )?;
+    fn preprocess_query(&mut self, pq_data: &PQData, metric: Metric) -> ANNResult<()> {
+        let table = pq_data.pq_table();
+        let expected_len = table.ncenters() * table.nchunks();
+        let dst = diskann_utils::views::MutMatrixView::try_from(
+            &mut self.aligned_pqtable_dist_scratch[..expected_len],
+            table.nchunks(),
+            table.ncenters(),
+        )
+        .bridge_err()?;
 
-    Ok(())
+        match metric {
+            // Prior to moving query preprocessing onto `PQScratch`, the disk index
+            // was hard-coded to use L2 distance for comparisons.
+            //
+            // We're keeping that behavior here - treating `Cosine` and `CosineNormalized`
+            // as L2 until a more thorough evaluation can be made.
+            Metric::L2 | Metric::Cosine | Metric::CosineNormalized => {
+                table.process_into::<diskann_quantization::distances::SquaredL2>(
+                    &self.query_scratch,
+                    dst,
+                );
+            }
+            Metric::InnerProduct => {
+                table.process_into::<diskann_quantization::distances::InnerProduct>(
+                    &self.query_scratch,
+                    dst,
+                );
+            }
+        }
+
+        Ok(())
+    }
 }