diff --git a/.github/scripts/ci_changes.py b/.github/scripts/ci_changes.py index 0bde4365..d4799c10 100644 --- a/.github/scripts/ci_changes.py +++ b/.github/scripts/ci_changes.py @@ -125,7 +125,7 @@ def classify(paths): ) elif path.startswith("python_bindings/"): selected.update(("python_build", "wheels", "benchmarks")) - if path.endswith(".py"): + if path.endswith((".py", ".pyi")): selected.add("python_quality") elif path.endswith(CPP_SUFFIXES): selected.update(("cpp_format", "cpp_tidy")) diff --git a/.github/scripts/test_ci_workflow_contracts.py b/.github/scripts/test_ci_workflow_contracts.py index dbbee0c3..b0837e91 100644 --- a/.github/scripts/test_ci_workflow_contracts.py +++ b/.github/scripts/test_ci_workflow_contracts.py @@ -6,6 +6,8 @@ import subprocess import sys import tempfile +import textwrap +import tomllib import unittest from pathlib import Path @@ -15,6 +17,62 @@ class WorkflowContractsTest(unittest.TestCase): + def test_installed_typing_checks_are_required(self): + config = tomllib.loads((ROOT / "pyproject.toml").read_text()) + test_dependencies = config["project"]["optional-dependencies"]["test"] + wheel_dependencies = config["tool"]["cibuildwheel"]["test-requires"] + checker = [value for value in test_dependencies if value.startswith("mypy==")] + self.assertEqual(len(checker), 1) + self.assertIn(checker[0], wheel_dependencies) + self.assertFalse( + any(value.startswith("mypy") for value in config["project"]["dependencies"]) + ) + workflow = (ROOT / ".github/workflows/python.yml").read_text() + build = workflow.split(" python-build:\n", 1)[1].split( + " python-result:\n", 1 + )[0] + self.assertIn('pip install ".[test]"', build) + self.assertIn("python -m pytest tests/python/test_editor_support.py", build) + self.assertNotIn("continue-on-error", build) + self.assertIn( + "pytest {project}/tests/python", + config["tool"]["cibuildwheel"]["test-command"], + ) + self.assertNotIn( + "--ignore={project}/tests/python/test_editor_support.py", + config["tool"]["cibuildwheel"]["test-command"], + ) + for path in ( + "python_bindings/__init__.pyi", + "tests/python/typing/public_api.py", + ): + with self.subTest(path=path): + self.assertLessEqual( + {"python_build", "python_quality", "wheels"}, classify([path]) + ) + # Exercise the actual aggregator script, so a failed selected matrix + # cannot be reported as a successful Python validation check. + result = workflow.split(" python-result:\n", 1)[1].split( + " python-windows:\n", 1 + )[0] + self.assertIn("needs: [changes, python-build]", result) + script = textwrap.dedent(result.split(" run: |\n", 1)[1]) + for matrix_result in ("success", "failure", "cancelled", "skipped"): + with self.subTest(matrix_result=matrix_result): + checked = subprocess.run( + ["bash", "-e", "-c", script], + env=dict( + os.environ, + CHANGES_RESULT="success", + PYTHON_SELECTED="true", + MATRIX_RESULT=matrix_result, + ), + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(checked.returncode == 0, matrix_result == "success") + def test_tidy_rejects_empty_scope_and_propagates_analysis_failure(self): with tempfile.TemporaryDirectory() as directory: work = Path(directory) diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 7020879a..08aa6f6a 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -68,6 +68,11 @@ jobs: - name: Build and install wheel run: pip install ".[test]" -v + - name: Check installed Python types and editor metadata + env: + OMP_NUM_THREADS: "2" + run: python -m pytest tests/python/test_editor_support.py --tb=short -q + # Repaired-wheel jobs run the full Python suite. - name: Smoke-test source installation run: >- diff --git a/pyproject.toml b/pyproject.toml index 8fa98544..c01aa53f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ Issues = "https://github.com/VectorDB-NTU/RaBitQ-Library/issues" Paper = "https://doi.org/10.1145/3725413" [project.optional-dependencies] -test = ["pytest>=7", "numpy"] +test = ["pytest>=7", "numpy", "mypy==2.4.0"] [tool.scikit-build] # Disable directory auto-discovery so we can map it via CMake's install target @@ -45,7 +45,7 @@ cmake.args = [ build = "cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64 cp314-manylinux_x86_64 cp311-manylinux_aarch64 cp312-manylinux_aarch64 cp313-manylinux_aarch64 cp314-manylinux_aarch64 cp311-win_amd64 cp312-win_amd64 cp313-win_amd64 cp314-win_amd64 cp311-macosx_arm64 cp312-macosx_arm64 cp313-macosx_arm64 cp314-macosx_arm64" build-verbosity = 1 environment = { CMAKE_ARGS = "-DRABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF" } -test-requires = ["pytest"] +test-requires = ["pytest", "mypy==2.4.0"] test-command = "pytest {project}/tests/python --ignore={project}/tests/python/test_compare_with_faiss.py --ignore={project}/tests/python/test_cpp_clustering_examples.py --tb=short -q" [tool.cibuildwheel.linux] diff --git a/python_bindings/CMakeLists.txt b/python_bindings/CMakeLists.txt index 68ce0f27..63175c21 100644 --- a/python_bindings/CMakeLists.txt +++ b/python_bindings/CMakeLists.txt @@ -58,7 +58,7 @@ endif() # 5. Map files to a 'rabitqlib' directory # inside the target wheel, completely bypassing your local layout. install(TARGETS _rabitqlib DESTINATION rabitqlib) -install(FILES __init__.py DESTINATION rabitqlib) +install(FILES __init__.py __init__.pyi py.typed DESTINATION rabitqlib) if(RABITQ_BUILD_BENCHMARKS) # Benchmark the production conversion helper without importing or replacing diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index 911763d3..059a94c1 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -31,10 +31,24 @@ def __init__(self, d, k, **kwargs): self._impl = None def train(self, x): - """Train on ``(n, d)`` data and return the last pre-update objective. - - ``final_obj`` measures the returned assignments against the final centroids. - A failed fit preserves the results of the previous successful fit. + """Train centroids and retain the final assignments. + + Parameters + ---------- + x : array_like of float, shape (n, d) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + + Returns + ------- + objective : float + Last pre-update objective. This can differ from final_obj, which uses + returned assignments and final centroids. + + Notes + ----- + A failed fit preserves the last successful results. Returned arrays are + copies. """ implementation = self._implementation_type(self.d, self.k, self.cp) objective = implementation.train(x) @@ -43,38 +57,137 @@ def train(self, x): @property def centroids(self): + """Final centroids, float32 array of shape (k, d), or None before training.""" return None if self._impl is None else self._impl.centroids @property def assignments(self): + """Final centroid IDs, uint32 array of shape (n,), or None before training.""" return None if self._impl is None else self._impl.assignments @property def distances(self): + """Final assignment distances, float32 array of shape (n,), or None before training.""" return None if self._impl is None else self._impl.distances @property def obj(self): + """Pre-update objectives, float64 array of shape (iterations,), or None before training.""" return None if self._impl is None else self._impl.obj @property def final_obj(self): + """Objective for final assignments and centroids, or None before training.""" return None if self._impl is None else self._impl.final_obj @property def iteration_stats(self): + """Per-iteration statistics, or None before training.""" return None if self._impl is None else self._impl.iteration_stats class QGKMeans(_KMeans): - """SymphonyQG graph assignment; recommended for many centroids.""" + """Cluster vectors using SymphonyQG graph assignment. + + Parameters + ---------- + d : int + Original vector dimension, from 64 through 65536. + k : int + Number of centroids; greater than graph_degree and no more than the + number of training vectors. + niter : int, optional + Maximum number of clustering iterations. Default is 25. + verbose : bool, optional + Print training progress when True. Default is False. + spherical : bool, optional + Normalize centroids after each update and use inner-product assignment. + Normalize input vectors yourself for cosine clustering. Default is + False. + seed : int, optional + Random seed. Default is 42. + min_points_per_centroid : int, optional + Warn below this many training vectors per centroid. Default is 39. + early_stop_threshold : float, optional + Stop when the absolute relative objective change is at most this + threshold. Default is 0.0. + num_threads : int, optional + Worker count; 0 selects the available CPU count. Default is 0. + final_assignment : FinalAssignmentMode, optional + Approximate uses the training assignment method; Exact uses exhaustive + assignment against final centroids. Default is + FinalAssignmentMode.Approximate. + graph_degree : int, optional + Centroid graph degree; a positive multiple of 32. Default is 32. + ef_build : int, optional + Positive centroid-graph construction search breadth. Default is 240. + ef_search : int, optional + Positive centroid-graph assignment search breadth. Default is 16. + graph_build_iterations : int, optional + Positive number of graph-building passes after initialization. Default + is 1. + quantization_bits : int, optional + Centroid storage: 0 for raw float32, or 4 or 8 for RaBitQ. Default is + 0. + + Notes + ----- + Training results are None until the first successful train call. Set cp + fields before training to customize subsequent runs. A failed train + preserves previous results. + """ _parameters_type = QGKMeansParameters _implementation_type = _QGKMeans class RaBitQKMeans(_KMeans): - """Flat RaBitQ assignment; recommended for few centroids.""" + """Cluster vectors using flat RaBitQ assignment. + + Parameters + ---------- + d : int + Original vector dimension, from 64 through 65536. + k : int + Number of centroids, between 1 and the number of training vectors. + niter : int, optional + Maximum number of clustering iterations. Default is 25. + verbose : bool, optional + Print training progress when True. Default is False. + spherical : bool, optional + Normalize centroids after each update and use inner-product assignment. + Normalize input vectors yourself for cosine clustering. Default is + False. + seed : int, optional + Random seed. Default is 42. + min_points_per_centroid : int, optional + Warn below this many training vectors per centroid. Default is 39. + early_stop_threshold : float, optional + Stop when the absolute relative objective change is at most this + threshold. Default is 0.0. + num_threads : int, optional + Worker count; 0 selects the available CPU count. Default is 0. + final_assignment : FinalAssignmentMode, optional + Approximate uses the training assignment method; Exact uses exhaustive + assignment against final centroids. Default is + FinalAssignmentMode.Approximate. + + Notes + ----- + Training results are None until the first successful train call. Set cp + fields before training to customize subsequent runs. A failed train + preserves previous results. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import RaBitQKMeans + >>> x = np.random.default_rng(42).normal(size=(64, 64)).astype(np.float32) + >>> model = RaBitQKMeans(64, 2, niter=2, num_threads=1) + >>> objective = model.train(x) + >>> model.centroids.shape + (2, 64) + """ _parameters_type = RaBitQKMeansParameters _implementation_type = _RaBitQKMeans diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi new file mode 100644 index 00000000..cbe23df8 --- /dev/null +++ b/python_bindings/__init__.pyi @@ -0,0 +1,1267 @@ +"""Public Python interfaces for RaBitQ indexes and clustering. + +Array shapes and numeric ranges are documented runtime constraints. +""" + +from collections.abc import Sequence +from typing import Any, ClassVar, Literal, SupportsIndex, TypeAlias + +import numpy as np +import numpy.typing as npt + +__all__ = [ + "FinalAssignmentMode", + "HnswIndex", + "IvfIndex", + "QGKMeans", + "QGKMeansIterationStats", + "QGKMeansParameters", + "RaBitQKMeans", + "RaBitQKMeansIterationStats", + "RaBitQKMeansParameters", + "SymqgIndex", +] + +_Metric: TypeAlias = Literal["l2", "ip", "innerproduct"] +_Initializer: TypeAlias = Literal["auto", "flat", "flat_rabitq", "hnsw"] +# Rank, ID ranges and Python bool (an int subtype) are validated at runtime. +_IntegerIds: TypeAlias = npt.NDArray[np.integer[Any]] | Sequence[int | np.integer[Any]] + +# pybind11 enums are not subclasses of enum.Enum. +class FinalAssignmentMode: + """Final centroid assignment policy. + + Approximate uses the selected clustering algorithm. SymphonyQG is a + compatibility alias for Approximate. Exact performs exhaustive assignment. + """ + + Approximate: ClassVar[FinalAssignmentMode] + SymphonyQG: ClassVar[FinalAssignmentMode] + Exact: ClassVar[FinalAssignmentMode] + __members__: ClassVar[dict[str, FinalAssignmentMode]] + + def __init__(self, value: SupportsIndex) -> None: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + def __int__(self) -> int: ... + def __index__(self) -> int: ... + +class QGKMeansIterationStats: + """Read-only statistics for one completed clustering iteration.""" + + @property + def iteration(self) -> int: + """Iteration number, starting at 1.""" + @property + def obj(self) -> float: + """Objective measured before the centroid update.""" + @property + def shift(self) -> float: + """Total squared centroid displacement during the update.""" + @property + def nsplit(self) -> int: + """Number of empty centroids repaired by splitting.""" + +RaBitQKMeansIterationStats = QGKMeansIterationStats + +class IvfIndex: + """An approximate nearest-neighbor index using IVF. + + Parameters + ---------- + dim : int + Original vector dimension, from 64 through 65536. + max_elements : int + Number of vectors required by the initial build. Later adds grow + storage. + num_clusters : int + Number of supplied centroids. + nbits : int + Total quantization bits, 1 through 9, or 32 for owned float32 vectors + with one-bit filtering and raw reranking. + metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. + initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional + Centroid routing strategy. Default is 'auto'. + + Notes + ----- + Build or load the index before searching or saving. Concurrent reads are + supported; an overlapping mutation raises RuntimeError. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import IvfIndex + >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + >>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) + >>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) + >>> ids, distances = index.search(x[:2], k=3, nprobe=1) + >>> ids.shape, distances.shape + ((2, 3), (2, 3)) + """ + + def __init__( + self, + dim: SupportsIndex, + max_elements: SupportsIndex, + num_clusters: SupportsIndex, + nbits: SupportsIndex, + metric: _Metric = "l2", + initializer: _Initializer = "auto", + ) -> None: + """An approximate nearest-neighbor index using IVF. + + Parameters + ---------- + dim : int + Original vector dimension, from 64 through 65536. + max_elements : int + Number of vectors required by the initial build. Later adds grow + storage. + num_clusters : int + Number of supplied centroids. + nbits : int + Total quantization bits, 1 through 9, or 32 for owned float32 vectors + with one-bit filtering and raw reranking. + metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. + initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional + Centroid routing strategy. Default is 'auto'. + + Notes + ----- + Build or load the index before searching or saving. Concurrent reads are + supported; an overlapping mutation raises RuntimeError. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import IvfIndex + >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + >>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) + >>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) + >>> ids, distances = index.search(x[:2], k=3, nprobe=1) + >>> ids.shape, distances.shape + ((2, 3), (2, 3)) + """ + + @property + def dim(self) -> int: + """Original vector dimension.""" + + @property + def max_elements(self) -> int: + """Current storage capacity.""" + + @property + def num_points(self) -> int: + """Number of live, searchable points.""" + + @property + def num_stored_points(self) -> int: + """Number of stored points, including removed points.""" + + @property + def is_built(self) -> bool: + """Whether the index has been built or loaded.""" + + @property + def metric(self) -> Literal["l2", "ip"]: + """Canonical metric name: 'l2' or 'ip'.""" + + @property + def nbits(self) -> int: + """Total storage bits per coordinate.""" + + @property + def num_clusters(self) -> int: + """Number of centroids used for quantization.""" + + def build( + self, + data: npt.ArrayLike, + centroids: npt.ArrayLike, + cluster_ids: _IntegerIds, + num_threads: SupportsIndex = 1, + fast_quantization: bool | np.bool_ = False, + ) -> None: + """Build the index from vectors and supplied centroid assignments. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + centroids : array_like of float, shape (num_clusters, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + cluster_ids : array_like of int, shape (n,) + Centroid ID per vector, in [0, num_clusters). Floating-point and + boolean IDs are rejected. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + + Notes + ----- + The initial data must satisfy n == max_elements. Centroids and data use + original coordinates. The input arrays need not be retained after this + synchronous call. + """ + + def add( + self, + data: npt.ArrayLike, + cluster_ids: _IntegerIds | None = None, + num_threads: SupportsIndex = 1, + fast_quantization: bool | np.bool_ = False, + ) -> npt.NDArray[np.uint32]: + """Append vectors using the existing centroids. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + cluster_ids : array_like of int, shape (n,), optional + Existing centroid IDs. None assigns each vector to its nearest + centroid. Centroids do not move. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + + Returns + ------- + ids : ndarray of uint32, shape (n,) + Assigned point IDs. + + Notes + ----- + Batch additions to reduce repeated cluster copying; reserve can preallocate + room. + """ + + def remove(self, ids: _IntegerIds) -> int: + """Exclude point IDs from subsequent search results. + + Parameters + ---------- + ids : array_like of int, shape (n,) + IDs in [0, num_stored_points). Floating-point and boolean IDs are + rejected. + + Returns + ------- + count : int + Number of newly removed points; repeats do not count twice. + + Notes + ----- + Removal does not reclaim storage and cannot be undone. + """ + + def search( + self, + queries: npt.ArrayLike, + k: SupportsIndex, + nprobe: SupportsIndex, + high_accuracy: bool | np.bool_ | None = None, + num_threads: SupportsIndex = 1, + ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + """Search a batch of query vectors. + + Parameters + ---------- + queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + k : int + Number of neighbors requested, between 1 and max_elements. + nprobe : int + Positive number of clusters to probe; capped at num_clusters. + high_accuracy : bool, optional + Enable HACC. None (default) selects it for 4-9 bits and standard + FastScan otherwise. Raw storage uses standard FastScan. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + + Returns + ------- + ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. + distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + + Notes + ----- + Returns (ids, distances), in that order. Missing neighbors have ID + 4294967295 and distance inf. Array shapes and value ranges are checked at + runtime, not by static type checkers. + """ + + def save(self, path: str) -> None: + """Save a built index to disk. + + Parameters + ---------- + path : str + Index file path. Pass str(path) for pathlib.Path objects. + + Notes + ----- + IVF with HNSW centroid routing also writes its required sidecar file. + """ + + @staticmethod + def load(path: str) -> IvfIndex: + """Load an index and its saved rotation from disk. + + Parameters + ---------- + path : str + Index file path. Pass str(path) for pathlib.Path objects. + + Returns + ------- + index : IvfIndex + A built index ready for search. + """ + + @property + def initializer(self) -> _Initializer: + """Resolved centroid routing strategy.""" + + def reserve(self, cluster_id: SupportsIndex, num_points: SupportsIndex) -> None: + """Preallocate storage for a cluster. + + Parameters + ---------- + cluster_id : int + Cluster ID in [0, num_clusters). + num_points : int + Nonnegative capacity to reserve for this cluster. + + Notes + ----- + Requires a built index. Reserved room is not saved. + """ + +class HnswIndex: + """An approximate nearest-neighbor index using HNSW. + + Parameters + ---------- + dim : int + Original vector dimension, from 64 through 65536. + max_elements : int + Storage capacity, including removed points. Use resize to grow it. + M : int, optional + Graph connectivity, at least 2. Default is 16. + ef_construction : int, optional + Construction search breadth. Default is 200. + nbits : int, optional + Total quantization bits, 1 through 9. Default is 8. + metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. + random_seed : int, optional + Graph random seed. Default is 100. + + Notes + ----- + Build or load the index before searching or saving. Concurrent reads are + supported; an overlapping mutation raises RuntimeError. + """ + + def __init__( + self, + dim: SupportsIndex, + max_elements: SupportsIndex, + M: SupportsIndex = 16, + ef_construction: SupportsIndex = 200, + nbits: SupportsIndex = 8, + metric: _Metric = "l2", + random_seed: SupportsIndex = 100, + ) -> None: + """An approximate nearest-neighbor index using HNSW. + + Parameters + ---------- + dim : int + Original vector dimension, from 64 through 65536. + max_elements : int + Storage capacity, including removed points. Use resize to grow it. + M : int, optional + Graph connectivity, at least 2. Default is 16. + ef_construction : int, optional + Construction search breadth. Default is 200. + nbits : int, optional + Total quantization bits, 1 through 9. Default is 8. + metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. + random_seed : int, optional + Graph random seed. Default is 100. + + Notes + ----- + Build or load the index before searching or saving. Concurrent reads are + supported; an overlapping mutation raises RuntimeError. + """ + + @property + def dim(self) -> int: + """Original vector dimension.""" + + @property + def max_elements(self) -> int: + """Current storage capacity.""" + + @property + def num_points(self) -> int: + """Number of live, searchable points.""" + + @property + def num_stored_points(self) -> int: + """Number of stored points, including removed points.""" + + @property + def is_built(self) -> bool: + """Whether the index has been built or loaded.""" + + @property + def metric(self) -> Literal["l2", "ip"]: + """Canonical metric name: 'l2' or 'ip'.""" + + @property + def nbits(self) -> int: + """Total storage bits per coordinate.""" + + @property + def num_clusters(self) -> int: + """Number of centroids used for quantization.""" + + def build( + self, + data: npt.ArrayLike, + centroids: npt.ArrayLike, + cluster_ids: _IntegerIds, + num_threads: SupportsIndex = 1, + fast_quantization: bool | np.bool_ = False, + ) -> None: + """Build the index from vectors and supplied centroid assignments. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + centroids : array_like of float, shape (num_clusters, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + cluster_ids : array_like of int, shape (n,) + Centroid ID per vector, in [0, num_clusters). Floating-point and + boolean IDs are rejected. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + + Notes + ----- + The initial data must satisfy 1 <= n <= max_elements. Centroids and data + use original coordinates. The input arrays need not be retained after this + synchronous call. + """ + + def add( + self, + data: npt.ArrayLike, + cluster_ids: _IntegerIds | None = None, + num_threads: SupportsIndex = 1, + fast_quantization: bool | np.bool_ = False, + ) -> npt.NDArray[np.uint32]: + """Append vectors using the existing centroids. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + cluster_ids : array_like of int, shape (n,), optional + Existing centroid IDs. None assigns each vector to its nearest + centroid. Centroids do not move. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + + Returns + ------- + ids : ndarray of uint32, shape (n,) + Assigned point IDs. + + Notes + ----- + The stored point count plus n must fit max_elements; resize first if + needed. Removed IDs are not reused. + """ + + def remove(self, ids: _IntegerIds) -> int: + """Exclude point IDs from subsequent search results. + + Parameters + ---------- + ids : array_like of int, shape (n,) + IDs in [0, num_stored_points). Floating-point and boolean IDs are + rejected. + + Returns + ------- + count : int + Number of newly removed points; repeats do not count twice. + + Notes + ----- + Removal does not reclaim storage and cannot be undone. + """ + + def search( + self, + queries: npt.ArrayLike, + k: SupportsIndex, + ef: SupportsIndex = 0, + num_threads: SupportsIndex = 1, + ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + """Search a batch of query vectors. + + Parameters + ---------- + queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + k : int + Number of neighbors requested, between 1 and max_elements. + ef : int, optional + Search breadth. 0 (default) selects max(k, 10). + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + + Returns + ------- + ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. + distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + + Notes + ----- + Returns (ids, distances), in that order. Missing neighbors have ID + 4294967295 and distance inf. Array shapes and value ranges are checked at + runtime, not by static type checkers. + """ + + def save(self, path: str) -> None: + """Save a built index to disk. + + Parameters + ---------- + path : str + Index file path. Pass str(path) for pathlib.Path objects. + + Notes + ----- + The saved state includes the rotation used to build the index. + """ + + @staticmethod + def load(path: str) -> HnswIndex: + """Load an index and its saved rotation from disk. + + Parameters + ---------- + path : str + Index file path. Pass str(path) for pathlib.Path objects. + + Returns + ------- + index : HnswIndex + A built index ready for search. + """ + + def resize(self, max_elements: SupportsIndex) -> None: + """Resize the storage capacity of a built index. + + Parameters + ---------- + max_elements : int + New capacity; must be at least num_stored_points. + """ + +class SymqgIndex: + """An approximate nearest-neighbor index using SymphonyQG. + + Parameters + ---------- + dim : int + Original vector dimension, from 64 through 65536. + max_degree : int + Graph degree bound; a positive multiple of 32. + metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. + quantization_bits : {0, 4, 8}, optional + Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 + (default). + + Notes + ----- + Build or load the index before searching or saving. Concurrent reads are + supported; an overlapping mutation raises RuntimeError. + """ + + def __init__( + self, + dim: SupportsIndex, + max_degree: SupportsIndex, + metric: _Metric = "l2", + quantization_bits: SupportsIndex = 0, + ) -> None: + """An approximate nearest-neighbor index using SymphonyQG. + + Parameters + ---------- + dim : int + Original vector dimension, from 64 through 65536. + max_degree : int + Graph degree bound; a positive multiple of 32. + metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. + quantization_bits : {0, 4, 8}, optional + Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 + (default). + + Notes + ----- + Build or load the index before searching or saving. Concurrent reads are + supported; an overlapping mutation raises RuntimeError. + """ + + @property + def dim(self) -> int: + """Original vector dimension.""" + + @property + def max_elements(self) -> int: + """Current storage capacity.""" + + @property + def num_points(self) -> int: + """Number of live, searchable points.""" + + @property + def num_stored_points(self) -> int: + """Number of stored points, including removed points.""" + + @property + def is_built(self) -> bool: + """Whether the index has been built or loaded.""" + + @property + def metric(self) -> Literal["l2", "ip"]: + """Canonical metric name: 'l2' or 'ip'.""" + + @property + def max_degree(self) -> int: + """Graph degree bound.""" + + @property + def quantization_bits(self) -> int: + """Storage bits: 0 for raw float32, or 4 or 8 for RaBitQ.""" + + def build( + self, + data: npt.ArrayLike, + ef_construction: SupportsIndex, + num_threads: SupportsIndex = 1, + init: Literal["pipnn", "random"] = "pipnn", + ) -> None: + """Build the graph from input vectors. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + ef_construction : int + Positive construction search breadth fitting uint32. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + init : {'pipnn', 'random'}, optional + Graph initialization method. Default is 'pipnn'. + + Notes + ----- + Requires n > max_degree. Input arrays need not be retained after this + synchronous call. + """ + + def search( + self, + queries: npt.ArrayLike, + k: SupportsIndex, + ef: SupportsIndex, + num_threads: SupportsIndex = 1, + ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + """Search a batch of query vectors. + + Parameters + ---------- + queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + k : int + Number of neighbors requested, between 1 and num_points. + ef : int + Search breadth. Must be positive and at least k. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + + Returns + ------- + ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. + distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + + Notes + ----- + Returns (ids, distances), in that order. Array shapes and value ranges are + checked at runtime, not by static type checkers. + """ + + def search_batch( + self, + queries: npt.ArrayLike, + k: SupportsIndex, + ef: SupportsIndex, + num_threads: SupportsIndex = 1, + ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: + """Search a batch of query vectors. + + Parameters + ---------- + queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + k : int + Number of neighbors requested, between 1 and num_points. + ef : int + Search breadth. Must be positive and at least k. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + + Returns + ------- + ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. + distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + + Notes + ----- + Returns (ids, distances), in that order. Array shapes and value ranges are + checked at runtime, not by static type checkers. + """ + + def save(self, path: str) -> None: + """Save a built index to disk. + + Parameters + ---------- + path : str + Index file path. Pass str(path) for pathlib.Path objects. + + Notes + ----- + The saved state includes the rotation used to build the index. + """ + + @staticmethod + def load(path: str) -> SymqgIndex: + """Load an index and its saved rotation from disk. + + Parameters + ---------- + path : str + Index file path. Pass str(path) for pathlib.Path objects. + + Returns + ------- + index : SymqgIndex + A built index ready for search. + """ + +class QGKMeans: + """Cluster vectors using SymphonyQG graph assignment. + + Parameters + ---------- + d : int + Original vector dimension, from 64 through 65536. + k : int + Number of centroids; greater than graph_degree and no more than the + number of training vectors. + niter : int, optional + Maximum number of clustering iterations. Default is 25. + verbose : bool, optional + Print training progress when True. Default is False. + spherical : bool, optional + Normalize centroids after each update and use inner-product assignment. + Normalize input vectors yourself for cosine clustering. Default is + False. + seed : int, optional + Random seed. Default is 42. + min_points_per_centroid : int, optional + Warn below this many training vectors per centroid. Default is 39. + early_stop_threshold : float, optional + Stop when the absolute relative objective change is at most this + threshold. Default is 0.0. + num_threads : int, optional + Worker count; 0 selects the available CPU count. Default is 0. + final_assignment : FinalAssignmentMode, optional + Approximate uses the training assignment method; Exact uses exhaustive + assignment against final centroids. Default is + FinalAssignmentMode.Approximate. + graph_degree : int, optional + Centroid graph degree; a positive multiple of 32. Default is 32. + ef_build : int, optional + Positive centroid-graph construction search breadth. Default is 240. + ef_search : int, optional + Positive centroid-graph assignment search breadth. Default is 16. + graph_build_iterations : int, optional + Positive number of graph-building passes after initialization. Default + is 1. + quantization_bits : int, optional + Centroid storage: 0 for raw float32, or 4 or 8 for RaBitQ. Default is + 0. + + Notes + ----- + Training results are None until the first successful train call. Set cp + fields before training to customize subsequent runs. A failed train + preserves previous results. + """ + + d: int + k: int + cp: QGKMeansParameters + + def __init__( + self, + d: SupportsIndex, + k: SupportsIndex, + *, + niter: SupportsIndex = 25, + verbose: bool | np.bool_ = False, + spherical: bool | np.bool_ = False, + seed: SupportsIndex = 42, + min_points_per_centroid: SupportsIndex = 39, + early_stop_threshold: float = 0.0, + num_threads: SupportsIndex = 0, + final_assignment: FinalAssignmentMode = FinalAssignmentMode.Approximate, + graph_degree: SupportsIndex = 32, + ef_build: SupportsIndex = 240, + ef_search: SupportsIndex = 16, + graph_build_iterations: SupportsIndex = 1, + quantization_bits: SupportsIndex = 0, + ) -> None: + """Cluster vectors using SymphonyQG graph assignment. + + Parameters + ---------- + d : int + Original vector dimension, from 64 through 65536. + k : int + Number of centroids; greater than graph_degree and no more than the + number of training vectors. + niter : int, optional + Maximum number of clustering iterations. Default is 25. + verbose : bool, optional + Print training progress when True. Default is False. + spherical : bool, optional + Normalize centroids after each update and use inner-product assignment. + Normalize input vectors yourself for cosine clustering. Default is + False. + seed : int, optional + Random seed. Default is 42. + min_points_per_centroid : int, optional + Warn below this many training vectors per centroid. Default is 39. + early_stop_threshold : float, optional + Stop when the absolute relative objective change is at most this + threshold. Default is 0.0. + num_threads : int, optional + Worker count; 0 selects the available CPU count. Default is 0. + final_assignment : FinalAssignmentMode, optional + Approximate uses the training assignment method; Exact uses exhaustive + assignment against final centroids. Default is + FinalAssignmentMode.Approximate. + graph_degree : int, optional + Centroid graph degree; a positive multiple of 32. Default is 32. + ef_build : int, optional + Positive centroid-graph construction search breadth. Default is 240. + ef_search : int, optional + Positive centroid-graph assignment search breadth. Default is 16. + graph_build_iterations : int, optional + Positive number of graph-building passes after initialization. Default + is 1. + quantization_bits : int, optional + Centroid storage: 0 for raw float32, or 4 or 8 for RaBitQ. Default is + 0. + + Notes + ----- + Training results are None until the first successful train call. Set cp + fields before training to customize subsequent runs. A failed train + preserves previous results. + """ + + def train(self, x: npt.ArrayLike) -> float: + """Train centroids and retain the final assignments. + + Parameters + ---------- + x : array_like of float, shape (n, d) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + + Returns + ------- + objective : float + Last pre-update objective. This can differ from final_obj, which uses + returned assignments and final centroids. + + Notes + ----- + A failed fit preserves the last successful results. Returned arrays are + copies. + """ + + @property + def centroids(self) -> npt.NDArray[np.float32] | None: + """Final centroids, float32 array of shape (k, d), or None before training.""" + + @property + def assignments(self) -> npt.NDArray[np.uint32] | None: + """Final centroid IDs, uint32 array of shape (n,), or None before training.""" + + @property + def distances(self) -> npt.NDArray[np.float32] | None: + """Final assignment distances, float32 array of shape (n,), or None before training.""" + + @property + def obj(self) -> npt.NDArray[np.float64] | None: + """Pre-update objectives, float64 array of shape (iterations,), or None before training.""" + + @property + def final_obj(self) -> float | None: + """Objective for final assignments and centroids, or None before training.""" + + @property + def iteration_stats(self) -> list[QGKMeansIterationStats] | None: + """Per-iteration statistics, or None before training.""" + +class QGKMeansParameters: + """Mutable training settings for QGKMeans. + + Notes + ----- + Create with no arguments; set attributes directly or use keyword arguments + on QGKMeans. + """ + + @property + def niter(self) -> int: + """Maximum number of clustering iterations. Default is 25.""" + @niter.setter + def niter(self, value: SupportsIndex) -> None: ... + @property + def verbose(self) -> bool: + """Print training progress when True. Default is False.""" + @verbose.setter + def verbose(self, value: bool | np.bool_) -> None: ... + @property + def spherical(self) -> bool: + """Normalize centroids after each update and use inner-product assignment. Normalize input vectors yourself for cosine clustering. Default is False.""" + @spherical.setter + def spherical(self, value: bool | np.bool_) -> None: ... + @property + def seed(self) -> int: + """Random seed. Default is 42.""" + @seed.setter + def seed(self, value: SupportsIndex) -> None: ... + @property + def min_points_per_centroid(self) -> int: + """Warn below this many training vectors per centroid. Default is 39.""" + @min_points_per_centroid.setter + def min_points_per_centroid(self, value: SupportsIndex) -> None: ... + + early_stop_threshold: float + """Stop when the absolute relative objective change is at most this threshold. Default is 0.0.""" + + @property + def num_threads(self) -> int: + """Worker count; 0 selects the available CPU count. Default is 0.""" + @num_threads.setter + def num_threads(self, value: SupportsIndex) -> None: ... + + final_assignment: FinalAssignmentMode + """Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is FinalAssignmentMode.Approximate.""" + + @property + def graph_degree(self) -> int: + """Centroid graph degree; a positive multiple of 32. Default is 32.""" + @graph_degree.setter + def graph_degree(self, value: SupportsIndex) -> None: ... + @property + def ef_build(self) -> int: + """Positive centroid-graph construction search breadth. Default is 240.""" + @ef_build.setter + def ef_build(self, value: SupportsIndex) -> None: ... + @property + def ef_search(self) -> int: + """Positive centroid-graph assignment search breadth. Default is 16.""" + @ef_search.setter + def ef_search(self, value: SupportsIndex) -> None: ... + @property + def graph_build_iterations(self) -> int: + """Positive number of graph-building passes after initialization. Default is 1.""" + @graph_build_iterations.setter + def graph_build_iterations(self, value: SupportsIndex) -> None: ... + @property + def quantization_bits(self) -> int: + """Centroid storage: 0 for raw float32, or 4 or 8 for RaBitQ. Default is 0.""" + @quantization_bits.setter + def quantization_bits(self, value: SupportsIndex) -> None: ... + def __init__(self) -> None: + """Create parameters with their documented defaults.""" + +class RaBitQKMeans: + """Cluster vectors using flat RaBitQ assignment. + + Parameters + ---------- + d : int + Original vector dimension, from 64 through 65536. + k : int + Number of centroids, between 1 and the number of training vectors. + niter : int, optional + Maximum number of clustering iterations. Default is 25. + verbose : bool, optional + Print training progress when True. Default is False. + spherical : bool, optional + Normalize centroids after each update and use inner-product assignment. + Normalize input vectors yourself for cosine clustering. Default is + False. + seed : int, optional + Random seed. Default is 42. + min_points_per_centroid : int, optional + Warn below this many training vectors per centroid. Default is 39. + early_stop_threshold : float, optional + Stop when the absolute relative objective change is at most this + threshold. Default is 0.0. + num_threads : int, optional + Worker count; 0 selects the available CPU count. Default is 0. + final_assignment : FinalAssignmentMode, optional + Approximate uses the training assignment method; Exact uses exhaustive + assignment against final centroids. Default is + FinalAssignmentMode.Approximate. + + Notes + ----- + Training results are None until the first successful train call. Set cp + fields before training to customize subsequent runs. A failed train + preserves previous results. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import RaBitQKMeans + >>> x = np.random.default_rng(42).normal(size=(64, 64)).astype(np.float32) + >>> model = RaBitQKMeans(64, 2, niter=2, num_threads=1) + >>> objective = model.train(x) + >>> model.centroids.shape + (2, 64) + """ + + d: int + k: int + cp: RaBitQKMeansParameters + + def __init__( + self, + d: SupportsIndex, + k: SupportsIndex, + *, + niter: SupportsIndex = 25, + verbose: bool | np.bool_ = False, + spherical: bool | np.bool_ = False, + seed: SupportsIndex = 42, + min_points_per_centroid: SupportsIndex = 39, + early_stop_threshold: float = 0.0, + num_threads: SupportsIndex = 0, + final_assignment: FinalAssignmentMode = FinalAssignmentMode.Approximate, + ) -> None: + """Cluster vectors using flat RaBitQ assignment. + + Parameters + ---------- + d : int + Original vector dimension, from 64 through 65536. + k : int + Number of centroids, between 1 and the number of training vectors. + niter : int, optional + Maximum number of clustering iterations. Default is 25. + verbose : bool, optional + Print training progress when True. Default is False. + spherical : bool, optional + Normalize centroids after each update and use inner-product assignment. + Normalize input vectors yourself for cosine clustering. Default is + False. + seed : int, optional + Random seed. Default is 42. + min_points_per_centroid : int, optional + Warn below this many training vectors per centroid. Default is 39. + early_stop_threshold : float, optional + Stop when the absolute relative objective change is at most this + threshold. Default is 0.0. + num_threads : int, optional + Worker count; 0 selects the available CPU count. Default is 0. + final_assignment : FinalAssignmentMode, optional + Approximate uses the training assignment method; Exact uses exhaustive + assignment against final centroids. Default is + FinalAssignmentMode.Approximate. + + Notes + ----- + Training results are None until the first successful train call. Set cp + fields before training to customize subsequent runs. A failed train + preserves previous results. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import RaBitQKMeans + >>> x = np.random.default_rng(42).normal(size=(64, 64)).astype(np.float32) + >>> model = RaBitQKMeans(64, 2, niter=2, num_threads=1) + >>> objective = model.train(x) + >>> model.centroids.shape + (2, 64) + """ + + def train(self, x: npt.ArrayLike) -> float: + """Train centroids and retain the final assignments. + + Parameters + ---------- + x : array_like of float, shape (n, d) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. + + Returns + ------- + objective : float + Last pre-update objective. This can differ from final_obj, which uses + returned assignments and final centroids. + + Notes + ----- + A failed fit preserves the last successful results. Returned arrays are + copies. + """ + + @property + def centroids(self) -> npt.NDArray[np.float32] | None: + """Final centroids, float32 array of shape (k, d), or None before training.""" + + @property + def assignments(self) -> npt.NDArray[np.uint32] | None: + """Final centroid IDs, uint32 array of shape (n,), or None before training.""" + + @property + def distances(self) -> npt.NDArray[np.float32] | None: + """Final assignment distances, float32 array of shape (n,), or None before training.""" + + @property + def obj(self) -> npt.NDArray[np.float64] | None: + """Pre-update objectives, float64 array of shape (iterations,), or None before training.""" + + @property + def final_obj(self) -> float | None: + """Objective for final assignments and centroids, or None before training.""" + + @property + def iteration_stats(self) -> list[QGKMeansIterationStats] | None: + """Per-iteration statistics, or None before training.""" + +class RaBitQKMeansParameters: + """Mutable training settings for RaBitQKMeans. + + Notes + ----- + Create with no arguments; set attributes directly or use keyword arguments + on RaBitQKMeans. + """ + + @property + def niter(self) -> int: + """Maximum number of clustering iterations. Default is 25.""" + @niter.setter + def niter(self, value: SupportsIndex) -> None: ... + @property + def verbose(self) -> bool: + """Print training progress when True. Default is False.""" + @verbose.setter + def verbose(self, value: bool | np.bool_) -> None: ... + @property + def spherical(self) -> bool: + """Normalize centroids after each update and use inner-product assignment. Normalize input vectors yourself for cosine clustering. Default is False.""" + @spherical.setter + def spherical(self, value: bool | np.bool_) -> None: ... + @property + def seed(self) -> int: + """Random seed. Default is 42.""" + @seed.setter + def seed(self, value: SupportsIndex) -> None: ... + @property + def min_points_per_centroid(self) -> int: + """Warn below this many training vectors per centroid. Default is 39.""" + @min_points_per_centroid.setter + def min_points_per_centroid(self, value: SupportsIndex) -> None: ... + + early_stop_threshold: float + """Stop when the absolute relative objective change is at most this threshold. Default is 0.0.""" + + @property + def num_threads(self) -> int: + """Worker count; 0 selects the available CPU count. Default is 0.""" + @num_threads.setter + def num_threads(self, value: SupportsIndex) -> None: ... + + final_assignment: FinalAssignmentMode + """Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is FinalAssignmentMode.Approximate.""" + + def __init__(self) -> None: + """Create parameters with their documented defaults.""" diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index fb404068..b2134586 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -290,7 +290,32 @@ class HnswIndex { void register_hnsw(py::module_& m) { using namespace rabitqlib::python_bindings; - py::class_(m, "HnswIndex") + py::class_( + m, "HnswIndex", R"doc(An approximate nearest-neighbor index using HNSW. + +Parameters +---------- +dim : int + Original vector dimension, from 64 through 65536. +max_elements : int + Storage capacity, including removed points. Use resize to grow it. +M : int, optional + Graph connectivity, at least 2. Default is 16. +ef_construction : int, optional + Construction search breadth. Default is 200. +nbits : int, optional + Total quantization bits, 1 through 9. Default is 8. +metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. +random_seed : int, optional + Graph random seed. Default is 100. + +Notes +----- +Build or load the index before searching or saving. Concurrent reads are +supported; an overlapping mutation raises RuntimeError.)doc" + ) .def( py::init(), py::arg("dim"), @@ -299,7 +324,31 @@ void register_hnsw(py::module_& m) { py::arg("ef_construction") = 200, py::arg("nbits") = 8, py::arg("metric") = "l2", - py::arg("random_seed") = 100 + py::arg("random_seed") = 100, + R"doc(An approximate nearest-neighbor index using HNSW. + +Parameters +---------- +dim : int + Original vector dimension, from 64 through 65536. +max_elements : int + Storage capacity, including removed points. Use resize to grow it. +M : int, optional + Graph connectivity, at least 2. Default is 16. +ef_construction : int, optional + Construction search breadth. Default is 200. +nbits : int, optional + Total quantization bits, 1 through 9. Default is 8. +metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. +random_seed : int, optional + Graph random seed. Default is 100. + +Notes +----- +Build or load the index before searching or saving. Concurrent reads are +supported; an overlapping mutation raises RuntimeError.)doc" ) .def( "build", @@ -308,7 +357,30 @@ void register_hnsw(py::module_& m) { py::arg("centroids"), py::arg("cluster_ids"), py::arg("num_threads") = 1, - py::arg("fast_quantization") = false + py::arg("fast_quantization") = false, + R"doc(Build the index from vectors and supplied centroid assignments. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +centroids : array_like of float, shape (num_clusters, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +cluster_ids : array_like of int, shape (n,) + Centroid ID per vector, in [0, num_clusters). Floating-point and + boolean IDs are rejected. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. +fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + +Notes +----- +The initial data must satisfy 1 <= n <= max_elements. Centroids and data +use original coordinates. The input arrays need not be retained after this +synchronous call.)doc" ) .def( "search", @@ -316,7 +388,34 @@ void register_hnsw(py::module_& m) { py::arg("queries"), py::arg("k"), py::arg("ef") = 0, - py::arg("num_threads") = 1 + py::arg("num_threads") = 1, + R"doc(Search a batch of query vectors. + +Parameters +---------- +queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +k : int + Number of neighbors requested, between 1 and max_elements. +ef : int, optional + Search breadth. 0 (default) selects max(k, 10). +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + +Returns +------- +ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. +distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + +Notes +----- +Returns (ids, distances), in that order. Missing neighbors have ID +4294967295 and distance inf. Array shapes and value ranges are checked at +runtime, not by static type checkers.)doc" ) .def( "add", @@ -324,26 +423,125 @@ void register_hnsw(py::module_& m) { py::arg("data"), py::arg("cluster_ids") = py::none(), py::arg("num_threads") = 1, - py::arg("fast_quantization") = false + py::arg("fast_quantization") = false, + R"doc(Append vectors using the existing centroids. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +cluster_ids : array_like of int, shape (n,), optional + Existing centroid IDs. None assigns each vector to its nearest + centroid. Centroids do not move. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. +fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + +Returns +------- +ids : ndarray of uint32, shape (n,) + Assigned point IDs. + +Notes +----- +The stored point count plus n must fit max_elements; resize first if +needed. Removed IDs are not reused.)doc" + ) + .def( + "resize", + &HnswIndex::resize, + py::arg("max_elements"), + R"doc(Resize the storage capacity of a built index. + +Parameters +---------- +max_elements : int + New capacity; must be at least num_stored_points.)doc" ) - .def("resize", &HnswIndex::resize, py::arg("max_elements")) .def( "remove", &HnswIndex::remove, py::arg("ids"), - "Exclude ids from later search results and return how many were newly " - "removed. Removed points stay in the graph and in num_stored_points, and " - "cannot " - "be restored." + R"doc(Exclude point IDs from subsequent search results. + +Parameters +---------- +ids : array_like of int, shape (n,) + IDs in [0, num_stored_points). Floating-point and boolean IDs are + rejected. + +Returns +------- +count : int + Number of newly removed points; repeats do not count twice. + +Notes +----- +Removal does not reclaim storage and cannot be undone.)doc" + ) + .def( + "save", + &HnswIndex::save, + py::arg("path"), + R"doc(Save a built index to disk. + +Parameters +---------- +path : str + Index file path. Pass str(path) for pathlib.Path objects. + +Notes +----- +The saved state includes the rotation used to build the index.)doc" + ) + .def_static( + "load", + &HnswIndex::load, + py::arg("path"), + R"doc(Load an index and its saved rotation from disk. + +Parameters +---------- +path : str + Index file path. Pass str(path) for pathlib.Path objects. + +Returns +------- +index : HnswIndex + A built index ready for search.)doc" + ) + .def_property_readonly( + "dim", &HnswIndex::dim, R"doc(Original vector dimension.)doc" + ) + .def_property_readonly( + "max_elements", &HnswIndex::max_elements, R"doc(Current storage capacity.)doc" + ) + .def_property_readonly( + "num_points", + &HnswIndex::num_points, + R"doc(Number of live, searchable points.)doc" + ) + .def_property_readonly( + "num_stored_points", + &HnswIndex::num_stored_points, + R"doc(Number of stored points, including removed points.)doc" ) - .def("save", &HnswIndex::save, py::arg("path")) - .def_static("load", &HnswIndex::load, py::arg("path")) - .def_property_readonly("dim", &HnswIndex::dim) - .def_property_readonly("max_elements", &HnswIndex::max_elements) - .def_property_readonly("num_points", &HnswIndex::num_points) - .def_property_readonly("num_stored_points", &HnswIndex::num_stored_points) - .def_property_readonly("nbits", &HnswIndex::nbits) - .def_property_readonly("num_clusters", &HnswIndex::num_clusters) - .def_property_readonly("is_built", &HnswIndex::is_built) - .def_property_readonly("metric", &HnswIndex::metric); + .def_property_readonly( + "nbits", &HnswIndex::nbits, R"doc(Total storage bits per coordinate.)doc" + ) + .def_property_readonly( + "num_clusters", + &HnswIndex::num_clusters, + R"doc(Number of centroids used for quantization.)doc" + ) + .def_property_readonly( + "is_built", + &HnswIndex::is_built, + R"doc(Whether the index has been built or loaded.)doc" + ) + .def_property_readonly( + "metric", &HnswIndex::metric, R"doc(Canonical metric name: 'l2' or 'ip'.)doc" + ); } diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index c84f9611..ed20c0dd 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -327,7 +327,43 @@ class IvfIndex { void register_ivf(py::module_& m) { using namespace rabitqlib::python_bindings; - py::class_(m, "IvfIndex") + py::class_( + m, "IvfIndex", R"doc(An approximate nearest-neighbor index using IVF. + +Parameters +---------- +dim : int + Original vector dimension, from 64 through 65536. +max_elements : int + Number of vectors required by the initial build. Later adds grow + storage. +num_clusters : int + Number of supplied centroids. +nbits : int + Total quantization bits, 1 through 9, or 32 for owned float32 vectors + with one-bit filtering and raw reranking. +metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. +initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional + Centroid routing strategy. Default is 'auto'. + +Notes +----- +Build or load the index before searching or saving. Concurrent reads are +supported; an overlapping mutation raises RuntimeError. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import IvfIndex +>>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) +>>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) +>>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) +>>> ids, distances = index.search(x[:2], k=3, nprobe=1) +>>> ids.shape, distances.shape +((2, 3), (2, 3)))doc" + ) .def( py::init< size_t, @@ -342,9 +378,41 @@ void register_ivf(py::module_& m) { py::arg("nbits"), py::arg("metric") = "l2", py::arg("initializer") = "auto", - "Create IVF with 1-9 quantization bits, or nbits=32 for owned raw-vector " - "reranking. initializer selects auto, flat, flat_rabitq, or hnsw centroid " - "routing." + R"doc(An approximate nearest-neighbor index using IVF. + +Parameters +---------- +dim : int + Original vector dimension, from 64 through 65536. +max_elements : int + Number of vectors required by the initial build. Later adds grow + storage. +num_clusters : int + Number of supplied centroids. +nbits : int + Total quantization bits, 1 through 9, or 32 for owned float32 vectors + with one-bit filtering and raw reranking. +metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. +initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional + Centroid routing strategy. Default is 'auto'. + +Notes +----- +Build or load the index before searching or saving. Concurrent reads are +supported; an overlapping mutation raises RuntimeError. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import IvfIndex +>>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) +>>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) +>>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) +>>> ids, distances = index.search(x[:2], k=3, nprobe=1) +>>> ids.shape, distances.shape +((2, 3), (2, 3)))doc" ) .def( "build", @@ -353,7 +421,30 @@ void register_ivf(py::module_& m) { py::arg("centroids"), py::arg("cluster_ids"), py::arg("num_threads") = 1, - py::arg("fast_quantization") = false + py::arg("fast_quantization") = false, + R"doc(Build the index from vectors and supplied centroid assignments. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +centroids : array_like of float, shape (num_clusters, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +cluster_ids : array_like of int, shape (n,) + Centroid ID per vector, in [0, num_clusters). Floating-point and + boolean IDs are rejected. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. +fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + +Notes +----- +The initial data must satisfy n == max_elements. Centroids and data use +original coordinates. The input arrays need not be retained after this +synchronous call.)doc" ) .def( "add", @@ -362,26 +453,69 @@ void register_ivf(py::module_& m) { py::arg("cluster_ids") = py::none(), py::arg("num_threads") = 1, py::arg("fast_quantization") = false, - "Append vectors without the original data and return their ids. Vectors " - "are quantized against the existing centroids, which do not move. Without " - "cluster_ids each vector goes to its nearest centroid. Each call copies " - "every cluster it adds to, unless reserve left room for the new vectors, " - "so add many vectors per call rather than one at a time." + R"doc(Append vectors using the existing centroids. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +cluster_ids : array_like of int, shape (n,), optional + Existing centroid IDs. None assigns each vector to its nearest + centroid. Centroids do not move. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. +fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + +Returns +------- +ids : ndarray of uint32, shape (n,) + Assigned point IDs. + +Notes +----- +Batch additions to reduce repeated cluster copying; reserve can preallocate +room.)doc" ) .def( "reserve", &IvfIndex::reserve, py::arg("cluster_id"), py::arg("num_points"), - "Preallocate room for a cluster to hold num_points vectors, so later adds " - "to it do not copy it. Reserved room is not saved." + R"doc(Preallocate storage for a cluster. + +Parameters +---------- +cluster_id : int + Cluster ID in [0, num_clusters). +num_points : int + Nonnegative capacity to reserve for this cluster. + +Notes +----- +Requires a built index. Reserved room is not saved.)doc" ) .def( "remove", &IvfIndex::remove, py::arg("ids"), - "Exclude ids from later search results and return how many were newly " - "removed. Removed vectors keep their storage and cannot be restored." + R"doc(Exclude point IDs from subsequent search results. + +Parameters +---------- +ids : array_like of int, shape (n,) + IDs in [0, num_stored_points). Floating-point and boolean IDs are + rejected. + +Returns +------- +count : int + Number of newly removed points; repeats do not count twice. + +Notes +----- +Removal does not reclaim storage and cannot be undone.)doc" ) .def( "search", @@ -390,17 +524,102 @@ void register_ivf(py::module_& m) { py::arg("k"), py::arg("nprobe"), py::arg("high_accuracy") = py::none(), - py::arg("num_threads") = 1 + py::arg("num_threads") = 1, + R"doc(Search a batch of query vectors. + +Parameters +---------- +queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +k : int + Number of neighbors requested, between 1 and max_elements. +nprobe : int + Positive number of clusters to probe; capped at num_clusters. +high_accuracy : bool, optional + Enable HACC. None (default) selects it for 4-9 bits and standard + FastScan otherwise. Raw storage uses standard FastScan. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + +Returns +------- +ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. +distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + +Notes +----- +Returns (ids, distances), in that order. Missing neighbors have ID +4294967295 and distance inf. Array shapes and value ranges are checked at +runtime, not by static type checkers.)doc" + ) + .def( + "save", + &IvfIndex::save, + py::arg("path"), + R"doc(Save a built index to disk. + +Parameters +---------- +path : str + Index file path. Pass str(path) for pathlib.Path objects. + +Notes +----- +IVF with HNSW centroid routing also writes its required sidecar file.)doc" + ) + .def_static( + "load", + &IvfIndex::load, + py::arg("path"), + R"doc(Load an index and its saved rotation from disk. + +Parameters +---------- +path : str + Index file path. Pass str(path) for pathlib.Path objects. + +Returns +------- +index : IvfIndex + A built index ready for search.)doc" + ) + .def_property_readonly("dim", &IvfIndex::dim, R"doc(Original vector dimension.)doc") + .def_property_readonly( + "max_elements", &IvfIndex::max_elements, R"doc(Current storage capacity.)doc" + ) + .def_property_readonly( + "num_points", + &IvfIndex::num_points, + R"doc(Number of live, searchable points.)doc" + ) + .def_property_readonly( + "num_stored_points", + &IvfIndex::num_stored_points, + R"doc(Number of stored points, including removed points.)doc" + ) + .def_property_readonly( + "num_clusters", + &IvfIndex::num_clusters, + R"doc(Number of centroids used for quantization.)doc" ) - .def("save", &IvfIndex::save, py::arg("path")) - .def_static("load", &IvfIndex::load, py::arg("path")) - .def_property_readonly("dim", &IvfIndex::dim) - .def_property_readonly("max_elements", &IvfIndex::max_elements) - .def_property_readonly("num_points", &IvfIndex::num_points) - .def_property_readonly("num_stored_points", &IvfIndex::num_stored_points) - .def_property_readonly("num_clusters", &IvfIndex::num_clusters) - .def_property_readonly("nbits", &IvfIndex::nbits) - .def_property_readonly("is_built", &IvfIndex::is_built) - .def_property_readonly("metric", &IvfIndex::metric) - .def_property_readonly("initializer", &IvfIndex::initializer); + .def_property_readonly( + "nbits", &IvfIndex::nbits, R"doc(Total storage bits per coordinate.)doc" + ) + .def_property_readonly( + "is_built", + &IvfIndex::is_built, + R"doc(Whether the index has been built or loaded.)doc" + ) + .def_property_readonly( + "metric", &IvfIndex::metric, R"doc(Canonical metric name: 'l2' or 'ip'.)doc" + ) + .def_property_readonly( + "initializer", + &IvfIndex::initializer, + R"doc(Resolved centroid routing strategy.)doc" + ); } diff --git a/python_bindings/kmeans_bindings.cpp b/python_bindings/kmeans_bindings.cpp index a889934c..53cc1d88 100644 --- a/python_bindings/kmeans_bindings.cpp +++ b/python_bindings/kmeans_bindings.cpp @@ -60,15 +60,43 @@ double train(Clustering& clustering, py::handle x) { template void bind_training_parameters(py::class_& binding) { - binding.def(py::init<>()) - .def_readwrite("niter", &Parameters::niter) - .def_readwrite("verbose", &Parameters::verbose) - .def_readwrite("spherical", &Parameters::spherical) - .def_readwrite("seed", &Parameters::seed) - .def_readwrite("min_points_per_centroid", &Parameters::min_points_per_centroid) - .def_readwrite("early_stop_threshold", &Parameters::early_stop_threshold) - .def_readwrite("num_threads", &Parameters::num_threads) - .def_readwrite("final_assignment", &Parameters::final_assignment); + binding.def(py::init<>(), "Create parameters with their documented defaults.") + .def_readwrite( + "niter", + &Parameters::niter, + R"doc(Maximum number of clustering iterations. Default is 25.)doc" + ) + .def_readwrite( + "verbose", + &Parameters::verbose, + R"doc(Print training progress when True. Default is False.)doc" + ) + .def_readwrite( + "spherical", + &Parameters::spherical, + R"doc(Normalize centroids after each update and use inner-product assignment. Normalize input vectors yourself for cosine clustering. Default is False.)doc" + ) + .def_readwrite("seed", &Parameters::seed, R"doc(Random seed. Default is 42.)doc") + .def_readwrite( + "min_points_per_centroid", + &Parameters::min_points_per_centroid, + R"doc(Warn below this many training vectors per centroid. Default is 39.)doc" + ) + .def_readwrite( + "early_stop_threshold", + &Parameters::early_stop_threshold, + R"doc(Stop when the absolute relative objective change is at most this threshold. Default is 0.0.)doc" + ) + .def_readwrite( + "num_threads", + &Parameters::num_threads, + R"doc(Worker count; 0 selects the available CPU count. Default is 0.)doc" + ) + .def_readwrite( + "final_assignment", + &Parameters::final_assignment, + R"doc(Approximate uses the training assignment method; Exact uses exhaustive assignment against final centroids. Default is FinalAssignmentMode.Approximate.)doc" + ); } template @@ -122,29 +150,87 @@ void register_kmeans(py::module_& m) { using rabitqlib::rabitqkmeans::RaBitQKMeans; using rabitqlib::rabitqkmeans::RaBitQKMeansParameters; - py::enum_(m, "FinalAssignmentMode") + py::enum_( + m, "FinalAssignmentMode", R"doc(Final centroid assignment policy. + +Approximate uses the selected clustering algorithm. SymphonyQG is a +compatibility alias for Approximate. Exact performs exhaustive assignment.)doc" + ) .value("Approximate", FinalAssignmentMode::Approximate) .value("SymphonyQG", FinalAssignmentMode::SymphonyQG) .value("Exact", FinalAssignmentMode::Exact); - py::class_ qg_parameters(m, "QGKMeansParameters"); + py::class_ qg_parameters( + m, "QGKMeansParameters", R"doc(Mutable training settings for QGKMeans. + +Notes +----- +Create with no arguments; set attributes directly or use keyword arguments +on QGKMeans.)doc" + ); bind_training_parameters(qg_parameters); - qg_parameters.def_readwrite("graph_degree", &QGKMeansParameters::graph_degree) - .def_readwrite("ef_build", &QGKMeansParameters::ef_build) - .def_readwrite("ef_search", &QGKMeansParameters::ef_search) + qg_parameters .def_readwrite( - "graph_build_iterations", &QGKMeansParameters::graph_build_iterations + "graph_degree", + &QGKMeansParameters::graph_degree, + R"doc(Centroid graph degree; a positive multiple of 32. Default is 32.)doc" ) - .def_readwrite("quantization_bits", &QGKMeansParameters::quantization_bits); - - py::class_ flat_parameters(m, "RaBitQKMeansParameters"); + .def_readwrite( + "ef_build", + &QGKMeansParameters::ef_build, + R"doc(Positive centroid-graph construction search breadth. Default is 240.)doc" + ) + .def_readwrite( + "ef_search", + &QGKMeansParameters::ef_search, + R"doc(Positive centroid-graph assignment search breadth. Default is 16.)doc" + ) + .def_readwrite( + "graph_build_iterations", + &QGKMeansParameters::graph_build_iterations, + R"doc(Positive number of graph-building passes after initialization. Default is 1.)doc" + ) + .def_readwrite( + "quantization_bits", + &QGKMeansParameters::quantization_bits, + R"doc(Centroid storage: 0 for raw float32, or 4 or 8 for RaBitQ. Default is 0.)doc" + ); + + py::class_ flat_parameters( + m, "RaBitQKMeansParameters", R"doc(Mutable training settings for RaBitQKMeans. + +Notes +----- +Create with no arguments; set attributes directly or use keyword arguments +on RaBitQKMeans.)doc" + ); bind_training_parameters(flat_parameters); - py::class_(m, "QGKMeansIterationStats") - .def_readonly("iteration", &QGKMeansIterationStats::iteration) - .def_readonly("obj", &QGKMeansIterationStats::obj) - .def_readonly("shift", &QGKMeansIterationStats::shift) - .def_readonly("nsplit", &QGKMeansIterationStats::nsplit); + py::class_( + m, + "QGKMeansIterationStats", + R"doc(Read-only statistics for one completed clustering iteration.)doc" + ) + .def_readonly( + "iteration", + &QGKMeansIterationStats::iteration, + R"doc(Iteration number, starting at 1.)doc" + ) + .def_readonly( + "obj", + &QGKMeansIterationStats::obj, + R"doc(Objective measured before the centroid update.)doc" + ) + .def_readonly( + "shift", + &QGKMeansIterationStats::shift, + R"doc(Total squared centroid displacement during the update.)doc" + ) + .def_readonly( + "nsplit", + &QGKMeansIterationStats::nsplit, + R"doc(Number of empty centroids repaired by splitting.)doc" + ); m.attr("RaBitQKMeansIterationStats") = m.attr("QGKMeansIterationStats"); bind_clustering(m, "_QGKMeans", "QGKMeans"); diff --git a/python_bindings/py.typed b/python_bindings/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 0dd09948..910b4dd4 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -211,13 +211,52 @@ class SymqgIndex { void register_symqg(py::module_& m) { using namespace rabitqlib::python_bindings; - py::class_(m, "SymqgIndex") + py::class_( + m, "SymqgIndex", R"doc(An approximate nearest-neighbor index using SymphonyQG. + +Parameters +---------- +dim : int + Original vector dimension, from 64 through 65536. +max_degree : int + Graph degree bound; a positive multiple of 32. +metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. +quantization_bits : {0, 4, 8}, optional + Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 + (default). + +Notes +----- +Build or load the index before searching or saving. Concurrent reads are +supported; an overlapping mutation raises RuntimeError.)doc" + ) .def( py::init(), py::arg("dim"), py::arg("max_degree"), py::arg("metric") = "l2", - py::arg("quantization_bits") = 0 + py::arg("quantization_bits") = 0, + R"doc(An approximate nearest-neighbor index using SymphonyQG. + +Parameters +---------- +dim : int + Original vector dimension, from 64 through 65536. +max_degree : int + Graph degree bound; a positive multiple of 32. +metric : {'l2', 'ip', 'innerproduct'}, optional + Squared L2 or inner-product distance. Default is 'l2'. Vectors are not + automatically normalized for cosine similarity. +quantization_bits : {0, 4, 8}, optional + Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 + (default). + +Notes +----- +Build or load the index before searching or saving. Concurrent reads are +supported; an overlapping mutation raises RuntimeError.)doc" ) .def( "build", @@ -225,7 +264,25 @@ void register_symqg(py::module_& m) { py::arg("data"), py::arg("ef_construction"), py::arg("num_threads") = 1, - py::arg("init") = "pipnn" + py::arg("init") = "pipnn", + R"doc(Build the graph from input vectors. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +ef_construction : int + Positive construction search breadth fitting uint32. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. +init : {'pipnn', 'random'}, optional + Graph initialization method. Default is 'pipnn'. + +Notes +----- +Requires n > max_degree. Input arrays need not be retained after this +synchronous call.)doc" ) .def( "search", @@ -233,7 +290,33 @@ void register_symqg(py::module_& m) { py::arg("queries"), py::arg("k"), py::arg("ef"), - py::arg("num_threads") = 1 + py::arg("num_threads") = 1, + R"doc(Search a batch of query vectors. + +Parameters +---------- +queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +k : int + Number of neighbors requested, between 1 and num_points. +ef : int + Search breadth. Must be positive and at least k. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + +Returns +------- +ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. +distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + +Notes +----- +Returns (ids, distances), in that order. Array shapes and value ranges are +checked at runtime, not by static type checkers.)doc" ) .def( "search_batch", @@ -241,16 +324,95 @@ void register_symqg(py::module_& m) { py::arg("queries"), py::arg("k"), py::arg("ef"), - py::arg("num_threads") = 1 + py::arg("num_threads") = 1, + R"doc(Search a batch of query vectors. + +Parameters +---------- +queries : array_like of float, shape (n_queries, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +k : int + Number of neighbors requested, between 1 and num_points. +ef : int + Search breadth. Must be positive and at least k. +num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + +Returns +------- +ids : ndarray of uint32, shape (n_queries, k) + Point IDs of the selected neighbors. +distances : ndarray of float32, shape (n_queries, k) + Lower values are better; squared L2 estimates for l2 and inner-product + distance estimates for ip. These are distances, not similarity scores. + +Notes +----- +Returns (ids, distances), in that order. Array shapes and value ranges are +checked at runtime, not by static type checkers.)doc" ) - .def("save", &SymqgIndex::save, py::arg("path")) - .def_static("load", &SymqgIndex::load, py::arg("path")) - .def_property_readonly("dim", &SymqgIndex::dim) - .def_property_readonly("max_degree", &SymqgIndex::max_degree) - .def_property_readonly("num_points", &SymqgIndex::num_points) - .def_property_readonly("num_stored_points", &SymqgIndex::num_stored_points) - .def_property_readonly("max_elements", &SymqgIndex::max_elements) - .def_property_readonly("is_built", &SymqgIndex::is_built) - .def_property_readonly("metric", &SymqgIndex::metric) - .def_property_readonly("quantization_bits", &SymqgIndex::quantization_bits); + .def( + "save", + &SymqgIndex::save, + py::arg("path"), + R"doc(Save a built index to disk. + +Parameters +---------- +path : str + Index file path. Pass str(path) for pathlib.Path objects. + +Notes +----- +The saved state includes the rotation used to build the index.)doc" + ) + .def_static( + "load", + &SymqgIndex::load, + py::arg("path"), + R"doc(Load an index and its saved rotation from disk. + +Parameters +---------- +path : str + Index file path. Pass str(path) for pathlib.Path objects. + +Returns +------- +index : SymqgIndex + A built index ready for search.)doc" + ) + .def_property_readonly( + "dim", &SymqgIndex::dim, R"doc(Original vector dimension.)doc" + ) + .def_property_readonly( + "max_degree", &SymqgIndex::max_degree, R"doc(Graph degree bound.)doc" + ) + .def_property_readonly( + "num_points", + &SymqgIndex::num_points, + R"doc(Number of live, searchable points.)doc" + ) + .def_property_readonly( + "num_stored_points", + &SymqgIndex::num_stored_points, + R"doc(Number of stored points, including removed points.)doc" + ) + .def_property_readonly( + "max_elements", &SymqgIndex::max_elements, R"doc(Current storage capacity.)doc" + ) + .def_property_readonly( + "is_built", + &SymqgIndex::is_built, + R"doc(Whether the index has been built or loaded.)doc" + ) + .def_property_readonly( + "metric", &SymqgIndex::metric, R"doc(Canonical metric name: 'l2' or 'ip'.)doc" + ) + .def_property_readonly( + "quantization_bits", + &SymqgIndex::quantization_bits, + R"doc(Storage bits: 0 for raw float32, or 4 or 8 for RaBitQ.)doc" + ); } diff --git a/tests/README.md b/tests/README.md index 503eef3c..64f0be54 100644 --- a/tests/README.md +++ b/tests/README.md @@ -321,3 +321,20 @@ wraps C allocations to sweep HNSW-RaBitQ constructor failures and verify ownersh these under ASan/UBSan as well as Release; do not disable allocator-mismatch checking. `compact_codes_failure_test` covers cold scratch allocation and growth failures in serial and OpenMP encoding, including the public FastScan batch entry points. + +## Python editor support + +Wheels include type stubs and `py.typed`; select the installed environment in your editor. +Keep [public declarations](../python_bindings/__init__.pyi) and runtime docstrings in sync, +following the [numpydoc style](https://numpydoc.readthedocs.io/en/latest/format.html). +Array shapes and value ranges remain runtime checks. + +After rebuilding and installing with the `test` extra, run: + +```bash +python -m pytest tests/python/test_editor_support.py -q +``` + +The suite checks installed declarations, signatures, documentation, examples, and +valid/invalid calls with strict mypy. CI requires it; missing mypy fails the test. +Keep the mypy pins in the `test` extra and wheel test requirements synchronized. diff --git a/tests/python/test_editor_support.py b/tests/python/test_editor_support.py new file mode 100644 index 00000000..df23aef5 --- /dev/null +++ b/tests/python/test_editor_support.py @@ -0,0 +1,255 @@ +"""Installed-package contracts for editor declarations and runtime help.""" + +import ast +import doctest +import re +import subprocess +import sys +from pathlib import Path + +import pytest +import rabitqlib + + +def installed_stub(): + package = Path(rabitqlib.__file__).parent + assert (package / "py.typed").is_file() + stub = package / "__init__.pyi" + assert stub.is_file() + return ast.parse(stub.read_text()) + + +def declarations(): + return { + node.name: node + for node in installed_stub().body + if isinstance(node, ast.ClassDef) + } + + +def test_installed_public_declarations(): + tree = installed_stub() + exports = next( + node + for node in tree.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "__all__" + for target in node.targets + ) + ) + assert ast.literal_eval(exports.value) == rabitqlib.__all__ + names = {node.name for node in tree.body if isinstance(node, ast.ClassDef)} + names.update( + target.id + for node in tree.body + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) + ) + assert set(rabitqlib.__all__) <= names + for name, cls in declarations().items(): + runtime = getattr(rabitqlib, name) + members = {node.name for node in cls.body if isinstance(node, ast.FunctionDef)} + for node in cls.body: + if isinstance(node, ast.AnnAssign): + assert isinstance(node.target, ast.Name) + members.add(node.target.id) + if name == "FinalAssignmentMode": + assert {"Approximate", "SymphonyQG", "Exact", "name", "value"} <= members + continue + public = {key for key in vars(runtime) if not key.startswith("_")} + assert public <= members, (name, public - members) + for member in members - {"d", "k", "cp"}: + assert hasattr(runtime, member), (name, member) + + +def normalized(text): + return " ".join(text.split()) + + +def test_editor_docs_match_runtime_help(): + for name, cls in declarations().items(): + runtime = getattr(rabitqlib, name) + assert normalized(ast.get_docstring(cls)) in normalized(runtime.__doc__) + for node in cls.body: + if not isinstance(node, ast.FunctionDef): + continue + documentation = ast.get_docstring(node) + if documentation is None: + continue + # Python wrappers inherit __init__(d, k, **kwargs); class help documents + # the supported keyword parameters and the stub expands them for IDEs. + if name in ("QGKMeans", "RaBitQKMeans") and node.name == "__init__": + continue + actual = getattr(runtime, node.name).__doc__ + assert normalized(documentation) in normalized(actual), (name, node.name) + + +@pytest.mark.parametrize("name", ["IvfIndex", "HnswIndex", "SymqgIndex"]) +def test_native_signature_names_and_defaults(name): + cls = declarations()[name] + runtime = getattr(rabitqlib, name) + for node in cls.body: + if not isinstance(node, ast.FunctionDef): + continue + if any( + isinstance(d, ast.Name) and d.id == "property" for d in node.decorator_list + ): + continue + signature = getattr(runtime, node.name).__doc__.splitlines()[0] + actual_names = re.findall(r"(?:\(|, )([A-Za-z_]\w*):", signature) + expected_names = [a.arg for a in node.args.args] + assert actual_names == expected_names, (name, node.name, signature) + default_args = node.args.args[len(node.args.args) - len(node.args.defaults) :] + for arg, value in zip(default_args, node.args.defaults): + actual = re.search(rf"\b{arg.arg}: [^=]+ = ([^,)]+)", signature) + assert actual is not None, (name, node.name, arg.arg) + assert ast.literal_eval(actual[1]) == ast.literal_eval(value) + + +@pytest.mark.parametrize("name", ["QGKMeans", "RaBitQKMeans"]) +def test_clustering_keyword_defaults(name): + cls = declarations()[name] + init = next( + n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "__init__" + ) + instance = getattr(rabitqlib, name)(64, 40) + for arg, default in zip(init.args.kwonlyargs, init.args.kw_defaults): + assert default is not None, (name, arg.arg) + if isinstance(default, ast.Attribute): + expected = getattr(rabitqlib.FinalAssignmentMode, default.attr) + else: + expected = ast.literal_eval(default) + assert getattr(instance.cp, arg.arg) == expected + for name in ( + "centroids", + "assignments", + "distances", + "obj", + "final_obj", + "iteration_stats", + ): + assert getattr(instance, name) is None + + +def test_documentation_examples(): + runner = doctest.DocTestRunner() + parser = doctest.DocTestParser() + for name in rabitqlib.__all__: + # Explicitly include native classes, which testmod otherwise skips + # because they belong to the extension module. + documentation = getattr(rabitqlib, name).__doc__ + runner.run(parser.get_doctest(documentation, {}, name, None, 0)) + assert runner.failures == 0 + assert runner.tries > 0 + + +def test_static_public_usage(tmp_path): + fixture = Path(__file__).parent / "typing" / "public_api.py" + result = subprocess.run( + [ + sys.executable, + "-m", + "mypy", + "--strict", + "--no-incremental", + "--cache-dir", + str(tmp_path / "cache"), + str(fixture), + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + # Removing ignores must expose the deliberately invalid calls, proving the + # package is typed rather than silently treated as Any. + invalid = tmp_path / "invalid.py" + invalid.write_text(re.sub(r" # type: ignore\[[^]]+\]", "", fixture.read_text())) + result = subprocess.run( + [ + sys.executable, + "-m", + "mypy", + "--strict", + "--no-incremental", + "--cache-dir", + str(tmp_path / "cache"), + str(invalid), + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 1, result.stdout + result.stderr + expected_lines = { + number + for number, line in enumerate(fixture.read_text().splitlines(), start=1) + if "# type: ignore[" in line + } + error_lines = { + int(number) + for number in re.findall( + rf"{re.escape(str(invalid))}:(\d+): error:", result.stdout + ) + } + assert error_lines == expected_lines, result.stdout + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "symqg"]) +def test_documented_search_result_types(kind): + import numpy as np + + data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + if kind == "symqg": + index = rabitqlib.SymqgIndex(64, 32) + index.build(data, 40) + result = index.search(data[:2], 3, ef=20) + else: + centroids = data[:2].copy() + labels = np.arange(80, dtype=np.uint32) % 2 + if kind == "ivf": + index = rabitqlib.IvfIndex(64, 80, 2, 4) + else: + index = rabitqlib.HnswIndex(64, 80) + index.build(data, centroids, labels) + if isinstance(index, rabitqlib.IvfIndex): + result = index.search(data[:2], 3, nprobe=2) + else: + result = index.search(data[:2], 3, ef=20) + ids, distances = result + assert ids.shape == distances.shape == (2, 3) + assert ids.dtype == np.uint32 + assert distances.dtype == np.float32 + + +@pytest.mark.parametrize("name", ["QGKMeans", "RaBitQKMeans"]) +def test_documented_training_result_types(name): + import numpy as np + + data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + model = getattr(rabitqlib, name)(64, 40, niter=2, num_threads=1) + assert isinstance(model.train(data), float) + assert model.centroids.shape == (40, 64) + assert model.centroids.dtype == np.float32 + assert model.assignments.shape == model.distances.shape == (80,) + assert model.assignments.dtype == np.uint32 + assert model.distances.dtype == np.float32 + assert model.obj.dtype == np.float64 + assert len(model.obj) == len(model.iteration_stats) + assert isinstance(model.final_obj, float) + + +def test_numpy_scalar_inputs_and_enum_members(): + import runpy + + fixture = Path(__file__).parent / "typing" / "public_api.py" + runpy.run_path(str(fixture))["check_numpy_scalars"]() + + +def test_valid_string_options_and_integer_ids(): + import runpy + + fixture = Path(__file__).parent / "typing" / "public_api.py" + runpy.run_path(str(fixture))["check_input_contracts"]() diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py new file mode 100644 index 00000000..2955a4c3 --- /dev/null +++ b/tests/python/typing/public_api.py @@ -0,0 +1,215 @@ +"""Static-only fixture: assert precise types and reject invalid usage.""" + +from typing import Literal, assert_type + +import numpy as np +import numpy.typing as npt +import rabitqlib +from rabitqlib import ( + FinalAssignmentMode, + HnswIndex, + IvfIndex, + QGKMeans, + QGKMeansIterationStats, + QGKMeansParameters, + RaBitQKMeans, + RaBitQKMeansIterationStats, + RaBitQKMeansParameters, + SymqgIndex, +) + + +def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: + assert_type(rabitqlib.__all__, list[str]) + ivf = IvfIndex(64, 64, 2, 4) + hnsw = HnswIndex(64, 128) + graph = SymqgIndex(64, 32) + for index in (ivf, hnsw): + index.build(x, x[:2], labels) + assert_type(index.add(x, labels), npt.NDArray[np.uint32]) + assert_type(index.remove([0, 1]), int) + assert_type(index.dim, int) + assert_type(index.is_built, bool) + index.save("index.bin") + ids, distances = ivf.search(x, k=5, nprobe=2) + assert_type(ids, npt.NDArray[np.uint32]) + assert_type(distances, npt.NDArray[np.float32]) + assert_type( + hnsw.search(x, 5), tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]] + ) + graph.build(x, 64) + assert_type( + graph.search_batch(x, 5, 32), + tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]], + ) + assert_type(IvfIndex.load("index.bin"), IvfIndex) + assert_type(HnswIndex.load("index.bin"), HnswIndex) + assert_type(SymqgIndex.load("index.bin"), SymqgIndex) + ivf.reserve(0, 128) + hnsw.resize(256) + qg = QGKMeans(64, 40, niter=2, quantization_bits=4) + flat = RaBitQKMeans(64, 2, final_assignment=FinalAssignmentMode.Exact) + assert_type(qg.cp, QGKMeansParameters) + assert_type(flat.cp, RaBitQKMeansParameters) + for model in (qg, flat): + assert_type(model.train(x), float) + assert_type(model.centroids, npt.NDArray[np.float32] | None) + assert_type(model.assignments, npt.NDArray[np.uint32] | None) + assert_type(model.distances, npt.NDArray[np.float32] | None) + assert_type(model.obj, npt.NDArray[np.float64] | None) + assert_type(model.final_obj, float | None) + if model.iteration_stats is not None: + assert_type(model.iteration_stats[0], QGKMeansIterationStats) + qg.cp.final_assignment = FinalAssignmentMode.Approximate + ivf.search(x, "five", 2) # type: ignore[arg-type] + ivf.dim = 16 # type: ignore[misc] + qg.cp.niter = "two" # type: ignore[assignment] + RaBitQKMeans(64, 2, graph_degree=32) # type: ignore[call-arg] + hnsw.search(x, 5, nprobe=2) # type: ignore[call-arg] + + +def stats_alias(stats: RaBitQKMeansIterationStats) -> QGKMeansIterationStats: + return stats + + +def check_parameter_reads( + parameters: QGKMeansParameters | RaBitQKMeansParameters, +) -> None: + # Read in a separate scope: Pyright narrows union-typed property writes to + # the assigned NumPy scalar, although the native setter normalizes it. + assert_type(parameters.niter, int) + assert_type(parameters.verbose, bool) + assert type(parameters.niter) is int + assert type(parameters.verbose) is bool + + +def check_numpy_scalars() -> None: + """The same positive cases run through mypy and the compiled bindings.""" + x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + labels = np.zeros(80, dtype=np.uint32) + centroids = x.mean(axis=0, keepdims=True) + ivf = IvfIndex(np.int64(64), np.int64(80), np.int64(1), np.int64(4)) + hnsw = HnswIndex(np.int64(64), np.int64(80), M=np.int64(16)) + graph = SymqgIndex(np.int64(64), np.int64(32), quantization_bits=np.int64(4)) + for index in (ivf, hnsw): + index.build( + x, + centroids, + labels, + num_threads=np.int64(1), + fast_quantization=np.bool_(False), + ) + assert_type(index.dim, int) + assert type(index.dim) is int + ivf.reserve(np.uint32(0), np.int64(100)) + ids, distances = ivf.search( + x[:2], k=np.int64(3), nprobe=np.int64(1), high_accuracy=np.bool_(True) + ) + assert_type(ids, npt.NDArray[np.uint32]) + assert_type(distances, npt.NDArray[np.float32]) + hnsw.resize(np.int64(100)) + hnsw.search(x[:2], np.int64(3), ef=np.int64(10)) + graph.build(x, np.int64(40), num_threads=np.int64(1)) + graph.search(x[:2], np.int64(3), np.int64(10)) + flat = RaBitQKMeans( + np.int64(64), + np.int64(2), + niter=np.int64(1), + spherical=np.bool_(False), + num_threads=np.int64(1), + ) + qg = QGKMeans( + np.int64(64), + np.int64(40), + niter=np.int64(1), + graph_degree=np.int64(32), + num_threads=np.int64(1), + ) + for model in (flat, qg): + model.cp.niter = np.int64(1) + model.cp.seed = np.uint32(42) + model.cp.min_points_per_centroid = np.int64(1) + model.cp.num_threads = np.int64(1) + model.cp.verbose = np.bool_(False) + model.cp.spherical = np.bool_(False) + check_parameter_reads(model.cp) + assert_type(model.train(x), float) + qg.cp.graph_degree = np.int64(32) + qg.cp.ef_build = np.int64(40) + qg.cp.ef_search = np.int64(16) + qg.cp.graph_build_iterations = np.int64(1) + qg.cp.quantization_bits = np.int64(4) + members = FinalAssignmentMode.__members__ + assert_type(members, dict[str, FinalAssignmentMode]) + assert_type(members["Exact"], FinalAssignmentMode) + assert members["Exact"] == FinalAssignmentMode.Exact + assert members["SymphonyQG"] == members["Approximate"] + assert FinalAssignmentMode(np.int64(1)) == members["Exact"] + + +def check_input_contracts() -> None: + """Valid literal options and integer ID containers also execute at runtime.""" + x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + centroids = x.mean(axis=0, keepdims=True) + signed = np.zeros(80, dtype=np.int64) + unsigned = np.zeros(80, dtype=np.uint32) + metrics: tuple[Literal["l2", "ip", "innerproduct"], ...] = ( + "l2", + "ip", + "innerproduct", + ) + for metric in metrics: + ivf = IvfIndex(64, 80, 1, 4, metric=metric) + HnswIndex(64, 80, metric=metric) + SymqgIndex(64, 32, metric=metric) + initializers: tuple[Literal["auto", "flat", "flat_rabitq", "hnsw"], ...] = ( + "auto", + "flat", + "flat_rabitq", + "hnsw", + ) + for initializer in initializers: + IvfIndex(64, 80, 1, 4, initializer=initializer) + for ids in (signed, unsigned, [0] * 80, (0,) * 80): + ivf = IvfIndex(64, 80, 1, 4) + ivf.build(x, centroids, ids) + ivf.add(x[:1], [np.int64(0)]) + ivf.add(x[:1], (np.uint32(0),)) + ivf.remove(np.array([0], dtype=np.int32)) + ivf.remove(np.array([1], dtype=np.uint64)) + ivf.remove([2]) + ivf.remove((3,)) + ivf.remove([]) + ivf.remove(()) + hnsw = HnswIndex(64, 82) + hnsw.build(x, centroids, ids) + hnsw.add(x[:1], [np.uint32(0)]) + hnsw.remove([0]) + hnsw.remove(np.array([1], dtype=np.int64)) + initializations: tuple[Literal["random", "pipnn"], ...] = ("random", "pipnn") + for init in initializations: + graph = SymqgIndex(64, 32) + graph.build(x, 40, init=init) + # Returned option values must be usable as inputs without casts. + IvfIndex(64, 80, 1, 4, metric=ivf.metric, initializer=ivf.initializer) + + +def reject_invalid_inputs(x: npt.NDArray[np.float32]) -> None: + float_ids: npt.NDArray[np.float64] = np.zeros(80, dtype=np.float64) + bool_ids: npt.NDArray[np.bool_] = np.array([True], dtype=np.bool_) + string_ids: npt.NDArray[np.str_] = np.array(["1"], dtype=np.str_) + ivf = IvfIndex(64, 80, 1, 4) + hnsw = HnswIndex(64, 80) + graph = SymqgIndex(64, 32) + IvfIndex(64, 80, 1, 4, metric="cosine") # type: ignore[arg-type] + HnswIndex(64, 80, metric="L2") # type: ignore[arg-type] + SymqgIndex(64, 32, metric="inner_product") # type: ignore[arg-type] + IvfIndex(64, 80, 1, 4, initializer="flat_rabit") # type: ignore[arg-type] + graph.build(x, 40, init="pippn") # type: ignore[arg-type] + ivf.build(x, x[:1], [0.0] * 80) # type: ignore[list-item] + hnsw.build(x, x[:1], float_ids) # type: ignore[arg-type] + ivf.add(x, float_ids) # type: ignore[arg-type] + hnsw.add(x, (0.0,)) # type: ignore[arg-type] + ivf.remove([1.5]) # type: ignore[list-item] + hnsw.remove(bool_ids) # type: ignore[arg-type] + ivf.remove(string_ids) # type: ignore[arg-type]