[FEA] expose python & C API for prefiltered brute force

cpp/include/cuvs/neighbors/brute_force.h

@@ -18,6 +18,7 @@
 
 #include <cuvs/core/c_api.h>
 #include <cuvs/distance/distance.h>
+#include <cuvs/neighbors/prefilters.h>
 #include <dlpack/dlpack.h>
 #include <stdint.h>
 
@@ -135,9 +136,13 @@ cuvsError_t cuvsBruteForceBuild(cuvsResources_t res,
  * DLManagedTensor dataset;
  * DLManagedTensor queries;
  * DLManagedTensor neighbors;
+ * DLManagedTensor bitmap;
+ *
+ * cuvsPrefilter prefilter{(uintptr_t)&bitmap, BITMAP};
  *
  * // Search the `index` built using `cuvsBruteForceBuild`
- * cuvsError_t search_status = cuvsBruteForceSearch(res, index, &queries, &neighbors, &distances);
+ * cuvsError_t search_status = cuvsBruteForceSearch(res, index, &queries, &neighbors, &distances,
+ * prefilter);
  *
  * // de-allocate `res`
  * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
@@ -148,12 +153,15 @@ cuvsError_t cuvsBruteForceBuild(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 cuvsPrefilter input prefilter that can be used
+              to filter queries and neighbors based on the given bitmap.
  */
 cuvsError_t cuvsBruteForceSearch(cuvsResources_t res,
                                  cuvsBruteForceIndex_t index,
                                  DLManagedTensor* queries,
                                  DLManagedTensor* neighbors,
-                                 DLManagedTensor* distances);
+                                 DLManagedTensor* distances,
+                                 cuvsPrefilter prefilter);
 /**
  * @}
  */

cpp/include/cuvs/neighbors/prefilters.h

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuvs/core/c_api.h>
+#include <cuvs/distance/distance.h>
+#include <dlpack/dlpack.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @defgroup prefilter Prefilter APIs
+ * @brief APIs related to prefilter functionality.
+ * @{
+ */
+
+/**
+ * @brief Enum to denote prefilter type.
+ */
+enum cuvsPrefilterType {
+  /* No filter */
+  NO_FILTER,
+  /* Filter an index with a bitset */
+  BITSET,
+  /* Filter an index with a bitmap */
+  BITMAP
+};
+
+/**
+ * @brief Struct to hold address of cuvs::neighbor::prefilter and its type
+ *
+ */
+typedef struct {
+  uintptr_t addr;
+  enum cuvsPrefilterType type;
+} cuvsPrefilter;
+
+/**
+ * @}
+ */
+
+#ifdef __cplusplus
+}
+#endif

cpp/src/neighbors/brute_force_c.cpp

@@ -27,6 +27,7 @@
 #include <cuvs/core/interop.hpp>
 #include <cuvs/neighbors/brute_force.h>
 #include <cuvs/neighbors/brute_force.hpp>
+#include <cuvs/neighbors/prefilters.h>
 
 namespace {
 
@@ -53,20 +54,38 @@ void _search(cuvsResources_t res,
              cuvsBruteForceIndex index,
              DLManagedTensor* queries_tensor,
              DLManagedTensor* neighbors_tensor,
-             DLManagedTensor* distances_tensor)
+             DLManagedTensor* distances_tensor,
+             cuvsPrefilter prefilter)
 {
   auto res_ptr   = reinterpret_cast<raft::resources*>(res);
   auto index_ptr = reinterpret_cast<cuvs::neighbors::brute_force::index<T>*>(index.addr);
 
   using queries_mdspan_type   = raft::device_matrix_view<T const, int64_t, raft::row_major>;
   using neighbors_mdspan_type = raft::device_matrix_view<int64_t, int64_t, raft::row_major>;
   using distances_mdspan_type = raft::device_matrix_view<float, int64_t, raft::row_major>;
-  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);
+  using prefilter_mds_type    = raft::device_vector_view<const uint32_t, int64_t>;
+  using prefilter_opt_type    = cuvs::core::bitmap_view<const uint32_t, int64_t>;
+
+  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);
+
+  std::optional<cuvs::core::bitmap_view<const uint32_t, int64_t>> filter_opt;
+
+  if (prefilter.type == NO_FILTER) {
+    filter_opt = std::nullopt;
+  } else {
+    auto prefilter_ptr  = reinterpret_cast<DLManagedTensor*>(prefilter.addr);
+    auto prefilter_mds  = cuvs::core::from_dlpack<prefilter_mds_type>(prefilter_ptr);
+    auto prefilter_view = prefilter_opt_type((const uint32_t*)prefilter_mds.data_handle(),
+                                             queries_mds.extent(0),
+                                             index_ptr->dataset().extent(0));
+
+    filter_opt = std::make_optional<prefilter_opt_type>(prefilter_view);
+  }
 
   cuvs::neighbors::brute_force::search(
-    *res_ptr, *index_ptr, queries_mds, neighbors_mds, distances_mds, std::nullopt);
+    *res_ptr, *index_ptr, queries_mds, neighbors_mds, distances_mds, filter_opt);
 }
 
 }  // namespace
@@ -120,7 +139,8 @@ extern "C" cuvsError_t cuvsBruteForceSearch(cuvsResources_t res,
                                             cuvsBruteForceIndex_t index_c_ptr,
                                             DLManagedTensor* queries_tensor,
                                             DLManagedTensor* neighbors_tensor,
-                                            DLManagedTensor* distances_tensor)
+                                            DLManagedTensor* distances_tensor,
+                                            cuvsPrefilter prefilter)
 {
   return cuvs::core::translate_exceptions([=] {
     auto queries   = queries_tensor->dl_tensor;
@@ -143,7 +163,7 @@ extern "C" cuvsError_t cuvsBruteForceSearch(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, index, queries_tensor, neighbors_tensor, distances_tensor);
+      _search<float>(res, index, queries_tensor, neighbors_tensor, distances_tensor, prefilter);
     } else {
       RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d",
                 queries.dtype.code,