From 1f8f8264c6137943b971f3aac97885cd5c94a3e5 Mon Sep 17 00:00:00 2001 From: LHT129 Date: Mon, 6 Jul 2026 15:17:44 +0800 Subject: [PATCH] feat(simq): add search statistics Populate Dataset statistics for SIMQ KNN and range search results, including coarse-search and rerank counters for observability. Fixes #2417 Signed-off-by: LHT129 Assisted-by: OpenCode:codex-1/gpt-5.5 --- src/algorithm/simq/simq.cpp | 94 +++++++++++++++++++++++++++++++++++-- src/algorithm/simq/simq.h | 6 ++- tests/test_simq.cpp | 67 ++++++++++++++++++++++++-- 3 files changed, 158 insertions(+), 9 deletions(-) diff --git a/src/algorithm/simq/simq.cpp b/src/algorithm/simq/simq.cpp index 9abc9ec058..82a0547dc4 100644 --- a/src/algorithm/simq/simq.cpp +++ b/src/algorithm/simq/simq.cpp @@ -16,6 +16,7 @@ #include "simq.h" #include +#include #include #include #include @@ -28,6 +29,7 @@ #include "index_feature_list.h" #include "inner_string_params.h" #include "metric_type.h" +#include "query_context.h" #include "storage/serialization.h" #include "storage/stream_reader.h" #include "storage/stream_writer.h" @@ -47,6 +49,38 @@ struct cluster_member_entry { float distance; }; +std::string +dump_simq_statistics(const SearchStatistics& stats, + uint64_t coarse_dist_cmp, + uint64_t coarse_probe_count, + uint64_t coarse_candidate_count, + uint64_t rerank_candidate_count, + uint64_t filtered_candidate_count, + uint64_t result_count, + bool limited_size_applied) { + auto json = JsonType::Parse(stats.Dump()); + json["simq_coarse_dist_cmp"].SetUint64(coarse_dist_cmp); + json["simq_coarse_probe_count"].SetUint64(coarse_probe_count); + json["simq_coarse_candidate_count"].SetUint64(coarse_candidate_count); + json["simq_rerank_candidate_count"].SetUint64(rerank_candidate_count); + json["simq_filtered_candidate_count"].SetUint64(filtered_candidate_count); + json["simq_result_count"].SetUint64(result_count); + json["simq_limited_size_applied"].SetBool(limited_size_applied); + return json.Dump(); +} + +uint64_t +read_dist_cmp(const DatasetPtr& result_ds) { + if (result_ds == nullptr) { + return 0; + } + auto values = result_ds->GetStatistics({"dist_cmp"}); + if (values.empty() || values[0].empty()) { + return 0; + } + return std::strtoull(values[0].c_str(), nullptr, 10); +} + class HGraphDynamicClustering { public: HGraphDynamicClustering(float init_cluster_ratio, @@ -603,7 +637,11 @@ SIMQ::split_cluster_incremental(InnerIdType cluster_idx) { // ───────────────────────────────────────────────────────────────────────────── std::vector> -SIMQ::coarse_search(const float* query_tokens, uint32_t query_token_count, int64_t coarse_k) const { +SIMQ::coarse_search(const float* query_tokens, + uint32_t query_token_count, + int64_t coarse_k, + uint64_t* coarse_dist_cmp, + uint64_t* coarse_probe_count) const { // All buffers are local — safe for concurrent searches under shared_lock. std::unordered_map score_map; score_map.reserve(static_cast(coarse_k) * static_cast(max_cluster_size_)); @@ -618,11 +656,17 @@ SIMQ::coarse_search(const float* query_tokens, uint32_t query_token_count, int64 if (actual_coarse_k <= 0) { continue; } + if (coarse_probe_count != nullptr) { + *coarse_probe_count += static_cast(actual_coarse_k); + } auto query_ds = Dataset::Make(); query_ds->NumElements(1)->Dim(dim_)->Float32Vectors(qt)->Owner(false); auto result_ds = rep_hgraph_->KnnSearch( query_ds, actual_coarse_k, R"({"hgraph": {"ef_search": 100}})", nullptr); + if (coarse_dist_cmp != nullptr) { + *coarse_dist_cmp += read_dist_cmp(result_ds); + } int64_t nres = result_ds->GetDim(); const auto* rdists = result_ds->GetDistances(); @@ -670,9 +714,12 @@ SIMQ::KnnSearch(const DatasetPtr& query, const std::string& parameters, const FilterPtr& filter) const { std::shared_lock lock(global_mutex_); + SearchStatistics stats; if (total_count_ == 0 || rep_hgraph_ == nullptr) { - return Dataset::Make(); + auto result = Dataset::Make(); + result->Statistics(dump_simq_statistics(stats, 0, 0, 0, 0, 0, 0, false)); + return result; } CHECK_ARGUMENT(query->GetNumElements() > 0, "simq search: query.num_elements must be > 0"); @@ -688,21 +735,29 @@ SIMQ::KnnSearch(const DatasetPtr& query, rerank_k = std::min(rerank_k, static_cast(total_count_)); k = std::min(k, static_cast(total_count_)); - auto coarse_results = coarse_search(query_mvs[0].vectors_, query_mvs[0].len_, coarse_k); + uint64_t coarse_dist_cmp = 0; + uint64_t coarse_probe_count = 0; + auto coarse_results = coarse_search( + query_mvs[0].vectors_, query_mvs[0].len_, coarse_k, &coarse_dist_cmp, &coarse_probe_count); + uint64_t coarse_candidate_count = coarse_results.size(); if (static_cast(coarse_results.size()) > rerank_k) { coarse_results.resize(rerank_k); } + uint64_t rerank_candidate_count = coarse_results.size(); // Exact MaxSim rerank via MultiVectorDataCell auto computer = mv_codes_->FactoryComputer(&query_mvs[0]); std::vector> reranked; reranked.reserve(coarse_results.size()); + uint64_t filtered_candidate_count = 0; for (auto& [doc_id, _] : coarse_results) { if (filter != nullptr && !filter->CheckValid(this->label_table_->GetLabelById(doc_id))) { + ++filtered_candidate_count; continue; } float dist = 0.0F; mv_codes_->Query(&dist, computer, &doc_id, 1); + ++stats.dist_cmp; reranked.emplace_back(dist, doc_id); } std::sort(reranked.begin(), reranked.end(), [](const auto& a, const auto& b) { @@ -715,6 +770,14 @@ SIMQ::KnnSearch(const DatasetPtr& query, dists[i] = reranked[i].first; ids[i] = this->label_table_->GetLabelById(reranked[i].second); } + result_ds->Statistics(dump_simq_statistics(stats, + coarse_dist_cmp, + coarse_probe_count, + coarse_candidate_count, + rerank_candidate_count, + filtered_candidate_count, + static_cast(result_count), + false)); return std::move(result_ds); } @@ -729,9 +792,12 @@ SIMQ::RangeSearch(const DatasetPtr& query, const FilterPtr& filter, int64_t limited_size) const { std::shared_lock lock(global_mutex_); + SearchStatistics stats; if (total_count_ == 0 || rep_hgraph_ == nullptr) { - return Dataset::Make(); + auto result = Dataset::Make(); + result->Statistics(dump_simq_statistics(stats, 0, 0, 0, 0, 0, 0, false)); + return result; } CHECK_ARGUMENT(query->GetNumElements() > 0, @@ -748,25 +814,35 @@ SIMQ::RangeSearch(const DatasetPtr& query, int64_t rerank_k = sp.rerank_k > 0 ? sp.rerank_k : default_rerank_k_; rerank_k = std::min(rerank_k, static_cast(total_count_)); - auto coarse_results = coarse_search(query_mvs[0].vectors_, query_mvs[0].len_, coarse_k); + uint64_t coarse_dist_cmp = 0; + uint64_t coarse_probe_count = 0; + auto coarse_results = coarse_search( + query_mvs[0].vectors_, query_mvs[0].len_, coarse_k, &coarse_dist_cmp, &coarse_probe_count); + uint64_t coarse_candidate_count = coarse_results.size(); if (static_cast(coarse_results.size()) > rerank_k) { coarse_results.resize(rerank_k); } + uint64_t rerank_candidate_count = coarse_results.size(); auto computer = mv_codes_->FactoryComputer(&query_mvs[0]); std::vector> in_range; + uint64_t filtered_candidate_count = 0; for (auto& [doc_id, _] : coarse_results) { if (filter != nullptr && !filter->CheckValid(this->label_table_->GetLabelById(doc_id))) { + ++filtered_candidate_count; continue; } float dist = 0.0F; mv_codes_->Query(&dist, computer, &doc_id, 1); + ++stats.dist_cmp; if (dist <= radius) { in_range.emplace_back(dist, doc_id); } } + bool limited_size_applied = false; if (limited_size >= 0 && static_cast(in_range.size()) > limited_size) { + limited_size_applied = true; std::nth_element(in_range.begin(), in_range.begin() + limited_size, in_range.end(), @@ -783,6 +859,14 @@ SIMQ::RangeSearch(const DatasetPtr& query, dists[i] = in_range[i].first; ids[i] = this->label_table_->GetLabelById(in_range[i].second); } + result_ds->Statistics(dump_simq_statistics(stats, + coarse_dist_cmp, + coarse_probe_count, + coarse_candidate_count, + rerank_candidate_count, + filtered_candidate_count, + static_cast(in_range.size()), + limited_size_applied)); return std::move(result_ds); } diff --git a/src/algorithm/simq/simq.h b/src/algorithm/simq/simq.h index b4435a81f0..adca9da986 100644 --- a/src/algorithm/simq/simq.h +++ b/src/algorithm/simq/simq.h @@ -96,7 +96,11 @@ class SIMQ : public InnerIndexInterface { build_rep_hgraph(const float* flat_vecs, int64_t dim); std::vector> - coarse_search(const float* query_tokens, uint32_t query_token_count, int64_t coarse_k) const; + coarse_search(const float* query_tokens, + uint32_t query_token_count, + int64_t coarse_k, + uint64_t* coarse_dist_cmp = nullptr, + uint64_t* coarse_probe_count = nullptr) const; void serialize_rep_hgraph(StreamWriter& writer) const; diff --git a/tests/test_simq.cpp b/tests/test_simq.cpp index 798e9f5dc3..90bfe8a703 100644 --- a/tests/test_simq.cpp +++ b/tests/test_simq.cpp @@ -48,6 +48,7 @@ #include #include #include +#include #include #include #include @@ -256,6 +257,59 @@ recall_at_k(const int64_t* returned, int64_t n_returned, const int64_t* gt, int6 return static_cast(hits) / static_cast(n_gt); } +struct SimqSearchStats { + uint64_t dist_cmp{0}; + uint64_t coarse_dist_cmp{0}; + uint64_t coarse_probe_count{0}; + uint64_t coarse_candidate_count{0}; + uint64_t rerank_candidate_count{0}; + uint64_t filtered_candidate_count{0}; + uint64_t result_count{0}; + std::string limited_size_applied; +}; + +static uint64_t +parse_u64(const std::string& value) { + REQUIRE(!value.empty()); + return std::strtoull(value.c_str(), nullptr, 10); +} + +static SimqSearchStats +get_simq_search_stats(const vsag::DatasetPtr& result) { + REQUIRE(result->GetStatistics() != "{}"); + auto values = result->GetStatistics({"dist_cmp", + "simq_coarse_dist_cmp", + "simq_coarse_probe_count", + "simq_coarse_candidate_count", + "simq_rerank_candidate_count", + "simq_filtered_candidate_count", + "simq_result_count", + "simq_limited_size_applied"}); + REQUIRE(values.size() == 8); + + SimqSearchStats stats; + stats.dist_cmp = parse_u64(values[0]); + stats.coarse_dist_cmp = parse_u64(values[1]); + stats.coarse_probe_count = parse_u64(values[2]); + stats.coarse_candidate_count = parse_u64(values[3]); + stats.rerank_candidate_count = parse_u64(values[4]); + stats.filtered_candidate_count = parse_u64(values[5]); + stats.result_count = parse_u64(values[6]); + stats.limited_size_applied = values[7]; + return stats; +} + +static void +require_simq_search_stats(const vsag::DatasetPtr& result) { + auto stats = get_simq_search_stats(result); + REQUIRE(stats.dist_cmp > 0); + REQUIRE(stats.coarse_dist_cmp > 0); + REQUIRE(stats.coarse_probe_count > 0); + REQUIRE(stats.coarse_candidate_count >= stats.rerank_candidate_count); + REQUIRE(stats.rerank_candidate_count >= stats.dist_cmp); + REQUIRE(stats.result_count == static_cast(result->GetDim())); +} + // ───────────────────────────────────────────────────────────────────────────── // Test cases // ───────────────────────────────────────────────────────────────────────────── @@ -315,6 +369,7 @@ TEST_CASE("SIMQ: build and knn search recall", "[simq][build][search]") { auto search_result = index->KnnSearch(one_query, TOP_K, search_param, vsag::FilterPtr(nullptr)); REQUIRE(search_result.has_value()); + require_simq_search_stats(search_result.value()); 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)); @@ -407,15 +462,17 @@ TEST_CASE("SIMQ: range search", "[simq][range_search]") { // Use KNN distances to pick a sensible radius: the worst distance in top-k auto knn_result = index->KnnSearch(one_query, TOP_K, search_param, vsag::FilterPtr(nullptr)); REQUIRE(knn_result.has_value()); + require_simq_search_stats(knn_result.value()); const float* knn_dists = knn_result.value()->GetDistances(); - int64_t knn_n = knn_result.value()->GetNumElements(); + int64_t knn_n = knn_result.value()->GetDim(); REQUIRE(knn_n > 0); float radius = knn_dists[knn_n - 1]; // worst distance among top-k results SECTION("all returned distances are within radius") { auto rr = index->RangeSearch(one_query, radius, search_param, vsag::FilterPtr(nullptr)); REQUIRE(rr.has_value()); - int64_t n = rr.value()->GetNumElements(); + require_simq_search_stats(rr.value()); + int64_t n = rr.value()->GetDim(); const float* rdists = rr.value()->GetDistances(); REQUIRE(n >= 1); for (int64_t i = 0; i < n; ++i) { @@ -428,7 +485,11 @@ TEST_CASE("SIMQ: range search", "[simq][range_search]") { auto rr = index->RangeSearch(one_query, radius, search_param, vsag::FilterPtr(nullptr), limited); REQUIRE(rr.has_value()); - REQUIRE(rr.value()->GetNumElements() <= limited); + require_simq_search_stats(rr.value()); + int64_t n = rr.value()->GetDim(); + REQUIRE(n <= limited); + auto stats = get_simq_search_stats(rr.value()); + REQUIRE(stats.limited_size_applied == "true"); } }