[diskann-garnet] FFI updates and additions

diskann-garnet/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "diskann-garnet"
-version = "4.0.4"
+version = "5.0.0"
 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.4</version>
+    <version>5.0.0</version>
     <readme>docs/README.md</readme>
     <authors>Microsoft</authors>
     <projectUrl>https://github.com/microsoft/DiskANN</projectUrl>

diskann-garnet/src/dyn_index.rs

@@ -11,6 +11,7 @@ use crate::{
 use diskann::{
     ANNResult,
     graph::{InplaceDeleteMethod, index::SearchStats, search},
+    neighbor::Neighbor,
     provider::DataProvider,
     utils::VectorRepr,
 };
@@ -69,7 +70,13 @@ pub(crate) trait DynIndex: Send + Sync {
 
     fn train_quantizer(&self, context: &Context) -> bool;
 
-    fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize);
+    fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize)
+    -> bool;
+
+    fn random_members(&self, context: &Context, count: u32, output: &mut SearchResults<'_>)
+    -> bool;
+
+    fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult<Vec<Neighbor<GarnetId>>>;
 }
 
 impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
@@ -182,9 +189,27 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
         self.inner.provider().train_quantizer(context)
     }
 
-    fn backfill_quant_vectors(&self, context: &Context, task_idx: usize, task_count: usize) {
+    fn backfill_quant_vectors(
+        &self,
+        context: &Context,
+        task_idx: usize,
+        task_count: usize,
+    ) -> bool {
         self.inner
             .provider()
-            .backfill_quant_vectors(context, task_idx, task_count);
+            .backfill_quant_vectors(context, task_idx, task_count)
+    }
+
+    fn random_members(
+        &self,
+        context: &Context,
+        count: u32,
+        output: &mut SearchResults<'_>,
+    ) -> bool {
+        self.inner.provider().random_members(context, count, output)
+    }
+
+    fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult<Vec<Neighbor<GarnetId>>> {
+        self.inner.provider().neighbors(context, id)
     }
 }

diskann-garnet/src/ffi_recall_tests.rs

@@ -175,6 +175,7 @@ mod tests {
                 callbacks.delete_callback(),
                 callbacks.rmw_callback(),
                 callbacks.filter_callback(),
+                callbacks.log_callback(),
                 &mut quant_needed,
             )
         };

diskann-garnet/src/ffi_tests.rs

@@ -60,6 +60,7 @@ mod tests {
                 callbacks.delete_callback(),
                 callbacks.rmw_callback(),
                 callbacks.filter_callback(),
+                callbacks.log_callback(),
                 &mut quant_needed,
             )
         };

diskann-garnet/src/fsm.rs

@@ -143,7 +143,7 @@ impl FreeSpaceMap {
         let block_key = Self::block_key(0);
         if this
             .callbacks
-            .exists_wid(&ctx.term(Term::Metadata), block_key)
+            .exists_wid(&ctx.term(Term::Metadata), block_key, BLOCK_SIZE_BYTES)
         {
             this.load_state(ctx)?;
         } else {
@@ -159,10 +159,11 @@ impl FreeSpaceMap {
     /// Load all state from Garnet by scanning the FSM blocks.
     fn load_state(&mut self, ctx: &Context) -> Result<(), FsmError> {
         let mut max_block_id = 0;
-        while self
-            .callbacks
-            .exists_wid(&ctx.term(Term::Metadata), Self::block_key(max_block_id))
-        {
+        while self.callbacks.exists_wid(
+            &ctx.term(Term::Metadata),
+            Self::block_key(max_block_id),
+            BLOCK_SIZE_BYTES,
+        ) {
             max_block_id += 1;
         }
 

diskann-garnet/src/garnet.rs

@@ -77,15 +77,16 @@ impl Context {
 impl ExecutionContext for Context {}
 
 pub(crate) type ReadCallback =
-    unsafe extern "C" fn(u64, u32, *const u8, usize, ReadDataCallback, *mut c_void);
+    unsafe extern "C" fn(u64, u32, u32, *const u8, usize, ReadDataCallback, *mut c_void);
 pub(crate) type WriteCallback =
     unsafe extern "C" fn(u64, *const u8, usize, *const u8, usize) -> bool;
 pub(crate) type DeleteCallback = unsafe extern "C" fn(u64, *const u8, usize) -> bool;
 pub(crate) type ReadModifyWriteCallback =
     unsafe extern "C" fn(u64, *const u8, usize, usize, RmwDataCallback, *mut c_void) -> bool;
 pub(crate) type ReadDataCallback = unsafe extern "C" fn(u32, *mut c_void, *const u8, usize);
 pub(crate) type RmwDataCallback = unsafe extern "C" fn(*mut c_void, *mut u8, usize);
-pub(crate) type FilterCallback = unsafe extern "C" fn(u64, u32) -> bool;
+pub(crate) type FilterCallback = unsafe extern "C" fn(u64, *const u8, usize) -> bool;
+pub(crate) type LogCallback = unsafe extern "C" fn(u64, *const u8, usize);
 
 #[derive(Copy, Clone)]
 pub(crate) struct Callbacks {
@@ -94,6 +95,7 @@ pub(crate) struct Callbacks {
     delete_callback: DeleteCallback,
     rmw_callback: ReadModifyWriteCallback,
     filter_callback: FilterCallback,
+    log_callback: LogCallback,
 }
 
 impl Callbacks {
@@ -103,13 +105,15 @@ impl Callbacks {
         delete_callback: DeleteCallback,
         rmw_callback: ReadModifyWriteCallback,
         filter_callback: FilterCallback,
+        log_callback: LogCallback,
     ) -> Self {
         Self {
             read_callback,
             write_callback,
             delete_callback,
             rmw_callback,
             filter_callback,
+            log_callback,
         }
     }
 
@@ -139,34 +143,39 @@ impl Callbacks {
     }
 
     #[cfg(test)]
-    pub(crate) fn exists_iid(&self, ctx: &Context, id: u32) -> bool {
+    pub(crate) fn log_callback(&self) -> LogCallback {
+        self.log_callback
+    }
+
+    #[cfg(test)]
+    pub(crate) fn exists_iid(&self, ctx: &Context, id: u32, length_hint: usize) -> bool {
         let key = [4, id];
         // SAFETY: Key bytes are preceded by 4 bytes of space.
-        unsafe { self.exists_raw(ctx, bytemuck::bytes_of(&key)) }
+        unsafe { self.exists_raw(ctx, bytemuck::bytes_of(&key), length_hint) }
     }
 
-    pub(crate) fn exists_wid(&self, ctx: &Context, key: u64) -> bool {
+    pub(crate) fn exists_wid(&self, ctx: &Context, key: u64, length_hint: usize) -> bool {
         // NOTE: the length is bit-shifted so that we have a u32 in the lower half of the u64.
         let mut key = [8 << 32, key];
         let key_bytes = bytemuck::bytes_of_mut(&mut key);
         // SAFETY: Key bytes are preceded by 8 bytes of extra space.
-        unsafe { self.exists_raw(ctx, &key_bytes[4..]) }
+        unsafe { self.exists_raw(ctx, &key_bytes[4..], length_hint) }
     }
 
     #[expect(
         dead_code,
         reason = "currently unused, but may be needed in the future"
     )]
-    pub(crate) fn exists_eid(&self, ctx: &Context, id: &GarnetId) -> bool {
+    pub(crate) fn exists_eid(&self, ctx: &Context, id: &GarnetId, length_hint: usize) -> bool {
         // SAFETY: GarnetId ensures there are 4 bytes preceding the key bytes.
-        unsafe { self.exists_raw(ctx, id) }
+        unsafe { self.exists_raw(ctx, id, length_hint) }
     }
 
     /// Check for a key's existance in Garnet.
     ///
     /// NOTE: The key bytes must be preceded by 4 valid bytes that Garnet can write into.
     /// This invariant must be checked by the caller.
-    unsafe fn exists_raw(&self, ctx: &Context, key: &[u8]) -> bool {
+    unsafe fn exists_raw(&self, ctx: &Context, key: &[u8], length_hint: usize) -> bool {
         let mut called = false;
         let mut cb = |_, _: &[u8]| {
             called = true;
@@ -176,6 +185,7 @@ impl Callbacks {
             (self.read_callback)(
                 ctx.inner,
                 1,
+                length_hint as u32,
                 key.as_ptr(),
                 key.len(),
                 make_read_call(&cb),
@@ -247,6 +257,7 @@ impl Callbacks {
     /// This invariant must be checked by the caller.
     #[must_use]
     unsafe fn read_single_raw(&self, ctx: &Context, key: &[u8], value: &mut [u8]) -> bool {
+        let length_hint = value.len() as u32;
         let mut found = false;
         let mut cb = |_, data: &[u8]| {
             found = true;
@@ -257,6 +268,7 @@ impl Callbacks {
             (self.read_callback)(
                 ctx.inner,
                 1,
+                length_hint,
                 key.as_ptr(),
                 key.len(),
                 make_read_call(&cb),
@@ -272,6 +284,7 @@ impl Callbacks {
         &self,
         ctx: &Context,
         ids: &[u32],
+        length_hint: usize,
         mut f: F,
     ) where
         F: FnMut(u32, &'a [T]),
@@ -284,6 +297,7 @@ impl Callbacks {
             (self.read_callback)(
                 ctx.inner,
                 ids.len() as u32 / 2,
+                length_hint as u32,
                 bytemuck::must_cast_slice::<_, u8>(ids).as_ptr(),
                 mem::size_of_val(ids),
                 make_read_call(&f),
@@ -316,11 +330,15 @@ impl Callbacks {
             result = Some(bytemuck::cast_slice::<u8, T>(data).to_owned());
         };
 
+        // NOTE: We hint the length as 8192 bytes, which will often overestimate. The only varsize
+        // things to read are the quant state and the external ID map. Quant state is
+        // maximum `117 + 6 * dim` bytes, which is several kilobytes in practice.
         // SAFETY: Key bytes are preceded by 4 bytes of extra space.
         unsafe {
             (self.read_callback)(
                 ctx.inner,
                 1,
+                8192,
                 bytemuck::bytes_of(&key).as_ptr(),
                 mem::size_of_val(&key),
                 make_read_call(&cb),
@@ -489,8 +507,27 @@ impl Callbacks {
 
     /// Evaluate the filter callback on an ID.
     #[must_use]
-    pub(crate) fn matches_filter(&self, ctx: &Context, id: u32) -> bool {
-        unsafe { (self.filter_callback)(ctx.inner, id) }
+    pub(crate) fn matches_filter(&self, ctx: &Context, data: &[u8]) -> bool {
+        unsafe {
+            (self.filter_callback)(
+                ctx.inner,
+                if data.is_empty() {
+                    std::ptr::null()
+                } else {
+                    data.as_ptr()
+                },
+                data.len(),
+            )
+        }
+    }
+
+    /// Log a message to Garnet.
+    ///
+    /// The context bits can be set with appropriate `Term` to flag which area the log message concerns.
+    pub(crate) fn log(&self, ctx: &Context, msg: &str) {
+        unsafe {
+            (self.log_callback)(ctx.inner, msg.as_ptr(), msg.len());
+        }
     }
 }