fix: correct multi-query top-k insertion in iterative filter (#51591)

`insert_helper` in `PhyIterativeFilterNode` indexed the per-query result
window by the **relative** count instead of the **absolute** window end
(`base_idx + count`). For `nq_index >= 1` the shift guard `count > pos`
is always false, so an out-of-order candidate (a closer distance
arriving after farther ones — which approximate vector iterators do
emit) overwrote an existing top-k entry instead of shifting it,
corrupting the top-k of **every query vector after the first** in a
multi-query filtered search.

Extracted the insertion into `Utils.h::topk_binsert` with the correct
absolute-window math, matching the already-correct math in
`PhyIterativeElementFilterNode::CollectResults`.

Deterministic regression test
`IterativeFilterInsert.MultiQueryWindowOrdering` inserts an out-of-order
stream into the second query window and asserts the correct sorted
result.

issue: #51385

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: xiaofanluan <xf@hjjaq.com>
Co-authored-by: xiaofanluan <xf@hjjaq.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
James
2026-07-20 22:32:42 +08:00
committed by GitHub
co-authored by xiaofanluan Claude Opus 4.8
parent b66f55fe67
commit f6bffd8de5
3 changed files with 122 additions and 69 deletions
@@ -88,47 +88,6 @@ PhyIterativeFilterNode::IsFinished() {
return is_finished_;
}
inline void
insert_helper(milvus::SearchResult& search_result,
int& topk,
const bool large_is_better,
const FixedVector<float>& distances,
const int64_t nq_index,
const int64_t unity_topk,
const int i,
int64_t doc_id,
std::optional<int32_t> elem_idx) {
auto pos = large_is_better
? find_binsert_position<true>(search_result.distances_,
nq_index * unity_topk,
nq_index * unity_topk + topk,
distances[i])
: find_binsert_position<false>(search_result.distances_,
nq_index * unity_topk,
nq_index * unity_topk + topk,
distances[i]);
if (topk > pos) {
std::memmove(&search_result.distances_[pos + 1],
&search_result.distances_[pos],
(topk - pos) * sizeof(float));
std::memmove(&search_result.seg_offsets_[pos + 1],
&search_result.seg_offsets_[pos],
(topk - pos) * sizeof(int64_t));
if (elem_idx.has_value()) {
std::memmove(&search_result.element_indices_[pos + 1],
&search_result.element_indices_[pos],
(topk - pos) * sizeof(int32_t));
}
}
search_result.seg_offsets_[pos] = doc_id;
if (elem_idx.has_value()) {
search_result.element_indices_[pos] = elem_idx.value();
}
search_result.distances_[pos] = distances[i];
++topk;
}
RowVectorPtr
PhyIterativeFilterNode::GetOutput() {
milvus::exec::checkCancellation(query_context_);
@@ -231,7 +190,7 @@ PhyIterativeFilterNode::GetOutput() {
for (auto& iterator : search_result.vector_iterators_.value()) {
EvalCtx eval_ctx(operator_context_->get_exec_context());
int topk = 0;
int64_t topk = 0;
while (iterator->HasNext() && topk < unity_topk) {
offsets.clear();
distances.clear();
@@ -319,15 +278,13 @@ PhyIterativeFilterNode::GetOutput() {
for (size_t i = 0; i < offsets.size(); ++i) {
auto [doc_id, elem_idx] = element_to_doc_mapping[i];
if (doc_eval_cache[doc_id]) {
insert_helper(search_result,
topk,
large_is_better,
distances,
nq_index,
unity_topk,
i,
doc_id,
elem_idx);
topk_binsert(search_result,
nq_index * unity_topk,
topk,
large_is_better,
distances[i],
doc_id,
elem_idx);
if (topk == unity_topk) {
break;
}
@@ -338,15 +295,13 @@ PhyIterativeFilterNode::GetOutput() {
Assert(bitsetview.size() == offsets.size());
for (auto i = 0; i < offsets.size(); ++i) {
if (bitsetview[i]) {
insert_helper(search_result,
topk,
large_is_better,
distances,
nq_index,
unity_topk,
i,
offsets[i],
std::nullopt);
topk_binsert(search_result,
nq_index * unity_topk,
topk,
large_is_better,
distances[i],
offsets[i],
std::nullopt);
if (topk == unity_topk) {
break;
}
@@ -357,15 +312,13 @@ PhyIterativeFilterNode::GetOutput() {
Assert(!element_level);
for (auto i = 0; i < offsets.size(); ++i) {
if (bitset[offsets[i]] > 0) {
insert_helper(search_result,
topk,
large_is_better,
distances,
nq_index,
unity_topk,
i,
offsets[i],
std::nullopt);
topk_binsert(search_result,
nq_index * unity_topk,
topk,
large_is_better,
distances[i],
offsets[i],
std::nullopt);
if (topk == unity_topk) {
break;
}
+54
View File
@@ -17,6 +17,8 @@
#pragma once
#include <cstddef>
#include <cstring>
#include <optional>
#include "common/OpContext.h"
#include "common/QueryInfo.h"
#include "common/QueryResult.h"
@@ -47,6 +49,58 @@ find_binsert_position(const std::vector<float>& distances,
return static_cast<size_t>(it - distances.begin());
}
// Insert one (distance, seg_offset[, elem_idx]) into the per-query result
// window [base_idx, base_idx + count) of search_result, keeping that window
// sorted best-first for the metric. Advances count. elem_idx == std::nullopt
// means non-element-level (element_indices_ is left untouched).
//
// Single owner of the top-k insertion invariant, shared by the iterative
// filter and iterative element filter nodes so the two cannot drift apart.
inline void
topk_binsert(SearchResult& search_result,
size_t base_idx,
int64_t& count,
bool large_is_better,
float distance,
int64_t seg_offset,
std::optional<int32_t> elem_idx) {
size_t pos = large_is_better
? find_binsert_position<true>(search_result.distances_,
base_idx,
base_idx + count,
distance)
: find_binsert_position<false>(search_result.distances_,
base_idx,
base_idx + count,
distance);
// The window is [base_idx, base_idx + count); shift only when inserting
// before its current end. Guard/size with the ABSOLUTE window end
// (base_idx + count), not the relative count — otherwise for nq_index >= 1
// (base_idx > 0) the shift is wrongly skipped and out-of-order insertions
// overwrite existing top-k entries.
if (count > 0 && pos < base_idx + count) {
size_t n = base_idx + count - pos;
std::memmove(&search_result.distances_[pos + 1],
&search_result.distances_[pos],
n * sizeof(float));
std::memmove(&search_result.seg_offsets_[pos + 1],
&search_result.seg_offsets_[pos],
n * sizeof(int64_t));
if (elem_idx.has_value()) {
std::memmove(&search_result.element_indices_[pos + 1],
&search_result.element_indices_[pos],
n * sizeof(int32_t));
}
}
search_result.seg_offsets_[pos] = seg_offset;
if (elem_idx.has_value()) {
search_result.element_indices_[pos] = elem_idx.value();
}
search_result.distances_[pos] = distance;
++count;
}
[[maybe_unused]] static bool
UseVectorIterator(const SearchInfo& search_info) {
return search_info.has_group_by() || search_info.iterative_filter_execution;
@@ -29,6 +29,7 @@
#include "common/TracerBase.h"
#include "common/Types.h"
#include "common/protobuf_utils.h"
#include "exec/operator/Utils.h"
#include "filemanager/InputStream.h"
#include "gtest/gtest.h"
#include "index/Index.h"
@@ -432,3 +433,48 @@ TEST(IterativeFilter, GrowingIndex) {
*search_result, *search_result2, topK, num_queries);
}
}
// Directly exercises the shared top-k insertion helper across multiple query
// windows. The per-query window starts at base_idx = nq_index * unity_topk;
// an out-of-order insertion (a better distance arriving after worse ones) must
// shift the entries already in that window. The pre-fix math guarded/sized the
// shift with the relative count instead of the absolute window end, so the
// shift was skipped for every query beyond the first (nq_index >= 1),
// corrupting their results. nq_index == 0 stayed correct, which is why the
// nq == 1 unit tests above never caught it.
TEST(IterativeFilterInsert, MultiQueryWindowOrdering) {
const int64_t unity_topk = 4;
const int64_t nq = 2;
const bool large_is_better = false; // L2: smaller distance ranks first
SearchResult sr;
sr.distances_.assign(nq * unity_topk, 0.0f);
sr.seg_offsets_.assign(nq * unity_topk, -1);
// Insert stream (distance, seg_offset): 1/100, 3/300, 2/200, 4/400.
// The 2/200 arrives after 3/300, so it must be shifted into the middle.
const std::vector<std::pair<float, int64_t>> stream = {
{1.0f, 100}, {3.0f, 300}, {2.0f, 200}, {4.0f, 400}};
const std::vector<int64_t> expected_offsets = {100, 200, 300, 400};
for (int64_t q = 0; q < nq; ++q) {
int64_t count = 0;
for (const auto& [dist, off] : stream) {
milvus::exec::topk_binsert(sr,
q * unity_topk,
count,
large_is_better,
dist,
off,
std::nullopt);
}
ASSERT_EQ(count, unity_topk);
for (int64_t k = 0; k < unity_topk; ++k) {
ASSERT_EQ(sr.seg_offsets_[q * unity_topk + k], expected_offsets[k])
<< "nq_index=" << q << " k=" << k;
ASSERT_FLOAT_EQ(sr.distances_[q * unity_topk + k],
static_cast<float>(k + 1))
<< "nq_index=" << q << " k=" << k;
}
}
}