diff --git a/src/VecSim/algorithms/svs/svs.h b/src/VecSim/algorithms/svs/svs.h index 7dc15dc0d..e5c9b6551 100644 --- a/src/VecSim/algorithms/svs/svs.h +++ b/src/VecSim/algorithms/svs/svs.h @@ -587,29 +587,58 @@ class SVSIndex : public VecSimIndexAbstract, fl auto processed_query_ptr = this->preprocessQuery(queryBlob); const void *processed_query = processed_query_ptr.get(); - auto query = svs::data::ConstSimpleDataView{ - static_cast(processed_query), 1, this->dim}; - auto result = svs::QueryResult{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{ + static_cast(processed_query), 1, this->dim}; + auto result = svs::QueryResult{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{static_cast(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. diff --git a/tests/benchmark/bm_initialization/bm_basics_svs_initialize_fp32.h b/tests/benchmark/bm_initialization/bm_basics_svs_initialize_fp32.h index bdbb7e007..225758018 100644 --- a/tests/benchmark/bm_initialization/bm_basics_svs_initialize_fp32.h +++ b/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"}); diff --git a/tests/benchmark/bm_vecsim_svs.h b/tests/benchmark/bm_vecsim_svs.h index 5acb882c0..b7628b920 100644 --- a/tests/benchmark/bm_vecsim_svs.h +++ b/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::RunGC(benchmark::State &st) { ASSERT_EQ(VecSimIndex_IndexSize(tiered_index), N_VECTORS - num_deletions); } +template +void BM_VecSimSVS::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