From 0b2aec10cd1ca0a81aaf1bc258e2d00d31acacce Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 15 Sep 2026 13:16:59 +0800 Subject: [PATCH] feat(symqg): use PiPNN for initialization by default - Encapsulate PiPNN initialization in QGBuilder with one refinement pass - Add init selection to C++ and Python APIs, retaining random initialization - Update documentation, examples, --- docs/docs/index/qg.md | 191 +++---- .../rabitqlib/index/symqg/detail/pipnn.hpp | 482 ++++++++++++++++++ include/rabitqlib/index/symqg/qg_builder.hpp | 195 +++++-- python_bindings/symqg_bindings.cpp | 27 +- sample/cpp/symqg_indexing.cpp | 41 +- sample/python/symqg_indexing.py | 16 +- tests/python/test_symqg.py | 51 ++ tests/unit/rabitqlib/index/pipnn_test.cpp | 120 +++++ tests/unit/rabitqlib/index/qg_test.cpp | 276 ++++++++++ 9 files changed, 1232 insertions(+), 167 deletions(-) create mode 100644 include/rabitqlib/index/symqg/detail/pipnn.hpp create mode 100644 tests/unit/rabitqlib/index/pipnn_test.cpp diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 7c20925..2d1e15b 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -1,133 +1,110 @@ # QG + RaBitQ (SymphonyQG) -[QG](https://medium.com/@masajiro.iwasaki/fusion-of-graph-based-indexing-and-product-quantization-for-ann-search-7d1f0336d0d0) -is a graph-based index originating from the -[NGT library](https://github.com/yahoojapan/NGT). This implementation comes -from the [SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) project. For -each vertex vanilla QG stores the raw vector, a fixed-size neighbor list, and batched -one-bit RaBitQ data for those neighbors. This layout uses more memory than the -raw vectors alone, but lets graph traversal estimate a group of neighbor -distances with FastScan while computing exact distances for visited vertices. - -QG-quant replaces each raw vector with an independent packed 4- or 8-bit RaBitQ -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. - -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 -`sample/cpp/symqg_querying.cpp` for complete programs. +[SymphonyQG](https://dl.acm.org/doi/abs/10.1145/3709730) combines the graph search +of [QG](https://medium.com/@masajiro.iwasaki/fusion-of-graph-based-indexing-and-product-quantization-for-ann-search-7d1f0336d0d0), +from [NGT](https://github.com/yahoojapan/NGT), with batched RaBitQ distance estimation. +FastScan estimates neighbor distances; visited vertices are scored using stored +raw vectors or quantized codes. + +Set `quantization_bits=0` for raw vectors (default), or `4`/`8` for QG-quant. +Quantized vectors share a rotated global centroid. Refinement temporarily +reconstructs source vectors and estimates distances to stored target codes; +reverse edges are rescored because these estimates are directional. Raw refinement +uses owned raw vectors. No full reconstructed-vector cache is retained. ## Index Construction -We build the QG by iteratively refining the graph structure. -Since the QG is more complicated than other indices, we need a QGBuilder to help us construct the index. +[PiPNN](https://dl.acm.org/doi/abs/10.1145/3770855.3817891) is the default initializer +for fast indexing, followed by one SymphonyQG refinement iteration. Random +initialization remains available and uses three iterations by default. -At the beginning, we need to intialize a QG and a QGBuilder by following construtor. -```cpp -QuantizedGraph::QuantizedGraph( - size_t num, - size_t dim, - size_t max_deg, - MetricType metric_type = METRIC_L2, - RotatorType rotator_type = RotatorType::FhtKacRotator, - size_t quantization_bits = 0 - ); - -QGBuilder::QGBuilder( - QuantizedGraph& index, - uint32_t ef_build, - const float* data, - size_t num_threads = std::numeric_limits::max() - ) -``` -- **num**: Number of vertices (vectors) in the dataset. -- **dim**: Dimension of the dataset. -- **max_deg**: Degree bound of QG, must be a multiple of 32. -- **quantization_bits**: `0` for vanilla QG, or `4`/`8` for QG-quant. -- **index**: Previously initialized QG. -- **ef_build**: Search window size during indexing. -- **data**: Pointer to the dataset, size of num * dim. -- **num_threads**: Number of threads to use (default: std::numeric_limits::max(), which auto-selects). -```cpp -size_t rows = 1000000; -size_t cols = 128; -size_t degree = 32; -size_t ef = 200; +### Python -std::vector data(rows * cols); // populate with the dataset +```python +from rabitqlib import SymqgIndex -QuantizedGraph qg(rows, cols, degree); +# data and queries are float32 arrays with shape (count, dim). +index = SymqgIndex(dim=data.shape[1], max_degree=32, metric="l2", quantization_bits=0) +index.build(data, ef_construction=200, num_threads=32, init="pipnn") +index.save("qg_example.index") -QGBuilder builder(qg, ef, data.data()); +loaded = SymqgIndex.load("qg_example.index") +ids, distances = loaded.search(queries, k=10, ef=100, num_threads=1) ``` -Then, we can use the builder to construct the index. Then we can save the index. +`init` defaults to `"pipnn"`; use `"random"` for random initialization. +Supported metrics are `"l2"` and `"ip"`. `max_degree` must be a multiple of 32 +and smaller than the point count. `ef_construction` controls the build search +window; `ef` controls the query search window. Python defaults to one thread. + +### C++ + +The C++ API uses `rabitqlib::symqg::QuantizedGraph` and `QGBuilder`: + ```cpp -builder.build(); // build index interatively +QuantizedGraph( + size_t num, size_t dim, size_t max_deg, + MetricType metric_type = METRIC_L2, + RotatorType rotator_type = RotatorType::FhtKacRotator, + size_t quantization_bits = 0 +); +QGBuilder( + QuantizedGraph& index, uint32_t ef_build, const float* data, + size_t num_threads = std::numeric_limits::max(), + QGInitialization init = QGInitialization::PiPNN +); +``` + +`data` contains `num * dim` floats; `max_deg` has the same constraints as Python's +`max_degree`. C++ defaults to all available threads. Pass +`QGInitialization::Random` as the final builder argument to use random initialization. +The builder handles initialization internally. -const char* index_file = "./qg_example.index"; -qg.save(index_file); // save index +```cpp +using namespace rabitqlib::symqg; + +// data contains rows * cols floats. +QuantizedGraph qg(rows, cols, 32); +{ + QGBuilder builder(qg, 200, data.data(), 32, QGInitialization::PiPNN); + builder.build(); +} // release builder scratch before saving +qg.save("qg_example.index"); ``` +The caller can release input vectors after the `QGBuilder` constructor returns. +Python retains the caller's array during `build`. Complete the build before +querying or saving. Initialization does not change query-distance conventions or +save/load formats. + ### Data Layout -Each indexed element is stored in the following layout. -``` -[Raw data vector] -[Batch data for QG] -[Edges] -``` +Each row contains: -For QG-quant, the first block becomes `[Packed 4/8-bit RaBitQ code + factors]`. -The index also stores one rotated global centroid shared by all rows. +```text +[Raw vector or packed 4/8-bit code + factors] +[One-bit neighbor codes + factors] +[Neighbor IDs] +``` -`Batch data for QG` contains one-bit codes and estimator factors for the -element's neighbors, organized in FastScan batches of 32. Consequently, -`max_deg` must be a multiple of 32. +Neighbor codes use FastScan batches of 32, which determines the degree alignment. +Quantized storage reduces each vector's size but not its neighborhood codes; +those codes and temporary refinement pools still consume substantial memory. ## Querying -For querying, code is pretty simple. -```cpp -void QuantizedGraph::search( - const T* __restrict__ query, - uint32_t k, - uint32_t* __restrict__ results, - T* __restrict__ dists); -``` -- **query**: Query vector. -- **k**: Top-k. -- **results**: Result buffer, size of k. -- **dists**: Distance buffer, size of k. -Then we can use a pre-constructed index to search. +C++ search accepts one vector in the original input dimension and writes `k` IDs +and distances: + ```cpp QuantizedGraph qg; -qg.load("./qg_example.index"); // load pre-constructed index - +qg.load("qg_example.index"); +qg.set_ef(100); -size_t ef = 100; -size_t topk = 10; -std::vector results(topk); // result buffer -std::vector dists(topk); // distance buffer -std::vector query(cols); // populate with a query vector - -qg.set_ef(ef); // set search window size -qg.search(query.data(), topk, results.data(), dists.data()); +std::vector ids(10); +std::vector distances(10); +qg.search(query.data(), 10, ids.data(), distances.data()); ``` + +See `sample/cpp/symqg_indexing.cpp`, `sample/cpp/symqg_querying.cpp`, and their +Python counterparts for complete examples. diff --git a/include/rabitqlib/index/symqg/detail/pipnn.hpp b/include/rabitqlib/index/symqg/detail/pipnn.hpp new file mode 100644 index 0000000..51a3fb1 --- /dev/null +++ b/include/rabitqlib/index/symqg/detail/pipnn.hpp @@ -0,0 +1,482 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/defines.hpp" +#include "rabitqlib/utils/buffer.hpp" +#include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" + +namespace rabitqlib::symqg::detail { + +// A temporary graph. No vectors, distances, or clustering scratch survive export. +struct InitialGraph { + std::vector offsets; + std::vector neighbors; +}; + +namespace pipnn_impl { +constexpr size_t kLeafSize = 1024; +constexpr size_t kTileSize = 2048; +constexpr size_t kHashBits = 12; +constexpr size_t kLocalDegree = 2; +constexpr float kAlpha = 1.1F; +using Bucket = std::vector; + +struct Candidate { + PID id; + float distance; + uint16_t hash; +}; + +struct BucketJob { + Bucket ids; + size_t depth = 0; + uint32_t seed = 555; +}; + +// Non-owning slices of a single build allocation. Keeping the large matrices out +// of worker heaps avoids retaining their high-water marks after seeding finishes. +struct Workspace { + float* points; + float* distances; + float* norms; + std::vector> candidates; +}; + +struct ScratchPool { + size_t point_size; + size_t distance_size; + size_t stride; + std::unique_ptr storage; + + ScratchPool(size_t count, size_t dim, size_t threads) + : point_size(std::min(count, kTileSize) * dim) + , distance_size(std::max( + std::min(count, kTileSize) * std::min(count, size_t{512}), + std::min(count, kLeafSize) * std::min(count, kLeafSize) + )) + , stride(point_size + distance_size + std::min(count, kTileSize)) + // Workers initialize their own slices before reading, preserving NUMA + // first-touch placement without a serial zero-fill of the entire pool. + , storage(new float[threads * stride]) {} + + Workspace worker(size_t id) { + float* base = storage.get() + id * stride; + return {base, base + point_size, base + point_size + distance_size, {}}; + } +}; + +inline void gather(const float* data, size_t dim, const Bucket& ids, Workspace& work) { + for (size_t i = 0; i < ids.size(); ++i) { + std::copy_n(data + ids[i] * dim, dim, work.points + i * dim); + } +} + +inline void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric) { + const auto count = static_cast(size); + RowMajorMatrixMap points(work.points, count, static_cast(dim)); + RowMajorMatrixMap distances(work.distances, count, count); + VectorMap norms(work.norms, count); + distances.setZero(); + distances.template selfadjointView().rankUpdate(points); + norms = points.rowwise().squaredNorm(); + for (Eigen::Index i = 0; i < count; ++i) { + for (Eigen::Index j = 0; j <= i; ++j) { + const float dot = distances(i, j); + distances(i, j) = + metric == METRIC_L2 ? std::max(0.0F, norms[i] + norms[j] - 2 * dot) : -dot; + } + } +} + +inline std::vector partition( + const float* data, + size_t dim, + const BucketJob& job, + MetricType metric, + size_t threads, + ScratchPool& scratch +) { + const auto& ids = job.ids; + const size_t leader_count = std::min( + ids.size(), std::clamp(job.depth == 0 ? 512 : ids.size() / 200, 3, 512) + ); + const size_t fanout = std::min( + leader_count, + job.depth == 0 ? size_t{10} + : job.depth == 1 ? size_t{3} + : size_t{1} + ); + std::mt19937 random(job.seed); + Bucket leaders; + leaders.reserve(leader_count); + std::sample(ids.begin(), ids.end(), std::back_inserter(leaders), leader_count, random); + RowMajorMatrix leader_data(leader_count, dim); + for (size_t i = 0; i < leader_count; ++i) { + std::copy_n(data + leaders[i] * dim, dim, leader_data.data() + i * dim); + } + const Vector leader_norms = leader_data.rowwise().squaredNorm(); + std::vector assignments(ids.size() * fanout); + const auto assign_tile = [&](size_t begin, Workspace& work) { + const size_t count = std::min(kTileSize, ids.size() - begin); + RowMajorMatrixMap points( + work.points, static_cast(count), static_cast(dim) + ); + RowMajorMatrixMap distances( + work.distances, + static_cast(count), + static_cast(leader_count) + ); + VectorMap norms(work.norms, static_cast(count)); + for (size_t i = 0; i < count; ++i) { + std::copy_n(data + ids[begin + i] * dim, dim, points.data() + i * dim); + } + distances.noalias() = points * leader_data.transpose(); + norms = points.rowwise().squaredNorm(); + work.candidates.resize(leader_count); + for (size_t i = 0; i < count; ++i) { + for (size_t j = 0; j < leader_count; ++j) { + const float dot = + distances(static_cast(i), static_cast(j)); + const float distance = + metric == METRIC_L2 + ? std::max(0.0F, norms.data()[i] + leader_norms.data()[j] - 2 * dot) + : -dot; + work.candidates[j] = {distance, static_cast(j)}; + } + std::partial_sort( + work.candidates.begin(), + work.candidates.begin() + static_cast(fanout), + work.candidates.end() + ); + for (size_t j = 0; j < fanout; ++j) { + assignments[(begin + i) * fanout + j] = work.candidates[j].second; + } + } + }; + if (job.depth == 0) { +#pragma omp parallel num_threads(threads) + { + auto work = scratch.worker(static_cast(omp_get_thread_num())); +#pragma omp for schedule(static) + for (size_t begin = 0; begin < ids.size(); begin += kTileSize) { + assign_tile(begin, work); + } + } + } else { + // Already inside the bucket-level team. Eigen must see that team rather + // than a serialized nested region when deciding its own parallelism. + auto work = scratch.worker(static_cast(omp_get_thread_num())); + for (size_t begin = 0; begin < ids.size(); begin += kTileSize) { + assign_tile(begin, work); + } + } + std::vector children(leader_count); + for (size_t i = 0; i < ids.size(); ++i) { + for (size_t j = 0; j < fanout; ++j) { + children[assignments[i * fanout + j]].push_back(ids[i]); + } + } + return children; +} + +inline std::vector cluster( + const float* data, + size_t count, + size_t dim, + MetricType metric, + size_t threads, + ScratchPool& scratch +) { + Bucket all(count); + std::iota(all.begin(), all.end(), PID{0}); + // The root parallelizes tiles; deeper levels parallelize independent buckets. + std::vector active; + active.push_back({std::move(all), 0, 555}); + std::vector leaves; + while (!active.empty()) { + std::vector> next(active.size()); + std::vector> finished(active.size()); +#pragma omp parallel for num_threads(active.front().depth == 0 ? 1 : threads) \ + schedule(dynamic) + for (size_t i = 0; i < active.size(); ++i) { + auto& job = active[i]; + if (job.ids.size() <= kLeafSize) { + finished[i].push_back(std::move(job.ids)); + continue; + } + auto children = partition(data, dim, job, metric, threads, scratch); + std::sort(children.begin(), children.end(), [](const auto& a, const auto& b) { + return a.size() < b.size(); + }); + Bucket small; + for (size_t j = 0; j < children.size(); ++j) { + auto& child = children[j]; + if (child.empty()) { + continue; + } + if (child.size() < kLeafSize / 10) { + small.insert(small.end(), child.begin(), child.end()); + if (small.size() < kLeafSize / 10) { + continue; + } + std::sort(small.begin(), small.end()); + small.erase(std::unique(small.begin(), small.end()), small.end()); + finished[i].push_back(std::move(small)); + small = Bucket(); + } else if (child.size() <= kLeafSize) { + finished[i].push_back(std::move(child)); + } else if (job.depth >= 4 || (job.depth >= 2 && child.size() * 5 > job.ids.size() * 4)) { + // Duplicates must not cause unbounded recursion/replication. + std::mt19937 random(job.seed + static_cast(j)); + std::shuffle(child.begin(), child.end(), random); + for (size_t begin = 0; begin < child.size(); begin += kLeafSize) { + finished[i].emplace_back( + child.begin() + static_cast(begin), + child.begin() + static_cast( + std::min(begin + kLeafSize, child.size()) + ) + ); + } + } else { + next[i].push_back( + {std::move(child), + job.depth + 1, + job.seed * 1664525U + static_cast(j) + 1013904223U} + ); + } + } + if (!small.empty()) { + std::sort(small.begin(), small.end()); + small.erase(std::unique(small.begin(), small.end()), small.end()); + finished[i].push_back(std::move(small)); + } + Bucket().swap(job.ids); + } + std::vector new_active; + for (size_t i = 0; i < active.size(); ++i) { + for (auto& leaf : finished[i]) { + leaves.push_back(std::move(leaf)); + } + for (auto& job : next[i]) { + new_active.push_back(std::move(job)); + } + } + active = std::move(new_active); + } + return leaves; +} +} // namespace pipnn_impl + +// Native PiPNN-style seeding: overlapping leader partitions, local dense kNN, +// directional hash reservoirs, and alpha pruning. Uses existing Eigen/OpenMP. +inline InitialGraph build_initial_graph( + const float* data, + size_t count, + size_t dim, + size_t degree, + MetricType metric = METRIC_L2, + size_t num_threads = std::numeric_limits::max() +) { + validate_metric_type(metric); + if (data == nullptr || dim == 0 || degree == 0 || degree >= count || + count >= buffer::kSearchBufferMaxPointCount || degree % 32 != 0) { + throw std::invalid_argument( + "PiPNN seed requires data, positive dim, and a degree multiple of 32 below " + "count" + ); + } + using namespace pipnn_impl; + const size_t threads = std::max(1, std::min(num_threads, total_threads())); + // Match QGBuilder's thread setting, including Eigen when threads == 1. + omp_set_num_threads(static_cast(threads)); + ScratchPool scratch(count, dim, threads); + auto leaves = cluster(data, count, dim, metric, threads, scratch); + RowMajorMatrix projections(dim, kHashBits); + std::mt19937 random(555); + std::normal_distribution normal; + for (Eigen::Index i = 0; i < projections.size(); ++i) { + projections.data()[i] = normal(random); + } + RowMajorMatrix sketches(count, kHashBits); +#pragma omp parallel for num_threads(threads) schedule(static) + for (size_t begin = 0; begin < count; begin += kTileSize) { + const size_t rows = std::min(kTileSize, count - begin); + ConstRowMajorMatrixMap points( + data + begin * dim, + static_cast(rows), + static_cast(dim) + ); + sketches + .middleRows(static_cast(begin), static_cast(rows)) + .noalias() = points * projections; + } + const size_t capacity = degree * 5 / 2; + std::vector table(count * capacity); + std::vector sizes(count, 0), worst(count, 0); + std::vector locks(count); + const auto merge = [&](PID source, PID target, float distance) { + uint16_t hash = 0; + for (size_t bit = 0; bit < kHashBits; ++bit) { + hash = static_cast( + (hash << 1) | (sketches(target, static_cast(bit)) > + sketches(source, static_cast(bit))) + ); + } + std::lock_guard lock(locks[source]); + auto* row = table.data() + source * capacity; + auto& size = sizes[source]; + auto& furthest = worst[source]; + if (size == capacity && distance >= row[furthest].distance) { + return; + } + auto* pos = std::lower_bound( + row, + row + size, + hash, + [](const Candidate& candidate, uint16_t key) { return candidate.hash < key; } + ); + const size_t slot = static_cast(pos - row); + if (slot < size && pos->hash == hash) { + if (distance >= pos->distance) { + return; + } + *pos = {target, distance, hash}; + } else if (size < capacity) { + std::move_backward(pos, row + size, row + size + 1); + *pos = {target, distance, hash}; + ++size; + } else { + const size_t insertion = slot > furthest ? slot - 1 : slot; + if (slot > furthest) { + std::move(row + furthest + 1, row + slot, row + furthest); + } else { + std::move_backward(row + slot, row + furthest, row + furthest + 1); + } + row[insertion] = {target, distance, hash}; + } + // Bounded rows keep this scan small; retain full-precision scores. + furthest = static_cast( + std::max_element( + row, + row + size, + [](const Candidate& a, const Candidate& b) { + return a.distance < b.distance; + } + ) - + row + ); + }; +#pragma omp parallel num_threads(threads) + { + auto work = scratch.worker(static_cast(omp_get_thread_num())); +#pragma omp for schedule(dynamic) + for (size_t leaf = 0; leaf < leaves.size(); ++leaf) { + const auto& ids = leaves[leaf]; + if (ids.size() < 2) { + continue; + } + gather(data, dim, ids, work); + pairwise(work, ids.size(), dim, metric); + for (size_t i = 0; i < ids.size(); ++i) { + work.candidates.clear(); + for (size_t j = 0; j < ids.size(); ++j) { + if (i != j) { + work.candidates.emplace_back( + work.distances[std::max(i, j) * ids.size() + std::min(i, j)], + ids[j] + ); + } + } + const size_t keep = std::min(kLocalDegree, work.candidates.size()); + std::partial_sort( + work.candidates.begin(), + work.candidates.begin() + static_cast(keep), + work.candidates.end() + ); + for (size_t j = 0; j < keep; ++j) { + const auto [distance, target] = work.candidates[j]; + merge(ids[i], target, distance); + merge(target, ids[i], distance); + } + } + } + } + std::vector().swap(leaves); + scratch.storage.reset(); + sketches.resize(0, 0); + InitialGraph result; + result.offsets.resize(count + 1, 0); + // Prune in-place: retain selected IDs in the front of each reservoir row. +#pragma omp parallel num_threads(threads) + { + std::vector candidates; + Bucket selected; +#pragma omp for schedule(dynamic) + for (size_t i = 0; i < count; ++i) { + const auto* row = table.data() + i * capacity; + candidates.assign(row, row + sizes[i]); + std::sort( + candidates.begin(), + candidates.end(), + [](const Candidate& a, const Candidate& b) { + return a.distance < b.distance || + (a.distance == b.distance && a.id < b.id); + } + ); + selected.clear(); + for (const auto& candidate : candidates) { + bool occluded = false; + for (PID accepted : selected) { + const float distance = + metric == METRIC_L2 + ? euclidean_sqr( + data + accepted * dim, data + candidate.id * dim, dim + ) + : dot_product_dis( + data + accepted * dim, data + candidate.id * dim, dim + ) - 1.0F; + if (distance <= candidate.distance / kAlpha) { + occluded = true; + break; + } + } + if (!occluded) { + selected.push_back(candidate.id); + } + if (selected.size() == degree) { + break; + } + } + for (size_t j = 0; j < selected.size(); ++j) { + table[i * capacity + j].id = selected[j]; + } + result.offsets[i + 1] = selected.size(); + } + } + std::partial_sum(result.offsets.begin(), result.offsets.end(), result.offsets.begin()); + result.neighbors.resize(result.offsets.back()); +#pragma omp parallel for num_threads(threads) + for (size_t i = 0; i < count; ++i) { + for (size_t j = 0; j < result.offsets[i + 1] - result.offsets[i]; ++j) { + result.neighbors[result.offsets[i] + j] = table[i * capacity + j].id; + } + } + return result; +} +} // namespace rabitqlib::symqg::detail diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index bf07598..c2e78d4 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -16,6 +16,7 @@ #include #include "rabitqlib/defines.hpp" +#include "rabitqlib/index/symqg/detail/pipnn.hpp" #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -24,12 +25,9 @@ namespace rabitqlib::symqg { constexpr size_t kMaxBsIter = 5; // max iter for binary search of pruning bar using CandidateList = std::vector>; +enum class QGInitialization { PiPNN, Random }; -/** - * @brief Builder of qg. Since we need to build the symphonyqg iteratively, which requires - * to record a lot of temp data, we use a separate class as a builder for this purpose. - * - */ +// Owns temporary state for SymphonyQG initialization and refinement. class QGBuilder { friend struct QGConstructionTestAccess; @@ -40,6 +38,7 @@ class QGBuilder { size_t num_nodes_; // num of data points size_t dim_; // dimension of data size_t degree_bound_; // degree bound for qg, multiple of 32 + bool seeded_ = true; static constexpr size_t kMaxCandidatePoolSize = 750; // max num of candidates for indexing static constexpr size_t kMaxPrunedSize = @@ -58,62 +57,108 @@ class QGBuilder { void graph_refine(); void iter(bool); - public: - explicit QGBuilder( - QuantizedGraph& index, - uint32_t ef_build, - const float* data, - size_t num_threads = std::numeric_limits::max() - ) + void initialize_storage(const float* data); + + QGBuilder(QuantizedGraph& index, uint32_t ef_build, size_t num_threads) : qg_{index} , ef_build_{ef_build} , num_threads_{std::max(1, std::min(num_threads, total_threads()))} , num_nodes_{qg_.num_vertices()} , dim_{qg_.dimension()} - , degree_bound_(qg_.degree_bound()) - , new_neighbors_(qg_.num_vertices()) - , pruned_neighbors_(qg_.num_vertices()) - , visited_list_( - num_threads_, - VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) - ) - , degrees_(qg_.num_vertices(), degree_bound_) { + , degree_bound_(qg_.degree_bound()) { omp_set_num_threads(static_cast(num_threads_)); + } - std::vector centroid = - compute_centroid(data, num_nodes_, dim_, num_threads_); - - qg_.set_quantization_centroid(centroid.data()); - qg_.copy_vectors(data); - - PID entry_point = 0; - if (qg_.is_quantized()) { - QuantizedQuery query( - qg_.centroid_.data(), - qg_.centroid_.data(), - qg_.padded_dim_, - qg_.metric_type_ + public: + explicit QGBuilder( + QuantizedGraph& index, + uint32_t ef_build, + const float* data, + size_t num_threads = std::numeric_limits::max(), + QGInitialization init = QGInitialization::PiPNN + ) + : QGBuilder(index, ef_build, num_threads) { + if (init == QGInitialization::PiPNN) { + auto seed = detail::build_initial_graph( + data, num_nodes_, dim_, degree_bound_, qg_.metric_type_, num_threads_ ); - 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; - } - } + initialize_seed(data, seed.offsets, seed.neighbors); + } else if (init == QGInitialization::Random) { + seeded_ = false; + initialize_storage(data); + random_init(); } else { - entry_point = exact_nn( - data, centroid.data(), num_nodes_, dim_, num_threads_, qg_.raw_dist_func_ - ); + throw std::invalid_argument("Unknown QG initialization"); } + } - qg_.set_ep(entry_point); + private: + void initialize_seed( + const float* data, + const std::vector& offsets, + const std::vector& neighbors + ) { + if (offsets.size() != num_nodes_ + 1 || offsets.front() != 0 || + offsets.back() != neighbors.size()) { + throw std::invalid_argument("Seed graph offsets must delimit every vertex"); + } + for (size_t i = 0; i < num_nodes_; ++i) { + if (offsets[i] > offsets[i + 1] || offsets[i + 1] > neighbors.size() || + offsets[i + 1] - offsets[i] > degree_bound_) { + throw std::invalid_argument( + "Seed graph row exceeds degree bound or has invalid offsets" + ); + } + for (size_t j = offsets[i]; j < offsets[i + 1]; ++j) { + if (neighbors[j] >= num_nodes_ || neighbors[j] == i) { + throw std::invalid_argument( + "Seed graph IDs must be in range and exclude self" + ); + } + } + } + initialize_storage(data); +#pragma omp parallel + { + CandidateList row; + row.reserve(degree_bound_); + std::vector ids; + ids.reserve(degree_bound_); +#pragma omp for schedule(dynamic) + for (size_t i = 0; i < num_nodes_; ++i) { + row.clear(); + ids.assign( + neighbors.begin() + static_cast(offsets[i]), + neighbors.begin() + static_cast(offsets[i + 1]) + ); + std::sort(ids.begin(), ids.end()); + ids.erase(std::unique(ids.begin(), ids.end()), ids.end()); + for (PID id : ids) { + // Encoding only consumes IDs. Score surviving seed candidates + // during search instead of retaining a second graph of scores. + row.emplace_back(id, 0.0F); + } + degrees_[i] = row.size(); + qg_.update_qg(i, row); + } + } + } - random_init(); + public: + // One complete search/prune/reverse-edge/degree-completion iteration. Call + // before serving a seeded graph: search() requires full final rows. + void refine() { iter(true); } + + // One refinement for PiPNN initialization, three passes for random init. + void build() { + if (seeded_) { + refine(); + } else { + build(3); + } } - void build(size_t num_iter = 3) { + void build(size_t num_iter) { if (num_iter < 2) { throw std::invalid_argument( "The number of QG build iterations must be at least 2" @@ -147,6 +192,42 @@ class QGBuilder { } }; +inline void QGBuilder::initialize_storage(const float* data) { + // Allocate refinement scratch only after PiPNN's dense workspace is released. + new_neighbors_.resize(num_nodes_); + pruned_neighbors_.resize(num_nodes_); + visited_list_ = std::vector( + num_threads_, + VisitedSet(num_nodes_, std::min(ef_build_ * ef_build_, num_nodes_ / 10)) + ); + degrees_.assign(num_nodes_, degree_bound_); + std::vector centroid = compute_centroid(data, num_nodes_, dim_, num_threads_); + + qg_.set_quantization_centroid(centroid.data()); + qg_.copy_vectors(data); + + 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); +} + inline void QGBuilder::add_pruned_edges( const CandidateList& result, const CandidateList& pruned_list, @@ -274,7 +355,24 @@ inline void QGBuilder::search_new_neighbors(bool refine) { vis.clear(); qg_.find_candidates(cur_id, ef_build_, candidates, vis, degrees_); - // add current neighbors + // Seeded construction keeps its initial edges only in QG. Materialize + // their scores on demand, after the caller can release the input/CSR. + if (new_neighbors_[cur_id].empty() && degrees_[cur_id] != 0) { + std::vector reconstructed; + std::optional> prepared; + const float* source = qg_.prepare_build_query(cur_id, reconstructed, prepared); + const PID* ids = qg_.get_neighbors(cur_id); + for (size_t j = 0; j < degrees_[cur_id]; ++j) { + if (ids[j] != cur_id && !vis.get(ids[j])) { + candidates.emplace_back( + ids[j], + qg_.point_distance(source, prepared ? &*prepared : nullptr, ids[j]) + ); + } + } + } + + // Add current neighbors retained by preceding iterations. for (auto& nei : new_neighbors_[cur_id]) { auto neighbor_id = nei.id; if (neighbor_id != cur_id && !vis.get(neighbor_id)) { @@ -291,6 +389,7 @@ inline void QGBuilder::search_new_neighbors(bool refine) { candidates.resize(min_size); // prune and update qg + new_neighbors_[cur_id].reserve(degree_bound_); heuristic_prune(cur_id, candidates, new_neighbors_[cur_id], refine); } } diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 064bb74..68afbfa 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -37,7 +37,19 @@ class SymqgIndex { } } - void build(py::handle data, size_t ef_construction, size_t num_threads = 1) { + void build( + py::handle data, + size_t ef_construction, + size_t num_threads = 1, + const std::string& init = "pipnn" + ) { + if (init != "random" && init != "pipnn") { + throw std::invalid_argument("init must be 'random' or 'pipnn'"); + } + if (ef_construction == 0 || + ef_construction > std::numeric_limits::max()) { + throw std::invalid_argument("ef_construction must be positive and fit uint32"); + } auto data_array = ensure_2d_array(data, "data"); if (static_cast(data_array.shape(1)) != dim_) { throw std::invalid_argument("data dimension does not match index dim"); @@ -53,8 +65,14 @@ class SymqgIndex { quantization_bits_ ); - rabitqlib::symqg::QGBuilder builder( - *index_, ef_construction, data_array.data(), num_threads + const auto initialization = init == "pipnn" ? symqg::QGInitialization::PiPNN + : symqg::QGInitialization::Random; + symqg::QGBuilder builder( + *index_, + static_cast(ef_construction), + data_array.data(), + num_threads, + initialization ); builder.build(); built_ = true; @@ -187,7 +205,8 @@ void register_symqg(py::module_& m) { &SymqgIndex::build, py::arg("data"), py::arg("ef_construction"), - py::arg("num_threads") = 1 + py::arg("num_threads") = 1, + py::arg("init") = "pipnn" ) .def( "search", diff --git a/sample/cpp/symqg_indexing.cpp b/sample/cpp/symqg_indexing.cpp index 170daca..6853bab 100644 --- a/sample/cpp/symqg_indexing.cpp +++ b/sample/cpp/symqg_indexing.cpp @@ -1,5 +1,9 @@ +#include #include #include +#include +#include +#include #include "rabitqlib/defines.hpp" #include "rabitqlib/index/symqg/qg.hpp" @@ -20,7 +24,9 @@ int run(int argc, char** argv) { << "arg3: ef for indexing \n" << "arg4: path for saving index\n" << "arg5: metric type (\"l2\" or \"ip\"), l2 by default\n" - << "arg6: vector quantization bits (0, 4, or 8), 0 by default\n"; + << "arg6: vector quantization bits (0, 4, or 8), 0 by default\n" + << "arg7: init (pipnn or random), pipnn by default\n" + << "arg8: construction threads, all available by default\n"; return 1; } @@ -47,8 +53,19 @@ int run(int argc, char** argv) { rabitqlib::load_vecs(data_file, data); + const std::string init = argc > 7 ? argv[7] : "pipnn"; + if (init != "random" && init != "pipnn") { + throw std::invalid_argument("Init must be random or pipnn"); + } + const bool pipnn = init == "pipnn"; + const size_t threads = argc > 8 ? std::stoul(argv[8]) : rabitqlib::total_threads(); + if (threads == 0 || ef == 0 || ef > std::numeric_limits::max()) { + throw std::invalid_argument( + "Threads and construction ef must be positive; ef must fit uint32" + ); + } rabitqlib::StopW stopw; - + rabitqlib::StopW stage; index_type qg( data.rows(), data.cols(), @@ -58,10 +75,22 @@ int run(int argc, char** argv) { quantization_bits ); - rabitqlib::symqg::QGBuilder builder(qg, ef, data.data()); - - // 3 iters, refine at last iter - builder.build(); + { + const auto initialization = pipnn ? rabitqlib::symqg::QGInitialization::PiPNN + : rabitqlib::symqg::QGInitialization::Random; + rabitqlib::symqg::QGBuilder builder( + qg, static_cast(ef), data.data(), threads, initialization + ); + std::cout << "Initialize and encode " << stage.get_elapsed_sec() << " secs\n"; + + // QG owns its vectors/codes now; release the input before refinement. + data = data_type(); + stage.reset(); + builder.build(); + std::cout << (pipnn ? "One refinement " : "Three iterations ") + << stage.get_elapsed_sec() << " secs\n"; + std::cout << "Average degree " << builder.avg_degree() << '\n'; + } // Release builder scratch before saving. auto milisecs = stopw.get_elapsed_mili(); diff --git a/sample/python/symqg_indexing.py b/sample/python/symqg_indexing.py index 4c2db6b..323abac 100644 --- a/sample/python/symqg_indexing.py +++ b/sample/python/symqg_indexing.py @@ -13,6 +13,7 @@ METRIC = "l2" # "l2" or "ip" QUANTIZATION_BITS = 0 # 0 for vanilla QG, or 4/8 for QG-quant NUM_THREADS = 16 # number of threads for build +INIT = "pipnn" # "pipnn" or "random" initialization of SymphonyQG # ────────────────────────────────────────────── @@ -26,7 +27,7 @@ def main(args=None) -> None: print( f"\nBuilding SymphonyQG index: n={n}, dim={dim}, MaxDegree={args.max_degree}, " f"ef={args.ef_construction}, metric={args.metric}, " - f"quantization_bits={args.quantization_bits}" + f"quantization_bits={args.quantization_bits}, init={args.init}" ) idx = SymqgIndex( @@ -37,7 +38,12 @@ def main(args=None) -> None: ) t0 = time() - idx.build(data, ef_construction=args.ef_construction, num_threads=args.num_threads) + idx.build( + data, + ef_construction=args.ef_construction, + num_threads=args.num_threads, + init=args.init, + ) print(f"Indexing time: {time() - t0:.2f}s") idx.save(args.index_file) @@ -80,6 +86,12 @@ def main(args=None) -> None: default=QUANTIZATION_BITS, help="Vector quantization bits: 0 for vanilla QG, or 4/8 for QG-quant", ) + parser.add_argument( + "--init", + choices=["pipnn", "random"], + default=INIT, + help="SymphonyQG initialization (default: pipnn)", + ) parser.add_argument( "--num-threads", dest="num_threads", diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 0e88356..20ef6a5 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -200,3 +200,54 @@ def test_constant_vectors_preserve_returned_distance(metric, tmp_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) + + +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("init", [None, "pipnn", "random"]) +def test_initialization_and_roundtrip(tmp_path, bits, metric, init): + rng = np.random.default_rng(17) + data = rng.standard_normal((160, 65)).astype(np.float32) + queries = rng.standard_normal((8, 65)).astype(np.float32) + index = SymqgIndex(65, max_degree=32, metric=metric, quantization_bits=bits) + kwargs = {} if init is None else {"init": init} + index.build(data, ef_construction=128, num_threads=2, **kwargs) + assert index.is_built + ids, distances = index.search(queries, k=10, ef=160, num_threads=2) + assert np.all(ids < len(data)) + assert np.isfinite(distances).all() + assert all(len(set(row)) == 10 for row in ids) + exact = ( + brute_force_knn(data, queries, 10)[0] + if metric == "l2" + else np.argsort(-(queries @ data.T), axis=1)[:, :10] + ) + assert recall_at_k(ids, exact, 10) >= (0.7 if bits else 0.95) + if bits == 0: + expected = ( + np.sum((queries[:, None, :] - data[ids]) ** 2, axis=2) + if metric == "l2" + else 1 - np.einsum("qd,qkd->qk", queries, data[ids]) + ) + np.testing.assert_allclose(distances, expected, rtol=1e-5, atol=1e-5) + path = str(tmp_path / "native_pipnn.index") + index.save(path) + restored = SymqgIndex.load(path) + restored_ids, restored_distances = restored.search(queries, k=10, ef=160) + np.testing.assert_array_equal(ids, restored_ids) + np.testing.assert_array_equal(distances, restored_distances) + + +def test_native_pipnn_rejects_invalid_build_arguments(base_data): + index = SymqgIndex(DIM, max_degree=32) + with pytest.raises(ValueError, match="init must be 'random' or 'pipnn'"): + index.build(base_data, ef_construction=64, init="unknown") + with pytest.raises( + ValueError, match="ef_construction must be positive and fit uint32" + ): + index.build(base_data, ef_construction=0, init="pipnn") + + +def test_default_initialization_is_pipnn(): + # pybind11 generates this signature from the actual Python argument defaults. + assert "init: str = 'pipnn'" in SymqgIndex.build.__doc__ diff --git a/tests/unit/rabitqlib/index/pipnn_test.cpp b/tests/unit/rabitqlib/index/pipnn_test.cpp new file mode 100644 index 0000000..aeffe09 --- /dev/null +++ b/tests/unit/rabitqlib/index/pipnn_test.cpp @@ -0,0 +1,120 @@ +#include "rabitqlib/index/symqg/detail/pipnn.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace rabitqlib::symqg::detail { +namespace { +TEST(PipnnTest, DenseDistancesMatchReference) { + constexpr size_t kCount = 1024, kDim = 65; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.17F); + } + pipnn_impl::ScratchPool scratch(kCount, kDim, 2); + // Reuse each worker's slice with changing matrix shapes, including the + // largest leaf, and check both distance conventions against scalar doubles. + for (size_t worker : {0U, 1U}) { + auto work = scratch.worker(worker); + for (size_t count : {33U, 1024U, 257U}) { + pipnn_impl::Bucket ids(count); + std::iota(ids.begin(), ids.end(), PID{0}); + pipnn_impl::gather(data.data(), kDim, ids, work); + for (auto metric : {METRIC_L2, METRIC_IP}) { + pipnn_impl::pairwise(work, count, kDim, metric); + for (size_t i = 0; i < count; ++i) { + for (size_t j = 0; j <= i; ++j) { + double expected = 0; + for (size_t d = 0; d < kDim; ++d) { + const double a = data[i * kDim + d], b = data[j * kDim + d]; + expected += metric == METRIC_L2 ? (a - b) * (a - b) : -a * b; + } + EXPECT_NEAR(work.distances[i * count + j], expected, 5e-5); + } + } + } + } + } +} + +TEST(PipnnTest, LocalSeedRetainsNearestNeighborsAndValidIds) { + constexpr size_t kCount = 97, kDim = 65; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto seed = build_initial_graph(data.data(), kCount, kDim, 32, metric, 2); + ASSERT_EQ(seed.offsets.size(), kCount + 1); + ASSERT_EQ(seed.offsets.front(), 0U); + ASSERT_EQ(seed.offsets.back(), seed.neighbors.size()); + for (size_t i = 0; i < kCount; ++i) { + std::vector row( + seed.neighbors.begin() + static_cast(seed.offsets[i]), + seed.neighbors.begin() + static_cast(seed.offsets[i + 1]) + ); + ASSERT_FALSE(row.empty()); + ASSERT_LE(row.size(), 32U); + std::sort(row.begin(), row.end()); + EXPECT_EQ(std::unique(row.begin(), row.end()), row.end()); + EXPECT_EQ(std::find(row.begin(), row.end(), i), row.end()); + EXPECT_LT(row.back(), kCount); + float best = std::numeric_limits::max(); + PID nearest = 0; + for (PID j = 0; j < kCount; ++j) { + if (i == j) { + continue; + } + const float distance = + metric == METRIC_L2 + ? euclidean_sqr( + data.data() + i * kDim, data.data() + j * kDim, kDim + ) + : dot_product_dis( + data.data() + i * kDim, data.data() + j * kDim, kDim + ); + if (distance < best) { + best = distance; + nearest = j; + } + } + EXPECT_NE(std::find(row.begin(), row.end(), nearest), row.end()); + } + } +} + +TEST(PipnnTest, OverlappingPartitionsTerminateForDuplicateVectors) { + constexpr size_t kCount = 1100, kDim = 7; + std::vector data(kCount * kDim, 0.5F); + pipnn_impl::ScratchPool scratch(kCount, kDim, 2); + auto leaves = pipnn_impl::cluster(data.data(), kCount, kDim, METRIC_L2, 2, scratch); + std::vector memberships(kCount, 0); + for (auto& leaf : leaves) { + EXPECT_LE(leaf.size(), pipnn_impl::kLeafSize); + std::sort(leaf.begin(), leaf.end()); + EXPECT_EQ(std::unique(leaf.begin(), leaf.end()), leaf.end()); + for (PID id : leaf) { + ASSERT_LT(id, kCount); + ++memberships[id]; + } + } + for (size_t membership : memberships) { + EXPECT_GE(membership, 1U); + EXPECT_LE(membership, 30U); + } +} + +TEST(PipnnTest, RejectsInvalidConfiguration) { + std::vector data(33 * 7, 0); + EXPECT_THROW(build_initial_graph(nullptr, 33, 7, 32), std::invalid_argument); + EXPECT_THROW(build_initial_graph(data.data(), 33, 0, 32), std::invalid_argument); + EXPECT_THROW(build_initial_graph(data.data(), 33, 7, 16), std::invalid_argument); + EXPECT_THROW(build_initial_graph(data.data(), 32, 7, 32), std::invalid_argument); +} +} // namespace +} // namespace rabitqlib::symqg::detail diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index e395e6d..2fd012c 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include "rabitqlib/index/symqg/qg_builder.hpp" @@ -15,6 +16,18 @@ namespace rabitqlib::symqg { struct QGConstructionTestAccess { + static QGBuilder from_graph( + QuantizedGraph& graph, + uint32_t ef, + const float* data, + const std::vector& offsets, + const std::vector& neighbors, + size_t threads + ) { + QGBuilder builder(graph, ef, threads); + builder.initialize_seed(data, offsets, neighbors); + return builder; + } static auto codes(const QuantizedGraph& graph) { std::vector result; for (PID id = 0; id < graph.num_points_; ++id) { @@ -26,6 +39,25 @@ struct QGConstructionTestAccess { static const auto& neighbors(const QGBuilder& builder) { return builder.new_neighbors_; } + static const auto& degrees(const QGBuilder& builder) { return builder.degrees_; } + static void retain_seed_scores(QGBuilder& builder) { + auto& graph = builder.qg_; + for (PID i = 0; i < builder.num_nodes_; ++i) { + std::vector reconstructed; + std::optional> prepared; + const float* source = graph.prepare_build_query(i, reconstructed, prepared); + for (size_t j = 0; j < builder.degrees_[i]; ++j) { + const PID id = graph.get_neighbors(i)[j]; + builder.new_neighbors_[i].emplace_back( + id, graph.point_distance(source, prepared ? &*prepared : nullptr, id) + ); + } + } + } + static void search(QGBuilder& builder) { builder.search_new_neighbors(false); } + static const PID* encoded_neighbors(const QuantizedGraph& graph, PID id) { + return graph.get_neighbors(id); + } static std::array estimate( const QuantizedGraph& graph, PID source, PID target ) { @@ -62,6 +94,18 @@ struct QGConstructionTestAccess { namespace { +static_assert( + !std::is_constructible_v< + QGBuilder, + QuantizedGraph&, + uint32_t, + const float*, + const std::vector&, + const std::vector&, + size_t>, + "Intermediate graph input must not be part of the public builder API" +); + TEST(QGEstimatorTest, MatchesExactDistancesForCollinearResiduals) { if (!cpu::has_avx2()) { GTEST_SKIP() << "FastScan requires AVX2/FMA"; @@ -148,6 +192,238 @@ TEST(QGConstructionTest, BuildsAndPrunesAfterInputReleaseUsingExistingCodes) { } } +TEST(QGConstructionTest, EncodedSeedSearchMatchesRetainedScores) { + if (!cpu::has_avx2()) { + GTEST_SKIP() << "FastScan requires AVX2/FMA"; + } + constexpr size_t kCount = 97, kDim = 65, kDegree = 32; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + std::vector offsets{0}; + std::vector edges; + for (PID i = 0; i < kCount; ++i) { + for (size_t j = 0; j < i % (kDegree + 1); ++j) { + edges.push_back((i + j + 1) % kCount); + } + offsets.push_back(edges.size()); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(::testing::Message() << metric << "/" << bits); + QuantizedGraph graph( + kCount, kDim, kDegree, metric, RotatorType::FhtKacRotator, bits + ); + // Both builders search the same immutable encoded graph/rotation. + auto compact = QGConstructionTestAccess::from_graph( + graph, 8, data.data(), offsets, edges, 2 + ); + auto retained = QGConstructionTestAccess::from_graph( + graph, 8, data.data(), offsets, edges, 2 + ); + QGConstructionTestAccess::retain_seed_scores(retained); + QGConstructionTestAccess::search(compact); + QGConstructionTestAccess::search(retained); + for (PID i = 0; i < kCount; ++i) { + const auto& actual = QGConstructionTestAccess::neighbors(compact)[i]; + const auto& expected = QGConstructionTestAccess::neighbors(retained)[i]; + ASSERT_EQ(actual.size(), expected.size()); + for (size_t j = 0; j < actual.size(); ++j) { + EXPECT_EQ(actual[j].id, expected[j].id); + EXPECT_FLOAT_EQ(actual[j].distance, expected[j].distance); + } + } + } + } +} + +TEST(QGConstructionTest, DefaultsToPipnnInitialization) { + if (!cpu::has_avx2()) { + GTEST_SKIP() << "FastScan requires AVX2/FMA"; + } + constexpr size_t kCount = 97, kDim = 65, kDegree = 32; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto seed = + detail::build_initial_graph(data.data(), kCount, kDim, kDegree, metric, 1); + for (size_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(::testing::Message() << metric << "/" << bits); + QuantizedGraph graph( + kCount, kDim, kDegree, metric, RotatorType::FhtKacRotator, bits + ); + QGBuilder builder(graph, 64, data.data(), 1); + for (PID i = 0; i < kCount; ++i) { + std::vector expected( + seed.neighbors.begin() + static_cast(seed.offsets[i]), + seed.neighbors.begin() + static_cast(seed.offsets[i + 1]) + ); + std::sort(expected.begin(), expected.end()); + ASSERT_EQ(QGConstructionTestAccess::degrees(builder)[i], expected.size()); + EXPECT_TRUE(QGConstructionTestAccess::neighbors(builder)[i].empty()); + const PID* actual = QGConstructionTestAccess::encoded_neighbors(graph, i); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), actual)); + } + builder.build(); + EXPECT_FLOAT_EQ(builder.avg_degree(), kDegree); + EXPECT_FALSE(builder.check_dup()); + } + } +} + +TEST(QGConstructionTest, SupportsExplicitRandomInitialization) { + if (!cpu::has_avx2()) { + GTEST_SKIP() << "FastScan requires AVX2/FMA"; + } + std::vector data(65 * 65, 0.5F); + QuantizedGraph graph(65, 65, 32); + QGBuilder builder(graph, 64, data.data(), 1, QGInitialization::Random); + for (const auto& row : QGConstructionTestAccess::neighbors(builder)) { + EXPECT_EQ(row.size(), 32U); + } + builder.build(); + EXPECT_FLOAT_EQ(builder.avg_degree(), 32); + EXPECT_FALSE(builder.check_dup()); + EXPECT_THROW(builder.build(0), std::invalid_argument); + EXPECT_THROW(builder.build(1), std::invalid_argument); + EXPECT_THROW( + (QGBuilder(graph, 64, data.data(), 1, static_cast(255))), + std::invalid_argument + ); +} + +TEST(QGConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { + if (!cpu::has_avx2()) { + GTEST_SKIP() << "FastScan requires AVX2/FMA"; + } + constexpr size_t kCount = 97, kDim = 65, kDegree = 64; + constexpr std::array kDegrees{0, 1, 31, 32, 33, 64}; + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {0U, 4U, 8U}) { + SCOPED_TRACE(::testing::Message() << metric << "/" << bits); + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + const std::vector query(data.begin(), data.begin() + kDim); + std::vector offsets{0}; + std::vector edges; + for (size_t i = 0; i < kCount; ++i) { + for (size_t j = 0; j < kDegrees[i % kDegrees.size()]; ++j) { + // Repeated IDs exercise import deduplication, including full rows. + edges.push_back((i + 1 + (i < kDegrees.size() ? j : j / 2)) % kCount); + } + offsets.push_back(edges.size()); + } + QuantizedGraph graph( + kCount, kDim, kDegree, metric, RotatorType::FhtKacRotator, bits + ); + auto builder = QGConstructionTestAccess::from_graph( + graph, 96, data.data(), offsets, edges, 2 + ); + const auto stored_vectors = QGConstructionTestAccess::codes(graph); + for (PID i = 0; i < kCount; ++i) { + std::vector expected( + edges.begin() + static_cast(offsets[i]), + edges.begin() + static_cast(offsets[i + 1]) + ); + std::sort(expected.begin(), expected.end()); + expected.erase( + std::unique(expected.begin(), expected.end()), expected.end() + ); + ASSERT_EQ(QGConstructionTestAccess::degrees(builder)[i], expected.size()); + const auto& row = QGConstructionTestAccess::neighbors(builder)[i]; + EXPECT_EQ(row.capacity(), 0U); + for (size_t j = 0; j < expected.size(); ++j) { + EXPECT_EQ( + QGConstructionTestAccess::encoded_neighbors(graph, i)[j], + expected[j] + ); + } + } + std::vector().swap(data); + std::vector().swap(offsets); + std::vector().swap(edges); + builder.refine(); + EXPECT_EQ(stored_vectors, QGConstructionTestAccess::codes(graph)); + EXPECT_FALSE(builder.check_dup()); + for (PID i = 0; i < kCount; ++i) { + const auto& row = QGConstructionTestAccess::neighbors(builder)[i]; + ASSERT_EQ(row.size(), kDegree); + EXPECT_EQ(QGConstructionTestAccess::degrees(builder)[i], kDegree); + for (size_t j = 0; j < row.size(); ++j) { + EXPECT_LT(row[j].id, kCount); + EXPECT_NE(row[j].id, i); + EXPECT_TRUE(std::isfinite(row[j].distance)); + EXPECT_EQ( + QGConstructionTestAccess::encoded_neighbors(graph, i)[j], row[j].id + ); + if (bits != 0) { + const auto expected = + QGConstructionTestAccess::estimate(graph, i, row[j].id); + EXPECT_NEAR(row[j].distance, expected[0], expected[1]); + } + } + } + graph.set_ef(96); + std::array ids{}, loaded_ids{}; + std::array distances{}, loaded_distances{}; + graph.search(query.data(), 10, ids.data(), distances.data()); + for (size_t j = 0; j < ids.size(); ++j) { + EXPECT_LT(ids[j], kCount); + EXPECT_TRUE(std::isfinite(distances[j])); + } + const std::string path = ::testing::TempDir() + "rabitq_qg_seeded.index"; + graph.save(path.c_str()); + QuantizedGraph loaded; + loaded.load(path.c_str()); + loaded.set_ef(96); + loaded.search(query.data(), 10, loaded_ids.data(), loaded_distances.data()); + EXPECT_EQ(ids, loaded_ids); + EXPECT_EQ(distances, loaded_distances); + EXPECT_EQ(graph.entry_point(), loaded.entry_point()); + std::remove(path.c_str()); + } + } +} + +TEST(QGConstructionTest, RejectsInvalidSeedStructureAndIds) { + constexpr size_t kCount = 33, kDim = 64; + std::vector data(kCount * kDim, 0.1F); + QuantizedGraph graph(kCount, kDim, 32); + const auto reject = [&](const std::vector& offsets, + const std::vector& edges, + const char* message) { + try { + auto builder = QGConstructionTestAccess::from_graph( + graph, 32, data.data(), offsets, edges, 1 + ); + FAIL() << "Invalid seed accepted"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), message); + } + }; + reject({}, {}, "Seed graph offsets must delimit every vertex"); + std::vector offsets(kCount + 1, 0); + offsets.back() = 1; + reject(offsets, {}, "Seed graph offsets must delimit every vertex"); + offsets.back() = 0; + offsets[1] = 1; + reject(offsets, {}, "Seed graph row exceeds degree bound or has invalid offsets"); + std::fill(offsets.begin() + 1, offsets.end(), 1); + reject(offsets, {kCount}, "Seed graph IDs must be in range and exclude self"); + reject(offsets, {0}, "Seed graph IDs must be in range and exclude self"); + std::fill(offsets.begin() + 1, offsets.end(), 33); + reject( + offsets, + std::vector(33, 1), + "Seed graph row exceeds degree bound or has invalid offsets" + ); +} + TEST(QuantizedGraphConfigurationTest, RejectsDegreeNotAlignedForFastScan) { EXPECT_THROW( (QuantizedGraph(64, 64, 16, METRIC_L2, RotatorType::MatrixRotator)),