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 diff --git a/src/numpy_vector_store/vector_store.py b/src/numpy_vector_store/vector_store.py index c9945c6..c1a8a1b 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. @@ -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, @@ -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. @@ -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, @@ -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. @@ -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, @@ -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): @@ -259,7 +266,7 @@ def _cosine_values( np.einsum( "ij,j->i", vectors, - query_norm, + query, dtype=np.float64, ) / vector_norms @@ -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( @@ -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) @@ -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, @@ -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] diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 9806a6d..5140390 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 {}]) @@ -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) @@ -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, @@ -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"] )