Switched to list for internal storage format for Vector and HalfVector [skip ci]

Author: ankane

pgvector/halfvec.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 import struct
-from typing import cast
 
 try:
     import numpy as np
@@ -12,50 +11,47 @@
 class HalfVector:
     def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
         if isinstance(value, list):
-            dim = len(value)
             try:
-                self._value = struct.pack(f'>HH{dim}e', dim, 0, *value)
-            except struct.error:
+                self._value = [float(v) for v in value]
+            except (TypeError, ValueError):
                 raise ValueError('expected list[float]')
         elif NUMPY_AVAILABLE and 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()
+            self._value = [float(v) for v in value]
         else:
             raise ValueError('expected list or ndarray')
 
     def __repr__(self) -> str:
-        return f'HalfVector({self.to_list()})'
+        return f'HalfVector({self._value})'
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
-            return self.to_binary() == other.to_binary()
+            return self._value == other._value
         return False
 
     def dimensions(self) -> int:
-        dim, = cast(tuple[int], struct.unpack_from('>H', self._value))
-        return dim
+        return len(self._value)
 
     def to_list(self) -> list[float]:
-        return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:]))
+        return self._value
 
     def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]:
-        return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4)
+        return np.array(self._value, dtype=np.float16)
 
     def to_text(self) -> str:
-        return f'[{",".join([str(v) for v in self.to_list()])}]'
+        return f'[{",".join([str(v) for v in self._value])}]'
 
     def to_binary(self) -> bytes:
-        return self._value
+        dim = len(self._value)
+        return struct.pack(f'>HH{dim}e', dim, 0, *self._value)
 
     @classmethod
     def from_text(cls, value: str) -> HalfVector:
-        return cls([float(v) for v in value[1:-1].split(',')])
+        vec = cls.__new__(cls)
+        vec._value = [float(v) for v in value[1:-1].split(',')]
+        return vec
 
     @classmethod
     def from_binary(cls, value: bytes) -> HalfVector:
@@ -68,7 +64,7 @@ def from_binary(cls, value: bytes) -> HalfVector:
             raise ValueError('expected unused to be 0')
 
         vec = cls.__new__(cls)
-        vec._value = value
+        vec._value = list(struct.unpack_from(f'>{dim}e', value[4:]))
         return vec
 
     @classmethod

pgvector/vector.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 import struct
-from typing import cast
 
 try:
     import numpy as np
@@ -12,50 +11,47 @@
 class Vector:
     def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
         if isinstance(value, list):
-            dim = len(value)
             try:
-                self._value = struct.pack(f'>HH{dim}f', dim, 0, *value)
-            except struct.error:
+                self._value = [float(v) for v in value]
+            except (TypeError, ValueError):
                 raise ValueError('expected list[float]')
         elif NUMPY_AVAILABLE and 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()
+            self._value = [float(v) for v in value]
         else:
             raise ValueError('expected list or ndarray')
 
     def __repr__(self) -> str:
-        return f'Vector({self.to_list()})'
+        return f'Vector({self._value})'
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
-            return self.to_binary() == other.to_binary()
+            return self._value == other._value
         return False
 
     def dimensions(self) -> int:
-        dim, = cast(tuple[int], struct.unpack_from('>H', self._value))
-        return dim
+        return len(self._value)
 
     def to_list(self) -> list[float]:
-        return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:]))
+        return self._value
 
     def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]:
-        return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4)
+        return np.array(self._value, dtype=np.float32)
 
     def to_text(self) -> str:
-        return f'[{",".join([str(v) for v in self.to_list()])}]'
+        return f'[{",".join([str(v) for v in self._value])}]'
 
     def to_binary(self) -> bytes:
-        return self._value
+        dim = len(self._value)
+        return struct.pack(f'>HH{dim}f', dim, 0, *self._value)
 
     @classmethod
     def from_text(cls, value: str) -> Vector:
-        return cls([float(v) for v in value[1:-1].split(',')])
+        vec = cls.__new__(cls)
+        vec._value = [float(v) for v in value[1:-1].split(',')]
+        return vec
 
     @classmethod
     def from_binary(cls, value: bytes) -> Vector:
@@ -68,7 +64,7 @@ def from_binary(cls, value: bytes) -> Vector:
             raise ValueError('expected unused to be 0')
 
         vec = cls.__new__(cls)
-        vec._value = value
+        vec._value = list(struct.unpack_from(f'>{dim}f', value[4:]))
         return vec
 
     @classmethod

tests/test_half_vector.py

@@ -11,7 +11,9 @@
 
 class TestHalfVector:
     def test_list(self) -> None:
-        assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3]
+        arr = [1.0, 2.0, 3.0]
+        assert HalfVector(arr).to_list() == arr
+        assert HalfVector(arr).to_list() is not arr
 
     def test_list_empty(self) -> None:
         assert HalfVector([]).to_list() == []
@@ -31,6 +33,7 @@ def test_ndarray(self) -> None:
         arr = np.array([1, 2, 3])
         assert HalfVector(arr).to_list() == [1, 2, 3]
         assert HalfVector(arr).to_numpy() is not arr
+        assert HalfVector(arr).to_numpy().dtype == np.float16
 
     def test_int(self) -> None:
         with pytest.raises(ValueError) as error:
@@ -48,13 +51,6 @@ def test_equality(self) -> None:
     def test_dimensions(self) -> None:
         assert HalfVector([1, 2, 3]).dimensions() == 3
 
-    @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
-    def test_to_numpy_readonly(self) -> None:
-        arr = HalfVector([1, 2, 3]).to_numpy()
-        with pytest.raises(ValueError) as error:
-            arr[0] = 4
-        assert str(error.value) == 'assignment destination is read-only'
-
     def test_from_text(self) -> None:
         vec = HalfVector.from_text('[1.5,2,3]')
         assert vec.to_list() == [1.5, 2, 3]

tests/test_vector.py

@@ -11,7 +11,9 @@
 
 class TestVector:
     def test_list(self) -> None:
-        assert Vector([1, 2, 3]).to_list() == [1, 2, 3]
+        arr = [1.0, 2.0, 3.0]
+        assert Vector(arr).to_list() == arr
+        assert Vector(arr).to_list() is not arr
 
     def test_list_empty(self) -> None:
         assert Vector([]).to_list() == []
@@ -31,6 +33,7 @@ def test_ndarray(self) -> None:
         arr = np.array([1, 2, 3])
         assert Vector(arr).to_list() == [1, 2, 3]
         assert Vector(arr).to_numpy() is not arr
+        assert Vector(arr).to_numpy().dtype == np.float32
 
     def test_int(self) -> None:
         with pytest.raises(ValueError) as error:
@@ -48,13 +51,6 @@ def test_equality(self) -> None:
     def test_dimensions(self) -> None:
         assert Vector([1, 2, 3]).dimensions() == 3
 
-    @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
-    def test_to_numpy_readonly(self) -> None:
-        arr = Vector([1, 2, 3]).to_numpy()
-        with pytest.raises(ValueError) as error:
-            arr[0] = 4
-        assert str(error.value) == 'assignment destination is read-only'
-
     def test_from_text(self) -> None:
         vec = Vector.from_text('[1.5,2,3]')
         assert vec.to_list() == [1.5, 2, 3]