From 45d3817691437cc9299e7d093c7aa979b5badbb7 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 22:43:01 +0800 Subject: [PATCH 1/9] fix: correct MatrixRotator serialization size and comments - use sizeof(T) when serializing MatrixRotator data - correct inaccurate and misspelled comments --- include/rabitqlib/fastscan/fastscan.hpp | 6 +++--- include/rabitqlib/index/ivf/initializer.hpp | 4 ++-- include/rabitqlib/index/lut.hpp | 4 ++-- include/rabitqlib/index/query.hpp | 4 ++-- include/rabitqlib/utils/rotator.hpp | 16 ++++++++-------- include/rabitqlib/utils/tools.hpp | 2 +- src/simd/pack_excode_kernels.hpp | 2 +- 7 files changed, 19 insertions(+), 19 deletions(-) diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index 6d7e30c..7fc285b 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -70,7 +70,7 @@ static inline void get_column( } /** - * @brief Pack quantization codes, store in blocks, the data orgnization is illustrated in + * @brief Pack quantization codes into blocks; the data organization is illustrated in * the link and kPerm0. Since we pack codes as 32-sized groups, if the num is not a multiple * of 32, we have to use some space for these absent data * @@ -97,8 +97,8 @@ inline void pack_codes( // each batch contain codes for 32 vectors for (size_t row = 0; row < num_rd; row += kBatchSize) { // get quantization codes for each column for each batch - // i.e., we get the codes for 8 dims of 32 vectors and re-orgnize the data layout - // based on the shuffle SIMD instruction used during quering + // i.e., we get the codes for 8 dims of 32 vectors and reorganize the data layout + // based on the shuffle SIMD instruction used during querying for (size_t i = 0; i < cols; ++i) { get_column(quantization_code, num, cols, row, i, col); for (size_t j = 0; j < 32; ++j) { diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 8980bf2..179c9ee 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -69,7 +69,7 @@ inline void parallel_for(size_t start, size_t end, size_t numThreads, Function f } /** - * @brief For ivf centroids, we need an intializer to get the candidate clusters. + * @brief For IVF centroids, an initializer finds the candidate clusters. */ class Initializer { protected: @@ -198,4 +198,4 @@ class HNSWInitializer : public Initializer { ~HNSWInitializer() override { delete alg_hnsw_; } }; -} // namespace rabitqlib::ivf \ No newline at end of file +} // namespace rabitqlib::ivf diff --git a/include/rabitqlib/index/lut.hpp b/include/rabitqlib/index/lut.hpp index 5690dcf..7d76829 100644 --- a/include/rabitqlib/index/lut.hpp +++ b/include/rabitqlib/index/lut.hpp @@ -14,7 +14,7 @@ template class Lut { static constexpr size_t kNumBits = 8; static constexpr size_t kNumBitsHacc = 16; - static_assert(std::is_floating_point_v, "T must be an floating type in Lut"); + static_assert(std::is_floating_point_v, "T must be a floating-point type in Lut"); private: size_t table_length_ = 0; @@ -62,4 +62,4 @@ class Lut { [[nodiscard]] T delta() const { return delta_; }; [[nodiscard]] T sum_vl() const { return sum_vl_lut_; }; }; -} // namespace rabitqlib \ No newline at end of file +} // namespace rabitqlib diff --git a/include/rabitqlib/index/query.hpp b/include/rabitqlib/index/query.hpp index 04ab11a..6e49bf2 100644 --- a/include/rabitqlib/index/query.hpp +++ b/include/rabitqlib/index/query.hpp @@ -42,8 +42,8 @@ class BatchQuery { [[nodiscard]] T g_add() const { return G_add_; } void set_g_add(T dist) { - // For L2, dist is the compute by euclidean_sqr() - // For IP, dist is computed by dot_product_dist() i.e., 1 - dot_product() + // 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; } diff --git a/include/rabitqlib/utils/rotator.hpp b/include/rabitqlib/utils/rotator.hpp index 138ba90..1611e7b 100644 --- a/include/rabitqlib/utils/rotator.hpp +++ b/include/rabitqlib/utils/rotator.hpp @@ -58,14 +58,14 @@ inline size_t padding_requirement(size_t dim, RotatorType type) { template class MatrixRotator : public Rotator { private: - RowMajorMatrix rand_mat_; // Rotation Maxtrix + RowMajorMatrix rand_mat_; // Rotation matrix public: explicit MatrixRotator(size_t dim, size_t padded_dim) : Rotator(dim, padded_dim), rand_mat_(dim, padded_dim) { RowMajorMatrix rand = random_gaussian_matrix(padded_dim, padded_dim); Eigen::HouseholderQR> qr(rand); - RowMajorMatrix q_inv = - qr.householderQ().transpose(); // inverse of orthogonal mat is its inverse + RowMajorMatrix q_inv = qr.householderQ().transpose( + ); // The inverse of an orthogonal matrix is its transpose. // the random matrix only need the first dim rows, since we just pad zeros for // the vector to be rotated to padded dimension @@ -84,27 +84,27 @@ class MatrixRotator : public Rotator { void load(std::ifstream& input) override { input.read( reinterpret_cast(rand_mat_.data()), - static_cast(sizeof(float) * this->dim_ * this->padded_dim_) + static_cast(sizeof(T) * this->dim_ * this->padded_dim_) ); } void save(std::ofstream& output) const override { output.write( reinterpret_cast(rand_mat_.data()), - (sizeof(float) * this->dim_ * this->padded_dim_) + (sizeof(T) * this->dim_ * this->padded_dim_) ); } void load(const char* data) override { - std::memcpy(rand_mat_.data(), data, sizeof(float) * this->dim_ * this->padded_dim_); + std::memcpy(rand_mat_.data(), data, sizeof(T) * this->dim_ * this->padded_dim_); } void save(char* data) const override { - std::memcpy(data, rand_mat_.data(), sizeof(float) * this->dim_ * this->padded_dim_); + std::memcpy(data, rand_mat_.data(), sizeof(T) * this->dim_ * this->padded_dim_); } size_t dump_bytes() const override { - return sizeof(float) * this->dim_ * this->padded_dim_; + return sizeof(T) * this->dim_ * this->padded_dim_; } void rotate(const T* vec, T* rotated_vec) const override { diff --git a/include/rabitqlib/utils/tools.hpp b/include/rabitqlib/utils/tools.hpp index 9dc1808..15ba9a0 100644 --- a/include/rabitqlib/utils/tools.hpp +++ b/include/rabitqlib/utils/tools.hpp @@ -23,7 +23,7 @@ inline void assert_floating() { ); } -// thread save rand int +// Generate a thread-safe random integer in the inclusive range [min, max]. template inline T rand_integer(T min, T max) { static thread_local std::mt19937 generator = [] { diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index 4cafa9f..b00ac5c 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -84,7 +84,7 @@ inline void packing_3bit_excode_intrinsics( inline void packing_4bit_excode_intrinsics( const uint8_t* o_raw, uint8_t* o_compact, size_t dim ) { - // although this part only requries SSE, computing inner product for this orgnization + // Although this part only requires SSE, computing inner products for this organization // requires AVX512F, similar for remaining functions // ! require dim % 16 == 0 for (size_t j = 0; j < dim; j += 16) { From 1d4716afa332c8d8c1eff4263e9a1a7104b81f56 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 22:47:50 +0800 Subject: [PATCH 2/9] perf: use stack storage for FastScan accumulators - replace dynamically allocated Eigen storage with fixed-size std::array buffers --- include/rabitqlib/index/estimator.hpp | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/include/rabitqlib/index/estimator.hpp b/include/rabitqlib/index/estimator.hpp index 083a1a2..ed121aa 100644 --- a/include/rabitqlib/index/estimator.hpp +++ b/include/rabitqlib/index/estimator.hpp @@ -33,13 +33,10 @@ inline void split_batch_estdist( ) { constexpr size_t kSafeChunkDim = 1024; ConstBatchDataMap cur_batch(batch_data, padded_dim); - RowMajorArray accu_arr(1, fastscan::kBatchSize); + std::array accu_values{}; + RowMajorArrayMap accu_arr(accu_values.data(), 1, fastscan::kBatchSize); const auto* codes_ptr = cur_batch.bin_code(); const auto* lut_ptr = q_obj.lut(); - for (size_t i = 0; i < fastscan::kBatchSize; ++i) { - accu_arr.data()[i] = 0; - } - if (use_hacc) { std::array accu_res; size_t remaining_dim = padded_dim; @@ -165,7 +162,7 @@ template inline void qg_batch_estdist( const char* batch_data, const BatchQuery& q_obj, size_t padded_dim, T* est_distance ) { - std::vector accu_res(fastscan::kBatchSize); + std::array accu_res{}; ConstQGBatchDataMap cur_batch(batch_data, padded_dim); From 73d96bff9ae55cde6422fcb1cd6ac6b819d4bf20 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 23:11:57 +0800 Subject: [PATCH 3/9] fix(ivf): harden storage lifecycle and Python validation - manage IVF initializers with unique ownership - safely replace buffers during reconstruction and loading - normalize OpenMP thread counts and avoid null extra-code pointer arithmetic - validate Python build and search inputs - use explicit sentinels for unfilled search results - add Python regressions for incomplete results and search-before-build --- include/rabitqlib/index/ivf/ivf.hpp | 68 +++++++++++++++++------------ python_bindings/ivf_bindings.cpp | 26 +++++++++-- tests/python/test_ivf.py | 20 +++++++++ 3 files changed, 81 insertions(+), 33 deletions(-) diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 55187f1..dc8600b 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -8,6 +8,8 @@ #include #include #include +#include +#include #include #include "rabitqlib/defines.hpp" @@ -27,18 +29,18 @@ namespace rabitqlib::ivf { class IVF { private: - Initializer* initer_ = nullptr; // initializer for find candidate cluster - char* batch_data_ = nullptr; // 1-bit code and factors - char* ex_data_ = nullptr; // code for remaining bits - PID* ids_ = nullptr; // PID of vectors (orgnized by clusters) - size_t num_; // num of data points - size_t dim_; // dimension of data points - size_t padded_dim_; // dimension after padding, - size_t num_cluster_; // num of centroids (clusters) - size_t ex_bits_; // total bits = ex_bits_ + 1 - RotatorType type_; // type of rotator - Rotator* rotator_ = nullptr; // Data Rotator - std::vector cluster_lst_; // List of clusters in ivf + std::unique_ptr initer_; // initializer for candidate clusters + char* batch_data_ = nullptr; // 1-bit code and factors + char* ex_data_ = nullptr; // code for remaining bits + PID* ids_ = nullptr; // PID of vectors (organized by clusters) + size_t num_ = 0; // num of data points + size_t dim_ = 0; // dimension of data points + size_t padded_dim_ = 0; // dimension after padding + size_t num_cluster_ = 0; // num of centroids (clusters) + size_t ex_bits_ = 0; // total bits = ex_bits_ + 1 + RotatorType type_ = RotatorType::FhtKacRotator; // type of rotator + Rotator* rotator_ = nullptr; // Data Rotator + std::vector cluster_lst_; // List of clusters in ivf MetricType metric_type_ = rabitqlib::METRIC_L2; // metric type float (*ip_func_)(const float*, const uint8_t*, size_t) = nullptr; @@ -66,10 +68,10 @@ class IVF { void init_clusters(const std::vector&); void free_memory() { - ::delete initer_; - std::free(batch_data_); - std::free(ex_data_); - std::free(ids_); + initer_.reset(); + std::free(std::exchange(batch_data_, nullptr)); + std::free(std::exchange(ex_data_, nullptr)); + std::free(std::exchange(ids_, nullptr)); } void search_cluster( @@ -87,7 +89,7 @@ class IVF { ) const; public: - explicit IVF() {} + explicit IVF() = default; explicit IVF( size_t, size_t, @@ -158,7 +160,7 @@ inline IVF::~IVF() { * * @param data Data objects (N*DIM) * @param centroids Centroid vectors (K*DIM) - * @param clustter_ids Cluster ID for each data objects + * @param cluster_ids Cluster ID for each data object */ inline void IVF::construct( const float* data, @@ -196,7 +198,7 @@ inline void IVF::construct( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } - num_threads = std::min(num_threads, rabitqlib::total_threads()); + num_threads = std::clamp(num_threads, size_t{1}, rabitqlib::total_threads()); /* Quantize each cluster */ #pragma omp parallel for schedule(dynamic) num_threads(num_threads) for (size_t i = 0; i < num_cluster_; ++i) { @@ -211,10 +213,13 @@ inline void IVF::construct( inline void IVF::allocate_memory(const std::vector& cluster_sizes) { std::cout << "Allocating memory for IVF...\n"; + free_memory(); + cluster_lst_.clear(); + if (num_cluster_ < 20000UL) { - this->initer_ = new FlatInitializer(padded_dim_, num_cluster_); + this->initer_ = std::make_unique(padded_dim_, num_cluster_); } else { - this->initer_ = new HNSWInitializer(padded_dim_, num_cluster_); + this->initer_ = std::make_unique(padded_dim_, num_cluster_); } this->batch_data_ = memory::align_allocate<64, char, true>(batch_data_bytes(cluster_sizes)); @@ -227,7 +232,7 @@ inline void IVF::allocate_memory(const std::vector& cluster_sizes) { } /** - * @brief intialize the cluster list: finding idx for all data + * @brief Initialize the cluster list by finding each cluster's storage offsets. */ inline void IVF::init_clusters(const std::vector& cluster_sizes) { this->cluster_lst_.reserve(num_cluster_); @@ -241,8 +246,9 @@ inline void IVF::init_clusters(const std::vector& cluster_sizes) { char* current_batch_data = batch_data_ + (BatchDataMap::data_bytes(padded_dim_) * added_batches); char* current_ex_data = - ex_data_ + - (added_vectors * ExDataMap::data_bytes(padded_dim_, ex_bits_)); + ex_bits_ > 0 ? ex_data_ + (added_vectors * + ExDataMap::data_bytes(padded_dim_, ex_bits_)) + : nullptr; PID* ids = ids_ + added_vectors; Cluster cur_cluster(num, current_batch_data, current_ex_data, ids); @@ -263,7 +269,7 @@ inline void IVF::quantize_cluster( ) { size_t num_points = IDs.size(); if (cp.num() != num_points) { - std::cerr << "Size of cluster and IDs are inequivalent\n"; + std::cerr << "Cluster size and ID count differ\n"; std::cerr << "Cluster: " << cp.num() << " IDs: " << num_points << '\n'; exit(1); } @@ -298,7 +304,9 @@ inline void IVF::quantize_cluster( ); batch_data += BatchDataMap::data_bytes(padded_dim_); - ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * n; + if (ex_bits_ > 0) { + ex_data += ExDataMap::data_bytes(padded_dim_, ex_bits_) * n; + } } } @@ -360,6 +368,7 @@ inline void IVF::load(const char* filename) { input.read(reinterpret_cast(&type_), sizeof(type_)); input.read(reinterpret_cast(&metric_type_), sizeof(metric_type_)); + delete rotator_; rotator_ = choose_rotator(dim_, type_, round_up_to_multiple(dim_, 64)); padded_dim_ = rotator_->size(); @@ -381,7 +390,6 @@ inline void IVF::load(const char* filename) { this->rotator_->load(input); /* Load data */ - free_memory(); allocate_memory(cluster_sizes); this->initer_->load(input, filename); input.read(batch_data_, static_cast(batch_data_bytes(cluster_sizes))); @@ -476,8 +484,10 @@ inline void IVF::search_cluster( ); batch_data += BatchDataMap::data_bytes(padded_dim_); - ex_data += - ExDataMap::data_bytes(padded_dim_, ex_bits_) * fastscan::kBatchSize; + if (ex_bits_ > 0) { + ex_data += + ExDataMap::data_bytes(padded_dim_, ex_bits_) * fastscan::kBatchSize; + } ids += fastscan::kBatchSize; } diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index defc06d..a9c2ea7 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -1,6 +1,6 @@ #include -#include +#include #include #include #include @@ -50,9 +50,18 @@ class IvfIndex { if (static_cast(data_array.shape(1)) != dim_) { throw std::invalid_argument("data dimension does not match index dim"); } + if (static_cast(data_array.shape(0)) != max_elements_) { + throw std::invalid_argument("number of data rows must match index max_elements" + ); + } if (static_cast(centroids_array.shape(1)) != dim_) { throw std::invalid_argument("centroid dimension does not match index dim"); } + if (static_cast(centroids_array.shape(0)) != num_clusters_) { + throw std::invalid_argument( + "number of centroid rows must match index num_clusters" + ); + } if (static_cast(cluster_ids_array.shape(0)) != static_cast(data_array.shape(0))) { throw std::invalid_argument( @@ -78,9 +87,18 @@ class IvfIndex { size_t num_threads = 1 ) { auto query_array = ensure_2d_array(queries, "queries"); + if (!built_) { + throw std::runtime_error("IvfIndex must be built or loaded before search"); + } if (static_cast(query_array.shape(1)) != dim_) { throw std::invalid_argument("query dimension does not match index dim"); } + if (k == 0 || k > max_elements_) { + throw std::invalid_argument("k must be between 1 and max_elements"); + } + if (nprobe == 0) { + throw std::invalid_argument("nprobe must be positive"); + } const size_t nq = static_cast(query_array.shape(0)); const auto shape = @@ -95,8 +113,8 @@ class IvfIndex { nq, num_threads, [&](size_t idx, size_t /*threadId*/) { - std::vector row_ids(k, 0); - std::vector row_dists(k, 0.0F); + std::vector row_ids(k, rabitqlib::kPidMax); + std::vector row_dists(k, std::numeric_limits::infinity()); index_->search( query_array.data() + (idx * dim_), @@ -199,4 +217,4 @@ void register_ivf(py::module_& m) { .def_property_readonly("nbits", &IvfIndex::nbits) .def_property_readonly("is_built", &IvfIndex::is_built) .def_property_readonly("metric", &IvfIndex::metric); -} \ No newline at end of file +} diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 607a6c5..55ceccf 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -79,6 +79,20 @@ def test_k_equals_one(built_ivf, query_data): assert dists.shape == (N_QUERIES, 1) +def test_unfilled_results_use_sentinels(): + data = np.zeros((4, DIM), dtype=np.float32) + data[1:] = 100.0 + centroids = np.stack((data[0], data[1])) + cluster_ids = np.array([0, 1, 1, 1], dtype=np.uint32) + idx = IvfIndex(DIM, 4, 2, nbits=4) + idx.build(data, centroids, cluster_ids) + + ids, dists = idx.search(data[:1], k=3, nprobe=1) + assert ids[0, 0] == 0 + np.testing.assert_array_equal(ids[0, 1:], np.iinfo(np.uint32).max) + assert np.all(np.isinf(dists[0, 1:])) + + # ── search correctness ──────────────────────────────────────────────────────── @@ -153,6 +167,12 @@ def test_wrong_query_dim_raises(built_ivf): built_ivf.search(bad_queries, k=1, nprobe=1) +def test_search_before_build_raises(query_data): + idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) + with pytest.raises(Exception): + idx.search(query_data[:1], k=1, nprobe=1) + + # ── save / load roundtrip ───────────────────────────────────────────────────── From f0e3bdedcddc225ace96c105fdd0ea22d5e942f9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 23:14:51 +0800 Subject: [PATCH 4/9] fix(symqg): validate graph and search configuration - require graph degrees to be positive multiples of 32 - reject degrees that cannot exclude the current vertex - enforce the search-buffer point-ID limit - normalize zero thread requests to one worker - validate Python k and ef search parameters - correct the graph-building iteration diagnostic - add configuration and rotator-lifecycle regression tests --- include/rabitqlib/index/symqg/qg.hpp | 30 ++++++++++++++++++-- include/rabitqlib/index/symqg/qg_builder.hpp | 4 +-- python_bindings/symqg_bindings.cpp | 14 +++++++-- tests/python/test_symqg.py | 5 ++++ tests/unit/rabitqlib/index/qg_test.cpp | 30 ++++++++++++++++++++ 5 files changed, 76 insertions(+), 7 deletions(-) create mode 100644 tests/unit/rabitqlib/index/qg_test.cpp diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 06d543a..4cd9dc1 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include "rabitqlib/defines.hpp" @@ -62,6 +63,8 @@ class QuantizedGraph { size_t row_offset_ = 0; // length of entire row size_t ef_ = 0; + void validate_configuration() const; + void initialize(); void copy_vectors(const T*); @@ -157,12 +160,32 @@ inline QuantizedGraph::QuantizedGraph( , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) , metric_type_(metric_type) , rotator_type_(rotator_type) { + validate_configuration(); initialize(); } +template +inline void QuantizedGraph::validate_configuration() const { + if (degree_bound_ == 0 || degree_bound_ % fastscan::kBatchSize != 0) { + throw std::invalid_argument( + "QuantizedGraph degree bound must be a positive multiple of 32" + ); + } + if (degree_bound_ >= num_points_) { + throw std::invalid_argument( + "QuantizedGraph degree bound must be smaller than the number of points" + ); + } + if (num_points_ > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument( + "QuantizedGraph point count exceeds the search-buffer ID limit" + ); + } +} + template inline QuantizedGraph::~QuantizedGraph() { - ::delete this->rotator_; + delete this->rotator_; } template @@ -225,6 +248,7 @@ inline void QuantizedGraph::load(const char* filename) { raw_dist_func_ = (metric_type_ == METRIC_IP) ? dot_product_dis : euclidean_sqr; + validate_configuration(); initialize(); /* Data */ @@ -335,7 +359,7 @@ inline void QuantizedGraph::search( } // scan a data row (including data vec and quantization codes for its neighbors) -// store estimated distance & return exact distnace for current vertex +// Store estimated neighbor distances; the caller computes the current vertex exactly. template void QuantizedGraph::scan_neighbors( const BatchQuery& q_obj, @@ -395,7 +419,7 @@ inline void QuantizedGraph::update_results( // initialize const offsets & data array template inline void QuantizedGraph::initialize() { - ::delete rotator_; + delete rotator_; rotator_ = choose_rotator(dim_, rotator_type_, round_up_to_multiple(dim_, 64)); padded_dim_ = rotator_->size(); diff --git a/include/rabitqlib/index/symqg/qg_builder.hpp b/include/rabitqlib/index/symqg/qg_builder.hpp index 229af3d..8bedd43 100644 --- a/include/rabitqlib/index/symqg/qg_builder.hpp +++ b/include/rabitqlib/index/symqg/qg_builder.hpp @@ -59,7 +59,7 @@ class QGBuilder { ) : qg_{index} , ef_build_{ef_build} - , num_threads_{std::min(num_threads, total_threads())} + , num_threads_{std::max(1, std::min(num_threads, total_threads()))} , num_nodes_{qg_.num_vertices()} , dim_{qg_.dimension()} , degree_bound_(qg_.degree_bound()) @@ -88,7 +88,7 @@ class QGBuilder { void build(size_t num_iter = 3) { if (num_iter < 2) { - std::cerr << "The number of iter for building qg should >= 3\n"; + std::cerr << "The number of iterations for building QG must be at least 2\n"; exit(1); } // for first iterations, we do not need to refine the graph structure diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 18480a1..4008e9b 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -17,7 +17,11 @@ namespace rabitqlib::python_bindings { class SymqgIndex { public: SymqgIndex(size_t dim, size_t max_degree, const std::string& metric = "l2") - : dim_(dim), max_degree_(max_degree), metric_(metric_from_string(metric)) {} + : dim_(dim), max_degree_(max_degree), metric_(metric_from_string(metric)) { + if (max_degree == 0 || max_degree % rabitqlib::fastscan::kBatchSize != 0) { + throw std::invalid_argument("max_degree must be a positive multiple of 32"); + } + } void build(py::handle data, size_t ef_construction, size_t num_threads = 1) { auto data_array = ensure_2d_array(data, "data"); @@ -42,6 +46,12 @@ class SymqgIndex { if (!built_) { throw std::runtime_error("SymqgIndex must be built or loaded before search"); } + if (k == 0 || k > num_points_) { + throw std::invalid_argument("k must be between 1 and num_points"); + } + if (ef == 0) { + throw std::invalid_argument("ef must be positive"); + } if (static_cast(query_array.shape(1)) != dim_) { throw std::invalid_argument("query dimension does not match index dim"); } @@ -152,4 +162,4 @@ void register_symqg(py::module_& m) { .def_property_readonly("num_points", &SymqgIndex::num_points) .def_property_readonly("is_built", &SymqgIndex::is_built) .def_property_readonly("metric", &SymqgIndex::metric); -} \ No newline at end of file +} diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 3603bcb..ff938ae 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -108,6 +108,11 @@ def test_search_before_build_raises(): idx.search(queries, k=1, ef=_EF) +def test_invalid_degree_raises(): + with pytest.raises(Exception): + SymqgIndex(DIM, max_degree=16) + + # ── save / load roundtrip ───────────────────────────────────────────────────── diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp new file mode 100644 index 0000000..02c2f69 --- /dev/null +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -0,0 +1,30 @@ +#include "rabitqlib/index/symqg/qg.hpp" + +#include + +#include + +namespace rabitqlib::symqg { +namespace { + +TEST(QuantizedGraphConfigurationTest, RejectsDegreeNotAlignedForFastScan) { + EXPECT_THROW( + (QuantizedGraph(64, 64, 16, METRIC_L2, RotatorType::MatrixRotator)), + std::invalid_argument + ); +} + +TEST(QuantizedGraphConfigurationTest, RejectsDegreeThatCannotExcludeSelf) { + EXPECT_THROW( + (QuantizedGraph(32, 64, 32, METRIC_L2, RotatorType::MatrixRotator)), + std::invalid_argument + ); +} + +TEST(QuantizedGraphLifecycleTest, DestroysConcreteRotatorThroughBasePointer) { + QuantizedGraph graph(33, 64, 32, METRIC_L2, RotatorType::MatrixRotator); + EXPECT_EQ(graph.num_vertices(), 33U); +} + +} // namespace +} // namespace rabitqlib::symqg From f93e6229947c9712dbe1497d1d919f8aaee51f64 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 23:16:07 +0800 Subject: [PATCH 5/9] fix(hnsw): harden Python binding validation - reject invalid data counts, empty centroids, and out-of-range cluster IDs - prevent searches before the index is built or loaded - validate requested neighbor counts - initialize unfilled results with explicit sentinels - add a regression test for search-before-build --- python_bindings/hnsw_bindings.cpp | 27 +++++++++++++++++++++++++++ tests/python/test_hnsw.py | 6 ++++++ 2 files changed, 33 insertions(+) diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 2550eef..32b6868 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -50,6 +51,12 @@ class HnswIndex { if (static_cast(data_array.shape(1)) != dim_) { throw std::invalid_argument("data dimension does not match index dim"); } + if (data_array.shape(0) == 0 || + static_cast(data_array.shape(0)) > max_elements_) { + throw std::invalid_argument( + "number of data rows must be between 1 and max_elements" + ); + } if (static_cast(centroids_array.shape(1)) != dim_) { throw std::invalid_argument("centroid dimension does not match index dim"); } @@ -61,6 +68,14 @@ class HnswIndex { } const size_t num_clusters = static_cast(centroids_array.shape(0)); + if (num_clusters == 0) { + throw std::invalid_argument("at least one centroid is required"); + } + for (ssize_t i = 0; i < cluster_ids_array.shape(0); ++i) { + if (cluster_ids_array.data()[i] >= num_clusters) { + throw std::invalid_argument("cluster_ids contains an out-of-range value"); + } + } num_clusters_ = num_clusters; // Ensure cluster_ids are writable for the C++ API by making a copy @@ -87,17 +102,29 @@ class HnswIndex { py::tuple search(py::handle queries, size_t k, size_t ef = 0, size_t num_threads = 1) { auto query_array = ensure_2d_array(queries, "queries"); + if (!built_) { + throw std::runtime_error("HnswIndex must be built or loaded before search"); + } if (dim_ != 0 && static_cast(query_array.shape(1)) != dim_) { throw std::invalid_argument("query dimension does not match index dim"); } if (ef == 0) { ef = std::max(k, 10); } + if (k == 0 || k > max_elements_) { + throw std::invalid_argument("k must be between 1 and max_elements"); + } const auto shape = std::vector{ static_cast(query_array.shape(0)), static_cast(k)}; auto ids = py::array_t(shape); auto dists = py::array_t(shape); + std::fill(ids.mutable_data(), ids.mutable_data() + ids.size(), rabitqlib::kPidMax); + std::fill( + dists.mutable_data(), + dists.mutable_data() + dists.size(), + std::numeric_limits::infinity() + ); auto ids_buf = ids.mutable_unchecked<2>(); auto dists_buf = dists.mutable_unchecked<2>(); diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index c9b96b8..6806c7b 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -115,6 +115,12 @@ def test_wrong_query_dim_raises(built_hnsw): built_hnsw.search(bad_queries, k=1) +def test_search_before_build_raises(query_data): + idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + with pytest.raises(Exception): + idx.search(query_data[:1], k=1) + + # ── save / load roundtrip ───────────────────────────────────────────────────── From b464b85a33aff3a2c0b2a3a3031e90cfccf55d78 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 23:33:10 +0800 Subject: [PATCH 6/9] fix(quantization): handle zero residuals and one-bit outputs - return finite factors for zero-residual one-bit and extra-bit quantization - handle zero-range scalar quantization without producing NaNs - fully initialize codes and factors for one-bit-only quantization - assert positive residual-code inner-product invariants - add regressions for degenerate inputs and poisoned output storage --- .../rabitqlib/quantization/rabitq_impl.hpp | 63 +++++++++--- .../rabitqlib/quantization/rabitq_test.cpp | 96 +++++++++++++++++++ 2 files changed, 147 insertions(+), 12 deletions(-) create mode 100644 tests/unit/rabitqlib/quantization/rabitq_test.cpp diff --git a/include/rabitqlib/quantization/rabitq_impl.hpp b/include/rabitqlib/quantization/rabitq_impl.hpp index cd91c4e..4d11461 100644 --- a/include/rabitqlib/quantization/rabitq_impl.hpp +++ b/include/rabitqlib/quantization/rabitq_impl.hpp @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -94,15 +95,27 @@ inline void one_bit_code_with_factor( T l2_sqr = l2norm_sqr(residual_arr.data(), dim); T l2_norm = std::sqrt(l2_sqr); + if (l2_sqr == 0) { + if (metric_type == METRIC_L2) { + f_add = 0; + } else if (metric_type == METRIC_IP) { + f_add = 1; + } else { + std::cerr << "Unsupported metric type in quantization\n" << std::flush; + exit(1); + } + f_rescale = 0; + f_error = 0; + return; + } + // dot product between residual and xu_cb T ip_resi_xucb = dot_product(residual_arr.data(), xu_cb.data(), dim); // dot product between centroid and xu_cb T ip_cent_xucb = dot_product(centroid, xu_cb.data(), dim); - // corner case - if (ip_resi_xucb == 0) { - ip_resi_xucb = std::numeric_limits::infinity(); - } + // A nonzero residual and its sign code have a strictly positive inner product. + assert(ip_resi_xucb > 0); // We use unnormalized vector to get error factor. To be more specific, // sqrt((1 - ^2) / ^2) / sqrt(dim - 1) = 3rd item in following @@ -450,6 +463,22 @@ inline void ex_bits_code_with_factor( // residual vector RowMajorArray residual_arr = data_arr - cent_arr; + const T l2_sqr = l2norm_sqr(residual_arr.data(), dim); + if (l2_sqr == 0) { + std::fill(ex_code, ex_code + dim, static_cast(0)); + if (metric_type == METRIC_L2) { + f_add_ex = 0; + } else if (metric_type == METRIC_IP) { + f_add_ex = 1; + } else { + std::cerr << "Unsupported metric for ex_bits_code()\n" << std::flush; + exit(1); + } + f_rescale_ex = 0; + f_error_ex = 0; + return; + } + T ipnorm_inv = ex_bits_code(residual_arr.data(), dim, ex_bits, ex_code, t_const); // get factors @@ -465,16 +494,13 @@ inline void ex_bits_code_with_factor( float cb = -(static_cast(1 << ex_bits) - 0.5F); RowMajorArray xu_cb = total_code.template cast() + cb; - T l2_sqr = l2norm_sqr(residual_arr.data(), dim); T l2_norm = std::sqrt(l2_sqr); T ip_resi_xucb = dot_product(residual_arr.data(), xu_cb.data(), dim); T ip_cent_xucb = dot_product(centroid, xu_cb.data(), dim); - // corner case - if (ip_resi_xucb == 0) { - ip_resi_xucb = std::numeric_limits::infinity(); - } + // A nonzero residual and its quantized code have a strictly positive inner product. + assert(ip_resi_xucb > 0); T tmp_error = l2_norm * kConstEpsilon * @@ -494,7 +520,7 @@ inline void ex_bits_code_with_factor( f_rescale_ex = ipnorm_inv * -l2_norm; f_error_ex = 1 * tmp_error; } else { - std::cerr << "Unsupport metric for ex_bits_code()\n" << std::flush; + std::cerr << "Unsupported metric for ex_bits_code()\n" << std::flush; exit(1); } } @@ -550,6 +576,13 @@ static inline void rabitq_scalar_impl( RowMajorArray residual_arr = rabitq_impl::one_bit::one_bit_code(data, centroid, dim, binary_code.data()); + if (l2norm_sqr(residual_arr.data(), dim) == 0) { + std::fill(total_code, total_code + dim, static_cast(0)); + delta = 0; + vl = 0; + return; + } + if (ex_bits > 0) { ex_bits::ex_bits_code( residual_arr.data(), dim, ex_bits, total_code, t_const @@ -558,7 +591,8 @@ static inline void rabitq_scalar_impl( // merge 2 one_bit code and ex_bits code for (size_t i = 0; i < dim; ++i) { - total_code[i] += static_cast(binary_code[i]) << ex_bits; + const TP sign_code = static_cast(binary_code[i]) << ex_bits; + total_code[i] = ex_bits > 0 ? total_code[i] + sign_code : sign_code; } float cb = -(static_cast(1 << ex_bits) - 0.5F); @@ -612,10 +646,15 @@ static inline void rabitq_full_impl( metric_type, t_const ); + } else { + one_bit::one_bit_code_with_factor( + data, centroid, dim, binary_code.data(), f_add, f_rescale, f_error, metric_type + ); } for (size_t i = 0; i < dim; ++i) { - total_code[i] += static_cast(binary_code[i]) << ex_bits; + const TP sign_code = static_cast(binary_code[i]) << ex_bits; + total_code[i] = ex_bits > 0 ? total_code[i] + sign_code : sign_code; } } } // namespace total_bits diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp new file mode 100644 index 0000000..22231be --- /dev/null +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -0,0 +1,96 @@ +#include "rabitqlib/quantization/rabitq.hpp" + +#include + +#include +#include +#include +#include + +namespace rabitqlib::quant { +namespace { + +TEST(RabitqDegenerateInputTest, OneBitFactorsAreFiniteForZeroResidual) { + constexpr size_t kDim = 64; + std::array data{}; + std::array centroid{}; + std::array code{}; + + for (MetricType metric : {METRIC_L2, METRIC_IP}) { + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + rabitq_impl::one_bit::one_bit_code_with_factor( + data.data(), + centroid.data(), + kDim, + code.data(), + f_add, + f_rescale, + f_error, + metric + ); + + EXPECT_TRUE(std::isfinite(f_add)); + EXPECT_TRUE(std::isfinite(f_rescale)); + EXPECT_FLOAT_EQ(f_error, 0.0F); + } +} + +TEST(RabitqDegenerateInputTest, ExtraBitFactorsAreFiniteForZeroResidual) { + constexpr size_t kDim = 64; + std::array data{}; + std::array centroid{}; + std::array code{}; + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + + rabitq_impl::ex_bits::ex_bits_code_with_factor( + data.data(), centroid.data(), kDim, 3, code.data(), f_add, f_rescale, f_error + ); + + EXPECT_TRUE(std::isfinite(f_add)); + EXPECT_TRUE(std::isfinite(f_rescale)); + EXPECT_FLOAT_EQ(f_error, 0.0F); +} + +TEST(RabitqDegenerateInputTest, ScalarQuantizationReconstructsZeroVector) { + constexpr size_t kDim = 64; + std::array data{}; + std::array code{}; + std::array reconstructed{}; + float delta = 1; + float vl = 1; + + quantize_scalar(data.data(), kDim, 4, code.data(), delta, vl); + reconstruct_vec(code.data(), delta, vl, kDim, reconstructed.data()); + + EXPECT_TRUE(std::isfinite(delta)); + EXPECT_TRUE(std::isfinite(vl)); + EXPECT_EQ(reconstructed, data); +} + +TEST(RabitqOneBitTest, FullQuantizationInitializesCodesAndFactors) { + constexpr size_t kDim = 64; + std::array data{}; + std::array code; + code.fill(0xFF); + data[0] = 1.0F; + float f_add = std::numeric_limits::quiet_NaN(); + float f_rescale = std::numeric_limits::quiet_NaN(); + float f_error = std::numeric_limits::quiet_NaN(); + + quantize_full_single(data.data(), kDim, 1, code.data(), f_add, f_rescale, f_error); + + EXPECT_EQ(code[0], 1); + for (size_t i = 1; i < kDim; ++i) { + EXPECT_EQ(code[i], 0); + } + EXPECT_TRUE(std::isfinite(f_add)); + EXPECT_TRUE(std::isfinite(f_rescale)); + EXPECT_TRUE(std::isfinite(f_error)); +} + +} // namespace +} // namespace rabitqlib::quant From b3d8917bb9fc4e6a381d812c6e35089739f1ecf0 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 23:35:39 +0800 Subject: [PATCH 7/9] fix(buffer): harden checked IDs and zero-capacity handling --- include/rabitqlib/utils/buffer.hpp | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/include/rabitqlib/utils/buffer.hpp b/include/rabitqlib/utils/buffer.hpp index f4daaad..786a6fc 100644 --- a/include/rabitqlib/utils/buffer.hpp +++ b/include/rabitqlib/utils/buffer.hpp @@ -1,12 +1,18 @@ #pragma once #include +#include #include #include "rabitqlib/defines.hpp" #include "rabitqlib/utils/memory.hpp" namespace rabitqlib::buffer { + +inline constexpr PID kSearchBufferCheckedMask = PID{1} + << (std::numeric_limits::digits - 1); +inline constexpr size_t kSearchBufferMaxPointCount = + static_cast(kSearchBufferCheckedMask); /** * @brief sorted linear buffer, used as beam set for graph-based ANN search. In symphonyqg, * the search buffer may contain duplicate id with different distances @@ -31,10 +37,10 @@ class SearchBuffer { } // set top bit to 1 as checked - static void set_checked(PID& data_id) { data_id |= (1 << 31); } + static void set_checked(PID& data_id) { data_id |= kSearchBufferCheckedMask; } [[nodiscard]] static auto is_checked(PID data_id) -> bool { - return static_cast(data_id >> 31); + return (data_id & kSearchBufferCheckedMask) != 0; } public: @@ -97,16 +103,21 @@ class SearchBuffer { } T top_dist() const { + if (capacity_ == 0) { + return std::numeric_limits::lowest(); + } return is_full() ? data_[size_ - 1].distance : std::numeric_limits::max(); } [[nodiscard]] auto is_full() const -> bool { return size_ == capacity_; } // judge if dist can be inserted into buffer - [[nodiscard]] auto is_full(T dist) const -> bool { return dist > top_dist(); } + [[nodiscard]] auto is_full(T dist) const -> bool { + return capacity_ == 0 || dist > top_dist(); + } const std::vector, memory::AlignedAllocator>>& data() { return data_; } }; -} // namespace rabitqlib::buffer \ No newline at end of file +} // namespace rabitqlib::buffer From 79c944f35b534993fd9b21621cd434e555380aa9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Mon, 24 Aug 2026 23:48:35 +0800 Subject: [PATCH 8/9] fix: honor requested thread counts in space utilities --- include/rabitqlib/utils/space.hpp | 33 ++++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index c0be95c..f6cc24e 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -154,10 +154,17 @@ template inline std::vector compute_centroid( const T* data, size_t num_points, size_t dim, size_t num_threads ) { - omp_set_num_threads(static_cast(num_threads)); - std::vector> all_results(num_threads, std::vector(dim, 0)); - -#pragma omp parallel for schedule(dynamic) + const auto thread_count = static_cast(std::max( + 1, + std::min( + {num_threads, + std::max(num_points, 1), + static_cast(std::numeric_limits::max())} + ) + )); + std::vector> all_results(thread_count, std::vector(dim, 0)); + +#pragma omp parallel for schedule(dynamic) num_threads(thread_count) for (size_t i = 0; i < num_points; ++i) { auto tid = omp_get_thread_num(); std::vector& cur_results = all_results[tid]; @@ -191,9 +198,17 @@ inline PID exact_nn( size_t num_threads, T (*dist_func)(const T*, const T*, size_t) ) { - std::vector> best_entries(num_threads); - -#pragma omp parallel for schedule(dynamic) + const auto thread_count = static_cast(std::max( + 1, + std::min( + {num_threads, + std::max(num_points, 1), + static_cast(std::numeric_limits::max())} + ) + )); + std::vector> best_entries(thread_count); + +#pragma omp parallel for schedule(dynamic) num_threads(thread_count) for (size_t i = 0; i < num_points; ++i) { auto tid = omp_get_thread_num(); AnnCandidate& cur_entry = best_entries[tid]; @@ -244,8 +259,8 @@ float ip64_fxu7_avx( // inner product between float type and int type vectors template inline TF ip_fxi(const TF* __restrict__ vec0, const TI* __restrict__ vec1, size_t dim) { - static_assert(std::is_floating_point_v, "TF must be an floating type"); - static_assert(std::is_integral_v, "TI must be an integeral type"); + static_assert(std::is_floating_point_v, "TF must be a floating-point type"); + static_assert(std::is_integral_v, "TI must be an integral type"); ConstVectorMap v0(vec0, dim); ConstVectorMap v1(vec1, dim); From 73c6edee1dd152eb5a3311f5989ece34dd494634 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Tue, 25 Aug 2026 00:00:23 +0800 Subject: [PATCH 9/9] test: use a valid degree for SymQG construction --- tests/python/test_import.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/test_import.py b/tests/python/test_import.py index 9508eb0..209b2c5 100644 --- a/tests/python/test_import.py +++ b/tests/python/test_import.py @@ -44,9 +44,9 @@ def test_hnsw_construct_explicit(): def test_symqg_construct(): - idx = SymqgIndex(64, max_degree=16) + idx = SymqgIndex(64, max_degree=32) assert idx.dim == 64 - assert idx.max_degree == 16 + assert idx.max_degree == 32 assert not idx.is_built