Add a filtering_rate hint to brute_force::search_params

cpp/include/cuvs/neighbors/brute_force.hpp

@@ -21,7 +21,22 @@ namespace brute_force {
 
 struct index_params : cuvs::neighbors::index_params {};
 
-struct search_params : cuvs::neighbors::search_params {};
+/** @brief Parameters for brute-force search. */
+struct search_params : cuvs::neighbors::search_params {
+  /**
+   * A hint for the fraction of entries the filter removes, i.e.
+   * `(n_queries * n_dataset - n_set_bits) / (n_queries * n_dataset)`; for a `bitset_filter` this is
+   * `(n_dataset - n_set_bits) / n_dataset`.
+   *
+   * - Negative (default): the rate is computed from the filter, which runs a popcount kernel and
+   *   synchronizes the stream on every filtered search.
+   * - In `[0.0, 1.0)`: the value is trusted and the popcount is skipped when the dense path is
+   *   selected. The sparse and gather paths still count the filter, because they need the exact
+   *   number of passing entries.
+   * - Any other value (including NaN) is rejected.
+   */
+  float filtering_rate = -1.0;
+};
 
 /**
  * @defgroup bruteforce_cpp_index Bruteforce index

cpp/src/neighbors/brute_force.cu

@@ -215,7 +215,7 @@ void index<T, DistT>::update_dataset(
               const cuvs::neighbors::filtering::base_filter& sample_filter)                     \
   {                                                                                             \
     detail::search<T, int64_t, DistT, raft::row_major>(                                         \
-      res, idx, queries, neighbors, distances, sample_filter);                                  \
+      res, params, idx, queries, neighbors, distances, sample_filter);                          \
   }                                                                                             \
   void search(raft::resources const& res,                                                       \
               const cuvs::neighbors::brute_force::index<T, DistT>& idx,                         \
@@ -225,7 +225,7 @@ void index<T, DistT>::update_dataset(
               const cuvs::neighbors::filtering::base_filter& sample_filter)                     \
   {                                                                                             \
     detail::search<T, int64_t, DistT, raft::row_major>(                                         \
-      res, idx, queries, neighbors, distances, sample_filter);                                  \
+      res, {}, idx, queries, neighbors, distances, sample_filter);                              \
   }                                                                                             \
   void search(raft::resources const& res,                                                       \
               const cuvs::neighbors::brute_force::search_params& params,                        \
@@ -236,7 +236,7 @@ void index<T, DistT>::update_dataset(
               const cuvs::neighbors::filtering::base_filter& sample_filter)                     \
   {                                                                                             \
     detail::search<T, int64_t, DistT, raft::col_major>(                                         \
-      res, idx, queries, neighbors, distances, sample_filter);                                  \
+      res, params, idx, queries, neighbors, distances, sample_filter);                          \
   }                                                                                             \
   void search(raft::resources const& res,                                                       \
               const cuvs::neighbors::brute_force::index<T, DistT>& idx,                         \
@@ -246,7 +246,7 @@ void index<T, DistT>::update_dataset(
               const cuvs::neighbors::filtering::base_filter& sample_filter)                     \
   {                                                                                             \
     detail::search<T, int64_t, DistT, raft::col_major>(                                         \
-      res, idx, queries, neighbors, distances, sample_filter);                                  \
+      res, {}, idx, queries, neighbors, distances, sample_filter);                              \
   }
 
 CUVS_INST_BFKNN(float, float);

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -736,6 +736,7 @@ void brute_force_search_gathered(
 template <typename T, typename IdxT, typename BitsT, typename DistanceT = float>
 void brute_force_search_filtered(
   raft::resources const& res,
+  const cuvs::neighbors::brute_force::search_params& params,
   const cuvs::neighbors::brute_force::index<T, DistanceT>& idx,
   raft::device_matrix_view<const T, IdxT, raft::row_major> queries,
   const cuvs::neighbors::filtering::base_filter* filter,
@@ -759,6 +760,15 @@ void brute_force_search_filtered(
                                     metric == cuvs::distance::DistanceType::CosineExpanded),
                "Index must has norms when using Euclidean, IP, and Cosine!");
 
+  // Negative means "auto-detect from the filter"; [0.0, 1.0) is trusted as-is. Anything else
+  // (including NaN, which compares false against both bounds) is out of contract.
+  const bool auto_filtering_rate = params.filtering_rate < 0.0f;
+  const bool use_hint            = params.filtering_rate >= 0.0f && params.filtering_rate < 1.0f;
+  RAFT_EXPECTS(auto_filtering_rate || use_hint,
+               "search_params::filtering_rate must be negative (auto-detect) or in [0.0, 1.0), "
+               "got %f",
+               params.filtering_rate);
+
   IdxT n_queries                                     = queries.extent(0);
   IdxT n_dataset                                     = idx.dataset().extent(0);
   IdxT dim                                           = idx.dataset().extent(1);
@@ -771,7 +781,10 @@ void brute_force_search_filtered(
                              const cuvs::core::bitset_view<BitsT, IdxT>>>
     filter_view;
 
-  IdxT nnz_h = 0;
+  IdxT nnz_h         = 0;
+  double selectivity = 0.0;
+  // Whether nnz_h / n_pass came from an actual popcount rather than the filtering_rate hint.
+  bool counted = false;
   // A bitset passes the same rows for every query, so it has a row count; a bitmap passes a
   // different set per query and has none.
   std::optional<IdxT> n_pass;
@@ -782,22 +795,51 @@ void brute_force_search_filtered(
     auto actual_filter =
       dynamic_cast<const cuvs::neighbors::filtering::bitmap_filter<BitsT, int64_t>*>(filter);
     filter_view.emplace(actual_filter->view());
-    nnz_h = actual_filter->view().count(res);
+    if (use_hint) {
+      selectivity = 1.0 - params.filtering_rate;
+    } else {
+      nnz_h       = actual_filter->view().count(res);
+      selectivity = static_cast<double>(nnz_h) /
+                    (static_cast<double>(n_queries) * static_cast<double>(n_dataset));
+      counted = true;
+    }
   } else if (filter_type == cuvs::neighbors::filtering::FilterType::Bitset) {
     auto actual_filter =
       dynamic_cast<const cuvs::neighbors::filtering::bitset_filter<BitsT, int64_t>*>(filter);
     filter_view.emplace(actual_filter->view());
-    n_pass = actual_filter->view().count(res);
-    nnz_h  = n_queries * (*n_pass);
+    if (use_hint) {
+      selectivity = 1.0 - params.filtering_rate;
+      // Path selection also needs a row count; the exact one is only fetched below, and only
+      // if the path it picks turns out to need it.
+      n_pass = static_cast<IdxT>(std::llround(selectivity * static_cast<double>(n_dataset)));
+    } else {
+      n_pass      = actual_filter->view().count(res);
+      nnz_h       = n_queries * (*n_pass);
+      selectivity = static_cast<double>(nnz_h) /
+                    (static_cast<double>(n_queries) * static_cast<double>(n_dataset));
+      counted = true;
+    }
   } else {
     RAFT_FAIL("Unsupported sample filter type");
   }
 
   std::visit([&](const auto& actual_view) { filter_data = actual_view.data(); }, *filter_view);
 
-  const double selectivity =
-    static_cast<double>(nnz_h) / (static_cast<double>(n_queries) * static_cast<double>(n_dataset));
-  const auto path = select_filtered_search_path(n_dataset, dim, selectivity, k, n_pass);
+  auto path = select_filtered_search_path(n_dataset, dim, selectivity, k, n_pass);
+
+  // Only the dense path can run on a hinted selectivity alone: sddmm sizes its CSR from the
+  // exact nnz, and gather sizes the compacted dataset from the exact row count. If we have to
+  // popcount after all, spend the exact numbers on the decision too.
+  if (!counted && path != filtered_search_path::dense) {
+    std::visit([&](const auto& actual_view) { nnz_h = actual_view.count(res); }, *filter_view);
+    if (n_pass) {
+      n_pass = nnz_h;
+      nnz_h  = n_queries * (*n_pass);
+    }
+    selectivity = static_cast<double>(nnz_h) /
+                  (static_cast<double>(n_queries) * static_cast<double>(n_dataset));
+    path = select_filtered_search_path(n_dataset, dim, selectivity, k, n_pass);
+  }
 
   if (path == filtered_search_path::gather) {
     auto bitset_view = std::get<const cuvs::core::bitset_view<BitsT, IdxT>>(*filter_view);
@@ -906,6 +948,7 @@ void brute_force_search_filtered(
 
 template <typename T, typename IdxT, typename DistT, typename LayoutT>
 void search(raft::resources const& res,
+            const cuvs::neighbors::brute_force::search_params& params,
             const cuvs::neighbors::brute_force::index<T, DistT>& idx,
             raft::device_matrix_view<const T, int64_t, LayoutT> queries,
             raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
@@ -926,7 +969,7 @@ void search(raft::resources const& res,
         dynamic_cast<const cuvs::neighbors::filtering::bitmap_filter<uint32_t, int64_t>&>(
           sample_filter_ref);
       return brute_force_search_filtered<T, int64_t, uint32_t, DistT>(
-        res, idx, queries, &sample_filter, neighbors, distances);
+        res, params, idx, queries, &sample_filter, neighbors, distances);
     } catch (const std::bad_cast&) {
     }
 
@@ -935,7 +978,7 @@ void search(raft::resources const& res,
         dynamic_cast<const cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>&>(
           sample_filter_ref);
       return brute_force_search_filtered<T, int64_t, uint32_t, DistT>(
-        res, idx, queries, &sample_filter, neighbors, distances);
+        res, params, idx, queries, &sample_filter, neighbors, distances);
     } catch (const std::bad_cast&) {
       RAFT_FAIL("Unsupported sample filter type");
     }

cpp/tests/neighbors/brute_force_prefiltered.cu

@@ -567,6 +567,88 @@ class PrefilteredBruteForceOnBitmapTest
                                                     true));
   }
 
+  // Same as Run(), but passes the filter's true filtering_rate as a hint to the params-taking
+  // search overload. Confirms results match auto-detection. The rate comes from the measured
+  // `nnz`, not `params.sparsity`: RMAT can draw the same edge twice, so fewer bits may be set.
+  void RunWithFilteringRateHint()
+  {
+    auto dataset_raw = raft::make_device_matrix_view<const value_t, index_t, raft::row_major>(
+      (const value_t*)dataset_d.data(), params.n_dataset, params.dim);
+
+    auto queries = raft::make_device_matrix_view<const value_t, index_t, raft::row_major>(
+      (const value_t*)queries_d.data(), params.n_queries, params.dim);
+
+    auto dataset = brute_force::build(handle, dataset_raw, params.metric);
+
+    auto filter = cuvs::core::bitmap_view<bitmap_t, index_t>(
+      (bitmap_t*)filter_d.data(), params.n_queries, params.n_dataset);
+
+    auto out_val = raft::make_device_matrix_view<dist_t, index_t, raft::row_major>(
+      out_val_d.data(), params.n_queries, params.top_k);
+    auto out_idx = raft::make_device_matrix_view<index_t, index_t, raft::row_major>(
+      out_idx_d.data(), params.n_queries, params.top_k);
+
+    cuvs::neighbors::brute_force::search_params search_params;
+    search_params.filtering_rate =
+      static_cast<float>(1.0 - static_cast<double>(nnz) / (static_cast<double>(params.n_queries) *
+                                                           static_cast<double>(params.n_dataset)));
+
+    brute_force::search(handle,
+                        search_params,
+                        dataset,
+                        queries,
+                        out_idx,
+                        out_val,
+                        cuvs::neighbors::filtering::bitmap_filter(filter));
+
+    raft::resource::sync_stream(handle);
+
+    ASSERT_TRUE(cuvs::neighbors::devArrMatchKnnPair(out_idx_expected_d.data(),
+                                                    out_idx.data_handle(),
+                                                    out_val_expected_d.data(),
+                                                    out_val.data_handle(),
+                                                    params.n_queries,
+                                                    params.top_k,
+                                                    0.001f,
+                                                    stream,
+                                                    true));
+  }
+
+  // A filtering_rate that is neither the negative auto-detect sentinel nor in [0.0, 1.0) is
+  // out of contract and must be rejected rather than silently steering path selection.
+  void RunWithInvalidFilteringRate()
+  {
+    auto dataset_raw = raft::make_device_matrix_view<const value_t, index_t, raft::row_major>(
+      (const value_t*)dataset_d.data(), params.n_dataset, params.dim);
+
+    auto queries = raft::make_device_matrix_view<const value_t, index_t, raft::row_major>(
+      (const value_t*)queries_d.data(), params.n_queries, params.dim);
+
+    auto dataset = brute_force::build(handle, dataset_raw, params.metric);
+
+    auto filter = cuvs::core::bitmap_view<bitmap_t, index_t>(
+      (bitmap_t*)filter_d.data(), params.n_queries, params.n_dataset);
+
+    auto out_val = raft::make_device_matrix_view<dist_t, index_t, raft::row_major>(
+      out_val_d.data(), params.n_queries, params.top_k);
+    auto out_idx = raft::make_device_matrix_view<index_t, index_t, raft::row_major>(
+      out_idx_d.data(), params.n_queries, params.top_k);
+
+    for (float invalid_rate : {1.0f, 1.5f, 50.0f, std::numeric_limits<float>::quiet_NaN()}) {
+      cuvs::neighbors::brute_force::search_params search_params;
+      search_params.filtering_rate = invalid_rate;
+
+      EXPECT_THROW(brute_force::search(handle,
+                                       search_params,
+                                       dataset,
+                                       queries,
+                                       out_idx,
+                                       out_val,
+                                       cuvs::neighbors::filtering::bitmap_filter(filter)),
+                   raft::logic_error);
+    }
+  }
+
  protected:
   raft::resources handle;
   cudaStream_t stream;
@@ -948,6 +1030,53 @@ class PrefilteredBruteForceOnBitsetTest
                                                     true));
   }
 
+  // Same as Run(), but passes the filter's true filtering_rate as a hint to the params-taking
+  // search overload. Confirms results match auto-detection. The rate comes from the measured
+  // `nnz`, not `params.sparsity`: RMAT can draw the same edge twice, so fewer bits may be set.
+  void RunWithFilteringRateHint()
+  {
+    auto dataset_raw = raft::make_device_matrix_view<const value_t, index_t, raft::row_major>(
+      (const value_t*)dataset_d.data(), params.n_dataset, params.dim);
+
+    auto queries = raft::make_device_matrix_view<const value_t, index_t, raft::row_major>(
+      (const value_t*)queries_d.data(), params.n_queries, params.dim);
+
+    auto dataset = brute_force::build(handle, dataset_raw, params.metric);
+
+    auto filter =
+      cuvs::core::bitset_view<bitset_t, index_t>((bitset_t*)filter_d.data(), params.n_dataset);
+
+    auto out_val = raft::make_device_matrix_view<dist_t, index_t, raft::row_major>(
+      out_val_d.data(), params.n_queries, params.top_k);
+    auto out_idx = raft::make_device_matrix_view<index_t, index_t, raft::row_major>(
+      out_idx_d.data(), params.n_queries, params.top_k);
+
+    cuvs::neighbors::brute_force::search_params search_params;
+    search_params.filtering_rate =
+      static_cast<float>(1.0 - static_cast<double>(nnz) / (static_cast<double>(params.n_queries) *
+                                                           static_cast<double>(params.n_dataset)));
+
+    brute_force::search(handle,
+                        search_params,
+                        dataset,
+                        queries,
+                        out_idx,
+                        out_val,
+                        cuvs::neighbors::filtering::bitset_filter(filter));
+
+    raft::resource::sync_stream(handle);
+
+    ASSERT_TRUE(cuvs::neighbors::devArrMatchKnnPair(out_idx_expected_d.data(),
+                                                    out_idx.data_handle(),
+                                                    out_val_expected_d.data(),
+                                                    out_val.data_handle(),
+                                                    params.n_queries,
+                                                    params.top_k,
+                                                    0.001f,
+                                                    stream,
+                                                    true));
+  }
+
  protected:
   raft::resources handle;
   cudaStream_t stream;
@@ -970,18 +1099,38 @@ class PrefilteredBruteForceOnBitsetTest
 using PrefilteredBruteForceTestOnBitmap_float_int64 =
   PrefilteredBruteForceOnBitmapTest<float, float, int64_t>;
 TEST_P(PrefilteredBruteForceTestOnBitmap_float_int64, Result) { Run(); }
+TEST_P(PrefilteredBruteForceTestOnBitmap_float_int64, ResultWithFilteringRateHint)
+{
+  RunWithFilteringRateHint();
+}
+TEST_P(PrefilteredBruteForceTestOnBitmap_float_int64, InvalidFilteringRateThrows)
+{
+  RunWithInvalidFilteringRate();
+}
 
 using PrefilteredBruteForceTestOnBitmap_half_int64 =
   PrefilteredBruteForceOnBitmapTest<half, float, int64_t>;
 TEST_P(PrefilteredBruteForceTestOnBitmap_half_int64, Result) { Run(); }
+TEST_P(PrefilteredBruteForceTestOnBitmap_half_int64, ResultWithFilteringRateHint)
+{
+  RunWithFilteringRateHint();
+}
 
 using PrefilteredBruteForceTestOnBitset_float_int64 =
   PrefilteredBruteForceOnBitsetTest<float, float, int64_t>;
 TEST_P(PrefilteredBruteForceTestOnBitset_float_int64, Result) { Run(); }
+TEST_P(PrefilteredBruteForceTestOnBitset_float_int64, ResultWithFilteringRateHint)
+{
+  RunWithFilteringRateHint();
+}
 
 using PrefilteredBruteForceTestOnBitset_half_int64 =
   PrefilteredBruteForceOnBitsetTest<half, float, int64_t>;
 TEST_P(PrefilteredBruteForceTestOnBitset_half_int64, Result) { Run(); }
+TEST_P(PrefilteredBruteForceTestOnBitset_half_int64, ResultWithFilteringRateHint)
+{
+  RunWithFilteringRateHint();
+}
 
 template <typename index_t>
 const std::vector<PrefilteredBruteForceInputs<index_t>> selectk_inputs = {

fern/pages/cpp_api/cpp-api-neighbors-brute-force.md

@@ -6,6 +6,25 @@ slug: api-reference/cpp-api-neighbors-brute-force
 
 _Source header: `cuvs/neighbors/brute_force.hpp`_
 
+## Types
+
+<a id="neighbors-brute-force-search-params"></a>
+### neighbors::brute_force::search_params
+
+Parameters for brute-force search.
+
+```cpp
+struct search_params : cuvs::neighbors::search_params {
+  float filtering_rate;
+};
+```
+
+**Fields**
+
+| Name | Type | Description |
+| --- | --- | --- |
+| `filtering_rate` | `float` | A hint for the fraction of entries the filter removes, i.e. `(n_queries * n_dataset - n_set_bits) / (n_queries * n_dataset)`; for a `bitset_filter` this is `(n_dataset - n_set_bits) / n_dataset`.<br /><br />- Negative (default): the rate is computed from the filter, which runs a popcount kernel and synchronizes the stream on every filtered search.<br />- In `[0.0, 1.0)`: the value is trusted and the popcount is skipped when the dense path is selected. The sparse and gather paths still count the filter, because they need the exact number of passing entries.<br />- Any other value (including NaN) is rejected. |
+
 ## Bruteforce index
 
 <a id="neighbors-brute-force-index"></a>