diff --git a/include/svs/concepts/distance.h b/include/svs/concepts/distance.h index ec240b88e..7877fab5d 100644 --- a/include/svs/concepts/distance.h +++ b/include/svs/concepts/distance.h @@ -25,6 +25,7 @@ #include #include +#include #include #include @@ -215,6 +216,12 @@ template auto comparator(const T& /*x*/) -> compare_t { return typename T::compare{}; } +/// Distance reported for a missing neighbor: it sorts after every real distance. +template float worst_distance(const T& x) { + return comparator(x)(0.0f, 1.0f) ? std::numeric_limits::infinity() + : -std::numeric_limits::infinity(); +} + ///// ///// Traits ///// diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h index 1fc1a8980..7e3b4a0b7 100644 --- a/include/svs/concurrent/dynamic_index.h +++ b/include/svs/concurrent/dynamic_index.h @@ -171,6 +171,59 @@ class ValidBuilder { const lib::SegmentedVector& 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 class FilteredValidBuilder { + public: + FilteredValidBuilder( + const lib::SegmentedVector& status, + const IDTranslator& translator, + const Filter& filter, + const void* owner + ) + : status_{status} + , translator_{&translator} + , filter_{&filter} + , owner_{owner} {} + + template + PredicatedSearchNeighbor operator()(I i, float distance) const { + bool invalid = + std::atomic_ref(const_cast(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, distance, !invalid); + } + + private: + const lib::SegmentedVector& status_; + const IDTranslator* translator_; + const Filter* filter_; + const void* owner_; +}; + template class MutableVamanaIndex { friend class MultiMutableVamanaIndex; @@ -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(-1)) { + continue; // Padding from a filtered search. + } auto internal = lib::narrow_cast(ids.at(i, j)); ids.at(i, j) = unsafe_translate_internal_id(internal); } @@ -882,6 +934,11 @@ class MutableVamanaIndex { // This is an internal method, mostly used to help implement the batch iterator. ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } + template + FilteredValidBuilder internal_search_builder(const Filter& filter) const { + return FilteredValidBuilder{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. @@ -889,6 +946,7 @@ class MutableVamanaIndex { /// 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 lock_storage_shared() const { + check_not_in_filter(this); std::shared_lock gatekeeper_lock{*compact_gatekeeper_}; return std::shared_lock(*compact_mutex_); } @@ -900,18 +958,28 @@ class MutableVamanaIndex { return std::shared_lock(*translator_mutex_); } + template 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* 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(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{std::span{&seed, 1}}, - internal_search_builder(), + builder, prefetch_parameters, cancel ); @@ -957,7 +1025,8 @@ 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 ) { // 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( + 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(-1)) { + FilterScope scope{this}; + if (!(*filter)(id)) { + results.index(i, j) = static_cast(-1); + } + } + // Also covers ids a concurrent erase left without a translation. + if (results.index(i, j) == static_cast(-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_}; diff --git a/include/svs/concurrent/greedy_search.h b/include/svs/concurrent/greedy_search.h index fe8d7fd59..0b0864f9e 100644 --- a/include/svs/concurrent/greedy_search.h +++ b/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. diff --git a/include/svs/core/translation.h b/include/svs/core/translation.h index 88b95311f..1ffa034bc 100644 --- a/include/svs/core/translation.h +++ b/include/svs/core/translation.h @@ -27,6 +27,7 @@ #include "svs/lib/narrow.h" #include "svs/lib/readwrite.h" #include "svs/lib/saveload.h" +#include "svs/lib/segmented_vector.h" #include "svs/lib/threads.h" #include "svs/third-party/fmt.h" @@ -35,9 +36,12 @@ #include "tsl/robin_map.h" // stl +#include #include #include +#include #include +#include namespace svs { @@ -50,6 +54,10 @@ class IDTranslator { tsl::robin_map::const_iterator; using value_type = typename const_iterator::value_type; + /// Marks an unmapped slot in the internal-to-external table. + static constexpr external_id_type invalid_external_id = + std::numeric_limits::max(); + // Construct the identity transformation of size `n`. struct Identity { Identity() = delete; @@ -67,17 +75,7 @@ class IDTranslator { /// /// @brief Return the number of translations. /// - size_t size() const { - if constexpr (checkbounds_v) { - const size_t e2i = external_to_internal_.size(); - const size_t i2e = internal_to_external_.size(); - - if (e2i != i2e) { - throw ANNEXCEPTION("Size mismatch! E2I is {} while I2E is {}!", e2i, i2e); - } - } - return external_to_internal_.size(); - } + size_t size() const { return external_to_internal_.size(); } /// /// @brief Insert the two ranges. @@ -124,6 +122,7 @@ class IDTranslator { if (!lib::all_unique(int_begin, int_end)) { throw ANNEXCEPTION("Internal IDs contain repeat elements!"); } + check_external_not_reserved(ext_begin, ext_end); check_external_free(ext_begin, ext_end); check_internal_free(int_begin, int_end); } @@ -141,7 +140,7 @@ class IDTranslator { template void insert_translation(Ext external_id, Int internal_id) { external_to_internal_[external_id] = lib::narrow(internal_id); - internal_to_external_[internal_id] = lib::narrow(external_id); + store_external(internal_id, lib::narrow(external_id)); } /// @@ -183,6 +182,8 @@ class IDTranslator { throw ANNEXCEPTION("Internal IDs contain repeat elements!"); } + check_external_not_reserved(ext_begin, ext_end); + // Validation pass for (auto e = ext_begin; e != ext_end; ++e) { auto found = external_to_internal_.find(*e); @@ -197,7 +198,7 @@ class IDTranslator { for (; e != ext_end; ++e, ++i) { auto found = external_to_internal_.find(*e); if (found != external_to_internal_.end()) { - internal_to_external_.erase(found->second); + store_external(found->second, invalid_external_id); } insert_translation(*e, *i); } @@ -218,7 +219,7 @@ class IDTranslator { /// @param e The internal ID to check. /// bool has_internal(internal_id_type e) const { - return internal_to_external_.contains(e); + return load_external(e) != invalid_external_id; } /// @@ -236,17 +237,18 @@ class IDTranslator { /// @param i The internal ID to translate to an external ID. /// external_id_type get_external(internal_id_type i) const { - return internal_to_external_.at(i); + auto e = load_external(i); + if (e == invalid_external_id) { + throw std::out_of_range("IDTranslator: internal ID not found"); + } + return e; } /// @brief Return the external ID, or a default if not found. external_id_type get_external_or(internal_id_type i, external_id_type default_val) const { - auto it = internal_to_external_.find(i); - if (it == internal_to_external_.end()) { - return default_val; - } - return it->second; + auto e = load_external(i); + return e == invalid_external_id ? default_val : e; } /// @brief Return the internal ID, or a default if not found. @@ -280,13 +282,12 @@ class IDTranslator { assert(has_internal(from)); assert(!has_internal(to)); - auto itr = internal_to_external_.find(from); - auto external = itr->second; + auto external = load_external(from); // Updating the internal-to-external ID is easy. external_to_internal_[external] = to; - internal_to_external_.erase(itr); - internal_to_external_.insert({to, external}); + store_external(from, invalid_external_id); + store_external(to, external); } /// @@ -314,7 +315,7 @@ class IDTranslator { auto internal = itr->second; // Updating the external-to-internal ID is easy. - internal_to_external_[internal] = to; + store_external(internal, to); external_to_internal_.erase(itr); external_to_internal_.insert({to, internal}); } @@ -340,7 +341,7 @@ class IDTranslator { for (const auto i : internal_ids) { auto e = get_external(i); - internal_to_external_.erase(i); + store_external(i, invalid_external_id); external_to_internal_.erase(e); } } @@ -366,7 +367,7 @@ class IDTranslator { for (const auto e : external_ids) { auto i = get_internal(e); - internal_to_external_.erase(i); + store_external(i, invalid_external_id); external_to_internal_.erase(e); } } @@ -377,9 +378,24 @@ class IDTranslator { /// @param begin Start forward iterator to the external IDs to check. /// @param end End forward iterator to the external IDs to check. /// + template + void check_external_not_reserved(const Begin& begin, const End& end) const { + check( + begin, + end, + [](auto e) { return std::cmp_equal(e, invalid_external_id); }, + "Reserved external" + ); + } + template void check_external_free(const Begin& begin, const End& end) const { - check(begin, end, external_to_internal_, "Index already contains external"); + check( + begin, + end, + [this](auto e) { return has_external(e); }, + "Index already contains external" + ); } /// @@ -393,7 +409,7 @@ class IDTranslator { check( begin, end, - external_to_internal_, + [this](auto e) { return has_external(e); }, "Index does not contain external", std::logical_not() ); @@ -407,7 +423,12 @@ class IDTranslator { /// template void check_internal_free(const Begin& begin, const End& end) const { - check(begin, end, internal_to_external_, "Index already contains internal"); + check( + begin, + end, + [this](auto i) { return has_internal(i); }, + "Index already contains internal" + ); } /// @@ -421,7 +442,7 @@ class IDTranslator { check( begin, end, - internal_to_external_, + [this](auto i) { return has_internal(i); }, "Index does not contain internal", std::logical_not() ); @@ -501,23 +522,43 @@ class IDTranslator { } private: - template + template void check( const Begin& begin, const End& end, - const Map& map, + const Contains& contains, const char* message, Modifier modify = lib::identity() ) const { for (auto i = begin; i != end; ++i) { - if (modify(map.contains(*i))) { + if (modify(contains(*i))) { throw ANNEXCEPTION("{} ID {}!", message, *i); } } } tsl::robin_map external_to_internal_{}; - tsl::robin_map internal_to_external_{}; + // Vector for lock-free reads; TODO: moves to #447 (same table), not in this PR. + lib::SegmentedVector internal_to_external_{}; + + external_id_type load_external(internal_id_type i) const { + if (i >= internal_to_external_.size()) { + return invalid_external_id; + } + return std::atomic_ref( + const_cast(internal_to_external_[i]) + ) + .load(std::memory_order_acquire); + } + + // Writers only; grows the table first when `i` is beyond it. + void store_external(size_t i, external_id_type e) { + if (i >= internal_to_external_.size()) { + internal_to_external_.resize(i + 1, invalid_external_id); + } + std::atomic_ref(internal_to_external_[i]) + .store(e, std::memory_order_release); + } }; } // namespace svs diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 8a9a64dd1..072ecb3c3 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -109,6 +109,26 @@ class ValidBuilder { const std::vector& status_; }; +// Separate template so the regular search compiles without any filter check. +template class FilteredValidBuilder { + public: + FilteredValidBuilder(const IDTranslator& translator, const Filter& filter) + : translator_{&translator} + , filter_{&filter} {} + + template + constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { + // delete_entries clears the translator entry, so a missing id means deleted. + auto id = translator_->get_external_or(i, IDTranslator::invalid_external_id); + bool invalid = id == IDTranslator::invalid_external_id || !(*filter_)(id); + return PredicatedSearchNeighbor(i, distance, !invalid); + } + + private: + const IDTranslator* translator_; + const Filter* filter_; +}; + namespace detail { template @@ -517,11 +537,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)); @@ -636,6 +652,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); } @@ -679,13 +698,28 @@ class MutableVamanaIndex { // This is an internal method, mostly used to help implement the batch iterator. ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } + template + FilteredValidBuilder internal_search_builder(const Filter& filter) const { + return FilteredValidBuilder{translator_, filter}; + } + + template 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* 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(filter); // Unused in regular search. + auto builder = [&] { + if constexpr (Filtered) { + return internal_search_builder(*filter); + } else { + return internal_search_builder(); + } + }(); // Perform the greedy search using the provided resources. greedy_search( graph_, @@ -695,7 +729,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - internal_search_builder(), + builder, prefetch_parameters, cancel ); @@ -727,7 +761,8 @@ 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 ) { threads::parallel_for( threadpool_, @@ -747,6 +782,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( + prefetch_parameters, cancel, filter + ), + *this, + cancel + ); + } + (); + return; + } extensions::per_thread_batch_search( data_, buffer, @@ -980,12 +1034,14 @@ class MutableVamanaIndex { return ids.size(); } + private: void delete_entry(size_t i) { SlotMetadata& meta = getindex(status_, i); assert(meta == SlotMetadata::Valid); meta = SlotMetadata::Deleted; } + public: bool is_deleted(size_t i) const { return status_[i] != SlotMetadata::Valid; } Idx entry_point() const { diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 96ffeaa06..a1b87ccc2 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -585,9 +585,15 @@ void svs_invoke( ); // Copy back results. - for (size_t j = 0; j < num_neighbors; ++j) { + auto num_results = std::min(num_neighbors, search_buffer.size()); + for (size_t j = 0; j < num_results; ++j) { result.set(search_buffer[j], i, j); } + // Fewer than k results (e.g. a selective filter): pad with invalid ids. + for (size_t j = num_results; j < num_neighbors; ++j) { + result.index(i, j) = static_cast(-1); + result.distance(i, j) = distance::worst_distance(distance); + } } } diff --git a/include/svs/index/vamana/greedy_search.h b/include/svs/index/vamana/greedy_search.h index f12c0129f..30225578f 100644 --- a/include/svs/index/vamana/greedy_search.h +++ b/include/svs/index/vamana/greedy_search.h @@ -197,7 +197,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)); + } } } } diff --git a/include/svs/index/vamana/index.h b/include/svs/index/vamana/index.h index 4d4a47f09..bce8d8b81 100644 --- a/include/svs/index/vamana/index.h +++ b/include/svs/index/vamana/index.h @@ -24,6 +24,7 @@ #include "svs/core/query_result.h" #include "svs/core/recall.h" #include "svs/index/vamana/calibrate.h" +#include "svs/index/vamana/dynamic_search_buffer.h" #include "svs/index/vamana/extensions.h" #include "svs/index/vamana/greedy_search.h" #include "svs/index/vamana/search_buffer.h" @@ -39,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -50,6 +52,22 @@ namespace svs::index::vamana { template class BatchIterator; +// Marks candidates failing the filter invalid. +template class FilteredNeighborBuilder { + public: + explicit FilteredNeighborBuilder(const Filter& filter) + : filter_{&filter} {} + + template + constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { + bool valid = (*filter_)(static_cast(i)); + return PredicatedSearchNeighbor(i, distance, valid); + } + + private: + const Filter* filter_; +}; + struct VamanaIndexParameters { // Members public: @@ -485,13 +503,23 @@ class VamanaIndex { // This is an internal method, mostly used to help implement the batch iterator. static NeighborBuilder internal_search_builder() { return NeighborBuilder(); } + template 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* 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(filter); // Unused in regular search. + auto builder = [&] { + if constexpr (Filtered) { + return FilteredNeighborBuilder>{*filter}; + } else { + return internal_search_builder(); + } + }(); greedy_search( graph_, data_, @@ -500,10 +528,14 @@ class VamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - internal_search_builder(), + builder, prefetch_parameters, cancel ); + if constexpr (Filtered) { + // Remove the invalid entries that remain in the buffer. + buffer.cleanup(); + } }; } @@ -578,12 +610,46 @@ class VamanaIndex { QueryResultView result, const Queries& queries, const search_parameters_type& search_parameters, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), + const std::function* filter = nullptr ) { threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { + if (filter != nullptr) { + [&]() __attribute__((noinline)) { + // The filtered search buffer must hold invalid entries. + auto search_buffer = MutableBuffer>{ + search_parameters.buffer_config_.get_total_capacity() < + result.n_neighbors() + ? SearchBufferConfig{result.n_neighbors()} + : search_parameters.buffer_config_, + distance::comparator(distance_), + search_parameters.search_buffer_visited_set_}; + auto prefetch_parameters = GreedySearchPrefetchParameters{ + search_parameters.prefetch_lookahead_, + search_parameters.prefetch_step_}; + auto scratch = + extensions::per_thread_batch_search_setup(data_, distance_); + extensions::per_thread_batch_search( + data_, + search_buffer, + scratch, + queries, + result, + threads::UnitRange{is}, + greedy_search_closure( + prefetch_parameters, cancel, filter + ), + *this, + cancel + ); + } + (); + return; + } + // The number of neighbors to store in the result. size_t num_neighbors = result.n_neighbors(); diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 261278d7b..5e1f53a0e 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -609,6 +609,21 @@ 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. + /// Rows with fewer passing entries than requested are padded with id ``size_t(-1)``. + /// ``filter`` must be thread-safe and must not call into the same index. + template + void filtered_search( + QueryResultView results, + data::ConstSimpleDataView queries, + const VamanaSearchParameters& search_parameters, + const std::function& filter + ) { + impl_->filtered_search( + results, AnonymousArray<2>(queries), search_parameters, filter + ); + } }; /// diff --git a/include/svs/orchestrators/manager.h b/include/svs/orchestrators/manager.h index b3f4d33d2..1625f256d 100644 --- a/include/svs/orchestrators/manager.h +++ b/include/svs/orchestrators/manager.h @@ -146,25 +146,11 @@ class ManagerImpl : public ManagerInterface { const search_parameters_type& search_parameters, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) override { - // See if we have a specialization for this particular query type. - // If so, invoke that specialization, otherwise throw - lib::match( - QueryTypes{}, - data.type(), - [&](lib::Type SVS_UNUSED(type)) { - const auto view = data::ConstSimpleDataView(data); - svs::index::search_batch_into_with( - implementation_, result, view, search_parameters, cancel - ); - }, - [&](svs::DataType data_type) { - throw ANNEXCEPTION( - "Unsupported datatype! Got: {}. Expected one of: {}.", - data_type, - fmt::join(QueryTypes::data_types(), ", ") - ); - } - ); + dispatch_queries(data, [&](const auto& view) { + svs::index::search_batch_into_with( + implementation_, result, view, search_parameters, cancel + ); + }); } // Data Interface @@ -191,6 +177,26 @@ class ManagerImpl : public ManagerInterface { Impl& impl() { return implementation_; } const Impl& impl() const { return implementation_; } + // Invoke `f` with the query view of the matching supported type, otherwise throw. + template + void dispatch_queries(const AnonymousArray<2>& data, F&& f) const { + lib::match( + QueryTypes{}, + data.type(), + [&](lib::Type SVS_UNUSED(type)) { + const auto view = data::ConstSimpleDataView(data); + f(view); + }, + [&](svs::DataType data_type) { + throw ANNEXCEPTION( + "Unsupported datatype! Got: {}. Expected one of: {}.", + data_type, + fmt::join(QueryTypes::data_types(), ", ") + ); + } + ); + } + private: friend class IndexManager; Impl implementation_; diff --git a/include/svs/orchestrators/vamana.h b/include/svs/orchestrators/vamana.h index 5787e8a7a..cada48419 100644 --- a/include/svs/orchestrators/vamana.h +++ b/include/svs/orchestrators/vamana.h @@ -35,6 +35,7 @@ // stdlib #include +#include #include #include #include @@ -108,6 +109,14 @@ class VamanaInterface { // through the wrong slot with no link-time diagnostic. virtual svs::index::vamana::MemoryBreakdown get_memory_breakdown() const = 0; virtual svs::index::vamana::VamanaIndexParameters parameters() const = 0; + + // Filtered search; keep last (ABI). + virtual void filtered_search( + svs::QueryResultView results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter + ) = 0; }; template @@ -283,6 +292,19 @@ class VamanaImpl : public manager::ManagerImpl { svs::index::vamana::VamanaIndexParameters parameters() const override { return impl().parameters(); } + + void filtered_search( + svs::QueryResultView results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter + ) override { + this->dispatch_queries(queries, [&](const auto& view) { + impl().search( + results, view, search_parameters, lib::Returns(lib::Const()), &filter + ); + }); + } }; ///// Forward declarations @@ -776,6 +798,21 @@ class Vamana : public manager::IndexManager { return impl_->parameters(); } + /// @brief Filtered search: only ids for which ``filter(id)`` is true are returned. + /// Rows with fewer passing entries than requested are padded with id ``size_t(-1)``. + /// ``filter`` must be thread-safe and must not call into the same index. + template + void filtered_search( + QueryResultView results, + data::ConstSimpleDataView queries, + const VamanaSearchParameters& search_parameters, + const std::function& filter + ) { + impl_->filtered_search( + results, AnonymousArray<2>(queries), search_parameters, filter + ); + } + private: template < manager::QueryTypeDefinition QueryTypes,