Improved internal storage format

Author: ankane

pgvector/halfvec.py

@@ -1,55 +1,63 @@
 from __future__ import annotations
 import numpy as np
-from struct import pack, unpack_from
+import struct
 
 
 class HalfVector:
     def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
-        # asarray still copies if same dtype
-        if not isinstance(value, np.ndarray) or value.dtype != '>f2':
-            value = np.asarray(value, dtype='>f2')
-
-        # for mypy
-        assert isinstance(value, np.ndarray)
-
-        if value.ndim != 1:
-            raise ValueError('expected ndim to be 1')
-
-        # atleast_1d for ty
-        self._value = np.atleast_1d(value)
+        if isinstance(value, list):
+            dim = len(value)
+            try:
+                self._value = struct.pack(f'>HH{dim}e', dim, 0, *value)
+            except struct.error as e:
+                raise ValueError('expected list[float]')
+        elif isinstance(value, np.ndarray):
+            if value.ndim != 1:
+                raise ValueError('expected ndim to be 1')
+
+            # asarray still copies if same dtype
+            if value.dtype != '>f2':
+                value = np.asarray(value, dtype='>f2')
+
+            self._value = struct.pack('>HH', value.shape[0], 0) + value.tobytes()
+        else:
+            raise ValueError('expected list or ndarray')
 
     def __repr__(self) -> str:
         return f'HalfVector({self.to_list()})'
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
-            return np.array_equal(self.to_numpy(), other.to_numpy())
+            return self.to_binary() == other.to_binary()
         return False
 
     def dimensions(self) -> int:
-        return len(self._value)
+        dim, = struct.unpack_from('>H', self._value)
+        return dim
 
     def to_list(self) -> list[float]:
-        return self._value.tolist()
+        return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:]))
 
     def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]:
         # TODO return native endian
-        return self._value
+        return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4)
 
     def to_text(self) -> str:
-        return '[' + ','.join([str(float(v)) for v in self._value]) + ']'
+        return '[' + ','.join([str(v) for v in self.to_list()]) + ']'
 
     def to_binary(self) -> bytes:
-        return pack('>HH', self.dimensions(), 0) + self._value.tobytes()
+        return self._value
 
     @classmethod
     def from_text(cls, value: str) -> HalfVector:
         return cls([float(v) for v in value[1:-1].split(',')])
 
     @classmethod
     def from_binary(cls, value: bytes) -> HalfVector:
-        dim, unused = unpack_from('>HH', value)
-        return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4))
+        # TODO check dimensions/length and unused
+        vec = cls.__new__(cls)
+        vec._value = value
+        return vec
 
     @classmethod
     def _to_db(cls, value: object, dim: int | None = None) -> str | None:

pgvector/vector.py

@@ -1,55 +1,63 @@
 from __future__ import annotations
 import numpy as np
-from struct import pack, unpack_from
+import struct
 
 
 class Vector:
     def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
-        # asarray still copies if same dtype
-        if not isinstance(value, np.ndarray) or value.dtype != '>f4':
-            value = np.asarray(value, dtype='>f4')
-
-        # for mypy
-        assert isinstance(value, np.ndarray)
-
-        if value.ndim != 1:
-            raise ValueError('expected ndim to be 1')
-
-        # atleast_1d for ty
-        self._value = np.atleast_1d(value)
+        if isinstance(value, list):
+            dim = len(value)
+            try:
+                self._value = struct.pack(f'>HH{dim}f', dim, 0, *value)
+            except struct.error as e:
+                raise ValueError('expected list[float]')
+        elif isinstance(value, np.ndarray):
+            if value.ndim != 1:
+                raise ValueError('expected ndim to be 1')
+
+            # asarray still copies if same dtype
+            if value.dtype != '>f4':
+                value = np.asarray(value, dtype='>f4')
+
+            self._value = struct.pack('>HH', value.shape[0], 0) + value.tobytes()
+        else:
+            raise ValueError('expected list or ndarray')
 
     def __repr__(self) -> str:
         return f'Vector({self.to_list()})'
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
-            return np.array_equal(self.to_numpy(), other.to_numpy())
+            return self.to_binary() == other.to_binary()
         return False
 
     def dimensions(self) -> int:
-        return len(self._value)
+        dim, = struct.unpack_from('>H', self._value)
+        return dim
 
     def to_list(self) -> list[float]:
-        return self._value.tolist()
+        return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:]))
 
     def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]:
         # TODO return native endian
-        return self._value
+        return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4)
 
     def to_text(self) -> str:
-        return '[' + ','.join([str(float(v)) for v in self._value]) + ']'
+        return '[' + ','.join([str(v) for v in self.to_list()]) + ']'
 
     def to_binary(self) -> bytes:
-        return pack('>HH', self.dimensions(), 0) + self._value.tobytes()
+        return self._value
 
     @classmethod
     def from_text(cls, value: str) -> Vector:
         return cls([float(v) for v in value[1:-1].split(',')])
 
     @classmethod
     def from_binary(cls, value: bytes) -> Vector:
-        dim, unused = unpack_from('>HH', value)
-        return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4))
+        # TODO check dimensions/length and unused
+        vec = cls.__new__(cls)
+        vec._value = value
+        return vec
 
     @classmethod
     def _to_db(cls, value: object, dim: int | None = None) -> str | None:

tests/test_half_vector.py

@@ -9,28 +9,24 @@ def test_list(self):
         assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3]
 
     def test_list_str(self):
-        with pytest.raises(ValueError, match='could not convert string to float'):
+        with pytest.raises(ValueError) as error:
             HalfVector([1, 'two', 3])  # ty: ignore[invalid-argument-type]
+        assert str(error.value) == 'expected list[float]'
 
     def test_ndarray(self):
         arr = np.array([1, 2, 3])
         assert HalfVector(arr).to_list() == [1, 2, 3]
         assert HalfVector(arr).to_numpy() is not arr
 
-    def test_ndarray_same_object(self):
-        arr = np.array([1, 2, 3], dtype='>f2')
-        assert HalfVector(arr).to_list() == [1, 2, 3]
-        assert HalfVector(arr).to_numpy() is arr
-
     def test_ndim_two(self):
         with pytest.raises(ValueError) as error:
             HalfVector([[1, 2], [3, 4]])  # ty: ignore[invalid-argument-type]
-        assert str(error.value) == 'expected ndim to be 1'
+        assert str(error.value) == 'expected list[float]'
 
     def test_ndim_zero(self):
         with pytest.raises(ValueError) as error:
             HalfVector(1)  # ty: ignore[invalid-argument-type]
-        assert str(error.value) == 'expected ndim to be 1'
+        assert str(error.value) == 'expected list or ndarray'
 
     def test_repr(self):
         assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])'

tests/test_vector.py

@@ -9,28 +9,24 @@ def test_list(self):
         assert Vector([1, 2, 3]).to_list() == [1, 2, 3]
 
     def test_list_str(self):
-        with pytest.raises(ValueError, match='could not convert string to float'):
+        with pytest.raises(ValueError) as error:
             Vector([1, 'two', 3])  # ty: ignore[invalid-argument-type]
+        assert str(error.value) == 'expected list[float]'
 
     def test_ndarray(self):
         arr = np.array([1, 2, 3])
         assert Vector(arr).to_list() == [1, 2, 3]
         assert Vector(arr).to_numpy() is not arr
 
-    def test_ndarray_same_object(self):
-        arr = np.array([1, 2, 3], dtype='>f4')
-        assert Vector(arr).to_list() == [1, 2, 3]
-        assert Vector(arr).to_numpy() is arr
-
     def test_ndim_two(self):
         with pytest.raises(ValueError) as error:
             Vector([[1, 2], [3, 4]])  # ty: ignore[invalid-argument-type]
-        assert str(error.value) == 'expected ndim to be 1'
+        assert str(error.value) == 'expected list[float]'
 
     def test_ndim_zero(self):
         with pytest.raises(ValueError) as error:
             Vector(1)  # ty: ignore[invalid-argument-type]
-        assert str(error.value) == 'expected ndim to be 1'
+        assert str(error.value) == 'expected list or ndarray'
 
     def test_repr(self):
         assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])'