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

ankane · ankane · commit f1cf31e0f7ae · 2024-05-16T18:32:32.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,7 +1,7 @@
-## 0.2.6 (unreleased)
+## 0.3.0 (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 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,4 +1,6 @@
 from .halfvec import Halfvec
+from .sparsevec import Sparsevec
 from .vector import Vector
+from ..utils import SparseVec
 
-__all__ = ['Vector', 'Halfvec']
+__all__ = ['Vector', 'Halfvec', 'Sparsevec']
diff --git a/pgvector/sqlalchemy/sparsevec.py b/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 SparseVec
+
+
+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 SparseVec.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(SparseVec.to_db(value, self.dim))
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return SparseVec.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
diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py
@@ -19,9 +19,13 @@ def to_dense(self):
             vec[i] = v
         return vec
 
-    def to_db(value):
+    def to_db(value, dim=None):
         if value is None:
             return value
+
+        if dim is not None and value.dim != dim:
+            raise ValueError('expected %d dimensions, not %d' % (dim, len(value)))
+
         return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(value.indices, value.values)]) + '}/' + str(value.dim)
 
     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, Halfvec
+from pgvector.sqlalchemy import Vector, Halfvec, Sparsevec, SparseVec
 import pytest
 from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer
 from sqlalchemy.exc import StatementError
@@ -21,6 +21,7 @@ class Item(Base):
     id = mapped_column(Integer, primary_key=True)
     embedding = mapped_column(Vector(3))
     half_embedding = mapped_column(Halfvec(3))
+    sparse_embedding = mapped_column(Sparsevec(3))
 
 
 Base.metadata.drop_all(engine)
@@ -44,7 +45,7 @@ def create_items():
     ]
     session = Session(engine)
     for i, v in enumerate(vectors):
-        session.add(Item(id=i + 1, embedding=v, half_embedding=v))
+        session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v)))
     session.commit()
 
 
