Added avg function with type casting to SQLAlchemy - #44

Co-authored-by: lucasgadams

Author: ankane

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 0.3.5 (unreleased)
 
+- Added `avg` function with type casting to SQLAlchemy
 - Added `globally` option for Psycopg 2
 
 ## 0.3.4 (2024-09-26)

pgvector/sqlalchemy/__init__.py

@@ -1,4 +1,5 @@
 from .bit import BIT
+from .functions import avg, sum
 from .halfvec import HALFVEC
 from .sparsevec import SPARSEVEC
 from .vector import VECTOR
@@ -12,5 +13,7 @@
     'BIT',
     'SPARSEVEC',
     'HalfVector',
-    'SparseVector'
+    'SparseVector',
+    'avg',
+    'sum'
 ]

pgvector/sqlalchemy/functions.py

@@ -0,0 +1,8 @@
+# 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'

tests/test_sqlalchemy.py

@@ -1,5 +1,5 @@
 import numpy as np
-from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector
+from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum
 import pytest
 from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer
 from sqlalchemy.exc import StatementError
@@ -339,41 +339,39 @@ def test_select_orm(self):
 
     def test_avg(self):
         with Session(engine) as session:
-            avg = session.query(func.avg(Item.embedding)).first()[0]
-            assert avg is None
+            res = session.query(avg(Item.embedding)).first()[0]
+            assert res is None
             session.add(Item(embedding=[1, 2, 3]))
             session.add(Item(embedding=[4, 5, 6]))
-            avg = session.query(func.avg(Item.embedding)).first()[0]
-            # does not type cast
-            assert avg == '[2.5,3.5,4.5]'
+            res = session.query(avg(Item.embedding)).first()[0]
+            assert np.array_equal(res, np.array([2.5, 3.5, 4.5]))
 
     def test_avg_orm(self):
         with Session(engine) as session:
-            avg = session.scalars(select(func.avg(Item.embedding))).first()
-            assert avg is None
+            res = session.scalars(select(avg(Item.embedding))).first()
+            assert res is None
             session.add(Item(embedding=[1, 2, 3]))
             session.add(Item(embedding=[4, 5, 6]))
-            avg = session.scalars(select(func.avg(Item.embedding))).first()
-            # does not type cast
-            assert avg == '[2.5,3.5,4.5]'
+            res = session.scalars(select(avg(Item.embedding))).first()
+            assert np.array_equal(res, np.array([2.5, 3.5, 4.5]))
 
     def test_sum(self):
         with Session(engine) as session:
-            sum = session.query(func.sum(Item.embedding)).first()[0]
-            assert sum is None
+            res = session.query(sum(Item.embedding)).first()[0]
+            assert res is None
             session.add(Item(embedding=[1, 2, 3]))
             session.add(Item(embedding=[4, 5, 6]))
-            sum = session.query(func.sum(Item.embedding)).first()[0]
-            assert np.array_equal(sum, np.array([5, 7, 9]))
+            res = session.query(sum(Item.embedding)).first()[0]
+            assert np.array_equal(res, np.array([5, 7, 9]))
 
     def test_sum_orm(self):
         with Session(engine) as session:
-            sum = session.scalars(select(func.sum(Item.embedding))).first()
-            assert sum is None
+            res = session.scalars(select(sum(Item.embedding))).first()
+            assert res is None
             session.add(Item(embedding=[1, 2, 3]))
             session.add(Item(embedding=[4, 5, 6]))
-            sum = session.scalars(select(func.sum(Item.embedding))).first()
-            assert np.array_equal(sum, np.array([5, 7, 9]))
+            res = session.scalars(select(sum(Item.embedding))).first()
+            assert np.array_equal(res, np.array([5, 7, 9]))
 
     def test_bad_dimensions(self):
         item = Item(embedding=[1, 2])

tests/test_sqlmodel.py

@@ -1,5 +1,5 @@
 import numpy as np
-from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector
+from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum
 import pytest
 from sqlalchemy import Column, Index
 from sqlalchemy.exc import StatementError
@@ -198,41 +198,39 @@ def test_select(self):
 
     def test_vector_avg(self):
         with Session(engine) as session:
-            avg = session.exec(select(func.avg(Item.embedding))).first()
-            assert avg is None
+            res = session.exec(select(avg(Item.embedding))).first()
+            assert res is None
             session.add(Item(embedding=[1, 2, 3]))
             session.add(Item(embedding=[4, 5, 6]))
-            avg = session.exec(select(func.avg(Item.embedding))).first()
-            # does not type cast
-            assert avg == '[2.5,3.5,4.5]'
+            res = session.exec(select(avg(Item.embedding))).first()
+            assert np.array_equal(res, np.array([2.5, 3.5, 4.5]))
 
     def test_vector_sum(self):
         with Session(engine) as session:
-            sum = session.exec(select(func.sum(Item.embedding))).first()
-            assert sum is None
+            res = session.exec(select(sum(Item.embedding))).first()
+            assert res is None
             session.add(Item(embedding=[1, 2, 3]))
             session.add(Item(embedding=[4, 5, 6]))
-            sum = session.exec(select(func.sum(Item.embedding))).first()
-            assert np.array_equal(sum, np.array([5, 7, 9]))
+            res = session.exec(select(sum(Item.embedding))).first()
+            assert np.array_equal(res, np.array([5, 7, 9]))
 
     def test_halfvec_avg(self):
         with Session(engine) as session:
-            avg = session.exec(select(func.avg(Item.half_embedding))).first()
-            assert avg is None
+            res = session.exec(select(avg(Item.half_embedding))).first()
+            assert res is None
             session.add(Item(half_embedding=[1, 2, 3]))
             session.add(Item(half_embedding=[4, 5, 6]))
-            avg = session.exec(select(func.avg(Item.half_embedding))).first()
-            # does not type cast
-            assert avg == '[2.5,3.5,4.5]'
+            res = session.exec(select(avg(Item.half_embedding))).first()
+            assert res.to_list() == [2.5, 3.5, 4.5]
 
     def test_halfvec_sum(self):
         with Session(engine) as session:
-            sum = session.exec(select(func.sum(Item.half_embedding))).first()
-            assert sum is None
+            res = session.exec(select(sum(Item.half_embedding))).first()
+            assert res is None
             session.add(Item(half_embedding=[1, 2, 3]))
             session.add(Item(half_embedding=[4, 5, 6]))
-            sum = session.exec(select(func.sum(Item.half_embedding))).first()
-            assert sum.to_list() == [5, 7, 9]
+            res = session.exec(select(sum(Item.half_embedding))).first()
+            assert res.to_list() == [5, 7, 9]
 
     def test_bad_dimensions(self):
         item = Item(embedding=[1, 2])