forked from pgvector/pgvector-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy path__init__.py
More file actions
49 lines (31 loc) · 1.27 KB
/
__init__.py
File metadata and controls
49 lines (31 loc) · 1.27 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
import numpy as np
from struct import pack, unpack
def from_db(value):
# could be ndarray if already cast by lower-level driver
if value is None or isinstance(value, np.ndarray):
return value
return np.array(value[1:-1].split(','), dtype=np.float32)
def from_db_binary(value):
if value is None:
return value
(dim, unused) = unpack('>HH', value[:4])
return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32)
def to_db(value, dim=None):
if value is None:
return value
if isinstance(value, np.ndarray):
if value.ndim != 1:
raise ValueError('expected ndim to be 1')
if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating):
raise ValueError('dtype must be numeric')
value = value.tolist()
if dim is not None and len(value) != dim:
raise ValueError('expected %d dimensions, not %d' % (dim, len(value)))
return '[' + ','.join([str(float(v)) for v in value]) + ']'
def to_db_binary(value):
if value is None:
return value
value = np.asarray(value, dtype='>f')
if value.ndim != 1:
raise ValueError('expected ndim to be 1')
return pack('>HH', value.shape[0], 0) + value.tobytes()