Added support for halfvec type to SQLAlchemy and SQLModel [skip ci]

ankane · ankane · commit b4b6c6b9f53a · 2024-05-16T18:28:59.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.2.6 (unreleased)
 
 - Added support for `halfvec` and `sparsevec` types to Django
+- Added support for `halfvec` type to SQLAlchemy and SQLModel
 - Added support for `halfvec` and `sparsevec` types to Psycopg 3
 - Added support for `halfvec` and `sparsevec` types to Psycopg 2
 - Added support for `halfvec` and `sparsevec` types to asyncpg
diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py
@@ -1,3 +1,4 @@
+from .halfvec import Halfvec
 from .vector import Vector
 
-__all__ = ['Vector']
+__all__ = ['Vector', 'Halfvec']
diff --git a/pgvector/sqlalchemy/halfvec.py b/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 HalfVec
+
+
+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 HalfVec.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(HalfVec.to_db(value, self.dim))
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return HalfVec.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
diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py
@@ -12,11 +12,15 @@ def __init__(self, value):
     def to_list(self):
         return list(self.value)
 
-    def to_db(value):
+    def to_db(value, dim=None):
         if value is None:
             return value
         if isinstance(value, HalfVec):
             value = value.value
+
+        if dim is not None and len(value) != dim:
+            raise ValueError('expected %d dimensions, not %d' % (dim, len(value)))
+
         return '[' + ','.join([str(float(v)) for v in value]) + ']'
 
     def to_db_binary(value):
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
@@ -1,5 +1,5 @@
 import numpy as np
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import Vector, Halfvec
 import pytest
 from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer
 from sqlalchemy.exc import StatementError
@@ -20,6 +20,7 @@ class Item(Base):
 
     id = mapped_column(Integer, primary_key=True)
     embedding = mapped_column(Vector(3))
+    half_embedding = mapped_column(Halfvec(3))
 
 
 Base.metadata.drop_all(engine)
@@ -43,7 +44,7 @@ def create_items():
     ]
     session = Session(engine)
     for i, v in enumerate(vectors):
-        session.add(Item(id=i + 1, embedding=v))
+        session.add(Item(id=i + 1, embedding=v, half_embedding=v))
     session.commit()
 
 
