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

@@ -21,18 +21,31 @@ use diskann_utils::views::{MatrixView, MutMatrixView};
 /// # Associated Types
 /// - [`Self::CompressorContext`]: An overloadable type that provides initialization parameters
 ///   for the compressor.
+/// - [`Self::Prepared`]: The ready-to-use compressor produced by [`Self::prepare`].
 ///
 /// # Methods
 /// - `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<T>: Sized + Sync
+/// - `prepare`: Returns a ready-to-use compressor.
+pub trait QuantCompressor<'a, T>: Sized
 where
     T: VectorRepr,
 {
-    type CompressorContext;
+    type CompressorContext: 'a;
+
+    type Prepared: PreparedCompressor + Sync;
+
+    fn new(context: &'a Self::CompressorContext) -> Self;
+
+    /// Returns an error if preparation fails.
+    fn prepare(&self) -> ANNResult<Self::Prepared>;
+}
 
-    fn new(context: &Self::CompressorContext) -> ANNResult<Self>;
+/// A quantizer that is ready to compress vectors.
+///
+/// # Methods
+/// - `compress`: Compresses a batch of vectors into the output buffer.
+/// - `compressed_bytes`: Returns the size in bytes of each compressed vector.
+pub trait PreparedCompressor {
     fn compress(&self, vector: MatrixView<f32>, output: MutMatrixView<u8>) -> ANNResult<()>;
     fn compressed_bytes(&self) -> usize;
 }

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

@@ -20,39 +20,38 @@ use tracing::info;
 
 use crate::{
     error::{diskann_error, ErrorKind},
-    storage::quant::compressor::QuantCompressor,
+    storage::quant::compressor::{PreparedCompressor, QuantCompressor},
 };
 
 /// [`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>,
+    Q: QuantCompressor<'a, T>,
 {
     pub quantizer: Q,
     pub data_path: String,
     pub compressed_data_path: String,
-    phantom: PhantomData<T>,
+    phantom: PhantomData<&'a T>,
 }
 
-impl<T, Q> QuantDataGenerator<T, Q>
+impl<'a, T, Q> QuantDataGenerator<'a, T, Q>
 where
     T: Copy + VectorRepr,
-    Q: QuantCompressor<T>,
+    Q: QuantCompressor<'a, 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: Q::new(quantizer_context),
             phantom: PhantomData,
-        })
+        }
     }
 
     /// This method reads the source data file, processes vectors in batches, compresses them
@@ -61,12 +60,13 @@ where
     /// The implementation is adapted from generate_quantized_data_internal in pq_construction.rs
     //
     /// # Processing Flow
-    /// 1. Opens the source data file and validates its metadata.
-    /// 2. Deletes any existing output.
-    /// 3. Creates or opens output compressed file and writes metadata header - [num_points as i32, compressed_vector_size as i32]
-    /// 4. Processes data in bounded blocks.
-    /// 5. Compresses each block in small batch sizes in parallel to (potentially) take advantage of batch compression with quantizer
-    /// 6. Writes compressed blocks to the output file.
+    /// 1. Prepares the quantizer (training or loading a codebook as needed).
+    /// 2. Opens the source data file and validates its metadata.
+    /// 3. Deletes any existing output.
+    /// 4. Creates or opens output compressed file and writes metadata header - [num_points as i32, compressed_vector_size as i32]
+    /// 5. Processes data in bounded blocks.
+    /// 6. Compresses each block in small batch sizes in parallel to (potentially) take advantage of batch compression with quantizer
+    /// 7. Writes compressed blocks to the output file.
     pub fn generate_data<Storage>(
         &self,
         storage_provider: &Storage, // Provider for reading source data and writing compressed results
@@ -93,6 +93,7 @@ where
             ));
         }
 
+        let compressor = self.quantizer.prepare()?;
         let compressed_path = self.compressed_data_path.as_str();
 
         if storage_provider.exists(compressed_path) {
@@ -105,10 +106,10 @@ 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())?
+        Metadata::new(num_points, compressor.compressed_bytes())?
             .write(&mut compressed_data_writer)?;
 
-        let compressed_size = self.quantizer.compressed_bytes();
+        let compressed_size = compressor.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 +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)| self.quantizer.compress(src, dst))?;
+                .try_for_each_in_pool(pool, |(src, dst)| compressor.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))?;
@@ -216,13 +217,20 @@ mod generator_tests {
             }
         }
     }
-    impl QuantCompressor<f32> for DummyCompressor {
+    impl<'a> QuantCompressor<'a, f32> for DummyCompressor {
         type CompressorContext = u32;
+        type Prepared = DummyCompressor;
+
+        fn new(context: &'a Self::CompressorContext) -> Self {
+            Self::new(*context)
+        }
 
-        fn new(context: &Self::CompressorContext) -> ANNResult<Self> {
-            Ok(Self::new(*context))
+        fn prepare(&self) -> ANNResult<Self::Prepared> {
+            Ok(Self::new(self.output_dim))
         }
+    }
 
+    impl PreparedCompressor for DummyCompressor {
         fn compress(
             &self,
             _vector: views::MatrixView<f32>,
@@ -276,20 +284,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,
+        output_dim: &'a u32,
         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, output_dim);
         let result = generator.generate_data(storage_provider, pool.as_ref(), max_block_size);
         (generator, result)
     }
@@ -308,7 +312,7 @@ mod generator_tests {
             compressed_path.clone(),
             &storage_provider,
             data_path,
-            output_dim,
+            &output_dim,
             10_000,
         );
 
@@ -350,10 +354,9 @@ mod generator_tests {
             compressed_path.clone(),
             &storage_provider,
             data_path,
-            4,
+            &4,
             10_000,
         );
-
         assert!(result.is_err());
         assert!(!storage_provider.exists(&compressed_path));
         Ok(())
@@ -364,7 +367,7 @@ mod generator_tests {
         let (storage_provider, data_path, compressed_path) = generate_data_files(1, 8)?;
 
         let (_, result) =
-            create_and_call_generator(compressed_path.clone(), &storage_provider, data_path, 4, 0);
+            create_and_call_generator(compressed_path.clone(), &storage_provider, data_path, &4, 0);
 
         assert!(result.is_err());
         assert!(!storage_provider.exists(&compressed_path));

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

@@ -7,8 +7,8 @@ mod generator;
 pub use generator::QuantDataGenerator;
 
 pub(crate) mod pq;
-pub use pq::pq_generation::{PQGeneration, PQGenerationContext};
+pub use pq::pq_generation::{PQCompressor, PQGeneration, PQGenerationContext};
 pub use pq::PQData;
 
 mod compressor;
-pub use compressor::QuantCompressor;
+pub use compressor::{PreparedCompressor, QuantCompressor};