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

include/svs/concepts/distance.h

@@ -25,6 +25,7 @@
 
 #include <concepts>
 #include <cstddef>
+#include <limits>
 #include <utility>
 #include <vector>
 
@@ -215,6 +216,12 @@ template <typename T> auto comparator(const T& /*x*/) -> compare_t<T> {
     return typename T::compare{};
 }
 
+/// Distance reported for a missing neighbor: it sorts after every real distance.
+template <typename T> float worst_distance(const T& x) {
+    return comparator(x)(0.0f, 1.0f) ? std::numeric_limits<float>::infinity()
+                                     : -std::numeric_limits<float>::infinity();
+}
+
 /////
 ///// Traits
 /////

include/svs/concurrent/dynamic_index.h

@@ -171,6 +171,59 @@ class ValidBuilder {
     const lib::SegmentedVector<SlotMetadata>& status_;
 };
 
+// The index whose user filter runs on this thread, so re-entering it fails fast.
+inline thread_local const void* filtering_index = nullptr;
+
+struct FilterScope {
+    explicit FilterScope(const void* index)
+        : previous_{filtering_index} {
+        filtering_index = index;
+    }
+    ~FilterScope() { filtering_index = previous_; }
+    const void* previous_;
+};
+
+inline void check_not_in_filter(const void* index) {
+    if (filtering_index == index) {
+        throw ANNEXCEPTION("filter must not call into the same index");
+    }
+}
+
+// Separate template so the regular search compiles without any filter check.
+template <typename Filter> class FilteredValidBuilder {
+  public:
+    FilteredValidBuilder(
+        const lib::SegmentedVector<SlotMetadata>& status,
+        const IDTranslator& translator,
+        const Filter& filter,
+        const void* owner
+    )
+        : status_{status}
+        , translator_{&translator}
+        , filter_{&filter}
+        , owner_{owner} {}
+
+    template <typename I>
+    PredicatedSearchNeighbor<I> operator()(I i, float distance) const {
+        bool invalid =
+            std::atomic_ref<SlotMetadata>(const_cast<SlotMetadata&>(getindex(status_, i)))
+                .load(std::memory_order_acquire) != SlotMetadata::Valid;
+        if (!invalid) {
+            // Lock-free read; an erased slot has no external id and fails.
+            size_t external = translator_->get_external_or(i, no_external_id);
+            FilterScope scope{owner_};
+            invalid = external == no_external_id || !(*filter_)(external);
+        }
+        return PredicatedSearchNeighbor<I>(i, distance, !invalid);
+    }
+
+  private:
+    const lib::SegmentedVector<SlotMetadata>& status_;
+    const IDTranslator* translator_;
+    const Filter* filter_;
+    const void* owner_;
+};
+
 template <graphs::MemoryGraph Graph, typename Data, typename Dist>
 class MutableVamanaIndex {
     friend class MultiMutableVamanaIndex<Graph, Data, Dist>;
@@ -606,11 +659,7 @@ class MutableVamanaIndex {
         size_t metadata_bytes = status_.capacity() * sizeof(SlotMetadata);
         metadata_bytes +=
             entry_point_.capacity() * sizeof(typename entry_point_type::value_type);
-        // The IDTranslator holds two tsl::robin_map instances (external->internal and
-        // internal->external), neither of which exposes its allocated byte count. We
-        // approximate the storage as the id pair held in each of the two directions. This
-        // ignores the maps' load-factor slack and control bytes, so it is an estimate of
-        // the hash-map overhead that is accurate to within a few percent.
+        // Estimate: one id pair per entry for each translation direction.
         metadata_bytes += 2 * translator_.size() *
                           (sizeof(IDTranslator::external_id_type) +
                            sizeof(IDTranslator::internal_id_type));
@@ -810,6 +859,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) = unsafe_translate_internal_id(internal);
                     }
@@ -882,13 +934,19 @@ class MutableVamanaIndex {
     // This is an internal method, mostly used to help implement the batch iterator.
     ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; }
 
+    template <typename Filter>
+    FilteredValidBuilder<Filter> internal_search_builder(const Filter& filter) const {
+        return FilteredValidBuilder<Filter>{status_, translator_, filter, this};
+    }
+
     /// @brief RAII shared lock guarding data_/graph_ against compact()'s shrink
     /// (which frees segments). Used by searches, iterators and mutations that
     /// may overlap each other but must exclude compaction.
     ///
     /// Hold this through internal-to-external ID translation. Acquire the translation
     /// lock inside this guard, preserving the compact -> translator lock order.
     [[nodiscard]] std::shared_lock<std::shared_mutex> lock_storage_shared() const {
+        check_not_in_filter(this);
         std::shared_lock gatekeeper_lock{*compact_gatekeeper_};
         return std::shared_lock<std::shared_mutex>(*compact_mutex_);
     }
@@ -900,18 +958,28 @@ class MutableVamanaIndex {
         return std::shared_lock<std::shared_mutex>(*translator_mutex_);
     }
 
+    template <bool Filtered = false>
     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)>* filter = nullptr
     ) const {
-        return [&, prefetch_parameters](
-                   const auto& query, auto& accessor, auto& distance, auto& buffer
-               ) {
+        return [&,
+                prefetch_parameters,
+                filter](const auto& query, auto& accessor, auto& distance, auto& buffer) {
+            static_cast<void>(filter); // Unused in regular search.
             const Idx seed = entry_point();
             if (seed == NO_ENTRY) {
                 buffer.clear();
                 return;
             }
+            auto builder = [&] {
+                if constexpr (Filtered) {
+                    return internal_search_builder(*filter);
+                } else {
+                    return internal_search_builder();
+                }
+            }();
             // Perform the greedy search using the provided resources.
             concurrent::greedy_search(
                 graph_,
@@ -921,7 +989,7 @@ class MutableVamanaIndex {
                 distance,
                 buffer,
                 vamana::EntryPointInitializer<Idx>{std::span<const Idx>{&seed, 1}},
-                internal_search_builder(),
+                builder,
                 prefetch_parameters,
                 cancel
             );
@@ -957,7 +1025,8 @@ 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)>* filter = nullptr
     ) {
         // Internal IDs must retain their meaning until translation has completed.
         auto compact_lock = lock_storage_shared();
@@ -982,6 +1051,25 @@ class MutableVamanaIndex {
                     auto scratch =
                         extensions::per_thread_batch_search_setup(data_, distance_);
 
+                    if (filter != nullptr) {
+                        [&]() __attribute__((noinline)) {
+                            extensions::per_thread_batch_search(
+                                data_,
+                                buffer,
+                                scratch,
+                                queries,
+                                results,
+                                threads::UnitRange{is},
+                                greedy_search_closure<true>(
+                                    prefetch_parameters, cancel, filter
+                                ),
+                                *this,
+                                cancel
+                            );
+                        }
+                        ();
+                        return;
+                    }
                     extensions::per_thread_batch_search(
                         data_,
                         buffer,
@@ -1005,6 +1093,26 @@ class MutableVamanaIndex {
         // After the search procedure, the indices in `results` are internal.
         // Perform one more pass to convert these to external ids.
         translate_to_external(results.indices());
+
+        if (filter != nullptr) {
+            // A concurrent remap may have changed the answer: re-check the final ids.
+            const float worst = distance::worst_distance(distance_);
+            for (size_t i = 0; i < results.n_queries(); ++i) {
+                for (size_t j = 0; j < results.n_neighbors(); ++j) {
+                    auto id = results.index(i, j);
+                    if (id != static_cast<I>(-1)) {
+                        FilterScope scope{this};
+                        if (!(*filter)(id)) {
+                            results.index(i, j) = static_cast<I>(-1);
+                        }
+                    }
+                    // Also covers ids a concurrent erase left without a translation.
+                    if (results.index(i, j) == static_cast<I>(-1)) {
+                        results.distance(i, j) = worst;
+                    }
+                }
+            }
+        }
     }
 
     ///
@@ -1400,6 +1508,7 @@ class MutableVamanaIndex {
     ///     improve performance but requires more working memory.
     ///
     void compact(Idx batch_size = 1'000) {
+        check_not_in_filter(this);
         std::lock_guard gatekeeper_lock{*compact_gatekeeper_};
         std::lock_guard compact_lock{*compact_mutex_};
 

include/svs/concurrent/greedy_search.h

@@ -146,7 +146,9 @@ void greedy_search(
                 // Compute distance and update search buffer.
                 auto dist =
                     distance::compute(distance_function, query, accessor(dataset, id));
-                search_buffer.insert(builder(id, dist));
+                if (!search_buffer.can_skip(dist)) {
+                    search_buffer.insert(builder(id, dist));
+                }
             }
 
             // Validate that no concurrent write occurred during the read.