From 1df9cc8f4b3673609a2e5741d0cb601acbebe8c2 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 14 Oct 2026 18:20:26 -0700 Subject: [PATCH 1/3] Add single-pass filtered top-k search for the dynamic Vamana index Filtered top-k in the C API ran an unfiltered search through BatchIterator, filtered each batch and widened the window until k results passed. At selective filters the search buffer grows to about k / selectivity entries, and with the visited set enabled the resumed search cannot re-admit candidates that were seen and then dropped, which lowers recall. Add index::vamana::two_heap_filtered_search: one graph traversal with a candidate min-heap, a result max-heap of non-deleted nodes that pass the filter (capped at max(k, search window)), and an exact visited set. Failing nodes are still expanded. DynamicVamana::two_heap_filtered_search exposes it through a new virtual appended last in DynamicVamanaInterface; it returns false for data types that need a custom search path (anything other than uncompressed and scalar-quantized data), and the C API then keeps using the batch iterator. The static index is unchanged. --- bindings/c/src/filtered_search.hpp | 21 ++ .../index/vamana/two_heap_filtered_search.h | 280 ++++++++++++++++++ include/svs/orchestrators/dynamic_vamana.h | 59 ++++ 3 files changed, 360 insertions(+) create mode 100644 include/svs/index/vamana/two_heap_filtered_search.h diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index 8e45ac691..e802d8d16 100644 --- a/bindings/c/src/filtered_search.hpp +++ b/bindings/c/src/filtered_search.hpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -31,6 +32,7 @@ #include #include #include +#include #include namespace svs::c_runtime { @@ -168,6 +170,25 @@ std::vector filtered_topk_search( assert(results.n_queries() == queries.size()); + // The dynamic index uses a single-pass search with a result heap of size + // max(k, search window) when its data type supports it. + if constexpr (std::is_same_v) { + std::vector found(queries.size(), 0); + const bool done = index.two_heap_filtered_search( + results.view(), + queries, + std::max(results.n_neighbors(), initial_batch_hint), + [id_filter](size_t id) { return id_filter->is_member(id); }, + found + ); + if (done) { + for (size_t i = 0; i < queries.size(); ++i) { + pad_result(results, i, found[i]); + } + return found; + } + } + // Filtered search: we need to estimate the batch size based on the filter rate and // the number of hits const auto num_neighbors = results.n_neighbors(); diff --git a/include/svs/index/vamana/two_heap_filtered_search.h b/include/svs/index/vamana/two_heap_filtered_search.h new file mode 100644 index 000000000..9db834e24 --- /dev/null +++ b/include/svs/index/vamana/two_heap_filtered_search.h @@ -0,0 +1,280 @@ +/* + * 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. + */ + +// Single-pass filtered top-k search for the dynamic Vamana index, following hnswlib's +// filtered `searchBaseLayerST`. +// +// - `candidates`: min-heap of nodes still to expand. +// - `results`: max-heap of nodes that pass the filter and are not deleted, capped at +// `ef`. Nodes that fail are expanded but never enter `results`. +// - The search stops when the best candidate is worse than the worst result and +// `results` is full. +// - The visited set is exact: one tag per slot, reset per query by a generation counter. +#pragma once + +#include "svs/concepts/data.h" +#include "svs/concepts/distance.h" +#include "svs/core/data/simple.h" +#include "svs/core/query_result.h" +#include "svs/index/vamana/extensions.h" +#include "svs/index/vamana/greedy_search.h" +#include "svs/lib/neighbor.h" +#include "svs/lib/prefetch.h" +#include "svs/lib/threads.h" +#include "svs/quantization/scalar/scalar.h" + +#include +#include +#include +#include +#include + +namespace svs::index::vamana { +namespace detail { + +// Per-thread scratch reused across queries. +template struct TwoHeapScratch { + // `tags[slot] == generation` means the slot was visited by the current query. + std::vector tags; + uint32_t generation = 0; + std::vector> candidates; + std::vector> results; +}; + +template TwoHeapScratch& two_heap_scratch() { + thread_local TwoHeapScratch scratch; + return scratch; +} + +// Searches one query. Writes up to `output.n_neighbors()` neighbors (external ids) into +// row `query_index` of `output` and returns how many were written. +template < + typename Index, + typename Graph, + typename Data, + typename Accessor, + typename DistFn, + typename Query, + typename Pred> +size_t two_heap_search_one( + const Index& index, + const Graph& graph, + const Data& data, + Accessor& accessor, + DistFn& distance, + std::span entry_points, + const Query& query, + size_t ef, + const Pred& pred, + GreedySearchPrefetchParameters prefetch_parameters, + QueryResultView& output, + size_t query_index +) { + using Idx = typename Graph::index_type; + using Entry = std::pair; + + auto& scratch = two_heap_scratch(); + auto& candidates = scratch.candidates; + auto& results = scratch.results; + candidates.clear(); + results.clear(); + + // Start a new generation; on wrap-around, clear all tags. + if (scratch.tags.size() < data.size()) { + scratch.tags.resize(data.size(), 0); + } + if (++scratch.generation == 0) { + std::fill(scratch.tags.begin(), scratch.tags.end(), 0); + scratch.generation = 1; + } + uint32_t* tags = scratch.tags.data(); + const uint32_t generation = scratch.generation; + + distance::maybe_fix_argument(distance, query); + // `cmp(a, b)` is true if distance `a` is strictly better than `b`. + // std heaps keep the "largest" element on top: best candidate, worst result. + const auto cmp = distance::comparator(distance); + const auto candidate_less = [&](const Entry& a, const Entry& b) { + return cmp(b.first, a.first); + }; + const auto result_less = [&](const Entry& a, const Entry& b) { + return cmp(a.first, b.first); + }; + + // Check deletion first: deleted slots may have no external id. + const auto passes = [&](Idx i) { + return !index.is_deleted(i) && pred(index.translate_internal_id(i)); + }; + const auto compute = [&](Idx i) { + return static_cast(distance::compute(distance, query, accessor(data, i))); + }; + // Keep `results` at most `ef` long by dropping the worst. + const auto add_result = [&](float dist, Idx i) { + results.emplace_back(dist, i); + std::push_heap(results.begin(), results.end(), result_less); + if (results.size() > ef) { + std::pop_heap(results.begin(), results.end(), result_less); + results.pop_back(); + } + }; + + // Entry points are always expanded; they are results only if they pass. + for (Idx ep : entry_points) { + tags[ep] = generation; + float dist = compute(ep); + candidates.emplace_back(dist, ep); + std::push_heap(candidates.begin(), candidates.end(), candidate_less); + if (passes(ep)) { + add_result(dist, ep); + } + } + + while (!candidates.empty()) { + const auto [candidate_dist, candidate_id] = candidates.front(); + // Best candidate is worse than the worst of a full result set. + if (results.size() == ef && cmp(results.front().first, candidate_dist)) { + break; + } + std::pop_heap(candidates.begin(), candidates.end(), candidate_less); + candidates.pop_back(); + + auto neighbors = graph.get_node(candidate_id); + const size_t num_neighbors = neighbors.size(); + // Same prefetching scheme as `greedy_search`. + auto prefetcher = lib::make_prefetcher( + lib::PrefetchParameters{ + prefetch_parameters.lookahead, prefetch_parameters.step}, + num_neighbors, + [&](size_t i) { accessor.prefetch(data, neighbors[i]); }, + [&](size_t i) { + if (i + 1 < num_neighbors) { + lib::prefetch_l0(tags + neighbors[i + 1]); + } + return tags[neighbors[i]] != generation; + } + ); + + prefetcher(); + for (auto id : neighbors) { + if (tags[id] == generation) { + continue; + } + tags[id] = generation; + prefetcher(); + + float dist = compute(id); + // Only nodes that could improve `results` are queued, and only those + // evaluate the filter. + if (results.size() < ef || cmp(dist, results.front().first)) { + candidates.emplace_back(dist, id); + std::push_heap(candidates.begin(), candidates.end(), candidate_less); + if (passes(id)) { + add_result(dist, id); + } + } + } + } + + // Best result first. + std::sort_heap(results.begin(), results.end(), result_less); + const size_t found = std::min(results.size(), output.n_neighbors()); + for (size_t j = 0; j < found; ++j) { + output.set( + Neighbor{ + index.translate_internal_id(results[j].second), results[j].first}, + query_index, + j + ); + } + return found; +} + +} // namespace detail + +/// Data types for which `two_heap_filtered_search` is valid: the default search path +/// (identity accessor, no reranking) must apply. +template inline constexpr bool two_heap_supported = false; +template +inline constexpr bool two_heap_supported> = true; +template +inline constexpr bool + two_heap_supported> = true; + +/// +/// @brief Filtered top-k search over a dynamic Vamana index. +/// +/// @param index The `MutableVamanaIndex` to search. +/// @param output Row `i` receives the neighbors (external ids) of query `i`. Rows with +/// fewer than `n_neighbors` results are not padded. +/// @param queries One datum per query. +/// @param ef Result-set capacity (the search window). Raised to `n_neighbors` if smaller. +/// @param pred Callable `bool(size_t external_id)`. Must be thread safe. +/// @param found Resized to `queries.size()`; receives the number of results per query. +/// @param prefetch_parameters Prefetch lookahead and step for neighbor vectors. +/// +/// Requires `two_heap_supported`. +/// +template +void two_heap_filtered_search( + Index& index, + QueryResultView output, + const Queries& queries, + size_t ef, + const Pred& pred, + std::vector& found, + GreedySearchPrefetchParameters prefetch_parameters = {} +) { + using Idx = typename Index::internal_id_type; + static_assert( + two_heap_supported, + "Data type needs a custom search path (e.g. reranking)." + ); + ef = std::max(ef, output.n_neighbors()); + found.assign(queries.size(), 0); + auto& pool = index.get_threadpool_handle(); + + index.experimental_escape_hatch([&](const auto& graph, + const auto& data, + const auto& distance, + std::span entry_points) { + threads::parallel_for( + pool, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + auto dist_fn = extensions::per_thread_batch_search_setup(data, distance); + auto accessor = data::GetDatumAccessor(); + for (auto i : is) { + found[i] = detail::two_heap_search_one( + index, + graph, + data, + accessor, + dist_fn, + entry_points, + queries.get_datum(i), + ef, + pred, + prefetch_parameters, + output, + i + ); + } + } + ); + }); +} + +} // namespace svs::index::vamana diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 7ebfaaaa5..36a57820b 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -17,12 +17,15 @@ #pragma once #include "svs/index/vamana/dynamic_index.h" +#include "svs/index/vamana/two_heap_filtered_search.h" #include "svs/orchestrators/manager.h" #include "svs/orchestrators/vamana.h" // stdlib +#include #include +#include namespace svs { @@ -56,6 +59,17 @@ class DynamicVamanaInterface : public VamanaInterface { // Check if the index has been consolidated - no `Deleted` entries remain. virtual bool is_consolidated() const = 0; + + // Single-pass filtered top-k; `filter` takes an external id. Returns false, leaving + // `results` untouched, if the data type does not support it. + // Keep last: new virtuals must be appended. + virtual bool two_heap_filtered_search( + svs::QueryResultView results, + AnonymousArray<2> queries, + size_t ef, + const std::function& filter, + std::vector& found + ) = 0; }; template @@ -109,6 +123,33 @@ class DynamicVamanaImpl : public VamanaImpl results, + AnonymousArray<2> queries, + size_t ef, + const std::function& filter, + std::vector& found + ) override { + if constexpr (index::vamana::two_heap_supported) { + lib::match( + QueryTypes{}, + queries.type(), + [&](lib::Type SVS_UNUSED(type)) { + const auto view = data::ConstSimpleDataView(queries); + index::vamana::two_heap_filtered_search( + impl(), results, view, ef, filter, found + ); + }, + [&](svs::DataType data_type) { + throw ANNEXCEPTION("Unsupported query datatype! Got: {}.", data_type); + } + ); + return true; + } else { + return false; + } + } }; // Forward Declaractions. @@ -506,6 +547,24 @@ class DynamicVamana : public manager::IndexManager { ); } + /// @brief Single-pass filtered top-k search, see ``two_heap_filtered_search.h``. + /// + /// @param filter Takes an external id. + /// @param found Receives the number of results per query. Rows are not padded. + /// @return False if the data type is unsupported, in which case nothing is written. + template + bool two_heap_filtered_search( + QueryResultView results, + data::ConstSimpleDataView queries, + size_t ef, + const std::function& filter, + std::vector& found + ) { + return impl_->two_heap_filtered_search( + results, AnonymousArray<2>(queries), ef, filter, found + ); + } + ///// Experimental Calibration template < data::ImmutableMemoryDataset Queries, From a4e7bb63780ff27ae305abd874e58a33b63f25e4 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 14 Oct 2026 18:20:26 -0700 Subject: [PATCH 2/3] Test filtered top-k recall, deletions and empty filter on the dynamic C API index Use random data in the existing dynamic filtered search test: the periodic generator produces a few distinct vectors, which leaves the graph disconnected. --- bindings/c/tests/c_api_dynamic_index.cpp | 187 ++++++++++++++++++++++- 1 file changed, 183 insertions(+), 4 deletions(-) diff --git a/bindings/c/tests/c_api_dynamic_index.cpp b/bindings/c/tests/c_api_dynamic_index.cpp index 372ef2158..9e798eecc 100644 --- a/bindings/c/tests/c_api_dynamic_index.cpp +++ b/bindings/c/tests/c_api_dynamic_index.cpp @@ -24,6 +24,12 @@ #include "c_api_test_utils.h" // Standard library +#include +#include +#include +#include +#include +#include #include CATCH_TEST_CASE("C API Dynamic Index", "[c_api][index][dynamic]") { @@ -892,11 +898,19 @@ CATCH_TEST_CASE( const size_t NUM_THREADS = 4; const size_t BLOCK_SIZE = 1024 * 1024; // 1 MB block size for testing - std::vector data; - std::vector queries; + // generate_test_data() yields only a few distinct vectors, which leaves the graph + // disconnected; use random data so every vector is reachable. + std::mt19937 rng(7); + std::uniform_real_distribution uniform(0.0f, 1.0f); + std::vector data(NUM_VECTORS * DIMENSION); + std::vector queries(NUM_QUERIES * DIMENSION); std::vector ids(NUM_VECTORS); - generate_test_data(data, NUM_VECTORS, DIMENSION); - generate_test_data(queries, NUM_QUERIES, DIMENSION); + for (auto& x : data) { + x = uniform(rng); + } + for (auto& x : queries) { + x = uniform(rng); + } for (size_t i = 0; i < NUM_VECTORS; ++i) { ids[i] = i; } @@ -1045,3 +1059,168 @@ CATCH_TEST_CASE( svs_error_free(error); } } + +namespace { + +// Accepts IDs that are multiples of the modulus stored in `self`. +bool filter_modulus(void* self, size_t id) { + return (id % *static_cast(self)) == 0; +} +bool filter_none(void* /*self*/, size_t /*id*/) { return false; } +float filter_rate_unknown(void* /*self*/) { return 0.0f; } + +// Exact top-k (squared L2) over the IDs that pass `keep`. +template +std::vector restricted_ground_truth( + const std::vector& data, + const float* query, + size_t dimension, + size_t k, + const Keep& keep +) { + std::vector> scored; + for (size_t id = 0; id < data.size() / dimension; ++id) { + if (keep(id)) { + scored.emplace_back( + euclidean_distance(query, data.data() + id * dimension, dimension), id + ); + } + } + const size_t n = std::min(k, scored.size()); + std::partial_sort(scored.begin(), scored.begin() + n, scored.end()); + std::vector top; + for (size_t i = 0; i < n; ++i) { + top.push_back(scored[i].second); + } + return top; +} + +} // namespace + +CATCH_TEST_CASE( + "C API Dynamic Filtered Search recall", "[c_api][index][dynamic][search][filter]" +) { + const size_t NUM_VECTORS = 10000; + const size_t NUM_QUERIES = 20; + const size_t DIMENSION = 32; + const size_t K = 10; + const size_t NUM_THREADS = 4; + const size_t BLOCK_SIZE = 1024 * 1024; + + std::mt19937 rng(123); + std::uniform_real_distribution uniform(0.0f, 1.0f); + std::vector data(NUM_VECTORS * DIMENSION); + std::vector queries(NUM_QUERIES * DIMENSION); + for (auto& x : data) { + x = uniform(rng); + } + for (auto& x : queries) { + x = uniform(rng); + } + std::vector ids(NUM_VECTORS); + for (size_t i = 0; i < 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, DIMENSION, algorithm, error + ); + CATCH_REQUIRE(builder != nullptr); + CATCH_REQUIRE(svs_index_builder_set_threadpool( + builder, SVS_THREADPOOL_KIND_NATIVE, NUM_THREADS, error + )); + svs_index_h index = svs_index_build_dynamic( + builder, data.data(), ids.data(), NUM_VECTORS, BLOCK_SIZE, error + ); + CATCH_REQUIRE(index != nullptr); + svs_search_params_h search_params = svs_search_params_create_vamana(100, error); + CATCH_REQUIRE(search_params != nullptr); + + std::unordered_set deleted; + + // Runs a filtered search and returns the mean recall against exact results restricted + // to the filter and to non-deleted IDs. Also checks that every returned ID is valid. + auto run = [&](svs_id_filter_interface& filter, auto keep) { + svs_search_results_t results = SVS_INIT_SEARCH_RESULTS(); + CATCH_REQUIRE(svs_index_search_topk( + index, queries.data(), NUM_QUERIES, K, &results, search_params, &filter, error + )); + CATCH_REQUIRE(svs_error_ok(error)); + CATCH_REQUIRE(results.num_queries == NUM_QUERIES); + + size_t hits = 0; + size_t expected = 0; + for (size_t q = 0; q < NUM_QUERIES; ++q) { + const size_t num_found = results.offsets[q + 1] - results.offsets[q]; + CATCH_REQUIRE(num_found <= K); + auto truth = restricted_ground_truth( + data, queries.data() + q * DIMENSION, DIMENSION, K, keep + ); + std::unordered_set truth_set(truth.begin(), truth.end()); + expected += truth.size(); + for (size_t j = 0; j < num_found; ++j) { + size_t id = results.indices[results.offsets[q] + j]; + CATCH_REQUIRE(keep(id)); + hits += truth_set.count(id); + } + } + svs_search_results_free(&results); + return static_cast(hits) / static_cast(expected); + }; + + for (size_t modulus : {2, 10, 100}) { + CATCH_SECTION("Recall with 1/" + std::to_string(modulus) + " selectivity") { + svs_id_filter_interface_ops ops = + SVS_INIT_ID_FILTER_OPS(filter_modulus, filter_rate_unknown); + svs_id_filter_interface filter = SVS_MAKE_INTERFACE(&modulus, ops); + auto keep = [&](size_t id) { return id % modulus == 0; }; + CATCH_REQUIRE(run(filter, keep) >= 0.95); + } + } + + CATCH_SECTION("Deleted points are never returned") { + // Delete half of the IDs that pass the filter. + std::vector to_delete; + for (size_t id = 0; id < NUM_VECTORS; id += 20) { + to_delete.push_back(id); + deleted.insert(id); + } + size_t deleted_count = 0; + CATCH_REQUIRE(svs_index_dynamic_delete_points( + index, to_delete.data(), to_delete.size(), &deleted_count, error + )); + CATCH_REQUIRE(deleted_count == to_delete.size()); + + size_t modulus = 10; + svs_id_filter_interface_ops ops = + SVS_INIT_ID_FILTER_OPS(filter_modulus, filter_rate_unknown); + svs_id_filter_interface filter = SVS_MAKE_INTERFACE(&modulus, ops); + auto keep = [&](size_t id) { return id % modulus == 0 && !deleted.count(id); }; + CATCH_REQUIRE(run(filter, keep) >= 0.95); + } + + CATCH_SECTION("Filter matching nothing returns no results") { + svs_id_filter_interface_ops ops = + SVS_INIT_ID_FILTER_OPS(filter_none, filter_rate_unknown); + svs_id_filter_interface filter = SVS_MAKE_INTERFACE(nullptr, ops); + + svs_search_results_t results = SVS_INIT_SEARCH_RESULTS(); + CATCH_REQUIRE(svs_index_search_topk( + index, queries.data(), NUM_QUERIES, K, &results, search_params, &filter, error + )); + CATCH_REQUIRE(svs_error_ok(error)); + for (size_t q = 0; q < NUM_QUERIES; ++q) { + CATCH_REQUIRE(results.offsets[q + 1] == results.offsets[q]); + } + svs_search_results_free(&results); + } + + svs_search_params_free(search_params); + svs_index_free(index); + svs_index_builder_free(builder); + svs_algorithm_free(algorithm); + svs_error_free(error); +} From 5bf53f8f060649c86015e4f0095eb59b7013fef7 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Thu, 15 Oct 2026 14:24:28 -0700 Subject: [PATCH 3/3] Use SVS search-window terminology in filtered top-k search --- bindings/c/src/filtered_search.hpp | 2 +- .../index/vamana/two_heap_filtered_search.h | 55 +++++++++++-------- include/svs/orchestrators/dynamic_vamana.h | 10 ++-- 3 files changed, 37 insertions(+), 30 deletions(-) diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index e802d8d16..c82134e1b 100644 --- a/bindings/c/src/filtered_search.hpp +++ b/bindings/c/src/filtered_search.hpp @@ -170,7 +170,7 @@ std::vector filtered_topk_search( assert(results.n_queries() == queries.size()); - // The dynamic index uses a single-pass search with a result heap of size + // The dynamic index uses a single-pass search with a filtered result set of size // max(k, search window) when its data type supports it. if constexpr (std::is_same_v) { std::vector found(queries.size(), 0); diff --git a/include/svs/index/vamana/two_heap_filtered_search.h b/include/svs/index/vamana/two_heap_filtered_search.h index 9db834e24..37ebe77dd 100644 --- a/include/svs/index/vamana/two_heap_filtered_search.h +++ b/include/svs/index/vamana/two_heap_filtered_search.h @@ -14,15 +14,18 @@ * limitations under the License. */ -// Single-pass filtered top-k search for the dynamic Vamana index, following hnswlib's -// filtered `searchBaseLayerST`. +// Single-pass filtered top-k search for the dynamic Vamana index, using the same stopping +// rule as hnswlib's filtered search. // // - `candidates`: min-heap of nodes still to expand. -// - `results`: max-heap of nodes that pass the filter and are not deleted, capped at -// `ef`. Nodes that fail are expanded but never enter `results`. +// - `results` (the filtered search buffer): max-heap of nodes that pass the filter and +// are not deleted, capped at the search window size. Nodes that fail the filter are +// expanded but never enter `results`. // - The search stops when the best candidate is worse than the worst result and // `results` is full. -// - The visited set is exact: one tag per slot, reset per query by a generation counter. +// - The visited set is an exact tag array (one tag per slot, reset per query by a +// generation counter). SVS's `VisitedFilter` is lossy and could let duplicates into +// the candidate heap, so it is not used here. #pragma once #include "svs/concepts/data.h" @@ -50,8 +53,8 @@ template struct TwoHeapScratch { // `tags[slot] == generation` means the slot was visited by the current query. std::vector tags; uint32_t generation = 0; - std::vector> candidates; - std::vector> results; + std::vector> candidates; + std::vector> results; }; template TwoHeapScratch& two_heap_scratch() { @@ -77,14 +80,14 @@ size_t two_heap_search_one( DistFn& distance, std::span entry_points, const Query& query, - size_t ef, + size_t search_window_size, const Pred& pred, GreedySearchPrefetchParameters prefetch_parameters, QueryResultView& output, size_t query_index ) { using Idx = typename Graph::index_type; - using Entry = std::pair; + using Entry = Neighbor; auto& scratch = two_heap_scratch(); auto& candidates = scratch.candidates; @@ -108,10 +111,10 @@ size_t two_heap_search_one( // std heaps keep the "largest" element on top: best candidate, worst result. const auto cmp = distance::comparator(distance); const auto candidate_less = [&](const Entry& a, const Entry& b) { - return cmp(b.first, a.first); + return cmp(b.distance(), a.distance()); }; const auto result_less = [&](const Entry& a, const Entry& b) { - return cmp(a.first, b.first); + return cmp(a.distance(), b.distance()); }; // Check deletion first: deleted slots may have no external id. @@ -121,11 +124,11 @@ size_t two_heap_search_one( const auto compute = [&](Idx i) { return static_cast(distance::compute(distance, query, accessor(data, i))); }; - // Keep `results` at most `ef` long by dropping the worst. + // Keep `results` at most `search_window_size` long by dropping the worst. const auto add_result = [&](float dist, Idx i) { - results.emplace_back(dist, i); + results.emplace_back(i, dist); std::push_heap(results.begin(), results.end(), result_less); - if (results.size() > ef) { + if (results.size() > search_window_size) { std::pop_heap(results.begin(), results.end(), result_less); results.pop_back(); } @@ -135,7 +138,7 @@ size_t two_heap_search_one( for (Idx ep : entry_points) { tags[ep] = generation; float dist = compute(ep); - candidates.emplace_back(dist, ep); + candidates.emplace_back(ep, dist); std::push_heap(candidates.begin(), candidates.end(), candidate_less); if (passes(ep)) { add_result(dist, ep); @@ -143,9 +146,11 @@ size_t two_heap_search_one( } while (!candidates.empty()) { - const auto [candidate_dist, candidate_id] = candidates.front(); + const auto candidate_dist = candidates.front().distance(); + const auto candidate_id = candidates.front().id(); // Best candidate is worse than the worst of a full result set. - if (results.size() == ef && cmp(results.front().first, candidate_dist)) { + if (results.size() == search_window_size && + cmp(results.front().distance(), candidate_dist)) { break; } std::pop_heap(candidates.begin(), candidates.end(), candidate_less); @@ -178,8 +183,9 @@ size_t two_heap_search_one( float dist = compute(id); // Only nodes that could improve `results` are queued, and only those // evaluate the filter. - if (results.size() < ef || cmp(dist, results.front().first)) { - candidates.emplace_back(dist, id); + if (results.size() < search_window_size || + cmp(dist, results.front().distance())) { + candidates.emplace_back(id, dist); std::push_heap(candidates.begin(), candidates.end(), candidate_less); if (passes(id)) { add_result(dist, id); @@ -194,7 +200,7 @@ size_t two_heap_search_one( for (size_t j = 0; j < found; ++j) { output.set( Neighbor{ - index.translate_internal_id(results[j].second), results[j].first}, + index.translate_internal_id(results[j].id()), results[j].distance()}, query_index, j ); @@ -220,7 +226,8 @@ inline constexpr bool /// @param output Row `i` receives the neighbors (external ids) of query `i`. Rows with /// fewer than `n_neighbors` results are not padded. /// @param queries One datum per query. -/// @param ef Result-set capacity (the search window). Raised to `n_neighbors` if smaller. +/// @param search_window_size Capacity of the filtered result set. Raised to `n_neighbors` +/// if smaller. /// @param pred Callable `bool(size_t external_id)`. Must be thread safe. /// @param found Resized to `queries.size()`; receives the number of results per query. /// @param prefetch_parameters Prefetch lookahead and step for neighbor vectors. @@ -232,7 +239,7 @@ void two_heap_filtered_search( Index& index, QueryResultView output, const Queries& queries, - size_t ef, + size_t search_window_size, const Pred& pred, std::vector& found, GreedySearchPrefetchParameters prefetch_parameters = {} @@ -242,7 +249,7 @@ void two_heap_filtered_search( two_heap_supported, "Data type needs a custom search path (e.g. reranking)." ); - ef = std::max(ef, output.n_neighbors()); + search_window_size = std::max(search_window_size, output.n_neighbors()); found.assign(queries.size(), 0); auto& pool = index.get_threadpool_handle(); @@ -265,7 +272,7 @@ void two_heap_filtered_search( dist_fn, entry_points, queries.get_datum(i), - ef, + search_window_size, pred, prefetch_parameters, output, diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 36a57820b..df4565f18 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -66,7 +66,7 @@ class DynamicVamanaInterface : public VamanaInterface { virtual bool two_heap_filtered_search( svs::QueryResultView results, AnonymousArray<2> queries, - size_t ef, + size_t search_window_size, const std::function& filter, std::vector& found ) = 0; @@ -127,7 +127,7 @@ class DynamicVamanaImpl : public VamanaImpl results, AnonymousArray<2> queries, - size_t ef, + size_t search_window_size, const std::function& filter, std::vector& found ) override { @@ -138,7 +138,7 @@ class DynamicVamanaImpl : public VamanaImpl(lib::Type SVS_UNUSED(type)) { const auto view = data::ConstSimpleDataView(queries); index::vamana::two_heap_filtered_search( - impl(), results, view, ef, filter, found + impl(), results, view, search_window_size, filter, found ); }, [&](svs::DataType data_type) { @@ -556,12 +556,12 @@ class DynamicVamana : public manager::IndexManager { bool two_heap_filtered_search( QueryResultView results, data::ConstSimpleDataView queries, - size_t ef, + size_t search_window_size, const std::function& filter, std::vector& found ) { return impl_->two_heap_filtered_search( - results, AnonymousArray<2>(queries), ef, filter, found + results, AnonymousArray<2>(queries), search_window_size, filter, found ); }