[SVS] Improve SVS topKQuery() by using "single-search" for non-multi

src/VecSim/algorithms/svs/svs.h

@@ -587,29 +587,58 @@ class SVSIndex : public VecSimIndexAbstract<svs_details::vecsim_dt<DataType>, fl
         auto processed_query_ptr = this->preprocessQuery(queryBlob);
         const void *processed_query = processed_query_ptr.get();
 
-        auto query = svs::data::ConstSimpleDataView<DataType>{
-            static_cast<const DataType *>(processed_query), 1, this->dim};
-        auto result = svs::QueryResult<size_t>{query.size(), k};
         auto sp = svs_details::joinSearchParams(impl_->get_search_parameters(), queryParams,
                                                 is_two_level_lvq);
 
         auto timeoutCtx = queryParams ? queryParams->timeoutCtx : nullptr;
         auto cancel = [timeoutCtx]() { return VECSIM_TIMEOUT(timeoutCtx); };
 
-        impl_->search(result.view(), query, sp, cancel);
-        if (cancel()) {
-            rep->code = VecSim_QueryReply_TimedOut;
-            return rep;
-        }
+        if constexpr (isMulti) {
+            auto query = svs::data::ConstSimpleDataView<DataType>{
+                static_cast<const DataType *>(processed_query), 1, this->dim};
+            auto result = svs::QueryResult<size_t>{query.size(), k};
+            impl_->search(result.view(), query, sp, cancel);
+            if (cancel()) {
+                rep->code = VecSim_QueryReply_TimedOut;
+                return rep;
+            }
+
+            assert(result.n_queries() == 1);
 
-        assert(result.n_queries() == 1);
+            const auto n_neighbors = result.n_neighbors();
+            rep->results.reserve(n_neighbors);
 
-        const auto n_neighbors = result.n_neighbors();
-        rep->results.reserve(n_neighbors);
+            for (size_t i = 0; i < n_neighbors; i++) {
+                rep->results.push_back(
+                    VecSimQueryResult{result.index(0, i), toVecSimDistance(result.distance(0, i))});
+            }
+        } else {
+            // Use the single-query search path of MutableVamanaIndex to avoid the
+            // parallel batch-search overhead for a single query.
+            auto scratch = impl_->scratchspace(sp);
+            // Ensure the scratch buffer can hold at least `k` neighbors.
+            if (scratch.buffer.target_capacity() < k) {
+                scratch.buffer.change_maxsize(k);
+            }
 
-        for (size_t i = 0; i < n_neighbors; i++) {
-            rep->results.push_back(
-                VecSimQueryResult{result.index(0, i), toVecSimDistance(result.distance(0, i))});
+            auto query = std::span<const DataType>{static_cast<const DataType *>(processed_query),
+                                                   this->dim};
+            impl_->search(query, scratch, cancel);
+            if (cancel()) {
+                rep->code = VecSim_QueryReply_TimedOut;
+                return rep;
+            }
+
+            const auto &buffer = scratch.buffer;
+            const auto n_neighbors = std::min(k, buffer.size());
+            rep->results.reserve(n_neighbors);
+
+            for (size_t i = 0; i < n_neighbors; i++) {
+                const auto& neighbor = buffer[i];
+                rep->results.push_back(
+                    VecSimQueryResult{impl_->translate_internal_id(neighbor.id()),
+                                      toVecSimDistance(neighbor.distance())});
+            }
         }
         // Workaround for VecSim merge_results() that expects results to be sorted
         // by score, then by id from both indices.

tests/benchmark/bm_initialization/bm_basics_svs_initialize_fp32.h

@@ -54,3 +54,17 @@ BENCHMARK_REGISTER_F(BM_VecSimSVS, BM_FUNC_NAME(BM_AddVectorsDuringTraining))
                    {2, 4}})
     ->ArgNames({"training_threshold", "thread_count"})
     ->UseRealTime();
+
+// TopK search on the loaded SVS index, using the search window_size and k defined below.
+// {window_size, k} (recall that always window_size >= k)
+BENCHMARK_TEMPLATE_DEFINE_F(BM_VecSimSVS, BM_FUNC_NAME(BM_TopK), DATA_TYPE_INDEX_T)
+(benchmark::State &st) { TopK_SVS(st); }
+BENCHMARK_REGISTER_F(BM_VecSimSVS, BM_FUNC_NAME(BM_TopK))
+    ->Unit(benchmark::kMillisecond)
+    ->Iterations(10)
+    ->Args({10, 10})
+    ->Args({200, 10})
+    ->Args({100, 100})
+    ->Args({200, 100})
+    ->Args({500, 500})
+    ->ArgNames({"window_size", "k"});

tests/benchmark/bm_vecsim_svs.h

@@ -54,6 +54,10 @@ class BM_VecSimSVS : public BM_VecSimGeneral {
     // Deletes an amount of labels from the index that triggers inplace consolidation.
     void RunGC(benchmark::State &st);
 
+    // TopK search benchmark. The search window size and k are provided by the benchmark
+    // registration.
+    void TopK_SVS(benchmark::State &st);
+
 private:
     static const char *svs_index_tar_file;
     static std::string base_path;
@@ -463,6 +467,29 @@ void BM_VecSimSVS<index_type_t>::RunGC(benchmark::State &st) {
     ASSERT_EQ(VecSimIndex_IndexSize(tiered_index), N_VECTORS - num_deletions);
 }
 
+template <typename index_type_t>
+void BM_VecSimSVS<index_type_t>::TopK_SVS(benchmark::State &st) {
+    // Search window size (Vamana graph accuracy/latency tuning) and number of results, as
+    // defined by the benchmark registration.
+    size_t window_size = st.range(0);
+    size_t k = st.range(1);
+
+    // Load the SVS index from file (update_threshold is irrelevant for a search-only benchmark).
+    auto *index = CreateSVSIndexFromFile(1);
+
+    VecSimQueryParams query_params = {.svsRuntimeParams = {.windowSize = window_size}};
+
+    size_t iter = 0;
+    for (auto _ : st) {
+        auto results = VecSimIndex_TopKQuery(index, test_vectors[iter % N_QUERIES].data(), k,
+                                             &query_params, BY_SCORE);
+        VecSimQueryReply_Free(results);
+        iter++;
+    }
+
+    VecSimIndex_Free(index);
+}
+
 #define UNIT_AND_ITERATIONS Unit(benchmark::kMillisecond)->Iterations(2)
 
 #if HAVE_SVS_LVQ