Skip to content

refactor: share binding helpers and clarify QG preparation #177

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
gouyt13 merged 1 commit into from
Oct 29, 2026
Merged

refactor: share binding helpers and clarify QG preparation #177

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

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 23 additions & 19 deletions include/rabitqlib/clustering/qgkmeans.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<size_t>(std::numeric_limits<std::ptrdiff_t>::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 {
Expand All @@ -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<size_t>(std::numeric_limits<std::ptrdiff_t>::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<float>
: euclidean_sqr<float>;
Expand Down
20 changes: 20 additions & 0 deletions python_bindings/bindings_common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,26 @@ ensure_1d_integer_array(py::handle value, const char* name) {
return ensure_1d_array<int64_t>(value, name);
}

// The caller has already validated rank/dtype. Removed points still occupy IDs.
inline std::vector<PID> checked_removal_ids(
const py::array_t<int64_t, py::array::c_style | py::array::forcecast>& ids,
size_t num_stored_points
) {
std::vector<PID> result;
result.reserve(static_cast<size_t>(ids.shape(0)));
const rabitqlib::detail::ConstPackedArrayView<int64_t> values(
reinterpret_cast<const char*>(ids.data())
);
for (py::ssize_t i = 0; i < ids.shape(0); ++i) {
const int64_t id = values[static_cast<size_t>(i)];
if (id < 0 || static_cast<uint64_t>(id) >= num_stored_points) {
throw std::invalid_argument("ids must be in [0, num_stored_points)");
}
result.push_back(static_cast<PID>(id));
}
return result;
}

// Validate before narrowing, and own the result for APIs that mutate assignments.
inline std::vector<PID> checked_cluster_ids(
py::handle value,
Expand Down
51 changes: 7 additions & 44 deletions python_bindings/hnsw_bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <stdexcept>
Expand All @@ -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;

Expand Down Expand Up @@ -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<rabitqlib::PID> to_remove;
to_remove.reserve(static_cast<size_t>(ids_array.shape(0)));
const rabitqlib::detail::ConstPackedArrayView<int64_t> values(
reinterpret_cast<const char*>(ids_array.data())
);
for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) {
const int64_t id = values[static_cast<size_t>(i)];
if (id < 0 || static_cast<uint64_t>(id) >= count) {
throw std::invalid_argument("ids must be in [0, num_stored_points)");
}
to_remove.push_back(static_cast<rabitqlib::PID>(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());
}
Expand Down Expand Up @@ -314,8 +299,8 @@ class HnswIndex {
void register_hnsw(py::module_& m) {
using namespace rabitqlib::python_bindings;

py::class_<HnswIndex>(
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
----------
Expand All @@ -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_<HnswIndex>(m, "HnswIndex", kIndexDoc)
.def(
py::init<size_t, size_t, size_t, size_t, size_t, const std::string&, size_t>(),
py::arg("dim"),
Expand All @@ -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",
Expand Down
59 changes: 7 additions & 52 deletions python_bindings/ivf_bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -199,18 +198,7 @@ class IvfIndex {
if (!built_) {
throw std::runtime_error("IvfIndex must be built or loaded before remove");
}
std::vector<rabitqlib::PID> to_remove;
to_remove.reserve(static_cast<size_t>(ids_array.shape(0)));
const rabitqlib::detail::ConstPackedArrayView<int64_t> values(
reinterpret_cast<const char*>(ids_array.data())
);
for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) {
const int64_t id = values[static_cast<size_t>(i)];
if (id < 0 || static_cast<uint64_t>(id) >= index_->num_stored_points()) {
throw std::invalid_argument("ids must be in [0, num_stored_points)");
}
to_remove.push_back(static_cast<rabitqlib::PID>(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());
}
Expand Down Expand Up @@ -355,8 +343,8 @@ class IvfIndex {
void register_ivf(py::module_& m) {
using namespace rabitqlib::python_bindings;

py::class_<IvfIndex>(
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
----------
Expand Down Expand Up @@ -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_<IvfIndex>(m, "IvfIndex", kIndexDoc)
.def(
py::init<
size_t,
Expand All @@ -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",
Expand Down
Loading
Loading