Split disk quantizer preparation from compression

diskann-disk/src/build/builder/build.rs

@@ -164,7 +164,7 @@ where
             self.index_writer.get_dataset_file(),
             self.pq_storage.get_compressed_data_path().into(),
             &quantizer_context,
-        )?;
+        );
         generator.generate_data(
             storage_provider,
             pool,

diskann-disk/src/storage/quant/compressor.rs

@@ -20,7 +20,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView};
 /// - [`CompressorContext`]: An overloadable type that provides initialization parameters for the compressor
 ///
 /// # Methods
-/// - `new`: Constructs a new compressor instance with the provided context.
+/// - `prepare`: Performs any setup needed before compression and returns a compressor.
 /// - `compress`: Compresses a batch of vectors into the output buffer.
 /// - `compressed_bytes`: Returns the size in bytes of each compressed vector
 pub trait QuantCompressor<T>: Sized + Sync
@@ -29,7 +29,7 @@ where
 {
     type CompressorContext;
 
-    fn new(context: &Self::CompressorContext) -> ANNResult<Self>;
+    fn prepare(context: &Self::CompressorContext) -> ANNResult<Self>;
     fn compress(&self, vector: MatrixView<f32>, output: MutMatrixView<u8>) -> ANNResult<()>;
     fn compressed_bytes(&self) -> usize;
 }

diskann-disk/src/storage/quant/generator.rs

@@ -25,34 +25,33 @@ use crate::{
 
 /// [`QuantDataGenerator`] orchestrates the process of reading vector data, applying quantization,
 /// and writing compressed results to storage in batches.
-pub struct QuantDataGenerator<T, Q>
+pub struct QuantDataGenerator<'a, T, Q>
 where
     T: Copy + VectorRepr,
     Q: QuantCompressor<T>,
 {
-    pub quantizer: Q,
+    quantizer_context: &'a Q::CompressorContext,
     pub data_path: String,
     pub compressed_data_path: String,
     phantom: PhantomData<T>,
 }
 
-impl<T, Q> QuantDataGenerator<T, Q>
+impl<'a, T, Q> QuantDataGenerator<'a, T, Q>
 where
     T: Copy + VectorRepr,
     Q: QuantCompressor<T>,
 {
     pub fn new(
         data_path: String,
         compressed_data_path: String,
-        quantizer_context: &Q::CompressorContext,
-    ) -> ANNResult<Self> {
-        let quantizer = Q::new(quantizer_context)?;
-        Ok(Self {
+        quantizer_context: &'a Q::CompressorContext,
+    ) -> Self {
+        Self {
             data_path,
             compressed_data_path,
-            quantizer,
+            quantizer_context,
             phantom: PhantomData,
-        })
+        }
     }
 
     /// This method reads the source data file, processes vectors in batches, compresses them
@@ -93,6 +92,8 @@ where
             ));
         }
 
+        let quantizer = Q::prepare(self.quantizer_context)?;
+        let compressed_size = quantizer.compressed_bytes();
         let compressed_path = self.compressed_data_path.as_str();
 
         if storage_provider.exists(compressed_path) {
@@ -105,10 +106,8 @@ where
         data_reader.seek(SeekFrom::Start((std::mem::size_of::<i32>() * 2) as u64))?;
 
         let mut compressed_data_writer = storage_provider.create_for_write(compressed_path)?;
-        Metadata::new(num_points, self.quantizer.compressed_bytes())?
-            .write(&mut compressed_data_writer)?;
+        Metadata::new(num_points, compressed_size)?.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;
 
@@ -163,7 +162,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)| self.quantizer.compress(src, dst))?;
+                .try_for_each_in_pool(pool, |(src, dst)| quantizer.compress(src, dst))?;
 
             let write_offset = start_index * compressed_size + std::mem::size_of::<i32>() * 2;
             compressed_data_writer.seek(SeekFrom::Start(write_offset as u64))?;
@@ -191,7 +190,10 @@ where
 
 #[cfg(test)]
 mod generator_tests {
-    use std::io::BufReader;
+    use std::{
+        io::BufReader,
+        sync::atomic::{AtomicUsize, Ordering},
+    };
 
     use diskann::utils::read_exact_into;
     use diskann_providers::storage::VirtualStorageProvider;
@@ -204,6 +206,24 @@ 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<u8>,
@@ -217,10 +237,11 @@ mod generator_tests {
         }
     }
     impl QuantCompressor<f32> for DummyCompressor {
-        type CompressorContext = u32;
+        type CompressorContext = DummyCompressorContext;
 
-        fn new(context: &Self::CompressorContext) -> ANNResult<Self> {
-            Ok(Self::new(*context))
+        fn prepare(context: &Self::CompressorContext) -> ANNResult<Self> {
+            context.prepare_calls.fetch_add(1, Ordering::SeqCst);
+            Ok(Self::new(context.output_dim))
         }
 
         fn compress(
@@ -276,20 +297,16 @@ mod generator_tests {
         Ok((storage_provider, data_path, compressed_path))
     }
 
-    fn create_and_call_generator<F: vfs::FileSystem>(
+    fn create_and_call_generator<'a, F: vfs::FileSystem>(
         compressed_path: String,
         storage_provider: &VirtualStorageProvider<F>,
         data_path: String,
-        output_dim: u32,
+        context: &'a DummyCompressorContext,
         max_block_size: usize,
-    ) -> (QuantDataGenerator<f32, DummyCompressor>, ANNResult<()>) {
+    ) -> (QuantDataGenerator<'a, f32, DummyCompressor>, ANNResult<()>) {
         let pool: diskann_providers::utils::RayonThreadPool = create_thread_pool_for_test();
-        let generator = QuantDataGenerator::<f32, DummyCompressor>::new(
-            data_path,
-            compressed_path,
-            &output_dim,
-        )
-        .unwrap();
+        let generator =
+            QuantDataGenerator::<f32, DummyCompressor>::new(data_path, compressed_path, context);
         let result = generator.generate_data(storage_provider, pool.as_ref(), max_block_size);
         (generator, result)
     }
@@ -304,15 +321,17 @@ mod generator_tests {
         #[case] output_dim: u32,
     ) -> ANNResult<()> {
         let (storage_provider, data_path, compressed_path) = generate_data_files(num_points, dim)?;
-        let (generator, result) = create_and_call_generator(
+        let context = DummyCompressorContext::new(output_dim);
+        let (_generator, result) = create_and_call_generator(
             compressed_path.clone(),
             &storage_provider,
             data_path,
-            output_dim,
+            &context,
             10_000,
         );
 
         result?;
+        assert_eq!(context.prepare_calls(), 1);
         assert!(storage_provider.exists(&compressed_path));
 
         let expected_size = num_points * output_dim as usize;
@@ -328,8 +347,9 @@ mod generator_tests {
         assert_eq!(metadata.ndims_u32(), output_dim);
         assert_eq!(metadata.npoints(), num_points);
 
+        let expected_code: Vec<u8> = (0..output_dim).map(|x| (x % 256) as u8).collect();
         data.chunks_exact(output_dim as usize)
-            .for_each(|chunk| assert_eq!(chunk, generator.quantizer.code.as_slice()));
+            .for_each(|chunk| assert_eq!(chunk, expected_code.as_slice()));
 
         Ok(())
     }
@@ -345,28 +365,37 @@ 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,
-            4,
+            &context,
             10_000,
         );
 
         assert!(result.is_err());
+        assert_eq!(context.prepare_calls(), 0);
         assert!(!storage_provider.exists(&compressed_path));
         Ok(())
     }
 
     #[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, 4, 0);
+        let (_, result) = create_and_call_generator(
+            compressed_path.clone(),
+            &storage_provider,
+            data_path,
+            &context,
+            0,
+        );
 
         assert!(result.is_err());
+        assert_eq!(context.prepare_calls(), 0);
         assert!(!storage_provider.exists(&compressed_path));
         Ok(())
     }