Converge PQ pivot loading

diskann-disk/src/storage/disk_index_reader.rs

@@ -6,7 +6,9 @@ use std::sync::Arc;
 
 use diskann::ANNResult;
 use diskann_providers::storage::StorageReadProvider;
-use diskann_providers::{storage::PQStorage, utils::load_metadata_from_file};
+use diskann_providers::{
+    model::FixedChunkPQTable, storage::PQStorage, utils::load_metadata_from_file,
+};
 
 use crate::search::pq::PQData;
 use tracing::info;
@@ -29,11 +31,7 @@ impl DiskIndexReader {
         storage_provider: &Storage,
     ) -> ANNResult<Self> {
         let pq_storage = PQStorage::new(&pq_pivot_path, &pq_compressed_data_path, None);
-        let pq_pivot_table = pq_storage.load_pq_pivots_bin::<Storage>(
-            &pq_pivot_path,
-            0, // Use 0 to infer num_pq_chunks from the file
-            storage_provider,
-        )?;
+        let pq_pivot_table: FixedChunkPQTable = pq_storage.load_pivots(storage_provider)?.into();
 
         // Auto-detect number of points from compressed PQ file metadata
         let metadata = load_metadata_from_file(storage_provider, &pq_compressed_data_path)?;

diskann-disk/src/storage/quant/pq/pq_generation.rs

@@ -8,12 +8,9 @@ use std::{marker::PhantomData, time::Instant};
 use diskann::utils::VectorRepr;
 use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider};
 use diskann_providers::{
-    model::{
-        pq::{accum_row_inplace, generate_pq_pivots},
-        GeneratePivotArguments,
-    },
+    model::{pq::generate_pq_pivots, GeneratePivotArguments},
     storage::PQStorage,
-    utils::{BridgeErr, RayonThreadPoolRef},
+    utils::RayonThreadPoolRef,
 };
 use diskann_quantization::{error::Format, product::TransposedTable, CompressInto};
 use diskann_utils::views::MatrixBase;
@@ -113,32 +110,28 @@ where
             .pq_storage
             .read_existing_pivot_metadata(context.storage_provider)?;
 
-        //Load the pivots
         let num_chunks = context.num_chunks;
-        let (mut full_pivot_data, centroid, chunk_offsets) =
-            context.pq_storage.load_existing_pivot_data(
-                &num_chunks,
-                &context.num_centers,
-                &full_dim,
-                context.storage_provider,
-            )?;
+        let table = context.pq_storage.load_pivots(context.storage_provider)?;
 
-        let mut full_pivot_data_mat = diskann_utils::views::MutMatrixView::try_from(
-            full_pivot_data.as_mut_slice(),
-            context.num_centers,
-            full_dim,
-        )
-        .bridge_err()?;
-
-        accum_row_inplace(full_pivot_data_mat.as_mut_view(), centroid.as_slice());
+        if table.nchunks() != num_chunks
+            || table.ncenters() != context.num_centers
+            || table.dim() != full_dim
+        {
+            return Err(diskann_error!(
+                ErrorKind::PQError,
+                "PQ pivot table mismatch: file has {} chunks, {} centers in {} dimensions but expected {} chunks, {} centers in {} dimensions.",
+                table.nchunks(),
+                table.ncenters(),
+                table.dim(),
+                num_chunks,
+                context.num_centers,
+                full_dim
+            ));
+        }
 
-        let table = TransposedTable::from_parts(
-            full_pivot_data_mat.as_view(),
-            diskann_quantization::views::ChunkOffsetsView::new(&chunk_offsets)
-                .bridge_err()?
-                .to_owned(),
-        )
-        .map_err(|err| diskann_error!(ErrorKind::PQError, "{}", Format(err)))?;
+        let table =
+            TransposedTable::from_parts(table.view_pivots(), table.view_offsets().to_owned())
+                .map_err(|err| diskann_error!(ErrorKind::PQError, "{}", Format(err)))?;
 
         Ok(Self {
             table,

diskann-providers/src/model/graph/provider/async_/experimental/multi_pq_async.rs

@@ -7,13 +7,14 @@ use std::sync::{Arc, Mutex};
 
 use arc_swap::{ArcSwap, Guard};
 use diskann::{ANNError, ANNResult, error::IntoANNResult, utils::VectorRepr};
+use diskann_quantization::CompressInto;
 use diskann_utils::lazy_format;
 use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric};
 use rand::{Rng, SeedableRng, rngs::StdRng};
 
-use crate::model::{
-    FixedChunkPQTable,
-    pq::{distance::multi, generate_pq_data_from_pivots_from_membuf},
+use crate::{
+    model::{FixedChunkPQTable, pq::distance::multi},
+    utils::BridgeErr,
 };
 
 /// The discriminant type for PQ vector versions.
@@ -155,17 +156,9 @@ impl TestMultiPQProviderAsync {
         };
 
         let mut quant_vector: Vec<u8> = vec![0; table.get_num_chunks()];
-        if generate_pq_data_from_pivots_from_membuf(
-            &vector_f32,
-            table.get_pq_table(),
-            table.get_num_centers(),
-            table.get_chunk_offsets(),
-            &mut quant_vector,
-        )
-        .is_err()
-        {
-            return Err(ANNError::message("Error in generating PQ data."));
-        }
+        table
+            .compress_into(vector_f32.as_slice(), &mut quant_vector)
+            .bridge_err()?;
 
         let new = Arc::new(VersionedPQVector::new(quant_vector, version));
         self.quant_vectors[id].swap(new);

diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs

@@ -229,7 +229,7 @@ impl FastMemoryQuantVectorProviderAsync {
 
     /// Load `self` from a pivots file and data file.
     ///
-    /// The pivots file follows the format in [`storage::PQStorage::load_pq_pivots_bin`] and
+    /// The pivots file follows the format in [`storage::PQStorage::load_pivots`] and
     /// the compressed code is saved in a canonical `.bin` format.
     ///
     /// See also: [`storage::bin::load_from_bin`].
@@ -245,8 +245,15 @@ impl FastMemoryQuantVectorProviderAsync {
             // We can use that information to load the pivots, then finish the rest
             // of initialization.
             let pq_storage = storage::PQStorage::new(pivots, data, None);
-            let table = pq_storage.load_pq_pivots_bin(pivots, pq_bytes, provider)?;
-            Ok(Self::new(metric, num_points, table))
+            let table = pq_storage.load_pivots(provider)?;
+            if table.nchunks() != pq_bytes {
+                return Err(ANNError::message(format!(
+                    "PQ pivot table mismatch: file has {} chunks but expected {} chunks.",
+                    table.nchunks(),
+                    pq_bytes
+                )));
+            }
+            Ok(Self::new(metric, num_points, table.into()))
         })
     }
 

diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs

@@ -162,7 +162,7 @@ impl MemoryQuantVectorProviderAsync {
 
     /// Load `self` from a pivots file and data file.
     ///
-    /// The pivots file follows the format in [`storage::PQStorage::load_pq_pivots_bin`] and
+    /// The pivots file follows the format in [`storage::PQStorage::load_pivots`] and
     /// the compressed code is saved in a canonical `.bin` format.
     ///
     /// See also: [`storage::bin::load_from_bin`].
@@ -178,8 +178,15 @@ impl MemoryQuantVectorProviderAsync {
             // We can use that information to load the pivots, then finish the rest
             // of initialization.
             let pq_storage = storage::PQStorage::new(pivots, data, None);
-            let table = pq_storage.load_pq_pivots_bin(pivots, pq_bytes, provider)?;
-            Ok(Self::new(metric, num_points, table))
+            let table = pq_storage.load_pivots(provider)?;
+            if table.nchunks() != pq_bytes {
+                return Err(ANNError::message(format!(
+                    "PQ pivot table mismatch: file has {} chunks but expected {} chunks.",
+                    table.nchunks(),
+                    pq_bytes
+                )));
+            }
+            Ok(Self::new(metric, num_points, table.into()))
         })
     }
 

diskann-providers/src/model/pq/fixed_chunk_pq_table.rs

@@ -430,6 +430,12 @@ impl FixedChunkPQTable {
     }
 }
 
+impl From<BasicTable> for FixedChunkPQTable {
+    fn from(table: BasicTable) -> Self {
+        Self { table }
+    }
+}
+
 // This goes against Rust's Orphan rule, so we cannot implement it directly.
 // However, we can use a wrapper type to implement the conversion.
 // This is a workaround to allow the conversion from `product::TableCompressionError` to
@@ -675,9 +681,8 @@ pub fn compute_pq_distance_for_pq_coordinates(
 mod fixed_chunk_pq_table_test {
     use core::ops::Range;
 
-    use crate::storage::{StorageReadProvider, VirtualStorageProvider};
+    use crate::storage::{PQStorage, VirtualStorageProvider};
     use approx::assert_relative_eq;
-    use diskann::error::ErrorContext;
     use diskann_utils::test_data_root;
     use diskann_vector::{
         PureDistanceFunction,
@@ -686,9 +691,18 @@ mod fixed_chunk_pq_table_test {
     use itertools::iproduct;
 
     use super::*;
-    use crate::{model::NUM_PQ_CENTROIDS, utils::read_bin_from};
+    use crate::model::NUM_PQ_CENTROIDS;
 
     const DIM: usize = 128;
+    const PQ_PIVOTS_PATH: &str = "/sift/siftsmall_learn_pq_pivots.bin";
+
+    fn load_test_pivots() -> FixedChunkPQTable {
+        let storage_provider = VirtualStorageProvider::new_overlay(test_data_root());
+        PQStorage::new(PQ_PIVOTS_PATH, "", None)
+            .load_pivots(&storage_provider)
+            .unwrap()
+            .into()
+    }
 
     #[test]
     fn constructor_errors() {
@@ -811,14 +825,8 @@ mod fixed_chunk_pq_table_test {
 
     #[test]
     fn load_pivot_test() {
-        let storage_provider = VirtualStorageProvider::new_overlay(test_data_root());
-        let pq_pivots_path: &str = "/sift/siftsmall_learn_pq_pivots.bin";
-        let (dim, pq_table, chunk_offsets) =
-            load_pq_pivots_bin(pq_pivots_path, &1, &storage_provider).unwrap();
-        let fixed_chunk_pq_table =
-            FixedChunkPQTable::new(dim, pq_table.into(), chunk_offsets.into()).unwrap();
+        let fixed_chunk_pq_table = load_test_pivots();
 
-        assert_eq!(dim, DIM);
         assert_eq!(fixed_chunk_pq_table.table.dim(), DIM);
         assert_eq!(fixed_chunk_pq_table.table.ncenters(), NUM_PQ_CENTROIDS);
 
@@ -838,14 +846,7 @@ mod fixed_chunk_pq_table_test {
 
     #[test]
     fn calculate_distances_tests() {
-        let storage_provider = VirtualStorageProvider::new_overlay(test_data_root());
-
-        let pq_pivots_path: &str = "/sift/siftsmall_learn_pq_pivots.bin";
-
-        let (dim, pq_table, chunk_offsets) =
-            load_pq_pivots_bin(pq_pivots_path, &1, &storage_provider).unwrap();
-        let fixed_chunk_pq_table =
-            FixedChunkPQTable::new(dim, pq_table.into(), chunk_offsets.into()).unwrap();
+        let fixed_chunk_pq_table = load_test_pivots();
 
         let query_vec: Vec<f32> = vec![
             32.39f32, 78.57f32, 50.32f32, 80.46f32, 6.47f32, 69.76f32, 94.2f32, 83.36f32, 5.8f32,
@@ -994,75 +995,6 @@ mod fixed_chunk_pq_table_test {
         }
     }
 
-    type LoadPQPivotResult = (usize, Vec<f32>, Vec<usize>);
-    fn load_pq_pivots_bin<StorageProvider: StorageReadProvider>(
-        pq_pivots_path: &str,
-        num_pq_chunks: &usize,
-        storage_provider: &StorageProvider,
-    ) -> ANNResult<LoadPQPivotResult> {
-        let mut reader = storage_provider
-            .open_reader(pq_pivots_path)
-            .with_context(|| format!("ERROR: Opening PQ k-means pivot file {}", pq_pivots_path))?;
-
-        let offsets = read_bin_from::<u64>(&mut reader, 0)?;
-        if offsets.nrows() != 4 {
-            return Err(ANNError::message(format!(
-                "Error reading pq_pivots file {}. \
-                 Offsets don't contain correct metadata, \
-                 # offsets = {}, but expecting 4.",
-                pq_pivots_path,
-                offsets.nrows()
-            )));
-        }
-        let file_offset_data = offsets.map(|x| x.into_usize());
-
-        let mut pivots = read_bin_from::<f32>(&mut reader, file_offset_data[(0, 0)])?;
-
-        if pivots.nrows() != NUM_PQ_CENTROIDS {
-            return Err(ANNError::message(format!(
-                "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.",
-                pq_pivots_path,
-                pivots.nrows(),
-                NUM_PQ_CENTROIDS
-            )));
-        }
-        let dim = pivots.ncols();
-
-        let centroids = read_bin_from::<f32>(&mut reader, file_offset_data[(1, 0)])?;
-        if centroids.nrows() != dim || centroids.ncols() != 1 {
-            return Err(ANNError::message(format!(
-                "Error reading pq_pivots file {}. file_dim = {}, \
-                 file_cols = {} but expecting {} entries in 1 dimension.",
-                pq_pivots_path,
-                centroids.nrows(),
-                centroids.ncols(),
-                dim
-            )));
-        }
-
-        pivots.row_iter_mut().for_each(|row| {
-            std::iter::zip(row.iter_mut(), centroids.as_slice().iter()).for_each(|(p, c)| *p += *c);
-        });
-
-        let chunk_offsets_m = read_bin_from::<u32>(&mut reader, file_offset_data[(2, 0)])?;
-        if chunk_offsets_m.nrows() != num_pq_chunks + 1 || chunk_offsets_m.ncols() != 1 {
-            return Err(ANNError::message(format!(
-                "Error reading pq_pivots file at chunk offsets; \
-                 file has nr={}, nc={} but expecting nr={} and nc=1.",
-                chunk_offsets_m.nrows(),
-                chunk_offsets_m.ncols(),
-                num_pq_chunks + 1
-            )));
-        }
-        let chunk_offsets = chunk_offsets_m.map(|x| x.into_usize());
-
-        Ok((
-            dim,
-            pivots.into_inner().into_vec(),
-            chunk_offsets.into_inner().into_vec(),
-        ))
-    }
-
     #[test]
     fn test_populate_chunk_distances() {
         let dim = 8;