From 5c4524164ecd48418890f5a903fb3ca4deb3af5e Mon Sep 17 00:00:00 2001 From: Tim VanReenen Date: Fri, 21 Aug 2026 11:57:20 -0400 Subject: [PATCH 1/5] Validate row selectors before search shortcuts --- src/numpy_vector_store/vector_store.py | 48 ++++++++++++++-------- tests/test_vector_store.py | 55 +++++++++++++++++++------- 2 files changed, 73 insertions(+), 30 deletions(-) diff --git a/src/numpy_vector_store/vector_store.py b/src/numpy_vector_store/vector_store.py index c9945c6..cdad30c 100644 --- a/src/numpy_vector_store/vector_store.py +++ b/src/numpy_vector_store/vector_store.py @@ -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. @@ -200,7 +202,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. @@ -224,7 +228,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. @@ -497,27 +503,35 @@ 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) + rows = np.asarray(within_rows, dtype=object) 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") + if any(np.asarray(row).ndim != 0 for row in rows): + raise ValueError("within_rows must be a 1D sequence of row indexes") - rows = rows.astype(np.intp, copy=False) - if np.any(rows < 0) or np.any(rows >= self._row_count): + 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 + + 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], @@ -534,13 +548,15 @@ 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 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] diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 9806a6d..2def2e5 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -326,20 +326,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) @@ -670,6 +656,47 @@ 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], 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)]) + @pytest.mark.parametrize( "search_name", ["cosine_search", "dot_search", "euclidean_search"] ) From e655a1368ae3805cdc5d992ae4a182fcf1f5ce79 Mon Sep 17 00:00:00 2001 From: Tim VanReenen Date: Fri, 21 Aug 2026 11:58:24 -0400 Subject: [PATCH 2/5] Apply zero-query rules before empty results --- src/numpy_vector_store/vector_store.py | 15 ++++++---- tests/test_vector_store.py | 38 ++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/src/numpy_vector_store/vector_store.py b/src/numpy_vector_store/vector_store.py index cdad30c..7cadcac 100644 --- a/src/numpy_vector_store/vector_store.py +++ b/src/numpy_vector_store/vector_store.py @@ -191,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, @@ -217,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, @@ -244,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, @@ -253,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): @@ -265,7 +266,7 @@ def _cosine_values( np.einsum( "ij,j->i", vectors, - query_norm, + query, dtype=np.float64, ) / vector_norms @@ -278,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( @@ -294,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) @@ -536,6 +535,7 @@ def _metric_search( [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, @@ -552,6 +552,9 @@ def _metric_search( 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 [] diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 2def2e5..bcd94b8 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -600,6 +600,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, @@ -697,6 +698,43 @@ def test_metric_searches_reject_out_of_bounds_rows_independent_of_state( with pytest.raises(IndexError, match="outside"): getattr(store, search_name)([1.0, 0.0], within_rows=[len(store)]) + @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"] ) From f1e414128f372fed4c19a46b05af057b474a2854 Mon Sep 17 00:00:00 2001 From: Tim VanReenen Date: Fri, 21 Aug 2026 11:59:35 -0400 Subject: [PATCH 3/5] Document search selector and zero-query rules --- README.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/README.md b/README.md index 61e749f..2dfb383 100644 --- a/README.md +++ b/README.md @@ -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 @@ -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 @@ -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 From 11eecbaf94767035272484d66de5faf1a3f51ba0 Mon Sep 17 00:00:00 2001 From: Tim VanReenen Date: Fri, 21 Aug 2026 13:43:47 -0400 Subject: [PATCH 4/5] Keep native row selector validation vectorized --- src/numpy_vector_store/vector_store.py | 24 ++++++++++++- tests/test_vector_store.py | 49 ++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/src/numpy_vector_store/vector_store.py b/src/numpy_vector_store/vector_store.py index 7cadcac..bdbe98c 100644 --- a/src/numpy_vector_store/vector_store.py +++ b/src/numpy_vector_store/vector_store.py @@ -3,6 +3,7 @@ import operator import os import stat +import warnings from collections.abc import Callable, Sequence from dataclasses import dataclass from pathlib import Path @@ -505,9 +506,30 @@ def _normalize_within_rows( self, within_rows: Sequence[SupportsIndex] | npt.NDArray[np.integer[Any]], ) -> npt.NDArray[np.intp]: - rows = np.asarray(within_rows, dtype=object) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="Creating an ndarray from ragged nested sequences", + ) + try: + rows = np.asarray(within_rows) + except ValueError: + rows = np.asarray(within_rows, dtype=object) + if rows.ndim != 1: raise ValueError("within_rows must be a 1D sequence of row indexes") + if len(rows) == 0: + 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") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index bcd94b8..88a6239 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -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 {}]) @@ -698,6 +708,45 @@ def test_metric_searches_reject_out_of_bounds_rows_independent_of_state( 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"), [ From 2dc64dd7cb015251f390724a5774230b87d94009 Mon Sep 17 00:00:00 2001 From: Tim VanReenen Date: Fri, 21 Aug 2026 13:50:49 -0400 Subject: [PATCH 5/5] Avoid global warning state during row validation --- src/numpy_vector_store/vector_store.py | 18 +++++++++--------- tests/test_vector_store.py | 1 + 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/numpy_vector_store/vector_store.py b/src/numpy_vector_store/vector_store.py index bdbe98c..c1a8a1b 100644 --- a/src/numpy_vector_store/vector_store.py +++ b/src/numpy_vector_store/vector_store.py @@ -3,7 +3,6 @@ import operator import os import stat -import warnings from collections.abc import Callable, Sequence from dataclasses import dataclass from pathlib import Path @@ -506,15 +505,16 @@ def _normalize_within_rows( self, within_rows: Sequence[SupportsIndex] | npt.NDArray[np.integer[Any]], ) -> npt.NDArray[np.intp]: - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message="Creating an ndarray from ragged nested sequences", + 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, ) - try: - rows = np.asarray(within_rows) - except ValueError: - rows = np.asarray(within_rows, dtype=object) if rows.ndim != 1: raise ValueError("within_rows must be a 1D sequence of row indexes") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 88a6239..5140390 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -675,6 +675,7 @@ def test_metric_searches_validate_common_inputs(self): [ ([[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"),