numpy-vector-store 0.2.0__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: numpy-vector-store
3
- Version: 0.2.0
3
+ Version: 0.3.0
4
4
  Summary: A fast, lightweight, and zero-setup in-memory vector store powered by NumPy
5
5
  Project-URL: Homepage, https://github.com/tvanreenen/numpy-vector-store
6
6
  Project-URL: Repository, https://github.com/tvanreenen/numpy-vector-store
@@ -27,7 +27,7 @@ Description-Content-Type: text/markdown
27
27
  A fast, lightweight, zero-setup in-memory vector store powered by NumPy.
28
28
 
29
29
  - **Tiny local vector search** for projects that do not need a vector database
30
- - **Fast exact cosine search** using vectorized NumPy operations
30
+ - **Fast exact vector search** using vectorized NumPy operations
31
31
  - **Simple typed API** returning `VectorHit(index, value, metadata)`
32
32
  - **Composable filtering** by passing prefiltered row indexes with `within_rows`
33
33
  - **Portable persistence** as trusted local `.npz` files with `vectors` + `metadata`
@@ -93,6 +93,58 @@ for hit in hits:
93
93
  dataclass, string, integer row ID, or any other Python object that fits your
94
94
  application.
95
95
 
96
+ ## Normalization
97
+
98
+ `VectorStore` defaults to `normalize=True`, which scales each stored vector to
99
+ length `1`. Normalization preserves vector direction while discarding magnitude:
100
+
101
+ ```python
102
+ [3.0, 4.0] -> [0.6, 0.8]
103
+ ```
104
+
105
+ This is the default because it makes cosine similarity fast and direction-only,
106
+ which is the common case for semantic embeddings. Use `normalize=False` when
107
+ vector length matters, such as when magnitude encodes strength, confidence,
108
+ counts, scale, or raw geometry.
109
+
110
+ Zero vectors are rejected in both modes because cosine similarity is undefined
111
+ for zero-norm vectors.
112
+
113
+ | Method | `normalize=True` default | `normalize=False` |
114
+ |---|---|---|
115
+ | `cosine_search` | True cosine similarity over stored unit vectors; fastest/default path for embeddings | True cosine similarity over raw vectors; computes vector norms during search |
116
+ | `dot_search` | Dot product of unit vectors, effectively equivalent to cosine similarity | True dot product over original vectors; use when magnitude should affect ranking |
117
+ | `euclidean_search` | Distance between normalized directions; useful only when direction-normalized distance is intended | True Euclidean distance over original vectors; use for geometric/feature-space nearest neighbors |
118
+ | `get` | Returns normalized vectors | Returns original vectors |
119
+ | `save` | Saves normalized vectors | Saves raw vectors |
120
+ | `load` | Loads and normalizes vectors | Loads vectors exactly as stored |
121
+
122
+ ## Search Methods
123
+
124
+ Use `cosine_search` for semantic embeddings and direction-only similarity:
125
+
126
+ ```python
127
+ hits = store.cosine_search(query, top_k=10, min_value=0.75)
128
+ ```
129
+
130
+ Use `dot_search` with `normalize=False` when larger-magnitude vectors should
131
+ rank higher:
132
+
133
+ ```python
134
+ store = VectorStore[dict[str, str]](dimensions=3, normalize=False)
135
+ store.add(vectors, metadata)
136
+ hits = store.dot_search(query, top_k=10, min_value=0.0)
137
+ ```
138
+
139
+ Use `euclidean_search` with `normalize=False` for raw coordinate or feature-space
140
+ nearest-neighbor search:
141
+
142
+ ```python
143
+ store = VectorStore[dict[str, str]](dimensions=3, normalize=False)
144
+ store.add(vectors, metadata)
145
+ hits = store.euclidean_search(query, top_k=10, max_value=1.5)
146
+ ```
147
+
96
148
  ## Prefiltering
97
149
 
98
150
  The store does not implement a metadata query language. To filter by metadata,
@@ -146,39 +198,45 @@ loaded = VectorStore[dict[str, str]](dimensions=1536, file_path="vectors.npz")
146
198
  loaded.load()
147
199
  ```
148
200
 
149
- Context manager usage auto-saves on exit:
201
+ If you save with `normalize=False`, load with `normalize=False` too:
150
202
 
151
203
  ```python
152
- with VectorStore[dict[str, str]](dimensions=1536, file_path="vectors.npz") as store:
153
- store.add(embeddings, metadata)
154
- ```
155
-
156
- Persistence uses a minimal NumPy `.npz` contract with `vectors` and `metadata`
157
- arrays. Vectors are normalized when added or loaded, and similarity search only
158
- normalizes the query vector. Loading validates shape, dimensions, row counts, and
159
- zero-norm vectors. It also uses `allow_pickle=True` for flexible Python metadata
160
- payloads, so only load files generated by your own application or another trusted
161
- local process. Loading untrusted `.npz` files is not a supported security model.
162
-
163
- ## Migrating from 0.1
204
+ store = VectorStore[dict[str, str]](
205
+ dimensions=1536,
206
+ file_path="raw-vectors.npz",
207
+ normalize=False,
208
+ )
209
+ store.add(raw_vectors, metadata)
210
+ store.save()
164
211
 
165
- The preferred 0.2 API is `add(...)` and `cosine_search(...)`.
212
+ loaded = VectorStore[dict[str, str]](
213
+ dimensions=1536,
214
+ file_path="raw-vectors.npz",
215
+ normalize=False,
216
+ )
217
+ loaded.load()
218
+ ```
166
219
 
167
- The 0.1 methods remain temporarily available, but emit `DeprecationWarning` and
168
- will be removed in a future 0.x release:
220
+ Context manager usage auto-saves on exit:
169
221
 
170
222
  ```python
171
- store.add_vectors(vectors_2d, metadata_array)
172
- results = store.search(query, top_k=3, score_cutoff=0.5)
173
- for index, value, metadata in results:
174
- ...
223
+ with VectorStore[dict[str, str]](dimensions=1536, file_path="vectors.npz") as store:
224
+ store.add(embeddings, metadata)
175
225
  ```
176
226
 
177
- Legacy `search(...)` keeps returning tuples. New `cosine_search(...)` returns
178
- `VectorHit` objects.
179
-
180
- `metadata_schema` was removed. For vectorized metadata filtering, keep metadata
181
- in a sidecar table and pass matching row indexes with `within_rows`.
227
+ Persistence uses a minimal NumPy `.npz` contract with `vectors` and `metadata`
228
+ arrays. The `.npz` file does not encode the `normalize` setting; choose the same
229
+ setting when loading that you used when saving. Loading validates shape,
230
+ dimensions, row counts, and zero-norm vectors. It also uses `allow_pickle=True`
231
+ for flexible Python metadata payloads, so only load files generated by your own
232
+ application or another trusted local process. Loading untrusted `.npz` files is
233
+ not a supported security model.
234
+
235
+ ## Compatibility
236
+
237
+ This project is still pre-1.0, so occasional breaking changes are expected while
238
+ the API stabilizes. Breaking changes are documented in GitHub release notes.
239
+ Deprecated APIs will keep warning for at least one point release before removal.
182
240
 
183
241
  ## Contributing
184
242
 
@@ -3,7 +3,7 @@
3
3
  A fast, lightweight, zero-setup in-memory vector store powered by NumPy.
4
4
 
5
5
  - **Tiny local vector search** for projects that do not need a vector database
6
- - **Fast exact cosine search** using vectorized NumPy operations
6
+ - **Fast exact vector search** using vectorized NumPy operations
7
7
  - **Simple typed API** returning `VectorHit(index, value, metadata)`
8
8
  - **Composable filtering** by passing prefiltered row indexes with `within_rows`
9
9
  - **Portable persistence** as trusted local `.npz` files with `vectors` + `metadata`
@@ -69,6 +69,58 @@ for hit in hits:
69
69
  dataclass, string, integer row ID, or any other Python object that fits your
70
70
  application.
71
71
 
72
+ ## Normalization
73
+
74
+ `VectorStore` defaults to `normalize=True`, which scales each stored vector to
75
+ length `1`. Normalization preserves vector direction while discarding magnitude:
76
+
77
+ ```python
78
+ [3.0, 4.0] -> [0.6, 0.8]
79
+ ```
80
+
81
+ This is the default because it makes cosine similarity fast and direction-only,
82
+ which is the common case for semantic embeddings. Use `normalize=False` when
83
+ vector length matters, such as when magnitude encodes strength, confidence,
84
+ counts, scale, or raw geometry.
85
+
86
+ Zero vectors are rejected in both modes because cosine similarity is undefined
87
+ for zero-norm vectors.
88
+
89
+ | Method | `normalize=True` default | `normalize=False` |
90
+ |---|---|---|
91
+ | `cosine_search` | True cosine similarity over stored unit vectors; fastest/default path for embeddings | True cosine similarity over raw vectors; computes vector norms during search |
92
+ | `dot_search` | Dot product of unit vectors, effectively equivalent to cosine similarity | True dot product over original vectors; use when magnitude should affect ranking |
93
+ | `euclidean_search` | Distance between normalized directions; useful only when direction-normalized distance is intended | True Euclidean distance over original vectors; use for geometric/feature-space nearest neighbors |
94
+ | `get` | Returns normalized vectors | Returns original vectors |
95
+ | `save` | Saves normalized vectors | Saves raw vectors |
96
+ | `load` | Loads and normalizes vectors | Loads vectors exactly as stored |
97
+
98
+ ## Search Methods
99
+
100
+ Use `cosine_search` for semantic embeddings and direction-only similarity:
101
+
102
+ ```python
103
+ hits = store.cosine_search(query, top_k=10, min_value=0.75)
104
+ ```
105
+
106
+ Use `dot_search` with `normalize=False` when larger-magnitude vectors should
107
+ rank higher:
108
+
109
+ ```python
110
+ store = VectorStore[dict[str, str]](dimensions=3, normalize=False)
111
+ store.add(vectors, metadata)
112
+ hits = store.dot_search(query, top_k=10, min_value=0.0)
113
+ ```
114
+
115
+ Use `euclidean_search` with `normalize=False` for raw coordinate or feature-space
116
+ nearest-neighbor search:
117
+
118
+ ```python
119
+ store = VectorStore[dict[str, str]](dimensions=3, normalize=False)
120
+ store.add(vectors, metadata)
121
+ hits = store.euclidean_search(query, top_k=10, max_value=1.5)
122
+ ```
123
+
72
124
  ## Prefiltering
73
125
 
74
126
  The store does not implement a metadata query language. To filter by metadata,
@@ -122,39 +174,45 @@ loaded = VectorStore[dict[str, str]](dimensions=1536, file_path="vectors.npz")
122
174
  loaded.load()
123
175
  ```
124
176
 
125
- Context manager usage auto-saves on exit:
177
+ If you save with `normalize=False`, load with `normalize=False` too:
126
178
 
127
179
  ```python
128
- with VectorStore[dict[str, str]](dimensions=1536, file_path="vectors.npz") as store:
129
- store.add(embeddings, metadata)
130
- ```
131
-
132
- Persistence uses a minimal NumPy `.npz` contract with `vectors` and `metadata`
133
- arrays. Vectors are normalized when added or loaded, and similarity search only
134
- normalizes the query vector. Loading validates shape, dimensions, row counts, and
135
- zero-norm vectors. It also uses `allow_pickle=True` for flexible Python metadata
136
- payloads, so only load files generated by your own application or another trusted
137
- local process. Loading untrusted `.npz` files is not a supported security model.
138
-
139
- ## Migrating from 0.1
180
+ store = VectorStore[dict[str, str]](
181
+ dimensions=1536,
182
+ file_path="raw-vectors.npz",
183
+ normalize=False,
184
+ )
185
+ store.add(raw_vectors, metadata)
186
+ store.save()
140
187
 
141
- The preferred 0.2 API is `add(...)` and `cosine_search(...)`.
188
+ loaded = VectorStore[dict[str, str]](
189
+ dimensions=1536,
190
+ file_path="raw-vectors.npz",
191
+ normalize=False,
192
+ )
193
+ loaded.load()
194
+ ```
142
195
 
143
- The 0.1 methods remain temporarily available, but emit `DeprecationWarning` and
144
- will be removed in a future 0.x release:
196
+ Context manager usage auto-saves on exit:
145
197
 
146
198
  ```python
147
- store.add_vectors(vectors_2d, metadata_array)
148
- results = store.search(query, top_k=3, score_cutoff=0.5)
149
- for index, value, metadata in results:
150
- ...
199
+ with VectorStore[dict[str, str]](dimensions=1536, file_path="vectors.npz") as store:
200
+ store.add(embeddings, metadata)
151
201
  ```
152
202
 
153
- Legacy `search(...)` keeps returning tuples. New `cosine_search(...)` returns
154
- `VectorHit` objects.
155
-
156
- `metadata_schema` was removed. For vectorized metadata filtering, keep metadata
157
- in a sidecar table and pass matching row indexes with `within_rows`.
203
+ Persistence uses a minimal NumPy `.npz` contract with `vectors` and `metadata`
204
+ arrays. The `.npz` file does not encode the `normalize` setting; choose the same
205
+ setting when loading that you used when saving. Loading validates shape,
206
+ dimensions, row counts, and zero-norm vectors. It also uses `allow_pickle=True`
207
+ for flexible Python metadata payloads, so only load files generated by your own
208
+ application or another trusted local process. Loading untrusted `.npz` files is
209
+ not a supported security model.
210
+
211
+ ## Compatibility
212
+
213
+ This project is still pre-1.0, so occasional breaking changes are expected while
214
+ the API stabilizes. Breaking changes are documented in GitHub release notes.
215
+ Deprecated APIs will keep warning for at least one point release before removal.
158
216
 
159
217
  ## Contributing
160
218
 
@@ -5,7 +5,7 @@ This package provides a lightweight vector store implementation for storing
5
5
  and searching vector embeddings using NumPy arrays.
6
6
  """
7
7
 
8
- __version__ = "0.2.0"
8
+ __version__ = "0.3.0"
9
9
  __author__ = "Tim VanReenen"
10
10
 
11
11
  from .vector_store import VectorHit, VectorStore
@@ -1,7 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
- import warnings
4
- from collections.abc import Sequence
3
+ from collections.abc import Callable, Sequence
5
4
  from dataclasses import dataclass
6
5
  from pathlib import Path
7
6
  from typing import Any, Generic, TypeVar
@@ -23,17 +22,19 @@ class VectorHit(Generic[TMetadata]):
23
22
 
24
23
  class VectorStore(Generic[TMetadata]):
25
24
  """
26
- A small in-memory vector store using NumPy cosine similarity.
25
+ A small in-memory vector store using NumPy exact vector search.
27
26
 
28
27
  Metadata is an opaque row payload returned with vector hits. The store does
29
28
  not implement a metadata query language; pass row indexes to within_rows for
30
- prefiltered similarity search.
29
+ prefiltered vector search.
31
30
  """
32
31
 
33
32
  def __init__(
34
33
  self,
35
34
  dimensions: int,
36
35
  file_path: str | Path | None = None,
36
+ *,
37
+ normalize: bool = True,
37
38
  ) -> None:
38
39
  """
39
40
  Initialize the vector store.
@@ -41,12 +42,14 @@ class VectorStore(Generic[TMetadata]):
41
42
  Args:
42
43
  dimensions: The number of dimensions for vectors to be stored.
43
44
  file_path: Optional path to save/load vectors from.
45
+ normalize: Whether to store vectors normalized to unit length.
44
46
  """
45
47
  if dimensions <= 0:
46
48
  raise ValueError("dimensions must be greater than 0")
47
49
 
48
50
  self.dimensions = dimensions
49
51
  self.file_path = Path(file_path) if file_path else None
52
+ self.normalize = normalize
50
53
  self.vectors: npt.NDArray[np.float32] = np.empty(
51
54
  (0, dimensions), dtype=np.float32
52
55
  )
@@ -59,8 +62,8 @@ class VectorStore(Generic[TMetadata]):
59
62
  """
60
63
  Add vectors and row metadata payloads.
61
64
 
62
- Vectors are normalized at insert time. Metadata items are stored as
63
- opaque payloads and returned with vector hits.
65
+ When normalize=True, vectors are normalized at insert time. Metadata
66
+ items are stored as opaque payloads and returned with vector hits.
64
67
  """
65
68
  vectors_2d = np.asarray(vectors, dtype=np.float32)
66
69
  if vectors_2d.ndim != 2:
@@ -78,40 +81,20 @@ class VectorStore(Generic[TMetadata]):
78
81
  if len(vectors_2d) != len(metadata_array):
79
82
  raise ValueError("Number of vectors must match number of metadata items")
80
83
 
81
- normalized_vectors = self._normalize_vectors(
84
+ vectors_to_store = self._prepare_vectors_for_storage(
82
85
  vectors_2d,
83
86
  error_message="Cannot add zero-norm vectors",
84
87
  )
85
88
 
86
89
  if len(self.vectors) == 0:
87
- self.vectors = normalized_vectors
90
+ self.vectors = vectors_to_store
88
91
  else:
89
- self.vectors = np.vstack([self.vectors, normalized_vectors]).astype(
92
+ self.vectors = np.vstack([self.vectors, vectors_to_store]).astype(
90
93
  np.float32, copy=False
91
94
  )
92
95
 
93
96
  self.metadata = np.append(self.metadata, metadata_array)
94
97
 
95
- def add_vectors(self, vectors_2d: np.ndarray, metadata_array: np.ndarray) -> None:
96
- """
97
- Deprecated API for adding vectors and metadata arrays.
98
-
99
- Use add(...) instead.
100
- """
101
- warnings.warn(
102
- "VectorStore.add_vectors() is deprecated and will be removed in a future 0.x release. "
103
- "Use VectorStore.add() instead.",
104
- DeprecationWarning,
105
- stacklevel=2,
106
- )
107
- if vectors_2d.ndim != 2:
108
- raise ValueError("vectors_2d must be a 2D NumPy array")
109
-
110
- if metadata_array.ndim != 1:
111
- raise ValueError("metadata_array must be a 1D NumPy array")
112
-
113
- self.add(vectors_2d, metadata_array.tolist())
114
-
115
98
  def load(self) -> None:
116
99
  """Load vectors from file if file_path is specified and exists."""
117
100
  if self._loaded or not self.file_path:
@@ -126,7 +109,7 @@ class VectorStore(Generic[TMetadata]):
126
109
  loaded_metadata = np.array(data["metadata"], copy=True)
127
110
 
128
111
  self._validate_loaded_arrays(loaded_vectors, loaded_metadata)
129
- loaded_vectors = self._normalize_vectors(
112
+ loaded_vectors = self._prepare_vectors_for_storage(
130
113
  loaded_vectors,
131
114
  error_message="Loaded vectors contain zero-norm vectors",
132
115
  )
@@ -164,87 +147,98 @@ class VectorStore(Generic[TMetadata]):
164
147
  min_value: Optional minimum cosine similarity value.
165
148
  within_rows: Optional row indexes to restrict search to.
166
149
  """
167
- query_vector = self._validate_query(query)
168
-
169
- if top_k <= 0:
170
- raise ValueError("top_k must be greater than 0")
171
-
172
- if len(self.vectors) == 0:
173
- return []
150
+ return self._metric_search(
151
+ query,
152
+ top_k=top_k,
153
+ within_rows=within_rows,
154
+ values_fn=self._cosine_values,
155
+ descending=True,
156
+ min_value=min_value,
157
+ max_value=None,
158
+ )
174
159
 
175
- row_indices = self._normalize_within_rows(within_rows)
176
- if len(row_indices) == 0:
177
- return []
160
+ def dot_search(
161
+ self,
162
+ query: npt.ArrayLike,
163
+ *,
164
+ top_k: int = 10,
165
+ min_value: float | None = None,
166
+ within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None = None,
167
+ ) -> list[VectorHit[TMetadata]]:
168
+ """
169
+ Return rows ranked by dot product.
178
170
 
179
- similarities = self._cosine_similarity_numpy(
180
- query_vector, self.vectors[row_indices]
171
+ With normalize=True, this is the dot product of unit vectors. With
172
+ normalize=False, this is the true dot product over original vectors.
173
+ """
174
+ return self._metric_search(
175
+ query,
176
+ top_k=top_k,
177
+ within_rows=within_rows,
178
+ values_fn=self._dot_values,
179
+ descending=True,
180
+ min_value=min_value,
181
+ max_value=None,
181
182
  )
182
183
 
183
- if min_value is not None:
184
- valid_local_indices = np.flatnonzero(similarities >= min_value)
185
- else:
186
- valid_local_indices = np.arange(len(similarities))
187
-
188
- if len(valid_local_indices) == 0:
189
- return []
184
+ def euclidean_search(
185
+ self,
186
+ query: npt.ArrayLike,
187
+ *,
188
+ top_k: int = 10,
189
+ max_value: float | None = None,
190
+ within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None = None,
191
+ ) -> list[VectorHit[TMetadata]]:
192
+ """
193
+ Return rows ranked by Euclidean distance.
190
194
 
191
- valid_similarities = similarities[valid_local_indices]
192
- result_count = min(top_k, len(valid_local_indices))
193
- if result_count < len(valid_local_indices):
194
- top_local_unsorted = np.argpartition(valid_similarities, -result_count)[
195
- -result_count:
196
- ]
197
- else:
198
- top_local_unsorted = np.arange(len(valid_similarities))
195
+ With normalize=True, this is distance between normalized directions.
196
+ With normalize=False, this is true Euclidean distance over original
197
+ vectors.
198
+ """
199
+ return self._metric_search(
200
+ query,
201
+ top_k=top_k,
202
+ within_rows=within_rows,
203
+ values_fn=self._euclidean_values,
204
+ descending=False,
205
+ min_value=None,
206
+ max_value=max_value,
207
+ )
199
208
 
200
- top_local_sorted = top_local_unsorted[
201
- np.argsort(valid_similarities[top_local_unsorted])[::-1]
202
- ]
203
- local_indices = valid_local_indices[top_local_sorted]
204
- original_indices = row_indices[local_indices]
209
+ def _cosine_values(
210
+ self, query: npt.NDArray[np.float32], vectors: npt.NDArray[np.float32]
211
+ ) -> npt.NDArray[np.float32]:
212
+ """Compute cosine similarity values."""
213
+ query_norm = self._normalize_query(query)
214
+ values = np.dot(vectors, query_norm)
205
215
 
206
- return [
207
- VectorHit(
208
- index=int(original_idx),
209
- value=float(similarities[local_idx]),
210
- metadata=self.metadata[original_idx],
211
- )
212
- for local_idx, original_idx in zip(
213
- local_indices, original_indices, strict=True
214
- )
215
- ]
216
+ if not self.normalize:
217
+ vector_norms = np.linalg.norm(vectors, axis=1)
218
+ if np.any(vector_norms == 0):
219
+ raise ValueError("Cannot cosine search zero-norm stored vectors")
220
+ values = values / vector_norms
216
221
 
217
- def search(
218
- self, query_vector: np.ndarray, top_k: int = 10, score_cutoff: float = 0.0
219
- ) -> list[tuple[int, float, TMetadata]]:
220
- """
221
- Deprecated tuple-returning similarity search.
222
+ return np.asarray(values, dtype=np.float32)
222
223
 
223
- Use cosine_search(...), which returns VectorHit objects and uses
224
- min_value instead.
225
- """
226
- warnings.warn(
227
- "VectorStore.search() is deprecated and will be removed in a future 0.x release. "
228
- "Use VectorStore.cosine_search() instead.",
229
- DeprecationWarning,
230
- stacklevel=2,
231
- )
232
- hits = self.cosine_search(query_vector, top_k=top_k, min_value=score_cutoff)
233
- return [(hit.index, hit.value, hit.metadata) for hit in hits]
224
+ def _dot_values(
225
+ self, query: npt.NDArray[np.float32], vectors: npt.NDArray[np.float32]
226
+ ) -> npt.NDArray[np.float32]:
227
+ """Compute dot product values."""
228
+ if self.normalize:
229
+ query = self._normalize_query(query)
230
+ return np.asarray(np.dot(vectors, query), dtype=np.float32)
234
231
 
235
- def _cosine_similarity_numpy(
232
+ def _euclidean_values(
236
233
  self, query: npt.NDArray[np.float32], vectors: npt.NDArray[np.float32]
237
234
  ) -> npt.NDArray[np.float32]:
238
- """Compute cosine similarity against normalized vectors."""
239
- query_magnitude = np.linalg.norm(query)
240
- if query_magnitude == 0:
241
- raise ValueError("Cannot search with zero-norm query vector")
242
- query_norm = query / query_magnitude
243
- similarities = np.asarray(np.dot(vectors, query_norm), dtype=np.float32)
244
- return similarities
235
+ """Compute Euclidean distance values."""
236
+ if self.normalize:
237
+ query = self._normalize_query(query)
238
+ return np.asarray(np.linalg.norm(vectors - query, axis=1), dtype=np.float32)
245
239
 
246
240
  def get(self, index: int) -> tuple[npt.NDArray[np.float32], TMetadata] | None:
247
- """Get a normalized vector and metadata payload by row index."""
241
+ """Get a stored vector and metadata payload by row index."""
248
242
  if 0 <= index < len(self.vectors):
249
243
  return (self.vectors[index], self.metadata[index])
250
244
  return None
@@ -272,13 +266,21 @@ class VectorStore(Generic[TMetadata]):
272
266
  ) -> npt.NDArray[Any]:
273
267
  return np.asarray(metadata, dtype=object)
274
268
 
275
- def _normalize_vectors(
269
+ def _prepare_vectors_for_storage(
276
270
  self, vectors: npt.NDArray[np.float32], *, error_message: str
277
271
  ) -> npt.NDArray[np.float32]:
278
- norms = np.linalg.norm(vectors, axis=1, keepdims=True)
272
+ norms = self._validate_non_zero_vectors(vectors, error_message=error_message)
273
+ if self.normalize:
274
+ return np.asarray(vectors / norms[:, np.newaxis], dtype=np.float32)
275
+ return vectors.astype(np.float32, copy=True)
276
+
277
+ def _validate_non_zero_vectors(
278
+ self, vectors: npt.NDArray[np.float32], *, error_message: str
279
+ ) -> npt.NDArray[np.float32]:
280
+ norms = np.linalg.norm(vectors, axis=1)
279
281
  if np.any(norms == 0):
280
282
  raise ValueError(error_message)
281
- return np.asarray(vectors / norms, dtype=np.float32)
283
+ return np.asarray(norms, dtype=np.float32)
282
284
 
283
285
  def _validate_query(self, query: npt.ArrayLike) -> npt.NDArray[np.float32]:
284
286
  query_vector = np.asarray(query, dtype=np.float32)
@@ -290,6 +292,14 @@ class VectorStore(Generic[TMetadata]):
290
292
  )
291
293
  return query_vector
292
294
 
295
+ def _normalize_query(
296
+ self, query: npt.NDArray[np.float32]
297
+ ) -> npt.NDArray[np.float32]:
298
+ query_magnitude = np.linalg.norm(query)
299
+ if query_magnitude == 0:
300
+ raise ValueError("Cannot search with zero-norm query vector")
301
+ return np.asarray(query / query_magnitude, dtype=np.float32)
302
+
293
303
  def _normalize_within_rows(
294
304
  self, within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None
295
305
  ) -> npt.NDArray[np.intp]:
@@ -309,6 +319,80 @@ class VectorStore(Generic[TMetadata]):
309
319
  raise IndexError("within_rows contains row indexes outside the store")
310
320
  return rows
311
321
 
322
+ def _metric_search(
323
+ self,
324
+ query: npt.ArrayLike,
325
+ *,
326
+ top_k: int,
327
+ within_rows: Sequence[int] | npt.NDArray[np.integer[Any]] | None,
328
+ values_fn: Callable[
329
+ [npt.NDArray[np.float32], npt.NDArray[np.float32]],
330
+ npt.NDArray[np.float32],
331
+ ],
332
+ descending: bool,
333
+ min_value: float | None,
334
+ max_value: float | None,
335
+ ) -> list[VectorHit[TMetadata]]:
336
+ query_vector = self._validate_query(query)
337
+
338
+ if top_k <= 0:
339
+ raise ValueError("top_k must be greater than 0")
340
+
341
+ if len(self.vectors) == 0:
342
+ return []
343
+
344
+ row_indices = self._normalize_within_rows(within_rows)
345
+ if len(row_indices) == 0:
346
+ return []
347
+
348
+ values = values_fn(query_vector, self.vectors[row_indices])
349
+
350
+ valid_local_indices = np.arange(len(values))
351
+ if min_value is not None:
352
+ valid_local_indices = valid_local_indices[
353
+ values[valid_local_indices] >= min_value
354
+ ]
355
+ if max_value is not None:
356
+ valid_local_indices = valid_local_indices[
357
+ values[valid_local_indices] <= max_value
358
+ ]
359
+
360
+ if len(valid_local_indices) == 0:
361
+ return []
362
+
363
+ valid_values = values[valid_local_indices]
364
+ result_count = min(top_k, len(valid_local_indices))
365
+ if result_count < len(valid_local_indices):
366
+ if descending:
367
+ top_valid_unsorted = np.argpartition(valid_values, -result_count)[
368
+ -result_count:
369
+ ]
370
+ else:
371
+ top_valid_unsorted = np.argpartition(valid_values, result_count - 1)[
372
+ :result_count
373
+ ]
374
+ else:
375
+ top_valid_unsorted = np.arange(len(valid_values))
376
+
377
+ sort_order = np.argsort(valid_values[top_valid_unsorted])
378
+ if descending:
379
+ sort_order = sort_order[::-1]
380
+
381
+ top_valid_sorted = top_valid_unsorted[sort_order]
382
+ local_indices = valid_local_indices[top_valid_sorted]
383
+ original_indices = row_indices[local_indices]
384
+
385
+ return [
386
+ VectorHit(
387
+ index=int(original_idx),
388
+ value=float(values[local_idx]),
389
+ metadata=self.metadata[original_idx],
390
+ )
391
+ for local_idx, original_idx in zip(
392
+ local_indices, original_indices, strict=True
393
+ )
394
+ ]
395
+
312
396
  def _validate_required_fields(self, files: set[str]) -> None:
313
397
  if "vectors" not in files:
314
398
  raise ValueError("Persisted vector store is missing vectors")
@@ -22,9 +22,16 @@ class TestVectorStore:
22
22
  """Test VectorStore initialization."""
23
23
  store = VectorStore(dimensions=3)
24
24
  assert store.dimensions == 3
25
+ assert store.normalize is True
25
26
  assert len(store.vectors) == 0
26
27
  assert len(store) == 0
27
28
 
29
+ def test_init_allows_raw_vector_storage(self):
30
+ """Test VectorStore can preserve raw vectors."""
31
+ store = VectorStore(dimensions=3, normalize=False)
32
+
33
+ assert store.normalize is False
34
+
28
35
  def test_init_rejects_non_positive_dimensions(self):
29
36
  """Test VectorStore rejects invalid dimensions."""
30
37
  with pytest.raises(ValueError, match="dimensions"):
@@ -47,6 +54,15 @@ class TestVectorStore:
47
54
 
48
55
  np.testing.assert_array_almost_equal(store.vectors[0], np.array([0.6, 0.8]))
49
56
 
57
+ def test_add_preserves_raw_vectors_when_normalize_false(self):
58
+ """Test add preserves original vectors when normalize=False."""
59
+ store = VectorStore[dict[str, str]](dimensions=2, normalize=False)
60
+
61
+ store.add([[3.0, 4.0]], [{"id": "v"}])
62
+
63
+ np.testing.assert_array_equal(store.vectors[0], np.array([3.0, 4.0]))
64
+ np.testing.assert_array_equal(store.get(0)[0], np.array([3.0, 4.0]))
65
+
50
66
  def test_add_rejects_wrong_dimension(self):
51
67
  """Test add rejects vectors with wrong dimensions."""
52
68
  store = VectorStore(dimensions=3)
@@ -82,6 +98,13 @@ class TestVectorStore:
82
98
  with pytest.raises(ValueError, match="zero-norm"):
83
99
  store.add([[0.0, 0.0, 0.0]], [{"id": "zero"}])
84
100
 
101
+ def test_add_rejects_zero_norm_vector_when_normalize_false(self):
102
+ """Test add rejects zero vectors because cosine search is always available."""
103
+ store = VectorStore(dimensions=3, normalize=False)
104
+
105
+ with pytest.raises(ValueError, match="zero-norm"):
106
+ store.add([[0.0, 0.0, 0.0]], [{"id": "zero"}])
107
+
85
108
  def test_cosine_search_returns_vector_hits(self):
86
109
  """Test preferred cosine_search returns typed vector hits."""
87
110
  store = VectorStore[dict[str, str]](dimensions=3)
@@ -217,6 +240,116 @@ class TestVectorStore:
217
240
 
218
241
  assert results == []
219
242
 
243
+ def test_cosine_search_with_raw_vectors(self):
244
+ """Test cosine_search computes true cosine when normalize=False."""
245
+ store = VectorStore[dict[str, str]](dimensions=2, normalize=False)
246
+ store.add(
247
+ [[10.0, 0.0], [0.0, 2.0], [3.0, 4.0]],
248
+ [{"id": "x"}, {"id": "y"}, {"id": "diag"}],
249
+ )
250
+
251
+ results = store.cosine_search([1.0, 0.0], top_k=3)
252
+
253
+ assert [hit.index for hit in results] == [0, 2, 1]
254
+ assert [hit.value for hit in results] == pytest.approx([1.0, 0.6, 0.0])
255
+
256
+ def test_dot_search_with_raw_vectors(self):
257
+ """Test dot_search ranks by true dot product when normalize=False."""
258
+ store = VectorStore[dict[str, str]](dimensions=2, normalize=False)
259
+ store.add(
260
+ [[10.0, 0.0], [0.0, 2.0], [3.0, 4.0]],
261
+ [{"id": "x"}, {"id": "y"}, {"id": "diag"}],
262
+ )
263
+
264
+ results = store.dot_search([1.0, 0.0], top_k=3)
265
+
266
+ assert all(isinstance(hit, VectorHit) for hit in results)
267
+ assert [hit.index for hit in results] == [0, 2, 1]
268
+ assert [hit.value for hit in results] == pytest.approx([10.0, 3.0, 0.0])
269
+
270
+ def test_dot_search_with_normalized_vectors(self):
271
+ """Test dot_search uses stored unit vectors when normalize=True."""
272
+ store = VectorStore[dict[str, str]](dimensions=2)
273
+ store.add(
274
+ [[10.0, 0.0], [3.0, 4.0], [0.0, 2.0]],
275
+ [{"id": "x"}, {"id": "diag"}, {"id": "y"}],
276
+ )
277
+
278
+ results = store.dot_search([2.0, 0.0], top_k=3, min_value=0.5)
279
+
280
+ assert [hit.index for hit in results] == [0, 1]
281
+ assert [hit.value for hit in results] == pytest.approx([1.0, 0.6])
282
+
283
+ def test_euclidean_search_with_raw_vectors(self):
284
+ """Test euclidean_search ranks by nearest raw vector when normalize=False."""
285
+ store = VectorStore[dict[str, str]](dimensions=2, normalize=False)
286
+ store.add(
287
+ [[1.0, 1.0], [4.0, 5.0], [2.0, 1.0]],
288
+ [{"id": "near"}, {"id": "far"}, {"id": "also-near"}],
289
+ )
290
+
291
+ results = store.euclidean_search([1.0, 2.0], top_k=3)
292
+
293
+ assert all(isinstance(hit, VectorHit) for hit in results)
294
+ assert [hit.index for hit in results] == [0, 2, 1]
295
+ assert [hit.value for hit in results] == pytest.approx(
296
+ [1.0, np.sqrt(2.0), np.sqrt(18.0)]
297
+ )
298
+
299
+ def test_euclidean_search_with_max_value(self):
300
+ """Test euclidean_search max_value filters by distance."""
301
+ store = VectorStore[dict[str, str]](dimensions=2, normalize=False)
302
+ store.add(
303
+ [[1.0, 1.0], [4.0, 5.0], [2.0, 1.0]],
304
+ [{"id": "near"}, {"id": "far"}, {"id": "also-near"}],
305
+ )
306
+
307
+ results = store.euclidean_search([1.0, 2.0], top_k=3, max_value=1.5)
308
+
309
+ assert [hit.index for hit in results] == [0, 2]
310
+
311
+ def test_euclidean_search_with_normalized_vectors(self):
312
+ """Test euclidean_search uses normalized stored vectors by default."""
313
+ store = VectorStore[dict[str, str]](dimensions=2)
314
+ store.add(
315
+ [[10.0, 0.0], [3.0, 4.0]],
316
+ [{"id": "x"}, {"id": "diag"}],
317
+ )
318
+
319
+ results = store.euclidean_search([2.0, 0.0], top_k=2)
320
+
321
+ assert [hit.index for hit in results] == [0, 1]
322
+ assert [hit.value for hit in results] == pytest.approx([0.0, np.sqrt(0.8)])
323
+
324
+ def test_metric_searches_support_within_rows(self):
325
+ """Test dot and Euclidean search support prefiltered row indexes."""
326
+ store = VectorStore[dict[str, str]](dimensions=2, normalize=False)
327
+ store.add(
328
+ [[10.0, 0.0], [0.0, 2.0], [3.0, 4.0]],
329
+ [{"id": "x"}, {"id": "y"}, {"id": "diag"}],
330
+ )
331
+
332
+ dot_results = store.dot_search([1.0, 0.0], within_rows=[1, 2])
333
+ euclidean_results = store.euclidean_search([1.0, 0.0], within_rows=[1, 2])
334
+
335
+ assert [hit.index for hit in dot_results] == [2, 1]
336
+ assert [hit.index for hit in euclidean_results] == [1, 2]
337
+
338
+ def test_metric_searches_validate_common_inputs(self):
339
+ """Test added search methods share common validation behavior."""
340
+ store = VectorStore(dimensions=2)
341
+ store.add([[1.0, 0.0]], [{"id": "x"}])
342
+
343
+ with pytest.raises(ValueError, match="top_k"):
344
+ store.dot_search([1.0, 0.0], top_k=0)
345
+
346
+ with pytest.raises(ValueError, match="Query vector dimension"):
347
+ store.euclidean_search([1.0])
348
+
349
+ assert store.dot_search([1.0, 0.0], within_rows=[]) == []
350
+ with pytest.raises(IndexError, match="outside"):
351
+ store.euclidean_search([1.0, 0.0], within_rows=[1])
352
+
220
353
  def test_get(self):
221
354
  """Test retrieving vector and metadata by index."""
222
355
  store = VectorStore(dimensions=2)
@@ -442,6 +575,27 @@ class TestVectorStore:
442
575
  finally:
443
576
  Path(file_path).unlink(missing_ok=True)
444
577
 
578
+ def test_load_preserves_raw_vectors_when_normalize_false(self):
579
+ """Test load preserves valid vectors when normalize=False."""
580
+ with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
581
+ file_path = tmp.name
582
+
583
+ try:
584
+ np.savez_compressed(
585
+ file_path,
586
+ vectors=np.array([[2.0, 0.0]], dtype=np.float32),
587
+ metadata=np.array([{"id": "test"}], dtype=object),
588
+ )
589
+
590
+ store = VectorStore(dimensions=2, file_path=file_path, normalize=False)
591
+ store.load()
592
+
593
+ np.testing.assert_array_equal(store.vectors[0], np.array([2.0, 0.0]))
594
+ assert store.cosine_search([1.0, 0.0])[0].value == pytest.approx(1.0)
595
+ assert store.dot_search([1.0, 0.0])[0].value == pytest.approx(2.0)
596
+ finally:
597
+ Path(file_path).unlink(missing_ok=True)
598
+
445
599
  def test_load_raises_on_zero_norm_vectors(self):
446
600
  """Test load rejects persisted zero-norm vectors."""
447
601
  with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
@@ -460,6 +614,24 @@ class TestVectorStore:
460
614
  finally:
461
615
  Path(file_path).unlink(missing_ok=True)
462
616
 
617
+ def test_load_raises_on_zero_norm_vectors_when_normalize_false(self):
618
+ """Test raw stores reject zero rows because cosine search is always available."""
619
+ with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
620
+ file_path = tmp.name
621
+
622
+ try:
623
+ np.savez_compressed(
624
+ file_path,
625
+ vectors=np.array([[0.0, 0.0]], dtype=np.float32),
626
+ metadata=np.array([{"id": "zero"}], dtype=object),
627
+ )
628
+
629
+ store = VectorStore(dimensions=2, file_path=file_path, normalize=False)
630
+ with pytest.raises(ValueError, match="zero-norm"):
631
+ store.load()
632
+ finally:
633
+ Path(file_path).unlink(missing_ok=True)
634
+
463
635
  def test_save_no_file_path(self):
464
636
  """Test save with no file path."""
465
637
  store = VectorStore(dimensions=2)
@@ -498,56 +670,24 @@ class TestVectorStore:
498
670
  finally:
499
671
  Path(file_path).unlink(missing_ok=True)
500
672
 
501
- def test_add_vectors_deprecated_compatibility_api(self):
502
- """Test deprecated add_vectors remains compatible temporarily."""
503
- store = VectorStore(dimensions=2)
504
- vectors_2d = np.array([[1.0, 2.0], [3.0, 4.0]])
505
- metadata_array = np.array([{"id": 1}, {"id": 2}])
506
-
507
- with pytest.deprecated_call(match="add_vectors"):
508
- store.add_vectors(vectors_2d, metadata_array)
509
-
510
- assert len(store) == 2
511
- assert store.get(0)[1] == {"id": 1}
512
-
513
- def test_add_vectors_deprecated_validation(self):
514
- """Test deprecated add_vectors still preserves validation."""
515
- store = VectorStore(dimensions=2)
516
-
517
- with pytest.raises(ValueError, match="2D"):
518
- with pytest.deprecated_call(match="add_vectors"):
519
- store.add_vectors(np.array([1.0, 2.0]), np.array([{"id": 1}]))
520
-
521
- with pytest.raises(ValueError, match="1D"):
522
- with pytest.deprecated_call(match="add_vectors"):
523
- store.add_vectors(np.array([[1.0, 2.0]]), np.array([[{"id": 1}]]))
524
-
525
- def test_search_deprecated_compatibility_api(self):
526
- """Test deprecated search remains compatible temporarily."""
527
- store = VectorStore(dimensions=3)
528
- store.add(
529
- [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
530
- [{"name": "x-axis"}, {"name": "y-axis"}],
531
- )
532
-
533
- with pytest.deprecated_call(match="search"):
534
- results = store.search(np.array([0.9, 0.1, 0.0]), top_k=2)
673
+ def test_save_and_load_raw_vectors_round_trip(self):
674
+ """Test normalize=False saves and loads raw vectors without mode metadata."""
675
+ with tempfile.NamedTemporaryFile(suffix=".npz", delete=False) as tmp:
676
+ file_path = tmp.name
535
677
 
536
- assert len(results) == 2
537
- assert results[0][0] == 0
538
- assert results[0][1] > 0.9
539
- assert results[0][2] == {"name": "x-axis"}
678
+ try:
679
+ store1 = VectorStore(dimensions=2, file_path=file_path, normalize=False)
680
+ store1.add([[3.0, 4.0]], [{"id": "raw"}])
681
+ store1.save()
540
682
 
541
- def test_search_deprecated_score_cutoff(self):
542
- """Test deprecated search still supports score_cutoff."""
543
- store = VectorStore(dimensions=3)
544
- store.add(
545
- [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
546
- [{"name": "x-axis"}, {"name": "y-axis"}],
547
- )
683
+ with np.load(file_path, allow_pickle=True) as data:
684
+ assert set(data.files) == {"vectors", "metadata"}
685
+ np.testing.assert_array_equal(data["vectors"][0], np.array([3.0, 4.0]))
548
686
 
549
- with pytest.deprecated_call(match="search"):
550
- results = store.search(np.array([0.9, 0.1, 0.0]), score_cutoff=0.8)
687
+ store2 = VectorStore(dimensions=2, file_path=file_path, normalize=False)
688
+ store2.load()
551
689
 
552
- assert len(results) == 1
553
- assert results[0][0] == 0
690
+ np.testing.assert_array_equal(store2.get(0)[0], np.array([3.0, 4.0]))
691
+ assert store2.get(0)[1] == {"id": "raw"}
692
+ finally:
693
+ Path(file_path).unlink(missing_ok=True)