diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index c5297ae9c..132c88026 100644 --- a/diskann-disk/src/storage/quant/pq/pq_generation.rs +++ b/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,27 @@ 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.pq_storage.get_pivot_data_path(), + Some(num_chunks), + 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.ncenters() != context.num_centers || table.dim() != full_dim { + return Err(diskann_error!( + ErrorKind::PQError, + "PQ pivot table mismatch: file has {} centers in {} dimensions but expected {} centers in {} dimensions.", + table.ncenters(), + table.dim(), + 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, diff --git a/diskann-providers/src/model/graph/provider/async_/experimental/multi_pq_async.rs b/diskann-providers/src/model/graph/provider/async_/experimental/multi_pq_async.rs index 012253507..d22b77802 100644 --- a/diskann-providers/src/model/graph/provider/async_/experimental/multi_pq_async.rs +++ b/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 = 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); diff --git a/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs b/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs index 6dbeeb89c..e12e1bc0a 100644 --- a/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs +++ b/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs @@ -144,6 +144,11 @@ impl FixedChunkPQTable { Ok(Self { table }) } + /// Wrap an already-validated basic PQ table. + pub fn from_basic_table(table: BasicTable) -> Self { + Self { table } + } + /// Get chunk number. pub fn get_num_chunks(&self) -> usize { self.table.nchunks() @@ -675,9 +680,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 +690,17 @@ 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_pq_pivots_bin(PQ_PIVOTS_PATH, 1, &storage_provider) + .unwrap() + } #[test] fn constructor_errors() { @@ -811,14 +823,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 +844,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 = vec![ 32.39f32, 78.57f32, 50.32f32, 80.46f32, 6.47f32, 69.76f32, 94.2f32, 83.36f32, 5.8f32, @@ -994,75 +993,6 @@ mod fixed_chunk_pq_table_test { } } - type LoadPQPivotResult = (usize, Vec, Vec); - fn load_pq_pivots_bin( - pq_pivots_path: &str, - num_pq_chunks: &usize, - storage_provider: &StorageProvider, - ) -> ANNResult { - 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::(&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::(&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::(&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::(&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; diff --git a/diskann-providers/src/model/pq/pq_construction.rs b/diskann-providers/src/model/pq/pq_construction.rs index 68fc1fed4..a95e43fa3 100644 --- a/diskann-providers/src/model/pq/pq_construction.rs +++ b/diskann-providers/src/model/pq/pq_construction.rs @@ -340,29 +340,29 @@ where let (num_points, dim) = Metadata::read(uncompressed_data_reader)?.into_dims(); - let mut full_pivot_data: Vec; - let centroid: Vec; - let chunk_offsets: Vec; let full_dim: usize; + let table; if !pq_storage.pivot_data_exist(storage_provider) { return Err(ANNError::message("ERROR: PQ k-means pivot file not found.")); } else { (_, full_dim) = pq_storage.read_existing_pivot_metadata(storage_provider)?; - (full_pivot_data, centroid, chunk_offsets) = pq_storage.load_existing_pivot_data( - &num_pq_chunks, - &num_centers, - &full_dim, + table = pq_storage.load_pivots( + pq_storage.get_pivot_data_path(), + Some(num_pq_chunks), storage_provider, )?; - } - // Instead of subtracting the center from each data set component, we instead - // add it to each center. - let mut full_pivot_data_mat = - MutMatrixView::try_from(full_pivot_data.as_mut_slice(), num_centers, full_dim) - .bridge_err()?; - accum_row_inplace(full_pivot_data_mat.as_mut_view(), centroid.as_slice()); + if table.ncenters() != num_centers || table.dim() != full_dim { + return Err(ANNError::message(format!( + "PQ pivot table mismatch: file has {} centers in {} dimensions but expected {} centers in {} dimensions.", + table.ncenters(), + table.dim(), + num_centers, + full_dim + ))); + } + } pq_storage.write_compressed_pivot_metadata::( num_points, @@ -387,13 +387,8 @@ where ))?; // The compression table. - let table = TransposedTable::from_parts( - full_pivot_data_mat.as_view(), - diskann_quantization::views::ChunkOffsetsView::new(&chunk_offsets) - .bridge_err()? - .to_owned(), - ) - .map_err(ANNError::new)?; + let table = TransposedTable::from_parts(table.view_pivots(), table.view_offsets().to_owned()) + .map_err(|err| ANNError::message(diskann_quantization::error::format(&err)))?; let mut buffer = vec![0.0; full_dim * block_size]; @@ -546,17 +541,20 @@ pub fn generate_pq_data_from_pivots_from_membuf_batch return Err(ANNError::message("Error: Invalid PQ buffer input size.")); } + let table = BasicTableView::new( + MatrixView::try_from(pivot_data, parameters.num_centers(), dim).bridge_err()?, + ChunkOffsetsView::new(offsets).bridge_err()?, + ) + .map_err(|err| ANNError::message(diskann_quantization::error::format(&err)))?; + pq_out .par_chunks_mut(num_pq_chunks) .zip(vector_data.par_chunks(dim)) - .try_for_each_in_pool(pool, |(pq_slice, vector_slice)| { - generate_pq_data_from_pivots_from_membuf( - vector_slice, - pivot_data, - parameters.num_centers(), - offsets, - pq_slice, - ) + .try_for_each_in_pool(pool, |(pq_slice, vector)| { + let data = vector.iter().map(|x| (*x).into()).collect::>(); + table + .compress_into(data.as_slice(), pq_slice) + .map_err(ANNError::new) }) } diff --git a/diskann-providers/src/storage/pq_storage.rs b/diskann-providers/src/storage/pq_storage.rs index 5b6449ef8..7bc06adf0 100644 --- a/diskann-providers/src/storage/pq_storage.rs +++ b/diskann-providers/src/storage/pq_storage.rs @@ -9,6 +9,7 @@ use diskann::{ ANNError, ANNResult, utils::{IntoUsize, VectorRepr}, }; +use diskann_quantization::{product::BasicTable, views::ChunkOffsetsBase}; use diskann_utils::{ io::{Metadata, read_bin, write_bin}, views::{Matrix, MatrixView}, @@ -179,66 +180,51 @@ impl PQStorage { where Storage: StorageReadProvider, { - // Load file offset data. File layout: offset table(4*1) -> pivot data(num_centers*dim) -> centroid(dim*1) -> chunk offsets(num_chunks+1*1) - let reader = &mut storage_provider.open_reader(&self.pivot_data_path)?; - - let offsets = read_bin_from::(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.", - &self.pivot_data_path, - offsets.nrows() - ))); - } - let file_offset_data = offsets.map(|x| x.into_usize()); - - info!(" Offset data: {:?}", file_offset_data.as_slice()); - - let pivots = read_bin_from::(reader, file_offset_data[(0, 0)])?; - if pivots.nrows() != *num_centers || pivots.ncols() != *dim { - return Err(ANNError::message(format!( - "Error reading pq_pivots file {}. file_num_centers = {}, \ - file_dim = {} but expecting {} centers in {} dimensions.", - &self.pivot_data_path, - pivots.nrows(), - pivots.ncols(), - num_centers, - dim - ))); - } - - let centroid_m = read_bin_from::(reader, file_offset_data[(1, 0)])?; - if centroid_m.nrows() != *dim || centroid_m.ncols() != 1 { - return Err(ANNError::message(format!( - "Error reading pq_pivots file {}. file_dim = {}, \ - file_cols = {} but expecting {} entries in 1 dimension.", - &self.pivot_data_path, - centroid_m.nrows(), - centroid_m.ncols(), - dim - ))); - } - - let chunk_offsets_m = read_bin_from::(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()); + let (pivots, centroid, chunk_offsets) = self.read_pivot_file( + &self.pivot_data_path, + Some(*num_pq_chunks), + Some(*num_centers), + Some(*dim), + storage_provider, + )?; + let table = + Self::pivot_data_into_basic_table(&self.pivot_data_path, pivots, chunk_offsets)?; + let (pivots, chunk_offsets) = table.into_parts(); Ok(( pivots.into_inner().into_vec(), - centroid_m.into_inner().into_vec(), + centroid.into_inner().into_vec(), chunk_offsets.into_inner().into_vec(), )) } + /// Load the effective PQ pivot table from a pivot file. + /// + /// If `expected_num_pq_chunks` is `None`, the chunk count is inferred from the + /// file. The loader verifies the pivot file layout and folds any stored legacy + /// centroid into the pivots. `BasicTable::new` validates the resulting table + /// invariants, including chunk-offset bounds and monotonicity. + pub fn load_pivots( + &self, + pq_pivots: &str, + expected_num_pq_chunks: Option, + storage_provider: &Storage, + ) -> ANNResult { + let (mut pivots, centroid, chunk_offsets) = self.read_pivot_file( + pq_pivots, + expected_num_pq_chunks, + None, + None, + storage_provider, + )?; + + if centroid.as_slice().iter().any(|c| *c != 0.0) { + accum_row_inplace(pivots.as_mut_view(), centroid.as_slice()) + } + + Self::pivot_data_into_basic_table(pq_pivots, pivots, chunk_offsets) + } + /// Load the compressed pq dataset from file. /// /// Returns a `num_points × num_pq_chunks` matrix of u8 codes. @@ -281,26 +267,63 @@ impl PQStorage { num_pq_chunks: usize, storage_provider: &Storage, ) -> ANNResult { + let table = self.load_pivots( + pq_pivots, + (num_pq_chunks != 0).then_some(num_pq_chunks), + storage_provider, + )?; + Ok(FixedChunkPQTable::from_basic_table(table)) + } + + fn read_pivot_file( + &self, + pq_pivots: &str, + expected_num_pq_chunks: Option, + expected_num_centers: Option, + expected_dim: Option, + storage_provider: &Storage, + ) -> ANNResult<(Matrix, Matrix, Matrix)> { if !storage_provider.exists(pq_pivots) { - return Err(ANNError::message("ERROR: PQ k-means pivot file not found.")); + return Err(ANNError::message(format!( + "ERROR: PQ k-means pivot file not found: {pq_pivots}." + ))); } info!("Loading PQ pivots from {}...", pq_pivots); let mut reader = storage_provider.open_reader(pq_pivots)?; + + // File layout: offset table(4*1) -> pivot data(num_centers*dim) -> + // centroid(dim*1) -> chunk offsets(num_chunks+1*1). + // + // The expected values here are file-format checks only. Structural PQ table + // invariants, such as chunk-offset monotonicity and pivot/offset dimension + // agreement, are validated by `ChunkOffsetsBase` and `BasicTable`. let offsets = read_bin_from::(&mut reader, 0)?; - if offsets.nrows() != 4 { + if offsets.nrows() != 4 || offsets.ncols() != 1 { return Err(ANNError::message(format!( "Error reading pq_pivots file {}. Offsets don't contain correct metadata, \ - # offsets = {}, but expecting 4.", + file has nr={}, nc={}, but expecting nr=4 and nc=1.", pq_pivots, - offsets.nrows() + offsets.nrows(), + offsets.ncols() ))); } let file_offset_data = offsets.map(|x| x.into_usize()); - let mut pivots = read_bin_from::(&mut reader, file_offset_data[(0, 0)])?; - if pivots.nrows() > NUM_PQ_CENTROIDS { + info!(" Offset data: {:?}", file_offset_data.as_slice()); + + let pivots = read_bin_from::(&mut reader, file_offset_data[(0, 0)])?; + if let Some(num_centers) = expected_num_centers { + if pivots.nrows() != num_centers { + return Err(ANNError::message(format!( + "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.", + pq_pivots, + pivots.nrows(), + num_centers + ))); + } + } else if pivots.nrows() > NUM_PQ_CENTROIDS { return Err(ANNError::message(format!( "Error reading pq_pivots file {}. file_num_centers = {}, but expecting {} centers.", pq_pivots, @@ -308,42 +331,71 @@ impl PQStorage { NUM_PQ_CENTROIDS ))); } - let dim = pivots.ncols(); - let centroids = read_bin_from::(&mut reader, file_offset_data[(1, 0)])?; - if centroids.nrows() != dim || centroids.ncols() != 1 { + if let Some(dim) = expected_dim + && pivots.ncols() != dim + { + return Err(ANNError::message(format!( + "Error reading pq_pivots file {}. file_dim = {} but expecting {} dimensions.", + pq_pivots, + pivots.ncols(), + dim + ))); + } + + let centroid = read_bin_from::(&mut reader, file_offset_data[(1, 0)])?; + if centroid.nrows() != pivots.ncols() || centroid.ncols() != 1 { return Err(ANNError::message(format!( "Error reading pq_pivots file {}. file_dim = {}, file_cols = {} \ but expecting {} entries in 1 dimension.", pq_pivots, - centroids.nrows(), - centroids.ncols(), - dim + centroid.nrows(), + centroid.ncols(), + pivots.ncols() ))); } let chunk_offsets_m = read_bin_from::(&mut reader, file_offset_data[(2, 0)])?; - if (chunk_offsets_m.nrows() != num_pq_chunks + 1 && num_pq_chunks as u32 != 0) - || chunk_offsets_m.ncols() != 1 + if let Some(num_pq_chunks) = expected_num_pq_chunks + && chunk_offsets_m.nrows() != num_pq_chunks + 1 { return Err(ANNError::message(format!( - "Error reading pq_pivots file at chunk offsets; file has nr={}, nc={} \ - but expecting nr={} and nc=1. The expected num_pq_chunks should be \ - passed as 0 if we want to infer.", + "Error reading pq_pivots file at chunk offsets; file has nr={}, but expecting nr={}.", chunk_offsets_m.nrows(), - chunk_offsets_m.ncols(), num_pq_chunks + 1 ))); } + if chunk_offsets_m.ncols() != 1 { + return Err(ANNError::message(format!( + "Error reading pq_pivots file at chunk offsets; file has nc={}, but expecting nc=1.", + chunk_offsets_m.ncols() + ))); + } let chunk_offsets = chunk_offsets_m.map(|x| x.into_usize()); - // If the centroid is non-zero, we need to add it to the pivots to restore the - // numeric behavior. - if centroids.as_slice().iter().any(|c| *c != 0.0) { - accum_row_inplace(pivots.as_mut_view(), centroids.as_slice()) - } + Ok((pivots, centroid, chunk_offsets)) + } + + fn pivot_data_into_basic_table( + pq_pivots: &str, + pivots: Matrix, + chunk_offsets: Matrix, + ) -> ANNResult { + let offsets = ChunkOffsetsBase::new(chunk_offsets.into_inner()).map_err(|err| { + ANNError::message(format!( + "Error constructing chunk offsets from pq_pivots file {}: {}", + pq_pivots, + diskann_quantization::error::format(&err) + )) + })?; - FixedChunkPQTable::new(dim, pivots.into_inner(), chunk_offsets.into_inner()) + BasicTable::new(pivots, offsets).map_err(|err| { + ANNError::message(format!( + "Error constructing PQ table from pq_pivots file {}: {}", + pq_pivots, + diskann_quantization::error::format(&err) + )) + }) } /// streams data from the file, and samples each vector with probability p_val @@ -377,6 +429,10 @@ impl PQStorage { pub fn get_compressed_data_path(&self) -> &str { &self.compressed_data_path } + + pub fn get_pivot_data_path(&self) -> &str { + &self.pivot_data_path + } } #[cfg(test)] @@ -393,6 +449,27 @@ mod pq_storage_tests { const PQ_PIVOT_PATH: &str = "/sift/siftsmall_learn_pq_pivots.bin"; const PQ_COMPRESSED_PATH: &str = "/sift/empty_pq_compressed.bin"; + fn write_test_pivots( + storage_provider: &VirtualStorageProvider, + pivot_path: &str, + num_centers: usize, + dim: usize, + centroid: Option<&[f32]>, + chunk_offsets: &[usize], + ) { + let pivots: Vec = (0..num_centers * dim).map(|i| i as f32).collect(); + PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .write_pivot_data( + &pivots, + centroid, + chunk_offsets, + num_centers, + dim, + storage_provider, + ) + .unwrap(); + } + #[test] fn new_test() { PQStorage::new(PQ_PIVOT_PATH, PQ_COMPRESSED_PATH, Some(DATA_FILE)); @@ -510,6 +587,231 @@ mod pq_storage_tests { assert_eq!(loaded_pivots, table.view_pivots().as_slice()); } + #[test] + fn load_pq_pivots_infer_chunks_loads_without_expected_count() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/infer_chunk_count_pivots.bin"; + + let num_centers = 3; + let dim = 4; + let pivots: Vec = (0..num_centers * dim).map(|i| i as f32).collect(); + let chunk_offsets = vec![0, 2, dim]; + + let pq_storage = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None); + pq_storage + .write_pivot_data( + &pivots, + None, + &chunk_offsets, + num_centers, + dim, + &storage_provider, + ) + .unwrap(); + + let table = pq_storage + .load_pivots(pivot_path, None, &storage_provider) + .unwrap(); + + assert_eq!(table.view_pivots().as_slice(), pivots); + } + + #[test] + fn load_pq_pivots_zero_chunk_count_infers_from_file() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/zero_chunk_count_pivots.bin"; + let pivots: Vec = (0..12).map(|i| i as f32).collect(); + + PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .write_pivot_data(&pivots, None, &[0, 2, 4], 3, 4, &storage_provider) + .unwrap(); + + let table = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_pq_pivots_bin(pivot_path, 0, &storage_provider) + .unwrap(); + + assert_eq!(table.view_pivots().as_slice(), pivots); + } + + #[test] + fn load_pivot_data_rejects_mismatched_shape() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/mismatched_shape_pivots.bin"; + + write_test_pivots(&storage_provider, pivot_path, 3, 4, None, &[0, 2, 4]); + let pq_storage = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None); + + assert!( + pq_storage + .load_existing_pivot_data(&2, &4, &4, &storage_provider) + .is_err() + ); + assert!( + pq_storage + .load_existing_pivot_data(&2, &3, &5, &storage_provider) + .is_err() + ); + } + + #[test] + fn load_pivot_data_rejects_invalid_centroid_and_chunk_count() { + let storage_provider = VirtualStorageProvider::new_memory(); + + let wrong_centroid_path = "/wrong_centroid_pivots.bin"; + write_test_pivots( + &storage_provider, + wrong_centroid_path, + 3, + 4, + Some(&[1.0, 2.0, 3.0]), + &[0, 2, 4], + ); + assert!( + PQStorage::new(wrong_centroid_path, PQ_COMPRESSED_PATH, None) + .load_existing_pivot_data(&2, &3, &4, &storage_provider) + .is_err() + ); + + let wrong_count_path = "/wrong_chunk_count_pivots.bin"; + write_test_pivots(&storage_provider, wrong_count_path, 3, 4, None, &[0, 4]); + assert!( + PQStorage::new(wrong_count_path, PQ_COMPRESSED_PATH, None) + .load_existing_pivot_data(&2, &3, &4, &storage_provider) + .is_err() + ); + } + + #[test] + fn load_pivot_data_rejects_invalid_chunk_bounds() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/legacy_chunk_bounds_pivots.bin"; + + write_test_pivots(&storage_provider, pivot_path, 3, 4, None, &[1, 4]); + + let err = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_existing_pivot_data(&1, &3, &4, &storage_provider) + .unwrap_err(); + + assert!(err.to_string().contains("offsets must begin at 0")); + } + + #[test] + fn load_pivot_data_rejects_chunk_offsets_dim_mismatch() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/chunk_offsets_dim_mismatch_pivots.bin"; + + write_test_pivots(&storage_provider, pivot_path, 3, 4, None, &[0, 2, 3]); + + let err = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_existing_pivot_data(&2, &3, &4, &storage_provider) + .unwrap_err(); + + assert!(err.to_string().contains("offsets expect 3")); + } + + #[test] + fn load_pq_pivots_infer_rejects_invalid_chunk_bounds() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/infer_wrong_chunk_bounds_pivots.bin"; + + write_test_pivots(&storage_provider, pivot_path, 3, 4, None, &[1, 4]); + + let err = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_pivots(pivot_path, None, &storage_provider) + .unwrap_err(); + + assert!(err.to_string().contains("offsets must begin at 0")); + } + + #[test] + fn load_pq_pivots_reports_chunk_offset_column_mismatch() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/wrong_chunk_offset_columns_pivots.bin"; + + { + let mut writer = storage_provider.create_for_write(pivot_path).unwrap(); + let mut cumul_bytes = [0usize; 4]; + cumul_bytes[0] = METADATA_SIZE; + + writer.seek(SeekFrom::Start(cumul_bytes[0] as u64)).unwrap(); + + let pivots = [0.0, 1.0, 2.0, 3.0]; + cumul_bytes[1] = cumul_bytes[0] + + write_bin( + MatrixView::try_from(pivots.as_slice(), 2, 2).unwrap(), + &mut writer, + ) + .unwrap(); + + let centroid = [0.0, 0.0]; + cumul_bytes[2] = cumul_bytes[1] + + write_bin(MatrixView::column_vector(centroid.as_slice()), &mut writer).unwrap(); + + let chunk_offsets = [0_u32, 2_u32]; + cumul_bytes[3] = cumul_bytes[2] + + write_bin( + MatrixView::try_from(chunk_offsets.as_slice(), 1, 2).unwrap(), + &mut writer, + ) + .unwrap(); + + let offsets: Vec = cumul_bytes.iter().map(|&offset| offset as u64).collect(); + write_bin_from( + MatrixView::column_vector(offsets.as_slice()), + &mut writer, + 0, + ) + .unwrap(); + } + + let err = PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_pivots(pivot_path, None, &storage_provider) + .unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("file has nc=2, but expecting nc=1")); + assert!(!message.contains("expecting nr=0")); + } + + #[test] + fn load_pq_pivots_rejects_too_many_centers() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/too_many_centers_pivots.bin"; + + write_test_pivots( + &storage_provider, + pivot_path, + NUM_PQ_CENTROIDS + 1, + 1, + None, + &[0, 1], + ); + + assert!( + PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_pivots(pivot_path, None, &storage_provider) + .is_err() + ); + } + + #[test] + fn load_pivot_data_rejects_malformed_offset_table() { + let storage_provider = VirtualStorageProvider::new_memory(); + let pivot_path = "/malformed_offsets_pivots.bin"; + + { + let mut writer = storage_provider.create_for_write(pivot_path).unwrap(); + let offsets = [METADATA_SIZE as u64, 0, 0]; + write_bin(MatrixView::column_vector(offsets.as_slice()), &mut writer).unwrap(); + } + + assert!( + PQStorage::new(pivot_path, PQ_COMPRESSED_PATH, None) + .load_pivots(pivot_path, None, &storage_provider) + .is_err() + ); + } + /// Write pivot data with a non-zero centroid, read it back, and verify that /// folding the centroid via `accum_row_inplace` produces the expected /// adjusted pivots. diff --git a/diskann-quantization/src/product/tables/basic.rs b/diskann-quantization/src/product/tables/basic.rs index 9469fa58d..822940645 100644 --- a/diskann-quantization/src/product/tables/basic.rs +++ b/diskann-quantization/src/product/tables/basic.rs @@ -5,7 +5,7 @@ use crate::traits::CompressInto; use crate::views::{ChunkOffsetsBase, ChunkOffsetsView}; -use diskann_utils::views::{DenseData, MatrixBase, MatrixView}; +use diskann_utils::views::{DenseData, MatrixBase, MatrixView, MutDenseData, MutMatrixView}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use thiserror::Error; @@ -86,6 +86,14 @@ where self.pivots.as_view() } + /// Return a mutable view over the pivot table. + pub fn view_pivots_mut(&mut self) -> MutMatrixView<'_, f32> + where + T: MutDenseData, + { + self.pivots.as_mut_view() + } + /// Return a view over the schema offsets. pub fn view_offsets(&self) -> ChunkOffsetsView<'_> { self.offsets.as_view() @@ -105,6 +113,11 @@ where pub fn dim(&self) -> usize { self.pivots.ncols() } + + /// Consume this table and return the underlying pivots and offsets. + pub fn into_parts(self) -> (MatrixBase, ChunkOffsetsBase) { + (self.pivots, self.offsets) + } } #[derive(Error, Debug)] diff --git a/diskann-quantization/src/views.rs b/diskann-quantization/src/views.rs index 04c4a0953..03ed01223 100644 --- a/diskann-quantization/src/views.rs +++ b/diskann-quantization/src/views.rs @@ -211,6 +211,11 @@ where pub fn as_slice(&self) -> &[usize] { self.offsets.as_slice() } + + /// Consume the offsets, returning the inner representation. + pub fn into_inner(self) -> T { + self.offsets + } } pub type ChunkOffsetsView<'a> = ChunkOffsetsBase<&'a [usize]>;