Add support for filtered search to the tiered_index code

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -27,18 +27,21 @@ dependencies:
 - go
 - graphviz
 - ipython
+- libboost-devel
 - libclang==20.1.4
 - libcublas-dev
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
 - librmm==25.8.*,>=0.0.0a0
 - make
+- maven
 - nccl>=2.19
 - ninja
 - numpy>=1.23,<3.0a0
 - numpydoc
 - openblas
+- openjdk=22.*
 - pre-commit
 - pylibraft==25.8.*,>=0.0.0a0
 - pytest-cov

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -27,18 +27,21 @@ dependencies:
 - go
 - graphviz
 - ipython
+- libboost-devel
 - libclang==20.1.4
 - libcublas-dev
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
 - librmm==25.8.*,>=0.0.0a0
 - make
+- maven
 - nccl>=2.19
 - ninja
 - numpy>=1.23,<3.0a0
 - numpydoc
 - openblas
+- openjdk=22.*
 - pre-commit
 - pylibraft==25.8.*,>=0.0.0a0
 - pytest-cov

cpp/include/cuvs/neighbors/cagra.h

@@ -590,7 +590,7 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res,
  * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
  *
  * // create an index with `cuvsCagraBuild`
- * cuvsCagraSerializeHnswlib(res, "/path/to/index", index);
+ * cuvsCagraSerializeToHnswlib(res, "/path/to/index", index);
  * @endcode
  *
  * @param[in] res cuvsResources_t opaque C handle

cpp/include/cuvs/neighbors/hnsw.h

@@ -430,7 +430,7 @@ cuvsError_t cuvsHnswSerialize(cuvsResources_t res, const char* filename, cuvsHns
  * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
  *
  * // create an index with `cuvsCagraBuild`
- * cuvsCagraSerializeHnswlib(res, "/path/to/index", index);
+ * cuvsCagraSerializeToHnswlib(res, "/path/to/index", index);
  *
  * // Load the serialized CAGRA index from file as an hnswlib index
  * // The index should have the same dtype as the one used to build CAGRA the index

cpp/include/cuvs/neighbors/tiered_index.h

@@ -246,6 +246,21 @@ cuvsError_t cuvsTieredIndexSearch(cuvsResources_t res,
 cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res,
                                   DLManagedTensor* new_vectors,
                                   cuvsTieredIndex_t index);
+/**
+ * @}
+ */
+/**
+ * @defgroup tiered_c_index_compact Tiered index compact
+ * @{
+ */
+/**
+ * @brief Compact the index
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[inout] index Tiered index to be compacted
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsTieredIndexCompact(cuvsResources_t res, cuvsTieredIndex_t index);
 /**
  * @}
  */

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -88,7 +88,8 @@ void tiled_brute_force_knn(const raft::resources& handle,
                            const uint32_t* filter_bits               = nullptr,
                            DistanceEpilogue distance_epilogue        = raft::identity_op(),
                            cuvs::neighbors::filtering::FilterType filter_type =
-                             cuvs::neighbors::filtering::FilterType::Bitmap)
+                             cuvs::neighbors::filtering::FilterType::Bitmap,
+                           size_t filter_col_offset = 0)
 {
   // Figure out the number of rows/cols to tile for
   size_t tile_rows = 0;
@@ -248,9 +249,9 @@ void tiled_brute_force_knn(const raft::resources& handle,
                          count,
                          count + current_query_size * current_centroid_size,
                          [=] __device__(IndexType idx) {
-                           IndexType row      = i + (idx / current_centroid_size);
-                           IndexType col      = j + (idx % current_centroid_size);
-                           IndexType g_idx    = row * n_cols + col;
+                           IndexType row   = i + (idx / current_centroid_size);
+                           IndexType col   = j + (idx % current_centroid_size) + filter_col_offset;
+                           IndexType g_idx = row * n_cols + col;
                            IndexType item_idx = (g_idx) >> 5;
                            uint32_t bit_idx   = (g_idx) & 31;
                            uint32_t filter    = filter_bits[item_idx];
@@ -596,12 +597,12 @@ void brute_force_search_filtered(
                  metric == cuvs::distance::DistanceType::L2Expanded ||
                  metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
                  metric == cuvs::distance::DistanceType::CosineExpanded,
-               "Only Euclidean, IP, and Cosine are supported!");
+               "Only Euclidean, IP, and Cosine distance are supported!");
 
   RAFT_EXPECTS(idx.has_norms() || !(metric == cuvs::distance::DistanceType::L2Expanded ||
                                     metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
                                     metric == cuvs::distance::DistanceType::CosineExpanded),
-               "Index must has norms when using Euclidean, IP, and Cosine!");
+               "Index must have norms when using Euclidean, IP, or Cosine distance!");
 
   IdxT n_queries                                     = queries.extent(0);
   IdxT n_dataset                                     = idx.dataset().extent(0);

cpp/src/neighbors/detail/tiered_index.cuh

@@ -33,6 +33,8 @@
 #include <cuvs/neighbors/knn_merge_parts.hpp>
 #include <cuvs/neighbors/tiered_index.hpp>
 
+#include "knn_brute_force.cuh"
+
 namespace cuvs::neighbors::tiered_index::detail {
 /**
   Storage for brute force based incremental indices
@@ -221,19 +223,62 @@ struct index_state {
         temp_distances.data_handle(), n_queries, k),
       sample_filter);
 
-    // search the bfknn index
     auto offset          = n_queries * k;
     auto bfknn_neighbors = raft::make_device_matrix_view<int64_t, int64_t>(
       temp_neighbors.data_handle() + offset, n_queries, k);
     auto bfknn_distances = raft::make_device_matrix_view<value_type, int64_t>(
       temp_distances.data_handle() + offset, n_queries, k);
-    brute_force::search(res,
-                        brute_force::search_params(),
-                        bfknn_index,
-                        queries,
-                        bfknn_neighbors,
-                        bfknn_distances,
-                        sample_filter);
+
+    switch (sample_filter.get_filter_type()) {
+      case filtering::FilterType::None: {
+        brute_force::search(res,
+                            brute_force::search_params(),
+                            bfknn_index,
+                            queries,
+                            bfknn_neighbors,
+                            bfknn_distances,
+                            sample_filter);
+        break;
+      }
+      case filtering::FilterType::Bitset: {
+        // We need to adjust the filter by the number of ann rows - which
+        // is a little tricky since this might not be aligned to the uint32_t
+        // bitset filter.  Use the detail api directly here which can support this
+        auto idx_norm =
+          bfknn_index.has_norms() ? const_cast<float*>(bfknn_index.norms().data_handle()) : nullptr;
+
+        auto actual_filter =
+          dynamic_cast<const cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>*>(
+            &sample_filter);
+        const uint32_t* filter_data = actual_filter->view().data();
+
+        neighbors::detail::tiled_brute_force_knn<float, int64_t, float>(
+          res,
+          queries.data_handle(),
+          bfknn_index.dataset().data_handle(),
+          n_queries,
+          bfknn_rows(),
+          storage->dim,
+          k,
+          bfknn_distances.data_handle(),
+          bfknn_neighbors.data_handle(),
+          build_params.metric,
+          2.0,
+          0,
+          0,
+          idx_norm,
+          nullptr,
+          filter_data,
+          raft::identity_op(),
+          filtering::FilterType::Bitset,
+          ann_rows());
+
+        break;
+      }
+      default: {
+        RAFT_FAIL("Only bitset filter is supported in tiered index");
+      }
+    }
 
     if (!distance::is_min_close(build_params.metric)) {
       // knn_merge_parts doesn't currently support InnerProduct distances etc

cpp/src/neighbors/tiered_index_c.cpp

@@ -158,6 +158,14 @@ void _extend(cuvsResources_t res, DLManagedTensor* new_vectors, cuvsTieredIndex
   tiered_index::extend(*res_ptr, vectors_mds, index_ptr);
 }
 
+template <typename UpstreamT>
+void _compact(cuvsResources_t res, cuvsTieredIndex index)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<tiered_index::index<UpstreamT>*>(index.addr);
+
+  tiered_index::compact(*res_ptr, index_ptr);
+}
 }  // namespace
 
 extern "C" cuvsError_t cuvsTieredIndexCreate(cuvsTieredIndex_t* index)
@@ -305,3 +313,25 @@ extern "C" cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res,
     }
   });
 }
+
+extern "C" cuvsError_t cuvsTieredIndexCompact(cuvsResources_t res, cuvsTieredIndex_t index_c_ptr)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto index = *index_c_ptr;
+    switch (index.algo) {
+      case CUVS_TIERED_INDEX_ALGO_CAGRA: {
+        _compact<cagra::index<float, uint32_t>>(res, index);
+        break;
+      }
+      case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: {
+        _compact<ivf_flat::index<float, int64_t>>(res, index);
+        break;
+      }
+      case CUVS_TIERED_INDEX_ALGO_IVF_PQ: {
+        _compact<ivf_pq::typed_index<float, int64_t>>(res, index);
+        break;
+      }
+      default: RAFT_FAIL("unsupported tiered index algorithm");
+    }
+  });
+}

dependencies.yaml

@@ -19,6 +19,7 @@ files:
       - depends_on_pylibraft
       - depends_on_nccl
       - docs
+      - java
       - rapids_build
       - run_py_cuvs
       - rust

python/cuvs/cuvs/neighbors/tiered_index/__init__.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 
 
-from .tiered_index import Index, IndexParams, build, extend, search
+from .tiered_index import Index, IndexParams, build, compact, extend, search
 
 __all__ = [
     "Index",
     "IndexParams",
     "build",
+    "compact",
     "extend",
     "search",
 ]

python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pxd

@@ -76,3 +76,6 @@ cdef extern from "cuvs/neighbors/tiered_index.h" nogil:
     cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res,
                                       DLManagedTensor* new_vectors,
                                       cuvsTieredIndex_t index)
+
+    cuvsError_t cuvsTieredIndexCompact(cuvsResources_t res,
+                                       cuvsTieredIndex_t index)

python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pyx

@@ -387,3 +387,33 @@ def extend(Index index, new_vectors, resources=None):
         ))
 
     return index
+
+
+@auto_sync_resources
+def compact(Index index, resources=None):
+    """
+    Compact the index
+
+    This function takes any data that has been added incrementally, and ensures
+    that it been added to the ANN index.
+
+    Parameters
+    ----------
+    index : tiered_index.Index
+        Trained tiered_index object.
+    {resources_docstring}
+
+    Returns
+    -------
+    index: py:class:`cuvs.neighbors.tiered_index.Index`
+    """
+
+    cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
+
+    with cuda_interruptible():
+        check_cuvs(cuvsTieredIndexCompact(
+            res,
+            index.index
+        ))
+
+    return index