C API: adaptive filtered top-k search by sampled pass rate

bindings/c/src/filtered_search.hpp

@@ -21,7 +21,9 @@
 
 #include <svs/concepts/data.h>
 #include <svs/core/query_result.h>
+#include <svs/index/vamana/search_buffer.h>
 #include <svs/lib/misc.h>
+#include <svs/lib/preprocessor.h>
 #include <svs/lib/threads.h>
 
 #include <algorithm>
@@ -79,25 +81,28 @@ inline bool hit_rate_sufficient(size_t total, size_t hits, float filter_rate) {
 /// @param goal The target number of hits to achieve - usually is K (from TopK).
 /// @param hint A hint for the batch size - usually based on prior knowledge.
 /// @param limit The maximum allowed batch size. E.g. index size.
+/// @param measured_pass_rate If not null, receives the sampled pass rate (hits / sampled
+/// ids; 0 if nothing was sampled). The sample is taken even when `filter_rate()` is 0.
 /// @return The estimated initial batch size, or 0 if the hit rate is insufficient.
 inline size_t estimate_initial_batch_size(
     const IDFilterInterface* id_filter,
     std::function<size_t()> sample_generator,
     size_t min_sample_size,
     size_t goal,
     size_t hint,
-    size_t limit
+    size_t limit,
+    float* measured_pass_rate = nullptr
 ) {
     assert(id_filter != nullptr);
     const auto filter_rate = id_filter->filter_rate();
-    if (filter_rate <= 0.0f) {
-        // If filter rate is 0.0 or negative, return the `hint` as the initial batch size -
-        // clamped to `limit` to avoid oversizing.
-        return std::min(hint, limit);
+    if (measured_pass_rate != nullptr) {
+        *measured_pass_rate = 0.0f;
     }
 
-    auto sample_size =
-        std::max({goal, min_sample_size, static_cast<size_t>(1.f / filter_rate)});
+    auto sample_size = std::max(goal, min_sample_size);
+    if (filter_rate > 0.0f) {
+        sample_size = std::max(sample_size, static_cast<size_t>(1.f / filter_rate));
+    }
     assert(sample_size > 0);
     size_t hits = 0;
     for (size_t i = 0; i < sample_size; ++i) {
@@ -111,6 +116,14 @@ inline size_t estimate_initial_batch_size(
             hits++;
         }
     }
+    if (measured_pass_rate != nullptr && sample_size > 0) {
+        *measured_pass_rate = static_cast<float>(hits) / static_cast<float>(sample_size);
+    }
+    if (filter_rate <= 0.0f) {
+        // No filter rate estimate from the caller: use the `hint` as the initial batch
+        // size - clamped to `limit` to avoid oversizing.
+        return std::min(hint, limit);
+    }
     // if hit rate is less than filter_rate, return 0 - which means we should not even start
     // the search
     if (!hit_rate_sufficient(sample_size, hits, filter_rate)) {
@@ -148,13 +161,21 @@ inline void set_empty_result(svs::QueryResult<size_t>& result) {
     std::fill(result.indices().begin(), result.indices().end(), static_cast<size_t>(-1));
 }
 
-// Perform a filtered nearest-neighbor search by iterating over batches of candidates and
-// keeping only those that pass the filter. The batch size is estimated adaptively based on
-// the observed hit rate. Results are written into `results`.
+// Perform a filtered nearest-neighbor search. The pass rate of the filter is measured on a
+// sample of ids, then:
+//  * pass rate >= GREEDY_MIN_PASS_RATE: one plain greedy search per query for
+//    max(hint, GREEDY_OVERSAMPLE * K / pass rate) neighbors, filtered in distance order.
+//    Queries that end up with fewer than K hits are redone with the batch loop below.
+//  * otherwise: iterate over batches of candidates and keep only those that pass the
+//    filter. The batch size is estimated adaptively based on the observed hit rate, and the
+//    iterator gets extra search buffer capacity that grows as the pass rate falls (60 to
+//    200 slots).
+// Results are written into `results`.
 // returns a vector containing the number of neighbors found for each query.
-template <typename IndexType>
+template <typename IndexType, typename SearchParams>
 std::vector<size_t> filtered_topk_search(
     IndexType& index,
+    const SearchParams& params,
     svs::QueryResult<size_t>& results,
     svs::data::ConstSimpleDataView<float> queries,
     size_t initial_batch_hint,
@@ -165,6 +186,19 @@ std::vector<size_t> filtered_topk_search(
     // between accuracy and performance. A larger sample size gives a more accurate
     // estimate of the filter hit rate, but takes longer to compute.
     const size_t MIN_SAMPLE_SIZE = 200;
+    // Sampled pass rate from which the greedy search is tried first.
+    const float GREEDY_MIN_PASS_RATE = 0.5f;
+    // The greedy search asks for this many times the expected number of candidates.
+    const float GREEDY_OVERSAMPLE = 1.2f;
+    // Extra search buffer capacity of the batch iterator (library default is 100):
+    // ceil(BATCH_EXTRA_PER_K * K / pass rate), clamped to [MIN, MAX]; MAX if nothing
+    // passed. With MIN < 55, a 10k-vector index at window 100 and 10 % pass rate lost
+    // recall.
+    const float BATCH_EXTRA_PER_K = 0.2f;
+    const size_t BATCH_EXTRA_MIN = 60;
+    const size_t BATCH_EXTRA_MAX = 200;
+    // Upper bound on neighbor slots of one greedy search call (limits memory).
+    const size_t GREEDY_MAX_SLOTS = size_t{1} << 22;
 
     assert(results.n_queries() == queries.size());
 
@@ -174,13 +208,15 @@ std::vector<size_t> filtered_topk_search(
     const auto index_size = index.size();
     std::vector<size_t> found_neighbors_per_query(queries.size(), 0);
 
+    float pass_rate = 0.0f;
     auto initial_batch_size = estimate_initial_batch_size(
         id_filter,
         sample_generator,
         MIN_SAMPLE_SIZE,
         num_neighbors,
         initial_batch_hint,
-        index_size
+        index_size,
+        &pass_rate
     );
     if (initial_batch_size == 0) {
         // If the batch size is 0, it means that the filter rate is too low than
@@ -190,49 +226,122 @@ std::vector<size_t> filtered_topk_search(
     }
 
     const auto filter_rate = id_filter->filter_rate();
+    const bool greedy_first = pass_rate >= GREEDY_MIN_PASS_RATE;
+    const size_t batch_extra_slots =
+        pass_rate > 0.0f
+            ? std::clamp(
+                  static_cast<size_t>(std::ceil(
+                      BATCH_EXTRA_PER_K * static_cast<float>(num_neighbors) / pass_rate
+                  )),
+                  BATCH_EXTRA_MIN,
+                  BATCH_EXTRA_MAX
+              )
+            : BATCH_EXTRA_MAX;
 
-    auto search_closure = [&](const auto& range, uint64_t SVS_UNUSED(tid)) {
-        for (auto i : range) {
-            auto query = queries.get_datum(i);
-            auto iterator = index.batch_iterator(query);
-            size_t found = 0;
-            size_t total_checked = 0;
-            auto batch_size = initial_batch_size;
-            do {
-                batch_size = estimate_batch_size(
-                    total_checked, found, num_neighbors, batch_size, index_size
-                );
-                iterator.next(batch_size);
-                total_checked += iterator.size();
-                for (auto& neighbor : iterator.results()) {
-                    if (id_filter->is_member(neighbor.id())) {
-                        results.set(neighbor, i, found);
-                        found++;
-                        if (found == num_neighbors) {
-                            break;
-                        }
+    // Batch iterator + filter loop for query `i`; returns the number of neighbors found.
+    auto batch_search = [&](size_t i, size_t extra_slots) {
+        auto query = queries.get_datum(i);
+        auto iterator = index.batch_iterator(query, extra_slots);
+        size_t found = 0;
+        size_t total_checked = 0;
+        auto batch_size = initial_batch_size;
+        do {
+            batch_size = estimate_batch_size(
+                total_checked, found, num_neighbors, batch_size, index_size
+            );
+            iterator.next(batch_size);
+            total_checked += iterator.size();
+            for (auto& neighbor : iterator.results()) {
+                if (id_filter->is_member(neighbor.id())) {
+                    results.set(neighbor, i, found);
+                    found++;
+                    if (found == num_neighbors) {
+                        break;
                     }
                 }
-                // TODO: clarify the contract here - should we return partial or no
-                // result if the hit rate is too low
-                if (found < num_neighbors &&
-                    !hit_rate_sufficient(total_checked, found, filter_rate)) {
-                    found = 0;
-                    break;
-                }
-            } while (found < num_neighbors && !iterator.done());
+            }
+            // TODO: clarify the contract here - should we return partial or no
+            // result if the hit rate is too low
+            if (found < num_neighbors &&
+                !hit_rate_sufficient(total_checked, found, filter_rate)) {
+                found = 0;
+                break;
+            }
+        } while (found < num_neighbors && !iterator.done());
 
-            // Pad results if not enough neighbors found
-            pad_result(results, i, found);
-            found_neighbors_per_query[i] = found;
-        }
+        // Pad results if not enough neighbors found
+        pad_result(results, i, found);
+        return found;
     };
 
-    svs::threads::parallel_for(
-        index.get_threadpool_handle(),
-        svs::threads::StaticPartition{queries.size()},
-        search_closure
+    if (!greedy_first) {
+        auto search_closure = [&](const auto& range, uint64_t SVS_UNUSED(tid)) {
+            for (auto i : range) {
+                found_neighbors_per_query[i] = batch_search(i, batch_extra_slots);
+            }
+        };
+        svs::threads::parallel_for(
+            index.get_threadpool_handle(),
+            svs::threads::StaticPartition{queries.size()},
+            search_closure
+        );
+        return found_neighbors_per_query;
+    }
+
+    // Window and capacity both equal the number of neighbors asked for.
+    const auto greedy_n = std::min(
+        std::max(
+            initial_batch_hint,
+            static_cast<size_t>(
+                std::ceil(GREEDY_OVERSAMPLE * static_cast<float>(num_neighbors) / pass_rate)
+            )
+        ),
+        index_size
     );
+    auto greedy_params = params;
+    greedy_params.buffer_config_ = svs::index::vamana::SearchBufferConfig{greedy_n};
+    const size_t chunk =
+        std::max(GREEDY_MAX_SLOTS / std::max(greedy_n, size_t{1}), size_t{1});
+    const size_t dim = queries.dimensions();
+
+    for (size_t first = 0; first < queries.size(); first += chunk) {
+        const size_t count = std::min(chunk, queries.size() - first);
+        auto greedy = svs::QueryResult<size_t>(count, greedy_n);
+        // Parallelized by the index's own thread pool.
+        index.search(
+            greedy.view(),
+            svs::data::ConstSimpleDataView<float>(queries.data() + first * dim, count, dim),
+            greedy_params
+        );
+
+        auto filter_closure = [&](const auto& range, uint64_t SVS_UNUSED(tid)) {
+            for (auto c : range) {
+                const auto i = first + c;
+                size_t found = 0;
+                for (size_t j = 0; j < greedy_n && found < num_neighbors; ++j) {
+                    const auto id = greedy.index(c, j);
+                    if (id != static_cast<size_t>(-1) && id_filter->is_member(id)) {
+                        results.set(
+                            svs::Neighbor<size_t>{id, greedy.distance(c, j)}, i, found
+                        );
+                        found++;
+                    }
+                }
+                if (found < num_neighbors) {
+                    // Not enough hits: redo this query with the batch loop.
+                    found = batch_search(i, svs::UNSIGNED_INTEGER_PLACEHOLDER);
+                } else {
+                    pad_result(results, i, found);
+                }
+                found_neighbors_per_query[i] = found;
+            }
+        };
+        svs::threads::parallel_for(
+            index.get_threadpool_handle(),
+            svs::threads::StaticPartition{count},
+            filter_closure
+        );
+    }
     return found_neighbors_per_query;
 }
 

bindings/c/src/index_vamana.cpp

@@ -83,7 +83,7 @@ std::pair<svs::QueryResult<size_t>, std::vector<size_t>> IndexVamana::search(
         std::max(num_neighbors, params.buffer_config_.get_search_window_size());
 
     auto found_neighbors_per_query = filtered_topk_search(
-        index, results, queries, batch_hint, id_filter, sample_generator
+        index, params, results, queries, batch_hint, id_filter, sample_generator
     );
     return {results, found_neighbors_per_query};
 }
@@ -166,7 +166,7 @@ std::pair<svs::QueryResult<size_t>, std::vector<size_t>> DynamicIndexVamana::sea
         std::max(num_neighbors, params.buffer_config_.get_search_window_size());
 
     auto found_neighbors_per_query = filtered_topk_search(
-        index, results, queries, batch_hint, id_filter, sample_generator
+        index, params, results, queries, batch_hint, id_filter, sample_generator
     );
     return {results, found_neighbors_per_query};
 }

bindings/c/tests/c_api_dynamic_index.cpp

@@ -21,6 +21,7 @@
 #include "catch2/catch_test_macros.hpp"
 
 // Test utilities
+#include "c_api_filtered_topk_checks.h"
 #include "c_api_test_utils.h"
 
 // Standard library
@@ -1045,3 +1046,35 @@ CATCH_TEST_CASE(
         svs_error_free(error);
     }
 }
+
+CATCH_TEST_CASE(
+    "C API Dynamic Filtered Search topK by pass rate",
+    "[c_api][index][dynamic][search][filter]"
+) {
+    filtered_topk_checks::Dataset ds;
+    std::vector<size_t> ids(ds.num_vectors);
+    for (size_t i = 0; i < ds.num_vectors; ++i) {
+        ids[i] = i;
+    }
+    svs_error_h error = svs_error_create();
+    svs_algorithm_h algorithm = svs_algorithm_create_vamana(32, 100, 100, error);
+    CATCH_REQUIRE(algorithm != nullptr);
+    svs_index_builder_h builder = svs_index_builder_create(
+        SVS_DISTANCE_METRIC_EUCLIDEAN, ds.dimension, algorithm, error
+    );
+    CATCH_REQUIRE(builder != nullptr);
+    CATCH_REQUIRE(
+        svs_index_builder_set_threadpool(builder, SVS_THREADPOOL_KIND_NATIVE, 4, error)
+    );
+    svs_index_h index = svs_index_build_dynamic(
+        builder, ds.data.data(), ids.data(), ds.num_vectors, 1024 * 1024, error
+    );
+    CATCH_REQUIRE(index != nullptr);
+
+    filtered_topk_checks::run(index, ds);
+
+    svs_index_free(index);
+    svs_index_builder_free(builder);
+    svs_algorithm_free(algorithm);
+    svs_error_free(error);
+}