diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index ec9ade194b..b2c0b47bac 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -157,14 +157,16 @@ where metric: self.index_configuration.dist_metric, }; + PQGeneration::::generate_pivots(&quantizer_context)?; + let generator = QuantDataGenerator::< Data::VectorDataType, PQGeneration, >::new( self.index_writer.get_dataset_file(), self.pq_storage.get_compressed_data_path().into(), - quantizer_context, - ); + &quantizer_context, + )?; generator.generate_data( storage_provider, pool, diff --git a/diskann-disk/src/storage/quant/compressor.rs b/diskann-disk/src/storage/quant/compressor.rs index 7b55634dff..acddd89940 100644 --- a/diskann-disk/src/storage/quant/compressor.rs +++ b/diskann-disk/src/storage/quant/compressor.rs @@ -20,8 +20,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; /// - [`CompressorContext`]: An overloadable type that provides initialization parameters for the compressor /// /// # Methods -/// - `new`: Constructs a compressor with the provided context. -/// - `generate`: Generates any data needed before compression. +/// - `new`: Constructs a new compressor instance with the provided context. /// - `compress`: Compresses a batch of vectors into the output buffer. /// - `compressed_bytes`: Returns the size in bytes of each compressed vector pub trait QuantCompressor: Sized + Sync @@ -30,8 +29,7 @@ where { type CompressorContext; - fn new(context: Self::CompressorContext) -> Self; - fn generate(&self) -> ANNResult<()>; + fn new(context: &Self::CompressorContext) -> ANNResult; fn compress(&self, vector: MatrixView, output: MutMatrixView) -> ANNResult<()>; fn compressed_bytes(&self) -> usize; } diff --git a/diskann-disk/src/storage/quant/generator.rs b/diskann-disk/src/storage/quant/generator.rs index 79913fd241..36af775bfb 100644 --- a/diskann-disk/src/storage/quant/generator.rs +++ b/diskann-disk/src/storage/quant/generator.rs @@ -44,15 +44,15 @@ where pub fn new( data_path: String, compressed_data_path: String, - quantizer_context: Q::CompressorContext, - ) -> Self { - let quantizer = Q::new(quantizer_context); - Self { + quantizer_context: &Q::CompressorContext, + ) -> ANNResult { + let quantizer = Q::new(quantizer_context)?; + Ok(Self { data_path, compressed_data_path, quantizer, phantom: PhantomData, - } + }) } /// This method reads the source data file, processes vectors in batches, compresses them @@ -93,7 +93,6 @@ where )); } - self.quantizer.generate()?; let compressed_path = self.compressed_data_path.as_str(); if storage_provider.exists(compressed_path) { @@ -220,12 +219,8 @@ mod generator_tests { impl QuantCompressor for DummyCompressor { type CompressorContext = u32; - fn new(context: Self::CompressorContext) -> Self { - Self::new(context) - } - - fn generate(&self) -> ANNResult<()> { - Ok(()) + fn new(context: &Self::CompressorContext) -> ANNResult { + Ok(Self::new(*context)) } fn compress( @@ -289,8 +284,12 @@ mod generator_tests { max_block_size: usize, ) -> (QuantDataGenerator, ANNResult<()>) { let pool: diskann_providers::utils::RayonThreadPool = create_thread_pool_for_test(); - let generator = - QuantDataGenerator::::new(data_path, compressed_path, output_dim); + let generator = QuantDataGenerator::::new( + data_path, + compressed_path, + &output_dim, + ) + .unwrap(); let result = generator.generate_data(storage_provider, pool.as_ref(), max_block_size); (generator, result) } diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index 3e64fbf43c..5aabf2254c 100644 --- a/diskann-disk/src/storage/quant/pq/pq_generation.rs +++ b/diskann-disk/src/storage/quant/pq/pq_generation.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use std::{marker::PhantomData, sync::OnceLock, time::Instant}; +use std::{marker::PhantomData, time::Instant}; use diskann::utils::VectorRepr; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; @@ -46,10 +46,10 @@ where T: VectorRepr, Storage: StorageReadProvider + StorageWriteProvider + 'a, { - context: PQGenerationContext<'a, Storage>, - table: OnceLock, + table: TransposedTable, num_chunks: usize, phantom_data: PhantomData, + phantom_storage: PhantomData<&'a Storage>, } impl<'a, T, Storage> PQGeneration<'a, T, Storage> @@ -120,28 +120,13 @@ where { type CompressorContext = PQGenerationContext<'a, Storage>; - fn new(context: Self::CompressorContext) -> Self { - let num_chunks = context.num_chunks; - Self { - context, - table: OnceLock::new(), - num_chunks, - phantom_data: PhantomData, - } - } - - fn generate(&self) -> diskann::ANNResult<()> { - if self.table.get().is_some() { - return Ok(()); - } - - let context = &self.context; + fn new(context: &Self::CompressorContext) -> diskann::ANNResult { Self::generate_pivots(context)?; + let (_, full_dim) = context .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( @@ -168,11 +153,11 @@ where ) .map_err(|err| diskann_error!(ErrorKind::PQError, "{}", Format(err)))?; - self.table.set(table).map_err(|_| { - diskann_error!( - ErrorKind::PQError, - "PQ compressor was generated concurrently" - ) + Ok(Self { + table, + num_chunks, + phantom_data: PhantomData, + phantom_storage: PhantomData, }) } @@ -182,13 +167,6 @@ where output: MatrixBase<&mut [u8]>, ) -> Result<(), diskann::ANNError> { self.table - .get() - .ok_or_else(|| { - diskann_error!( - ErrorKind::PQError, - "PQ compressor must be generated before compression" - ) - })? .compress_into(vector, output) .map_err(|err| diskann_error!(ErrorKind::PQError, "{}", Format(err))) } @@ -294,24 +272,22 @@ mod pq_generation_tests { assert!(!storage_provider.exists(pivot_file_name)); - let compressor = PQGeneration::::new(context); - assert!(!storage_provider.exists(pivot_file_name)); - - let result = compressor.generate(); + let result = PQGeneration::::generate_pivots(&context); assert!(result.is_ok()); assert!(storage_provider.exists(pivot_file_name)); + let compressor = PQGeneration::::new(&context).unwrap(); + assert_eq!(compressor.num_chunks, num_chunks); assert_eq!(compressor.compressed_bytes(), num_chunks); - let table = compressor.table.get().unwrap(); - assert_eq!(table.dim(), dim); - assert_eq!(table.ncenters(), num_centers); - assert_eq!(table.nchunks(), num_chunks); + assert_eq!(compressor.table.dim(), dim); + assert_eq!(compressor.table.ncenters(), num_centers); + assert_eq!(compressor.table.nchunks(), num_chunks); } #[rstest] - fn generate_creates_missing_pivots() { + fn new_preserves_missing_pivot_generation_fallback() { let storage_provider = VirtualStorageProvider::new_memory(); storage_provider .filesystem() @@ -342,10 +318,8 @@ mod pq_generation_tests { Some(data_path), ); - let compressor = PQGeneration::::new(context); - let result = compressor.generate(); - - assert!(result.is_ok()); + let compressor = PQGeneration::::new(&context); + assert!(compressor.is_ok()); assert!(storage_provider.exists(pivot_file_name)); } @@ -370,14 +344,14 @@ mod pq_generation_tests { "".to_string(), None, ); - let compressor = PQGeneration::::new(context); - let result = compressor.generate(); + let compressor = PQGeneration::::new(&context); - if let Err(x) = result.as_ref() { + if let Err(x) = compressor.as_ref() { println!("Error creating compressor: {x}"); }; - assert!(result.is_ok()); + assert!(compressor.is_ok()); + let compressor = compressor.unwrap(); let data_matrix = read_bin::(&mut storage_provider.open_reader(TEST_PQ_DATA_PATH).unwrap()).unwrap(); @@ -423,7 +397,7 @@ mod pq_generation_tests { "".to_string(), None, ); - let result = PQGeneration::::new(context).generate(); + let result = PQGeneration::::new(&context); assert!(result.is_err()); } }