[diskann-garnet] Fix Q8 quantizer state loading

diskann-garnet/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "diskann-garnet"
-version = "4.0.1"
+version = "4.0.2"
 edition = "2024"
 authors.workspace = true
 license.workspace = true

diskann-garnet/diskann-garnet.nuspec

@@ -2,7 +2,7 @@
 <package>
   <metadata>
     <id>diskann-garnet</id>
-    <version>4.0.1</version>
+    <version>4.0.2</version>
     <readme>docs/README.md</readme>
     <authors>Microsoft</authors>
     <projectUrl>https://github.com/microsoft/DiskANN</projectUrl>

diskann-garnet/src/provider.rs

@@ -212,8 +212,25 @@ impl<T: VectorRepr> GarnetProvider<T> {
                     return Err(GarnetProviderError::InvalidQuantizer);
                 }
 
-                let quantizer = Box::new(quantization::MinMax8Bit::new(dim, metric_type)?)
-                    as Box<dyn GarnetQuantizer>;
+                let quantizer = if let Some(quant_state) =
+                    callbacks.read_varsize_iid::<u8>(&context.term(Term::Metadata), QUANT_STATE_KEY)
+                {
+                    quantization::MinMax8Bit::new_from_bytes(metric_type, &quant_state)?
+                } else {
+                    let quantizer = quantization::MinMax8Bit::new(dim, metric_type)?;
+
+                    if !callbacks.write_iid(
+                        &context.term(Term::Metadata),
+                        QUANT_STATE_KEY,
+                        &quantizer.serialize()?,
+                    ) {
+                        return Err(GarnetError::Write.into());
+                    }
+
+                    quantizer
+                };
+
+                let quantizer = Box::new(quantizer) as Box<dyn GarnetQuantizer>;
                 let canonical_bytes = quantizer.bytes();
 
                 // NOTE: Q8 needs no training, so it always starts with backfill
@@ -2096,4 +2113,70 @@ mod tests {
         // quantized
         assert!(store.full_reads() < store.quant_reads());
     }
+
+    /// Test that restarts during phase three quant bootstrap work.
+    /// Phase three starts once backfill is complete and lasts for the remaining
+    /// life of the index.
+    #[test]
+    fn restart_q8() {
+        let store = Store::new();
+        let ctx = Context::new(0);
+        let index = create_2d_f32_index(VectorQuantType::Q8, Metric::L2, &store, &ctx);
+        let provider = index.inner.provider();
+
+        let mut rng = rand::rng();
+
+        let mut first_insert = true;
+        for id in 0..10 {
+            let v = [rng.random(), rng.random()];
+
+            if first_insert {
+                provider.maybe_set_start_point(&ctx, &v).unwrap();
+                first_insert = false;
+            }
+
+            DynIndex::insert(
+                &index,
+                &ctx,
+                &GarnetId::from(bytemuck::bytes_of::<u32>(&id)),
+                bytemuck::cast_slice::<f32, u8>(&v),
+            )
+            .unwrap();
+        }
+
+        // Drop and re-create the index, keeping the same backing store
+        let index = create_2d_f32_index(VectorQuantType::Q8, Metric::L2, &store, &ctx);
+        let provider = index.inner.provider();
+
+        // There should be saved quant state.
+        assert!(
+            provider
+                .callbacks
+                .exists_iid(&ctx.term(Term::Metadata), QUANT_STATE_KEY),
+            "quant state missing"
+        );
+
+        // Vectors should serialize identically
+
+        let quantizer = provider.quantizer.as_ref().expect("missing quantizer");
+        let mut fv = vec![0f32; 2];
+        let mut orig_qv = vec![0u8; quantizer.bytes()];
+        let mut qv = vec![0u8; quantizer.bytes()];
+
+        assert!(provider.callbacks.read_single_iid(
+            &ctx.term(Term::Vector),
+            1,
+            bytemuck::cast_slice_mut::<f32, u8>(&mut fv),
+        ));
+
+        quantizer.compress(&fv, &mut qv).unwrap();
+
+        assert!(
+            provider
+                .callbacks
+                .read_single_iid(&ctx.term(Term::Quantized), 1, &mut orig_qv)
+        );
+
+        assert_eq!(orig_qv, qv, "quant vectors mismatched");
+    }
 }

diskann-garnet/src/quantization.rs

@@ -5,7 +5,7 @@ use diskann_quantization::{
     CompressInto,
     algorithms::{Transform, TransformKind, transforms::NewTransformError},
     alloc::{GlobalAllocator, Poly, ScopedAllocator},
-    minmax,
+    minmax::{self, MinMaxQuantizer},
     num::POSITIVE_ONE_F32,
     spherical::{
         self, Data, PreScale, SphericalQuantizer, SupportedMetric,
@@ -241,6 +241,7 @@ impl MinMax8Bit {
             Some(d) => d,
             None => return Err(GarnetQuantizerError::ZeroDim),
         };
+
         let mut rng = rand::rng();
         let transform = Transform::new(
             TransformKind::DoubleHadamard {
@@ -257,6 +258,15 @@ impl MinMax8Bit {
             inner: minmax::MinMaxQuantizer::new(transform, grid_scale),
         })
     }
+
+    pub(crate) fn new_from_bytes(
+        metric: Metric,
+        bytes: &[u8],
+    ) -> Result<Self, GarnetQuantizerError> {
+        let inner = MinMaxQuantizer::try_deserialize(bytes)
+            .map_err(|e| GarnetQuantizerError::Deserialization(Box::new(e)))?;
+        Ok(Self { metric, inner })
+    }
 }
 
 impl GarnetQuantizer for MinMax8Bit {
@@ -306,7 +316,9 @@ impl GarnetQuantizer for MinMax8Bit {
     }
 
     fn serialize(&self) -> Result<Poly<[u8], GlobalAllocator>, GarnetQuantizerError> {
-        Err(GarnetQuantizerError::UnsupportedSerialization)
+        self.inner
+            .serialize(GlobalAllocator)
+            .map_err(|e| GarnetQuantizerError::Alloc(Box::new(e)))
     }
 
     fn deserialize(&self, _state: &[u8]) -> Result<(), GarnetQuantizerError> {

diskann-quantization/build.rs

@@ -34,7 +34,11 @@ mod compile {
             )
         });
         let out_dir = env::var_os("OUT_DIR").expect("This should be set by Cargo");
-        let input_files = ["transforms_v001.fbs", "spherical_v001.fbs"];
+        let input_files = [
+            "transforms_v001.fbs",
+            "spherical_v001.fbs",
+            "minmax_v001.fbs",
+        ];
 
         // Let `cargo` know that we want to rerun under the following conditions:
         //
@@ -55,7 +59,7 @@ mod compile {
         //
         // We first remove the current contents of the generated `flatbuffers` directory
         // to ensure we don't leave cruft behind.
-        let output_dirs = ["transforms", "spherical"];
+        let output_dirs = ["transforms", "spherical", "minmax"];
 
         let generated_folder = "./src/flatbuffers";
 

diskann-quantization/schemas/minmax_v001.fbs

@@ -0,0 +1,12 @@
+include "transforms_v001.fbs";
+
+namespace Minmax;
+
+file_identifier "mq00";
+
+table Quantizer {
+    transform:Transforms.Transform (required);
+    grid_scale:float32;
+}
+
+root_type Quantizer;

diskann-quantization/src/flatbuffers.rs

@@ -78,6 +78,11 @@ pub(crate) mod spherical {
     import_schema!(quantizer_generated, "spherical/quantizer_generated.rs");
 }
 
+pub(crate) mod minmax {
+    use super::*;
+    import_schema!(quantizer_generated, "minmax/quantizer_generated.rs");
+}
+
 /// Create a `FlatBufferBuilder` and pass it to the closure.
 ///
 /// After the closure runs, finish the serialization with the returned offset and pass