Improved internal representation of Bit class

Author: ankane

pgvector/bit.py

@@ -5,51 +5,58 @@
 
 class Bit:
     def __init__(self, value):
-        if isinstance(value, str):
-            self._value = self.from_text(value)._value
-        elif isinstance(value, bytes):
-            self._value = np.unpackbits(np.frombuffer(value, dtype=np.uint8)).astype(bool)
+        if isinstance(value, bytes):
+            self._len = 8 * len(value)
+            self._data = value
         else:
-            value = np.asarray(value)
+            if isinstance(value, str):
+                value = [v != '0' for v in value]
+            else:
+                value = np.asarray(value)
 
-            if value.dtype != np.bool:
-                warn('expected elements to be boolean', stacklevel=2)
-                value = value.astype(bool)
+                if value.dtype != np.bool:
+                    warn('expected elements to be boolean', stacklevel=2)
+                    value = value.astype(bool)
 
-            if value.ndim != 1:
-                raise ValueError('expected ndim to be 1')
+                if value.ndim != 1:
+                    raise ValueError('expected ndim to be 1')
 
-            self._value = value
+            self._len = len(value)
+            self._data = np.packbits(value).tobytes()
 
     def __repr__(self):
         return f'Bit({self.to_text()})'
 
     def __eq__(self, other):
         if isinstance(other, self.__class__):
-            return np.array_equal(self.to_numpy(), other.to_numpy())
+            return self._len == other._len and self._data == other._data
         return False
 
     def to_list(self):
-        return self._value.tolist()
+        return self.to_numpy().tolist()
 
     def to_numpy(self):
-        return self._value
+        return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool)
 
     def to_text(self):
-        return ''.join(self._value.astype(np.uint8).astype(str))
+        return ''.join(format(v, '08b') for v in self._data)[:self._len]
 
     def to_binary(self):
-        return pack('>i', len(self._value)) + np.packbits(self._value).tobytes()
+        return pack('>i', self._len) + self._data
 
     @classmethod
     def from_text(cls, value):
-        return cls(np.asarray([v != '0' for v in value], dtype=bool))
+        return cls(str(value))
 
     @classmethod
     def from_binary(cls, value):
-        count = unpack_from('>i', value)[0]
-        buf = np.frombuffer(value, dtype=np.uint8, offset=4)
-        return cls(np.unpackbits(buf, count=count).astype(bool))
+        if not isinstance(value, bytes):
+            raise ValueError('expected bytes')
+
+        bit = cls.__new__(cls)
+        bit._len = unpack_from('>i', value)[0]
+        bit._data = value[4:]
+        return bit
 
     @classmethod
     def _to_db(cls, value):

tests/test_bit.py

@@ -25,6 +25,11 @@ def test_bytes(self):
         assert Bit(b'\xff\x00').to_list() == [True, True, True, True, True, True, True, True, False, False, False, False, False, False, False, False]
         assert Bit(b'\xfe\x07').to_list() == [True, True, True, True, True, True, True, False, False, False, False, False, False, True, True, True]
 
+    def test_ndarray(self):
+        arr = np.array([True, False, True])
+        assert Bit(arr).to_list() == [True, False, True]
+        assert np.array_equal(Bit(arr).to_numpy(), arr)
+
     def test_ndarray_uint8(self):
         arr = np.array([254, 7, 0], dtype=np.uint8)
         with pytest.warns(UserWarning, match='expected elements to be boolean'):
@@ -35,11 +40,6 @@ def test_ndarray_uint16(self):
         with pytest.warns(UserWarning, match='expected elements to be boolean'):
             assert Bit(arr).to_text() == '110'
 
-    def test_ndarray_same_object(self):
-        arr = np.array([True, False, True])
-        assert Bit(arr).to_list() == [True, False, True]
-        assert Bit(arr).to_numpy() is arr
-
     def test_ndim_two(self):
         with pytest.raises(ValueError) as error:
             Bit([[True, False], [True, False]])