pgvector 0.3.5__py3-none-any.whl

pgvector/utils/sparsevec.py

@@ -0,0 +1,156 @@
+import numpy as np
+from struct import pack, unpack_from
+
+NO_DEFAULT = object()
+
+
+class SparseVector:
+    def __init__(self, value, dimensions=NO_DEFAULT, /):
+        if value.__class__.__module__.startswith('scipy.sparse.'):
+            if dimensions is not NO_DEFAULT:
+                raise ValueError('extra argument')
+
+            self._from_sparse(value)
+        elif isinstance(value, dict):
+            if dimensions is NO_DEFAULT:
+                raise ValueError('missing dimensions')
+
+            self._from_dict(value, dimensions)
+        else:
+            if dimensions is not NO_DEFAULT:
+                raise ValueError('extra argument')
+
+            self._from_dense(value)
+
+    def __repr__(self):
+        elements = dict(zip(self._indices, self._values))
+        return f'SparseVector({elements}, {self._dim})'
+
+    def dimensions(self):
+        return self._dim
+
+    def indices(self):
+        return self._indices
+
+    def values(self):
+        return self._values
+
+    def to_coo(self):
+        from scipy.sparse import coo_array
+
+        coords = ([0] * len(self._indices), self._indices)
+        return coo_array((self._values, coords), shape=(1, self._dim))
+
+    def to_list(self):
+        vec = [0.0] * self._dim
+        for i, v in zip(self._indices, self._values):
+            vec[i] = v
+        return vec
+
+    def to_numpy(self):
+        vec = np.repeat(0.0, self._dim).astype(np.float32)
+        for i, v in zip(self._indices, self._values):
+            vec[i] = v
+        return vec
+
+    def to_text(self):
+        return '{' + ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + '}/' + str(int(self._dim))
+
+    def to_binary(self):
+        nnz = len(self._indices)
+        return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values)
+
+    def _from_dict(self, d, dim):
+        elements = [(i, v) for i, v in d.items() if v != 0]
+        elements.sort()
+
+        self._dim = int(dim)
+        self._indices = [int(v[0]) for v in elements]
+        self._values = [float(v[1]) for v in elements]
+
+    def _from_sparse(self, value):
+        value = value.tocoo()
+
+        if value.ndim == 1:
+            self._dim = value.shape[0]
+        elif value.ndim == 2 and value.shape[0] == 1:
+            self._dim = value.shape[1]
+        else:
+            raise ValueError('expected ndim to be 1')
+
+        if hasattr(value, 'coords'):
+            # scipy 1.13+
+            self._indices = value.coords[0].tolist()
+        else:
+            self._indices = value.col.tolist()
+        self._values = value.data.tolist()
+
+    def _from_dense(self, value):
+        self._dim = len(value)
+        self._indices = [i for i, v in enumerate(value) if v != 0]
+        self._values = [float(value[i]) for i in self._indices]
+
+    @classmethod
+    def from_text(cls, value):
+        elements, dim = value.split('/', 2)
+        indices = []
+        values = []
+        # split on empty string returns single element list
+        if len(elements) > 2:
+            for e in elements[1:-1].split(','):
+                i, v = e.split(':', 2)
+                indices.append(int(i) - 1)
+                values.append(float(v))
+        return cls._from_parts(int(dim), indices, values)
+
+    @classmethod
+    def from_binary(cls, value):
+        dim, nnz, unused = unpack_from('>iii', value)
+        indices = unpack_from(f'>{nnz}i', value, 12)
+        values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4)
+        return cls._from_parts(int(dim), indices, values)
+
+    @classmethod
+    def _from_parts(cls, dim, indices, values):
+        vec = cls.__new__(cls)
+        vec._dim = dim
+        vec._indices = indices
+        vec._values = values
+        return vec
+
+    @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)

pgvector/utils/vector.py

@@ -0,0 +1,78 @@
+import numpy as np
+from struct import pack, unpack_from
+
+
+class Vector:
+    def __init__(self, value):
+        # asarray still copies if same dtype
+        if not isinstance(value, np.ndarray) or value.dtype != '>f4':
+            value = np.asarray(value, dtype='>f4')
+
+        if value.ndim != 1:
+            raise ValueError('expected ndim to be 1')
+
+        self._value = value
+
+    def __repr__(self):
+        return f'Vector({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='>f4', 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, np.ndarray):
+            return value
+
+        return cls.from_text(value).to_numpy().astype(np.float32)
+
+    @classmethod
+    def _from_db_binary(cls, value):
+        if value is None or isinstance(value, np.ndarray):
+            return value
+
+        return cls.from_binary(value).to_numpy().astype(np.float32)

pgvector-0.3.5.dist-info/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2021-2024 Andrew Kane
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.