Skip to content

Commit 1e59bf6

Browse files
authoredMar 19, 2025··
Solve batch iterator reranking issue (#95)
This PR addresses the reranking issue in the batch iterator. When the dataset requires reranking after a search, it is necessary to restore the distances using the primary dataset before proceeding to search for the next batch. Failing to do so may result in the search buffer being mixed with both reranking results and primary dataset results.
1 parent 1878951 commit 1e59bf6

File tree

1 file changed

+16
-0
lines changed

1 file changed

+16
-0
lines changed
 

‎include/svs/index/vamana/iterator.h

+16
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,22 @@ class BatchIterator {
334334
) {
335335
auto search_closure =
336336
[&](const auto& query, const auto& accessor, auto& d, auto& buffer) {
337+
constexpr vamana::extensions::UsesReranking<
338+
std::remove_const_t<std::remove_reference_t<decltype(data)>>>
339+
uses_reranking{};
340+
if constexpr (uses_reranking()) {
341+
distance::maybe_fix_argument(d, query);
342+
// recompute search buffer using only primary dataset
343+
for (size_t j = 0, jmax = buffer.size(); j < jmax; ++j) {
344+
auto& neighbor = buffer[j];
345+
auto id = neighbor.id();
346+
auto new_distance =
347+
distance::compute(d, query, data.get_primary(id));
348+
neighbor.set_distance(new_distance);
349+
}
350+
buffer.sort();
351+
}
352+
337353
vamana::greedy_search(
338354
graph,
339355
data,

0 commit comments

Comments
 (0)
Please sign in to comment.