diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index 8e45ac691..dd216a591 100644 --- a/bindings/c/src/filtered_search.hpp +++ b/bindings/c/src/filtered_search.hpp @@ -21,7 +21,9 @@ #include #include +#include #include +#include #include #include @@ -79,6 +81,8 @@ 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, @@ -86,18 +90,19 @@ inline size_t estimate_initial_batch_size( 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(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(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(hits) / static_cast(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& result) { std::fill(result.indices().begin(), result.indices().end(), static_cast(-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 +template std::vector filtered_topk_search( IndexType& index, + const SearchParams& params, svs::QueryResult& results, svs::data::ConstSimpleDataView queries, size_t initial_batch_hint, @@ -165,6 +186,19 @@ std::vector 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 filtered_topk_search( const auto index_size = index.size(); std::vector 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 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(std::ceil( + BATCH_EXTRA_PER_K * static_cast(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( + std::ceil(GREEDY_OVERSAMPLE * static_cast(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(count, greedy_n); + // Parallelized by the index's own thread pool. + index.search( + greedy.view(), + svs::data::ConstSimpleDataView(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(-1) && id_filter->is_member(id)) { + results.set( + svs::Neighbor{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; } diff --git a/bindings/c/src/index_vamana.cpp b/bindings/c/src/index_vamana.cpp index bcaa00c3d..7d11d4811 100644 --- a/bindings/c/src/index_vamana.cpp +++ b/bindings/c/src/index_vamana.cpp @@ -83,7 +83,7 @@ std::pair, std::vector> 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, std::vector> 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}; } diff --git a/bindings/c/tests/c_api_dynamic_index.cpp b/bindings/c/tests/c_api_dynamic_index.cpp index 372ef2158..3c56df714 100644 --- a/bindings/c/tests/c_api_dynamic_index.cpp +++ b/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 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); +} diff --git a/bindings/c/tests/c_api_filtered_topk_checks.h b/bindings/c/tests/c_api_filtered_topk_checks.h new file mode 100644 index 000000000..cba9cadff --- /dev/null +++ b/bindings/c/tests/c_api_filtered_topk_checks.h @@ -0,0 +1,189 @@ +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +#pragma once + +// Checks of the filtered top-k search paths that depend on the filter pass rate. Shared by +// the static and the dynamic index tests: both call the same filtered search. + +#include "svs/c/svs_c.h" + +#include "catch2/catch_test_macros.hpp" + +#include +#include +#include +#include +#include + +namespace filtered_topk_checks { + +// Filter: ids that are multiples of the modulus stored in `self` (pass rate 1 / modulus). +inline bool is_multiple(void* self, size_t id) { + return id % *static_cast(self) == 0; +} +inline float rate_unknown(void*) { return 0.0f; } +inline float rate_half(void*) { return 0.5f; } +inline float rate_one_percent(void*) { return 0.01f; } + +struct Dataset { + size_t num_vectors = 5000; + size_t num_queries = 50; + size_t dimension = 32; + size_t k = 10; + std::vector data; + std::vector queries; + + Dataset() { + // Random data: the periodic test generator produces a disconnected graph. + std::mt19937 rng(7); + std::uniform_real_distribution uniform(0.0f, 1.0f); + data.resize(num_vectors * dimension); + queries.resize(num_queries * dimension); + for (auto& x : data) { + x = uniform(rng); + } + for (auto& x : queries) { + x = uniform(rng); + } + } + + // Exact top-k ids of query `q` among the ids with id % modulus == 0. + std::vector ground_truth(size_t q, size_t modulus) const { + std::vector> dist; + const float* query = queries.data() + q * dimension; + for (size_t id = 0; id < num_vectors; id += modulus) { + float d = 0; + for (size_t j = 0; j < dimension; ++j) { + float diff = query[j] - data[id * dimension + j]; + d += diff * diff; + } + dist.emplace_back(d, id); + } + std::partial_sort(dist.begin(), dist.begin() + k, dist.end()); + std::vector top; + for (size_t i = 0; i < k; ++i) { + top.push_back(dist[i].second); + } + return top; + } +}; + +struct Outcome { + double recall = 0; + size_t min_found = 0; // fewest results returned for any query + size_t max_found = 0; // most results returned for any query +}; + +// Runs a top-k search with the filter "id % modulus == 0" (modulus 0: no filter) and +// compares with the exact filtered results. Checks that every returned id passes the +// filter and, if `check_order`, that distances are non-decreasing (results gathered from +// several iterator batches are sorted within a batch only). +inline Outcome search( + svs_index_h index, + const Dataset& ds, + svs_search_params_h params, + size_t modulus, + float (*filter_rate)(void*), + bool check_order +) { + svs_error_h error = svs_error_create(); + svs_id_filter_interface_ops ops = SVS_INIT_ID_FILTER_OPS(is_multiple, *filter_rate); + svs_id_filter_interface filter = SVS_MAKE_INTERFACE(&modulus, ops); + svs_search_results_t results = SVS_INIT_SEARCH_RESULTS(); + CATCH_REQUIRE(svs_index_search_topk( + index, + ds.queries.data(), + ds.num_queries, + ds.k, + &results, + params, + modulus == 0 ? nullptr : &filter, + error + )); + CATCH_REQUIRE(svs_error_ok(error)); + CATCH_REQUIRE(results.num_queries == ds.num_queries); + + Outcome out; + out.min_found = ds.k; + size_t hits = 0; + size_t expected = 0; + for (size_t q = 0; q < ds.num_queries; ++q) { + const size_t begin = results.offsets[q]; + const size_t found = results.offsets[q + 1] - begin; + CATCH_REQUIRE(found <= ds.k); + out.min_found = std::min(out.min_found, found); + out.max_found = std::max(out.max_found, found); + auto truth = ds.ground_truth(q, std::max(modulus, 1)); + std::unordered_set truth_set(truth.begin(), truth.end()); + expected += truth.size(); + for (size_t j = 0; j < found; ++j) { + const size_t id = results.indices[begin + j]; + CATCH_REQUIRE((modulus == 0 || id % modulus == 0)); + if (check_order && j > 0) { + CATCH_REQUIRE( + results.distances[begin + j] >= results.distances[begin + j - 1] + ); + } + hits += truth_set.count(id); + } + } + svs_search_results_free(&results); + svs_error_free(error); + out.recall = static_cast(hits) / static_cast(expected); + return out; +} + +// `index` holds ds.data with ids 0..num_vectors-1. +inline void run(svs_index_h index, const Dataset& ds) { + svs_error_h error = svs_error_create(); + svs_search_params_h params = svs_search_params_create_vamana(100, error); + CATCH_REQUIRE(params != nullptr); + + // Reference: the same search without a filter. + const auto unfiltered = search(index, ds, params, 0, rate_unknown, true); + CATCH_REQUIRE(unfiltered.min_found == ds.k); + + // Pass rate 1/2 and 1 (unknown filter rate): the greedy oversampled search. Every query + // gets k results and the recall is comparable to the unfiltered search. + for (size_t modulus : {1, 2}) { + const auto out = search(index, ds, params, modulus, rate_unknown, true); + CATCH_REQUIRE(out.min_found == ds.k); + CATCH_REQUIRE(out.recall >= unfiltered.recall - 0.05); + } + + // Pass rates 1/10 and 1/100: the batch iterator. A window of 400 gathers enough hits in + // the first batch, so this does not depend on how a search resumes with a larger + // window. Every query still gets k results at 1/10; at 1/100 all returned ids must + // pass the filter (checked in search()). + svs_search_params_h wide = svs_search_params_create_vamana(400, error); + CATCH_REQUIRE(wide != nullptr); + const auto tenth = search(index, ds, wide, 10, rate_unknown, false); + CATCH_REQUIRE(tenth.min_found == ds.k); + CATCH_REQUIRE(tenth.recall >= 0.8); + const auto hundredth = search(index, ds, wide, 100, rate_unknown, false); + CATCH_REQUIRE(hundredth.max_found <= ds.k); + svs_search_params_free(wide); + + // A positive filter rate above the actual pass rate (1/20) stops the search early with + // an empty result; a rate below it does not. + const auto stopped = search(index, ds, params, 20, rate_half, false); + CATCH_REQUIRE(stopped.max_found == 0); + const auto allowed = search(index, ds, params, 20, rate_one_percent, false); + CATCH_REQUIRE(allowed.min_found == ds.k); + + svs_search_params_free(params); + svs_error_free(error); +} + +} // namespace filtered_topk_checks diff --git a/bindings/c/tests/c_api_index.cpp b/bindings/c/tests/c_api_index.cpp index 46fd1215c..054afb6f1 100644 --- a/bindings/c/tests/c_api_index.cpp +++ b/bindings/c/tests/c_api_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 @@ -1548,3 +1549,28 @@ CATCH_TEST_CASE("C API Filtered Search topK", "[c_api][index][search][filter]") svs_error_free(error); } } + +CATCH_TEST_CASE( + "C API Filtered Search topK by pass rate", "[c_api][index][search][filter]" +) { + filtered_topk_checks::Dataset ds; + 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(builder, ds.data.data(), ds.num_vectors, 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); +}