Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion docs/docs/index/usage.md
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,11 @@ IVF-RaBitQ grows automatically and has no spare logical capacity after build or
HNSW-RaBitQ's `resize` changes only capacity. SymphonyQG has no incremental updates;
all three counts are equal after build.

`repr(index)` shows dimension, metric, quantization bits, live point count and build
state. IVF/HNSW also show stored count, capacity and centroid count; IVF includes
the resolved initializer, and SymphonyQG includes its degree bound. The summary
reads only metadata and follows the same concurrency rules as property reads.

## Estimate storage

The following estimates describe the main float32 index storage, not total process
Expand Down Expand Up @@ -104,7 +109,7 @@ construction, updates, and file I/O. The following contract applies to

| Operation on the same index | Concurrent access |
| --- | --- |
| `search`, `search_batch`, `save`, and property reads | May run together |
| `search`, `search_batch`, `save`, `repr`, and property reads | May run together |
| `build`, `add`, `remove`, `reserve`, and `resize`, where available | Require exclusive access |

A conflicting call immediately raises
Expand Down
9 changes: 9 additions & 0 deletions python_bindings/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,9 @@ class IvfIndex:
((2, 3), (2, 3))
"""

def __repr__(self) -> str:
"""Return a one-line summary of configuration, counts and build state."""

@property
def dim(self) -> int:
"""Original vector dimension."""
Expand Down Expand Up @@ -422,6 +425,9 @@ class HnswIndex:
supported; an overlapping mutation raises RuntimeError.
"""

def __repr__(self) -> str:
"""Return a one-line summary of configuration, counts and build state."""

@property
def dim(self) -> int:
"""Original vector dimension."""
Expand Down Expand Up @@ -664,6 +670,9 @@ class SymqgIndex:
supported; an overlapping mutation raises RuntimeError.
"""

def __repr__(self) -> str:
"""Return a one-line summary of configuration, counts and build state."""

@property
def dim(self) -> int:
"""Original vector dimension."""
Expand Down
18 changes: 18 additions & 0 deletions python_bindings/hnsw_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,19 @@ class HnswIndex {
return wrapper;
}

[[nodiscard]] std::string repr() const {
const auto operation = access_.read();
return "HnswIndex(dim=" + std::to_string(index_->dimension()) + ", metric='" +
metric_to_string(index_->metric_type()) +
"', nbits=" + std::to_string(index_->nbits()) +
", num_points=" + std::to_string(built_ ? index_->num_points() : 0) +
", num_stored_points=" +
std::to_string(built_ ? index_->num_stored_points() : 0) +
", max_elements=" + std::to_string(index_->max_elements()) +
", num_clusters=" + std::to_string(index_->num_clusters()) +
", is_built=" + (built_ ? "True" : "False") + ")";
}

[[nodiscard]] size_t dim() const {
const auto operation = access_.read();
return index_->dimension();
Expand Down Expand Up @@ -521,6 +534,11 @@ Returns
index : HnswIndex
A built index ready for search.)doc"
)
.def(
"__repr__",
&HnswIndex::repr,
"Return a one-line summary of configuration, counts and build state."
)
.def_property_readonly(
"dim", &HnswIndex::dim, R"doc(Original vector dimension.)doc"
)
Expand Down
18 changes: 18 additions & 0 deletions python_bindings/ivf_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -284,6 +284,19 @@ class IvfIndex {
return wrapper;
}

[[nodiscard]] std::string repr() const {
const auto operation = access_.read();
return "IvfIndex(dim=" + std::to_string(index_->dimension()) + ", metric='" +
metric_to_string(index_->metric_type()) +
"', nbits=" + std::to_string(index_->nbits()) +
", num_points=" + std::to_string(index_->num_points()) +
", num_stored_points=" + std::to_string(index_->num_stored_points()) +
", max_elements=" + std::to_string(index_->max_elements()) +
", num_clusters=" + std::to_string(index_->num_clusters()) +
", initializer='" + initializer_to_string(index_->initializer_type()) +
"', is_built=" + (built_ ? "True" : "False") + ")";
}

[[nodiscard]] size_t dim() const {
const auto operation = access_.read();
return index_->dimension();
Expand Down Expand Up @@ -597,6 +610,11 @@ Returns
index : IvfIndex
A built index ready for search.)doc"
)
.def(
"__repr__",
&IvfIndex::repr,
"Return a one-line summary of configuration, counts and build state."
)
.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"
Expand Down
15 changes: 15 additions & 0 deletions python_bindings/symqg_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,16 @@ class SymqgIndex {
return wrapper;
}

[[nodiscard]] std::string repr() const {
const auto operation = access_.read();
return "SymqgIndex(dim=" + std::to_string(dim_) + ", metric='" +
metric_to_string(metric_) +
"', quantization_bits=" + std::to_string(quantization_bits_) +
", num_points=" + std::to_string(index_ ? index_->num_points() : 0) +
", max_degree=" + std::to_string(max_degree_) +
", is_built=" + (built_ ? "True" : "False") + ")";
}

[[nodiscard]] size_t dim() const {
const auto operation = access_.read();
return dim_;
Expand Down Expand Up @@ -383,6 +393,11 @@ Returns
index : SymqgIndex
A built index ready for search.)doc"
)
.def(
"__repr__",
&SymqgIndex::repr,
"Return a one-line summary of configuration, counts and build state."
)
.def_property_readonly(
"dim", &SymqgIndex::dim, R"doc(Original vector dimension.)doc"
)
Expand Down
4 changes: 4 additions & 0 deletions tests/python/test_concurrent_indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,12 +105,14 @@ def __array__(self, dtype=None, copy=None):
def test_active_search_allows_reads_and_rejects_mutations(kind, tmp_path):
index, data, build = make_index(kind)
queries = PausedArray(data[:3])
summary = repr(index)
options = parameters(kind, True)
with ThreadPoolExecutor(max_workers=1) as pool:
running = pool.submit(index.search, queries, 5, **options)
try:
assert queries.entered.wait(timeout=10)
assert index.dim == 65
assert repr(index) == summary
index.search(data[:1], 5, **options)
index.save(str(tmp_path / "reading.index"))
mutations = [build]
Expand Down Expand Up @@ -145,6 +147,8 @@ def test_active_mutation_rejects_reads_and_other_mutations(kind, tmp_path):
lambda: index.save(str(tmp_path / "busy.index")),
lambda: index.dim,
lambda: index.is_built,
lambda: repr(index),
lambda: str(index),
build,
]
if kind != "qg":
Expand Down
103 changes: 103 additions & 0 deletions tests/python/test_index_repr.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
"""Concise metadata summaries across index lifecycle and count changes."""

import numpy as np
import pytest
from rabitqlib import HnswIndex, IvfIndex, SymqgIndex

SUMMARIES = {
"ivf": (
"IvfIndex(dim=64, metric='ip', nbits={bits}, num_points={live}, "
"num_stored_points={stored}, max_elements={capacity}, "
"num_clusters={clusters}, initializer='flat', is_built={built})"
),
"hnsw": (
"HnswIndex(dim=64, metric='ip', nbits={bits}, num_points={live}, "
"num_stored_points={stored}, max_elements={capacity}, "
"num_clusters={clusters}, is_built={built})"
),
"symqg": (
"SymqgIndex(dim=64, metric='ip', quantization_bits={bits}, "
"num_points={live}, max_degree=32, is_built={built})"
),
}


@pytest.fixture
def vectors():
return np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32)


def unbuilt_index(kind, bits):
if kind == "ivf":
return IvfIndex(64, 80, 1, bits, metric="innerproduct")
if kind == "hnsw":
return HnswIndex(64, 96, nbits=bits, metric="innerproduct")
return SymqgIndex(64, 32, metric="innerproduct", quantization_bits=bits)


def build_index(index, vectors):
if isinstance(index, SymqgIndex):
index.build(vectors, 40)
else:
index.build(
vectors, vectors.mean(axis=0, keepdims=True), np.zeros(80, dtype=np.uint32)
)


@pytest.mark.parametrize(
"kind,bits", [("ivf", 4), ("ivf", 32), ("hnsw", 8), ("symqg", 0), ("symqg", 8)]
)
def test_repr_lifecycle(kind, bits, vectors, tmp_path):
index = unbuilt_index(kind, bits)
capacity = 96 if kind == "hnsw" else 80
before = SUMMARIES[kind].format(
bits=bits,
live=0,
stored=0,
capacity=capacity,
clusters=0 if kind == "hnsw" else 1,
built=False,
)
assert repr(index) == str(index) == before

build_index(index, vectors)
after = SUMMARIES[kind].format(
bits=bits, live=80, stored=80, capacity=capacity, clusters=1, built=True
)
reference = index.search(vectors[:2], 3, 32)
assert repr(index) == str(index) == after
assert repr(index) == after
result = index.search(vectors[:2], 3, 32)
np.testing.assert_array_equal(result[0], reference[0])
np.testing.assert_array_equal(result[1], reference[1])

path = tmp_path / "summary.index"
index.save(path)
restored = type(index).load(path)
assert repr(restored) == str(restored) == after


@pytest.mark.parametrize("kind", ["ivf", "hnsw"])
def test_repr_tracks_live_stored_and_capacity(kind, vectors, tmp_path):
index = unbuilt_index(kind, 4)
assert isinstance(index, (IvfIndex, HnswIndex))
build_index(index, vectors)
assert index.remove([0]) == 1
capacity = 96 if kind == "hnsw" else 80
assert repr(index) == SUMMARIES[kind].format(
bits=4, live=79, stored=80, capacity=capacity, clusters=1, built=True
)

index.add(vectors[:1])
if isinstance(index, HnswIndex):
index.resize(128)
capacity = 128
else:
capacity = 81
expected = SUMMARIES[kind].format(
bits=4, live=80, stored=81, capacity=capacity, clusters=1, built=True
)
assert repr(index) == expected
path = tmp_path / "updated.index"
index.save(path)
assert repr(type(index).load(path)) == expected
2 changes: 2 additions & 0 deletions tests/python/typing/public_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None:
ivf = IvfIndex(64, 64, 2, 4)
hnsw = HnswIndex(64, 128)
graph = SymqgIndex(64, 32)
for index in (ivf, hnsw, graph):
assert_type(index.__repr__(), str)
for index in (ivf, hnsw):
index.build(x, x[:2], labels)
assert_type(index.add(x, labels), npt.NDArray[np.uint32])
Expand Down
Loading