Skip to content

feat(python): train, set_centroids and add for IvfIndex and HnswIndex - #204

Merged
gouyt13 merged 2 commits into
mainfrom
feat/python-train-api
Nov 1, 2026
Merged

gouyt13 merged 2 commits into
mainfrom
feat/python-train-api

Conversation

@gouyt13

@gouyt13 gouyt13 commented Nov 1, 2026

Copy link
Copy Markdown
Member

Summary

Second of three PRs moving the Python indexes to a Faiss-style train → add → search workflow, on top of the C++ training support from #203. IvfIndex and HnswIndex now train their own centroids, so callers no longer run k-means and pass centroids and assignments to build. SymphonyQG follows in PR 3.

New API

  • IvfIndex(dim, num_clusters, *, nbits, metric="l2", initializer="auto", seed=None) and HnswIndex(dim, *, nbits, num_clusters=16, M=16, ef_construction=200, metric="l2", seed=100). nbits is required; HNSW's old default of 8 is gone. Everything except dim (and IVF's num_clusters) is keyword-only.
  • seed=None for IVF draws one random seed at construction, used for the rotation and cp.seed, and reported as index.seed. This keeps 0.7's random default rotation while making it reproducible. Loaded indexes report None.
  • index.cp is a RaBitQKMeansParameters with index defaults: niter=25, max_points_per_centroid=256, spherical=True for inner product, and seed from the index.
  • train(x, *, num_threads=None): None uses cp.num_threads (default 0); an explicit value overrides it for that call only. IVF uses RaBitQKMeans up to 2000 clusters and QGKMeans above that, a provisional, unmeasured threshold; HNSW always uses RaBitQKMeans.
  • set_centroids(centroids) takes shape (num_clusters, dim). is_trained, len(index) (live points), and add(data, *, cluster_ids=None, num_threads=0, fast_quantization=False).
  • Retraining is allowed while the index stores no points and raises RuntimeError afterwards. Operations before training raise "<Index> must be trained or loaded before <operation>". A trained index without points cannot be saved.
  • HNSW: capacity grows geometrically on add (resize still preallocates), and search bounds k by num_stored_points instead of the capacity, as IVF already did. The first add into an empty index links with exact vectors, like the 0.7 build.

0.8 compatibility layer (removed in 0.9)

  • The pybind classes expose only the 0.8 API. As Faiss does for its SWIG classes, rabitqlib/__init__.py patches the native classes at import time in one marked section: __init__ (with the original kept), build, is_built, and add. The public types stay the native extension types, and static load keeps returning them.
  • Accepted with a DeprecationWarning attributed to the caller's line:
    • IvfIndex(dim, max_elements, num_clusters, nbits, ...) and max_elements=;
    • HnswIndex(dim, max_elements, ...) and random_seed=;
    • build() and is_built;
    • add options passed by position.
  • Old forms keep their 0.7 defaults: build uses one thread and legacy HNSW keeps nbits=8. Mixed or ambiguous calls raise TypeError, which says how the call was read.
  • The stubs declare every old form as an @overload marked @deprecated, so editors flag them.

Behavior changes for existing code

  • add(x) without num_threads now uses every CPU, including in 0.7 code, with no warning. Default HNSW builds are therefore not reproducible; pass num_threads=1. search still defaults to one thread.
  • HNSW capacity grows in add and in the deprecated build, so they no longer raise when max_elements is exceeded. Before training, num_clusters reports the configured count (16 for old-form indexes) instead of 0.
  • State messages say "trained" instead of "built", repr shows is_trained, and HNSW search rejects k > num_stored_points.

Other changes

  • C++: IVF::set_centroids rotates centroids with parallel_for, as construct does. ivf.hpp documents that leaving off the seed in the new constructor selects the point-count constructor.
  • Docs: the README and Quick Start use train + add without k-means. index/ivf.md, hnsw.md, usage.md and clustering.md are updated, and a new Migrating to 0.8 page with a table of old and new calls is added to the nav.
  • Samples: the IVF/HNSW indexing samples train by default and still accept --clusters.

Deliberately left for later

  • Version pins and the "describes 0.7.2" text stay until the 0.8.0 release, after the third PR, even though the docs already show the 0.8 API.
  • benchmarks/run.py still uses the old forms; it will be migrated at release.
  • Existing Python tests stay on the 0.7 forms, where they double as coverage of the compatibility layer.

No saved-index format changes and no new dependencies. Files from earlier releases load as trained indexes; an untrained, empty HNSW file from 0.7 loads with the default of 16 clusters.

Testing

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

  • Formatting and static analysis: ./scripts/check-format.sh, ./scripts/check-python.sh, full ./scripts/check-tidy.sh (library and bindings, clang-tidy 23) and ./scripts/check-isa-symbols.sh passed.
  • CTest (--no-tests=error): 485/485 passed; one test skipped because OMP_PROC_BIND was not set. IvfTrainTest.SetCentroidsAndAddMatchConstructByteForByte now sets centroids on every thread and still matches construct byte for byte.
  • Python, after reinstalling into the project environment and confirming the imported extension and __init__.py are the new build:
    • pytest: 958 passed; 2 skipped (macOS/ARM wheel-only tests).
    • Editor, mypy and Pyright checks: 19 passed. A new mypy check confirms that all seven deprecated forms are reported with --enable-error-code deprecated.
    • The ~990 warnings are all the compatibility layer's own DeprecationWarnings from tests that still use the 0.7 forms.
  • New tests:
    • test_train_api.py: signatures and keyword-only arguments, argument errors, seeds, cp defaults per metric, and that train matches set_centroids with RaBitQKMeans centroids from the same cp. Also the num_threads override, retraining rules, centroid shape, operations before training, saving without points, len, HNSW capacity growth and k bound, and that a failed HNSW add keeps the capacity. A single-threaded first add produces the same file as build.
    • test_compat_08.py: every 0.7 form warns once, at the caller's line. Covers old defaults, mixed forms raising TypeError, build/is_built, positional add, failed build leaving the index unchanged, and build on a trained HNSW index.
    • test_txtai_compat.py: txtai's rabitqlib backend call by call (constructors, build, search, add, remove, max_elements, num_clusters, save/load), allowing only DeprecationWarnings.
    • test_examples.py: the indexing samples train without a clustering file, with DeprecationWarning treated as an error.
    • The three bug-fix tests (failed build, failed add, build on a trained HNSW index) were first run against the previous build and failed there.
  • Old forms give the same results as main: main (6ba97a8) and this branch were both built in scratch copies. In both copies, the IVF point-count rotation seed was pinned by a temporary patch that is not part of this PR. 302 arrays from the 0.7 call forms were compared bit for bit, and none differ. They cover L2/IP, IVF 1/4/9/32 bits with flat and HNSW routing, HNSW 1/4/8 bits, build, search (including HACC and threads), add, remove, resize, saved file bytes and loaded searches.
  • Docs: mkdocs build --strict passed. Every changed Python snippet (README, Quick Start, HNSW workflow, IVF training) was executed with DeprecationWarning as an error.
  • Manual check: an untrained, empty HNSW file written by the pre-feat(index): train centroids inside IVF and HNSW indexes #203 C++ library loads with 16 clusters, trains and accepts add.
  • No performance claims; the 2000-cluster threshold and niter will be measured separately.

IvfIndex and HnswIndex now follow a train -> add -> search workflow, so
callers no longer run k-means and pass centroids and assignments to build.

- IvfIndex(dim, num_clusters, *, nbits, metric, initializer, seed=None)
  and HnswIndex(dim, *, nbits, num_clusters=16, M, ef_construction,
  metric, seed=100). nbits is required; seed=None draws one seed, which
  index.seed reports.
- index.cp (RaBitQKMeansParameters: niter 25, max_points_per_centroid
  256, spherical for inner product, seed from the index),
  train(x, *, num_threads=None), set_centroids(centroids), is_trained,
  len(index).
- add(data, *, cluster_ids=None, num_threads=0, fast_quantization=False).
  train and add default to every CPU; search keeps one thread. HNSW
  capacity grows geometrically on add, and search bounds k by the stored
  point count.
- A 0.8 compatibility layer in rabitqlib/__init__.py patches the native
  classes, as Faiss does for its SWIG classes: the 0.7 constructor forms,
  max_elements, random_seed, build(), is_built and positional add options
  still work, raise DeprecationWarning at the caller, and give the same
  results as 0.7. Stubs mark them @deprecated.
- IVF set_centroids rotates centroids in parallel, as construct does.
- Docs, samples and a 0.7 -> 0.8 migration page use the new API.
@github-actions github-actions Bot added the release:added Release category managed from the PR title label Nov 1, 2026
…uments

The compatibility layer passes inspect.Signature.bind's message through.
Python 3.11 reports a missing keyword-only argument as "missing a required
argument", while 3.12+ says "missing a required keyword-only argument", so
the wheel jobs failed on 3.11. Match both.
@gouyt13
gouyt13 merged commit a4d0c8a into main Nov 1, 2026
39 checks passed
@gouyt13
gouyt13 deleted the feat/python-train-api branch November 1, 2026 10:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

release:added Release category managed from the PR title

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant