Filtered top-k search on the dynamic Vamana index via the search builder

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<SlotMetadata>& 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<SlotMetadata>& status,
+        const IDTranslator* translator = nullptr,
+        const std::function<bool(size_t)>* pred = nullptr
+    )
+        : status_{status}
+        , translator_{translator}
+        , pred_{pred} {}
 
     template <typename I>
     constexpr PredicatedSearchNeighbor<I> 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>(i, distance, !invalid);
     }
 
   private:
     const std::vector<SlotMetadata>& status_;
+    const IDTranslator* translator_;
+    const std::function<bool(size_t)>* pred_;
 };
 
 template <graphs::MemoryGraph Graph, typename Data, typename Dist>
@@ -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<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);
                     }
@@ -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<bool(size_t)>* pred = nullptr
+    ) const {
+        return ValidBuilder{status_, &translator_, pred};
+    }
 
     auto greedy_search_closure(
         GreedySearchPrefetchParameters prefetch_parameters,
-        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
     ) 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<Idx>{lib::as_const_span(entry_point_)},
-                internal_search_builder(),
+                internal_search_builder(pred),
                 prefetch_parameters,
                 cancel
             );
@@ -623,8 +643,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()},
@@ -643,17 +667,25 @@ class MutableVamanaIndex {
 
                 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<decltype(search)>{search, pred});
+                }
             }
         );
 

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,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<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);
         }
     }

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<size_t> results,
+        AnonymousArray<2> queries,
+        const index::vamana::VamanaSearchParameters& search_parameters,
+        const std::function<bool(size_t)>& filter
+    ) = 0;
 };
 
+// Filtered search is only enabled for uncompressed data.
+template <typename T> inline constexpr bool is_simple_data = false;
+template <typename T, size_t Extent, typename Alloc>
+inline constexpr bool is_simple_data<data::SimpleData<T, Extent, Alloc>> = true;
+
 template <lib::TypeList QueryTypes, typename Impl>
 class DynamicVamanaImpl : public VamanaImpl<QueryTypes, Impl, DynamicVamanaInterface> {
   public:
@@ -109,6 +122,46 @@ class DynamicVamanaImpl : public VamanaImpl<QueryTypes, Impl, DynamicVamanaInter
 
     // Check if the index has been consolidated - no `Deleted` entries remain.
     bool is_consolidated() const override { return impl().is_consolidated(); }
+
+    bool filtered_search(
+        svs::QueryResultView<size_t> results,
+        AnonymousArray<2> queries,
+        const index::vamana::VamanaSearchParameters& search_parameters,
+        const std::function<bool(size_t)>& 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<typename Impl::data_type> &&
+            requires(
+                Impl & idx,
+                svs::QueryResultView<size_t> r,
+                data::ConstSimpleDataView<float> q,
+                const index::vamana::VamanaSearchParameters& sp,
+                const std::function<bool(size_t)>* pred
+            ) { idx.search(r, q, sp, lib::Returns(lib::Const<false>()), pred); };
+        if constexpr (supported) {
+            lib::match(
+                QueryTypes{},
+                queries.type(),
+                [&]<typename T>(lib::Type<T> SVS_UNUSED(type)) {
+                    impl().search(
+                        results,
+                        data::ConstSimpleDataView<T>(queries),
+                        search_parameters,
+                        lib::Returns(lib::Const<false>()),
+                        &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<DynamicVamanaInterface> {
 
     /// @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 <typename QueryType>
+    bool filtered_search(
+        QueryResultView<size_t> results,
+        data::ConstSimpleDataView<QueryType> queries,
+        const VamanaSearchParameters& search_parameters,
+        const std::function<bool(size_t)>& filter
+    ) {
+        return impl_->filtered_search(
+            results, AnonymousArray<2>(queries), search_parameters, filter
+        );
+    }
 };
 
 ///