Inline filtering for diskann-garnet

Author: Jack Moffitt

Cargo.lock

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

diskann-garnet/Cargo.toml

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

diskann-garnet/diskann-garnet.nuspec

@@ -6,7 +6,6 @@
 use crate::{
     SearchResults,
     garnet::{Context, GarnetId},
-    labels::GarnetQueryLabelProvider,
     provider::{DynamicQuantization, GarnetProvider},
 };
 use diskann::{
@@ -15,10 +14,7 @@ use diskann::{
     provider::DataProvider,
     utils::VectorRepr,
 };
-use diskann_providers::{
-    index::wrapped_async::DiskANNIndex, model::graph::provider::layers::BetaFilter,
-};
-use std::sync::Arc;
+use diskann_providers::index::wrapped_async::DiskANNIndex;
 
 /// Type-erased version of `DiskANNIndex<GarnetProvider>`.
 /// All vector data is passed as untyped byte slices.
@@ -31,17 +27,31 @@ pub(crate) trait DynIndex: Send + Sync {
         &self,
         context: &Context,
         data: &[u8],
-        params: &search::Knn,
-        filter: Option<(&GarnetQueryLabelProvider, f32)>,
+        params: search::Knn,
         output: &mut SearchResults<'_>,
     ) -> ANNResult<SearchStats>;
 
     fn search_element(
         &self,
         context: &Context,
         id: &GarnetId,
-        params: &search::Knn,
-        filter: Option<(&GarnetQueryLabelProvider, f32)>,
+        params: search::Knn,
+        output: &mut SearchResults<'_>,
+    ) -> ANNResult<SearchStats>;
+
+    fn filtered_search_vector(
+        &self,
+        context: &Context,
+        data: &[u8],
+        params: search::InlineFilterSearch,
+        output: &mut SearchResults<'_>,
+    ) -> ANNResult<SearchStats>;
+
+    fn filtered_search_element(
+        &self,
+        context: &Context,
+        id: &GarnetId,
+        params: search::InlineFilterSearch,
         output: &mut SearchResults<'_>,
     ) -> ANNResult<SearchStats>;
 
@@ -84,32 +94,50 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
         &self,
         context: &Context,
         data: &[u8],
-        params: &search::Knn,
-        filter: Option<(&GarnetQueryLabelProvider, f32)>,
+        params: search::Knn,
         output: &mut SearchResults<'_>,
     ) -> ANNResult<SearchStats> {
         let query = bytemuck::cast_slice::<u8, T>(data);
-        if let Some((labels, beta)) = filter {
-            let beta_filter = BetaFilter::new(DynamicQuantization, Arc::new(labels.clone()), beta);
-            self.search(*params, &beta_filter, context, query, output)
-        } else {
-            self.search(*params, &DynamicQuantization, context, query, output)
-        }
+        self.search(params, &DynamicQuantization, context, query, output)
     }
 
     fn search_element(
         &self,
         context: &Context,
         id: &GarnetId,
-        params: &search::Knn,
-        filter: Option<(&GarnetQueryLabelProvider, f32)>,
+        params: search::Knn,
+        output: &mut SearchResults<'_>,
+    ) -> ANNResult<SearchStats> {
+        // Look up internal ID
+        let iid = self.inner.provider().to_internal_id(context, id)?;
+        let data = self.inner.provider().get_full_vector(context, iid)?;
+        let data_bytes = bytemuck::cast_slice::<T, u8>(&data);
+        self.search_vector(context, data_bytes, params, output)
+    }
+
+    fn filtered_search_vector(
+        &self,
+        context: &Context,
+        data: &[u8],
+        params: search::InlineFilterSearch,
+        output: &mut SearchResults<'_>,
+    ) -> ANNResult<SearchStats> {
+        let query = bytemuck::cast_slice::<u8, T>(data);
+        self.search(params, &DynamicQuantization, context, query, output)
+    }
+
+    fn filtered_search_element(
+        &self,
+        context: &Context,
+        id: &GarnetId,
+        params: search::InlineFilterSearch,
         output: &mut SearchResults<'_>,
     ) -> ANNResult<SearchStats> {
         // Look up internal ID
         let iid = self.inner.provider().to_internal_id(context, id)?;
         let data = self.inner.provider().get_full_vector(context, iid)?;
         let data_bytes = bytemuck::cast_slice::<T, u8>(&data);
-        self.search_vector(context, data_bytes, params, filter, output)
+        self.filtered_search_vector(context, data_bytes, params, output)
     }
 
     fn remove(&self, context: &Context, id: &GarnetId) -> ANNResult<()> {

diskann-garnet/src/dyn_index.rs

@@ -173,6 +173,7 @@ mod tests {
                 callbacks.write_callback(),
                 callbacks.delete_callback(),
                 callbacks.rmw_callback(),
+                callbacks.filter_callback(),
             )
         };
         assert!(!index_ptr.is_null());

diskann-garnet/src/ffi_recall_tests.rs

@@ -60,6 +60,7 @@ mod tests {
                 callbacks.write_callback(),
                 callbacks.delete_callback(),
                 callbacks.rmw_callback(),
+                callbacks.filter_callback(),
             )
         };
 
@@ -759,76 +760,6 @@ mod tests {
         }
     }
 
-    #[test]
-    fn search_with_bitmap_all_match() {
-        let store = Store;
-        let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant);
-
-        unsafe {
-            assert_eq!(
-                insert_f32_vector(&ctx, index_ptr, 10, &[1.0, 0.0]),
-                InsertResult::Success
-            );
-            assert_eq!(
-                insert_f32_vector(&ctx, index_ptr, 20, &[0.0, 1.0]),
-                InsertResult::Success
-            );
-            assert_eq!(
-                insert_f32_vector(&ctx, index_ptr, 30, &[1.0, 1.0]),
-                InsertResult::Success
-            );
-
-            // Bitmap with bits 1,2,3 set (internal IDs for the 3 inserted vectors;
-            // internal ID 0 is the start point)
-            let bitmap: [u8; 8] = [0b00001110, 0, 0, 0, 0, 0, 0, 0];
-            let (ids, _dists) = do_search(&ctx, index_ptr, &[1.0, 0.0], 3, Some(&bitmap));
-            // Start point (internal ID 0) is filtered out from results,
-            // so we may get fewer than k results.
-            assert!(ids.len() >= 2, "should return at least 2 matching vectors");
-            assert_eq!(ids[0], 10, "closest should still be id=10");
-
-            drop_index(ctx.get(), index_ptr);
-        }
-    }
-
-    #[test]
-    fn search_with_bitmap_partial_match() {
-        let store = Store;
-        let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant);
-
-        unsafe {
-            // Internal ID 0 -> EID 10, vector [1,0]
-            assert_eq!(
-                insert_f32_vector(&ctx, index_ptr, 10, &[1.0, 0.0]),
-                InsertResult::Success
-            );
-            // Internal ID 1 -> EID 20, vector [0,1]
-            assert_eq!(
-                insert_f32_vector(&ctx, index_ptr, 20, &[0.0, 1.0]),
-                InsertResult::Success
-            );
-            // Internal ID 2 -> EID 30, vector [1,1]
-            assert_eq!(
-                insert_f32_vector(&ctx, index_ptr, 30, &[1.0, 1.0]),
-                InsertResult::Success
-            );
-
-            // Bitmap with only bit 2 set (internal ID 2 = EID 20, second inserted vector)
-            let bitmap: [u8; 8] = [0b00000100, 0, 0, 0, 0, 0, 0, 0];
-            // Query close to EID 20's vector [0,1] to ensure it appears in results
-            let (ids, _dists) = do_search(&ctx, index_ptr, &[0.0, 1.0], 3, Some(&bitmap));
-            // BetaFilter biases toward matching vectors by scaling their distances.
-            assert!(!ids.is_empty(), "should return at least one result");
-            // EID 20 should appear since it's the closest to query AND matches the filter
-            assert!(
-                ids.contains(&20),
-                "filtered vector EID 20 should be in results"
-            );
-
-            drop_index(ctx.get(), index_ptr);
-        }
-    }
-
     #[test]
     fn search_with_null_bitmap_same_as_unfiltered() {
         let store = Store;

diskann-garnet/src/ffi_tests.rs

@@ -85,13 +85,15 @@ 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;
 
 #[derive(Copy, Clone)]
 pub(crate) struct Callbacks {
     read_callback: ReadCallback,
     write_callback: WriteCallback,
     delete_callback: DeleteCallback,
     rmw_callback: ReadModifyWriteCallback,
+    filter_callback: FilterCallback,
 }
 
 impl Callbacks {
@@ -100,12 +102,14 @@ impl Callbacks {
         write_callback: WriteCallback,
         delete_callback: DeleteCallback,
         rmw_callback: ReadModifyWriteCallback,
+        filter_callback: FilterCallback,
     ) -> Self {
         Self {
             read_callback,
             write_callback,
             delete_callback,
             rmw_callback,
+            filter_callback,
         }
     }
 
@@ -129,6 +133,11 @@ impl Callbacks {
         self.rmw_callback
     }
 
+    #[cfg(test)]
+    pub(crate) fn filter_callback(&self) -> FilterCallback {
+        self.filter_callback
+    }
+
     #[cfg(test)]
     pub(crate) fn exists_iid(&self, ctx: &Context, id: u32) -> bool {
         let key = [4, id];
@@ -477,6 +486,12 @@ 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) }
+    }
 }
 
 unsafe extern "C" fn read_call<'a, F, T>(index: u32, ptr: *mut c_void, data: *const u8, len: usize)