Skip to content

Make search preconditions independent of store state #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tvanreenen merged 5 commits into from
Aug 21, 2026
Merged

Make search preconditions independent of store state #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5c45241
Validate row selectors before search shortcuts
tvanreenen Aug 21, 2026
e655a13
Apply zero-query rules before empty results
tvanreenen Aug 21, 2026
f1e4141
Document search selector and zero-query rules
tvanreenen Aug 21, 2026
11eecba
Keep native row selector validation vectorized
tvanreenen Aug 21, 2026
2dc64dd
Avoid global warning state during row validation
tvanreenen Aug 21, 2026
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
15 changes: 15 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,12 @@ integer or floating-point values and NumPy integer or floating scalars.
Booleans, strings, complex numbers, and arrays are not threshold scalars and are
rejected.

`within_rows` must be a one-dimensional sequence of unique integer row indexes.
Python and NumPy integers are accepted; booleans and non-integer values are not.
Malformed shapes and duplicate indexes raise `ValueError`, while an index
outside the current store raises `IndexError`. These checks still run when the
store is empty.

For these scalar inputs, an inappropriate type raises `TypeError` and a
supported type with an invalid value raises `ValueError`. Row selectors outside
the store raise `IndexError`, and filesystem operations continue to raise the
Expand Down Expand Up @@ -166,6 +172,11 @@ search. Because cosine similarity is undefined for zero vectors,
`cosine_search` raises an error when its selected rows include one; use
`within_rows` to exclude zero rows when needed.

Cosine search also requires a non-zero query. Dot-product and Euclidean searches
require one when `normalize=True`, because they normalize the query before
comparison. With `normalize=False`, both methods accept a zero query. These
rules apply even when the store or `within_rows` selection is empty.

### Numerical inputs

Stored vectors use `float32` to keep the store compact. Vectors and queries must
Expand Down Expand Up @@ -240,6 +251,10 @@ rows = [
hits = store.cosine_search(query, top_k=10, within_rows=rows)
```

Each stored row may appear at most once in `within_rows`; duplicate indexes are
rejected rather than producing duplicate hits. An empty sequence returns no
hits, but it does not bypass validation of the query or other search arguments.

Searches without `within_rows` compute directly against the stored vector matrix
and do not make a full copy of it. A filtered search gathers the selected rows
into a temporary matrix, so its additional memory use scales with the number of
Expand Down
83 changes: 62 additions & 21 deletions src/numpy_vector_store/vector_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,9 @@ def cosine_search(
*,
top_k: SupportsIndex = 10,
min_value: _RealScalar | None = None,
within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None = None,
within_rows: Sequence[SupportsIndex]
| npt.NDArray[np.integer[Any]]
| None = None,
) -> list[VectorHit[TMetadata]]:
"""
Return the most similar rows using cosine similarity.
Expand All @@ -189,6 +191,7 @@ def cosine_search(
top_k=top_k,
within_rows=within_rows,
values_fn=self._cosine_values,
normalize_query=True,
descending=True,
min_value=min_value,
max_value=None,
Expand All @@ -200,7 +203,9 @@ def dot_search(
*,
top_k: SupportsIndex = 10,
min_value: _RealScalar | None = None,
within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None = None,
within_rows: Sequence[SupportsIndex]
| npt.NDArray[np.integer[Any]]
| None = None,
) -> list[VectorHit[TMetadata]]:
"""
Return rows ranked by dot product.
Expand All @@ -213,6 +218,7 @@ def dot_search(
top_k=top_k,
within_rows=within_rows,
values_fn=self._dot_values,
normalize_query=self._normalize,
descending=True,
min_value=min_value,
max_value=None,
Expand All @@ -224,7 +230,9 @@ def euclidean_search(
*,
top_k: SupportsIndex = 10,
max_value: _RealScalar | None = None,
within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None = None,
within_rows: Sequence[SupportsIndex]
| npt.NDArray[np.integer[Any]]
| None = None,
) -> list[VectorHit[TMetadata]]:
"""
Return rows ranked by Euclidean distance.
Expand All @@ -238,6 +246,7 @@ def euclidean_search(
top_k=top_k,
within_rows=within_rows,
values_fn=self._euclidean_values,
normalize_query=self._normalize,
descending=False,
min_value=None,
max_value=max_value,
Expand All @@ -247,10 +256,8 @@ def _cosine_values(
self, query: npt.NDArray[np.float32], vectors: npt.NDArray[np.float32]
) -> npt.NDArray[np.float32]:
"""Compute cosine similarity values."""
query_norm = self._normalize_query(query)

if self._normalize:
values = np.dot(vectors, query_norm)
values = np.dot(vectors, query)
else:
vector_norms = self._row_norms(vectors)
if np.any(vector_norms == 0):
Expand All @@ -259,7 +266,7 @@ def _cosine_values(
np.einsum(
"ij,j->i",
vectors,
query_norm,
query,
dtype=np.float64,
)
/ vector_norms
Expand All @@ -272,7 +279,6 @@ def _dot_values(
) -> npt.NDArray[np.float64]:
"""Compute dot product values."""
if self._normalize:
query = self._normalize_query(query)
values = np.dot(vectors, query)
else:
values = np.einsum(
Expand All @@ -288,7 +294,6 @@ def _euclidean_values(
) -> npt.NDArray[np.float64]:
"""Compute Euclidean distance values."""
if self._normalize:
query = self._normalize_query(query)
differences = vectors - query
else:
differences = np.empty(vectors.shape, dtype=np.float64)
Expand Down Expand Up @@ -497,31 +502,62 @@ def _normalize_query(
)

def _normalize_within_rows(
self, within_rows: Sequence[int] | npt.NDArray[np.integer[Any]]
self,
within_rows: Sequence[SupportsIndex] | npt.NDArray[np.integer[Any]],
) -> npt.NDArray[np.intp]:
rows = np.asarray(within_rows)
if isinstance(within_rows, np.ndarray):
rows = np.asarray(within_rows)
else:
has_nested_rows = any(
isinstance(row, (list, tuple, np.ndarray)) for row in within_rows
)
rows = np.asarray(
within_rows,
dtype=object if has_nested_rows else None,
)

if rows.ndim != 1:
raise ValueError("within_rows must be a 1D sequence of row indexes")
if len(rows) == 0:
return np.array([], dtype=np.intp)
if not np.issubdtype(rows.dtype, np.integer):
raise ValueError("within_rows must contain integer row indexes")
return np.empty(0, dtype=np.intp)

if np.issubdtype(rows.dtype, np.integer):
if len(np.unique(rows)) != len(rows):
raise ValueError("within_rows must contain unique row indexes")
if np.any(rows < 0) or np.any(rows >= self._row_count):
raise IndexError("within_rows contains row indexes outside the store")
return rows.astype(np.intp, copy=False)

if rows.dtype != np.dtype(object):
raise TypeError("within_rows must contain integer row indexes")
if any(np.asarray(row).ndim != 0 for row in rows):
raise ValueError("within_rows must be a 1D sequence of row indexes")

try:
normalized_rows = [
self._validate_integer(row, name="within_rows row index")
for row in rows
]
except TypeError:
raise TypeError("within_rows must contain integer row indexes") from None

rows = rows.astype(np.intp, copy=False)
if np.any(rows < 0) or np.any(rows >= self._row_count):
if len(set(normalized_rows)) != len(normalized_rows):
raise ValueError("within_rows must contain unique row indexes")
if any(row < 0 or row >= self._row_count for row in normalized_rows):
raise IndexError("within_rows contains row indexes outside the store")
return rows
return np.asarray(normalized_rows, dtype=np.intp)

def _metric_search(
self,
query: npt.ArrayLike,
*,
top_k: SupportsIndex,
within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None,
within_rows: Sequence[SupportsIndex] | npt.NDArray[np.integer[Any]] | None,
values_fn: Callable[
[npt.NDArray[np.float32], npt.NDArray[np.float32]],
npt.NDArray[np.float32] | npt.NDArray[np.float64],
],
normalize_query: bool,
descending: bool,
min_value: _RealScalar | None,
max_value: _RealScalar | None,
Expand All @@ -534,13 +570,18 @@ def _metric_search(
min_value = self._validate_search_threshold(min_value, name="min_value")
max_value = self._validate_search_threshold(max_value, name="max_value")

row_indices = None
if within_rows is not None:
row_indices = self._normalize_within_rows(within_rows)

if normalize_query:
query_vector = self._normalize_query(query_vector)

if self._row_count == 0:
return []

row_indices = None
selected_vectors = self._vectors[: self._row_count]
if within_rows is not None:
row_indices = self._normalize_within_rows(within_rows)
if row_indices is not None:
if len(row_indices) == 0:
return []
selected_vectors = self._vectors[row_indices]
Expand Down
143 changes: 129 additions & 14 deletions tests/test_vector_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,16 @@ class MetadataRecord:
label: str


@dataclass
class IndexValue:
"""Custom integer-index value used to exercise selector fallback behavior."""

value: int

def __index__(self):
return self.value


def add_single_vector(store, vector, metadata=None):
"""Helper function to add a single vector."""
store.add(np.atleast_2d(vector), [metadata or {}])
Expand Down Expand Up @@ -326,20 +336,6 @@ def test_cosine_search_empty_within_rows(self):

assert store.cosine_search([1.0, 0.0, 0.0], within_rows=[]) == []

def test_cosine_search_rejects_invalid_within_rows(self):
"""Test within_rows validates shape, dtype, and bounds."""
store = VectorStore[dict[str, str]](dimensions=3)
store.add([[1.0, 0.0, 0.0]], [{"id": "x"}])

with pytest.raises(ValueError, match="1D"):
store.cosine_search([1.0, 0.0, 0.0], within_rows=[[0]])

with pytest.raises(ValueError, match="integer"):
store.cosine_search([1.0, 0.0, 0.0], within_rows=[0.0])

with pytest.raises(IndexError, match="outside"):
store.cosine_search([1.0, 0.0, 0.0], within_rows=[1])

def test_cosine_search_rejects_wrong_query_dimensions(self):
"""Test cosine_search rejects wrong query dimensions."""
store = VectorStore(dimensions=3)
Expand Down Expand Up @@ -614,6 +610,7 @@ def capture_vectors(query, vectors):
top_k=1,
within_rows=None,
values_fn=capture_vectors,
normalize_query=False,
descending=True,
min_value=None,
max_value=None,
Expand Down Expand Up @@ -670,6 +667,124 @@ def test_metric_searches_validate_common_inputs(self):
)
assert [hit.index for hit in results] == [0]

@pytest.mark.parametrize(
"search_name", ["cosine_search", "dot_search", "euclidean_search"]
)
@pytest.mark.parametrize(
("within_rows", "exception", "message"),
[
([[0]], ValueError, "1D"),
([[0], [0, 1]], ValueError, "1D"),
([0, [0, 1]], ValueError, "1D"),
([0.0], TypeError, "integer"),
([True], TypeError, "integer"),
([np.bool_(True)], TypeError, "integer"),
([0, 0], ValueError, "unique"),
],
)
@pytest.mark.parametrize("populated", [False, True])
def test_metric_searches_reject_invalid_row_selectors_independent_of_state(
self, search_name, within_rows, exception, message, populated
):
"""Test malformed and duplicate row selectors never bypass validation."""
store = VectorStore(dimensions=2)
if populated:
store.add([[1.0, 0.0]], [{"id": "x"}])

with pytest.raises(exception, match=message):
getattr(store, search_name)([1.0, 0.0], within_rows=within_rows)

@pytest.mark.parametrize(
"search_name", ["cosine_search", "dot_search", "euclidean_search"]
)
@pytest.mark.parametrize("populated", [False, True])
def test_metric_searches_reject_out_of_bounds_rows_independent_of_state(
self, search_name, populated
):
"""Test row bounds are checked before an empty search can return."""
store = VectorStore(dimensions=2)
if populated:
store.add([[1.0, 0.0]], [{"id": "x"}])

with pytest.raises(IndexError, match="outside"):
getattr(store, search_name)([1.0, 0.0], within_rows=[len(store)])

def test_native_integer_row_selectors_remain_zero_copy(self):
"""Test common NumPy selectors stay on the vectorized validation path."""
store = VectorStore(dimensions=2)
store.add([[1.0, 0.0], [0.0, 1.0]], ["x", "y"])
rows = np.array([1, 0], dtype=np.intp)

normalized_rows = store._normalize_within_rows(rows)

assert np.shares_memory(normalized_rows, rows)

def test_row_selectors_accept_generic_integer_index_values(self):
"""Test object-valued selectors use the integer-index fallback."""
store = VectorStore(dimensions=2)
store.add([[1.0, 0.0], [0.0, 1.0]], ["x", "y"])

results = store.cosine_search(
[1.0, 0.0], within_rows=[IndexValue(0), IndexValue(1)]
)

assert [hit.index for hit in results] == [0, 1]

@pytest.mark.parametrize(
("within_rows", "exception", "message"),
[
([IndexValue(0), object()], TypeError, "integer"),
([IndexValue(0), IndexValue(0)], ValueError, "unique"),
([IndexValue(2)], IndexError, "outside"),
],
)
def test_generic_integer_row_selectors_preserve_failure_contracts(
self, within_rows, exception, message
):
"""Test fallback validation matches native selector failures."""
store = VectorStore(dimensions=2)
store.add([[1.0, 0.0], [0.0, 1.0]], ["x", "y"])

with pytest.raises(exception, match=message):
store.cosine_search([1.0, 0.0], within_rows=within_rows)

@pytest.mark.parametrize(
("search_name", "normalize"),
[
("cosine_search", True),
("cosine_search", False),
("dot_search", True),
("euclidean_search", True),
],
)
@pytest.mark.parametrize("within_rows", [None, []])
@pytest.mark.parametrize("populated", [False, True])
def test_metric_searches_reject_zero_queries_independent_of_state(
self, search_name, normalize, within_rows, populated
):
"""Test query normalization cannot be bypassed by an empty search."""
store = VectorStore(dimensions=2, normalize=normalize)
if populated:
store.add([[1.0, 0.0]], [{"id": "x"}])

with pytest.raises(ValueError, match="zero-norm"):
getattr(store, search_name)([0.0, 0.0], within_rows=within_rows)

@pytest.mark.parametrize("search_name", ["dot_search", "euclidean_search"])
@pytest.mark.parametrize("within_rows", [None, []])
@pytest.mark.parametrize("populated", [False, True])
def test_raw_metric_searches_accept_zero_queries_independent_of_state(
self, search_name, within_rows, populated
):
"""Test raw dot and Euclidean searches keep zero-query semantics."""
store = VectorStore(dimensions=2, normalize=False)
if populated:
store.add([[1.0, 0.0]], [{"id": "x"}])

results = getattr(store, search_name)([0.0, 0.0], within_rows=within_rows)

assert len(results) == (1 if populated and within_rows is None else 0)

@pytest.mark.parametrize(
"search_name", ["cosine_search", "dot_search", "euclidean_search"]
)
Expand Down