From 63039943b8e911fcbaaaf27d4e83742f5191a55b Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sun, 13 Sep 2026 17:12:57 +0800 Subject: [PATCH] Build quantized QG using stored RaBitQ codes --- docs/docs/index/qg.md | 21 +++-- include/rabitqlib/index/symqg/qg.hpp | 43 +++++---- include/rabitqlib/index/symqg/qg_builder.hpp | 65 +++++++++++--- tests/python/test_symqg.py | 27 ++++++ tests/unit/rabitqlib/index/qg_test.cpp | 91 ++++++++++++++++++++ 5 files changed, 212 insertions(+), 35 deletions(-) diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 8ff9f2d..7c20925 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -14,12 +14,21 @@ code. All codes use the dataset's global centroid and are not combined with the one-bit neighbor codes. Pass `quantization_bits` as `4` or `8` to select QG-quant, or leave it at `0` for vanilla QG. -During graph construction, candidate discovery uses a floating-point source -vector against the stored qg-quant candidate codes. Those estimated distances -are retained for candidate ordering and source-to-candidate pruning terms, while -candidate-to-candidate pruning comparisons and graph refinement use the available -raw build vectors. The raw vectors are not retained in the completed qg-quant -index. +After computing the centroid and encoding the input, quantized graph construction +uses only the existing stored RaBitQ representation. A source is temporarily +reconstructed in rotated coordinates, then the existing RaBitQ estimator scores +target codes. This applies to candidate discovery, initial edges, pairwise pruning, +refinement, and fallback edges. Reverse edges are rescored in their own direction +because these estimates need not be symmetric. Entry-point selection scores the +stored codes against the centroid. Neighbor FastScan encoding continues to use +code reconstructions. + +No raw input pointer or full reconstructed-vector cache is retained. The caller +can release the input after the synchronous C++ `QGBuilder` constructor finishes. +The existing pruning rules operate on the estimated distances; this does not make +them exact distances between reconstructed vectors. External queries, the 4/8-bit +encoding and correction factors, neighbor FastScan layout, and version-1 file +format are unchanged. Raw QG continues to construct using its owned raw vectors. Memory and performance depend on the dimension, degree, build window, and search window. See `sample/cpp/symqg_indexing.cpp` and diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 24b2f5f..aa5656b 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -62,6 +62,7 @@ class QuantizedQuery { template class QuantizedGraph { friend class QGBuilder; + friend struct QGConstructionTestAccess; private: size_t num_points_ = 0; // num points @@ -72,9 +73,8 @@ class QuantizedGraph { PID entry_point_ = 0; // Entry point of graph MetricType metric_type_ = MetricType::METRIC_L2; RotatorType rotator_type_ = RotatorType::FhtKacRotator; - size_t quantization_bits_ = 0; // 0: raw vectors, 4/8: packed RaBitQ vectors - const T* build_data_ = nullptr; // non-owning, used only while QGBuilder runs - std::vector centroid_; // rotated global centroid for qg-quant + size_t quantization_bits_ = 0; // 0: raw vectors, 4/8: packed RaBitQ vectors + std::vector centroid_; // rotated global centroid for qg-quant ex_ipfunc quantized_ip_func_ = nullptr; Array< @@ -111,13 +111,6 @@ class QuantizedGraph { return reinterpret_cast(&data_.at(row_offset_ * data_id)); } - [[nodiscard]] const T* get_build_vector(PID data_id) const { - if (quantization_bits_ == 0) { - return get_vector(data_id); - } - return build_data_ + (dim_ * data_id); - } - [[nodiscard]] char* get_quantized_vector(PID data_id) { return &data_.at(row_offset_ * data_id); } @@ -128,6 +121,9 @@ class QuantizedGraph { void prepare_query(const T*, std::vector&, std::optional>&) const; + const T* prepare_build_query(PID, std::vector&, std::optional>&) + const; + T point_distance(const T*, const QuantizedQuery*, PID) const; T quantized_distance(const QuantizedQuery&, PID) const; @@ -263,7 +259,6 @@ inline void QuantizedGraph::validate_configuration() const { template inline void QuantizedGraph::copy_vectors(const T* data) { - build_data_ = data; if (quantization_bits_ != 0) { if constexpr (!std::is_same_v) { throw std::logic_error("qg-quant currently requires float data"); @@ -670,6 +665,22 @@ inline void QuantizedGraph::reconstruct_quantized_vector(PID data_id, T* reco ); } +// Construction sources come from owned raw rows or from the existing RaBitQ codes. +// Reconstructed sources are already rotated; never pass them through prepare_query. +template +inline const T* QuantizedGraph::prepare_build_query( + PID id, std::vector& rotated, std::optional>& prepared +) const { + if (is_quantized()) { + rotated.resize(padded_dim_); + reconstruct_quantized_vector(id, rotated.data()); + prepared.emplace(rotated.data(), centroid_.data(), padded_dim_, metric_type_); + return rotated.data(); + } + prepared.reset(); + return get_vector(id); +} + // find candidate neighbors for cur_id, exclude the vertex itself template inline void QuantizedGraph::find_candidates( @@ -679,10 +690,12 @@ inline void QuantizedGraph::find_candidates( VisitedSet& vis, const std::vector& degrees ) const { - const T* query = get_build_vector(cur_id); std::vector rotated_query(padded_dim_); std::optional> quantized_query; - prepare_query(query, rotated_query, quantized_query); + const T* query = prepare_build_query(cur_id, rotated_query, quantized_query); + if (!is_quantized()) { + rotator_->rotate(query, rotated_query.data()); + } BatchQuery q_obj(rotated_query.data(), padded_dim_); // insert entry point to initialize search buffer @@ -729,7 +742,7 @@ inline void QuantizedGraph::update_qg( std::vector rotated_centroid(padded_dim_); for (size_t i = 0; i < cur_degree; ++i) { if (quantization_bits_ == 0) { - const T* neighbor_vec = get_build_vector(new_neighbors[i].id); + const T* neighbor_vec = get_vector(new_neighbors[i].id); this->rotator_->rotate(neighbor_vec, &rotated_data[i * padded_dim_]); } else { reconstruct_quantized_vector( @@ -738,7 +751,7 @@ inline void QuantizedGraph::update_qg( } } if (quantization_bits_ == 0) { - this->rotator_->rotate(get_build_vector(cur_id), rotated_centroid.data()); + this->rotator_->rotate(get_vector(cur_id), rotated_centroid.data()); } else { reconstruct_quantized_vector(cur_id, rotated_centroid.data()); } diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 19f86a4..bf07598 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -30,6 +31,8 @@ using CandidateList = std::vector>; * */ class QGBuilder { + friend struct QGConstructionTestAccess; + private: QuantizedGraph& qg_; size_t ef_build_; // size of search pool for indexing @@ -83,17 +86,33 @@ class QGBuilder { qg_.set_quantization_centroid(centroid.data()); qg_.copy_vectors(data); - PID entry_point = exact_nn( - data, centroid.data(), num_nodes_, dim_, num_threads_, qg_.raw_dist_func_ - ); + PID entry_point = 0; + if (qg_.is_quantized()) { + QuantizedQuery query( + qg_.centroid_.data(), + qg_.centroid_.data(), + qg_.padded_dim_, + qg_.metric_type_ + ); + float best = std::numeric_limits::max(); + for (PID id = 0; id < num_nodes_; ++id) { + const float distance = qg_.quantized_distance(query, id); + if (distance < best) { + best = distance; + entry_point = id; + } + } + } else { + entry_point = exact_nn( + data, centroid.data(), num_nodes_, dim_, num_threads_, qg_.raw_dist_func_ + ); + } qg_.set_ep(entry_point); random_init(); } - ~QGBuilder() { qg_.build_data_ = nullptr; } - void build(size_t num_iter = 3) { if (num_iter < 2) { throw std::invalid_argument( @@ -144,10 +163,12 @@ inline void QGBuilder::add_pruned_edges( nei_set.emplace(nei.id); } + std::vector reconstructed; + std::optional> prepared; while (new_result.size() < degree_bound_ && start < pruned_list.size()) { const auto& cur = pruned_list[start]; bool occlude = false; - const float* cur_data = qg_.get_build_vector(cur.id); + const float* cur_data = qg_.prepare_build_query(cur.id, reconstructed, prepared); float dik_sqr = cur.distance; if (nei_set.find(cur.id) != nei_set.end()) { @@ -160,7 +181,7 @@ inline void QGBuilder::add_pruned_edges( break; } float djk_sqr = - qg_.raw_dist_func_(qg_.get_build_vector(nei.id), cur_data, dim_); + qg_.point_distance(cur_data, prepared ? &*prepared : nullptr, nei.id); float cosine = (dik_sqr + dij_sqr - djk_sqr) / (2 * std::sqrt(dij_sqr * dik_sqr)); if (cosine > threshold) { @@ -199,6 +220,8 @@ inline void QGBuilder::heuristic_prune( ); // bool vector to record if this neighbor is pruned size_t start = 0; // start position + std::vector reconstructed; + std::optional> prepared; while (pruned_results.size() < degree_bound_ && start < poolsize) { auto candidate_id = pool[start].id; @@ -209,7 +232,8 @@ inline void QGBuilder::heuristic_prune( } pruned_results.emplace_back(pool[start]); // add current candidate to result - const float* data_j = qg_.get_build_vector(candidate_id); + const float* data_j = + qg_.prepare_build_query(candidate_id, reconstructed, prepared); // i : current vertex // j : neighbor added in this iter @@ -219,7 +243,8 @@ inline void QGBuilder::heuristic_prune( continue; } float dik = pool[k].distance; - auto djk = qg_.raw_dist_func_(data_j, qg_.get_build_vector(pool[k].id), dim_); + auto djk = + qg_.point_distance(data_j, prepared ? &*prepared : nullptr, pool[k].id); if (djk < dik) { if (refine && pruned_neighbors_[cur_id].size() < kMaxPrunedSize) { @@ -305,6 +330,15 @@ inline void QGBuilder::add_reverse_edges(bool refine) { #pragma omp parallel for schedule(dynamic) for (PID data_id = 0; data_id < num_nodes_; ++data_id) { CandidateList& tmp_pool = reverse_buffer[data_id]; + if (qg_.is_quantized() && !tmp_pool.empty()) { + // RaBitQ estimates are directional: score destination -> source afresh. + std::vector reconstructed; + std::optional> prepared; + qg_.prepare_build_query(data_id, reconstructed, prepared); + for (auto& candidate : tmp_pool) { + candidate.distance = qg_.quantized_distance(*prepared, candidate.id); + } + } tmp_pool.reserve(tmp_pool.size() + degree_bound_); tmp_pool.insert( tmp_pool.end(), new_neighbors_[data_id].begin(), new_neighbors_[data_id].end() @@ -328,12 +362,14 @@ inline void QGBuilder::random_init() { } } - const float* cur_data = qg_.get_build_vector(i); + std::vector reconstructed; + std::optional> prepared; + const float* cur_data = qg_.prepare_build_query(i, reconstructed, prepared); new_neighbors_[i].reserve(degree_bound_); for (PID cur_neigh : neighbor_set) { new_neighbors_[i].emplace_back( cur_neigh, - qg_.raw_dist_func_(cur_data, qg_.get_build_vector(cur_neigh), dim_) + qg_.point_distance(cur_data, prepared ? &*prepared : nullptr, cur_neigh) ); } @@ -388,14 +424,15 @@ inline void QGBuilder::graph_refine() { for (auto& neighbor : new_result) { ids.emplace(neighbor.id); } + std::vector reconstructed; + std::optional> prepared; + const float* source = qg_.prepare_build_query(i, reconstructed, prepared); while (new_result.size() < degree_bound_) { PID rand_id = rand_integer(0, static_cast(num_nodes_) - 1); if (rand_id != static_cast(i) && ids.find(rand_id) == ids.end()) { new_result.emplace_back( rand_id, - qg_.raw_dist_func_( - qg_.get_build_vector(rand_id), qg_.get_build_vector(i), dim_ - ) + qg_.point_distance(source, prepared ? &*prepared : nullptr, rand_id) ); ids.emplace(rand_id); } diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 566dca7..3c321df 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -1,5 +1,7 @@ """Tests for SymqgIndex: construction, search, properties, error handling, save/load.""" +import struct + import numpy as np import pytest from conftest import DIM, N_QUERIES, N_VECTORS, brute_force_knn, recall_at_k @@ -155,3 +157,28 @@ def test_save_load_roundtrip(built_symqg, query_data, tmp_path): ids_load, dists_load = loaded.search(query_data, k=_TOPK, ef=_EF) np.testing.assert_array_equal(ids_orig, ids_load) np.testing.assert_allclose(dists_orig, dists_load, rtol=1e-5) + + +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_code_only_build_preserves_v1_format(bits, metric, tmp_path): + rng = np.random.default_rng(851) + data = rng.normal(size=(65, 65)).astype(np.float32) + queries = rng.normal(size=(3, 65)).astype(np.float32) + index = SymqgIndex(65, max_degree=32, metric=metric, quantization_bits=bits) + index.build(data, ef_construction=64) + data.fill(np.nan) + path = tmp_path / "code_only.index" + index.save(str(path)) + payload = path.read_bytes() + assert struct.unpack_from(" #include #include +#include #include #include #include @@ -12,8 +13,98 @@ #include "rabitqlib/index/symqg/qg_builder.hpp" namespace rabitqlib::symqg { +struct QGConstructionTestAccess { + static auto codes(const QuantizedGraph& graph) { + std::vector result; + for (PID id = 0; id < graph.num_points_; ++id) { + const char* code = graph.get_quantized_vector(id); + result.insert(result.end(), code, code + graph.batch_data_offset_); + } + return result; + } + static const auto& neighbors(const QGBuilder& builder) { + return builder.new_neighbors_; + } + static std::array estimate( + const QuantizedGraph& graph, PID source, PID target + ) { + std::vector query(graph.padded_dim_); + graph.reconstruct_quantized_vector(source, query.data()); + ConstExDataMap code( + graph.get_quantized_vector(target), graph.padded_dim_, graph.quantization_bits_ + ); + const double midpoint = ((1U << graph.quantization_bits_) - 1) / 2.0; + double dot = 0, add = code.f_add_ex(), magnitude = std::abs(add); + for (size_t d = 0; d < query.size(); ++d) { + const auto byte = graph.quantization_bits_ == 8 + ? code.ex_code()[d] + : code.ex_code()[(d / 16) * 8 + d % 8]; + const auto value = graph.quantization_bits_ == 8 + ? byte + : ((d % 16 < 8) ? byte & 15U : byte >> 4); + dot += query[d] * (value - midpoint); + // The existing kernel subtracts uncentered float sums; include both + // operands in its rounding bound when the final distance is small. + magnitude += std::abs(code.f_rescale_ex() * query[d]) * (value + midpoint); + const double residual = static_cast(query[d]) - graph.centroid_[d]; + const double term = graph.metric_type_ == METRIC_L2 + ? residual * residual + : -static_cast(query[d]) * graph.centroid_[d]; + add += term; + magnitude += std::abs(term); + } + return { + add + code.f_rescale_ex() * dot, + 8 * std::numeric_limits::epsilon() * std::max(1.0, magnitude)}; + } +}; + namespace { +TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { + constexpr size_t kCount = 65, kDim = 65; + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {4U, 8U}) { + for (bool constant : {false, true}) { + SCOPED_TRACE( + ::testing::Message() << metric << "/" << bits << "/" << constant + ); + std::vector data(kCount * kDim, 2.0F); + if (!constant) { + for (size_t i = 0; i < data.size(); ++i) { + data[i] += std::sin(static_cast(i) * 0.13F) * + static_cast(1 + (i / kDim) % 4); + } + } + QuantizedGraph graph( + kCount, kDim, 32, metric, RotatorType::FhtKacRotator, bits + ); + QGBuilder builder(graph, 64, data.data(), 1); + const auto codes = QGConstructionTestAccess::codes(graph); + std::fill( + data.begin(), data.end(), std::numeric_limits::quiet_NaN() + ); + std::vector().swap(data); + builder.build(); + EXPECT_EQ(QGConstructionTestAccess::codes(graph), codes); + EXPECT_FALSE(builder.check_dup()); + EXPECT_FLOAT_EQ(builder.avg_degree(), 32); + const auto& neighbors = QGConstructionTestAccess::neighbors(builder); + for (PID id = 0; id < kCount; ++id) { + ASSERT_EQ(neighbors[id].size(), 32U); + for (const auto& neighbor : neighbors[id]) { + EXPECT_NE(neighbor.id, id); + ASSERT_LT(neighbor.id, kCount); + const auto expected = + QGConstructionTestAccess::estimate(graph, id, neighbor.id); + EXPECT_NEAR(neighbor.distance, expected[0], expected[1]); + } + } + } + } + } +} + TEST(QuantizedGraphConfigurationTest, RejectsDegreeNotAlignedForFastScan) { EXPECT_THROW( (QuantizedGraph(64, 64, 16, METRIC_L2, RotatorType::MatrixRotator)),