aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/pgvector/utils/halfvec.py
diff options
context:
space:
mode:
Diffstat (limited to '.venv/lib/python3.12/site-packages/pgvector/utils/halfvec.py')
-rw-r--r--.venv/lib/python3.12/site-packages/pgvector/utils/halfvec.py78
1 files changed, 78 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/pgvector/utils/halfvec.py b/.venv/lib/python3.12/site-packages/pgvector/utils/halfvec.py
new file mode 100644
index 00000000..e1e50514
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/pgvector/utils/halfvec.py
@@ -0,0 +1,78 @@
+import numpy as np
+from struct import pack, unpack_from
+
+
+class HalfVector:
+ def __init__(self, value):
+ # asarray still copies if same dtype
+ if not isinstance(value, np.ndarray) or value.dtype != '>f2':
+ value = np.asarray(value, dtype='>f2')
+
+ if value.ndim != 1:
+ raise ValueError('expected ndim to be 1')
+
+ self._value = value
+
+ def __repr__(self):
+ return f'HalfVector({self.to_list()})'
+
+ def dimensions(self):
+ return len(self._value)
+
+ def to_list(self):
+ return self._value.tolist()
+
+ def to_numpy(self):
+ return self._value
+
+ def to_text(self):
+ return '[' + ','.join([str(float(v)) for v in self._value]) + ']'
+
+ def to_binary(self):
+ return pack('>HH', self.dimensions(), 0) + self._value.tobytes()
+
+ @classmethod
+ def from_text(cls, value):
+ return cls([float(v) for v in value[1:-1].split(',')])
+
+ @classmethod
+ def from_binary(cls, value):
+ dim, unused = unpack_from('>HH', value)
+ return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4))
+
+ @classmethod
+ def _to_db(cls, value, dim=None):
+ if value is None:
+ return value
+
+ if not isinstance(value, cls):
+ value = cls(value)
+
+ if dim is not None and value.dimensions() != dim:
+ raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions()))
+
+ return value.to_text()
+
+ @classmethod
+ def _to_db_binary(cls, value):
+ if value is None:
+ return value
+
+ if not isinstance(value, cls):
+ value = cls(value)
+
+ return value.to_binary()
+
+ @classmethod
+ def _from_db(cls, value):
+ if value is None or isinstance(value, cls):
+ return value
+
+ return cls.from_text(value)
+
+ @classmethod
+ def _from_db_binary(cls, value):
+ if value is None or isinstance(value, cls):
+ return value
+
+ return cls.from_binary(value)