[diskann-garnet] Implement BIN and Q8 quantizers

diskann-garnet/Cargo.toml

@@ -14,10 +14,11 @@ bytemuck.workspace = true
 crossbeam = "0.8.4"
 dashmap = { workspace = true, features = ["inline"] }
 diskann.workspace = true
-diskann-quantization.workspace = true
+diskann-quantization = { workspace = true, features = ["flatbuffers"] }
 diskann-providers.workspace = true
+diskann-utils.workspace = true
 diskann-vector.workspace = true
 foldhash = "0.2.0"
+rand.workspace = true
 thiserror.workspace = true
-tokio.workspace = true
-diskann-utils.workspace = true
+tokio = { workspace = true, features = ["sync"] }

diskann-garnet/diskann-garnet.nuspec

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

diskann-garnet/src/alloc.rs

@@ -9,7 +9,7 @@ use std::ptr::NonNull;
 /// Custom allocator that over-aligns to 8 bytes. This is needed since Garnet will hand us byte slices for f32 data
 /// that may be unaligned, so we need an allocator to make owned, aligned byte containers.
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct AlignToEight;
+pub struct AlignToEight;
 
 unsafe impl AllocatorCore for AlignToEight {
     #[inline]

diskann-garnet/src/dyn_index.rs

@@ -7,7 +7,7 @@ use crate::{
     SearchResults,
     garnet::{Context, GarnetId},
     labels::GarnetQueryLabelProvider,
-    provider::{self, GarnetProvider},
+    provider::{self, DynamicQuantization, GarnetProvider},
 };
 use diskann::{
     ANNError, ANNResult,
@@ -16,8 +16,7 @@ use diskann::{
     utils::VectorRepr,
 };
 use diskann_providers::{
-    index::wrapped_async::DiskANNIndex,
-    model::graph::provider::{async_::common::FullPrecision, layers::BetaFilter},
+    index::wrapped_async::DiskANNIndex, model::graph::provider::layers::BetaFilter,
 };
 use std::sync::Arc;
 
@@ -55,6 +54,10 @@ pub trait DynIndex: Send + Sync {
     fn internal_id_exists(&self, context: &Context, id: u32) -> bool;
 
     fn external_id_exists(&self, context: &Context, id: &GarnetId) -> bool;
+
+    fn train_quantizer(&self, context: &Context) -> bool;
+
+    fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize);
 }
 
 impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
@@ -63,7 +66,7 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
     /// The data slice here must be aligned to `T` or this will panic.
     fn insert(&self, context: &Context, id: &GarnetId, data: &[u8]) -> ANNResult<()> {
         self.insert(
-            FullPrecision,
+            DynamicQuantization,
             context,
             id,
             bytemuck::cast_slice::<u8, T>(data),
@@ -87,10 +90,10 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
     ) -> ANNResult<SearchStats> {
         let query = bytemuck::cast_slice::<u8, T>(data);
         if let Some((labels, beta)) = filter {
-            let beta_filter = BetaFilter::new(FullPrecision, Arc::new(labels.clone()), beta);
+            let beta_filter = BetaFilter::new(DynamicQuantization, Arc::new(labels.clone()), beta);
             self.search(*params, &beta_filter, context, query, output)
         } else {
-            self.search(*params, &FullPrecision, context, query, output)
+            self.search(*params, &DynamicQuantization, context, query, output)
         }
     }
 
@@ -105,23 +108,22 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
         let rt = tokio::runtime::Builder::new_current_thread()
             .build()
             .map_err(|e| ANNError::new(diskann::ANNErrorKind::Opaque, e))?;
-        let mut accessor: provider::FullAccessor<'_, T> =
-            <FullPrecision as SearchStrategy<_, _>>::search_accessor(
-                &FullPrecision,
+        let mut accessor: provider::DynamicAccessor<'_, T> =
+            <DynamicQuantization as SearchStrategy<_, _>>::search_accessor(
+                &DynamicQuantization,
                 self.inner.provider(),
                 context,
             )?;
 
         // Look up internal ID
         let iid = self.inner.provider().to_internal_id(context, id)?;
         let data = rt.block_on(accessor.get_element(iid))?;
-        let data_bytes = bytemuck::cast_slice::<T, u8>(&data);
-        self.search_vector(context, data_bytes, params, filter, output)
+        self.search_vector(context, &data, params, filter, output)
     }
 
     fn remove(&self, context: &Context, id: &GarnetId) -> ANNResult<()> {
         self.inplace_delete(
-            FullPrecision,
+            DynamicQuantization,
             context,
             id,
             3,
@@ -147,4 +149,14 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
     fn external_id_exists(&self, context: &Context, id: &GarnetId) -> bool {
         self.inner.provider().vector_id_exists(context, id)
     }
+
+    fn train_quantizer(&self, context: &Context) -> bool {
+        self.inner.provider().train_quantizer(context)
+    }
+
+    fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize) {
+        self.inner
+            .provider()
+            .backfill_quant_vectors(context, task_idx, task_count);
+    }
 }

diskann-garnet/src/ffi_recall_tests.rs

@@ -10,8 +10,8 @@ mod tests {
     use diskann_vector::distance::{Cosine, Metric, SquaredL2};
 
     use crate::{
-        VectorQuantType, VectorValueType, create_index, drop_index, garnet::Context, insert,
-        search_vector, test_utils::Store,
+        VectorQuantType, create_index, drop_index, garnet::Context, insert, search_vector,
+        test_utils::Store,
     };
 
     /// Helper to insert a vector with a string external ID and FP32 data.
@@ -25,16 +25,15 @@ mod tests {
         let vector_bytes: &[u8] = bytemuck::cast_slice(vector);
         unsafe {
             insert(
-                ctx.0,
+                ctx.get(),
                 index_ptr,
                 id_bytes.as_ptr(),
                 id_bytes.len(),
-                VectorValueType::FP32,
                 vector_bytes.as_ptr(),
                 vector.len(),
                 b"".as_ptr(),
                 0,
-            )
+            ) > 0
         }
     }
 
@@ -156,14 +155,14 @@ mod tests {
     ) -> f64 {
         store.clear();
         let callbacks = store.callbacks();
-        let ctx = Context(0);
+        let ctx = Context::new(0);
 
         let reduce_dimensions = 0;
         let l_build = 100;
         let max_degree = 32;
         let index_ptr = unsafe {
             create_index(
-                ctx.0,
+                ctx.get(),
                 dimensions,
                 reduce_dimensions,
                 VectorQuantType::NoQuant,
@@ -204,9 +203,8 @@ mod tests {
 
             let count = unsafe {
                 search_vector(
-                    ctx.0,
+                    ctx.get(),
                     index_ptr,
-                    VectorValueType::FP32,
                     query_bytes.as_ptr(),
                     vec.len(),
                     delta,
@@ -237,7 +235,7 @@ mod tests {
             total_expected += expected_ids.len();
         }
 
-        unsafe { drop_index(ctx.0, index_ptr) };
+        unsafe { drop_index(ctx.get(), index_ptr) };
 
         total_matches as f64 / total_expected as f64
     }