Add filtering for CAGRA to C API

cpp/include/cuvs/neighbors/cagra.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cuvs/core/c_api.h>
+#include <cuvs/neighbors/common.h>
 #include <dlpack/dlpack.h>
 #include <stdbool.h>
 #include <stdint.h>
@@ -385,13 +386,16 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
  * @param[in] queries DLManagedTensor* queries dataset to search
  * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries
  * @param[out] distances DLManagedTensor* output `k` distances for queries
+ * @param[in] prefilter cuvsFilter input prefilter that can be used
+              to filter queries and neighbors based on the given bitset.
  */
 cuvsError_t cuvsCagraSearch(cuvsResources_t res,
                             cuvsCagraSearchParams_t params,
                             cuvsCagraIndex_t index,
                             DLManagedTensor* queries,
                             DLManagedTensor* neighbors,
-                            DLManagedTensor* distances);
+                            DLManagedTensor* distances,
+                            cuvsFilter filter);
 
 /**
  * @}

cpp/src/neighbors/cagra_c.cpp

@@ -28,6 +28,7 @@
 #include <cuvs/core/interop.hpp>
 #include <cuvs/neighbors/cagra.h>
 #include <cuvs/neighbors/cagra.hpp>
+#include <cuvs/neighbors/common.h>
 
 #include <fstream>
 
@@ -91,7 +92,8 @@ void _search(cuvsResources_t res,
              cuvsCagraIndex index,
              DLManagedTensor* queries_tensor,
              DLManagedTensor* neighbors_tensor,
-             DLManagedTensor* distances_tensor)
+             DLManagedTensor* distances_tensor,
+             cuvsFilter filter)
 {
   auto res_ptr   = reinterpret_cast<raft::resources*>(res);
   auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index.addr);
@@ -117,8 +119,26 @@ void _search(cuvsResources_t res,
   auto queries_mds            = cuvs::core::from_dlpack<queries_mdspan_type>(queries_tensor);
   auto neighbors_mds          = cuvs::core::from_dlpack<neighbors_mdspan_type>(neighbors_tensor);
   auto distances_mds          = cuvs::core::from_dlpack<distances_mdspan_type>(distances_tensor);
-  cuvs::neighbors::cagra::search(
-    *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds);
+  if (filter.type == NO_FILTER) {
+    cuvs::neighbors::cagra::search(
+      *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds);
+  } else if (filter.type == BITSET) {
+    using filter_mdspan_type    = raft::device_vector_view<std::uint32_t, int64_t, raft::row_major>;
+    auto removed_indices_tensor = reinterpret_cast<DLManagedTensor*>(filter.addr);
+    auto removed_indices = cuvs::core::from_dlpack<filter_mdspan_type>(removed_indices_tensor);
+    cuvs::core::bitset_view<std::uint32_t, int64_t> removed_indices_bitset(
+      removed_indices, index_ptr->dataset().extent(0));
+    auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset);
+    cuvs::neighbors::cagra::search(*res_ptr,
+                                   search_params,
+                                   *index_ptr,
+                                   queries_mds,
+                                   neighbors_mds,
+                                   distances_mds,
+                                   bitset_filter_obj);
+  } else {
+    RAFT_FAIL("Unsupported filter type: BITMAP");
+  }
 }
 
 template <typename T>
@@ -213,7 +233,8 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
                                        cuvsCagraIndex_t index_c_ptr,
                                        DLManagedTensor* queries_tensor,
                                        DLManagedTensor* neighbors_tensor,
-                                       DLManagedTensor* distances_tensor)
+                                       DLManagedTensor* distances_tensor,
+                                       cuvsFilter filter)
 {
   return cuvs::core::translate_exceptions([=] {
     auto queries   = queries_tensor->dl_tensor;
@@ -236,11 +257,14 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
     RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries");
 
     if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) {
-      _search<float>(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor);
+      _search<float>(
+        res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter);
     } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) {
-      _search<int8_t>(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor);
+      _search<int8_t>(
+        res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter);
     } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) {
-      _search<uint8_t>(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor);
+      _search<uint8_t>(
+        res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter);
     } else {
       RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d",
                 queries.dtype.code,

cpp/test/neighbors/ann_cagra_c.cu

@@ -34,9 +34,14 @@ float queries[4][2] = {{0.48216683, 0.0428398},
                        {0.51260436, 0.2643005},
                        {0.05198065, 0.5789965}};
 
+uint32_t filter[1] = {0b1001};  // index 1 and 2 are removed
+
 uint32_t neighbors_exp[4] = {3, 0, 3, 1};
 float distances_exp[4]    = {0.03878258, 0.12472608, 0.04776672, 0.15224178};
 
+uint32_t neighbors_exp_filtered[4] = {3, 0, 3, 0};
+float distances_exp_filtered[4]    = {0.03878258, 0.12472608, 0.04776672, 0.59063464};
+
 TEST(CagraC, BuildSearch)
 {
   // create cuvsResources_t
@@ -109,10 +114,15 @@ TEST(CagraC, BuildSearch)
   distances_tensor.dl_tensor.shape              = distances_shape;
   distances_tensor.dl_tensor.strides            = nullptr;
 
+  cuvsFilter filter;
+  filter.type = NO_FILTER;
+  filter.addr = (uintptr_t)NULL;
+
   // search index
   cuvsCagraSearchParams_t search_params;
   cuvsCagraSearchParamsCreate(&search_params);
-  cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor);
+  cuvsCagraSearch(
+    res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter);
 
   // verify output
   ASSERT_TRUE(
@@ -126,3 +136,113 @@ TEST(CagraC, BuildSearch)
   cuvsCagraIndexDestroy(index);
   cuvsResourcesDestroy(res);
 }
+
+TEST(CagraC, BuildSearchFiltered)
+{
+  // create cuvsResources_t
+  cuvsResources_t res;
+  cuvsResourcesCreate(&res);
+  cudaStream_t stream;
+  cuvsStreamGet(res, &stream);
+
+  // create dataset DLTensor
+  DLManagedTensor dataset_tensor;
+  dataset_tensor.dl_tensor.data               = dataset;
+  dataset_tensor.dl_tensor.device.device_type = kDLCPU;
+  dataset_tensor.dl_tensor.ndim               = 2;
+  dataset_tensor.dl_tensor.dtype.code         = kDLFloat;
+  dataset_tensor.dl_tensor.dtype.bits         = 32;
+  dataset_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t dataset_shape[2]                    = {4, 2};
+  dataset_tensor.dl_tensor.shape              = dataset_shape;
+  dataset_tensor.dl_tensor.strides            = nullptr;
+
+  // create index
+  cuvsCagraIndex_t index;
+  cuvsCagraIndexCreate(&index);
+
+  // build index
+  cuvsCagraIndexParams_t build_params;
+  cuvsCagraIndexParamsCreate(&build_params);
+  cuvsCagraBuild(res, build_params, &dataset_tensor, index);
+
+  // create queries DLTensor
+  rmm::device_uvector<float> queries_d(4 * 2, stream);
+  raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream);
+
+  DLManagedTensor queries_tensor;
+  queries_tensor.dl_tensor.data               = queries_d.data();
+  queries_tensor.dl_tensor.device.device_type = kDLCUDA;
+  queries_tensor.dl_tensor.ndim               = 2;
+  queries_tensor.dl_tensor.dtype.code         = kDLFloat;
+  queries_tensor.dl_tensor.dtype.bits         = 32;
+  queries_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t queries_shape[2]                    = {4, 2};
+  queries_tensor.dl_tensor.shape              = queries_shape;
+  queries_tensor.dl_tensor.strides            = nullptr;
+
+  // create neighbors DLTensor
+  rmm::device_uvector<uint32_t> neighbors_d(4, stream);
+
+  DLManagedTensor neighbors_tensor;
+  neighbors_tensor.dl_tensor.data               = neighbors_d.data();
+  neighbors_tensor.dl_tensor.device.device_type = kDLCUDA;
+  neighbors_tensor.dl_tensor.ndim               = 2;
+  neighbors_tensor.dl_tensor.dtype.code         = kDLUInt;
+  neighbors_tensor.dl_tensor.dtype.bits         = 32;
+  neighbors_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t neighbors_shape[2]                    = {4, 1};
+  neighbors_tensor.dl_tensor.shape              = neighbors_shape;
+  neighbors_tensor.dl_tensor.strides            = nullptr;
+
+  // create distances DLTensor
+  rmm::device_uvector<float> distances_d(4, stream);
+
+  DLManagedTensor distances_tensor;
+  distances_tensor.dl_tensor.data               = distances_d.data();
+  distances_tensor.dl_tensor.device.device_type = kDLCUDA;
+  distances_tensor.dl_tensor.ndim               = 2;
+  distances_tensor.dl_tensor.dtype.code         = kDLFloat;
+  distances_tensor.dl_tensor.dtype.bits         = 32;
+  distances_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t distances_shape[2]                    = {4, 1};
+  distances_tensor.dl_tensor.shape              = distances_shape;
+  distances_tensor.dl_tensor.strides            = nullptr;
+
+  // create filter DLTensor
+  rmm::device_uvector<uint32_t> filter_d(1, stream);
+  raft::copy(filter_d.data(), filter, 1, stream);
+
+  cuvsFilter filter;
+
+  DLManagedTensor filter_tensor;
+  filter_tensor.dl_tensor.data               = filter_d.data();
+  filter_tensor.dl_tensor.device.device_type = kDLCUDA;
+  filter_tensor.dl_tensor.ndim               = 1;
+  filter_tensor.dl_tensor.dtype.code         = kDLUInt;
+  filter_tensor.dl_tensor.dtype.bits         = 32;
+  filter_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t filter_shape[1]                    = {1};
+  filter_tensor.dl_tensor.shape              = filter_shape;
+  filter_tensor.dl_tensor.strides            = nullptr;
+
+  filter.type = BITSET;
+  filter.addr = (uintptr_t)&filter_tensor;
+
+  // search index
+  cuvsCagraSearchParams_t search_params;
+  cuvsCagraSearchParamsCreate(&search_params);
+  cuvsCagraSearch(
+    res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter);
+  // verify output
+  ASSERT_TRUE(cuvs::devArrMatchHost(
+    neighbors_exp_filtered, neighbors_d.data(), 4, cuvs::Compare<uint32_t>()));
+  ASSERT_TRUE(cuvs::devArrMatchHost(
+    distances_exp_filtered, distances_d.data(), 4, cuvs::CompareApprox<float>(0.001f)));
+
+  // de-allocate index and res
+  cuvsCagraSearchParamsDestroy(search_params);
+  cuvsCagraIndexParamsDestroy(build_params);
+  cuvsCagraIndexDestroy(index);
+  cuvsResourcesDestroy(res);
+}

examples/c/src/cagra_c_example.c

@@ -67,9 +67,9 @@ void cagra_build_search_simple() {
   // Allocate memory for `queries`, `neighbors` and `distances` output
   uint32_t *neighbors;
   float *distances, *queries_d;
-  cuvsRMMAlloc(res, (void**) &queries_d, sizeof(float) * n_queries * n_cols);
-  cuvsRMMAlloc(res, (void**) &neighbors, sizeof(uint32_t) * n_queries * topk);
-  cuvsRMMAlloc(res, (void**) &distances, sizeof(float) * n_queries * topk);
+  cuvsRMMAlloc(res, (void **)&queries_d, sizeof(float) * n_queries * n_cols);
+  cuvsRMMAlloc(res, (void **)&neighbors, sizeof(uint32_t) * n_queries * topk);
+  cuvsRMMAlloc(res, (void **)&distances, sizeof(float) * n_queries * topk);
 
   // Use DLPack to represent `queries`, `neighbors` and `distances` as tensors
   cudaMemcpy(queries_d, queries, sizeof(float) * 4 * 2, cudaMemcpyDefault);
@@ -111,8 +111,12 @@ void cagra_build_search_simple() {
   cuvsCagraSearchParams_t search_params;
   cuvsCagraSearchParamsCreate(&search_params);
 
+  cuvsFilter filter;
+  filter.type = NO_FILTER;
+  filter.addr = (uintptr_t)NULL;
+
   cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor,
-                  &distances_tensor);
+                  &distances_tensor, filter);
 
   // print results
   uint32_t *neighbors_h =

python/cuvs/cuvs/neighbors/cagra/cagra.pxd

@@ -28,7 +28,7 @@ from libcpp cimport bool
 
 from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t
 from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor
-
+from cuvs.neighbors.filters.filters cimport cuvsFilter
 
 cdef extern from "cuvs/neighbors/cagra.h" nogil:
 
@@ -113,7 +113,8 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil:
                                 cuvsCagraIndex_t index,
                                 DLManagedTensor* queries,
                                 DLManagedTensor* neighbors,
-                                DLManagedTensor* distances) except +
+                                DLManagedTensor* distances,
+                                cuvsFilter filter) except +
 
     cuvsError_t cuvsCagraSerialize(cuvsResources_t res,
                                    const char * filename,

python/cuvs/cuvs/neighbors/cagra/cagra.pyx

@@ -47,6 +47,8 @@ from libc.stdint cimport (
 
 from cuvs.common.exceptions import check_cuvs
 
+from cuvs.neighbors.filters import no_filter
+
 
 cdef class CompressionParams:
     """
@@ -480,7 +482,8 @@ def search(SearchParams search_params,
            k,
            neighbors=None,
            distances=None,
-           resources=None):
+           resources=None,
+           filter=None):
     """
     Find the k nearest neighbors for each query.
 
@@ -553,6 +556,9 @@ def search(SearchParams search_params,
     _check_input_array(distances_cai, [np.dtype('float32')],
                        exp_rows=n_queries, exp_cols=k)
 
+    if filter is None:
+        filter = no_filter()
+
     cdef cuvsCagraSearchParams* params = &search_params.params
     cdef cydlpack.DLManagedTensor* queries_dlpack = \
         cydlpack.dlpack_c(queries_cai)
@@ -569,7 +575,8 @@ def search(SearchParams search_params,
             index.index,
             queries_dlpack,
             neighbors_dlpack,
-            distances_dlpack
+            distances_dlpack,
+            filter.prefilter
         ))
 
     return (distances, neighbors)

python/cuvs/cuvs/neighbors/filters/__init__.py

@@ -13,6 +13,6 @@
 # limitations under the License.
 
 
-from .filters import Prefilter, from_bitmap, no_filter
+from .filters import Prefilter, from_bitmap, from_bitset, no_filter
 
-__all__ = ["no_filter", "from_bitmap", "Prefilter"]
+__all__ = ["no_filter", "from_bitmap", "from_bitset", "Prefilter"]

python/cuvs/cuvs/neighbors/filters/filters.pyx

@@ -95,3 +95,52 @@ def from_bitmap(bitmap):
     filter.addr = <uintptr_t> bitmap_dlpack
 
     return Prefilter(filter, parent=bitmap)
+
+def from_bitset(bitset):
+    """
+    Create a pre-filter from an array with type of uint32.
+
+    Parameters
+    ----------
+    bitmap : numpy.ndarray
+        An array with type of `uint32` where each bit in the array corresponds
+        to if a sample and query pair is greenlit (not filtered) or filtered.
+        The array is row-major, meaning the bits are ordered by rows first.
+        Each bit in a `uint32` element represents a different sample-query
+        pair.
+
+        - Bit value of 1: The sample-query pair is greenlit (allowed).
+        - Bit value of 0: The sample-query pair is filtered.
+
+    Returns
+    -------
+    filter : cuvs.neighbors.filters.Prefilter
+        An instance of `Prefilter` that can be used to filter neighbors
+        based on the given bitmap.
+    {resources_docstring}
+
+    Examples
+    --------
+
+    >>> import cupy as cp
+    >>> import numpy as np
+    >>> from cuvs.neighbors import filters
+    >>>
+    >>> n_samples = 50000
+    >>> n_queries = 1000
+    >>>
+    >>> n_bitmap = np.ceil(n_samples * n_queries / 32).astype(int)
+    >>> bitmap = cp.random.randint(1, 100, size=(n_bitmap,), dtype=cp.uint32)
+    >>> prefilter = filters.from_bitmap(bitmap)
+    """
+    bitset_cai = wrap_array(bitset)
+    _check_input_array(bitset_cai, [np.dtype('uint32')])
+
+    cdef cydlpack.DLManagedTensor* bitset_dlpack = \
+        cydlpack.dlpack_c(bitset_cai)
+
+    cdef cuvsFilter filter
+    filter.type = BITSET
+    filter.addr = <uintptr_t> bitset_dlpack
+
+    return Prefilter(filter, parent=bitset)