Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
rfsaliev wants to merge 1 commit into
base: main
Choose a base branch
from
+43 −14
Draft

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 43 additions & 14 deletions src/VecSim/algorithms/svs/svs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down