From 5314f5c82b24ef3b5c650c371d74e7694af3e8bb Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 12 Nov 2026 12:52:46 +0800 Subject: [PATCH 1/2] test(python): migrate ten test files off the 0.7 call forms Rewrite index setup in ten Python test files with the 0.8 API, so they pass under -W error::DeprecationWarning and keep working after the 0.7 forms are removed in 0.9 (#213): - constructors take num_clusters, nbits and seed by keyword, and SymphonyQG takes ef_construction and nbits=32 for raw vectors; - build() becomes set_centroids() and add(), or add() for SymphonyQG, with num_threads=1 where build() used its single-thread default; - search_batch() becomes search(), and is_built becomes len(). test_concurrent_indexes used build() as a repeatable mutation, which 0.8 does not have. Active operations are now checked to reject add, train and set_centroids, and SymphonyQG's released access is checked through a second add, which takes write access before it rejects the existing graph. No other test checked train or set_centroids for BUSY. test_array_alignment reproduces the IVF rotation with seed= instead of rebuilding a loaded reference, since 0.8 rejects set_centroids on an index that stores points. test_index_repr expects 0.8 capacities and fixes the HNSW num_clusters at construction. Related to #214. --- tests/python/test_array_alignment.py | 20 +++--- tests/python/test_build_cluster_ids.py | 32 +++++---- tests/python/test_concurrent_indexes.py | 86 ++++++++++++++--------- tests/python/test_index_repr.py | 30 ++++---- tests/python/test_kmeans_sampling.py | 5 +- tests/python/test_large_dimensions.py | 26 ++++--- tests/python/test_persistence_paths.py | 13 ++-- tests/python/test_quantization_factors.py | 5 +- tests/python/test_remove_arrays.py | 7 +- tests/python/test_thread_counts.py | 14 ++-- 10 files changed, 137 insertions(+), 101 deletions(-) diff --git a/tests/python/test_array_alignment.py b/tests/python/test_array_alignment.py index bb083f5d..d9b5d9dc 100644 --- a/tests/python/test_array_alignment.py +++ b/tests/python/test_array_alignment.py @@ -46,12 +46,12 @@ def test_indexes_accept_unaligned_build_and_query_arrays(kind, offset, tmp_path) def make_index(): if kind == "ivf": - return IvfIndex(65, len(data), 2, 4) + return IvfIndex(65, 2, nbits=4, seed=42) if kind == "hnsw": return HnswIndex( - 65, len(data), M=8, ef_construction=32, nbits=4, random_seed=42 + 65, num_clusters=2, M=8, ef_construction=32, nbits=4, seed=42 ) - return SymqgIndex(65, 32) + return SymqgIndex(65, max_degree=32, ef_construction=32) if kind == "ivf": parameters = {"nprobe": 2} @@ -60,7 +60,8 @@ def make_index(): if kind != "qg": reference = make_index() - reference.build(data, centroids, labels, num_threads=1) + reference.set_centroids(centroids) + reference.add(data, cluster_ids=labels, num_threads=1) reference_path = tmp_path / "aligned.index" reference.save(str(reference_path)) expected_bytes = reference_path.read_bytes() @@ -71,11 +72,11 @@ def make_index(): if kind != "qg": inputs += [(data, unaligned_centroids), (unaligned_data, unaligned_centroids)] for vectors, centers in inputs: - # IVF samples a rotation at construction; rebuild a loaded reference to - # reuse it. HNSW's explicit seed reproduces its rotation and serial graph. - index = IvfIndex.load(str(reference_path)) if kind == "ivf" else make_index() + # Explicit seeds reproduce the IVF rotation and HNSW's rotation and serial + # graph. + index = make_index() if kind == "qg": - index.build(vectors, ef_construction=32, num_threads=1) + index.add(vectors, num_threads=1) ids, distances = index.search(queries, k=len(data), **parameters) for row in ids: np.testing.assert_array_equal(np.sort(row), np.arange(len(data))) @@ -88,7 +89,8 @@ def make_index(): ) expected = index.search(queries, k=5, num_threads=1, **parameters) else: - index.build(vectors, centers, labels, num_threads=1) + index.set_centroids(centers) + index.add(vectors, cluster_ids=labels, num_threads=1) actual_path = tmp_path / "unaligned.index" index.save(str(actual_path)) assert actual_path.read_bytes() == expected_bytes diff --git a/tests/python/test_build_cluster_ids.py b/tests/python/test_build_cluster_ids.py index 61d8844d..251c19d5 100644 --- a/tests/python/test_build_cluster_ids.py +++ b/tests/python/test_build_cluster_ids.py @@ -1,4 +1,4 @@ -"""Build must validate original assignments before narrowing them to point IDs.""" +"""`add` must validate original assignments before narrowing them to point IDs.""" import numpy as np import pytest @@ -10,33 +10,35 @@ "bad", [2**32, -(2**32), 0.5, -1, 1, np.uint64(2**64 - 1), True, "0"], ) -def test_build_rejects_invalid_cluster_ids(kind, bad): +def test_add_rejects_invalid_cluster_ids(kind, bad): data = np.ones((2, 64), dtype=np.float32) centroids = np.zeros((1, 64), dtype=np.float32) index = ( - IvfIndex(64, 2, 1, 4) + IvfIndex(64, 1, nbits=4) if kind == "ivf" - else HnswIndex(64, 2, M=2, ef_construction=10, nbits=4) + else HnswIndex(64, num_clusters=1, M=2, ef_construction=10, nbits=4) ) + index.set_centroids(centroids) ids = np.full(2, bad) with pytest.raises(ValueError, match="cluster_ids"): - index.build(data, centroids, ids) - assert not index.is_built - index.build(data, centroids, np.zeros(2, dtype=np.int64)) - assert index.is_built + index.add(data, cluster_ids=ids, num_threads=1) + assert len(index) == 0 + index.add(data, cluster_ids=np.zeros(2, dtype=np.int64), num_threads=1) + assert len(index) == 2 @pytest.mark.parametrize("kind", ["ivf", "hnsw"]) @pytest.mark.parametrize("layout", ["native", "strided", "swapped", "unaligned"]) -def test_build_uint32_cluster_ids_validate_and_preserve_input(kind, layout): +def test_add_uint32_cluster_ids_validate_and_preserve_input(kind, layout): data = np.ones((2, 64), dtype=np.float32) data[1] = 2 centroids = data.copy() index = ( - IvfIndex(64, 2, 2, 4) + IvfIndex(64, 2, nbits=4) if kind == "ivf" - else HnswIndex(64, 2, M=2, ef_construction=10, nbits=4) + else HnswIndex(64, num_clusters=2, M=2, ef_construction=10, nbits=4) ) + index.set_centroids(centroids) if layout == "strided": ids = np.zeros(4, dtype=np.uint32)[::2] elif layout == "swapped": @@ -47,11 +49,11 @@ def test_build_uint32_cluster_ids_validate_and_preserve_input(kind, layout): ids = np.zeros(2, dtype=np.uint32) ids[1] = np.iinfo(np.uint32).max with pytest.raises(ValueError, match="cluster_ids"): - index.build(data, centroids, ids, num_threads=1) - assert not index.is_built + index.add(data, cluster_ids=ids, num_threads=1) + assert len(index) == 0 ids[1] = 1 ids.flags.writeable = False before = ids.copy() - index.build(data, centroids, ids, num_threads=1) - assert index.is_built + index.add(data, cluster_ids=ids, num_threads=1) + assert len(index) == 2 np.testing.assert_array_equal(ids, before) diff --git a/tests/python/test_concurrent_indexes.py b/tests/python/test_concurrent_indexes.py index 63c21c3f..787caf04 100644 --- a/tests/python/test_concurrent_indexes.py +++ b/tests/python/test_concurrent_indexes.py @@ -18,33 +18,41 @@ ("ivf_flat_rabitq", 4), ("hnsw", 1), ("hnsw", 9), - ("qg", 0), + ("qg", 32), ("qg", 4), ("qg", 8), ] -def make_index(kind, bits=4, metric="l2"): +def make_index(kind, bits=4, metric="l2", fill=True): rng = np.random.default_rng(713) data = rng.standard_normal((129, 65)).astype(np.float32) labels = (np.arange(len(data)) % 5).astype(np.uint32) centroids = np.stack([data[labels == i].mean(axis=0) for i in range(5)]) if kind.startswith("ivf"): initializer = kind.removeprefix("ivf_") if "_" in kind else "flat" - index = IvfIndex(65, len(data), 5, bits, metric, initializer) + index = IvfIndex(65, 5, nbits=bits, metric=metric, initializer=initializer) elif kind == "hnsw": - index = HnswIndex(65, len(data) + 8, 8, 64, bits, metric, 42) + index = HnswIndex( + 65, + num_clusters=5, + M=8, + ef_construction=64, + nbits=bits, + metric=metric, + seed=42, + ) else: - index = SymqgIndex(65, 32, metric, bits) - - def build(vectors=data): + index = SymqgIndex( + 65, max_degree=32, ef_construction=64, metric=metric, nbits=bits + ) + if fill: if kind == "qg": - index.build(vectors, 64, 1) + index.add(data, num_threads=1) else: - index.build(vectors, centroids, labels, 1) - - build() - return index, data, build + index.set_centroids(centroids) + index.add(data, cluster_ids=labels, num_threads=1) + return index, data def parameters(kind, wide): @@ -56,7 +64,7 @@ def parameters(kind, wide): @pytest.mark.parametrize("kind,bits", CASES) @pytest.mark.parametrize("metric", ["l2", "ip"]) def test_concurrent_search_matches_serial_and_loaded(kind, bits, metric, tmp_path): - index, _, _ = make_index(kind, bits, metric) + index, _ = make_index(kind, bits, metric) # Exercise conversions and ownership of a temporary contiguous float32 copy. queries = np.random.default_rng(18).standard_normal((17, 130))[:, ::2] expected = [index.search(queries, 5, **parameters(kind, wide)) for wide in (0, 1)] @@ -69,12 +77,7 @@ def search(worker): wide = worker % 2 start.wait(timeout=10) for count in (0, 1, 17, 17): - method = ( - current.search_batch - if kind == "qg" and worker >= 2 - else current.search - ) - actual = method( + actual = current.search( queries[:count], 5, num_threads=1 + worker // 2, @@ -104,7 +107,7 @@ def __array__(self, dtype=None, copy=None): @pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) def test_active_search_allows_reads_and_rejects_mutations(kind, tmp_path): - index, data, build = make_index(kind) + index, data = make_index(kind) queries = PausedArray(data[:3]) summary = repr(index) options = parameters(kind, True) @@ -116,9 +119,13 @@ def test_active_search_allows_reads_and_rejects_mutations(kind, tmp_path): assert repr(index) == summary index.search(data[:1], 5, **options) index.save(str(tmp_path / "reading.index")) - mutations = [build] + mutations = [lambda: index.add(data[:1])] if kind != "qg": - mutations += [lambda: index.add(data[:1]), lambda: index.remove([0])] + mutations += [ + lambda: index.remove([0]), + lambda: index.train(data), + lambda: index.set_centroids(data[:5]), + ] if kind == "hnsw": mutations.append(lambda: index.resize(256)) for mutate in mutations: @@ -130,47 +137,56 @@ def test_active_search_allows_reads_and_rejects_mutations(kind, tmp_path): # Failed conflicts must not leave the access state busy. if kind == "hnsw": index.resize(256) + elif kind == "qg": + # add takes write access before it rejects a second graph. + with pytest.raises(RuntimeError, match="does not support incremental addition"): + index.add(data) else: - build() + index.add(data[:1]) index.search(data[:1], 5, **options) @pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) def test_active_mutation_rejects_reads_and_other_mutations(kind, tmp_path): - index, data, build = make_index(kind) - vectors = PausedArray(data[:1] if kind == "hnsw" else data) + # SymphonyQG adds all vectors at once, so its index starts empty. + index, data = make_index(kind, fill=kind != "qg") + vectors = PausedArray(data if kind == "qg" else data[:1]) with ThreadPoolExecutor(max_workers=1) as pool: - running = pool.submit(index.add if kind == "hnsw" else build, vectors) + running = pool.submit(index.add, vectors, num_threads=1) try: assert vectors.entered.wait(timeout=10) calls = [ lambda: index.search(data[:1], 5, **parameters(kind, True)), lambda: index.save(str(tmp_path / "busy.index")), lambda: index.dim, - lambda: index.is_built, + lambda: len(index), lambda: repr(index), lambda: str(index), - build, + lambda: index.add(data[:1]), ] if kind != "qg": - calls += [lambda: index.add(data[:1]), lambda: index.remove([0])] + calls += [ + lambda: index.remove([0]), + lambda: index.train(data), + lambda: index.set_centroids(data[:5]), + ] if kind == "hnsw": calls.append(lambda: index.resize(256)) for call in calls: with pytest.raises(RuntimeError, match=f"^{BUSY}$"): call() - other, other_data, _ = make_index(kind) + other, other_data = make_index(kind) other.search(other_data[:1], 5, **parameters(kind, True)) finally: vectors.resume.set() running.result(timeout=10) - assert index.is_built + assert len(index) == (len(data) if kind == "qg" else len(data) + 1) index.search(data[:1], 5, **parameters(kind, True)) @pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) def test_search_releases_gil(kind): - index, data, _ = make_index(kind) + index, data = make_index(kind) queries = np.tile(data, (32, 1)) ready, progress = Event(), Event() active = stopped = False @@ -208,12 +224,14 @@ def python_worker(): def test_native_exception_reacquires_gil_and_releases_access(): - index, data, build = make_index("qg") + index, data = make_index("qg") queries = data[:2].copy() queries[1, 0] = np.nan # The native search checks coordinates after releasing the GIL. with pytest.raises(ValueError, match="must contain only finite values"): index.search(queries, 5, ef=64) - build() + # add takes write access before it rejects a second graph. + with pytest.raises(RuntimeError, match="does not support incremental addition"): + index.add(data) ids, distances = index.search(data[:2], 5, ef=64) assert ids.shape == distances.shape == (2, 5) diff --git a/tests/python/test_index_repr.py b/tests/python/test_index_repr.py index 667c5217..cc7c32ac 100644 --- a/tests/python/test_index_repr.py +++ b/tests/python/test_index_repr.py @@ -16,7 +16,7 @@ "num_clusters={clusters}, is_trained={built})" ), "symqg": ( - "SymqgIndex(dim=64, metric='ip', nbits={nbits}, " + "SymqgIndex(dim=64, metric='ip', nbits={bits}, " "num_points={live}, max_degree=32)" ), } @@ -29,42 +29,44 @@ def vectors(): def unbuilt_index(kind, bits): if kind == "ivf": - return IvfIndex(64, 80, 1, bits, metric="innerproduct") + return IvfIndex(64, 1, nbits=bits, metric="innerproduct") if kind == "hnsw": - return HnswIndex(64, 96, nbits=bits, metric="innerproduct") - return SymqgIndex(64, 32, metric="innerproduct", quantization_bits=bits) + index = HnswIndex(64, nbits=bits, num_clusters=1, metric="innerproduct") + # Keep capacity distinct from the 80 stored points. + index.resize(96) + return index + return SymqgIndex( + 64, max_degree=32, ef_construction=40, metric="innerproduct", nbits=bits + ) def build_index(index, vectors): if isinstance(index, SymqgIndex): - index.build(vectors, 40) + index.add(vectors, num_threads=1) else: - index.build( - vectors, vectors.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32) - ) + index.set_centroids(vectors.mean(axis=0, keepdims=True)) + index.add(vectors, cluster_ids=np.zeros(80, dtype=np.uint32), num_threads=1) @pytest.mark.parametrize( - "kind,bits", [("ivf", 4), ("ivf", 32), ("hnsw", 8), ("symqg", 0), ("symqg", 8)] + "kind,bits", [("ivf", 4), ("ivf", 32), ("hnsw", 8), ("symqg", 32), ("symqg", 8)] ) def test_repr_lifecycle(kind, bits, vectors, tmp_path): index = unbuilt_index(kind, bits) - capacity = 96 if kind == "hnsw" else 80 before = SUMMARIES[kind].format( bits=bits, - nbits=32 if bits == 0 else bits, live=0, stored=0, - capacity=capacity, - clusters=16 if kind == "hnsw" else 1, + capacity=96 if kind == "hnsw" else 0, + clusters=1, built=False, ) assert repr(index) == str(index) == before build_index(index, vectors) + capacity = 96 if kind == "hnsw" else 80 after = SUMMARIES[kind].format( bits=bits, - nbits=32 if bits == 0 else bits, live=80, stored=80, capacity=capacity, diff --git a/tests/python/test_kmeans_sampling.py b/tests/python/test_kmeans_sampling.py index ee08fc28..09c4de9c 100644 --- a/tests/python/test_kmeans_sampling.py +++ b/tests/python/test_kmeans_sampling.py @@ -255,8 +255,9 @@ def test_sampled_results_build_an_ivf_index(model_type): x = clustered(4 * clusters, 23) model = make_model(model_type, max_points_per_centroid=1) model.train(x) - index = IvfIndex(DIM, len(x), clusters, 4) - index.build(x, model.centroids, model.assignments) + index = IvfIndex(DIM, clusters, nbits=4) + index.set_centroids(model.centroids) + index.add(x, cluster_ids=model.assignments, num_threads=1) ids, _ = index.search(x[:3], 5, nprobe=clusters) for query, neighbors in enumerate(ids): assert query in neighbors diff --git a/tests/python/test_large_dimensions.py b/tests/python/test_large_dimensions.py index d03b56ee..29765daf 100644 --- a/tests/python/test_large_dimensions.py +++ b/tests/python/test_large_dimensions.py @@ -14,21 +14,29 @@ def test_large_dimension_round_trip(tmp_path, dim, metric, kind): data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32) data /= np.linalg.norm(data, axis=1, keepdims=True) if kind == "ivf": - index = IvfIndex(dim, count, 1, nbits=4, metric=metric) + index = IvfIndex(dim, 1, nbits=4, metric=metric) search_args = {"nprobe": 1} elif kind == "hnsw": index = HnswIndex( - dim, count, M=count, ef_construction=count, nbits=4, metric=metric + dim, + num_clusters=1, + M=count, + ef_construction=count, + nbits=4, + metric=metric, ) search_args = {"ef": count} else: - index = SymqgIndex(dim, max_degree=32, quantization_bits=4, metric=metric) + index = SymqgIndex( + dim, max_degree=32, ef_construction=count, nbits=4, metric=metric + ) search_args = {"ef": count} if kind == "qg": - index.build(data, ef_construction=count, num_threads=2) + index.add(data, num_threads=2) else: - index.build(data, data.mean(axis=0, keepdims=True), np.zeros(count, np.uint32)) + index.set_centroids(data.mean(axis=0, keepdims=True)) + index.add(data, cluster_ids=np.zeros(count, np.uint32), num_threads=1) queries = data[:3] ids, distances = index.search(queries, k=1, **search_args) np.testing.assert_array_equal(ids[:, 0], np.arange(3)) @@ -47,9 +55,9 @@ def test_rejects_dimensions_above_fht_limit(kind): dim = 65537 with pytest.raises(ValueError, match="Unsupported dimension for FhtKacRotator"): if kind == "ivf": - IvfIndex(dim, 40, 1, nbits=4) + IvfIndex(dim, 1, nbits=4) elif kind == "hnsw": - HnswIndex(dim, 40, nbits=4) + HnswIndex(dim, nbits=4) else: - index = SymqgIndex(dim, max_degree=32, quantization_bits=4) - index.build(np.zeros((40, dim), dtype=np.float32), ef_construction=40) + index = SymqgIndex(dim, max_degree=32, ef_construction=40, nbits=4) + index.add(np.zeros((40, dim), dtype=np.float32), num_threads=1) diff --git a/tests/python/test_persistence_paths.py b/tests/python/test_persistence_paths.py index 912c2dda..b03a7923 100644 --- a/tests/python/test_persistence_paths.py +++ b/tests/python/test_persistence_paths.py @@ -10,18 +10,19 @@ def test_unicode_path_round_trip(tmp_path, kind, unicode_directory): data = np.random.default_rng(42).standard_normal((64, 64)).astype(np.float32) if kind.startswith("ivf"): - index = IvfIndex(64, 64, 1, nbits=32 if kind == "ivf_raw" else 4) + index = IvfIndex(64, 1, nbits=32 if kind == "ivf_raw" else 4) elif kind == "hnsw": - index = HnswIndex(64, 64, M=4, nbits=4) + index = HnswIndex(64, num_clusters=1, M=4, nbits=4) else: - bits = {"qg": 0, "qg4": 4, "qg8": 8}[kind] - index = SymqgIndex(64, max_degree=32, quantization_bits=bits) + bits = {"qg": 32, "qg4": 4, "qg8": 8}[kind] + index = SymqgIndex(64, max_degree=32, ef_construction=64, nbits=bits) if kind.startswith("qg"): - index.build(data, ef_construction=64) + index.add(data, num_threads=1) search_args = {"ef": 64} else: - index.build(data, data.mean(axis=0, keepdims=True), np.zeros(64, np.uint32)) + index.set_centroids(data.mean(axis=0, keepdims=True)) + index.add(data, cluster_ids=np.zeros(64, np.uint32), num_threads=1) search_args = {"nprobe": 1} if kind.startswith("ivf") else {"ef": 64} directory = tmp_path / "\u6d4b\u8bd5_\U0001f680" if unicode_directory else tmp_path diff --git a/tests/python/test_quantization_factors.py b/tests/python/test_quantization_factors.py index aebbb247..15cffdac 100644 --- a/tests/python/test_quantization_factors.py +++ b/tests/python/test_quantization_factors.py @@ -11,8 +11,9 @@ def test_collinear_ivf_query_keeps_finite_distances(nbits, metric, high_accuracy): data = np.full((1, 512), 0.1, dtype=np.float32) centroid = np.zeros_like(data) - index = IvfIndex(data.shape[1], len(data), 1, nbits=nbits, metric=metric) - index.build(data, centroid, np.zeros(len(data), dtype=np.uint32), num_threads=1) + index = IvfIndex(data.shape[1], 1, nbits=nbits, metric=metric) + index.set_centroids(centroid) + index.add(data, cluster_ids=np.zeros(len(data), dtype=np.uint32), num_threads=1) # The C++ regression fixes the rotated residual to trigger negative roundoff. # This exercises the normal Python boundary with its randomly sampled rotation. diff --git a/tests/python/test_remove_arrays.py b/tests/python/test_remove_arrays.py index d53b4c6b..e746561b 100644 --- a/tests/python/test_remove_arrays.py +++ b/tests/python/test_remove_arrays.py @@ -14,11 +14,12 @@ def test_remove_integer_views_preserve_input_and_validate_atomically( count, dim = 8, 64 data = np.random.default_rng(42).normal(size=(count, dim)).astype(np.float32) index = ( - IvfIndex(dim, count, 1, 4) + IvfIndex(dim, 1, nbits=4) if kind == "ivf" - else HnswIndex(dim, count, M=2, ef_construction=10, nbits=4) + else HnswIndex(dim, num_clusters=1, M=2, ef_construction=10, nbits=4) ) - index.build(data, np.zeros((1, dim), np.float32), np.zeros(count, np.uint32)) + index.set_centroids(np.zeros((1, dim), np.float32)) + index.add(data, cluster_ids=np.zeros(count, np.uint32), num_threads=1) search_args = {"nprobe": 1} if kind == "ivf" else {"ef": count} before = index.search(data, k=1, **search_args) diff --git a/tests/python/test_thread_counts.py b/tests/python/test_thread_counts.py index 4896d8a9..888febcc 100644 --- a/tests/python/test_thread_counts.py +++ b/tests/python/test_thread_counts.py @@ -19,23 +19,23 @@ def test_build_and_search_thread_counts(kind, metric, build_threads): data = rng.standard_normal((96, 64)).astype(np.float32) data /= np.linalg.norm(data, axis=1, keepdims=True) if kind == "ivf": - index = IvfIndex(64, len(data), 1, nbits=8, metric=metric) + index = IvfIndex(64, 1, nbits=8, metric=metric) search_args = {"nprobe": 1} elif kind == "hnsw": index = HnswIndex( - 64, len(data), M=16, ef_construction=96, nbits=8, metric=metric + 64, num_clusters=1, M=16, ef_construction=96, nbits=8, metric=metric ) search_args = {"ef": 96} else: - index = SymqgIndex(64, max_degree=32, metric=metric) + index = SymqgIndex(64, max_degree=32, ef_construction=96, metric=metric) search_args = {"ef": 96} if kind == "qg": - index.build(data, ef_construction=96, num_threads=build_threads) + index.add(data, num_threads=build_threads) else: - index.build( + index.set_centroids(data.mean(axis=0, keepdims=True)) + index.add( data, - data.mean(axis=0, keepdims=True), - np.zeros(len(data), dtype=np.uint32), + cluster_ids=np.zeros(len(data), dtype=np.uint32), num_threads=build_threads, ) From 06016bdbab4221873e8c613c030e92b50c7a2600 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Thu, 12 Nov 2026 13:29:12 +0800 Subject: [PATCH 2/2] test(python): stop requiring every graph point in two tests Apply the #216 fixes to the two migrated tests that failed in CI on this branch, each time on a graph build that left one point unreachable: - test_array_alignment builds SymphonyQG with seed=42 and one thread, as the IVF and HNSW cases already do. It compares the saved bytes and the k=5 results of aligned and unaligned inputs, and drops the k=len(data) search that required every point to be reachable. The exact-distance check now runs on the k=5 reference results. - test_thread_counts keeps exact self-matches for IVF, but requires a self-recall of at least 95% for HNSW and SymphonyQG, whose parallel or unseeded builds can leave a point unreachable. Searches on the same graph must still agree across thread counts. Related to #216. --- tests/python/test_array_alignment.py | 61 ++++++++++++++-------------- tests/python/test_thread_counts.py | 7 +++- 2 files changed, 37 insertions(+), 31 deletions(-) diff --git a/tests/python/test_array_alignment.py b/tests/python/test_array_alignment.py index d9b5d9dc..968df5bd 100644 --- a/tests/python/test_array_alignment.py +++ b/tests/python/test_array_alignment.py @@ -51,49 +51,50 @@ def make_index(): return HnswIndex( 65, num_clusters=2, M=8, ef_construction=32, nbits=4, seed=42 ) - return SymqgIndex(65, max_degree=32, ef_construction=32) + return SymqgIndex(65, max_degree=32, ef_construction=32, seed=42) + + def fill(index, vectors, centers): + if kind == "qg": + index.add(vectors, num_threads=1) + else: + index.set_centroids(centers) + index.add(vectors, cluster_ids=labels, num_threads=1) if kind == "ivf": parameters = {"nprobe": 2} else: parameters = {"ef": 65} - if kind != "qg": - reference = make_index() - reference.set_centroids(centroids) - reference.add(data, cluster_ids=labels, num_threads=1) - reference_path = tmp_path / "aligned.index" - reference.save(str(reference_path)) - expected_bytes = reference_path.read_bytes() - expected = reference.search(queries, k=5, num_threads=1, **parameters) + # Explicit seeds and one build thread make every index reproducible. + reference = make_index() + fill(reference, data, centroids) + reference_path = tmp_path / "aligned.index" + reference.save(str(reference_path)) + expected_bytes = reference_path.read_bytes() + expected = reference.search(queries, k=5, num_threads=1, **parameters) + if kind == "qg": + # SymphonyQG stores raw vectors, so its distances are exact. + exact = np.sum( + (queries.astype(np.float64)[:, None] - data.astype(np.float64)) ** 2, + axis=2, + ) + np.testing.assert_allclose( + expected[1], + np.take_along_axis(exact, expected[0], axis=1), + rtol=2e-6, + atol=1e-5, + ) # Isolate each conversion so a damaged centroid cannot hide behind data damage. inputs = [(unaligned_data, centroids)] if kind != "qg": inputs += [(data, unaligned_centroids), (unaligned_data, unaligned_centroids)] for vectors, centers in inputs: - # Explicit seeds reproduce the IVF rotation and HNSW's rotation and serial - # graph. index = make_index() - if kind == "qg": - index.add(vectors, num_threads=1) - ids, distances = index.search(queries, k=len(data), **parameters) - for row in ids: - np.testing.assert_array_equal(np.sort(row), np.arange(len(data))) - exact = np.sum( - (queries.astype(np.float64)[:, None] - data.astype(np.float64)) ** 2, - axis=2, - ) - np.testing.assert_allclose( - distances, np.take_along_axis(exact, ids, axis=1), rtol=2e-6, atol=1e-5 - ) - expected = index.search(queries, k=5, num_threads=1, **parameters) - else: - index.set_centroids(centers) - index.add(vectors, cluster_ids=labels, num_threads=1) - actual_path = tmp_path / "unaligned.index" - index.save(str(actual_path)) - assert actual_path.read_bytes() == expected_bytes + fill(index, vectors, centers) + actual_path = tmp_path / "unaligned.index" + index.save(str(actual_path)) + assert actual_path.read_bytes() == expected_bytes for values in (queries, unaligned_queries): actual = index.search(values, k=5, num_threads=1, **parameters) for result, reference_result in zip(actual, expected, strict=True): diff --git a/tests/python/test_thread_counts.py b/tests/python/test_thread_counts.py index 888febcc..d3f5267a 100644 --- a/tests/python/test_thread_counts.py +++ b/tests/python/test_thread_counts.py @@ -43,7 +43,12 @@ def test_build_and_search_thread_counts(kind, metric, build_threads): expected_ids, expected_distances = index.search( data, k=5, num_threads=1, **search_args ) - np.testing.assert_array_equal(expected_ids[:, 0], np.arange(len(data))) + self_recall = np.mean(expected_ids[:, 0] == np.arange(len(data))) + if kind == "ivf": + assert self_recall == 1 + else: + # A parallel or unseeded graph build can leave a point unreachable (#216). + assert self_recall >= 0.95 assert np.isfinite(expected_distances).all() for threads in (0, MAX_THREADS): ids, distances = index.search(data, k=5, num_threads=threads, **search_args)