Add cancel (#409)

GitOrigin-RevId: 98418f02d3916800f722fb97b1755086973895ad

Author: dian-lun-lin

include/svs/index/flat/flat.h

@@ -212,7 +212,6 @@ class FlatIndex {
     /// Return the logical number of dimensions of the indexed vectors.
     size_t dimensions() const { return data_.dimensions(); }
 
-    ///
     /// @anchor flat_class_search_mutating
     /// @brief Fill the result with the ``num_neighbors`` nearest neighbors for each query.
     ///
@@ -225,6 +224,9 @@ class FlatIndex {
     /// @param result The result data structure to populate.
     ///     Row `i` in the result corresponds to the neighbors for the `i`th query.
     ///     Neighbors within each row are ordered from nearest to furthest.
+    /// @param cancel A predicate called during the search to determine if the search should be cancelled.
+    //      Return ``true`` if the search should be cancelled. This functor must implement ``bool operator()()``.
+    //      Note: This predicate should be thread-safe as it can be called concurrently by different threads during the search.
     /// @param predicate A predicate functor that can be used to exclude certain dataset
     ///     elements from consideration. This functor must implement
     ///     ``bool operator()(size_t)`` where the ``size_t`` argument is an index in
@@ -260,6 +262,7 @@ class FlatIndex {
         QueryResultView<size_t> result,
         const data::ConstSimpleDataView<QueryType>& queries,
         const search_parameters_type& search_parameters,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>()),
         Pred predicate = lib::Returns(lib::Const<true>())
     ) {
         const size_t data_max_size = data_.size();
@@ -276,12 +279,18 @@ class FlatIndex {
 
         size_t start = 0;
         while (start < data_.size()) {
+            // Check if request to cancel the search
+            if (cancel()) {
+                scratch.cleanup();
+                return;
+            }
             size_t stop = std::min(data_max_size, start + data_batch_size);
             search_subset(
                 queries,
                 threads::UnitRange(start, stop),
                 scratch,
                 search_parameters,
+                cancel,
                 predicate
             );
             start = stop;
@@ -311,6 +320,7 @@ class FlatIndex {
         const threads::UnitRange<size_t>& data_indices,
         sorter_type& scratch,
         const search_parameters_type& search_parameters,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>()),
         Pred predicate = lib::Returns(lib::Const<true>())
     ) {
         // Process all queries.
@@ -331,6 +341,7 @@ class FlatIndex {
                     threads::UnitRange(query_indices),
                     scratch,
                     distances,
+                    cancel,
                     predicate
                 );
             }
@@ -352,6 +363,7 @@ class FlatIndex {
         const threads::UnitRange<size_t>& query_indices,
         sorter_type& scratch,
         distance::BroadcastDistance<DistFull>& distance_functors,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>()),
         Pred predicate = lib::Returns(lib::Const<true>())
     ) {
         assert(distance_functors.size() >= query_indices.size());
@@ -365,6 +377,11 @@ class FlatIndex {
         }
 
         for (auto data_index : data_indices) {
+            // Check if request to cancel the search
+            if (cancel()) {
+                return;
+            }
+
             // Skip this index if it doesn't pass the predicate.
             if (!predicate(data_index)) {
                 continue;

include/svs/index/index.h

@@ -41,20 +41,24 @@ void search_batch_into_with(
     Index& index,
     svs::QueryResultView<I> result,
     const Queries& queries,
-    const search_parameters_t<Index>& search_parameters
+    const search_parameters_t<Index>& search_parameters,
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     // Assert pre-conditions.
     assert(result.n_queries() == queries.size());
-    index.search(result, queries, search_parameters);
+    index.search(result, queries, search_parameters, cancel);
 }
 
 // Apply default search parameters
 template <typename Index, std::integral I, data::ImmutableMemoryDataset Queries>
 void search_batch_into(
-    Index& index, svs::QueryResultView<I> result, const Queries& queries
+    Index& index,
+    svs::QueryResultView<I> result,
+    const Queries& queries,
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     svs::index::search_batch_into_with(
-        index, result, queries, index.get_search_parameters()
+        index, result, queries, index.get_search_parameters(), cancel
     );
 }
 
@@ -64,19 +68,26 @@ svs::QueryResult<size_t> search_batch_with(
     Index& index,
     const Queries& queries,
     size_t num_neighbors,
-    const search_parameters_t<Index>& search_parameters
+    const search_parameters_t<Index>& search_parameters,
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     auto result = svs::QueryResult<size_t>{queries.size(), num_neighbors};
-    svs::index::search_batch_into_with(index, result.view(), queries, search_parameters);
+    svs::index::search_batch_into_with(
+        index, result.view(), queries, search_parameters, cancel
+    );
     return result;
 }
 
 // Obtain default search parameters.
 template <typename Index, data::ImmutableMemoryDataset Queries>
-svs::QueryResult<size_t>
-search_batch(Index& index, const Queries& queries, size_t num_neighbors) {
+svs::QueryResult<size_t> search_batch(
+    Index& index,
+    const Queries& queries,
+    size_t num_neighbors,
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
+) {
     return svs::index::search_batch_with(
-        index, queries, num_neighbors, index.get_search_parameters()
+        index, queries, num_neighbors, index.get_search_parameters(), cancel
     );
 }
 } // namespace svs::index

include/svs/index/inverted/memory_based.h

@@ -376,7 +376,8 @@ template <typename Index, typename Cluster> class InvertedIndex {
     void search(
         QueryResultView<Idx> results,
         const Queries& queries,
-        const search_parameters_type& search_parameters
+        const search_parameters_type& search_parameters,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
     ) {
         threads::run(
             threadpool_,
@@ -396,9 +397,11 @@ template <typename Index, typename Cluster> class InvertedIndex {
 
                 for (auto i : is) {
                     buffer.clear();
+
                     auto&& query = queries.get_datum(i);
                     // Primary Index Search
-                    index_.search(query, scratch);
+                    index_.search(query, scratch, cancel);
+
                     auto& d = scratch.scratch;
                     auto compare = distance::comparator(d);
 
@@ -411,6 +414,11 @@ template <typename Index, typename Cluster> class InvertedIndex {
                     );
 
                     for (size_t j = 0, jmax = scratch_buffer.size(); j < jmax; ++j) {
+                        // Check if request to cancel the search
+                        // TODO: this cancel may also be inside on_leaves()
+                        if (cancel()) {
+                            return;
+                        }
                         auto candidate = scratch_buffer[j];
                         if (!compare(candidate.distance(), cutoff_distance)) {
                             break;

include/svs/index/vamana/dynamic_index.h

@@ -402,7 +402,7 @@ class MutableVamanaIndex {
     ///
     size_t dimensions() const { return data_.dimensions(); }
 
-    auto greedy_search_closure(GreedySearchPrefetchParameters prefetch_parameters) const {
+    auto greedy_search_closure(GreedySearchPrefetchParameters prefetch_parameters, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())) const {
         return [&, prefetch_parameters](
                    const auto& query, auto& accessor, auto& distance, auto& buffer
                ) {
@@ -416,7 +416,8 @@ class MutableVamanaIndex {
                 buffer,
                 entry_point_,
                 ValidBuilder{status_},
-                prefetch_parameters
+                prefetch_parameters,
+                cancel
             );
             // Take a pass over the search buffer to remove any deleted elements that
             // might remain.
@@ -426,7 +427,10 @@ class MutableVamanaIndex {
 
     template <typename I, data::ImmutableMemoryDataset Queries>
     void search(
-        QueryResultView<I> results, const Queries& queries, const search_parameters_type& sp
+        QueryResultView<I> results,
+        const Queries& queries,
+        const search_parameters_type& sp,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
     ) {
         threads::run(
             threadpool_,
@@ -452,11 +456,17 @@ class MutableVamanaIndex {
                     queries,
                     results,
                     threads::UnitRange{is},
-                    greedy_search_closure(prefetch_parameters)
+                    greedy_search_closure(prefetch_parameters, cancel),
+                    cancel
                 );
             }
         );
 
+        // Check if request to cancel the search
+        if (cancel()) {
+            return;
+        }
+
         // 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());

include/svs/index/vamana/extensions.h

@@ -411,15 +411,19 @@ struct VamanaSingleSearchType {
         SearchBuffer& search_buffer,
         Scratch& scratch,
         const Query& query,
-        const Search& search
+        const Search& search,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
     ) const {
-        svs::svs_invoke(*this, data, search_buffer, scratch, query, search);
+        svs::svs_invoke(*this, data, search_buffer, scratch, query, search, cancel);
     }
 };
 
 /// Customization point object for processing single queries.
 inline constexpr VamanaSingleSearchType single_search{};
 
+/// In this function, cancel does not need to be called since search will call cancel
+/// However, lvq requires cancel to be called to skip reranking
+/// Here we have cancel to make the interface consistent with lvq
 template <
     typename Data,
     typename SearchBuffer,
@@ -432,8 +436,13 @@ SVS_FORCE_INLINE void svs_invoke(
     SearchBuffer& search_buffer,
     Distance& distance,
     const Query& query,
-    const Search& search
+    const Search& search,
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
+    // Check if request to cancel the search
+    if (cancel()) {
+        return;
+    }
     // Perform graph search.
     auto accessor = data::GetDatumAccessor();
     search(query, accessor, distance, search_buffer);
@@ -482,10 +491,19 @@ struct VamanaPerThreadBatchSearchType {
         const Queries& queries,
         QueryResultView<I>& result,
         threads::UnitRange<size_t> thread_indices,
-        const Search& search
+        const Search& search,
+        const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
     ) const {
         svs::svs_invoke(
-            *this, data, search_buffer, scratch, queries, result, thread_indices, search
+            *this,
+            data,
+            search_buffer,
+            scratch,
+            queries,
+            result,
+            thread_indices,
+            search,
+            cancel
         );
     }
 };
@@ -509,13 +527,18 @@ void svs_invoke(
     const Queries& queries,
     QueryResultView<I>& result,
     threads::UnitRange<size_t> thread_indices,
-    const Search& search
+    const Search& search,
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     // Fallback implementation
     size_t num_neighbors = result.n_neighbors();
     for (auto i : thread_indices) {
+        // Check if request to cancel the search
+        if (cancel()) {
+            return;
+        }
         // Perform search - results will be queued in the search buffer.
-        single_search(dataset, search_buffer, distance, queries.get_datum(i), search);
+        single_search(dataset, search_buffer, distance, queries.get_datum(i), search, cancel);
 
         // Copy back results.
         for (size_t j = 0; j < num_neighbors; ++j) {

include/svs/index/vamana/greedy_search.h

@@ -78,7 +78,8 @@ void greedy_search(
     const Ep& entry_points,
     const Builder& builder,
     Tracker& search_tracker,
-    GreedySearchPrefetchParameters prefetch_parameters = {}
+    GreedySearchPrefetchParameters prefetch_parameters = {},
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     using I = typename Graph::index_type;
 
@@ -102,6 +103,10 @@ void greedy_search(
     // Main search routine.
     search_buffer.sort();
     while (!search_buffer.done()) {
+        // Check if request to cancel the search
+        if (cancel()) {
+            return;
+        }
         // Get the next unvisited vertex.
         const auto& node = search_buffer.next();
         auto node_id = node.id();
@@ -169,7 +174,8 @@ void greedy_search(
     Buffer& search_buffer,
     const Ep& entry_points,
     const Builder& builder = NeighborBuilder(),
-    GreedySearchPrefetchParameters prefetch_parameters = {}
+    GreedySearchPrefetchParameters prefetch_parameters = {},
+    const lib::DefaultPredicate& cancel = lib::Returns(lib::Const<false>())
 ) {
     auto null_tracker = NullTracker{};
     greedy_search(
@@ -182,7 +188,8 @@ void greedy_search(
         entry_points,
         builder,
         null_tracker,
-        prefetch_parameters
+        prefetch_parameters,
+        cancel
     );
 }
 } // namespace svs::index::vamana