diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index 5989692f..d008d759 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -184,25 +184,7 @@ class QGAssigner { void clear_hints() { previous_labels_.clear(); } private: - void assign_with_cache( - const float* centroids, - const float* x, - size_t n, - PID* labels, - float* distances, - detail::QuerySumCache& query_sums - ) { - if (centroids == nullptr || x == nullptr || labels == nullptr || - distances == nullptr) { - throw std::invalid_argument( - "QGAssigner input and output pointers must not be null" - ); - } - - if (n > static_cast(std::numeric_limits::max()) / d_) { - throw std::invalid_argument("QGAssigner input dimensions overflow ptrdiff_t"); - } - + void rebuild_centroid_graph(const float* centroids) { if (builder_) { builder_->reset(centroids); } else { @@ -227,6 +209,28 @@ class QGAssigner { } else { builder_->build(niter_); } + } + + void assign_with_cache( + const float* centroids, + const float* x, + size_t n, + PID* labels, + float* distances, + detail::QuerySumCache& query_sums + ) { + if (centroids == nullptr || x == nullptr || labels == nullptr || + distances == nullptr) { + throw std::invalid_argument( + "QGAssigner input and output pointers must not be null" + ); + } + + if (n > static_cast(std::numeric_limits::max()) / d_) { + throw std::invalid_argument("QGAssigner input dimensions overflow ptrdiff_t"); + } + + rebuild_centroid_graph(centroids); const auto distance_func = graph_->metric_type() == METRIC_IP ? dot_product_dis : euclidean_sqr; diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index e9d05cd0..85dc1783 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -225,6 +225,26 @@ ensure_1d_integer_array(py::handle value, const char* name) { return ensure_1d_array(value, name); } +// The caller has already validated rank/dtype. Removed points still occupy IDs. +inline std::vector checked_removal_ids( + const py::array_t& ids, + size_t num_stored_points +) { + std::vector result; + result.reserve(static_cast(ids.shape(0))); + const rabitqlib::detail::ConstPackedArrayView values( + reinterpret_cast(ids.data()) + ); + for (py::ssize_t i = 0; i < ids.shape(0); ++i) { + const int64_t id = values[static_cast(i)]; + if (id < 0 || static_cast(id) >= num_stored_points) { + throw std::invalid_argument("ids must be in [0, num_stored_points)"); + } + result.push_back(static_cast(id)); + } + return result; +} + // Validate before narrowing, and own the result for APIs that mutate assignments. inline std::vector checked_cluster_ids( py::handle value, diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 318b97ae..4d38191a 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -7,7 +7,6 @@ #include #include -#include #include #include #include @@ -18,7 +17,6 @@ #include "bindings_common.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" #include "rabitqlib/index/types.hpp" -#include "rabitqlib/utils/packed.hpp" namespace py = pybind11; @@ -156,20 +154,7 @@ class HnswIndex { if (!built_) { throw std::runtime_error("HnswIndex must be built or loaded before remove"); } - // Removed labels remain allocated; validate against the full stored ID range. - const size_t count = index_->num_stored_points(); - std::vector to_remove; - to_remove.reserve(static_cast(ids_array.shape(0))); - const rabitqlib::detail::ConstPackedArrayView values( - reinterpret_cast(ids_array.data()) - ); - for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) { - const int64_t id = values[static_cast(i)]; - if (id < 0 || static_cast(id) >= count) { - throw std::invalid_argument("ids must be in [0, num_stored_points)"); - } - to_remove.push_back(static_cast(id)); - } + auto to_remove = checked_removal_ids(ids_array, index_->num_stored_points()); py::gil_scoped_release release; return index_->remove(to_remove.data(), to_remove.size()); } @@ -314,8 +299,8 @@ class HnswIndex { void register_hnsw(py::module_& m) { using namespace rabitqlib::python_bindings; - py::class_( - m, "HnswIndex", R"doc(An approximate nearest-neighbor index using HNSW. + static constexpr char kIndexDoc[] = + R"doc(An approximate nearest-neighbor index using HNSW. Parameters ---------- @@ -338,8 +323,9 @@ random_seed : int, optional Notes ----- Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc" - ) +supported; an overlapping mutation raises RuntimeError.)doc"; + + py::class_(m, "HnswIndex", kIndexDoc) .def( py::init(), py::arg("dim"), @@ -349,30 +335,7 @@ supported; an overlapping mutation raises RuntimeError.)doc" py::arg("nbits") = 8, py::arg("metric") = "l2", py::arg("random_seed") = 100, - R"doc(An approximate nearest-neighbor index using HNSW. - -Parameters ----------- -dim : int - Original vector dimension, from 64 through 65536. -max_elements : int - Storage capacity, including removed points. Use resize to grow it. -M : int, optional - Graph connectivity, at least 2. Default is 16. -ef_construction : int, optional - Construction search breadth. Default is 200. -nbits : int, optional - Total quantization bits, 1 through 9. Default is 8. -metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. -random_seed : int, optional - Graph random seed. Default is 100. - -Notes ------ -Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc" + kIndexDoc ) .def( "build", diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 20b81065..0f96da2c 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -21,7 +21,6 @@ #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" #include "rabitqlib/index/types.hpp" -#include "rabitqlib/utils/packed.hpp" namespace py = pybind11; @@ -199,18 +198,7 @@ class IvfIndex { if (!built_) { throw std::runtime_error("IvfIndex must be built or loaded before remove"); } - std::vector to_remove; - to_remove.reserve(static_cast(ids_array.shape(0))); - const rabitqlib::detail::ConstPackedArrayView values( - reinterpret_cast(ids_array.data()) - ); - for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) { - const int64_t id = values[static_cast(i)]; - if (id < 0 || static_cast(id) >= index_->num_stored_points()) { - throw std::invalid_argument("ids must be in [0, num_stored_points)"); - } - to_remove.push_back(static_cast(id)); - } + auto to_remove = checked_removal_ids(ids_array, index_->num_stored_points()); py::gil_scoped_release release; return index_->remove(to_remove.data(), to_remove.size()); } @@ -355,8 +343,8 @@ class IvfIndex { void register_ivf(py::module_& m) { using namespace rabitqlib::python_bindings; - py::class_( - m, "IvfIndex", R"doc(An approximate nearest-neighbor index using IVF. + static constexpr char kIndexDoc[] = + R"doc(An approximate nearest-neighbor index using IVF. Parameters ---------- @@ -390,8 +378,9 @@ Examples >>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) >>> ids, distances = index.search(x[:2], k=3, nprobe=1) >>> ids.shape, distances.shape -((2, 3), (2, 3)))doc" - ) +((2, 3), (2, 3)))doc"; + + py::class_(m, "IvfIndex", kIndexDoc) .def( py::init< size_t, @@ -406,41 +395,7 @@ Examples py::arg("nbits"), py::arg("metric") = "l2", py::arg("initializer") = "auto", - R"doc(An approximate nearest-neighbor index using IVF. - -Parameters ----------- -dim : int - Original vector dimension, from 64 through 65536. -max_elements : int - Number of vectors required by the initial build. Later adds grow - storage. -num_clusters : int - Number of supplied centroids. -nbits : int - Total quantization bits, 1 through 9, or 32 for owned float32 vectors - with one-bit filtering and raw reranking. -metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. -initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional - Centroid routing strategy. Default is 'auto'. - -Notes ------ -Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError. - -Examples --------- ->>> import numpy as np ->>> from rabitqlib import IvfIndex ->>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) ->>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) ->>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) ->>> ids, distances = index.search(x[:2], k=3, nprobe=1) ->>> ids.shape, distances.shape -((2, 3), (2, 3)))doc" + kIndexDoc ) .def( "build", diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 94bcfdc8..7f49e0f1 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -219,8 +219,8 @@ class SymqgIndex { void register_symqg(py::module_& m) { using namespace rabitqlib::python_bindings; - py::class_( - m, "SymqgIndex", R"doc(An approximate nearest-neighbor index using SymphonyQG. + static constexpr char kIndexDoc[] = + R"doc(An approximate nearest-neighbor index using SymphonyQG. Parameters ---------- @@ -238,33 +238,43 @@ quantization_bits : {0, 4, 8}, optional Notes ----- Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc" - ) - .def( - py::init(), - py::arg("dim"), - py::arg("max_degree"), - py::arg("metric") = "l2", - py::arg("quantization_bits") = 0, - R"doc(An approximate nearest-neighbor index using SymphonyQG. +supported; an overlapping mutation raises RuntimeError.)doc"; + + static constexpr char kSearchDoc[] = R"doc(Search a batch of query vectors. Parameters ---------- -dim : int - Original vector dimension, from 64 through 65536. -max_degree : int - Graph degree bound; a positive multiple of 32. -metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. -quantization_bits : {0, 4, 8}, optional - Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 - (default). +queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +k : int + Number of neighbors requested, between 1 and num_points. +ef : int + Search breadth. Must be positive and at least k. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + +Returns +------- +ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. +distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. Notes ----- -Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError.)doc" +Returns (ids, distances), in that order. Array shapes and value ranges are +checked at runtime, not by static type checkers.)doc"; + + py::class_(m, "SymqgIndex", kIndexDoc) + .def( + py::init(), + py::arg("dim"), + py::arg("max_degree"), + py::arg("metric") = "l2", + py::arg("quantization_bits") = 0, + kIndexDoc ) .def( "build", @@ -299,32 +309,7 @@ synchronous call.)doc" py::arg("k"), py::arg("ef"), py::arg("num_threads") = 1, - R"doc(Search a batch of query vectors. - -Parameters ----------- -queries : array_like of float, shape (n_queries, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -k : int - Number of neighbors requested, between 1 and num_points. -ef : int - Search breadth. Must be positive and at least k. -num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. - -Returns -------- -ids : ndarray of uint32, shape (n_queries, k) - Point IDs of the selected neighbors. -distances : ndarray of float32, shape (n_queries, k) - Lower values are better; squared L2 estimates for l2 and inner-product - distance estimates for ip. These are distances, not similarity scores. - -Notes ------ -Returns (ids, distances), in that order. Array shapes and value ranges are -checked at runtime, not by static type checkers.)doc" + kSearchDoc ) .def( "search_batch", @@ -333,32 +318,7 @@ checked at runtime, not by static type checkers.)doc" py::arg("k"), py::arg("ef"), py::arg("num_threads") = 1, - R"doc(Search a batch of query vectors. - -Parameters ----------- -queries : array_like of float, shape (n_queries, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -k : int - Number of neighbors requested, between 1 and num_points. -ef : int - Search breadth. Must be positive and at least k. -num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. - -Returns -------- -ids : ndarray of uint32, shape (n_queries, k) - Point IDs of the selected neighbors. -distances : ndarray of float32, shape (n_queries, k) - Lower values are better; squared L2 estimates for l2 and inner-product - distance estimates for ip. These are distances, not similarity scores. - -Notes ------ -Returns (ids, distances), in that order. Array shapes and value ranges are -checked at runtime, not by static type checkers.)doc" + kSearchDoc ) .def( "save", diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 12920132..518f751f 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -39,6 +39,12 @@ namespace rabitqlib::symqg { +namespace { +// Native-endian numeric tag; little-endian files contain the bytes "QTIBARGQ". +constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; +constexpr uint32_t kFormatVersion = 1; +} // namespace + QuantizedQuery::QuantizedQuery( const float* rotated_query, const float* centroid, @@ -261,9 +267,6 @@ void QuantizedGraph::save(const char* filename) const { } output.exceptions(std::ios::badbit | std::ios::failbit); - // Native-endian numeric tag; little-endian files contain the bytes "QTIBARGQ". - constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; - constexpr uint32_t kFormatVersion = 1; if (quantization_bits_ != 0) { output.write(reinterpret_cast(&kFormatMagic), sizeof(kFormatMagic)); output.write( @@ -333,9 +336,6 @@ void QuantizedGraph::load(const char* filename) { read_exact(&value, sizeof(value), field); }; - // Native-endian numeric tag; little-endian files contain the bytes "QTIBARGQ". - constexpr uint64_t kFormatMagic = 0x5147524142495451ULL; - constexpr uint32_t kFormatVersion = 1; uint64_t magic = 0; read_value(magic, "format marker"); if (magic == kFormatMagic) {