Add support for filtered search to the tiered_index code

c/include/cuvs/neighbors/cagra.h

@@ -783,7 +783,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

c/include/cuvs/neighbors/hnsw.h

@@ -419,7 +419,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

c/include/cuvs/neighbors/tiered_index.h

@@ -239,6 +239,22 @@ cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res,
  * @}
  */
 
+/**
+ * @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);
+/**
+ * @}
+ */
+
 /**
  * @defgroup tiered_c_index_merge Tiered index merge
  * @{
@@ -258,7 +274,6 @@ cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res,
                                  cuvsTieredIndex_t* indices,
                                  size_t num_indices,
                                  cuvsTieredIndex_t output_index);
-
 /**
  * @}
  */

c/src/neighbors/tiered_index.cpp

@@ -204,6 +204,14 @@ void _merge(cuvsResources_t res,
   output_index->algo  = indices[0]->algo;
 }
 
+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)
@@ -352,6 +360,28 @@ 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");
+    }
+  });
+}
+
 extern "C" cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res,
                                             cuvsTieredIndexParams_t params,
                                             cuvsTieredIndex_t* indices,

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

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

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

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

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

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

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

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

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -77,7 +77,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;
@@ -237,9 +238,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];
@@ -585,12 +586,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

@@ -22,6 +22,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
@@ -210,19 +212,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

dependencies.yaml

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

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

@@ -2,12 +2,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-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

@@ -65,3 +65,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

@@ -376,3 +376,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