Migrate disk PQ flat scan to flat API

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

@@ -6,9 +6,11 @@
 //! Product quantization types and functions used during disk-based search.
 
 mod pq_scratch;
-pub use pq_scratch::PQScratch;
+pub(crate) use pq_scratch::{PQBatchScratch, PQQueryComputerArgs, PQQueryComputerStorage};
+pub use pq_scratch::{PQQueryComputer, PQScratch};
 
 pub(crate) use crate::storage::quant::pq::PQData;
 
 mod quantizer_preprocess;
+pub(crate) use quantizer_preprocess::prepare_query;
 pub use quantizer_preprocess::quantizer_preprocess;

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

@@ -7,14 +7,143 @@
 use diskann::ANNResult;
 
 use diskann_quantization::alloc::{AlignedAllocator, Poly};
+use diskann_utils::object_pool::{ObjectPool, PoolOption, TryAsPooled};
+use diskann_vector::PreprocessedDistanceFunction;
+use std::sync::Arc;
 
 use crate::error::{diskann_error, ErrorKind};
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct PQQueryComputerArgs {
+    dim: usize,
+    num_pq_chunks: usize,
+    num_centers: usize,
+}
+
+impl PQQueryComputerArgs {
+    pub(crate) fn new(dim: usize, num_pq_chunks: usize, num_centers: usize) -> Self {
+        Self {
+            dim,
+            num_pq_chunks,
+            num_centers,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct PQQueryComputerStorage {
+    aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>,
+    query_scratch: Vec<f32>,
+    num_pq_chunks: usize,
+    num_centers: usize,
+}
+
+impl TryAsPooled<PQQueryComputerArgs> for PQQueryComputerStorage {
+    type Error = diskann::ANNError;
+
+    fn try_create(args: PQQueryComputerArgs) -> Result<Self, Self::Error> {
+        let aligned_pqtable_dist_scratch = Poly::broadcast(
+            0f32,
+            args.num_centers * args.num_pq_chunks,
+            AlignedAllocator::A128,
+        )
+        .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
+
+        Ok(Self {
+            aligned_pqtable_dist_scratch,
+            query_scratch: vec![0.0; args.dim],
+            num_pq_chunks: args.num_pq_chunks,
+            num_centers: args.num_centers,
+        })
+    }
+
+    fn try_modify(&mut self, args: PQQueryComputerArgs) -> Result<(), Self::Error> {
+        if self.query_scratch.len() != args.dim
+            || self.num_pq_chunks != args.num_pq_chunks
+            || self.num_centers != args.num_centers
+        {
+            *self = Self::try_create(args)?;
+        }
+        Ok(())
+    }
+}
+
+/// Opaque preprocessed query state created internally by disk search strategies.
+#[derive(Debug)]
+pub struct PQQueryComputer {
+    storage: PoolOption<PQQueryComputerStorage>,
+}
+
+impl PQQueryComputer {
+    /// Create an empty query computer for the given PQ schema.
+    #[cfg(test)]
+    pub(crate) fn new(dim: usize, num_pq_chunks: usize, num_centers: usize) -> ANNResult<Self> {
+        Ok(Self {
+            storage: PoolOption::try_non_pooled_create(PQQueryComputerArgs::new(
+                dim,
+                num_pq_chunks,
+                num_centers,
+            ))?,
+        })
+    }
+
+    pub(crate) fn pooled(
+        pool: &Arc<ObjectPool<PQQueryComputerStorage>>,
+        args: PQQueryComputerArgs,
+    ) -> ANNResult<Self> {
+        Ok(Self {
+            storage: PoolOption::try_pooled(pool, args)?,
+        })
+    }
+
+    /// Copy a full-precision query into the preprocessing buffer.
+    pub(crate) fn set(&mut self, query: &[f32]) -> ANNResult<()> {
+        let dim = self.storage.query_scratch.len();
+        if query.len() != dim {
+            return Err(diskann_error!(
+                ErrorKind::DimensionMismatchError,
+                "PQQueryComputer::set: expected query of length {dim}, got {}",
+                query.len()
+            ));
+        }
+        self.storage.query_scratch.copy_from_slice(query);
+        Ok(())
+    }
+
+    pub(crate) fn lookup_table(&self) -> &[f32] {
+        &self.storage.aligned_pqtable_dist_scratch
+    }
+
+    pub(super) fn preprocessing_buffers(&mut self) -> (&[f32], &mut [f32]) {
+        let storage = &mut *self.storage;
+        (
+            &storage.query_scratch,
+            &mut storage.aligned_pqtable_dist_scratch,
+        )
+    }
+}
+
+impl PreprocessedDistanceFunction<&[u8], f32> for PQQueryComputer {
+    fn evaluate_similarity(&self, code: &[u8]) -> f32 {
+        assert_eq!(
+            code.len(),
+            self.storage.num_pq_chunks,
+            "PQ code has the wrong number of chunks",
+        );
+        code.iter()
+            .enumerate()
+            .map(|(chunk, &center)| {
+                self.storage.aligned_pqtable_dist_scratch
+                    [chunk * self.storage.num_centers + center as usize]
+            })
+            .sum()
+    }
+}
+
 #[derive(Debug)]
 /// PQ scratch
 pub struct PQScratch {
-    /// Aligned pq table distance scratch, the length must be at least [256 * NCHUNKS]. 256 is the number of PQ centroids.
-    /// This is used to store the distance between each chunk in the query vector to each centroid, which is why the length is num of centroids * num of chunks
+    /// Aligned PQ table distance scratch.
     pub aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>,
 
     /// Aligned dist scratch, must be at least diskann MAX_DEGREE
@@ -25,9 +154,7 @@ pub struct PQScratch {
     /// This is used to store the pq coordinates of the candidate vectors.
     pub aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>,
 
-    /// Query scratch buffer stored as `f32`, sized by the PQ table's logical dimension.
-    /// `set` populates it from a caller-provided `&[f32]`; `PQTable::preprocess_query` can
-    /// then rotate or otherwise preprocess it.
+    /// Query scratch buffer stored as `f32`.
     pub query_scratch: Vec<f32>,
 }
 
@@ -45,18 +172,17 @@ impl PQScratch {
         let aligned_pq_coord_scratch =
             Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128)
                 .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
+        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
+            .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
         let aligned_pqtable_dist_scratch =
             Poly::broadcast(0f32, num_centers * num_pq_chunks, AlignedAllocator::A128)
                 .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
-        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
-            .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
-        let query_scratch = vec![0.0f32; dim];
 
         Ok(Self {
             aligned_pqtable_dist_scratch,
             aligned_dist_scratch,
             aligned_pq_coord_scratch,
-            query_scratch,
+            query_scratch: vec![0.0; dim],
         })
     }
 
@@ -80,22 +206,86 @@ impl PQScratch {
         Ok(())
     }
 
-    /// Return the largest number of PQ vectors whose distances can be computed using this
-    /// scratch data structure.
+    /// Return the largest number of PQ vectors that fit in the batch scratch.
+    #[cfg(test)]
+    pub(crate) fn max_vectors(&self) -> usize {
+        self.aligned_dist_scratch.len()
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct PQBatchScratch {
+    pub(crate) aligned_dist_scratch: Poly<[f32], AlignedAllocator>,
+    pub(crate) aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>,
+}
+
+impl PQBatchScratch {
+    pub(crate) fn new(graph_degree: usize, num_pq_chunks: usize) -> ANNResult<Self> {
+        let aligned_pq_coord_scratch =
+            Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128)
+                .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
+        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
+            .map_err(|e| diskann_error!(ErrorKind::IndexError, e))?;
+        Ok(Self {
+            aligned_dist_scratch,
+            aligned_pq_coord_scratch,
+        })
+    }
+
     pub(crate) fn max_vectors(&self) -> usize {
         self.aligned_dist_scratch.len()
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use diskann_quantization::num::PowerOfTwo;
+    use diskann_utils::object_pool::ObjectPool;
+    use diskann_vector::PreprocessedDistanceFunction;
     use rstest::rstest;
 
-    use super::PQScratch;
+    use super::{PQQueryComputer, PQQueryComputerArgs, PQQueryComputerStorage, PQScratch};
 
     use crate::error::{error_kind, ErrorKind};
 
+    #[test]
+    fn query_computer_scores_pq_code() {
+        let mut computer = PQQueryComputer::new(2, 2, 3).unwrap();
+        computer
+            .preprocessing_buffers()
+            .1
+            .copy_from_slice(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
+
+        assert_eq!(computer.evaluate_similarity(&[1, 2]), 6.0);
+    }
+
+    #[test]
+    fn query_computer_reuses_pooled_storage() {
+        let args = PQQueryComputerArgs::new(2, 2, 3);
+        let pool = Arc::new(ObjectPool::<PQQueryComputerStorage>::try_new(args, 0, None).unwrap());
+        assert!(pool.is_empty());
+
+        let first_ptr;
+        {
+            let computer = PQQueryComputer::pooled(&pool, args).unwrap();
+            assert!(pool.is_empty());
+            first_ptr = computer.lookup_table().as_ptr();
+        }
+        assert_eq!(pool.len(), 1);
+
+        let second_ptr;
+        {
+            let computer = PQQueryComputer::pooled(&pool, args).unwrap();
+            assert!(pool.is_empty());
+            second_ptr = computer.lookup_table().as_ptr();
+        }
+        assert_eq!(pool.len(), 1);
+
+        assert_eq!(first_ptr, second_ptr);
+    }
+
     #[rstest]
     #[case(512, 8, 128, 256)] // default test case
     #[case(59, 16, 37, 41)] // not multiple of 256
@@ -120,7 +310,6 @@ mod tests {
             (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
             0
         );
-
         assert_eq!(pq_scratch.max_vectors(), graph_degree);
 
         // Test set() method

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

@@ -9,21 +9,21 @@ use diskann_vector::distance::Metric;
 use diskann_providers::model::compute_pq_distance;
 use diskann_providers::utils::BridgeErr;
 
-use super::{PQData, PQScratch};
+use super::{PQData, PQQueryComputer, 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,
+fn preprocess_query(
+    query: &[f32],
+    lookup_table: &mut [f32],
     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],
+        &mut lookup_table[..expected_len],
         table.nchunks(),
         table.ncenters(),
     )
@@ -36,21 +36,40 @@ pub fn quantizer_preprocess(
         // 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,
-            );
+            table.process_into::<diskann_quantization::distances::SquaredL2>(query, dst);
         }
         Metric::InnerProduct => {
-            table.process_into::<diskann_quantization::distances::InnerProduct>(
-                &pq_scratch.query_scratch,
-                dst,
-            );
+            table.process_into::<diskann_quantization::distances::InnerProduct>(query, dst);
         }
     }
 
-    // Compute the pq distance between query vector to all the vertex in the pq
-    // calculation id scratch.
+    Ok(())
+}
+
+pub(crate) fn prepare_query(
+    computer: &mut PQQueryComputer,
+    pq_data: &PQData,
+    metric: Metric,
+    query: &[f32],
+) -> ANNResult<()> {
+    computer.set(query)?;
+    let (query, lookup_table) = computer.preprocessing_buffers();
+    preprocess_query(query, lookup_table, pq_data, metric)
+}
+
+pub fn quantizer_preprocess(
+    pq_scratch: &mut PQScratch,
+    pq_data: &PQData,
+    metric: Metric,
+    id_to_calculate_pq_distance: &[u32],
+) -> ANNResult<()> {
+    preprocess_query(
+        &pq_scratch.query_scratch,
+        &mut pq_scratch.aligned_pqtable_dist_scratch,
+        pq_data,
+        metric,
+    )?;
+
     compute_pq_distance(
         id_to_calculate_pq_distance,
         pq_data.get_num_chunks(),