Repository navigation
feat(python): train, set_centroids and add for IvfIndex and HnswIndex - #204
Merged
Merged
Conversation
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.
…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.
This was referenced Nov 1, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Second of three PRs moving the Python indexes to a Faiss-style
train→add→searchworkflow, on top of the C++ training support from #203.IvfIndexandHnswIndexnow train their own centroids, so callers no longer run k-means and pass centroids and assignments tobuild. SymphonyQG follows in PR 3.New API
IvfIndex(dim, num_clusters, *, nbits, metric="l2", initializer="auto", seed=None)andHnswIndex(dim, *, nbits, num_clusters=16, M=16, ef_construction=200, metric="l2", seed=100).nbitsis required; HNSW's old default of 8 is gone. Everything exceptdim(and IVF'snum_clusters) is keyword-only.seed=Nonefor IVF draws one random seed at construction, used for the rotation andcp.seed, and reported asindex.seed. This keeps 0.7's random default rotation while making it reproducible. Loaded indexes reportNone.index.cpis aRaBitQKMeansParameterswith index defaults:niter=25,max_points_per_centroid=256,spherical=Truefor inner product, andseedfrom the index.train(x, *, num_threads=None):Noneusescp.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), andadd(data, *, cluster_ids=None, num_threads=0, fast_quantization=False).RuntimeErrorafterwards. Operations before training raise "<Index>must be trained or loaded before<operation>". A trained index without points cannot be saved.add(resizestill preallocates), andsearchboundskbynum_stored_pointsinstead of the capacity, as IVF already did. The firstaddinto an empty index links with exact vectors, like the 0.7build.0.8 compatibility layer (removed in 0.9)
rabitqlib/__init__.pypatches the native classes at import time in one marked section:__init__(with the original kept),build,is_built, andadd. The public types stay the native extension types, and staticloadkeeps returning them.DeprecationWarningattributed to the caller's line:IvfIndex(dim, max_elements, num_clusters, nbits, ...)andmax_elements=;HnswIndex(dim, max_elements, ...)andrandom_seed=;build()andis_built;addoptions passed by position.builduses one thread and legacy HNSW keepsnbits=8. Mixed or ambiguous calls raiseTypeError, which says how the call was read.@overloadmarked@deprecated, so editors flag them.Behavior changes for existing code
add(x)withoutnum_threadsnow uses every CPU, including in 0.7 code, with no warning. Default HNSW builds are therefore not reproducible; passnum_threads=1.searchstill defaults to one thread.addand in the deprecatedbuild, so they no longer raise whenmax_elementsis exceeded. Before training,num_clustersreports the configured count (16 for old-form indexes) instead of 0.reprshowsis_trained, and HNSWsearchrejectsk > num_stored_points.Other changes
IVF::set_centroidsrotates centroids withparallel_for, asconstructdoes.ivf.hppdocuments that leaving off the seed in the new constructor selects the point-count constructor.train+addwithout k-means.index/ivf.md,hnsw.md,usage.mdandclustering.mdare updated, and a new Migrating to 0.8 page with a table of old and new calls is added to the nav.--clusters.Deliberately left for later
benchmarks/run.pystill uses the old forms; it will be migrated at release.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:./scripts/check-format.sh,./scripts/check-python.sh, full./scripts/check-tidy.sh(library and bindings, clang-tidy 23) and./scripts/check-isa-symbols.shpassed.--no-tests=error): 485/485 passed; one test skipped becauseOMP_PROC_BINDwas not set.IvfTrainTest.SetCentroidsAndAddMatchConstructByteForBytenow sets centroids on every thread and still matchesconstructbyte for byte.__init__.pyare the new build:--enable-error-code deprecated.DeprecationWarnings from tests that still use the 0.7 forms.test_train_api.py: signatures and keyword-only arguments, argument errors, seeds,cpdefaults per metric, and thattrainmatchesset_centroidswithRaBitQKMeanscentroids from the samecp. Also thenum_threadsoverride, retraining rules, centroid shape, operations before training, saving without points,len, HNSW capacity growth andkbound, and that a failed HNSWaddkeeps the capacity. A single-threaded firstaddproduces the same file asbuild.test_compat_08.py: every 0.7 form warns once, at the caller's line. Covers old defaults, mixed forms raisingTypeError,build/is_built, positionaladd, failedbuildleaving the index unchanged, andbuildon 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 onlyDeprecationWarnings.test_examples.py: the indexing samples train without a clustering file, withDeprecationWarningtreated as an error.build, failedadd,buildon a trained HNSW index) were first run against the previous build and failed there.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.mkdocs build --strictpassed. Every changed Python snippet (README, Quick Start, HNSW workflow, IVF training) was executed withDeprecationWarningas an error.add.niterwill be measured separately.