From e9f2531e476c1ca84f1e7e1d321de0aa659a70e2 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Fri, 13 Nov 2026 13:15:48 +0800 Subject: [PATCH] test(python): migrate the last seven test files off the 0.7 call forms Migrate test_ivf, test_hnsw, test_symqg, test_input_boundaries, test_ivf_initializers, test_portable_indexes and test_import to the 0.8 API: constructors without max_elements, and set_centroids plus add in place of build. Calls that relied on build's single thread pass num_threads=1, and HNSW passes nbits and num_clusters explicitly. Tests of 0.7-only behavior move to test_compat_08.py, wrapped in quietly: build's row-count diagnostics, SymphonyQG rebuilds and build's default init, HNSW's rejected rebuild and zero capacity, and IVF's rebuild of a loaded index. test_invalid_quantization_bits_raises is deleted, because test_compat_08.py already covers it. Where 0.8 behaves differently, the expectations follow it: set_centroids and HNSW add report their own messages, an untrained IvfIndex trains on its first add unless cluster IDs are given, and SymphonyQG's access-guard check uses add, its only write. IvfIndex.set_centroids builds the HNSW routing graph with every CPU, so the 128-centroid routing test checks self-recall in aggregate. Every test file except the four exempt ones now passes with -W error::DeprecationWarning. --- tests/python/test_compat_08.py | 161 +++++++++++++++++++++- tests/python/test_hnsw.py | 156 ++++++++++++--------- tests/python/test_import.py | 18 +-- tests/python/test_input_boundaries.py | 84 +++++------- tests/python/test_ivf.py | 188 +++++++++++++------------- tests/python/test_ivf_initializers.py | 50 ++++--- tests/python/test_portable_indexes.py | 13 +- tests/python/test_symqg.py | 163 +++++++--------------- 8 files changed, 475 insertions(+), 358 deletions(-) diff --git a/tests/python/test_compat_08.py b/tests/python/test_compat_08.py index 288a5f35..b677431c 100644 --- a/tests/python/test_compat_08.py +++ b/tests/python/test_compat_08.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from conftest import DIM, N_VECTORS +from conftest import DIM, N_CLUSTERS, N_VECTORS from rabitqlib import HnswIndex, IvfIndex, SymqgIndex @@ -315,3 +315,162 @@ def test_symqg_build_draws_a_new_seed_unless_one_was_given(base_data): for _ in range(2): quietly(lambda: seeded.build(base_data, 50)) assert seeded.seed == 3 + + +# Moved from test_input_boundaries.py (#214): build's row count has no 0.8 form. +@pytest.mark.parametrize("index_kind,rows", [("ivf", 79), ("hnsw", 0)]) +def test_build_row_count_diagnostics(index_kind, rows): + if index_kind == "ivf": + index = quietly(lambda: IvfIndex(64, 80, 1, 4)) + expected = "number of data rows must match index max_elements; expected 80, got shape (79, 64)" + else: + # HNSW capacity grows, so only an empty build is rejected. + index = quietly(lambda: HnswIndex(64, 80)) + expected = "number of data rows must be at least 1; got shape (0, 64)" + with pytest.raises(ValueError) as error: + quietly( + lambda: index.build( + np.zeros((rows, 64)), + np.zeros((1, 64)), + np.zeros(rows, dtype=np.uint32), + ) + ) + assert str(error.value) == expected + + +# Moved from test_symqg.py (#214): build replacing a graph has no 0.8 form. +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("replacement_points", [0, 32]) +def test_symqg_failed_rebuild_preserves_index( + base_data, query_data, tmp_path, bits, metric, replacement_points +): + data = base_data[:40] + queries = query_data[:3] + index = quietly( + lambda: SymqgIndex(DIM, max_degree=32, metric=metric, quantization_bits=bits) + ) + quietly(lambda: index.build(data, ef_construction=50)) + expected_ids, expected_distances = index.search(queries, k=5, ef=50) + before_path = tmp_path / "before.index" + index.save(str(before_path)) + + with pytest.raises( + ValueError, + match="QuantizedGraph degree bound must be smaller than the number of points", + ): + quietly(lambda: index.build(data[:replacement_points], ef_construction=50)) + + assert quietly(lambda: index.is_built) + assert index.num_points == len(data) + assert index.dim == DIM + assert index.max_degree == 32 + assert index.metric == metric + assert quietly(lambda: index.quantization_bits) == bits + ids, distances = index.search(queries, k=5, ef=50) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) + + after_path = tmp_path / "after.index" + index.save(str(after_path)) + assert after_path.read_bytes() == before_path.read_bytes() + restored = SymqgIndex.load(str(after_path)) + assert restored.num_points == len(data) + restored_ids, restored_distances = restored.search(queries, k=5, ef=50) + np.testing.assert_array_equal(restored_ids, expected_ids) + np.testing.assert_array_equal(restored_distances, expected_distances) + + +# Moved from test_symqg.py (#214). +@pytest.mark.parametrize("bits", [0, 4, 8]) +def test_symqg_successful_rebuild_replaces_index(base_data, tmp_path, bits): + # Seeded, and build uses one thread, so every self-query below must succeed. + index = quietly( + lambda: SymqgIndex(DIM, max_degree=32, quantization_bits=bits, seed=42) + ) + quietly(lambda: index.build(base_data[:40], ef_construction=50)) + replacement = base_data[40:88] + 10 + quietly(lambda: index.build(replacement, ef_construction=50)) + + assert quietly(lambda: index.is_built) + assert index.num_points == len(replacement) + assert quietly(lambda: index.quantization_bits) == bits + ids, distances = index.search(replacement[:3], k=1, ef=50) + np.testing.assert_array_equal(ids[:, 0], np.arange(3)) + + path = str(tmp_path / "rebuilt.index") + index.save(path) + restored = SymqgIndex.load(path) + assert restored.num_points == len(replacement) + restored_ids, restored_distances = restored.search(replacement[:3], k=1, ef=50) + np.testing.assert_array_equal(restored_ids, ids) + np.testing.assert_array_equal(restored_distances, distances) + + +# Moved from test_symqg.py (#214). +def test_symqg_build_defaults_to_pipnn(): + assert inspect.signature(SymqgIndex.build).parameters["init"].default == "pipnn" + + +# Moved from test_hnsw.py (#214): build on a built index has no 0.8 form. +@pytest.mark.parametrize("reload", [False, True]) +def test_hnsw_rejected_rebuild_preserves_index(base_data, clusters, tmp_path, reload): + idx = quietly(lambda: HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4)) + centroids, cluster_ids = clusters + quietly(lambda: idx.build(base_data, centroids, cluster_ids)) + if reload: + path = str(tmp_path / "hnsw.index") + idx.save(path) + idx = HnswIndex.load(path) + before = idx.search(base_data[:10], k=3, ef=50) + with pytest.raises( + RuntimeError, match="^HNSW index is already trained, constructed or loaded$" + ): + quietly( + lambda: idx.build( + base_data, centroids[:1], np.zeros(N_VECTORS, dtype=np.uint32) + ) + ) + assert quietly(lambda: idx.is_built) + assert idx.num_clusters == N_CLUSTERS + after = idx.search(base_data[:10], k=3, ef=50) + for actual, expected in zip(after, before, strict=True): + np.testing.assert_array_equal(actual, expected) + + +# Moved from test_hnsw.py (#214): the 0.8 constructor takes no max_elements. +def test_hnsw_zero_capacity_raises(): + with pytest.raises( + ValueError, match="^HNSW max_elements must be greater than zero$" + ): + quietly(lambda: HnswIndex(max_elements=0, dim=64)) + + +# Moved from test_ivf.py (#214): build on a loaded index rebuilds it, which has no +# 0.8 form. C++ IvfConstructionTest covers the same encoding. +@pytest.mark.parametrize("bits", [1, 4, 32]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_ivf_rebuild_matches_incremental_encoding_with_empty_clusters( + tmp_path, bits, metric +): + rng = np.random.default_rng(731) + data = rng.standard_normal((130, 65), dtype=np.float32) + centers = rng.standard_normal((4, 65), dtype=np.float32) + labels = np.full(130, 3, dtype=np.uint32) + labels[0] = 1 + labels[1:33] = 2 + incremental = quietly(lambda: IvfIndex(65, 1, 4, nbits=bits, metric=metric)) + quietly(lambda: incremental.build(data[:1], centers, labels[:1], num_threads=1)) + incremental.add(data[1:], cluster_ids=labels[1:], num_threads=1) + path = tmp_path / "incremental.index" + incremental.save(str(path)) + expected_bytes = path.read_bytes() + expected_ids, expected_distances = incremental.search(data[:3], k=10, nprobe=4) + rebuilt = IvfIndex.load(str(path)) + for threads in (1, 4): + quietly(lambda: rebuilt.build(data, centers, labels, num_threads=threads)) + rebuilt.save(str(path)) + assert path.read_bytes() == expected_bytes + ids, distances = rebuilt.search(data[:3], k=10, nprobe=4) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index c478fcb5..d57c8397 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -15,12 +15,25 @@ # ── fixtures ────────────────────────────────────────────────────────────────── +def _built( + data, centroids, cluster_ids, num_threads=1, fast_quantization=False, **options +): + """An index over data with the given centroids; the first add links the points.""" + index = HnswIndex(data.shape[1], num_clusters=len(centroids), **options) + index.set_centroids(centroids) + index.add( + data, + cluster_ids=cluster_ids, + num_threads=num_threads, + fast_quantization=fast_quantization, + ) + return index + + @pytest.fixture(scope="module") def built_hnsw(base_data, clusters): - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters - idx.build(base_data, centroids, cluster_ids) - return idx + return _built(base_data, centroids, cluster_ids, M=8, ef_construction=50, nbits=4) # ── construction ────────────────────────────────────────────────────────────── @@ -29,32 +42,11 @@ def built_hnsw(base_data, clusters): @pytest.mark.parametrize("m", [0, 1]) def test_rejects_too_small_m(m): with pytest.raises(ValueError, match="^HNSW M must be at least 2$"): - HnswIndex(DIM, N_VECTORS, M=m) - - -@pytest.mark.parametrize("reload", [False, True]) -def test_rejected_rebuild_preserves_index(base_data, clusters, tmp_path, reload): - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) - centroids, cluster_ids = clusters - idx.build(base_data, centroids, cluster_ids) - if reload: - path = str(tmp_path / "hnsw.index") - idx.save(path) - idx = HnswIndex.load(path) - before = idx.search(base_data[:10], k=3, ef=_EF) - with pytest.raises( - RuntimeError, match="^HNSW index is already trained, constructed or loaded$" - ): - idx.build(base_data, centroids[:1], np.zeros(N_VECTORS, dtype=np.uint32)) - assert idx.is_built - assert idx.num_clusters == N_CLUSTERS - after = idx.search(base_data[:10], k=3, ef=_EF) - for actual, expected in zip(after, before, strict=True): - np.testing.assert_array_equal(actual, expected) + HnswIndex(DIM, M=m) def test_properties(built_hnsw): - assert built_hnsw.is_built + assert len(built_hnsw) == N_VECTORS assert built_hnsw.dim == DIM assert built_hnsw.nbits == 4 assert built_hnsw.metric == "l2" @@ -71,10 +63,9 @@ def test_seed_reproduces_serial_build_and_search(tmp_path): snapshots = [] results = [] for run, seed in enumerate((42, 42, 43)): - index = HnswIndex( - 65, len(data), M=8, ef_construction=32, nbits=4, random_seed=seed + index = _built( + data, centroids, clusters, M=8, ef_construction=32, nbits=4, seed=seed ) - index.build(data, centroids, clusters, num_threads=1) path = tmp_path / f"seed-{run}.index" index.save(str(path)) snapshots.append(path.read_bytes()) @@ -89,9 +80,16 @@ def test_parallel_build_preserves_self_retrieval(base_data, clusters): # M equals the point count, so the base layer holds 2M >= N - 1 links and never # prunes a reverse link: every point stays reachable under any thread # interleaving. This checks parallel construction, not sparse-graph recall. - idx = HnswIndex(DIM, N_VECTORS, M=N_VECTORS, ef_construction=N_VECTORS, nbits=4) centroids, cluster_ids = clusters - idx.build(base_data, centroids, cluster_ids, num_threads=8) + idx = _built( + base_data, + centroids, + cluster_ids, + num_threads=8, + M=N_VECTORS, + ef_construction=N_VECTORS, + nbits=4, + ) ids, distances = idx.search(base_data, k=1, ef=N_VECTORS) np.testing.assert_array_equal(ids[:, 0], np.arange(N_VECTORS)) assert np.isfinite(distances).all() @@ -129,10 +127,17 @@ def test_self_retrieval(built_hnsw, base_data, clusters, fast_quantization): """Each database vector must be its own nearest neighbor at high ef.""" idx = built_hnsw if fast_quantization: - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters - idx.build(base_data, centroids, cluster_ids, fast_quantization=True) - assert idx.is_built + idx = _built( + base_data, + centroids, + cluster_ids, + fast_quantization=True, + M=8, + ef_construction=50, + nbits=4, + ) + assert len(idx) == N_VECTORS probes = base_data[:10] ids, _ = idx.search(probes, k=1, ef=200) for i in range(10): @@ -152,8 +157,9 @@ def test_recall_vs_brute_force(built_hnsw, base_data, query_data): def test_wrong_data_dim_raises(clusters): - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters + idx = HnswIndex(DIM, M=8, ef_construction=50, nbits=4, num_clusters=len(centroids)) + idx.set_centroids(centroids) bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) with pytest.raises( ValueError, @@ -163,16 +169,20 @@ def test_wrong_data_dim_raises(clusters): ) + "$", ): - idx.build(bad_data, centroids, cluster_ids) + idx.add(bad_data, cluster_ids=cluster_ids, num_threads=1) def test_out_of_range_cluster_id_raises(base_data, clusters): - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters + idx = HnswIndex(DIM, M=8, ef_construction=50, nbits=4, num_clusters=len(centroids)) + idx.set_centroids(centroids) invalid_ids = cluster_ids.copy() invalid_ids[0] = len(centroids) - with pytest.raises(ValueError, match="cluster_ids contains an out-of-range value"): - idx.build(base_data, centroids, invalid_ids) + with pytest.raises( + ValueError, match=re.escape("cluster_ids must be in [0, num_clusters)") + ): + idx.add(base_data, cluster_ids=invalid_ids, num_threads=1) + assert len(idx) == 0 def test_wrong_query_dim_raises(built_hnsw): @@ -189,7 +199,7 @@ def test_wrong_query_dim_raises(built_hnsw): def test_search_before_build_raises(query_data): - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + idx = HnswIndex(DIM, M=8, ef_construction=50, nbits=4) with pytest.raises( RuntimeError, match="^HnswIndex must be trained or loaded before search$" ): @@ -215,7 +225,7 @@ def test_save_load_roundtrip(built_hnsw, query_data, tmp_path): built_hnsw.save(path) loaded = HnswIndex.load(path) - assert loaded.is_built + assert len(loaded) == len(built_hnsw) assert loaded.dim == built_hnsw.dim assert loaded.nbits == built_hnsw.nbits @@ -240,9 +250,11 @@ def test_load_rejects_count_larger_than_capacity(built_hnsw, tmp_path): @pytest.fixture def partial_hnsw(base_data, clusters): """Built from the first 400 vectors, with room for the remaining 100.""" - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) centroids, cluster_ids = clusters - idx.build(base_data[:400], centroids, cluster_ids[:400]) + idx = HnswIndex(DIM, M=8, ef_construction=50, nbits=4, num_clusters=len(centroids)) + idx.set_centroids(centroids) + idx.resize(N_VECTORS) + idx.add(base_data[:400], cluster_ids=cluster_ids[:400], num_threads=1) return idx @@ -272,8 +284,14 @@ def test_add_into_a_dense_graph_retrieves_every_point(base_data, clusters): # M equals the point count, so the base layer holds 2M >= N - 1 links and never # prunes a reverse link: even a parallel add keeps every point reachable. centroids, cluster_ids = clusters - idx = HnswIndex(DIM, N_VECTORS, M=N_VECTORS, ef_construction=N_VECTORS, nbits=4) - idx.build(base_data[:400], centroids, cluster_ids[:400]) + idx = _built( + base_data[:400], + centroids, + cluster_ids[:400], + M=N_VECTORS, + ef_construction=N_VECTORS, + nbits=4, + ) idx.add(base_data[400:], cluster_ids=cluster_ids[400:]) ids, _ = idx.search(base_data[400:], k=1, ef=N_VECTORS) np.testing.assert_array_equal(ids[:, 0], np.arange(400, N_VECTORS)) @@ -315,21 +333,22 @@ def test_add_rejects_wrong_dimension(partial_hnsw): def test_add_before_build_trains_first(base_data): - idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + idx = HnswIndex(DIM, M=8, ef_construction=50, nbits=4) np.testing.assert_array_equal(idx.add(base_data[:4]), np.arange(4)) assert idx.is_trained assert idx.num_clusters == 4 # fewer rows than clusters: one centroid per row with pytest.raises(RuntimeError, match="trained or loaded"): - HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4).add( + HnswIndex(DIM, M=8, ef_construction=50, nbits=4).add( base_data[:4], cluster_ids=np.zeros(4, dtype=np.uint32) ) def test_resize_grows_capacity_and_allows_more(base_data, clusters): centroids, cluster_ids = clusters - idx = HnswIndex(DIM, 400, M=8, ef_construction=50, nbits=4) - idx.build(base_data[:400], centroids, cluster_ids[:400]) - assert idx.max_elements == 400 + idx = _built( + base_data[:400], centroids, cluster_ids[:400], M=8, ef_construction=50, nbits=4 + ) + assert idx.max_elements == 400 # the first add allocates exactly its rows idx.resize(N_VECTORS) assert idx.max_elements == N_VECTORS @@ -341,13 +360,6 @@ def test_resize_grows_capacity_and_allows_more(base_data, clusters): np.testing.assert_array_equal(ids[:, 0], np.arange(10)) -def test_zero_capacity_raises(): - with pytest.raises( - ValueError, match="^HNSW max_elements must be greater than zero$" - ): - HnswIndex(max_elements=0, dim=64) - - def test_resize_below_element_count_raises(partial_hnsw): with pytest.raises(ValueError, match="element count"): partial_hnsw.resize(10) @@ -361,8 +373,14 @@ def test_resize_across_allocation_threshold_preserves_search(tmp_path, nbits, me data = rng.standard_normal((97, 65)).astype(np.float32) queries = rng.standard_normal((7, 65)).astype(np.float32) centroids = data.mean(axis=0, keepdims=True) - index = HnswIndex(65, len(data), M=8, nbits=nbits, metric=metric) - index.build(data, centroids, np.zeros(len(data), dtype=np.uint32)) + index = _built( + data, + centroids, + np.zeros(len(data), dtype=np.uint32), + M=8, + nbits=nbits, + metric=metric, + ) expected_ids, expected_distances = index.search(queries, k=7, ef=97) # The level-0 allocation moves from below 2 MiB to above it and back. @@ -402,12 +420,19 @@ def test_added_points_survive_save_and_load( def _removable(nbits=4, metric="l2", count=120, dim=64, seed=61): - """A small built index with room for as many points again, plus those points.""" + """A small built index, its points, and as many new points to add.""" rng = np.random.default_rng(seed) data = rng.standard_normal((count * 2, dim)).astype(np.float32) centroids = data[:count].mean(axis=0, keepdims=True) - idx = HnswIndex(dim, count, M=8, ef_construction=50, nbits=nbits, metric=metric) - idx.build(data[:count], centroids, np.zeros(count, dtype=np.uint32)) + idx = _built( + data[:count], + centroids, + np.zeros(count, dtype=np.uint32), + M=8, + ef_construction=50, + nbits=nbits, + metric=metric, + ) return idx, data[:count], data[count:] @@ -503,7 +528,7 @@ def test_remove_rejects_invalid_ids_without_removing_anything(): np.testing.assert_array_equal(before[1], after[1]) assert idx.remove([0, 1]) == 2 - unbuilt = HnswIndex(64, 4, M=8, ef_construction=50, nbits=4) + unbuilt = HnswIndex(64, M=8, ef_construction=50, nbits=4) with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.remove([0]) @@ -573,8 +598,9 @@ def test_legacy_padding_round_trip_and_add(tmp_path): def test_huge_ef_is_clamped_to_the_stored_point_count(): data = np.random.default_rng(5).normal(size=(80, 64)).astype(np.float32) - index = HnswIndex(64, 80, nbits=4) - index.build(data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32)) + index = _built( + data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32), nbits=4 + ) exhaustive_ids, exhaustive_distances = index.search(data[:3], k=5, ef=80) ids, distances = index.search(data[:3], k=5, ef=2**40) np.testing.assert_array_equal(ids, exhaustive_ids) diff --git a/tests/python/test_import.py b/tests/python/test_import.py index 806da025..099129ec 100644 --- a/tests/python/test_import.py +++ b/tests/python/test_import.py @@ -9,21 +9,21 @@ [ ( IvfIndex, - (64, 500, 5), + (64, 5), {"nbits": 4, "metric": "l2"}, {"num_clusters": 5, "nbits": 4, "metric": "l2"}, ), - (IvfIndex, (64, 500, 5), {"nbits": 4, "metric": "ip"}, {"metric": "ip"}), - (HnswIndex, (64, 500), {}, {}), + (IvfIndex, (64, 5), {"nbits": 4, "metric": "ip"}, {"metric": "ip"}), + (HnswIndex, (64,), {}, {}), ( HnswIndex, - (64, 500), + (64,), { "M": 8, "ef_construction": 50, "nbits": 4, "metric": "ip", - "random_seed": 7, + "seed": 7, }, {"metric": "ip", "nbits": 4}, ), @@ -34,7 +34,7 @@ def test_constructor_contract(cls, args, kwargs, expected): idx = cls(*args, **kwargs) assert idx.dim == 64 - assert not idx.is_built + assert len(idx) == 0 for name, value in expected.items(): assert getattr(idx, name) == value @@ -42,9 +42,9 @@ def test_constructor_contract(cls, args, kwargs, expected): @pytest.mark.parametrize( "cls,args", [ - (IvfIndex, (64, 500, 5, 4)), - (HnswIndex, (64, 500)), - (SymqgIndex, (64, 16)), + (IvfIndex, (64, 5)), + (HnswIndex, (64,)), + (SymqgIndex, (64,)), ], ) def test_invalid_metric_raises(cls, args): diff --git a/tests/python/test_input_boundaries.py b/tests/python/test_input_boundaries.py index 1d8caab3..1178e913 100644 --- a/tests/python/test_input_boundaries.py +++ b/tests/python/test_input_boundaries.py @@ -23,17 +23,16 @@ def __fspath__(self): def built_index(request): data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) if request.param == "symqg": - index = SymqgIndex(64, 32) - index.build(data, 40) + index = SymqgIndex(64, max_degree=32, ef_construction=40) + index.add(data, num_threads=1) else: if request.param.startswith("ivf"): initializer = "hnsw" if request.param == "ivf_hnsw" else "flat" - index = IvfIndex(64, 80, 1, 4, initializer=initializer) + index = IvfIndex(64, 1, nbits=4, initializer=initializer) else: - index = HnswIndex(64, 80, nbits=4) - index.build( - data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32) - ) + index = HnswIndex(64, nbits=4, num_clusters=1) + index.set_centroids(data.mean(axis=0, keepdims=True)) + index.add(data, cluster_ids=np.zeros(80, dtype=np.uint32), num_threads=1) return index, data[:2] @@ -215,10 +214,10 @@ def __fspath__(self): elif isinstance(index, HnswIndex): index.resize(index.max_elements) else: - data = ( - np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) - ) - index.build(data, 40) + # add is SymphonyQG's only write. It takes the write guard and then + # refuses a built graph; a guard still held would raise "busy". + with pytest.raises(RuntimeError, match="incremental addition"): + index.add(np.zeros((1, 64), dtype=np.float32)) return str(tmp_path / "callback.index") index.save(UpdatingPath()) @@ -227,37 +226,23 @@ def __fspath__(self): np.testing.assert_array_equal(restored.search(queries, 3, 32)[0], expected[0]) -@pytest.mark.parametrize("index_kind,rows", [("ivf", 79), ("hnsw", 0)]) -def test_build_row_count_diagnostics(index_kind, rows): - if index_kind == "ivf": - index = IvfIndex(64, 80, 1, 4) - expected = "number of data rows must match index max_elements; expected 80, got shape (79, 64)" - else: - # HNSW capacity grows, so only an empty build is rejected. - index = HnswIndex(64, 80) - expected = "number of data rows must be at least 1; got shape (0, 64)" - with pytest.raises(ValueError) as error: - index.build( - np.zeros((rows, 64)), np.zeros((1, 64)), np.zeros(rows, dtype=np.uint32) - ) - assert str(error.value) == expected - - @pytest.mark.parametrize("index_kind,dtype", [("ivf", np.uint32), ("hnsw", np.int64)]) def test_cluster_id_error_identifies_position( index_kind, dtype: type[np.int64] | type[np.uint32] ): if index_kind == "ivf": - index = IvfIndex(64, 80, 1, 4) - prefix = "cluster_ids must be in [0, num_clusters)" + index = IvfIndex(64, 1, nbits=4) else: - index = HnswIndex(64, 80) - prefix = "cluster_ids contains an out-of-range value" + index = HnswIndex(64, nbits=8, num_clusters=1) + index.set_centroids(np.zeros((1, 64))) ids = np.zeros(80, dtype=dtype) ids[-1] = 1 with pytest.raises(ValueError) as error: - index.build(np.zeros((80, 64)), np.zeros((1, 64)), ids) - assert str(error.value) == prefix + "; expected [0, 1), invalid entry at index 79" + index.add(np.zeros((80, 64)), cluster_ids=ids, num_threads=1) + assert str(error.value) == ( + "cluster_ids must be in [0, num_clusters); expected [0, 1), " + "invalid entry at index 79" + ) @pytest.mark.parametrize("built_index", ["ivf_hnsw"], indirect=True) @@ -298,19 +283,9 @@ def test_nonfinite_or_oversized_coordinates_are_rejected(built_index, value, mes index.search(bad_queries, 3, 32) extra = np.random.default_rng(7).normal(size=(2, 64)).astype(np.float32) extra[0, 0] = value - if isinstance(index, IvfIndex): - with pytest.raises(ValueError, match=message): - index.add(extra) - elif isinstance(index, HnswIndex): - # The shared fixture is at capacity, which is checked before coordinates. - data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) - spare = HnswIndex(64, 82, nbits=4) - spare.build( - data, data.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32) - ) + if isinstance(index, (IvfIndex, HnswIndex)): with pytest.raises(ValueError, match=message): - spare.add(extra) - assert spare.num_points == 80 + index.add(extra, num_threads=1) assert index.num_points == before ids, distances = index.search(queries, 3, 32) assert ids.shape == (2, 3) and np.all(np.isfinite(distances)) @@ -323,17 +298,22 @@ def test_build_rejects_bad_coordinates_without_building(index_kind, value): bad_data = data.copy() bad_data[5, 3] = value if index_kind == "symqg": - index = SymqgIndex(64, 32) + index = SymqgIndex(64, max_degree=32, ef_construction=40) with pytest.raises(ValueError): - index.build(bad_data, 40) + index.add(bad_data, num_threads=1) else: centroid = data.mean(axis=0, keepdims=True) ids = np.zeros(80, dtype=np.uint32) - index = IvfIndex(64, 80, 1, 4) if index_kind == "ivf" else HnswIndex(64, 80) - with pytest.raises(ValueError): - index.build(bad_data, centroid, ids) + index = ( + IvfIndex(64, 1, nbits=4) + if index_kind == "ivf" + else HnswIndex(64, nbits=8, num_clusters=1) + ) bad_centroid = centroid.copy() bad_centroid[0, 0] = value with pytest.raises(ValueError): - index.build(data, bad_centroid, ids) - assert not index.is_built + index.set_centroids(bad_centroid) + index.set_centroids(centroid) + with pytest.raises(ValueError): + index.add(bad_data, cluster_ids=ids, num_threads=1) + assert len(index) == 0 diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 59764f96..777b9823 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -15,11 +15,31 @@ # ── fixtures ────────────────────────────────────────────────────────────────── +def _built( + data, + centroids, + cluster_ids, + nbits, + num_threads=1, + fast_quantization=False, + **options, +): + """An index over data with the given centroids and cluster assignments.""" + index = IvfIndex(data.shape[1], len(centroids), nbits=nbits, **options) + index.set_centroids(centroids) + index.add( + data, + cluster_ids=cluster_ids, + num_threads=num_threads, + fast_quantization=fast_quantization, + ) + return index + + @pytest.fixture(scope="module") def built_ivf(base_data, clusters): - idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) centroids, cluster_ids = clusters - idx.build(base_data, centroids, cluster_ids) + idx = _built(base_data, centroids, cluster_ids, nbits=4) return idx @@ -27,7 +47,7 @@ def built_ivf(base_data, clusters): def test_properties(built_ivf): - assert built_ivf.is_built + assert len(built_ivf) == N_VECTORS assert built_ivf.dim == DIM assert built_ivf.num_clusters == N_CLUSTERS assert built_ivf.nbits == 4 @@ -62,8 +82,7 @@ def test_batch_search_matches_individual_queries(bits, metric): queries = rng.standard_normal((13, 65)).astype(np.float32) centroids = data[:4].copy() labels = np.arange(len(data), dtype=np.uint32) % 3 - index = IvfIndex(65, len(data), 4, bits, metric=metric) - index.build(data, centroids, labels) + index = _built(data, centroids, labels, nbits=bits, metric=metric) for hacc in [None, False, True]: for nprobe in [1, 99]: kwargs = dict(k=7, nprobe=nprobe, high_accuracy=hacc) @@ -94,8 +113,7 @@ def test_unfilled_results_use_sentinels(): data[1:] = 100.0 centroids = np.stack((data[0], data[1])) cluster_ids = np.array([0, 1, 1, 1], dtype=np.uint32) - idx = IvfIndex(DIM, 4, 2, nbits=4) - idx.build(data, centroids, cluster_ids) + idx = _built(data, centroids, cluster_ids, nbits=4) ids, dists = idx.search(data[:1], k=3, nprobe=1) assert ids[0, 0] == 0 @@ -111,10 +129,9 @@ def test_self_retrieval(built_ivf, base_data, clusters, fast_quantization): """Probing all clusters: each database vector must be its own nearest neighbor.""" idx = built_ivf if fast_quantization: - idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) centroids, cluster_ids = clusters - idx.build(base_data, centroids, cluster_ids, fast_quantization=True) - assert idx.is_built + idx = _built(base_data, centroids, cluster_ids, nbits=4, fast_quantization=True) + assert len(idx) == N_VECTORS probes = base_data[:10] ids, _ = idx.search(probes, k=1, nprobe=_NPROBE_ALL) for i in range(10): @@ -139,8 +156,7 @@ def test_inner_product_routing_uses_inner_product(): query = np.zeros((1, DIM), dtype=np.float32) query[0, 0] = 1.0 - idx = IvfIndex(DIM, 2, 2, nbits=32, metric="ip") - idx.build(data, centroids, cluster_ids, num_threads=1) + idx = _built(data, centroids, cluster_ids, nbits=32, metric="ip") ids, distances = idx.search(query, k=1, nprobe=1, num_threads=1) assert ids[0, 0] == 1 @@ -158,13 +174,8 @@ def test_full_quantization_centroid_distances_roundtrip(nbits, metric, tmp_path) data = rng.standard_normal((33, 129)).astype(np.float32) data[0] = 0 centroid = np.zeros((1, 129), dtype=np.float32) - idx = IvfIndex(129, len(data), 1, nbits=nbits, metric=metric) - idx.build( - data, - centroid, - np.zeros(len(data), dtype=np.uint32), - fast_quantization=False, - num_threads=1, + idx = _built( + data, centroid, np.zeros(len(data), dtype=np.uint32), nbits=nbits, metric=metric ) ids, distances = idx.search(centroid, k=len(data), nprobe=1, num_threads=1) @@ -190,8 +201,9 @@ def test_full_quantization_centroid_distances_roundtrip(nbits, metric, tmp_path) def test_wrong_data_dim_raises(clusters): - idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) + idx = IvfIndex(DIM, N_CLUSTERS, nbits=4) centroids, cluster_ids = clusters + idx.set_centroids(centroids) bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) with pytest.raises( ValueError, @@ -201,33 +213,34 @@ def test_wrong_data_dim_raises(clusters): ) + "$", ): - idx.build(bad_data, centroids, cluster_ids) + idx.add(bad_data, cluster_ids=cluster_ids, num_threads=1) -def test_wrong_centroid_dim_raises(base_data, clusters): - idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) - _, cluster_ids = clusters +def test_wrong_centroid_dim_raises(): + idx = IvfIndex(DIM, N_CLUSTERS, nbits=4) bad_centroids = np.zeros((N_CLUSTERS, DIM + 1), dtype=np.float32) with pytest.raises( ValueError, match="^" + re.escape( - f"centroid dimension does not match index dim; expected last dimension {DIM}, got shape {bad_centroids.shape}" + "centroids must have shape (num_clusters, dim); " + f"expected ({N_CLUSTERS}, {DIM}), got shape {bad_centroids.shape}" ) + "$", ): - idx.build(base_data, bad_centroids, cluster_ids) + idx.set_centroids(bad_centroids) def test_wrong_cluster_ids_length_raises(base_data, clusters): - idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) + idx = IvfIndex(DIM, N_CLUSTERS, nbits=4) centroids, _ = clusters + idx.set_centroids(centroids) bad_ids = np.zeros(N_VECTORS + 5, dtype=np.uint32) with pytest.raises( ValueError, match=f"^cluster_ids length must match number of rows in data; expected {N_VECTORS} IDs, got {len(bad_ids)}$", ): - idx.build(base_data, centroids, bad_ids) + idx.add(base_data, cluster_ids=bad_ids, num_threads=1) def test_wrong_query_dim_raises(built_ivf): @@ -244,7 +257,7 @@ def test_wrong_query_dim_raises(built_ivf): def test_search_before_build_raises(query_data): - idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) + idx = IvfIndex(DIM, N_CLUSTERS, nbits=4) with pytest.raises( RuntimeError, match="^IvfIndex must be trained or loaded before search$" ): @@ -259,7 +272,7 @@ def test_save_load_roundtrip(built_ivf, query_data, tmp_path): built_ivf.save(path) loaded = IvfIndex.load(path) - assert loaded.is_built + assert len(loaded) == len(built_ivf) assert loaded.dim == built_ivf.dim assert loaded.num_clusters == built_ivf.num_clusters assert loaded.nbits == built_ivf.nbits @@ -281,14 +294,15 @@ def test_raw_reranking(metric, high_accuracy, fast_quantization, tmp_path): data = rng.standard_normal((65, dim)).astype(np.float32) data[0] = 0 # zero residual, padded dimension, and tail batches queries = rng.standard_normal((3, dim)).astype(np.float32) - idx = IvfIndex(dim, len(data), 3, 32, metric) cluster_ids = np.arange(len(data), dtype=np.uint32) % 2 # empty third cluster - idx.build( + idx = _built( data, np.zeros((3, dim), dtype=np.float32), cluster_ids, - fast_quantization=fast_quantization, + nbits=32, + metric=metric, num_threads=2, + fast_quantization=fast_quantization, ) # Use a wider oracle: float32 BLAS reduction error near cancellation can # exceed the tolerance even when the returned float32 distance is accurate. @@ -329,8 +343,13 @@ def test_raw_reranking(metric, high_accuracy, fast_quantization, tmp_path): ) order = np.argsort(cluster_ids, kind="stable") np.testing.assert_array_equal(stored.reshape(original.shape), original[order]) - one_bit = IvfIndex(dim, len(data), 3, 1, metric) - one_bit.build(original, np.zeros((3, dim), dtype=np.float32), cluster_ids) + one_bit = _built( + original, + np.zeros((3, dim), dtype=np.float32), + cluster_ids, + nbits=1, + metric=metric, + ) one_bit_path = tmp_path / "one-bit.index" one_bit.save(str(one_bit_path)) assert len(payload) - one_bit_path.stat().st_size == raw_bytes @@ -397,8 +416,7 @@ def test_legacy_ivf_fixture(tmp_path, filename, dim, bits, padded, nonzero_dista ) def test_raw_index_rejects_invalid_files(tmp_path, damage): data = np.zeros((2, 65), dtype=np.float32) - idx = IvfIndex(65, 2, 1, 32) - idx.build(data, data[:1], np.zeros(2, dtype=np.uint32)) + idx = _built(data, data[:1], np.zeros(2, dtype=np.uint32), nbits=32) path = tmp_path / "invalid.index" idx.save(str(path)) payload = bytearray(path.read_bytes()) @@ -433,7 +451,7 @@ def test_raw_index_rejects_invalid_files(tmp_path, damage): @pytest.mark.parametrize("nbits", [0, 10, 31, 33]) def test_invalid_ivf_bits(nbits): with pytest.raises(ValueError, match="IVF bits must be in"): - IvfIndex(3, 2, 1, nbits) + IvfIndex(3, 1, nbits=nbits) @pytest.mark.parametrize("nbits", [1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) @@ -441,9 +459,11 @@ def test_automatic_high_accuracy(nbits, tmp_path): rng = np.random.default_rng(29) data = rng.standard_normal((65, 65)).astype(np.float32) queries = rng.standard_normal((3, 65)).astype(np.float32) - index = IvfIndex(65, len(data), 1, nbits) - index.build( - data, np.zeros((1, 65), dtype=np.float32), np.zeros(len(data), dtype=np.uint32) + index = _built( + data, + np.zeros((1, 65), dtype=np.float32), + np.zeros(len(data), dtype=np.uint32), + nbits=nbits, ) path = tmp_path / "auto.index" index.save(str(path)) @@ -463,8 +483,9 @@ def test_hnsw_sidecar_save_failures(tmp_path, failure): pytest.skip("requires a device that rejects writes") data = np.zeros((2, 64), dtype=np.float32) data[1, 0] = 1 - idx = IvfIndex(64, 2, 2, nbits=32, initializer="hnsw") - idx.build(data, data, np.arange(2, dtype=np.uint32)) + idx = _built( + data, data, np.arange(2, dtype=np.uint32), nbits=32, initializer="hnsw" + ) expected = idx.search(data, 1, 1) path = tmp_path / "sidecar_failure.index" sidecar = Path(str(path) + ".hnsw") @@ -495,13 +516,17 @@ def test_hnsw_centroid_routing_roundtrip(tmp_path): count, dim = 128, 64 centroids = rng.standard_normal((count, dim)).astype(np.float32) labels = np.arange(count, dtype=np.uint32) - idx = IvfIndex(dim, count, count, nbits=1, initializer="hnsw") - idx.build(centroids, centroids, labels, num_threads=4) - selected = np.array([0, 17, 63, count - 1]) - queries = centroids[selected] + idx = _built( + centroids, centroids, labels, nbits=1, num_threads=4, initializer="hnsw" + ) + # set_centroids builds the routing graph with every CPU, and a parallel build may + # leave a centroid unreachable: check every centroid's self-recall in aggregate. + every, every_distances = idx.search(centroids, k=1, nprobe=4, num_threads=4) + found = every[:, 0] == labels + assert np.count_nonzero(found) >= 0.97 * count + np.testing.assert_allclose(every_distances[found], 0, rtol=0, atol=1e-5) + queries = centroids[[0, 17, 63, count - 1]] ids, distances = idx.search(queries, k=1, nprobe=4, num_threads=4) - np.testing.assert_array_equal(ids[:, 0], selected) - np.testing.assert_allclose(distances, 0, rtol=0, atol=1e-5) path = str(tmp_path / "hnsw_centroid.index") idx.save(path) @@ -519,8 +544,17 @@ def test_hnsw_inner_product_routing_roundtrip(tmp_path): data = centroids[:1].copy() query = np.zeros((4, dim), dtype=np.float32) query[:, 0] = 1.0 - idx = IvfIndex(dim, 1, count, nbits=32, metric="ip", initializer="hnsw") - idx.build(data, centroids, np.array([0], dtype=np.uint32), num_threads=4) + # Centroid 0 is every other centroid's nearest under IP, so the routing graph's + # pruning always keeps it and every query reaches it. + idx = _built( + data, + centroids, + np.array([0], dtype=np.uint32), + nbits=32, + num_threads=4, + metric="ip", + initializer="hnsw", + ) path = str(tmp_path / "hnsw_ip_centroid.index") ids, distances = idx.search(query, k=1, nprobe=1, num_threads=4) @@ -547,8 +581,7 @@ def _dynamic_index(nbits=4, metric="l2", count=120, dim=65, clusters=3, seed=61) gap = data[:, None, :] - centroids[None, :, :] assigned = np.argmin(np.einsum("ncd,ncd->nc", gap, gap), axis=1) assigned = assigned.astype(np.uint32) - idx = IvfIndex(dim, count, clusters, nbits, metric) - idx.build(data[:count], centroids, assigned[:count]) + idx = _built(data[:count], centroids, assigned[:count], nbits=nbits, metric=metric) return idx, data[:count], data[count:], assigned[count:] @@ -593,10 +626,10 @@ def test_add_chunks_match_a_single_add(tmp_path): base = tmp_path / "base.index" idx.save(str(base)) whole = IvfIndex.load(str(base)) - whole.add(new, assigned) + whole.add(new, cluster_ids=assigned) parts = IvfIndex.load(str(base)) for chunk in np.array_split(np.arange(len(new)), [1, 33, 34, 90]): - ids = parts.add(new[chunk], assigned[chunk]) + ids = parts.add(new[chunk], cluster_ids=assigned[chunk]) assert len(ids) == len(chunk) a, b = tmp_path / "a.index", tmp_path / "b.index" whole.save(str(a)) @@ -610,12 +643,12 @@ def test_reserved_adds_match_a_single_add(tmp_path, nbits): base = tmp_path / "base.index" idx.save(str(base)) whole = IvfIndex.load(str(base)) - whole.add(new, assigned) + whole.add(new, cluster_ids=assigned) parts = IvfIndex.load(str(base)) for cluster in range(parts.num_clusters): parts.reserve(cluster, len(old) + len(new)) for chunk in np.array_split(np.arange(len(new)), [1, 33, 34, 90]): - parts.add(new[chunk], assigned[chunk]) + parts.add(new[chunk], cluster_ids=assigned[chunk]) a, b = tmp_path / "a.index", tmp_path / "b.index" whole.save(str(a)) parts.save(str(b)) @@ -634,7 +667,7 @@ def test_reserve_rejects_invalid_arguments(): idx.reserve(0, 0) assert idx.max_elements == len(old) - unbuilt = IvfIndex(65, 4, 1, 4) + unbuilt = IvfIndex(65, 1, nbits=4) with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.reserve(0, 1) @@ -654,17 +687,18 @@ def test_add_rejects_invalid_arguments(): with pytest.raises(ValueError, match="2D"): idx.add(new[0]) with pytest.raises(ValueError, match="length"): - idx.add(new, assigned[:-1]) + idx.add(new, cluster_ids=assigned[:-1]) for bad in (-1, 3, 2**32): # 2**32 would wrap to cluster 0 if cast to uint32 first cluster_ids = assigned.astype(np.int64) cluster_ids[0] = bad with pytest.raises(ValueError, match="cluster_ids"): - idx.add(new, cluster_ids) + idx.add(new, cluster_ids=cluster_ids) assert idx.max_elements == len(old) - unbuilt = IvfIndex(65, 4, 1, 4) + # Without cluster IDs, the first add trains the index; with them, it needs centroids. + unbuilt = IvfIndex(65, 1, nbits=4) with pytest.raises(RuntimeError, match="trained or loaded"): - unbuilt.add(new[:2]) + unbuilt.add(new[:2], cluster_ids=[0, 0]) # IvfRemoveTest owns the full width x metric x scan-path matrix in C++. @@ -742,7 +776,7 @@ def test_remove_rejects_invalid_ids_without_removing_anything(): np.testing.assert_array_equal(before[0], after[0]) np.testing.assert_array_equal(before[1], after[1]) - unbuilt = IvfIndex(65, 4, 1, 4) + unbuilt = IvfIndex(65, 1, nbits=4) with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.remove([0]) @@ -790,31 +824,3 @@ def test_remove_and_add_accept_any_integer_dtype_and_empty_input(): assert idx.remove(np.array([5], dtype=np.int16)) == 1 assert len(idx.add(new[:3], cluster_ids=[0, 1, 2])) == 3 assert len(idx.add(new[:3], cluster_ids=np.array([0, 1, 2], dtype=np.uint16))) == 3 - - -@pytest.mark.parametrize("bits", [1, 4, 32]) -@pytest.mark.parametrize("metric", ["l2", "ip"]) -def test_rebuild_matches_incremental_encoding_with_empty_clusters( - tmp_path, bits, metric -): - rng = np.random.default_rng(731) - data = rng.standard_normal((130, 65), dtype=np.float32) - centers = rng.standard_normal((4, 65), dtype=np.float32) - labels = np.full(130, 3, dtype=np.uint32) - labels[0] = 1 - labels[1:33] = 2 - incremental = IvfIndex(65, 1, 4, nbits=bits, metric=metric) - incremental.build(data[:1], centers, labels[:1], num_threads=1) - incremental.add(data[1:], cluster_ids=labels[1:], num_threads=1) - path = tmp_path / "incremental.index" - incremental.save(str(path)) - expected_bytes = path.read_bytes() - expected_ids, expected_distances = incremental.search(data[:3], k=10, nprobe=4) - rebuilt = IvfIndex.load(str(path)) - for threads in (1, 4): - rebuilt.build(data, centers, labels, num_threads=threads) - rebuilt.save(str(path)) - assert path.read_bytes() == expected_bytes - ids, distances = rebuilt.search(data[:3], k=10, nprobe=4) - np.testing.assert_array_equal(ids, expected_ids) - np.testing.assert_array_equal(distances, expected_distances) diff --git a/tests/python/test_ivf_initializers.py b/tests/python/test_ivf_initializers.py index 93393aba..96d8e1a9 100644 --- a/tests/python/test_ivf_initializers.py +++ b/tests/python/test_ivf_initializers.py @@ -15,8 +15,11 @@ def test_hnsw_parallel_build_routes_and_reloads(tmp_path, metric, threads): count, dim = 1024, 65 data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32) data /= np.linalg.norm(data, axis=1, keepdims=True) - index = IvfIndex(dim, count, count, 4, metric, initializer="hnsw") - index.build(data, data, np.arange(count, dtype=np.uint32), num_threads=threads) + index = IvfIndex(dim, count, nbits=4, metric=metric, initializer="hnsw") + # set_centroids builds the routing graph with every CPU; `threads` applies to + # add and search. + index.set_centroids(data) + index.add(data, cluster_ids=np.arange(count, dtype=np.uint32), num_threads=threads) # A parallel build may leave a centroid unreachable in the routing graph, so # check the aggregate self-recall of every point (one point per cluster). every, _ = index.search(data, 1, 8, num_threads=threads) @@ -40,8 +43,9 @@ def test_hnsw_concurrent_routing_windows_match_serial(tmp_path, metric): rng = np.random.default_rng(42) data = rng.standard_normal((count, dim)).astype(np.float32) queries = rng.standard_normal((4, dim)).astype(np.float32) - index = IvfIndex(dim, count, count, 1, metric, initializer="hnsw") - index.build(data, data, np.arange(count, dtype=np.uint32), num_threads=1) + index = IvfIndex(dim, count, nbits=1, metric=metric, initializer="hnsw") + index.set_centroids(data) + index.add(data, cluster_ids=np.arange(count, dtype=np.uint32), num_threads=1) probes = (1, 16, 400, 500) expected = [index.search(queries, 5, probe, num_threads=1) for probe in probes] path = tmp_path / "concurrent.index" @@ -62,8 +66,9 @@ def test_hnsw_concurrent_routing_windows_match_serial(tmp_path, metric): def test_hnsw_rejects_corrupt_sidecar(tmp_path, metric, damage): count, dim = 64, 64 data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32) - index = IvfIndex(dim, count, count, 4, metric, initializer="hnsw") - index.build(data, data, np.arange(count, dtype=np.uint32), num_threads=1) + index = IvfIndex(dim, count, nbits=4, metric=metric, initializer="hnsw") + index.set_centroids(data) + index.add(data, cluster_ids=np.arange(count, dtype=np.uint32), num_threads=1) expected = index.search(data[:2], 1, 4, num_threads=1) path = tmp_path / "centroids.index" index.save(str(path)) @@ -113,10 +118,13 @@ def test_initializer_roundtrip_and_updates(tmp_path, metric, bits, initializer): dim, count = 65, 33 data = rng.standard_normal((count + 1, dim)).astype(np.float32) data /= np.linalg.norm(data, axis=1, keepdims=True) - index = IvfIndex(dim, count, count, bits, metric, initializer=initializer) + index = IvfIndex(dim, count, nbits=bits, metric=metric, initializer=initializer) resolved = "flat" if initializer == "auto" else initializer assert index.initializer == resolved - index.build(data[:count], data[:count], np.arange(count, dtype=np.uint32), 2) + index.set_centroids(data[:count]) + index.add( + data[:count], cluster_ids=np.arange(count, dtype=np.uint32), num_threads=2 + ) queries = data[[0, 17, 32]] expected = index.search(queries, 3, 7, num_threads=2) np.testing.assert_array_equal(expected[0][:, 0], [0, 17, 32]) @@ -162,9 +170,9 @@ def test_initializer_roundtrip_and_updates(tmp_path, metric, bits, initializer): ], ) def test_auto_initializer_thresholds(clusters, expected): - assert IvfIndex(64, 1, clusters, 1).initializer == expected + assert IvfIndex(64, clusters, nbits=1).initializer == expected for mode in ("flat", "flat_rabitq", "hnsw"): - assert IvfIndex(64, 1, clusters, 1, initializer=mode).initializer == mode + assert IvfIndex(64, clusters, nbits=1, initializer=mode).initializer == mode # Historical routing depends only on the cluster count, and both storage widths' @@ -188,8 +196,9 @@ def historical_ivf(request): np.random.default_rng(43).standard_normal((clusters, 64)).astype(np.float32) ) historical = "flat" if clusters < 20000 else "hnsw" - index = IvfIndex(64, 3, clusters, bits, initializer=historical) - index.build(centroids[:3], centroids, np.arange(3, dtype=np.uint32), 2) + index = IvfIndex(64, clusters, nbits=bits, initializer=historical) + index.set_centroids(centroids) + index.add(centroids[:3], cluster_ids=np.arange(3, dtype=np.uint32), num_threads=2) return index, centroids[:3].copy() @@ -221,15 +230,16 @@ def test_legacy_formats_infer_historical_initializer( def test_initializer_rejects_invalid_argument(): with pytest.raises(ValueError, match="initializer must be"): - IvfIndex(64, 3, 1, 1, initializer="unknown") + IvfIndex(64, 1, nbits=1, initializer="unknown") @pytest.mark.parametrize("initializer", ["flat", "flat_rabitq"]) def test_explicit_flat_routing_above_historical_threshold(tmp_path, initializer): centroids = np.zeros((20000, 64), dtype=np.float32) centroids[17, 0] = 1 - index = IvfIndex(64, 1, len(centroids), 32, initializer=initializer) - index.build(centroids[17:18], centroids, np.array([17], dtype=np.uint32), 2) + index = IvfIndex(64, len(centroids), nbits=32, initializer=initializer) + index.set_centroids(centroids) + index.add(centroids[17:18], cluster_ids=np.array([17], np.uint32), num_threads=2) path = tmp_path / "large.index" index.save(str(path)) assert not Path(str(path) + ".hnsw").exists() @@ -243,8 +253,9 @@ def test_explicit_flat_routing_above_historical_threshold(tmp_path, initializer) @pytest.mark.parametrize("damage", ["auto", "unknown", "truncated_initializer"]) def test_initializer_rejects_invalid_persistence(tmp_path, damage): data = np.zeros((3, 65), dtype=np.float32) - index = IvfIndex(65, 3, 1, 1, initializer="flat_rabitq") - index.build(data, data[:1], np.zeros(3, dtype=np.uint32)) + index = IvfIndex(65, 1, nbits=1, initializer="flat_rabitq") + index.set_centroids(data[:1]) + index.add(data, cluster_ids=np.zeros(3, dtype=np.uint32), num_threads=1) path = tmp_path / "bad.index" index.save(str(path)) payload = bytearray(path.read_bytes()) @@ -262,8 +273,9 @@ def test_initializer_rejects_invalid_persistence(tmp_path, damage): def test_auto_flat_rabitq_persists_actual_routing(tmp_path, clusters): centroids = np.zeros((clusters, 64), dtype=np.float32) centroids[17, 0] = 1 - index = IvfIndex(64, 1, clusters, 32) - index.build(centroids[17:18], centroids, np.array([17], dtype=np.uint32), 2) + index = IvfIndex(64, clusters, nbits=32) + index.set_centroids(centroids) + index.add(centroids[17:18], cluster_ids=np.array([17], np.uint32), num_threads=2) path = tmp_path / "auto.index" index.save(str(path)) assert not Path(str(path) + ".hnsw").exists() diff --git a/tests/python/test_portable_indexes.py b/tests/python/test_portable_indexes.py index 35476f1e..8bed52a1 100644 --- a/tests/python/test_portable_indexes.py +++ b/tests/python/test_portable_indexes.py @@ -11,17 +11,17 @@ def check_index_round_trip(tmp_path, metric, dim, kind, bits): # Unit vectors make self-query the exact winner for both metrics. data /= np.linalg.norm(data, axis=1, keepdims=True) if kind == "ivf": - index = IvfIndex(dim, len(data), 1, nbits=bits, metric=metric) + index = IvfIndex(dim, 1, nbits=bits, metric=metric) search_args = {"nprobe": 1} elif kind == "hnsw": index = HnswIndex( dim, - len(data), + num_clusters=1, M=16, ef_construction=96, nbits=bits, metric=metric, - random_seed=42, + seed=42, ) search_args = {"ef": 96} else: @@ -38,9 +38,8 @@ def check_index_round_trip(tmp_path, metric, dim, kind, bits): if kind == "qg": index.add(data, num_threads=1) else: - index.build( - data, data.mean(axis=0, keepdims=True), np.zeros(len(data), np.uint32) - ) + index.set_centroids(data.mean(axis=0, keepdims=True)) + index.add(data, cluster_ids=np.zeros(len(data), np.uint32), num_threads=1) queries = data[:5].copy() ids, distances = index.search(queries, k=5, **search_args) assert ids.shape == distances.shape == (5, 5) @@ -74,7 +73,7 @@ def check_index_round_trip(tmp_path, metric, dim, kind, bits): loaded = type(index).load(str(path)) assert loaded.dim == dim assert loaded.metric == metric - assert (loaded.quantization_bits if kind == "qg" else loaded.nbits) == bits + assert loaded.nbits == (bits or 32 if kind == "qg" else bits) actual_ids, actual_distances = loaded.search(queries, k=5, **search_args) np.testing.assert_array_equal(actual_ids, ids) np.testing.assert_array_equal(actual_distances, distances) diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 4d2be889..3fbc22cf 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -32,12 +32,12 @@ def built_symqg(base_data): def test_properties(built_symqg): - assert built_symqg.is_built + assert len(built_symqg) == N_VECTORS assert built_symqg.dim == DIM assert built_symqg.max_degree == _MAX_DEGREE assert built_symqg.num_points == N_VECTORS assert built_symqg.metric == "l2" - assert built_symqg.quantization_bits == 0 + assert built_symqg.nbits == 32 # ── search output shape and dtype ───────────────────────────────────────────── @@ -62,15 +62,17 @@ def test_batch_search_matches_individual_queries(tmp_path, bits, metric): queries = rng.standard_normal((19, 130)).astype(np.float32)[:, ::2] assert not queries.flags.c_contiguous queries[4] = queries[0] - index = SymqgIndex(65, 32, metric=metric, quantization_bits=bits) - index.build(data, ef_construction=48, num_threads=1) + index = SymqgIndex( + 65, max_degree=32, metric=metric, nbits=bits or 32, ef_construction=48 + ) + index.add(data, num_threads=1) path = str(tmp_path / "batch.index") index.save(path) restored = SymqgIndex.load(path) - # C++ owns the full ef/thread matrix. These scenarios retain the two aliases, - # serial/automatic/parallel scheduling, changing windows, and strided inputs - # while sharing a single build for each metric/storage pair. + # C++ owns the full ef/thread matrix. These scenarios retain serial, automatic + # and parallel scheduling, changing windows, and strided inputs while sharing a + # single build for each metric/storage pair. for k, ef, threads in [ (7, 64, 1), (1, 1, 4), @@ -84,29 +86,27 @@ def test_batch_search_matches_individual_queries(tmp_path, bits, metric): ] expected_ids = np.concatenate([result[0] for result in individual]) expected_distances = np.concatenate([result[1] for result in individual]) - for method in ("search", "search_batch"): - search = getattr(restored, method) - ids, distances = search(queries, k=k, ef=ef, num_threads=threads) - np.testing.assert_array_equal(ids, expected_ids) - np.testing.assert_array_equal( - distances.view(np.uint32), expected_distances.view(np.uint32) - ) - empty_ids, empty_distances = search( - queries[:0], k=k, ef=ef, num_threads=threads - ) - assert empty_ids.shape == empty_distances.shape == (0, k) - - -def test_search_batch_validates_inputs(built_symqg, query_data): + ids, distances = restored.search(queries, k=k, ef=ef, num_threads=threads) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal( + distances.view(np.uint32), expected_distances.view(np.uint32) + ) + empty_ids, empty_distances = restored.search( + queries[:0], k=k, ef=ef, num_threads=threads + ) + assert empty_ids.shape == empty_distances.shape == (0, k) + + +def test_search_validates_inputs(built_symqg, query_data): with pytest.raises(ValueError, match="queries must be a 2D NumPy array"): - built_symqg.search_batch(query_data[0], k=1, ef=_EF) + built_symqg.search(query_data[0], k=1, ef=_EF) with pytest.raises(ValueError, match="query dimension does not match index dim"): - built_symqg.search_batch(query_data[:, :-1], k=1, ef=_EF) + built_symqg.search(query_data[:, :-1], k=1, ef=_EF) with pytest.raises(ValueError, match=r"k must be between 1 and 2\*\*31"): - built_symqg.search_batch(query_data, k=0, ef=_EF) + built_symqg.search(query_data, k=0, ef=_EF) unbuilt = SymqgIndex(DIM, max_degree=_MAX_DEGREE) with pytest.raises(RuntimeError, match="add them or load an index before search"): - unbuilt.search_batch(query_data, k=1, ef=_EF) + unbuilt.search(query_data, k=1, ef=_EF) # ── search correctness ──────────────────────────────────────────────────────── @@ -133,7 +133,7 @@ def test_recall_vs_brute_force(built_symqg, base_data, query_data): def test_wrong_data_dim_raises(): - idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) + idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE, ef_construction=_EF_BUILD) bad_data = np.zeros((N_VECTORS, DIM + 1), dtype=np.float32) with pytest.raises( ValueError, @@ -143,7 +143,7 @@ def test_wrong_data_dim_raises(): ) + "$", ): - idx.build(bad_data, ef_construction=_EF_BUILD) + idx.add(bad_data, num_threads=1) def test_wrong_query_dim_raises(built_symqg): @@ -192,71 +192,6 @@ def test_search_raises_ef_to_k_and_defaults_it(built_symqg, query_data): np.testing.assert_array_equal(built_symqg.search(query_data, 3)[0], default[0]) -@pytest.mark.parametrize("bits", [0, 4, 8]) -@pytest.mark.parametrize("metric", ["l2", "ip"]) -@pytest.mark.parametrize("replacement_points", [0, _MAX_DEGREE]) -def test_failed_rebuild_preserves_index( - base_data, query_data, tmp_path, bits, metric, replacement_points -): - data = base_data[:40] - queries = query_data[:3] - index = SymqgIndex( - DIM, max_degree=_MAX_DEGREE, metric=metric, quantization_bits=bits - ) - index.build(data, ef_construction=_EF_BUILD) - expected_ids, expected_distances = index.search(queries, k=5, ef=_EF) - before_path = tmp_path / "before.index" - index.save(str(before_path)) - - with pytest.raises( - ValueError, - match="QuantizedGraph degree bound must be smaller than the number of points", - ): - index.build(data[:replacement_points], ef_construction=_EF_BUILD) - - assert index.is_built - assert index.num_points == len(data) - assert index.dim == DIM - assert index.max_degree == _MAX_DEGREE - assert index.metric == metric - assert index.quantization_bits == bits - ids, distances = index.search(queries, k=5, ef=_EF) - np.testing.assert_array_equal(ids, expected_ids) - np.testing.assert_array_equal(distances, expected_distances) - - after_path = tmp_path / "after.index" - index.save(str(after_path)) - assert after_path.read_bytes() == before_path.read_bytes() - restored = SymqgIndex.load(str(after_path)) - assert restored.num_points == len(data) - restored_ids, restored_distances = restored.search(queries, k=5, ef=_EF) - np.testing.assert_array_equal(restored_ids, expected_ids) - np.testing.assert_array_equal(restored_distances, expected_distances) - - -@pytest.mark.parametrize("bits", [0, 4, 8]) -def test_successful_rebuild_replaces_index(base_data, tmp_path, bits): - # Seeded, and build uses one thread, so every self-query below must succeed. - index = SymqgIndex(DIM, max_degree=_MAX_DEGREE, quantization_bits=bits, seed=42) - index.build(base_data[:40], ef_construction=_EF_BUILD) - replacement = base_data[40:88] + 10 - index.build(replacement, ef_construction=_EF_BUILD) - - assert index.is_built - assert index.num_points == len(replacement) - assert index.quantization_bits == bits - ids, distances = index.search(replacement[:3], k=1, ef=_EF) - np.testing.assert_array_equal(ids[:, 0], np.arange(3)) - - path = str(tmp_path / "rebuilt.index") - index.save(path) - restored = SymqgIndex.load(path) - assert restored.num_points == len(replacement) - restored_ids, restored_distances = restored.search(replacement[:3], k=1, ef=_EF) - np.testing.assert_array_equal(restored_ids, ids) - np.testing.assert_array_equal(restored_distances, distances) - - def test_invalid_degree_raises(): with pytest.raises( ValueError, match="^max_degree must be a positive multiple of 32$" @@ -264,11 +199,6 @@ def test_invalid_degree_raises(): SymqgIndex(DIM, max_degree=16) -def test_invalid_quantization_bits_raises(): - with pytest.raises(ValueError, match="^quantization_bits must be 0, 4, or 8$"): - SymqgIndex(DIM, max_degree=_MAX_DEGREE, quantization_bits=6) - - def test_save_propagates_write_failure(built_symqg): if not os.path.exists("/dev/full"): pytest.skip("requires a failing output device") @@ -278,9 +208,9 @@ def test_save_propagates_write_failure(built_symqg): @pytest.mark.parametrize("bits", [4, 8]) def test_qg_quant_build_search_and_roundtrip(bits, base_data, query_data, tmp_path): - idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE, quantization_bits=bits) - idx.build(base_data, ef_construction=_EF_BUILD) - assert idx.quantization_bits == bits + idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE, nbits=bits, ef_construction=_EF_BUILD) + idx.add(base_data, num_threads=1) + assert idx.nbits == bits ids, dists = idx.search(query_data, k=5, ef=_EF) assert np.all(ids < N_VECTORS) @@ -289,7 +219,7 @@ def test_qg_quant_build_search_and_roundtrip(bits, base_data, query_data, tmp_pa path = str(tmp_path / f"symqg_quant{bits}.index") idx.save(path) loaded = SymqgIndex.load(path) - assert loaded.quantization_bits == bits + assert loaded.nbits == bits loaded_ids, loaded_dists = loaded.search(query_data, k=5, ef=_EF) np.testing.assert_array_equal(loaded_ids, ids) np.testing.assert_allclose(loaded_dists, dists, rtol=1e-5) @@ -303,7 +233,7 @@ def test_save_load_roundtrip(built_symqg, query_data, tmp_path): built_symqg.save(path) loaded = SymqgIndex.load(path) - assert loaded.is_built + assert len(loaded) == len(built_symqg) assert loaded.dim == built_symqg.dim assert loaded.max_degree == built_symqg.max_degree assert loaded.num_points == built_symqg.num_points @@ -320,8 +250,8 @@ def test_code_only_build_preserves_v1_format(bits, metric, tmp_path): rng = np.random.default_rng(851) data = rng.normal(size=(65, 65)).astype(np.float32) queries = rng.normal(size=(3, 65)).astype(np.float32) - index = SymqgIndex(65, max_degree=32, metric=metric, quantization_bits=bits) - index.build(data, ef_construction=64) + index = SymqgIndex(65, max_degree=32, metric=metric, nbits=bits, ef_construction=64) + index.add(data, num_threads=1) data.fill(np.nan) path = tmp_path / "code_only.index" index.save(str(path)) @@ -346,8 +276,8 @@ def test_constant_vectors_preserve_returned_distance(metric, tmp_path): """Neighbor corrections must not shift the vertex distances returned to Python.""" data = np.full((65, 64), 0.125, dtype=np.float32) queries = np.full((2, 64), 0.0625, dtype=np.float32) - index = SymqgIndex(64, max_degree=32, metric=metric) - index.build(data, ef_construction=64) + index = SymqgIndex(64, max_degree=32, metric=metric, ef_construction=64) + index.add(data, num_threads=1) ids, distances = index.search(queries, k=5, ef=64) expected = 0.5 if metric == "ip" else 0.25 np.testing.assert_allclose(distances, expected, rtol=0, atol=1e-6) @@ -366,10 +296,17 @@ def test_initialization_and_roundtrip(tmp_path, bits, metric, init): rng = np.random.default_rng(17) data = rng.standard_normal((160, 65)).astype(np.float32) queries = rng.standard_normal((8, 65)).astype(np.float32) - index = SymqgIndex(65, max_degree=32, metric=metric, quantization_bits=bits) kwargs = {} if init is None else {"init": init} - index.build(data, ef_construction=128, num_threads=2, **kwargs) - assert index.is_built + index = SymqgIndex( + 65, + max_degree=32, + metric=metric, + nbits=bits or 32, + ef_construction=128, + **kwargs, + ) + index.add(data, num_threads=2) + assert len(index) == len(data) ids, distances = index.search(queries, k=10, ef=160, num_threads=2) assert np.all(ids < len(data)) assert np.isfinite(distances).all() @@ -395,17 +332,15 @@ def test_initialization_and_roundtrip(tmp_path, bits, metric, init): np.testing.assert_array_equal(distances, restored_distances) -def test_native_pipnn_rejects_invalid_build_arguments(base_data): - index = SymqgIndex(DIM, max_degree=32) +def test_native_pipnn_rejects_invalid_build_arguments(): with pytest.raises(ValueError, match="init must be 'random' or 'pipnn'"): - index.build(base_data, ef_construction=64, init="unknown") + SymqgIndex(DIM, max_degree=32, ef_construction=64, init="unknown") with pytest.raises( ValueError, match="ef_construction must be positive and fit uint32" ): - index.build(base_data, ef_construction=0, init="pipnn") + SymqgIndex(DIM, max_degree=32, ef_construction=0, init="pipnn") def test_default_initialization_is_pipnn(): assert inspect.signature(SymqgIndex.__init__).parameters["init"].default == "pipnn" assert SymqgIndex(DIM).init == "pipnn" - assert inspect.signature(SymqgIndex.build).parameters["init"].default == "pipnn"