pgvector 0.3.5__tar.gz → 0.3.6__tar.gz

{pgvector-0.3.5 → pgvector-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector
-Version: 0.3.5
+Version: 0.3.6
 Summary: pgvector support for Python
 Author-email: Andrew Kane <andrew@ankane.org>
 License: MIT
@@ -197,7 +197,7 @@ Average vectors
 ```python
 from pgvector.sqlalchemy import avg
 
-session.scalars(select(func.avg(Item.embedding))).first()
+session.scalars(select(avg(Item.embedding))).first()
 ```
 
 Also supports `sum`
@@ -279,7 +279,7 @@ Average vectors
 ```python
 from pgvector.sqlalchemy import avg
 
-session.exec(select(func.avg(Item.embedding))).first()
+session.exec(select(avg(Item.embedding))).first()
 ```
 
 Also supports `sum`

{pgvector-0.3.5 → pgvector-0.3.6}/README.md

@@ -185,7 +185,7 @@ Average vectors
 ```python
 from pgvector.sqlalchemy import avg
 
-session.scalars(select(func.avg(Item.embedding))).first()
+session.scalars(select(avg(Item.embedding))).first()
 ```
 
 Also supports `sum`
@@ -267,7 +267,7 @@ Average vectors
 ```python
 from pgvector.sqlalchemy import avg
 
-session.exec(select(func.avg(Item.embedding))).first()
+session.exec(select(avg(Item.embedding))).first()
 ```
 
 Also supports `sum`

{pgvector-0.3.5 → pgvector-0.3.6}/pgvector/psycopg2/halfvec.py

@@ -1,4 +1,4 @@
-from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type
 from ..utils import HalfVector
 
 
@@ -14,7 +14,12 @@ def cast_halfvec(value, cur):
     return HalfVector._from_db(value)
 
 
-def register_halfvec_info(oid, scope):
+def register_halfvec_info(oid, array_oid, scope):
     halfvec = new_type((oid,), 'HALFVEC', cast_halfvec)
     register_type(halfvec, scope)
+
+    if array_oid is not None:
+        halfvecarray = new_array_type((array_oid,), 'HALFVECARRAY', halfvec)
+        register_type(halfvecarray, scope)
+
     register_adapter(HalfVector, HalfvecAdapter)

{pgvector-0.3.5 → pgvector-0.3.6}/pgvector/psycopg2/register.py

@@ -7,22 +7,23 @@ 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):
+# TODO make arrays True by defalt in 0.4.0
+def register_vector(conn_or_curs=None, globally=True, arrays=False):
     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'))")
+    cur.execute("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('_vector'), to_regtype('halfvec'), to_regtype('_halfvec'), to_regtype('sparsevec'), 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)
+    register_vector_info(type_info['vector'], type_info['_vector'] if arrays else None, scope)
 
     if 'halfvec' in type_info:
-        register_halfvec_info(type_info['halfvec'], scope)
+        register_halfvec_info(type_info['halfvec'], type_info['_halfvec'] if arrays else None, scope)
 
     if 'sparsevec' in type_info:
-        register_sparsevec_info(type_info['sparsevec'], scope)
+        register_sparsevec_info(type_info['sparsevec'], type_info['_sparsevec'] if arrays else None, scope)

{pgvector-0.3.5 → pgvector-0.3.6}/pgvector/psycopg2/sparsevec.py

@@ -1,4 +1,4 @@
-from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type
 from ..utils import SparseVector
 
 
@@ -14,7 +14,12 @@ def cast_sparsevec(value, cur):
     return SparseVector._from_db(value)
 
 
-def register_sparsevec_info(oid, scope):
+def register_sparsevec_info(oid, array_oid, scope):
     sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec)
     register_type(sparsevec, scope)
+
+    if array_oid is not None:
+        sparsevecarray = new_array_type((array_oid,), 'SPARSEVECARRAY', sparsevec)
+        register_type(sparsevecarray, scope)
+
     register_adapter(SparseVector, SparsevecAdapter)

{pgvector-0.3.5 → pgvector-0.3.6}/pgvector/psycopg2/vector.py

@@ -1,5 +1,5 @@
 import numpy as np
-from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type
 from ..utils import Vector
 
 
@@ -15,7 +15,12 @@ def cast_vector(value, cur):
     return Vector._from_db(value)
 
 
-def register_vector_info(oid, scope):
+def register_vector_info(oid, array_oid, scope):
     vector = new_type((oid,), 'VECTOR', cast_vector)
     register_type(vector, scope)
+
+    if array_oid is not None:
+        vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector)
+        register_type(vectorarray, scope)
+
     register_adapter(np.ndarray, VectorAdapter)

{pgvector-0.3.5 → pgvector-0.3.6}/pgvector.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pgvector
-Version: 0.3.5
+Version: 0.3.6
 Summary: pgvector support for Python
 Author-email: Andrew Kane <andrew@ankane.org>
 License: MIT
@@ -197,7 +197,7 @@ Average vectors
 ```python
 from pgvector.sqlalchemy import avg
 
-session.scalars(select(func.avg(Item.embedding))).first()
+session.scalars(select(avg(Item.embedding))).first()
 ```
 
 Also supports `sum`
@@ -279,7 +279,7 @@ Average vectors
 ```python
 from pgvector.sqlalchemy import avg
 
-session.exec(select(func.avg(Item.embedding))).first()
+session.exec(select(avg(Item.embedding))).first()
 ```
 
 Also supports `sum`

{pgvector-0.3.5 → pgvector-0.3.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgvector"
-version = "0.3.5"
+version = "0.3.6"
 description = "pgvector support for Python"
 readme = "README.md"
 authors = [

{pgvector-0.3.5 → pgvector-0.3.6}/tests/test_asyncpg.py

@@ -94,6 +94,24 @@ class TestAsyncpg:
 
         await conn.close()
 
+    @pytest.mark.asyncio
+    async def test_vector_array(self):
+        conn = await asyncpg.connect(database='pgvector_python_test')
+        await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
+        await conn.execute('DROP TABLE IF EXISTS asyncpg_items')
+        await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])')
+
+        await register_vector(conn)
+
+        embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings[0], embeddings[1])
+
+        res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
+        assert np.array_equal(res[0]['embeddings'][0], embeddings[0])
+        assert np.array_equal(res[0]['embeddings'][1], embeddings[1])
+
+        await conn.close()
+
     @pytest.mark.asyncio
     async def test_pool(self):
         async def init(conn):

{pgvector-0.3.5 → pgvector-0.3.6}/tests/test_django.py

@@ -1,8 +1,10 @@
 import django
 from django.conf import settings
+from django.contrib.postgres.fields import ArrayField
 from django.core import serializers
 from django.db import connection, migrations, models
-from django.db.models import Avg, Sum
+from django.db.models import Avg, Sum, FloatField, DecimalField
+from django.db.models.functions import Cast
 from django.db.migrations.loader import MigrationLoader
 from django.forms import ModelForm
 from math import sqrt
@@ -46,6 +48,9 @@ class Item(models.Model):
     half_embedding = HalfVectorField(dimensions=3, null=True, blank=True)
     binary_embedding = BitField(length=3, null=True, blank=True)
     sparse_embedding = SparseVectorField(dimensions=3, null=True, blank=True)
+    embeddings = ArrayField(VectorField(dimensions=3), null=True, blank=True)
+    double_embedding = ArrayField(FloatField(), null=True, blank=True)
+    numeric_embedding = ArrayField(DecimalField(max_digits=20, decimal_places=10), null=True, blank=True)
 
     class Meta:
         app_label = 'django_app'
@@ -82,6 +87,9 @@ class Migration(migrations.Migration):
                 ('half_embedding', pgvector.django.HalfVectorField(dimensions=3, null=True, blank=True)),
                 ('binary_embedding', pgvector.django.BitField(length=3, null=True, blank=True)),
                 ('sparse_embedding', pgvector.django.SparseVectorField(dimensions=3, null=True, blank=True)),
+                ('embeddings', ArrayField(pgvector.django.VectorField(dimensions=3), null=True, blank=True)),
+                ('double_embedding', ArrayField(FloatField(), null=True, blank=True)),
+                ('numeric_embedding', ArrayField(DecimalField(max_digits=20, decimal_places=10), null=True, blank=True)),
             ],
         ),
         migrations.AddIndex(
@@ -433,3 +441,35 @@ class TestDjango:
         assert Item.objects.first().half_embedding is None
         assert Item.objects.first().binary_embedding is None
         assert Item.objects.first().sparse_embedding is None
+
+    def test_vector_array(self):
+        Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]).save()
+
+        with connection.cursor() as cursor:
+            from pgvector.psycopg import register_vector
+            register_vector(cursor.connection)
+
+            # this fails if the driver does not cast arrays
+            item = Item.objects.get(pk=1)
+            assert item.embeddings[0].tolist() == [1, 2, 3]
+            assert item.embeddings[1].tolist() == [4, 5, 6]
+
+    def test_double_array(self):
+        Item(id=1, double_embedding=[1, 1, 1]).save()
+        Item(id=2, double_embedding=[2, 2, 2]).save()
+        Item(id=3, double_embedding=[1, 1, 2]).save()
+        distance = L2Distance(Cast('double_embedding', VectorField()), [1, 1, 1])
+        items = Item.objects.annotate(distance=distance).order_by(distance)
+        assert [v.id for v in items] == [1, 3, 2]
+        assert [v.distance for v in items] == [0, 1, sqrt(3)]
+        assert items[1].double_embedding == [1, 1, 2]
+
+    def test_numeric_array(self):
+        Item(id=1, numeric_embedding=[1, 1, 1]).save()
+        Item(id=2, numeric_embedding=[2, 2, 2]).save()
+        Item(id=3, numeric_embedding=[1, 1, 2]).save()
+        distance = L2Distance(Cast('numeric_embedding', VectorField()), [1, 1, 1])
+        items = Item.objects.annotate(distance=distance).order_by(distance)
+        assert [v.id for v in items] == [1, 3, 2]
+        assert [v.distance for v in items] == [0, 1, sqrt(3)]
+        assert items[1].numeric_embedding == [1, 1, 2]

{pgvector-0.3.5 → pgvector-0.3.6}/tests/test_peewee.py

@@ -199,3 +199,25 @@ class TestPeewee:
         Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]})
         Item.get_or_create(embedding=np.array([4, 5, 6]))
         Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9]))
+
+    def test_vector_array(self):
+        from playhouse.postgres_ext import PostgresqlExtDatabase, ArrayField
+
+        ext_db = PostgresqlExtDatabase('pgvector_python_test')
+
+        class ExtItem(BaseModel):
+            embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False)
+
+            class Meta:
+                database = ext_db
+                table_name = 'peewee_ext_item'
+
+        ext_db.connect()
+        ext_db.drop_tables([ExtItem])
+        ext_db.create_tables([ExtItem])
+
+        # fails with column "embeddings" is of type vector[] but expression is of type text[]
+        # ExtItem.create(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])
+        # item = ExtItem.get_by_id(1)
+        # assert np.array_equal(item.embeddings[0], np.array([1, 2, 3]))
+        # assert np.array_equal(item.embeddings[1], np.array([4, 5, 6]))

{pgvector-0.3.5 → pgvector-0.3.6}/tests/test_psycopg2.py

@@ -1,5 +1,5 @@
 import numpy as np
-from pgvector.psycopg2 import register_vector, SparseVector
+from pgvector.psycopg2 import register_vector, HalfVector, SparseVector
 import psycopg2
 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor
 
@@ -9,9 +9,9 @@ conn.autocommit = True
 cur = conn.cursor()
 cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
 cur.execute('DROP TABLE IF EXISTS psycopg2_items')
-cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))')
+cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector[], half_embeddings halfvec[], sparse_embeddings sparsevec[])')
 
-register_vector(cur, globally=False)
+register_vector(cur, globally=False, arrays=True)
 
 
 class TestPsycopg2:
@@ -55,6 +55,33 @@ class TestPsycopg2:
         assert res[0][0].to_list() == [1.5, 2, 3]
         assert res[1][0] is None
 
+    def test_vector_array(self):
+        embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        cur.execute('INSERT INTO psycopg2_items (embeddings) VALUES (%s::vector[])', (embeddings,))
+
+        cur.execute('SELECT embeddings FROM psycopg2_items ORDER BY id')
+        res = cur.fetchone()
+        assert np.array_equal(res[0][0], embeddings[0])
+        assert np.array_equal(res[0][1], embeddings[1])
+
+    def test_halfvec_array(self):
+        embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])]
+        cur.execute('INSERT INTO psycopg2_items (half_embeddings) VALUES (%s::halfvec[])', (embeddings,))
+
+        cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id')
+        res = cur.fetchone()
+        assert res[0][0].to_list() == [1.5, 2, 3]
+        assert res[0][1].to_list() == [4.5, 5, 6]
+
+    def test_sparsevec_array(self):
+        embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])]
+        cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,))
+
+        cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id')
+        res = cur.fetchone()
+        assert res[0][0].to_list() == [1.5, 2, 3]
+        assert res[0][1].to_list() == [4.5, 5, 6]
+
     def test_cursor_factory(self):
         for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]:
             conn = psycopg2.connect(dbname='pgvector_python_test')

{pgvector-0.3.5 → pgvector-0.3.6}/tests/test_sqlalchemy.py

@@ -1,7 +1,7 @@
 import numpy as np
 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 import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY
 from sqlalchemy.exc import StatementError
 from sqlalchemy.ext.automap import automap_base
 from sqlalchemy.orm import declarative_base, Session
@@ -31,6 +31,7 @@ class Item(Base):
     half_embedding = mapped_column(HALFVEC(3))
     binary_embedding = mapped_column(BIT(3))
     sparse_embedding = mapped_column(SPARSEVEC(3))
+    embeddings = mapped_column(ARRAY(VECTOR(3)))
 
 
 Base.metadata.drop_all(engine)
@@ -70,7 +71,8 @@ class TestSqlalchemy:
             Column('embedding', VECTOR(3)),
             Column('half_embedding', HALFVEC(3)),
             Column('binary_embedding', BIT(3)),
-            Column('sparse_embedding', SPARSEVEC(3))
+            Column('sparse_embedding', SPARSEVEC(3)),
+            Column('embeddings', ARRAY(VECTOR(3)))
         )
 
         metadata.drop_all(engine)
@@ -422,6 +424,20 @@ class TestSqlalchemy:
         item = session.query(AutoItem).first()
         assert item.embedding.tolist() == [1, 2, 3]
 
+    def test_vector_array(self):
+        session = Session(engine)
+        session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
+        session.commit()
+
+        with engine.connect() as connection:
+            from pgvector.psycopg2 import register_vector
+            register_vector(connection.connection.dbapi_connection, globally=False, arrays=True)
+
+            # this fails if the driver does not cast arrays
+            item = Session(bind=connection).get(Item, 1)
+            assert item.embeddings[0].tolist() == [1, 2, 3]
+            assert item.embeddings[1].tolist() == [4, 5, 6]
+
     @pytest.mark.asyncio
     @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
     async def test_async(self):