feat: add threshold filtering to knn search - #2531
Conversation
wxyucs
added
kind/feature
|
/label status/waiting-for-review |
|
Automated pull request review completed. Review effort: Submitted 4 inline comments. |
Merge Protections🟢 All 3 merge protections satisfied — ready to merge. Show 3 satisfied protections🟢 Require kind label
🟢 Require version label
🟢 Require linked issue for feature/bug PRs
|
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: medium (246 changed lines across 13 files).
Submitted 3 inline comments.
Reviewed commit 9b6cf96.
9b6cf96 to
e718e38
Compare
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (397 changed lines across 15 files).
Submitted 7 inline comments.
Reviewed commit e718e38.
e718e38 to
d05a524
Compare
LHT129
left a comment
LHT129
left a comment
There was a problem hiding this comment.
This update addresses all six issues raised in the previous review round:
- IVF reorder threshold — Threshold is now applied after exact-distance reorder, with
FilterDatasetByThresholdcapping attopk_. The candidate pool is retained at full size when threshold is present. - Non-finite threshold validation —
ValidateSearchThresholdrejects NaN/Inf inSearchWithRequestpaths, andindex_impl.hvalidates insideSAFE_CALLfor consistent error handling. - Allocator preservation —
FilterDatasetByThresholdaccepts anAllocator*, allocates through it, and setsOwner(true, allocator). HNSW and SINDI pass their caller allocator; a regression test verifies balanced deallocation. - route_buckets_only threshold — IVF now explicitly rejects threshold with
disable_bucket_scanviaCHECK_ARGUMENT. - IVF reasoning report —
AttachReasoningReportis called afterFilterDatasetByThresholdin the reorder path. - Pyramid/SIMQ threshold — Both now parse and apply threshold via
FilterDatasetByThresholdwith their allocators and k-capping.
Additional observations:
- The
#include <unordered_set>addition intest_brute_force.cppfixes a missing-include bug in the existing code (it was used but not included). ParseSearchThresholdnow validates JSON type viaIsNumber()before callingGetFloat(), preventing silent coercion of string values.- The
SetOFFFirstUsed()call in HGraph when threshold filtering empties the result is a correct fix to maintain iterator filter state consistency. FilterDatasetByThresholdcorrectly preserves statistics and reasoning from the input dataset.
One minor suggestion was posted inline about the two-pass iteration in FilterDatasetByThreshold.
0990cfd to
44a70f2
Compare
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (651 changed lines across 23 files).
Submitted 4 inline comments.
Reviewed commit 44a70f2.
LHT129
left a comment
LHT129
left a comment
There was a problem hiding this comment.
Review Summary
This PR adds inclusive top-level distance threshold filtering for KNN search across all maintained indexes (BruteForce, HGraph, IVF, Pyramid, SINDI, SIMQ). The implementation is thorough and well-tested with 628 additions across 23 files.
What was reviewed
- New
src/utils/search_threshold.hwithParseSearchThreshold,ValidateSearchThreshold,FilterDatasetByThreshold, andAllocateThresholdArray - Threshold integration in all 6 maintained indexes
- Non-finite distance filtering in
MemmoveHeap,BasicSearcher,ParallelSearcher, and HGraph - JSON parsing support (
JsonWrapper::IsNumber) SearchRequest::threshold_fieldIndexImplvalidation layer across allKnnSearchoverloads andSearchWithRequest- Documentation updates (EN + ZH)
- Comprehensive tests for BruteForce, HGraph, IVF, Pyramid, and SIMQ
Key observations
- The non-finite distance defense is layered at multiple levels (heap Push, searcher entry points, HGraph result filtering) — good defense-in-depth.
- The HGraph iterator path correctly consumes threshold-rejected candidates via
SetPointso pagination state is maintained. - IVF reorder path correctly retains the full candidate pool for exact-distance threshold filtering.
FilterDatasetByThresholdexception safety was improved from earlier rounds — allocations are now attached to the owning Dataset before subsequent allocations.
Minor suggestions
FilterDatasetByThresholditerates 3 times over the input (count, copy ids+dists, copy extra-infos). The second and third passes could be merged.- When
result_count == 0with extra info present, the filtered Dataset hasExtraInfoSize > 0butExtraInfos == nullptr. Downstream consumers should be null-safe on this pointer.
Overall
The implementation is solid, well-tested, and addresses all critical issues from previous review rounds. No blocking concerns.
44a70f2 to
c429a73
Compare
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (728 changed lines across 22 files).
Submitted 2 inline comments.
Reviewed commit c429a73.
c429a73 to
632fa56
Compare
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (787 changed lines across 22 files).
Submitted 4 inline comments.
Reviewed commit 632fa56.
LHT129
left a comment
LHT129
left a comment
There was a problem hiding this comment.
Reviewed the latest commit (632fa56). All previously flagged issues have been properly addressed:
- JSON type validation before GetFloat (IsNumber check)
- Route-only IVF explicitly rejects threshold with INVALID_ARGUMENT
- IVF reasoning report attached after threshold filtering
- Pyramid and SIMQ threshold filtering implemented with regressions
- IndexImpl validates threshold JSON before empty-index shortcuts
- Non-finite distance handling in searchers (basic and parallel)
The implementation is thorough with good test coverage across all six maintained indexes. No new issues found.
Signed-off-by: Xiangyu Wang <wxy407827@antgroup.com> Assisted-by: OpenAI Codex:GPT-5
632fa56 to
b35bbc3
Compare
| The result contains up to `k` neighbors sorted by ascending distance to the query. | ||
|
|
||
| ## Threshold Filtering | ||
|
|
There was a problem hiding this comment.
[note] The documentation references "This PR's threshold support" which ties the docs to a specific pull request. User-facing documentation should describe the feature itself without referencing the PR that added it. Consider rephrasing to something like "Threshold filtering is supported for the following maintained indexes: BruteForce, HGraph, IVF, Pyramid, SINDI, and SIMQ."
The same issue appears in the Chinese documentation.
|
|
||
| 搜索请求至多返回 k 个结果，这些结果按照最近邻和查询向量的距离升序排序。输出的结果类似于： | ||
|
|
||
| ## 按阈值过滤 |
There was a problem hiding this comment.
[note] Same issue as the English docs: "本 PR 的阈值支持范围仅包括" references the PR itself. Consider rephrasing to describe the feature directly, e.g. "阈值过滤支持以下仍在维护的索引：BruteForce、HGraph、IVF、Pyramid、SINDI 和 SIMQ。"
vsag-bot
left a comment
vsag-bot
left a comment
There was a problem hiding this comment.
Automated inline review completed.
Review effort: high (1177 changed lines across 27 files).
Submitted 4 inline comments.
Reviewed commit b35bbc3.
| for (const auto& record : finite_records) { | ||
| search_result->Push(record); | ||
| } | ||
| filter_search_result_by_threshold(search_result, request.threshold_, ctx.alloc); |
There was a problem hiding this comment.
[P2] Filter non-finite HGraph candidates before ef pruning
This filtering happens only after the default single-threaded graph search has bounded candidates to ef/topk. With IP vectors {FLT_MAX, 0}, query FLT_MAX, k = ef_search = 1, and threshold 1, the computed -inf candidate occupies the sole result slot and the qualifying distance-1 candidate is discarded, so this returns empty. Exclude non-finite candidates during bounded result selection while retaining them only as traversal seeds.
| for (uint32_t i = 0; i < count_no_visited; i++) { | ||
| dist = line_dists[i]; | ||
| const auto cur_id = to_be_visited_id[i]; | ||
| if (not std::isfinite(dist)) { |
There was a problem hiding this comment.
[P2] Keep NaN intermediate nodes as traversal seeds
This branch enqueues infinities for traversal but drops NaN nodes entirely. In a topology finite entry -> NaN-valued bridge -> finite target, the public HGraph iterator can no longer reach the target, even without a threshold; the previous path enqueued that bridge. Normalize NaN to a safe traversal priority while excluding it from the result heap.
| if (std::isinf(dist)) { | ||
| if constexpr (mode == KNN_SEARCH) { | ||
| if (not is_id_allowed || is_id_allowed->CheckValid(cur_id)) { | ||
| top_candidates->Push(dist, cur_id); |
There was a problem hiding this comment.
[P2] Maintain the ef bound when retaining infinite results
This push bypasses the top_candidates->Size() > ef eviction and lower-bound update used by the finite branch. Once more than ef connected nodes have overflowed L2 distances, the heap remains oversized, the equality-based termination condition is disabled, and iterator search can traverse the full graph while retaining O(N) candidates. Preserve the infinity as a traversal seed/result without violating the ef bound; the parallel implementation has the same pattern.
| for (const auto& [distance, _] : reranked) { | ||
| if (not threshold.has_value() or | ||
| (std::isfinite(distance) and distance <= threshold.value())) { | ||
| ++result_count; |
There was a problem hiding this comment.
[P2] Handle zero k before counting SIMQ results
For k == 0, result_count is incremented to 1 before the equality check, so it can never equal zero and the loop counts every reranked candidate. SIMQ therefore returns all candidates, whereas the previous min(k, size) path returned none. Validate that k is positive or retain an explicit zero-result fast path.
Summary
Add inclusive top-level distance threshold filtering for KNN in the maintained indexes only:
HNSW and DiskANN are deprecated and explicitly unsupported non-goals for this PR. Their source, test, and documentation paths are absent from the full PR diff.
Threshold filtering excludes non-finite values before bounded result selection and allows eligible finite candidates to backfill. Non-reorder IVF filters during bucket selection, while reordered IVF filters exact distances after reranking. HGraph iterator search consumes rejected pages and retained discard state internally. Ordinary KNN preserves infinite results and may use infinite intermediate graph nodes as traversal seeds. Threshold scratch storage uses the active request allocator.
Validation
clang-format-15, directclang-tidy-15, bounded debug build, andgit diff --checkpassed.Fixes: #2081