pgvector 0.3.5__py3-none-any.whl

pgvector/psycopg/sparsevec.py

@@ -0,0 +1,53 @@
+from psycopg.adapt import Loader, Dumper
+from psycopg.pq import Format
+from ..utils import SparseVector
+
+
+class SparseVectorDumper(Dumper):
+
+    format = Format.TEXT
+
+    def dump(self, obj):
+        return SparseVector._to_db(obj).encode('utf8')
+
+
+class SparseVectorBinaryDumper(SparseVectorDumper):
+
+    format = Format.BINARY
+
+    def dump(self, obj):
+        return SparseVector._to_db_binary(obj)
+
+
+class SparseVectorLoader(Loader):
+
+    format = Format.TEXT
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return SparseVector._from_db(data.decode('utf8'))
+
+
+class SparseVectorBinaryLoader(SparseVectorLoader):
+
+    format = Format.BINARY
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return SparseVector._from_db_binary(data)
+
+
+def register_sparsevec_info(context, info):
+    info.register(context)
+
+    # add oid to anonymous class for set_types
+    text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid})
+    binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid})
+
+    adapters = context.adapters
+    adapters.register_dumper(SparseVector, text_dumper)
+    adapters.register_dumper(SparseVector, binary_dumper)
+    adapters.register_loader(info.oid, SparseVectorLoader)
+    adapters.register_loader(info.oid, SparseVectorBinaryLoader)

pgvector/psycopg/vector.py

@@ -0,0 +1,58 @@
+import psycopg
+from psycopg.adapt import Loader, Dumper
+from psycopg.pq import Format
+from ..utils import Vector
+
+
+class VectorDumper(Dumper):
+
+    format = Format.TEXT
+
+    def dump(self, obj):
+        return Vector._to_db(obj).encode('utf8')
+
+
+class VectorBinaryDumper(VectorDumper):
+
+    format = Format.BINARY
+
+    def dump(self, obj):
+        return Vector._to_db_binary(obj)
+
+
+class VectorLoader(Loader):
+
+    format = Format.TEXT
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return Vector._from_db(data.decode('utf8'))
+
+
+class VectorBinaryLoader(VectorLoader):
+
+    format = Format.BINARY
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return Vector._from_db_binary(data)
+
+
+def register_vector_info(context, info):
+    if info is None:
+        raise psycopg.ProgrammingError('vector type not found in the database')
+    info.register(context)
+
+    # add oid to anonymous class for set_types
+    text_dumper = type('', (VectorDumper,), {'oid': info.oid})
+    binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid})
+
+    adapters = context.adapters
+    adapters.register_dumper('numpy.ndarray', text_dumper)
+    adapters.register_dumper('numpy.ndarray', binary_dumper)
+    adapters.register_dumper(Vector, text_dumper)
+    adapters.register_dumper(Vector, binary_dumper)
+    adapters.register_loader(info.oid, VectorLoader)
+    adapters.register_loader(info.oid, VectorBinaryLoader)

pgvector/psycopg2/__init__.py

@@ -0,0 +1,8 @@
+from .register import register_vector
+from ..utils import HalfVector, SparseVector
+
+__all__ = [
+    'register_vector',
+    'HalfVector',
+    'SparseVector'
+]

pgvector/psycopg2/halfvec.py

@@ -0,0 +1,20 @@
+from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from ..utils import HalfVector
+
+
+class HalfvecAdapter:
+    def __init__(self, value):
+        self._value = value
+
+    def getquoted(self):
+        return adapt(HalfVector._to_db(self._value)).getquoted()
+
+
+def cast_halfvec(value, cur):
+    return HalfVector._from_db(value)
+
+
+def register_halfvec_info(oid, scope):
+    halfvec = new_type((oid,), 'HALFVEC', cast_halfvec)
+    register_type(halfvec, scope)
+    register_adapter(HalfVector, HalfvecAdapter)

pgvector/psycopg2/register.py

@@ -0,0 +1,28 @@
+import psycopg2
+from psycopg2.extensions import cursor
+from .halfvec import register_halfvec_info
+from .sparsevec import register_sparsevec_info
+from .vector import register_vector_info
+
+
+# TODO make globally False by default in 0.4.0
+# note: register_adapter is always global
+def register_vector(conn_or_curs=None, globally=True):
+    conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection
+    cur = conn.cursor(cursor_factory=cursor)
+    scope = None if globally else conn_or_curs
+
+    # use to_regtype to get first matching type in search path
+    cur.execute("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")
+    type_info = dict(cur.fetchall())
+
+    if 'vector' not in type_info:
+        raise psycopg2.ProgrammingError('vector type not found in the database')
+
+    register_vector_info(type_info['vector'], scope)
+
+    if 'halfvec' in type_info:
+        register_halfvec_info(type_info['halfvec'], scope)
+
+    if 'sparsevec' in type_info:
+        register_sparsevec_info(type_info['sparsevec'], scope)

pgvector/psycopg2/sparsevec.py

@@ -0,0 +1,20 @@
+from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from ..utils import SparseVector
+
+
+class SparsevecAdapter:
+    def __init__(self, value):
+        self._value = value
+
+    def getquoted(self):
+        return adapt(SparseVector._to_db(self._value)).getquoted()
+
+
+def cast_sparsevec(value, cur):
+    return SparseVector._from_db(value)
+
+
+def register_sparsevec_info(oid, scope):
+    sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec)
+    register_type(sparsevec, scope)
+    register_adapter(SparseVector, SparsevecAdapter)

pgvector/psycopg2/vector.py

@@ -0,0 +1,21 @@
+import numpy as np
+from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from ..utils import Vector
+
+
+class VectorAdapter:
+    def __init__(self, value):
+        self._value = value
+
+    def getquoted(self):
+        return adapt(Vector._to_db(self._value)).getquoted()
+
+
+def cast_vector(value, cur):
+    return Vector._from_db(value)
+
+
+def register_vector_info(oid, scope):
+    vector = new_type((oid,), 'VECTOR', cast_vector)
+    register_type(vector, scope)
+    register_adapter(np.ndarray, VectorAdapter)

pgvector/sqlalchemy/__init__.py

@@ -0,0 +1,19 @@
+from .bit import BIT
+from .functions import avg, sum
+from .halfvec import HALFVEC
+from .sparsevec import SPARSEVEC
+from .vector import VECTOR
+from .vector import VECTOR as Vector
+from ..utils import HalfVector, SparseVector
+
+__all__ = [
+    'Vector',
+    'VECTOR',
+    'HALFVEC',
+    'BIT',
+    'SPARSEVEC',
+    'HalfVector',
+    'SparseVector',
+    'avg',
+    'sum'
+]

pgvector/sqlalchemy/bit.py

@@ -0,0 +1,26 @@
+from sqlalchemy.dialects.postgresql.base import ischema_names
+from sqlalchemy.types import UserDefinedType, Float
+
+
+class BIT(UserDefinedType):
+    cache_ok = True
+
+    def __init__(self, length=None):
+        super(UserDefinedType, self).__init__()
+        self.length = length
+
+    def get_col_spec(self, **kw):
+        if self.length is None:
+            return 'BIT'
+        return 'BIT(%d)' % self.length
+
+    class comparator_factory(UserDefinedType.Comparator):
+        def hamming_distance(self, other):
+            return self.op('<~>', return_type=Float)(other)
+
+        def jaccard_distance(self, other):
+            return self.op('<%>', return_type=Float)(other)
+
+
+# for reflection
+ischema_names['bit'] = BIT

pgvector/sqlalchemy/functions.py

@@ -0,0 +1,14 @@
+# https://docs.sqlalchemy.org/en/20/core/functions.html
+# include sum for a consistent API
+from sqlalchemy.sql.functions import ReturnTypeFromArgs, sum
+
+
+class avg(ReturnTypeFromArgs):
+    inherit_cache = True
+    package = 'pgvector'
+
+
+__all__ = [
+    'avg',
+    'sum'
+]

pgvector/sqlalchemy/halfvec.py

@@ -0,0 +1,51 @@
+from sqlalchemy.dialects.postgresql.base import ischema_names
+from sqlalchemy.types import UserDefinedType, Float, String
+from ..utils import HalfVector
+
+
+class HALFVEC(UserDefinedType):
+    cache_ok = True
+    _string = String()
+
+    def __init__(self, dim=None):
+        super(UserDefinedType, self).__init__()
+        self.dim = dim
+
+    def get_col_spec(self, **kw):
+        if self.dim is None:
+            return 'HALFVEC'
+        return 'HALFVEC(%d)' % self.dim
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return HalfVector._to_db(value, self.dim)
+        return process
+
+    def literal_processor(self, dialect):
+        string_literal_processor = self._string._cached_literal_processor(dialect)
+
+        def process(value):
+            return string_literal_processor(HalfVector._to_db(value, self.dim))
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return HalfVector._from_db(value)
+        return process
+
+    class comparator_factory(UserDefinedType.Comparator):
+        def l2_distance(self, other):
+            return self.op('<->', return_type=Float)(other)
+
+        def max_inner_product(self, other):
+            return self.op('<#>', return_type=Float)(other)
+
+        def cosine_distance(self, other):
+            return self.op('<=>', return_type=Float)(other)
+
+        def l1_distance(self, other):
+            return self.op('<+>', return_type=Float)(other)
+
+
+# for reflection
+ischema_names['halfvec'] = HALFVEC

pgvector/sqlalchemy/sparsevec.py

@@ -0,0 +1,51 @@
+from sqlalchemy.dialects.postgresql.base import ischema_names
+from sqlalchemy.types import UserDefinedType, Float, String
+from ..utils import SparseVector
+
+
+class SPARSEVEC(UserDefinedType):
+    cache_ok = True
+    _string = String()
+
+    def __init__(self, dim=None):
+        super(UserDefinedType, self).__init__()
+        self.dim = dim
+
+    def get_col_spec(self, **kw):
+        if self.dim is None:
+            return 'SPARSEVEC'
+        return 'SPARSEVEC(%d)' % self.dim
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return SparseVector._to_db(value, self.dim)
+        return process
+
+    def literal_processor(self, dialect):
+        string_literal_processor = self._string._cached_literal_processor(dialect)
+
+        def process(value):
+            return string_literal_processor(SparseVector._to_db(value, self.dim))
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return SparseVector._from_db(value)
+        return process
+
+    class comparator_factory(UserDefinedType.Comparator):
+        def l2_distance(self, other):
+            return self.op('<->', return_type=Float)(other)
+
+        def max_inner_product(self, other):
+            return self.op('<#>', return_type=Float)(other)
+
+        def cosine_distance(self, other):
+            return self.op('<=>', return_type=Float)(other)
+
+        def l1_distance(self, other):
+            return self.op('<+>', return_type=Float)(other)
+
+
+# for reflection
+ischema_names['sparsevec'] = SPARSEVEC

pgvector/sqlalchemy/vector.py

@@ -0,0 +1,51 @@
+from sqlalchemy.dialects.postgresql.base import ischema_names
+from sqlalchemy.types import UserDefinedType, Float, String
+from ..utils import Vector
+
+
+class VECTOR(UserDefinedType):
+    cache_ok = True
+    _string = String()
+
+    def __init__(self, dim=None):
+        super(UserDefinedType, self).__init__()
+        self.dim = dim
+
+    def get_col_spec(self, **kw):
+        if self.dim is None:
+            return 'VECTOR'
+        return 'VECTOR(%d)' % self.dim
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return Vector._to_db(value, self.dim)
+        return process
+
+    def literal_processor(self, dialect):
+        string_literal_processor = self._string._cached_literal_processor(dialect)
+
+        def process(value):
+            return string_literal_processor(Vector._to_db(value, self.dim))
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return Vector._from_db(value)
+        return process
+
+    class comparator_factory(UserDefinedType.Comparator):
+        def l2_distance(self, other):
+            return self.op('<->', return_type=Float)(other)
+
+        def max_inner_product(self, other):
+            return self.op('<#>', return_type=Float)(other)
+
+        def cosine_distance(self, other):
+            return self.op('<=>', return_type=Float)(other)
+
+        def l1_distance(self, other):
+            return self.op('<+>', return_type=Float)(other)
+
+
+# for reflection
+ischema_names['vector'] = VECTOR

pgvector/utils/__init__.py

@@ -0,0 +1,11 @@
+from .bit import Bit
+from .halfvec import HalfVector
+from .sparsevec import SparseVector
+from .vector import Vector
+
+__all__ = [
+    'Vector',
+    'HalfVector',
+    'Bit',
+    'SparseVector'
+]

pgvector/utils/bit.py

@@ -0,0 +1,61 @@
+import numpy as np
+from struct import pack, unpack_from
+
+
+class Bit:
+    def __init__(self, value):
+        if isinstance(value, str):
+            self._value = self.from_text(value)._value
+        else:
+            # TODO change in 0.4.0
+            # TODO raise if dtype not bool or uint8
+            # if isinstance(value, np.ndarray) and value.dtype == np.uint8:
+            #     value = np.unpackbits(value)
+            # else:
+            #     value = np.asarray(value, dtype=bool)
+
+            value = np.asarray(value, dtype=bool)
+
+            if value.ndim != 1:
+                raise ValueError('expected ndim to be 1')
+
+            self._value = value
+
+    def __repr__(self):
+        return f'Bit({self.to_text()})'
+
+    def to_list(self):
+        return self._value.tolist()
+
+    def to_numpy(self):
+        return self._value
+
+    def to_text(self):
+        return ''.join(self._value.astype(np.uint8).astype(str))
+
+    def to_binary(self):
+        return pack('>i', len(self._value)) + np.packbits(self._value).tobytes()
+
+    @classmethod
+    def from_text(cls, value):
+        return cls(np.asarray([v != '0' for v in value], dtype=bool))
+
+    @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))
+
+    @classmethod
+    def _to_db(cls, value):
+        if not isinstance(value, cls):
+            raise ValueError('expected bit')
+
+        return value.to_text()
+
+    @classmethod
+    def _to_db_binary(cls, value):
+        if not isinstance(value, cls):
+            raise ValueError('expected bit')
+
+        return value.to_binary()

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)