Added tests for halfvec and sparsevec types with SQLAlchemy and asyncpg [skip ci]

Author: ankane

tests/test_sqlalchemy.py

@@ -550,6 +550,29 @@ async def test_asyncpg_vector(self):
 
         await engine.dispose()
 
+    @pytest.mark.asyncio
+    @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
+    async def test_asyncpg_halfvec(self):
+        import asyncpg
+
+        engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test')
+        async_session = async_sessionmaker(engine, expire_on_commit=False)
+
+        # TODO do not throw error when types are registered
+        # @event.listens_for(engine.sync_engine, "connect")
+        # def connect(dbapi_connection, connection_record):
+        #     from pgvector.asyncpg import register_vector
+        #     dbapi_connection.run_async(register_vector)
+
+        async with async_session() as session:
+            async with session.begin():
+                embedding = [1, 2, 3]
+                session.add(Item(id=1, half_embedding=embedding))
+                item = await session.get(Item, 1)
+                assert item.half_embedding.to_list() == embedding
+
+        await engine.dispose()
+
     @pytest.mark.asyncio
     @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
     async def test_asyncpg_bit(self):
@@ -566,3 +589,26 @@ async def test_asyncpg_bit(self):
                 assert item.binary_embedding == embedding
 
         await engine.dispose()
+
+    @pytest.mark.asyncio
+    @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
+    async def test_asyncpg_sparsevec(self):
+        import asyncpg
+
+        engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test')
+        async_session = async_sessionmaker(engine, expire_on_commit=False)
+
+        # TODO do not throw error when types are registered
+        # @event.listens_for(engine.sync_engine, "connect")
+        # def connect(dbapi_connection, connection_record):
+        #     from pgvector.asyncpg import register_vector
+        #     dbapi_connection.run_async(register_vector)
+
+        async with async_session() as session:
+            async with session.begin():
+                embedding = [1, 2, 3]
+                session.add(Item(id=1, sparse_embedding=embedding))
+                item = await session.get(Item, 1)
+                assert item.sparse_embedding.to_list() == embedding
+
+        await engine.dispose()