[gouyt13]

Summary

Adds optional training-set sampling to RaBitQKMeans and QGKMeans, and an assign method that maps new vectors to trained centroids.

Sampling (max_points_per_centroid, default 0)

• Default 0 trains on all vectors, as before. Faiss uses 256.
• When positive and n > k * max_points_per_centroid, centroid initialization and the Lloyd iterations use k * max_points_per_centroid vectors drawn without replacement. seed selects them through a random stream separate from initialization.
• The final centroids then assign all n vectors with the configured final_assignment. assignments and distances keep shape (n,), so IvfIndex.build(data, centroids, assignments) and similar calls are unchanged.
• obj and iteration_stats describe the sample; distances and final_obj cover all points. The docs say the two are not directly comparable.
• Training that does not sample is bit-for-bit unchanged.
• Sampling copies the sample, so it raises peak memory (documented).

assign

• C++: assign(n, x, assignments, distances) const writes to caller-owned buffers.
• Python: assign(x) returns (uint32 (m,), float32 (m,)).
• It uses final_assignment and the other parameters of the last successful fit.
• It uses no labels or search hints from training, so each vector's result is independent of the other vectors in the call:
  • RaBitQKMeans stores the center its training points were encoded around, and encodes new vectors around it.
  • QGKMeans caches a graph over the centroids. It is rebuilt only when the centroids or graph settings change, and searched read-only with per-call scratch.
• A sampled fit's final pass is this same assignment, so assign reproduces its labels. It also reproduces training labels in Exact mode and for RaBitQKMeans. Unsampled approximate QGKMeans training uses search hints, so its labels can differ (documented).
• Input checks:
  • x and the public centroids are validated like training input.
  • Before a successful train(), it raises RuntimeError / std::runtime_error with "<Name> has not been trained", including when C++ callers wrote centroids directly.

Compatibility

• New public API: the C++ assign method, and a max_points_per_centroid field added at the end of both parameter structs, so positional aggregate initialization keeps its meaning.
• Python: a new keyword argument and cp attribute; the type stubs and docstrings are updated.
• No saved-index format changes and no new dependencies.

Testing

All checks ran on this commit (rebased on cc2e4f2), on Linux x86-64, GCC 14.3, RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF:

• CTest (--no-tests=error): 468/468 passed. One test (ThreadCount.BoundParallelForUsesRuntimePlacementAndPropagatesErrors) skipped because OMP_PROC_BIND was not set.

• Static analysis and formatting: full ./scripts/check-tidy.sh passed, including kmeans_bindings.cpp. Focused clang-tidy on the changed test files passed. ./scripts/check-format.sh and ./scripts/check-python.sh passed.

• Python (after reinstalling into the project environment and confirming the imported extension was the new build):

  • pytest: 888 passed; 2 skipped (macOS/ARM wheel-only tests).
  • C++ clustering example tests with RABITQ_REQUIRE_CPP_EXAMPLES=1: 18 passed.
  • Editor, mypy and Pyright checks: 18 passed.

• Docs: mkdocs build --strict passed.

• Compatibility: the same program was built against the main headers and this branch's headers and compared over 112 configurations: both classes, both final modes, L2 and spherical, QG bits 0/4/8, 1 and 4 threads, retraining on own centroids, and n = k * 256 exactly. Outputs are bit-identical single-threaded. With 4 threads everything except iteration_stats[].shift matches; shift already varies run to run on main because it comes from an OpenMP + reduction.

• New tests (kmeans_sampling_test.cpp, test_kmeans_sampling.py), for both classes and both final modes:

  • the sampling boundary at n = k*m and k*m + 1;
  • sampled training equals unsampled training on the drawn sample;
  • all n points are assigned;
  • same-seed reproducibility;
  • validation of points outside the sample;
  • retraining on own centroids;
  • assign errors and validation, including invalid centroids;
  • assign reproduces training labels in the cases listed above;
  • batch independence;
  • QG graph-cache invalidation when centroids or settings change.

  Mutating the threshold (both off-by-one directions), removing centroid validation, encoding around the batch mean, or disabling cache invalidation each makes the corresponding tests fail.

Not run: macOS, Windows and ARM builds, which CI covers. Performance is not claimed in this PR; a benchmark of training time and downstream IVF recall with and without sampling is planned for the next release.