xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test: stabilize DiskANN and IVF searcher tests (#637)

This commit is contained in:
ZeFeng Yin 2026-08-03 10:39:17 +08:00 committed by GitHub
parent ad36b75966
commit be6d7b7cfb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 41 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
tests/core/algorithm/diskann/diskann_searcher_test.cc
View File

@ -16,6 +16,8 @@
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <cstring>
#include <unordered_set>
#include <ailego/math/distance.h>
#include <gtest/gtest.h>
#include <zvec/ailego/container/vector.h>
@ -126,8 +128,6 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) {
NumericalVector<float> vec(dim);
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
size_t topk = 200;
uint64_t knnTotalTime = 0;
uint64_t linearTotalTime = 0;
int totalHits = 0;
int totalCnts = 0;
int topk1Hits = 0;
@ -185,14 +185,8 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) {
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
auto t1 = Realtime::MicroSeconds();
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
auto t2 = Realtime::MicroSeconds();
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
auto t3 = Realtime::MicroSeconds();
knnTotalTime += t2 - t1;
linearTotalTime += t3 - t2;
auto &knnResult = knnCtx->result();
// TODO: check
@ -215,11 +209,9 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) {
float recall = totalHits * step * step * 1.0f / totalCnts;
float topk1Recall = topk1Hits * step * 1.0f / doc_cnt;
float cost = linearTotalTime * 1.0f / knnTotalTime;
EXPECT_GT(recall, 0.90f);
EXPECT_GT(topk1Recall, 0.80f);
EXPECT_GT(cost, 2.0f);
}
TEST_F(DiskAnnSearcherTest, TestNodeCache) {
@ -294,8 +286,6 @@ TEST_F(DiskAnnSearcherTest, TestNodeCache) {
NumericalVector<float> vec(dim);
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
size_t topk = 200;
uint64_t knnTotalTime = 0;
uint64_t linearTotalTime = 0;
int totalHits = 0;
int totalCnts = 0;
int topk1Hits = 0;
@ -308,14 +298,8 @@ TEST_F(DiskAnnSearcherTest, TestNodeCache) {
for (size_t j = 0; j < dim; ++j) {
vec[j] = i + 0.1f;
}
auto t1 = Realtime::MicroSeconds();
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
auto t2 = Realtime::MicroSeconds();
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
auto t3 = Realtime::MicroSeconds();
knnTotalTime += t2 - t1;
linearTotalTime += t3 - t2;
auto &knnResult = knnCtx->result();
// TODO: check
@ -338,11 +322,9 @@ TEST_F(DiskAnnSearcherTest, TestNodeCache) {
float recall = totalHits * step * step * 1.0f / totalCnts;
float topk1Recall = topk1Hits * step * 1.0f / doc_cnt;
float cost = linearTotalTime * 1.0f / knnTotalTime;
EXPECT_GT(recall, 0.90f);
EXPECT_GT(topk1Recall, 0.80f);
EXPECT_GT(cost, 2.0f);
}
TEST_F(DiskAnnSearcherTest, TestFilter) {
@ -433,9 +415,6 @@ TEST_F(DiskAnnSearcherTest, TestFilter) {
auto &knnResult = knnCtx->result();
ASSERT_EQ(topk, knnResult.size());
ASSERT_EQ(50UL, knnResult[0].key());
ASSERT_EQ(51UL, knnResult[1].key());
ASSERT_EQ(49UL, knnResult[2].key());
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
@ -461,9 +440,13 @@ TEST_F(DiskAnnSearcherTest, TestFilter) {
auto &knnResult = knnCtx->result();
ASSERT_EQ(topk, knnResult.size());
ASSERT_EQ(52UL, knnResult[0].key());
ASSERT_EQ(48UL, knnResult[1].key());
ASSERT_EQ(53UL, knnResult[2].key());
std::unordered_set<uint64_t> knn_keys;
for (const auto &result : knnResult) {
ASSERT_TRUE(knn_keys.emplace(result.key()).second);
EXPECT_NE(50UL, result.key());
EXPECT_NE(51UL, result.key());
EXPECT_NE(49UL, result.key());
}
linearCtx->set_filter(filterFunc);
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
@ -473,6 +456,13 @@ TEST_F(DiskAnnSearcherTest, TestFilter) {
ASSERT_EQ(52UL, linearResult[0].key());
ASSERT_EQ(48UL, linearResult[1].key());
ASSERT_EQ(53UL, linearResult[2].key());
size_t hit_count = 0;
for (const auto &result : linearResult) {
hit_count += knn_keys.count(result.key());
}
const float recall = static_cast<float>(hit_count) / topk;
EXPECT_GT(recall, 0.90f);
}
}
@ -689,7 +679,9 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) {
std::string vec_value;
ASSERT_EQ(0, searcher->get_vector(i, linearCtx, vec_value));
float vector_value = *(const float *)(vec_value.data());
ASSERT_GE(vec_value.size(), sizeof(float));
float vector_value = 0.0f;
std::memcpy(&vector_value, vec_value.data(), sizeof(vector_value));
ASSERT_EQ(vector_value, i);
}
@ -722,9 +714,15 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) {
ASSERT_EQ(topk, linearResult.size());
ASSERT_EQ(i, linearResult[0].key());
ASSERT_NE(knnResult[0].vector_string(), "");
float vector_value = *((float *)(knnResult[0].vector_string().data()));
ASSERT_EQ(vector_value, i);
const auto &vector_string = knnResult[0].vector_string();
ASSERT_GE(vector_string.size(), sizeof(float));
// DiskAnn is approximate, so the first KNN result is not guaranteed to
// be the exact query vector on every graph build. Verify that the fetched
// payload belongs to the returned key instead.
std::string expected_vector;
ASSERT_EQ(0, searcher->get_vector(knnResult[0].key(), linearCtx,
expected_vector));
ASSERT_EQ(vector_string, expected_vector);
}
}

14
tests/core/algorithm/ivf/ivf_searcher_test.cc
View File

@ -2552,11 +2552,23 @@ TEST_F(IVFSearcherTest, TestQuantizedPerCentroid) {
const IndexDocumentList &result = context->result(0);
EXPECT_EQ((size_t)topk, result.size());
uint64_t max_rank_error = 0;
uint64_t total_rank_error = 0;
for (size_t i = 0; i < topk; ++i) {
ASSERT_NEAR((uint64_t)(total - 1) - i, result[i].key(), 150);
const uint64_t expected_key = (uint64_t)(total - 1) - i;
const uint64_t actual_key = result[i].key();
const uint64_t rank_error = expected_key > actual_key
? expected_key - actual_key
: actual_key - expected_key;
max_rank_error = std::max(max_rank_error, rank_error);
total_rank_error += rank_error;
float expect = (float)result[i].key() * 500.0f * dimension_;
ASSERT_NEAR(expect, std::abs(result[i].score()), expect * 0.2 + 500000);
}
// K-means initialization and int8 score ties can vary by platform. Check
// the overall approximate ranking quality without rejecting one outlier.
EXPECT_LE(max_rank_error, 200U);
EXPECT_LE(total_rank_error, topk * 30U);
}
// batch bf serch