Skip to content

Commit e771e69

Browse files
committed
Restored fast path for NumPy
1 parent 0b8fbd7 commit e771e69

2 files changed

Lines changed: 11 additions & 3 deletions

File tree

pgvector/psycopg/vector.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,7 @@ class VectorBinaryDumper(VectorDumper):
2828
format = Format.BINARY
2929

3030
def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None:
31-
if not isinstance(obj, Vector):
32-
obj = Vector(obj)
33-
return obj.to_binary()
31+
return Vector._to_db_binary(obj)
3432

3533

3634
class VectorLoader(Loader):

pgvector/vector.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,16 @@ def _to_db_binary(cls, value: object) -> bytes | None:
8686
if value is None:
8787
return value
8888

89+
# fast path for NumPy
90+
if NUMPY_AVAILABLE and isinstance(value, np.ndarray):
91+
if value.ndim != 1:
92+
raise ValueError('expected ndim to be 1')
93+
94+
if value.dtype != '>f4':
95+
value = np.asarray(value, dtype='>f4')
96+
97+
return struct.pack('>HH', len(value), 0) + value.tobytes()
98+
8999
if not isinstance(value, cls):
90100
value = cls(value) # type: ignore
91101

0 commit comments

Comments
 (0)