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.
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/PKG-INFO +85 -27
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/README.md +84 -26
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/src/numpy_vector_store/__init__.py +1 -1
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/src/numpy_vector_store/vector_store.py +184 -100
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/tests/test_vector_store.py +189 -49
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/.github/FUNDING.yml +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/.github/workflows/publish-pypi.yml +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/.github/workflows/publish-testpypi.yml +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/.gitignore +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/LICENSE +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/justfile +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/pyproject.toml +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/tests/__init__.py +0 -0
- {numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: numpy-vector-store
|
|
3
|
-
Version: 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
|
|
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
|
-
|
|
201
|
+
If you save with `normalize=False`, load with `normalize=False` too:
|
|
150
202
|
|
|
151
203
|
```python
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
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
|
-
|
|
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
|
-
|
|
168
|
-
will be removed in a future 0.x release:
|
|
220
|
+
Context manager usage auto-saves on exit:
|
|
169
221
|
|
|
170
222
|
```python
|
|
171
|
-
|
|
172
|
-
|
|
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
|
-
|
|
178
|
-
`
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
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
|
|
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
|
-
|
|
177
|
+
If you save with `normalize=False`, load with `normalize=False` too:
|
|
126
178
|
|
|
127
179
|
```python
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
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
|
-
|
|
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
|
-
|
|
144
|
-
will be removed in a future 0.x release:
|
|
196
|
+
Context manager usage auto-saves on exit:
|
|
145
197
|
|
|
146
198
|
```python
|
|
147
|
-
|
|
148
|
-
|
|
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
|
-
|
|
154
|
-
`
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
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.
|
|
8
|
+
__version__ = "0.3.0"
|
|
9
9
|
__author__ = "Tim VanReenen"
|
|
10
10
|
|
|
11
11
|
from .vector_store import VectorHit, VectorStore
|
{numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/src/numpy_vector_store/vector_store.py
RENAMED
|
@@ -1,7 +1,6 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
-
import
|
|
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
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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 =
|
|
90
|
+
self.vectors = vectors_to_store
|
|
88
91
|
else:
|
|
89
|
-
self.vectors = np.vstack([self.vectors,
|
|
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.
|
|
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
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
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
|
-
|
|
176
|
-
|
|
177
|
-
|
|
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
|
-
|
|
180
|
-
|
|
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
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
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
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
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
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
|
|
204
|
-
|
|
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
|
-
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
210
|
-
|
|
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
|
-
|
|
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
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
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
|
|
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
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
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
|
|
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
|
|
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 =
|
|
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(
|
|
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
|
|
502
|
-
"""Test
|
|
503
|
-
|
|
504
|
-
|
|
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
|
-
|
|
537
|
-
|
|
538
|
-
|
|
539
|
-
|
|
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
|
-
|
|
542
|
-
|
|
543
|
-
|
|
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
|
-
|
|
550
|
-
|
|
687
|
+
store2 = VectorStore(dimensions=2, file_path=file_path, normalize=False)
|
|
688
|
+
store2.load()
|
|
551
689
|
|
|
552
|
-
|
|
553
|
-
|
|
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)
|
|
File without changes
|
|
File without changes
|
{numpy_vector_store-0.2.0 → numpy_vector_store-0.3.0}/.github/workflows/publish-testpypi.yml
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|