diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index f7f58a19c..81a7f0215 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include #include @@ -388,13 +389,16 @@ cuvsError_t cuvsCagraBuild(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 cuvsFilter input prefilter that can be used + to filter queries and neighbors based on the given bitset. */ cuvsError_t cuvsCagraSearch(cuvsResources_t res, cuvsCagraSearchParams_t params, cuvsCagraIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, - DLManagedTensor* distances); + DLManagedTensor* distances, + cuvsFilter filter); /** * @} diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 02b7a566e..20cbae7d0 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include @@ -92,7 +93,8 @@ void _search(cuvsResources_t res, cuvsCagraIndex 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); @@ -118,8 +120,26 @@ void _search(cuvsResources_t res, auto queries_mds = cuvs::core::from_dlpack(queries_tensor); auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); auto distances_mds = cuvs::core::from_dlpack(distances_tensor); - cuvs::neighbors::cagra::search( - *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); + if (filter.type == NO_FILTER) { + cuvs::neighbors::cagra::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->dataset().extent(0)); + auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); + cuvs::neighbors::cagra::search(*res_ptr, + search_params, + *index_ptr, + queries_mds, + neighbors_mds, + distances_mds, + bitset_filter_obj); + } else { + RAFT_FAIL("Unsupported filter type: BITMAP"); + } } template @@ -214,7 +234,8 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, cuvsCagraIndex_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; @@ -237,11 +258,14 @@ extern "C" cuvsError_t cuvsCagraSearch(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(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); + _search( + 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); + _search( + 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); + _search( + 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/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index 599d2d842..7315813cc 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -34,9 +34,14 @@ float queries[4][2] = {{0.48216683, 0.0428398}, {0.51260436, 0.2643005}, {0.05198065, 0.5789965}}; +uint32_t filter[1] = {0b1001}; // index 1 and 2 are removed + uint32_t neighbors_exp[4] = {3, 0, 3, 1}; float distances_exp[4] = {0.03878258, 0.12472608, 0.04776672, 0.15224178}; +uint32_t neighbors_exp_filtered[4] = {3, 0, 3, 0}; +float distances_exp_filtered[4] = {0.03878258, 0.12472608, 0.04776672, 0.59063464}; + TEST(CagraC, BuildSearch) { // create cuvsResources_t @@ -109,10 +114,15 @@ TEST(CagraC, BuildSearch) distances_tensor.dl_tensor.shape = distances_shape; distances_tensor.dl_tensor.strides = nullptr; + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + // search index cuvsCagraSearchParams_t search_params; cuvsCagraSearchParamsCreate(&search_params); - cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + cuvsCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); // verify output ASSERT_TRUE( @@ -126,3 +136,113 @@ TEST(CagraC, BuildSearch) cuvsCagraIndexDestroy(index); cuvsResourcesDestroy(res); } + +TEST(CagraC, BuildSearchFiltered) +{ + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + cudaStream_t stream; + cuvsStreamGet(res, &stream); + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create index + cuvsCagraIndex_t index; + cuvsCagraIndexCreate(&index); + + // build index + cuvsCagraIndexParams_t build_params; + cuvsCagraIndexParamsCreate(&build_params); + cuvsCagraBuild(res, build_params, &dataset_tensor, index); + + // create queries DLTensor + rmm::device_uvector queries_d(4 * 2, stream); + raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream); + + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries_d.data(); + queries_tensor.dl_tensor.device.device_type = kDLCUDA; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + rmm::device_uvector neighbors_d(4, stream); + + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_d.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCUDA; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 32; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + rmm::device_uvector distances_d(4, stream); + + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_d.data(); + distances_tensor.dl_tensor.device.device_type = kDLCUDA; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + // create filter DLTensor + rmm::device_uvector filter_d(1, stream); + raft::copy(filter_d.data(), filter, 1, stream); + + cuvsFilter filter; + + DLManagedTensor filter_tensor; + filter_tensor.dl_tensor.data = filter_d.data(); + filter_tensor.dl_tensor.device.device_type = kDLCUDA; + filter_tensor.dl_tensor.ndim = 1; + filter_tensor.dl_tensor.dtype.code = kDLUInt; + filter_tensor.dl_tensor.dtype.bits = 32; + filter_tensor.dl_tensor.dtype.lanes = 1; + int64_t filter_shape[1] = {1}; + filter_tensor.dl_tensor.shape = filter_shape; + filter_tensor.dl_tensor.strides = nullptr; + + filter.type = BITSET; + filter.addr = (uintptr_t)&filter_tensor; + + // search index + cuvsCagraSearchParams_t search_params; + cuvsCagraSearchParamsCreate(&search_params); + cuvsCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); + // verify output + ASSERT_TRUE(cuvs::devArrMatchHost( + neighbors_exp_filtered, neighbors_d.data(), 4, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + distances_exp_filtered, distances_d.data(), 4, cuvs::CompareApprox(0.001f))); + + // de-allocate index and res + cuvsCagraSearchParamsDestroy(search_params); + cuvsCagraIndexParamsDestroy(build_params); + cuvsCagraIndexDestroy(index); + cuvsResourcesDestroy(res); +} diff --git a/examples/c/src/cagra_c_example.c b/examples/c/src/cagra_c_example.c index fdcbbf571..83957aabc 100644 --- a/examples/c/src/cagra_c_example.c +++ b/examples/c/src/cagra_c_example.c @@ -67,9 +67,9 @@ void cagra_build_search_simple() { // Allocate memory for `queries`, `neighbors` and `distances` output uint32_t *neighbors; float *distances, *queries_d; - cuvsRMMAlloc(res, (void**) &queries_d, sizeof(float) * n_queries * n_cols); - cuvsRMMAlloc(res, (void**) &neighbors, sizeof(uint32_t) * n_queries * topk); - cuvsRMMAlloc(res, (void**) &distances, sizeof(float) * n_queries * topk); + cuvsRMMAlloc(res, (void **)&queries_d, sizeof(float) * n_queries * n_cols); + cuvsRMMAlloc(res, (void **)&neighbors, sizeof(uint32_t) * n_queries * topk); + cuvsRMMAlloc(res, (void **)&distances, sizeof(float) * n_queries * topk); // Use DLPack to represent `queries`, `neighbors` and `distances` as tensors cudaMemcpy(queries_d, queries, sizeof(float) * 4 * 2, cudaMemcpyDefault); @@ -111,8 +111,12 @@ void cagra_build_search_simple() { cuvsCagraSearchParams_t search_params; cuvsCagraSearchParamsCreate(&search_params); + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, - &distances_tensor); + &distances_tensor, filter); // print results uint32_t *neighbors_h = diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index a0f811480..fba7e3d1e 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -29,6 +29,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/cagra.h" nogil: @@ -115,7 +116,8 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsCagraIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, - DLManagedTensor* distances) except + + DLManagedTensor* distances, + cuvsFilter filter) except + cuvsError_t cuvsCagraSerialize(cuvsResources_t res, const char * filename, diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index fd55905cf..f62563f61 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -48,6 +48,7 @@ from libc.stdint cimport ( ) from cuvs.common.exceptions import check_cuvs +from cuvs.neighbors.filters import no_filter cdef class CompressionParams: @@ -484,7 +485,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. @@ -503,6 +505,9 @@ 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 @@ -557,6 +562,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 cuvsCagraSearchParams* params = &search_params.params cdef cydlpack.DLManagedTensor* queries_dlpack = \ cydlpack.dlpack_c(queries_cai) @@ -573,7 +581,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/neighbors/filters/__init__.py b/python/cuvs/cuvs/neighbors/filters/__init__.py index 2ad118965..0ddf809c9 100644 --- a/python/cuvs/cuvs/neighbors/filters/__init__.py +++ b/python/cuvs/cuvs/neighbors/filters/__init__.py @@ -13,6 +13,6 @@ # limitations under the License. -from .filters import Prefilter, from_bitmap, no_filter +from .filters import Prefilter, from_bitmap, from_bitset, no_filter -__all__ = ["no_filter", "from_bitmap", "Prefilter"] +__all__ = ["no_filter", "from_bitmap", "from_bitset", "Prefilter"] diff --git a/python/cuvs/cuvs/neighbors/filters/filters.pyx b/python/cuvs/cuvs/neighbors/filters/filters.pyx index 9bc2a905c..7bc6f9dae 100644 --- a/python/cuvs/cuvs/neighbors/filters/filters.pyx +++ b/python/cuvs/cuvs/neighbors/filters/filters.pyx @@ -22,7 +22,7 @@ from libc.stdint cimport uintptr_t from cuvs.common cimport cydlpack from cuvs.neighbors.common import _check_input_array -from .filters cimport BITMAP, NO_FILTER, cuvsFilter +from .filters cimport BITMAP, BITSET, NO_FILTER, cuvsFilter from pylibraft.common.cai_wrapper import wrap_array @@ -95,3 +95,50 @@ def from_bitmap(bitmap): filter.addr = bitmap_dlpack return Prefilter(filter, parent=bitmap) + +def from_bitset(bitset): + """ + Create a pre-filter from an array with type of uint32. + + Parameters + ---------- + bitset : numpy.ndarray + An array with type of `uint32` where each bit in the array corresponds + to if a sample is greenlit (not filtered) or filtered. + Each bit in a `uint32` element represents a different sample of the dataset. + + - Bit value of 1: The sample is greenlit (allowed). + - Bit value of 0: The sample pair is filtered. + + Returns + ------- + filter : cuvs.neighbors.filters.Prefilter + An instance of `Prefilter` that can be used to filter neighbors + based on the given bitset. + {resources_docstring} + + Examples + -------- + + >>> import cupy as cp + >>> import numpy as np + >>> from cuvs.neighbors import filters + >>> + >>> n_samples = 50000 + >>> n_queries = 1000 + >>> + >>> n_bitset = np.ceil(n_samples / 32).astype(int) + >>> bitset = cp.random.randint(1, 100, size=(n_bitset,), dtype=cp.uint32) + >>> prefilter = filters.from_bitset(bitset) + """ + bitset_cai = wrap_array(bitset) + _check_input_array(bitset_cai, [np.dtype('uint32')]) + + cdef cydlpack.DLManagedTensor* bitset_dlpack = \ + cydlpack.dlpack_c(bitset_cai) + + cdef cuvsFilter filter + filter.type = BITSET + filter.addr = bitset_dlpack + + return Prefilter(filter, parent=bitset) diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index d3b03a5d0..650eb7b5d 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -16,10 +16,11 @@ import numpy as np import pytest from pylibraft.common import device_ndarray +from scipy.spatial.distance import cdist from sklearn.neighbors import NearestNeighbors from sklearn.preprocessing import normalize -from cuvs.neighbors import cagra +from cuvs.neighbors import cagra, filters from cuvs.test.ann_utils import calc_recall, generate_data @@ -168,6 +169,113 @@ 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("n_rows", [1000, 5000]) +@pytest.mark.parametrize("n_cols", [10, 50]) +@pytest.mark.parametrize("n_queries", [10, 100]) +@pytest.mark.parametrize("k", [10, 20]) +@pytest.mark.parametrize("sparsity", [0.2, 0.5]) +def test_filtered_cagra( + n_rows, + n_cols, + n_queries, + k, + sparsity, +): + 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( + metric="euclidean", + intermediate_graph_degree=64, + graph_degree=32, + build_algo="nn_descent", + ) + index = cagra.build(build_params, dataset_device) + + prefilter = 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=prefilter, + ) + + # 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] + + # Calculate reference values with sklearn on filtered dataset + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric="euclidean" + ) + nn_skl.fit(filtered_dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + # Get actual results + actual_indices = out_idx_device.copy_to_host() + actual_distances = out_dist_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" + ) + + # Verify filtering - no filtered indices should be in results + filtered_indices = np.where(bool_filter)[0] + for i in range(n_queries): + assert not np.intersect1d(filtered_indices, actual_indices[i]).size + + # Now compare with sklearn results + recall = calc_recall(mapped_actual_indices, skl_idx) + + assert recall > 0.7 + + def test_cagra_index_params(params): # Note that inner_product tests use normalized input which we cannot # represent in int8, therefore we test only sqeuclidean metric here. diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index 959959f60..bf316b4d7 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/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::cuvsCagraSearch( res.0, params.0, @@ -85,6 +90,7 @@ impl Index { queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), + prefilter, )) } } @@ -167,7 +173,8 @@ mod tests { #[test] fn test_cagra_compression() { use crate::cagra::CompressionParams; - let build_params = IndexParams::new().unwrap() + let build_params = IndexParams::new() + .unwrap() .set_compression(CompressionParams::new().unwrap()); test_cagra(build_params); }