diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9336f82e1..ca4afba3a 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -94,19 +94,32 @@ inline constexpr std::string_view name(SlotMetadata metadata) { class ValidBuilder { public: - ValidBuilder(const std::vector& status) - : status_{status} {} + // Optional `pred` (with `translator`): an entry is also invalid if `pred(external_id)` + // is false. Invalid entries are navigated through but never returned. + ValidBuilder( + const std::vector& status, + const IDTranslator* translator = nullptr, + const std::function* pred = nullptr + ) + : status_{status} + , translator_{translator} + , pred_{pred} {} template constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { bool invalid = getindex(status_, i) == SlotMetadata::Deleted; // This neighbor should be skipped if the metadata corresponding to the given index // marks this slot as deleted. + if (!invalid && pred_ != nullptr) { + invalid = !(*pred_)(translator_->get_external(i)); + } return PredicatedSearchNeighbor(i, distance, !invalid); } private: const std::vector& status_; + const IDTranslator* translator_; + const std::function* pred_; }; template @@ -532,6 +545,9 @@ class MutableVamanaIndex { [&](const auto is, uint64_t /*tid*/) { for (auto i : is) { for (size_t j = 0, jmax = getsize<1>(ids); j < jmax; ++j) { + if (ids.at(i, j) == static_cast(-1)) { + continue; // Padding from a filtered search. + } auto internal = lib::narrow_cast(ids.at(i, j)); ids.at(i, j) = translate_internal_id(internal); } @@ -573,13 +589,17 @@ class MutableVamanaIndex { // Return a `greedy_search` compatible builder for this index. // This is an internal method, mostly used to help implement the batch iterator. - ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } + ValidBuilder internal_search_builder(const std::function* pred = nullptr + ) const { + return ValidBuilder{status_, &translator_, pred}; + } auto greedy_search_closure( GreedySearchPrefetchParameters prefetch_parameters, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), + const std::function* pred = nullptr ) const { - return [&, prefetch_parameters]( + return [&, prefetch_parameters, pred]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { // Perform the greedy search using the provided resources. @@ -591,7 +611,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - internal_search_builder(), + internal_search_builder(pred), prefetch_parameters, cancel ); @@ -623,8 +643,12 @@ class MutableVamanaIndex { QueryResultView results, const Queries& queries, const search_parameters_type& sp, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), + const std::function* pred = nullptr ) { + // Optional `pred`: only entries whose external id passes are returned. It is called + // from all search threads. Rows with fewer passing entries than requested are + // padded with id `size_t(-1)` and the worst score. threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, @@ -637,23 +661,33 @@ class MutableVamanaIndex { : sp.buffer_config_, distance::comparator(distance_), sp.search_buffer_visited_set_}; + // Filtered: visit entries from a min-heap during the search. + buffer.enable_heap(pred != nullptr); auto prefetch_parameters = GreedySearchPrefetchParameters{ sp.prefetch_lookahead_, sp.prefetch_step_}; auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); - extensions::per_thread_batch_search( - data_, - buffer, - scratch, - queries, - results, - threads::UnitRange{is}, - greedy_search_closure(prefetch_parameters, cancel), - *this, - cancel - ); + auto run = [&](const auto& search) { + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + search, + *this, + cancel + ); + }; + auto search = greedy_search_closure(prefetch_parameters, cancel, pred); + if (pred == nullptr) { + run(search); + } else { + run(extensions::FilteredSearch{search, pred}); + } } ); diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 9e8f304e3..23c31a6db 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -95,6 +95,14 @@ template > class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; + // Optional (`enable_heap`): while searching, entries to visit are taken from this + // min-heap (best on top), and `candidates_` keeps only the best valid entries, sorted. + // `seen_`: ids already queued in this search (exact duplicate check). + bool use_heap_ = false; + bool heap_done_ = false; + vector_type heap_{}; + std::unordered_set seen_{}; + value_type popped_{}; public: MutableBuffer() = default; @@ -197,6 +205,9 @@ template > class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); + heap_.clear(); + seen_.clear(); + heap_done_ = false; best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -223,6 +234,9 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } + /// @brief Visit entries from a min-heap during search (search window = capacity). + void enable_heap(bool enable = true) { use_heap_ = enable; } + /// @brief Return the number of valid elements currently contained in the buffer. size_t valid() const { return valid_; } @@ -267,6 +281,12 @@ template > class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { + if (heap_mode()) { + // Nothing left to visit, or the best one is worse than the worst result. + return heap_.empty() || + (full() && + compare_(candidates_.back().distance(), heap_.front().distance())); + } // Until we've reached the target number of valid elements, we have to ignore // the state of `roi_end_`. return best_unvisited_ == ((valid() < target_valid_) ? size() : roi_end_); @@ -282,6 +302,13 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { + if (heap_mode()) { + std::pop_heap(heap_.begin(), heap_.end(), heap_compare()); + popped_ = heap_.back(); + heap_.pop_back(); + popped_.set_visited(); + return popped_; + } // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); node.set_visited(); @@ -302,6 +329,10 @@ template > class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { + if (heap_mode()) { + heap_insert(neighbor); + return; + } // Follow the contract of not appending more valid elements than the capacity // allows. bool valid = neighbor.valid(); @@ -453,6 +484,9 @@ template > class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { + if (heap_mode()) { + return heap_insert(neighbor); + } if (can_skip(neighbor.distance())) { return size(); } @@ -505,6 +539,9 @@ template > class MutableBuffer { // Sort all stored elements in the buffer. void sort() { + if (heap_mode()) { + return; // `candidates_` is kept sorted on every insert. + } assert(valid() <= valid_capacity_); // Put the elements in ascending order. @@ -554,6 +591,13 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { + if (heap_mode()) { + // `candidates_` already holds only valid entries; continue as a sorted list + // (e.g. for the top-up). + heap_.clear(); + seen_.clear(); + heap_done_ = true; + } auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); candidates_.resize(new_end - begin()); @@ -619,6 +663,43 @@ template > class MutableBuffer { } private: + bool heap_mode() const { return use_heap_ && !heap_done_; } + + // Min-heap order: the best entry on top. + auto heap_compare() const { + return [this](const value_type& a, const value_type& b) { + return compare_(b.distance(), a.distance()); + }; + } + + size_t heap_insert(value_type neighbor) { + // Skip entries worse than the worst result when full, and ids already queued. + if (can_skip(neighbor.distance()) || !seen_.insert(neighbor.id()).second) { + return size(); + } + heap_.push_back(neighbor); + std::push_heap(heap_.begin(), heap_.end(), heap_compare()); + if (!neighbor.valid()) { + return size(); + } + // Valid: keep it among the best `valid_capacity_` results. + auto pos = std::upper_bound( + begin(), + end(), + neighbor.distance(), + [&](const float& d, const value_type& other) { + return compare_(d, other.distance()); + } + ); + size_t i = pos - begin(); + candidates_.insert(pos, neighbor); + if (candidates_.size() > valid_capacity_) { + candidates_.pop_back(); + } + valid_ = candidates_.size(); + return i; + } + /// Return how many more valid candidates exist than required. /// If the number of valid candidates is *less* than the target, a negative number /// is returned. diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 96ffeaa06..883286b27 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -27,6 +27,9 @@ #include "svs/lib/preprocessor.h" #include "svs/lib/threads.h" +#include +#include + namespace svs::index::vamana::extensions { ///// @@ -432,12 +435,18 @@ struct VamanaSingleSearchType { /// with additional results. template void check_and_supplement_search_buffer( - const Index& index, SearchBuffer& search_buffer, const Query& query + const Index& index, + SearchBuffer& search_buffer, + const Query& query, + const std::function* pred = nullptr ) { if (search_buffer.valid() < search_buffer.target_window() && search_buffer.valid() < index.size()) { search_buffer.sort(); for (auto external_id : index.external_ids()) { + if (pred != nullptr && !(*pred)(external_id)) { + continue; + } auto internal_id = index.translate_external_id(external_id); auto dist = index.get_distance(external_id, query); auto builder = index.internal_search_builder(); @@ -449,6 +458,15 @@ void check_and_supplement_search_buffer( } } +/// A search closure that also carries the filter of a filtered search. +template struct FilteredSearch { + Search search; + const std::function* pred; + template void operator()(Args&&... args) const { + search(std::forward(args)...); + } +}; + /// Customization point object for processing single queries. inline constexpr VamanaSingleSearchType single_search{}; @@ -481,7 +499,11 @@ SVS_FORCE_INLINE void svs_invoke( search(query, accessor, distance, search_buffer); if constexpr (Index::needs_id_translation) { - check_and_supplement_search_buffer(index, search_buffer, query); + const std::function* pred = nullptr; + if constexpr (requires { search.pred; }) { + pred = search.pred; // Filtered search: keep filtered-out ids out of the top-up. + } + check_and_supplement_search_buffer(index, search_buffer, query, pred); } } @@ -586,6 +608,16 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { + if constexpr (requires { search.pred; }) { + if (j >= search_buffer.size()) { // Filtered: fewer than k passed. + result.index(i, j) = static_cast(-1); + result.distance(i, j) = + distance::comparator(distance)(0.0f, 1.0f) + ? std::numeric_limits::infinity() + : -std::numeric_limits::infinity(); + continue; + } + } result.set(search_buffer[j], i, j); } } diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 067c7c486..9619bdc7d 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -56,8 +56,21 @@ class DynamicVamanaInterface : public VamanaInterface { // Check if the index has been consolidated - no `Deleted` entries remain. virtual bool is_consolidated() const = 0; + + // Filtered search; false if the data type does not support it. Keep last (ABI). + virtual bool filtered_search( + svs::QueryResultView results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter + ) = 0; }; +// Filtered search is only enabled for uncompressed data. +template inline constexpr bool is_simple_data = false; +template +inline constexpr bool is_simple_data> = true; + template class DynamicVamanaImpl : public VamanaImpl { public: @@ -109,6 +122,46 @@ class DynamicVamanaImpl : public VamanaImpl results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter + ) override { + // Indexes whose `search` has no filter argument (e.g. the concurrent index) report + // "unsupported" like other data types. + constexpr bool supported = + is_simple_data && + requires( + Impl & idx, + svs::QueryResultView r, + data::ConstSimpleDataView q, + const index::vamana::VamanaSearchParameters& sp, + const std::function* pred + ) { idx.search(r, q, sp, lib::Returns(lib::Const()), pred); }; + if constexpr (supported) { + lib::match( + QueryTypes{}, + queries.type(), + [&](lib::Type SVS_UNUSED(type)) { + impl().search( + results, + data::ConstSimpleDataView(queries), + search_parameters, + lib::Returns(lib::Const()), + &filter + ); + }, + [&](svs::DataType data_type) { + throw ANNEXCEPTION("Unsupported query datatype! Got: {}.", data_type); + } + ); + return true; + } else { + return false; + } + } }; // Forward Declaractions. @@ -609,6 +662,20 @@ class DynamicVamana : public manager::IndexManager { /// @brief Check if the index has been consolidated - no `Deleted` entries remain. bool is_consolidated() const { return impl_->is_consolidated(); } + + /// @brief Filtered search: only ids for which ``filter(id)`` is true are returned. + /// Returns false if the data type does not support it. + template + bool filtered_search( + QueryResultView results, + data::ConstSimpleDataView queries, + const VamanaSearchParameters& search_parameters, + const std::function& filter + ) { + return impl_->filtered_search( + results, AnonymousArray<2>(queries), search_parameters, filter + ); + } }; /// diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 94daddcf1..44ed215d7 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -23,11 +23,14 @@ #include #include #include +#include #include #include +#include #include #include #include +#include // svs #include "svs/core/recall.h" @@ -464,3 +467,222 @@ CATCH_TEST_CASE("MutableVamana Index Memory Usage", "[graph_index][dynamic_index const size_t usage = index.get_memory_breakdown().total(); CATCH_REQUIRE(usage == expected_total_bytes); } + +CATCH_TEST_CASE( + "MutableVamana Index Filtered Search", "[graph_index][dynamic_index][filtered]" +) { + const size_t num_threads = 2; + const size_t n = 3000; + const size_t dim = 32; + const size_t num_queries = 50; + const size_t k = 10; + using Distance = svs::distance::DistanceL2; + + std::mt19937 rng(42); + std::normal_distribution normal; + auto data = svs::data::BlockedData(n, dim); + std::vector> raw(n, std::vector(dim)); + for (size_t i = 0; i < n; ++i) { + for (auto& x : raw[i]) { + x = normal(rng); + } + data.set_datum(i, raw[i]); + } + auto queries = svs::data::SimpleData(num_queries, dim); + for (size_t i = 0; i < num_queries; ++i) { + std::vector q(dim); + for (auto& x : q) { + x = normal(rng); + } + queries.set_datum(i, q); + } + + // External ids differ from internal ids. + std::vector ids(n); + std::iota(ids.begin(), ids.end(), 7); + svs::index::vamana::VamanaBuildParameters parameters{1.2, 32, 64, 64, 32, true}; + auto index = svs::index::vamana::MutableVamanaIndex( + parameters, std::move(data), ids, Distance(), num_threads + ); + + auto sp = svs::index::vamana::VamanaSearchParameters{}; + sp.buffer_config_ = svs::index::vamana::SearchBufferConfig{50}; + + // Runs the filtered search; `found` = number of non-padding ids per query. + auto search_filtered = [&](auto& results, + const std::function& pred, + std::vector& found) { + index.search( + results.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + &pred + ); + found.assign(num_queries, 0); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + found[q] += results.index(q, j) != static_cast(-1); + } + } + }; + + // Brute force over live ids passing `pred`. + std::unordered_set deleted; + auto brute_force = [&](size_t q, auto&& pred) { + std::vector> all; + for (size_t i = 0; i < n; ++i) { + size_t e = ids[i]; + if (deleted.contains(e) || !pred(e)) { + continue; + } + float d = 0; + for (size_t j = 0; j < dim; ++j) { + float diff = queries.get_datum(q)[j] - raw[i][j]; + d += diff * diff; + } + all.emplace_back(d, e); + } + std::sort(all.begin(), all.end()); + std::vector out; + for (size_t i = 0; i < std::min(k, all.size()); ++i) { + out.push_back(all[i].second); + } + return out; + }; + + auto check = [&](auto&& pred, bool expect_full) { + auto results = svs::QueryResult(num_queries, k); + std::vector found; + search_filtered(results, pred, found); + CATCH_REQUIRE(found.size() == num_queries); + size_t hits = 0, total = 0; + for (size_t q = 0; q < num_queries; ++q) { + auto gt = brute_force(q, pred); + CATCH_REQUIRE(found[q] <= k); + if (expect_full) { + CATCH_REQUIRE(found[q] == k); + } + for (size_t j = 0; j < found[q]; ++j) { + size_t id = results.index(q, j); + CATCH_REQUIRE(pred(id)); + CATCH_REQUIRE(!deleted.contains(id)); + hits += std::count(gt.begin(), gt.end(), id); + } + for (size_t j = found[q]; j < k; ++j) { + CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); + CATCH_REQUIRE( + results.distance(q, j) == std::numeric_limits::infinity() + ); + } + total += gt.size(); + } + return double(hits) / double(total); + }; + + auto mod10 = [](size_t e) { return e % 10 == 0; }; + + CATCH_SECTION("recall against brute force") { + CATCH_REQUIRE(check(mod10, true) >= 0.95); + } + + CATCH_SECTION("deleted ids are never returned") { + std::vector to_delete; + for (size_t i = 0; i < n; i += 3) { + to_delete.push_back(ids[i]); + } + index.delete_entries(to_delete); + deleted.insert(to_delete.begin(), to_delete.end()); + CATCH_REQUIRE(check(mod10, true) >= 0.9); + CATCH_REQUIRE(check([](size_t) { return true; }, true) >= 0.9); + } + + CATCH_SECTION("fewer than k passing ids are padded") { + auto few = [](size_t e) { return e == 100 || e == 200 || e == 300; }; + check(few, false); + auto results = svs::QueryResult(num_queries, k); + std::vector found; + search_filtered(results, few, found); + for (size_t q = 0; q < num_queries; ++q) { + CATCH_REQUIRE(found[q] == 3); + } + auto none = [](size_t) { return false; }; + search_filtered(results, none, found); + for (size_t q = 0; q < num_queries; ++q) { + CATCH_REQUIRE(found[q] == 0); + } + } + + CATCH_SECTION("unfiltered search is unchanged by the optional argument") { + index.delete_entries(std::vector{ids[0], ids[1], ids[2]}); + auto plain = svs::QueryResult(num_queries, k); + auto with_null = svs::QueryResult(num_queries, k); + index.search(plain.view(), queries.cview(), sp); + index.search( + with_null.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + nullptr + ); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE(plain.index(q, j) == with_null.index(q, j)); + CATCH_REQUIRE(plain.distance(q, j) == with_null.distance(q, j)); + } + } + } + + CATCH_SECTION("inner product pads with the worst score (-inf)") { + auto data_ip = svs::data::BlockedData(n, dim); + for (size_t i = 0; i < n; ++i) { + data_ip.set_datum(i, raw[i]); + } + svs::index::vamana::VamanaBuildParameters parameters_ip{0.95, 32, 64, 64, 32, true}; + auto index_ip = svs::index::vamana::MutableVamanaIndex( + parameters_ip, std::move(data_ip), ids, svs::distance::DistanceIP(), num_threads + ); + std::function few = [](size_t e) { return e == 100 || e == 200; }; + auto results = svs::QueryResult(num_queries, k); + index_ip.search( + results.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + &few + ); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 2; j < k; ++j) { + CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); + CATCH_REQUIRE( + results.distance(q, j) == -std::numeric_limits::infinity() + ); + } + } + } + + CATCH_SECTION("orchestrator entry point") { + auto dyn = svs::DynamicVamana::build( + parameters, + [&] { + auto d = svs::data::SimpleData(n, dim); + for (size_t i = 0; i < n; ++i) { + d.set_datum(i, raw[i]); + } + return d; + }(), + ids, + Distance(), + num_threads + ); + auto results = svs::QueryResult(num_queries, k); + CATCH_REQUIRE(dyn.filtered_search( + results.view(), queries.cview(), sp, std::function(mod10) + )); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE(mod10(results.index(q, j))); + } + } + } +} diff --git a/tests/svs/index/vamana/search_buffer.cpp b/tests/svs/index/vamana/search_buffer.cpp index 8ecb698a5..cdf860975 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -906,6 +906,45 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") { CATCH_REQUIRE(b.done()); } } + + // With `enable_heap`, entries are visited from a min-heap and `candidates_` keeps only + // the best valid entries, sorted. + CATCH_SECTION("Heap") { + buffer.enable_heap(); + buffer.insert({10, 5, false}); + buffer.insert({0, 100, true}); + buffer.insert({11, 1, false}); + buffer.insert({1, 50, true}); + buffer.insert({11, 1, false}); // duplicate: ignored + CATCH_REQUIRE(buffer.size() == 2); + CATCH_REQUIRE(buffer.valid() == 2); + + CATCH_REQUIRE(buffer.next().id() == 11); + CATCH_REQUIRE(buffer.next().id() == 10); + buffer.insert({2, 40, true}); + buffer.insert({3, 30, true}); + CATCH_REQUIRE(buffer.full()); + // Full: worse than the worst valid entry (100) is skipped. + buffer.insert({4, 200, true}); + buffer.insert({12, 150, false}); + // A better valid entry pushes out the worst valid one (id 0). + buffer.insert({5, 20, true}); + CATCH_REQUIRE(buffer.size() == 4); + + CATCH_REQUIRE(buffer.next().id() == 5); + CATCH_REQUIRE(buffer.next().id() == 3); + CATCH_REQUIRE(buffer.next().id() == 2); + CATCH_REQUIRE(buffer.next().id() == 1); + // Only id 0 (100) is left in the heap; it is worse than the worst result (50). + CATCH_REQUIRE(buffer.done()); + + buffer.cleanup(); + CATCH_REQUIRE(buffer.size() == 4); + CATCH_REQUIRE(eq(buffer[0], {5, 20, true})); + CATCH_REQUIRE(eq(buffer[1], {3, 30, true})); + CATCH_REQUIRE(eq(buffer[2], {2, 40, true})); + CATCH_REQUIRE(eq(buffer[3], {1, 50, true})); + } } CATCH_TEST_CASE("Fuzzing Mutable", "[core][search_buffer]") {