[rfsaliev]

Predicated search and range search for the Vamana indexes

Summary

Adds a pluggable stop predicate to the dynamic search buffer (MutableBuffer) and uses it to provide two new search modes on all three Vamana indexes (VamanaIndex, MutableVamanaIndex, concurrent::MutableVamanaIndex):

• predicated_search: graph search that terminates when a caller-supplied predicate says so, instead of the fixed "search window fully visited" rule.
• range_search(radius): returns all neighbors within a distance bound instead of a fixed number of neighbors.

Cancellation (the existing cancel predicate used for timeouts) is folded into the stop predicate for the new paths. A cancelled query now ends gracefully with the best results found so far instead of returning undefined rows.

Existing search() paths are unchanged in behavior and performance.

Motivation

• Range search bounds results by distance to the query, which the top-K window logic cannot express.
• Upcoming filtered search needs extra stop criteria, e.g. stop early when the share of candidates passing the filter drops below a threshold. That is easiest to express as a predicate over the buffer state.
• cancel was checked at several redundant points per query and, on cancel, left result rows partially filled with stale IDs.

Changes

Search buffer (include/svs/index/vamana/dynamic_search_buffer.h)

• Stop predicate as a template parameter. MutableBuffer<Idx, Cmp, Stop = DefaultStop>, so there is no runtime cost. done() is now best_unvisited == size() || stop(MutableBufferState{...}).
• MutableBufferState is the predicate input:
  • valid: valid results currently in the buffer;
  • total_candidates: all candidates offered to the buffer;
  • valid_visited: valid entries before the best unvisited one;
  • best_unvisited_distance;
  • search_window.
• Predicates:
  • DefaultStop reproduces the previous done() exactly (valid_visited >= search_window).
  • RangeStop<Cmp>{radius} visits at least the search window, then continues while the best unvisited candidate is within the radius. Both L2 (less) and inner product (greater) are handled through the comparator.
  • CancelableStop<Stop> wraps any predicate with a cancel callback.
• total_candidates is opt-in. It is counted only for predicates that declare static constexpr bool counts_total_candidates = true. Counting it unconditionally cost about 1–2% top-K QPS in the per-candidate insert path.
• roi_end_ is replaced by valid_visited_. The new counter holds the number of valid entries in [0, best_unvisited_).
  • It is provably equivalent for the default stop rule; the existing fuzz test's reference model already defines done() this way.
  • It is meaningful to arbitrary predicates.
  • It removes the 5A/5B region-of-interest bookkeeping from unsafe_insert.
  • next() now scans up to size(), so a search that continues past the old window never re-expands a visited node.
• Unbounded capacity. MutableBuffer::unbounded_capacity (SIZE_MAX). In the new APIs, a zero total capacity in the search parameters also means unbounded.
• Small fixes and cleanups:
  • sort() now recomputes best_unvisited_ in all cases; previously it was left stale when fewer than search_window valid entries were present.
  • change_maxsize() keeps the cursor in range when shrinking.
  • slack(), which wrapped for an unbounded capacity, is replaced by direct comparisons.
  • The concurrent buffer's sort() override became identical to the base version and was removed.

Shared implementation (include/svs/index/vamana/predicated_search.h, new)

• predicated_buffer_config() maps search parameters to a buffer configuration (zero capacity means unbounded).
• PredicatedNeighborBuilder marks every candidate valid (used by the static index).
• detail::predicated_batch_search and detail::range_batch_search hold the batch loops shared by all three indexes. Each index supplies only its scratchspace, single-query search, ID translation and an optional per-query translation lock.

Index APIs

Each of VamanaIndex, MutableVamanaIndex and concurrent::MutableVamanaIndex gains:

template <typename Stop> auto predicated_scratchspace(const VamanaSearchParameters& sp, Stop stop) const;
template <typename Query, typename Stop> void predicated_search(const Query& q, predicated_scratchspace_type<Stop>& scratch) const;
template <typename I, typename Queries, typename Stop>
void predicated_search(QueryResultView<I> results, const Queries& queries, const VamanaSearchParameters& sp,
                       const Stop& stop, const lib::DefaultPredicate& cancel = /* never */);
template <typename Queries>
std::vector<std::vector<Neighbor<size_t>>> range_search(const Queries& queries, float radius,
                       const VamanaSearchParameters& sp, const lib::DefaultPredicate& cancel = /* never */);

• Batch predicated_search writes up to k results per query and pads missing ones with sentinel neighbors.
• range_search returns in-radius neighbors with external IDs, sorted best to worst. A non-zero buffer capacity caps the number of results per query; zero means unbounded. The radius filter runs after any dataset-specific reranking done by an extensions::single_search override, so it applies to the final distances.
• Static index: uses MutableBuffer with all candidates valid, because the static SearchBuffer has no stop predicate.
• Concurrent index:
  • The batch paths hold lock_storage_shared() for the whole batch, like search().
  • A private LockedSearch adapter runs the per-query search without taking that lock again. Recursive shared acquisition is the documented hazard.
  • Each query's results are translated under lock_for_translation().
  • Entries erased by a concurrent consolidation keep no_external_id in batch rows, as search() does, and are dropped from range results.
  • It uses the index's own supplement_search_buffer top-up and seqlock greedy_search.

Cancellation

• greedy_search (both the plain and the concurrent version) now takes a templated Cancel that defaults to a compile-time false. Existing callers pass const lib::DefaultPredicate& and behave exactly as before.
• The new paths check cancel once per hop, inside CancelableStop, plus once per query. That is num_queries + 1 fewer calls than search(); the test asserts this exactly.
• On cancel: queries in flight stop at their next hop and keep the best results found so far, which are cleaned, reranked and translated. Queries not yet started get sentinel rows in batch predicated_search, or empty vectors in range_search.
• The single-query predicated_search(query, scratch) has no cancel parameter. Wrap the stop predicate in CancelableStop instead.

Compatibility

• Additive API. MutableBuffer's new template parameter is defaulted, so existing code compiles unchanged.
• Source-level changes to MutableBuffer members:
  • unsafe_insert(neighbor, pos, offset) loses its now-unused offset argument, matching SearchBuffer::unsafe_insert.
  • The protected slack() is removed.
• best_unvisited() can now advance past the old window boundary. Termination and results are unchanged.

Performance

Measured locally: dynamic index, 30k × 96 random vectors with 10% deleted, 5k queries, 32 threads, window and capacity 10–200. Each figure is the median of 6–8 interleaved runs.

Comparison	QPS change	Results
Regular search() vs. main	within ±0.9% (noise)	identical
predicated_search with DefaultStop vs. search()	1.0–1.6% faster	identical
Unconditional total_candidates counting (rejected)	1–2% slower	identical

The last row is why the counter is opt-in.

Testing

• search_buffer.cpp:
  • Existing MutableBuffer and fuzz tests pass unchanged.
  • New: fuzzing with unbounded capacity, and a consistency check of best_unvisited/valid_visited against the reference model.
  • New: RangeStop for less and greater, CancelableStop (including opt-in forwarding), and a filter-rate stopper built on valid / total_candidates.
• vamana/index.cpp (static) and vamana/dynamic_index_2.cpp (dynamic), for L2 and inner product:
  • DefaultStop gives results identical to search().
  • Range search: recall ≥ 0.95 against brute force, all results within the radius, bounded mode respects the capacity, unbounded mode returns everything for a large radius, and deleted points are never returned.
  • Cancellation: exact cancel call counts, sorted partial rows followed by sentinels, and empty or partial range results.
• concurrent/dynamic_index_2.cpp:
  • The same functional checks.
  • Range and predicated search running concurrently with delete_entries, consolidate and compact. This section passed 20 of 20 repeated runs; it has not been run under ThreadSanitizer.
• The full test suite passes.

Notes for reviewers

• Depends on a separate PR fixing a logger-sink use-after-free in tests/svs/index/vamana/dynamic_index_2.cpp: two tests attach a global log sink that captures a local vector and never remove it. The added tests change heap layout enough to turn that latent bug into a SIGSEGV in a later test. The PR fix(test): add scope guard for logger sink cleanup in dynamic index tests #471 should merge first.
• Hot-path files: dynamic_search_buffer.h and both greedy_search.h.
• Follow-ups, not in this PR:
  • Remove the redundant cancel() check at the entry of the default extensions::single_search.
  • Move regular search() onto the graceful-cancel model.
  • Expose the new APIs through the orchestrators and the C/runtime bindings, and switch the runtime range_search from the batch iterator to range_search.