Derive the filtering rate in filtered multi-partition CAGRA search

cpp/src/neighbors/cagra.cuh

@@ -420,6 +420,16 @@ void search_with_filtering(raft::resources const& res,
     res, params, idx, queries, neighbors, distances, sample_filter);
 }
 
+namespace detail {
+/** Clamp a filtering rate derived from a filter to the range the search plan accepts. */
+inline float clamp_filtering_rate(float filtering_rate)
+{
+  const float min_filtering_rate = 0.0f;
+  const float max_filtering_rate = 0.999f;
+  return std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate);
+}
+}  // namespace detail
+
 template <typename T,
           typename IdxT,
           cuvs::neighbors::ann_dataset_view DatasetViewT,
@@ -451,10 +461,7 @@ void search(raft::resources const& res,
     if (params.filtering_rate < 0.0) {
       const auto num_set_bits = sample_filter.bitset_view_.count(res);
       auto filtering_rate = (float)(idx.dataset().n_rows() - num_set_bits) / idx.dataset().n_rows();
-      const float min_filtering_rate = 0.0;
-      const float max_filtering_rate = 0.999;
-      params_copy.filtering_rate =
-        std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate);
+      params_copy.filtering_rate = detail::clamp_filtering_rate(filtering_rate);
     }
     auto sample_filter_copy = sample_filter;
     return search_with_filtering<T, IdxT, decltype(sample_filter_copy), OutputIdxT, DatasetViewT>(
@@ -503,13 +510,7 @@ void search(raft::resources const& res,
       dynamic_cast<const cuvs::neighbors::filtering::udf_filter&>(sample_filter_ref);
     search_params params_copy = params;
     if (params.filtering_rate < 0.0) {
-      const float min_filtering_rate = 0.0f;
-      const float max_filtering_rate = 0.999f;
-      params_copy.filtering_rate =
-        sample_filter.filtering_rate < 0.0f
-          ? 0.0f
-          : std::min(std::max(sample_filter.filtering_rate, min_filtering_rate),
-                     max_filtering_rate);
+      params_copy.filtering_rate = detail::clamp_filtering_rate(sample_filter.filtering_rate);
     }
     auto sample_filter_copy = sample_filter;
     return search_with_filtering<T, IdxT, decltype(sample_filter_copy), OutputIdxT, DatasetViewT>(
@@ -606,10 +607,35 @@ void search(
                                           cuvs::neighbors::filtering::none_sample_filter>(
       res, params, indices, queries, partition_ids, neighbors, distances, partition_bitsets);
   } else {
+    // As in the single-partition bitset path, derive the filtering rate from the bitsets when the
+    // caller did not set one, so that the search plan enlarges itopk_size for selective filters.
+    // Without this, a selective filter leaves too few valid candidates and the search can return
+    // fewer than k neighbors. One plan serves every partition, so size it for the most selective
+    // partition that still accepts at least one row; a partition that accepts nothing contributes
+    // no neighbors whatever the itopk_size, and an unfiltered partition has a rate of zero.
+    search_params params_copy = params;
+    if (params.filtering_rate < 0.0) {
+      float filtering_rate = 0.0f;
+      for (size_t i = 0; i < partition_bitsets.size() && i < indices.size(); i++) {
+        const auto& v = partition_bitsets[i];
+        if (v.data() == nullptr || v.size() == 0) { continue; }
+        const int64_t n_rows = indices[i]->dataset().n_rows();
+        if (n_rows == 0) { continue; }
+        // Count only this partition's rows; the bitset may be padded past n_rows.
+        const auto num_set_bits =
+          cuvs::core::bitset_view<std::uint32_t, int64_t>(const_cast<std::uint32_t*>(v.data()),
+                                                          std::min<int64_t>(n_rows, v.size()))
+            .count(res);
+        if (num_set_bits == 0) { continue; }
+        filtering_rate =
+          std::max(filtering_rate, static_cast<float>(n_rows - num_set_bits) / n_rows);
+      }
+      params_copy.filtering_rate = detail::clamp_filtering_rate(filtering_rate);
+    }
     using bitset_filter_t = cuvs::neighbors::filtering::bitset_filter<std::uint32_t, int64_t>;
     cagra::detail::search_multi_partition<T, OutputIdxT, IdxT, float, bitset_filter_t>(
       res,
-      params,
+      params_copy,
       indices,
       queries,
       partition_ids,

cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu

@@ -7,6 +7,8 @@
 
 #include "../ann_cagra.cuh"
 
+#include <algorithm>
+#include <limits>
 #include <utility>
 #include <vector>
 
@@ -135,6 +137,152 @@ TEST(AnnCagraMultiPartition, MultiKernelRejected)
                                        search_params);
 }
 
+namespace {
+// Runs a filtered multi-partition MULTI_CTA search where every partition keeps only ~4% of its rows
+// (still more than k each) and reports the number of queries that returned fewer than k neighbors
+// and the recall against an exact search over the kept rows.
+void search_with_selective_filter(float filtering_rate, int& short_queries, double& recall)
+{
+  raft::resources handle;
+  auto stream = raft::resource::get_cuda_stream(handle);
+
+  constexpr int num_partitions = 4, part_size = 500, dim = 128, n_queries = 500, k = 10;
+  // Keep one row in `keep_every` (~4% of each partition, i.e. ~21 rows per partition).
+  constexpr int keep_every = 24;
+  constexpr int n_rows     = num_partitions * part_size;
+
+  rmm::device_uvector<float> database(static_cast<size_t>(n_rows) * dim, stream);
+  rmm::device_uvector<float> queries(static_cast<size_t>(n_queries) * dim, stream);
+  raft::random::RngState r(1234ULL);
+  InitDataset(handle, database.data(), n_rows, dim, cuvs::distance::DistanceType::L2Expanded, r);
+  InitDataset(handle, queries.data(), n_queries, dim, cuvs::distance::DistanceType::L2Expanded, r);
+  raft::resource::sync_stream(handle);
+
+  cagra::index_params index_params;
+  index_params.graph_degree              = 32;
+  index_params.intermediate_graph_degree = 64;
+  index_params.graph_build_params        = graph_build_params::nn_descent_params(
+    index_params.intermediate_graph_degree, index_params.metric);
+
+  std::vector<cagra::index<float, std::uint32_t>> part_indices;
+  // An index only holds a view, so any padded copy has to outlive it.
+  std::vector<cuvs::neighbors::test::padded_device_matrix_for_cagra<float>> part_padded;
+  part_padded.reserve(num_partitions);
+  std::vector<cuvs::core::bitset<uint32_t, int64_t>> part_bitsets;
+  part_bitsets.reserve(num_partitions);
+  std::vector<cuvs::core::bitset_view<uint32_t, int64_t>> partition_bitsets;
+  std::vector<int64_t> removed_local;
+  for (int64_t i = 0; i < part_size; i++) {
+    if (i % keep_every != 0) { removed_local.push_back(i); }
+  }
+  auto removed =
+    raft::make_device_vector<int64_t, int64_t>(handle, static_cast<int64_t>(removed_local.size()));
+  raft::update_device(removed.data_handle(), removed_local.data(), removed_local.size(), stream);
+  for (int p = 0; p < num_partitions; p++) {
+    auto view = raft::make_device_matrix_view<const float, int64_t>(
+      database.data() + static_cast<size_t>(p) * part_size * dim, part_size, dim);
+    part_padded.emplace_back(handle, view);
+    auto const& padded = part_padded.back().view;
+    part_indices.push_back(cagra::build(handle, index_params, padded));
+    auto& part_index = part_indices.back();
+    part_index       = cagra::update_dataset(handle, std::move(part_index), padded);
+    part_bitsets.emplace_back(handle, removed.view(), part_size);
+    partition_bitsets.push_back(part_bitsets.back().view());
+  }
+  std::vector<const cagra::index<float, std::uint32_t>*> index_ptrs;
+  for (auto& idx : part_indices) {
+    index_ptrs.push_back(&idx);
+  }
+
+  const size_t out_size = static_cast<size_t>(n_queries) * k;
+  rmm::device_uvector<uint32_t> partition_ids_dev(out_size, stream);
+  rmm::device_uvector<uint32_t> neighbors_dev(out_size, stream);
+  rmm::device_uvector<float> distances_dev(out_size, stream);
+
+  // A small itopk_size, as used for small k, on the MULTI_CTA algorithm that AUTO picks for a
+  // handful of queries over a handful of partitions.
+  cagra::search_params search_params;
+  search_params.algo           = search_algo::MULTI_CTA;
+  search_params.itopk_size     = k;
+  search_params.filtering_rate = filtering_rate;
+
+  cagra::search(
+    handle,
+    search_params,
+    index_ptrs,
+    raft::make_device_matrix_view<const float, int64_t>(queries.data(), n_queries, dim),
+    raft::make_device_matrix_view<uint32_t, int64_t>(partition_ids_dev.data(), n_queries, k),
+    raft::make_device_matrix_view<uint32_t, int64_t>(neighbors_dev.data(), n_queries, k),
+    raft::make_device_matrix_view<float, int64_t>(distances_dev.data(), n_queries, k),
+    partition_bitsets);
+
+  std::vector<uint32_t> partition_ids(out_size);
+  std::vector<uint32_t> neighbors(out_size);
+  std::vector<float> distances(out_size);
+  std::vector<float> host_database(static_cast<size_t>(n_rows) * dim);
+  std::vector<float> host_queries(static_cast<size_t>(n_queries) * dim);
+  raft::update_host(partition_ids.data(), partition_ids_dev.data(), out_size, stream);
+  raft::update_host(neighbors.data(), neighbors_dev.data(), out_size, stream);
+  raft::update_host(distances.data(), distances_dev.data(), out_size, stream);
+  raft::update_host(host_database.data(), database.data(), host_database.size(), stream);
+  raft::update_host(host_queries.data(), queries.data(), host_queries.size(), stream);
+  raft::resource::sync_stream(handle);
+
+  // Exact top-k over the rows that pass the filter, for recall.
+  std::vector<int64_t> kept_rows;
+  for (int64_t g = 0; g < n_rows; g++) {
+    if ((g % part_size) % keep_every == 0) { kept_rows.push_back(g); }
+  }
+  ASSERT_GT(kept_rows.size(), static_cast<size_t>(k));
+
+  short_queries  = 0;
+  size_t matches = 0;
+  for (int q = 0; q < n_queries; q++) {
+    std::vector<std::pair<float, int64_t>> exact;
+    for (auto g : kept_rows) {
+      float d = 0;
+      for (int j = 0; j < dim; j++) {
+        const float diff = host_queries[q * dim + j] - host_database[g * dim + j];
+        d += diff * diff;
+      }
+      exact.emplace_back(d, g);
+    }
+    std::partial_sort(exact.begin(), exact.begin() + k, exact.end());
+    std::vector<int64_t> expected;
+    for (int j = 0; j < k; j++) {
+      expected.push_back(exact[j].second);
+    }
+
+    int valid = 0;
+    for (int j = 0; j < k; j++) {
+      const size_t i = static_cast<size_t>(q) * k + j;
+      if (distances[i] == std::numeric_limits<float>::max()) { continue; }
+      ASSERT_LT(partition_ids[i], static_cast<uint32_t>(num_partitions));
+      ASSERT_LT(neighbors[i], static_cast<uint32_t>(part_size));
+      ASSERT_EQ(neighbors[i] % keep_every, 0u) << "filtered-out row returned";
+      valid++;
+      const int64_t g = static_cast<int64_t>(partition_ids[i]) * part_size + neighbors[i];
+      if (std::find(expected.begin(), expected.end(), g) != expected.end()) { matches++; }
+    }
+    if (valid < k) { short_queries++; }
+  }
+  recall = static_cast<double>(matches) / out_size;
+}
+}  // namespace
+
+// A selective per-partition filter must not leave the search short of neighbors. Single-partition
+// search derives the filtering rate from the bitset and enlarges itopk_size accordingly; the
+// multi-partition search must do the same, or MULTI_CTA can return fewer than k valid neighbors
+// even though every partition holds more than k rows that pass the filter.
+TEST(AnnCagraMultiPartition, SelectiveFilterReturnsKNeighbors)
+{
+  int short_queries = 0;
+  double recall     = 0;
+  search_with_selective_filter(-1.0f, short_queries, recall);
+  EXPECT_EQ(short_queries, 0) << "queries returning fewer than k neighbors";
+  EXPECT_GE(recall, 0.95);
+}
+
 // The shared plan descriptor and the cross-partition select_k direction are derived from
 // indices[0], so all partitions must share one metric; a mismatch must be rejected.
 TEST(AnnCagraMultiPartition, MixedMetricRejected)

java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java

@@ -114,7 +114,9 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException {
 
     // Use the first vector (from document 0) as query
     float[] queryVector = dataset[0];
-    int topK = random.nextInt(5, TOP_K_LIMIT);
+    // Only even-numbered documents have vectors, so never ask for more results than exist.
+    int vectorCount = (datasetSize + 1) / 2;
+    int topK = Math.min(random.nextInt(5, TOP_K_LIMIT), vectorCount);
 
     GPUKnnFloatVectorQuery query =
         new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1);