[diskann-garnet] Implement continue_search()

diskann-garnet/Cargo.toml

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

diskann-garnet/diskann-garnet.nuspec

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

diskann-garnet/src/dyn_index.rs

@@ -71,6 +71,9 @@ pub(crate) trait DynIndex: Send + Sync {
     /// Return an approximate count of vectors in the index
     fn approximate_count(&self) -> u64;
 
+    /// Return the maximum degree of the index graph
+    fn max_degree(&self) -> usize;
+
     /// Set a start point if one doesn't already exist.
     /// If there is already a start point, this is a no-op.
     fn maybe_set_start_point(&self, context: &Context, data: &[u8]) -> ANNResult<()>;
@@ -101,6 +104,11 @@ pub(crate) trait DynIndex: Send + Sync {
 
     /// Returns the neighbors of and distances from the target vector
     fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult<Vec<Neighbor<GarnetId>>>;
+
+    /// Log a message to Garnet. The context term can be used to scope the log
+    /// message to an area (e.g. `Term::Quantized` for quantization related
+    /// messages).
+    fn log(&self, context: &Context, msg: &str);
 }
 
 impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
@@ -194,6 +202,10 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
         self.inner.provider().max_internal_id() as u64
     }
 
+    fn max_degree(&self) -> usize {
+        self.inner.provider().max_degree()
+    }
+
     fn maybe_set_start_point(&self, context: &Context, data: &[u8]) -> ANNResult<()> {
         self.inner
             .provider()
@@ -236,4 +248,8 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
     fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult<Vec<Neighbor<GarnetId>>> {
         self.inner.provider().neighbors(context, id)
     }
+
+    fn log(&self, context: &Context, msg: &str) {
+        self.inner.provider().log(context, msg);
+    }
 }

diskann-garnet/src/ffi_recall_tests.rs

@@ -10,8 +10,8 @@ mod tests {
     use diskann_vector::distance::{Cosine, Metric, SquaredL2};
 
     use crate::{
-        VectorQuantType, create_index, drop_index, garnet::Context, insert, search_vector,
-        test_utils::Store,
+        Continuation, 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.
@@ -197,13 +197,13 @@ mod tests {
         let delta = 2.0_f32;
         let search_exploration_factor = 200_u32;
         let max_filtering_effort = 0_usize;
-        let continuation = ptr::null_mut();
 
         for vec in vectors {
             let query_bytes: &[u8] = bytemuck::cast_slice(vec);
             let max_id_size = mem::size_of::<u32>() + max_id_len;
             let mut output_id_buffer = vec![0u8; k * max_id_size];
             let mut output_dists = vec![0f32; k];
+            let mut continuation = ptr::null_mut();
 
             let count = unsafe {
                 search_vector(
@@ -221,7 +221,7 @@ mod tests {
                     output_dists.as_mut_ptr(),
                     output_dists.len(),
                     1,
-                    continuation,
+                    &mut continuation,
                 )
             };
             assert!(count >= 0, "search failed");
@@ -238,6 +238,10 @@ mod tests {
             );
             total_matches += matches;
             total_expected += expected_ids.len();
+
+            if !continuation.is_null() {
+                unsafe { drop(Continuation::from_ptr(continuation)) };
+            }
         }
 
         unsafe { drop_index(ctx.get(), index_ptr) };

diskann-garnet/src/ffi_tests.rs

@@ -10,8 +10,9 @@ mod tests {
     use rand::{Rng, seq::SliceRandom};
 
     use crate::{
-        Index, InsertResult, VectorQuantType, backfill_quant_vectors, build_quant_table, card,
-        check_external_id_valid, check_internal_id_valid, create_index, drop_index,
+        Continuation, Index, InsertResult, VectorQuantType, backfill_quant_vectors,
+        build_quant_table, card, check_external_id_valid, check_internal_id_valid, create_index,
+        drop_index,
         garnet::{Context, Term},
         insert,
         quantization::{GarnetQuantizer, Spherical1Bit},
@@ -465,6 +466,7 @@ mod tests {
         let qv = &[0.0f32, 0.0];
         let mut output_id_buffer = vec![0u8; 2 * (mem::size_of::<u64>() + mem::size_of::<u32>())];
         let mut output_dists = vec![0f32; 2];
+        let mut continuation = ptr::null_mut();
 
         let count = unsafe {
             search_vector(
@@ -482,11 +484,12 @@ mod tests {
                 output_dists.as_mut_ptr(),
                 output_dists.len(),
                 1,
-                ptr::null_mut(),
+                &mut continuation,
             )
         };
 
         assert_eq!(count, 2);
+        assert!(continuation.is_null());
 
         let mut output_ids = vec![];
         let mut offset = 0;
@@ -573,6 +576,7 @@ mod tests {
 
         let mut output_id_buffer = vec![0u8; 2 * (mem::size_of::<u64>() + mem::size_of::<u32>())];
         let mut output_dists = vec![0f32; 2];
+        let mut continuation = ptr::null_mut();
 
         let count = unsafe {
             crate::search_element(
@@ -590,11 +594,12 @@ mod tests {
                 output_dists.as_mut_ptr(),
                 output_dists.len(),
                 1,
-                ptr::null_mut(),
+                &mut continuation,
             )
         };
 
         assert_eq!(count, 2);
+        assert!(continuation.is_null());
 
         let mut output_ids = vec![];
         let mut offset = 0;
@@ -633,6 +638,48 @@ mod tests {
         }
     }
 
+    #[test]
+    fn search_element_writes_continuation_output() {
+        let store = Store::new();
+        let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant);
+        let id = 1u32;
+        assert_eq!(
+            insert_f32_vector(&ctx, index_ptr, id, &[0.0, 1.0]),
+            InsertResult::Success
+        );
+
+        let mut output_ids: [u8; 0] = [];
+        let mut output_distances: [f32; 1] = [0f32];
+        let mut continuation = ptr::null_mut();
+        let count = unsafe {
+            crate::search_element(
+                ctx.get(),
+                index_ptr,
+                bytemuck::bytes_of(&id).as_ptr(),
+                mem::size_of::<u32>(),
+                1.0,
+                10,
+                ptr::null(),
+                0,
+                0,
+                output_ids.as_mut_ptr(),
+                output_ids.len(),
+                output_distances.as_mut_ptr(),
+                output_distances.len(),
+                1,
+                &mut continuation,
+            )
+        };
+
+        assert_eq!(count, 0);
+        assert!(!continuation.is_null());
+
+        unsafe {
+            drop(Continuation::from_ptr(continuation));
+            drop_index(ctx.get(), index_ptr);
+        }
+    }
+
     #[test]
     fn continue_search() {
         let store = Store::new();
@@ -694,6 +741,7 @@ mod tests {
             Some(b) => (b.as_ptr(), b.len()),
             None => (ptr::null(), 0),
         };
+        let mut continuation = ptr::null_mut();
 
         let count = unsafe {
             search_vector(
@@ -711,10 +759,14 @@ mod tests {
                 output_dists.as_mut_ptr(),
                 output_dists.len(),
                 1,
-                ptr::null_mut(),
+                &mut continuation,
             )
         };
 
+        if !continuation.is_null() {
+            unsafe { drop(Continuation::from_ptr(continuation)) };
+        }
+
         assert!(count >= 0, "search failed with {count}");
         let count = count as usize;
 
@@ -736,6 +788,49 @@ mod tests {
         (ids, output_dists)
     }
 
+    #[test]
+    fn search_vector_writes_continuation_output() {
+        let store = Store::new();
+        let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant);
+        let vector = [0.0f32, 1.0];
+        assert_eq!(
+            insert_f32_vector(&ctx, index_ptr, 1, &vector),
+            InsertResult::Success
+        );
+
+        let query_bytes = bytemuck::cast_slice(&vector);
+        let mut output_ids: [u8; 0] = [];
+        let mut output_distances: [f32; 1] = [0f32];
+        let mut continuation = ptr::null_mut();
+        let count = unsafe {
+            search_vector(
+                ctx.get(),
+                index_ptr,
+                query_bytes.as_ptr(),
+                vector.len(),
+                1.0,
+                10,
+                ptr::null(),
+                0,
+                0,
+                output_ids.as_mut_ptr(),
+                output_ids.len(),
+                output_distances.as_mut_ptr(),
+                output_distances.len(),
+                1,
+                &mut continuation,
+            )
+        };
+
+        assert_eq!(count, 0);
+        assert!(!continuation.is_null());
+
+        unsafe {
+            drop(Continuation::from_ptr(continuation));
+            drop_index(ctx.get(), index_ptr);
+        }
+    }
+
     #[test]
     fn search_without_filter() {
         let store = Store::new();