diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index 5db76bb1a..ec9ade194 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -163,7 +163,7 @@ where >::new( self.index_writer.get_dataset_file(), self.pq_storage.get_compressed_data_path().into(), - &quantizer_context, + quantizer_context, ); generator.generate_data( storage_provider, diff --git a/diskann-disk/src/storage/quant/compressor.rs b/diskann-disk/src/storage/quant/compressor.rs index 2de3cdfbf..7b55634df 100644 --- a/diskann-disk/src/storage/quant/compressor.rs +++ b/diskann-disk/src/storage/quant/compressor.rs @@ -20,7 +20,8 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; /// - [`CompressorContext`]: An overloadable type that provides initialization parameters for the compressor /// /// # Methods -/// - `prepare`: Performs any setup needed before compression and returns a compressor. +/// - `new`: Constructs a compressor with the provided context. +/// - `generate`: Generates any data needed before compression. /// - `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 @@ -29,7 +30,8 @@ where { type CompressorContext; - fn prepare(context: &Self::CompressorContext) -> ANNResult; + fn new(context: Self::CompressorContext) -> Self; + fn generate(&self) -> 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 0ebae5a42..79913fd24 100644 --- a/diskann-disk/src/storage/quant/generator.rs +++ b/diskann-disk/src/storage/quant/generator.rs @@ -25,18 +25,18 @@ use crate::{ /// [`QuantDataGenerator`] orchestrates the process of reading vector data, applying quantization, /// and writing compressed results to storage in batches. -pub struct QuantDataGenerator<'a, T, Q> +pub struct QuantDataGenerator where T: Copy + VectorRepr, Q: QuantCompressor, { - quantizer_context: &'a Q::CompressorContext, + pub quantizer: Q, pub data_path: String, pub compressed_data_path: String, phantom: PhantomData, } -impl<'a, T, Q> QuantDataGenerator<'a, T, Q> +impl QuantDataGenerator where T: Copy + VectorRepr, Q: QuantCompressor, @@ -44,12 +44,13 @@ where pub fn new( data_path: String, compressed_data_path: String, - quantizer_context: &'a Q::CompressorContext, + quantizer_context: Q::CompressorContext, ) -> Self { + let quantizer = Q::new(quantizer_context); Self { data_path, compressed_data_path, - quantizer_context, + quantizer, phantom: PhantomData, } } @@ -92,8 +93,7 @@ where )); } - let quantizer = Q::prepare(self.quantizer_context)?; - let compressed_size = quantizer.compressed_bytes(); + self.quantizer.generate()?; let compressed_path = self.compressed_data_path.as_str(); if storage_provider.exists(compressed_path) { @@ -106,8 +106,10 @@ where data_reader.seek(SeekFrom::Start((std::mem::size_of::() * 2) as u64))?; let mut compressed_data_writer = storage_provider.create_for_write(compressed_path)?; - Metadata::new(num_points, compressed_size)?.write(&mut compressed_data_writer)?; + Metadata::new(num_points, self.quantizer.compressed_bytes())? + .write(&mut compressed_data_writer)?; + let compressed_size = self.quantizer.compressed_bytes(); let block_size = std::cmp::min(num_points, max_block_size); let num_blocks = num_points / block_size + !num_points.is_multiple_of(block_size) as usize; @@ -162,7 +164,7 @@ where base_block .par_window_iter(BATCH_SIZE) .zip_eq(compressed_block.par_window_iter_mut(BATCH_SIZE)) - .try_for_each_in_pool(pool, |(src, dst)| quantizer.compress(src, dst))?; + .try_for_each_in_pool(pool, |(src, dst)| self.quantizer.compress(src, dst))?; let write_offset = start_index * compressed_size + std::mem::size_of::() * 2; compressed_data_writer.seek(SeekFrom::Start(write_offset as u64))?; @@ -190,10 +192,7 @@ where #[cfg(test)] mod generator_tests { - use std::{ - io::BufReader, - sync::atomic::{AtomicUsize, Ordering}, - }; + use std::io::BufReader; use diskann::utils::read_exact_into; use diskann_providers::storage::VirtualStorageProvider; @@ -206,24 +205,6 @@ mod generator_tests { use vfs::{FileSystem, MemoryFS}; use super::*; - pub struct DummyCompressorContext { - pub output_dim: u32, - pub prepare_calls: AtomicUsize, - } - - impl DummyCompressorContext { - pub fn new(output_dim: u32) -> Self { - Self { - output_dim, - prepare_calls: AtomicUsize::new(0), - } - } - - pub fn prepare_calls(&self) -> usize { - self.prepare_calls.load(Ordering::SeqCst) - } - } - pub struct DummyCompressor { pub output_dim: u32, pub code: Vec, @@ -237,11 +218,14 @@ mod generator_tests { } } impl QuantCompressor for DummyCompressor { - type CompressorContext = DummyCompressorContext; + type CompressorContext = u32; + + fn new(context: Self::CompressorContext) -> Self { + Self::new(context) + } - fn prepare(context: &Self::CompressorContext) -> ANNResult { - context.prepare_calls.fetch_add(1, Ordering::SeqCst); - Ok(Self::new(context.output_dim)) + fn generate(&self) -> ANNResult<()> { + Ok(()) } fn compress( @@ -297,16 +281,16 @@ mod generator_tests { Ok((storage_provider, data_path, compressed_path)) } - fn create_and_call_generator<'a, F: vfs::FileSystem>( + fn create_and_call_generator( compressed_path: String, storage_provider: &VirtualStorageProvider, data_path: String, - context: &'a DummyCompressorContext, + output_dim: u32, max_block_size: usize, - ) -> (QuantDataGenerator<'a, f32, DummyCompressor>, ANNResult<()>) { + ) -> (QuantDataGenerator, ANNResult<()>) { let pool: diskann_providers::utils::RayonThreadPool = create_thread_pool_for_test(); let generator = - QuantDataGenerator::::new(data_path, compressed_path, context); + QuantDataGenerator::::new(data_path, compressed_path, output_dim); let result = generator.generate_data(storage_provider, pool.as_ref(), max_block_size); (generator, result) } @@ -321,17 +305,15 @@ mod generator_tests { #[case] output_dim: u32, ) -> ANNResult<()> { let (storage_provider, data_path, compressed_path) = generate_data_files(num_points, dim)?; - let context = DummyCompressorContext::new(output_dim); - let (_generator, result) = create_and_call_generator( + let (generator, result) = create_and_call_generator( compressed_path.clone(), &storage_provider, data_path, - &context, + output_dim, 10_000, ); result?; - assert_eq!(context.prepare_calls(), 1); assert!(storage_provider.exists(&compressed_path)); let expected_size = num_points * output_dim as usize; @@ -347,9 +329,8 @@ mod generator_tests { assert_eq!(metadata.ndims_u32(), output_dim); assert_eq!(metadata.npoints(), num_points); - let expected_code: Vec = (0..output_dim).map(|x| (x % 256) as u8).collect(); data.chunks_exact(output_dim as usize) - .for_each(|chunk| assert_eq!(chunk, expected_code.as_slice())); + .for_each(|chunk| assert_eq!(chunk, generator.quantizer.code.as_slice())); Ok(()) } @@ -365,18 +346,15 @@ mod generator_tests { let data_path = "/test_data/empty.bin".to_string(); let compressed_path = "/test_data/empty_compressed.bin".to_string(); Metadata::new(0, 8)?.write(&mut storage_provider.create_for_write(data_path.as_str())?)?; - let context = DummyCompressorContext::new(4); let (_, result) = create_and_call_generator( compressed_path.clone(), &storage_provider, data_path, - &context, + 4, 10_000, ); - assert!(result.is_err()); - assert_eq!(context.prepare_calls(), 0); assert!(!storage_provider.exists(&compressed_path)); Ok(()) } @@ -384,18 +362,11 @@ mod generator_tests { #[test] fn generate_data_rejects_zero_chunk_size() -> ANNResult<()> { let (storage_provider, data_path, compressed_path) = generate_data_files(1, 8)?; - let context = DummyCompressorContext::new(4); - let (_, result) = create_and_call_generator( - compressed_path.clone(), - &storage_provider, - data_path, - &context, - 0, - ); + let (_, result) = + create_and_call_generator(compressed_path.clone(), &storage_provider, data_path, 4, 0); assert!(result.is_err()); - assert_eq!(context.prepare_calls(), 0); assert!(!storage_provider.exists(&compressed_path)); Ok(()) } diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index 0ff9f445d..3e64fbf43 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, time::Instant}; +use std::{marker::PhantomData, sync::OnceLock, time::Instant}; use diskann::utils::VectorRepr; use diskann_providers::storage::{StorageReadProvider, StorageWriteProvider}; @@ -46,10 +46,10 @@ where T: VectorRepr, Storage: StorageReadProvider + StorageWriteProvider + 'a, { - table: TransposedTable, + context: PQGenerationContext<'a, Storage>, + table: OnceLock, num_chunks: usize, phantom_data: PhantomData, - phantom_storage: PhantomData<&'a Storage>, } impl<'a, T, Storage> PQGeneration<'a, T, Storage> @@ -120,9 +120,23 @@ where { type CompressorContext = PQGenerationContext<'a, Storage>; - fn prepare(context: &Self::CompressorContext) -> diskann::ANNResult { - Self::generate_pivots(context)?; + 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; + Self::generate_pivots(context)?; let (_, full_dim) = context .pq_storage .read_existing_pivot_metadata(context.storage_provider)?; @@ -154,11 +168,11 @@ where ) .map_err(|err| diskann_error!(ErrorKind::PQError, "{}", Format(err)))?; - Ok(Self { - table, - num_chunks, - phantom_data: PhantomData, - phantom_storage: PhantomData, + self.table.set(table).map_err(|_| { + diskann_error!( + ErrorKind::PQError, + "PQ compressor was generated concurrently" + ) }) } @@ -168,6 +182,13 @@ 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))) } @@ -273,21 +294,24 @@ mod pq_generation_tests { assert!(!storage_provider.exists(pivot_file_name)); - let compressor = PQGeneration::::prepare(&context); - assert!(compressor.is_ok()); + let compressor = PQGeneration::::new(context); + assert!(!storage_provider.exists(pivot_file_name)); + + let result = compressor.generate(); + assert!(result.is_ok()); assert!(storage_provider.exists(pivot_file_name)); - let compressor = compressor.unwrap(); assert_eq!(compressor.num_chunks, num_chunks); assert_eq!(compressor.compressed_bytes(), num_chunks); - assert_eq!(compressor.table.dim(), dim); - assert_eq!(compressor.table.ncenters(), num_centers); - assert_eq!(compressor.table.nchunks(), 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); } #[rstest] - fn prepare_generates_missing_pivots() { + fn generate_creates_missing_pivots() { let storage_provider = VirtualStorageProvider::new_memory(); storage_provider .filesystem() @@ -318,9 +342,10 @@ mod pq_generation_tests { Some(data_path), ); - let compressor = PQGeneration::::prepare(&context); + let compressor = PQGeneration::::new(context); + let result = compressor.generate(); - assert!(compressor.is_ok()); + assert!(result.is_ok()); assert!(storage_provider.exists(pivot_file_name)); } @@ -345,19 +370,20 @@ mod pq_generation_tests { "".to_string(), None, ); - let compressor = PQGeneration::::prepare(&context); + let compressor = PQGeneration::::new(context); + let result = compressor.generate(); - if let Err(x) = compressor.as_ref() { + if let Err(x) = result.as_ref() { println!("Error creating compressor: {x}"); }; - assert!(compressor.is_ok()); + assert!(result.is_ok()); let data_matrix = read_bin::(&mut storage_provider.open_reader(TEST_PQ_DATA_PATH).unwrap()).unwrap(); let npts = data_matrix.nrows(); let mut compressed_mat = vec![0_u8; num_chunks * npts]; - let result = compressor.unwrap().compress( + let result = compressor.compress( data_matrix.as_view(), MutMatrixView::try_from(&mut compressed_mat, npts, num_chunks).unwrap(), ); @@ -397,7 +423,7 @@ mod pq_generation_tests { "".to_string(), None, ); - let compressor = PQGeneration::::prepare(&context); - assert!(compressor.is_err()); + let result = PQGeneration::::new(context).generate(); + assert!(result.is_err()); } }