diff --git a/docs/docs/en/src/guide/knn_search.md b/docs/docs/en/src/guide/knn_search.md index 5bfd79d422..eabb3abc01 100644 --- a/docs/docs/en/src/guide/knn_search.md +++ b/docs/docs/en/src/guide/knn_search.md @@ -87,3 +87,31 @@ for (int64_t i = 0; i < result->GetDim(); ++i) { ``` The result contains up to `k` neighbors sorted by ascending distance to the query. + +## Threshold Filtering + +This PR's threshold support is limited to the maintained indexes: BruteForce, HGraph, IVF, +Pyramid, SINDI, and SIMQ. HNSW and DiskANN are deprecated and are explicit unsupported +non-goals; this option does not add or change behavior for either legacy index. + +KNN search accepts an optional top-level `threshold` in the JSON search parameters: + +```json +{ + "threshold": 4.0, + "hgraph": { "ef_search": 64 } +} +``` + +When present, the result contains at most `k` neighbors and every returned distance is less than +or equal to `threshold`. The boundary is inclusive, results remain sorted by ascending distance, +and fewer than `k` results—or no results—may be returned. When omitted, KNN behavior is unchanged. +The value uses the index metric's returned distance: squared L2 for `l2`, `1 - inner_product` for +`ip`, and `1 - cosine_similarity` for `cosine`. Therefore negative thresholds can be meaningful +for `ip`; the threshold must be a finite JSON number. + +The same option is available through `SearchRequest::threshold_` on BruteForce, HGraph, and IVF; +Pyramid, SINDI, and SIMQ support the option through their KNN JSON parameters. These are the six +maintained indexes supported by this feature. It only limits KNN results; use `RangeSearch` and +its `radius` argument for range-search requests. Existing index-specific search parameters can be +supplied alongside `threshold`. diff --git a/docs/docs/zh/src/guide/knn_search.md b/docs/docs/zh/src/guide/knn_search.md index ebcf03f6bd..f6dc7e2009 100644 --- a/docs/docs/zh/src/guide/knn_search.md +++ b/docs/docs/zh/src/guide/knn_search.md @@ -92,6 +92,30 @@ BruteForce 索引支持用 Build 和 Add 方法写入数据，这里我们用 Ad 搜索请求至多返回 k 个结果，这些结果按照最近邻和查询向量的距离升序排序。输出的结果类似于： +## 按阈值过滤 + +本 PR 的阈值支持范围仅包括仍在维护的索引：BruteForce、HGraph、IVF、Pyramid、SINDI 和 +SIMQ。HNSW 和 DiskANN 已弃用，是明确不在范围内的目标；本选项不会为这两个旧索引新增或修改行为。 + +KNN 搜索参数支持在 JSON 顶层设置可选的 `threshold`： + +```json +{ + "threshold": 4.0, + "hgraph": { "ef_search": 64 } +} +``` + +设置后，结果至多返回 `k` 个邻居，并且每个返回距离都小于或等于 `threshold`。 +边界值会被包含，结果仍按距离升序排列，因此可能返回少于 `k` 个结果，甚至没有结果。 +不设置该字段时，KNN 行为保持不变。阈值使用索引返回的距离语义：`l2` 为 L2 平方距离， +`ip` 为 `1 - inner_product`，`cosine` 为 `1 - cosine_similarity`。因此 `ip` 可以使用负阈值； +阈值必须是有限的 JSON 数字。 + +通过统一请求接口时，BruteForce、HGraph 和 IVF 支持使用 `SearchRequest::threshold_`； +Pyramid、SINDI 和 SIMQ 通过 KNN JSON 参数支持该选项。这六种索引是本功能支持的仍在维护的索引。 +该选项只限制 KNN 结果；范围搜索请使用 `RangeSearch` 及其 `radius` 参数。`threshold` 可以和现有索引专用搜索参数同时传入。 + ```bash results: 6519: 13.855 @@ -105,4 +129,3 @@ results: 8703: 16.1161 5583: 16.1256 ``` - diff --git a/include/vsag/search_request.h b/include/vsag/search_request.h index 705cd3b99d..7cfed34ee0 100644 --- a/include/vsag/search_request.h +++ b/include/vsag/search_request.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include @@ -183,6 +184,14 @@ class SearchRequest { * Default is empty (no reasoning enabled). */ std::vector expected_labels_{}; + + /** + * @brief Optional inclusive distance threshold for KNN search mode + * @details When set, at most topk_ results are returned and every returned + * distance is less than or equal to this threshold. An unset value + * preserves the default KNN behavior. + */ + std::optional threshold_{std::nullopt}; }; } // namespace vsag diff --git a/src/algorithm/bruteforce/bruteforce.cpp b/src/algorithm/bruteforce/bruteforce.cpp index e81235cfb1..f4926cea36 100644 --- a/src/algorithm/bruteforce/bruteforce.cpp +++ b/src/algorithm/bruteforce/bruteforce.cpp @@ -15,6 +15,7 @@ #include "bruteforce.h" #include +#include #include #include #include @@ -36,6 +37,7 @@ #include "storage/serialization_tags.h" #include "storage/tlv_section.h" #include "typing.h" +#include "utils/search_threshold.h" #include "utils/slow_task_timer.h" #include "utils/util_functions.h" namespace vsag { @@ -360,6 +362,7 @@ BruteForce::KnnSearch(const DatasetPtr& query, req.query_ = query; req.topk_ = k; req.params_str_ = parameters; + req.threshold_ = ParseSearchThreshold(parameters); if (filter != nullptr) { req.filter_ = filter; } @@ -368,6 +371,7 @@ BruteForce::KnnSearch(const DatasetPtr& query, DatasetPtr BruteForce::SearchWithRequest(const SearchRequest& request) const { + ValidateSearchThreshold(request.threshold_); std::shared_lock read_lock(this->global_mutex_); auto computer = this->make_search_computer(request.query_); @@ -470,7 +474,9 @@ BruteForce::SearchWithRequest(const SearchRequest& request) const { if (is_range and dist > radius) { continue; } - cur_heap->Push(dist, i); + if (not request.threshold_.has_value() || std::isfinite(dist)) { + cur_heap->Push(dist, i); + } } else { if (reasoning != nullptr) { reasoning->RecordFilterReject(i); @@ -506,6 +512,13 @@ BruteForce::SearchWithRequest(const SearchRequest& request) const { } } + if (not is_range) { + filter_search_result_by_threshold( + heap, + request.threshold_, + select_query_allocator(request.search_allocator_, this->allocator_)); + } + // Collect result inner IDs before pack_knn_result_with_extra_info consumes the heap, // so we can call MarkResult for reasoning analysis. Vector result_inner_ids(this->allocator_); diff --git a/src/algorithm/hgraph/hgraph_search.cpp b/src/algorithm/hgraph/hgraph_search.cpp index cc3f117c78..16f36fa047 100644 --- a/src/algorithm/hgraph/hgraph_search.cpp +++ b/src/algorithm/hgraph/hgraph_search.cpp @@ -14,6 +14,8 @@ #include +#include + #include "attr/argparse.h" #include "dataset_impl.h" #include "hgraph.h" // IWYU pragma: keep @@ -21,6 +23,7 @@ #include "impl/filter/iterator_filter.h" #include "impl/heap/standard_heap.h" #include "impl/reasoning/search_reasoning.h" +#include "utils/search_threshold.h" #include "utils/util_functions.h" namespace vsag { @@ -52,6 +55,7 @@ HGraph::KnnSearch(const DatasetPtr& query, req.topk_ = k; req.filter_ = filter; req.params_str_ = parameters; + req.threshold_ = ParseSearchThreshold(parameters); req.search_allocator_ = allocator; return this->SearchWithRequest(req); } @@ -76,6 +80,7 @@ HGraph::KnnSearch(const DatasetPtr& query, this->validate_knn_args(query, k); auto params = HGraphSearchParameters::FromJson(parameters); + const auto threshold = ParseSearchThreshold(parameters); ctx.rabitq_error_rate = params.rabitq_error_rate; CHECK_ARGUMENT( // NOLINT params.ef_search >= 1, @@ -109,7 +114,6 @@ HGraph::KnnSearch(const DatasetPtr& query, } auto* iter_filter_ctx = static_cast(iter_ctx); - auto search_result = DistanceHeap::MakeInstanceBySize(ctx.alloc, k); const auto* query_data = get_data(query); // Note: brute_force_threshold is intentionally not applied here. The // iterator KnnSearch API pages results across multiple calls via @@ -117,106 +121,133 @@ HGraph::KnnSearch(const DatasetPtr& query, // that pagination state itself or be wasted on subsequent calls. The // non-iterator KnnSearch overload (which delegates to SearchWithRequest) // still benefits from the brute-force fallback. - if (is_last_filter) { - while (!iter_filter_ctx->Empty()) { - uint32_t cur_inner_id = iter_filter_ctx->GetTopID(); - float cur_dist = iter_filter_ctx->GetTopDist(); - search_result->Push(cur_dist, cur_inner_id); - iter_filter_ctx->PopDiscard(); - } - } else { - InnerSearchParam search_param; - search_param.ep = this->entry_point_id_; - search_param.topk = 1; - search_param.ef = 1; - search_param.is_inner_id_allowed = nullptr; - search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search; - if (search_param.ep == INVALID_ENTRY_POINT) { - return make_empty_dataset_with_stats(); - } - if (iter_filter_ctx->IsFirstUsed()) { - for (auto i = static_cast(this->route_graphs_.size() - 1); i >= 0; --i) { - auto result = this->search_one_graph(query_data, - this->route_graphs_[i], - this->basic_flatten_codes_, - search_param, - (VisitedListPtr) nullptr, - &ctx); - search_param.ep = result->Top().second; + while (true) { + auto search_result = DistanceHeap::MakeInstanceBySize(ctx.alloc, k); + if (is_last_filter) { + while (!iter_filter_ctx->Empty()) { + uint32_t cur_inner_id = iter_filter_ctx->GetTopID(); + float cur_dist = iter_filter_ctx->GetTopDist(); + search_result->Push(cur_dist, cur_inner_id); + iter_filter_ctx->PopDiscard(); + } + } else { + InnerSearchParam search_param; + search_param.ep = this->entry_point_id_; + search_param.topk = 1; + search_param.ef = 1; + search_param.is_inner_id_allowed = nullptr; + search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search; + if (search_param.ep == INVALID_ENTRY_POINT) { + return make_empty_dataset_with_stats(); + } + if (iter_filter_ctx->IsFirstUsed()) { + for (auto i = static_cast(this->route_graphs_.size() - 1); i >= 0; --i) { + auto result = this->search_one_graph(query_data, + this->route_graphs_[i], + this->basic_flatten_codes_, + search_param, + (VisitedListPtr) nullptr, + &ctx); + if (result->Empty()) { + return make_empty_dataset_with_stats(); + } + search_param.ep = result->Top().second; + } } - } - - search_param.ef = std::max(params.ef_search, k); - search_param.is_inner_id_allowed = ft; - search_param.topk = static_cast(search_param.ef); - search_param.parallel_search_thread_count = params.parallel_search_thread_count; - search_param.enable_reorder = params.enable_reorder; - search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search; - search_param.skip_ratio = params.skip_ratio; - search_param.skip_strategy_type = params.skip_strategy_type; - - DistanceRecordVector rabitq_lower_bound_candidates(ctx.alloc); - auto* rabitq_lower_bound_candidates_ptr = - search_param.enable_rabitq_one_bit_search and use_reorder_ and - search_param.enable_reorder and reorder_by_base_ - ? &rabitq_lower_bound_candidates - : nullptr; - - search_result = this->search_one_graph(query_data, - this->bottom_graph_, - this->basic_flatten_codes_, - search_param, - iter_filter_ctx, - &ctx, - rabitq_lower_bound_candidates_ptr); - if (use_reorder_ and search_param.enable_reorder) { - this->reorder(query_data, - this->get_reorder_codes(), - search_result, - k, - iter_filter_ctx, - ctx, - rabitq_lower_bound_candidates_ptr); - } else if (search_param.enable_reorder and params.rabitq_one_bit_search) { - this->reorder( - query_data, this->basic_flatten_codes_, search_result, k, iter_filter_ctx, ctx); + search_param.ef = std::max(params.ef_search, k); + search_param.is_inner_id_allowed = ft; + search_param.topk = static_cast(search_param.ef); + search_param.parallel_search_thread_count = params.parallel_search_thread_count; + search_param.enable_reorder = params.enable_reorder; + search_param.enable_rabitq_one_bit_search = params.rabitq_one_bit_search; + search_param.skip_ratio = params.skip_ratio; + search_param.skip_strategy_type = params.skip_strategy_type; + + DistanceRecordVector rabitq_lower_bound_candidates(ctx.alloc); + auto* rabitq_lower_bound_candidates_ptr = + search_param.enable_rabitq_one_bit_search and use_reorder_ and + search_param.enable_reorder and reorder_by_base_ + ? &rabitq_lower_bound_candidates + : nullptr; + + search_result = this->search_one_graph(query_data, + this->bottom_graph_, + this->basic_flatten_codes_, + search_param, + iter_filter_ctx, + &ctx, + rabitq_lower_bound_candidates_ptr); + + if (use_reorder_ and search_param.enable_reorder) { + this->reorder(query_data, + this->get_reorder_codes(), + search_result, + k, + iter_filter_ctx, + ctx, + rabitq_lower_bound_candidates_ptr); + } else if (search_param.enable_reorder and params.rabitq_one_bit_search) { + this->reorder( + query_data, this->basic_flatten_codes_, search_result, k, iter_filter_ctx, ctx); + } } - } - while (search_result->Size() > k) { - auto curr = search_result->Top(); - iter_filter_ctx->AddDiscardNode(curr.first, curr.second); - search_result->Pop(); - } + if (threshold.has_value()) { + DistanceRecordVector valid_records(ctx.alloc); + valid_records.reserve(search_result->Size()); + while (not search_result->Empty()) { + const auto record = search_result->Top(); + search_result->Pop(); + if (std::isfinite(record.first) and record.first <= threshold.value()) { + valid_records.push_back(record); + } else { + iter_filter_ctx->SetPoint(record.second); + } + } + for (const auto& record : valid_records) { + search_result->Push(record); + } + } + while (search_result->Size() > k) { + auto curr = search_result->Top(); + iter_filter_ctx->AddDiscardNode(curr.first, curr.second); + search_result->Pop(); + } - // return an empty dataset directly if searcher returns nothing - if (search_result->Empty()) { - return DatasetImpl::MakeEmptyDataset(); - } - auto count = static_cast(search_result->Size()); - auto [dataset_results, dists, ids] = create_fast_dataset(count, ctx.alloc); - char* extra_infos = nullptr; - if (extra_info_size_ > 0) { - extra_infos = - static_cast(ctx.alloc->Allocate(extra_info_size_ * search_result->Size())); - dataset_results->ExtraInfos(extra_infos) - ->ExtraInfoSize(static_cast(extra_info_size_)); - } - for (int64_t j = count - 1; j >= 0; --j) { - dists[j] = search_result->Top().first; - ids[j] = this->label_table_->GetLabelById(search_result->Top().second); - iter_filter_ctx->SetPoint(search_result->Top().second); - if (extra_infos != nullptr) { - this->extra_infos_->GetExtraInfoById(search_result->Top().second, - extra_infos + extra_info_size_ * j); + // An empty page is terminal to iterator callers, so consume retained traversal state + // internally until an eligible result is found or the discard heap is exhausted. + if (search_result->Empty()) { + iter_filter_ctx->SetOFFFirstUsed(); + if (not iter_filter_ctx->Empty()) { + continue; + } + return DatasetImpl::MakeEmptyDataset(); } - search_result->Pop(); - } - iter_filter_ctx->SetOFFFirstUsed(); + auto count = static_cast(search_result->Size()); + auto [dataset_results, dists, ids] = create_fast_dataset(count, ctx.alloc); + char* extra_infos = nullptr; + if (extra_info_size_ > 0) { + extra_infos = + static_cast(ctx.alloc->Allocate(extra_info_size_ * search_result->Size())); + dataset_results->ExtraInfos(extra_infos) + ->ExtraInfoSize(static_cast(extra_info_size_)); + } + for (int64_t j = count - 1; j >= 0; --j) { + dists[j] = search_result->Top().first; + ids[j] = this->label_table_->GetLabelById(search_result->Top().second); + iter_filter_ctx->SetPoint(search_result->Top().second); + if (extra_infos != nullptr) { + this->extra_infos_->GetExtraInfoById(search_result->Top().second, + extra_infos + extra_info_size_ * j); + } + search_result->Pop(); + } + iter_filter_ctx->SetOFFFirstUsed(); - dataset_results->Statistics(stats.Dump()); - return std::move(dataset_results); + dataset_results->Statistics(stats.Dump()); + return std::move(dataset_results); + } } template @@ -370,6 +401,7 @@ HGraph::RangeSearch(const DatasetPtr& query, [[nodiscard]] DatasetPtr HGraph::SearchWithRequest(const SearchRequest& request) const { + ValidateSearchThreshold(request.threshold_); SearchStatistics stats; QueryContext ctx{.alloc = this->allocator_, .stats = &stats}; if (request.search_allocator_ != nullptr) { @@ -456,6 +488,10 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { for (auto i = static_cast(this->route_graphs_.size() - 1); i >= 0; --i) { auto result = this->search_one_graph( raw_query, this->route_graphs_[i], this->basic_flatten_codes_, search_param, vt, &ctx); + if (result->Empty()) { + this->pool_->ReturnOne(vt); + return make_empty_dataset_with_stats(); + } search_param.ep = result->Top().second; } @@ -580,7 +616,32 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { return result; } - // KNN mode: trim by k + // Threshold KNN mode excludes unordered distances; ordinary KNN preserves its legacy + // non-finite-distance behavior while the searcher uses an ordered infinity sentinel. + bool query_is_finite = true; + if (request.query_ != nullptr and request.query_->GetFloat32Vectors() != nullptr) { + const auto query_size = request.query_->GetNumElements() * request.query_->GetDim(); + for (int64_t i = 0; i < query_size; ++i) { + if (not std::isfinite(request.query_->GetFloat32Vectors()[i])) { + query_is_finite = false; + break; + } + } + } + DistanceRecordVector finite_records(ctx.alloc); + finite_records.reserve(search_result->Size()); + while (not search_result->Empty()) { + const auto record = search_result->Top(); + search_result->Pop(); + if ((not request.threshold_.has_value() and query_is_finite) || + std::isfinite(record.first)) { + finite_records.push_back(record); + } + } + for (const auto& record : finite_records) { + search_result->Push(record); + } + filter_search_result_by_threshold(search_result, request.threshold_, ctx.alloc); while (search_result->Size() > static_cast(k)) { search_result->Pop(); } diff --git a/src/algorithm/inner_index_interface.cpp b/src/algorithm/inner_index_interface.cpp index 6c3daf2e75..8c47497480 100644 --- a/src/algorithm/inner_index_interface.cpp +++ b/src/algorithm/inner_index_interface.cpp @@ -1130,6 +1130,27 @@ InnerIndexInterface::pack_knn_result(DistHeapPtr& heap, Allocator* allocator) co return std::move(dataset_results); } +void +InnerIndexInterface::filter_search_result_by_threshold(DistHeapPtr& result, + const std::optional& threshold, + Allocator* allocator) { + if (not threshold.has_value() or result == nullptr) { + return; + } + DistanceRecordVector valid_records(allocator); + valid_records.reserve(result->Size()); + while (not result->Empty()) { + const auto record = result->Top(); + result->Pop(); + if (std::isfinite(record.first) and record.first <= threshold.value()) { + valid_records.push_back(record); + } + } + for (const auto& record : valid_records) { + result->Push(record); + } +} + DatasetPtr InnerIndexInterface::pack_knn_result_with_extra_info(DistHeapPtr& heap, Allocator* allocator) const { diff --git a/src/algorithm/inner_index_interface.h b/src/algorithm/inner_index_interface.h index a9dcc4ba15..77afd87801 100644 --- a/src/algorithm/inner_index_interface.h +++ b/src/algorithm/inner_index_interface.h @@ -653,6 +653,11 @@ class InnerIndexInterface { void validate_range_args(const DatasetPtr& query, float radius, int64_t limited_size) const; + static void + filter_search_result_by_threshold(DistHeapPtr& result, + const std::optional& threshold, + Allocator* allocator); + public: LabelTablePtr label_table_{nullptr}; mutable std::shared_mutex label_lookup_mutex_{}; // lock for label_lookup_ & labels_ diff --git a/src/algorithm/ivf/flat_bucket_searcher.cpp b/src/algorithm/ivf/flat_bucket_searcher.cpp index 80d41bd106..45174487a5 100644 --- a/src/algorithm/ivf/flat_bucket_searcher.cpp +++ b/src/algorithm/ivf/flat_bucket_searcher.cpp @@ -14,6 +14,7 @@ #include "flat_bucket_searcher.h" +#include #include #include "attr/executor/executor.h" @@ -67,6 +68,10 @@ FlatBucketSearcher::Search(BucketIdType bucket_id, if (reasoning_ctx != nullptr) { reasoning_ctx->RecordVisit(origin_id, dist[j], 0); } + if (param.distance_threshold.has_value() and + (not std::isfinite(dist[j]) or dist[j] > param.distance_threshold.value())) { + continue; + } if (attr_ft != nullptr and not attr_ft->CheckValid(j)) { if (reasoning_ctx != nullptr) { reasoning_ctx->RecordFilterReject(origin_id); diff --git a/src/algorithm/ivf/ivf.cpp b/src/algorithm/ivf/ivf.cpp index 313e7daeba..83d1068d3d 100644 --- a/src/algorithm/ivf/ivf.cpp +++ b/src/algorithm/ivf/ivf.cpp @@ -43,6 +43,7 @@ #include "storage/stream_reader.h" #include "storage/stream_writer.h" #include "storage/tlv_section.h" +#include "utils/search_threshold.h" #include "utils/util_functions.h" #include "vsag_exception.h" @@ -566,6 +567,7 @@ IVF::KnnSearch(const DatasetPtr& query, req.query_ = query; req.topk_ = k; req.params_str_ = parameters; + req.threshold_ = ParseSearchThreshold(parameters); if (filter != nullptr) { req.filter_ = filter; } @@ -1085,8 +1087,6 @@ IVF::reorder(int64_t topk, auto reorder_heap = reorder_->Reorder(input, query, topk, ctx); auto dataset_results = this->pack_knn_result(reorder_heap, ctx.alloc); - this->AttachReasoningReport(dataset_results, reasoning_ctx); - return dataset_results; } @@ -1283,6 +1283,7 @@ IVF::check_merge_illegal(const vsag::MergeUnit& unit) const { DatasetPtr IVF::SearchWithRequest(const SearchRequest& request) const { + ValidateSearchThreshold(request.threshold_); SearchStatistics stats; QueryContext ctx{.alloc = request.search_allocator_, .stats = &stats}; @@ -1298,6 +1299,8 @@ IVF::SearchWithRequest(const SearchRequest& request) const { CHECK_ARGUMENT(query->GetFloat32Vectors() != nullptr, "query float32 vectors cannot be null"); CHECK_ARGUMENT(query->GetDim() == this->dim_, "query dimension must match index dimension"); + CHECK_ARGUMENT(not request.threshold_.has_value(), + "threshold filtering is not supported with disable_bucket_scan"); auto result = this->route_buckets_only(query, param, ctx); result->Statistics(stats.Dump()); return result; @@ -1368,6 +1371,7 @@ IVF::SearchWithRequest(const SearchRequest& request) const { auto result = reorder( k, search_result, query->GetFloat32Vectors(), param, ctx, reasoning_ctx.get()); result->Statistics(stats.Dump()); + this->AttachReasoningReport(result, reasoning_ctx.get()); return result; } auto dataset_results = this->pack_knn_result(search_result, ctx.alloc); @@ -1384,18 +1388,29 @@ IVF::SearchWithRequest(const SearchRequest& request) const { param.factor > 0.0F, fmt::format("factor must be positive when use_reorder is true, got {}", param.factor)); param.topk = static_cast(param.factor * static_cast(request.topk_)); + if (request.threshold_.has_value()) { + param.topk = std::max(param.topk, request.topk_); + } + } + const bool reorder_enabled = use_reorder_ and param.enable_reorder; + if (not reorder_enabled) { + param.distance_threshold = request.threshold_; } auto search_result = this->search(query, param, ctx, reasoning_ctx.get()); - if (use_reorder_ and param.enable_reorder) { - auto result = reorder(request.topk_, + if (reorder_enabled) { + auto result = reorder(request.threshold_.has_value() ? param.topk : request.topk_, search_result, query->GetFloat32Vectors(), param, ctx, reasoning_ctx.get()); + result = FilterDatasetByThreshold(result, request.threshold_, ctx.alloc, request.topk_); + AttachReasoningReport(result, reasoning_ctx.get()); result->Statistics(stats.Dump()); return result; } + filter_search_result_by_threshold( + search_result, request.threshold_, select_query_allocator(ctx.alloc, this->allocator_)); if (search_result == nullptr || search_result->Empty()) { auto dataset_results = DatasetImpl::MakeEmptyDataset(); this->AttachReasoningReport(dataset_results, reasoning_ctx.get()); @@ -1417,7 +1432,7 @@ IVF::AttachReasoningReport(const DatasetPtr& dataset_results, if (reasoning_ctx == nullptr) { return; } - auto count = dataset_results->GetNumElements(); + auto count = dataset_results->GetDim(); if (count > 0 and dataset_results->GetIds() != nullptr) { Vector result_inner_ids(static_cast(count), this->allocator_); { diff --git a/src/algorithm/pyramid/pyramid.cpp b/src/algorithm/pyramid/pyramid.cpp index f1324b93d9..49b178a43e 100644 --- a/src/algorithm/pyramid/pyramid.cpp +++ b/src/algorithm/pyramid/pyramid.cpp @@ -29,6 +29,7 @@ #include "storage/serialization.h" #include "storage/serialization_tags.h" #include "storage/tlv_section.h" +#include "utils/search_threshold.h" #include "utils/slow_task_timer.h" #include "utils/util_functions.h" namespace vsag { @@ -260,6 +261,7 @@ Pyramid::KnnSearch(const DatasetPtr& query, SearchStatistics stats; QueryContext ctx{.stats = &stats}; + const auto threshold = ParseSearchThreshold(parameters); auto parsed_param = PyramidSearchParameters::FromJson(parameters); CHECK_ARGUMENT(k > 0, fmt::format("k({}) must be greater than 0", k)); CHECK_ARGUMENT(parsed_param.hierarchy_op == PyramidSearchParameters::HierarchyOp::SINGLE, @@ -273,7 +275,7 @@ Pyramid::KnnSearch(const DatasetPtr& query, InnerSearchParam search_param; search_param.ef = std::max(parsed_param.ef_search, static_cast(k)); search_param.radius = std::numeric_limits::max(); - search_param.topk = k; + search_param.topk = threshold.has_value() ? static_cast(search_param.ef) : k; search_param.search_mode = KNN_SEARCH; search_param.parallel_search_thread_count = parsed_param.parallel_search_thread_count; if (this->support_duplicate_) { @@ -295,7 +297,7 @@ Pyramid::KnnSearch(const DatasetPtr& query, parsed_param.hierarchies.empty() ? "" : parsed_param.hierarchies[0]; auto result = this->search_impl(query, search_func, search_param, ctx, hierarchy_name); result->Statistics(stats.Dump()); - return result; + return FilterDatasetByThreshold(result, threshold, allocator_, k); } DatasetPtr diff --git a/src/algorithm/simq/simq.cpp b/src/algorithm/simq/simq.cpp index 35b81e88d8..eddc75229c 100644 --- a/src/algorithm/simq/simq.cpp +++ b/src/algorithm/simq/simq.cpp @@ -35,6 +35,7 @@ #include "storage/stream_reader.h" #include "storage/stream_writer.h" #include "typing.h" +#include "utils/search_threshold.h" #include "utils/util_functions.h" namespace vsag { @@ -716,6 +717,7 @@ SIMQ::KnnSearch(const DatasetPtr& query, const FilterPtr& filter) const { std::shared_lock lock(global_mutex_); SearchStatistics stats; + const auto threshold = ParseSearchThreshold(parameters); if (total_count_ == 0 || rep_hgraph_ == nullptr) { auto result = Dataset::Make(); @@ -780,11 +782,28 @@ SIMQ::KnnSearch(const DatasetPtr& query, return a.first < b.first; }); - int64_t result_count = std::min(k, static_cast(reranked.size())); + int64_t result_count = 0; + for (const auto& [distance, _] : reranked) { + if (not threshold.has_value() or + (std::isfinite(distance) and distance <= threshold.value())) { + ++result_count; + if (result_count == k) { + break; + } + } + } auto [result_ds, dists, ids] = create_fast_dataset(result_count, allocator_); - for (int64_t i = 0; i < result_count; ++i) { - dists[i] = reranked[i].first; - ids[i] = this->label_table_->GetLabelById(reranked[i].second); + int64_t result_index = 0; + for (const auto& [distance, inner_id] : reranked) { + if (threshold.has_value() and + (not std::isfinite(distance) or distance > threshold.value())) { + continue; + } + dists[result_index] = distance; + ids[result_index] = this->label_table_->GetLabelById(inner_id); + if (++result_index == result_count) { + break; + } } result_ds->Statistics(dump_simq_statistics(stats, coarse_dist_cmp, @@ -792,9 +811,9 @@ SIMQ::KnnSearch(const DatasetPtr& query, coarse_candidate_count, rerank_candidate_count, filtered_candidate_count, - static_cast(result_count), + static_cast(result_ds->GetDim()), false)); - return std::move(result_ds); + return result_ds; } // ───────────────────────────────────────────────────────────────────────────── diff --git a/src/algorithm/sindi/sindi.cpp b/src/algorithm/sindi/sindi.cpp index 855d860fa4..14061e017f 100644 --- a/src/algorithm/sindi/sindi.cpp +++ b/src/algorithm/sindi/sindi.cpp @@ -38,6 +38,7 @@ #include "storage/serialization.h" #include "storage/serialization_tags.h" #include "storage/tlv_section.h" +#include "utils/search_threshold.h" #include "utils/util_functions.h" #include "vsag/allocator.h" #include "vsag/options.h" @@ -600,6 +601,7 @@ SINDI::KnnSearch(const DatasetPtr& query, // search parameter SINDISearchParameter search_param; search_param.FromJson(JsonType::Parse(parameters)); + const auto threshold = ParseSearchThreshold(parameters); CHECK_ARGUMENT(search_param.n_candidate <= SPARSE_AMPLIFICATION_FACTOR * k, fmt::format("n_candidate ({}) should be less than {} * k ({})", search_param.n_candidate, @@ -607,7 +609,7 @@ SINDI::KnnSearch(const DatasetPtr& query, k)); InnerSearchParam inner_param; inner_param.ef = std::max(static_cast(search_param.n_candidate), k); - inner_param.topk = k; + inner_param.topk = threshold.has_value() ? static_cast(inner_param.ef) : k; inner_param.is_inner_id_allowed = this->create_search_filter(filter); @@ -623,12 +625,15 @@ SINDI::KnnSearch(const DatasetPtr& query, auto computer = std::make_shared(effective_query, search_param, allocator_); const SparseVector* rerank_query = (remap_term_ids_ && use_reorder_) ? &sparse_query : nullptr; + DatasetPtr result; if (immutable_data_ != nullptr) { - return immutable_search_impl( + result = immutable_search_impl( + computer, inner_param, allocator, UseTermListsHeapInsert(search_param), rerank_query); + } else { + result = search_impl( computer, inner_param, allocator, UseTermListsHeapInsert(search_param), rerank_query); } - return search_impl( - computer, inner_param, allocator, UseTermListsHeapInsert(search_param), rerank_query); + return FilterDatasetByThreshold(result, threshold, allocator, k); } std::optional diff --git a/src/impl/inner_search_param.h b/src/impl/inner_search_param.h index 8ea012c96e..0bcdd1a1fc 100644 --- a/src/impl/inner_search_param.h +++ b/src/impl/inner_search_param.h @@ -17,6 +17,7 @@ #include #include +#include #include "typing.h" #include "utils/filter_search_skip_strategy.h" @@ -57,6 +58,7 @@ class InnerSearchParam { float factor{2.0F}; bool enable_reorder{true}; float first_order_scan_ratio{1.0F}; + std::optional distance_threshold{std::nullopt}; std::vector executors; // deal with duplicate ids diff --git a/src/impl/searcher/basic_searcher.cpp b/src/impl/searcher/basic_searcher.cpp index f2d19d5914..0a6d426d8f 100644 --- a/src/impl/searcher/basic_searcher.cpp +++ b/src/impl/searcher/basic_searcher.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -206,6 +207,11 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, vl->Set(cur_inner_id); if (iter_ctx->CheckPoint(cur_inner_id)) { flatten->Query(&cur_dist, computer, &cur_inner_id, 1, ctx); + if (std::isnan(cur_dist) or + (mode == InnerSearchMode::RANGE_SEARCH and not std::isfinite(cur_dist))) { + iter_ctx->PopDiscard(); + continue; + } // Sign convention: top_candidates stores positive distances (nearest = smallest); // candidate_set is a max-heap, so distances are negated (nearest = largest, popped first). top_candidates->Push(cur_dist, cur_inner_id); @@ -237,10 +243,13 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, flatten->Query(&dist, computer, &ep, 1, ctx); } if (not is_id_allowed || is_id_allowed->CheckValid(ep)) { - top_candidates->Push(dist, ep); - lower_bound = top_candidates->Top().first; + if (std::isfinite(dist) or (mode == InnerSearchMode::KNN_SEARCH and std::isinf(dist))) { + top_candidates->Push(dist, ep); + lower_bound = top_candidates->Top().first; + } } - candidate_set->Push(-dist, ep); + const auto entry_priority = std::isnan(dist) ? -std::numeric_limits::max() : -dist; + candidate_set->Push(entry_priority, ep); vl->Set(ep); } @@ -300,6 +309,20 @@ BasicSearcher::search_impl(const GraphInterfacePtr& graph, 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)) { + 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); + } + } + if (iter_ctx->CheckPoint(cur_id)) { + candidate_set->Push(std::numeric_limits::max(), cur_id); + flatten->Prefetch(cur_id); + } + } + continue; + } const bool id_allowed = not is_id_allowed || is_id_allowed->CheckValid(cur_id); if constexpr (mode == KNN_SEARCH) { if (collect_rabitq_lower_bound and lower_bound_dists[i] < lower_bound and diff --git a/src/impl/searcher/basic_searcher_test.cpp b/src/impl/searcher/basic_searcher_test.cpp index 089ac31b81..abf728eb15 100644 --- a/src/impl/searcher/basic_searcher_test.cpp +++ b/src/impl/searcher/basic_searcher_test.cpp @@ -517,3 +517,49 @@ TEST_CASE("BasicSearcher iterator drain path handles sign and lower_bound correc delete iter_ctx; } + +TEST_CASE("BasicSearcher traverses through an infinite-distance bridge", + "[ut][BasicSearcher][nonfinite]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + IndexCommonParam common; + common.dim_ = 1; + common.allocator_ = allocator; + common.metric_ = MetricType::METRIC_TYPE_L2SQR; + + constexpr const char* param_temp = R"({{"type": "{}"}})"; + auto quantizer_param = QuantizerParameter::GetQuantizerParameterByJson( + JsonType::Parse(fmt::format(param_temp, "fp32"))); + auto io_param = + IOParameter::GetIOParameterByJson(JsonType::Parse(fmt::format(param_temp, "memory_io"))); + auto flatten = + std::make_shared, MemoryIO>>( + quantizer_param, io_param, common); + flatten->SetQuantizer( + std::make_shared>(1, allocator.get())); + flatten->SetIO(std::make_unique(allocator.get())); + std::vector vectors = {10.0F, std::numeric_limits::max(), 1.0F}; + std::vector ids = {0, 1, 2}; + flatten->Train(vectors.data(), ids.size()); + flatten->BatchInsertVector(vectors.data(), ids.size(), ids.data()); + + auto graph = + std::make_shared(std::vector>{{1}, {2}, {}}); + auto pool = std::make_shared(1, allocator.get(), ids.size(), allocator.get()); + InnerSearchParam param; + param.ep = 0; + param.ef = 3; + param.topk = 2; + float query = 0.0F; + auto vl = pool->TakeOne(); + QueryContext* ctx = nullptr; + auto result = + BasicSearcher(common).Search(graph, flatten, vl, &query, param, LabelTablePtr{}, ctx); + pool->ReturnOne(vl); + + bool found_target = false; + while (not result->Empty()) { + found_target = found_target or result->Top().second == 2; + result->Pop(); + } + REQUIRE(found_target); +} diff --git a/src/impl/searcher/parallel_searcher.cpp b/src/impl/searcher/parallel_searcher.cpp index 8a12618ca1..29cbd75c5c 100644 --- a/src/impl/searcher/parallel_searcher.cpp +++ b/src/impl/searcher/parallel_searcher.cpp @@ -16,6 +16,7 @@ #include "parallel_searcher.h" #include +#include #include #include #include @@ -174,8 +175,10 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, flatten->Query(&dist, computer, &ep, 1, ctx); } if (check_func(ep)) { - top_candidates->Push(dist, ep); - lower_bound = top_candidates->Top().first; + if (std::isfinite(dist) or (mode == InnerSearchMode::KNN_SEARCH and std::isinf(dist))) { + top_candidates->Push(dist, ep); + lower_bound = top_candidates->Top().first; + } } if constexpr (mode == InnerSearchMode::RANGE_SEARCH) { if (dist > inner_search_param.radius and not top_candidates->Empty()) { @@ -185,7 +188,8 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, if (dist < THRESHOLD_ERROR) { inner_search_param.duplicate_id = ep; } - candidate_set->Push(-dist, ep); + const auto entry_priority = std::isnan(dist) ? -std::numeric_limits::max() : -dist; + candidate_set->Push(entry_priority, ep); vl->Set(ep); auto num_threads = inner_search_param.parallel_search_thread_count - 1; @@ -308,6 +312,17 @@ ParallelSearcher::search_impl(const GraphInterfacePtr& graph, for (uint64_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)) { + if (std::isinf(dist)) { + if constexpr (mode == KNN_SEARCH) { + if (check_func(cur_id)) { + top_candidates->Push(dist, cur_id); + } + } + candidate_set->Push(std::numeric_limits::max(), cur_id); + } + continue; + } if constexpr (mode == KNN_SEARCH) { if (collect_rabitq_lower_bound and lower_bound_dists[i] < lower_bound and check_func(cur_id)) { diff --git a/src/impl/searcher/parallel_searcher_test.cpp b/src/impl/searcher/parallel_searcher_test.cpp index 0da03e1c3d..d458a2fa0d 100644 --- a/src/impl/searcher/parallel_searcher_test.cpp +++ b/src/impl/searcher/parallel_searcher_test.cpp @@ -176,3 +176,49 @@ TEST_CASE("Parallel search with HNSW", "[ut][ParallelSearcher]") { } } } + +TEST_CASE("ParallelSearcher traverses through an infinite-distance bridge", + "[ut][ParallelSearcher][nonfinite]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + IndexCommonParam common; + common.dim_ = 1; + common.allocator_ = allocator; + common.metric_ = MetricType::METRIC_TYPE_L2SQR; + + constexpr const char* param_temp = R"({{"type": "{}"}})"; + auto quantizer_param = QuantizerParameter::GetQuantizerParameterByJson( + JsonType::Parse(fmt::format(param_temp, "fp32"))); + auto io_param = + IOParameter::GetIOParameterByJson(JsonType::Parse(fmt::format(param_temp, "memory_io"))); + auto flatten = + std::make_shared, MemoryIO>>( + quantizer_param, io_param, common); + flatten->SetQuantizer( + std::make_shared>(1, allocator.get())); + flatten->SetIO(std::make_unique(allocator.get())); + std::vector vectors = {10.0F, std::numeric_limits::max(), 1.0F}; + std::vector ids = {0, 1, 2}; + flatten->Train(vectors.data(), ids.size()); + flatten->BatchInsertVector(vectors.data(), ids.size(), ids.data()); + + auto graph = + std::make_shared(std::vector>{{1}, {2}, {}}); + auto pool = std::make_shared(1, allocator.get(), ids.size(), allocator.get()); + InnerSearchParam param; + param.ep = 0; + param.ef = 3; + param.topk = 2; + param.parallel_search_thread_count = 2; + float query = 0.0F; + auto vl = pool->TakeOne(); + auto result = ParallelSearcher(common, SafeThreadPool::FactoryDefaultThreadPool()) + .Search(graph, flatten, vl, &query, param); + pool->ReturnOne(vl); + + bool found_target = false; + while (not result->Empty()) { + found_target = found_target or result->Top().second == 2; + result->Pop(); + } + REQUIRE(found_target); +} diff --git a/src/index/index_impl.h b/src/index/index_impl.h index c58af1666a..f0451c7242 100644 --- a/src/index/index_impl.h +++ b/src/index/index_impl.h @@ -19,6 +19,7 @@ #include "algorithm/inner_index_interface.h" #include "common.h" #include "index_common_param.h" +#include "utils/search_threshold.h" #include "vsag/index.h" namespace vsag { @@ -302,6 +303,10 @@ class IndexImpl : public Index { int64_t k, const std::string& parameters, BitsetPtr invalid = nullptr) const override { + auto threshold_validation = ValidateThresholdParameters(parameters); + if (not threshold_validation.has_value()) { + return tl::unexpected(threshold_validation.error()); + } CHECK_QUERY_RETURN_EMPTY_DATASET(query); if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { return DatasetImpl::MakeEmptyDataset(); @@ -314,6 +319,10 @@ class IndexImpl : public Index { int64_t k, const std::string& parameters, const std::function& filter) const override { + auto threshold_validation = ValidateThresholdParameters(parameters); + if (not threshold_validation.has_value()) { + return tl::unexpected(threshold_validation.error()); + } CHECK_QUERY_RETURN_EMPTY_DATASET(query); if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { return DatasetImpl::MakeEmptyDataset(); @@ -326,6 +335,10 @@ class IndexImpl : public Index { int64_t k, const std::string& parameters, const FilterPtr& filter) const override { + auto threshold_validation = ValidateThresholdParameters(parameters); + if (not threshold_validation.has_value()) { + return tl::unexpected(threshold_validation.error()); + } CHECK_QUERY_RETURN_EMPTY_DATASET(query); if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { return DatasetImpl::MakeEmptyDataset(); @@ -335,6 +348,10 @@ class IndexImpl : public Index { tl::expected KnnSearch(const DatasetPtr& query, int64_t k, SearchParam& search_param) const override { + auto threshold_validation = ValidateThresholdParameters(search_param.parameters); + if (not threshold_validation.has_value()) { + return tl::unexpected(threshold_validation.error()); + } CHECK_QUERY_RETURN_EMPTY_DATASET(query); if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(search_param.parameters)) { return DatasetImpl::MakeEmptyDataset(); @@ -360,6 +377,10 @@ class IndexImpl : public Index { const FilterPtr& filter, IteratorContext*& iter_ctx, bool is_last_filter) const override { + auto threshold_validation = ValidateThresholdParameters(parameters); + if (not threshold_validation.has_value()) { + return tl::unexpected(threshold_validation.error()); + } CHECK_QUERY_RETURN_EMPTY_DATASET(query); if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(parameters)) { return DatasetImpl::MakeEmptyDataset(); @@ -476,10 +497,10 @@ class IndexImpl : public Index { [[nodiscard]] tl::expected SearchWithRequest(const SearchRequest& request) const override { - if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(request.params_str_)) { - return DatasetImpl::MakeEmptyDataset(); - } - SAFE_CALL(return this->inner_index_->SearchWithRequest(request)); + SAFE_CALL(ValidateSearchThreshold(request.threshold_); + if (GetNumElements() == 0 && !this->ShouldSkipEmptyCheck(request.params_str_)) { + return DatasetImpl::MakeEmptyDataset(); + } return this->inner_index_->SearchWithRequest(request)); } tl::expected @@ -557,6 +578,18 @@ class IndexImpl : public Index { } private: + tl::expected + ValidateThresholdParameters(const std::string& parameters) const { + try { + ParseSearchThreshold(parameters); + return {}; + } catch (const VsagException& e) { + return tl::unexpected(e.error_); + } catch (const std::exception& e) { + return tl::unexpected(Error(ErrorType::UNKNOWN_ERROR, e.what())); + } + } + bool ShouldSkipEmptyCheck(const std::string& params_str) const { if (GetNumElements() != 0 || params_str.empty()) { diff --git a/src/json_wrapper.cpp b/src/json_wrapper.cpp index c341777b5e..659b109475 100644 --- a/src/json_wrapper.cpp +++ b/src/json_wrapper.cpp @@ -92,6 +92,11 @@ JsonWrapper::IsNumberUnsigned() const { return json_->is_number_unsigned(); } +bool +JsonWrapper::IsNumber() const { + return json_->is_number(); +} + bool JsonWrapper::IsString() const { return json_->is_string(); diff --git a/src/json_wrapper.h b/src/json_wrapper.h index fef2900cc4..db3325b95b 100644 --- a/src/json_wrapper.h +++ b/src/json_wrapper.h @@ -44,6 +44,9 @@ class JsonWrapper { bool IsNumberUnsigned() const; + bool + IsNumber() const; + bool IsString() const; diff --git a/src/utils/search_threshold.h b/src/utils/search_threshold.h new file mode 100644 index 0000000000..34c3e31065 --- /dev/null +++ b/src/utils/search_threshold.h @@ -0,0 +1,125 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "json_types.h" +#include "vsag/dataset.h" + +namespace vsag { + +inline constexpr const char* SEARCH_THRESHOLD = "threshold"; + +inline std::optional +ParseSearchThreshold(const std::string& parameters) { + if (parameters.empty()) { + return std::nullopt; + } + const auto json = JsonType::Parse(parameters); + if (not json.Contains(SEARCH_THRESHOLD)) { + return std::nullopt; + } + CHECK_ARGUMENT(json[SEARCH_THRESHOLD].IsNumber(), "search threshold must be a number"); + const auto threshold = json[SEARCH_THRESHOLD].GetFloat(); + CHECK_ARGUMENT(std::isfinite(threshold), "search threshold must be finite"); + return threshold; +} + +inline void +ValidateSearchThreshold(const std::optional& threshold) { + if (threshold.has_value()) { + CHECK_ARGUMENT(std::isfinite(threshold.value()), "search threshold must be finite"); + } +} + +template +inline T* +AllocateThresholdArray(uint64_t count, Allocator* allocator) { + if (count == 0) { + return nullptr; + } + if (allocator != nullptr) { + auto* result = static_cast(allocator->Allocate(sizeof(T) * count)); + if (result == nullptr) { + throw std::bad_alloc(); + } + return result; + } + return new T[count]; +} + +inline DatasetPtr +FilterDatasetByThreshold(const DatasetPtr& input, + const std::optional& threshold, + Allocator* allocator = nullptr, + int64_t max_results = -1) { + if (not threshold.has_value()) { + return input; + } + int64_t result_count = 0; + for (int64_t i = 0; i < input->GetDim(); ++i) { + if (std::isfinite(input->GetDistances()[i]) and + input->GetDistances()[i] <= threshold.value()) { + ++result_count; + if (max_results > 0 and result_count == max_results) { + break; + } + } + } + auto result = Dataset::Make(); + result->NumElements(1)->Owner(true, allocator); + auto* result_ids = AllocateThresholdArray(result_count, allocator); + result->Dim(result_count)->Ids(result_ids); + auto* result_distances = AllocateThresholdArray(result_count, allocator); + result->Distances(result_distances); + const auto extra_size = input->GetExtraInfoSize(); + const auto* input_extra_infos = input->GetExtraInfos(); + char* extra_infos = nullptr; + if (result_count > 0 and extra_size > 0 and input_extra_infos != nullptr) { + extra_infos = AllocateThresholdArray(static_cast(result_count) * extra_size, + allocator); + result->ExtraInfos(extra_infos)->ExtraInfoSize(extra_size); + } + int64_t result_index = 0; + for (int64_t i = 0; i < input->GetDim() and result_index < result_count; ++i) { + if (std::isfinite(input->GetDistances()[i]) and + input->GetDistances()[i] <= threshold.value()) { + result_ids[result_index] = input->GetIds()[i]; + result_distances[result_index] = input->GetDistances()[i]; + if (extra_infos != nullptr) { + std::memcpy(extra_infos + result_index * extra_size, + input_extra_infos + i * extra_size, + extra_size); + } + ++result_index; + } + } + if (result_count == 0) { + result->ExtraInfos(nullptr)->ExtraInfoSize(0); + } else if (extra_infos == nullptr and extra_size > 0 and input_extra_infos != nullptr) { + result->ExtraInfos(nullptr)->ExtraInfoSize(extra_size); + } + result->Statistics(input->GetStatistics())->Reasoning(input->GetReasoning()); + return result; +} + +} // namespace vsag diff --git a/tests/test_brute_force.cpp b/tests/test_brute_force.cpp index b940ed9a72..085fafd8bc 100644 --- a/tests/test_brute_force.cpp +++ b/tests/test_brute_force.cpp @@ -23,11 +23,13 @@ #include #include #include +#include #include "functest.h" #include "storage/serialization_tags.h" #include "storage/streaming_serialization_test_utils.h" #include "test_index.h" +#include "utils/search_threshold.h" #include "vsag/constants.h" #include "vsag/options.h" #include "vsag/search_request.h" @@ -51,6 +53,7 @@ class TrackingAllocator : public vsag::Allocator { if (ptr != nullptr) { allocations_[ptr] = size; allocated_bytes_ += size; + allocation_count_ += 1; } return ptr; } @@ -103,6 +106,12 @@ class TrackingAllocator : public vsag::Allocator { return allocated_bytes_; } + uint64_t + AllocationCount() const { + std::scoped_lock lock(mutex_); + return allocation_count_; + } + void SetAllocationLimit(uint64_t limit) { std::scoped_lock lock(mutex_); @@ -113,6 +122,7 @@ class TrackingAllocator : public vsag::Allocator { mutable std::mutex mutex_; std::unordered_map allocations_; uint64_t allocated_bytes_{0}; + uint64_t allocation_count_{0}; uint64_t allocation_limit_{std::numeric_limits::max()}; }; @@ -1287,6 +1297,190 @@ TEST_CASE("(PR) BruteForce SearchWithRequest Reasoning", "[ft][bruteforce][reaso REQUIRE(empty_result.value()->GetReasoning().find("filter_rejected") != std::string::npos); } +TEST_CASE("(PR) BruteForce KnnSearch threshold filtering", "[ft][bruteforce][threshold][pr]") { + using namespace fixtures; + + auto param = BruteForceTestIndex::GenerateBruteForceBuildParametersString("l2", 1, "fp32"); + auto index = TestIndex::TestFactory(BruteForceTestIndex::name, param, true); + auto base = vsag::Dataset::Make(); + int64_t ids[] = {10, 11, 12, 13}; + float vectors[] = {0.0F, 1.0F, 2.0F, 3.0F}; + base->NumElements(4)->Dim(1)->Ids(ids)->Float32Vectors(vectors)->Owner(false); + REQUIRE(index->Build(base).has_value()); + + auto query = vsag::Dataset::Make(); + float query_vector[] = {0.0F}; + query->NumElements(1)->Dim(1)->Float32Vectors(query_vector)->Owner(false); + + auto baseline = index->KnnSearch(query, 4, "{}").value(); + auto no_threshold = index->KnnSearch(query, 4, R"({"threshold": 100.0})").value(); + REQUIRE(no_threshold->GetDim() == baseline->GetDim()); + for (int64_t i = 0; i < baseline->GetDim(); ++i) { + REQUIRE(no_threshold->GetIds()[i] == baseline->GetIds()[i]); + REQUIRE(no_threshold->GetDistances()[i] == baseline->GetDistances()[i]); + } + + auto filtered = index->KnnSearch(query, 3, R"({"threshold": 4.0})").value(); + REQUIRE(filtered->GetDim() == 3); + REQUIRE(filtered->GetIds()[0] == 10); + REQUIRE(filtered->GetIds()[1] == 11); + REQUIRE(filtered->GetIds()[2] == 12); + REQUIRE(filtered->GetDistances()[2] == 4.0F); + + auto empty = index->KnnSearch(query, 3, R"({"threshold": -0.1})").value(); + REQUIRE(empty->GetDim() == 0); + + auto overflow_index = TestIndex::TestFactory(BruteForceTestIndex::name, param, true); + float overflow_vector[] = {-std::numeric_limits::max()}; + auto overflow_base = vsag::Dataset::Make(); + overflow_base->NumElements(1)->Dim(1)->Ids(ids)->Float32Vectors(overflow_vector)->Owner(false); + REQUIRE(overflow_index->Build(overflow_base).has_value()); + float overflow_query_value = std::numeric_limits::max(); + auto overflow_query = vsag::Dataset::Make(); + overflow_query->NumElements(1)->Dim(1)->Float32Vectors(&overflow_query_value)->Owner(false); + auto overflow_result = overflow_index->KnnSearch(overflow_query, 1, "{}").value(); + REQUIRE(overflow_result->GetDim() == 1); + REQUIRE(std::isinf(overflow_result->GetDistances()[0])); + auto overflow_filtered = + overflow_index->KnnSearch(overflow_query, 1, R"({"threshold": 0.0})").value(); + REQUIRE(overflow_filtered->GetDim() == 0); + + vsag::SearchRequest request; + request.query_ = query; + request.topk_ = 4; + request.threshold_ = 1.0F; + TrackingAllocator request_allocator; + request.search_allocator_ = &request_allocator; + auto request_result = index->SearchWithRequest(request).value(); + REQUIRE(request_result->GetDim() == 2); + REQUIRE(request_result->GetDistances()[0] <= request_result->GetDistances()[1]); + REQUIRE(request_result->GetDistances()[1] == 1.0F); + REQUIRE(request_allocator.AllocationCount() > 0); + REQUIRE(request_allocator.AllocatedBytes() == 0); + + for (const auto threshold : {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + request.threshold_ = threshold; + REQUIRE_FALSE(index->SearchWithRequest(request).has_value()); + } + REQUIRE_FALSE(index->KnnSearch(query, 1, R"({"threshold":"bad"})").has_value()); + + auto nan_index = TestIndex::TestFactory(BruteForceTestIndex::name, param, true); + float nan_vectors[] = {std::numeric_limits::quiet_NaN(), 0.0F}; + auto nan_base = vsag::Dataset::Make(); + nan_base->NumElements(2)->Dim(1)->Ids(ids)->Float32Vectors(nan_vectors)->Owner(false); + REQUIRE(nan_index->Build(nan_base).has_value()); + auto nan_filtered = nan_index->KnnSearch(query, 1, R"({"threshold": 0.0})").value(); + REQUIRE(nan_filtered->GetDim() == 1); + REQUIRE(nan_filtered->GetIds()[0] == 11); + REQUIRE(nan_filtered->GetDistances()[0] == 0.0F); + + auto ip_param = BruteForceTestIndex::GenerateBruteForceBuildParametersString("ip", 1, "fp32"); + auto ip_index = TestIndex::TestFactory(BruteForceTestIndex::name, ip_param, true); + float ip_vectors[] = {1.0F, 0.5F, 0.0F}; + auto ip_base = vsag::Dataset::Make(); + ip_base->NumElements(3)->Dim(1)->Ids(ids)->Float32Vectors(ip_vectors)->Owner(false); + REQUIRE(ip_index->Build(ip_base).has_value()); + auto ip_query = vsag::Dataset::Make(); + float ip_query_vector[] = {1.0F}; + ip_query->NumElements(1)->Dim(1)->Float32Vectors(ip_query_vector)->Owner(false); + auto ip_filtered = ip_index->KnnSearch(ip_query, 3, R"({"threshold": 0.5})").value(); + REQUIRE(ip_filtered->GetDim() == 2); + REQUIRE(ip_filtered->GetDistances()[0] == 0.0F); + REQUIRE(ip_filtered->GetDistances()[1] == 0.5F); + + auto empty_index = TestIndex::TestFactory(BruteForceTestIndex::name, param, true); + auto malformed_empty = empty_index->KnnSearch(query, 1, R"({"threshold":"bad"})"); + REQUIRE_FALSE(malformed_empty.has_value()); + auto nonfinite_empty = empty_index->KnnSearch(query, 1, R"({"threshold":1e100})"); + REQUIRE_FALSE(nonfinite_empty.has_value()); +} + +TEST_CASE("(PR) SearchRequest preserves legacy aggregate initialization", + "[ut][search_request][compatibility][pr]") { + vsag::SearchRequest request{nullptr, + vsag::SearchMode::RANGE_SEARCH, + 7, + 1.5F, + 3, + "{}", + true, + "attr", + true, + nullptr, + true, + nullptr, + nullptr, + true, + nullptr, + false, + {42}}; + + REQUIRE(request.mode_ == vsag::SearchMode::RANGE_SEARCH); + REQUIRE(request.topk_ == 7); + REQUIRE(request.radius_ == 1.5F); + REQUIRE(request.limited_size_ == 3); + REQUIRE(request.params_str_ == "{}"); + REQUIRE(request.enable_attribute_filter_); + REQUIRE(request.attribute_filter_str_ == "attr"); + REQUIRE(request.enable_filter_); + REQUIRE(request.enable_bitset_filter_); + REQUIRE(request.enable_iterator_search_); + REQUIRE_FALSE(request.is_last_search_); + REQUIRE(request.expected_labels_ == std::vector{42}); + REQUIRE_FALSE(request.threshold_.has_value()); +} + +TEST_CASE("(PR) Threshold filtering preserves allocator ownership", "[ft][threshold][pr]") { + int64_t ids[] = {1, 2, 3}; + float distances[] = {0.0F, 1.0F, 2.0F}; + const char extra_info[] = "aabbcc"; + auto input = vsag::Dataset::Make(); + input->NumElements(1) + ->Dim(3) + ->Ids(ids) + ->Distances(distances) + ->ExtraInfoSize(2) + ->ExtraInfos(extra_info) + ->Owner(false); + + TrackingAllocator allocator; + { + auto result = vsag::FilterDatasetByThreshold(input, 1.0F, &allocator); + REQUIRE(result->GetDim() == 2); + REQUIRE(result->GetIds()[0] == 1); + REQUIRE(result->GetIds()[1] == 2); + REQUIRE(std::memcmp(result->GetExtraInfos(), "aabb", 4) == 0); + REQUIRE(allocator.AllocatedBytes() > 0); + } + REQUIRE(allocator.AllocatedBytes() == 0); + + auto empty = vsag::FilterDatasetByThreshold(input, -1.0F, &allocator); + REQUIRE(empty->GetDim() == 0); + REQUIRE(empty->GetExtraInfoSize() == 0); + REQUIRE(empty->GetExtraInfos() == nullptr); + + int64_t non_finite_ids[] = {7, 8, 9}; + float non_finite_distances[] = {-std::numeric_limits::infinity(), 0.0F, 1.0F}; + auto non_finite_input = vsag::Dataset::Make(); + non_finite_input->NumElements(1) + ->Dim(3) + ->Ids(non_finite_ids) + ->Distances(non_finite_distances) + ->Owner(false); + auto finite_only = vsag::FilterDatasetByThreshold(non_finite_input, 1.0F, &allocator, 1); + REQUIRE(finite_only->GetDim() == 1); + REQUIRE(finite_only->GetIds()[0] == 8); + REQUIRE(finite_only->GetDistances()[0] == 0.0F); + finite_only.reset(); + REQUIRE(allocator.AllocatedBytes() == 0); + + allocator.SetAllocationLimit(sizeof(int64_t) * 2); + REQUIRE_THROWS_AS(vsag::FilterDatasetByThreshold(input, 1.0F, &allocator), std::bad_alloc); + REQUIRE(allocator.AllocatedBytes() == 0); +} + TEST_CASE("(PR) BruteForce Reasoning Found Verification", "[ft][bruteforce][reasoning][pr]") { using namespace fixtures; diff --git a/tests/test_hgraph.cpp b/tests/test_hgraph.cpp index 6a6243efc9..8c660e885f 100644 --- a/tests/test_hgraph.cpp +++ b/tests/test_hgraph.cpp @@ -25,6 +25,7 @@ #include #include "functest.h" +#include "impl/filter/iterator_filter.h" #include "inner_string_params.h" #include "storage/serialization_tags.h" #include "storage/streaming_serialization_test_utils.h" @@ -3298,6 +3299,117 @@ TEST_CASE("(PR) HGraph brute_force_threshold default is no-op", REQUIRE_FALSE(negative_ef.has_value()); } +TEST_CASE("(PR) HGraph threshold iterator consumes rejected pages", + "[ft][hgraph][threshold][iterator][pr]") { + constexpr int64_t dim = 1; + constexpr int64_t base_count = 32; + std::string params = R"({ + "dtype":"float32", "metric_type":"l2", "dim":1, + "index_param":{"base_quantization_type":"fp32","max_degree":16, + "ef_construction":64,"use_reorder":false} + })"; + auto index = vsag::Factory::CreateIndex("hgraph", params).value(); + std::vector vectors(base_count); + std::vector ids(base_count); + for (int64_t i = 0; i < base_count; ++i) { + vectors[i] = static_cast(i); + ids[i] = i; + } + auto base = vsag::Dataset::Make(); + base->NumElements(base_count) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->Owner(false); + REQUIRE(index->Build(base).has_value()); + auto query = vsag::Dataset::Make(); + float query_value = 0.0F; + query->NumElements(1)->Dim(dim)->Float32Vectors(&query_value)->Owner(false); + + vsag::IteratorContext* iter_ctx = nullptr; + const auto search_params = R"({"hgraph":{"ef_search":8},"threshold":-1.0})"; + auto first = + index->KnnSearch(query, 1, search_params, vsag::FilterPtr(nullptr), iter_ctx, false); + REQUIRE(first.has_value()); + REQUIRE(first.value()->GetDim() == 0); + auto* filter_ctx = static_cast(iter_ctx); + REQUIRE(filter_ctx->GetDiscardElementNum() == 0); + delete iter_ctx; +} + +TEST_CASE("(PR) HGraph ignores non-finite entry distances", + "[ft][hgraph][threshold][nonfinite][pr]") { + const auto params = R"({ + "dtype":"float32", "metric_type":"l2", "dim":1, + "index_param":{"base_quantization_type":"fp32","max_degree":16, + "ef_construction":32,"use_reorder":false} + })"; + auto index = vsag::Factory::CreateIndex("hgraph", params).value(); + std::vector vectors = {0.0F, std::numeric_limits::max()}; + std::vector ids = {0, 1}; + auto base = vsag::Dataset::Make(); + base->NumElements(2)->Dim(1)->Ids(ids.data())->Float32Vectors(vectors.data())->Owner(false); + REQUIRE(index->Build(base).has_value()); + + float overflow_query_value = std::numeric_limits::max(); + auto overflow_query = vsag::Dataset::Make(); + overflow_query->NumElements(1)->Dim(1)->Float32Vectors(&overflow_query_value)->Owner(false); + auto overflow_result = index->KnnSearch(overflow_query, 1, R"({"hgraph":{"ef_search":8}})"); + REQUIRE(overflow_result.has_value()); + REQUIRE(overflow_result.value()->GetDim() == 1); + REQUIRE(overflow_result.value()->GetIds()[0] == 1); + REQUIRE(overflow_result.value()->GetDistances()[0] == 0.0F); + + float nan_query_value = std::numeric_limits::quiet_NaN(); + auto nan_query = vsag::Dataset::Make(); + nan_query->NumElements(1)->Dim(1)->Float32Vectors(&nan_query_value)->Owner(false); + auto nan_result = index->KnnSearch(nan_query, 1, R"({"hgraph":{"ef_search":8}})"); + REQUIRE(nan_result.has_value()); + for (int64_t i = 0; i < nan_result.value()->GetDim(); ++i) { + REQUIRE(std::isfinite(nan_result.value()->GetDistances()[i])); + } +} + +TEST_CASE("(PR) HGraph iterator preserves infinity without threshold", + "[ft][hgraph][threshold][iterator][nonfinite][pr]") { + const auto params = R"({ + "dtype":"float32", "metric_type":"ip", "dim":1, + "index_param":{"base_quantization_type":"fp32","max_degree":16, + "ef_construction":32,"use_reorder":false} + })"; + auto index = vsag::Factory::CreateIndex("hgraph", params).value(); + float vector = std::numeric_limits::max(); + int64_t id = 42; + auto base = vsag::Dataset::Make(); + base->NumElements(1)->Dim(1)->Ids(&id)->Float32Vectors(&vector)->Owner(false); + REQUIRE(index->Build(base).has_value()); + + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(1)->Float32Vectors(&vector)->Owner(false); + const auto search_params = R"({"hgraph":{"ef_search":8}})"; + + vsag::IteratorContext* iter_ctx = nullptr; + auto ordinary = + index->KnnSearch(query, 1, search_params, vsag::FilterPtr(nullptr), iter_ctx, false); + REQUIRE(ordinary.has_value()); + REQUIRE(ordinary.value()->GetDim() == 1); + REQUIRE(ordinary.value()->GetIds()[0] == id); + REQUIRE(std::isinf(ordinary.value()->GetDistances()[0])); + REQUIRE(ordinary.value()->GetDistances()[0] < 0.0F); + delete iter_ctx; + + iter_ctx = nullptr; + auto threshold = index->KnnSearch(query, + 1, + R"({"hgraph":{"ef_search":8},"threshold":0.0})", + vsag::FilterPtr(nullptr), + iter_ctx, + false); + REQUIRE(threshold.has_value()); + REQUIRE(threshold.value()->GetDim() == 0); + delete iter_ctx; +} + TEST_CASE("HGraph ExportCache + ImportCache + Build acceleration smoke test", "[ft][hgraph][cache][pr]") { // End-to-end smoke test for the cache-accelerated Build path: diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index d2faff76c0..e030613586 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -1687,6 +1687,77 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, } } +TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, + "IVF reorder applies threshold to exact distances", + "[ft][search][ivf][threshold][pr]") { + constexpr int64_t dim = 16; + auto param = IVFTestIndex::GenerateIVFBuildParametersString("l2", dim, "sq8,fp32", 10); + auto index = TestIndex::TestFactory(IVFTestIndex::name, param, true); + auto dataset = IVFTestIndex::pool.GetDatasetAndCreate(dim, 200, "l2"); + TestIndex::TestBuildIndex(index, dataset, true); + + auto query = vsag::Dataset::Make(); + query->NumElements(1) + ->Dim(dim) + ->Float32Vectors(dataset->base_->GetFloat32Vectors()) + ->Owner(false); + auto result = index->KnnSearch( + query, 2, R"({"ivf":{"scan_buckets_count":10,"factor":4.0},"threshold":0.0})"); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() >= 1); + REQUIRE(result.value()->GetDim() <= 2); + REQUIRE(result.value()->GetIds()[0] == dataset->base_->GetIds()[0]); + for (int64_t i = 0; i < result.value()->GetDim(); ++i) { + REQUIRE(result.value()->GetDistances()[i] <= 0.0F); + if (i > 0) { + REQUIRE(result.value()->GetDistances()[i - 1] <= result.value()->GetDistances()[i]); + } + } + + vsag::SearchRequest reasoning_request; + reasoning_request.query_ = query; + reasoning_request.topk_ = 2; + reasoning_request.params_str_ = R"({"ivf":{"scan_buckets_count":10,"factor":4.0}})"; + reasoning_request.threshold_ = -1.0F; + reasoning_request.expected_labels_ = {dataset->base_->GetIds()[0]}; + auto reasoning_result = index->SearchWithRequest(reasoning_request); + REQUIRE(reasoning_result.has_value()); + REQUIRE(reasoning_result.value()->GetDim() == 0); + REQUIRE(reasoning_result.value()->GetReasoning().find("0/1 expected labels found") != + std::string::npos); + + auto post_filter_baseline = index->KnnSearch( + query, 2, R"({"ivf":{"scan_buckets_count":10,"factor":4.0},"threshold":1e9})"); + REQUIRE(post_filter_baseline.has_value()); + REQUIRE(post_filter_baseline.value()->GetDim() == 2); + reasoning_request.threshold_ = 1e9F; + reasoning_request.expected_labels_ = {post_filter_baseline.value()->GetIds()[1]}; + auto post_filter_reasoning = index->SearchWithRequest(reasoning_request); + REQUIRE(post_filter_reasoning.has_value()); + REQUIRE(post_filter_reasoning.value()->GetReasoning().find("1/1 expected labels found") != + std::string::npos); +} + +TEST_CASE_PERSISTENT_FIXTURE(fixtures::IVFTestIndex, + "IVF threshold filters before top-k selection", + "[ft][search][ivf][threshold][nonfinite][pr]") { + auto param = IVFTestIndex::GenerateIVFBuildParametersString("ip", 1, "fp32", 1, "random"); + auto index = TestIndex::TestFactory(IVFTestIndex::name, param, true); + std::vector vectors = {std::numeric_limits::max(), 0.0F}; + std::vector ids = {10, 20}; + auto base = vsag::Dataset::Make(); + base->NumElements(2)->Dim(1)->Ids(ids.data())->Float32Vectors(vectors.data())->Owner(false); + REQUIRE(index->Build(base).has_value()); + + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(1)->Float32Vectors(vectors.data())->Owner(false); + auto result = index->KnnSearch(query, 1, R"({"ivf":{"scan_buckets_count":1},"threshold":1.0})"); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 1); + REQUIRE(result.value()->GetIds()[0] == 20); + REQUIRE(result.value()->GetDistances()[0] == 1.0F); +} + // RejectAllFilter for testing empty results class RejectAllFilter : public vsag::Filter { public: diff --git a/tests/test_pyramid.cpp b/tests/test_pyramid.cpp index 9896f636b0..1fe491689b 100644 --- a/tests/test_pyramid.cpp +++ b/tests/test_pyramid.cpp @@ -450,6 +450,11 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::PyramidTestIndex, index->KnnSearch(query, topk, GeneratePyramidSearchParametersString(ef_search)); REQUIRE(search_result.has_value()); REQUIRE(search_result.value()->GetDim() == topk); + + auto threshold_result = + index->KnnSearch(query, topk, R"({"threshold": 0.0, "pyramid": {"ef_search": 5}})"); + REQUIRE(threshold_result.has_value()); + REQUIRE(threshold_result.value()->GetDim() == 1); } TEST_CASE_PERSISTENT_FIXTURE(fixtures::PyramidTestIndex, diff --git a/tests/test_simq.cpp b/tests/test_simq.cpp index 90bfe8a703..e31206a9e9 100644 --- a/tests/test_simq.cpp +++ b/tests/test_simq.cpp @@ -371,6 +371,15 @@ TEST_CASE("SIMQ: build and knn search recall", "[simq][build][search]") { REQUIRE(search_result.has_value()); require_simq_search_stats(search_result.value()); + if (q == 0) { + auto threshold_result = index->KnnSearch( + one_query, TOP_K, R"({"threshold": -1000000.0, "simq": {"coarse_k": 10}})"); + REQUIRE(threshold_result.has_value()); + REQUIRE(threshold_result.value()->GetDim() == 0); + auto threshold_stats = get_simq_search_stats(threshold_result.value()); + REQUIRE(threshold_stats.result_count == 0); + } + auto* ret_ids = search_result.value()->GetIds(); float r = recall_at_k(ret_ids, TOP_K, ds.gt_ids[q].data(), static_cast(TOP_K)); total_recall += r; diff --git a/tests/test_sindi.cpp b/tests/test_sindi.cpp index 1c944d64a9..efe5e40893 100644 --- a/tests/test_sindi.cpp +++ b/tests/test_sindi.cpp @@ -15,6 +15,7 @@ #include #include +#include #include "functest.h" #include "test_index.h" @@ -209,6 +210,39 @@ TEST_CASE_PERSISTENT_FIXTURE(fixtures::SINDITestIndex, TestIndexStatus(index); } +TEST_CASE_PERSISTENT_FIXTURE(fixtures::SINDITestIndex, + "SINDI threshold backfills after non-finite candidates", + "[ft][search][sindi][threshold][nonfinite][pr]") { + fixtures::SINDIParam param; + param.use_reorder = true; + param.sparse_value_quant_type = "fp32"; + auto index = TestFactory("sindi", GenerateBuildParameter(param), true); + + std::vector term_ids = {1}; + std::vector overflow_value = {std::numeric_limits::max()}; + std::vector finite_value = {0.25F}; + vsag::SparseVector base_vectors[2]; + base_vectors[0] = {1, term_ids.data(), overflow_value.data()}; + base_vectors[1] = {1, term_ids.data(), finite_value.data()}; + int64_t ids[] = {10, 20}; + auto base = vsag::Dataset::Make(); + base->NumElements(2)->Ids(ids)->SparseVectors(base_vectors)->Owner(false); + REQUIRE(index->Build(base).has_value()); + + std::vector query_value = {2.0F}; + vsag::SparseVector query_vector{1, term_ids.data(), query_value.data()}; + auto query = vsag::Dataset::Make(); + query->NumElements(1)->SparseVectors(&query_vector)->Owner(false); + auto result = index->KnnSearch(query, + 1, + R"({"sindi":{"n_candidate":2,"query_prune_ratio":0.0, + "term_prune_ratio":0.0},"threshold":1.0})"); + REQUIRE(result.has_value()); + REQUIRE(result.value()->GetDim() == 1); + REQUIRE(result.value()->GetIds()[0] == 20); + REQUIRE(result.value()->GetDistances()[0] == 0.5F); +} + TEST_CASE_PERSISTENT_FIXTURE(fixtures::SINDITestIndex, "SINDI Analyze", "[ft][analyze][sindi]") { fixtures::SINDIParam param; param.use_reorder = GENERATE(true, false);