Skip to content

Build quantized QG using stored RaBitQ codes #101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gouyt13 merged 1 commit into from
Sep 13, 2026
Merged

Build quantized QG using stored RaBitQ codes #101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions docs/docs/index/qg.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
43 changes: 28 additions & 15 deletions include/rabitqlib/index/symqg/qg.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ class QuantizedQuery {
template <typename T = float>
class QuantizedGraph {
friend class QGBuilder;
friend struct QGConstructionTestAccess;

private:
size_t num_points_ = 0; // num points
Expand All @@ -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<T> centroid_; // rotated global centroid for qg-quant
size_t quantization_bits_ = 0; // 0: raw vectors, 4/8: packed RaBitQ vectors
std::vector<T> centroid_; // rotated global centroid for qg-quant
ex_ipfunc quantized_ip_func_ = nullptr;

Array<
Expand Down Expand Up @@ -111,13 +111,6 @@ class QuantizedGraph {
return reinterpret_cast<const T*>(&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);
}
Expand All @@ -128,6 +121,9 @@ class QuantizedGraph {

void prepare_query(const T*, std::vector<T>&, std::optional<QuantizedQuery<T>>&) const;

const T* prepare_build_query(PID, std::vector<T>&, std::optional<QuantizedQuery<T>>&)
const;

T point_distance(const T*, const QuantizedQuery<T>*, PID) const;

T quantized_distance(const QuantizedQuery<T>&, PID) const;
Expand Down Expand Up @@ -263,7 +259,6 @@ inline void QuantizedGraph<T>::validate_configuration() const {

template <typename T>
inline void QuantizedGraph<T>::copy_vectors(const T* data) {
build_data_ = data;
if (quantization_bits_ != 0) {
if constexpr (!std::is_same_v<T, float>) {
throw std::logic_error("qg-quant currently requires float data");
Expand Down Expand Up @@ -670,6 +665,22 @@ inline void QuantizedGraph<T>::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 <typename T>
inline const T* QuantizedGraph<T>::prepare_build_query(
PID id, std::vector<T>& rotated, std::optional<QuantizedQuery<T>>& 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 <typename T>
inline void QuantizedGraph<T>::find_candidates(
Expand All @@ -679,10 +690,12 @@ inline void QuantizedGraph<T>::find_candidates(
VisitedSet& vis,
const std::vector<uint32_t>& degrees
) const {
const T* query = get_build_vector(cur_id);
std::vector<T> rotated_query(padded_dim_);
std::optional<QuantizedQuery<T>> 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<T> q_obj(rotated_query.data(), padded_dim_);

// insert entry point to initialize search buffer
Expand Down Expand Up @@ -729,7 +742,7 @@ inline void QuantizedGraph<T>::update_qg(
std::vector<T> 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(
Expand All @@ -738,7 +751,7 @@ inline void QuantizedGraph<T>::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());
}
Expand Down
65 changes: 51 additions & 14 deletions include/rabitqlib/index/symqg/qg_builder.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <limits>
#include <mutex>
#include <numeric>
#include <optional>
#include <stdexcept>
#include <unordered_set>
#include <vector>
Expand All @@ -30,6 +31,8 @@ using CandidateList = std::vector<AnnCandidate<float>>;
*
*/
class QGBuilder {
friend struct QGConstructionTestAccess;

private:
QuantizedGraph<float>& qg_;
size_t ef_build_; // size of search pool for indexing
Expand Down Expand Up @@ -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<float> query(
qg_.centroid_.data(),
qg_.centroid_.data(),
qg_.padded_dim_,
qg_.metric_type_
);
float best = std::numeric_limits<float>::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(
Expand Down Expand Up @@ -144,10 +163,12 @@ inline void QGBuilder::add_pruned_edges(
nei_set.emplace(nei.id);
}

std::vector<float> reconstructed;
std::optional<QuantizedQuery<float>> 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()) {
Expand All @@ -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) {
Expand Down Expand Up @@ -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<float> reconstructed;
std::optional<QuantizedQuery<float>> prepared;
while (pruned_results.size() < degree_bound_ && start < poolsize) {
auto candidate_id = pool[start].id;

Expand All @@ -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
Expand All @@ -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) {
Expand Down Expand Up @@ -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<float> reconstructed;
std::optional<QuantizedQuery<float>> 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()
Expand All @@ -328,12 +362,14 @@ inline void QGBuilder::random_init() {
}
}

const float* cur_data = qg_.get_build_vector(i);
std::vector<float> reconstructed;
std::optional<QuantizedQuery<float>> 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)
);
}

Expand Down Expand Up @@ -388,14 +424,15 @@ inline void QGBuilder::graph_refine() {
for (auto& neighbor : new_result) {
ids.emplace(neighbor.id);
}
std::vector<float> reconstructed;
std::optional<QuantizedQuery<float>> prepared;
const float* source = qg_.prepare_build_query(i, reconstructed, prepared);
while (new_result.size() < degree_bound_) {
PID rand_id = rand_integer<PID>(0, static_cast<PID>(num_nodes_) - 1);
if (rand_id != static_cast<PID>(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);
}
Expand Down
27 changes: 27 additions & 0 deletions tests/python/test_symqg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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("<QI", payload) == (0x5147524142495451, 1)
# Existing v1 header, centroid, ExDataMap row, neighbor batches/IDs, FHT state.
row_bytes = 128 * bits // 8 + 8 + 128 * 4 + 256 + 128
assert len(payload) == 58 + 128 * 4 + 65 * row_bytes + 128 // 2
ids, distances = index.search(queries, 10, 64)
assert np.isfinite(distances).all()
assert all(len(set(row)) == 10 for row in ids)
loaded = SymqgIndex.load(str(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)
Loading