From de28f011102e096501871dd4d1a014e7b12b7e68 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 14 Sep 2026 19:27:08 +0800 Subject: [PATCH] Fix SymphonyQG inner-product query correction --- include/rabitqlib/index/query.hpp | 11 ++++--- include/rabitqlib/index/symqg/qg.hpp | 20 ++++++------ tests/python/test_symqg.py | 18 +++++++++++ tests/unit/rabitqlib/index/qg_test.cpp | 43 ++++++++++++++++++++++++++ 4 files changed, 78 insertions(+), 14 deletions(-) diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index 1c31313..c5c6a38 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -22,9 +22,13 @@ class BatchQuery { Lut lookup_table_; T G_add_ = 0; T G_k1xSumq_ = 0; // G_k1xSumq + MetricType metric_type_; public: - explicit BatchQuery(const T* rotated_query, size_t padded_dim) { + explicit BatchQuery( + const T* rotated_query, size_t padded_dim, MetricType metric_type = METRIC_L2 + ) + : metric_type_(metric_type) { lookup_table_ = std::move(Lut(rotated_query, padded_dim)); float c_1 = -((1 << 1) - 1) / 2.F; @@ -44,9 +48,8 @@ class BatchQuery { [[nodiscard]] T g_add() const { return G_add_; } void set_g_add(T dist) { - // For L2, dist is computed by euclidean_sqr(). - // For IP, dist is computed by dot_product_dis(), i.e. 1 - dot_product(). - G_add_ = dist; + // dist is squared L2 or 1 - dot_product; IP encoder factors already include 1. + G_add_ = metric_type_ == METRIC_IP ? dist - T{1} : dist; } [[nodiscard]] const uint8_t* lut() const { return lookup_table_.lut(); } diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index aa5656b..73da6a0 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -441,7 +441,7 @@ inline void QuantizedGraph::search( std::vector rotated_query(padded_dim_); std::optional> quantized_query; prepare_query(query, rotated_query, quantized_query); - BatchQuery q_obj(rotated_query.data(), padded_dim_); + BatchQuery q_obj(rotated_query.data(), padded_dim_, metric_type_); buffer::SearchBuffer search_pool(ef_); // init search buffer @@ -466,14 +466,14 @@ inline void QuantizedGraph::search( } vis->set(cur_node); - q_obj.set_g_add( - point_distance(query, quantized_query ? &*quantized_query : nullptr, cur_node) - ); + const T vertex_distance = + point_distance(query, quantized_query ? &*quantized_query : nullptr, cur_node); + q_obj.set_g_add(vertex_distance); scan_neighbors( q_obj, cur_node, est_dist.data(), search_pool, *vis, this->degree_bound_ ); - res_pool.insert(cur_node, q_obj.g_add()); + res_pool.insert(cur_node, vertex_distance); } update_results(res_pool, *vis, query, quantized_query ? &*quantized_query : nullptr); @@ -696,7 +696,7 @@ inline void QuantizedGraph::find_candidates( if (!is_quantized()) { rotator_->rotate(query, rotated_query.data()); } - BatchQuery q_obj(rotated_query.data(), padded_dim_); + BatchQuery q_obj(rotated_query.data(), padded_dim_, metric_type_); // insert entry point to initialize search buffer buffer::SearchBuffer tmp_pool(search_ef); @@ -711,12 +711,12 @@ inline void QuantizedGraph::find_candidates( } vis.set(cur_candi); auto cur_degree = degrees[cur_candi]; - q_obj.set_g_add( - point_distance(query, quantized_query ? &*quantized_query : nullptr, cur_candi) - ); + const T vertex_distance = + point_distance(query, quantized_query ? &*quantized_query : nullptr, cur_candi); + q_obj.set_g_add(vertex_distance); scan_neighbors(q_obj, cur_candi, est_dist.data(), tmp_pool, vis, cur_degree); if (cur_candi != cur_id) { - results.emplace_back(cur_candi, q_obj.g_add()); + results.emplace_back(cur_candi, vertex_distance); } } } diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 3c321df..0e88356 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -182,3 +182,21 @@ def test_code_only_build_preserves_v1_format(bits, metric, tmp_path): loaded_ids, loaded_distances = loaded.search(queries, 10, 64) np.testing.assert_array_equal(loaded_ids, ids) np.testing.assert_array_equal(loaded_distances, distances) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_constant_vectors_preserve_returned_distance(metric, tmp_path): + """Neighbor corrections must not shift the vertex distances returned to Python.""" + data = np.full((65, 64), 0.125, dtype=np.float32) + queries = np.full((2, 64), 0.0625, dtype=np.float32) + index = SymqgIndex(64, max_degree=32, metric=metric) + index.build(data, ef_construction=64) + ids, distances = index.search(queries, k=5, ef=64) + expected = 0.5 if metric == "ip" else 0.25 + np.testing.assert_allclose(distances, expected, rtol=0, atol=1e-6) + assert np.all(ids < len(data)) + path = str(tmp_path / "constant.symqg") + index.save(path) + loaded_ids, loaded_distances = SymqgIndex.load(path).search(queries, k=5, ef=64) + np.testing.assert_array_equal(loaded_ids, ids) + np.testing.assert_allclose(loaded_distances, expected, rtol=0, atol=1e-6) diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index b0c49fc..e395e6d 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -11,6 +11,7 @@ #include #include "rabitqlib/index/symqg/qg_builder.hpp" +#include "rabitqlib/utils/cpu_features.hpp" namespace rabitqlib::symqg { struct QGConstructionTestAccess { @@ -61,6 +62,48 @@ struct QGConstructionTestAccess { namespace { +TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { + if (!cpu::has_avx2()) { + GTEST_SKIP() << "FastScan requires AVX2/FMA"; + } + for (size_t dim : {64U, 1024U, 1088U}) { + for (auto metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(::testing::Message() << dim << "/" << metric); + std::vector centroid(dim, 0.125F), query(dim, 0.0625F); + std::vector data(fastscan::kBatchSize * dim); + for (size_t i = 0; i < fastscan::kBatchSize; ++i) { + // Collinear residuals have exact codes, including the zero-residual case. + const float value = 0.125F + (static_cast(i % 3) - 1.0F) * 0.03125F; + std::fill_n(data.data() + i * dim, dim, value); + } + std::vector batch(QGBatchDataMap::data_bytes(dim)); + quant::quantize_qg_batch( + data.data(), + centroid.data(), + fastscan::kBatchSize, + dim, + batch.data(), + metric + ); + const auto distance = + metric == METRIC_IP ? dot_product_dis : euclidean_sqr; + const float vertex_distance = distance(query.data(), centroid.data(), dim); + BatchQuery q_obj(query.data(), dim, metric); + q_obj.set_g_add(vertex_distance); + std::array estimates{}; + qg_batch_estdist(batch.data(), q_obj, dim, estimates.data()); + EXPECT_FLOAT_EQ( + q_obj.g_add(), metric == METRIC_IP ? vertex_distance - 1 : vertex_distance + ); + for (size_t i = 0; i < fastscan::kBatchSize; ++i) { + EXPECT_NEAR( + estimates[i], distance(query.data(), data.data() + i * dim, dim), 1e-5F + ); + } + } + } +} + TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { constexpr size_t kCount = 65, kDim = 65; for (auto metric : {METRIC_L2, METRIC_IP}) {