Handle corner in dynamic index with insufficient valid search results (#164)

This PR addresses a rare scenario in dynamic indexes where the search
buffer may not be populated with enough results. This typically occurs
when a large number of vectors have been deleted, resulting in a
sparsely connected graph. To ensure robustness, the buffer is now
supplemented with additional results when needed.

This is a corner case and is not expected to occur frequently, but
handling it improves stability and consistency in edge conditions.

Author: ibhati

include/svs/index/vamana/dynamic_index.h

@@ -491,7 +491,8 @@ class MutableVamanaIndex {
             scratch.buffer,
             scratch.scratch,
             query,
-            greedy_search_closure(scratch.prefetch_parameters, cancel)
+            greedy_search_closure(scratch.prefetch_parameters, cancel),
+            *this
         );
     }
 
@@ -514,7 +515,7 @@ class MutableVamanaIndex {
                     sp.prefetch_lookahead_, sp.prefetch_step_};
 
                 // Legalize search buffer for this search.
-                if (buffer.target() < num_neighbors) {
+                if (buffer.target_capacity() < num_neighbors) {
                     buffer.change_maxsize(num_neighbors);
                 }
                 auto scratch = extensions::per_thread_batch_search_setup(data_, distance_);
@@ -527,6 +528,7 @@ class MutableVamanaIndex {
                     results,
                     threads::UnitRange{is},
                     greedy_search_closure(prefetch_parameters, cancel),
+                    *this,
                     cancel
                 );
             }

include/svs/index/vamana/dynamic_search_buffer.h

@@ -193,11 +193,14 @@ template <typename Idx, typename Cmp = std::less<>> class MutableBuffer {
     /// @brief Return the number of valid elements currently contained in the buffer.
     size_t valid() const { return valid_; }
 
-    /// @brief Return the target number of valid candidates.
-    size_t target() const { return valid_capacity_; }
+    /// @brief Return the target valid capacity as equivalent to buffer capacity
+    size_t target_capacity() const { return valid_capacity_; }
+
+    /// @brief Return the target valid candidates as equivalent to the search window
+    size_t target_window() const { return target_valid_; }
 
     /// @brief Return whether or not the buffer contains its target number of candidates.
-    bool full() const { return valid() == target(); }
+    bool full() const { return valid() == target_capacity(); }
 
     /// @brief Return the candidate at index `i`.
     ///
@@ -577,7 +580,7 @@ template <typename Idx, typename Cmp = std::less<>> class MutableBuffer {
     /// If the number of valid candidates is *less* than the target, a negative number
     /// is returned.
     int64_t slack() const {
-        return lib::narrow_cast<int64_t>(valid()) - lib::narrow_cast<int64_t>(target());
+        return lib::narrow_cast<int64_t>(valid()) - lib::narrow_cast<int64_t>(target_capacity());
     }
 
     /// Return the index of the first preceding valid candidate beginning at the provided
@@ -597,7 +600,7 @@ template <typename Idx, typename Cmp = std::less<>> class MutableBuffer {
 template <typename Idx, typename Cmp>
 std::ostream& operator<<(std::ostream& io, const MutableBuffer<Idx, Cmp>& buffer) {
     return io << "MutableBuffer<" << datatype_v<Idx> << ">("
-              << "target_valid = " << buffer.target()
+              << "target_valid = " << buffer.target_capacity()
               << ", best_unvisited = " << buffer.best_unvisited()
               << ", valid = " << buffer.valid() << ", size = " << buffer.size() << ")";
 }

include/svs/index/vamana/extensions.h

@@ -410,19 +410,44 @@ struct VamanaSingleSearchType {
         typename SearchBuffer,
         typename Scratch,
         typename Query,
-        typename Search>
+        typename Search,
+        typename Index>
     void operator()(
         const Data& data,
         SearchBuffer& search_buffer,
         Scratch& scratch,
         const Query& query,
         const Search& search,
+        const Index& index,
         const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
     ) const {
-        svs::svs_invoke(*this, data, search_buffer, scratch, query, search, cancel);
+        svs::svs_invoke(*this, data, search_buffer, scratch, query, search, index, cancel);
     }
 };
 
+/// In rare cases, the search buffer may not be filled with enough results.
+/// This can occur in dynamic indexes when many vectors have been deleted
+/// and the graph becomes sparsely connected. It's a corner case and should
+/// not happen frequently, but when it does, we may need to supplement the buffer
+/// 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
+) {
+    if (search_buffer.valid() < search_buffer.target_window() &&
+        search_buffer.valid() < index.size()) {
+        for (auto external_id : index.external_ids()) {
+            auto internal_id = index.translate_external_id(external_id);
+            auto dist = index.get_distance(external_id, query);
+            auto builder = index.internal_search_builder();
+            search_buffer.insert(builder(internal_id, dist));
+            if (search_buffer.valid() >= search_buffer.target_window()) {
+                break;
+            }
+        }
+    }
+}
+
 /// Customization point object for processing single queries.
 inline constexpr VamanaSingleSearchType single_search{};
 
@@ -434,14 +459,16 @@ template <
     typename SearchBuffer,
     typename Distance,
     typename Query,
-    typename Search>
+    typename Search,
+    typename Index>
 SVS_FORCE_INLINE void svs_invoke(
     svs::tag_t<single_search>,
     const Data& SVS_UNUSED(dataset),
     SearchBuffer& search_buffer,
     Distance& distance,
     const Query& query,
     const Search& search,
+    const Index& index,
     const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     // Check if request to cancel the search
@@ -451,6 +478,10 @@ SVS_FORCE_INLINE void svs_invoke(
     // Perform graph search.
     auto accessor = data::GetDatumAccessor();
     search(query, accessor, distance, search_buffer);
+
+    if constexpr (Index::needs_id_translation) {
+        check_and_supplement_search_buffer(index, search_buffer, query);
+    }
 }
 
 ///
@@ -488,7 +519,8 @@ struct VamanaPerThreadBatchSearchType {
         typename Scratch,
         data::ImmutableMemoryDataset Queries,
         std::integral I,
-        typename Search>
+        typename Search,
+        typename Index>
     SVS_FORCE_INLINE void operator()(
         const Data& data,
         SearchBuffer& search_buffer,
@@ -497,6 +529,7 @@ struct VamanaPerThreadBatchSearchType {
         QueryResultView<I>& result,
         threads::UnitRange<size_t> thread_indices,
         const Search& search,
+        const Index& index,
         const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
     ) const {
         svs::svs_invoke(
@@ -508,6 +541,7 @@ struct VamanaPerThreadBatchSearchType {
             result,
             thread_indices,
             search,
+            index,
             cancel
         );
     }
@@ -523,7 +557,8 @@ template <
     typename Distance,
     typename Queries,
     std::integral I,
-    typename Search>
+    typename Search,
+    typename Index>
 void svs_invoke(
     svs::tag_t<per_thread_batch_search>,
     const Data& dataset,
@@ -533,6 +568,7 @@ void svs_invoke(
     QueryResultView<I>& result,
     threads::UnitRange<size_t> thread_indices,
     const Search& search,
+    const Index& index,
     const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     // Fallback implementation
@@ -544,7 +580,7 @@ void svs_invoke(
         }
         // Perform search - results will be queued in the search buffer.
         single_search(
-            dataset, search_buffer, distance, queries.get_datum(i), search, cancel
+            dataset, search_buffer, distance, queries.get_datum(i), search, index, cancel
         );
 
         // Copy back results.

include/svs/index/vamana/index.h

@@ -510,7 +510,8 @@ class VamanaIndex {
             scratch.buffer,
             scratch.scratch,
             query,
-            greedy_search_closure(scratch.prefetch_parameters, cancel)
+            greedy_search_closure(scratch.prefetch_parameters, cancel),
+            *this
         );
     }
 
@@ -592,6 +593,7 @@ class VamanaIndex {
                     result,
                     threads::UnitRange{is},
                     greedy_search_closure(prefetch_parameters, cancel),
+                    *this,
                     cancel
                 );
             }

include/svs/index/vamana/iterator.h

@@ -303,7 +303,8 @@ template <typename Index, typename QueryType> class BatchIterator {
                 scratchspace_.buffer,
                 scratchspace_.scratch,
                 lib::as_const_span(query_),
-                search_closure
+                search_closure,
+                *parent_
             );
         });
 

tests/svs/index/vamana/search_buffer.cpp

@@ -736,7 +736,7 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") {
         CATCH_SECTION("Full Buffer") {
             // We should be able to add elements to the buffer.
             // Valid elements should only be appended until 4 have been added.
-            CATCH_REQUIRE(b.target() == 4);
+            CATCH_REQUIRE(b.target_capacity() == 4);
             CATCH_REQUIRE(b.size() == 0);
             CATCH_REQUIRE(b.valid() == 0);
             CATCH_REQUIRE(!b.full());