diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h index 5c6162041..c46a9b2cd 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.h +++ b/cpp/include/cuvs/neighbors/ivf_flat.h @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -267,13 +268,17 @@ cuvsError_t cuvsIvfFlatBuild(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] filter cuvsFilter input filter that can be used + to filter queries and neighbors based on the given bitset. */ cuvsError_t cuvsIvfFlatSearch(cuvsResources_t res, cuvsIvfFlatSearchParams_t search_params, cuvsIvfFlatIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, - DLManagedTensor* distances); + DLManagedTensor* distances, + cuvsFilter filter); + /** * @} */ diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp index 2acc6b678..b38f1808d 100644 --- a/cpp/src/neighbors/ivf_flat_c.cpp +++ b/cpp/src/neighbors/ivf_flat_c.cpp @@ -67,7 +67,8 @@ void _search(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* queries_tensor, DLManagedTensor* neighbors_tensor, - DLManagedTensor* distances_tensor) + DLManagedTensor* distances_tensor, + cuvsFilter filter) { auto res_ptr = reinterpret_cast(res); auto index_ptr = reinterpret_cast*>(index.addr); @@ -82,8 +83,27 @@ void _search(cuvsResources_t res, auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); auto distances_mds = cuvs::core::from_dlpack(distances_tensor); - cuvs::neighbors::ivf_flat::search( - *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); + if (filter.type == NO_FILTER) { + cuvs::neighbors::ivf_flat::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; + auto removed_indices_tensor = reinterpret_cast(filter.addr); + auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); + cuvs::core::bitset_view removed_indices_bitset(removed_indices, + index_ptr->size()); + auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); + cuvs::neighbors::ivf_flat::search(*res_ptr, + search_params, + *index_ptr, + queries_mds, + neighbors_mds, + distances_mds, + bitset_filter_obj); + + } else { + RAFT_FAIL("Unsupported filter type: BITMAP"); + } } template @@ -179,7 +199,9 @@ extern "C" cuvsError_t cuvsIvfFlatSearch(cuvsResources_t res, cuvsIvfFlatIndex_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; @@ -203,13 +225,13 @@ extern "C" cuvsError_t cuvsIvfFlatSearch(cuvsResources_t res, if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); } else { RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", queries.dtype.code, diff --git a/cpp/tests/neighbors/run_ivf_flat_c.c b/cpp/tests/neighbors/run_ivf_flat_c.c index 8cd79c91f..d58f11bf3 100644 --- a/cpp/tests/neighbors/run_ivf_flat_c.c +++ b/cpp/tests/neighbors/run_ivf_flat_c.c @@ -91,12 +91,16 @@ void run_ivf_flat(int64_t n_rows, distances_tensor.dl_tensor.shape = distances_shape; distances_tensor.dl_tensor.strides = NULL; + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + // search index cuvsIvfFlatSearchParams_t search_params; cuvsIvfFlatSearchParamsCreate(&search_params); search_params->n_probes = n_probes; cuvsIvfFlatSearch( - res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); // de-allocate index and res cuvsIvfFlatSearchParamsDestroy(search_params); diff --git a/examples/c/src/ivf_flat_c_example.c b/examples/c/src/ivf_flat_c_example.c index 2121ca35e..510d624f3 100644 --- a/examples/c/src/ivf_flat_c_example.c +++ b/examples/c/src/ivf_flat_c_example.c @@ -67,8 +67,12 @@ void ivf_flat_build_search_simple(cuvsResources_t *res, DLManagedTensor * datase search_params->n_probes = 50; // Search the `index` built using `ivfFlatBuild` + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + cuvsError_t search_status = cuvsIvfFlatSearch(*res, search_params, index, - queries_tensor, &neighbors_tensor, &distances_tensor); + queries_tensor, &neighbors_tensor, &distances_tensor, filter); if (build_status != CUVS_SUCCESS) { printf("%s.\n", cuvsGetLastErrorText()); } @@ -165,8 +169,11 @@ void ivf_flat_build_extend_search(cuvsResources_t *res, DLManagedTensor * trains search_params->n_probes = 10; // Search the `index` built using `ivfFlatBuild` + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; cuvsError_t search_status = cuvsIvfFlatSearch(*res, search_params, index, - queries_tensor, &neighbors_tensor, &distances_tensor); + queries_tensor, &neighbors_tensor, &distances_tensor, filter); if (search_status != CUVS_SUCCESS) { printf("%s.\n", cuvsGetLastErrorText()); exit(-1); diff --git a/go/ivf_flat/ivf_flat.go b/go/ivf_flat/ivf_flat.go index 3330eb95e..61e17172f 100644 --- a/go/ivf_flat/ivf_flat.go +++ b/go/ivf_flat/ivf_flat.go @@ -67,6 +67,10 @@ func SearchIndex[T any](Resources cuvs.Resource, params *SearchParams, index *Iv if !index.trained { return errors.New("index needs to be built before calling search") } + prefilter := C.cuvsFilter{ + addr: 0, + _type: C.NO_FILTER, + } - return cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsIvfFlatSearch(C.cuvsResources_t(Resources.Resource), params.params, index.index, (*C.DLManagedTensor)(unsafe.Pointer(queries.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(neighbors.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(distances.C_tensor))))) + return cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsIvfFlatSearch(C.cuvsResources_t(Resources.Resource), params.params, index.index, (*C.DLManagedTensor)(unsafe.Pointer(queries.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(neighbors.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(distances.C_tensor)), prefilter))) } diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd index 96bc557e6..f2bd6a9b1 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd @@ -21,6 +21,7 @@ from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor from cuvs.distance_type cimport cuvsDistanceType +from cuvs.neighbors.filters.filters cimport cuvsFilter cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: @@ -71,7 +72,8 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: cuvsIvfFlatIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, - DLManagedTensor* distances) except + + DLManagedTensor* distances, + cuvsFilter filter) except + cuvsError_t cuvsIvfFlatSerialize(cuvsResources_t res, const char * filename, diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index 7a169e1a0..437499c1e 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -34,6 +34,7 @@ from pylibraft.common.interruptible import cuda_interruptible from cuvs.distance import DISTANCE_TYPES from cuvs.neighbors.common import _check_input_array +from cuvs.neighbors.filters import no_filter from libc.stdint cimport ( int8_t, @@ -274,7 +275,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. @@ -293,6 +295,8 @@ def search(SearchParams search_params, distances : Optional CUDA array interface compliant matrix shape (n_queries, k) If supplied, the distances to the neighbors will be written here in-place. (default None) + filter: Optional cuvs.neighbors.cuvsFilter can be used to filter + neighbors based on a given bitset. (default None) {resources_docstring} Examples @@ -339,6 +343,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 cuvsIvfFlatSearchParams* params = search_params.params cdef cuvsError_t search_status cdef cydlpack.DLManagedTensor* queries_dlpack = \ @@ -356,7 +363,8 @@ def search(SearchParams search_params, index.index, queries_dlpack, neighbors_dlpack, - distances_dlpack + distances_dlpack, + filter.prefilter )) return (distances, neighbors) diff --git a/python/cuvs/cuvs/tests/ann_utils.py b/python/cuvs/cuvs/tests/ann_utils.py index 60db7f327..b8f5d0bb0 100644 --- a/python/cuvs/cuvs/tests/ann_utils.py +++ b/python/cuvs/cuvs/tests/ann_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2024, NVIDIA CORPORATION. +# Copyright (c) 2023-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. @@ -13,6 +13,10 @@ # limitations under the License. import numpy as np +from pylibraft.common import device_ndarray +from sklearn.neighbors import NearestNeighbors + +from cuvs.neighbors import filters def generate_data(shape, dtype): @@ -33,3 +37,89 @@ def calc_recall(ann_idx, true_nn_idx): n += np.intersect1d(ann_idx[i, :], true_nn_idx[i, :]).size recall = n / ann_idx.size return recall + + +def create_sparse_bitset(n_size, sparsity): + bits_per_uint32 = 32 + num_bits = n_size + num_uint32s = (num_bits + bits_per_uint32 - 1) // bits_per_uint32 + num_ones = int(num_bits * sparsity) + + array = np.zeros(num_uint32s, dtype=np.uint32) + indices = np.random.choice(num_bits, num_ones, replace=False) + + for index in indices: + i = index // bits_per_uint32 + bit_position = index % bits_per_uint32 + array[i] |= 1 << bit_position + + return array + + +def run_filtered_search_test( + search_module, + sparsity, + n_rows=10000, + n_cols=10, + n_queries=10, + k=10, +): + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + bitset = create_sparse_bitset(n_rows, sparsity) + + dataset_device = device_ndarray(dataset) + queries_device = device_ndarray(queries) + bitset_device = device_ndarray(bitset) + + build_params = search_module.IndexParams() + index = search_module.build(build_params, dataset_device) + + filter_ = filters.from_bitset(bitset_device) + + search_params = search_module.SearchParams() + ret_distances, ret_indices = search_module.search( + search_params, + index, + queries_device, + k, + filter=filter_, + ) + + # Convert bitset to bool array for validation + bitset_as_uint8 = bitset.view(np.uint8) + bool_filter = np.unpackbits(bitset_as_uint8) + bool_filter = bool_filter.reshape(-1, 4, 8) + bool_filter = np.flip(bool_filter, axis=2) + bool_filter = bool_filter.reshape(-1)[:n_rows] + bool_filter = np.logical_not(bool_filter) # Flip so True means filtered + + # Get filtered dataset for reference calculation + non_filtered_mask = ~bool_filter + filtered_dataset = dataset[non_filtered_mask] + + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric="euclidean" + ) + nn_skl.fit(filtered_dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + actual_indices = ret_indices.copy_to_host() + + filtered_idx_map = ( + np.cumsum(~bool_filter) - 1 + ) # -1 because cumsum starts at 1 + + # Map ANN indices to filtered space + mapped_actual_indices = np.take( + filtered_idx_map, actual_indices, mode="clip" + ) + + filtered_indices = np.where(bool_filter)[0] + for i in range(n_queries): + assert not np.intersect1d(filtered_indices, actual_indices[i]).size + + recall = calc_recall(mapped_actual_indices, skl_idx) + + assert recall > 0.7 diff --git a/python/cuvs/cuvs/tests/test_cagra.py b/python/cuvs/cuvs/tests/test_cagra.py index 831ef11f7..f3de488da 100644 --- a/python/cuvs/cuvs/tests/test_cagra.py +++ b/python/cuvs/cuvs/tests/test_cagra.py @@ -19,8 +19,12 @@ from sklearn.neighbors import NearestNeighbors from sklearn.preprocessing import normalize -from cuvs.neighbors import cagra, filters -from cuvs.tests.ann_utils import calc_recall, generate_data +from cuvs.neighbors import cagra +from cuvs.tests.ann_utils import ( + calc_recall, + generate_data, + run_filtered_search_test, +) def run_cagra_build_search_test( @@ -139,97 +143,9 @@ def test_cagra_dataset_dtype_host_device( ) -def create_sparse_bitset(n_size, sparsity): - bits_per_uint32 = 32 - num_bits = n_size - num_uint32s = (num_bits + bits_per_uint32 - 1) // bits_per_uint32 - num_ones = int(num_bits * sparsity) - - array = np.zeros(num_uint32s, dtype=np.uint32) - indices = np.random.choice(num_bits, num_ones, replace=False) - - for index in indices: - i = index // bits_per_uint32 - bit_position = index % bits_per_uint32 - array[i] |= 1 << bit_position - - return array - - @pytest.mark.parametrize("sparsity", [0.2, 0.5, 0.7, 1.0]) -def test_filtered_cagra( - sparsity, - n_rows=10000, - n_cols=10, - n_queries=10, - k=10, -): - dataset = generate_data((n_rows, n_cols), np.float32) - queries = generate_data((n_queries, n_cols), np.float32) - - bitset = create_sparse_bitset(n_rows, sparsity) - - dataset_device = device_ndarray(dataset) - queries_device = device_ndarray(queries) - bitset_device = device_ndarray(bitset) - - build_params = cagra.IndexParams() - index = cagra.build(build_params, dataset_device) - - filter_ = filters.from_bitset(bitset_device) - - out_idx = np.zeros((n_queries, k), dtype=np.uint32) - out_dist = np.zeros((n_queries, k), dtype=np.float32) - out_idx_device = device_ndarray(out_idx) - out_dist_device = device_ndarray(out_dist) - - search_params = cagra.SearchParams() - ret_distances, ret_indices = cagra.search( - search_params, - index, - queries_device, - k, - neighbors=out_idx_device, - distances=out_dist_device, - filter=filter_, - ) - - # Convert bitset to bool array for validation - bitset_as_uint8 = bitset.view(np.uint8) - bool_filter = np.unpackbits(bitset_as_uint8) - bool_filter = bool_filter.reshape(-1, 4, 8) - bool_filter = np.flip(bool_filter, axis=2) - bool_filter = bool_filter.reshape(-1)[:n_rows] - bool_filter = np.logical_not(bool_filter) # Flip so True means filtered - - # Get filtered dataset for reference calculation - non_filtered_mask = ~bool_filter - filtered_dataset = dataset[non_filtered_mask] - - nn_skl = NearestNeighbors( - n_neighbors=k, algorithm="brute", metric="euclidean" - ) - nn_skl.fit(filtered_dataset) - skl_idx = nn_skl.kneighbors(queries, return_distance=False) - - actual_indices = out_idx_device.copy_to_host() - - filtered_idx_map = ( - np.cumsum(~bool_filter) - 1 - ) # -1 because cumsum starts at 1 - - # Map CAGRA indices to filtered space - mapped_actual_indices = np.take( - filtered_idx_map, actual_indices, mode="clip" - ) - - filtered_indices = np.where(bool_filter)[0] - for i in range(n_queries): - assert not np.intersect1d(filtered_indices, actual_indices[i]).size - - recall = calc_recall(mapped_actual_indices, skl_idx) - - assert recall > 0.7 +def test_filtered_cagra(sparsity): + run_filtered_search_test(cagra, sparsity) @pytest.mark.parametrize( diff --git a/python/cuvs/cuvs/tests/test_ivf_flat.py b/python/cuvs/cuvs/tests/test_ivf_flat.py index c3ec0252a..6b89041fa 100644 --- a/python/cuvs/cuvs/tests/test_ivf_flat.py +++ b/python/cuvs/cuvs/tests/test_ivf_flat.py @@ -20,7 +20,11 @@ from sklearn.preprocessing import normalize from cuvs.neighbors import ivf_flat -from cuvs.tests.ann_utils import calc_recall, generate_data +from cuvs.tests.ann_utils import ( + calc_recall, + generate_data, + run_filtered_search_test, +) def run_ivf_flat_build_search_test( @@ -129,3 +133,8 @@ def test_extend(dtype): dtype=dtype, add_data_on_build=False, ) + + +@pytest.mark.parametrize("sparsity", [0.5, 0.7, 1.0]) +def test_filtered_ivf_flat(sparsity): + run_filtered_search_test(ivf_flat, sparsity) diff --git a/rust/cuvs/src/ivf_flat/index.rs b/rust/cuvs/src/ivf_flat/index.rs index b1462a0e6..f66904142 100644 --- a/rust/cuvs/src/ivf_flat/index.rs +++ b/rust/cuvs/src/ivf_flat/index.rs @@ -78,6 +78,11 @@ impl Index { distances: &ManagedTensor, ) -> Result<()> { unsafe { + let prefilter = ffi::cuvsFilter { + addr: 0, + type_: ffi::cuvsFilterType::NO_FILTER, + }; + check_cuvs(ffi::cuvsIvfFlatSearch( res.0, params.0, @@ -85,6 +90,7 @@ impl Index { queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), + prefilter, )) } }