From 7562c5a4e281ce3828632ee2a234e3ba7c0e8a49 Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Wed, 11 Jun 2025 15:51:32 -0700 Subject: [PATCH 1/4] Add support for filtered search to the tiered_index code The tiered_index code wasn't working with filtered search - since the filter passed to the brute force portion wasn't translating for the number of rows seen already. Fix and add a python level unittest that would have caught this. Also: * expose 'tiered_index.compact' to python * add maven as a dependency of the java bindings --- cpp/include/cuvs/neighbors/cagra.h | 2 +- cpp/include/cuvs/neighbors/hnsw.h | 2 +- cpp/include/cuvs/neighbors/tiered_index.h | 15 +++ cpp/src/neighbors/detail/knn_brute_force.cuh | 13 +- cpp/src/neighbors/detail/tiered_index.cuh | 61 ++++++++-- cpp/src/neighbors/tiered_index_c.cpp | 30 +++++ dependencies.yaml | 1 + .../cuvs/neighbors/tiered_index/__init__.py | 3 +- .../neighbors/tiered_index/tiered_index.pxd | 3 + .../neighbors/tiered_index/tiered_index.pyx | 30 +++++ python/cuvs/cuvs/tests/test_tiered_index.py | 112 ++++++++++++++++++ 11 files changed, 255 insertions(+), 17 deletions(-) create mode 100644 python/cuvs/cuvs/tests/test_tiered_index.py diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 5959124870..35165c6f90 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/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 diff --git a/cpp/include/cuvs/neighbors/hnsw.h b/cpp/include/cuvs/neighbors/hnsw.h index b5f129e96d..096fe2d6dc 100644 --- a/cpp/include/cuvs/neighbors/hnsw.h +++ b/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 diff --git a/cpp/include/cuvs/neighbors/tiered_index.h b/cpp/include/cuvs/neighbors/tiered_index.h index a5bf9b3b9c..dc55b8cfa0 100644 --- a/cpp/include/cuvs/neighbors/tiered_index.h +++ b/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); /** * @} */ diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index 9202fe1c89..2045879fa2 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/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; @@ -261,9 +262,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]; @@ -609,12 +610,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); diff --git a/cpp/src/neighbors/detail/tiered_index.cuh b/cpp/src/neighbors/detail/tiered_index.cuh index 2251448aab..bf0ebceac7 100644 --- a/cpp/src/neighbors/detail/tiered_index.cuh +++ b/cpp/src/neighbors/detail/tiered_index.cuh @@ -32,6 +32,8 @@ #include #include +#include "knn_brute_force.cuh" + namespace cuvs::neighbors::tiered_index::detail { /** Storage for brute force based incremental indices @@ -227,19 +229,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( temp_neighbors.data_handle() + offset, n_queries, k); auto bfknn_distances = raft::make_device_matrix_view( 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(bfknn_index.norms().data_handle()) : nullptr; + + auto actual_filter = + dynamic_cast*>( + &sample_filter); + const uint32_t* filter_data = actual_filter->view().data(); + + neighbors::detail::tiled_brute_force_knn( + 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 diff --git a/cpp/src/neighbors/tiered_index_c.cpp b/cpp/src/neighbors/tiered_index_c.cpp index 196b199e10..a6882d6089 100644 --- a/cpp/src/neighbors/tiered_index_c.cpp +++ b/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 +void _compact(cuvsResources_t res, cuvsTieredIndex index) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(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>(res, index); + break; + } + case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { + _compact>(res, index); + break; + } + case CUVS_TIERED_INDEX_ALGO_IVF_PQ: { + _compact>(res, index); + break; + } + default: RAFT_FAIL("unsupported tiered index algorithm"); + } + }); +} diff --git a/dependencies.yaml b/dependencies.yaml index c7a2daa9e7..2050467874 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -120,6 +120,7 @@ files: - depends_on_libraft - depends_on_nccl - java + - maven - rapids_build py_build_libcuvs: output: pyproject diff --git a/python/cuvs/cuvs/neighbors/tiered_index/__init__.py b/python/cuvs/cuvs/neighbors/tiered_index/__init__.py index d6440431f3..03b5c58afe 100644 --- a/python/cuvs/cuvs/neighbors/tiered_index/__init__.py +++ b/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", ] diff --git a/python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pxd b/python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pxd index 9c0e35c609..060a9c6f8c 100644 --- a/python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pxd +++ b/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) diff --git a/python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pyx b/python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pyx index 53766f99c8..8365f8496b 100644 --- a/python/cuvs/cuvs/neighbors/tiered_index/tiered_index.pyx +++ b/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 = resources.get_c_obj() + + with cuda_interruptible(): + check_cuvs(cuvsTieredIndexCompact( + res, + index.index + )) + + return index diff --git a/python/cuvs/cuvs/tests/test_tiered_index.py b/python/cuvs/cuvs/tests/test_tiered_index.py new file mode 100644 index 0000000000..c885cbc161 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_tiered_index.py @@ -0,0 +1,112 @@ +# Copyright (c) 2025, 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. +# + +import numpy as np +import pytest +from pylibraft.common import device_ndarray + +from cuvs.neighbors import ( + brute_force, + cagra, + filters, + ivf_flat, + ivf_pq, + tiered_index, +) +from cuvs.tests.ann_utils import calc_recall, create_sparse_bitset + + +@pytest.mark.parametrize("n_dataset_rows", [1024, 10000]) +@pytest.mark.parametrize("n_query_rows", [10]) +@pytest.mark.parametrize("n_cols", [10]) +@pytest.mark.parametrize("k", [8, 16]) +@pytest.mark.parametrize("dtype", ["float32"]) +@pytest.mark.parametrize( + "metric", + [ + "sqeuclidean", + "inner_product", + ], +) +@pytest.mark.parametrize( + "algo", + [ + "cagra", + "ivf_flat", + "ivf_pq", + ], +) +@pytest.mark.parametrize("filter_type", ["bitset_filter", "no_filter"]) +def test_tiered_index( + n_dataset_rows, n_query_rows, n_cols, k, dtype, metric, algo, filter_type +): + dataset = np.random.random_sample((n_dataset_rows, n_cols)).astype(dtype) + queries = np.random.random_sample((n_query_rows, n_cols)).astype(dtype) + + indices = np.zeros((n_query_rows, k), dtype="int64") + distances = np.zeros((n_query_rows, k), dtype="float32") + + dataset_device = device_ndarray(dataset) + queries_device = device_ndarray(queries) + indices_device = device_ndarray(indices) + distances_device = device_ndarray(distances) + + # build with half the dataset, then extend with the other half + dataset_1_device = device_ndarray(dataset[: n_dataset_rows // 2, :]) + dataset_2_device = device_ndarray(dataset[n_dataset_rows // 2 :, :]) + + build_params = tiered_index.IndexParams( + metric=metric, algo=algo, min_ann_rows=1000 + ) + index = tiered_index.build(build_params, dataset_1_device) + index = tiered_index.extend(index, dataset_2_device) + + if filter_type == "bitset_filter": + sparsity = 0.5 + bitset = create_sparse_bitset(n_dataset_rows, sparsity) + bitset_device = device_ndarray(bitset) + prefilter = filters.from_bitset(bitset_device) + + # compact the index until we fully support filtered search here + # index = tiered_index.compact(index) + else: + prefilter = filters.no_filter() + + if algo == "cagra": + search_params = cagra.SearchParams() + elif algo == "ivf_flat": + search_params = ivf_flat.SearchParams(n_probes=64) + elif algo == "ivf_pq": + search_params = ivf_pq.SearchParams(n_probes=64) + + ret_distances, ret_indices = tiered_index.search( + search_params, + index, + queries_device, + k, + neighbors=indices_device, + distances=distances_device, + filter=prefilter, + ) + + bfknn_index = brute_force.build(dataset_device, metric) + groundtruth_neighbors, groundtruth_indices = brute_force.search( + bfknn_index, queries_device, k, prefilter=prefilter + ) + + ret_indices = ret_indices.copy_to_host() + groundtruth_indices = groundtruth_indices.copy_to_host() + recall = calc_recall(ret_indices, groundtruth_indices) + assert recall > 0.7 From e6271ba56eca10d7e23cae287d97cb0e97d4acaa Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Wed, 11 Jun 2025 17:26:58 -0700 Subject: [PATCH 2/4] Add java to 'all' conda environment --- conda/environments/all_cuda-128_arch-aarch64.yaml | 3 +++ conda/environments/all_cuda-128_arch-x86_64.yaml | 3 +++ dependencies.yaml | 2 +- 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/conda/environments/all_cuda-128_arch-aarch64.yaml b/conda/environments/all_cuda-128_arch-aarch64.yaml index c1adb7062f..89960ae93c 100644 --- a/conda/environments/all_cuda-128_arch-aarch64.yaml +++ b/conda/environments/all_cuda-128_arch-aarch64.yaml @@ -27,6 +27,7 @@ dependencies: - go - graphviz - ipython +- libboost-devel - libclang==20.1.4 - libcublas-dev - libcurand-dev @@ -34,11 +35,13 @@ dependencies: - 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 diff --git a/conda/environments/all_cuda-128_arch-x86_64.yaml b/conda/environments/all_cuda-128_arch-x86_64.yaml index ec6a42d8b4..373a279443 100644 --- a/conda/environments/all_cuda-128_arch-x86_64.yaml +++ b/conda/environments/all_cuda-128_arch-x86_64.yaml @@ -27,6 +27,7 @@ dependencies: - go - graphviz - ipython +- libboost-devel - libclang==20.1.4 - libcublas-dev - libcurand-dev @@ -34,11 +35,13 @@ dependencies: - 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 diff --git a/dependencies.yaml b/dependencies.yaml index 2050467874..df73d012b7 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -19,6 +19,7 @@ files: - depends_on_pylibraft - depends_on_nccl - docs + - java - rapids_build - run_py_cuvs - rust @@ -120,7 +121,6 @@ files: - depends_on_libraft - depends_on_nccl - java - - maven - rapids_build py_build_libcuvs: output: pyproject From 32be4d010c70b651c16d65fb2556ef627b5e279d Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Wed, 19 Nov 2025 14:58:27 -0800 Subject: [PATCH 3/4] style fix --- python/cuvs/cuvs/tests/test_tiered_index.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/python/cuvs/cuvs/tests/test_tiered_index.py b/python/cuvs/cuvs/tests/test_tiered_index.py index c885cbc161..f326061ffa 100644 --- a/python/cuvs/cuvs/tests/test_tiered_index.py +++ b/python/cuvs/cuvs/tests/test_tiered_index.py @@ -1,17 +1,5 @@ -# Copyright (c) 2025, 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. -# +# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 import numpy as np import pytest From d26b609861834d73db4fdb32550befb115ffe3bd Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Wed, 19 Nov 2025 21:29:54 -0800 Subject: [PATCH 4/4] fix merge error --- c/include/cuvs/neighbors/tiered_index.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/c/include/cuvs/neighbors/tiered_index.h b/c/include/cuvs/neighbors/tiered_index.h index 2cc27ba920..c210e2bc8b 100644 --- a/c/include/cuvs/neighbors/tiered_index.h +++ b/c/include/cuvs/neighbors/tiered_index.h @@ -251,6 +251,9 @@ cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res, * @return cuvsError_t */ cuvsError_t cuvsTieredIndexCompact(cuvsResources_t res, cuvsTieredIndex_t index); +/** + * @} + */ /** * @defgroup tiered_c_index_merge Tiered index merge