Filtered search inside greedy search, with a two-heap search buffer

bindings/c/src/filtered_search.hpp

@@ -23,6 +23,7 @@
 #include <svs/core/query_result.h>
 #include <svs/lib/misc.h>
 #include <svs/lib/threads.h>
+#include <svs/orchestrators/dynamic_vamana.h>
 
 #include <algorithm>
 #include <cassert>
@@ -31,6 +32,7 @@
 #include <functional>
 #include <limits>
 #include <stdexcept>
+#include <type_traits>
 #include <vector>
 
 namespace svs::c_runtime {
@@ -168,6 +170,29 @@ std::vector<size_t> 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<IndexType, svs::DynamicVamana>) {
+        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<size_t> 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<size_t>(-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();

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<size_t>(-1)) {
+                            continue; // Padding from a filtered search.
+                        }
                         auto internal = lib::narrow_cast<Idx>(ids.at(i, j));
                         ids.at(i, j) = translate_internal_id(internal);
                     }
@@ -623,8 +626,12 @@ class MutableVamanaIndex {
         QueryResultView<I> results,
         const Queries& queries,
         const search_parameters_type& sp,
-        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>()),
+        const std::function<bool(size_t)>* 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 +644,37 @@ class MutableVamanaIndex {
                         : sp.buffer_config_,
                     distance::comparator(distance_),
                     sp.search_buffer_visited_set_};
+                // Filtered: the buffer checks the filter itself (two heaps).
+                std::function<bool(Idx)> passes;
+                if (pred != nullptr) {
+                    passes = [&](Idx i) { return (*pred)(translator_.get_external(i)); };
+                    buffer.enable_heap(&passes);
+                }
 
                 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);
+                if (pred == nullptr) {
+                    run(search);
+                } else {
+                    run(extensions::FilteredSearch<decltype(search)>{search, pred});
+                }
             }
         );
 

include/svs/index/vamana/dynamic_search_buffer.h

@@ -95,6 +95,18 @@ template <typename Idx, typename Cmp = std::less<>> class MutableBuffer {
     vector_type candidates_{};
     // An optional visited filter.
     std::optional<set_type> 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).
+    // `passes_`: the filter, called with an internal id.
+    bool use_heap_ = false;
+    const std::function<bool(Idx)>* passes_ = nullptr;
+    bool heap_done_ = false;
+    vector_type to_visit_{};
+    std::unordered_set<Idx> seen_{};
+    value_type popped_{};
 
   public:
     MutableBuffer() = default;
@@ -197,6 +209,9 @@ template <typename Idx, typename Cmp = std::less<>> 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 +238,13 @@ template <typename Idx, typename Cmp = std::less<>> 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 `*passes` accepts its internal id.
+    void enable_heap(const std::function<bool(Idx)>* passes) {
+        use_heap_ = passes != nullptr;
+        passes_ = passes;
+    }
+
     /// @brief Return the number of valid elements currently contained in the buffer.
     size_t valid() const { return valid_; }
 
@@ -267,6 +289,12 @@ template <typename Idx, typename Cmp = std::less<>> 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 +310,13 @@ template <typename Idx, typename Cmp = std::less<>> 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 +337,10 @@ template <typename Idx, typename Cmp = std::less<>> 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 +492,9 @@ template <typename Idx, typename Cmp = std::less<>> 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 +547,9 @@ template <typename Idx, typename Cmp = std::less<>> 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 +599,14 @@ template <typename Idx, typename Cmp = std::less<>> 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 +672,45 @@ template <typename Idx, typename Cmp = std::less<>> class MutableBuffer {
     }
 
   private:
+    bool heap_mode() const { 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() || !(*passes_)(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.

include/svs/index/vamana/extensions.h

@@ -27,6 +27,9 @@
 #include "svs/lib/preprocessor.h"
 #include "svs/lib/threads.h"
 
+#include <functional>
+#include <limits>
+
 namespace svs::index::vamana::extensions {
 
 /////
@@ -432,12 +435,18 @@ struct VamanaSingleSearchType {
 /// with additional results.
 template <typename Index, typename SearchBuffer, typename Query>
 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<bool(size_t)>* 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 <typename Search> struct FilteredSearch {
+    Search search;
+    const std::function<bool(size_t)>* pred;
+    template <typename... Args> void operator()(Args&&... args) const {
+        search(std::forward<Args>(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<bool(size_t)>* 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,15 @@ 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<I>(-1);
+                    result.distance(i, j) = distance::comparator(distance)(0.0f, 1.0f)
+                                                ? std::numeric_limits<float>::infinity()
+                                                : -std::numeric_limits<float>::infinity();
+                    continue;
+                }
+            }
             result.set(search_buffer[j], i, j);
         }
     }
@@ -607,6 +638,23 @@ template <typename Data> struct UsesReranking {
 
 template <typename Data> inline constexpr UsesReranking<Data> calibration_uses_reranking{};
 
+/// Whether `DynamicVamana::filtered_search` is supported for `Data`.
+/// Data types opt in by providing an `svs_invoke` overload.
+template <typename Data> struct SupportsFilteredSearch {
+    constexpr bool operator()() const {
+        if constexpr (svs::svs_invocable<SupportsFilteredSearch>) {
+            return svs::svs_invoke(*this);
+        } else {
+            return false;
+        }
+    }
+};
+
+template <typename T, size_t Extent, typename Alloc>
+constexpr bool svs_invoke(SupportsFilteredSearch<data::SimpleData<T, Extent, Alloc>>) {
+    return true;
+}
+
 /////
 ///// Reconstruct Vector
 /////