Fix SymphonyQG inner-product query correction

include/rabitqlib/index/query.hpp

@@ -22,9 +22,13 @@ class BatchQuery {
     Lut<T> 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<T>(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(); }

include/rabitqlib/index/symqg/qg.hpp

@@ -441,7 +441,7 @@ inline void QuantizedGraph<T>::search(
     std::vector<T> rotated_query(padded_dim_);
     std::optional<QuantizedQuery<T>> quantized_query;
     prepare_query(query, rotated_query, quantized_query);
-    BatchQuery<T> q_obj(rotated_query.data(), padded_dim_);
+    BatchQuery<T> q_obj(rotated_query.data(), padded_dim_, metric_type_);
 
     buffer::SearchBuffer<T> search_pool(ef_);
     // init search buffer
@@ -466,14 +466,14 @@ inline void QuantizedGraph<T>::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<T>::find_candidates(
     if (!is_quantized()) {
         rotator_->rotate(query, rotated_query.data());
     }
-    BatchQuery<T> q_obj(rotated_query.data(), padded_dim_);
+    BatchQuery<T> 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<T>::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);
         }
     }
 }

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)

tests/unit/rabitqlib/index/qg_test.cpp

@@ -11,6 +11,7 @@
 #include <vector>
 
 #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<float> centroid(dim, 0.125F), query(dim, 0.0625F);
+            std::vector<float> 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<float>(i % 3) - 1.0F) * 0.03125F;
+                std::fill_n(data.data() + i * dim, dim, value);
+            }
+            std::vector<char> batch(QGBatchDataMap<float>::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<float> : euclidean_sqr<float>;
+            const float vertex_distance = distance(query.data(), centroid.data(), dim);
+            BatchQuery<float> q_obj(query.data(), dim, metric);
+            q_obj.set_g_add(vertex_distance);
+            std::array<float, fastscan::kBatchSize> 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}) {