mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
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:
co-authored by
xiaofanluan
Claude Opus 4.8
parent
b66f55fe67
commit
f6bffd8de5
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user