diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index 8e45ac691..76956ea25 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,29 @@ std::vector filtered_topk_search( assert(results.n_queries() == queries.size()); + // The dynamic index filters inside the search itself when its data type supports it. + // The search window is max(k, initial_batch_hint), the same as the C API's unfiltered + // search (hint = max(k, window)); other parameters come from the index. + if constexpr (std::is_same_v) { + auto params = index.get_search_parameters(); + params.buffer_config_ = svs::index::vamana::SearchBufferConfig{ + std::max(results.n_neighbors(), initial_batch_hint)}; + if (index.filtered_search( + results.view(), + queries, + params, + [id_filter](size_t id) { return id_filter->is_member(id); } + )) { + std::vector found(queries.size(), 0); + for (size_t i = 0; i < queries.size(); ++i) { + for (size_t j = 0; j < results.n_neighbors(); ++j) { + found[i] += (results.index(i, j) != static_cast(-1)); + } + } + 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/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9336f82e1..534154e26 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -532,6 +532,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); } @@ -623,37 +626,66 @@ 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* filter = nullptr ) { + // Optional `filter`: 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()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { - size_t num_neighbors = results.n_neighbors(); - auto buffer = search_buffer_type{ - // Legalize search buffer for this search. - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - - 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 body = [&](std::bool_constant) { + size_t num_neighbors = results.n_neighbors(); + auto buffer = MutableBuffer, Filtered>{ + // Legalize search buffer for this search. + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + // Filtered: the buffer checks the filter itself (two heaps). + std::function internal_filter; + if constexpr (Filtered) { + internal_filter = [&](Idx i) { + return (*filter)(translator_.get_external(i)); + }; + buffer.enable_heap(&internal_filter); + } + + auto prefetch_parameters = GreedySearchPrefetchParameters{ + sp.prefetch_lookahead_, sp.prefetch_step_}; + + auto scratch = + extensions::per_thread_batch_search_setup(data_, distance_); + + 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); + if constexpr (Filtered) { + run(extensions::FilteredSearch{search, filter}); + } else { + run(search); + } + }; + if (filter != nullptr) { + [&]() __attribute__((noinline)) { body(std::true_type{}); } + (); + } else { + body(std::false_type{}); + } } ); diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 9e8f304e3..769a68e00 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -38,7 +38,8 @@ namespace svs::index::vamana { /// The search buffer will still navigate through those entries, but it won't include /// them in its final result. /// -template > class MutableBuffer { +template , bool Filtered = false> +class MutableBuffer { public: // Type Aliases using index_type = Idx; @@ -95,6 +96,18 @@ template > class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; + // Optional (`enable_heap`): two heaps while searching. `to_visit_` is a min-heap of + // entries still to visit (best on top). During heap mode `candidates_` holds the + // results: a max-heap of the best valid entries that pass the filter (worst on top). + // `cleanup()` sorts `candidates_` once at the end. + // `seen_`: ids already queued in this search (exact duplicate check). + // `internal_filter_`: the filter, called with an internal id. + bool use_heap_ = false; + const std::function* internal_filter_ = nullptr; + bool heap_done_ = false; + vector_type to_visit_{}; + std::unordered_set seen_{}; + value_type popped_{}; public: MutableBuffer() = default; @@ -197,6 +210,9 @@ template > class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); + to_visit_.clear(); + seen_.clear(); + heap_done_ = false; best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -223,6 +239,13 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } + /// @brief Use two heaps during search (search window = capacity). A neighbor enters the + /// results only if `*internal_filter` accepts its internal id. + void enable_heap(const std::function* internal_filter) { + use_heap_ = internal_filter != nullptr; + internal_filter_ = internal_filter; + } + /// @brief Return the number of valid elements currently contained in the buffer. size_t valid() const { return valid_; } @@ -267,6 +290,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 to_visit_.empty() || + (full() && + compare_(candidates_.front().distance(), to_visit_.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 +311,13 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { + if (heap_mode()) { + std::pop_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + popped_ = to_visit_.back(); + to_visit_.pop_back(); + popped_.set_visited(); + return popped_; + } // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); node.set_visited(); @@ -302,6 +338,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 +493,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 +548,9 @@ template > class MutableBuffer { // Sort all stored elements in the buffer. void sort() { + if (heap_mode()) { + return; // Both heaps are kept in order on every insert. + } assert(valid() <= valid_capacity_); // Put the elements in ascending order. @@ -554,6 +600,14 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { + if (heap_mode()) { + // `candidates_` holds only valid entries: sort once, then continue as a sorted + // list (e.g. for the top-up). + std::sort(begin(), end(), compare_); + to_visit_.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 +673,52 @@ template > class MutableBuffer { } private: + bool heap_mode() const { + if constexpr (!Filtered) { + return false; + } else { + return use_heap_ && !heap_done_; + } + } + + // `to_visit_`: the best entry on top. + auto to_visit_compare() const { + return [this](const value_type& a, const value_type& b) { + return compare_(b.distance(), a.distance()); + }; + } + + // `candidates_`: the worst entry on top. + auto result_compare() const { + return [this](const value_type& a, const value_type& b) { + return compare_(a.distance(), b.distance()); + }; + } + + size_t heap_insert(value_type neighbor) { + // Skip entries worse than the worst result when full, and ids already queued. + if ((full() && compare_(candidates_.front().distance(), neighbor.distance())) || + !seen_.insert(neighbor.id()).second) { + return size(); + } + to_visit_.push_back(neighbor); + std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + // Only valid entries that pass the filter become results; the rest are only + // visited. + if (!neighbor.valid() || !(*internal_filter_)(neighbor.id())) { + return size(); + } + // Keep it among the best `valid_capacity_` results (drop the worst). + candidates_.push_back(neighbor); + std::push_heap(begin(), end(), result_compare()); + if (candidates_.size() > valid_capacity_) { + std::pop_heap(begin(), end(), result_compare()); + candidates_.pop_back(); + } + valid_ = candidates_.size(); + return size(); + } + /// 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. @@ -641,8 +741,9 @@ template > class MutableBuffer { } }; -template -std::ostream& operator<<(std::ostream& io, const MutableBuffer& buffer) { +template +std::ostream& +operator<<(std::ostream& io, const MutableBuffer& buffer) { return io << "MutableBuffer<" << datatype_v << ">(" << "target_valid = " << buffer.target_capacity() << ", best_unvisited = " << buffer.best_unvisited() diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 96ffeaa06..9ec594486 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* filter = 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 (filter != nullptr && !(*filter)(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* filter; + 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,12 @@ 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* filter = nullptr; + if constexpr (requires { search.filter; }) { + filter = + search.filter; // Filtered search: keep filtered-out ids out of the top-up. + } + check_and_supplement_search_buffer(index, search_buffer, query, filter); } } @@ -586,6 +609,15 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { + if constexpr (requires { search.filter; }) { + 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..1fb30f041 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -56,6 +56,14 @@ 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; }; template @@ -109,6 +117,44 @@ class DynamicVamanaImpl : public VamanaImpl results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter + ) override { + // Filtered search is available for every data type. Indexes whose `search` has no + // filter argument (e.g. the concurrent index) report "unsupported". + constexpr bool supported = requires( + Impl & idx, + svs::QueryResultView r, + data::ConstSimpleDataView q, + const index::vamana::VamanaSearchParameters& sp, + const std::function* filter + ) { idx.search(r, q, sp, lib::Returns(lib::Const()), filter); }; + 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 +655,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 + ); + } }; ///