From 3ddd4db6444945411302e03ab7c0566879f30fcd Mon Sep 17 00:00:00 2001 From: Marc-Antoine Parent Date: Sun, 22 May 2022 21:12:04 -0400 Subject: [PATCH 001/879] distance operations are comparisons, returning floats (#2) --- pgvector/sqlalchemy/__init__.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 6956691..1cf5b66 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,4 +1,4 @@ -from sqlalchemy.types import UserDefinedType +from sqlalchemy.types import UserDefinedType, Float from ..utils import from_db, to_db __all__ = ['Vector'] @@ -28,10 +28,10 @@ def process(value): class comparator_factory(UserDefinedType.Comparator): def l2_distance(self, other): - return self.op('<->')(other) + return self.op('<->', return_type=Float)(other) def max_inner_product(self, other): - return self.op('<#>')(other) + return self.op('<#>', return_type=Float)(other) def cosine_distance(self, other): - return self.op('<=>')(other) + return self.op('<=>', return_type=Float)(other) From 494413b7e970d052390a118477969e68eaa89139 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 22 May 2022 18:14:52 -0700 Subject: [PATCH 002/879] Added test for select with SQLAlchemy --- tests/test_sqlalchemy.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 11d551a..b3488f7 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -105,6 +105,12 @@ def test_cosine_distance(self): items = session.query(Item).order_by(Item.factors.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] + def test_select(self): + with Session(engine) as session: + session.add(Item(factors=[2, 3, 3])) + item = session.query(Item.factors.l2_distance([1, 1, 1])).first() + assert item[0] == 3 + def test_bad_dimensions(self): item = Item(factors=[1, 2]) session = Session(engine) From 7ed26cacbc807abd4a7dd91fed07fb5d60ec5ddb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 22 May 2022 18:15:02 -0700 Subject: [PATCH 003/879] Updated changelog [skip ci] --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 16617ca..5d149c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.1.6 (unreleased) + +- Fixed return type for distance operators with SQLAlchemy + ## 0.1.5 (2022-01-14) - Fixed `operator does not exist` error with Django From b8da873e4a489cb6591aebaa7d6904ee5e120762 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 22 May 2022 18:16:29 -0700 Subject: [PATCH 004/879] Version bump to 0.1.6 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d149c6..3789c24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.1.6 (unreleased) +## 0.1.6 (2022-05-22) - Fixed return type for distance operators with SQLAlchemy diff --git a/setup.py b/setup.py index 0f4d7da..82af1aa 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.1.5', + version='0.1.6', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From c26fbae3ff5d6a34bf2a98ef69a60e553c0f2f4c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Aug 2022 00:18:04 -0700 Subject: [PATCH 005/879] Test with pgvector 0.2.7 on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e821057..3adb090 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -13,7 +13,7 @@ jobs: - run: | sudo apt-get update && sudo apt-get install postgresql-server-dev-14 cd /tmp - git clone --branch v0.2.4 https://github.com/pgvector/pgvector.git + git clone --branch v0.2.7 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From dcacba5472ad182e948a0503c4cffb7f226184f5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Oct 2022 14:48:57 -0700 Subject: [PATCH 006/879] Test with pgvector 0.3.0 on CI --- .github/workflows/build.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 3adb090..1bd6c14 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -4,8 +4,8 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 - run: pip install -r requirements.txt - uses: ankane/setup-postgres@v1 with: @@ -13,7 +13,7 @@ jobs: - run: | sudo apt-get update && sudo apt-get install postgresql-server-dev-14 cd /tmp - git clone --branch v0.2.7 https://github.com/pgvector/pgvector.git + git clone --branch v0.3.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From 34fb52a2286a53b57317c2c215310371de47211b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Oct 2022 14:50:21 -0700 Subject: [PATCH 007/879] Added Python version to CI --- .github/workflows/build.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 1bd6c14..3de7aec 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,6 +6,8 @@ jobs: steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 + with: + python-version: '3.10' - run: pip install -r requirements.txt - uses: ankane/setup-postgres@v1 with: From c2b253ea6e2fcc59f7770c65108518e0b7681651 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Nov 2022 15:31:05 -0700 Subject: [PATCH 008/879] Updated pgvector to 0.3.1 on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 3de7aec..13ab885 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -15,7 +15,7 @@ jobs: - run: | sudo apt-get update && sudo apt-get install postgresql-server-dev-14 cd /tmp - git clone --branch v0.3.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.3.1 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From f26c843187e95e656926fede637f9a4606adcb5e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 3 Nov 2022 15:24:36 -0700 Subject: [PATCH 009/879] Switched to dev-files option on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 13ab885..531c01d 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -12,8 +12,8 @@ jobs: - uses: ankane/setup-postgres@v1 with: database: pgvector_python_test + dev-files: true - run: | - sudo apt-get update && sudo apt-get install postgresql-server-dev-14 cd /tmp git clone --branch v0.3.1 https://github.com/pgvector/pgvector.git cd pgvector From cea62b529fae9706cde10652c8d0f79112414709 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 22 Nov 2022 13:58:02 -0800 Subject: [PATCH 010/879] Updated pgvector to 0.3.2 on CI --- .github/workflows/build.yml | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 531c01d..fadf5d4 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -4,19 +4,19 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 - - uses: actions/setup-python@v4 - with: - python-version: '3.10' - - run: pip install -r requirements.txt - - uses: ankane/setup-postgres@v1 - with: - database: pgvector_python_test - dev-files: true - - run: | - cd /tmp - git clone --branch v0.3.1 https://github.com/pgvector/pgvector.git - cd pgvector - make - sudo make install - - run: pytest + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: '3.10' + - run: pip install -r requirements.txt + - uses: ankane/setup-postgres@v1 + with: + database: pgvector_python_test + dev-files: true + - run: | + cd /tmp + git clone --branch v0.3.2 https://github.com/pgvector/pgvector.git + cd pgvector + make + sudo make install + - run: pytest From acaa5e5c388c67f17012f706ed11377632c8be6e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 13 Jan 2023 12:26:13 -0800 Subject: [PATCH 011/879] Updated pgvector to 0.4.0 on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index fadf5d4..6ba84b4 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -15,7 +15,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.3.2 https://github.com/pgvector/pgvector.git + git clone --branch v0.4.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From e7ff1468a1f0651e7b8ec8319e0b1405bc280f47 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Mar 2023 18:36:42 -0800 Subject: [PATCH 012/879] Updated examples to use embedding --- README.md | 36 ++++++++++++++++++------------------ tests/test_asyncpg.py | 16 ++++++++-------- tests/test_django.py | 22 +++++++++++----------- tests/test_psycopg.py | 36 ++++++++++++++++++------------------ tests/test_psycopg2.py | 8 ++++---- tests/test_sqlalchemy.py | 40 ++++++++++++++++++++-------------------- 6 files changed, 79 insertions(+), 79 deletions(-) diff --git a/README.md b/README.md index 5e4406e..5797721 100644 --- a/README.md +++ b/README.md @@ -50,13 +50,13 @@ Add a vector field from pgvector.django import VectorField class Item(models.Model): - factors = VectorField(dimensions=3) + embedding = VectorField(dimensions=3) ``` Insert a vector ```python -item = Item(factors=[1, 2, 3]) +item = Item(embedding=[1, 2, 3]) item.save() ``` @@ -65,7 +65,7 @@ Get the nearest neighbors to a vector ```python from pgvector.django import L2Distance -Item.objects.order_by(L2Distance('factors', [3, 1, 2]))[:5] +Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5] ``` Also supports `MaxInnerProduct` and `CosineDistance` @@ -80,7 +80,7 @@ class Item(models.Model): indexes = [ IvfflatIndex( name='my_index', - fields=['factors'], + fields=['embedding'], lists=100, opclasses=['vector_l2_ops'] ) @@ -97,13 +97,13 @@ Add a vector column from pgvector.sqlalchemy import Vector class Item(Base): - factors = Column(Vector(3)) + embedding = Column(Vector(3)) ``` Insert a vector ```python -item = Item(factors=[1, 2, 3]) +item = Item(embedding=[1, 2, 3]) session.add(item) session.commit() ``` @@ -111,7 +111,7 @@ session.commit() Get the nearest neighbors to a vector ```python -session.query(Item).order_by(Item.factors.l2_distance([3, 1, 2])).limit(5).all() +session.query(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5).all() ``` Also supports `max_inner_product` and `cosine_distance` @@ -119,10 +119,10 @@ Also supports `max_inner_product` and `cosine_distance` Add an approximate index ```python -index = Index('my_index', Item.factors, +index = Index('my_index', Item.embedding, postgresql_using='ivfflat', postgresql_with={'lists': 100}, - postgresql_ops={'factors': 'vector_l2_ops'} + postgresql_ops={'embedding': 'vector_l2_ops'} ) index.create(engine) ``` @@ -142,14 +142,14 @@ register_vector(conn) Insert a vector ```python -factors = np.array([1, 2, 3]) -conn.execute('INSERT INTO item (factors) VALUES (%s)', (factors,)) +embedding = np.array([1, 2, 3]) +conn.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) ``` Get the nearest neighbors to a vector ```python -conn.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,)).fetchall() +conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() ``` ## Psycopg 2 @@ -165,14 +165,14 @@ register_vector(conn) Insert a vector ```python -factors = np.array([1, 2, 3]) -cur.execute('INSERT INTO item (factors) VALUES (%s)', (factors,)) +embedding = np.array([1, 2, 3]) +cur.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) ``` Get the nearest neighbors to a vector ```python -cur.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,)) +cur.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)) cur.fetchall() ``` @@ -189,14 +189,14 @@ await register_vector(conn) Insert a vector ```python -factors = np.array([1, 2, 3]) -await conn.execute('INSERT INTO item (factors) VALUES ($1)', factors) +embedding = np.array([1, 2, 3]) +await conn.execute('INSERT INTO item (embedding) VALUES ($1)', embedding) ``` Get the nearest neighbors to a vector ```python -await conn.fetch('SELECT * FROM item ORDER BY factors <-> $1 LIMIT 5', factors) +await conn.fetch('SELECT * FROM item ORDER BY embedding <-> $1 LIMIT 5', embedding) ``` ## History diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 12ab806..7f7dc80 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -11,22 +11,22 @@ async def test_works(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 item') - await conn.execute('CREATE TABLE item (id bigserial primary key, factors vector(3))') + await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') await register_vector(conn) - factors = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO item (factors) VALUES ($1), (NULL)", factors) + embedding = np.array([1.5, 2, 3]) + await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM item ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 - assert np.array_equal(res[0]['factors'], factors) - assert res[0]['factors'].dtype == np.float32 - assert res[1]['factors'] is None + assert np.array_equal(res[0]['embedding'], embedding) + assert res[0]['embedding'].dtype == np.float32 + assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT factors::text FROM item ORDER BY id LIMIT 1") - assert text_res[0]['factors'] == '[1.5,2,3]' + text_res = await conn.fetch("SELECT embedding::text FROM item ORDER BY id LIMIT 1") + assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() diff --git a/tests/test_django.py b/tests/test_django.py index 95e25f2..b645b3c 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -18,14 +18,14 @@ class Item(models.Model): - factors = VectorField(dimensions=3) + embedding = VectorField(dimensions=3) class Meta: app_label = 'myapp' indexes = [ IvfflatIndex( name='my_index', - fields=['factors'], + fields=['embedding'], lists=100, opclasses=['vector_l2_ops'] ) @@ -44,12 +44,12 @@ class Migration(migrations.Migration): name='Item', fields=[ ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('factors', pgvector.django.VectorField(dimensions=3)), + ('embedding', pgvector.django.VectorField(dimensions=3)), ], ), migrations.AddIndex( model_name='item', - index=pgvector.django.IvfflatIndex(fields=['factors'], lists=1, name='my_index', opclasses=['vector_l2_ops']), + index=pgvector.django.IvfflatIndex(fields=['embedding'], lists=1, name='my_index', opclasses=['vector_l2_ops']), ) ] @@ -72,7 +72,7 @@ def create_items(): [1, 1, 2] ] for i, v in enumerate(vectors): - item = Item(id=i + 1, factors=v) + item = Item(id=i + 1, embedding=v) item.save() @@ -81,24 +81,24 @@ def setup_method(self, test_method): Item.objects.all().delete() def test_works(self): - item = Item(id=1, factors=[1, 2, 3]) + item = Item(id=1, embedding=[1, 2, 3]) item.save() item = Item.objects.get(pk=1) assert item.id == 1 - assert np.array_equal(item.factors, np.array([1, 2, 3])) - assert item.factors.dtype == np.float32 + assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert item.embedding.dtype == np.float32 def test_l2_distance(self): create_items() - items = Item.objects.order_by(L2Distance('factors', [1, 1, 1])) + items = Item.objects.order_by(L2Distance('embedding', [1, 1, 1])) assert [v.id for v in items] == [1, 3, 2] def test_max_inner_product(self): create_items() - items = Item.objects.order_by(MaxInnerProduct('factors', [1, 1, 1])) + items = Item.objects.order_by(MaxInnerProduct('embedding', [1, 1, 1])) assert [v.id for v in items] == [2, 3, 1] def test_cosine_distance(self): create_items() - items = Item.objects.order_by(CosineDistance('factors', [1, 1, 1])) + items = Item.objects.order_by(CosineDistance('embedding', [1, 1, 1])) assert [v.id for v in items] == [1, 2, 3] diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 3fe1edc..2ebbcd9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -7,7 +7,7 @@ conn.execute('CREATE EXTENSION IF NOT EXISTS vector') conn.execute('DROP TABLE IF EXISTS item') -conn.execute('CREATE TABLE item (id bigserial primary key, factors vector(3))') +conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') register_vector(conn) @@ -17,37 +17,37 @@ def setup_method(self, test_method): conn.execute('DELETE FROM item') def test_works(self): - factors = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO item (factors) VALUES (%s), (NULL)', (factors,)) + embedding = np.array([1.5, 2, 3]) + conn.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,)) res = conn.execute('SELECT * FROM item ORDER BY id').fetchall() - assert np.array_equal(res[0][1], factors) + assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None def test_binary_format(self): - factors = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector', (factors,)).fetchone()[0] - assert np.array_equal(res, factors) + embedding = np.array([1.5, 2, 3]) + res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] + assert np.array_equal(res, embedding) def test_text_format(self): - factors = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %t::vector', (factors,)).fetchone()[0] - assert np.array_equal(res, factors) + embedding = np.array([1.5, 2, 3]) + res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] + assert np.array_equal(res, embedding) def test_binary_format_correct(self): - factors = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector::text', (factors,)).fetchone()[0] + embedding = np.array([1.5, 2, 3]) + res = conn.execute('SELECT %b::vector::text', (embedding,)).fetchone()[0] assert res == '[1.5,2,3]' def test_text_format_non_contiguous(self): - factors = np.flipud(np.array([1.5, 2, 3])) - assert not factors.data.contiguous - res = conn.execute('SELECT %t::vector', (factors,)).fetchone()[0] + embedding = np.flipud(np.array([1.5, 2, 3])) + assert not embedding.data.contiguous + res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, np.array([3, 2, 1.5])) def test_binary_format_non_contiguous(self): - factors = np.flipud(np.array([1.5, 2, 3])) - assert not factors.data.contiguous - res = conn.execute('SELECT %b::vector', (factors,)).fetchone()[0] + embedding = np.flipud(np.array([1.5, 2, 3])) + assert not embedding.data.contiguous + res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, np.array([3, 2, 1.5])) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index df86706..ac899ed 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -8,7 +8,7 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') cur.execute('DROP TABLE IF EXISTS item') -cur.execute('CREATE TABLE item (id bigserial primary key, factors vector(3))') +cur.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') register_vector(cur) @@ -18,11 +18,11 @@ def setup_method(self, test_method): cur.execute('DELETE FROM item') def test_works(self): - factors = np.array([1.5, 2, 3]) - cur.execute('INSERT INTO item (factors) VALUES (%s), (NULL)', (factors,)) + embedding = np.array([1.5, 2, 3]) + cur.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT * FROM item ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][1], factors) + assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index b3488f7..ca07fa2 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -17,7 +17,7 @@ class Item(Base): __tablename__ = 'orm_item' id = Column(Integer, primary_key=True) - factors = Column(Vector(3)) + embedding = Column(Vector(3)) Base.metadata.drop_all(engine) @@ -32,7 +32,7 @@ def create_items(): ] session = Session(engine) for i, v in enumerate(vectors): - session.add(Item(id=i + 1, factors=v)) + session.add(Item(id=i + 1, embedding=v)) session.commit() @@ -49,7 +49,7 @@ def test_core(self): 'core_item', metadata, Column('id', Integer, primary_key=True), - Column('factors', Vector(3)) + Column('embedding', Vector(3)) ) metadata.drop_all(engine) @@ -57,16 +57,16 @@ def test_core(self): index = Index( 'my_core_index', - item_table.c.factors, + item_table.c.embedding, postgresql_using='ivfflat', postgresql_with={'lists': 1}, - postgresql_ops={'factors': 'vector_l2_ops'} + postgresql_ops={'embedding': 'vector_l2_ops'} ) index.create(engine) def test_orm(self): - item = Item(factors=np.array([1.5, 2, 3])) - item2 = Item(factors=[4, 5, 6]) + item = Item(embedding=np.array([1.5, 2, 3])) + item2 = Item(embedding=[4, 5, 6]) item3 = Item() session = Session(engine) @@ -81,52 +81,52 @@ def test_orm(self): assert items[0].id == 1 assert items[1].id == 2 assert items[2].id == 3 - assert np.array_equal(items[0].factors, np.array([1.5, 2, 3])) - assert items[0].factors.dtype == np.float32 - assert np.array_equal(items[1].factors, np.array([4, 5, 6])) - assert items[1].factors.dtype == np.float32 - assert items[2].factors is None + assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) + assert items[0].embedding.dtype == np.float32 + assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) + assert items[1].embedding.dtype == np.float32 + assert items[2].embedding is None def test_l2_distance(self): create_items() with Session(engine) as session: - items = session.query(Item).order_by(Item.factors.l2_distance([1, 1, 1])).all() + items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] def test_max_inner_product(self): create_items() with Session(engine) as session: - items = session.query(Item).order_by(Item.factors.max_inner_product([1, 1, 1])).all() + items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] def test_cosine_distance(self): create_items() with Session(engine) as session: - items = session.query(Item).order_by(Item.factors.cosine_distance([1, 1, 1])).all() + items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] def test_select(self): with Session(engine) as session: - session.add(Item(factors=[2, 3, 3])) - item = session.query(Item.factors.l2_distance([1, 1, 1])).first() + session.add(Item(embedding=[2, 3, 3])) + item = session.query(Item.embedding.l2_distance([1, 1, 1])).first() assert item[0] == 3 def test_bad_dimensions(self): - item = Item(factors=[1, 2]) + item = Item(embedding=[1, 2]) session = Session(engine) session.add(item) with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() def test_bad_ndim(self): - item = Item(factors=np.array([[1, 2, 3]])) + item = Item(embedding=np.array([[1, 2, 3]])) session = Session(engine) session.add(item) with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() def test_bad_dtype(self): - item = Item(factors=np.array(['one', 'two', 'three'])) + item = Item(embedding=np.array(['one', 'two', 'three'])) session = Session(engine) session.add(item) with pytest.raises(StatementError, match='dtype must be numeric'): From 4984f92aa7defa42680dfdf1e5e7ab873a8d061e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 11:59:53 -0700 Subject: [PATCH 013/879] Updated example for SQLAlchemy 2 [skip ci] --- examples/surprise_recs.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/examples/surprise_recs.py b/examples/surprise_recs.py index 1d957c3..ce80eaa 100644 --- a/examples/surprise_recs.py +++ b/examples/surprise_recs.py @@ -1,9 +1,9 @@ from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, text, Column, Integer -from sqlalchemy.orm import declarative_base, Session +from sqlalchemy import create_engine, select, text, Integer +from sqlalchemy.orm import declarative_base, mapped_column, Session from surprise import Dataset, SVD -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example', future=True) +engine = create_engine('postgresql+psycopg2://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() @@ -14,15 +14,15 @@ class User(Base): __tablename__ = 'user' - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) + id = mapped_column(Integer, primary_key=True) + factors = mapped_column(Vector(20)) class Item(Base): __tablename__ = 'item' - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) + id = mapped_column(Integer, primary_key=True) + factors = mapped_column(Vector(20)) Base.metadata.drop_all(engine) @@ -41,10 +41,10 @@ class Item(Base): session.bulk_insert_mappings(Item, items) session.commit() -user = session.query(User).get(1) -items = session.query(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5).all() +user = session.get(User, 1) +items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) print('user-based recs:', [item.id for item in items]) -item = session.query(Item).get(50) -items = session.query(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5).all() +item = session.get(Item, 50) +items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.id for item in items]) From b35853e2c8df2c333341d1b3a3797795ff99bdcf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 12:07:35 -0700 Subject: [PATCH 014/879] Fixed warning in example [skip ci] --- examples/pytorch_image_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index 3551e16..21d421e 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -26,7 +26,7 @@ # load pretrained model -model = torchvision.models.resnet18(pretrained=True) +model = torchvision.models.resnet18(weights='DEFAULT') model.fc = torch.nn.Identity() model.eval() From 141ca8f94199acfb8c73d1702a6e00f85551f8af Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 12:07:47 -0700 Subject: [PATCH 015/879] Added requirements.txt for examples [skip ci] --- examples/requirements.txt | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 examples/requirements.txt diff --git a/examples/requirements.txt b/examples/requirements.txt new file mode 100644 index 0000000..9a36c0b --- /dev/null +++ b/examples/requirements.txt @@ -0,0 +1,5 @@ +implicit +lightfm +torch +torchvision +scikit-surprise From 3fdb077033a74c23ae8a15573400c81105521eb7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 12:13:32 -0700 Subject: [PATCH 016/879] Updated example for SQLAlchemy 2 [skip ci] --- examples/lightfm_recs.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/examples/lightfm_recs.py b/examples/lightfm_recs.py index 9b9bc5f..77f7d3a 100644 --- a/examples/lightfm_recs.py +++ b/examples/lightfm_recs.py @@ -1,10 +1,10 @@ from lightfm import LightFM from lightfm.datasets import fetch_movielens from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, text, Column, Float, Integer, String -from sqlalchemy.orm import declarative_base, Session +from sqlalchemy import create_engine, select, text, Float, Integer, String +from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example', future=True) +engine = create_engine('postgresql+psycopg2://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() @@ -15,17 +15,17 @@ class User(Base): __tablename__ = 'user' - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) + id = mapped_column(Integer, primary_key=True) + factors = mapped_column(Vector(20)) class Item(Base): __tablename__ = 'item' - id = Column(Integer, primary_key=True) - title = Column(String) - factors = Column(Vector(20)) - bias = Column(Float) + id = mapped_column(Integer, primary_key=True) + title = mapped_column(String) + factors = mapped_column(Vector(20)) + bias = mapped_column(Float) Base.metadata.drop_all(engine) @@ -46,11 +46,11 @@ class Item(Base): session.bulk_insert_mappings(Item, items) session.commit() -user = session.query(User).get(1) +user = session.get(User, 1) # subtract item bias for negative inner product -items = session.query(Item).order_by(Item.factors.max_inner_product(user.factors) - Item.bias).limit(5).all() +items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors) - Item.bias).limit(5)) print('user-based recs:', [item.title for item in items]) -item = session.query(Item).filter(Item.title == 'Star Wars (1977)').first() -items = session.query(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5).all() +item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() +items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) From 6ff2e0750d167a3c010a459e592ca4da830382fd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 12:13:52 -0700 Subject: [PATCH 017/879] Added comment [skip ci] --- examples/lightfm_recs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/lightfm_recs.py b/examples/lightfm_recs.py index 77f7d3a..970756c 100644 --- a/examples/lightfm_recs.py +++ b/examples/lightfm_recs.py @@ -51,6 +51,7 @@ class Item(Base): items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors) - Item.bias).limit(5)) print('user-based recs:', [item.title for item in items]) +# broken due to https://github.com/lyst/lightfm/issues/682 item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) From e6c62b18942e7a20dd29acfa4cb5a22c85e0e1d6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 12:18:41 -0700 Subject: [PATCH 018/879] Updated example for SQLAlchemy 2 [skip ci] --- examples/implicit_recs.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/examples/implicit_recs.py b/examples/implicit_recs.py index 10a1fe0..46923a1 100644 --- a/examples/implicit_recs.py +++ b/examples/implicit_recs.py @@ -1,10 +1,10 @@ import implicit from implicit.datasets.movielens import get_movielens from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, text, Column, Integer, String -from sqlalchemy.orm import declarative_base, Session +from sqlalchemy import create_engine, select, text, Integer, String +from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example', future=True) +engine = create_engine('postgresql+psycopg2://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() @@ -15,16 +15,16 @@ class User(Base): __tablename__ = 'user' - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) + id = mapped_column(Integer, primary_key=True) + factors = mapped_column(Vector(20)) class Item(Base): __tablename__ = 'item' - id = Column(Integer, primary_key=True) - title = Column(String) - factors = Column(Vector(20)) + id = mapped_column(Integer, primary_key=True) + title = mapped_column(String) + factors = mapped_column(Vector(20)) Base.metadata.drop_all(engine) @@ -35,17 +35,17 @@ class Item(Base): model.fit(ratings) users = [dict(id=i, factors=factors) for i, factors in enumerate(model.user_factors)] -items = [dict(id=i, title=titles[i].decode('utf-8'), factors=factors) for i, factors in enumerate(model.item_factors)] +items = [dict(id=i, title=titles[i], factors=factors) for i, factors in enumerate(model.item_factors)] session = Session(engine) session.bulk_insert_mappings(User, users) session.bulk_insert_mappings(Item, items) session.commit() -user = session.query(User).get(1) -items = session.query(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5).all() +user = session.get(User, 1) +items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) print('user-based recs:', [item.title for item in items]) -item = session.query(Item).filter(Item.title == 'Star Wars (1977)').first() -items = session.query(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5).all() +item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() +items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) From 509cfa7176c5f9205a69be588c700c426edaa5d1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 12:24:20 -0700 Subject: [PATCH 019/879] Updated examples for SQLAlchemy 2 [skip ci] --- README.md | 4 ++-- tests/test_sqlalchemy.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 5797721..9f40eb2 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,7 @@ Add a vector column from pgvector.sqlalchemy import Vector class Item(Base): - embedding = Column(Vector(3)) + embedding = mapped_column(Vector(3)) ``` Insert a vector @@ -111,7 +111,7 @@ session.commit() Get the nearest neighbors to a vector ```python -session.query(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5).all() +session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) ``` Also supports `max_inner_product` and `cosine_distance` diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index ca07fa2..52fe09b 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -3,9 +3,9 @@ import pytest from sqlalchemy import create_engine, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError -from sqlalchemy.orm import declarative_base, Session +from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test', future=True) +engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with engine.connect() as con: con.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) con.commit() @@ -16,8 +16,8 @@ class Item(Base): __tablename__ = 'orm_item' - id = Column(Integer, primary_key=True) - embedding = Column(Vector(3)) + id = mapped_column(Integer, primary_key=True) + embedding = mapped_column(Vector(3)) Base.metadata.drop_all(engine) From a7ac6d744bc335283980f1448255c0c027cce9ed Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 23:23:10 -0700 Subject: [PATCH 020/879] Added sentence embeddings example [skip ci] --- README.md | 1 + examples/requirements.txt | 1 + examples/sentence_embeddings.py | 43 +++++++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+) create mode 100644 examples/sentence_embeddings.py diff --git a/README.md b/README.md index 9f40eb2..27f12d1 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ And follow the instructions for your database library: Or check out some examples: +- [Sentence embeddings](examples/sentence_embeddings.py) with SentenceTransformers - [Image search](examples/pytorch_image_search.py) with PyTorch - [Implicit feedback recommendations](examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](examples/surprise_recs.py) with Surprise diff --git a/examples/requirements.txt b/examples/requirements.txt index 9a36c0b..72c47f1 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -3,3 +3,4 @@ lightfm torch torchvision scikit-surprise +sentence-transformers diff --git a/examples/sentence_embeddings.py b/examples/sentence_embeddings.py new file mode 100644 index 0000000..a117604 --- /dev/null +++ b/examples/sentence_embeddings.py @@ -0,0 +1,43 @@ +from pgvector.sqlalchemy import Vector +from sentence_transformers import SentenceTransformer +from sqlalchemy import create_engine, select, text, Integer, String, Text +from sqlalchemy.orm import declarative_base, mapped_column, Session + +engine = create_engine('postgresql+psycopg://localhost/pgvector_example') +with engine.connect() as conn: + conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) + conn.commit() + +Base = declarative_base() + + +class Document(Base): + __tablename__ = 'document' + + id = mapped_column(Integer, primary_key=True) + content = mapped_column(Text) + embedding = mapped_column(Vector(384)) + + +Base.metadata.drop_all(engine) +Base.metadata.create_all(engine) + +sentences = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] + +model = SentenceTransformer('all-MiniLM-L6-v2') +embeddings = model.encode(sentences) + +documents = [dict(content=sentences[i], embedding=embedding) for i, embedding in enumerate(embeddings)] + +session = Session(engine) +session.bulk_insert_mappings(Document, documents) +session.commit() + +doc = session.get(Document, 1) +neighbors = session.scalars(select(Document).filter(Document.id != doc.id).order_by(Document.embedding.cosine_distance(doc.embedding)).limit(5)) +for neighbor in neighbors: + print(neighbor.content) From db8743d2efe19485532a9f3e71b8d92d5047ad70 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 23:25:15 -0700 Subject: [PATCH 021/879] Updated examples to use Psycopg 3 [skip ci] --- examples/implicit_recs.py | 2 +- examples/lightfm_recs.py | 2 +- examples/surprise_recs.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/implicit_recs.py b/examples/implicit_recs.py index 46923a1..1565f82 100644 --- a/examples/implicit_recs.py +++ b/examples/implicit_recs.py @@ -4,7 +4,7 @@ from sqlalchemy import create_engine, select, text, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example') +engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() diff --git a/examples/lightfm_recs.py b/examples/lightfm_recs.py index 970756c..b83e04b 100644 --- a/examples/lightfm_recs.py +++ b/examples/lightfm_recs.py @@ -4,7 +4,7 @@ from sqlalchemy import create_engine, select, text, Float, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example') +engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() diff --git a/examples/surprise_recs.py b/examples/surprise_recs.py index ce80eaa..a562bf7 100644 --- a/examples/surprise_recs.py +++ b/examples/surprise_recs.py @@ -3,7 +3,7 @@ from sqlalchemy.orm import declarative_base, mapped_column, Session from surprise import Dataset, SVD -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example') +engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() From 8a480d796e9eb2c5e129f0a1a004a41dbce6e173 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 23:25:44 -0700 Subject: [PATCH 022/879] Updated CI --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 6ba84b4..e327f34 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -7,7 +7,7 @@ jobs: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 with: - python-version: '3.10' + python-version: 3.11 - run: pip install -r requirements.txt - uses: ankane/setup-postgres@v1 with: @@ -15,7 +15,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.4.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.4.1 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From 16e7474a92093953fc39f14c40a67e580275c5ca Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 23:28:05 -0700 Subject: [PATCH 023/879] Improved examples [skip ci] --- examples/implicit_recs.py | 4 ++-- examples/lightfm_recs.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/implicit_recs.py b/examples/implicit_recs.py index 1565f82..390fc36 100644 --- a/examples/implicit_recs.py +++ b/examples/implicit_recs.py @@ -34,8 +34,8 @@ class Item(Base): model = implicit.als.AlternatingLeastSquares(factors=20) model.fit(ratings) -users = [dict(id=i, factors=factors) for i, factors in enumerate(model.user_factors)] -items = [dict(id=i, title=titles[i], factors=factors) for i, factors in enumerate(model.item_factors)] +users = [dict(factors=factors) for i, factors in enumerate(model.user_factors)] +items = [dict(title=titles[i], factors=factors) for i, factors in enumerate(model.item_factors)] session = Session(engine) session.bulk_insert_mappings(User, users) diff --git a/examples/lightfm_recs.py b/examples/lightfm_recs.py index b83e04b..e5247d6 100644 --- a/examples/lightfm_recs.py +++ b/examples/lightfm_recs.py @@ -38,8 +38,8 @@ class Item(Base): user_biases, user_factors = model.get_user_representations() item_biases, item_factors = model.get_item_representations() -users = [dict(id=i, factors=factors) for i, factors in enumerate(user_factors)] -items = [dict(id=i, title=data['item_labels'][i], factors=factors, bias=item_biases[i].item()) for i, factors in enumerate(item_factors)] +users = [dict(factors=factors) for i, factors in enumerate(user_factors)] +items = [dict(title=data['item_labels'][i], factors=factors, bias=item_biases[i].item()) for i, factors in enumerate(item_factors)] session = Session(engine) session.bulk_insert_mappings(User, users) From 49f5a6b5e32acf13cf0ffa16dac8887feddf7949 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 23:28:28 -0700 Subject: [PATCH 024/879] Updated readme [skip ci] --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index 27f12d1..0e3033a 100644 --- a/README.md +++ b/README.md @@ -2,8 +2,6 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Great for online recommendations :tada: - Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), and [asyncpg](https://github.com/MagicStack/asyncpg) [![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions) From 6ad8e33160321be6c11f85e9ea52bc9ab66289f9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 Mar 2023 23:54:21 -0700 Subject: [PATCH 025/879] Added OpenAI embeddings example [skip ci] --- README.md | 1 + examples/openai_embeddings.py | 42 +++++++++++++++++++++++++++++++++++ examples/requirements.txt | 1 + 3 files changed, 44 insertions(+) create mode 100644 examples/openai_embeddings.py diff --git a/README.md b/README.md index 0e3033a..97f43cd 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ And follow the instructions for your database library: Or check out some examples: +- [Embeddings](examples/openai_embeddings.py) with OpenAI - [Sentence embeddings](examples/sentence_embeddings.py) with SentenceTransformers - [Image search](examples/pytorch_image_search.py) with PyTorch - [Implicit feedback recommendations](examples/implicit_recs.py) with Implicit diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py new file mode 100644 index 0000000..fb26d24 --- /dev/null +++ b/examples/openai_embeddings.py @@ -0,0 +1,42 @@ +import openai +from pgvector.sqlalchemy import Vector +from sentence_transformers import SentenceTransformer +from sqlalchemy import create_engine, select, text, Integer, String, Text +from sqlalchemy.orm import declarative_base, mapped_column, Session + +engine = create_engine('postgresql+psycopg://localhost/pgvector_example') +with engine.connect() as conn: + conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) + conn.commit() + +Base = declarative_base() + + +class Document(Base): + __tablename__ = 'document' + + id = mapped_column(Integer, primary_key=True) + content = mapped_column(Text) + embedding = mapped_column(Vector(1536)) + + +Base.metadata.drop_all(engine) +Base.metadata.create_all(engine) + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] + +embeddings = [v['embedding'] for v in openai.Embedding.create(input=input, model='text-embedding-ada-002')['data']] +documents = [dict(content=input[i], embedding=embedding) for i, embedding in enumerate(embeddings)] + +session = Session(engine) +session.bulk_insert_mappings(Document, documents) +session.commit() + +doc = session.get(Document, 1) +neighbors = session.scalars(select(Document).filter(Document.id != doc.id).order_by(Document.embedding.max_inner_product(doc.embedding)).limit(5)) +for neighbor in neighbors: + print(neighbor.content) diff --git a/examples/requirements.txt b/examples/requirements.txt index 72c47f1..4be7ae1 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -1,5 +1,6 @@ implicit lightfm +openai torch torchvision scikit-surprise From 21a8a3d2a70aa9c545332322728c88121bc3884e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Apr 2023 00:03:20 -0700 Subject: [PATCH 026/879] Updated examples to use latest bulk insert API [skip ci] --- examples/implicit_recs.py | 7 +++---- examples/lightfm_recs.py | 7 +++---- examples/openai_embeddings.py | 5 ++--- examples/sentence_embeddings.py | 5 ++--- examples/surprise_recs.py | 7 +++---- 5 files changed, 13 insertions(+), 18 deletions(-) diff --git a/examples/implicit_recs.py b/examples/implicit_recs.py index 390fc36..f70eb8c 100644 --- a/examples/implicit_recs.py +++ b/examples/implicit_recs.py @@ -1,7 +1,7 @@ import implicit from implicit.datasets.movielens import get_movielens from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, select, text, Integer, String +from sqlalchemy import create_engine, insert, select, text, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -38,9 +38,8 @@ class Item(Base): items = [dict(title=titles[i], factors=factors) for i, factors in enumerate(model.item_factors)] session = Session(engine) -session.bulk_insert_mappings(User, users) -session.bulk_insert_mappings(Item, items) -session.commit() +session.execute(insert(User), users) +session.execute(insert(Item), items) user = session.get(User, 1) items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) diff --git a/examples/lightfm_recs.py b/examples/lightfm_recs.py index e5247d6..fcb9027 100644 --- a/examples/lightfm_recs.py +++ b/examples/lightfm_recs.py @@ -1,7 +1,7 @@ from lightfm import LightFM from lightfm.datasets import fetch_movielens from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, select, text, Float, Integer, String +from sqlalchemy import create_engine, insert, select, text, Float, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -42,9 +42,8 @@ class Item(Base): items = [dict(title=data['item_labels'][i], factors=factors, bias=item_biases[i].item()) for i, factors in enumerate(item_factors)] session = Session(engine) -session.bulk_insert_mappings(User, users) -session.bulk_insert_mappings(Item, items) -session.commit() +session.execute(insert(User), users) +session.execute(insert(Item), items) user = session.get(User, 1) # subtract item bias for negative inner product diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index fb26d24..795828f 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -1,7 +1,7 @@ import openai from pgvector.sqlalchemy import Vector from sentence_transformers import SentenceTransformer -from sqlalchemy import create_engine, select, text, Integer, String, Text +from sqlalchemy import create_engine, insert, select, text, Integer, String, Text from sqlalchemy.orm import declarative_base, mapped_column, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -33,8 +33,7 @@ class Document(Base): documents = [dict(content=input[i], embedding=embedding) for i, embedding in enumerate(embeddings)] session = Session(engine) -session.bulk_insert_mappings(Document, documents) -session.commit() +session.execute(insert(Document), documents) doc = session.get(Document, 1) neighbors = session.scalars(select(Document).filter(Document.id != doc.id).order_by(Document.embedding.max_inner_product(doc.embedding)).limit(5)) diff --git a/examples/sentence_embeddings.py b/examples/sentence_embeddings.py index a117604..f6d0775 100644 --- a/examples/sentence_embeddings.py +++ b/examples/sentence_embeddings.py @@ -1,6 +1,6 @@ from pgvector.sqlalchemy import Vector from sentence_transformers import SentenceTransformer -from sqlalchemy import create_engine, select, text, Integer, String, Text +from sqlalchemy import create_engine, insert, select, text, Integer, String, Text from sqlalchemy.orm import declarative_base, mapped_column, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -34,8 +34,7 @@ class Document(Base): documents = [dict(content=sentences[i], embedding=embedding) for i, embedding in enumerate(embeddings)] session = Session(engine) -session.bulk_insert_mappings(Document, documents) -session.commit() +session.execute(insert(Document), documents) doc = session.get(Document, 1) neighbors = session.scalars(select(Document).filter(Document.id != doc.id).order_by(Document.embedding.cosine_distance(doc.embedding)).limit(5)) diff --git a/examples/surprise_recs.py b/examples/surprise_recs.py index a562bf7..bd7d18d 100644 --- a/examples/surprise_recs.py +++ b/examples/surprise_recs.py @@ -1,5 +1,5 @@ from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, select, text, Integer +from sqlalchemy import create_engine, insert, select, text, Integer from sqlalchemy.orm import declarative_base, mapped_column, Session from surprise import Dataset, SVD @@ -37,9 +37,8 @@ class Item(Base): items = [dict(id=trainset.to_raw_iid(i), factors=algo.qi[i]) for i in trainset.all_items()] session = Session(engine) -session.bulk_insert_mappings(User, users) -session.bulk_insert_mappings(Item, items) -session.commit() +session.execute(insert(User), users) +session.execute(insert(Item), items) user = session.get(User, 1) items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) From a874dd60413a43fa45914f6276fe9e3559a17527 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Apr 2023 10:21:11 -0700 Subject: [PATCH 027/879] Added example of getting distance with Django - #5 --- README.md | 6 ++++++ tests/test_django.py | 13 ++++++++++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 97f43cd..5fcdc4b 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,12 @@ Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5] Also supports `MaxInnerProduct` and `CosineDistance` +Get the distance + +```python +Item.objects.annotate(distance=L2Distance('embedding', [3, 1, 2])) +``` + Add an approximate index ```python diff --git a/tests/test_django.py b/tests/test_django.py index b645b3c..2314822 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -2,6 +2,7 @@ from django.conf import settings from django.db import connection, migrations, models from django.db.migrations.loader import MigrationLoader +from math import sqrt import numpy as np import pgvector.django from pgvector.django import VectorExtension, VectorField, IvfflatIndex, L2Distance, MaxInnerProduct, CosineDistance @@ -90,15 +91,21 @@ def test_works(self): def test_l2_distance(self): create_items() - items = Item.objects.order_by(L2Distance('embedding', [1, 1, 1])) + distance = L2Distance('embedding', [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)] def test_max_inner_product(self): create_items() - items = Item.objects.order_by(MaxInnerProduct('embedding', [1, 1, 1])) + distance = MaxInnerProduct('embedding', [1, 1, 1]) + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [-6, -4, -3] def test_cosine_distance(self): create_items() - items = Item.objects.order_by(CosineDistance('embedding', [1, 1, 1])) + distance = CosineDistance('embedding', [1, 1, 1]) + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] + assert [v.distance for v in items] == [0, 0, 0.05719095841793653] From e3c8a0fe0f3b6a5fe02bc94f08e783bc012bf9d9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Apr 2023 11:34:24 -0700 Subject: [PATCH 028/879] Added example of getting items within a certain distance with Django - closes #5 --- README.md | 6 ++++++ tests/test_django.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/README.md b/README.md index 5fcdc4b..295812d 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,12 @@ Get the distance Item.objects.annotate(distance=L2Distance('embedding', [3, 1, 2])) ``` +Get items within a certain distance + +```python +Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance__lt=5) +``` + Add an approximate index ```python diff --git a/tests/test_django.py b/tests/test_django.py index 2314822..39ed689 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -109,3 +109,9 @@ def test_cosine_distance(self): items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + + def test_filter(self): + create_items() + distance = L2Distance('embedding', [1, 1, 1]) + items = Item.objects.alias(distance=distance).filter(distance__lt=1) + assert [v.id for v in items] == [1] From d7d4d35a8f8a098ce91bc65c3a2ef277f396cca3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Apr 2023 18:30:26 -0700 Subject: [PATCH 029/879] Added SQLModel to docs - closes #6 --- README.md | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 295812d..e9e3265 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), and [asyncpg](https://github.com/MagicStack/asyncpg) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), and [asyncpg](https://github.com/MagicStack/asyncpg) [![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions) @@ -18,6 +18,7 @@ And follow the instructions for your database library: - [Django](#django) - [SQLAlchemy](#sqlalchemy) +- [SQLModel](#sqlmodel) - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) @@ -141,6 +142,34 @@ index.create(engine) Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +## SQLModel + +Add a vector column + +```python +from pgvector.sqlalchemy import Vector +from sqlalchemy import Column + +class Item(SQLModel, table=True): + embedding: List[float] = Field(default=None, sa_column=Column(Vector(3))) +``` + +Insert a vector + +```python +item = Item(embedding=[1, 2, 3]) +session.add(item) +session.commit() +``` + +Get the nearest neighbors to a vector + +```python +session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) +``` + +Also supports `max_inner_product` and `cosine_distance` + ## Psycopg 3 Register the vector type with your connection From 91f34f23fb61e7674b5509ea094c7a8608f1075c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Apr 2023 18:48:22 -0700 Subject: [PATCH 030/879] Updated example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index e9e3265..1779b8d 100644 --- a/README.md +++ b/README.md @@ -151,7 +151,7 @@ from pgvector.sqlalchemy import Vector from sqlalchemy import Column class Item(SQLModel, table=True): - embedding: List[float] = Field(default=None, sa_column=Column(Vector(3))) + embedding: List[float] = Field(sa_column=Column(Vector(3))) ``` Insert a vector From a0b7897080e91b62008efae260f2395f1d1bdc87 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 3 May 2023 19:12:14 -0700 Subject: [PATCH 031/879] Fixed binary copy for psycopg3 --- CHANGELOG.md | 4 ++++ pgvector/psycopg/__init__.py | 8 ++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3789c24..71e5d84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.1.7 (unreleased) + +- Fixed binary copy for psycopg3 + ## 0.1.6 (2022-05-22) - Fixed return type for distance operators with SQLAlchemy diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index ecb8fe6..69e7f68 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -49,8 +49,12 @@ def register_vector(context): raise psycopg.ProgrammingError('vector type not found in the database') info.register(context) + # add oid to anonymous class + text_dumper = type('', (VectorDumper,), {'oid': info.oid}) + binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) + adapters = context.adapters - adapters.register_dumper('numpy.ndarray', VectorDumper) - adapters.register_dumper('numpy.ndarray', VectorBinaryDumper) + adapters.register_dumper('numpy.ndarray', text_dumper) + adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) From f2d5dab2e5a663e10544122cdea0170c18810600 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 3 May 2023 19:18:05 -0700 Subject: [PATCH 032/879] Added tests for copy --- CHANGELOG.md | 2 +- tests/test_psycopg.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 71e5d84..616993b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## 0.1.7 (unreleased) -- Fixed binary copy for psycopg3 +- Fixed `set_types` for psycopg3 ## 0.1.6 (2022-05-22) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 2ebbcd9..c21bcad 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -51,3 +51,22 @@ def test_binary_format_non_contiguous(self): assert not embedding.data.contiguous res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, np.array([3, 2, 1.5])) + + def test_text_copy(self): + embedding = np.array([1.5, 2, 3]) + cur = conn.cursor() + with cur.copy("COPY item (embedding) FROM STDIN") as copy: + copy.write_row([embedding]) + + def test_binary_copy(self): + embedding = np.array([1.5, 2, 3]) + cur = conn.cursor() + with cur.copy("COPY item (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.write_row([embedding]) + + def test_binary_copy_set_types(self): + embedding = np.array([1.5, 2, 3]) + cur = conn.cursor() + with cur.copy("COPY item (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.set_types(['int8', 'vector']) + copy.write_row([1, embedding]) From dcf931ca7149ef01f2478473c63868985fb54222 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 3 May 2023 19:18:55 -0700 Subject: [PATCH 033/879] Updated comment [skip ci] --- pgvector/psycopg/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 69e7f68..a9659aa 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -49,7 +49,7 @@ def register_vector(context): raise psycopg.ProgrammingError('vector type not found in the database') info.register(context) - # add oid to anonymous class + # add oid to anonymous class for set_types text_dumper = type('', (VectorDumper,), {'oid': info.oid}) binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) From 8ea0541d048cd6af71c371272c1a5acf91a2cd78 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 May 2023 20:48:45 -0700 Subject: [PATCH 034/879] Test with multiple Python versions on CI --- .github/workflows/build.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e327f34..7de1a2c 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -3,11 +3,15 @@ on: [push, pull_request] jobs: build: runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python: [3.11, 3.6] steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 with: - python-version: 3.11 + python-version: ${{ matrix.python }} - run: pip install -r requirements.txt - uses: ankane/setup-postgres@v1 with: From 007467b5bd73f115d4813101a7f9da9b4203cb31 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 May 2023 20:51:59 -0700 Subject: [PATCH 035/879] Fixed CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7de1a2c..e6edcfb 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -2,7 +2,7 @@ name: build on: [push, pull_request] jobs: build: - runs-on: ubuntu-latest + runs-on: ubuntu-20.04 strategy: fail-fast: false matrix: From 866f038acf22e5ef2227f666373bfcd7b7c3e944 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 May 2023 21:12:42 -0700 Subject: [PATCH 036/879] Added register_vector_async for psycopg3 Co-authored-by: Alexander Koshelev --- CHANGELOG.md | 1 + README.md | 8 ++++++++ pgvector/psycopg/__init__.py | 9 +++++++++ tests/test_psycopg.py | 23 ++++++++++++++++++++++- 4 files changed, 40 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 616993b..59b7307 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.1.7 (unreleased) +- Added `register_vector_async` for psycopg3 - Fixed `set_types` for psycopg3 ## 0.1.6 (2022-05-22) diff --git a/README.md b/README.md index 1779b8d..7aa9919 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,14 @@ from pgvector.psycopg import register_vector register_vector(conn) ``` +For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use [unreleased] + +```python +from pgvector.psycopg import register_vector_async + +await register_vector_async(conn) +``` + Insert a vector ```python diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index a9659aa..bd398ec 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -45,6 +45,15 @@ def load(self, data): def register_vector(context): info = TypeInfo.fetch(context, 'vector') + register_vector_info(context, info) + + +async def register_vector_async(context): + info = await TypeInfo.fetch(context, 'vector') + register_vector_info(context, info) + + +def register_vector_info(context, info): if info is None: raise psycopg.ProgrammingError('vector type not found in the database') info.register(context) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index c21bcad..e2f52b5 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,6 +1,7 @@ import numpy as np -from pgvector.psycopg import register_vector +from pgvector.psycopg import register_vector, register_vector_async import psycopg +import pytest conn = psycopg.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -70,3 +71,23 @@ def test_binary_copy_set_types(self): with cur.copy("COPY item (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector']) copy.write_row([1, embedding]) + + @pytest.mark.asyncio + async def test_async(self): + conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) + + await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + await conn.execute('DROP TABLE IF EXISTS item') + await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') + + await register_vector_async(conn) + + embedding = np.array([1.5, 2, 3]) + await conn.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,)) + + async with conn.cursor() as cur: + await cur.execute('SELECT * FROM item ORDER BY id') + res = await cur.fetchall() + assert np.array_equal(res[0][1], embedding) + assert res[0][1].dtype == np.float32 + assert res[1][1] is None From 6a05297516474b7cb0810c80c7a9ce39d93c40fb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 May 2023 21:21:12 -0700 Subject: [PATCH 037/879] Fixed CI --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e6edcfb..62683dd 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -2,11 +2,11 @@ name: build on: [push, pull_request] jobs: build: - runs-on: ubuntu-20.04 + runs-on: ubuntu-latest strategy: fail-fast: false matrix: - python: [3.11, 3.6] + python: [3.11, 3.7] steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 From 57226caaf3e02085b2f90d8ead66655a468be664 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 11 May 2023 09:26:36 -0700 Subject: [PATCH 038/879] Added example of getting items within a certain distance with SQLAlchemy - closes #11 --- README.md | 6 ++++++ tests/test_sqlalchemy.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/README.md b/README.md index 7aa9919..da9560e 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,12 @@ session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).lim Also supports `max_inner_product` and `cosine_distance` +Get items within a certain distance + +```python +session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) +``` + Add an approximate index ```python diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 52fe09b..3d4d2af 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -105,6 +105,12 @@ def test_cosine_distance(self): items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] + def test_filter(self): + create_items() + with Session(engine) as session: + items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() + assert [v.id for v in items] == [1] + def test_select(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) From 8885197e0e4be7c674887a89e57638137a3285c0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 11 May 2023 09:39:54 -0700 Subject: [PATCH 039/879] Added example of getting items within a certain distance with SQLAlchemy ORM --- tests/test_sqlalchemy.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 3d4d2af..e3ca027 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -111,6 +111,12 @@ def test_filter(self): items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() assert [v.id for v in items] == [1] + def test_filter_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) + assert [v.id for v in items] == [1] + def test_select(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) From 1e2aca4fe74c1ef33b2d6fcac9fd7f110ace7c06 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 11 May 2023 14:43:29 -0700 Subject: [PATCH 040/879] Added more SQLAlchemy ORM tests --- tests/test_sqlalchemy.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index e3ca027..63b53e6 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -93,18 +93,36 @@ def test_l2_distance(self): items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] + def test_l2_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_max_inner_product(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] + def test_max_inner_product_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) + assert [v.id for v in items] == [2, 3, 1] + def test_cosine_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] + def test_cosine_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 2, 3] + def test_filter(self): create_items() with Session(engine) as session: From 6f50e90db7efb8d133d061a364ee8096586ba778 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 11 May 2023 14:50:07 -0700 Subject: [PATCH 041/879] Added example of getting distance with SQLAlchemy --- README.md | 6 ++++++ tests/test_sqlalchemy.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/README.md b/README.md index da9560e..5337ec7 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,12 @@ session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).lim Also supports `max_inner_product` and `cosine_distance` +Get the distance + +```python +session.scalars(select(Item.embedding.l2_distance([3, 1, 2]))) +``` + Get items within a certain distance ```python diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 63b53e6..f0f103e 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -141,6 +141,12 @@ def test_select(self): item = session.query(Item.embedding.l2_distance([1, 1, 1])).first() assert item[0] == 3 + def test_select_orm(self): + with Session(engine) as session: + session.add(Item(embedding=[2, 3, 3])) + item = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() + assert item[0] == 3 + def test_bad_dimensions(self): item = Item(embedding=[1, 2]) session = Session(engine) From 7a2f45cbeba69bed9b6895da676136559e6f245e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 11 May 2023 14:55:35 -0700 Subject: [PATCH 042/879] Version bump to 0.1.7 [skip ci] --- CHANGELOG.md | 2 +- README.md | 2 +- setup.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 59b7307..9917993 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.1.7 (unreleased) +## 0.1.7 (2023-05-11) - Added `register_vector_async` for psycopg3 - Fixed `set_types` for psycopg3 diff --git a/README.md b/README.md index 5337ec7..6dc8ab7 100644 --- a/README.md +++ b/README.md @@ -192,7 +192,7 @@ from pgvector.psycopg import register_vector register_vector(conn) ``` -For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use [unreleased] +For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use ```python from pgvector.psycopg import register_vector_async diff --git a/setup.py b/setup.py index 82af1aa..2b0c709 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.1.6', + version='0.1.7', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From f6273e9ae6703d3b71d6813f27b7886bf23acdeb Mon Sep 17 00:00:00 2001 From: Kevin Wood Date: Sat, 13 May 2023 14:29:36 -0500 Subject: [PATCH 043/879] Fix Django serialization (#12) --- pgvector/django/__init__.py | 3 +++ tests/test_django.py | 11 +++++++++++ 2 files changed, 14 insertions(+) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 02e77cd..4fac58d 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -39,6 +39,9 @@ def to_python(self, value): def get_prep_value(self, value): return to_db(value) + def value_to_string(self, obj): + return to_db(self.value_from_object(obj)) + class IvfflatIndex(PostgresIndex): suffix = 'ivfflat' diff --git a/tests/test_django.py b/tests/test_django.py index 39ed689..bca8271 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,4 +1,5 @@ import django +from django.core import serializers from django.conf import settings from django.db import connection, migrations, models from django.db.migrations.loader import MigrationLoader @@ -6,6 +7,7 @@ import numpy as np import pgvector.django from pgvector.django import VectorExtension, VectorField, IvfflatIndex, L2Distance, MaxInnerProduct, CosineDistance +from unittest import mock settings.configure( DATABASES={ @@ -115,3 +117,12 @@ def test_filter(self): distance = L2Distance('embedding', [1, 1, 1]) items = Item.objects.alias(distance=distance).filter(distance__lt=1) assert [v.id for v in items] == [1] + + def test_serialization(self): + create_items() + items = Item.objects.all() + data = serializers.serialize('json', items) + with mock.patch('django.core.serializers.python.apps.get_model') as get_model: + get_model.return_value = Item + for obj in serializers.deserialize('json', data, ): + obj.save() From 1bad2e6f779bfaeb68652d32df3a9656d37f04f1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:31:52 -0700 Subject: [PATCH 044/879] Use get_prep_value to match Django docs [skip ci] --- pgvector/django/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 4fac58d..d0da5a1 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -40,7 +40,7 @@ def get_prep_value(self, value): return to_db(value) def value_to_string(self, obj): - return to_db(self.value_from_object(obj)) + return self.get_prep_value(self.value_from_object(obj)) class IvfflatIndex(PostgresIndex): From 2a9cc90a6a5e82ed3d5d55d53664819e08e92c70 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:32:10 -0700 Subject: [PATCH 045/879] Updated link [skip ci] --- pgvector/django/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index d0da5a1..5c1e51d 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -11,7 +11,7 @@ def __init__(self): self.name = 'vector' -# https://docs.djangoproject.com/en/3.2/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ class VectorField(Field): description = 'Vector' From c9de0942aaa5986deb972e4ae5e98a2f2b400733 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:32:22 -0700 Subject: [PATCH 046/879] Fixed order [skip ci] --- tests/test_django.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index bca8271..8307938 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,6 +1,6 @@ import django -from django.core import serializers from django.conf import settings +from django.core import serializers from django.db import connection, migrations, models from django.db.migrations.loader import MigrationLoader from math import sqrt From 38e4453b6c2fdb2003e81d068a5ad83fbf10ae2c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:32:44 -0700 Subject: [PATCH 047/879] Updated license year [skip ci] --- LICENSE.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/LICENSE.txt b/LICENSE.txt index 5474bcb..b3134ac 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2022 Andrew Kane +Copyright (c) 2021-2023 Andrew Kane Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal From 47925f7e8192ecd9526505294f0cff245420f6b7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:34:45 -0700 Subject: [PATCH 048/879] Updated changelog [skip ci] --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9917993..8d16e30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.1.8 (unreleased) + +- Fixed serialization with Django + ## 0.1.7 (2023-05-11) - Added `register_vector_async` for psycopg3 From 181a232b95c97a80a6dfb97ad1f186ae916a12fd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:38:23 -0700 Subject: [PATCH 049/879] Simplified code --- tests/test_django.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 8307938..7c5acda 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -124,5 +124,5 @@ def test_serialization(self): data = serializers.serialize('json', items) with mock.patch('django.core.serializers.python.apps.get_model') as get_model: get_model.return_value = Item - for obj in serializers.deserialize('json', data, ): + for obj in serializers.deserialize('json', data): obj.save() From 9d76928fd8cc93408e442bbb31ef7465743040e9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 13 May 2023 12:40:40 -0700 Subject: [PATCH 050/879] Test multiple formats --- tests/test_django.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 7c5acda..d6298ea 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -121,8 +121,9 @@ def test_filter(self): def test_serialization(self): create_items() items = Item.objects.all() - data = serializers.serialize('json', items) - with mock.patch('django.core.serializers.python.apps.get_model') as get_model: - get_model.return_value = Item - for obj in serializers.deserialize('json', data): - obj.save() + for format in ['json', 'xml']: + data = serializers.serialize(format, items) + with mock.patch('django.core.serializers.python.apps.get_model') as get_model: + get_model.return_value = Item + for obj in serializers.deserialize(format, data): + obj.save() From 60fc9c66ca27a7b870558ad1a850c601e016b0fc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 20 May 2023 17:31:44 -0700 Subject: [PATCH 051/879] Version bump to 0.1.8 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d16e30..3c2835e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.1.8 (unreleased) +## 0.1.8 (2023-05-20) - Fixed serialization with Django diff --git a/setup.py b/setup.py index 2b0c709..5945a3d 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.1.7', + version='0.1.8', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From b95d4a8e8ef0a17bc399819a8b0d2c37999078c0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 28 May 2023 12:59:26 -0700 Subject: [PATCH 052/879] Improved examples [skip ci] --- examples/pytorch_image_search.py | 5 ++--- tests/test_psycopg.py | 3 +-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index 21d421e..d0e3ec9 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -10,8 +10,7 @@ # establish connection -conn = psycopg.connect(dbname='pgvector_example') -conn.autocommit = True +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') register_vector(conn) @@ -38,7 +37,7 @@ def generate_embeddings(inputs): # generate, save, and index embeddings if seed: conn.execute('DROP TABLE IF EXISTS image') - conn.execute('CREATE TABLE image (id bigserial primary key, embedding vector(512))') + conn.execute('CREATE TABLE image (id bigserial PRIMARY KEY, embedding vector(512))') for data in tqdm(dataloader): embeddings = generate_embeddings(data[0]) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index e2f52b5..73c54d9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -3,8 +3,7 @@ import psycopg import pytest -conn = psycopg.connect(dbname='pgvector_python_test') -conn.autocommit = True +conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') conn.execute('DROP TABLE IF EXISTS item') From ccbb1a9daabb4e5f4c992f98efbfa716cdb96fe3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 31 May 2023 16:01:56 -0700 Subject: [PATCH 053/879] Added hybrid search example [skip ci] --- README.md | 1 + examples/hybrid_search.py | 73 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) create mode 100644 examples/hybrid_search.py diff --git a/README.md b/README.md index 6dc8ab7..7235041 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Or check out some examples: - [Embeddings](examples/openai_embeddings.py) with OpenAI - [Sentence embeddings](examples/sentence_embeddings.py) with SentenceTransformers +- [Hybrid search](examples/hybrid_search.py) with SentenceTransformers - [Image search](examples/pytorch_image_search.py) with PyTorch - [Implicit feedback recommendations](examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](examples/surprise_recs.py) with Surprise diff --git a/examples/hybrid_search.py b/examples/hybrid_search.py new file mode 100644 index 0000000..c2ebd89 --- /dev/null +++ b/examples/hybrid_search.py @@ -0,0 +1,73 @@ +# good resources +# https://qdrant.tech/articles/hybrid-search/ +# https://www.sbert.net/examples/applications/semantic-search/README.html + +import asyncio +import itertools +from pgvector.psycopg import register_vector_async +import psycopg +from sentence_transformers import CrossEncoder, SentenceTransformer + +sentences = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +query = 'growling bear' + + +async def create_schema(conn): + await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + await register_vector_async(conn) + + await conn.execute('DROP TABLE IF EXISTS document') + await conn.execute('CREATE TABLE document (id bigserial PRIMARY KEY, content text, embedding vector(384))') + await conn.execute("CREATE INDEX ON document USING GIN (to_tsvector('english', content))") + + +async def insert_data(conn): + model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') + embeddings = model.encode(sentences) + + sql = 'INSERT INTO document (content, embedding) VALUES ' + ', '.join(['(%s, %s)' for _ in embeddings]) + params = list(itertools.chain(*zip(sentences, embeddings))) + await conn.execute(sql, params) + + +async def semantic_search(conn, query): + model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') + embedding = model.encode(query) + + async with conn.cursor() as cur: + await cur.execute('SELECT id, content FROM document ORDER BY embedding <=> %s LIMIT 5', (embedding,)) + return await cur.fetchall() + + +async def keyword_search(conn, query): + async with conn.cursor() as cur: + await cur.execute("SELECT id, content FROM document, plainto_tsquery('english', %s) query WHERE to_tsvector('english', content) @@ query ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC LIMIT 5", (query,)) + return await cur.fetchall() + + +def rerank(query, results): + # deduplicate + results = set(itertools.chain(*results)) + + # re-rank + encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') + scores = encoder.predict([(query, item[1]) for item in results]) + return [v for _, v in sorted(zip(scores, results), reverse=True)] + + +async def main(): + conn = await psycopg.AsyncConnection.connect(dbname='pgvector_example', autocommit=True) + await create_schema(conn) + await insert_data(conn) + + # perform queries in parallel + results = await asyncio.gather(semantic_search(conn, query), keyword_search(conn, query)) + results = rerank(query, results) + print(results) + + +asyncio.run(main()) From 34db96b3c1cb420fc1350b6540dcb10dd765ef68 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 28 Jun 2023 18:30:12 -0400 Subject: [PATCH 054/879] Added pool example and test for asyncpg - closes #18 --- README.md | 9 +++++++++ tests/test_asyncpg.py | 22 ++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/README.md b/README.md index 7235041..c6693c0 100644 --- a/README.md +++ b/README.md @@ -248,6 +248,15 @@ from pgvector.asyncpg import register_vector await register_vector(conn) ``` +or your pool + +```python +async def init(conn): + await register_vector(conn) + +pool = await asyncpg.create_pool(..., init=register_vector) +``` + Insert a vector ```python diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 7f7dc80..2634310 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -30,3 +30,25 @@ async def test_works(self): assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() + + @pytest.mark.asyncio + async def test_pool(self): + async def init(conn): + await register_vector(conn) + + pool = await asyncpg.create_pool(database='pgvector_python_test', init=init) + + async with pool.acquire() as conn: + await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + await conn.execute('DROP TABLE IF EXISTS item') + await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') + + embedding = np.array([1.5, 2, 3]) + await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding) + + res = await conn.fetch("SELECT * FROM item ORDER BY id") + assert res[0]['id'] == 1 + assert res[1]['id'] == 2 + assert np.array_equal(res[0]['embedding'], embedding) + assert res[0]['embedding'].dtype == np.float32 + assert res[1]['embedding'] is None From 62a08391b4a49d3f11b79659b24d2c97e9707f08 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 28 Jun 2023 18:33:47 -0400 Subject: [PATCH 055/879] Fixed example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c6693c0..6406c85 100644 --- a/README.md +++ b/README.md @@ -254,7 +254,7 @@ or your pool async def init(conn): await register_vector(conn) -pool = await asyncpg.create_pool(..., init=register_vector) +pool = await asyncpg.create_pool(..., init=init) ``` Insert a vector From 38173ede3f4d1b59b0e42a8bcc3a10735e00b914 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 2 Jul 2023 11:31:03 -0400 Subject: [PATCH 056/879] Dropped support for Python < 3.8 [skip ci] --- .github/workflows/build.yml | 2 +- CHANGELOG.md | 4 ++++ setup.py | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 62683dd..9e7fd5d 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.11, 3.7] + python: [3.11, 3.8] steps: - uses: actions/checkout@v3 - uses: actions/setup-python@v4 diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c2835e..bc616d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.0 (unreleased) + +- Dropped support for Python < 3.8 + ## 0.1.8 (2023-05-20) - Fixed serialization with Django diff --git a/setup.py b/setup.py index 5945a3d..64c6678 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,7 @@ 'pgvector.sqlalchemy', 'pgvector.utils' ], - python_requires='>=3.6', + python_requires='>=3.8', install_requires=[ 'numpy' ], From 7682bb816b4411d9d0a06687fd5a1f52bd1b810c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 2 Jul 2023 11:32:54 -0400 Subject: [PATCH 057/879] Updated pgvector on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 9e7fd5d..5d88016 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.4.1 https://github.com/pgvector/pgvector.git + git clone --branch v0.4.4 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From 5611b36c2bb65a02b189fef409b139aca87f3186 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 23 Jul 2023 09:12:03 -0700 Subject: [PATCH 058/879] Fixed form validation with Django - closes #21 --- CHANGELOG.md | 1 + pgvector/django/__init__.py | 6 ++++++ tests/test_django.py | 11 +++++++++++ 3 files changed, 18 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index bc616d1..97175b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.2.0 (unreleased) +- Fixed form validation with Django - Dropped support for Python < 3.8 ## 0.1.8 (2023-05-20) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 5c1e51d..2a1d101 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -42,6 +42,12 @@ def get_prep_value(self, value): def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) + def validate(self, value, model_instance): + super().validate(value.tolist(), model_instance) + + def run_validators(self, value): + super().run_validators(value.tolist()) + class IvfflatIndex(PostgresIndex): suffix = 'ivfflat' diff --git a/tests/test_django.py b/tests/test_django.py index d6298ea..ab02851 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -3,6 +3,7 @@ from django.core import serializers from django.db import connection, migrations, models from django.db.migrations.loader import MigrationLoader +from django.forms import ModelForm from math import sqrt import numpy as np import pgvector.django @@ -79,6 +80,12 @@ def create_items(): item.save() +class ItemForm(ModelForm): + class Meta: + model = Item + fields = ['embedding'] + + class TestDjango: def setup_method(self, test_method): Item.objects.all().delete() @@ -127,3 +134,7 @@ def test_serialization(self): get_model.return_value = Item for obj in serializers.deserialize(format, data): obj.save() + + def test_form(self): + form = ItemForm(data={'embedding': [1, 2, 3]}) + assert form.is_valid() From 305621e5dfd27a98f3ee855d437924a6782383ff Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 23 Jul 2023 20:06:36 -0700 Subject: [PATCH 059/879] Version bump to 0.2.0 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 97175b4..3eb67c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.0 (unreleased) +## 0.2.0 (2023-07-23) - Fixed form validation with Django - Dropped support for Python < 3.8 diff --git a/setup.py b/setup.py index 64c6678..79e777d 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.1.8', + version='0.2.0', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From d650dfef74d28134d2a9a6e38f3562e79f1224b8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 31 Jul 2023 18:59:17 -0700 Subject: [PATCH 060/879] Added more Django form tests --- tests/test_django.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index ab02851..3969d02 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -138,3 +138,19 @@ def test_serialization(self): def test_form(self): form = ItemForm(data={'embedding': [1, 2, 3]}) assert form.is_valid() + assert 'value="[1, 2, 3]"' in form.as_div() + + def test_form_instance(self): + Item(id=1, embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = ItemForm(instance=item) + # TODO fix + # assert 'value="[1. 2. 3.]"' in form.as_div() + + def test_form_save(self): + Item(id=1, embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = ItemForm(instance=item, data={'embedding': [4, 5, 6]}) + assert form.is_valid() + assert form.save() + assert [4, 5, 6], Item.objects.get(pk=1).embedding From 4ea06c041fdfdd67568401ce215f9b3b3143d29f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 31 Jul 2023 19:03:43 -0700 Subject: [PATCH 061/879] Improved test [skip ci] --- tests/test_django.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 3969d02..f34b392 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -145,7 +145,7 @@ def test_form_instance(self): item = Item.objects.get(pk=1) form = ItemForm(instance=item) # TODO fix - # assert 'value="[1. 2. 3.]"' in form.as_div() + # assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() def test_form_save(self): Item(id=1, embedding=[1, 2, 3]).save() From 557983b36b147e21fc298ae65c8ddae5d142ff63 Mon Sep 17 00:00:00 2001 From: Jerin Peter George Date: Mon, 31 Jul 2023 21:06:25 -0700 Subject: [PATCH 062/879] Enhances DjangoAdmin support (#22) --- pgvector/django/__init__.py | 4 ++++ pgvector/django/forms.py | 14 ++++++++++++++ pgvector/django/widgets.py | 11 +++++++++++ tests/test_django.py | 12 +++++++++--- 4 files changed, 38 insertions(+), 3 deletions(-) create mode 100644 pgvector/django/forms.py create mode 100644 pgvector/django/widgets.py diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 2a1d101..816e1d6 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,6 +1,7 @@ from django.contrib.postgres.operations import CreateExtension from django.contrib.postgres.indexes import PostgresIndex from django.db.models import Field, FloatField, Func, Value +from .forms import VectorFormField from ..utils import from_db, to_db __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance'] @@ -48,6 +49,9 @@ def validate(self, value, model_instance): def run_validators(self, value): super().run_validators(value.tolist()) + def formfield(self, **kwargs): + return super().formfield(form_class=VectorFormField, **kwargs) + class IvfflatIndex(PostgresIndex): suffix = 'ivfflat' diff --git a/pgvector/django/forms.py b/pgvector/django/forms.py new file mode 100644 index 0000000..128f8cc --- /dev/null +++ b/pgvector/django/forms.py @@ -0,0 +1,14 @@ +from django import forms +from .widgets import VectorWidget + + +class VectorFormField(forms.CharField): + widget = VectorWidget + + def has_changed(self, initial, data): + try: + initial = initial.tolist() + except AttributeError: + # initial could be None + pass + return super().has_changed(initial, data) diff --git a/pgvector/django/widgets.py b/pgvector/django/widgets.py new file mode 100644 index 0000000..a2bbb94 --- /dev/null +++ b/pgvector/django/widgets.py @@ -0,0 +1,11 @@ +from django import forms + + +class VectorWidget(forms.TextInput): + def format_value(self, value): + try: + value = value.tolist() + except AttributeError: + # value could be None + pass + return super().format_value(value) diff --git a/tests/test_django.py b/tests/test_django.py index f34b392..2b8b55a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -136,10 +136,16 @@ def test_serialization(self): obj.save() def test_form(self): - form = ItemForm(data={'embedding': [1, 2, 3]}) + form = ItemForm(data={'embedding': "[1, 2, 3]"}) assert form.is_valid() assert 'value="[1, 2, 3]"' in form.as_div() + def test_form_has_changed(self): + Item(id=1, embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = ItemForm(instance=item, data={'embedding': "[1, 2, 4]"}) + assert form.has_changed() + def test_form_instance(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) @@ -150,7 +156,7 @@ def test_form_instance(self): def test_form_save(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - form = ItemForm(instance=item, data={'embedding': [4, 5, 6]}) + form = ItemForm(instance=item, data={'embedding': "[4, 5, 6]"}) assert form.is_valid() assert form.save() - assert [4, 5, 6], Item.objects.get(pk=1).embedding + assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() From a715c68346616731beeb860b0b935a52b6e8eeee Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 31 Jul 2023 21:09:42 -0700 Subject: [PATCH 063/879] Updated tests --- tests/test_django.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 2b8b55a..0b594fa 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -136,27 +136,21 @@ def test_serialization(self): obj.save() def test_form(self): - form = ItemForm(data={'embedding': "[1, 2, 3]"}) + form = ItemForm(data={'embedding': '[1, 2, 3]'}) assert form.is_valid() assert 'value="[1, 2, 3]"' in form.as_div() - def test_form_has_changed(self): - Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) - form = ItemForm(instance=item, data={'embedding': "[1, 2, 4]"}) - assert form.has_changed() - def test_form_instance(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = ItemForm(instance=item) - # TODO fix - # assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() def test_form_save(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - form = ItemForm(instance=item, data={'embedding': "[4, 5, 6]"}) + form = ItemForm(instance=item, data={'embedding': '[4, 5, 6]'}) assert form.is_valid() assert form.save() + assert form.has_changed() assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() From ad5fcb7f2d36500b106b6cc9429fb804f7a23640 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 31 Jul 2023 21:10:16 -0700 Subject: [PATCH 064/879] Updated test [skip ci] --- tests/test_django.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 0b594fa..b9bf00f 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -150,7 +150,7 @@ def test_form_save(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = ItemForm(instance=item, data={'embedding': '[4, 5, 6]'}) + assert form.has_changed() assert form.is_valid() assert form.save() - assert form.has_changed() assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() From ae61511ec85fc91a3eadcff37113b670235fb6af Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 31 Jul 2023 21:12:15 -0700 Subject: [PATCH 065/879] Updated changelog [skip ci] --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3eb67c1..9931185 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.1 (unreleased) + +- Fixed form issues with Django + ## 0.2.0 (2023-07-23) - Fixed form validation with Django From e6ca2c27c6a6fa8c183eea675fb8c8593403bad6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 31 Jul 2023 21:20:19 -0700 Subject: [PATCH 066/879] Version bump to 0.2.1 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9931185..88801b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.1 (unreleased) +## 0.2.1 (2023-07-31) - Fixed form issues with Django diff --git a/setup.py b/setup.py index 79e777d..49d8396 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.2.0', + version='0.2.1', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From 6867e8b451ba4175ae53482243371f3824f0d0a3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 30 Aug 2023 20:14:27 -0700 Subject: [PATCH 067/879] Added HnswIndex for Django --- .github/workflows/build.yml | 2 +- CHANGELOG.md | 4 ++++ pgvector/django/__init__.py | 25 +++++++++++++++++++++++++ tests/test_django.py | 11 +++++++++-- 4 files changed, 39 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 5d88016..d711cd5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.4.4 https://github.com/pgvector/pgvector.git + git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install diff --git a/CHANGELOG.md b/CHANGELOG.md index 88801b9..fd948be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.2 (unreleased) + +- Added `HnswIndex` for Django + ## 0.2.1 (2023-07-31) - Fixed form issues with Django diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 816e1d6..954133f 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -73,6 +73,31 @@ def get_with_params(self): return with_params +class HnswIndex(PostgresIndex): + suffix = 'hnsw' + + def __init__(self, *expressions, m=None, ef_construction=None, **kwargs): + self.m = m + self.ef_construction = ef_construction + super().__init__(*expressions, **kwargs) + + def deconstruct(self): + path, args, kwargs = super().deconstruct() + if self.m is not None: + kwargs['m'] = self.m + if self.ef_construction is not None: + kwargs['ef_construction'] = self.ef_construction + return path, args, kwargs + + def get_with_params(self): + with_params = [] + if self.m is not None: + with_params.append('m = %d' % self.m) + if self.ef_construction is not None: + with_params.append('ef_construction = %d' % self.ef_construction) + return with_params + + class DistanceBase(Func): output_field = FloatField() diff --git a/tests/test_django.py b/tests/test_django.py index b9bf00f..d762141 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -7,7 +7,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, IvfflatIndex, L2Distance, MaxInnerProduct, CosineDistance +from pgvector.django import VectorExtension, VectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance from unittest import mock settings.configure( @@ -28,10 +28,17 @@ class Meta: app_label = 'myapp' indexes = [ IvfflatIndex( - name='my_index', + name='ivfflat_idx', fields=['embedding'], lists=100, opclasses=['vector_l2_ops'] + ), + HnswIndex( + name='hnsw_idx', + fields=['embedding'], + m=16, + ef_construction=100, + opclasses=['vector_l2_ops'] ) ] From ff9012eb03ce5a71143c5c57a898db290a5c176e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Sep 2023 12:59:20 -0700 Subject: [PATCH 068/879] Added support for Peewee - closes #26 and closes #27 Co-authored-by: John-Craig Borman --- CHANGELOG.md | 1 + README.md | 40 +++++++++++++++++++++- pgvector/peewee/__init__.py | 31 +++++++++++++++++ requirements.txt | 1 + setup.py | 1 + tests/test_peewee.py | 68 +++++++++++++++++++++++++++++++++++++ 6 files changed, 141 insertions(+), 1 deletion(-) create mode 100644 pgvector/peewee/__init__.py create mode 100644 tests/test_peewee.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fd948be..32a43e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.2.2 (unreleased) +- Added support for Peewee - Added `HnswIndex` for Django ## 0.2.1 (2023-07-31) diff --git a/README.md b/README.md index 6406c85..9229e30 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), and [asyncpg](https://github.com/MagicStack/asyncpg) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions) @@ -22,6 +22,7 @@ And follow the instructions for your database library: - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) +- [Peewee](#peewee) [unreleased] Or check out some examples: @@ -270,6 +271,43 @@ Get the nearest neighbors to a vector await conn.fetch('SELECT * FROM item ORDER BY embedding <-> $1 LIMIT 5', embedding) ``` +## Peewee + +Add a vector column + +```python +from pgvector.peewee import VectorField + +class Item(BaseModel): + embedding = VectorField(dimensions=3) +``` + +Insert a vector + +```python +item = Item.create(embedding=[1, 2, 3]) +``` + +Get the nearest neighbors to a vector + +```python +Item.select().order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5) +``` + +Also supports `max_inner_product` and `cosine_distance` + +Get the distance + +```python +Item.select(Item.embedding.l2_distance([3, 1, 2]).alias('distance')) +``` + +Get items within a certain distance + +```python +Item.select().where(Item.embedding.l2_distance([3, 1, 2]) < 5) +``` + ## History View the [changelog](https://github.com/pgvector/pgvector-python/blob/master/CHANGELOG.md) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py new file mode 100644 index 0000000..4014091 --- /dev/null +++ b/pgvector/peewee/__init__.py @@ -0,0 +1,31 @@ +from peewee import Expression, Field, Value +from ..utils import from_db, to_db + + +class VectorField(Field): + field_type = 'vector' + + def __init__(self, dimensions=None, *args, **kwargs): + self.dimensions = dimensions + super(VectorField, self).__init__(*args, **kwargs) + + def get_modifiers(self): + return self.dimensions and [self.dimensions] or None + + def db_value(self, value): + return to_db(value) + + def python_value(self, value): + return from_db(value) + + def _distance(self, op, vector): + return Expression(lhs=self, op=op, rhs=Value(vector, converter=to_db, unpack=False)) + + def l2_distance(self, vector): + return self._distance('<->', vector) + + def max_inner_product(self, vector): + return self._distance('<#>', vector) + + def cosine_distance(self, vector): + return self._distance('<=>', vector) diff --git a/requirements.txt b/requirements.txt index 712f909..e111e7d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ asyncpg Django numpy +peewee psycopg[binary] psycopg2-binary pytest diff --git a/setup.py b/setup.py index 49d8396..407762e 100644 --- a/setup.py +++ b/setup.py @@ -16,6 +16,7 @@ packages=[ 'pgvector.asyncpg', 'pgvector.django', + 'pgvector.peewee', 'pgvector.psycopg', 'pgvector.psycopg2', 'pgvector.sqlalchemy', diff --git a/tests/test_peewee.py b/tests/test_peewee.py new file mode 100644 index 0000000..b9017ab --- /dev/null +++ b/tests/test_peewee.py @@ -0,0 +1,68 @@ +from math import sqrt +import numpy as np +from peewee import Model, PostgresqlDatabase +from pgvector.peewee import VectorField + +db = PostgresqlDatabase('pgvector_python_test') + + +class BaseModel(Model): + class Meta: + database = db + + +class Item(BaseModel): + embedding = VectorField(dimensions=3) + + +db.connect() +db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') +db.drop_tables([Item]) +db.create_tables([Item]) + + +def create_items(): + vectors = [ + [1, 1, 1], + [2, 2, 2], + [1, 1, 2] + ] + for i, v in enumerate(vectors): + Item.create(id=i + 1, embedding=v) + + +class TestPeewee: + def setup_method(self, test_method): + Item.truncate_table() + + def test_works(self): + Item.create(id=1, embedding=[1, 2, 3]) + item = Item.get_by_id(1) + assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert item.embedding.dtype == np.float32 + + def test_l2_distance(self): + create_items() + distance = Item.embedding.l2_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, sqrt(3)] + + def test_max_inner_product(self): + create_items() + distance = Item.embedding.max_inner_product([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [-6, -4, -3] + + def test_cosine_distance(self): + create_items() + distance = Item.embedding.cosine_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 2, 3] + assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + + def test_where(self): + create_items() + items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) + assert [v.id for v in items] == [1] From a7ce7bb7a73a33083ccf4bd0373b065915a75a7d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Sep 2023 14:42:58 -0700 Subject: [PATCH 069/879] Added index docs for Peewee [skip ci] --- README.md | 8 ++++++++ tests/test_peewee.py | 2 ++ 2 files changed, 10 insertions(+) diff --git a/README.md b/README.md index 9229e30..ce20460 100644 --- a/README.md +++ b/README.md @@ -308,6 +308,14 @@ Get items within a certain distance Item.select().where(Item.embedding.l2_distance([3, 1, 2]) < 5) ``` +Add an approximate index + +```python +Item.add_index('embedding vector_l2_ops', using='hnsw') +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## History View the [changelog](https://github.com/pgvector/pgvector-python/blob/master/CHANGELOG.md) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index b9017ab..8041d1d 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -15,6 +15,8 @@ class Item(BaseModel): embedding = VectorField(dimensions=3) +Item.add_index('embedding vector_l2_ops', using='hnsw') + db.connect() db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') db.drop_tables([Item]) From 9c0c00f8ae9931891928a88201b9b6fa3541dd0b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 13:32:06 -0700 Subject: [PATCH 070/879] Added test and docs for HNSW --- README.md | 15 +++++++++++++++ tests/test_sqlalchemy.py | 15 ++++++++++++--- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index ce20460..91db1a2 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,14 @@ class Item(models.Model): fields=['embedding'], lists=100, opclasses=['vector_l2_ops'] + ), + # or + HnswIndex( + name='my_index', + fields=['embedding'], + m=16, + ef_construction=64, + opclasses=['vector_l2_ops'] ) ] ``` @@ -151,6 +159,13 @@ index = Index('my_index', Item.embedding, postgresql_with={'lists': 100}, postgresql_ops={'embedding': 'vector_l2_ops'} ) +# or +index = Index('my_index', Item.embedding, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) + index.create(engine) ``` diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f0f103e..14e48c0 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -55,14 +55,23 @@ def test_core(self): metadata.drop_all(engine) metadata.create_all(engine) - index = Index( - 'my_core_index', + ivfflat_index = Index( + 'ivfflat_core_index', item_table.c.embedding, postgresql_using='ivfflat', postgresql_with={'lists': 1}, postgresql_ops={'embedding': 'vector_l2_ops'} ) - index.create(engine) + ivfflat_index.create(engine) + + hnsw_index = Index( + 'hnsw_core_index', + item_table.c.embedding, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} + ) + hnsw_index.create(engine) def test_orm(self): item = Item(embedding=np.array([1.5, 2, 3])) From e63c63561ef0b319348c107724e875b42891bf25 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 13:37:36 -0700 Subject: [PATCH 071/879] Version bump to 0.2.2 [skip ci] --- CHANGELOG.md | 2 +- README.md | 2 +- setup.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 32a43e8..f1e1cbf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.2 (unreleased) +## 0.2.2 (2023-08-08) - Added support for Peewee - Added `HnswIndex` for Django diff --git a/README.md b/README.md index 91db1a2..b1eec40 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ And follow the instructions for your database library: - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) -- [Peewee](#peewee) [unreleased] +- [Peewee](#peewee) Or check out some examples: diff --git a/setup.py b/setup.py index 407762e..601d6c9 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.2.1', + version='0.2.2', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From 3cb04790c75ce9adeea0972a9240e1c1143eee14 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 13:38:47 -0700 Subject: [PATCH 072/879] Fixed example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b1eec40..ea461b2 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance_ Add an approximate index ```python -from pgvector.django import IvfflatIndex +from pgvector.django import IvfflatIndex, HnswIndex class Item(models.Model): class Meta: From ae3085a2570b74992d0b290641c91796b017593d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 15:19:09 -0700 Subject: [PATCH 073/879] Added HNSW index to __all__ --- pgvector/django/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 954133f..282edef 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -4,7 +4,7 @@ from .forms import VectorFormField from ..utils import from_db, to_db -__all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance'] +__all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance'] class VectorExtension(CreateExtension): From 0e26fe1ba22563214f541141277a9ee65ce8ecab Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 15:34:40 -0700 Subject: [PATCH 074/879] Converted OpenAI example to psycopg [skip ci] --- examples/openai_embeddings.py | 37 ++++++++++++----------------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index 795828f..671ffcf 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -1,27 +1,15 @@ import openai -from pgvector.sqlalchemy import Vector +from pgvector.psycopg import register_vector +import psycopg from sentence_transformers import SentenceTransformer -from sqlalchemy import create_engine, insert, select, text, Integer, String, Text -from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg://localhost/pgvector_example') -with engine.connect() as conn: - conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - conn.commit() +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) -Base = declarative_base() +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) - -class Document(Base): - __tablename__ = 'document' - - id = mapped_column(Integer, primary_key=True) - content = mapped_column(Text) - embedding = mapped_column(Vector(1536)) - - -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) +conn.execute('DROP TABLE IF EXISTS document') +conn.execute('CREATE TABLE document (id bigserial PRIMARY KEY, content text, embedding vector(1536))') input = [ 'The dog is barking', @@ -30,12 +18,11 @@ class Document(Base): ] embeddings = [v['embedding'] for v in openai.Embedding.create(input=input, model='text-embedding-ada-002')['data']] -documents = [dict(content=input[i], embedding=embedding) for i, embedding in enumerate(embeddings)] -session = Session(engine) -session.execute(insert(Document), documents) +for content, embedding in zip(input, embeddings): + conn.execute('INSERT INTO document (content, embedding) VALUES (%s, %s)', (content, embedding)) -doc = session.get(Document, 1) -neighbors = session.scalars(select(Document).filter(Document.id != doc.id).order_by(Document.embedding.max_inner_product(doc.embedding)).limit(5)) +document_id = 2 +neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() for neighbor in neighbors: - print(neighbor.content) + print(neighbor[0]) From db17185ffaaea4ba0c8e62eb639362f39f5ab295 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 15:35:44 -0700 Subject: [PATCH 075/879] Improved example [skip ci] --- examples/openai_embeddings.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index 671ffcf..abcc787 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -17,12 +17,13 @@ 'The bear is growling' ] -embeddings = [v['embedding'] for v in openai.Embedding.create(input=input, model='text-embedding-ada-002')['data']] +response = openai.Embedding.create(input=input, model='text-embedding-ada-002') +embeddings = [v['embedding'] for v in response['data']] for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO document (content, embedding) VALUES (%s, %s)', (content, embedding)) -document_id = 2 +document_id = 1 neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() for neighbor in neighbors: print(neighbor[0]) From 607376856b35f9c4b480289a7cf577d769ebffba Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 15:37:33 -0700 Subject: [PATCH 076/879] Updated example [skip ci] --- examples/openai_embeddings.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index abcc787..afd7273 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -8,8 +8,8 @@ conn.execute('CREATE EXTENSION IF NOT EXISTS vector') register_vector(conn) -conn.execute('DROP TABLE IF EXISTS document') -conn.execute('CREATE TABLE document (id bigserial PRIMARY KEY, content text, embedding vector(1536))') +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))') input = [ 'The dog is barking', @@ -21,7 +21,7 @@ embeddings = [v['embedding'] for v in response['data']] for content, embedding in zip(input, embeddings): - conn.execute('INSERT INTO document (content, embedding) VALUES (%s, %s)', (content, embedding)) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() From 105f74a59dc43312d3a2aff4eeae9cc914002168 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 15:38:59 -0700 Subject: [PATCH 077/879] Removed unneeded import from example [skip ci] --- examples/openai_embeddings.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index afd7273..08926f0 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -1,7 +1,6 @@ import openai from pgvector.psycopg import register_vector import psycopg -from sentence_transformers import SentenceTransformer conn = psycopg.connect(dbname='pgvector_example', autocommit=True) From e4c735983d7c972a3261340d776972f1049f7b0e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 15:40:16 -0700 Subject: [PATCH 078/879] Converted SentenceTransformers example to psycopg [skip ci] --- examples/sentence_embeddings.py | 38 +++++++++++---------------------- 1 file changed, 12 insertions(+), 26 deletions(-) diff --git a/examples/sentence_embeddings.py b/examples/sentence_embeddings.py index f6d0775..4b8b96a 100644 --- a/examples/sentence_embeddings.py +++ b/examples/sentence_embeddings.py @@ -1,26 +1,14 @@ -from pgvector.sqlalchemy import Vector +from pgvector.psycopg import register_vector +import psycopg from sentence_transformers import SentenceTransformer -from sqlalchemy import create_engine, insert, select, text, Integer, String, Text -from sqlalchemy.orm import declarative_base, mapped_column, Session -engine = create_engine('postgresql+psycopg://localhost/pgvector_example') -with engine.connect() as conn: - conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - conn.commit() +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) -Base = declarative_base() +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) - -class Document(Base): - __tablename__ = 'document' - - id = mapped_column(Integer, primary_key=True) - content = mapped_column(Text) - embedding = mapped_column(Vector(384)) - - -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') sentences = [ 'The dog is barking', @@ -31,12 +19,10 @@ class Document(Base): model = SentenceTransformer('all-MiniLM-L6-v2') embeddings = model.encode(sentences) -documents = [dict(content=sentences[i], embedding=embedding) for i, embedding in enumerate(embeddings)] - -session = Session(engine) -session.execute(insert(Document), documents) +for content, embedding in zip(sentences, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) -doc = session.get(Document, 1) -neighbors = session.scalars(select(Document).filter(Document.id != doc.id).order_by(Document.embedding.cosine_distance(doc.embedding)).limit(5)) +document_id = 1 +neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() for neighbor in neighbors: - print(neighbor.content) + print(neighbor[0]) From 8bc6a3da2d370d182045cc22b804d5c840be9f2a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 8 Sep 2023 16:15:18 -0700 Subject: [PATCH 079/879] Updated example [skip ci] --- examples/sentence_embeddings.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/sentence_embeddings.py b/examples/sentence_embeddings.py index 4b8b96a..d4e7f96 100644 --- a/examples/sentence_embeddings.py +++ b/examples/sentence_embeddings.py @@ -10,16 +10,16 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') -sentences = [ +input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] model = SentenceTransformer('all-MiniLM-L6-v2') -embeddings = model.encode(sentences) +embeddings = model.encode(input) -for content, embedding in zip(sentences, embeddings): +for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 From 615a21c993c19a5025616b8ae19ca53f9d5d530f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 9 Sep 2023 11:10:47 -0700 Subject: [PATCH 080/879] Fixed date in changelog [skip ci] --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f1e1cbf..5b73471 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.2 (2023-08-08) +## 0.2.2 (2023-09-08) - Added support for Peewee - Added `HnswIndex` for Django From dd42b1815061ffd5f1d03eb4d279bba3399b4e41 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 9 Sep 2023 11:33:37 -0700 Subject: [PATCH 081/879] Updated example [skip ci] --- examples/pytorch_image_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index d0e3ec9..480e746 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -46,7 +46,7 @@ def generate_embeddings(inputs): params = [embedding for embedding in embeddings] conn.execute(sql, params) - conn.execute('CREATE INDEX ON image USING ivfflat (embedding vector_cosine_ops)') + conn.execute('CREATE INDEX ON image USING hnsw (embedding vector_cosine_ops)') def show_images(dataset_images): From 20a78b3627516b2560a49870501ffb8459249b73 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 9 Sep 2023 11:41:34 -0700 Subject: [PATCH 082/879] Improved example [skip ci] --- examples/pytorch_image_search.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index 480e746..3f7ba92 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -25,13 +25,15 @@ # load pretrained model +device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') model = torchvision.models.resnet18(weights='DEFAULT') model.fc = torch.nn.Identity() +model.to(device) model.eval() def generate_embeddings(inputs): - return model(inputs).detach().numpy() + return model(inputs.to(device)).detach().cpu().numpy() # generate, save, and index embeddings @@ -53,7 +55,8 @@ def show_images(dataset_images): grid = torchvision.utils.make_grid(dataset_images) img = (grid / 2 + 0.5).permute(1, 2, 0).numpy() plt.imshow(img) - plt.waitforbuttonpress() + plt.draw() + plt.waitforbuttonpress(timeout=3) # load 5 random unseen images From c8c45b926fecdbd5de16b74379091d9eb13ea3bd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 9 Sep 2023 11:44:18 -0700 Subject: [PATCH 083/879] Improved example [skip ci] --- examples/pytorch_image_search.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index 3f7ba92..7cb587e 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -41,6 +41,7 @@ def generate_embeddings(inputs): conn.execute('DROP TABLE IF EXISTS image') conn.execute('CREATE TABLE image (id bigserial PRIMARY KEY, embedding vector(512))') + print('Generating embeddings') for data in tqdm(dataloader): embeddings = generate_embeddings(data[0]) @@ -48,6 +49,7 @@ def generate_embeddings(inputs): params = [embedding for embedding in embeddings] conn.execute(sql, params) + print('Creating index') conn.execute('CREATE INDEX ON image USING hnsw (embedding vector_cosine_ops)') From 8123b94a3b1342b6f4bec52bc53d59fa84a24540 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 9 Sep 2023 11:50:07 -0700 Subject: [PATCH 084/879] Updated example [skip ci] --- examples/pytorch_image_search.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index 7cb587e..0cc9af0 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -49,9 +49,6 @@ def generate_embeddings(inputs): params = [embedding for embedding in embeddings] conn.execute(sql, params) - print('Creating index') - conn.execute('CREATE INDEX ON image USING hnsw (embedding vector_cosine_ops)') - def show_images(dataset_images): grid = torchvision.utils.make_grid(dataset_images) From d686807348cb2b73492abc19db33ffedc70cbd62 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 15 Sep 2023 11:01:32 -0700 Subject: [PATCH 085/879] Added tests for aggregates with Django --- tests/test_django.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index d762141..0f539f4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -2,6 +2,7 @@ from django.conf import settings from django.core import serializers from django.db import connection, migrations, models +from django.db.models import Avg, Sum from django.db.migrations.loader import MigrationLoader from django.forms import ModelForm from math import sqrt @@ -132,6 +133,22 @@ def test_filter(self): items = Item.objects.alias(distance=distance).filter(distance__lt=1) assert [v.id for v in items] == [1] + def test_avg(self): + avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] + assert avg is None + Item(embedding=[1, 2, 3]).save() + Item(embedding=[4, 5, 6]).save() + avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] + assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + + def test_sum(self): + avg = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] + assert avg is None + Item(embedding=[1, 2, 3]).save() + Item(embedding=[4, 5, 6]).save() + avg = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] + assert np.array_equal(avg, np.array([5, 7, 9])) + def test_serialization(self): create_items() items = Item.objects.all() From 1a568d70dacd17776a3b151480f310dfe0d6d039 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 15 Sep 2023 15:44:35 -0700 Subject: [PATCH 086/879] Added test for get_or_create for Peewee - #31 --- tests/test_peewee.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 8041d1d..0c21e45 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -68,3 +68,7 @@ def test_where(self): create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) assert [v.id for v in items] == [1] + + def test_get_or_create(self): + Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) + Item.get_or_create(embedding=np.array([4, 5, 6])) From 695557e09310a45dba5b9b1fdd0416a941ed5671 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 15 Sep 2023 15:48:22 -0700 Subject: [PATCH 087/879] Simplified Peewee code --- pgvector/peewee/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 4014091..148f8b3 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -19,7 +19,7 @@ def python_value(self, value): return from_db(value) def _distance(self, op, vector): - return Expression(lhs=self, op=op, rhs=Value(vector, converter=to_db, unpack=False)) + return Expression(lhs=self, op=op, rhs=self.to_value(vector)) def l2_distance(self, vector): return self._distance('<->', vector) From 01159f5f81339d1978345f602993452aabe5575f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 15 Sep 2023 15:53:26 -0700 Subject: [PATCH 088/879] Added test for AsIs --- tests/test_peewee.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 0c21e45..755710b 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,6 +1,6 @@ from math import sqrt import numpy as np -from peewee import Model, PostgresqlDatabase +from peewee import AsIs, Model, PostgresqlDatabase from pgvector.peewee import VectorField db = PostgresqlDatabase('pgvector_python_test') @@ -72,3 +72,4 @@ def test_where(self): def test_get_or_create(self): 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=AsIs([7, 8, 9])) From fb571b573f17507ee3f17eb5219dc150f76a8688 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 15 Sep 2023 15:56:24 -0700 Subject: [PATCH 089/879] Updated get_or_create test --- tests/test_peewee.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 755710b..93cb037 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,6 +1,6 @@ from math import sqrt import numpy as np -from peewee import AsIs, Model, PostgresqlDatabase +from peewee import Model, PostgresqlDatabase from pgvector.peewee import VectorField db = PostgresqlDatabase('pgvector_python_test') @@ -72,4 +72,4 @@ def test_where(self): def test_get_or_create(self): 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=AsIs([7, 8, 9])) + Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) From b421a54336acf5def714146cb5e397921d5756fe Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 15 Sep 2023 16:44:56 -0700 Subject: [PATCH 090/879] Added example for averaging vectors with Django [skip ci] --- README.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/README.md b/README.md index ea461b2..d1ea4ba 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,16 @@ Get items within a certain distance Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance__lt=5) ``` +Average vectors + +```python +from django.db.models import Avg + +Item.objects.aggregate(Avg('embedding')) +``` + +Also supports `Sum` + Add an approximate index ```python From 2d8736bb49bfb6969dda9a15f2471b0f0ca26995 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 18 Sep 2023 15:37:45 -0700 Subject: [PATCH 091/879] Added docs for enabling extension [skip ci] --- README.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index d1ea4ba..ed60cb3 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ Or check out some examples: ## Django -Create the extension +Enable the extension ```python from pgvector.django import VectorExtension @@ -211,6 +211,12 @@ Also supports `max_inner_product` and `cosine_distance` ## Psycopg 3 +Enable the extension + +```python +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + Register the vector type with your connection ```python @@ -242,6 +248,12 @@ conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding, ## Psycopg 2 +Enable the extension + +```python +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + Register the vector type with your connection or cursor ```python @@ -266,6 +278,12 @@ cur.fetchall() ## asyncpg +Enable the extension + +```python +await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + Register the vector type with your connection ```python From a5953a92df09207f21c94715eddf6b1656930ea1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 18 Sep 2023 15:39:11 -0700 Subject: [PATCH 092/879] Improved example [skip ci] --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index ed60cb3..f7efff0 100644 --- a/README.md +++ b/README.md @@ -251,7 +251,8 @@ conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding, Enable the extension ```python -conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +cur = conn.cursor() +cur.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` Register the vector type with your connection or cursor From 454ec2d7a237d45ce3fb9c637d9b472fe4ea66f2 Mon Sep 17 00:00:00 2001 From: Keith Date: Sun, 24 Sep 2023 20:24:17 -0700 Subject: [PATCH 093/879] Add tests for Django clean and get_or_create functions and fix related issues (#34) --- pgvector/django/__init__.py | 3 +++ tests/test_django.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 282edef..af27f7d 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,3 +1,4 @@ +import numpy as np from django.contrib.postgres.operations import CreateExtension from django.contrib.postgres.indexes import PostgresIndex from django.db.models import Field, FloatField, Func, Value @@ -35,6 +36,8 @@ def from_db_value(self, value, expression, connection): return from_db(value) def to_python(self, value): + if isinstance(value, list): + return np.array(value, dtype=np.float32) return from_db(value) def get_prep_value(self, value): diff --git a/tests/test_django.py b/tests/test_django.py index 0f539f4..ea6dbf9 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -178,3 +178,10 @@ def test_form_save(self): assert form.is_valid() assert form.save() assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() + + def test_clean(self): + item = Item(id=1, embedding=[1, 2, 3]) + item.full_clean() + + def test_get_or_create(self): + Item.objects.get_or_create(embedding=[1, 2, 3]) From 5e34237cff657b07b03c6f134f28c0b4e16b0262 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 24 Sep 2023 20:48:27 -0700 Subject: [PATCH 094/879] Fixed null values for Django - fixes #35 --- CHANGELOG.md | 4 ++++ pgvector/django/__init__.py | 9 +++++++-- tests/test_django.py | 6 +++++- 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b73471..2355d1a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.3 (unreleased) + +- Fixed null values for Django + ## 0.2.2 (2023-09-08) - Added support for Peewee diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index af27f7d..c42d36a 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -16,6 +16,7 @@ def __init__(self): # https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ class VectorField(Field): description = 'Vector' + empty_strings_allowed = False def __init__(self, *args, dimensions=None, **kwargs): self.dimensions = dimensions @@ -47,10 +48,14 @@ def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) def validate(self, value, model_instance): - super().validate(value.tolist(), model_instance) + if isinstance(value, np.ndarray): + value = value.tolist() + super().validate(value, model_instance) def run_validators(self, value): - super().run_validators(value.tolist()) + if isinstance(value, np.ndarray): + value = value.tolist() + super().run_validators(value) def formfield(self, **kwargs): return super().formfield(form_class=VectorFormField, **kwargs) diff --git a/tests/test_django.py b/tests/test_django.py index ea6dbf9..28123b4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -56,7 +56,7 @@ class Migration(migrations.Migration): name='Item', fields=[ ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('embedding', pgvector.django.VectorField(dimensions=3)), + ('embedding', pgvector.django.VectorField(dimensions=3, null=True)), ], ), migrations.AddIndex( @@ -185,3 +185,7 @@ def test_clean(self): def test_get_or_create(self): Item.objects.get_or_create(embedding=[1, 2, 3]) + + def test_missing(self): + Item().save() + assert Item.objects.first().embedding is None From c66b1f5780cf4524590ac3ef28106a8ec11d6125 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 24 Sep 2023 20:50:19 -0700 Subject: [PATCH 095/879] Improved code --- pgvector/django/__init__.py | 2 +- pgvector/django/forms.py | 6 ++---- pgvector/django/widgets.py | 6 ++---- 3 files changed, 5 insertions(+), 9 deletions(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index c42d36a..ab250e6 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,7 +1,7 @@ -import numpy as np from django.contrib.postgres.operations import CreateExtension from django.contrib.postgres.indexes import PostgresIndex from django.db.models import Field, FloatField, Func, Value +import numpy as np from .forms import VectorFormField from ..utils import from_db, to_db diff --git a/pgvector/django/forms.py b/pgvector/django/forms.py index 128f8cc..3748236 100644 --- a/pgvector/django/forms.py +++ b/pgvector/django/forms.py @@ -1,4 +1,5 @@ from django import forms +import numpy as np from .widgets import VectorWidget @@ -6,9 +7,6 @@ class VectorFormField(forms.CharField): widget = VectorWidget def has_changed(self, initial, data): - try: + if isinstance(initial, np.ndarray): initial = initial.tolist() - except AttributeError: - # initial could be None - pass return super().has_changed(initial, data) diff --git a/pgvector/django/widgets.py b/pgvector/django/widgets.py index a2bbb94..731d632 100644 --- a/pgvector/django/widgets.py +++ b/pgvector/django/widgets.py @@ -1,11 +1,9 @@ from django import forms +import numpy as np class VectorWidget(forms.TextInput): def format_value(self, value): - try: + if isinstance(value, np.ndarray): value = value.tolist() - except AttributeError: - # value could be None - pass return super().format_value(value) From 1d7ff8853f5e2c198f6a07791d660d5b412d93de Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 24 Sep 2023 20:51:38 -0700 Subject: [PATCH 096/879] Updated changelog [skip ci] --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2355d1a..9066920 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.2.3 (unreleased) -- Fixed null values for Django +- Fixed null values with Django +- Fixed `full_clean` with Django ## 0.2.2 (2023-09-08) From 8ff7d3ffe5e81f02acb3128cd822dc64d494f2e0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 25 Sep 2023 18:28:57 -0700 Subject: [PATCH 097/879] Version bump to 0.2.3 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9066920..a8dc947 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.3 (unreleased) +## 0.2.3 (2023-09-25) - Fixed null values with Django - Fixed `full_clean` with Django diff --git a/setup.py b/setup.py index 601d6c9..227b6cf 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.2.2', + version='0.2.3', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From 1d8f0bc813cdc6e8caa9ac66af5de03a023ed69b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 20:43:45 -0700 Subject: [PATCH 098/879] Fixed variable name [skip ci] --- tests/test_django.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 28123b4..5b9c386 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -142,12 +142,12 @@ def test_avg(self): assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) def test_sum(self): - avg = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] - assert avg is None + sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] + assert sum is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() - avg = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] - assert np.array_equal(avg, np.array([5, 7, 9])) + sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] + assert np.array_equal(sum, np.array([5, 7, 9])) def test_serialization(self): create_items() From caed7d6aeb892c54a4fda4674da9eb08a5a0c0d6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 20:44:47 -0700 Subject: [PATCH 099/879] Added tests for aggregates with SQLAlchemy --- tests/test_sqlalchemy.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 14e48c0..932257f 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -4,6 +4,7 @@ from sqlalchemy import create_engine, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.sql import func engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with engine.connect() as con: @@ -144,6 +145,24 @@ def test_filter_orm(self): items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] + def test_avg(self): + with Session(engine) as session: + avg = session.query(func.avg(Item.embedding)).first()[0] + assert avg 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] + assert np.array_equal(avg, 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 + 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])) + def test_select(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) From 446859c905f1e5d477ea9195703c3f773ec0f763 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 20:49:47 -0700 Subject: [PATCH 100/879] Added tests for aggregates with Peewee --- tests/test_peewee.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 93cb037..38acb7c 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,6 +1,6 @@ from math import sqrt import numpy as np -from peewee import Model, PostgresqlDatabase +from peewee import Model, PostgresqlDatabase, fn from pgvector.peewee import VectorField db = PostgresqlDatabase('pgvector_python_test') @@ -69,6 +69,22 @@ def test_where(self): items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) assert [v.id for v in items] == [1] + def test_avg(self): + avg = Item.select(fn.avg(Item.embedding)).scalar() + assert avg is None + Item.create(embedding=[1, 2, 3]) + Item.create(embedding=[4, 5, 6]) + avg = Item.select(fn.avg(Item.embedding)).scalar() + assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + + def test_sum(self): + sum = Item.select(fn.sum(Item.embedding)).scalar() + assert sum is None + Item.create(embedding=[1, 2, 3]) + Item.create(embedding=[4, 5, 6]) + sum = Item.select(fn.sum(Item.embedding)).scalar() + assert np.array_equal(sum, np.array([5, 7, 9])) + def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) Item.get_or_create(embedding=np.array([4, 5, 6])) From d0e13cc4d81ef4bebd15216b6cbc37b9106a93d7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 20:52:49 -0700 Subject: [PATCH 101/879] Moved tests [skip ci] --- tests/test_sqlalchemy.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 932257f..646b908 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -145,6 +145,18 @@ def test_filter_orm(self): items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] + def test_select(self): + with Session(engine) as session: + session.add(Item(embedding=[2, 3, 3])) + item = session.query(Item.embedding.l2_distance([1, 1, 1])).first() + assert item[0] == 3 + + def test_select_orm(self): + with Session(engine) as session: + session.add(Item(embedding=[2, 3, 3])) + item = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() + assert item[0] == 3 + def test_avg(self): with Session(engine) as session: avg = session.query(func.avg(Item.embedding)).first()[0] @@ -163,18 +175,6 @@ def test_sum(self): sum = session.query(func.sum(Item.embedding)).first()[0] assert np.array_equal(sum, np.array([5, 7, 9])) - def test_select(self): - with Session(engine) as session: - session.add(Item(embedding=[2, 3, 3])) - item = session.query(Item.embedding.l2_distance([1, 1, 1])).first() - assert item[0] == 3 - - def test_select_orm(self): - with Session(engine) as session: - session.add(Item(embedding=[2, 3, 3])) - item = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() - assert item[0] == 3 - def test_bad_dimensions(self): item = Item(embedding=[1, 2]) session = Session(engine) From a3b8e846c99749bf2e51cfdf030b4db1575dcf8a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 20:53:47 -0700 Subject: [PATCH 102/879] Added more docs on aggregates [skip ci] --- README.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/README.md b/README.md index f7efff0..d2a5f14 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,16 @@ Get items within a certain distance session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) ``` +Average vectors + +```python +from sqlalchemy.sql import func + +session.query(func.avg(Item.embedding)).first()[0] +``` + +Also supports `sum` + Add an approximate index ```python @@ -352,6 +362,16 @@ Get items within a certain distance Item.select().where(Item.embedding.l2_distance([3, 1, 2]) < 5) ``` +Average vectors + +```python +from peewee import fn + +Item.select(fn.avg(Item.embedding)).scalar() +``` + +Also supports `sum` + Add an approximate index ```python From bd36d645c8b82e7b2d76208527099de1628e355a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 21:02:19 -0700 Subject: [PATCH 103/879] Use session [skip ci] --- tests/test_sqlalchemy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 646b908..21b46ac 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -7,9 +7,9 @@ from sqlalchemy.sql import func engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -with engine.connect() as con: - con.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - con.commit() +with Session(engine) as session: + session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) + session.commit() Base = declarative_base() From 8c151f5aab2b0af2a672ac839a43b94ef5473e06 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 1 Oct 2023 21:02:31 -0700 Subject: [PATCH 104/879] Updated readme [skip ci] --- README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/README.md b/README.md index d2a5f14..351be4a 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,12 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc ## SQLAlchemy +Enable the extension + +```python +session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) +``` + Add a vector column ```python @@ -193,6 +199,12 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc ## SQLModel +Enable the extension + +```python +session.exec('CREATE EXTENSION IF NOT EXISTS vector') +``` + Add a vector column ```python From cf8dd1c69397b1aaefa638a0e5755c7042295c3c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Oct 2023 10:16:19 -0700 Subject: [PATCH 105/879] Updated table name in example [skip ci] --- examples/hybrid_search.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/hybrid_search.py b/examples/hybrid_search.py index c2ebd89..c4e6eea 100644 --- a/examples/hybrid_search.py +++ b/examples/hybrid_search.py @@ -20,16 +20,16 @@ async def create_schema(conn): await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') await register_vector_async(conn) - await conn.execute('DROP TABLE IF EXISTS document') - await conn.execute('CREATE TABLE document (id bigserial PRIMARY KEY, content text, embedding vector(384))') - await conn.execute("CREATE INDEX ON document USING GIN (to_tsvector('english', content))") + await conn.execute('DROP TABLE IF EXISTS documents') + await conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') + await conn.execute("CREATE INDEX ON documents USING GIN (to_tsvector('english', content))") async def insert_data(conn): model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') embeddings = model.encode(sentences) - sql = 'INSERT INTO document (content, embedding) VALUES ' + ', '.join(['(%s, %s)' for _ in embeddings]) + sql = 'INSERT INTO documents (content, embedding) VALUES ' + ', '.join(['(%s, %s)' for _ in embeddings]) params = list(itertools.chain(*zip(sentences, embeddings))) await conn.execute(sql, params) @@ -39,13 +39,13 @@ async def semantic_search(conn, query): embedding = model.encode(query) async with conn.cursor() as cur: - await cur.execute('SELECT id, content FROM document ORDER BY embedding <=> %s LIMIT 5', (embedding,)) + await cur.execute('SELECT id, content FROM documents ORDER BY embedding <=> %s LIMIT 5', (embedding,)) return await cur.fetchall() async def keyword_search(conn, query): async with conn.cursor() as cur: - await cur.execute("SELECT id, content FROM document, plainto_tsquery('english', %s) query WHERE to_tsvector('english', content) @@ query ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC LIMIT 5", (query,)) + await cur.execute("SELECT id, content FROM documents, plainto_tsquery('english', %s) query WHERE to_tsvector('english', content) @@ query ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC LIMIT 5", (query,)) return await cur.fetchall() From 188d8a4f479dc49a9b1f18507ad8ceb0beba9a05 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Oct 2023 10:35:44 -0700 Subject: [PATCH 106/879] Added example of Reciprocal Rank Fusion in SQL [skip ci] --- examples/hybrid_search_rrf.py | 50 +++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 examples/hybrid_search_rrf.py diff --git a/examples/hybrid_search_rrf.py b/examples/hybrid_search_rrf.py new file mode 100644 index 0000000..d7fd307 --- /dev/null +++ b/examples/hybrid_search_rrf.py @@ -0,0 +1,50 @@ +from pgvector.psycopg import register_vector +import psycopg +from sentence_transformers import SentenceTransformer + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') +conn.execute("CREATE INDEX ON documents USING GIN (to_tsvector('english', content))") + +sentences = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') +embeddings = model.encode(sentences) +for content, embedding in zip(sentences, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) + +sql = """ +WITH semantic_search AS ( + SELECT id, RANK () OVER (ORDER BY embedding <=> %(embedding)s) AS rank + FROM documents + ORDER BY embedding <=> %(embedding)s + LIMIT 20 +), +keyword_search AS ( + SELECT id, RANK () OVER (ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC) + FROM documents, plainto_tsquery('english', %(query)s) query + WHERE to_tsvector('english', content) @@ query + ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC + LIMIT 20 +) +SELECT + COALESCE(semantic_search.id, keyword_search.id) AS id, + COALESCE(1.0 / (%(k)s + semantic_search.rank), 0.0) + COALESCE(1.0 / (%(k)s + keyword_search.rank), 0.0) AS score +FROM semantic_search +FULL OUTER JOIN keyword_search ON semantic_search.id = keyword_search.id +ORDER BY score DESC +LIMIT 10 +""" +query = 'growling bear' +embedding = model.encode(query) +k = 60 +results = conn.execute(sql, {'query': query, 'embedding': embedding, 'k': k}).fetchall() +print(results) From 22c822584342d2fdb33e08219c30ee000bb280f8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Oct 2023 10:46:18 -0700 Subject: [PATCH 107/879] Updated example [skip ci] --- examples/hybrid_search_rrf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hybrid_search_rrf.py b/examples/hybrid_search_rrf.py index d7fd307..ffb914e 100644 --- a/examples/hybrid_search_rrf.py +++ b/examples/hybrid_search_rrf.py @@ -41,7 +41,7 @@ FROM semantic_search FULL OUTER JOIN keyword_search ON semantic_search.id = keyword_search.id ORDER BY score DESC -LIMIT 10 +LIMIT 5 """ query = 'growling bear' embedding = model.encode(query) From 4c3e8b7973b6568ddd64f9d961dc03b4f55e573e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Oct 2023 10:51:42 -0700 Subject: [PATCH 108/879] Updated example [skip ci] --- examples/hybrid_search_rrf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/hybrid_search_rrf.py b/examples/hybrid_search_rrf.py index ffb914e..601bd5a 100644 --- a/examples/hybrid_search_rrf.py +++ b/examples/hybrid_search_rrf.py @@ -37,7 +37,8 @@ ) SELECT COALESCE(semantic_search.id, keyword_search.id) AS id, - COALESCE(1.0 / (%(k)s + semantic_search.rank), 0.0) + COALESCE(1.0 / (%(k)s + keyword_search.rank), 0.0) AS score + COALESCE(1.0 / (%(k)s + semantic_search.rank), 0.0) + + COALESCE(1.0 / (%(k)s + keyword_search.rank), 0.0) AS score FROM semantic_search FULL OUTER JOIN keyword_search ON semantic_search.id = keyword_search.id ORDER BY score DESC From 93f14a88aebf65c70194dddb7129a41bea06f88a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Oct 2023 10:55:28 -0700 Subject: [PATCH 109/879] Updated example [skip ci] --- examples/hybrid_search_rrf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/hybrid_search_rrf.py b/examples/hybrid_search_rrf.py index 601bd5a..74d60bf 100644 --- a/examples/hybrid_search_rrf.py +++ b/examples/hybrid_search_rrf.py @@ -48,4 +48,5 @@ embedding = model.encode(query) k = 60 results = conn.execute(sql, {'query': query, 'embedding': embedding, 'k': k}).fetchall() -print(results) +for row in results: + print('document:', row[0], 'RRF score:', row[1]) From 48133b15d1423e6a41d5286e91e72b2259044185 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 26 Oct 2023 13:41:15 -0700 Subject: [PATCH 110/879] Updated order [skip ci] --- README.md | 142 +++++++++++++++++++++++++++--------------------------- 1 file changed, 71 insertions(+), 71 deletions(-) diff --git a/README.md b/README.md index 351be4a..c910b37 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) +Supports [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions) @@ -16,11 +16,11 @@ pip install pgvector And follow the instructions for your database library: +- [Psycopg 3](#psycopg-3) +- [Psycopg 2](#psycopg-2) - [Django](#django) - [SQLAlchemy](#sqlalchemy) - [SQLModel](#sqlmodel) -- [Psycopg 3](#psycopg-3) -- [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) - [Peewee](#peewee) @@ -34,6 +34,74 @@ Or check out some examples: - [Explicit feedback recommendations](examples/surprise_recs.py) with Surprise - [Recommendations](examples/lightfm_recs.py) with LightFM +## Psycopg 3 + +Enable the extension + +```python +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the vector type with your connection + +```python +from pgvector.psycopg import register_vector + +register_vector(conn) +``` + +For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use + +```python +from pgvector.psycopg import register_vector_async + +await register_vector_async(conn) +``` + +Insert a vector + +```python +embedding = np.array([1, 2, 3]) +conn.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) +``` + +Get the nearest neighbors to a vector + +```python +conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() +``` + +## Psycopg 2 + +Enable the extension + +```python +cur = conn.cursor() +cur.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the vector type with your connection or cursor + +```python +from pgvector.psycopg2 import register_vector + +register_vector(conn) +``` + +Insert a vector + +```python +embedding = np.array([1, 2, 3]) +cur.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) +``` + +Get the nearest neighbors to a vector + +```python +cur.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)) +cur.fetchall() +``` + ## Django Enable the extension @@ -231,74 +299,6 @@ session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit( Also supports `max_inner_product` and `cosine_distance` -## Psycopg 3 - -Enable the extension - -```python -conn.execute('CREATE EXTENSION IF NOT EXISTS vector') -``` - -Register the vector type with your connection - -```python -from pgvector.psycopg import register_vector - -register_vector(conn) -``` - -For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use - -```python -from pgvector.psycopg import register_vector_async - -await register_vector_async(conn) -``` - -Insert a vector - -```python -embedding = np.array([1, 2, 3]) -conn.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) -``` - -Get the nearest neighbors to a vector - -```python -conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() -``` - -## Psycopg 2 - -Enable the extension - -```python -cur = conn.cursor() -cur.execute('CREATE EXTENSION IF NOT EXISTS vector') -``` - -Register the vector type with your connection or cursor - -```python -from pgvector.psycopg2 import register_vector - -register_vector(conn) -``` - -Insert a vector - -```python -embedding = np.array([1, 2, 3]) -cur.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) -``` - -Get the nearest neighbors to a vector - -```python -cur.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)) -cur.fetchall() -``` - ## asyncpg Enable the extension From 03a760b14fa6f0f395230bd177251dd1f7f5c277 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 26 Oct 2023 13:44:22 -0700 Subject: [PATCH 111/879] Updated readme [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c910b37..8b540b1 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ cur.fetchall() ## Django -Enable the extension +Create a migration to enable the extension ```python from pgvector.django import VectorExtension From eb474015492df4d4573f92b180428cac60d56f11 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 26 Oct 2023 13:45:44 -0700 Subject: [PATCH 112/879] Updated readme [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8b540b1..11e8526 100644 --- a/README.md +++ b/README.md @@ -115,7 +115,7 @@ class Migration(migrations.Migration): ] ``` -Add a vector field +Add a vector field to your model ```python from pgvector.django import VectorField From e50b2e295fe9719a6f110f580c1c537f90bf1ef1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 26 Oct 2023 13:49:04 -0700 Subject: [PATCH 113/879] Revert "Updated order [skip ci]" This reverts commit 48133b15d1423e6a41d5286e91e72b2259044185. --- README.md | 142 +++++++++++++++++++++++++++--------------------------- 1 file changed, 71 insertions(+), 71 deletions(-) diff --git a/README.md b/README.md index 11e8526..37030e0 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions) @@ -16,11 +16,11 @@ pip install pgvector And follow the instructions for your database library: -- [Psycopg 3](#psycopg-3) -- [Psycopg 2](#psycopg-2) - [Django](#django) - [SQLAlchemy](#sqlalchemy) - [SQLModel](#sqlmodel) +- [Psycopg 3](#psycopg-3) +- [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) - [Peewee](#peewee) @@ -34,74 +34,6 @@ Or check out some examples: - [Explicit feedback recommendations](examples/surprise_recs.py) with Surprise - [Recommendations](examples/lightfm_recs.py) with LightFM -## Psycopg 3 - -Enable the extension - -```python -conn.execute('CREATE EXTENSION IF NOT EXISTS vector') -``` - -Register the vector type with your connection - -```python -from pgvector.psycopg import register_vector - -register_vector(conn) -``` - -For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use - -```python -from pgvector.psycopg import register_vector_async - -await register_vector_async(conn) -``` - -Insert a vector - -```python -embedding = np.array([1, 2, 3]) -conn.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) -``` - -Get the nearest neighbors to a vector - -```python -conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() -``` - -## Psycopg 2 - -Enable the extension - -```python -cur = conn.cursor() -cur.execute('CREATE EXTENSION IF NOT EXISTS vector') -``` - -Register the vector type with your connection or cursor - -```python -from pgvector.psycopg2 import register_vector - -register_vector(conn) -``` - -Insert a vector - -```python -embedding = np.array([1, 2, 3]) -cur.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) -``` - -Get the nearest neighbors to a vector - -```python -cur.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)) -cur.fetchall() -``` - ## Django Create a migration to enable the extension @@ -299,6 +231,74 @@ session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit( Also supports `max_inner_product` and `cosine_distance` +## Psycopg 3 + +Enable the extension + +```python +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the vector type with your connection + +```python +from pgvector.psycopg import register_vector + +register_vector(conn) +``` + +For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use + +```python +from pgvector.psycopg import register_vector_async + +await register_vector_async(conn) +``` + +Insert a vector + +```python +embedding = np.array([1, 2, 3]) +conn.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) +``` + +Get the nearest neighbors to a vector + +```python +conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() +``` + +## Psycopg 2 + +Enable the extension + +```python +cur = conn.cursor() +cur.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the vector type with your connection or cursor + +```python +from pgvector.psycopg2 import register_vector + +register_vector(conn) +``` + +Insert a vector + +```python +embedding = np.array([1, 2, 3]) +cur.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) +``` + +Get the nearest neighbors to a vector + +```python +cur.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)) +cur.fetchall() +``` + ## asyncpg Enable the extension From 0c88c2152918ed4d7abb83d93694ac966074a15b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Nov 2023 13:18:06 -0700 Subject: [PATCH 114/879] Added Reciprocal Rank Fusion example to readme [skip ci] --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 37030e0..5362a1d 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,8 @@ Or check out some examples: - [Embeddings](examples/openai_embeddings.py) with OpenAI - [Sentence embeddings](examples/sentence_embeddings.py) with SentenceTransformers -- [Hybrid search](examples/hybrid_search.py) with SentenceTransformers +- [Hybrid search](examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) +- [Hybrid search](examples/hybrid_search.py) with SentenceTransformers (cross-encoder) - [Image search](examples/pytorch_image_search.py) with PyTorch - [Implicit feedback recommendations](examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](examples/surprise_recs.py) with Surprise From 416b8456a19b7625ea80cf768083ccbc70ff7fb2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Nov 2023 13:19:58 -0700 Subject: [PATCH 115/879] Use absolute urls for examples for PyPI [skip ci] --- README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 5362a1d..1540410 100644 --- a/README.md +++ b/README.md @@ -26,14 +26,14 @@ And follow the instructions for your database library: Or check out some examples: -- [Embeddings](examples/openai_embeddings.py) with OpenAI -- [Sentence embeddings](examples/sentence_embeddings.py) with SentenceTransformers -- [Hybrid search](examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) -- [Hybrid search](examples/hybrid_search.py) with SentenceTransformers (cross-encoder) -- [Image search](examples/pytorch_image_search.py) with PyTorch -- [Implicit feedback recommendations](examples/implicit_recs.py) with Implicit -- [Explicit feedback recommendations](examples/surprise_recs.py) with Surprise -- [Recommendations](examples/lightfm_recs.py) with LightFM +- [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai_embeddings.py) with OpenAI +- [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_embeddings.py) with SentenceTransformers +- [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) +- [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) +- [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch +- [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit +- [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise +- [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM ## Django From 7997a4b2b98f318fb289902eb64e922b770c25ce Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 18 Nov 2023 11:47:03 -0800 Subject: [PATCH 116/879] Updated OpenAI example [skip ci] --- examples/openai_embeddings.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index 08926f0..ec766d5 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -1,4 +1,4 @@ -import openai +from openai import OpenAI from pgvector.psycopg import register_vector import psycopg @@ -16,8 +16,9 @@ 'The bear is growling' ] -response = openai.Embedding.create(input=input, model='text-embedding-ada-002') -embeddings = [v['embedding'] for v in response['data']] +client = OpenAI() +response = client.embeddings.create(input=input, model='text-embedding-ada-002') +embeddings = [v.embedding for v in response.data] for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) From bab9e4b2319329e61caa203fd08bff5f42e7e7f6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 21 Nov 2023 09:22:51 -0800 Subject: [PATCH 117/879] Improved test --- tests/test_django.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 5b9c386..b68b5f2 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -61,7 +61,11 @@ class Migration(migrations.Migration): ), migrations.AddIndex( model_name='item', - index=pgvector.django.IvfflatIndex(fields=['embedding'], lists=1, name='my_index', opclasses=['vector_l2_ops']), + index=pgvector.django.IvfflatIndex(fields=['embedding'], lists=1, name='ivfflat_idx', opclasses=['vector_l2_ops']), + ), + migrations.AddIndex( + model_name='item', + index=pgvector.django.HnswIndex(fields=['embedding'], m=16, ef_construction=64, name='hnsw_idx', opclasses=['vector_l2_ops']), ) ] From 0705e3abbfe248f3ef288078405328cc3f62ad5d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 21 Nov 2023 10:44:08 -0800 Subject: [PATCH 118/879] Added more tests for aggregates with SQLAlchemy - #44 --- tests/test_sqlalchemy.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 21b46ac..c837840 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -166,6 +166,15 @@ def test_avg(self): avg = session.query(func.avg(Item.embedding)).first()[0] assert np.array_equal(avg, 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 + session.add(Item(embedding=[1, 2, 3])) + session.add(Item(embedding=[4, 5, 6])) + avg = session.scalars(select(func.avg(Item.embedding))).first() + assert np.array_equal(avg, 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] @@ -175,6 +184,15 @@ def test_sum(self): sum = session.query(func.sum(Item.embedding)).first()[0] assert np.array_equal(sum, 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 + 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])) + def test_bad_dimensions(self): item = Item(embedding=[1, 2]) session = Session(engine) From a03a9a6e61f3bf65f82b6ea94d493edaa338aadc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 21 Nov 2023 10:47:17 -0800 Subject: [PATCH 119/879] Updated SQLAlchemy example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 1540410..173c227 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ Average vectors ```python from sqlalchemy.sql import func -session.query(func.avg(Item.embedding)).first()[0] +session.scalars(select(func.avg(Item.embedding))).first() ``` Also supports `sum` From 7ab418e705f3dbd4c58c05f35c02b10b49568fbb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 21 Nov 2023 11:03:51 -0800 Subject: [PATCH 120/879] Added test for asyncio with SQLAlchemy - #44 --- requirements.txt | 2 +- tests/test_sqlalchemy.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index e111e7d..b4ce55e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,4 @@ psycopg[binary] psycopg2-binary pytest pytest-asyncio -SQLAlchemy +SQLAlchemy[asyncio] diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index c837840..9808f0e 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -3,6 +3,7 @@ import pytest from sqlalchemy import create_engine, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base, mapped_column, Session from sqlalchemy.sql import func @@ -213,3 +214,17 @@ def test_bad_dtype(self): session.add(item) with pytest.raises(StatementError, match='dtype must be numeric'): session.commit() + + @pytest.mark.asyncio + async def test_async(self): + engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + session.add(Item(embedding=[1, 2, 3])) + session.add(Item(embedding=[4, 5, 6])) + avg = await session.scalars(select(func.avg(Item.embedding))) + assert avg.first() == '[2.5,3.5,4.5]' + + await engine.dispose() From 4e02d6f7d62eba2e8febf731693accbf82925704 Mon Sep 17 00:00:00 2001 From: KellyRousselHoomano <86828998+KellyRousselHoomano@users.noreply.github.com> Date: Wed, 22 Nov 2023 19:16:26 +0100 Subject: [PATCH 121/879] Registering Vector Type with SQLAlchemy Registries (#45) --- pgvector/sqlalchemy/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 1cf5b66..ec2d85a 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,6 +1,6 @@ from sqlalchemy.types import UserDefinedType, Float from ..utils import from_db, to_db - +from sqlalchemy.dialects import postgresql __all__ = ['Vector'] @@ -35,3 +35,6 @@ def max_inner_product(self, other): def cosine_distance(self, other): return self.op('<=>', return_type=Float)(other) + +# Register Vector type to PostgreSQL's reflection subsystem +postgresql.base.ischema_names['vector'] = Vector \ No newline at end of file From 478d0ab4bb64c3e8b3795f0af812a758736af5d7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:18:40 -0800 Subject: [PATCH 122/879] Fixed style lint [skip ci] --- pgvector/sqlalchemy/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index ec2d85a..775518e 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,6 +1,7 @@ +from sqlalchemy.dialects import postgresql from sqlalchemy.types import UserDefinedType, Float from ..utils import from_db, to_db -from sqlalchemy.dialects import postgresql + __all__ = ['Vector'] @@ -36,5 +37,6 @@ def max_inner_product(self, other): def cosine_distance(self, other): return self.op('<=>', return_type=Float)(other) + # Register Vector type to PostgreSQL's reflection subsystem -postgresql.base.ischema_names['vector'] = Vector \ No newline at end of file +postgresql.base.ischema_names['vector'] = Vector From dd793e8f91006c60a39e1e9954cda13675bd110b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:19:25 -0800 Subject: [PATCH 123/879] Added test for reflection with SQLAlchemy --- tests/test_sqlalchemy.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 9808f0e..32c341c 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,7 +1,7 @@ import numpy as np from pgvector.sqlalchemy import Vector import pytest -from sqlalchemy import create_engine, select, text, MetaData, Table, Column, Index, Integer +from sqlalchemy import create_engine, inspect, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base, mapped_column, Session @@ -215,6 +215,10 @@ def test_bad_dtype(self): with pytest.raises(StatementError, match='dtype must be numeric'): session.commit() + def test_inspect(self): + columns = inspect(engine).get_columns('orm_item') + assert isinstance(columns[1]['type'], Vector) + @pytest.mark.asyncio async def test_async(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') From 323e525911c6fb776602c8ca903b4d3e36f77bdf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:25:17 -0800 Subject: [PATCH 124/879] Updated changelog [skip ci] --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a8dc947..96ac6a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.4 (unreleased) + +- Improved reflection with SQLAlchemy + ## 0.2.3 (2023-09-25) - Fixed null values with Django From ef5de46dcb7b2bc4c3deb5e6d2eca0cadeb9e88c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:28:04 -0800 Subject: [PATCH 125/879] Updated CI --- .github/workflows/build.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d711cd5..768cc3c 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,9 +6,9 @@ jobs: strategy: fail-fast: false matrix: - python: [3.11, 3.8] + python: [3.12, 3.8] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 - uses: actions/setup-python@v4 with: python-version: ${{ matrix.python }} @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From 30d29e7da71de5035f4f0fbdabfec88dcb96fdf8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:41:36 -0800 Subject: [PATCH 126/879] Added tests for SQLModel --- requirements.txt | 3 +- tests/test_sqlmodel.py | 89 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 tests/test_sqlmodel.py diff --git a/requirements.txt b/requirements.txt index b4ce55e..da0ef44 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ psycopg[binary] psycopg2-binary pytest pytest-asyncio -SQLAlchemy[asyncio] +SQLAlchemy[asyncio]>=2 +sqlmodel>=0.0.12 diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py new file mode 100644 index 0000000..c80817e --- /dev/null +++ b/tests/test_sqlmodel.py @@ -0,0 +1,89 @@ +import numpy as np +from pgvector.sqlalchemy import Vector +import pytest +from sqlalchemy import Column +from sqlalchemy.exc import StatementError +from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text +from typing import List, Optional + +engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +with Session(engine) as session: + session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) + + +class Item(SQLModel, table=True): + __tablename__ = 'sqlmodel_item' + + id: Optional[int] = Field(default=None, primary_key=True) + embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Vector(3))) + + +SQLModel.metadata.drop_all(engine) +SQLModel.metadata.create_all(engine) + + +def create_items(): + vectors = [ + [1, 1, 1], + [2, 2, 2], + [1, 1, 2] + ] + session = Session(engine) + for i, v in enumerate(vectors): + session.add(Item(id=i + 1, embedding=v)) + session.commit() + + +class TestSqlmodel: + def setup_method(self, test_method): + with Session(engine) as session: + session.exec(delete(Item)) + session.commit() + + def test_orm(self): + item = Item(embedding=[1.5, 2, 3]) + item2 = Item(embedding=[4, 5, 6]) + item3 = Item() + + session = Session(engine) + session.add(item) + session.add(item2) + session.add(item3) + session.commit() + + stmt = select(Item) + with Session(engine) as session: + items = session.exec(stmt).all() + assert items[0].id == 1 + assert items[1].id == 2 + assert items[2].id == 3 + assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) + assert items[0].embedding.dtype == np.float32 + assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) + assert items[1].embedding.dtype == np.float32 + assert items[2].embedding is None + + def test_l2_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + + def test_max_inner_product(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) + assert [v.id for v in items] == [2, 3, 1] + + def test_cosine_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 2, 3] + + def test_bad_dimensions(self): + item = Item(embedding=[1, 2]) + session = Session(engine) + session.add(item) + with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + session.commit() From 696e81bc12c9f183ce0c4a898d84d08014003df3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:44:42 -0800 Subject: [PATCH 127/879] Updated example for latest SQLModel [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 173c227..4937fde 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc Enable the extension ```python -session.exec('CREATE EXTENSION IF NOT EXISTS vector') +session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) ``` Add a vector column From 38dea36c1a34496b2657b366da27b02b0f3e0c4a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 10:54:59 -0800 Subject: [PATCH 128/879] Improved docs [skip ci] --- README.md | 18 ++++++++++++++++++ tests/test_asyncpg.py | 4 ++-- tests/test_psycopg.py | 4 ++-- tests/test_psycopg2.py | 2 +- 4 files changed, 23 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 4937fde..637e7a1 100644 --- a/README.md +++ b/README.md @@ -256,6 +256,12 @@ from pgvector.psycopg import register_vector_async await register_vector_async(conn) ``` +Create a table + +```python +conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +``` + Insert a vector ```python @@ -286,6 +292,12 @@ from pgvector.psycopg2 import register_vector register_vector(conn) ``` +Create a table + +```python +cur.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +``` + Insert a vector ```python @@ -325,6 +337,12 @@ async def init(conn): pool = await asyncpg.create_pool(..., init=init) ``` +Create a table + +```python +await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +``` + Insert a vector ```python diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 2634310..33b8f96 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -11,7 +11,7 @@ async def test_works(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 item') - await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') + await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') await register_vector(conn) @@ -41,7 +41,7 @@ async def init(conn): async with pool.acquire() as conn: await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') await conn.execute('DROP TABLE IF EXISTS item') - await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') + await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') embedding = np.array([1.5, 2, 3]) await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 73c54d9..3e9ced9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -7,7 +7,7 @@ conn.execute('CREATE EXTENSION IF NOT EXISTS vector') conn.execute('DROP TABLE IF EXISTS item') -conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') +conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') register_vector(conn) @@ -77,7 +77,7 @@ async def test_async(self): await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') await conn.execute('DROP TABLE IF EXISTS item') - await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') + await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') await register_vector_async(conn) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index ac899ed..94ffba3 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -8,7 +8,7 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') cur.execute('DROP TABLE IF EXISTS item') -cur.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))') +cur.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') register_vector(cur) From 96862702c838c63e700f70632c9366a376b930f3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 22 Nov 2023 11:00:58 -0800 Subject: [PATCH 129/879] Updated table names [skip ci] --- README.md | 18 +++++++++--------- examples/pytorch_image_search.py | 8 ++++---- tests/test_asyncpg.py | 18 +++++++++--------- tests/test_psycopg.py | 24 ++++++++++++------------ tests/test_psycopg2.py | 10 +++++----- 5 files changed, 39 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 637e7a1..2b5f8b9 100644 --- a/README.md +++ b/README.md @@ -259,20 +259,20 @@ await register_vector_async(conn) Create a table ```python -conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') ``` Insert a vector ```python embedding = np.array([1, 2, 3]) -conn.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) +conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) ``` Get the nearest neighbors to a vector ```python -conn.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() +conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() ``` ## Psycopg 2 @@ -295,20 +295,20 @@ register_vector(conn) Create a table ```python -cur.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') ``` Insert a vector ```python embedding = np.array([1, 2, 3]) -cur.execute('INSERT INTO item (embedding) VALUES (%s)', (embedding,)) +cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) ``` Get the nearest neighbors to a vector ```python -cur.execute('SELECT * FROM item ORDER BY embedding <-> %s LIMIT 5', (embedding,)) +cur.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)) cur.fetchall() ``` @@ -340,20 +340,20 @@ pool = await asyncpg.create_pool(..., init=init) Create a table ```python -await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') ``` Insert a vector ```python embedding = np.array([1, 2, 3]) -await conn.execute('INSERT INTO item (embedding) VALUES ($1)', embedding) +await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding) ``` Get the nearest neighbors to a vector ```python -await conn.fetch('SELECT * FROM item ORDER BY embedding <-> $1 LIMIT 5', embedding) +await conn.fetch('SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5', embedding) ``` ## Peewee diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index 0cc9af0..e05571c 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -38,14 +38,14 @@ def generate_embeddings(inputs): # generate, save, and index embeddings if seed: - conn.execute('DROP TABLE IF EXISTS image') - conn.execute('CREATE TABLE image (id bigserial PRIMARY KEY, embedding vector(512))') + conn.execute('DROP TABLE IF EXISTS images') + conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, embedding vector(512))') print('Generating embeddings') for data in tqdm(dataloader): embeddings = generate_embeddings(data[0]) - sql = 'INSERT INTO image (embedding) VALUES ' + ','.join(['(%s)' for _ in embeddings]) + sql = 'INSERT INTO images (embedding) VALUES ' + ','.join(['(%s)' for _ in embeddings]) params = [embedding for embedding in embeddings] conn.execute(sql, params) @@ -67,5 +67,5 @@ def show_images(dataset_images): # generate and query embeddings embeddings = generate_embeddings(images) for image, embedding in zip(images, embeddings): - result = conn.execute('SELECT id FROM image ORDER BY embedding <=> %s LIMIT 15', (embedding,)).fetchall() + result = conn.execute('SELECT id FROM images ORDER BY embedding <=> %s LIMIT 15', (embedding,)).fetchall() show_images([image] + [dataset[row[0] - 1][0] for row in result]) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 33b8f96..f06ac3d 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -10,15 +10,15 @@ class TestAsyncpg: async def test_works(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 item') - await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') + await conn.execute('DROP TABLE IF EXISTS items') + await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') await register_vector(conn) embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM item ORDER BY id") + res = await conn.fetch("SELECT * FROM items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert np.array_equal(res[0]['embedding'], embedding) @@ -26,7 +26,7 @@ async def test_works(self): assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM item ORDER BY id LIMIT 1") + text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @@ -40,13 +40,13 @@ async def init(conn): async with pool.acquire() as conn: await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') - await conn.execute('DROP TABLE IF EXISTS item') - await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') + await conn.execute('DROP TABLE IF EXISTS items') + await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM item ORDER BY id") + res = await conn.fetch("SELECT * FROM items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert np.array_equal(res[0]['embedding'], embedding) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 3e9ced9..926bf45 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -6,21 +6,21 @@ conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') -conn.execute('DROP TABLE IF EXISTS item') -conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +conn.execute('DROP TABLE IF EXISTS items') +conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') register_vector(conn) class TestPsycopg: def setup_method(self, test_method): - conn.execute('DELETE FROM item') + conn.execute('DELETE FROM items') def test_works(self): embedding = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,)) + conn.execute('INSERT INTO items (embedding) VALUES (%s), (NULL)', (embedding,)) - res = conn.execute('SELECT * FROM item ORDER BY id').fetchall() + res = conn.execute('SELECT * FROM items ORDER BY id').fetchall() assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None @@ -55,19 +55,19 @@ def test_binary_format_non_contiguous(self): def test_text_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY item (embedding) FROM STDIN") as copy: + with cur.copy("COPY items (embedding) FROM STDIN") as copy: copy.write_row([embedding]) def test_binary_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY item (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy("COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.write_row([embedding]) def test_binary_copy_set_types(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY item (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy("COPY items (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector']) copy.write_row([1, embedding]) @@ -76,16 +76,16 @@ async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') - await conn.execute('DROP TABLE IF EXISTS item') - await conn.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') + await conn.execute('DROP TABLE IF EXISTS items') + await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') await register_vector_async(conn) embedding = np.array([1.5, 2, 3]) - await conn.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,)) + await conn.execute('INSERT INTO items (embedding) VALUES (%s), (NULL)', (embedding,)) async with conn.cursor() as cur: - await cur.execute('SELECT * FROM item ORDER BY id') + await cur.execute('SELECT * FROM items ORDER BY id') res = await cur.fetchall() assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 94ffba3..925d60f 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -7,21 +7,21 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') -cur.execute('DROP TABLE IF EXISTS item') -cur.execute('CREATE TABLE item (id bigserial PRIMARY KEY, embedding vector(3))') +cur.execute('DROP TABLE IF EXISTS items') +cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') register_vector(cur) class TestPsycopg2: def setup_method(self, test_method): - cur.execute('DELETE FROM item') + cur.execute('DELETE FROM items') def test_works(self): embedding = np.array([1.5, 2, 3]) - cur.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,)) + cur.execute('INSERT INTO items (embedding) VALUES (%s), (NULL)', (embedding,)) - cur.execute('SELECT * FROM item ORDER BY id') + cur.execute('SELECT * FROM items ORDER BY id') res = cur.fetchall() assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 From a8c5d8479a066c6d941ee7277561745c422aa9b3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 24 Nov 2023 09:23:59 -0800 Subject: [PATCH 130/879] Updated import --- pgvector/sqlalchemy/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 775518e..e6c128a 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,4 +1,4 @@ -from sqlalchemy.dialects import postgresql +from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float from ..utils import from_db, to_db @@ -38,5 +38,5 @@ def cosine_distance(self, other): return self.op('<=>', return_type=Float)(other) -# Register Vector type to PostgreSQL's reflection subsystem -postgresql.base.ischema_names['vector'] = Vector +# for reflection +ischema_names['vector'] = Vector From d86ea90da3817a015cdac18264de8e79275b3071 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 24 Nov 2023 09:26:23 -0800 Subject: [PATCH 131/879] Version bump to 0.2.4 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 96ac6a4..69bdcdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.4 (unreleased) +## 0.2.4 (2023-11-24) - Improved reflection with SQLAlchemy diff --git a/setup.py b/setup.py index 227b6cf..4b24a00 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.2.3', + version='0.2.4', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From e9232f9f5e098f670e6f723b078f24ae506688fd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 29 Nov 2023 21:43:13 -0800 Subject: [PATCH 132/879] Added Citus example [skip ci] --- README.md | 1 + examples/citus.py | 49 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) create mode 100644 examples/citus.py diff --git a/README.md b/README.md index 2b5f8b9..94ff51e 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ Or check out some examples: - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise - [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM +- [Horizontal scaling](https://github.com/pgvector/pgvector-python/blob/master/examples/citus.py) with Citus ## Django diff --git a/examples/citus.py b/examples/citus.py new file mode 100644 index 0000000..915c25f --- /dev/null +++ b/examples/citus.py @@ -0,0 +1,49 @@ +import numpy as np +from pgvector.psycopg import register_vector +import psycopg + +# generate random data +rows = 100000 +dimensions = 128 +embeddings = np.random.rand(rows, dimensions) +categories = np.random.randint(100, size=rows).tolist() +queries = np.random.rand(10, dimensions) + +# enable extensions +conn = psycopg.connect(dbname='pgvector_citus', autocommit=True) +conn.execute('CREATE EXTENSION IF NOT EXISTS citus') +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + +# GUC variables set on the session do not propagate to Citus workers +# https://github.com/citusdata/citus/issues/462 +# you can either: +# 1. set them on the system, user, or database and reconnect +# 2. set them for a transaction with SET LOCAL +conn.execute("ALTER DATABASE pgvector_citus SET maintenance_work_mem = '512MB'") +conn.execute('ALTER DATABASE pgvector_citus SET hnsw.ef_search = 20') +conn.close() + +# reconnect for updated GUC variables to take effect +conn = psycopg.connect(dbname='pgvector_citus', autocommit=True) +register_vector(conn) + +print('Creating distributed table') +conn.execute('DROP TABLE IF EXISTS items') +conn.execute('CREATE TABLE items (id bigserial, embedding vector(%d), category_id bigint, PRIMARY KEY (id, category_id))' % dimensions) +conn.execute('SET citus.shard_count = 4') +conn.execute("SELECT create_distributed_table('items', 'category_id')") + +print('Loading data in parallel') +with conn.cursor().copy('COPY items (embedding, category_id) FROM STDIN WITH (FORMAT BINARY)') as copy: + copy.set_types(['vector', 'bigint']) + + for i in range(rows): + copy.write_row([embeddings[i], categories[i]]) + +print('Creating index in parallel') +conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') + +print('Running distributed queries') +for query in queries: + items = conn.execute('SELECT id FROM items ORDER BY embedding <-> %s LIMIT 10', (query,)).fetchall() + print([r[0] for r in items]) From 907f67836708a3533b08a1827cba5a9f72561f54 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 18 Jan 2024 16:39:25 -0800 Subject: [PATCH 133/879] Added literal binds support for SQLAlchemy - closes #51 --- CHANGELOG.md | 4 ++++ pgvector/sqlalchemy/__init__.py | 9 ++++++++- tests/test_sqlalchemy.py | 4 ++++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 69bdcdb..cebbf99 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.5 (unreleased) + +- Added literal binds support for SQLAlchemy + ## 0.2.4 (2023-11-24) - Improved reflection with SQLAlchemy diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index e6c128a..8b798e4 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,5 +1,5 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float +from sqlalchemy.types import UserDefinedType, Float, String from ..utils import from_db, to_db __all__ = ['Vector'] @@ -7,6 +7,7 @@ class Vector(UserDefinedType): cache_ok = True + _string = String() def __init__(self, dim=None): super(UserDefinedType, self).__init__() @@ -22,6 +23,12 @@ def process(value): return 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(to_db(value, self.dim)) + return process + def result_processor(self, dialect, coltype): def process(value): return from_db(value) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 32c341c..0c4681e 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -219,6 +219,10 @@ def test_inspect(self): columns = inspect(engine).get_columns('orm_item') assert isinstance(columns[1]['type'], Vector) + def test_literal_binds(self): + sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(compile_kwargs={'literal_binds': True}) + assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) + @pytest.mark.asyncio async def test_async(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') From aa7214eb9a138c547c35d27423dcbaf5c887d8a8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 18 Jan 2024 16:49:07 -0800 Subject: [PATCH 134/879] Fixed CI --- pgvector/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 pgvector/__init__.py diff --git a/pgvector/__init__.py b/pgvector/__init__.py new file mode 100644 index 0000000..e69de29 From c6fb53b1db9e9f275c380009741cb7fa1724e772 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 18 Jan 2024 18:58:13 -0800 Subject: [PATCH 135/879] Improved test --- tests/test_sqlalchemy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0c4681e..0f40e43 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -220,7 +220,7 @@ def test_inspect(self): assert isinstance(columns[1]['type'], Vector) def test_literal_binds(self): - sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(compile_kwargs={'literal_binds': True}) + sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) @pytest.mark.asyncio From 2bb8c3c2dbe4bb6620db246cfab1b7f8f4ddbb43 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jan 2024 18:22:06 -0800 Subject: [PATCH 136/879] Updated OpenAI model [skip ci] --- examples/openai_embeddings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/openai_embeddings.py b/examples/openai_embeddings.py index ec766d5..ebed3d0 100644 --- a/examples/openai_embeddings.py +++ b/examples/openai_embeddings.py @@ -17,7 +17,7 @@ ] client = OpenAI() -response = client.embeddings.create(input=input, model='text-embedding-ada-002') +response = client.embeddings.create(input=input, model='text-embedding-3-small') embeddings = [v.embedding for v in response.data] for content, embedding in zip(input, embeddings): From d5f764f1798fe4b4708c05224f2eb2bd14ae5a6a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 7 Feb 2024 11:34:21 -0800 Subject: [PATCH 137/879] Version bump to 0.2.5 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cebbf99..5bfc2dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.2.5 (unreleased) +## 0.2.5 (2024-02-07) - Added literal binds support for SQLAlchemy diff --git a/setup.py b/setup.py index 4b24a00..dff986d 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.2.4', + version='0.2.5', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From bfa5954f3e731894b51d119dc9d23f3f16645870 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 7 Feb 2024 11:35:36 -0800 Subject: [PATCH 138/879] Fixed lint [skip ci] --- pgvector/sqlalchemy/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 8b798e4..6220082 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -25,6 +25,7 @@ def process(value): def literal_processor(self, dialect): string_literal_processor = self._string._cached_literal_processor(dialect) + def process(value): return string_literal_processor(to_db(value, self.dim)) return process From 8730940ea235e16ac30c9cb8f27ed41a47aa0ca5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 17 Feb 2024 10:22:15 -0800 Subject: [PATCH 139/879] Updated badge [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 94ff51e..54f74be 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) -[![Build Status](https://github.com/pgvector/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector-python/actions) +[![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) ## Installation From ee12393e9e18910bc6f9115053ce11efdc229134 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 17 Feb 2024 10:31:45 -0800 Subject: [PATCH 140/879] Updated pgvector on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 768cc3c..c39241c 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git + git clone --branch v0.6.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From 913fbb0dcb1e3e98f0b748cc0ad09c8ddd7fac05 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 18 Feb 2024 12:07:13 -0800 Subject: [PATCH 141/879] Updated readme [skip ci] --- README.md | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 54f74be..f8b0d9b 100644 --- a/README.md +++ b/README.md @@ -100,23 +100,23 @@ Also supports `Sum` Add an approximate index ```python -from pgvector.django import IvfflatIndex, HnswIndex +from pgvector.django import HnswIndex, IvfflatIndex class Item(models.Model): class Meta: indexes = [ - IvfflatIndex( + HnswIndex( name='my_index', fields=['embedding'], - lists=100, + m=16, + ef_construction=64, opclasses=['vector_l2_ops'] ), # or - HnswIndex( + IvfflatIndex( name='my_index', fields=['embedding'], - m=16, - ef_construction=64, + lists=100, opclasses=['vector_l2_ops'] ) ] @@ -183,14 +183,14 @@ Add an approximate index ```python index = Index('my_index', Item.embedding, - postgresql_using='ivfflat', - postgresql_with={'lists': 100}, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} ) # or index = Index('my_index', Item.embedding, - postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_using='ivfflat', + postgresql_with={'lists': 100}, postgresql_ops={'embedding': 'vector_l2_ops'} ) From 2836dd3d12f07869fbc9655d9751a6910be745c3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 22 Feb 2024 20:09:38 -0800 Subject: [PATCH 142/879] Added more SQLModel examples and tests --- README.md | 22 ++++++++++++++++++++++ tests/test_sqlmodel.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/README.md b/README.md index f8b0d9b..475fb70 100644 --- a/README.md +++ b/README.md @@ -233,6 +233,28 @@ session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit( Also supports `max_inner_product` and `cosine_distance` +Get the distance + +```python +session.exec(select(Item.embedding.l2_distance([3, 1, 2]))) +``` + +Get items within a certain distance + +```python +session.exec(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) +``` + +Average vectors + +```python +from sqlalchemy.sql import func + +session.exec(select(func.avg(Item.embedding))).first() +``` + +Also supports `sum` + ## Psycopg 3 Enable the extension diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index c80817e..9dccedc 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -3,6 +3,7 @@ import pytest from sqlalchemy import Column from sqlalchemy.exc import StatementError +from sqlalchemy.sql import func from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text from typing import List, Optional @@ -81,6 +82,36 @@ def test_cosine_distance(self): items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] + def test_filter(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) + assert [v.id for v in items] == [1] + + def test_select(self): + with Session(engine) as session: + session.add(Item(embedding=[2, 3, 3])) + item = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() + assert item[0] == 3 + + def test_avg(self): + with Session(engine) as session: + avg = session.exec(select(func.avg(Item.embedding))).first() + assert avg 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() + assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + + def test_sum(self): + with Session(engine) as session: + sum = session.exec(select(func.sum(Item.embedding))).first() + assert sum 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])) + def test_bad_dimensions(self): item = Item(embedding=[1, 2]) session = Session(engine) From 69e1e92f5f9ae7f6c27a784f271fe0a8e61221b1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 22 Feb 2024 20:12:27 -0800 Subject: [PATCH 143/879] Improved examples in readme [skip ci] --- README.md | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/README.md b/README.md index 475fb70..0cb6928 100644 --- a/README.md +++ b/README.md @@ -298,6 +298,16 @@ Get the nearest neighbors to a vector conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() ``` +Add an approximate index + +```python +conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') +# or +conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## Psycopg 2 Enable the extension @@ -335,6 +345,16 @@ cur.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding, cur.fetchall() ``` +Add an approximate index + +```python +cur.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') +# or +cur.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## asyncpg Enable the extension @@ -379,6 +399,16 @@ Get the nearest neighbors to a vector await conn.fetch('SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5', embedding) ``` +Add an approximate index + +```python +await conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') +# or +await conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## Peewee Add a vector column From b2dafb6d385bea9c51dea55f7223b3b63bf84bed Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 22 Feb 2024 20:26:21 -0800 Subject: [PATCH 144/879] Improved readme and tests --- README.md | 34 ++++++++++++++++++++++++++++++++-- tests/test_sqlalchemy.py | 9 +++++++++ tests/test_sqlmodel.py | 11 ++++++++++- 3 files changed, 51 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 0cb6928..2122299 100644 --- a/README.md +++ b/README.md @@ -182,13 +182,17 @@ Also supports `sum` Add an approximate index ```python -index = Index('my_index', Item.embedding, +index = Index( + 'my_index', + Item.embedding, postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} ) # or -index = Index('my_index', Item.embedding, +index = Index( + 'my_index', + Item.embedding, postgresql_using='ivfflat', postgresql_with={'lists': 100}, postgresql_ops={'embedding': 'vector_l2_ops'} @@ -255,6 +259,32 @@ session.exec(select(func.avg(Item.embedding))).first() Also supports `sum` +Add an approximate index + +```python +from sqlalchemy import Index + +index = Index( + 'my_index', + Item.embedding, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) +# or +index = Index( + 'my_index', + Item.embedding, + postgresql_using='ivfflat', + postgresql_with={'lists': 100}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) + +index.create(engine) +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## Psycopg 3 Enable the extension diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0f40e43..858e58c 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -25,6 +25,15 @@ class Item(Base): Base.metadata.drop_all(engine) Base.metadata.create_all(engine) +index = Index( + 'orm_index', + Item.embedding, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) +index.create(engine) + def create_items(): vectors = [ diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 9dccedc..4e30d2b 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,7 +1,7 @@ import numpy as np from pgvector.sqlalchemy import Vector import pytest -from sqlalchemy import Column +from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError from sqlalchemy.sql import func from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text @@ -22,6 +22,15 @@ class Item(SQLModel, table=True): SQLModel.metadata.drop_all(engine) SQLModel.metadata.create_all(engine) +index = Index( + 'sqlmodel_index', + Item.embedding, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) +index.create(engine) + def create_items(): vectors = [ From 65263193b10aaa3b10c59ef82fe92c4df4f2d46f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 23 Feb 2024 11:14:40 -0800 Subject: [PATCH 145/879] Improved tests [skip ci] --- tests/test_sqlalchemy.py | 8 ++++---- tests/test_sqlmodel.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 858e58c..441b8fd 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -158,14 +158,14 @@ def test_filter_orm(self): def test_select(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - item = session.query(Item.embedding.l2_distance([1, 1, 1])).first() - assert item[0] == 3 + items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() + assert items[0] == 3 def test_select_orm(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - item = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() - assert item[0] == 3 + items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() + assert items[0] == 3 def test_avg(self): with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 4e30d2b..71fa801 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -100,8 +100,8 @@ def test_filter(self): def test_select(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - item = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() - assert item[0] == 3 + items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() + assert items[0] == 3 def test_avg(self): with Session(engine) as session: From 2866df999e558eeb5397214143acf859a1ae2a4f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 29 Feb 2024 21:15:58 -0800 Subject: [PATCH 146/879] Added bulk loading example [skip ci] --- README.md | 1 + examples/bulk_loading.py | 42 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 examples/bulk_loading.py diff --git a/README.md b/README.md index 2122299..8aa0411 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,7 @@ Or check out some examples: - [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise - [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM - [Horizontal scaling](https://github.com/pgvector/pgvector-python/blob/master/examples/citus.py) with Citus +- [Bulk loading](https://github.com/pgvector/pgvector-python/blob/master/examples/bulk_loading.py) with `COPY` ## Django diff --git a/examples/bulk_loading.py b/examples/bulk_loading.py new file mode 100644 index 0000000..8fa1928 --- /dev/null +++ b/examples/bulk_loading.py @@ -0,0 +1,42 @@ +import numpy as np +from pgvector.psycopg import register_vector +import psycopg + +# generate random data +rows = 1000000 +dimensions = 128 +embeddings = np.random.rand(rows, dimensions) + +# enable extension +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +# create table +conn.execute('DROP TABLE IF EXISTS items') +conn.execute(f'CREATE TABLE items (id bigserial, embedding vector({dimensions}))') + +# load data in batches +cur = conn.cursor() +batches = len(embeddings) // 10000 +print(f'Loading {len(embeddings)} rows over {batches} batches') +for batch in np.array_split(embeddings, batches): + # show progress + print('.', end='', flush=True) + + with cur.copy('COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: + copy.set_types(['vector']) + + for embedding in batch: + copy.write_row([embedding]) + +print('\nSuccess!') + +# create any indexes *after* loading initial data (skipping for this example) +# print('Creating index') +# conn.execute("SET maintenance_work_mem = '8GB'") +# conn.execute("SET max_parallel_maintenance_workers = 7") +# conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)') + +# update planner statistics for good measure +conn.execute('ANALYZE items') From 9f55a60892fa8e9e14d946fc5c26c4f2027f0c07 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 29 Feb 2024 21:27:13 -0800 Subject: [PATCH 147/879] Added comment [skip ci] --- examples/bulk_loading.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/bulk_loading.py b/examples/bulk_loading.py index 8fa1928..1fe546b 100644 --- a/examples/bulk_loading.py +++ b/examples/bulk_loading.py @@ -25,6 +25,8 @@ print('.', end='', flush=True) with cur.copy('COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: + # use set_types for binary copy + # https://www.psycopg.org/psycopg3/docs/basic/copy.html#binary-copy copy.set_types(['vector']) for embedding in batch: From abe0560564e78e41875150e736007a23326d4fdb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 3 Mar 2024 14:26:39 -0800 Subject: [PATCH 148/879] Improved example [skip ci] --- examples/bulk_loading.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/bulk_loading.py b/examples/bulk_loading.py index 1fe546b..2c2ef44 100644 --- a/examples/bulk_loading.py +++ b/examples/bulk_loading.py @@ -16,11 +16,10 @@ conn.execute('DROP TABLE IF EXISTS items') conn.execute(f'CREATE TABLE items (id bigserial, embedding vector({dimensions}))') -# load data in batches +# load data +print(f'Loading {len(embeddings)} rows') cur = conn.cursor() -batches = len(embeddings) // 10000 -print(f'Loading {len(embeddings)} rows over {batches} batches') -for batch in np.array_split(embeddings, batches): +for batch in np.array_split(embeddings, len(embeddings) // 10000): # show progress print('.', end='', flush=True) From ddffb8068014c5bdcdcbaf2dcbb3107d2dc98241 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 3 Mar 2024 21:18:12 -0800 Subject: [PATCH 149/879] Improved example [skip ci] --- examples/bulk_loading.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/examples/bulk_loading.py b/examples/bulk_loading.py index 2c2ef44..957be4c 100644 --- a/examples/bulk_loading.py +++ b/examples/bulk_loading.py @@ -19,17 +19,21 @@ # load data print(f'Loading {len(embeddings)} rows') cur = conn.cursor() -for batch in np.array_split(embeddings, len(embeddings) // 10000): - # show progress - print('.', end='', flush=True) +with cur.copy('COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: + # use set_types for binary copy + # https://www.psycopg.org/psycopg3/docs/basic/copy.html#binary-copy + copy.set_types(['vector']) - with cur.copy('COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: - # use set_types for binary copy - # https://www.psycopg.org/psycopg3/docs/basic/copy.html#binary-copy - copy.set_types(['vector']) + for i, embedding in enumerate(embeddings): + # show progress + if i % 10000 == 0: + print('.', end='', flush=True) - for embedding in batch: - copy.write_row([embedding]) + copy.write_row([embedding]) + + # flush data + while conn.pgconn.flush() == 1: + pass print('\nSuccess!') From da00e9b896f6017af51cd80376cc69b2baf2b098 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Mar 2024 22:43:30 -0800 Subject: [PATCH 150/879] Added more insert tests for SQLAlchemy --- tests/test_sqlalchemy.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 441b8fd..f595e1b 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,7 +1,7 @@ import numpy as np from pgvector.sqlalchemy import Vector import pytest -from sqlalchemy import create_engine, inspect, select, text, MetaData, Table, Column, Index, Integer +from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base, mapped_column, Session @@ -232,6 +232,12 @@ def test_literal_binds(self): sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) + def test_insert(self): + session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) + + def test_insert_bulk(self): + session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) + @pytest.mark.asyncio async def test_async(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') From dd84bc1a7f8afa5910b8660c660fb4add3ad4341 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Mar 2024 22:46:33 -0800 Subject: [PATCH 151/879] Added another insert test for SQLAlchemy --- tests/test_sqlalchemy.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f595e1b..b79f835 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -238,6 +238,9 @@ def test_insert(self): def test_insert_bulk(self): session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) + def test_insert_text(self): + session.execute(text('INSERT INTO orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) + @pytest.mark.asyncio async def test_async(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') From 1b81759673df531879b5d9081025c0b356acff50 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Mar 2024 15:14:37 -0700 Subject: [PATCH 152/879] Updated example [skip ci] --- examples/bulk_loading.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/bulk_loading.py b/examples/bulk_loading.py index 957be4c..1df8fe7 100644 --- a/examples/bulk_loading.py +++ b/examples/bulk_loading.py @@ -38,10 +38,11 @@ print('\nSuccess!') # create any indexes *after* loading initial data (skipping for this example) -# print('Creating index') -# conn.execute("SET maintenance_work_mem = '8GB'") -# conn.execute("SET max_parallel_maintenance_workers = 7") -# conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)') +if False: + print('Creating index') + conn.execute("SET maintenance_work_mem = '8GB'") + conn.execute('SET max_parallel_maintenance_workers = 7') + conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)') # update planner statistics for good measure conn.execute('ANALYZE items') From bc213809d99dd2fa46ef9537a80e2567c5f5c202 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 22 Mar 2024 17:08:08 -0700 Subject: [PATCH 153/879] Improved example [skip ci] --- examples/citus.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/citus.py b/examples/citus.py index 915c25f..d448204 100644 --- a/examples/citus.py +++ b/examples/citus.py @@ -40,6 +40,9 @@ for i in range(rows): copy.write_row([embeddings[i], categories[i]]) + while conn.pgconn.flush() == 1: + pass + print('Creating index in parallel') conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') From 8bd28ca5ef9737ed179c6c0d2caeb69be3d4d124 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 24 Mar 2024 10:36:39 -0700 Subject: [PATCH 154/879] Improved example [skip ci] --- examples/pytorch_image_search.py | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/examples/pytorch_image_search.py b/examples/pytorch_image_search.py index e05571c..b4cf131 100644 --- a/examples/pytorch_image_search.py +++ b/examples/pytorch_image_search.py @@ -8,13 +8,11 @@ seed = True - # establish connection conn = psycopg.connect(dbname='pgvector_example', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') register_vector(conn) - # load images transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), @@ -23,7 +21,6 @@ dataset = torchvision.datasets.CIFAR10(root=tempfile.gettempdir(), train=True, download=True, transform=transform) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1000) - # load pretrained model device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') model = torchvision.models.resnet18(weights='DEFAULT') @@ -36,7 +33,7 @@ def generate_embeddings(inputs): return model(inputs.to(device)).detach().cpu().numpy() -# generate, save, and index embeddings +# generate and store embeddings if seed: conn.execute('DROP TABLE IF EXISTS images') conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, embedding vector(512))') @@ -49,23 +46,24 @@ def generate_embeddings(inputs): params = [embedding for embedding in embeddings] conn.execute(sql, params) - -def show_images(dataset_images): - grid = torchvision.utils.make_grid(dataset_images) - img = (grid / 2 + 0.5).permute(1, 2, 0).numpy() - plt.imshow(img) - plt.draw() - plt.waitforbuttonpress(timeout=3) - - # load 5 random unseen images queryset = torchvision.datasets.CIFAR10(root=tempfile.gettempdir(), train=False, download=True, transform=transform) queryloader = torch.utils.data.DataLoader(queryset, batch_size=5, shuffle=True) images = next(iter(queryloader))[0] - # generate and query embeddings +results = [] embeddings = generate_embeddings(images) for image, embedding in zip(images, embeddings): - result = conn.execute('SELECT id FROM images ORDER BY embedding <=> %s LIMIT 15', (embedding,)).fetchall() - show_images([image] + [dataset[row[0] - 1][0] for row in result]) + result = conn.execute('SELECT id FROM images ORDER BY embedding <=> %s LIMIT 5', (embedding,)).fetchall() + nearest_images = [dataset[row[0] - 1][0] for row in result] + results.append([image] + nearest_images) + +# show images +fig, axs = plt.subplots(len(results), len(results[0])) +for i, result in enumerate(results): + for j, image in enumerate(result): + ax = axs[i, j] + ax.imshow((image / 2 + 0.5).permute(1, 2, 0).numpy()) + ax.set_axis_off() +plt.show(block=True) From 751dc1de90aaa713218fcc7b771d331c0fd4e50b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 24 Mar 2024 11:04:46 -0700 Subject: [PATCH 155/879] Added hash image search example [skip ci] --- README.md | 1 + examples/hash_image_search.py | 43 +++++++++++++++++++++++++++++++++++ examples/requirements.txt | 1 + 3 files changed, 45 insertions(+) create mode 100644 examples/hash_image_search.py diff --git a/README.md b/README.md index 8aa0411..19acac5 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,7 @@ Or check out some examples: - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch +- [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise - [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM diff --git a/examples/hash_image_search.py b/examples/hash_image_search.py new file mode 100644 index 0000000..d30027d --- /dev/null +++ b/examples/hash_image_search.py @@ -0,0 +1,43 @@ +from datasets import load_dataset +import matplotlib.pyplot as plt +import psycopg +from imagehash import phash + + +def hash_image(img): + return ''.join(['1' if v else '0' for v in phash(img).hash.flatten()]) + + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('DROP TABLE IF EXISTS images') +conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, hash bit(64))') + +print('Loading dataset') +dataset = load_dataset('mnist') + +print('Generating hashes') +images = [{'hash': hash_image(row['image'])} for row in dataset['train']] + +print('Storing hashes') +cur = conn.cursor() +with cur.copy('COPY images (hash) FROM STDIN') as copy: + for image in images: + copy.write_row([image['hash']]) + +print('Querying hashes') +results = [] +for i in range(5): + image = dataset['test'][i]['image'] + result = conn.execute('SELECT id FROM images ORDER BY bit_count(hash # %s) LIMIT 5', (hash_image(image),)).fetchall() + nearest_images = [dataset['train'][row[0] - 1]['image'] for row in result] + results.append([image] + nearest_images) + +print('Showing results (first column is query image)') +fig, axs = plt.subplots(len(results), len(results[0])) +for i, result in enumerate(results): + for j, image in enumerate(result): + ax = axs[i, j] + ax.imshow(image) + ax.set_axis_off() +plt.show(block=True) diff --git a/examples/requirements.txt b/examples/requirements.txt index 4be7ae1..b919260 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -1,3 +1,4 @@ +imagehash implicit lightfm openai From 78569e0156d3674588b948d24723ad3db3762b89 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 24 Mar 2024 11:10:06 -0700 Subject: [PATCH 156/879] Updated dependencies for examples [skip ci] --- examples/requirements.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/requirements.txt b/examples/requirements.txt index b919260..6832670 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -1,6 +1,8 @@ +datasets imagehash implicit lightfm +matplotlib openai torch torchvision From ad32dc0f00e905c7a6a245ed20bcdbffcfe4e627 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 26 Mar 2024 12:27:47 -0700 Subject: [PATCH 157/879] Fixed warnings on CI and updated pgvector --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c39241c..0724ffe 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -9,7 +9,7 @@ jobs: python: [3.12, 3.8] steps: - uses: actions/checkout@v4 - - uses: actions/setup-python@v4 + - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python }} - run: pip install -r requirements.txt @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.6.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.6.2 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From a362227f1ff0a669942f13a4a0836b66df00713a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Apr 2024 11:07:34 -0700 Subject: [PATCH 158/879] Moved code to separate files --- pgvector/psycopg/__init__.py | 57 ++---------------------------------- pgvector/psycopg/vector.py | 55 ++++++++++++++++++++++++++++++++++ pgvector/utils/__init__.py | 50 +------------------------------ pgvector/utils/vector.py | 49 +++++++++++++++++++++++++++++++ 4 files changed, 108 insertions(+), 103 deletions(-) create mode 100644 pgvector/psycopg/vector.py create mode 100644 pgvector/utils/vector.py diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index bd398ec..37a5de2 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,48 +1,13 @@ import psycopg -from psycopg.adapt import Loader, Dumper -from psycopg.pq import Format from psycopg.types import TypeInfo +from .vector import * + +# TODO remove in 0.3.0 from ..utils import from_db, from_db_binary, to_db, to_db_binary __all__ = ['register_vector'] -class VectorDumper(Dumper): - - format = Format.TEXT - - def dump(self, obj): - return to_db(obj).encode("utf8") - - -class VectorBinaryDumper(VectorDumper): - - format = Format.BINARY - - def dump(self, obj): - return to_db_binary(obj) - - -class VectorLoader(Loader): - - format = Format.TEXT - - def load(self, data): - if isinstance(data, memoryview): - data = bytes(data) - return from_db(data.decode("utf8")) - - -class VectorBinaryLoader(VectorLoader): - - format = Format.BINARY - - def load(self, data): - if isinstance(data, memoryview): - data = bytes(data) - return from_db_binary(data) - - def register_vector(context): info = TypeInfo.fetch(context, 'vector') register_vector_info(context, info) @@ -51,19 +16,3 @@ def register_vector(context): async def register_vector_async(context): info = await TypeInfo.fetch(context, 'vector') register_vector_info(context, info) - - -def register_vector_info(context, info): - if info is None: - raise psycopg.ProgrammingError('vector type not found in the database') - info.register(context) - - # add oid to anonymous class for set_types - text_dumper = type('', (VectorDumper,), {'oid': info.oid}) - binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) - - adapters = context.adapters - adapters.register_dumper('numpy.ndarray', text_dumper) - adapters.register_dumper('numpy.ndarray', binary_dumper) - adapters.register_loader(info.oid, VectorLoader) - adapters.register_loader(info.oid, VectorBinaryLoader) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py new file mode 100644 index 0000000..2289298 --- /dev/null +++ b/pgvector/psycopg/vector.py @@ -0,0 +1,55 @@ +from psycopg.adapt import Loader, Dumper +from psycopg.pq import Format +from ..utils import from_db, from_db_binary, to_db, to_db_binary + + +class VectorDumper(Dumper): + + format = Format.TEXT + + def dump(self, obj): + return to_db(obj).encode("utf8") + + +class VectorBinaryDumper(VectorDumper): + + format = Format.BINARY + + def dump(self, obj): + return to_db_binary(obj) + + +class VectorLoader(Loader): + + format = Format.TEXT + + def load(self, data): + if isinstance(data, memoryview): + data = bytes(data) + return from_db(data.decode("utf8")) + + +class VectorBinaryLoader(VectorLoader): + + format = Format.BINARY + + def load(self, data): + if isinstance(data, memoryview): + data = bytes(data) + return from_db_binary(data) + + +def register_vector_info(context, info): + if info is None: + raise psycopg.ProgrammingError('vector type not found in the database') + info.register(context) + + # add oid to anonymous class for set_types + text_dumper = type('', (VectorDumper,), {'oid': info.oid}) + binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) + + adapters = context.adapters + adapters.register_dumper('numpy.ndarray', text_dumper) + adapters.register_dumper('numpy.ndarray', binary_dumper) + adapters.register_loader(info.oid, VectorLoader) + adapters.register_loader(info.oid, VectorBinaryLoader) diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index 5640b60..e58cbab 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1,49 +1 @@ -import numpy as np -from struct import pack, unpack - - -def from_db(value): - # could be ndarray if already cast by lower-level driver - if value is None or isinstance(value, np.ndarray): - return value - - return np.array(value[1:-1].split(','), dtype=np.float32) - - -def from_db_binary(value): - if value is None: - return value - - (dim, unused) = unpack('>HH', value[:4]) - return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32) - - -def to_db(value, dim=None): - if value is None: - return value - - if isinstance(value, np.ndarray): - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating): - raise ValueError('dtype must be numeric') - - value = value.tolist() - - 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): - if value is None: - return value - - value = np.asarray(value, dtype='>f') - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - return pack('>HH', value.shape[0], 0) + value.tobytes() +from .vector import * diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py new file mode 100644 index 0000000..5640b60 --- /dev/null +++ b/pgvector/utils/vector.py @@ -0,0 +1,49 @@ +import numpy as np +from struct import pack, unpack + + +def from_db(value): + # could be ndarray if already cast by lower-level driver + if value is None or isinstance(value, np.ndarray): + return value + + return np.array(value[1:-1].split(','), dtype=np.float32) + + +def from_db_binary(value): + if value is None: + return value + + (dim, unused) = unpack('>HH', value[:4]) + return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32) + + +def to_db(value, dim=None): + if value is None: + return value + + if isinstance(value, np.ndarray): + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating): + raise ValueError('dtype must be numeric') + + value = value.tolist() + + 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): + if value is None: + return value + + value = np.asarray(value, dtype='>f') + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + return pack('>HH', value.shape[0], 0) + value.tobytes() From 3dd2067fbcc2255c065ed8cc1b5dc7127739c396 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Apr 2024 12:04:06 -0700 Subject: [PATCH 159/879] Improved code --- pgvector/utils/vector.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 5640b60..d831dac 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -1,5 +1,5 @@ import numpy as np -from struct import pack, unpack +from struct import pack, unpack_from def from_db(value): @@ -14,7 +14,7 @@ def from_db_binary(value): if value is None: return value - (dim, unused) = unpack('>HH', value[:4]) + dim, unused = unpack_from('>HH', value) return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32) From 70fc5dfdea1c72867996f30077e95767a630b8d9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Apr 2024 12:14:02 -0700 Subject: [PATCH 160/879] Use consistent style [skip ci] --- pgvector/psycopg/vector.py | 4 ++-- pgvector/sqlalchemy/__init__.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 2289298..aa15b92 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -8,7 +8,7 @@ class VectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return to_db(obj).encode("utf8") + return to_db(obj).encode('utf8') class VectorBinaryDumper(VectorDumper): @@ -26,7 +26,7 @@ class VectorLoader(Loader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return from_db(data.decode("utf8")) + return from_db(data.decode('utf8')) class VectorBinaryLoader(VectorLoader): diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 6220082..d788770 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -15,8 +15,8 @@ def __init__(self, dim=None): def get_col_spec(self, **kw): if self.dim is None: - return "VECTOR" - return "VECTOR(%d)" % self.dim + return 'VECTOR' + return 'VECTOR(%d)' % self.dim def bind_processor(self, dialect): def process(value): From 5ba96ffa9f4ef751e8385aa7684cd2764f5434e7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Apr 2024 12:26:59 -0700 Subject: [PATCH 161/879] Added test for bit [skip ci] --- tests/test_psycopg.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 926bf45..a974808 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -71,6 +71,10 @@ def test_binary_copy_set_types(self): copy.set_types(['int8', 'vector']) copy.write_row([1, embedding]) + def test_bit(self): + res = conn.execute('SELECT %s::bit(3)', ('101',)).fetchone()[0] + assert res == '101' + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) From 3c51f396b3dd8390beae372e4a982dda2c01ae9a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 May 2024 16:13:29 -0700 Subject: [PATCH 162/879] Added support for halfvec and sparsevec types to Psycopg 3 --- .github/workflows/build.yml | 2 +- CHANGELOG.md | 4 +++ pgvector/psycopg/__init__.py | 22 +++++++++++++- pgvector/psycopg/halfvec.py | 57 +++++++++++++++++++++++++++++++++++ pgvector/psycopg/sparsevec.py | 57 +++++++++++++++++++++++++++++++++++ pgvector/utils/__init__.py | 2 ++ pgvector/utils/halfvec.py | 29 ++++++++++++++++++ pgvector/utils/sparsevec.py | 46 ++++++++++++++++++++++++++++ tests/test_psycopg.py | 43 ++++++++++++++++++++++++-- 9 files changed, 258 insertions(+), 4 deletions(-) create mode 100644 pgvector/psycopg/halfvec.py create mode 100644 pgvector/psycopg/sparsevec.py create mode 100644 pgvector/utils/halfvec.py create mode 100644 pgvector/utils/sparsevec.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0724ffe..f8bcaa3 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.6.2 https://github.com/pgvector/pgvector.git + git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bfc2dc..cf51599 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.2.6 (unreleased) + +- Added support for `halfvec` and `sparsevec` types to Psycopg 3 + ## 0.2.5 (2024-02-07) - Added literal binds support for SQLAlchemy diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 37a5de2..3228afc 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,8 +1,12 @@ import psycopg from psycopg.types import TypeInfo -from .vector import * +from .halfvec import register_halfvec_info +from .sparsevec import register_sparsevec_info +from .vector import register_vector_info +from ..utils import HalfVec, SparseVec # TODO remove in 0.3.0 +from .vector import * from ..utils import from_db, from_db_binary, to_db, to_db_binary __all__ = ['register_vector'] @@ -12,7 +16,23 @@ def register_vector(context): info = TypeInfo.fetch(context, 'vector') register_vector_info(context, info) + info = TypeInfo.fetch(context, 'halfvec') + if info is not None: + register_halfvec_info(context, info) + + info = TypeInfo.fetch(context, 'sparsevec') + if info is not None: + register_sparsevec_info(context, info) + async def register_vector_async(context): info = await TypeInfo.fetch(context, 'vector') register_vector_info(context, info) + + info = await TypeInfo.fetch(context, 'halfvec') + if info is not None: + register_halfvec_info(context, info) + + info = await TypeInfo.fetch(context, 'sparsevec') + if info is not None: + register_sparsevec_info(context, info) diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py new file mode 100644 index 0000000..3a159d6 --- /dev/null +++ b/pgvector/psycopg/halfvec.py @@ -0,0 +1,57 @@ +from psycopg.adapt import Loader, Dumper +from psycopg.pq import Format +from ..utils import HalfVec + + +class HalfVecDumper(Dumper): + + format = Format.TEXT + + def dump(self, obj): + return obj.to_db().encode('utf8') + + +class HalfVecBinaryDumper(HalfVecDumper): + + format = Format.BINARY + + def dump(self, obj): + return obj.to_db_binary() + + +class HalfVecLoader(Loader): + + format = Format.TEXT + + def load(self, data): + if data is None: + return None + if isinstance(data, memoryview): + data = bytes(data) + return HalfVec.from_db(data.decode('utf8')) + + +class HalfVecBinaryLoader(HalfVecLoader): + + format = Format.BINARY + + def load(self, data): + if data is None: + return None + if isinstance(data, memoryview): + data = bytes(data) + return HalfVec.from_db_binary(data) + + +def register_halfvec_info(context, info): + info.register(context) + + # add oid to anonymous class for set_types + text_dumper = type('', (HalfVecDumper,), {'oid': info.oid}) + binary_dumper = type('', (HalfVecBinaryDumper,), {'oid': info.oid}) + + adapters = context.adapters + adapters.register_dumper(HalfVec, text_dumper) + adapters.register_dumper(HalfVec, binary_dumper) + adapters.register_loader(info.oid, HalfVecLoader) + adapters.register_loader(info.oid, HalfVecBinaryLoader) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py new file mode 100644 index 0000000..18f6e1a --- /dev/null +++ b/pgvector/psycopg/sparsevec.py @@ -0,0 +1,57 @@ +from psycopg.adapt import Loader, Dumper +from psycopg.pq import Format +from ..utils import SparseVec + + +class SparseVecDumper(Dumper): + + format = Format.TEXT + + def dump(self, obj): + return obj.to_db().encode('utf8') + + +class SparseVecBinaryDumper(SparseVecDumper): + + format = Format.BINARY + + def dump(self, obj): + return obj.to_db_binary() + + +class SparseVecLoader(Loader): + + format = Format.TEXT + + def load(self, data): + if data is None: + return None + if isinstance(data, memoryview): + data = bytes(data) + return SparseVec.from_db(data.decode('utf8')) + + +class SparseVecBinaryLoader(SparseVecLoader): + + format = Format.BINARY + + def load(self, data): + if data is None: + return None + if isinstance(data, memoryview): + data = bytes(data) + return SparseVec.from_db_binary(data) + + +def register_sparsevec_info(context, info): + info.register(context) + + # add oid to anonymous class for set_types + text_dumper = type('', (SparseVecDumper,), {'oid': info.oid}) + binary_dumper = type('', (SparseVecBinaryDumper,), {'oid': info.oid}) + + adapters = context.adapters + adapters.register_dumper(SparseVec, text_dumper) + adapters.register_dumper(SparseVec, binary_dumper) + adapters.register_loader(info.oid, SparseVecLoader) + adapters.register_loader(info.oid, SparseVecBinaryLoader) diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index e58cbab..897862b 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1 +1,3 @@ +from .halfvec import * +from .sparsevec import * from .vector import * diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py new file mode 100644 index 0000000..262d3ad --- /dev/null +++ b/pgvector/utils/halfvec.py @@ -0,0 +1,29 @@ +from struct import pack, unpack_from + + +class HalfVec: + def __init__(self, value): + # TODO support np.array + if not isinstance(value, (list, tuple)): + raise ValueError('expected list or tuple') + + self.value = value + + def to_list(self): + return list(self.value) + + def to_db(self): + return '[' + ','.join([str(float(v)) for v in self.value]) + ']' + + def to_db_binary(self): + return pack(f'>HH{len(self.value)}e', len(self.value), 0, *self.value) + + def from_db(value): + return HalfVec([float(v) for v in value[1:-1].split(',')]) + + def from_db_binary(value): + dim, unused = unpack_from('>HH', value) + return HalfVec(unpack_from(f'>{dim}e', value, 4)) + + def __repr__(self): + return f'HalfVec({self.value})' diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py new file mode 100644 index 0000000..72ec0e4 --- /dev/null +++ b/pgvector/utils/sparsevec.py @@ -0,0 +1,46 @@ +from struct import pack, unpack_from + + +class SparseVec: + def __init__(self, dim, indices, values): + self.dim = dim + self.indices = indices + self.values = values + + def from_dense(value): + dim = len(value) + indices = [i for i, v in enumerate(value) if v != 0] + values = [value[i] for i in indices] + return SparseVec(dim, indices, values) + + def to_dense(self): + vec = [0] * self.dim + for i, v in zip(self.indices, self.values): + vec[i] = v + return vec + + def to_db(self): + return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(self.indices, self.values)]) + '}/' + str(self.dim) + + def to_db_binary(self): + nnz = len(self.indices) + return pack(f'>iii{nnz}i{nnz}f', self.dim, nnz, 0, *self.indices, *self.values) + + def from_db(value): + elements, dim = value.split('/') + indices = [] + values = [] + for e in elements[1:-1].split(','): + i, v = e.split(':') + indices.append(int(i) - 1) + values.append(float(v)) + return SparseVec(int(dim), indices, values) + + def from_db_binary(value): + dim, nnz, unused = unpack_from('>iii', value) + indices = unpack_from(f'>{nnz}i', value, 12) + values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) + return SparseVec(int(dim), indices, values) + + def __repr__(self): + return f'SparseVec({self.dim}, {self.indices}, {self.values})' diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index a974808..9a7f611 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg import register_vector, register_vector_async +from pgvector.psycopg import register_vector, register_vector_async, HalfVec, SparseVec import psycopg import pytest @@ -27,7 +27,7 @@ def test_works(self): def test_binary_format(self): embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] + res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] assert np.array_equal(res, embedding) def test_text_format(self): @@ -71,6 +71,45 @@ def test_binary_copy_set_types(self): copy.set_types(['int8', 'vector']) copy.write_row([1, embedding]) + def test_halfvec(self): + conn.execute('DROP TABLE IF EXISTS half_items') + conn.execute('CREATE TABLE half_items (id bigserial PRIMARY KEY, embedding halfvec(3))') + + embedding = HalfVec([1.5, 2, 3]) + conn.execute('INSERT INTO half_items (embedding) VALUES (%s)', (embedding,)) + + res = conn.execute('SELECT * FROM half_items ORDER BY id').fetchall() + + def test_halfvec_binary_format(self): + embedding = HalfVec([1.5, 2, 3]) + res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] + assert res.to_list() == [1.5, 2, 3] + + def test_halfvec_text_format(self): + embedding = HalfVec([1.5, 2, 3]) + res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] + assert res.to_list() == [1.5, 2, 3] + + def test_sparsevec(self): + conn.execute('DROP TABLE IF EXISTS sparse_items') + conn.execute('CREATE TABLE sparse_items (id bigserial PRIMARY KEY, embedding sparsevec(6))') + + embedding = SparseVec.from_dense([0, 1.5, 0, 2, 0, 3]) + conn.execute('INSERT INTO sparse_items (embedding) VALUES (%s)', (embedding,)) + + res = conn.execute('SELECT * FROM sparse_items ORDER BY id').fetchall() + assert res[0][1].to_dense() == [0, 1.5, 0, 2, 0, 3] + + def test_sparsevec_binary_format(self): + embedding = SparseVec.from_dense([1.5, 2, 3]) + res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] + assert res.to_dense() == [1.5, 2, 3] + + def test_sparsevec_text_format(self): + embedding = SparseVec.from_dense([1.5, 2, 3]) + res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] + assert res.to_dense() == [1.5, 2, 3] + def test_bit(self): res = conn.execute('SELECT %s::bit(3)', ('101',)).fetchone()[0] assert res == '101' From 62c69fca1b63faa7500350cc8044bdce78cca96c Mon Sep 17 00:00:00 2001 From: Domenico Date: Thu, 16 May 2024 17:54:41 +0200 Subject: [PATCH 163/879] L1 distance support for sqlalchemy and sqlmodel (#69) --- pgvector/sqlalchemy/__init__.py | 3 +++ tests/test_sqlalchemy.py | 12 ++++++++++++ tests/test_sqlmodel.py | 6 ++++++ 3 files changed, 21 insertions(+) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index d788770..9bd6ccc 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -39,6 +39,9 @@ class comparator_factory(UserDefinedType.Comparator): def l2_distance(self, other): return self.op('<->', return_type=Float)(other) + def l1_distance(self, other): + return self.op('<+>', return_type=Float)(other) + def max_inner_product(self, other): return self.op('<#>', return_type=Float)(other) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index b79f835..e5ce4ce 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -119,6 +119,18 @@ def test_l2_distance_orm(self): items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 2])).all() + assert [v.id for v in items] == [3, 1, 2] + + def test_l1_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 2]))) + assert [v.id for v in items] == [3, 1, 2] + def test_max_inner_product(self): create_items() with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 71fa801..fbd4410 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -79,6 +79,12 @@ def test_l2_distance(self): items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_max_inner_product(self): create_items() with Session(engine) as session: From 5a530ee5478adb1b63e1d9edbaa817f87a22a067 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 11:57:22 -0400 Subject: [PATCH 164/879] Moved l1_distance function and tests to be consistent with other libraries --- CHANGELOG.md | 1 + pgvector/sqlalchemy/__init__.py | 6 +++--- tests/test_sqlalchemy.py | 24 ++++++++++++------------ tests/test_sqlmodel.py | 12 ++++++------ 4 files changed, 22 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cf51599..2fe2ce1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.2.6 (unreleased) - Added support for `halfvec` and `sparsevec` types to Psycopg 3 +- Added `l1_distance` for SQLAlchemy and SQLModel ## 0.2.5 (2024-02-07) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 9bd6ccc..e035e79 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -39,15 +39,15 @@ class comparator_factory(UserDefinedType.Comparator): def l2_distance(self, other): return self.op('<->', return_type=Float)(other) - def l1_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['vector'] = Vector diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index e5ce4ce..5db608f 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -119,18 +119,6 @@ def test_l2_distance_orm(self): items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_l1_distance(self): - create_items() - with Session(engine) as session: - items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 2])).all() - assert [v.id for v in items] == [3, 1, 2] - - def test_l1_distance_orm(self): - create_items() - with Session(engine) as session: - items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 2]))) - assert [v.id for v in items] == [3, 1, 2] - def test_max_inner_product(self): create_items() with Session(engine) as session: @@ -155,6 +143,18 @@ def test_cosine_distance_orm(self): items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] + def test_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_l1_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_filter(self): create_items() with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index fbd4410..8d349ea 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -79,12 +79,6 @@ def test_l2_distance(self): items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_l1_distance(self): - create_items() - with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) - assert [v.id for v in items] == [1, 3, 2] - def test_max_inner_product(self): create_items() with Session(engine) as session: @@ -97,6 +91,12 @@ def test_cosine_distance(self): items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] + def test_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_filter(self): create_items() with Session(engine) as session: From 2a0eeb47fe9b818495f091001f2b412d66f31d23 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 12:00:39 -0400 Subject: [PATCH 165/879] Added L1Distance for Django --- CHANGELOG.md | 1 + pgvector/django/__init__.py | 7 ++++++- tests/test_django.py | 9 ++++++++- 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fe2ce1..b9087a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.2.6 (unreleased) - Added support for `halfvec` and `sparsevec` types to Psycopg 3 +- Added `L1Distance` for Django - Added `l1_distance` for SQLAlchemy and SQLModel ## 0.2.5 (2024-02-07) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index ab250e6..e4c8e21 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -5,7 +5,7 @@ from .forms import VectorFormField from ..utils import from_db, to_db -__all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance'] +__all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] class VectorExtension(CreateExtension): @@ -128,3 +128,8 @@ class MaxInnerProduct(DistanceBase): class CosineDistance(DistanceBase): function = '' arg_joiner = ' <=> ' + + +class L1Distance(DistanceBase): + function = '' + arg_joiner = ' <+> ' diff --git a/tests/test_django.py b/tests/test_django.py index b68b5f2..b6274cd 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance +from pgvector.django import VectorExtension, VectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance from unittest import mock settings.configure( @@ -131,6 +131,13 @@ def test_cosine_distance(self): assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + def test_l1_distance(self): + create_items() + distance = L1Distance('embedding', [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, 3] + def test_filter(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) From f489f56414da1b505175fbfab68ca9e94f56860b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 12:02:48 -0400 Subject: [PATCH 166/879] Added l1_distance for Peewee --- CHANGELOG.md | 2 +- pgvector/peewee/__init__.py | 3 +++ tests/test_peewee.py | 7 +++++++ 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b9087a1..b1108c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ - Added support for `halfvec` and `sparsevec` types to Psycopg 3 - Added `L1Distance` for Django -- Added `l1_distance` for SQLAlchemy and SQLModel +- Added `l1_distance` for SQLAlchemy, SQLModel, and Peewee ## 0.2.5 (2024-02-07) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 148f8b3..756959d 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -29,3 +29,6 @@ def max_inner_product(self, vector): def cosine_distance(self, vector): return self._distance('<=>', vector) + + def l1_distance(self, vector): + return self._distance('<+>', vector) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 38acb7c..becbb58 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -64,6 +64,13 @@ def test_cosine_distance(self): assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + def test_l1_distance(self): + create_items() + distance = Item.embedding.l1_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, 3] + def test_where(self): create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) From 7add3381143806bbdb7fa78d8ec886112a89ceab Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 12:21:01 -0400 Subject: [PATCH 167/879] Updated license year [skip ci] --- LICENSE.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/LICENSE.txt b/LICENSE.txt index b3134ac..d205f4e 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2023 Andrew Kane +Copyright (c) 2021-2024 Andrew Kane Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal From 03e8e3f7ec1de13d50590e7f0c5c5c88c1696658 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 12:22:05 -0400 Subject: [PATCH 168/879] Updated changelog [skip ci] --- CHANGELOG.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b1108c2..1d2ae09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,8 +37,8 @@ ## 0.1.7 (2023-05-11) -- Added `register_vector_async` for psycopg3 -- Fixed `set_types` for psycopg3 +- Added `register_vector_async` for Psycopg 3 +- Fixed `set_types` for Psycopg 3 ## 0.1.6 (2022-05-22) @@ -51,12 +51,12 @@ ## 0.1.4 (2021-10-12) -- Updated psycopg3 integration for 3.0 release (no longer experimental) +- Updated Psycopg 3 integration for 3.0 release (no longer experimental) ## 0.1.3 (2021-06-22) - Added support for asyncpg -- Added experimental support for psycopg3 +- Added experimental support for Psycopg 3 ## 0.1.2 (2021-06-13) From 03f7ea82df9f2c195e08754b6e7b3df8ad1460a3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 16:11:51 -0400 Subject: [PATCH 169/879] Added support for halfvec and sparsevec types to asyncpg --- CHANGELOG.md | 1 + pgvector/asyncpg/__init__.py | 17 +++++++++++- tests/test_asyncpg.py | 52 ++++++++++++++++++++++++++++++++++-- 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d2ae09..591a3d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.2.6 (unreleased) - Added support for `halfvec` and `sparsevec` types to Psycopg 3 +- Added support for `halfvec` and `sparsevec` types to asyncpg - Added `L1Distance` for Django - Added `l1_distance` for SQLAlchemy, SQLModel, and Peewee diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 5797327..9603ac5 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,4 +1,4 @@ -from ..utils import from_db, from_db_binary, to_db, to_db_binary +from ..utils import from_db, from_db_binary, to_db, to_db_binary, HalfVec, SparseVec __all__ = ['register_vector'] @@ -10,3 +10,18 @@ async def register_vector(conn): decoder=from_db_binary, format='binary' ) + + await conn.set_type_codec( + 'halfvec', + encoder=lambda v: HalfVec(v).to_db_binary(), + decoder=HalfVec.from_db_binary, + format='binary' + ) + + await conn.set_type_codec( + 'sparsevec', + # TODO fix + encoder=lambda v: isinstance(v, SparseVec) and v.to_db_binary(), + decoder=SparseVec.from_db_binary, + format='binary' + ) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index f06ac3d..28cbacd 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,13 +1,13 @@ import asyncio import asyncpg import numpy as np -from pgvector.asyncpg import register_vector +from pgvector.asyncpg import register_vector, SparseVec import pytest class TestAsyncpg: @pytest.mark.asyncio - async def test_works(self): + async def test_vector(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 items') @@ -31,6 +31,54 @@ async def test_works(self): await conn.close() + @pytest.mark.asyncio + async def test_halfvec(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 items') + await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding halfvec(3))') + + await register_vector(conn) + + embedding = [1.5, 2, 3] + await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + + res = await conn.fetch("SELECT * FROM items ORDER BY id") + assert res[0]['id'] == 1 + assert res[1]['id'] == 2 + assert res[0]['embedding'].to_list() == [1.5, 2, 3] + assert res[1]['embedding'] is None + + # ensures binary format is correct + text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + assert text_res[0]['embedding'] == '[1.5,2,3]' + + await conn.close() + + @pytest.mark.asyncio + async def test_sparsevec(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 items') + await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding sparsevec(3))') + + await register_vector(conn) + + embedding = SparseVec.from_dense([1.5, 2, 3]) + await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + + res = await conn.fetch("SELECT * FROM items ORDER BY id") + assert res[0]['id'] == 1 + assert res[1]['id'] == 2 + assert res[0]['embedding'].to_dense() == [1.5, 2, 3] + assert res[1]['embedding'] is None + + # ensures binary format is correct + text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + assert text_res[0]['embedding'] == '{1:1.5,2:2,3:3}/3' + + await conn.close() + @pytest.mark.asyncio async def test_pool(self): async def init(conn): From 8ef3e0dd2a7af8a76c7c35b2e06d2739335569dc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 16:16:06 -0400 Subject: [PATCH 170/879] Moved VectorField to separate file [skip ci] --- pgvector/peewee/__init__.py | 34 +--------------------------------- pgvector/peewee/vector.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 33 deletions(-) create mode 100644 pgvector/peewee/vector.py diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 756959d..ec4b615 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -1,34 +1,2 @@ -from peewee import Expression, Field, Value from ..utils import from_db, to_db - - -class VectorField(Field): - field_type = 'vector' - - def __init__(self, dimensions=None, *args, **kwargs): - self.dimensions = dimensions - super(VectorField, self).__init__(*args, **kwargs) - - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None - - def db_value(self, value): - return to_db(value) - - def python_value(self, value): - return from_db(value) - - def _distance(self, op, vector): - return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - - def l2_distance(self, vector): - return self._distance('<->', vector) - - def max_inner_product(self, vector): - return self._distance('<#>', vector) - - def cosine_distance(self, vector): - return self._distance('<=>', vector) - - def l1_distance(self, vector): - return self._distance('<+>', vector) +from .vector import VectorField diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py new file mode 100644 index 0000000..756959d --- /dev/null +++ b/pgvector/peewee/vector.py @@ -0,0 +1,34 @@ +from peewee import Expression, Field, Value +from ..utils import from_db, to_db + + +class VectorField(Field): + field_type = 'vector' + + def __init__(self, dimensions=None, *args, **kwargs): + self.dimensions = dimensions + super(VectorField, self).__init__(*args, **kwargs) + + def get_modifiers(self): + return self.dimensions and [self.dimensions] or None + + def db_value(self, value): + return to_db(value) + + def python_value(self, value): + return from_db(value) + + def _distance(self, op, vector): + return Expression(lhs=self, op=op, rhs=self.to_value(vector)) + + def l2_distance(self, vector): + return self._distance('<->', vector) + + def max_inner_product(self, vector): + return self._distance('<#>', vector) + + def cosine_distance(self, vector): + return self._distance('<=>', vector) + + def l1_distance(self, vector): + return self._distance('<+>', vector) From 3ee6b94571189e5721170c7d02fba1cbece0ab33 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 16:21:40 -0400 Subject: [PATCH 171/879] Added support for halfvec type to Peewee [skip ci] --- CHANGELOG.md | 1 + pgvector/peewee/__init__.py | 1 + pgvector/peewee/halfvec.py | 35 +++++++++++++++++++++++++++++++++++ tests/test_peewee.py | 14 +++++++++++--- 4 files changed, 48 insertions(+), 3 deletions(-) create mode 100644 pgvector/peewee/halfvec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 591a3d1..c9826d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ - Added support for `halfvec` and `sparsevec` types to Psycopg 3 - Added support for `halfvec` and `sparsevec` types to asyncpg +- Added support for `halfvec` type to Peewee - Added `L1Distance` for Django - Added `l1_distance` for SQLAlchemy, SQLModel, and Peewee diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index ec4b615..c79cbf9 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -1,2 +1,3 @@ from ..utils import from_db, to_db +from .halfvec import HalfvecField from .vector import VectorField diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py new file mode 100644 index 0000000..8d8d176 --- /dev/null +++ b/pgvector/peewee/halfvec.py @@ -0,0 +1,35 @@ +from peewee import Expression, Field, Value +# TODO use halfvec functions +from ..utils import from_db, to_db + + +class HalfvecField(Field): + field_type = 'halfvec' + + def __init__(self, dimensions=None, *args, **kwargs): + self.dimensions = dimensions + super(HalfvecField, self).__init__(*args, **kwargs) + + def get_modifiers(self): + return self.dimensions and [self.dimensions] or None + + def db_value(self, value): + return to_db(value) + + def python_value(self, value): + return from_db(value) + + def _distance(self, op, vector): + return Expression(lhs=self, op=op, rhs=self.to_value(vector)) + + def l2_distance(self, vector): + return self._distance('<->', vector) + + def max_inner_product(self, vector): + return self._distance('<#>', vector) + + def cosine_distance(self, vector): + return self._distance('<=>', vector) + + def l1_distance(self, vector): + return self._distance('<+>', vector) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index becbb58..970ee41 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField +from pgvector.peewee import VectorField, HalfvecField db = PostgresqlDatabase('pgvector_python_test') @@ -12,7 +12,8 @@ class Meta: class Item(BaseModel): - embedding = VectorField(dimensions=3) + embedding = VectorField(dimensions=3, null=True) + half_embedding = HalfvecField(dimensions=3, null=True) Item.add_index('embedding vector_l2_ops', using='hnsw') @@ -30,7 +31,7 @@ def create_items(): [1, 1, 2] ] for i, v in enumerate(vectors): - Item.create(id=i + 1, embedding=v) + Item.create(id=i + 1, embedding=v, half_embedding=v) class TestPeewee: @@ -71,6 +72,13 @@ def test_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_halfvec_l2_distance(self): + create_items() + distance = Item.half_embedding.l2_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_where(self): create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) From bd9d7d68d6b6a81fcbc10b2b7804fa775057eee0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 16:28:42 -0400 Subject: [PATCH 172/879] Added support for sparsevec type to Peewee [skip ci] --- CHANGELOG.md | 2 +- pgvector/peewee/__init__.py | 3 ++- pgvector/peewee/halfvec.py | 9 +++++---- pgvector/peewee/sparsevec.py | 37 ++++++++++++++++++++++++++++++++++++ tests/test_peewee.py | 12 ++++++++++-- 5 files changed, 55 insertions(+), 8 deletions(-) create mode 100644 pgvector/peewee/sparsevec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c9826d3..c48f5ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ - Added support for `halfvec` and `sparsevec` types to Psycopg 3 - Added support for `halfvec` and `sparsevec` types to asyncpg -- Added support for `halfvec` type to Peewee +- Added support for `halfvec` and `sparsevec` types to Peewee - Added `L1Distance` for Django - Added `l1_distance` for SQLAlchemy, SQLModel, and Peewee diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index c79cbf9..106170f 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -1,3 +1,4 @@ -from ..utils import from_db, to_db +from ..utils import from_db, to_db, SparseVec from .halfvec import HalfvecField +from .sparsevec import SparsevecField from .vector import VectorField diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 8d8d176..896bffd 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,6 +1,5 @@ from peewee import Expression, Field, Value -# TODO use halfvec functions -from ..utils import from_db, to_db +from ..utils import HalfVec class HalfvecField(Field): @@ -14,10 +13,12 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return to_db(value) + return HalfVec(value).to_db() def python_value(self, value): - return from_db(value) + if value is None: + return value + return HalfVec.from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py new file mode 100644 index 0000000..e171e63 --- /dev/null +++ b/pgvector/peewee/sparsevec.py @@ -0,0 +1,37 @@ +from peewee import Expression, Field, Value +from ..utils import SparseVec + + +class SparsevecField(Field): + field_type = 'sparsevec' + + def __init__(self, dimensions=None, *args, **kwargs): + self.dimensions = dimensions + super(SparsevecField, self).__init__(*args, **kwargs) + + def get_modifiers(self): + return self.dimensions and [self.dimensions] or None + + def db_value(self, value): + # TODO improve + return value.to_db() + + def python_value(self, value): + if value is None: + return value + return SparseVec.from_db(value) + + def _distance(self, op, vector): + return Expression(lhs=self, op=op, rhs=self.to_value(vector)) + + def l2_distance(self, vector): + return self._distance('<->', vector) + + def max_inner_product(self, vector): + return self._distance('<#>', vector) + + def cosine_distance(self, vector): + return self._distance('<=>', vector) + + def l1_distance(self, vector): + return self._distance('<+>', vector) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 970ee41..e429f8c 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField, HalfvecField +from pgvector.peewee import VectorField, HalfvecField, SparsevecField, SparseVec db = PostgresqlDatabase('pgvector_python_test') @@ -14,6 +14,7 @@ class Meta: class Item(BaseModel): embedding = VectorField(dimensions=3, null=True) half_embedding = HalfvecField(dimensions=3, null=True) + sparse_embedding = SparsevecField(dimensions=3, null=True) Item.add_index('embedding vector_l2_ops', using='hnsw') @@ -31,7 +32,7 @@ def create_items(): [1, 1, 2] ] for i, v in enumerate(vectors): - Item.create(id=i + 1, embedding=v, half_embedding=v) + Item.create(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v)) class TestPeewee: @@ -79,6 +80,13 @@ def test_halfvec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_sparsevec_l2_distance(self): + create_items() + distance = Item.sparse_embedding.l2_distance(SparseVec.from_dense([1, 1, 1])) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_where(self): create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) From afb416f947a21c08738857da68dfc7cca724926a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 16:31:04 -0400 Subject: [PATCH 173/879] Moved Vector to separate file for SQLAlchemy [skip ci] --- pgvector/sqlalchemy/__init__.py | 52 +-------------------------------- pgvector/sqlalchemy/vector.py | 51 ++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 51 deletions(-) create mode 100644 pgvector/sqlalchemy/vector.py diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index e035e79..65193c6 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,53 +1,3 @@ -from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import from_db, to_db +from .vector import Vector __all__ = ['Vector'] - - -class Vector(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 'VECTOR' - return 'VECTOR(%d)' % self.dim - - def bind_processor(self, dialect): - def process(value): - return 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(to_db(value, self.dim)) - return process - - def result_processor(self, dialect, coltype): - def process(value): - return 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['vector'] = Vector diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py new file mode 100644 index 0000000..01f1f24 --- /dev/null +++ b/pgvector/sqlalchemy/vector.py @@ -0,0 +1,51 @@ +from sqlalchemy.dialects.postgresql.base import ischema_names +from sqlalchemy.types import UserDefinedType, Float, String +from ..utils import from_db, to_db + + +class Vector(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 'VECTOR' + return 'VECTOR(%d)' % self.dim + + def bind_processor(self, dialect): + def process(value): + return 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(to_db(value, self.dim)) + return process + + def result_processor(self, dialect, coltype): + def process(value): + return 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['vector'] = Vector From 97e28a58b97e5fd347cd4e64a260792b634f239f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:27:21 -0400 Subject: [PATCH 174/879] Improved code --- pgvector/asyncpg/__init__.py | 5 ++--- pgvector/peewee/halfvec.py | 4 +--- pgvector/peewee/sparsevec.py | 5 +---- pgvector/psycopg/halfvec.py | 8 ++------ pgvector/psycopg/sparsevec.py | 8 ++------ pgvector/utils/halfvec.py | 22 +++++++++++++++++----- pgvector/utils/sparsevec.py | 18 +++++++++++++----- 7 files changed, 38 insertions(+), 32 deletions(-) diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 9603ac5..c8fc6b2 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -13,15 +13,14 @@ async def register_vector(conn): await conn.set_type_codec( 'halfvec', - encoder=lambda v: HalfVec(v).to_db_binary(), + encoder=HalfVec.to_db_binary, decoder=HalfVec.from_db_binary, format='binary' ) await conn.set_type_codec( 'sparsevec', - # TODO fix - encoder=lambda v: isinstance(v, SparseVec) and v.to_db_binary(), + encoder=SparseVec.to_db_binary, decoder=SparseVec.from_db_binary, format='binary' ) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 896bffd..60edd83 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -13,11 +13,9 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return HalfVec(value).to_db() + return HalfVec.to_db(value) def python_value(self, value): - if value is None: - return value return HalfVec.from_db(value) def _distance(self, op, vector): diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index e171e63..c44d4fe 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -13,12 +13,9 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - # TODO improve - return value.to_db() + return SparseVec.to_db(value) def python_value(self, value): - if value is None: - return value return SparseVec.from_db(value) def _distance(self, op, vector): diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 3a159d6..023586a 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -8,7 +8,7 @@ class HalfVecDumper(Dumper): format = Format.TEXT def dump(self, obj): - return obj.to_db().encode('utf8') + return HalfVec.to_db(obj).encode('utf8') class HalfVecBinaryDumper(HalfVecDumper): @@ -16,7 +16,7 @@ class HalfVecBinaryDumper(HalfVecDumper): format = Format.BINARY def dump(self, obj): - return obj.to_db_binary() + return HalfVec.to_db_binary(obj) class HalfVecLoader(Loader): @@ -24,8 +24,6 @@ class HalfVecLoader(Loader): format = Format.TEXT def load(self, data): - if data is None: - return None if isinstance(data, memoryview): data = bytes(data) return HalfVec.from_db(data.decode('utf8')) @@ -36,8 +34,6 @@ class HalfVecBinaryLoader(HalfVecLoader): format = Format.BINARY def load(self, data): - if data is None: - return None if isinstance(data, memoryview): data = bytes(data) return HalfVec.from_db_binary(data) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 18f6e1a..a5dec21 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -8,7 +8,7 @@ class SparseVecDumper(Dumper): format = Format.TEXT def dump(self, obj): - return obj.to_db().encode('utf8') + return SparseVec.to_db(obj).encode('utf8') class SparseVecBinaryDumper(SparseVecDumper): @@ -16,7 +16,7 @@ class SparseVecBinaryDumper(SparseVecDumper): format = Format.BINARY def dump(self, obj): - return obj.to_db_binary() + return SparseVec.to_db_binary(obj) class SparseVecLoader(Loader): @@ -24,8 +24,6 @@ class SparseVecLoader(Loader): format = Format.TEXT def load(self, data): - if data is None: - return None if isinstance(data, memoryview): data = bytes(data) return SparseVec.from_db(data.decode('utf8')) @@ -36,8 +34,6 @@ class SparseVecBinaryLoader(SparseVecLoader): format = Format.BINARY def load(self, data): - if data is None: - return None if isinstance(data, memoryview): data = bytes(data) return SparseVec.from_db_binary(data) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 262d3ad..36b1e4c 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -12,16 +12,28 @@ def __init__(self, value): def to_list(self): return list(self.value) - def to_db(self): - return '[' + ','.join([str(float(v)) for v in self.value]) + ']' - - def to_db_binary(self): - return pack(f'>HH{len(self.value)}e', len(self.value), 0, *self.value) + def to_db(value): + if value is None: + return value + if isinstance(value, HalfVec): + value = value.value + return '[' + ','.join([str(float(v)) for v in value]) + ']' + + def to_db_binary(value): + if value is None: + return value + if isinstance(value, HalfVec): + value = value.value + return pack(f'>HH{len(value)}e', len(value), 0, *value) def from_db(value): + if value is None: + return value return HalfVec([float(v) for v in value[1:-1].split(',')]) def from_db_binary(value): + if value is None: + return value dim, unused = unpack_from('>HH', value) return HalfVec(unpack_from(f'>{dim}e', value, 4)) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 72ec0e4..7e4cb6b 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -19,14 +19,20 @@ def to_dense(self): vec[i] = v return vec - def to_db(self): - return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(self.indices, self.values)]) + '}/' + str(self.dim) + def to_db(value): + if value is None: + return value + return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(value.indices, value.values)]) + '}/' + str(value.dim) - def to_db_binary(self): - nnz = len(self.indices) - return pack(f'>iii{nnz}i{nnz}f', self.dim, nnz, 0, *self.indices, *self.values) + def to_db_binary(value): + if value is None: + return value + nnz = len(value.indices) + return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) def from_db(value): + if value is None: + return value elements, dim = value.split('/') indices = [] values = [] @@ -37,6 +43,8 @@ def from_db(value): return SparseVec(int(dim), indices, values) def from_db_binary(value): + if value is None: + return value dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) From 4f07066ad0e999dcf45019f6857fcc733d89c592 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:31:57 -0400 Subject: [PATCH 175/879] Moved VectorAdapter to separate file for Psycopg 2 [skip ci] --- pgvector/psycopg2/__init__.py | 19 ++----------------- pgvector/psycopg2/vector.py | 21 +++++++++++++++++++++ 2 files changed, 23 insertions(+), 17 deletions(-) create mode 100644 pgvector/psycopg2/vector.py diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 94cc081..b51be42 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,23 +1,10 @@ -import numpy as np import psycopg2 -from psycopg2.extensions import adapt, new_type, register_adapter, register_type +from .vector import register_vector_info from ..utils import from_db, to_db __all__ = ['register_vector'] -class VectorAdapter(object): - def __init__(self, vector): - self._vector = vector - - def getquoted(self): - return adapt(to_db(self._vector)).getquoted() - - -def cast_vector(value, cur): - return from_db(value) - - def register_vector(conn_or_curs=None): cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs @@ -27,6 +14,4 @@ def register_vector(conn_or_curs=None): except psycopg2.errors.UndefinedObject: raise psycopg2.ProgrammingError('vector type not found in the database') - vector = new_type((oid,), 'VECTOR', cast_vector) - register_type(vector) - register_adapter(np.ndarray, VectorAdapter) + register_vector_info(oid) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py new file mode 100644 index 0000000..f314948 --- /dev/null +++ b/pgvector/psycopg2/vector.py @@ -0,0 +1,21 @@ +import numpy as np +from psycopg2.extensions import adapt, new_type, register_adapter, register_type +from ..utils import from_db, to_db + + +class VectorAdapter(object): + def __init__(self, vector): + self._vector = vector + + def getquoted(self): + return adapt(to_db(self._vector)).getquoted() + + +def cast_vector(value, cur): + return from_db(value) + + +def register_vector_info(oid): + vector = new_type((oid,), 'VECTOR', cast_vector) + register_type(vector) + register_adapter(np.ndarray, VectorAdapter) From 73ccb192bfb443d51674c82c1083e03a1a2c06a3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:40:55 -0400 Subject: [PATCH 176/879] Added support for halfvec type to Psycopg 2 [skip ci] --- CHANGELOG.md | 1 + pgvector/psycopg2/__init__.py | 9 +++++++-- pgvector/psycopg2/halfvec.py | 20 ++++++++++++++++++++ pgvector/utils/halfvec.py | 4 ++-- pgvector/utils/sparsevec.py | 4 ++-- tests/test_psycopg2.py | 21 +++++++++++++++------ 6 files changed, 47 insertions(+), 12 deletions(-) create mode 100644 pgvector/psycopg2/halfvec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c48f5ad..2c93317 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.2.6 (unreleased) - Added support for `halfvec` and `sparsevec` types to Psycopg 3 +- Added support for `halfvec` type to Psycopg 2 - Added support for `halfvec` and `sparsevec` types to asyncpg - Added support for `halfvec` and `sparsevec` types to Peewee - Added `L1Distance` for Django diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index b51be42..bde04e5 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,4 +1,5 @@ import psycopg2 +from .halfvec import register_halfvec_info from .vector import register_vector_info from ..utils import from_db, to_db @@ -10,8 +11,12 @@ def register_vector(conn_or_curs=None): try: cur.execute('SELECT NULL::vector') - oid = cur.description[0][1] + register_vector_info(cur.description[0][1]) except psycopg2.errors.UndefinedObject: raise psycopg2.ProgrammingError('vector type not found in the database') - register_vector_info(oid) + try: + cur.execute('SELECT NULL::halfvec') + register_halfvec_info(cur.description[0][1]) + except psycopg2.errors.UndefinedObject: + pass diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py new file mode 100644 index 0000000..dc87ba1 --- /dev/null +++ b/pgvector/psycopg2/halfvec.py @@ -0,0 +1,20 @@ +from psycopg2.extensions import adapt, new_type, register_adapter, register_type +from ..utils import HalfVec + + +class HalfvecAdapter(object): + def __init__(self, halfvec): + self._halfvec = halfvec + + def getquoted(self): + return adapt(HalfVec.to_db(self._halfvec)).getquoted() + + +def cast_halfvec(value, cur): + return HalfVec.from_db(value) + + +def register_halfvec_info(oid): + halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) + register_type(halfvec) + register_adapter(HalfVec, HalfvecAdapter) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 36b1e4c..627c1c6 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -27,12 +27,12 @@ def to_db_binary(value): return pack(f'>HH{len(value)}e', len(value), 0, *value) def from_db(value): - if value is None: + if value is None or isinstance(value, HalfVec): return value return HalfVec([float(v) for v in value[1:-1].split(',')]) def from_db_binary(value): - if value is None: + if value is None or isinstance(value, HalfVec): return value dim, unused = unpack_from('>HH', value) return HalfVec(unpack_from(f'>{dim}e', value, 4)) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 7e4cb6b..f7f6607 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -31,7 +31,7 @@ def to_db_binary(value): return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) def from_db(value): - if value is None: + if value is None or isinstance(value, SparseVec): return value elements, dim = value.split('/') indices = [] @@ -43,7 +43,7 @@ def from_db(value): return SparseVec(int(dim), indices, values) def from_db_binary(value): - if value is None: + if value is None or isinstance(value, SparseVec): return value dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 925d60f..1cc9669 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -7,22 +7,31 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') -cur.execute('DROP TABLE IF EXISTS items') -cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +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))') register_vector(cur) class TestPsycopg2: def setup_method(self, test_method): - cur.execute('DELETE FROM items') + cur.execute('DELETE FROM psycopg2_items') - def test_works(self): + def test_vector(self): embedding = np.array([1.5, 2, 3]) - cur.execute('INSERT INTO items (embedding) VALUES (%s), (NULL)', (embedding,)) + cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) - cur.execute('SELECT * FROM items ORDER BY id') + cur.execute('SELECT * FROM psycopg2_items ORDER BY id') res = cur.fetchall() assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None + + def test_halfvec(self): + embedding = [1.5, 2, 3] + cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT id, half_embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][1].to_list() == [1.5, 2, 3] + assert res[1][1] is None From 34c891e4131b13f434c3318d822288311a28c831 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:45:12 -0400 Subject: [PATCH 177/879] Added support for sparsevec type to Psycopg 2 [skip ci] --- CHANGELOG.md | 2 +- pgvector/psycopg2/__init__.py | 9 ++++++++- pgvector/psycopg2/sparsevec.py | 20 ++++++++++++++++++++ tests/test_psycopg2.py | 27 ++++++++++++++++++--------- 4 files changed, 47 insertions(+), 11 deletions(-) create mode 100644 pgvector/psycopg2/sparsevec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c93317..d4b1502 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ ## 0.2.6 (unreleased) - Added support for `halfvec` and `sparsevec` types to Psycopg 3 -- Added support for `halfvec` type to Psycopg 2 +- Added support for `halfvec` and `sparsevec` types to Psycopg 2 - Added support for `halfvec` and `sparsevec` types to asyncpg - Added support for `halfvec` and `sparsevec` types to Peewee - Added `L1Distance` for Django diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index bde04e5..16b7531 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,7 +1,8 @@ import psycopg2 from .halfvec import register_halfvec_info +from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import from_db, to_db +from ..utils import from_db, to_db, SparseVec __all__ = ['register_vector'] @@ -20,3 +21,9 @@ def register_vector(conn_or_curs=None): register_halfvec_info(cur.description[0][1]) except psycopg2.errors.UndefinedObject: pass + + try: + cur.execute('SELECT NULL::sparsevec') + register_sparsevec_info(cur.description[0][1]) + except psycopg2.errors.UndefinedObject: + pass diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py new file mode 100644 index 0000000..6603b2f --- /dev/null +++ b/pgvector/psycopg2/sparsevec.py @@ -0,0 +1,20 @@ +from psycopg2.extensions import adapt, new_type, register_adapter, register_type +from ..utils import SparseVec + + +class SparsevecAdapter(object): + def __init__(self, value): + self._value = value + + def getquoted(self): + return adapt(SparseVec.to_db(self._value)).getquoted() + + +def cast_sparsevec(value, cur): + return SparseVec.from_db(value) + + +def register_sparsevec_info(oid): + sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) + register_type(sparsevec) + register_adapter(SparseVec, SparsevecAdapter) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 1cc9669..b5ac50a 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg2 import register_vector +from pgvector.psycopg2 import register_vector, SparseVec import psycopg2 conn = psycopg2.connect(dbname='pgvector_python_test') @@ -8,7 +8,7 @@ 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))') +cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), sparse_embedding sparsevec(3))') register_vector(cur) @@ -21,17 +21,26 @@ def test_vector(self): embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) - cur.execute('SELECT * FROM psycopg2_items ORDER BY id') + cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][1], embedding) - assert res[0][1].dtype == np.float32 - assert res[1][1] is None + assert np.array_equal(res[0][0], embedding) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None def test_halfvec(self): embedding = [1.5, 2, 3] cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) - cur.execute('SELECT id, half_embedding FROM psycopg2_items ORDER BY id') + cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][1].to_list() == [1.5, 2, 3] - assert res[1][1] is None + assert res[0][0].to_list() == [1.5, 2, 3] + assert res[1][0] is None + + def test_sparsevec(self): + embedding = SparseVec.from_dense([1.5, 2, 3]) + cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][0].to_dense() == [1.5, 2, 3] + assert res[1][0] is None From 57c1b3fca849c1c1ed781356e261dd248c024b06 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:46:05 -0400 Subject: [PATCH 178/879] Improved code [skip ci] --- pgvector/psycopg2/halfvec.py | 6 +++--- pgvector/psycopg2/vector.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index dc87ba1..8d974b3 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -3,11 +3,11 @@ class HalfvecAdapter(object): - def __init__(self, halfvec): - self._halfvec = halfvec + def __init__(self, value): + self._value = value def getquoted(self): - return adapt(HalfVec.to_db(self._halfvec)).getquoted() + return adapt(HalfVec.to_db(self._value)).getquoted() def cast_halfvec(value, cur): diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index f314948..8bcef06 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -4,11 +4,11 @@ class VectorAdapter(object): - def __init__(self, vector): - self._vector = vector + def __init__(self, value): + self._value = value def getquoted(self): - return adapt(to_db(self._vector)).getquoted() + return adapt(to_db(self._value)).getquoted() def cast_vector(value, cur): From 00805065e599dafb1ee0af11ec9dc0c9ee9e28a9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:48:16 -0400 Subject: [PATCH 179/879] Improved code [skip ci] --- pgvector/asyncpg/__init__.py | 6 +++--- pgvector/utils/vector.py | 7 +++++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index c8fc6b2..9c6b660 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,4 +1,4 @@ -from ..utils import from_db, from_db_binary, to_db, to_db_binary, HalfVec, SparseVec +from ..utils import Vector, HalfVec, SparseVec __all__ = ['register_vector'] @@ -6,8 +6,8 @@ async def register_vector(conn): await conn.set_type_codec( 'vector', - encoder=to_db_binary, - decoder=from_db_binary, + encoder=Vector.to_db_binary, + decoder=Vector.from_db_binary, format='binary' ) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index d831dac..bf3e165 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -47,3 +47,10 @@ def to_db_binary(value): raise ValueError('expected ndim to be 1') return pack('>HH', value.shape[0], 0) + value.tobytes() + + +class Vector: + from_db = from_db + from_db_binary = from_db_binary + to_db = to_db + to_db_binary = to_db_binary From 36e20b300795ea1909994c561969026f77e94aaa Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:54:28 -0400 Subject: [PATCH 180/879] Improved code [skip ci] --- pgvector/django/__init__.py | 10 ++--- pgvector/peewee/__init__.py | 2 +- pgvector/peewee/vector.py | 6 +-- pgvector/psycopg/__init__.py | 3 -- pgvector/psycopg/vector.py | 10 ++--- pgvector/psycopg2/__init__.py | 2 +- pgvector/psycopg2/vector.py | 6 +-- pgvector/sqlalchemy/vector.py | 8 ++-- pgvector/utils/vector.py | 69 +++++++++++++++-------------------- 9 files changed, 52 insertions(+), 64 deletions(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index e4c8e21..abcf6dd 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -3,7 +3,7 @@ from django.db.models import Field, FloatField, Func, Value import numpy as np from .forms import VectorFormField -from ..utils import from_db, to_db +from ..utils import Vector __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] @@ -34,15 +34,15 @@ def db_type(self, connection): return 'vector(%d)' % self.dimensions def from_db_value(self, value, expression, connection): - return from_db(value) + return Vector.from_db(value) def to_python(self, value): if isinstance(value, list): return np.array(value, dtype=np.float32) - return from_db(value) + return Vector.from_db(value) def get_prep_value(self, value): - return to_db(value) + return Vector.to_db(value) def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) @@ -111,7 +111,7 @@ class DistanceBase(Func): def __init__(self, expression, vector, **extra): if not hasattr(vector, 'resolve_expression'): - vector = Value(to_db(vector)) + vector = Value(Vector.to_db(vector)) super().__init__(expression, vector, **extra) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 106170f..665c385 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -1,4 +1,4 @@ -from ..utils import from_db, to_db, SparseVec +from ..utils import SparseVec from .halfvec import HalfvecField from .sparsevec import SparsevecField from .vector import VectorField diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 756959d..3c11e74 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,5 +1,5 @@ from peewee import Expression, Field, Value -from ..utils import from_db, to_db +from ..utils import Vector class VectorField(Field): @@ -13,10 +13,10 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return to_db(value) + return Vector.to_db(value) def python_value(self, value): - return from_db(value) + return Vector.from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 3228afc..b5df2d1 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -5,9 +5,6 @@ from .vector import register_vector_info from ..utils import HalfVec, SparseVec -# TODO remove in 0.3.0 -from .vector import * -from ..utils import from_db, from_db_binary, to_db, to_db_binary __all__ = ['register_vector'] diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index aa15b92..fe606a6 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,6 +1,6 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import from_db, from_db_binary, to_db, to_db_binary +from ..utils import Vector class VectorDumper(Dumper): @@ -8,7 +8,7 @@ class VectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return to_db(obj).encode('utf8') + return Vector.to_db(obj).encode('utf8') class VectorBinaryDumper(VectorDumper): @@ -16,7 +16,7 @@ class VectorBinaryDumper(VectorDumper): format = Format.BINARY def dump(self, obj): - return to_db_binary(obj) + return Vector.to_db_binary(obj) class VectorLoader(Loader): @@ -26,7 +26,7 @@ class VectorLoader(Loader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return from_db(data.decode('utf8')) + return Vector.from_db(data.decode('utf8')) class VectorBinaryLoader(VectorLoader): @@ -36,7 +36,7 @@ class VectorBinaryLoader(VectorLoader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return from_db_binary(data) + return Vector.from_db_binary(data) def register_vector_info(context, info): diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 16b7531..764c0f7 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -2,7 +2,7 @@ from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import from_db, to_db, SparseVec +from ..utils import SparseVec __all__ = ['register_vector'] diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 8bcef06..c0b066a 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,6 +1,6 @@ import numpy as np from psycopg2.extensions import adapt, new_type, register_adapter, register_type -from ..utils import from_db, to_db +from ..utils import Vector class VectorAdapter(object): @@ -8,11 +8,11 @@ def __init__(self, value): self._value = value def getquoted(self): - return adapt(to_db(self._value)).getquoted() + return adapt(Vector.to_db(self._value)).getquoted() def cast_vector(value, cur): - return from_db(value) + return Vector.from_db(value) def register_vector_info(oid): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 01f1f24..056dcef 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import from_db, to_db +from ..utils import Vector as Vec class Vector(UserDefinedType): @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return to_db(value, self.dim) + return Vec.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(to_db(value, self.dim)) + return string_literal_processor(Vec.to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return from_db(value) + return Vec.from_db(value) return process class comparator_factory(UserDefinedType.Comparator): diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index bf3e165..c172b5e 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -2,55 +2,46 @@ from struct import pack, unpack_from -def from_db(value): - # could be ndarray if already cast by lower-level driver - if value is None or isinstance(value, np.ndarray): - return value - - return np.array(value[1:-1].split(','), dtype=np.float32) - - -def from_db_binary(value): - if value is None: - return value +class Vector: + def from_db(value): + # could be ndarray if already cast by lower-level driver + if value is None or isinstance(value, np.ndarray): + return value - dim, unused = unpack_from('>HH', value) - return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32) + return np.array(value[1:-1].split(','), dtype=np.float32) + def from_db_binary(value): + if value is None or isinstance(value, np.ndarray): + return value -def to_db(value, dim=None): - if value is None: - return value + dim, unused = unpack_from('>HH', value) + return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32) - if isinstance(value, np.ndarray): - if value.ndim != 1: - raise ValueError('expected ndim to be 1') + def to_db(value, dim=None): + if value is None: + return value - if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating): - raise ValueError('dtype must be numeric') + if isinstance(value, np.ndarray): + if value.ndim != 1: + raise ValueError('expected ndim to be 1') - value = value.tolist() + if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating): + raise ValueError('dtype must be numeric') - if dim is not None and len(value) != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) + value = value.tolist() - return '[' + ','.join([str(float(v)) for v in 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): - if value is None: - return value + def to_db_binary(value): + if value is None: + return value - value = np.asarray(value, dtype='>f') - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - return pack('>HH', value.shape[0], 0) + value.tobytes() + value = np.asarray(value, dtype='>f') + if value.ndim != 1: + raise ValueError('expected ndim to be 1') -class Vector: - from_db = from_db - from_db_binary = from_db_binary - to_db = to_db - to_db_binary = to_db_binary + return pack('>HH', value.shape[0], 0) + value.tobytes() From 5061cae02c461cf11980a51110884b74d99e49cc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 17:58:13 -0400 Subject: [PATCH 181/879] Moved indexes to separate file for Django [skip ci] --- pgvector/django/__init__.py | 47 +------------------------------------ pgvector/django/indexes.py | 46 ++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 46 deletions(-) create mode 100644 pgvector/django/indexes.py diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index abcf6dd..04d013d 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,8 +1,8 @@ from django.contrib.postgres.operations import CreateExtension -from django.contrib.postgres.indexes import PostgresIndex from django.db.models import Field, FloatField, Func, Value import numpy as np from .forms import VectorFormField +from .indexes import IvfflatIndex, HnswIndex from ..utils import Vector __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] @@ -61,51 +61,6 @@ def formfield(self, **kwargs): return super().formfield(form_class=VectorFormField, **kwargs) -class IvfflatIndex(PostgresIndex): - suffix = 'ivfflat' - - def __init__(self, *expressions, lists=None, **kwargs): - self.lists = lists - super().__init__(*expressions, **kwargs) - - def deconstruct(self): - path, args, kwargs = super().deconstruct() - if self.lists is not None: - kwargs['lists'] = self.lists - return path, args, kwargs - - def get_with_params(self): - with_params = [] - if self.lists is not None: - with_params.append('lists = %d' % self.lists) - return with_params - - -class HnswIndex(PostgresIndex): - suffix = 'hnsw' - - def __init__(self, *expressions, m=None, ef_construction=None, **kwargs): - self.m = m - self.ef_construction = ef_construction - super().__init__(*expressions, **kwargs) - - def deconstruct(self): - path, args, kwargs = super().deconstruct() - if self.m is not None: - kwargs['m'] = self.m - if self.ef_construction is not None: - kwargs['ef_construction'] = self.ef_construction - return path, args, kwargs - - def get_with_params(self): - with_params = [] - if self.m is not None: - with_params.append('m = %d' % self.m) - if self.ef_construction is not None: - with_params.append('ef_construction = %d' % self.ef_construction) - return with_params - - class DistanceBase(Func): output_field = FloatField() diff --git a/pgvector/django/indexes.py b/pgvector/django/indexes.py new file mode 100644 index 0000000..5bec0eb --- /dev/null +++ b/pgvector/django/indexes.py @@ -0,0 +1,46 @@ +from django.contrib.postgres.indexes import PostgresIndex + + +class IvfflatIndex(PostgresIndex): + suffix = 'ivfflat' + + def __init__(self, *expressions, lists=None, **kwargs): + self.lists = lists + super().__init__(*expressions, **kwargs) + + def deconstruct(self): + path, args, kwargs = super().deconstruct() + if self.lists is not None: + kwargs['lists'] = self.lists + return path, args, kwargs + + def get_with_params(self): + with_params = [] + if self.lists is not None: + with_params.append('lists = %d' % self.lists) + return with_params + + +class HnswIndex(PostgresIndex): + suffix = 'hnsw' + + def __init__(self, *expressions, m=None, ef_construction=None, **kwargs): + self.m = m + self.ef_construction = ef_construction + super().__init__(*expressions, **kwargs) + + def deconstruct(self): + path, args, kwargs = super().deconstruct() + if self.m is not None: + kwargs['m'] = self.m + if self.ef_construction is not None: + kwargs['ef_construction'] = self.ef_construction + return path, args, kwargs + + def get_with_params(self): + with_params = [] + if self.m is not None: + with_params.append('m = %d' % self.m) + if self.ef_construction is not None: + with_params.append('ef_construction = %d' % self.ef_construction) + return with_params From 3415860599853caad00243060412776e247d3746 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:00:52 -0400 Subject: [PATCH 182/879] Moved VectorField to separate file for Django [skip ci] --- pgvector/django/__init__.py | 50 +---------------------------------- pgvector/django/vector.py | 52 +++++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 49 deletions(-) create mode 100644 pgvector/django/vector.py diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 04d013d..06adf01 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,8 +1,8 @@ from django.contrib.postgres.operations import CreateExtension from django.db.models import Field, FloatField, Func, Value import numpy as np -from .forms import VectorFormField from .indexes import IvfflatIndex, HnswIndex +from .vector import VectorField from ..utils import Vector __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] @@ -13,54 +13,6 @@ def __init__(self): self.name = 'vector' -# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ -class VectorField(Field): - description = 'Vector' - empty_strings_allowed = False - - def __init__(self, *args, dimensions=None, **kwargs): - self.dimensions = dimensions - super().__init__(*args, **kwargs) - - def deconstruct(self): - name, path, args, kwargs = super().deconstruct() - if self.dimensions is not None: - kwargs['dimensions'] = self.dimensions - return name, path, args, kwargs - - def db_type(self, connection): - if self.dimensions is None: - return 'vector' - return 'vector(%d)' % self.dimensions - - def from_db_value(self, value, expression, connection): - return Vector.from_db(value) - - def to_python(self, value): - if isinstance(value, list): - return np.array(value, dtype=np.float32) - return Vector.from_db(value) - - def get_prep_value(self, value): - return Vector.to_db(value) - - def value_to_string(self, obj): - return self.get_prep_value(self.value_from_object(obj)) - - def validate(self, value, model_instance): - if isinstance(value, np.ndarray): - value = value.tolist() - super().validate(value, model_instance) - - def run_validators(self, value): - if isinstance(value, np.ndarray): - value = value.tolist() - super().run_validators(value) - - def formfield(self, **kwargs): - return super().formfield(form_class=VectorFormField, **kwargs) - - class DistanceBase(Func): output_field = FloatField() diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py new file mode 100644 index 0000000..05789fa --- /dev/null +++ b/pgvector/django/vector.py @@ -0,0 +1,52 @@ +from django.db.models import Field +import numpy as np +from .forms import VectorFormField +from ..utils import Vector + + +# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ +class VectorField(Field): + description = 'Vector' + empty_strings_allowed = False + + def __init__(self, *args, dimensions=None, **kwargs): + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def deconstruct(self): + name, path, args, kwargs = super().deconstruct() + if self.dimensions is not None: + kwargs['dimensions'] = self.dimensions + return name, path, args, kwargs + + def db_type(self, connection): + if self.dimensions is None: + return 'vector' + return 'vector(%d)' % self.dimensions + + def from_db_value(self, value, expression, connection): + return Vector.from_db(value) + + def to_python(self, value): + if isinstance(value, list): + return np.array(value, dtype=np.float32) + return Vector.from_db(value) + + def get_prep_value(self, value): + return Vector.to_db(value) + + def value_to_string(self, obj): + return self.get_prep_value(self.value_from_object(obj)) + + def validate(self, value, model_instance): + if isinstance(value, np.ndarray): + value = value.tolist() + super().validate(value, model_instance) + + def run_validators(self, value): + if isinstance(value, np.ndarray): + value = value.tolist() + super().run_validators(value) + + def formfield(self, **kwargs): + return super().formfield(form_class=VectorFormField, **kwargs) From 18d1b65687e868eeb1959864c70b344d485798e2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:04:03 -0400 Subject: [PATCH 183/879] Improved code [skip ci] --- pgvector/django/__init__.py | 33 +-------------------------------- pgvector/django/forms.py | 12 ------------ pgvector/django/functions.py | 30 ++++++++++++++++++++++++++++++ pgvector/django/vector.py | 18 +++++++++++++++++- pgvector/django/widgets.py | 9 --------- 5 files changed, 48 insertions(+), 54 deletions(-) delete mode 100644 pgvector/django/forms.py create mode 100644 pgvector/django/functions.py delete mode 100644 pgvector/django/widgets.py diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 06adf01..a32d329 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,9 +1,7 @@ from django.contrib.postgres.operations import CreateExtension -from django.db.models import Field, FloatField, Func, Value -import numpy as np +from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance from .indexes import IvfflatIndex, HnswIndex from .vector import VectorField -from ..utils import Vector __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] @@ -11,32 +9,3 @@ class VectorExtension(CreateExtension): def __init__(self): self.name = 'vector' - - -class DistanceBase(Func): - output_field = FloatField() - - def __init__(self, expression, vector, **extra): - if not hasattr(vector, 'resolve_expression'): - vector = Value(Vector.to_db(vector)) - super().__init__(expression, vector, **extra) - - -class L2Distance(DistanceBase): - function = '' - arg_joiner = ' <-> ' - - -class MaxInnerProduct(DistanceBase): - function = '' - arg_joiner = ' <#> ' - - -class CosineDistance(DistanceBase): - function = '' - arg_joiner = ' <=> ' - - -class L1Distance(DistanceBase): - function = '' - arg_joiner = ' <+> ' diff --git a/pgvector/django/forms.py b/pgvector/django/forms.py deleted file mode 100644 index 3748236..0000000 --- a/pgvector/django/forms.py +++ /dev/null @@ -1,12 +0,0 @@ -from django import forms -import numpy as np -from .widgets import VectorWidget - - -class VectorFormField(forms.CharField): - widget = VectorWidget - - def has_changed(self, initial, data): - if isinstance(initial, np.ndarray): - initial = initial.tolist() - return super().has_changed(initial, data) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py new file mode 100644 index 0000000..2638d67 --- /dev/null +++ b/pgvector/django/functions.py @@ -0,0 +1,30 @@ +from django.db.models import FloatField, Func, Value +from ..utils import Vector + +class DistanceBase(Func): + output_field = FloatField() + + def __init__(self, expression, vector, **extra): + if not hasattr(vector, 'resolve_expression'): + vector = Value(Vector.to_db(vector)) + super().__init__(expression, vector, **extra) + + +class L2Distance(DistanceBase): + function = '' + arg_joiner = ' <-> ' + + +class MaxInnerProduct(DistanceBase): + function = '' + arg_joiner = ' <#> ' + + +class CosineDistance(DistanceBase): + function = '' + arg_joiner = ' <=> ' + + +class L1Distance(DistanceBase): + function = '' + arg_joiner = ' <+> ' diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 05789fa..063b5ff 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -1,6 +1,6 @@ +from django import forms from django.db.models import Field import numpy as np -from .forms import VectorFormField from ..utils import Vector @@ -50,3 +50,19 @@ def run_validators(self, value): def formfield(self, **kwargs): return super().formfield(form_class=VectorFormField, **kwargs) + + +class VectorWidget(forms.TextInput): + def format_value(self, value): + if isinstance(value, np.ndarray): + value = value.tolist() + return super().format_value(value) + + +class VectorFormField(forms.CharField): + widget = VectorWidget + + def has_changed(self, initial, data): + if isinstance(initial, np.ndarray): + initial = initial.tolist() + return super().has_changed(initial, data) diff --git a/pgvector/django/widgets.py b/pgvector/django/widgets.py deleted file mode 100644 index 731d632..0000000 --- a/pgvector/django/widgets.py +++ /dev/null @@ -1,9 +0,0 @@ -from django import forms -import numpy as np - - -class VectorWidget(forms.TextInput): - def format_value(self, value): - if isinstance(value, np.ndarray): - value = value.tolist() - return super().format_value(value) From 7dbef6894d583901ab7a7b9135f7aea043b0792b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:06:06 -0400 Subject: [PATCH 184/879] Moved VectorExtension to separate file for Django [skip ci] --- pgvector/django/__init__.py | 7 +------ pgvector/django/extensions.py | 6 ++++++ pgvector/django/functions.py | 1 + 3 files changed, 8 insertions(+), 6 deletions(-) create mode 100644 pgvector/django/extensions.py diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index a32d329..051090a 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,11 +1,6 @@ -from django.contrib.postgres.operations import CreateExtension +from .extensions import VectorExtension from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance from .indexes import IvfflatIndex, HnswIndex from .vector import VectorField __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] - - -class VectorExtension(CreateExtension): - def __init__(self): - self.name = 'vector' diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py new file mode 100644 index 0000000..0573f72 --- /dev/null +++ b/pgvector/django/extensions.py @@ -0,0 +1,6 @@ +from django.contrib.postgres.operations import CreateExtension + + +class VectorExtension(CreateExtension): + def __init__(self): + self.name = 'vector' diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 2638d67..4f696ff 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,6 +1,7 @@ from django.db.models import FloatField, Func, Value from ..utils import Vector + class DistanceBase(Func): output_field = FloatField() From 2d49011020173ca4f5e41514c4b26a81672d00b2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:12:55 -0400 Subject: [PATCH 185/879] Added .pytest_cache to .gitignore [skip ci] --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 1d3b727..f7ff659 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ venv/ .cache/ *.pyc __pycache__ +.pytest_cache/ From 828508ae6785400370263755f8ec01bc6b07e26a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:19:21 -0400 Subject: [PATCH 186/879] Added support for halfvec type to Django [skip ci] --- CHANGELOG.md | 1 + pgvector/django/__init__.py | 3 ++- pgvector/django/halfvec.py | 35 +++++++++++++++++++++++++++++++++++ tests/test_django.py | 17 +++++++++++++---- 4 files changed, 51 insertions(+), 5 deletions(-) create mode 100644 pgvector/django/halfvec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d4b1502..792dbab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.2.6 (unreleased) +- Added support for `halfvec` type to Django - 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/django/__init__.py b/pgvector/django/__init__.py index 051090a..5530c52 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,6 +1,7 @@ from .extensions import VectorExtension from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance +from .halfvec import HalfvecField from .indexes import IvfflatIndex, HnswIndex from .vector import VectorField -__all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] +__all__ = ['VectorExtension', 'VectorField', 'HalfvecField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py new file mode 100644 index 0000000..ec08a42 --- /dev/null +++ b/pgvector/django/halfvec.py @@ -0,0 +1,35 @@ +from django.db.models import Field +from ..utils import HalfVec + + +# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ +class HalfvecField(Field): + description = 'Halfvec' + empty_strings_allowed = False + + def __init__(self, *args, dimensions=None, **kwargs): + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def deconstruct(self): + name, path, args, kwargs = super().deconstruct() + if self.dimensions is not None: + kwargs['dimensions'] = self.dimensions + return name, path, args, kwargs + + def db_type(self, connection): + if self.dimensions is None: + return 'halfvec' + return 'halfvec(%d)' % self.dimensions + + def from_db_value(self, value, expression, connection): + return HalfVec.from_db(value) + + def to_python(self, value): + return HalfVec.from_db(value) + + def get_prep_value(self, value): + return HalfVec.to_db(value) + + def value_to_string(self, obj): + return self.get_prep_value(self.value_from_object(obj)) diff --git a/tests/test_django.py b/tests/test_django.py index b6274cd..b181d5c 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance +from pgvector.django import VectorExtension, VectorField, HalfvecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance from unittest import mock settings.configure( @@ -23,7 +23,8 @@ class Item(models.Model): - embedding = VectorField(dimensions=3) + embedding = VectorField(dimensions=3, null=True, blank=True) + half_embedding = HalfvecField(dimensions=3, null=True, blank=True) class Meta: app_label = 'myapp' @@ -56,7 +57,8 @@ class Migration(migrations.Migration): name='Item', fields=[ ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), - ('embedding', pgvector.django.VectorField(dimensions=3, null=True)), + ('embedding', pgvector.django.VectorField(dimensions=3, null=True, blank=True)), + ('half_embedding', pgvector.django.HalfvecField(dimensions=3, null=True, blank=True)), ], ), migrations.AddIndex( @@ -102,7 +104,7 @@ class TestDjango: def setup_method(self, test_method): Item.objects.all().delete() - def test_works(self): + def test_vector(self): item = Item(id=1, embedding=[1, 2, 3]) item.save() item = Item.objects.get(pk=1) @@ -110,6 +112,13 @@ def test_works(self): assert np.array_equal(item.embedding, np.array([1, 2, 3])) assert item.embedding.dtype == np.float32 + def test_halfvec(self): + item = Item(id=1, half_embedding=[1, 2, 3]) + item.save() + item = Item.objects.get(pk=1) + assert item.id == 1 + assert item.half_embedding.to_list() == [1, 2, 3] + def test_l2_distance(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) From 70c1f5813db4f0c312f39cac8a546c65e520c5ef Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:23:36 -0400 Subject: [PATCH 187/879] Added support for sparsevec type to Django [skip ci] --- CHANGELOG.md | 2 +- pgvector/django/__init__.py | 4 +++- pgvector/django/sparsevec.py | 35 +++++++++++++++++++++++++++++++++++ tests/test_django.py | 11 ++++++++++- 4 files changed, 49 insertions(+), 3 deletions(-) create mode 100644 pgvector/django/sparsevec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 792dbab..d5c9642 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## 0.2.6 (unreleased) -- Added support for `halfvec` type to Django +- Added support for `halfvec` and `sparsevec` types to Django - 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/django/__init__.py b/pgvector/django/__init__.py index 5530c52..2281ed9 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -2,6 +2,8 @@ from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance from .halfvec import HalfvecField from .indexes import IvfflatIndex, HnswIndex +from .sparsevec import SparsevecField from .vector import VectorField +from ..utils import SparseVec -__all__ = ['VectorExtension', 'VectorField', 'HalfvecField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] +__all__ = ['VectorExtension', 'VectorField', 'HalfvecField', 'SparsevecField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py new file mode 100644 index 0000000..91b9cf8 --- /dev/null +++ b/pgvector/django/sparsevec.py @@ -0,0 +1,35 @@ +from django.db.models import Field +from ..utils import SparseVec + + +# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ +class SparsevecField(Field): + description = 'Sparsevec' + empty_strings_allowed = False + + def __init__(self, *args, dimensions=None, **kwargs): + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def deconstruct(self): + name, path, args, kwargs = super().deconstruct() + if self.dimensions is not None: + kwargs['dimensions'] = self.dimensions + return name, path, args, kwargs + + def db_type(self, connection): + if self.dimensions is None: + return 'sparsevec' + return 'sparsevec(%d)' % self.dimensions + + def from_db_value(self, value, expression, connection): + return SparseVec.from_db(value) + + def to_python(self, value): + return SparseVec.from_db(value) + + def get_prep_value(self, value): + return SparseVec.to_db(value) + + def value_to_string(self, obj): + return self.get_prep_value(self.value_from_object(obj)) diff --git a/tests/test_django.py b/tests/test_django.py index b181d5c..271c2cc 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfvecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance +from pgvector.django import VectorExtension, VectorField, HalfvecField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, SparseVec from unittest import mock settings.configure( @@ -25,6 +25,7 @@ class Item(models.Model): embedding = VectorField(dimensions=3, null=True, blank=True) half_embedding = HalfvecField(dimensions=3, null=True, blank=True) + sparse_embedding = SparsevecField(dimensions=3, null=True, blank=True) class Meta: app_label = 'myapp' @@ -59,6 +60,7 @@ class Migration(migrations.Migration): ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), ('embedding', pgvector.django.VectorField(dimensions=3, null=True, blank=True)), ('half_embedding', pgvector.django.HalfvecField(dimensions=3, null=True, blank=True)), + ('sparse_embedding', pgvector.django.SparsevecField(dimensions=3, null=True, blank=True)), ], ), migrations.AddIndex( @@ -119,6 +121,13 @@ def test_halfvec(self): assert item.id == 1 assert item.half_embedding.to_list() == [1, 2, 3] + def test_sparsevec(self): + item = Item(id=1, sparse_embedding=SparseVec.from_dense([1, 2, 3])) + item.save() + item = Item.objects.get(pk=1) + assert item.id == 1 + assert item.sparse_embedding.to_dense() == [1, 2, 3] + def test_l2_distance(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) From b4b6c6b9f53a6574bb08773677e10472614f12b0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:28:59 -0400 Subject: [PATCH 188/879] Added support for halfvec type to SQLAlchemy and SQLModel [skip ci] --- CHANGELOG.md | 1 + pgvector/sqlalchemy/__init__.py | 3 +- pgvector/sqlalchemy/halfvec.py | 51 +++++++++++++++++++++++++++++++++ pgvector/utils/halfvec.py | 6 +++- tests/test_sqlalchemy.py | 5 ++-- 5 files changed, 62 insertions(+), 4 deletions(-) create mode 100644 pgvector/sqlalchemy/halfvec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d5c9642..1331f28 100644 --- 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 index 65193c6..d4c0179 100644 --- 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 new file mode 100644 index 0000000..694d14f --- /dev/null +++ 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 index 627c1c6..8b760cd 100644 --- 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 index 5db608f..9d5d168 100644 --- 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() From f1cf31e0f7ae7b3106f85cd06ffdc8725fc89e74 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:32:32 -0400 Subject: [PATCH 189/879] Added support for sparsevec type to SQLAlchemy and SQLModel [skip ci] --- CHANGELOG.md | 4 +-- pgvector/sqlalchemy/__init__.py | 4 ++- pgvector/sqlalchemy/sparsevec.py | 51 ++++++++++++++++++++++++++++++++ pgvector/utils/sparsevec.py | 6 +++- tests/test_sqlalchemy.py | 5 ++-- 5 files changed, 64 insertions(+), 6 deletions(-) create mode 100644 pgvector/sqlalchemy/sparsevec.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1331f28..eb6fab7 100644 --- 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 index d4c0179..a7db132 100644 --- 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 new file mode 100644 index 0000000..6e3c691 --- /dev/null +++ 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 index f7f6607..9aae64e 100644 --- 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 index 9d5d168..e6b869a 100644 --- 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() From 19054225385229ddf54f579010d509f608e10fbb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:33:52 -0400 Subject: [PATCH 190/879] Added new types to SQLModel tests [skip ci] --- tests/test_sqlmodel.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 8d349ea..2fe2930 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import Vector, Halfvec, Sparsevec, SparseVec import pytest from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError @@ -17,6 +17,8 @@ class Item(SQLModel, table=True): id: Optional[int] = Field(default=None, primary_key=True) embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Vector(3))) + half_embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Halfvec(3))) + sparse_embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Sparsevec(3))) SQLModel.metadata.drop_all(engine) @@ -40,7 +42,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, sparse_embedding=SparseVec.from_dense(v))) session.commit() From c662eac0f5f0a112a6712293bd1b9d3676ea391c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:47:03 -0400 Subject: [PATCH 191/879] Added support for sparsevec to distance functions in Django [skip ci] --- pgvector/django/functions.py | 9 +++++++-- tests/test_django.py | 16 +++++++++++++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 4f696ff..fff2735 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,5 +1,5 @@ from django.db.models import FloatField, Func, Value -from ..utils import Vector +from ..utils import Vector, HalfVec, SparseVec class DistanceBase(Func): @@ -7,7 +7,12 @@ class DistanceBase(Func): def __init__(self, expression, vector, **extra): if not hasattr(vector, 'resolve_expression'): - vector = Value(Vector.to_db(vector)) + if isinstance(vector, HalfVec): + vector = Value(HalfVec.to_db(vector)) + elif isinstance(vector, SparseVec): + vector = Value(SparseVec.to_db(vector)) + else: + vector = Value(Vector.to_db(vector)) super().__init__(expression, vector, **extra) diff --git a/tests/test_django.py b/tests/test_django.py index 271c2cc..1c33b43 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -92,7 +92,7 @@ def create_items(): [1, 1, 2] ] for i, v in enumerate(vectors): - item = Item(id=i + 1, embedding=v) + item = Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v)) item.save() @@ -156,6 +156,20 @@ def test_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_halfvec_l2_distance(self): + create_items() + distance = L2Distance('half_embedding', [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)] + + def test_sparsevec_l2_distance(self): + create_items() + distance = L2Distance('sparse_embedding', SparseVec.from_dense([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)] + def test_filter(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) From 82c976f96c1bbfdbbb78cdc2a014e229b9e51b2d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:48:45 -0400 Subject: [PATCH 192/879] Improved code and test [skip ci] --- pgvector/django/__init__.py | 17 +++++++++++++++-- tests/test_django.py | 4 ++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 2281ed9..f03fad3 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -4,6 +4,19 @@ from .indexes import IvfflatIndex, HnswIndex from .sparsevec import SparsevecField from .vector import VectorField -from ..utils import SparseVec +from ..utils import HalfVec, SparseVec -__all__ = ['VectorExtension', 'VectorField', 'HalfvecField', 'SparsevecField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance'] +__all__ = [ + 'VectorExtension', + 'VectorField', + 'HalfvecField', + 'SparsevecField', + 'IvfflatIndex', + 'HnswIndex', + 'L2Distance', + 'MaxInnerProduct', + 'CosineDistance', + 'L1Distance', + 'HalfVec', + 'SparseVec' +] diff --git a/tests/test_django.py b/tests/test_django.py index 1c33b43..98e4adf 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfvecField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, SparseVec +from pgvector.django import VectorExtension, VectorField, HalfvecField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HalfVec, SparseVec from unittest import mock settings.configure( @@ -158,7 +158,7 @@ def test_l1_distance(self): def test_halfvec_l2_distance(self): create_items() - distance = L2Distance('half_embedding', [1, 1, 1]) + distance = L2Distance('half_embedding', HalfVec([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)] From f399d3f99519acd489e466af76fd241698594b83 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:53:21 -0400 Subject: [PATCH 193/879] Improved typing in example and test [skip ci] --- README.md | 2 +- tests/test_sqlmodel.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 19acac5..28703a6 100644 --- a/README.md +++ b/README.md @@ -220,7 +220,7 @@ from pgvector.sqlalchemy import Vector from sqlalchemy import Column class Item(SQLModel, table=True): - embedding: List[float] = Field(sa_column=Column(Vector(3))) + embedding: Any = Field(sa_column=Column(Vector(3))) ``` Insert a vector diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 2fe2930..777b6bc 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -5,7 +5,7 @@ from sqlalchemy.exc import StatementError from sqlalchemy.sql import func from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text -from typing import List, Optional +from typing import Any, List, Optional engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: @@ -16,9 +16,9 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Vector(3))) - half_embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Halfvec(3))) - sparse_embedding: Optional[List[float]] = Field(default=None, sa_column=Column(Sparsevec(3))) + embedding: Optional[Any] = Field(default=None, sa_column=Column(Vector(3))) + half_embedding: Optional[Any] = Field(default=None, sa_column=Column(Halfvec(3))) + sparse_embedding: Optional[Any] = Field(default=None, sa_column=Column(Sparsevec(3))) SQLModel.metadata.drop_all(engine) From 66af6db211f695271f95f5868b1ed86d15f0290c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 18:53:50 -0400 Subject: [PATCH 194/879] Removed unused import [skip ci] --- tests/test_sqlmodel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 777b6bc..2c81715 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -5,7 +5,7 @@ from sqlalchemy.exc import StatementError from sqlalchemy.sql import func from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text -from typing import Any, List, Optional +from typing import Any, Optional engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: From 91900d6150480b36806ebc5a83391524a1d9d8df Mon Sep 17 00:00:00 2001 From: Cdingram Date: Thu, 16 May 2024 17:19:48 -0600 Subject: [PATCH 195/879] Fix for blank django form fields when VectorField is not required (#68) --- pgvector/django/vector.py | 5 +++++ tests/test_django.py | 8 ++++++++ 2 files changed, 13 insertions(+) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 063b5ff..f6ef0f2 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -66,3 +66,8 @@ def has_changed(self, initial, data): if isinstance(initial, np.ndarray): initial = initial.tolist() return super().has_changed(initial, data) + + def to_python(self, value): + if isinstance(value, str) and value == '': + return None + return super().to_python(value) diff --git a/tests/test_django.py b/tests/test_django.py index 98e4adf..24be5f9 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -222,6 +222,14 @@ def test_form_save(self): assert form.save() assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() + def test_form_save_missing(self): + Item(id=1).save() + item = Item.objects.get(pk=1) + form = ItemForm(instance=item, data={'embedding': ''}) + assert form.is_valid() + assert form.save() + assert Item.objects.get(pk=1).embedding is None + def test_clean(self): item = Item(id=1, embedding=[1, 2, 3]) item.full_clean() From 9a74ba38eedf03c373474cc371df957724dcbc5f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 19:21:04 -0400 Subject: [PATCH 196/879] Improved field descriptions for Django [skip ci] --- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index ec08a42..79103ad 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -4,7 +4,7 @@ # https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ class HalfvecField(Field): - description = 'Halfvec' + description = 'Half vector' empty_strings_allowed = False def __init__(self, *args, dimensions=None, **kwargs): diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 91b9cf8..1f39981 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -4,7 +4,7 @@ # https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ class SparsevecField(Field): - description = 'Sparsevec' + description = 'Sparse vector' empty_strings_allowed = False def __init__(self, *args, dimensions=None, **kwargs): From e583e442539b4185960d99231bdb14c127c988db Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 19:25:35 -0400 Subject: [PATCH 197/879] Updated URLs [skip ci] --- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 79103ad..2fa9e52 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -2,7 +2,7 @@ from ..utils import HalfVec -# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ class HalfvecField(Field): description = 'Half vector' empty_strings_allowed = False diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 1f39981..b7f4c79 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -2,7 +2,7 @@ from ..utils import SparseVec -# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ class SparsevecField(Field): description = 'Sparse vector' empty_strings_allowed = False diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index f6ef0f2..30fd99f 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -4,7 +4,7 @@ from ..utils import Vector -# https://docs.djangoproject.com/en/4.2/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ class VectorField(Field): description = 'Vector' empty_strings_allowed = False From bf53f41e8caf077e2dab65eb3e20c186f9794994 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 19:38:05 -0400 Subject: [PATCH 198/879] Added support for bit type to Django [skip ci] --- CHANGELOG.md | 4 ++-- pgvector/django/__init__.py | 6 +++++- pgvector/django/bit.py | 21 +++++++++++++++++++++ pgvector/django/functions.py | 19 +++++++++++++++++++ tests/test_django.py | 22 +++++++++++++++++++++- 5 files changed, 68 insertions(+), 4 deletions(-) create mode 100644 pgvector/django/bit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index eb6fab7..69be6fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,12 @@ ## 0.3.0 (unreleased) -- Added support for `halfvec` and `sparsevec` types to Django +- Added support for `halfvec`, `bit`, and `sparsevec` types to Django - 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 - Added support for `halfvec` and `sparsevec` types to Peewee -- Added `L1Distance` for Django +- Added `L1Distance`, `HammingDistance`, and `JaccardDistance` for Django - Added `l1_distance` for SQLAlchemy, SQLModel, and Peewee ## 0.2.5 (2024-02-07) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index f03fad3..677d77e 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,5 +1,6 @@ +from .bit import BitField from .extensions import VectorExtension -from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance +from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from .halfvec import HalfvecField from .indexes import IvfflatIndex, HnswIndex from .sparsevec import SparsevecField @@ -10,6 +11,7 @@ 'VectorExtension', 'VectorField', 'HalfvecField', + 'BitField', 'SparsevecField', 'IvfflatIndex', 'HnswIndex', @@ -17,6 +19,8 @@ 'MaxInnerProduct', 'CosineDistance', 'L1Distance', + 'HammingDistance', + 'JaccardDistance', 'HalfVec', 'SparseVec' ] diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py new file mode 100644 index 0000000..941d694 --- /dev/null +++ b/pgvector/django/bit.py @@ -0,0 +1,21 @@ +from django.db.models import Field + + +# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +class BitField(Field): + description = 'Bit string' + + def __init__(self, *args, length=None, **kwargs): + self.length = length + super().__init__(*args, **kwargs) + + def deconstruct(self): + name, path, args, kwargs = super().deconstruct() + if self.length is not None: + kwargs['length'] = self.length + return name, path, args, kwargs + + def db_type(self, connection): + if self.length is None: + return 'bit' + return 'bit(%d)' % self.length diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index fff2735..fe6f2c5 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -16,6 +16,15 @@ def __init__(self, expression, vector, **extra): super().__init__(expression, vector, **extra) +class BitDistanceBase(Func): + output_field = FloatField() + + def __init__(self, expression, vector, **extra): + if not hasattr(vector, 'resolve_expression'): + vector = Value(vector) + super().__init__(expression, vector, **extra) + + class L2Distance(DistanceBase): function = '' arg_joiner = ' <-> ' @@ -34,3 +43,13 @@ class CosineDistance(DistanceBase): class L1Distance(DistanceBase): function = '' arg_joiner = ' <+> ' + + +class HammingDistance(BitDistanceBase): + function = '' + arg_joiner = ' <~> ' + + +class JaccardDistance(BitDistanceBase): + function = '' + arg_joiner = ' <%%> ' diff --git a/tests/test_django.py b/tests/test_django.py index 24be5f9..6dc91e0 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfvecField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HalfVec, SparseVec +from pgvector.django import VectorExtension, VectorField, HalfvecField, BitField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVec, SparseVec from unittest import mock settings.configure( @@ -25,6 +25,7 @@ class Item(models.Model): embedding = VectorField(dimensions=3, null=True, blank=True) half_embedding = HalfvecField(dimensions=3, null=True, blank=True) + binary_embedding = BitField(length=3, null=True, blank=True) sparse_embedding = SparsevecField(dimensions=3, null=True, blank=True) class Meta: @@ -60,6 +61,7 @@ class Migration(migrations.Migration): ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), ('embedding', pgvector.django.VectorField(dimensions=3, null=True, blank=True)), ('half_embedding', pgvector.django.HalfvecField(dimensions=3, null=True, blank=True)), + ('binary_embedding', pgvector.django.BitField(length=3, null=True, blank=True)), ('sparse_embedding', pgvector.django.SparsevecField(dimensions=3, null=True, blank=True)), ], ), @@ -170,6 +172,24 @@ def test_sparsevec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_bit_hamming_distance(self): + Item(id=1, binary_embedding='000').save() + Item(id=2, binary_embedding='101').save() + Item(id=3, binary_embedding='111').save() + distance = HammingDistance('binary_embedding', '101') + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [0, 1, 2] + + def test_bit_jaccard_distance(self): + Item(id=1, binary_embedding='000').save() + Item(id=2, binary_embedding='101').save() + Item(id=3, binary_embedding='111').save() + distance = JaccardDistance('binary_embedding', '101') + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [2, 3, 1] + # assert [v.distance for v in items] == [0, 1/3, 1] + def test_filter(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) From 054eb4aed441d2752497531da08f668cf1c32325 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 19:49:48 -0400 Subject: [PATCH 199/879] Added support for bit type to SQLAlchemy and SQLModel [skip ci] --- CHANGELOG.md | 5 +++-- pgvector/sqlalchemy/__init__.py | 3 ++- pgvector/sqlalchemy/bit.py | 26 ++++++++++++++++++++++++++ tests/test_sqlalchemy.py | 28 ++++++++++++++++++++++++++-- 4 files changed, 57 insertions(+), 5 deletions(-) create mode 100644 pgvector/sqlalchemy/bit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 69be6fd..74a0f2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,13 +1,14 @@ ## 0.3.0 (unreleased) - Added support for `halfvec`, `bit`, and `sparsevec` types to Django -- Added support for `halfvec` and `sparsevec` types to SQLAlchemy and SQLModel +- Added support for `halfvec`, `bit`, 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 - Added support for `halfvec` and `sparsevec` types to Peewee - Added `L1Distance`, `HammingDistance`, and `JaccardDistance` for Django -- Added `l1_distance` for SQLAlchemy, SQLModel, and Peewee +- Added `l1_distance`, `hamming_distance`, and `jaccard_distance` for SQLAlchemy and SQLModel +- Added `l1_distance` for Peewee ## 0.2.5 (2024-02-07) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index a7db132..6e5a206 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,6 +1,7 @@ +from .bit import Bit from .halfvec import Halfvec from .sparsevec import Sparsevec from .vector import Vector from ..utils import SparseVec -__all__ = ['Vector', 'Halfvec', 'Sparsevec'] +__all__ = ['Vector', 'Halfvec', 'Bit', 'Sparsevec'] diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py new file mode 100644 index 0000000..f71c1d0 --- /dev/null +++ b/pgvector/sqlalchemy/bit.py @@ -0,0 +1,26 @@ +from sqlalchemy.dialects.postgresql.base import ischema_names +from sqlalchemy.types import UserDefinedType, Float + + +class Bit(UserDefinedType): + cache_ok = True + + def __init__(self, length=None): + super(UserDefinedType, self).__init__() + self.length = length + + def get_col_spec(self, **kw): + if self.length is None: + return 'BIT' + return 'BIT(%d)' % self.length + + class comparator_factory(UserDefinedType.Comparator): + def hamming_distance(self, other): + return self.op('<~>', return_type=Float)(other) + + def jaccard_distance(self, other): + return self.op('<%>', return_type=Float)(other) + + +# for reflection +ischema_names['bit'] = Bit diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index e6b869a..0003be4 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.sqlalchemy import Vector, Halfvec, Sparsevec, SparseVec +from pgvector.sqlalchemy import Vector, Halfvec, Bit, 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)) + binary_embedding = mapped_column(Bit(3)) sparse_embedding = mapped_column(Sparsevec(3)) @@ -62,7 +63,10 @@ def test_core(self): 'core_item', metadata, Column('id', Integer, primary_key=True), - Column('embedding', Vector(3)) + Column('embedding', Vector(3)), + Column('half_embedding', Halfvec(3)), + Column('binary_embedding', Bit(3)), + Column('sparse_embedding', Sparsevec(3)) ) metadata.drop_all(engine) @@ -157,6 +161,26 @@ def test_l1_distance_orm(self): items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_bit_hamming_distance(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='000')) + session.add(Item(id=2, binary_embedding='101')) + session.add(Item(id=3, binary_embedding='111')) + session.commit() + with Session(engine) as session: + items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() + assert [v.id for v in items] == [2, 3, 1] + + def test_bit_jaccard_distance(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='000')) + session.add(Item(id=2, binary_embedding='101')) + session.add(Item(id=3, binary_embedding='111')) + session.commit() + with Session(engine) as session: + items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() + assert [v.id for v in items] == [2, 3, 1] + def test_filter(self): create_items() with Session(engine) as session: From 29781ed5801fbdda121f3ab8c51bea9614131886 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:25:30 -0400 Subject: [PATCH 200/879] Added support for bit type to Peewee [skip ci] --- CHANGELOG.md | 2 +- pgvector/peewee/__init__.py | 12 +++++++++++- pgvector/peewee/bit.py | 21 +++++++++++++++++++++ tests/test_peewee.py | 21 ++++++++++++++++++++- 4 files changed, 53 insertions(+), 3 deletions(-) create mode 100644 pgvector/peewee/bit.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 74a0f2a..a37fb09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ - 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 -- Added support for `halfvec` and `sparsevec` types to Peewee +- Added support for `halfvec`, `bit`, and `sparsevec` types to Peewee - Added `L1Distance`, `HammingDistance`, and `JaccardDistance` for Django - Added `l1_distance`, `hamming_distance`, and `jaccard_distance` for SQLAlchemy and SQLModel - Added `l1_distance` for Peewee diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 665c385..ef0df67 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -1,4 +1,14 @@ -from ..utils import SparseVec +from .bit import FixedBitField from .halfvec import HalfvecField from .sparsevec import SparsevecField from .vector import VectorField +from ..utils import HalfVec, SparseVec + +__all__ = [ + 'VectorField', + 'HalfvecField', + 'FixedBitField', + 'SparsevecField', + 'HalfVec', + 'SparseVec' +] diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py new file mode 100644 index 0000000..8a3bfbc --- /dev/null +++ b/pgvector/peewee/bit.py @@ -0,0 +1,21 @@ +from peewee import Expression, Field, Value + + +class FixedBitField(Field): + field_type = 'bit' + + def __init__(self, max_length=None, *args, **kwargs): + self.max_length = max_length + super(FixedBitField, self).__init__(*args, **kwargs) + + def get_modifiers(self): + return self.max_length and [self.max_length] or None + + def _distance(self, op, vector): + return Expression(lhs=self, op=op, rhs=self.to_value(vector)) + + def hamming_distance(self, vector): + return self._distance('<~>', vector) + + def jaccard_distance(self, vector): + return self._distance('<%%>', vector) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index e429f8c..7b073e9 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField, HalfvecField, SparsevecField, SparseVec +from pgvector.peewee import VectorField, HalfvecField, FixedBitField, SparsevecField, SparseVec db = PostgresqlDatabase('pgvector_python_test') @@ -14,6 +14,7 @@ class Meta: class Item(BaseModel): embedding = VectorField(dimensions=3, null=True) half_embedding = HalfvecField(dimensions=3, null=True) + binary_embedding = FixedBitField(max_length=3, null=True) sparse_embedding = SparsevecField(dimensions=3, null=True) @@ -87,6 +88,24 @@ def test_sparsevec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_bit_hamming_distance(self): + Item.create(id=1, binary_embedding='000') + Item.create(id=2, binary_embedding='101') + Item.create(id=3, binary_embedding='111') + distance = Item.binary_embedding.hamming_distance('101') + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [0, 1, 2] + + def test_bit_jaccard_distance(self): + Item.create(id=1, binary_embedding='000') + Item.create(id=2, binary_embedding='101') + Item.create(id=3, binary_embedding='111') + distance = Item.binary_embedding.jaccard_distance('101') + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + # assert [v.distance for v in items] == [0, 1/3, 1] + def test_where(self): create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) From 70a88bc4e59b90e94d1bc83b066be3554b57ae59 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:26:12 -0400 Subject: [PATCH 201/879] Updated changelog [skip ci] --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a37fb09..14db629 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ - Added support for `halfvec`, `bit`, and `sparsevec` types to Peewee - Added `L1Distance`, `HammingDistance`, and `JaccardDistance` for Django - Added `l1_distance`, `hamming_distance`, and `jaccard_distance` for SQLAlchemy and SQLModel -- Added `l1_distance` for Peewee +- Added `l1_distance`, `hamming_distance`, and `jaccard_distance` for Peewee ## 0.2.5 (2024-02-07) From 66467240d5f47fbc2a6af755472453e6062e2286 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:32:42 -0400 Subject: [PATCH 202/879] Added test for bit with asyncpg [skip ci] --- tests/test_asyncpg.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 28cbacd..bc7a6e4 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -55,6 +55,30 @@ async def test_halfvec(self): await conn.close() + @pytest.mark.asyncio + async def test_bit(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 items') + await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding bit(3))') + + await register_vector(conn) + + embedding = asyncpg.BitString.from_int(5, 3) + await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + + res = await conn.fetch("SELECT * FROM items ORDER BY id") + assert res[0]['id'] == 1 + assert res[1]['id'] == 2 + assert res[0]['embedding'].to_int() == 5 + assert res[1]['embedding'] is None + + # ensures binary format is correct + text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + assert text_res[0]['embedding'] == '101' + + await conn.close() + @pytest.mark.asyncio async def test_sparsevec(self): conn = await asyncpg.connect(database='pgvector_python_test') From 9f8cf3b7b17ec4b9646e4ff35dd871dffbcda43c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:34:17 -0400 Subject: [PATCH 203/879] Improved test [skip ci] --- tests/test_asyncpg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index bc7a6e4..06187ef 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -64,7 +64,7 @@ async def test_bit(self): await register_vector(conn) - embedding = asyncpg.BitString.from_int(5, 3) + embedding = asyncpg.BitString.from_int(5, length=3) await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM items ORDER BY id") From 0275ac502a21683b90dfe8bf4f6a62ce5701db90 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:38:03 -0400 Subject: [PATCH 204/879] Added test for bit with Psycopg 2 [skip ci] --- tests/test_psycopg2.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index b5ac50a..7302296 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -8,7 +8,7 @@ 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), 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))') register_vector(cur) @@ -36,6 +36,15 @@ def test_halfvec(self): assert res[0][0].to_list() == [1.5, 2, 3] assert res[1][0] is None + def test_bit(self): + embedding = '101' + cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT binary_embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][0] == '101' + assert res[1][0] is None + def test_sparsevec(self): embedding = SparseVec.from_dense([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) From 8cd68b16f7cc0f29cde08ae4878ef49822180466 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:40:47 -0400 Subject: [PATCH 205/879] Improved tests for bit with Psycopg 3 [skip ci] --- tests/test_psycopg.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 9a7f611..0fa0789 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -114,6 +114,14 @@ def test_bit(self): res = conn.execute('SELECT %s::bit(3)', ('101',)).fetchone()[0] assert res == '101' + def test_bit_binary_format(self): + res = conn.execute('SELECT %b::bit(3)', ('101',), binary=True).fetchone()[0] + assert res == b'\x00\x00\x00\x03\xa0' + + def test_bit_text_format(self): + res = conn.execute('SELECT %t::bit(3)', ('101',)).fetchone()[0] + assert res == '101' + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) From ca5036dbc73b509a4b5f9c43cbf2e78d1e24f21b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:48:10 -0400 Subject: [PATCH 206/879] Improved test names [skip ci] --- tests/test_django.py | 8 ++++---- tests/test_peewee.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 6dc91e0..a32a774 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -130,28 +130,28 @@ def test_sparsevec(self): assert item.id == 1 assert item.sparse_embedding.to_dense() == [1, 2, 3] - def test_l2_distance(self): + def test_vector_l2_distance(self): create_items() distance = L2Distance('embedding', [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)] - def test_max_inner_product(self): + def test_vector_max_inner_product(self): create_items() distance = MaxInnerProduct('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_cosine_distance(self): + def test_vector_cosine_distance(self): create_items() distance = CosineDistance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_l1_distance(self): + def test_vector_l1_distance(self): create_items() distance = L1Distance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 7b073e9..9b21c0d 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -40,34 +40,34 @@ class TestPeewee: def setup_method(self, test_method): Item.truncate_table() - def test_works(self): + def test_vector(self): Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) assert np.array_equal(item.embedding, np.array([1, 2, 3])) assert item.embedding.dtype == np.float32 - def test_l2_distance(self): + def test_vector_l2_distance(self): create_items() distance = Item.embedding.l2_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_max_inner_product(self): + def test_vector_max_inner_product(self): create_items() distance = Item.embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_cosine_distance(self): + def test_vector_cosine_distance(self): create_items() distance = Item.embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_l1_distance(self): + def test_vector_l1_distance(self): create_items() distance = Item.embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) From ebe99b511d9eccf3fae5dd6195b95484818e4fb9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:51:29 -0400 Subject: [PATCH 207/879] Added support for NumPy arrays to HalfVec [skip ci] --- pgvector/utils/halfvec.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 8b760cd..75fca83 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -1,9 +1,12 @@ +import numpy as np from struct import pack, unpack_from class HalfVec: def __init__(self, value): - # TODO support np.array + if isinstance(value, np.ndarray): + value = value.tolist() + if not isinstance(value, (list, tuple)): raise ValueError('expected list or tuple') From 89b9b85a1c912d9f0562cf83153452231bc71a45 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:54:06 -0400 Subject: [PATCH 208/879] Improved Psycopg 3 tests [skip ci] --- tests/test_psycopg.py | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 0fa0789..306b222 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -6,47 +6,47 @@ conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') -conn.execute('DROP TABLE IF EXISTS items') -conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +conn.execute('DROP TABLE IF EXISTS psycopg_items') +conn.execute('CREATE TABLE psycopg_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))') register_vector(conn) class TestPsycopg: def setup_method(self, test_method): - conn.execute('DELETE FROM items') + conn.execute('DELETE FROM psycopg_items') - def test_works(self): + def test_vector(self): embedding = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO items (embedding) VALUES (%s), (NULL)', (embedding,)) + conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) - res = conn.execute('SELECT * FROM items ORDER BY id').fetchall() + res = conn.execute('SELECT * FROM psycopg_items ORDER BY id').fetchall() assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None - def test_binary_format(self): + def test_vector_binary_format(self): embedding = np.array([1.5, 2, 3]) res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] assert np.array_equal(res, embedding) - def test_text_format(self): + def test_vector_text_format(self): embedding = np.array([1.5, 2, 3]) res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, embedding) - def test_binary_format_correct(self): + def test_vector_binary_format_correct(self): embedding = np.array([1.5, 2, 3]) res = conn.execute('SELECT %b::vector::text', (embedding,)).fetchone()[0] assert res == '[1.5,2,3]' - def test_text_format_non_contiguous(self): + def test_vector_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, np.array([3, 2, 1.5])) - def test_binary_format_non_contiguous(self): + def test_vector_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] @@ -55,19 +55,19 @@ def test_binary_format_non_contiguous(self): def test_text_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY items (embedding) FROM STDIN") as copy: - copy.write_row([embedding]) + with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: + copy.write_row([embedding, HalfVec(embedding), '101', SparseVec.from_dense(embedding)]) def test_binary_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy("COPY psycopg_items (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.write_row([embedding]) def test_binary_copy_set_types(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY items (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy("COPY psycopg_items (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector']) copy.write_row([1, embedding]) @@ -127,16 +127,16 @@ async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') - await conn.execute('DROP TABLE IF EXISTS items') - await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') + await conn.execute('DROP TABLE IF EXISTS psycopg_items') + await conn.execute('CREATE TABLE psycopg_items (id bigserial PRIMARY KEY, embedding vector(3))') await register_vector_async(conn) embedding = np.array([1.5, 2, 3]) - await conn.execute('INSERT INTO items (embedding) VALUES (%s), (NULL)', (embedding,)) + await conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) async with conn.cursor() as cur: - await cur.execute('SELECT * FROM items ORDER BY id') + await cur.execute('SELECT * FROM psycopg_items ORDER BY id') res = await cur.fetchall() assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 From cb755fa5c6728e293512f8a6787ee4e325968aa3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 20:55:10 -0400 Subject: [PATCH 209/879] Improved test [skip ci] --- tests/test_psycopg.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 306b222..f674a66 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -67,9 +67,9 @@ def test_binary_copy(self): def test_binary_copy_set_types(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY psycopg_items (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.set_types(['int8', 'vector']) - copy.write_row([1, embedding]) + with cur.copy("COPY psycopg_items (id, embedding, half_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.set_types(['int8', 'vector', 'halfvec', 'sparsevec']) + copy.write_row([1, embedding, HalfVec(embedding), SparseVec.from_dense(embedding)]) def test_halfvec(self): conn.execute('DROP TABLE IF EXISTS half_items') From e1843c20722b89547b8ac44857271439796ead95 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:05:49 -0400 Subject: [PATCH 210/879] Added tests for bit for SQLModel --- tests/test_sqlmodel.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 2c81715..d0a6ccd 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.sqlalchemy import Vector, Halfvec, Sparsevec, SparseVec +from pgvector.sqlalchemy import Vector, Halfvec, Bit, Sparsevec, SparseVec import pytest from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError @@ -18,6 +18,7 @@ class Item(SQLModel, table=True): id: Optional[int] = Field(default=None, primary_key=True) embedding: Optional[Any] = Field(default=None, sa_column=Column(Vector(3))) half_embedding: Optional[Any] = Field(default=None, sa_column=Column(Halfvec(3))) + binary_embedding: Optional[Any] = Field(default=None, sa_column=Column(Bit(3))) sparse_embedding: Optional[Any] = Field(default=None, sa_column=Column(Sparsevec(3))) @@ -99,6 +100,26 @@ def test_l1_distance(self): items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_bit_hamming_distance(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='000')) + session.add(Item(id=2, binary_embedding='101')) + session.add(Item(id=3, binary_embedding='111')) + session.commit() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + + def test_bit_jaccard_distance(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='000')) + session.add(Item(id=2, binary_embedding='101')) + session.add(Item(id=3, binary_embedding='111')) + session.commit() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + def test_filter(self): create_items() with Session(engine) as session: From ca571f0607d80ead8306102b2b824f03b3225d17 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:15:22 -0400 Subject: [PATCH 211/879] Added support for passing dense vectors to sparsevec columns and functions [skip ci] --- pgvector/utils/sparsevec.py | 10 ++++++++++ tests/test_sqlmodel.py | 20 ++++++++++++++++---- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 9aae64e..afcef65 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -1,3 +1,4 @@ +import numpy as np from struct import pack, unpack_from @@ -8,6 +9,8 @@ def __init__(self, dim, indices, values): self.values = values def from_dense(value): + if isinstance(value, np.ndarray): + value = value.tolist() dim = len(value) indices = [i for i, v in enumerate(value) if v != 0] values = [value[i] for i in indices] @@ -23,6 +26,9 @@ def to_db(value, dim=None): if value is None: return value + if isinstance(value, (list, np.ndarray)): + value = SparseVec.from_dense(value) + if dim is not None and value.dim != dim: raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) @@ -31,6 +37,10 @@ def to_db(value, dim=None): def to_db_binary(value): if value is None: return value + + if isinstance(value, (list, np.ndarray)): + value = SparseVec.from_dense(value) + nnz = len(value.indices) return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index d0a6ccd..068d801 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -76,30 +76,42 @@ def test_orm(self): assert items[1].embedding.dtype == np.float32 assert items[2].embedding is None - def test_l2_distance(self): + def test_vector_l2_distance(self): create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_max_inner_product(self): + def test_vector_max_inner_product(self): create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_cosine_distance(self): + def test_vector_cosine_distance(self): create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_l1_distance(self): + def test_vector_l1_distance(self): create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_halfvec_l2_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + + def test_sparsevec_l2_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_bit_hamming_distance(self): session = Session(engine) session.add(Item(id=1, binary_embedding='000')) From 92cc17fb4c0f09614aed6198bd309f57ca6bfa1a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:22:13 -0400 Subject: [PATCH 212/879] Improved SparseVec code [skip ci] --- pgvector/utils/sparsevec.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index afcef65..040ef2c 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -2,6 +2,15 @@ from struct import pack, unpack_from +def to_db_value(value): + if isinstance(value, SparseVec): + return value + elif isinstance(value, (list, np.ndarray)): + return SparseVec.from_dense(value) + else: + raise ValueError('expected sparsevec') + + class SparseVec: def __init__(self, dim, indices, values): self.dim = dim @@ -26,11 +35,10 @@ def to_db(value, dim=None): if value is None: return value - if isinstance(value, (list, np.ndarray)): - value = SparseVec.from_dense(value) + value = to_db_value(value) if dim is not None and value.dim != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) + raise ValueError('expected %d dimensions, not %d' % (dim, value.dim)) return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(value.indices, value.values)]) + '}/' + str(value.dim) @@ -38,9 +46,7 @@ def to_db_binary(value): if value is None: return value - if isinstance(value, (list, np.ndarray)): - value = SparseVec.from_dense(value) - + value = to_db_value(value) nnz = len(value.indices) return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) From 62e5befc993000b6e632c2f73912f051b236a7cd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:28:07 -0400 Subject: [PATCH 213/879] Added more tests for SQLModel [skip ci] --- tests/test_sqlmodel.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 068d801..1a5b7de 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -106,12 +106,48 @@ def test_halfvec_l2_distance(self): items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_halfvec_max_inner_product(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) + assert [v.id for v in items] == [2, 3, 1] + + def test_halfvec_cosine_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 2, 3] + + def test_halfvec_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_sparsevec_l2_distance(self): create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_sparsevec_max_inner_product(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) + assert [v.id for v in items] == [2, 3, 1] + + def test_sparsevec_cosine_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 2, 3] + + def test_sparsevec_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_bit_hamming_distance(self): session = Session(engine) session.add(Item(id=1, binary_embedding='000')) From 705f27051eda69077ee54fa300e4b367ab73e406 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:30:31 -0400 Subject: [PATCH 214/879] Added more tests for SQLModel [skip ci] --- tests/test_sqlmodel.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 1a5b7de..58ac2c9 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -180,7 +180,7 @@ def test_select(self): items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() assert items[0] == 3 - def test_avg(self): + def test_vector_avg(self): with Session(engine) as session: avg = session.exec(select(func.avg(Item.embedding))).first() assert avg is None @@ -189,7 +189,7 @@ def test_avg(self): avg = session.exec(select(func.avg(Item.embedding))).first() assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) - def test_sum(self): + def test_vector_sum(self): with Session(engine) as session: sum = session.exec(select(func.sum(Item.embedding))).first() assert sum is None @@ -198,6 +198,24 @@ def test_sum(self): sum = session.exec(select(func.sum(Item.embedding))).first() assert np.array_equal(sum, 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 + 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() + assert avg.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 + 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] + def test_bad_dimensions(self): item = Item(embedding=[1, 2]) session = Session(engine) From 34c2f92985e37b142232a49e73808dd76eb69d10 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:32:21 -0400 Subject: [PATCH 215/879] Added more tests for SQLAlchemy [skip ci] --- tests/test_sqlalchemy.py | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0003be4..887859c 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -113,54 +113,66 @@ def test_orm(self): assert items[1].embedding.dtype == np.float32 assert items[2].embedding is None - def test_l2_distance(self): + def test_vector_l2_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_l2_distance_orm(self): + def test_vector_l2_distance_orm(self): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_max_inner_product(self): + def test_vector_max_inner_product(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_max_inner_product_orm(self): + def test_vector_max_inner_product_orm(self): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_cosine_distance(self): + def test_vector_cosine_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_cosine_distance_orm(self): + def test_vector_cosine_distance_orm(self): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_l1_distance(self): + def test_vector_l1_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_l1_distance_orm(self): + def test_vector_l1_distance_orm(self): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_halfvec_l2_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_sparsevec_l2_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + def test_bit_hamming_distance(self): session = Session(engine) session.add(Item(id=1, binary_embedding='000')) From 637f318b4a58f1182e488217957cd817bab6c9b5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:36:18 -0400 Subject: [PATCH 216/879] Added more tests for Peewee [skip ci] --- tests/test_peewee.py | 62 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 60 insertions(+), 2 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 9b21c0d..c9bd085 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -81,6 +81,27 @@ def test_halfvec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_halfvec_max_inner_product(self): + create_items() + distance = Item.half_embedding.max_inner_product([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [-6, -4, -3] + + def test_halfvec_cosine_distance(self): + create_items() + distance = Item.half_embedding.cosine_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 2, 3] + assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + + def test_halfvec_l1_distance(self): + create_items() + distance = Item.half_embedding.l1_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, 3] + def test_sparsevec_l2_distance(self): create_items() distance = Item.sparse_embedding.l2_distance(SparseVec.from_dense([1, 1, 1])) @@ -88,6 +109,27 @@ def test_sparsevec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_sparsevec_max_inner_product(self): + create_items() + distance = Item.sparse_embedding.max_inner_product([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [-6, -4, -3] + + def test_sparsevec_cosine_distance(self): + create_items() + distance = Item.sparse_embedding.cosine_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 2, 3] + assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + + def test_sparsevec_l1_distance(self): + create_items() + distance = Item.sparse_embedding.l1_distance([1, 1, 1]) + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, 3] + def test_bit_hamming_distance(self): Item.create(id=1, binary_embedding='000') Item.create(id=2, binary_embedding='101') @@ -111,7 +153,7 @@ def test_where(self): items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) assert [v.id for v in items] == [1] - def test_avg(self): + def test_vector_avg(self): avg = Item.select(fn.avg(Item.embedding)).scalar() assert avg is None Item.create(embedding=[1, 2, 3]) @@ -119,7 +161,7 @@ def test_avg(self): avg = Item.select(fn.avg(Item.embedding)).scalar() assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) - def test_sum(self): + def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding)).scalar() assert sum is None Item.create(embedding=[1, 2, 3]) @@ -127,6 +169,22 @@ def test_sum(self): sum = Item.select(fn.sum(Item.embedding)).scalar() assert np.array_equal(sum, np.array([5, 7, 9])) + def test_halfvec_avg(self): + avg = Item.select(fn.avg(Item.half_embedding)).scalar() + assert avg is None + Item.create(half_embedding=[1, 2, 3]) + Item.create(half_embedding=[4, 5, 6]) + avg = Item.select(fn.avg(Item.half_embedding)).scalar() + assert avg.to_list() == [2.5, 3.5, 4.5] + + def test_halfvec_sum(self): + sum = Item.select(fn.sum(Item.half_embedding)).scalar() + assert sum is None + Item.create(half_embedding=[1, 2, 3]) + Item.create(half_embedding=[4, 5, 6]) + sum = Item.select(fn.sum(Item.half_embedding)).scalar() + assert sum.to_list() == [5, 7, 9] + def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) Item.get_or_create(embedding=np.array([4, 5, 6])) From 84c25d843d49d52024df490fbb31e470995583a3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:40:16 -0400 Subject: [PATCH 217/879] Added more tests for Django [skip ci] --- tests/test_django.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index a32a774..2281d8a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -165,6 +165,27 @@ def test_halfvec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_halfvec_max_inner_product(self): + create_items() + distance = MaxInnerProduct('half_embedding', HalfVec([1, 1, 1])) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [-6, -4, -3] + + def test_halfvec_cosine_distance(self): + create_items() + distance = CosineDistance('half_embedding', HalfVec([1, 1, 1])) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [1, 2, 3] + assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + + def test_halfvec_l1_distance(self): + create_items() + distance = L1Distance('half_embedding', HalfVec([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, 3] + def test_sparsevec_l2_distance(self): create_items() distance = L2Distance('sparse_embedding', SparseVec.from_dense([1, 1, 1])) @@ -172,6 +193,27 @@ def test_sparsevec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_sparsevec_max_inner_product(self): + create_items() + distance = MaxInnerProduct('sparse_embedding', SparseVec.from_dense([1, 1, 1])) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [-6, -4, -3] + + def test_sparsevec_cosine_distance(self): + create_items() + distance = CosineDistance('sparse_embedding', SparseVec.from_dense([1, 1, 1])) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [1, 2, 3] + assert [v.distance for v in items] == [0, 0, 0.05719095841793653] + + def test_sparsevec_l1_distance(self): + create_items() + distance = L1Distance('sparse_embedding', SparseVec.from_dense([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, 3] + def test_bit_hamming_distance(self): Item(id=1, binary_embedding='000').save() Item(id=2, binary_embedding='101').save() From ae6384b2b395721b2c9009bb20bc463e522ae70d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:42:01 -0400 Subject: [PATCH 218/879] Added more tests for Django [skip ci] --- tests/test_django.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 2281d8a..12d9c37 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -238,7 +238,7 @@ def test_filter(self): items = Item.objects.alias(distance=distance).filter(distance__lt=1) assert [v.id for v in items] == [1] - def test_avg(self): + def test_vector_avg(self): avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert avg is None Item(embedding=[1, 2, 3]).save() @@ -246,7 +246,7 @@ def test_avg(self): avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) - def test_sum(self): + def test_vector_sum(self): sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert sum is None Item(embedding=[1, 2, 3]).save() @@ -254,6 +254,22 @@ def test_sum(self): sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert np.array_equal(sum, np.array([5, 7, 9])) + def test_halfvec_avg(self): + avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] + assert avg is None + Item(half_embedding=[1, 2, 3]).save() + Item(half_embedding=[4, 5, 6]).save() + avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] + assert avg.to_list() == [2.5, 3.5, 4.5] + + def test_halfvec_sum(self): + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] + assert sum is None + Item(half_embedding=[1, 2, 3]).save() + Item(half_embedding=[4, 5, 6]).save() + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] + assert sum.to_list() == [5, 7, 9] + def test_serialization(self): create_items() items = Item.objects.all() From 16508ecc16977a14953d80e3f69348cf55af1a5e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:44:27 -0400 Subject: [PATCH 219/879] Added bit test for Django [skip ci] --- tests/test_django.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index 12d9c37..a8b1adf 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -123,6 +123,13 @@ def test_halfvec(self): assert item.id == 1 assert item.half_embedding.to_list() == [1, 2, 3] + def test_bit(self): + item = Item(id=1, binary_embedding='101') + item.save() + item = Item.objects.get(pk=1) + assert item.id == 1 + assert item.binary_embedding == '101' + def test_sparsevec(self): item = Item(id=1, sparse_embedding=SparseVec.from_dense([1, 2, 3])) item.save() From 5c8711537e8202bc438dffde15e7d0560367c96f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:48:08 -0400 Subject: [PATCH 220/879] Updated table name for asyncpg tests [skip ci] --- tests/test_asyncpg.py | 48 +++++++++++++++++++++---------------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 06187ef..7405631 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -10,15 +10,15 @@ class TestAsyncpg: async def test_vector(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 items') - await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') await register_vector(conn) embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM items ORDER BY id") + res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert np.array_equal(res[0]['embedding'], embedding) @@ -26,7 +26,7 @@ async def test_vector(self): assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @@ -35,22 +35,22 @@ async def test_vector(self): async def test_halfvec(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 items') - await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding halfvec(3))') + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') await register_vector(conn) embedding = [1.5, 2, 3] - await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM items ORDER BY id") + res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert res[0]['embedding'].to_list() == [1.5, 2, 3] assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @@ -59,22 +59,22 @@ async def test_halfvec(self): async def test_bit(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 items') - await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding bit(3))') + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') await register_vector(conn) embedding = asyncpg.BitString.from_int(5, length=3) - await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM items ORDER BY id") + res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert res[0]['embedding'].to_int() == 5 assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") assert text_res[0]['embedding'] == '101' await conn.close() @@ -83,22 +83,22 @@ async def test_bit(self): async def test_sparsevec(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 items') - await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding sparsevec(3))') + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') await register_vector(conn) embedding = SparseVec.from_dense([1.5, 2, 3]) - await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM items ORDER BY id") + res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert res[0]['embedding'].to_dense() == [1.5, 2, 3] assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM items ORDER BY id LIMIT 1") + text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") assert text_res[0]['embedding'] == '{1:1.5,2:2,3:3}/3' await conn.close() @@ -112,13 +112,13 @@ async def init(conn): async with pool.acquire() as conn: await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') - await conn.execute('DROP TABLE IF EXISTS items') - await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO items (embedding) VALUES ($1), (NULL)", embedding) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) - res = await conn.fetch("SELECT * FROM items ORDER BY id") + res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['id'] == 1 assert res[1]['id'] == 2 assert np.array_equal(res[0]['embedding'], embedding) From a761245b0a967cd6f2d230c7fb4f3c514a822060 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 16 May 2024 21:52:30 -0400 Subject: [PATCH 221/879] Updated table name for Peewee tests [skip ci] --- tests/test_peewee.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index c9bd085..4e7eb44 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -17,6 +17,9 @@ class Item(BaseModel): binary_embedding = FixedBitField(max_length=3, null=True) sparse_embedding = SparsevecField(dimensions=3, null=True) + class Meta: + table_name = 'peewee_item' + Item.add_index('embedding vector_l2_ops', using='hnsw') From 92fe82326e7f6262471a63fa42ef8c4d059b44a1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 12:58:21 -0400 Subject: [PATCH 222/879] Renamed HalfVec to HalfVector and SparseVec to SparseVector --- pgvector/asyncpg/__init__.py | 10 +++++----- pgvector/django/__init__.py | 6 +++--- pgvector/django/functions.py | 10 +++++----- pgvector/django/halfvec.py | 8 ++++---- pgvector/django/sparsevec.py | 8 ++++---- pgvector/peewee/__init__.py | 6 +++--- pgvector/peewee/halfvec.py | 6 +++--- pgvector/peewee/sparsevec.py | 6 +++--- pgvector/psycopg/__init__.py | 2 +- pgvector/psycopg/halfvec.py | 30 +++++++++++++++--------------- pgvector/psycopg/sparsevec.py | 30 +++++++++++++++--------------- pgvector/psycopg2/__init__.py | 2 +- pgvector/psycopg2/halfvec.py | 8 ++++---- pgvector/psycopg2/sparsevec.py | 8 ++++---- pgvector/sqlalchemy/__init__.py | 2 +- pgvector/sqlalchemy/halfvec.py | 8 ++++---- pgvector/sqlalchemy/sparsevec.py | 8 ++++---- pgvector/utils/halfvec.py | 16 ++++++++-------- pgvector/utils/sparsevec.py | 18 +++++++++--------- tests/test_asyncpg.py | 4 ++-- tests/test_django.py | 22 +++++++++++----------- tests/test_peewee.py | 6 +++--- tests/test_psycopg.py | 18 +++++++++--------- tests/test_psycopg2.py | 4 ++-- tests/test_sqlalchemy.py | 4 ++-- tests/test_sqlmodel.py | 4 ++-- 26 files changed, 127 insertions(+), 127 deletions(-) diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 9c6b660..2252c7a 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,4 +1,4 @@ -from ..utils import Vector, HalfVec, SparseVec +from ..utils import Vector, HalfVector, SparseVector __all__ = ['register_vector'] @@ -13,14 +13,14 @@ async def register_vector(conn): await conn.set_type_codec( 'halfvec', - encoder=HalfVec.to_db_binary, - decoder=HalfVec.from_db_binary, + encoder=HalfVector.to_db_binary, + decoder=HalfVector.from_db_binary, format='binary' ) await conn.set_type_codec( 'sparsevec', - encoder=SparseVec.to_db_binary, - decoder=SparseVec.from_db_binary, + encoder=SparseVector.to_db_binary, + decoder=SparseVector.from_db_binary, format='binary' ) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 677d77e..08407c7 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -5,7 +5,7 @@ from .indexes import IvfflatIndex, HnswIndex from .sparsevec import SparsevecField from .vector import VectorField -from ..utils import HalfVec, SparseVec +from ..utils import HalfVector, SparseVector __all__ = [ 'VectorExtension', @@ -21,6 +21,6 @@ 'L1Distance', 'HammingDistance', 'JaccardDistance', - 'HalfVec', - 'SparseVec' + 'HalfVector', + 'SparseVector' ] diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index fe6f2c5..9b11869 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,5 +1,5 @@ from django.db.models import FloatField, Func, Value -from ..utils import Vector, HalfVec, SparseVec +from ..utils import Vector, HalfVector, SparseVector class DistanceBase(Func): @@ -7,10 +7,10 @@ class DistanceBase(Func): def __init__(self, expression, vector, **extra): if not hasattr(vector, 'resolve_expression'): - if isinstance(vector, HalfVec): - vector = Value(HalfVec.to_db(vector)) - elif isinstance(vector, SparseVec): - vector = Value(SparseVec.to_db(vector)) + if isinstance(vector, HalfVector): + vector = Value(HalfVector.to_db(vector)) + elif isinstance(vector, SparseVector): + vector = Value(SparseVector.to_db(vector)) else: vector = Value(Vector.to_db(vector)) super().__init__(expression, vector, **extra) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 2fa9e52..884f706 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -1,5 +1,5 @@ from django.db.models import Field -from ..utils import HalfVec +from ..utils import HalfVector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ @@ -23,13 +23,13 @@ def db_type(self, connection): return 'halfvec(%d)' % self.dimensions def from_db_value(self, value, expression, connection): - return HalfVec.from_db(value) + return HalfVector.from_db(value) def to_python(self, value): - return HalfVec.from_db(value) + return HalfVector.from_db(value) def get_prep_value(self, value): - return HalfVec.to_db(value) + return HalfVector.to_db(value) def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index b7f4c79..1f4c21f 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -1,5 +1,5 @@ from django.db.models import Field -from ..utils import SparseVec +from ..utils import SparseVector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ @@ -23,13 +23,13 @@ def db_type(self, connection): return 'sparsevec(%d)' % self.dimensions def from_db_value(self, value, expression, connection): - return SparseVec.from_db(value) + return SparseVector.from_db(value) def to_python(self, value): - return SparseVec.from_db(value) + return SparseVector.from_db(value) def get_prep_value(self, value): - return SparseVec.to_db(value) + return SparseVector.to_db(value) def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index ef0df67..8f54401 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -2,13 +2,13 @@ from .halfvec import HalfvecField from .sparsevec import SparsevecField from .vector import VectorField -from ..utils import HalfVec, SparseVec +from ..utils import HalfVector, SparseVector __all__ = [ 'VectorField', 'HalfvecField', 'FixedBitField', 'SparsevecField', - 'HalfVec', - 'SparseVec' + 'HalfVector', + 'SparseVector' ] diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 60edd83..e663d72 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,5 +1,5 @@ from peewee import Expression, Field, Value -from ..utils import HalfVec +from ..utils import HalfVector class HalfvecField(Field): @@ -13,10 +13,10 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return HalfVec.to_db(value) + return HalfVector.to_db(value) def python_value(self, value): - return HalfVec.from_db(value) + return HalfVector.from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index c44d4fe..272c9c6 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -1,5 +1,5 @@ from peewee import Expression, Field, Value -from ..utils import SparseVec +from ..utils import SparseVector class SparsevecField(Field): @@ -13,10 +13,10 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return SparseVec.to_db(value) + return SparseVector.to_db(value) def python_value(self, value): - return SparseVec.from_db(value) + return SparseVector.from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index b5df2d1..e47d699 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -3,7 +3,7 @@ from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import HalfVec, SparseVec +from ..utils import HalfVector, SparseVector __all__ = ['register_vector'] diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 023586a..6ca232d 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,53 +1,53 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import HalfVec +from ..utils import HalfVector -class HalfVecDumper(Dumper): +class HalfVectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return HalfVec.to_db(obj).encode('utf8') + return HalfVector.to_db(obj).encode('utf8') -class HalfVecBinaryDumper(HalfVecDumper): +class HalfVectorBinaryDumper(HalfVectorDumper): format = Format.BINARY def dump(self, obj): - return HalfVec.to_db_binary(obj) + return HalfVector.to_db_binary(obj) -class HalfVecLoader(Loader): +class HalfVectorLoader(Loader): format = Format.TEXT def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return HalfVec.from_db(data.decode('utf8')) + return HalfVector.from_db(data.decode('utf8')) -class HalfVecBinaryLoader(HalfVecLoader): +class HalfVectorBinaryLoader(HalfVectorLoader): format = Format.BINARY def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return HalfVec.from_db_binary(data) + return HalfVector.from_db_binary(data) def register_halfvec_info(context, info): info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (HalfVecDumper,), {'oid': info.oid}) - binary_dumper = type('', (HalfVecBinaryDumper,), {'oid': info.oid}) + text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid}) + binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid}) adapters = context.adapters - adapters.register_dumper(HalfVec, text_dumper) - adapters.register_dumper(HalfVec, binary_dumper) - adapters.register_loader(info.oid, HalfVecLoader) - adapters.register_loader(info.oid, HalfVecBinaryLoader) + adapters.register_dumper(HalfVector, text_dumper) + adapters.register_dumper(HalfVector, binary_dumper) + adapters.register_loader(info.oid, HalfVectorLoader) + adapters.register_loader(info.oid, HalfVectorBinaryLoader) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index a5dec21..634d83b 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,53 +1,53 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import SparseVec +from ..utils import SparseVector -class SparseVecDumper(Dumper): +class SparseVectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return SparseVec.to_db(obj).encode('utf8') + return SparseVector.to_db(obj).encode('utf8') -class SparseVecBinaryDumper(SparseVecDumper): +class SparseVectorBinaryDumper(SparseVectorDumper): format = Format.BINARY def dump(self, obj): - return SparseVec.to_db_binary(obj) + return SparseVector.to_db_binary(obj) -class SparseVecLoader(Loader): +class SparseVectorLoader(Loader): format = Format.TEXT def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return SparseVec.from_db(data.decode('utf8')) + return SparseVector.from_db(data.decode('utf8')) -class SparseVecBinaryLoader(SparseVecLoader): +class SparseVectorBinaryLoader(SparseVectorLoader): format = Format.BINARY def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return SparseVec.from_db_binary(data) + return SparseVector.from_db_binary(data) def register_sparsevec_info(context, info): info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (SparseVecDumper,), {'oid': info.oid}) - binary_dumper = type('', (SparseVecBinaryDumper,), {'oid': info.oid}) + text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid}) + binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid}) adapters = context.adapters - adapters.register_dumper(SparseVec, text_dumper) - adapters.register_dumper(SparseVec, binary_dumper) - adapters.register_loader(info.oid, SparseVecLoader) - adapters.register_loader(info.oid, SparseVecBinaryLoader) + adapters.register_dumper(SparseVector, text_dumper) + adapters.register_dumper(SparseVector, binary_dumper) + adapters.register_loader(info.oid, SparseVectorLoader) + adapters.register_loader(info.oid, SparseVectorBinaryLoader) diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 764c0f7..9d0473a 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -2,7 +2,7 @@ from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import SparseVec +from ..utils import SparseVector __all__ = ['register_vector'] diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 8d974b3..e5ec111 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,5 +1,5 @@ from psycopg2.extensions import adapt, new_type, register_adapter, register_type -from ..utils import HalfVec +from ..utils import HalfVector class HalfvecAdapter(object): @@ -7,14 +7,14 @@ def __init__(self, value): self._value = value def getquoted(self): - return adapt(HalfVec.to_db(self._value)).getquoted() + return adapt(HalfVector.to_db(self._value)).getquoted() def cast_halfvec(value, cur): - return HalfVec.from_db(value) + return HalfVector.from_db(value) def register_halfvec_info(oid): halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) register_type(halfvec) - register_adapter(HalfVec, HalfvecAdapter) + register_adapter(HalfVector, HalfvecAdapter) diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 6603b2f..4bcdd92 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,5 +1,5 @@ from psycopg2.extensions import adapt, new_type, register_adapter, register_type -from ..utils import SparseVec +from ..utils import SparseVector class SparsevecAdapter(object): @@ -7,14 +7,14 @@ def __init__(self, value): self._value = value def getquoted(self): - return adapt(SparseVec.to_db(self._value)).getquoted() + return adapt(SparseVector.to_db(self._value)).getquoted() def cast_sparsevec(value, cur): - return SparseVec.from_db(value) + return SparseVector.from_db(value) def register_sparsevec_info(oid): sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) register_type(sparsevec) - register_adapter(SparseVec, SparsevecAdapter) + register_adapter(SparseVector, SparsevecAdapter) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 6e5a206..fde5752 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -2,6 +2,6 @@ from .halfvec import Halfvec from .sparsevec import Sparsevec from .vector import Vector -from ..utils import SparseVec +from ..utils import SparseVector __all__ = ['Vector', 'Halfvec', 'Bit', 'Sparsevec'] diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 694d14f..cafc6c6 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import HalfVec +from ..utils import HalfVector class Halfvec(UserDefinedType): @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return HalfVec.to_db(value, self.dim) + return HalfVector.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 string_literal_processor(HalfVector.to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return HalfVec.from_db(value) + return HalfVector.from_db(value) return process class comparator_factory(UserDefinedType.Comparator): diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 6e3c691..499dff0 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import SparseVec +from ..utils import SparseVector class Sparsevec(UserDefinedType): @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return SparseVec.to_db(value, self.dim) + return SparseVector.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 string_literal_processor(SparseVector.to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return SparseVec.from_db(value) + return SparseVector.from_db(value) return process class comparator_factory(UserDefinedType.Comparator): diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 75fca83..b977291 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -2,7 +2,7 @@ from struct import pack, unpack_from -class HalfVec: +class HalfVector: def __init__(self, value): if isinstance(value, np.ndarray): value = value.tolist() @@ -18,7 +18,7 @@ def to_list(self): def to_db(value, dim=None): if value is None: return value - if isinstance(value, HalfVec): + if isinstance(value, HalfVector): value = value.value if dim is not None and len(value) != dim: @@ -29,20 +29,20 @@ def to_db(value, dim=None): def to_db_binary(value): if value is None: return value - if isinstance(value, HalfVec): + if isinstance(value, HalfVector): value = value.value return pack(f'>HH{len(value)}e', len(value), 0, *value) def from_db(value): - if value is None or isinstance(value, HalfVec): + if value is None or isinstance(value, HalfVector): return value - return HalfVec([float(v) for v in value[1:-1].split(',')]) + return HalfVector([float(v) for v in value[1:-1].split(',')]) def from_db_binary(value): - if value is None or isinstance(value, HalfVec): + if value is None or isinstance(value, HalfVector): return value dim, unused = unpack_from('>HH', value) - return HalfVec(unpack_from(f'>{dim}e', value, 4)) + return HalfVector(unpack_from(f'>{dim}e', value, 4)) def __repr__(self): - return f'HalfVec({self.value})' + return f'HalfVector({self.value})' diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 040ef2c..0c6885c 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -3,15 +3,15 @@ def to_db_value(value): - if isinstance(value, SparseVec): + if isinstance(value, SparseVector): return value elif isinstance(value, (list, np.ndarray)): - return SparseVec.from_dense(value) + return SparseVector.from_dense(value) else: raise ValueError('expected sparsevec') -class SparseVec: +class SparseVector: def __init__(self, dim, indices, values): self.dim = dim self.indices = indices @@ -23,7 +23,7 @@ def from_dense(value): dim = len(value) indices = [i for i, v in enumerate(value) if v != 0] values = [value[i] for i in indices] - return SparseVec(dim, indices, values) + return SparseVector(dim, indices, values) def to_dense(self): vec = [0] * self.dim @@ -51,7 +51,7 @@ def to_db_binary(value): return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) def from_db(value): - if value is None or isinstance(value, SparseVec): + if value is None or isinstance(value, SparseVector): return value elements, dim = value.split('/') indices = [] @@ -60,15 +60,15 @@ def from_db(value): i, v = e.split(':') indices.append(int(i) - 1) values.append(float(v)) - return SparseVec(int(dim), indices, values) + return SparseVector(int(dim), indices, values) def from_db_binary(value): - if value is None or isinstance(value, SparseVec): + if value is None or isinstance(value, SparseVector): return value dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return SparseVec(int(dim), indices, values) + return SparseVector(int(dim), indices, values) def __repr__(self): - return f'SparseVec({self.dim}, {self.indices}, {self.values})' + return f'SparseVector({self.dim}, {self.indices}, {self.values})' diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 7405631..474ba92 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,7 +1,7 @@ import asyncio import asyncpg import numpy as np -from pgvector.asyncpg import register_vector, SparseVec +from pgvector.asyncpg import register_vector, SparseVector import pytest @@ -88,7 +88,7 @@ async def test_sparsevec(self): await register_vector(conn) - embedding = SparseVec.from_dense([1.5, 2, 3]) + embedding = SparseVector.from_dense([1.5, 2, 3]) await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") diff --git a/tests/test_django.py b/tests/test_django.py index a8b1adf..586e9ad 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfvecField, BitField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVec, SparseVec +from pgvector.django import VectorExtension, VectorField, HalfvecField, BitField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVector, SparseVector from unittest import mock settings.configure( @@ -94,7 +94,7 @@ def create_items(): [1, 1, 2] ] for i, v in enumerate(vectors): - item = Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v)) + item = Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v)) item.save() @@ -131,7 +131,7 @@ def test_bit(self): assert item.binary_embedding == '101' def test_sparsevec(self): - item = Item(id=1, sparse_embedding=SparseVec.from_dense([1, 2, 3])) + item = Item(id=1, sparse_embedding=SparseVector.from_dense([1, 2, 3])) item.save() item = Item.objects.get(pk=1) assert item.id == 1 @@ -167,56 +167,56 @@ def test_vector_l1_distance(self): def test_halfvec_l2_distance(self): create_items() - distance = L2Distance('half_embedding', HalfVec([1, 1, 1])) + distance = L2Distance('half_embedding', HalfVector([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)] def test_halfvec_max_inner_product(self): create_items() - distance = MaxInnerProduct('half_embedding', HalfVec([1, 1, 1])) + distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_halfvec_cosine_distance(self): create_items() - distance = CosineDistance('half_embedding', HalfVec([1, 1, 1])) + distance = CosineDistance('half_embedding', HalfVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_halfvec_l1_distance(self): create_items() - distance = L1Distance('half_embedding', HalfVec([1, 1, 1])) + distance = L1Distance('half_embedding', HalfVector([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, 3] def test_sparsevec_l2_distance(self): create_items() - distance = L2Distance('sparse_embedding', SparseVec.from_dense([1, 1, 1])) + distance = L2Distance('sparse_embedding', SparseVector.from_dense([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)] def test_sparsevec_max_inner_product(self): create_items() - distance = MaxInnerProduct('sparse_embedding', SparseVec.from_dense([1, 1, 1])) + distance = MaxInnerProduct('sparse_embedding', SparseVector.from_dense([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_sparsevec_cosine_distance(self): create_items() - distance = CosineDistance('sparse_embedding', SparseVec.from_dense([1, 1, 1])) + distance = CosineDistance('sparse_embedding', SparseVector.from_dense([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_sparsevec_l1_distance(self): create_items() - distance = L1Distance('sparse_embedding', SparseVec.from_dense([1, 1, 1])) + distance = L1Distance('sparse_embedding', SparseVector.from_dense([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, 3] diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 4e7eb44..5c8ee5a 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField, HalfvecField, FixedBitField, SparsevecField, SparseVec +from pgvector.peewee import VectorField, HalfvecField, FixedBitField, SparsevecField, SparseVector db = PostgresqlDatabase('pgvector_python_test') @@ -36,7 +36,7 @@ def create_items(): [1, 1, 2] ] for i, v in enumerate(vectors): - Item.create(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v)) + Item.create(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v)) class TestPeewee: @@ -107,7 +107,7 @@ def test_halfvec_l1_distance(self): def test_sparsevec_l2_distance(self): create_items() - distance = Item.sparse_embedding.l2_distance(SparseVec.from_dense([1, 1, 1])) + distance = Item.sparse_embedding.l2_distance(SparseVector.from_dense([1, 1, 1])) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index f674a66..aeca9d9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg import register_vector, register_vector_async, HalfVec, SparseVec +from pgvector.psycopg import register_vector, register_vector_async, HalfVector, SparseVector import psycopg import pytest @@ -56,7 +56,7 @@ def test_text_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: - copy.write_row([embedding, HalfVec(embedding), '101', SparseVec.from_dense(embedding)]) + copy.write_row([embedding, HalfVector(embedding), '101', SparseVector.from_dense(embedding)]) def test_binary_copy(self): embedding = np.array([1.5, 2, 3]) @@ -69,24 +69,24 @@ def test_binary_copy_set_types(self): cur = conn.cursor() with cur.copy("COPY psycopg_items (id, embedding, half_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector', 'halfvec', 'sparsevec']) - copy.write_row([1, embedding, HalfVec(embedding), SparseVec.from_dense(embedding)]) + copy.write_row([1, embedding, HalfVector(embedding), SparseVector.from_dense(embedding)]) def test_halfvec(self): conn.execute('DROP TABLE IF EXISTS half_items') conn.execute('CREATE TABLE half_items (id bigserial PRIMARY KEY, embedding halfvec(3))') - embedding = HalfVec([1.5, 2, 3]) + embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO half_items (embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT * FROM half_items ORDER BY id').fetchall() def test_halfvec_binary_format(self): - embedding = HalfVec([1.5, 2, 3]) + embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] assert res.to_list() == [1.5, 2, 3] def test_halfvec_text_format(self): - embedding = HalfVec([1.5, 2, 3]) + embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] assert res.to_list() == [1.5, 2, 3] @@ -94,19 +94,19 @@ def test_sparsevec(self): conn.execute('DROP TABLE IF EXISTS sparse_items') conn.execute('CREATE TABLE sparse_items (id bigserial PRIMARY KEY, embedding sparsevec(6))') - embedding = SparseVec.from_dense([0, 1.5, 0, 2, 0, 3]) + embedding = SparseVector.from_dense([0, 1.5, 0, 2, 0, 3]) conn.execute('INSERT INTO sparse_items (embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT * FROM sparse_items ORDER BY id').fetchall() assert res[0][1].to_dense() == [0, 1.5, 0, 2, 0, 3] def test_sparsevec_binary_format(self): - embedding = SparseVec.from_dense([1.5, 2, 3]) + embedding = SparseVector.from_dense([1.5, 2, 3]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] assert res.to_dense() == [1.5, 2, 3] def test_sparsevec_text_format(self): - embedding = SparseVec.from_dense([1.5, 2, 3]) + embedding = SparseVector.from_dense([1.5, 2, 3]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] assert res.to_dense() == [1.5, 2, 3] diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 7302296..cc7e5c0 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg2 import register_vector, SparseVec +from pgvector.psycopg2 import register_vector, SparseVector import psycopg2 conn = psycopg2.connect(dbname='pgvector_python_test') @@ -46,7 +46,7 @@ def test_bit(self): assert res[1][0] is None def test_sparsevec(self): - embedding = SparseVec.from_dense([1.5, 2, 3]) + embedding = SparseVector.from_dense([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 887859c..9378c82 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.sqlalchemy import Vector, Halfvec, Bit, Sparsevec, SparseVec +from pgvector.sqlalchemy import Vector, Halfvec, Bit, Sparsevec, SparseVector import pytest from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError @@ -46,7 +46,7 @@ def create_items(): ] session = Session(engine) for i, v in enumerate(vectors): - session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v))) + session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v))) session.commit() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 58ac2c9..86f0110 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.sqlalchemy import Vector, Halfvec, Bit, Sparsevec, SparseVec +from pgvector.sqlalchemy import Vector, Halfvec, Bit, Sparsevec, SparseVector import pytest from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError @@ -43,7 +43,7 @@ def create_items(): ] session = Session(engine) for i, v in enumerate(vectors): - session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v))) + session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v))) session.commit() From 2a036aebcfec5289876383721eb2bf6fb6f4ba6e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 14:21:35 -0400 Subject: [PATCH 223/879] Renamed HalfvecField to HalfVectorField and SparsevecField to SparseVectorField --- pgvector/django/__init__.py | 8 ++++---- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/peewee/__init__.py | 8 ++++---- pgvector/peewee/halfvec.py | 4 ++-- pgvector/peewee/sparsevec.py | 4 ++-- tests/test_django.py | 10 +++++----- tests/test_peewee.py | 6 +++--- 8 files changed, 22 insertions(+), 22 deletions(-) diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 08407c7..09978a9 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,18 +1,18 @@ from .bit import BitField from .extensions import VectorExtension from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance -from .halfvec import HalfvecField +from .halfvec import HalfVectorField from .indexes import IvfflatIndex, HnswIndex -from .sparsevec import SparsevecField +from .sparsevec import SparseVectorField from .vector import VectorField from ..utils import HalfVector, SparseVector __all__ = [ 'VectorExtension', 'VectorField', - 'HalfvecField', + 'HalfVectorField', 'BitField', - 'SparsevecField', + 'SparseVectorField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 884f706..1c3bae5 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -3,7 +3,7 @@ # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ -class HalfvecField(Field): +class HalfVectorField(Field): description = 'Half vector' empty_strings_allowed = False diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 1f4c21f..3a06574 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -3,7 +3,7 @@ # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ -class SparsevecField(Field): +class SparseVectorField(Field): description = 'Sparse vector' empty_strings_allowed = False diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 8f54401..945e0dc 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -1,14 +1,14 @@ from .bit import FixedBitField -from .halfvec import HalfvecField -from .sparsevec import SparsevecField +from .halfvec import HalfVectorField +from .sparsevec import SparseVectorField from .vector import VectorField from ..utils import HalfVector, SparseVector __all__ = [ 'VectorField', - 'HalfvecField', + 'HalfVectorField', 'FixedBitField', - 'SparsevecField', + 'SparseVectorField', 'HalfVector', 'SparseVector' ] diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index e663d72..e30dcd0 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -2,12 +2,12 @@ from ..utils import HalfVector -class HalfvecField(Field): +class HalfVectorField(Field): field_type = 'halfvec' def __init__(self, dimensions=None, *args, **kwargs): self.dimensions = dimensions - super(HalfvecField, self).__init__(*args, **kwargs) + super(HalfVectorField, self).__init__(*args, **kwargs) def get_modifiers(self): return self.dimensions and [self.dimensions] or None diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 272c9c6..b6d3b91 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -2,12 +2,12 @@ from ..utils import SparseVector -class SparsevecField(Field): +class SparseVectorField(Field): field_type = 'sparsevec' def __init__(self, dimensions=None, *args, **kwargs): self.dimensions = dimensions - super(SparsevecField, self).__init__(*args, **kwargs) + super(SparseVectorField, self).__init__(*args, **kwargs) def get_modifiers(self): return self.dimensions and [self.dimensions] or None diff --git a/tests/test_django.py b/tests/test_django.py index 586e9ad..b5c67e4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -8,7 +8,7 @@ from math import sqrt import numpy as np import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfvecField, BitField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVector, SparseVector +from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVector, SparseVector from unittest import mock settings.configure( @@ -24,9 +24,9 @@ class Item(models.Model): embedding = VectorField(dimensions=3, null=True, blank=True) - half_embedding = HalfvecField(dimensions=3, null=True, blank=True) + half_embedding = HalfVectorField(dimensions=3, null=True, blank=True) binary_embedding = BitField(length=3, null=True, blank=True) - sparse_embedding = SparsevecField(dimensions=3, null=True, blank=True) + sparse_embedding = SparseVectorField(dimensions=3, null=True, blank=True) class Meta: app_label = 'myapp' @@ -60,9 +60,9 @@ class Migration(migrations.Migration): fields=[ ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), ('embedding', pgvector.django.VectorField(dimensions=3, null=True, blank=True)), - ('half_embedding', pgvector.django.HalfvecField(dimensions=3, null=True, blank=True)), + ('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.SparsevecField(dimensions=3, null=True, blank=True)), + ('sparse_embedding', pgvector.django.SparseVectorField(dimensions=3, null=True, blank=True)), ], ), migrations.AddIndex( diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 5c8ee5a..d83ef6f 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField, HalfvecField, FixedBitField, SparsevecField, SparseVector +from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField, SparseVector db = PostgresqlDatabase('pgvector_python_test') @@ -13,9 +13,9 @@ class Meta: class Item(BaseModel): embedding = VectorField(dimensions=3, null=True) - half_embedding = HalfvecField(dimensions=3, null=True) + half_embedding = HalfVectorField(dimensions=3, null=True) binary_embedding = FixedBitField(max_length=3, null=True) - sparse_embedding = SparsevecField(dimensions=3, null=True) + sparse_embedding = SparseVectorField(dimensions=3, null=True) class Meta: table_name = 'peewee_item' From 26b69ca54f054cf75e20fa3d1de7c90cb3550bc2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 16:15:27 -0400 Subject: [PATCH 224/879] Added Bit class for Psycopg 3 --- pgvector/psycopg/__init__.py | 9 ++++++++- pgvector/psycopg/bit.py | 31 +++++++++++++++++++++++++++++++ pgvector/utils/__init__.py | 1 + pgvector/utils/bit.py | 34 ++++++++++++++++++++++++++++++++++ tests/test_psycopg.py | 31 ++++++++++++++++++------------- 5 files changed, 92 insertions(+), 14 deletions(-) create mode 100644 pgvector/psycopg/bit.py create mode 100644 pgvector/utils/bit.py diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index e47d699..90ff8b7 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,9 +1,10 @@ import psycopg from psycopg.types import TypeInfo +from .bit import register_bit_info from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import HalfVector, SparseVector +from ..utils import Bit, HalfVector, SparseVector __all__ = ['register_vector'] @@ -13,6 +14,9 @@ def register_vector(context): info = TypeInfo.fetch(context, 'vector') register_vector_info(context, info) + info = TypeInfo.fetch(context, 'bit') + register_bit_info(context, info) + info = TypeInfo.fetch(context, 'halfvec') if info is not None: register_halfvec_info(context, info) @@ -26,6 +30,9 @@ async def register_vector_async(context): info = await TypeInfo.fetch(context, 'vector') register_vector_info(context, info) + info = await TypeInfo.fetch(context, 'bit') + register_bit_info(context, info) + info = await TypeInfo.fetch(context, 'halfvec') if info is not None: register_halfvec_info(context, info) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py new file mode 100644 index 0000000..80cfac0 --- /dev/null +++ b/pgvector/psycopg/bit.py @@ -0,0 +1,31 @@ +from psycopg.adapt import Loader, Dumper +from psycopg.pq import Format +from ..utils import Bit + + +class BitDumper(Dumper): + + format = Format.TEXT + + def dump(self, obj): + return Bit.to_db(obj).encode('utf8') + + +class BitBinaryDumper(BitDumper): + + format = Format.BINARY + + def dump(self, obj): + return Bit.to_db_binary(obj) + + +def register_bit_info(context, info): + info.register(context) + + # add oid to anonymous class for set_types + text_dumper = type('', (BitDumper,), {'oid': info.oid}) + binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid}) + + adapters = context.adapters + adapters.register_dumper(Bit, text_dumper) + adapters.register_dumper(Bit, binary_dumper) diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index 897862b..69e8cfb 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1,3 +1,4 @@ +from .bit import * from .halfvec import * from .sparsevec import * from .vector import * diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py new file mode 100644 index 0000000..ef639ba --- /dev/null +++ b/pgvector/utils/bit.py @@ -0,0 +1,34 @@ +import numpy as np +from struct import pack, unpack_from + + +class Bit: + def __init__(self, value): + if isinstance(value, bytes): + count = unpack_from('>i', value)[0] + buf = np.frombuffer(value[4:], dtype=np.uint8) + self.value = np.unpackbits(buf, count=count).astype(bool) + elif isinstance(value, str): + self.value = np.array([v != '0' for v in value], dtype=bool) + else: + self.value = np.array(value, dtype=bool) + + def to_db(value): + if not isinstance(value, Bit): + raise ValueError('expected bit') + + value = value.value + return ''.join(value.astype(np.uint8).astype(str)) + + def to_db_binary(value): + if not isinstance(value, Bit): + raise ValueError('expected bit') + + value = value.value + return pack('>i', len(value)) + np.packbits(value).tobytes() + + def __str__(self): + return self.__class__.to_db(self) + + def __repr__(self): + return f'Bit({self})' diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index aeca9d9..cba62fd 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg import register_vector, register_vector_async, HalfVector, SparseVector +from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector import psycopg import pytest @@ -90,6 +90,23 @@ def test_halfvec_text_format(self): res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] assert res.to_list() == [1.5, 2, 3] + def test_bit(self): + embedding = Bit([False, True, False, True, False, False, False, False, True]) + res = conn.execute('SELECT %s::bit(9)', (embedding,)).fetchone()[0] + assert res == '010100001' + assert str(Bit(res)) == '010100001' + + def test_bit_binary_format(self): + embedding = Bit([False, True, False, True, False, False, False, False, True]) + res = conn.execute('SELECT %b::bit(9)', (embedding,), binary=True).fetchone()[0] + assert str(Bit(res)) == '010100001' + assert repr(Bit(res)) == 'Bit(010100001)' + + def test_bit_text_format(self): + embedding = Bit([False, True, False, True, False, False, False, False, True]) + res = conn.execute('SELECT %t::bit(9)', (embedding,)).fetchone()[0] + assert res == '010100001' + def test_sparsevec(self): conn.execute('DROP TABLE IF EXISTS sparse_items') conn.execute('CREATE TABLE sparse_items (id bigserial PRIMARY KEY, embedding sparsevec(6))') @@ -110,18 +127,6 @@ def test_sparsevec_text_format(self): res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] assert res.to_dense() == [1.5, 2, 3] - def test_bit(self): - res = conn.execute('SELECT %s::bit(3)', ('101',)).fetchone()[0] - assert res == '101' - - def test_bit_binary_format(self): - res = conn.execute('SELECT %b::bit(3)', ('101',), binary=True).fetchone()[0] - assert res == b'\x00\x00\x00\x03\xa0' - - def test_bit_text_format(self): - res = conn.execute('SELECT %t::bit(3)', ('101',)).fetchone()[0] - assert res == '101' - @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) From bb060c87a8d23c67f78a5897e855a02dceecfa0d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 20:34:28 -0400 Subject: [PATCH 225/879] Moved copy tests [skip ci] --- tests/test_psycopg.py | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index cba62fd..c823eac 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -52,25 +52,6 @@ def test_vector_binary_format_non_contiguous(self): res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, np.array([3, 2, 1.5])) - def test_text_copy(self): - embedding = np.array([1.5, 2, 3]) - cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: - copy.write_row([embedding, HalfVector(embedding), '101', SparseVector.from_dense(embedding)]) - - def test_binary_copy(self): - embedding = np.array([1.5, 2, 3]) - cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.write_row([embedding]) - - def test_binary_copy_set_types(self): - embedding = np.array([1.5, 2, 3]) - cur = conn.cursor() - with cur.copy("COPY psycopg_items (id, embedding, half_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.set_types(['int8', 'vector', 'halfvec', 'sparsevec']) - copy.write_row([1, embedding, HalfVector(embedding), SparseVector.from_dense(embedding)]) - def test_halfvec(self): conn.execute('DROP TABLE IF EXISTS half_items') conn.execute('CREATE TABLE half_items (id bigserial PRIMARY KEY, embedding halfvec(3))') @@ -127,6 +108,25 @@ def test_sparsevec_text_format(self): res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] assert res.to_dense() == [1.5, 2, 3] + def test_text_copy(self): + embedding = np.array([1.5, 2, 3]) + cur = conn.cursor() + with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: + copy.write_row([embedding, HalfVector(embedding), '101', SparseVector.from_dense(embedding)]) + + def test_binary_copy(self): + embedding = np.array([1.5, 2, 3]) + cur = conn.cursor() + with cur.copy("COPY psycopg_items (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.write_row([embedding]) + + def test_binary_copy_set_types(self): + embedding = np.array([1.5, 2, 3]) + cur = conn.cursor() + with cur.copy("COPY psycopg_items (id, embedding, half_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.set_types(['int8', 'vector', 'halfvec', 'sparsevec']) + copy.write_row([1, embedding, HalfVector(embedding), SparseVector.from_dense(embedding)]) + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) From 950551ac3813af132849c5070f89ccc195bf67d1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 20:35:04 -0400 Subject: [PATCH 226/879] Added bit type to binary copy test --- tests/test_psycopg.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index c823eac..790d71e 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -123,9 +123,9 @@ def test_binary_copy(self): def test_binary_copy_set_types(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY psycopg_items (id, embedding, half_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.set_types(['int8', 'vector', 'halfvec', 'sparsevec']) - copy.write_row([1, embedding, HalfVector(embedding), SparseVector.from_dense(embedding)]) + with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) + copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector.from_dense(embedding)]) @pytest.mark.asyncio async def test_async(self): From e341c4d623ecffad8df23301f2e000b26f5aa126 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 20:43:00 -0400 Subject: [PATCH 227/879] Improved Psycopg 3 tests --- tests/test_psycopg.py | 46 +++++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 790d71e..c06dacd 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -20,10 +20,10 @@ def test_vector(self): embedding = np.array([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) - res = conn.execute('SELECT * FROM psycopg_items ORDER BY id').fetchall() - assert np.array_equal(res[0][1], embedding) - assert res[0][1].dtype == np.float32 - assert res[1][1] is None + res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() + assert np.array_equal(res[0][0], embedding) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None def test_vector_binary_format(self): embedding = np.array([1.5, 2, 3]) @@ -53,13 +53,11 @@ def test_vector_binary_format_non_contiguous(self): assert np.array_equal(res, np.array([3, 2, 1.5])) def test_halfvec(self): - conn.execute('DROP TABLE IF EXISTS half_items') - conn.execute('CREATE TABLE half_items (id bigserial PRIMARY KEY, embedding halfvec(3))') - embedding = HalfVector([1.5, 2, 3]) - conn.execute('INSERT INTO half_items (embedding) VALUES (%s)', (embedding,)) + conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT * FROM half_items ORDER BY id').fetchall() + res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + assert res.to_list() == [1.5, 2, 3] def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) @@ -72,10 +70,11 @@ def test_halfvec_text_format(self): assert res.to_list() == [1.5, 2, 3] def test_bit(self): - embedding = Bit([False, True, False, True, False, False, False, False, True]) - res = conn.execute('SELECT %s::bit(9)', (embedding,)).fetchone()[0] - assert res == '010100001' - assert str(Bit(res)) == '010100001' + embedding = Bit([True, False, True]) + conn.execute('INSERT INTO psycopg_items (binary_embedding) VALUES (%s)', (embedding,)) + + res = conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + assert res == '101' def test_bit_binary_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) @@ -87,26 +86,25 @@ def test_bit_text_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) res = conn.execute('SELECT %t::bit(9)', (embedding,)).fetchone()[0] assert res == '010100001' + assert str(Bit(res)) == '010100001' + assert repr(Bit(res)) == 'Bit(010100001)' def test_sparsevec(self): - conn.execute('DROP TABLE IF EXISTS sparse_items') - conn.execute('CREATE TABLE sparse_items (id bigserial PRIMARY KEY, embedding sparsevec(6))') - - embedding = SparseVector.from_dense([0, 1.5, 0, 2, 0, 3]) - conn.execute('INSERT INTO sparse_items (embedding) VALUES (%s)', (embedding,)) + embedding = SparseVector.from_dense([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT * FROM sparse_items ORDER BY id').fetchall() - assert res[0][1].to_dense() == [0, 1.5, 0, 2, 0, 3] + res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + assert res.to_dense() == [1.5, 2, 3] def test_sparsevec_binary_format(self): - embedding = SparseVector.from_dense([1.5, 2, 3]) + embedding = SparseVector.from_dense([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] - assert res.to_dense() == [1.5, 2, 3] + assert res.to_dense() == [1.5, 0, 2, 0, 3, 0] def test_sparsevec_text_format(self): - embedding = SparseVector.from_dense([1.5, 2, 3]) + embedding = SparseVector.from_dense([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] - assert res.to_dense() == [1.5, 2, 3] + assert res.to_dense() == [1.5, 0, 2, 0, 3, 0] def test_text_copy(self): embedding = np.array([1.5, 2, 3]) From 0d6260a328f3395456b619a80905a67d3bb160de Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 20:44:30 -0400 Subject: [PATCH 228/879] Improved test [skip ci] --- tests/test_psycopg.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index c06dacd..a11b73c 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -115,8 +115,8 @@ def test_text_copy(self): def test_binary_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.write_row([embedding]) + with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector.from_dense(embedding)]) def test_binary_copy_set_types(self): embedding = np.array([1.5, 2, 3]) From a8e971a9f7e1f22783b0bed2c837664434e1bfaa Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 20:59:28 -0400 Subject: [PATCH 229/879] Improved tests [skip ci] --- tests/test_django.py | 19 +++++-------------- tests/test_peewee.py | 18 +++++------------- 2 files changed, 10 insertions(+), 27 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index b5c67e4..2652e39 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -88,14 +88,9 @@ class Migration(migrations.Migration): def create_items(): - vectors = [ - [1, 1, 1], - [2, 2, 2], - [1, 1, 2] - ] - for i, v in enumerate(vectors): - item = Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v)) - item.save() + Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1])).save() + Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2])).save() + Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2])).save() class ItemForm(ModelForm): @@ -222,18 +217,14 @@ def test_sparsevec_l1_distance(self): assert [v.distance for v in items] == [0, 1, 3] def test_bit_hamming_distance(self): - Item(id=1, binary_embedding='000').save() - Item(id=2, binary_embedding='101').save() - Item(id=3, binary_embedding='111').save() + create_items() distance = HammingDistance('binary_embedding', '101') items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] def test_bit_jaccard_distance(self): - Item(id=1, binary_embedding='000').save() - Item(id=2, binary_embedding='101').save() - Item(id=3, binary_embedding='111').save() + create_items() distance = JaccardDistance('binary_embedding', '101') items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] diff --git a/tests/test_peewee.py b/tests/test_peewee.py index d83ef6f..a5f3d97 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -30,13 +30,9 @@ class Meta: def create_items(): - vectors = [ - [1, 1, 1], - [2, 2, 2], - [1, 1, 2] - ] - for i, v in enumerate(vectors): - Item.create(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v)) + Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1])) + Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2])) + Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2])) class TestPeewee: @@ -134,18 +130,14 @@ def test_sparsevec_l1_distance(self): assert [v.distance for v in items] == [0, 1, 3] def test_bit_hamming_distance(self): - Item.create(id=1, binary_embedding='000') - Item.create(id=2, binary_embedding='101') - Item.create(id=3, binary_embedding='111') + create_items() distance = Item.binary_embedding.hamming_distance('101') items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] def test_bit_jaccard_distance(self): - Item.create(id=1, binary_embedding='000') - Item.create(id=2, binary_embedding='101') - Item.create(id=3, binary_embedding='111') + create_items() distance = Item.binary_embedding.jaccard_distance('101') items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] From ad7a1863e208c4507f0d96f71aa4c16edee6c777 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:01:12 -0400 Subject: [PATCH 230/879] Improved tests [skip ci] --- tests/test_sqlalchemy.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 9378c82..c0fdb5c 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -39,14 +39,10 @@ class Item(Base): def create_items(): - vectors = [ - [1, 1, 1], - [2, 2, 2], - [1, 1, 2] - ] session = Session(engine) - for i, v in enumerate(vectors): - session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v))) + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2]))) session.commit() @@ -174,21 +170,13 @@ def test_sparsevec_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] def test_bit_hamming_distance(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='000')) - session.add(Item(id=2, binary_embedding='101')) - session.add(Item(id=3, binary_embedding='111')) - session.commit() + create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] def test_bit_jaccard_distance(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='000')) - session.add(Item(id=2, binary_embedding='101')) - session.add(Item(id=3, binary_embedding='111')) - session.commit() + create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] From ff129253971b2923ddf53178663236ac5f8f2a4f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:02:58 -0400 Subject: [PATCH 231/879] Improved tests [skip ci] --- tests/test_sqlmodel.py | 22 +++++----------------- 1 file changed, 5 insertions(+), 17 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 86f0110..c527a18 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -36,14 +36,10 @@ class Item(SQLModel, table=True): def create_items(): - vectors = [ - [1, 1, 1], - [2, 2, 2], - [1, 1, 2] - ] session = Session(engine) - for i, v in enumerate(vectors): - session.add(Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVector.from_dense(v))) + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2]))) session.commit() @@ -149,21 +145,13 @@ def test_sparsevec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] def test_bit_hamming_distance(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='000')) - session.add(Item(id=2, binary_embedding='101')) - session.add(Item(id=3, binary_embedding='111')) - session.commit() + create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) assert [v.id for v in items] == [2, 3, 1] def test_bit_jaccard_distance(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='000')) - session.add(Item(id=2, binary_embedding='101')) - session.add(Item(id=3, binary_embedding='111')) - session.commit() + create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) assert [v.id for v in items] == [2, 3, 1] From ce48133aacc9073adedff88863912402f40eeeea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:07:50 -0400 Subject: [PATCH 232/879] Improved table names [skip ci] --- tests/test_django.py | 6 +++--- tests/test_sqlalchemy.py | 14 +++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 2652e39..a0c837c 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -29,7 +29,7 @@ class Item(models.Model): sparse_embedding = SparseVectorField(dimensions=3, null=True, blank=True) class Meta: - app_label = 'myapp' + app_label = 'django_app' indexes = [ IvfflatIndex( name='ivfflat_idx', @@ -77,9 +77,9 @@ class Migration(migrations.Migration): # probably a better way to do this -migration = Migration('initial', 'myapp') +migration = Migration('initial', 'django_app') loader = MigrationLoader(connection, replace_migrations=False) -loader.graph.add_node(('myapp', migration.name), migration) +loader.graph.add_node(('django_app', migration.name), migration) sql_statements = loader.collect_sql([(migration, False)]) with connection.cursor() as cursor: diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index c0fdb5c..b03e6a0 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -16,7 +16,7 @@ class Item(Base): - __tablename__ = 'orm_item' + __tablename__ = 'sqlalchemy_orm_item' id = mapped_column(Integer, primary_key=True) embedding = mapped_column(Vector(3)) @@ -29,7 +29,7 @@ class Item(Base): Base.metadata.create_all(engine) index = Index( - 'orm_index', + 'sqlalchemy_orm_index', Item.embedding, postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, @@ -56,7 +56,7 @@ def test_core(self): metadata = MetaData() item_table = Table( - 'core_item', + 'sqlalchemy_core_item', metadata, Column('id', Integer, primary_key=True), Column('embedding', Vector(3)), @@ -69,7 +69,7 @@ def test_core(self): metadata.create_all(engine) ivfflat_index = Index( - 'ivfflat_core_index', + 'sqlalchemy_core_ivfflat_index', item_table.c.embedding, postgresql_using='ivfflat', postgresql_with={'lists': 1}, @@ -78,7 +78,7 @@ def test_core(self): ivfflat_index.create(engine) hnsw_index = Index( - 'hnsw_core_index', + 'sqlalchemy_core_hnsw_index', item_table.c.embedding, postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, @@ -263,7 +263,7 @@ def test_bad_dtype(self): session.commit() def test_inspect(self): - columns = inspect(engine).get_columns('orm_item') + columns = inspect(engine).get_columns('sqlalchemy_orm_item') assert isinstance(columns[1]['type'], Vector) def test_literal_binds(self): @@ -277,7 +277,7 @@ def test_insert_bulk(self): session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) def test_insert_text(self): - session.execute(text('INSERT INTO orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) + session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) @pytest.mark.asyncio async def test_async(self): From 2fb168420685a5ec679e28810e2796bb2b963f37 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:08:23 -0400 Subject: [PATCH 233/879] Fixed table name [skip ci] --- tests/test_django.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index a0c837c..9523c65 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -83,7 +83,7 @@ class Migration(migrations.Migration): sql_statements = loader.collect_sql([(migration, False)]) with connection.cursor() as cursor: - cursor.execute("DROP TABLE IF EXISTS myapp_item") + cursor.execute("DROP TABLE IF EXISTS django_app_item") cursor.execute('\n'.join(sql_statements)) From ab048e93906f9709ab7f7325c41bac392f1dfd03 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:18:47 -0400 Subject: [PATCH 234/879] Removed unneeded tests [skip ci] --- tests/test_asyncpg.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 474ba92..4f0a4bd 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -19,8 +19,6 @@ async def test_vector(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['id'] == 1 - assert res[1]['id'] == 2 assert np.array_equal(res[0]['embedding'], embedding) assert res[0]['embedding'].dtype == np.float32 assert res[1]['embedding'] is None @@ -44,8 +42,6 @@ async def test_halfvec(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['id'] == 1 - assert res[1]['id'] == 2 assert res[0]['embedding'].to_list() == [1.5, 2, 3] assert res[1]['embedding'] is None @@ -68,8 +64,6 @@ async def test_bit(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['id'] == 1 - assert res[1]['id'] == 2 assert res[0]['embedding'].to_int() == 5 assert res[1]['embedding'] is None @@ -92,8 +86,6 @@ async def test_sparsevec(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['id'] == 1 - assert res[1]['id'] == 2 assert res[0]['embedding'].to_dense() == [1.5, 2, 3] assert res[1]['embedding'] is None @@ -119,8 +111,6 @@ async def init(conn): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['id'] == 1 - assert res[1]['id'] == 2 assert np.array_equal(res[0]['embedding'], embedding) assert res[0]['embedding'].dtype == np.float32 assert res[1]['embedding'] is None From 1566eac74e42b8d512bd3ce30c7bfc275995237e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:37:26 -0400 Subject: [PATCH 235/879] Added more tests [skip ci] --- tests/test_peewee.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index a5f3d97..1146234 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -73,6 +73,11 @@ def test_vector_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_halfvec(self): + Item.create(id=1, half_embedding=[1, 2, 3]) + item = Item.get_by_id(1) + assert item.half_embedding.to_list() == [1, 2, 3] + def test_halfvec_l2_distance(self): create_items() distance = Item.half_embedding.l2_distance([1, 1, 1]) @@ -101,6 +106,11 @@ def test_halfvec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_sparsevec(self): + Item.create(id=1, sparse_embedding=[1, 2, 3]) + item = Item.get_by_id(1) + assert item.sparse_embedding.to_dense() == [1, 2, 3] + def test_sparsevec_l2_distance(self): create_items() distance = Item.sparse_embedding.l2_distance(SparseVector.from_dense([1, 1, 1])) @@ -129,6 +139,11 @@ def test_sparsevec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_bit(self): + Item.create(id=1, binary_embedding='101') + item = Item.get_by_id(1) + assert item.binary_embedding == '101' + def test_bit_hamming_distance(self): create_items() distance = Item.binary_embedding.hamming_distance('101') From b2ab8617a1d02c1259bc7e939047f39d1de47e02 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:39:47 -0400 Subject: [PATCH 236/879] Improved tests [skip ci] --- tests/test_django.py | 40 ++++++++++++++++------------------------ 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 9523c65..9f38761 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -104,34 +104,11 @@ def setup_method(self, test_method): Item.objects.all().delete() def test_vector(self): - item = Item(id=1, embedding=[1, 2, 3]) - item.save() + Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert item.id == 1 assert np.array_equal(item.embedding, np.array([1, 2, 3])) assert item.embedding.dtype == np.float32 - def test_halfvec(self): - item = Item(id=1, half_embedding=[1, 2, 3]) - item.save() - item = Item.objects.get(pk=1) - assert item.id == 1 - assert item.half_embedding.to_list() == [1, 2, 3] - - def test_bit(self): - item = Item(id=1, binary_embedding='101') - item.save() - item = Item.objects.get(pk=1) - assert item.id == 1 - assert item.binary_embedding == '101' - - def test_sparsevec(self): - item = Item(id=1, sparse_embedding=SparseVector.from_dense([1, 2, 3])) - item.save() - item = Item.objects.get(pk=1) - assert item.id == 1 - assert item.sparse_embedding.to_dense() == [1, 2, 3] - def test_vector_l2_distance(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) @@ -160,6 +137,11 @@ def test_vector_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_halfvec(self): + Item(id=1, half_embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + assert item.half_embedding.to_list() == [1, 2, 3] + def test_halfvec_l2_distance(self): create_items() distance = L2Distance('half_embedding', HalfVector([1, 1, 1])) @@ -188,6 +170,11 @@ def test_halfvec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_sparsevec(self): + Item(id=1, sparse_embedding=SparseVector.from_dense([1, 2, 3])).save() + item = Item.objects.get(pk=1) + assert item.sparse_embedding.to_dense() == [1, 2, 3] + def test_sparsevec_l2_distance(self): create_items() distance = L2Distance('sparse_embedding', SparseVector.from_dense([1, 1, 1])) @@ -216,6 +203,11 @@ def test_sparsevec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_bit(self): + Item(id=1, binary_embedding='101').save() + item = Item.objects.get(pk=1) + assert item.binary_embedding == '101' + def test_bit_hamming_distance(self): create_items() distance = HammingDistance('binary_embedding', '101') From 0250d419cf5e2ab483a8c04bf38ed8dc5c27654e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:45:37 -0400 Subject: [PATCH 237/879] Improved tests [skip ci] --- tests/test_sqlalchemy.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index b03e6a0..56a2539 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -109,6 +109,13 @@ def test_orm(self): assert items[1].embedding.dtype == np.float32 assert items[2].embedding is None + def test_vector(self): + session = Session(engine) + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.embedding.tolist() == [1, 2, 3] + def test_vector_l2_distance(self): create_items() with Session(engine) as session: @@ -157,18 +164,39 @@ def test_vector_l1_distance_orm(self): items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_halfvec(self): + session = Session(engine) + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding.to_list() == [1, 2, 3] + def test_halfvec_l2_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] + def test_sparsevec(self): + session = Session(engine) + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding.to_dense() == [1, 2, 3] + def test_sparsevec_l2_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] + def test_bit(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' + def test_bit_hamming_distance(self): create_items() with Session(engine) as session: From 505af704163cce9dcbd56b1491b435e36d282f18 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:48:56 -0400 Subject: [PATCH 238/879] Improved tests [skip ci] --- tests/test_sqlalchemy.py | 96 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 56a2539..174b5a0 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -177,6 +177,48 @@ def test_halfvec_l2_distance(self): items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] + def test_halfvec_l2_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + + def test_halfvec_max_inner_product(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all() + assert [v.id for v in items] == [2, 3, 1] + + def test_halfvec_max_inner_product_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) + assert [v.id for v in items] == [2, 3, 1] + + def test_halfvec_cosine_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 2, 3] + + def test_halfvec_cosine_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 2, 3] + + def test_halfvec_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_halfvec_l1_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_sparsevec(self): session = Session(engine) session.add(Item(id=1, sparse_embedding=[1, 2, 3])) @@ -190,6 +232,48 @@ def test_sparsevec_l2_distance(self): items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] + def test_sparsevec_l2_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + + def test_sparsevec_max_inner_product(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all() + assert [v.id for v in items] == [2, 3, 1] + + def test_sparsevec_max_inner_product_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) + assert [v.id for v in items] == [2, 3, 1] + + def test_sparsevec_cosine_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 2, 3] + + def test_sparsevec_cosine_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 2, 3] + + def test_sparsevec_l1_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_sparsevec_l1_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) + assert [v.id for v in items] == [1, 3, 2] + def test_bit(self): session = Session(engine) session.add(Item(id=1, binary_embedding='101')) @@ -203,12 +287,24 @@ def test_bit_hamming_distance(self): items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] + def test_bit_hamming_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + def test_bit_jaccard_distance(self): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] + def test_bit_jaccard_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + def test_filter(self): create_items() with Session(engine) as session: From 2c7b585afa976de6bcf2831eea24cf1cbda610f0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:51:07 -0400 Subject: [PATCH 239/879] Improved tests [skip ci] --- tests/test_sqlmodel.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index c527a18..c3716c9 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -72,6 +72,13 @@ def test_orm(self): assert items[1].embedding.dtype == np.float32 assert items[2].embedding is None + def test_vector(self): + session = Session(engine) + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.embedding.tolist() == [1, 2, 3] + def test_vector_l2_distance(self): create_items() with Session(engine) as session: @@ -96,6 +103,13 @@ def test_vector_l1_distance(self): items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_halfvec(self): + session = Session(engine) + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding.to_list() == [1, 2, 3] + def test_halfvec_l2_distance(self): create_items() with Session(engine) as session: @@ -120,6 +134,13 @@ def test_halfvec_l1_distance(self): items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_sparsevec(self): + session = Session(engine) + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding.to_dense() == [1, 2, 3] + def test_sparsevec_l2_distance(self): create_items() with Session(engine) as session: @@ -144,6 +165,13 @@ def test_sparsevec_l1_distance(self): items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_bit(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' + def test_bit_hamming_distance(self): create_items() with Session(engine) as session: From 22025d4b8d87f191f500f49395b3908cc85995f9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 20 May 2024 21:54:57 -0400 Subject: [PATCH 240/879] Moved tests [skip ci] --- tests/test_django.py | 38 ++++++++++++------------ tests/test_peewee.py | 38 ++++++++++++------------ tests/test_sqlalchemy.py | 62 ++++++++++++++++++++-------------------- tests/test_sqlmodel.py | 38 ++++++++++++------------ 4 files changed, 88 insertions(+), 88 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 9f38761..b40a47a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -170,6 +170,25 @@ def test_halfvec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_bit(self): + Item(id=1, binary_embedding='101').save() + item = Item.objects.get(pk=1) + assert item.binary_embedding == '101' + + def test_bit_hamming_distance(self): + create_items() + distance = HammingDistance('binary_embedding', '101') + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [0, 1, 2] + + def test_bit_jaccard_distance(self): + create_items() + distance = JaccardDistance('binary_embedding', '101') + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [2, 3, 1] + # assert [v.distance for v in items] == [0, 1/3, 1] + def test_sparsevec(self): Item(id=1, sparse_embedding=SparseVector.from_dense([1, 2, 3])).save() item = Item.objects.get(pk=1) @@ -203,25 +222,6 @@ def test_sparsevec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_bit(self): - Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) - assert item.binary_embedding == '101' - - def test_bit_hamming_distance(self): - create_items() - distance = HammingDistance('binary_embedding', '101') - items = Item.objects.annotate(distance=distance).order_by(distance) - assert [v.id for v in items] == [2, 3, 1] - assert [v.distance for v in items] == [0, 1, 2] - - def test_bit_jaccard_distance(self): - create_items() - distance = JaccardDistance('binary_embedding', '101') - items = Item.objects.annotate(distance=distance).order_by(distance) - assert [v.id for v in items] == [2, 3, 1] - # assert [v.distance for v in items] == [0, 1/3, 1] - def test_filter(self): create_items() distance = L2Distance('embedding', [1, 1, 1]) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 1146234..deb571f 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -106,6 +106,25 @@ def test_halfvec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] + def test_bit(self): + Item.create(id=1, binary_embedding='101') + item = Item.get_by_id(1) + assert item.binary_embedding == '101' + + def test_bit_hamming_distance(self): + create_items() + distance = Item.binary_embedding.hamming_distance('101') + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + assert [v.distance for v in items] == [0, 1, 2] + + def test_bit_jaccard_distance(self): + create_items() + distance = Item.binary_embedding.jaccard_distance('101') + items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) + assert [v.id for v in items] == [2, 3, 1] + # assert [v.distance for v in items] == [0, 1/3, 1] + def test_sparsevec(self): Item.create(id=1, sparse_embedding=[1, 2, 3]) item = Item.get_by_id(1) @@ -139,25 +158,6 @@ def test_sparsevec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_bit(self): - Item.create(id=1, binary_embedding='101') - item = Item.get_by_id(1) - assert item.binary_embedding == '101' - - def test_bit_hamming_distance(self): - create_items() - distance = Item.binary_embedding.hamming_distance('101') - items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) - assert [v.id for v in items] == [2, 3, 1] - assert [v.distance for v in items] == [0, 1, 2] - - def test_bit_jaccard_distance(self): - create_items() - distance = Item.binary_embedding.jaccard_distance('101') - items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) - assert [v.id for v in items] == [2, 3, 1] - # assert [v.distance for v in items] == [0, 1/3, 1] - def test_where(self): create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 174b5a0..1fff283 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -219,6 +219,37 @@ def test_halfvec_l1_distance_orm(self): items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_bit(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' + + def test_bit_hamming_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() + assert [v.id for v in items] == [2, 3, 1] + + def test_bit_hamming_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + + def test_bit_jaccard_distance(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() + assert [v.id for v in items] == [2, 3, 1] + + def test_bit_jaccard_distance_orm(self): + create_items() + with Session(engine) as session: + items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + def test_sparsevec(self): session = Session(engine) session.add(Item(id=1, sparse_embedding=[1, 2, 3])) @@ -274,37 +305,6 @@ def test_sparsevec_l1_distance_orm(self): items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' - - def test_bit_hamming_distance(self): - create_items() - with Session(engine) as session: - items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() - assert [v.id for v in items] == [2, 3, 1] - - def test_bit_hamming_distance_orm(self): - create_items() - with Session(engine) as session: - items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) - assert [v.id for v in items] == [2, 3, 1] - - def test_bit_jaccard_distance(self): - create_items() - with Session(engine) as session: - items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() - assert [v.id for v in items] == [2, 3, 1] - - def test_bit_jaccard_distance_orm(self): - create_items() - with Session(engine) as session: - items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) - assert [v.id for v in items] == [2, 3, 1] - def test_filter(self): create_items() with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index c3716c9..8563bec 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -134,6 +134,25 @@ def test_halfvec_l1_distance(self): items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_bit(self): + session = Session(engine) + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' + + def test_bit_hamming_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + + def test_bit_jaccard_distance(self): + create_items() + with Session(engine) as session: + items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) + assert [v.id for v in items] == [2, 3, 1] + def test_sparsevec(self): session = Session(engine) session.add(Item(id=1, sparse_embedding=[1, 2, 3])) @@ -165,25 +184,6 @@ def test_sparsevec_l1_distance(self): items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' - - def test_bit_hamming_distance(self): - create_items() - with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) - assert [v.id for v in items] == [2, 3, 1] - - def test_bit_jaccard_distance(self): - create_items() - with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) - assert [v.id for v in items] == [2, 3, 1] - def test_filter(self): create_items() with Session(engine) as session: From 68c0c4d3a242f5df0514f7582cdf8c1ed49f884b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 12:52:58 -0400 Subject: [PATCH 241/879] Added sparse search example [skip ci] --- README.md | 1 + examples/sparse_search.py | 54 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 examples/sparse_search.py diff --git a/README.md b/README.md index 28703a6..710525a 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ Or check out some examples: - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_embeddings.py) with SentenceTransformers - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) +- [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search.py) with Transformers (unreleased) - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit diff --git a/examples/sparse_search.py b/examples/sparse_search.py new file mode 100644 index 0000000..7e786e3 --- /dev/null +++ b/examples/sparse_search.py @@ -0,0 +1,54 @@ +# good resources +# https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/ +# https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1 + +import numpy as np +from pgvector.psycopg import register_vector, SparseVector +import psycopg +import torch +from transformers import AutoModelForMaskedLM, AutoTokenizer + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding sparsevec(30522))') + +model_id = 'opensearch-project/opensearch-neural-sparse-encoding-v1' +model = AutoModelForMaskedLM.from_pretrained(model_id) +tokenizer = AutoTokenizer.from_pretrained(model_id) +special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()] + + +def fetch_embeddings(input): + feature = tokenizer( + input, + padding=True, + truncation=True, + return_tensors='pt', + return_token_type_ids=False + ) + output = model(**feature)[0] + + values, _ = torch.max(output * feature['attention_mask'].unsqueeze(-1), dim=1) + values = torch.log(1 + torch.relu(values)) + values[:, special_token_ids] = 0 + return values.detach().cpu().numpy() + + +# note: works much better with longer content +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +embeddings = fetch_embeddings(input) +for content, embedding in zip(input, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector.from_dense(embedding))) + +query = 'forest' +query_embedding = fetch_embeddings([query])[0] +result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector.from_dense(query_embedding),)).fetchall() +for row in result: + print(row[0]) From 44855c65d295d4e5380fe49d3ef978bf372915c7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 14:30:39 -0400 Subject: [PATCH 242/879] Added support for Vector class to Psycopg 3 [skip ci] --- pgvector/psycopg/__init__.py | 2 +- pgvector/psycopg/vector.py | 2 ++ pgvector/utils/vector.py | 14 ++++++++++++++ tests/test_psycopg.py | 12 +++++++++++- 4 files changed, 28 insertions(+), 2 deletions(-) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 90ff8b7..d9e2a95 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -4,7 +4,7 @@ from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import Bit, HalfVector, SparseVector +from ..utils import Bit, HalfVector, SparseVector, Vector __all__ = ['register_vector'] diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index fe606a6..4875fe2 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -51,5 +51,7 @@ def register_vector_info(context, info): adapters = context.adapters adapters.register_dumper('numpy.ndarray', text_dumper) adapters.register_dumper('numpy.ndarray', binary_dumper) + adapters.register_dumper(Vector, text_dumper) + adapters.register_dumper(Vector, binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index c172b5e..3f44e0d 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -3,6 +3,15 @@ class Vector: + def __init__(self, value): + if isinstance(value, np.ndarray): + value = value.tolist() + + if not isinstance(value, (list, tuple)): + raise ValueError('expected list or tuple') + + self.value = value + def from_db(value): # could be ndarray if already cast by lower-level driver if value is None or isinstance(value, np.ndarray): @@ -29,6 +38,8 @@ def to_db(value, dim=None): raise ValueError('dtype must be numeric') value = value.tolist() + elif isinstance(value, Vector): + value = value.value if dim is not None and len(value) != dim: raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) @@ -39,6 +50,9 @@ def to_db_binary(value): if value is None: return value + if isinstance(value, Vector): + value = value.value + value = np.asarray(value, dtype='>f') if value.ndim != 1: diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index a11b73c..956a25d 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector +from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector, Vector import psycopg import pytest @@ -52,6 +52,16 @@ def test_vector_binary_format_non_contiguous(self): res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] assert np.array_equal(res, np.array([3, 2, 1.5])) + def test_vector_class_binary_format(self): + embedding = Vector([1.5, 2, 3]) + res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] + assert np.array_equal(res, np.array([1.5, 2, 3])) + + def test_vector_class_text_format(self): + embedding = Vector([1.5, 2, 3]) + res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] + assert np.array_equal(res, np.array([1.5, 2, 3])) + def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) From e5a8284486c09fb92445e06faec412458247e670 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 14:51:10 -0400 Subject: [PATCH 243/879] Improved __all__ [skip ci] --- pgvector/asyncpg/__init__.py | 6 +++++- pgvector/psycopg/__init__.py | 8 +++++++- pgvector/psycopg2/__init__.py | 8 ++++++-- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 2252c7a..67d6f5c 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,6 +1,10 @@ from ..utils import Vector, HalfVector, SparseVector -__all__ = ['register_vector'] +__all__ = [ + 'register_vector', + 'HalfVector', + 'SparseVector' +] async def register_vector(conn): diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index d9e2a95..20ff68f 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -7,7 +7,13 @@ from ..utils import Bit, HalfVector, SparseVector, Vector -__all__ = ['register_vector'] +__all__ = [ + 'register_vector', + 'Vector', + 'HalfVector', + 'Bit', + 'SparseVector' +] def register_vector(context): diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 9d0473a..7f34d7c 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -2,9 +2,13 @@ from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -from ..utils import SparseVector +from ..utils import HalfVector, SparseVector -__all__ = ['register_vector'] +__all__ = [ + 'register_vector', + 'HalfVector', + 'SparseVector' +] def register_vector(conn_or_curs=None): From 6f8a2cfdd849ea917272efefea69f1163e209ba2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 14:53:27 -0400 Subject: [PATCH 244/879] Improved code [skip ci] --- pgvector/asyncpg/__init__.py | 24 +-------------------- pgvector/asyncpg/register.py | 24 +++++++++++++++++++++ pgvector/psycopg/__init__.py | 41 ++---------------------------------- pgvector/psycopg/register.py | 38 +++++++++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 62 deletions(-) create mode 100644 pgvector/asyncpg/register.py create mode 100644 pgvector/psycopg/register.py diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 67d6f5c..f4e8754 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,3 +1,4 @@ +from .register import register_vector from ..utils import Vector, HalfVector, SparseVector __all__ = [ @@ -5,26 +6,3 @@ 'HalfVector', 'SparseVector' ] - - -async def register_vector(conn): - await conn.set_type_codec( - 'vector', - encoder=Vector.to_db_binary, - decoder=Vector.from_db_binary, - format='binary' - ) - - await conn.set_type_codec( - 'halfvec', - encoder=HalfVector.to_db_binary, - decoder=HalfVector.from_db_binary, - format='binary' - ) - - await conn.set_type_codec( - 'sparsevec', - encoder=SparseVector.to_db_binary, - decoder=SparseVector.from_db_binary, - format='binary' - ) diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py new file mode 100644 index 0000000..4664218 --- /dev/null +++ b/pgvector/asyncpg/register.py @@ -0,0 +1,24 @@ +from ..utils import Vector, HalfVector, SparseVector + + +async def register_vector(conn): + await conn.set_type_codec( + 'vector', + encoder=Vector.to_db_binary, + decoder=Vector.from_db_binary, + format='binary' + ) + + await conn.set_type_codec( + 'halfvec', + encoder=HalfVector.to_db_binary, + decoder=HalfVector.from_db_binary, + format='binary' + ) + + await conn.set_type_codec( + 'sparsevec', + encoder=SparseVector.to_db_binary, + decoder=SparseVector.from_db_binary, + format='binary' + ) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 20ff68f..9007c37 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,48 +1,11 @@ -import psycopg -from psycopg.types import TypeInfo -from .bit import register_bit_info -from .halfvec import register_halfvec_info -from .sparsevec import register_sparsevec_info -from .vector import register_vector_info +from .register import register_vector, register_vector_async from ..utils import Bit, HalfVector, SparseVector, Vector - __all__ = [ 'register_vector', + 'register_vector_async', 'Vector', 'HalfVector', 'Bit', 'SparseVector' ] - - -def register_vector(context): - info = TypeInfo.fetch(context, 'vector') - register_vector_info(context, info) - - info = TypeInfo.fetch(context, 'bit') - register_bit_info(context, info) - - info = TypeInfo.fetch(context, 'halfvec') - if info is not None: - register_halfvec_info(context, info) - - info = TypeInfo.fetch(context, 'sparsevec') - if info is not None: - register_sparsevec_info(context, info) - - -async def register_vector_async(context): - info = await TypeInfo.fetch(context, 'vector') - register_vector_info(context, info) - - info = await TypeInfo.fetch(context, 'bit') - register_bit_info(context, info) - - info = await TypeInfo.fetch(context, 'halfvec') - if info is not None: - register_halfvec_info(context, info) - - info = await TypeInfo.fetch(context, 'sparsevec') - if info is not None: - register_sparsevec_info(context, info) diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py new file mode 100644 index 0000000..7f54a31 --- /dev/null +++ b/pgvector/psycopg/register.py @@ -0,0 +1,38 @@ +import psycopg +from psycopg.types import TypeInfo +from .bit import register_bit_info +from .halfvec import register_halfvec_info +from .sparsevec import register_sparsevec_info +from .vector import register_vector_info + + +def register_vector(context): + info = TypeInfo.fetch(context, 'vector') + register_vector_info(context, info) + + info = TypeInfo.fetch(context, 'bit') + register_bit_info(context, info) + + info = TypeInfo.fetch(context, 'halfvec') + if info is not None: + register_halfvec_info(context, info) + + info = TypeInfo.fetch(context, 'sparsevec') + if info is not None: + register_sparsevec_info(context, info) + + +async def register_vector_async(context): + info = await TypeInfo.fetch(context, 'vector') + register_vector_info(context, info) + + info = await TypeInfo.fetch(context, 'bit') + register_bit_info(context, info) + + info = await TypeInfo.fetch(context, 'halfvec') + if info is not None: + register_halfvec_info(context, info) + + info = await TypeInfo.fetch(context, 'sparsevec') + if info is not None: + register_sparsevec_info(context, info) From db250ee221ebf8aab8933ba37a48c1c6826a33a8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 15:00:05 -0400 Subject: [PATCH 245/879] Improved code [skip ci] --- pgvector/psycopg2/__init__.py | 27 +-------------------------- pgvector/psycopg2/register.py | 26 ++++++++++++++++++++++++++ pgvector/sqlalchemy/__init__.py | 11 +++++++++-- 3 files changed, 36 insertions(+), 28 deletions(-) create mode 100644 pgvector/psycopg2/register.py diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 7f34d7c..7c95295 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,7 +1,4 @@ -import psycopg2 -from .halfvec import register_halfvec_info -from .sparsevec import register_sparsevec_info -from .vector import register_vector_info +from .register import register_vector from ..utils import HalfVector, SparseVector __all__ = [ @@ -9,25 +6,3 @@ 'HalfVector', 'SparseVector' ] - - -def register_vector(conn_or_curs=None): - cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs - - try: - cur.execute('SELECT NULL::vector') - register_vector_info(cur.description[0][1]) - except psycopg2.errors.UndefinedObject: - raise psycopg2.ProgrammingError('vector type not found in the database') - - try: - cur.execute('SELECT NULL::halfvec') - register_halfvec_info(cur.description[0][1]) - except psycopg2.errors.UndefinedObject: - pass - - try: - cur.execute('SELECT NULL::sparsevec') - register_sparsevec_info(cur.description[0][1]) - except psycopg2.errors.UndefinedObject: - pass diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py new file mode 100644 index 0000000..0ffd461 --- /dev/null +++ b/pgvector/psycopg2/register.py @@ -0,0 +1,26 @@ +import psycopg2 +from .halfvec import register_halfvec_info +from .sparsevec import register_sparsevec_info +from .vector import register_vector_info + + +def register_vector(conn_or_curs=None): + cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs + + try: + cur.execute('SELECT NULL::vector') + register_vector_info(cur.description[0][1]) + except psycopg2.errors.UndefinedObject: + raise psycopg2.ProgrammingError('vector type not found in the database') + + try: + cur.execute('SELECT NULL::halfvec') + register_halfvec_info(cur.description[0][1]) + except psycopg2.errors.UndefinedObject: + pass + + try: + cur.execute('SELECT NULL::sparsevec') + register_sparsevec_info(cur.description[0][1]) + except psycopg2.errors.UndefinedObject: + pass diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index fde5752..355a2d7 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -2,6 +2,13 @@ from .halfvec import Halfvec from .sparsevec import Sparsevec from .vector import Vector -from ..utils import SparseVector +from ..utils import HalfVector, SparseVector -__all__ = ['Vector', 'Halfvec', 'Bit', 'Sparsevec'] +__all__ = [ + 'Vector', + 'Halfvec', + 'Bit', + 'Sparsevec', + 'HalfVector', + 'SparseVector' +] From e0c84e4c0f8def287aa45327bf809fc2b8760abe Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 15:03:54 -0400 Subject: [PATCH 246/879] Improved code [skip ci] --- pgvector/utils/sparsevec.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 0c6885c..6860740 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -2,15 +2,6 @@ from struct import pack, unpack_from -def to_db_value(value): - if isinstance(value, SparseVector): - return value - elif isinstance(value, (list, np.ndarray)): - return SparseVector.from_dense(value) - else: - raise ValueError('expected sparsevec') - - class SparseVector: def __init__(self, dim, indices, values): self.dim = dim @@ -35,7 +26,7 @@ def to_db(value, dim=None): if value is None: return value - value = to_db_value(value) + value = __class__.to_db_value(value) if dim is not None and value.dim != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim)) @@ -46,10 +37,18 @@ def to_db_binary(value): if value is None: return value - value = to_db_value(value) + value = __class__.to_db_value(value) nnz = len(value.indices) return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) + def to_db_value(value): + if isinstance(value, SparseVector): + return value + elif isinstance(value, (list, np.ndarray)): + return SparseVector.from_dense(value) + else: + raise ValueError('expected sparsevec') + def from_db(value): if value is None or isinstance(value, SparseVector): return value From 1aec1c0e17a0dbe0617436b45626dee84a84a4d7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 15:07:31 -0400 Subject: [PATCH 247/879] Improved code [skip ci] --- pgvector/utils/bit.py | 12 ++++++------ pgvector/utils/halfvec.py | 6 +++--- pgvector/utils/sparsevec.py | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index ef639ba..88dfe4a 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -13,6 +13,12 @@ def __init__(self, value): else: self.value = np.array(value, dtype=bool) + def __str__(self): + return self.__class__.to_db(self) + + def __repr__(self): + return f'Bit({self})' + def to_db(value): if not isinstance(value, Bit): raise ValueError('expected bit') @@ -26,9 +32,3 @@ def to_db_binary(value): value = value.value return pack('>i', len(value)) + np.packbits(value).tobytes() - - def __str__(self): - return self.__class__.to_db(self) - - def __repr__(self): - return f'Bit({self})' diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index b977291..734bb25 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -15,6 +15,9 @@ def __init__(self, value): def to_list(self): return list(self.value) + def __repr__(self): + return f'HalfVector({self.value})' + def to_db(value, dim=None): if value is None: return value @@ -43,6 +46,3 @@ def from_db_binary(value): return value dim, unused = unpack_from('>HH', value) return HalfVector(unpack_from(f'>{dim}e', value, 4)) - - def __repr__(self): - return f'HalfVector({self.value})' diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 6860740..79ac80c 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -22,6 +22,9 @@ def to_dense(self): vec[i] = v return vec + def __repr__(self): + return f'SparseVector({self.dim}, {self.indices}, {self.values})' + def to_db(value, dim=None): if value is None: return value @@ -68,6 +71,3 @@ def from_db_binary(value): indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) return SparseVector(int(dim), indices, values) - - def __repr__(self): - return f'SparseVector({self.dim}, {self.indices}, {self.values})' From 739afcd7c790edf62d67f4391a0919c678e73ac1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 23 May 2024 15:07:56 -0400 Subject: [PATCH 248/879] Revert "Improved code [skip ci]" This reverts commit e0c84e4c0f8def287aa45327bf809fc2b8760abe. --- pgvector/utils/sparsevec.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 79ac80c..cbd37c7 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -2,6 +2,15 @@ from struct import pack, unpack_from +def to_db_value(value): + if isinstance(value, SparseVector): + return value + elif isinstance(value, (list, np.ndarray)): + return SparseVector.from_dense(value) + else: + raise ValueError('expected sparsevec') + + class SparseVector: def __init__(self, dim, indices, values): self.dim = dim @@ -29,7 +38,7 @@ def to_db(value, dim=None): if value is None: return value - value = __class__.to_db_value(value) + value = to_db_value(value) if dim is not None and value.dim != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim)) @@ -40,18 +49,10 @@ def to_db_binary(value): if value is None: return value - value = __class__.to_db_value(value) + value = to_db_value(value) nnz = len(value.indices) return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) - def to_db_value(value): - if isinstance(value, SparseVector): - return value - elif isinstance(value, (list, np.ndarray)): - return SparseVector.from_dense(value) - else: - raise ValueError('expected sparsevec') - def from_db(value): if value is None or isinstance(value, SparseVector): return value From c13d9bc5d52a401917290c21d7188003655a0fd6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 31 May 2024 19:23:32 -0700 Subject: [PATCH 249/879] Improved code [skip ci] --- pgvector/psycopg2/halfvec.py | 2 +- pgvector/psycopg2/sparsevec.py | 2 +- pgvector/psycopg2/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index e5ec111..b480911 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -2,7 +2,7 @@ from ..utils import HalfVector -class HalfvecAdapter(object): +class HalfvecAdapter: def __init__(self, value): self._value = value diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 4bcdd92..31e86f3 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -2,7 +2,7 @@ from ..utils import SparseVector -class SparsevecAdapter(object): +class SparsevecAdapter: def __init__(self, value): self._value = value diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index c0b066a..4fdb077 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -3,7 +3,7 @@ from ..utils import Vector -class VectorAdapter(object): +class VectorAdapter: def __init__(self, value): self._value = value From ba45eb23dd2eda0e8e15c4b5874c2cba5a577f81 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 14:55:50 -0700 Subject: [PATCH 250/879] Switched to uppercase for SQLAlchemy data types --- pgvector/sqlalchemy/__init__.py | 16 +++++++++------- pgvector/sqlalchemy/bit.py | 4 ++-- pgvector/sqlalchemy/halfvec.py | 4 ++-- pgvector/sqlalchemy/sparsevec.py | 4 ++-- pgvector/sqlalchemy/vector.py | 12 ++++++------ tests/test_sqlalchemy.py | 20 ++++++++++---------- tests/test_sqlmodel.py | 10 +++++----- 7 files changed, 36 insertions(+), 34 deletions(-) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 355a2d7..67b1d16 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,14 +1,16 @@ -from .bit import Bit -from .halfvec import Halfvec -from .sparsevec import Sparsevec -from .vector import Vector +from .bit import BIT +from .halfvec import HALFVEC +from .sparsevec import SPARSEVEC +from .vector import VECTOR +from .vector import VECTOR as Vector from ..utils import HalfVector, SparseVector __all__ = [ 'Vector', - 'Halfvec', - 'Bit', - 'Sparsevec', + 'VECTOR' + 'HALFVEC', + 'BIT', + 'SPARSEVEC', 'HalfVector', 'SparseVector' ] diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index f71c1d0..0f83f3c 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -2,7 +2,7 @@ from sqlalchemy.types import UserDefinedType, Float -class Bit(UserDefinedType): +class BIT(UserDefinedType): cache_ok = True def __init__(self, length=None): @@ -23,4 +23,4 @@ def jaccard_distance(self, other): # for reflection -ischema_names['bit'] = Bit +ischema_names['bit'] = BIT diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index cafc6c6..1c5f8f1 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -3,7 +3,7 @@ from ..utils import HalfVector -class Halfvec(UserDefinedType): +class HALFVEC(UserDefinedType): cache_ok = True _string = String() @@ -48,4 +48,4 @@ def l1_distance(self, other): # for reflection -ischema_names['halfvec'] = Halfvec +ischema_names['halfvec'] = HALFVEC diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 499dff0..ca5e9aa 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -3,7 +3,7 @@ from ..utils import SparseVector -class Sparsevec(UserDefinedType): +class SPARSEVEC(UserDefinedType): cache_ok = True _string = String() @@ -48,4 +48,4 @@ def l1_distance(self, other): # for reflection -ischema_names['sparsevec'] = Sparsevec +ischema_names['sparsevec'] = SPARSEVEC diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 056dcef..b2ce7f7 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,9 +1,9 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import Vector as Vec +from ..utils import Vector -class Vector(UserDefinedType): +class VECTOR(UserDefinedType): cache_ok = True _string = String() @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return Vec.to_db(value, self.dim) + return Vector.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(Vec.to_db(value, self.dim)) + return string_literal_processor(Vector.to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return Vec.from_db(value) + return Vector.from_db(value) return process class comparator_factory(UserDefinedType.Comparator): @@ -48,4 +48,4 @@ def l1_distance(self, other): # for reflection -ischema_names['vector'] = Vector +ischema_names['vector'] = VECTOR diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 1fff283..52b3e36 100644 --- a/tests/test_sqlalchemy.py +++ b/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 import pytest from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError @@ -19,10 +19,10 @@ class Item(Base): __tablename__ = 'sqlalchemy_orm_item' id = mapped_column(Integer, primary_key=True) - embedding = mapped_column(Vector(3)) - half_embedding = mapped_column(Halfvec(3)) - binary_embedding = mapped_column(Bit(3)) - sparse_embedding = mapped_column(Sparsevec(3)) + embedding = mapped_column(VECTOR(3)) + half_embedding = mapped_column(HALFVEC(3)) + binary_embedding = mapped_column(BIT(3)) + sparse_embedding = mapped_column(SPARSEVEC(3)) Base.metadata.drop_all(engine) @@ -59,10 +59,10 @@ def test_core(self): 'sqlalchemy_core_item', metadata, Column('id', Integer, primary_key=True), - Column('embedding', Vector(3)), - Column('half_embedding', Halfvec(3)), - Column('binary_embedding', Bit(3)), - Column('sparse_embedding', Sparsevec(3)) + Column('embedding', VECTOR(3)), + Column('half_embedding', HALFVEC(3)), + Column('binary_embedding', BIT(3)), + Column('sparse_embedding', SPARSEVEC(3)) ) metadata.drop_all(engine) @@ -388,7 +388,7 @@ def test_bad_dtype(self): def test_inspect(self): columns = inspect(engine).get_columns('sqlalchemy_orm_item') - assert isinstance(columns[1]['type'], Vector) + assert isinstance(columns[1]['type'], VECTOR) def test_literal_binds(self): sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 8563bec..1b6baaf 100644 --- a/tests/test_sqlmodel.py +++ b/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 import pytest from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError @@ -16,10 +16,10 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[Any] = Field(default=None, sa_column=Column(Vector(3))) - half_embedding: Optional[Any] = Field(default=None, sa_column=Column(Halfvec(3))) - binary_embedding: Optional[Any] = Field(default=None, sa_column=Column(Bit(3))) - sparse_embedding: Optional[Any] = Field(default=None, sa_column=Column(Sparsevec(3))) + embedding: Optional[Any] = Field(default=None, sa_column=Column(VECTOR(3))) + half_embedding: Optional[Any] = Field(default=None, sa_column=Column(HALFVEC(3))) + binary_embedding: Optional[Any] = Field(default=None, sa_column=Column(BIT(3))) + sparse_embedding: Optional[Any] = Field(default=None, sa_column=Column(SPARSEVEC(3))) SQLModel.metadata.drop_all(engine) From 4010076f7a7dc1641210da10acff2c17e73137b5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 15:57:45 -0700 Subject: [PATCH 251/879] Updated readme [skip ci] --- README.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 710525a..cf39d7e 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,8 @@ class Item(models.Model): embedding = VectorField(dimensions=3) ``` +Also supports `HalfVectorField` (unreleased), `BitField` (unreleased), and `SparseVectorField` (unreleased) + Insert a vector ```python @@ -76,7 +78,7 @@ from pgvector.django import L2Distance Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5] ``` -Also supports `MaxInnerProduct` and `CosineDistance` +Also supports `MaxInnerProduct`, `CosineDistance`, `L1Distance` (unreleased), `HammingDistance` (unreleased), and `JaccardDistance` (unreleased) Get the distance @@ -144,6 +146,8 @@ class Item(Base): embedding = mapped_column(Vector(3)) ``` +Also supports `HALFVEC` (unreleased), `BIT` (unreleased), and `SPARSEVEC` (unreleased) + Insert a vector ```python @@ -158,7 +162,7 @@ Get the nearest neighbors to a vector session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) ``` -Also supports `max_inner_product` and `cosine_distance` +Also supports `max_inner_product`, `cosine_distance`, `l1_distance` (unreleased), `hamming_distance` (unreleased), and `jaccard_distance` (unreleased) Get the distance @@ -224,6 +228,8 @@ class Item(SQLModel, table=True): embedding: Any = Field(sa_column=Column(Vector(3))) ``` +Also supports `HALFVEC` (unreleased), `BIT` (unreleased), and `SPARSEVEC` (unreleased) + Insert a vector ```python @@ -238,7 +244,7 @@ Get the nearest neighbors to a vector session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) ``` -Also supports `max_inner_product` and `cosine_distance` +Also supports `max_inner_product`, `cosine_distance`, `l1_distance` (unreleased), `hamming_distance` (unreleased), and `jaccard_distance` (unreleased) Get the distance @@ -453,6 +459,8 @@ class Item(BaseModel): embedding = VectorField(dimensions=3) ``` +Also supports `HalfVectorField` (unreleased), `FixedBitField` (unreleased), and `SparseVectorField` (unreleased) + Insert a vector ```python @@ -465,7 +473,7 @@ Get the nearest neighbors to a vector Item.select().order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5) ``` -Also supports `max_inner_product` and `cosine_distance` +Also supports `max_inner_product`, `cosine_distance`, `l1_distance` (unreleased), `hamming_distance` (unreleased), and `jaccard_distance` (unreleased) Get the distance From c1a7312a3509932b2c39e730a7906fd44c86015c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:04:29 -0700 Subject: [PATCH 252/879] Improved bit code [skip ci] --- pgvector/utils/bit.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 88dfe4a..4c0b449 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -7,28 +7,30 @@ def __init__(self, value): if isinstance(value, bytes): count = unpack_from('>i', value)[0] buf = np.frombuffer(value[4:], dtype=np.uint8) - self.value = np.unpackbits(buf, count=count).astype(bool) + self._value = np.unpackbits(buf, count=count).astype(bool) elif isinstance(value, str): - self.value = np.array([v != '0' for v in value], dtype=bool) + self._value = np.array([v != '0' for v in value], dtype=bool) else: - self.value = np.array(value, dtype=bool) + self._value = np.array(value, dtype=bool) def __str__(self): - return self.__class__.to_db(self) + return ''.join(self._value.astype(np.uint8).astype(str)) def __repr__(self): return f'Bit({self})' + def to_binary(self): + value = self._value + return pack('>i', len(value)) + np.packbits(value).tobytes() + def to_db(value): if not isinstance(value, Bit): raise ValueError('expected bit') - value = value.value - return ''.join(value.astype(np.uint8).astype(str)) + return str(value) def to_db_binary(value): if not isinstance(value, Bit): raise ValueError('expected bit') - value = value.value - return pack('>i', len(value)) + np.packbits(value).tobytes() + return value.to_binary() From e93e54ee10759ff608d33c68cd10bb5dd6d9319b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:06:21 -0700 Subject: [PATCH 253/879] Improved bit code [skip ci] --- pgvector/utils/bit.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 4c0b449..878273d 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -14,11 +14,14 @@ def __init__(self, value): self._value = np.array(value, dtype=bool) def __str__(self): - return ''.join(self._value.astype(np.uint8).astype(str)) + return self.to_text() def __repr__(self): return f'Bit({self})' + def to_text(self): + return ''.join(self._value.astype(np.uint8).astype(str)) + def to_binary(self): value = self._value return pack('>i', len(value)) + np.packbits(value).tobytes() @@ -27,7 +30,7 @@ def to_db(value): if not isinstance(value, Bit): raise ValueError('expected bit') - return str(value) + return value.to_text() def to_db_binary(value): if not isinstance(value, Bit): From 1efe8154772c9b154a21164f48a343ecf197a1b8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:10:24 -0700 Subject: [PATCH 254/879] Improved type code [skip ci] --- pgvector/utils/halfvec.py | 10 +++++----- pgvector/utils/sparsevec.py | 22 +++++++++++----------- pgvector/utils/vector.py | 6 +++--- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 734bb25..7f04088 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -10,19 +10,19 @@ def __init__(self, value): if not isinstance(value, (list, tuple)): raise ValueError('expected list or tuple') - self.value = value + self._value = value def to_list(self): - return list(self.value) + return list(self._value) def __repr__(self): - return f'HalfVector({self.value})' + return f'HalfVector({self._value})' def to_db(value, dim=None): if value is None: return value if isinstance(value, HalfVector): - value = value.value + value = value._value if dim is not None and len(value) != dim: raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) @@ -33,7 +33,7 @@ def to_db_binary(value): if value is None: return value if isinstance(value, HalfVector): - value = value.value + value = value._value return pack(f'>HH{len(value)}e', len(value), 0, *value) def from_db(value): diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index cbd37c7..1e09327 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -13,9 +13,9 @@ def to_db_value(value): class SparseVector: def __init__(self, dim, indices, values): - self.dim = dim - self.indices = indices - self.values = values + self._dim = dim + self._indices = indices + self._values = values def from_dense(value): if isinstance(value, np.ndarray): @@ -26,13 +26,13 @@ def from_dense(value): return SparseVector(dim, indices, values) def to_dense(self): - vec = [0] * self.dim - for i, v in zip(self.indices, self.values): + vec = [0] * self._dim + for i, v in zip(self._indices, self._values): vec[i] = v return vec def __repr__(self): - return f'SparseVector({self.dim}, {self.indices}, {self.values})' + return f'SparseVector({self._dim}, {self._indices}, {self._values})' def to_db(value, dim=None): if value is None: @@ -40,18 +40,18 @@ def to_db(value, dim=None): value = to_db_value(value) - if dim is not None and value.dim != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dim)) + if dim is not None and value._dim != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value._dim)) - return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(value.indices, value.values)]) + '}/' + str(value.dim) + return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(value._indices, value._values)]) + '}/' + str(value._dim) def to_db_binary(value): if value is None: return value value = to_db_value(value) - nnz = len(value.indices) - return pack(f'>iii{nnz}i{nnz}f', value.dim, nnz, 0, *value.indices, *value.values) + nnz = len(value._indices) + return pack(f'>iii{nnz}i{nnz}f', value._dim, nnz, 0, *value._indices, *value._values) def from_db(value): if value is None or isinstance(value, SparseVector): diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 3f44e0d..dd67700 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -10,7 +10,7 @@ def __init__(self, value): if not isinstance(value, (list, tuple)): raise ValueError('expected list or tuple') - self.value = value + self._value = value def from_db(value): # could be ndarray if already cast by lower-level driver @@ -39,7 +39,7 @@ def to_db(value, dim=None): value = value.tolist() elif isinstance(value, Vector): - value = value.value + value = value._value if dim is not None and len(value) != dim: raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) @@ -51,7 +51,7 @@ def to_db_binary(value): return value if isinstance(value, Vector): - value = value.value + value = value._value value = np.asarray(value, dtype='>f') From a4bf4aed85d75046eb7e28a0a4feff231c46a96b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:13:52 -0700 Subject: [PATCH 255/879] Improved HalfVector code [skip ci] --- pgvector/utils/halfvec.py | 32 +++++++++++++++++++++----------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 7f04088..9f617ec 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -12,29 +12,39 @@ def __init__(self, value): self._value = value - def to_list(self): - return list(self._value) - def __repr__(self): return f'HalfVector({self._value})' + def to_text(self): + return '[' + ','.join([str(float(v)) for v in self._value]) + ']' + + def to_binary(self): + return pack(f'>HH{len(self._value)}e', len(self._value), 0, *self._value) + + def dim(self): + return len(self._value) + + def to_list(self): + return list(self._value) + def to_db(value, dim=None): if value is None: return value - if isinstance(value, HalfVector): - value = value._value + if not isinstance(value, HalfVector): + value = HalfVector(value) - if dim is not None and len(value) != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, len(value))) + if dim is not None and value.dim() != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) - return '[' + ','.join([str(float(v)) for v in value]) + ']' + return value.to_text() def to_db_binary(value): if value is None: return value - if isinstance(value, HalfVector): - value = value._value - return pack(f'>HH{len(value)}e', len(value), 0, *value) + if not isinstance(value, HalfVector): + value = HalfVector(value) + + return value.to_binary() def from_db(value): if value is None or isinstance(value, HalfVector): From 00ad5852619151edc2d616d996151c2c0e91b807 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:16:13 -0700 Subject: [PATCH 256/879] Improved HalfVector code [skip ci] --- pgvector/utils/halfvec.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 9f617ec..e596fd6 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -27,6 +27,13 @@ def dim(self): def to_list(self): return list(self._value) + def from_text(value): + return HalfVector([float(v) for v in value[1:-1].split(',')]) + + def from_binary(value): + dim, unused = unpack_from('>HH', value) + return HalfVector(unpack_from(f'>{dim}e', value, 4)) + def to_db(value, dim=None): if value is None: return value @@ -49,10 +56,9 @@ def to_db_binary(value): def from_db(value): if value is None or isinstance(value, HalfVector): return value - return HalfVector([float(v) for v in value[1:-1].split(',')]) + return __class__.from_text(value) def from_db_binary(value): if value is None or isinstance(value, HalfVector): return value - dim, unused = unpack_from('>HH', value) - return HalfVector(unpack_from(f'>{dim}e', value, 4)) + return __class__.from_binary(value) From 63ac583a8bac60548eaacaad0f48c45ae5152207 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:21:11 -0700 Subject: [PATCH 257/879] Improved SparseVector code [skip ci] --- pgvector/utils/bit.py | 2 ++ pgvector/utils/halfvec.py | 2 ++ pgvector/utils/sparsevec.py | 51 ++++++++++++++++++++++++------------- 3 files changed, 38 insertions(+), 17 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 878273d..7fc75bf 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -26,6 +26,8 @@ def to_binary(self): value = self._value return pack('>i', len(value)) + np.packbits(value).tobytes() + # TODO move rest + def to_db(value): if not isinstance(value, Bit): raise ValueError('expected bit') diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index e596fd6..11ae7d1 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -34,6 +34,8 @@ def from_binary(value): dim, unused = unpack_from('>HH', value) return HalfVector(unpack_from(f'>{dim}e', value, 4)) + # TODO move rest + def to_db(value, dim=None): if value is None: return value diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 1e09327..1b42fef 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -17,6 +17,9 @@ def __init__(self, dim, indices, values): self._indices = indices self._values = values + def __repr__(self): + return f'SparseVector({self._dim}, {self._indices}, {self._values})' + def from_dense(value): if isinstance(value, np.ndarray): value = value.tolist() @@ -31,8 +34,30 @@ def to_dense(self): vec[i] = v return vec - def __repr__(self): - return f'SparseVector({self._dim}, {self._indices}, {self._values})' + def to_text(self): + return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(self._indices, self._values)]) + '}/' + str(self._dim) + + def to_binary(self): + nnz = len(self._indices) + return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) + + def from_text(value): + elements, dim = value.split('/') + indices = [] + values = [] + for e in elements[1:-1].split(','): + i, v = e.split(':') + indices.append(int(i) - 1) + values.append(float(v)) + return SparseVector(int(dim), indices, values) + + def from_binary(value): + dim, nnz, unused = unpack_from('>iii', value) + indices = unpack_from(f'>{nnz}i', value, 12) + values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) + return SparseVector(int(dim), indices, values) + + # TODO move rest def to_db(value, dim=None): if value is None: @@ -43,32 +68,24 @@ def to_db(value, dim=None): if dim is not None and value._dim != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value._dim)) - return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(value._indices, value._values)]) + '}/' + str(value._dim) + return value.to_text() def to_db_binary(value): if value is None: return value value = to_db_value(value) - nnz = len(value._indices) - return pack(f'>iii{nnz}i{nnz}f', value._dim, nnz, 0, *value._indices, *value._values) + + return value.to_binary() def from_db(value): if value is None or isinstance(value, SparseVector): return value - elements, dim = value.split('/') - indices = [] - values = [] - for e in elements[1:-1].split(','): - i, v = e.split(':') - indices.append(int(i) - 1) - values.append(float(v)) - return SparseVector(int(dim), indices, values) + + return __class__.from_text(value) def from_db_binary(value): if value is None or isinstance(value, SparseVector): return value - dim, nnz, unused = unpack_from('>iii', value) - indices = unpack_from(f'>{nnz}i', value, 12) - values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return SparseVector(int(dim), indices, values) + + return __class__.from_binary(value) From 585045976942fbfcabfc420f8c10bb906cc68b58 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 16:26:13 -0700 Subject: [PATCH 258/879] Improved code [skip ci] --- pgvector/utils/halfvec.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 11ae7d1..4a14aeb 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -39,6 +39,7 @@ def from_binary(value): def to_db(value, dim=None): if value is None: return value + if not isinstance(value, HalfVector): value = HalfVector(value) @@ -50,6 +51,7 @@ def to_db(value, dim=None): def to_db_binary(value): if value is None: return value + if not isinstance(value, HalfVector): value = HalfVector(value) @@ -58,9 +60,11 @@ def to_db_binary(value): def from_db(value): if value is None or isinstance(value, HalfVector): return value + return __class__.from_text(value) def from_db_binary(value): if value is None or isinstance(value, HalfVector): return value + return __class__.from_binary(value) From 3b4643ba1ada45f9ad26759eff37d944c790af22 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 17:12:46 -0700 Subject: [PATCH 259/879] Improved HalfVector code [skip ci] --- pgvector/utils/halfvec.py | 30 ++++++++++++++++-------------- tests/test_psycopg.py | 2 ++ 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 4a14aeb..590929e 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -4,35 +4,37 @@ class HalfVector: def __init__(self, value): - if isinstance(value, np.ndarray): - value = value.tolist() + value = np.asarray(value, dtype=np.float16) - if not isinstance(value, (list, tuple)): - raise ValueError('expected list or tuple') + if value.ndim != 1: + raise ValueError('expected ndim to be 1') self._value = value def __repr__(self): - return f'HalfVector({self._value})' - - def to_text(self): - return '[' + ','.join([str(float(v)) for v in self._value]) + ']' - - def to_binary(self): - return pack(f'>HH{len(self._value)}e', len(self._value), 0, *self._value) + return f'HalfVector({self.to_list()})' def dim(self): - return len(self._value) + return self._value.shape[0] def to_list(self): - return list(self._value) + return self._value.tolist() + + def to_numpy(self): + return self._value + + def to_text(self): + return '[' + ','.join([str(v) for v in self._value]) + ']' + + def to_binary(self): + return pack('>HH', self.dim(), 0) + np.array(self._value, dtype='>f2').tobytes() def from_text(value): return HalfVector([float(v) for v in value[1:-1].split(',')]) def from_binary(value): dim, unused = unpack_from('>HH', value) - return HalfVector(unpack_from(f'>{dim}e', value, 4)) + return HalfVector(np.frombuffer(value, dtype='>f2', count=dim, offset=4).astype(dtype=np.float16)) # TODO move rest diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 956a25d..a692ce0 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -73,11 +73,13 @@ def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] assert res.to_list() == [1.5, 2, 3] + assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] assert res.to_list() == [1.5, 2, 3] + assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) def test_bit(self): embedding = Bit([True, False, True]) From 7af59e7e73932ecf0ee20fec335fee9e175591cb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 17:17:55 -0700 Subject: [PATCH 260/879] Reduced copying for HalfVector --- pgvector/utils/halfvec.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 590929e..52436b8 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -4,7 +4,7 @@ class HalfVector: def __init__(self, value): - value = np.asarray(value, dtype=np.float16) + value = np.asarray(value, dtype='>f2') if value.ndim != 1: raise ValueError('expected ndim to be 1') @@ -27,14 +27,14 @@ def to_text(self): return '[' + ','.join([str(v) for v in self._value]) + ']' def to_binary(self): - return pack('>HH', self.dim(), 0) + np.array(self._value, dtype='>f2').tobytes() + return pack('>HH', self.dim(), 0) + self._value.tobytes() def from_text(value): return HalfVector([float(v) for v in value[1:-1].split(',')]) def from_binary(value): dim, unused = unpack_from('>HH', value) - return HalfVector(np.frombuffer(value, dtype='>f2', count=dim, offset=4).astype(dtype=np.float16)) + return HalfVector(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) # TODO move rest From 6f8bcfe6d5f442c786124c856ae9168896093f62 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 17:32:57 -0700 Subject: [PATCH 261/879] Reduced copying for HalfVector [skip ci] --- pgvector/utils/halfvec.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 52436b8..57ca4df 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -4,7 +4,9 @@ class HalfVector: def __init__(self, value): - value = np.asarray(value, dtype='>f2') + # asarray still copies if same dtype + if not isinstance(value, np.ndarray) or value.dtype != '>f2': + value = np.asarray(value, dtype='>f2') if value.ndim != 1: raise ValueError('expected ndim to be 1') From d1a77b5a29852d460f693caf0c3b045155977df7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 17:39:13 -0700 Subject: [PATCH 262/879] Added tests for HalfVector class --- tests/test_half_vector.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 tests/test_half_vector.py diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py new file mode 100644 index 0000000..d059458 --- /dev/null +++ b/tests/test_half_vector.py @@ -0,0 +1,26 @@ +import numpy as np +from pgvector.utils import HalfVector +import pytest + + +class TestHalfVector: + def test_list(self): + assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] + + def test_list_str(self): + with pytest.raises(ValueError) as error: + HalfVector([1, 'two', 3]) + assert str(error.value) == "could not convert string to float: 'two'" + + def test_tuple(self): + assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3] + + def test_ndarray(self): + arr = np.array([1, 2, 3]) + assert HalfVector(arr).to_list() == [1, 2, 3] + assert HalfVector(arr).to_numpy() is not arr + + def test_ndarray_same_object(self): + arr = np.array([1, 2, 3], dtype='>f2') + assert HalfVector(arr).to_list() == [1, 2, 3] + assert HalfVector(arr).to_numpy() is arr From 85ea3573184fad34bfd2e16df2abccf6a16ece7f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 17:43:27 -0700 Subject: [PATCH 263/879] Added tests for ndim [skip ci] --- tests/test_half_vector.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index d059458..9114c2c 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -24,3 +24,13 @@ def test_ndarray_same_object(self): arr = np.array([1, 2, 3], dtype='>f2') assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is arr + + def test_ndim_two(self): + with pytest.raises(ValueError) as error: + HalfVector([[1, 2], [3, 4]]) + assert str(error.value) == 'expected ndim to be 1' + + def test_ndim_zero(self): + with pytest.raises(ValueError) as error: + HalfVector(1) + assert str(error.value) == 'expected ndim to be 1' From f75cf2feadb23db05ff52a937bf9d4a1f8114be5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 18:03:17 -0700 Subject: [PATCH 264/879] Improved Vector logic [skip ci] --- pgvector/utils/vector.py | 77 +++++++++++++++++++++++---------------- tests/test_half_vector.py | 3 +- tests/test_sqlalchemy.py | 2 +- 3 files changed, 47 insertions(+), 35 deletions(-) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index dd67700..adecda3 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -4,58 +4,71 @@ class Vector: def __init__(self, value): - if isinstance(value, np.ndarray): - value = value.tolist() + # asarray still copies if same dtype + if not isinstance(value, np.ndarray) or value.dtype != '>f4': + value = np.asarray(value, dtype='>f4') - if not isinstance(value, (list, tuple)): - raise ValueError('expected list or tuple') + if value.ndim != 1: + raise ValueError('expected ndim to be 1') self._value = value - def from_db(value): - # could be ndarray if already cast by lower-level driver - if value is None or isinstance(value, np.ndarray): - return value + def __repr__(self): + return f'Vector({self.to_list()})' - return np.array(value[1:-1].split(','), dtype=np.float32) + def dim(self): + return self._value.shape[0] - def from_db_binary(value): - if value is None or isinstance(value, np.ndarray): - return value + def to_list(self): + return self._value.tolist() + + def to_numpy(self): + return self._value + + def to_text(self): + return '[' + ','.join([str(v) for v in self._value]) + ']' + def to_binary(self): + return pack('>HH', self.dim(), 0) + self._value.tobytes() + + def from_text(value): + return Vector([float(v) for v in value[1:-1].split(',')]) + + def from_binary(value): dim, unused = unpack_from('>HH', value) - return np.frombuffer(value, dtype='>f', count=dim, offset=4).astype(dtype=np.float32) + return Vector(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) + + # TODO move rest def to_db(value, dim=None): if value is None: return value - if isinstance(value, np.ndarray): - if value.ndim != 1: - raise ValueError('expected ndim to be 1') + if not isinstance(value, Vector): + value = Vector(value) - if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating): - raise ValueError('dtype must be numeric') + if dim is not None and value.dim() != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) - value = value.tolist() - elif isinstance(value, Vector): - 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]) + ']' + return value.to_text() def to_db_binary(value): if value is None: return value - if isinstance(value, Vector): - value = value._value + if not isinstance(value, Vector): + value = Vector(value) - value = np.asarray(value, dtype='>f') + return value.to_binary() - if value.ndim != 1: - raise ValueError('expected ndim to be 1') + def from_db(value): + if value is None or isinstance(value, np.ndarray): + return value + + return __class__.from_text(value).to_numpy().astype(np.float32) + + def from_db_binary(value): + if value is None or isinstance(value, np.ndarray): + return value - return pack('>HH', value.shape[0], 0) + value.tobytes() + return __class__.from_binary(value).to_numpy().astype(np.float32) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 9114c2c..3a3f3ca 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -8,9 +8,8 @@ def test_list(self): assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] def test_list_str(self): - with pytest.raises(ValueError) as error: + with pytest.raises(ValueError, match='could not convert string to float') as error: HalfVector([1, 'two', 3]) - assert str(error.value) == "could not convert string to float: 'two'" def test_tuple(self): assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3] diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 52b3e36..92bc89a 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -383,7 +383,7 @@ def test_bad_dtype(self): item = Item(embedding=np.array(['one', 'two', 'three'])) session = Session(engine) session.add(item) - with pytest.raises(StatementError, match='dtype must be numeric'): + with pytest.raises(StatementError, match='could not convert string to float'): session.commit() def test_inspect(self): From 7aef0f1ccc683fddf1cd5e4367782632b1f93eb4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 18:59:20 -0700 Subject: [PATCH 265/879] Renamed to_dense to to_list and added to_numpy to SparseVector --- pgvector/utils/sparsevec.py | 10 ++++++++-- tests/test_asyncpg.py | 2 +- tests/test_django.py | 2 +- tests/test_peewee.py | 2 +- tests/test_psycopg.py | 8 +++++--- tests/test_psycopg2.py | 2 +- tests/test_sqlalchemy.py | 2 +- tests/test_sqlmodel.py | 2 +- 8 files changed, 19 insertions(+), 11 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 1b42fef..611bb16 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -28,8 +28,14 @@ def from_dense(value): values = [value[i] for i in indices] return SparseVector(dim, indices, values) - def to_dense(self): - vec = [0] * self._dim + def to_list(self): + vec = [0.0] * self._dim + for i, v in zip(self._indices, self._values): + vec[i] = v + return vec + + def to_numpy(self): + vec = np.repeat(0.0, self._dim).astype(np.float32) for i, v in zip(self._indices, self._values): vec[i] = v return vec diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 4f0a4bd..3bfc888 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -86,7 +86,7 @@ async def test_sparsevec(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['embedding'].to_dense() == [1.5, 2, 3] + assert res[0]['embedding'].to_list() == [1.5, 2, 3] assert res[1]['embedding'] is None # ensures binary format is correct diff --git a/tests/test_django.py b/tests/test_django.py index b40a47a..1c611d9 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -192,7 +192,7 @@ def test_bit_jaccard_distance(self): def test_sparsevec(self): Item(id=1, sparse_embedding=SparseVector.from_dense([1, 2, 3])).save() item = Item.objects.get(pk=1) - assert item.sparse_embedding.to_dense() == [1, 2, 3] + assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() diff --git a/tests/test_peewee.py b/tests/test_peewee.py index deb571f..1455303 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -128,7 +128,7 @@ def test_bit_jaccard_distance(self): def test_sparsevec(self): Item.create(id=1, sparse_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.sparse_embedding.to_dense() == [1, 2, 3] + assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index a692ce0..17483f7 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -106,17 +106,19 @@ def test_sparsevec(self): conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] - assert res.to_dense() == [1.5, 2, 3] + assert res.to_list() == [1.5, 2, 3] def test_sparsevec_binary_format(self): embedding = SparseVector.from_dense([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] - assert res.to_dense() == [1.5, 0, 2, 0, 3, 0] + assert res.to_list() == [1.5, 0, 2, 0, 3, 0] + assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_sparsevec_text_format(self): embedding = SparseVector.from_dense([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] - assert res.to_dense() == [1.5, 0, 2, 0, 3, 0] + assert res.to_list() == [1.5, 0, 2, 0, 3, 0] + assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_text_copy(self): embedding = np.array([1.5, 2, 3]) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index cc7e5c0..f18405f 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -51,5 +51,5 @@ def test_sparsevec(self): cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0].to_dense() == [1.5, 2, 3] + assert res[0][0].to_list() == [1.5, 2, 3] assert res[1][0] is None diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 92bc89a..687a0d0 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -255,7 +255,7 @@ def test_sparsevec(self): session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.sparse_embedding.to_dense() == [1, 2, 3] + assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 1b6baaf..90f7e21 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -158,7 +158,7 @@ def test_sparsevec(self): session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.sparse_embedding.to_dense() == [1, 2, 3] + assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() From 323b0ab3ee21c09134e2477a982f8b732fe6c94b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:01:50 -0700 Subject: [PATCH 266/879] Addd test for __repr__ [skip ci] --- tests/test_half_vector.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 3a3f3ca..c60433d 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -33,3 +33,6 @@ def test_ndim_zero(self): with pytest.raises(ValueError) as error: HalfVector(1) assert str(error.value) == 'expected ndim to be 1' + + def test_repr(self): + assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' From 241b087e797e84a3a585fb5b6261a6e706602a87 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:03:09 -0700 Subject: [PATCH 267/879] Added tests for Vector class [skip ci] --- tests/test_vector.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tests/test_vector.py diff --git a/tests/test_vector.py b/tests/test_vector.py new file mode 100644 index 0000000..f424019 --- /dev/null +++ b/tests/test_vector.py @@ -0,0 +1,38 @@ +import numpy as np +from pgvector.utils import Vector +import pytest + + +class TestVector: + def test_list(self): + assert Vector([1, 2, 3]).to_list() == [1, 2, 3] + + def test_list_str(self): + with pytest.raises(ValueError, match='could not convert string to float') as error: + Vector([1, 'two', 3]) + + def test_tuple(self): + assert Vector((1, 2, 3)).to_list() == [1, 2, 3] + + def test_ndarray(self): + arr = np.array([1, 2, 3]) + assert Vector(arr).to_list() == [1, 2, 3] + assert Vector(arr).to_numpy() is not arr + + def test_ndarray_same_object(self): + arr = np.array([1, 2, 3], dtype='>f4') + assert Vector(arr).to_list() == [1, 2, 3] + assert Vector(arr).to_numpy() is arr + + def test_ndim_two(self): + with pytest.raises(ValueError) as error: + Vector([[1, 2], [3, 4]]) + assert str(error.value) == 'expected ndim to be 1' + + def test_ndim_zero(self): + with pytest.raises(ValueError) as error: + Vector(1) + assert str(error.value) == 'expected ndim to be 1' + + def test_repr(self): + assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' From e3708976aad6bc2e7e980b139d2e64df6582a1f8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:05:19 -0700 Subject: [PATCH 268/879] Added tests for Bit class [skip ci] --- tests/test_bit.py | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 tests/test_bit.py diff --git a/tests/test_bit.py b/tests/test_bit.py new file mode 100644 index 0000000..35dbc82 --- /dev/null +++ b/tests/test_bit.py @@ -0,0 +1,11 @@ +import numpy as np +from pgvector.utils import Bit +import pytest + + +class TestBit: + def test_list(self): + assert str(Bit([True, False, True])) == '101' + + def test_str(self): + assert str(Bit('101')) == '101' From a64323d6e6919caa8b509415b1acc555fd94732f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:08:33 -0700 Subject: [PATCH 269/879] Added test for SparseVector class [skip ci] --- tests/test_sparse_vector.py | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 tests/test_sparse_vector.py diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py new file mode 100644 index 0000000..488908c --- /dev/null +++ b/tests/test_sparse_vector.py @@ -0,0 +1,9 @@ +import numpy as np +from pgvector.utils import SparseVector +import pytest + + +class TestSparseVector: + def test_from_dense(self): + assert SparseVector.from_dense([1, 2, 3]).to_list() == [1, 2, 3] + assert SparseVector.from_dense([1, 2, 3]).to_numpy().tolist() == [1, 2, 3] From 140e444742de020982cba3ef908388c9a8b29680 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:12:14 -0700 Subject: [PATCH 270/879] Improved Bit code [skip ci] --- pgvector/utils/bit.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 7fc75bf..1ecbb4c 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -5,13 +5,11 @@ class Bit: def __init__(self, value): if isinstance(value, bytes): - count = unpack_from('>i', value)[0] - buf = np.frombuffer(value[4:], dtype=np.uint8) - self._value = np.unpackbits(buf, count=count).astype(bool) + self._value = __class__.from_binary(value)._value elif isinstance(value, str): - self._value = np.array([v != '0' for v in value], dtype=bool) + self._value = __class__.from_text(value)._value else: - self._value = np.array(value, dtype=bool) + self._value = np.asarray(value, dtype=bool) def __str__(self): return self.to_text() @@ -26,6 +24,14 @@ def to_binary(self): value = self._value return pack('>i', len(value)) + np.packbits(value).tobytes() + def from_text(value): + return Bit(np.asarray([v != '0' for v in value], dtype=bool)) + + def from_binary(value): + count = unpack_from('>i', value)[0] + buf = np.frombuffer(value[4:], dtype=np.uint8) + return Bit(np.unpackbits(buf, count=count).astype(bool)) + # TODO move rest def to_db(value): From 343ee3e4e5584d83419f16ffed1fbd18a8a69fbc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:19:15 -0700 Subject: [PATCH 271/879] Added more methods to Bit class [skip ci] --- pgvector/utils/bit.py | 9 +++++++++ tests/test_bit.py | 10 +++++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 1ecbb4c..ce7dd2d 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -17,6 +17,15 @@ def __str__(self): def __repr__(self): return f'Bit({self})' + def dim(self): + return self._value.shape[0] + + def to_list(self): + return self._value.tolist() + + def to_numpy(self): + return self._value + def to_text(self): return ''.join(self._value.astype(np.uint8).astype(str)) diff --git a/tests/test_bit.py b/tests/test_bit.py index 35dbc82..8de6a2e 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -5,7 +5,15 @@ class TestBit: def test_list(self): - assert str(Bit([True, False, True])) == '101' + assert Bit([True, False, True]).to_list() == [True, False, True] + + def test_tuple(self): + assert Bit((True, False, True)).to_list() == [True, False, True] def test_str(self): assert str(Bit('101')) == '101' + + def test_ndarray_same_object(self): + arr = np.array([True, False, True], dtype=bool) + assert Bit(arr).to_list() == [True, False, True] + assert Bit(arr).to_numpy() is arr From 22c81a4d833bd33e84ad71fd679a4954b0efba71 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:19:34 -0700 Subject: [PATCH 272/879] Simplified test [skip ci] --- tests/test_bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 8de6a2e..1f744c6 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -14,6 +14,6 @@ def test_str(self): assert str(Bit('101')) == '101' def test_ndarray_same_object(self): - arr = np.array([True, False, True], dtype=bool) + arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert Bit(arr).to_numpy() is arr From 8fac4979aac664da6f93fd690a2290c16e9b8869 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:20:47 -0700 Subject: [PATCH 273/879] Added test for __repr__ [skip ci] --- tests/test_bit.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index 1f744c6..e2006b9 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -17,3 +17,6 @@ def test_ndarray_same_object(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert Bit(arr).to_numpy() is arr + + def test_repr(self): + assert repr(Bit([True, False, True])) == 'Bit(101)' From 7e17e9e3091abb168b89289acc36dff5531a6e50 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:22:11 -0700 Subject: [PATCH 274/879] Added test for __repr__ [skip ci] --- tests/test_sparse_vector.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 488908c..e4af365 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -7,3 +7,6 @@ class TestSparseVector: def test_from_dense(self): assert SparseVector.from_dense([1, 2, 3]).to_list() == [1, 2, 3] assert SparseVector.from_dense([1, 2, 3]).to_numpy().tolist() == [1, 2, 3] + + def test_repr(self): + assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])' From db122fe5dc532a63322b2d10da7037525e459b34 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:31:53 -0700 Subject: [PATCH 275/879] Added check for ndim for Bit [skip ci] --- pgvector/utils/bit.py | 7 ++++++- tests/test_bit.py | 10 ++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index ce7dd2d..d92c327 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -9,7 +9,12 @@ def __init__(self, value): elif isinstance(value, str): self._value = __class__.from_text(value)._value else: - self._value = np.asarray(value, dtype=bool) + value = np.asarray(value, dtype=bool) + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + self._value = value def __str__(self): return self.to_text() diff --git a/tests/test_bit.py b/tests/test_bit.py index e2006b9..001ecae 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -18,5 +18,15 @@ def test_ndarray_same_object(self): assert Bit(arr).to_list() == [True, False, True] assert Bit(arr).to_numpy() is arr + def test_ndim_two(self): + with pytest.raises(ValueError) as error: + Bit([[True, False], [True, False]]) + assert str(error.value) == 'expected ndim to be 1' + + def test_ndim_zero(self): + with pytest.raises(ValueError) as error: + Bit(True) + assert str(error.value) == 'expected ndim to be 1' + def test_repr(self): assert repr(Bit([True, False, True])) == 'Bit(101)' From 33ec0f2ff007736e2fd808a8c7d1b84f290d6072 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:39:25 -0700 Subject: [PATCH 276/879] Removed from_binary path from Bit constructor [skip ci] --- pgvector/utils/bit.py | 4 +--- tests/test_psycopg.py | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index d92c327..da7a532 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -4,9 +4,7 @@ class Bit: def __init__(self, value): - if isinstance(value, bytes): - self._value = __class__.from_binary(value)._value - elif isinstance(value, str): + if isinstance(value, str): self._value = __class__.from_text(value)._value else: value = np.asarray(value, dtype=bool) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 17483f7..8da32f6 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -91,8 +91,8 @@ def test_bit(self): def test_bit_binary_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) res = conn.execute('SELECT %b::bit(9)', (embedding,), binary=True).fetchone()[0] - assert str(Bit(res)) == '010100001' - assert repr(Bit(res)) == 'Bit(010100001)' + assert str(Bit.from_binary(res)) == '010100001' + assert repr(Bit.from_binary(res)) == 'Bit(010100001)' def test_bit_text_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) From a3dc9166cbc1653f6bb7928c04ba5a0e30c8ea1b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:44:52 -0700 Subject: [PATCH 277/879] Use consistent __str__ representation [skip ci] --- pgvector/utils/bit.py | 5 +---- tests/test_bit.py | 3 ++- tests/test_half_vector.py | 1 + tests/test_psycopg.py | 2 -- tests/test_sparse_vector.py | 1 + tests/test_vector.py | 1 + 6 files changed, 6 insertions(+), 7 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index da7a532..9deea1f 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -14,11 +14,8 @@ def __init__(self, value): self._value = value - def __str__(self): - return self.to_text() - def __repr__(self): - return f'Bit({self})' + return f'Bit({self.to_text()})' def dim(self): return self._value.shape[0] diff --git a/tests/test_bit.py b/tests/test_bit.py index 001ecae..c42bd05 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -11,7 +11,7 @@ def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] def test_str(self): - assert str(Bit('101')) == '101' + assert Bit('101').to_list() == [True, False, True] def test_ndarray_same_object(self): arr = np.array([True, False, True]) @@ -30,3 +30,4 @@ def test_ndim_zero(self): def test_repr(self): assert repr(Bit([True, False, True])) == 'Bit(101)' + assert str(Bit([True, False, True])) == 'Bit(101)' diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index c60433d..683c6d5 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -36,3 +36,4 @@ def test_ndim_zero(self): def test_repr(self): assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' + assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 8da32f6..2de1ec7 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -91,14 +91,12 @@ def test_bit(self): def test_bit_binary_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) res = conn.execute('SELECT %b::bit(9)', (embedding,), binary=True).fetchone()[0] - assert str(Bit.from_binary(res)) == '010100001' assert repr(Bit.from_binary(res)) == 'Bit(010100001)' def test_bit_text_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) res = conn.execute('SELECT %t::bit(9)', (embedding,)).fetchone()[0] assert res == '010100001' - assert str(Bit(res)) == '010100001' assert repr(Bit(res)) == 'Bit(010100001)' def test_sparsevec(self): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index e4af365..49dbf6f 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -10,3 +10,4 @@ def test_from_dense(self): def test_repr(self): assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])' + assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])' diff --git a/tests/test_vector.py b/tests/test_vector.py index f424019..8e1bd67 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -36,3 +36,4 @@ def test_ndim_zero(self): def test_repr(self): assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' + assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' From f59434f938d44f4aa1cae8c965d695bc1e75e346 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:50:02 -0700 Subject: [PATCH 278/879] Added dim method to SparseVector [skip ci] --- pgvector/utils/sparsevec.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 611bb16..bc9a981 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -28,6 +28,9 @@ def from_dense(value): values = [value[i] for i in indices] return SparseVector(dim, indices, values) + def dim(self): + return self._dim + def to_list(self): vec = [0.0] * self._dim for i, v in zip(self._indices, self._values): @@ -71,8 +74,8 @@ def to_db(value, dim=None): value = to_db_value(value) - if dim is not None and value._dim != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value._dim)) + if dim is not None and value.dim() != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) return value.to_text() From 9f5f4eb6cc200c0047bb348068d5b01bcb16c6f4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 19:59:22 -0700 Subject: [PATCH 279/879] Use class name [skip ci] --- pgvector/utils/halfvec.py | 4 ++-- pgvector/utils/sparsevec.py | 4 ++-- pgvector/utils/vector.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 57ca4df..d408816 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -65,10 +65,10 @@ def from_db(value): if value is None or isinstance(value, HalfVector): return value - return __class__.from_text(value) + return HalfVector.from_text(value) def from_db_binary(value): if value is None or isinstance(value, HalfVector): return value - return __class__.from_binary(value) + return HalfVector.from_binary(value) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index bc9a981..ff39b97 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -91,10 +91,10 @@ def from_db(value): if value is None or isinstance(value, SparseVector): return value - return __class__.from_text(value) + return SparseVector.from_text(value) def from_db_binary(value): if value is None or isinstance(value, SparseVector): return value - return __class__.from_binary(value) + return SparseVector.from_binary(value) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index adecda3..2294f4f 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -65,10 +65,10 @@ def from_db(value): if value is None or isinstance(value, np.ndarray): return value - return __class__.from_text(value).to_numpy().astype(np.float32) + return Vector.from_text(value).to_numpy().astype(np.float32) def from_db_binary(value): if value is None or isinstance(value, np.ndarray): return value - return __class__.from_binary(value).to_numpy().astype(np.float32) + return Vector.from_binary(value).to_numpy().astype(np.float32) From b1d2cdb17c8abf25925b5ee0625ef9e77c5d25b2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 20:02:37 -0700 Subject: [PATCH 280/879] Improved SparseVector code [skip ci] --- pgvector/utils/sparsevec.py | 9 +++++---- tests/test_sparse_vector.py | 4 ++-- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index ff39b97..e435009 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -13,9 +13,10 @@ def to_db_value(value): class SparseVector: def __init__(self, dim, indices, values): - self._dim = dim - self._indices = indices - self._values = values + # TODO improve + self._dim = int(dim) + self._indices = [int(i) for i in indices] + self._values = [float(v) for v in values] def __repr__(self): return f'SparseVector({self._dim}, {self._indices}, {self._values})' @@ -25,7 +26,7 @@ def from_dense(value): value = value.tolist() dim = len(value) indices = [i for i, v in enumerate(value) if v != 0] - values = [value[i] for i in indices] + values = [float(value[i]) for i in indices] return SparseVector(dim, indices, values) def dim(self): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 49dbf6f..04fcf88 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -9,5 +9,5 @@ def test_from_dense(self): assert SparseVector.from_dense([1, 2, 3]).to_numpy().tolist() == [1, 2, 3] def test_repr(self): - assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])' - assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])' + assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])' + assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])' From 6c98554e13060d33168a453d42d6055f3cec740a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 20:13:31 -0700 Subject: [PATCH 281/879] Added casting [skip ci] --- pgvector/utils/halfvec.py | 2 +- pgvector/utils/sparsevec.py | 2 +- pgvector/utils/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index d408816..16626f6 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -26,7 +26,7 @@ def to_numpy(self): return self._value def to_text(self): - return '[' + ','.join([str(v) for v in self._value]) + ']' + return '[' + ','.join([str(float(v)) for v in self._value]) + ']' def to_binary(self): return pack('>HH', self.dim(), 0) + self._value.tobytes() diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index e435009..4a8ebd4 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -45,7 +45,7 @@ def to_numpy(self): return vec def to_text(self): - return '{' + ','.join([f'{i + 1}:{v}' for i, v in zip(self._indices, self._values)]) + '}/' + str(self._dim) + return '{' + ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + '}/' + str(int(self._dim)) def to_binary(self): nnz = len(self._indices) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 2294f4f..6264c93 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -26,7 +26,7 @@ def to_numpy(self): return self._value def to_text(self): - return '[' + ','.join([str(v) for v in self._value]) + ']' + return '[' + ','.join([str(float(v)) for v in self._value]) + ']' def to_binary(self): return pack('>HH', self.dim(), 0) + self._value.tobytes() From 1d82171235e9d55905f0962e4997380306914c60 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 20:17:56 -0700 Subject: [PATCH 282/879] Removed dim method from Bit [skip ci] --- pgvector/utils/bit.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 9deea1f..c99ddbf 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -17,9 +17,6 @@ def __init__(self, value): def __repr__(self): return f'Bit({self.to_text()})' - def dim(self): - return self._value.shape[0] - def to_list(self): return self._value.tolist() From 87249e92ad2d244da7272dc70b59af874139e4a0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 20:20:27 -0700 Subject: [PATCH 283/879] Added tests for dim methods [skip ci] --- pgvector/utils/bit.py | 3 +-- pgvector/utils/halfvec.py | 2 +- pgvector/utils/vector.py | 2 +- tests/test_half_vector.py | 3 +++ tests/test_sparse_vector.py | 3 +++ tests/test_vector.py | 3 +++ 6 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index c99ddbf..26c12ca 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -27,8 +27,7 @@ def to_text(self): return ''.join(self._value.astype(np.uint8).astype(str)) def to_binary(self): - value = self._value - return pack('>i', len(value)) + np.packbits(value).tobytes() + return pack('>i', len(self._value)) + np.packbits(self._value).tobytes() def from_text(value): return Bit(np.asarray([v != '0' for v in value], dtype=bool)) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 16626f6..596f04c 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -17,7 +17,7 @@ def __repr__(self): return f'HalfVector({self.to_list()})' def dim(self): - return self._value.shape[0] + return len(self._value) def to_list(self): return self._value.tolist() diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 6264c93..c21f109 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -17,7 +17,7 @@ def __repr__(self): return f'Vector({self.to_list()})' def dim(self): - return self._value.shape[0] + return len(self._value) def to_list(self): return self._value.tolist() diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 683c6d5..b2811b2 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -37,3 +37,6 @@ def test_ndim_zero(self): def test_repr(self): assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' + + def test_dim(self): + assert HalfVector([1, 2, 3]).dim() == 3 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 04fcf88..34884fb 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -11,3 +11,6 @@ def test_from_dense(self): def test_repr(self): assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])' assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])' + + def test_dim(self): + assert SparseVector.from_dense([1, 2, 3]).dim() == 3 diff --git a/tests/test_vector.py b/tests/test_vector.py index 8e1bd67..d6ecb5b 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -37,3 +37,6 @@ def test_ndim_zero(self): def test_repr(self): assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' + + def test_dim(self): + assert Vector([1, 2, 3]).dim() == 3 From 4bf9d32dcaa4d37d57d45441968342434c6a3f97 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 1 Jun 2024 20:51:15 -0700 Subject: [PATCH 284/879] Improved code [skip ci] --- pgvector/utils/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 26c12ca..78533ef 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -34,7 +34,7 @@ def from_text(value): def from_binary(value): count = unpack_from('>i', value)[0] - buf = np.frombuffer(value[4:], dtype=np.uint8) + buf = np.frombuffer(value, dtype=np.uint8, offset=4) return Bit(np.unpackbits(buf, count=count).astype(bool)) # TODO move rest From 240e4135bfe4a63bdb32c2c137bfe0ca49165c37 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 3 Jun 2024 18:40:39 -0700 Subject: [PATCH 285/879] Added underscores to private functions [skip ci] --- pgvector/asyncpg/register.py | 12 ++++++------ pgvector/django/functions.py | 6 +++--- pgvector/django/halfvec.py | 6 +++--- pgvector/django/sparsevec.py | 6 +++--- pgvector/django/vector.py | 6 +++--- pgvector/peewee/halfvec.py | 4 ++-- pgvector/peewee/sparsevec.py | 4 ++-- pgvector/peewee/vector.py | 4 ++-- pgvector/psycopg/bit.py | 4 ++-- pgvector/psycopg/halfvec.py | 8 ++++---- pgvector/psycopg/sparsevec.py | 8 ++++---- pgvector/psycopg/vector.py | 8 ++++---- pgvector/psycopg2/halfvec.py | 4 ++-- pgvector/psycopg2/sparsevec.py | 4 ++-- pgvector/psycopg2/vector.py | 4 ++-- pgvector/sqlalchemy/halfvec.py | 6 +++--- pgvector/sqlalchemy/sparsevec.py | 6 +++--- pgvector/sqlalchemy/vector.py | 6 +++--- pgvector/utils/bit.py | 6 ++---- pgvector/utils/halfvec.py | 10 ++++------ pgvector/utils/sparsevec.py | 31 ++++++++++++++----------------- pgvector/utils/vector.py | 10 ++++------ 22 files changed, 77 insertions(+), 86 deletions(-) diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 4664218..7148ffa 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -4,21 +4,21 @@ async def register_vector(conn): await conn.set_type_codec( 'vector', - encoder=Vector.to_db_binary, - decoder=Vector.from_db_binary, + encoder=Vector._to_db_binary, + decoder=Vector._from_db_binary, format='binary' ) await conn.set_type_codec( 'halfvec', - encoder=HalfVector.to_db_binary, - decoder=HalfVector.from_db_binary, + encoder=HalfVector._to_db_binary, + decoder=HalfVector._from_db_binary, format='binary' ) await conn.set_type_codec( 'sparsevec', - encoder=SparseVector.to_db_binary, - decoder=SparseVector.from_db_binary, + encoder=SparseVector._to_db_binary, + decoder=SparseVector._from_db_binary, format='binary' ) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 9b11869..da9fbf8 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -8,11 +8,11 @@ class DistanceBase(Func): def __init__(self, expression, vector, **extra): if not hasattr(vector, 'resolve_expression'): if isinstance(vector, HalfVector): - vector = Value(HalfVector.to_db(vector)) + vector = Value(HalfVector._to_db(vector)) elif isinstance(vector, SparseVector): - vector = Value(SparseVector.to_db(vector)) + vector = Value(SparseVector._to_db(vector)) else: - vector = Value(Vector.to_db(vector)) + vector = Value(Vector._to_db(vector)) super().__init__(expression, vector, **extra) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 1c3bae5..a7921b7 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -23,13 +23,13 @@ def db_type(self, connection): return 'halfvec(%d)' % self.dimensions def from_db_value(self, value, expression, connection): - return HalfVector.from_db(value) + return HalfVector._from_db(value) def to_python(self, value): - return HalfVector.from_db(value) + return HalfVector._from_db(value) def get_prep_value(self, value): - return HalfVector.to_db(value) + return HalfVector._to_db(value) def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 3a06574..4ec734f 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -23,13 +23,13 @@ def db_type(self, connection): return 'sparsevec(%d)' % self.dimensions def from_db_value(self, value, expression, connection): - return SparseVector.from_db(value) + return SparseVector._from_db(value) def to_python(self, value): - return SparseVector.from_db(value) + return SparseVector._from_db(value) def get_prep_value(self, value): - return SparseVector.to_db(value) + return SparseVector._to_db(value) def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 30fd99f..a89d540 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -25,15 +25,15 @@ def db_type(self, connection): return 'vector(%d)' % self.dimensions def from_db_value(self, value, expression, connection): - return Vector.from_db(value) + return Vector._from_db(value) def to_python(self, value): if isinstance(value, list): return np.array(value, dtype=np.float32) - return Vector.from_db(value) + return Vector._from_db(value) def get_prep_value(self, value): - return Vector.to_db(value) + return Vector._to_db(value) def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index e30dcd0..bed7d1f 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -13,10 +13,10 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return HalfVector.to_db(value) + return HalfVector._to_db(value) def python_value(self, value): - return HalfVector.from_db(value) + return HalfVector._from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index b6d3b91..8bba5cf 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -13,10 +13,10 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return SparseVector.to_db(value) + return SparseVector._to_db(value) def python_value(self, value): - return SparseVector.from_db(value) + return SparseVector._from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 3c11e74..a9ebf87 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -13,10 +13,10 @@ def get_modifiers(self): return self.dimensions and [self.dimensions] or None def db_value(self, value): - return Vector.to_db(value) + return Vector._to_db(value) def python_value(self, value): - return Vector.from_db(value) + return Vector._from_db(value) def _distance(self, op, vector): return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index 80cfac0..605c6f0 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -8,7 +8,7 @@ class BitDumper(Dumper): format = Format.TEXT def dump(self, obj): - return Bit.to_db(obj).encode('utf8') + return Bit._to_db(obj).encode('utf8') class BitBinaryDumper(BitDumper): @@ -16,7 +16,7 @@ class BitBinaryDumper(BitDumper): format = Format.BINARY def dump(self, obj): - return Bit.to_db_binary(obj) + return Bit._to_db_binary(obj) def register_bit_info(context, info): diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 6ca232d..351d2cb 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -8,7 +8,7 @@ class HalfVectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return HalfVector.to_db(obj).encode('utf8') + return HalfVector._to_db(obj).encode('utf8') class HalfVectorBinaryDumper(HalfVectorDumper): @@ -16,7 +16,7 @@ class HalfVectorBinaryDumper(HalfVectorDumper): format = Format.BINARY def dump(self, obj): - return HalfVector.to_db_binary(obj) + return HalfVector._to_db_binary(obj) class HalfVectorLoader(Loader): @@ -26,7 +26,7 @@ class HalfVectorLoader(Loader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return HalfVector.from_db(data.decode('utf8')) + return HalfVector._from_db(data.decode('utf8')) class HalfVectorBinaryLoader(HalfVectorLoader): @@ -36,7 +36,7 @@ class HalfVectorBinaryLoader(HalfVectorLoader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return HalfVector.from_db_binary(data) + return HalfVector._from_db_binary(data) def register_halfvec_info(context, info): diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 634d83b..435fd06 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -8,7 +8,7 @@ class SparseVectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return SparseVector.to_db(obj).encode('utf8') + return SparseVector._to_db(obj).encode('utf8') class SparseVectorBinaryDumper(SparseVectorDumper): @@ -16,7 +16,7 @@ class SparseVectorBinaryDumper(SparseVectorDumper): format = Format.BINARY def dump(self, obj): - return SparseVector.to_db_binary(obj) + return SparseVector._to_db_binary(obj) class SparseVectorLoader(Loader): @@ -26,7 +26,7 @@ class SparseVectorLoader(Loader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return SparseVector.from_db(data.decode('utf8')) + return SparseVector._from_db(data.decode('utf8')) class SparseVectorBinaryLoader(SparseVectorLoader): @@ -36,7 +36,7 @@ class SparseVectorBinaryLoader(SparseVectorLoader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return SparseVector.from_db_binary(data) + return SparseVector._from_db_binary(data) def register_sparsevec_info(context, info): diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 4875fe2..55e31a6 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -8,7 +8,7 @@ class VectorDumper(Dumper): format = Format.TEXT def dump(self, obj): - return Vector.to_db(obj).encode('utf8') + return Vector._to_db(obj).encode('utf8') class VectorBinaryDumper(VectorDumper): @@ -16,7 +16,7 @@ class VectorBinaryDumper(VectorDumper): format = Format.BINARY def dump(self, obj): - return Vector.to_db_binary(obj) + return Vector._to_db_binary(obj) class VectorLoader(Loader): @@ -26,7 +26,7 @@ class VectorLoader(Loader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return Vector.from_db(data.decode('utf8')) + return Vector._from_db(data.decode('utf8')) class VectorBinaryLoader(VectorLoader): @@ -36,7 +36,7 @@ class VectorBinaryLoader(VectorLoader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return Vector.from_db_binary(data) + return Vector._from_db_binary(data) def register_vector_info(context, info): diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index b480911..e3c0cdf 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -7,11 +7,11 @@ def __init__(self, value): self._value = value def getquoted(self): - return adapt(HalfVector.to_db(self._value)).getquoted() + return adapt(HalfVector._to_db(self._value)).getquoted() def cast_halfvec(value, cur): - return HalfVector.from_db(value) + return HalfVector._from_db(value) def register_halfvec_info(oid): diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 31e86f3..7cdf38c 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -7,11 +7,11 @@ def __init__(self, value): self._value = value def getquoted(self): - return adapt(SparseVector.to_db(self._value)).getquoted() + return adapt(SparseVector._to_db(self._value)).getquoted() def cast_sparsevec(value, cur): - return SparseVector.from_db(value) + return SparseVector._from_db(value) def register_sparsevec_info(oid): diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 4fdb077..c895f86 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -8,11 +8,11 @@ def __init__(self, value): self._value = value def getquoted(self): - return adapt(Vector.to_db(self._value)).getquoted() + return adapt(Vector._to_db(self._value)).getquoted() def cast_vector(value, cur): - return Vector.from_db(value) + return Vector._from_db(value) def register_vector_info(oid): diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 1c5f8f1..639f77b 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return HalfVector.to_db(value, self.dim) + return HalfVector._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(HalfVector.to_db(value, self.dim)) + return string_literal_processor(HalfVector._to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return HalfVector.from_db(value) + return HalfVector._from_db(value) return process class comparator_factory(UserDefinedType.Comparator): diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index ca5e9aa..370f5d1 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return SparseVector.to_db(value, self.dim) + return SparseVector._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(SparseVector.to_db(value, self.dim)) + return string_literal_processor(SparseVector._to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return SparseVector.from_db(value) + return SparseVector._from_db(value) return process class comparator_factory(UserDefinedType.Comparator): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index b2ce7f7..f57a045 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -18,19 +18,19 @@ def get_col_spec(self, **kw): def bind_processor(self, dialect): def process(value): - return Vector.to_db(value, self.dim) + return Vector._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(Vector.to_db(value, self.dim)) + return string_literal_processor(Vector._to_db(value, self.dim)) return process def result_processor(self, dialect, coltype): def process(value): - return Vector.from_db(value) + return Vector._from_db(value) return process class comparator_factory(UserDefinedType.Comparator): diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 78533ef..da342a0 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -37,15 +37,13 @@ def from_binary(value): buf = np.frombuffer(value, dtype=np.uint8, offset=4) return Bit(np.unpackbits(buf, count=count).astype(bool)) - # TODO move rest - - def to_db(value): + def _to_db(value): if not isinstance(value, Bit): raise ValueError('expected bit') return value.to_text() - def to_db_binary(value): + def _to_db_binary(value): if not isinstance(value, Bit): raise ValueError('expected bit') diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index 596f04c..e2836c8 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -38,9 +38,7 @@ def from_binary(value): dim, unused = unpack_from('>HH', value) return HalfVector(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) - # TODO move rest - - def to_db(value, dim=None): + def _to_db(value, dim=None): if value is None: return value @@ -52,7 +50,7 @@ def to_db(value, dim=None): return value.to_text() - def to_db_binary(value): + def _to_db_binary(value): if value is None: return value @@ -61,13 +59,13 @@ def to_db_binary(value): return value.to_binary() - def from_db(value): + def _from_db(value): if value is None or isinstance(value, HalfVector): return value return HalfVector.from_text(value) - def from_db_binary(value): + def _from_db_binary(value): if value is None or isinstance(value, HalfVector): return value diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 4a8ebd4..ace268a 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -2,15 +2,6 @@ from struct import pack, unpack_from -def to_db_value(value): - if isinstance(value, SparseVector): - return value - elif isinstance(value, (list, np.ndarray)): - return SparseVector.from_dense(value) - else: - raise ValueError('expected sparsevec') - - class SparseVector: def __init__(self, dim, indices, values): # TODO improve @@ -67,34 +58,40 @@ def from_binary(value): values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) return SparseVector(int(dim), indices, values) - # TODO move rest - - def to_db(value, dim=None): + def _to_db(value, dim=None): if value is None: return value - value = to_db_value(value) + value = __class__._to_db_value(value) if dim is not None and value.dim() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) return value.to_text() - def to_db_binary(value): + def _to_db_binary(value): if value is None: return value - value = to_db_value(value) + value = __class__._to_db_value(value) return value.to_binary() - def from_db(value): + def _to_db_value(value): + if isinstance(value, SparseVector): + return value + elif isinstance(value, (list, np.ndarray)): + return SparseVector.from_dense(value) + else: + raise ValueError('expected sparsevec') + + def _from_db(value): if value is None or isinstance(value, SparseVector): return value return SparseVector.from_text(value) - def from_db_binary(value): + def _from_db_binary(value): if value is None or isinstance(value, SparseVector): return value diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index c21f109..4c70b34 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -38,9 +38,7 @@ def from_binary(value): dim, unused = unpack_from('>HH', value) return Vector(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) - # TODO move rest - - def to_db(value, dim=None): + def _to_db(value, dim=None): if value is None: return value @@ -52,7 +50,7 @@ def to_db(value, dim=None): return value.to_text() - def to_db_binary(value): + def _to_db_binary(value): if value is None: return value @@ -61,13 +59,13 @@ def to_db_binary(value): return value.to_binary() - def from_db(value): + def _from_db(value): if value is None or isinstance(value, np.ndarray): return value return Vector.from_text(value).to_numpy().astype(np.float32) - def from_db_binary(value): + def _from_db_binary(value): if value is None or isinstance(value, np.ndarray): return value From f1ce5f3f11d351c1fca19b9b7ff3ed99f6f1db70 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 3 Jun 2024 18:42:36 -0700 Subject: [PATCH 286/879] Removed unneeded code [skip ci] --- pgvector/utils/sparsevec.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index ace268a..0fde7f1 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -13,8 +13,6 @@ def __repr__(self): return f'SparseVector({self._dim}, {self._indices}, {self._values})' def from_dense(value): - if isinstance(value, np.ndarray): - value = value.tolist() dim = len(value) indices = [i for i, v in enumerate(value) if v != 0] values = [float(value[i]) for i in indices] From b43b58ff7536dd3cdecad5aaa8c5c18a7746827a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 14:17:55 -0700 Subject: [PATCH 287/879] Use classmethod decorator - closes #72 Co-authored-by: "domenico.cinque" --- pgvector/utils/bit.py | 10 ++++++---- pgvector/utils/halfvec.py | 38 +++++++++++++++++++++---------------- pgvector/utils/sparsevec.py | 31 +++++++++++++++++------------- pgvector/utils/vector.py | 24 +++++++++++++---------- 4 files changed, 60 insertions(+), 43 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index da342a0..b5ede20 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -37,14 +37,16 @@ def from_binary(value): buf = np.frombuffer(value, dtype=np.uint8, offset=4) return Bit(np.unpackbits(buf, count=count).astype(bool)) - def _to_db(value): - if not isinstance(value, Bit): + @classmethod + def _to_db(cls, value): + if not isinstance(value, cls): raise ValueError('expected bit') return value.to_text() - def _to_db_binary(value): - if not isinstance(value, Bit): + @classmethod + def _to_db_binary(cls, value): + if not isinstance(value, cls): raise ValueError('expected bit') return value.to_binary() diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index e2836c8..d01b381 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -31,42 +31,48 @@ def to_text(self): def to_binary(self): return pack('>HH', self.dim(), 0) + self._value.tobytes() - def from_text(value): - return HalfVector([float(v) for v in value[1:-1].split(',')]) + @classmethod + def from_text(cls, value): + return cls([float(v) for v in value[1:-1].split(',')]) - def from_binary(value): + @classmethod + def from_binary(cls, value): dim, unused = unpack_from('>HH', value) - return HalfVector(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) + return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) - def _to_db(value, dim=None): + @classmethod + def _to_db(cls, value, dim=None): if value is None: return value - if not isinstance(value, HalfVector): - value = HalfVector(value) + if not isinstance(value, cls): + value = cls(value) if dim is not None and value.dim() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) return value.to_text() - def _to_db_binary(value): + @classmethod + def _to_db_binary(cls, value): if value is None: return value - if not isinstance(value, HalfVector): - value = HalfVector(value) + if not isinstance(value, cls): + value = cls(value) return value.to_binary() - def _from_db(value): - if value is None or isinstance(value, HalfVector): + @classmethod + def _from_db(cls, value): + if value is None or isinstance(value, cls): return value - return HalfVector.from_text(value) + return cls.from_text(value) - def _from_db_binary(value): - if value is None or isinstance(value, HalfVector): + @classmethod + def _from_db_binary(cls, value): + if value is None or isinstance(value, cls): return value - return HalfVector.from_binary(value) + return cls.from_binary(value) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 0fde7f1..18e510d 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -56,41 +56,46 @@ def from_binary(value): values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) return SparseVector(int(dim), indices, values) - def _to_db(value, dim=None): + @classmethod + def _to_db(cls, value, dim=None): if value is None: return value - value = __class__._to_db_value(value) + value = cls._to_db_value(value) if dim is not None and value.dim() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) return value.to_text() - def _to_db_binary(value): + @classmethod + def _to_db_binary(cls, value): if value is None: return value - value = __class__._to_db_value(value) + value = cls._to_db_value(value) return value.to_binary() - def _to_db_value(value): - if isinstance(value, SparseVector): + @classmethod + def _to_db_value(cls, value): + if isinstance(value, cls): return value elif isinstance(value, (list, np.ndarray)): - return SparseVector.from_dense(value) + return cls.from_dense(value) else: raise ValueError('expected sparsevec') - def _from_db(value): - if value is None or isinstance(value, SparseVector): + @classmethod + def _from_db(cls, value): + if value is None or isinstance(value, cls): return value - return SparseVector.from_text(value) + return cls.from_text(value) - def _from_db_binary(value): - if value is None or isinstance(value, SparseVector): + @classmethod + def _from_db_binary(cls, value): + if value is None or isinstance(value, cls): return value - return SparseVector.from_binary(value) + return cls.from_binary(value) diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 4c70b34..1ddba4c 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -38,35 +38,39 @@ def from_binary(value): dim, unused = unpack_from('>HH', value) return Vector(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) - def _to_db(value, dim=None): + @classmethod + def _to_db(cls, value, dim=None): if value is None: return value - if not isinstance(value, Vector): - value = Vector(value) + if not isinstance(value, cls): + value = cls(value) if dim is not None and value.dim() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) return value.to_text() - def _to_db_binary(value): + @classmethod + def _to_db_binary(cls, value): if value is None: return value - if not isinstance(value, Vector): - value = Vector(value) + if not isinstance(value, cls): + value = cls(value) return value.to_binary() - def _from_db(value): + @classmethod + def _from_db(cls, value): if value is None or isinstance(value, np.ndarray): return value - return Vector.from_text(value).to_numpy().astype(np.float32) + return cls.from_text(value).to_numpy().astype(np.float32) - def _from_db_binary(value): + @classmethod + def _from_db_binary(cls, value): if value is None or isinstance(value, np.ndarray): return value - return Vector.from_binary(value).to_numpy().astype(np.float32) + return cls.from_binary(value).to_numpy().astype(np.float32) From 3a4118ca3687a5ff366087924b0aacb523bab66e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 15:26:07 -0700 Subject: [PATCH 288/879] Improved SparseVector tests [skip ci] --- tests/test_sparse_vector.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 34884fb..92c6432 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -5,12 +5,12 @@ class TestSparseVector: def test_from_dense(self): - assert SparseVector.from_dense([1, 2, 3]).to_list() == [1, 2, 3] - assert SparseVector.from_dense([1, 2, 3]).to_numpy().tolist() == [1, 2, 3] + assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] def test_repr(self): - assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])' - assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])' + assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' + assert str(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' def test_dim(self): - assert SparseVector.from_dense([1, 2, 3]).dim() == 3 + assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).dim() == 6 From fb8edf9da0fc1d2bab273e283689ec2dc93cdbe5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 15:29:56 -0700 Subject: [PATCH 289/879] Use classmethod decorator for more methods --- pgvector/utils/bit.py | 10 ++++++---- pgvector/utils/sparsevec.py | 10 ++++++---- pgvector/utils/vector.py | 10 ++++++---- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index b5ede20..3b27db3 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -29,13 +29,15 @@ def to_text(self): def to_binary(self): return pack('>i', len(self._value)) + np.packbits(self._value).tobytes() - def from_text(value): - return Bit(np.asarray([v != '0' for v in value], dtype=bool)) + @classmethod + def from_text(cls, value): + return cls(np.asarray([v != '0' for v in value], dtype=bool)) - def from_binary(value): + @classmethod + def from_binary(cls, value): count = unpack_from('>i', value)[0] buf = np.frombuffer(value, dtype=np.uint8, offset=4) - return Bit(np.unpackbits(buf, count=count).astype(bool)) + return cls(np.unpackbits(buf, count=count).astype(bool)) @classmethod def _to_db(cls, value): diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 18e510d..8d3f866 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -40,7 +40,8 @@ def to_binary(self): nnz = len(self._indices) return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) - def from_text(value): + @classmethod + def from_text(cls, value): elements, dim = value.split('/') indices = [] values = [] @@ -48,13 +49,14 @@ def from_text(value): i, v = e.split(':') indices.append(int(i) - 1) values.append(float(v)) - return SparseVector(int(dim), indices, values) + return cls(int(dim), indices, values) - def from_binary(value): + @classmethod + def from_binary(cls, value): dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return SparseVector(int(dim), indices, values) + return cls(int(dim), indices, values) @classmethod def _to_db(cls, value, dim=None): diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 1ddba4c..7a0f6bc 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -31,12 +31,14 @@ def to_text(self): def to_binary(self): return pack('>HH', self.dim(), 0) + self._value.tobytes() - def from_text(value): - return Vector([float(v) for v in value[1:-1].split(',')]) + @classmethod + def from_text(cls, value): + return cls([float(v) for v in value[1:-1].split(',')]) - def from_binary(value): + @classmethod + def from_binary(cls, value): dim, unused = unpack_from('>HH', value) - return Vector(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) + return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) @classmethod def _to_db(cls, value, dim=None): From 63bc2507e47a0a87cc941044a40529b9b4b677cd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 15:30:54 -0700 Subject: [PATCH 290/879] Use classmethod decorator for from_dense [skip ci] --- pgvector/utils/sparsevec.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 8d3f866..962938b 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -12,11 +12,12 @@ def __init__(self, dim, indices, values): def __repr__(self): return f'SparseVector({self._dim}, {self._indices}, {self._values})' - def from_dense(value): + @classmethod + def from_dense(cls, value): dim = len(value) indices = [i for i, v in enumerate(value) if v != 0] values = [float(value[i]) for i in indices] - return SparseVector(dim, indices, values) + return cls(dim, indices, values) def dim(self): return self._dim From df3427dace3cf9de6fe231229c319b3c1b958e58 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 15:36:51 -0700 Subject: [PATCH 291/879] Added from_coordinates method to SparseVector --- pgvector/utils/sparsevec.py | 10 ++++++++++ tests/test_sparse_vector.py | 4 ++++ 2 files changed, 14 insertions(+) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 962938b..9254e54 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -12,6 +12,16 @@ def __init__(self, dim, indices, values): def __repr__(self): return f'SparseVector({self._dim}, {self._indices}, {self._values})' + @classmethod + def from_coordinates(cls, coordinates, dim): + if isinstance(coordinates, dict): + coordinates = coordinates.items() + elements = [(i, v) for i, v in coordinates] + elements.sort() + indices = [int(v[0]) for v in elements] + values = [float(v[1]) for v in elements] + return cls(dim, indices, values) + @classmethod def from_dense(cls, value): dim = len(value) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 92c6432..5fb4b7e 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -8,6 +8,10 @@ def test_from_dense(self): assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0] assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] + def test_from_coordinates(self): + assert SparseVector.from_coordinates({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_coordinates([(0, 1), (2, 2), (4, 3)], 6).to_list() == [1, 0, 2, 0, 3, 0] + def test_repr(self): assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' assert str(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' From 08a73cdf7dc062ed11c0b8b0227a7def850ae144 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 16:13:47 -0700 Subject: [PATCH 292/879] Removed support for array from from_coordinates [skip ci] --- pgvector/utils/sparsevec.py | 4 +--- tests/test_sparse_vector.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 9254e54..afcd856 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -14,9 +14,7 @@ def __repr__(self): @classmethod def from_coordinates(cls, coordinates, dim): - if isinstance(coordinates, dict): - coordinates = coordinates.items() - elements = [(i, v) for i, v in coordinates] + elements = [(i, v) for i, v in coordinates.items()] elements.sort() indices = [int(v[0]) for v in elements] values = [float(v[1]) for v in elements] diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 5fb4b7e..f6f96e0 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -10,7 +10,6 @@ def test_from_dense(self): def test_from_coordinates(self): assert SparseVector.from_coordinates({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_coordinates([(0, 1), (2, 2), (4, 3)], 6).to_list() == [1, 0, 2, 0, 3, 0] def test_repr(self): assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' From 999af56edf4ee95bf54783c242dec8834d9daf32 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 16:15:23 -0700 Subject: [PATCH 293/879] Improved SparseVector test [skip ci] --- tests/test_sparse_vector.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index f6f96e0..f622108 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -7,6 +7,7 @@ class TestSparseVector: def test_from_dense(self): assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0] assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_dense(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] def test_from_coordinates(self): assert SparseVector.from_coordinates({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] From c708d8f0eaedd4608109aca79bd97bd08b5c4528 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 17:47:58 -0700 Subject: [PATCH 294/879] Updated function name [skip ci] --- pgvector/utils/sparsevec.py | 4 ++-- tests/test_sparse_vector.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index afcd856..0a3df3b 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -13,8 +13,8 @@ def __repr__(self): return f'SparseVector({self._dim}, {self._indices}, {self._values})' @classmethod - def from_coordinates(cls, coordinates, dim): - elements = [(i, v) for i, v in coordinates.items()] + def from_dict(cls, d, dim): + elements = [(i, v) for i, v in d.items()] elements.sort() indices = [int(v[0]) for v in elements] values = [float(v[1]) for v in elements] diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index f622108..0f69b8a 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -10,7 +10,7 @@ def test_from_dense(self): assert SparseVector.from_dense(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] def test_from_coordinates(self): - assert SparseVector.from_coordinates({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_dict({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] def test_repr(self): assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' From 1b5dc44cc8cbe5a7885989da6f0cc0a173b2f474 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 18:04:02 -0700 Subject: [PATCH 295/879] Added from_scipy method to SparseVector --- pgvector/utils/sparsevec.py | 12 ++++++++++++ requirements.txt | 1 + tests/test_sparse_vector.py | 6 ++++++ 3 files changed, 19 insertions(+) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 0a3df3b..b1a53d7 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -20,6 +20,18 @@ def from_dict(cls, d, dim): values = [float(v[1]) for v in elements] return cls(dim, indices, values) + @classmethod + def from_scipy(cls, value): + value = value.tocoo() + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + dim = value.shape[0] + indices = value.coords[0].tolist() + values = value.data.tolist() + return cls(dim, indices, values) + @classmethod def from_dense(cls, value): dim = len(value) diff --git a/requirements.txt b/requirements.txt index da0ef44..c1e11f3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,5 +6,6 @@ psycopg[binary] psycopg2-binary pytest pytest-asyncio +scipy SQLAlchemy[asyncio]>=2 sqlmodel>=0.0.12 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 0f69b8a..f79d62a 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,6 +1,7 @@ import numpy as np from pgvector.utils import SparseVector import pytest +from scipy.sparse import coo_array class TestSparseVector: @@ -12,6 +13,11 @@ def test_from_dense(self): def test_from_coordinates(self): assert SparseVector.from_dict({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] + def test_from_scipy(self): + arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) + assert SparseVector.from_scipy(arr).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_scipy(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] + def test_repr(self): assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' assert str(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' From 550aa6403daf8aa0becac1b3e652aab92407e844 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 18:05:49 -0700 Subject: [PATCH 296/879] Renamed from_scipy method to from_sparse [skip ci] --- pgvector/utils/sparsevec.py | 2 +- tests/test_sparse_vector.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index b1a53d7..7de4607 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -21,7 +21,7 @@ def from_dict(cls, d, dim): return cls(dim, indices, values) @classmethod - def from_scipy(cls, value): + def from_sparse(cls, value): value = value.tocoo() if value.ndim != 1: diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index f79d62a..f447442 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -13,10 +13,10 @@ def test_from_dense(self): def test_from_coordinates(self): assert SparseVector.from_dict({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] - def test_from_scipy(self): + def test_from_sparse(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) - assert SparseVector.from_scipy(arr).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_scipy(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_sparse(arr).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_sparse(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] def test_repr(self): assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' From d858fcdfaee47c709566be8ff0b00decbb9d4810 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 18:21:40 -0700 Subject: [PATCH 297/879] Fixed from_sparse method for scipy < 1.13 --- pgvector/utils/sparsevec.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 7de4607..c4aa4f1 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -24,11 +24,18 @@ def from_dict(cls, d, dim): def from_sparse(cls, value): value = value.tocoo() - if value.ndim != 1: + if value.ndim == 1: + dim = value.shape[0] + elif value.ndim == 2 and value.shape[0] == 1: + dim = value.shape[1] + else: raise ValueError('expected ndim to be 1') - dim = value.shape[0] - indices = value.coords[0].tolist() + if hasattr(value, 'coords'): + # scipy 1.13+ + indices = value.coords[0].tolist() + else: + indices = value.col.tolist() values = value.data.tolist() return cls(dim, indices, values) From b3851e5ac1ac0dcda7221c66d568e889fb22516f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 4 Jun 2024 18:49:32 -0700 Subject: [PATCH 298/879] Improved code [skip ci] --- pgvector/utils/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 3b27db3..8a6f2ac 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -5,7 +5,7 @@ class Bit: def __init__(self, value): if isinstance(value, str): - self._value = __class__.from_text(value)._value + self._value = self.from_text(value)._value else: value = np.asarray(value, dtype=bool) From 33eb182c8cbb638f4b6266e0518b14d0beed21d8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Jun 2024 13:50:55 -0700 Subject: [PATCH 299/879] Added more tests [skip ci] --- tests/test_sparse_vector.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index f447442..c799fa3 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -16,6 +16,8 @@ def test_from_coordinates(self): def test_from_sparse(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) assert SparseVector.from_sparse(arr).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_sparse(arr.tocsc()).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector.from_sparse(arr.tocsr()).to_list() == [1, 0, 2, 0, 3, 0] assert SparseVector.from_sparse(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] def test_repr(self): From 8eaad5c3e83d2f2156d1d550a898db07f4c55740 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Jun 2024 13:56:12 -0700 Subject: [PATCH 300/879] Added to_dict method to SparseVector --- pgvector/utils/sparsevec.py | 3 +++ tests/test_sparse_vector.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index c4aa4f1..f33e496 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -49,6 +49,9 @@ def from_dense(cls, value): def dim(self): return self._dim + def to_dict(self): + return {i: v for i, v in zip(self._indices, self._values)} + def to_list(self): vec = [0.0] * self._dim for i, v in zip(self._indices, self._values): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index c799fa3..c1258dc 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -26,3 +26,6 @@ def test_repr(self): def test_dim(self): assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).dim() == 6 + + def test_to_dict(self): + assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3} From d0a50651bfdaf4ede4feed284f34a22932f52bb7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Jun 2024 13:59:10 -0700 Subject: [PATCH 301/879] Fixed CI --- tests/test_sparse_vector.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index c1258dc..73986b3 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -16,8 +16,6 @@ def test_from_coordinates(self): def test_from_sparse(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) assert SparseVector.from_sparse(arr).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_sparse(arr.tocsc()).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_sparse(arr.tocsr()).to_list() == [1, 0, 2, 0, 3, 0] assert SparseVector.from_sparse(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] def test_repr(self): From f36aa11a32e8c09196b52061466b3acd66ecd188 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Jun 2024 14:08:41 -0700 Subject: [PATCH 302/879] Simplified code [skip ci] --- pgvector/utils/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index f33e496..7112736 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -50,7 +50,7 @@ def dim(self): return self._dim def to_dict(self): - return {i: v for i, v in zip(self._indices, self._values)} + return dict(zip(self._indices, self._values)) def to_list(self): vec = [0.0] * self._dim From 8be15b29e90082b01a173503624659437035cde5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Jun 2024 14:54:21 -0700 Subject: [PATCH 303/879] Added to_coo method to SparseVector --- pgvector/utils/sparsevec.py | 6 ++++++ tests/test_sparse_vector.py | 3 +++ 2 files changed, 9 insertions(+) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 7112736..ebfa0a9 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -52,6 +52,12 @@ def dim(self): def to_dict(self): return dict(zip(self._indices, self._values)) + def to_coo(self): + from scipy.sparse import coo_array + + coords = ([0] * len(self._indices), self._indices) + return coo_array((self._values, coords), shape=(1, self._dim)) + def to_list(self): vec = [0.0] * self._dim for i, v in zip(self._indices, self._values): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 73986b3..c4565fb 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -27,3 +27,6 @@ def test_dim(self): def test_to_dict(self): assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3} + + def test_to_coo(self): + assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] From 3fa840e8046410fdefcea5e613314058c6a419b1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 11 Jun 2024 10:41:55 -0700 Subject: [PATCH 304/879] Added test for Automap with SQLAlchemy - #74 --- tests/test_sqlalchemy.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 687a0d0..1c0fb80 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -3,6 +3,7 @@ import pytest from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError +from sqlalchemy.ext.automap import automap_base from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base, mapped_column, Session from sqlalchemy.sql import func @@ -403,6 +404,16 @@ def test_insert_bulk(self): def test_insert_text(self): session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) + def test_automap(self): + metadata = MetaData() + metadata.reflect(engine, only=['sqlalchemy_orm_item']) + AutoBase = automap_base(metadata=metadata) + AutoBase.prepare() + AutoItem = AutoBase.classes.sqlalchemy_orm_item + session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) + item = session.query(AutoItem).first() + assert item.embedding.tolist() == [1, 2, 3] + @pytest.mark.asyncio async def test_async(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') From 0c8e1f0608f407e006a7514e95db60ebbcc8cd97 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 09:43:17 -0700 Subject: [PATCH 305/879] Improved Django tests --- tests/test_django.py | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 1c611d9..4fe4000 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -93,12 +93,18 @@ def create_items(): Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2])).save() -class ItemForm(ModelForm): +class VectorForm(ModelForm): class Meta: model = Item fields = ['embedding'] +class BitForm(ModelForm): + class Meta: + model = Item + fields = ['binary_embedding'] + + class TestDjango: def setup_method(self, test_method): Item.objects.all().delete() @@ -270,36 +276,41 @@ def test_serialization(self): for obj in serializers.deserialize(format, data): obj.save() - def test_form(self): - form = ItemForm(data={'embedding': '[1, 2, 3]'}) + def test_vector_form(self): + form = VectorForm(data={'embedding': '[1, 2, 3]'}) assert form.is_valid() assert 'value="[1, 2, 3]"' in form.as_div() - def test_form_instance(self): + def test_vector_form_instance(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - form = ItemForm(instance=item) + form = VectorForm(instance=item) assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() - def test_form_save(self): + def test_vector_form_save(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - form = ItemForm(instance=item, data={'embedding': '[4, 5, 6]'}) + form = VectorForm(instance=item, data={'embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() - def test_form_save_missing(self): + def test_vector_form_save_missing(self): Item(id=1).save() item = Item.objects.get(pk=1) - form = ItemForm(instance=item, data={'embedding': ''}) + form = VectorForm(instance=item, data={'embedding': ''}) assert form.is_valid() assert form.save() assert Item.objects.get(pk=1).embedding is None + def test_bit_form(self): + form = BitForm(data={'binary_embedding': '101'}) + assert form.is_valid() + assert 'value="101"' in form.as_div() + def test_clean(self): - item = Item(id=1, embedding=[1, 2, 3]) + item = Item(id=1, embedding=[1, 2, 3], binary_embedding='101') item.full_clean() def test_get_or_create(self): @@ -308,3 +319,6 @@ def test_get_or_create(self): def test_missing(self): Item().save() assert Item.objects.first().embedding is None + assert Item.objects.first().half_embedding is None + assert Item.objects.first().binary_embedding is None + assert Item.objects.first().sparse_embedding is None From b2d1913c7cfd888687e56ee6dc9243a90b03cff9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 09:47:05 -0700 Subject: [PATCH 306/879] Improved Django tests [skip ci] --- pgvector/utils/sparsevec.py | 4 ++-- tests/test_django.py | 13 ++++++++++++- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index ebfa0a9..96bbcdf 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -79,11 +79,11 @@ def to_binary(self): @classmethod def from_text(cls, value): - elements, dim = value.split('/') + elements, dim = value.split('/', 2) indices = [] values = [] for e in elements[1:-1].split(','): - i, v = e.split(':') + i, v = e.split(':', 2) indices.append(int(i) - 1) values.append(float(v)) return cls(int(dim), indices, values) diff --git a/tests/test_django.py b/tests/test_django.py index 4fe4000..734c1ff 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -105,6 +105,12 @@ class Meta: fields = ['binary_embedding'] +class SparseVectorForm(ModelForm): + class Meta: + model = Item + fields = ['sparse_embedding'] + + class TestDjango: def setup_method(self, test_method): Item.objects.all().delete() @@ -309,8 +315,13 @@ def test_bit_form(self): assert form.is_valid() assert 'value="101"' in form.as_div() + def test_sparsevec_form(self): + form = SparseVectorForm(data={'sparse_embedding': '{1:1,3:2,5:3}/6'}) + assert form.is_valid() + assert 'value="{1:1,3:2,5:3}/6"' in form.as_div() + def test_clean(self): - item = Item(id=1, embedding=[1, 2, 3], binary_embedding='101') + item = Item(id=1, embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector.from_dense([1, 2, 3])) item.full_clean() def test_get_or_create(self): From 404458990ebc64f2e49a5236bb2e6aed03d3c9da Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 09:52:17 -0700 Subject: [PATCH 307/879] Improved Django tests [skip ci] --- tests/test_django.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 734c1ff..7f563d0 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -315,10 +315,41 @@ def test_bit_form(self): assert form.is_valid() assert 'value="101"' in form.as_div() + def test_bit_form_instance(self): + Item(id=1, binary_embedding='101').save() + item = Item.objects.get(pk=1) + form = BitForm(instance=item) + assert 'value="101"' in form.as_div() + + def test_bit_form_save(self): + Item(id=1, binary_embedding='101').save() + item = Item.objects.get(pk=1) + form = BitForm(instance=item, data={'binary_embedding': '010'}) + assert form.has_changed() + assert form.is_valid() + assert form.save() + assert '010' == Item.objects.get(pk=1).binary_embedding + def test_sparsevec_form(self): - form = SparseVectorForm(data={'sparse_embedding': '{1:1,3:2,5:3}/6'}) + form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) + assert form.is_valid() + assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() + + def test_sparsevec_form_instance(self): + Item(id=1, sparse_embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = SparseVectorForm(instance=item) + # TODO fix + # assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + + def test_sparsevec_form_save(self): + Item(id=1, sparse_embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = SparseVectorForm(instance=item, data={'sparse_embedding': '{1:4,2:5,3:6}/3'}) + assert form.has_changed() assert form.is_valid() - assert 'value="{1:1,3:2,5:3}/6"' in form.as_div() + assert form.save() + assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list() def test_clean(self): item = Item(id=1, embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector.from_dense([1, 2, 3])) From 47138b0aad47f196b242ea6405f9eea7c69bf256 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 09:54:43 -0700 Subject: [PATCH 308/879] Added test for halfvec form [skip ci] --- tests/test_django.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index 7f563d0..920bae4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -99,6 +99,12 @@ class Meta: fields = ['embedding'] +class HalfVectorForm(ModelForm): + class Meta: + model = Item + fields = ['half_embedding'] + + class BitForm(ModelForm): class Meta: model = Item @@ -310,6 +316,11 @@ def test_vector_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).embedding is None + def test_halfvec_form(self): + form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) + assert form.is_valid() + assert 'value="[1, 2, 3]"' in form.as_div() + def test_bit_form(self): form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() From 58718785a747ac7ca7c574c1778b50a7d76e1c4d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 10:02:53 -0700 Subject: [PATCH 309/879] Added more tests for halfvec form [skip ci] --- tests/test_django.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index 920bae4..7712d7a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -321,6 +321,22 @@ def test_halfvec_form(self): assert form.is_valid() assert 'value="[1, 2, 3]"' in form.as_div() + def test_halfvec_form_instance(self): + Item(id=1, half_embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = HalfVectorForm(instance=item) + # TODO fix + # assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + + def test_halfvec_form_save(self): + Item(id=1, half_embedding=[1, 2, 3]).save() + item = Item.objects.get(pk=1) + form = HalfVectorForm(instance=item, data={'half_embedding': '[4, 5, 6]'}) + assert form.has_changed() + assert form.is_valid() + assert form.save() + assert [4, 5, 6] == Item.objects.get(pk=1).half_embedding.to_list() + def test_bit_form(self): form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() From 2dc31ebc907ba585897b761eed3c5eb2140d6707 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 10:04:47 -0700 Subject: [PATCH 310/879] Added test for halfvec form [skip ci] --- tests/test_django.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index 7712d7a..c2ca8e3 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -337,6 +337,15 @@ def test_halfvec_form_save(self): assert form.save() assert [4, 5, 6] == Item.objects.get(pk=1).half_embedding.to_list() + def test_halfvec_form_save_missing(self): + Item(id=1).save() + item = Item.objects.get(pk=1) + form = HalfVectorForm(instance=item, data={'half_embedding': ''}) + assert form.is_valid() + # TODO fix + # assert form.save() + assert Item.objects.get(pk=1).half_embedding is None + def test_bit_form(self): form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() From 0686c58feffd7fb26b509cc83618f4b5aed18cdf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 10:31:17 -0700 Subject: [PATCH 311/879] Improved to_python to HalfVectorField --- pgvector/django/halfvec.py | 7 ++++++- tests/test_django.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index a7921b7..7ea6fbb 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -26,7 +26,12 @@ def from_db_value(self, value, expression, connection): return HalfVector._from_db(value) def to_python(self, value): - return HalfVector._from_db(value) + if value is None or isinstance(value, HalfVector): + return value + elif isinstance(value, str): + return HalfVector._from_db(value) + else: + return HalfVector(value) def get_prep_value(self, value): return HalfVector._to_db(value) diff --git a/tests/test_django.py b/tests/test_django.py index c2ca8e3..108bd85 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -388,7 +388,7 @@ def test_sparsevec_form_save(self): assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list() def test_clean(self): - item = Item(id=1, embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector.from_dense([1, 2, 3])) + item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector.from_dense([1, 2, 3])) item.full_clean() def test_get_or_create(self): From 1e2006a8138b674fe3d4a6080396b081096c9d84 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 15:00:14 -0700 Subject: [PATCH 312/879] Fixed HalfVectorField forms for Django - fixes #75 --- pgvector/django/halfvec.py | 20 ++++++++++++++++++++ tests/test_django.py | 6 ++---- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 7ea6fbb..6b59a7f 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -1,3 +1,4 @@ +from django import forms from django.db.models import Field from ..utils import HalfVector @@ -38,3 +39,22 @@ def get_prep_value(self, value): def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) + + def formfield(self, **kwargs): + return super().formfield(form_class=HalfVectorFormField, **kwargs) + + +class HalfVectorWidget(forms.TextInput): + def format_value(self, value): + if isinstance(value, HalfVector): + value = value.to_list() + return super().format_value(value) + + +class HalfVectorFormField(forms.CharField): + widget = HalfVectorWidget + + def to_python(self, value): + if isinstance(value, str) and value == '': + return None + return super().to_python(value) diff --git a/tests/test_django.py b/tests/test_django.py index 108bd85..e76f77a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -325,8 +325,7 @@ def test_halfvec_form_instance(self): Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item) - # TODO fix - # assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() def test_halfvec_form_save(self): Item(id=1, half_embedding=[1, 2, 3]).save() @@ -342,8 +341,7 @@ def test_halfvec_form_save_missing(self): item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': ''}) assert form.is_valid() - # TODO fix - # assert form.save() + assert form.save() assert Item.objects.get(pk=1).half_embedding is None def test_bit_form(self): From 8d6da7465055102b246cbed9f1dd01e3f8edabf5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 15:08:20 -0700 Subject: [PATCH 313/879] Improved SparseVectorField forms for Django --- pgvector/django/sparsevec.py | 11 +++++++++++ tests/test_django.py | 8 ++++++++ 2 files changed, 19 insertions(+) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 4ec734f..6633b23 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -1,3 +1,4 @@ +from django import forms from django.db.models import Field from ..utils import SparseVector @@ -33,3 +34,13 @@ def get_prep_value(self, value): def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) + + def formfield(self, **kwargs): + return super().formfield(form_class=SparseVectorFormField, **kwargs) + + +class SparseVectorFormField(forms.CharField): + def to_python(self, value): + if isinstance(value, str) and value == '': + return None + return super().to_python(value) diff --git a/tests/test_django.py b/tests/test_django.py index e76f77a..5692d2a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -385,6 +385,14 @@ def test_sparsevec_form_save(self): assert form.save() assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list() + def test_sparesevec_form_save_missing(self): + Item(id=1).save() + item = Item.objects.get(pk=1) + form = SparseVectorForm(instance=item, data={'sparse_embedding': ''}) + assert form.is_valid() + assert form.save() + assert Item.objects.get(pk=1).sparse_embedding is None + def test_clean(self): item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector.from_dense([1, 2, 3])) item.full_clean() From 6d41b794b64e91a760cdb6268e1a8d87870b09e4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 15:10:30 -0700 Subject: [PATCH 314/879] Improved SparseVectorField forms for Django [skip ci] --- pgvector/django/sparsevec.py | 9 +++++++++ tests/test_django.py | 4 ++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 6633b23..d0d2d07 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -39,7 +39,16 @@ def formfield(self, **kwargs): return super().formfield(form_class=SparseVectorFormField, **kwargs) +class SparseVectorWidget(forms.TextInput): + def format_value(self, value): + if isinstance(value, SparseVector): + value = value.to_text() + return super().format_value(value) + + class SparseVectorFormField(forms.CharField): + widget = SparseVectorWidget + def to_python(self, value): if isinstance(value, str) and value == '': return None diff --git a/tests/test_django.py b/tests/test_django.py index 5692d2a..bd229fe 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -373,8 +373,8 @@ def test_sparsevec_form_instance(self): Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item) - # TODO fix - # assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + # TODO improve + assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() def test_sparsevec_form_save(self): Item(id=1, sparse_embedding=[1, 2, 3]).save() From d4469b97327241afb2e88f3b12bcb38a8def362b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 12 Jun 2024 15:16:50 -0700 Subject: [PATCH 315/879] Improved BitField forms for Django [skip ci] --- pgvector/django/bit.py | 11 +++++++++++ tests/test_django.py | 8 ++++++++ 2 files changed, 19 insertions(+) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 941d694..2cc847a 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -1,3 +1,4 @@ +from django import forms from django.db.models import Field @@ -19,3 +20,13 @@ def db_type(self, connection): if self.length is None: return 'bit' return 'bit(%d)' % self.length + + def formfield(self, **kwargs): + return super().formfield(form_class=BitFormField, **kwargs) + + +class BitFormField(forms.CharField): + def to_python(self, value): + if isinstance(value, str) and value == '': + return None + return super().to_python(value) diff --git a/tests/test_django.py b/tests/test_django.py index bd229fe..186cf3a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -364,6 +364,14 @@ def test_bit_form_save(self): assert form.save() assert '010' == Item.objects.get(pk=1).binary_embedding + def test_bit_form_save_missing(self): + Item(id=1).save() + item = Item.objects.get(pk=1) + form = BitForm(instance=item, data={'binary_embedding': ''}) + assert form.is_valid() + assert form.save() + assert Item.objects.get(pk=1).binary_embedding is None + def test_sparsevec_form(self): form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) assert form.is_valid() From aba997c0e3a3f2df706ffb082ac17d6f2b6f1a28 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 14 Jun 2024 10:29:02 -0700 Subject: [PATCH 316/879] Updated test name [skip ci] --- tests/test_sparse_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index c4565fb..1286dbb 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -10,7 +10,7 @@ def test_from_dense(self): assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] assert SparseVector.from_dense(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] - def test_from_coordinates(self): + def test_from_dict(self): assert SparseVector.from_dict({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] def test_from_sparse(self): From 50def9c996b0a90d7d4e4dccf6807b03a179ff13 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 16:22:41 -0700 Subject: [PATCH 317/879] Simplified SparseVector construction --- examples/sparse_search.py | 4 +- pgvector/utils/sparsevec.py | 83 ++++++++++++++++++++----------------- tests/test_asyncpg.py | 2 +- tests/test_django.py | 18 ++++---- tests/test_peewee.py | 8 ++-- tests/test_psycopg.py | 12 +++--- tests/test_psycopg2.py | 2 +- tests/test_sparse_vector.py | 37 ++++++++++++----- tests/test_sqlalchemy.py | 6 +-- tests/test_sqlmodel.py | 6 +-- 10 files changed, 100 insertions(+), 78 deletions(-) diff --git a/examples/sparse_search.py b/examples/sparse_search.py index 7e786e3..6ce33e8 100644 --- a/examples/sparse_search.py +++ b/examples/sparse_search.py @@ -45,10 +45,10 @@ def fetch_embeddings(input): ] embeddings = fetch_embeddings(input) for content, embedding in zip(input, embeddings): - conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector.from_dense(embedding))) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector(embedding))) query = 'forest' query_embedding = fetch_embeddings([query])[0] -result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector.from_dense(query_embedding),)).fetchall() +result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector(query_embedding),)).fetchall() for row in result: print(row[0]) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 96bbcdf..e6d7f87 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -3,48 +3,54 @@ class SparseVector: - def __init__(self, dim, indices, values): - # TODO improve - self._dim = int(dim) - self._indices = [int(i) for i in indices] - self._values = [float(v) for v in values] + def __init__(self, value, dimensions=None): + if value.__class__.__module__ == 'scipy.sparse._arrays': + if dimensions is not None: + raise ValueError('dimensions not allowed') + + self._from_sparse(value) + elif isinstance(value, dict): + self._from_dict(value, dimensions) + else: + if dimensions is not None: + raise ValueError('dimensions not allowed') + + self._from_dense(value) def __repr__(self): - return f'SparseVector({self._dim}, {self._indices}, {self._values})' + return f'SparseVector({self.to_dict()}, {self.dim()})' + + def _from_dict(self, d, dim): + if dim is None: + raise ValueError('dimensions required') - @classmethod - def from_dict(cls, d, dim): elements = [(i, v) for i, v in d.items()] elements.sort() - indices = [int(v[0]) for v in elements] - values = [float(v[1]) for v in elements] - return cls(dim, indices, values) + self._dim = int(dim) + self._indices = [int(v[0]) for v in elements] + self._values = [float(v[1]) for v in elements] - @classmethod - def from_sparse(cls, value): + def _from_sparse(self, value): value = value.tocoo() if value.ndim == 1: - dim = value.shape[0] + self._dim = value.shape[0] elif value.ndim == 2 and value.shape[0] == 1: - dim = value.shape[1] + self._dim = value.shape[1] else: raise ValueError('expected ndim to be 1') if hasattr(value, 'coords'): # scipy 1.13+ - indices = value.coords[0].tolist() + self._indices = value.coords[0].tolist() else: - indices = value.col.tolist() - values = value.data.tolist() - return cls(dim, indices, values) + self._indices = value.col.tolist() + self._values = value.data.tolist() - @classmethod - def from_dense(cls, value): - dim = len(value) - indices = [i for i, v in enumerate(value) if v != 0] - values = [float(value[i]) for i in indices] - return cls(dim, indices, values) + def _from_dense(self, value): + self._dim = len(value) + self._indices = [i for i, v in enumerate(value) if v != 0] + self._values = [float(value[i]) for i in self._indices] def dim(self): return self._dim @@ -86,21 +92,30 @@ def from_text(cls, value): i, v = e.split(':', 2) indices.append(int(i) - 1) values.append(float(v)) - return cls(int(dim), indices, values) + return cls._from_parts(int(dim), indices, values) @classmethod def from_binary(cls, value): dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return cls(int(dim), indices, values) + return cls._from_parts(int(dim), indices, values) + + @classmethod + def _from_parts(cls, dim, indices, values): + vec = cls.__new__(cls) + vec._dim = dim + vec._indices = indices + vec._values = values + return vec @classmethod def _to_db(cls, value, dim=None): if value is None: return value - value = cls._to_db_value(value) + if not isinstance(value, cls): + value = cls(value) if dim is not None and value.dim() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) @@ -112,19 +127,11 @@ def _to_db_binary(cls, value): if value is None: return value - value = cls._to_db_value(value) + if not isinstance(value, cls): + value = cls(value) return value.to_binary() - @classmethod - def _to_db_value(cls, value): - if isinstance(value, cls): - return value - elif isinstance(value, (list, np.ndarray)): - return cls.from_dense(value) - else: - raise ValueError('expected sparsevec') - @classmethod def _from_db(cls, value): if value is None or isinstance(value, cls): diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 3bfc888..829883e 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -82,7 +82,7 @@ async def test_sparsevec(self): await register_vector(conn) - embedding = SparseVector.from_dense([1.5, 2, 3]) + embedding = SparseVector([1.5, 2, 3]) await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") diff --git a/tests/test_django.py b/tests/test_django.py index 186cf3a..421966f 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -88,9 +88,9 @@ class Migration(migrations.Migration): def create_items(): - Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1])).save() - Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2])).save() - Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2])).save() + Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])).save() + Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])).save() + Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])).save() class VectorForm(ModelForm): @@ -208,34 +208,34 @@ def test_bit_jaccard_distance(self): # assert [v.distance for v in items] == [0, 1/3, 1] def test_sparsevec(self): - Item(id=1, sparse_embedding=SparseVector.from_dense([1, 2, 3])).save() + Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() item = Item.objects.get(pk=1) assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() - distance = L2Distance('sparse_embedding', SparseVector.from_dense([1, 1, 1])) + distance = L2Distance('sparse_embedding', SparseVector([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)] def test_sparsevec_max_inner_product(self): create_items() - distance = MaxInnerProduct('sparse_embedding', SparseVector.from_dense([1, 1, 1])) + distance = MaxInnerProduct('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_sparsevec_cosine_distance(self): create_items() - distance = CosineDistance('sparse_embedding', SparseVector.from_dense([1, 1, 1])) + distance = CosineDistance('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_sparsevec_l1_distance(self): create_items() - distance = L1Distance('sparse_embedding', SparseVector.from_dense([1, 1, 1])) + distance = L1Distance('sparse_embedding', SparseVector([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, 3] @@ -402,7 +402,7 @@ def test_sparesevec_form_save_missing(self): assert Item.objects.get(pk=1).sparse_embedding is None def test_clean(self): - item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector.from_dense([1, 2, 3])) + item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector([1, 2, 3])) item.full_clean() def test_get_or_create(self): diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 1455303..0882890 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -30,9 +30,9 @@ class Meta: def create_items(): - Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1])) - Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2])) - Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2])) + Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])) + Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])) + Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])) class TestPeewee: @@ -132,7 +132,7 @@ def test_sparsevec(self): def test_sparsevec_l2_distance(self): create_items() - distance = Item.sparse_embedding.l2_distance(SparseVector.from_dense([1, 1, 1])) + distance = Item.sparse_embedding.l2_distance(SparseVector([1, 1, 1])) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 2de1ec7..79ac190 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -100,20 +100,20 @@ def test_bit_text_format(self): assert repr(Bit(res)) == 'Bit(010100001)' def test_sparsevec(self): - embedding = SparseVector.from_dense([1.5, 2, 3]) + embedding = SparseVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] assert res.to_list() == [1.5, 2, 3] def test_sparsevec_binary_format(self): - embedding = SparseVector.from_dense([1.5, 0, 2, 0, 3, 0]) + embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] assert res.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_sparsevec_text_format(self): - embedding = SparseVector.from_dense([1.5, 0, 2, 0, 3, 0]) + embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] assert res.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) @@ -122,20 +122,20 @@ def test_text_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: - copy.write_row([embedding, HalfVector(embedding), '101', SparseVector.from_dense(embedding)]) + copy.write_row([embedding, HalfVector(embedding), '101', SparseVector(embedding)]) def test_binary_copy(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector.from_dense(embedding)]) + copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) def test_binary_copy_set_types(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) - copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector.from_dense(embedding)]) + copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) @pytest.mark.asyncio async def test_async(self): diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index f18405f..54da6a7 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -46,7 +46,7 @@ def test_bit(self): assert res[1][0] is None def test_sparsevec(self): - embedding = SparseVector.from_dense([1.5, 2, 3]) + embedding = SparseVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 1286dbb..ae38e2a 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -6,27 +6,42 @@ class TestSparseVector: def test_from_dense(self): - assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_dense(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] + assert SparseVector(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] + + def test_from_dense_dimensions(self): + with pytest.raises(ValueError) as error: + SparseVector([1, 0, 2, 0, 3, 0], 6) + assert str(error.value) == 'dimensions not allowed' def test_from_dict(self): - assert SparseVector.from_dict({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] + + def test_from_dict_no_dimensions(self): + with pytest.raises(ValueError) as error: + SparseVector({0: 1, 2: 2, 4: 3}) + assert str(error.value) == 'dimensions required' def test_from_sparse(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) - assert SparseVector.from_sparse(arr).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector.from_sparse(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector(arr).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] + + def test_from_sparse_dimensions(self): + with pytest.raises(ValueError) as error: + SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) + assert str(error.value) == 'dimensions not allowed' def test_repr(self): - assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' - assert str(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])' + assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' + assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' def test_dim(self): - assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).dim() == 6 + assert SparseVector([1, 0, 2, 0, 3, 0]).dim() == 6 def test_to_dict(self): - assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3} + assert SparseVector([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3} def test_to_coo(self): - assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] + assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 1c0fb80..edce3dc 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -41,9 +41,9 @@ class Item(Base): def create_items(): session = Session(engine) - session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1]))) - session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2]))) - session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2]))) + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) session.commit() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 90f7e21..5685ce6 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -37,9 +37,9 @@ class Item(SQLModel, table=True): def create_items(): session = Session(engine) - session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector.from_dense([1, 1, 1]))) - session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector.from_dense([2, 2, 2]))) - session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector.from_dense([1, 1, 2]))) + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) session.commit() From 7088b11380ca5e8ab57cdf7f4e5a1ce180dfd1c3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 16:26:23 -0700 Subject: [PATCH 318/879] Renamed dim method to dimensions [skip ci] --- pgvector/utils/halfvec.py | 8 ++++---- pgvector/utils/sparsevec.py | 8 ++++---- pgvector/utils/vector.py | 8 ++++---- tests/test_half_vector.py | 4 ++-- tests/test_sparse_vector.py | 4 ++-- tests/test_vector.py | 4 ++-- 6 files changed, 18 insertions(+), 18 deletions(-) diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index d01b381..e1e5051 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -16,7 +16,7 @@ def __init__(self, value): def __repr__(self): return f'HalfVector({self.to_list()})' - def dim(self): + def dimensions(self): return len(self._value) def to_list(self): @@ -29,7 +29,7 @@ def to_text(self): return '[' + ','.join([str(float(v)) for v in self._value]) + ']' def to_binary(self): - return pack('>HH', self.dim(), 0) + self._value.tobytes() + return pack('>HH', self.dimensions(), 0) + self._value.tobytes() @classmethod def from_text(cls, value): @@ -48,8 +48,8 @@ def _to_db(cls, value, dim=None): if not isinstance(value, cls): value = cls(value) - if dim is not None and value.dim() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) + if dim is not None and value.dimensions() != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) return value.to_text() diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index e6d7f87..d2fff13 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -18,7 +18,7 @@ def __init__(self, value, dimensions=None): self._from_dense(value) def __repr__(self): - return f'SparseVector({self.to_dict()}, {self.dim()})' + return f'SparseVector({self.to_dict()}, {self.dimensions()})' def _from_dict(self, d, dim): if dim is None: @@ -52,7 +52,7 @@ def _from_dense(self, value): self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] - def dim(self): + def dimensions(self): return self._dim def to_dict(self): @@ -117,8 +117,8 @@ def _to_db(cls, value, dim=None): if not isinstance(value, cls): value = cls(value) - if dim is not None and value.dim() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) + if dim is not None and value.dimensions() != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) return value.to_text() diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 7a0f6bc..3fa2f35 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -16,7 +16,7 @@ def __init__(self, value): def __repr__(self): return f'Vector({self.to_list()})' - def dim(self): + def dimensions(self): return len(self._value) def to_list(self): @@ -29,7 +29,7 @@ def to_text(self): return '[' + ','.join([str(float(v)) for v in self._value]) + ']' def to_binary(self): - return pack('>HH', self.dim(), 0) + self._value.tobytes() + return pack('>HH', self.dimensions(), 0) + self._value.tobytes() @classmethod def from_text(cls, value): @@ -48,8 +48,8 @@ def _to_db(cls, value, dim=None): if not isinstance(value, cls): value = cls(value) - if dim is not None and value.dim() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dim())) + if dim is not None and value.dimensions() != dim: + raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) return value.to_text() diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index b2811b2..2a0d3a3 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -38,5 +38,5 @@ def test_repr(self): assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' - def test_dim(self): - assert HalfVector([1, 2, 3]).dim() == 3 + def test_dimensions(self): + assert HalfVector([1, 2, 3]).dimensions() == 3 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index ae38e2a..6a49f57 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -37,8 +37,8 @@ def test_repr(self): assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' - def test_dim(self): - assert SparseVector([1, 0, 2, 0, 3, 0]).dim() == 6 + def test_dimensions(self): + assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 def test_to_dict(self): assert SparseVector([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3} diff --git a/tests/test_vector.py b/tests/test_vector.py index d6ecb5b..90d6a9b 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -38,5 +38,5 @@ def test_repr(self): assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' - def test_dim(self): - assert Vector([1, 2, 3]).dim() == 3 + def test_dimensions(self): + assert Vector([1, 2, 3]).dimensions() == 3 From 88872a91a3fac052381ea6fd00d8068e066aaf47 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 16:29:21 -0700 Subject: [PATCH 319/879] Added indices and values methods to SparseVector [skip ci] --- pgvector/utils/sparsevec.py | 10 +++++++--- tests/test_sparse_vector.py | 7 +++++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index d2fff13..041e727 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -18,7 +18,8 @@ def __init__(self, value, dimensions=None): self._from_dense(value) def __repr__(self): - return f'SparseVector({self.to_dict()}, {self.dimensions()})' + elements = dict(zip(self._indices, self._values)) + return f'SparseVector({elements}, {self.dimensions()})' def _from_dict(self, d, dim): if dim is None: @@ -55,8 +56,11 @@ def _from_dense(self, value): def dimensions(self): return self._dim - def to_dict(self): - return dict(zip(self._indices, self._values)) + def indices(self): + return self._indices + + def values(self): + return self._values def to_coo(self): from scipy.sparse import coo_array diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 6a49f57..f03e270 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -40,8 +40,11 @@ def test_repr(self): def test_dimensions(self): assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 - def test_to_dict(self): - assert SparseVector([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3} + def test_indices(self): + assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] + + def test_values(self): + assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] def test_to_coo(self): assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] From 6fa839d32de6f8a69ca01896bbc45bb48afa18af Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 16:29:57 -0700 Subject: [PATCH 320/879] Improved test [skip ci] --- tests/test_sparse_vector.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index f03e270..59c3844 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -16,7 +16,8 @@ def test_from_dense_dimensions(self): assert str(error.value) == 'dimensions not allowed' def test_from_dict(self): - assert SparseVector({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6).to_list() == [1, 0, 2, 0, 3, 0] + assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] def test_from_dict_no_dimensions(self): with pytest.raises(ValueError) as error: From dfa3e98c4a0829bd9d606acb3139007865c36e62 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 16:32:55 -0700 Subject: [PATCH 321/879] Moved methods [skip ci] --- pgvector/utils/sparsevec.py | 64 ++++++++++++++++++------------------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 041e727..2060b2b 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -21,38 +21,6 @@ def __repr__(self): elements = dict(zip(self._indices, self._values)) return f'SparseVector({elements}, {self.dimensions()})' - def _from_dict(self, d, dim): - if dim is None: - raise ValueError('dimensions required') - - elements = [(i, v) for i, v in d.items()] - elements.sort() - self._dim = int(dim) - self._indices = [int(v[0]) for v in elements] - self._values = [float(v[1]) for v in elements] - - def _from_sparse(self, value): - value = value.tocoo() - - if value.ndim == 1: - self._dim = value.shape[0] - elif value.ndim == 2 and value.shape[0] == 1: - self._dim = value.shape[1] - else: - raise ValueError('expected ndim to be 1') - - if hasattr(value, 'coords'): - # scipy 1.13+ - self._indices = value.coords[0].tolist() - else: - self._indices = value.col.tolist() - self._values = value.data.tolist() - - def _from_dense(self, value): - self._dim = len(value) - self._indices = [i for i, v in enumerate(value) if v != 0] - self._values = [float(value[i]) for i in self._indices] - def dimensions(self): return self._dim @@ -87,6 +55,38 @@ def to_binary(self): nnz = len(self._indices) return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) + def _from_dict(self, d, dim): + if dim is None: + raise ValueError('dimensions required') + + elements = [(i, v) for i, v in d.items()] + elements.sort() + self._dim = int(dim) + self._indices = [int(v[0]) for v in elements] + self._values = [float(v[1]) for v in elements] + + def _from_sparse(self, value): + value = value.tocoo() + + if value.ndim == 1: + self._dim = value.shape[0] + elif value.ndim == 2 and value.shape[0] == 1: + self._dim = value.shape[1] + else: + raise ValueError('expected ndim to be 1') + + if hasattr(value, 'coords'): + # scipy 1.13+ + self._indices = value.coords[0].tolist() + else: + self._indices = value.col.tolist() + self._values = value.data.tolist() + + def _from_dense(self, value): + self._dim = len(value) + self._indices = [i for i, v in enumerate(value) if v != 0] + self._values = [float(value[i]) for i in self._indices] + @classmethod def from_text(cls, value): elements, dim = value.split('/', 2) From 0266a2ee38e3812bc66ba2a96d9b8bd73b58f42e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 18:39:25 -0700 Subject: [PATCH 322/879] Fixed dict constructor for SparseVector [skip ci] --- pgvector/utils/sparsevec.py | 2 +- tests/test_sparse_vector.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 2060b2b..7d1839f 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -59,7 +59,7 @@ def _from_dict(self, d, dim): if dim is None: raise ValueError('dimensions required') - elements = [(i, v) for i, v in d.items()] + elements = [(i, v) for i, v in d.items() if v != 0] elements.sort() self._dim = int(dim) self._indices = [int(v[0]) for v in elements] diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 59c3844..1fc70ba 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -16,8 +16,9 @@ def test_from_dense_dimensions(self): assert str(error.value) == 'dimensions not allowed' def test_from_dict(self): - assert SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] + vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] def test_from_dict_no_dimensions(self): with pytest.raises(ValueError) as error: From 1bc26e40545dfe6719354f48d34d7f403c6cf627 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 18:40:02 -0700 Subject: [PATCH 323/879] Improved code [skip ci] --- pgvector/utils/sparsevec.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 7d1839f..5d62554 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -61,6 +61,7 @@ def _from_dict(self, d, dim): elements = [(i, v) for i, v in d.items() if v != 0] elements.sort() + self._dim = int(dim) self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] From ac8680e0e5851e30ac8e843a2bcc8187d37a5fc0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 18:43:34 -0700 Subject: [PATCH 324/879] Require dimensions to be positional [skip ci] --- pgvector/utils/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 5d62554..c21e0d3 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -3,7 +3,7 @@ class SparseVector: - def __init__(self, value, dimensions=None): + def __init__(self, value, dimensions=None, /): if value.__class__.__module__ == 'scipy.sparse._arrays': if dimensions is not None: raise ValueError('dimensions not allowed') From ae8402fbc122a7d496007bb3bb62b611956025d2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 18:49:26 -0700 Subject: [PATCH 325/879] Improved tests [skip ci] --- tests/test_sparse_vector.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 1fc70ba..023493b 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -5,16 +5,19 @@ class TestSparseVector: - def test_from_dense(self): - assert SparseVector([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0] - assert SparseVector(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] + def test_from_list(self): + vec = SparseVector([1, 0, 2, 0, 3, 0]) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.to_numpy().tolist() == [1, 0, 2, 0, 3, 0] - def test_from_dense_dimensions(self): + def test_from_list_dimensions(self): with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) assert str(error.value) == 'dimensions not allowed' + def test_from_numpy(self): + assert SparseVector(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] + def test_from_dict(self): vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert vec.to_list() == [1, 0, 2, 0, 3, 0] @@ -25,12 +28,12 @@ def test_from_dict_no_dimensions(self): SparseVector({0: 1, 2: 2, 4: 3}) assert str(error.value) == 'dimensions required' - def test_from_sparse(self): + def test_from_scipy(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) assert SparseVector(arr).to_list() == [1, 0, 2, 0, 3, 0] assert SparseVector(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] - def test_from_sparse_dimensions(self): + def test_from_scipy_dimensions(self): with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) assert str(error.value) == 'dimensions not allowed' From 2b89e1688dc196bdd851d56331598783a9456fa5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 20 Jun 2024 18:54:55 -0700 Subject: [PATCH 326/879] Improved tests [skip ci] --- tests/test_sparse_vector.py | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 023493b..064f115 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -5,39 +5,49 @@ class TestSparseVector: - def test_from_list(self): + def test_list(self): vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.to_numpy().tolist() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] - def test_from_list_dimensions(self): + def test_list_dimensions(self): with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) assert str(error.value) == 'dimensions not allowed' - def test_from_numpy(self): - assert SparseVector(np.array([1, 0, 2, 0, 3, 0])).to_list() == [1, 0, 2, 0, 3, 0] + def test_ndarray(self): + vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] - def test_from_dict(self): + def test_dict(self): vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_from_dict_no_dimensions(self): + def test_dict_no_dimensions(self): with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) assert str(error.value) == 'dimensions required' - def test_from_scipy(self): + def test_coo_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) - assert SparseVector(arr).to_list() == [1, 0, 2, 0, 3, 0] - assert SparseVector(arr.todok()).to_list() == [1, 0, 2, 0, 3, 0] + vec = SparseVector(arr) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] - def test_from_scipy_dimensions(self): + def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) assert str(error.value) == 'dimensions not allowed' + def test_dok_array(self): + arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() + vec = SparseVector(arr) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + def test_repr(self): assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' From 0cf5ca5adb0abf3487a23698e91b8afee8e9e486 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 25 Jun 2024 16:46:14 -0700 Subject: [PATCH 327/879] Improved code --- pgvector/utils/sparsevec.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index c21e0d3..a451fc6 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -1,18 +1,23 @@ import numpy as np from struct import pack, unpack_from +NO_DEFAULT = object() + class SparseVector: - def __init__(self, value, dimensions=None, /): + def __init__(self, value, dimensions=NO_DEFAULT, /): if value.__class__.__module__ == 'scipy.sparse._arrays': - if dimensions is not None: + if dimensions is not NO_DEFAULT: raise ValueError('dimensions not allowed') self._from_sparse(value) elif isinstance(value, dict): + if dimensions is NO_DEFAULT: + raise ValueError('dimensions required') + self._from_dict(value, dimensions) else: - if dimensions is not None: + if dimensions is not NO_DEFAULT: raise ValueError('dimensions not allowed') self._from_dense(value) @@ -56,9 +61,6 @@ def to_binary(self): return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) def _from_dict(self, d, dim): - if dim is None: - raise ValueError('dimensions required') - elements = [(i, v) for i, v in d.items() if v != 0] elements.sort() From 1451cc60b995ac304f972284741a7a8cc28b1171 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 25 Jun 2024 17:04:54 -0700 Subject: [PATCH 328/879] Fixed check --- pgvector/utils/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index a451fc6..632962f 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -6,7 +6,7 @@ class SparseVector: def __init__(self, value, dimensions=NO_DEFAULT, /): - if value.__class__.__module__ == 'scipy.sparse._arrays': + if value.__class__.__module__.startswith('scipy.sparse.'): if dimensions is not NO_DEFAULT: raise ValueError('dimensions not allowed') From e005b45a29c985c5908d9fbf3ceb6d6bfb725f71 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 25 Jun 2024 17:07:28 -0700 Subject: [PATCH 329/879] Improved code [skip ci] --- pgvector/utils/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 632962f..0c35d25 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -24,7 +24,7 @@ def __init__(self, value, dimensions=NO_DEFAULT, /): def __repr__(self): elements = dict(zip(self._indices, self._values)) - return f'SparseVector({elements}, {self.dimensions()})' + return f'SparseVector({elements}, {self._dim})' def dimensions(self): return self._dim From 12370bb0b4d7e42c6ebee30f0eed360bf3c65f1d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 25 Jun 2024 17:45:39 -0700 Subject: [PATCH 330/879] Improved error messages --- pgvector/utils/sparsevec.py | 6 +++--- tests/test_sparse_vector.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 0c35d25..1c81876 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -8,17 +8,17 @@ class SparseVector: def __init__(self, value, dimensions=NO_DEFAULT, /): if value.__class__.__module__.startswith('scipy.sparse.'): if dimensions is not NO_DEFAULT: - raise ValueError('dimensions not allowed') + raise ValueError('extra argument') self._from_sparse(value) elif isinstance(value, dict): if dimensions is NO_DEFAULT: - raise ValueError('dimensions required') + raise ValueError('missing dimensions') self._from_dict(value, dimensions) else: if dimensions is not NO_DEFAULT: - raise ValueError('dimensions not allowed') + raise ValueError('extra argument') self._from_dense(value) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 064f115..643e841 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -14,7 +14,7 @@ def test_list(self): def test_list_dimensions(self): with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) - assert str(error.value) == 'dimensions not allowed' + assert str(error.value) == 'extra argument' def test_ndarray(self): vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) @@ -29,7 +29,7 @@ def test_dict(self): def test_dict_no_dimensions(self): with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) - assert str(error.value) == 'dimensions required' + assert str(error.value) == 'missing dimensions' def test_coo_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) @@ -40,7 +40,7 @@ def test_coo_array(self): def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) - assert str(error.value) == 'dimensions not allowed' + assert str(error.value) == 'extra argument' def test_dok_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() From 3eadd987c169247a61dbc3a025552e81cd02c61e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 25 Jun 2024 18:01:38 -0700 Subject: [PATCH 331/879] Version bump to 0.3.0 [skip ci] --- CHANGELOG.md | 2 +- README.md | 18 +++++++++--------- setup.py | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14db629..999a859 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.0 (unreleased) +## 0.3.0 (2024-06-25) - Added support for `halfvec`, `bit`, and `sparsevec` types to Django - Added support for `halfvec`, `bit`, and `sparsevec` types to SQLAlchemy and SQLModel diff --git a/README.md b/README.md index cf39d7e..670c6ea 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ Or check out some examples: - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_embeddings.py) with SentenceTransformers - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) -- [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search.py) with Transformers (unreleased) +- [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search.py) with Transformers - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit @@ -61,7 +61,7 @@ class Item(models.Model): embedding = VectorField(dimensions=3) ``` -Also supports `HalfVectorField` (unreleased), `BitField` (unreleased), and `SparseVectorField` (unreleased) +Also supports `HalfVectorField`, `BitField`, and `SparseVectorField` Insert a vector @@ -78,7 +78,7 @@ from pgvector.django import L2Distance Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5] ``` -Also supports `MaxInnerProduct`, `CosineDistance`, `L1Distance` (unreleased), `HammingDistance` (unreleased), and `JaccardDistance` (unreleased) +Also supports `MaxInnerProduct`, `CosineDistance`, `L1Distance`, `HammingDistance`, and `JaccardDistance` Get the distance @@ -146,7 +146,7 @@ class Item(Base): embedding = mapped_column(Vector(3)) ``` -Also supports `HALFVEC` (unreleased), `BIT` (unreleased), and `SPARSEVEC` (unreleased) +Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` Insert a vector @@ -162,7 +162,7 @@ Get the nearest neighbors to a vector session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) ``` -Also supports `max_inner_product`, `cosine_distance`, `l1_distance` (unreleased), `hamming_distance` (unreleased), and `jaccard_distance` (unreleased) +Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` Get the distance @@ -228,7 +228,7 @@ class Item(SQLModel, table=True): embedding: Any = Field(sa_column=Column(Vector(3))) ``` -Also supports `HALFVEC` (unreleased), `BIT` (unreleased), and `SPARSEVEC` (unreleased) +Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` Insert a vector @@ -244,7 +244,7 @@ Get the nearest neighbors to a vector session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) ``` -Also supports `max_inner_product`, `cosine_distance`, `l1_distance` (unreleased), `hamming_distance` (unreleased), and `jaccard_distance` (unreleased) +Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` Get the distance @@ -459,7 +459,7 @@ class Item(BaseModel): embedding = VectorField(dimensions=3) ``` -Also supports `HalfVectorField` (unreleased), `FixedBitField` (unreleased), and `SparseVectorField` (unreleased) +Also supports `HalfVectorField`, `FixedBitField`, and `SparseVectorField` Insert a vector @@ -473,7 +473,7 @@ Get the nearest neighbors to a vector Item.select().order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5) ``` -Also supports `max_inner_product`, `cosine_distance`, `l1_distance` (unreleased), `hamming_distance` (unreleased), and `jaccard_distance` (unreleased) +Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` Get the distance diff --git a/setup.py b/setup.py index dff986d..37e0f68 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.2.5', + version='0.3.0', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From a782ac6669a50a8b8ee2f89e29eff91f22367d3a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 27 Jun 2024 19:58:59 -0700 Subject: [PATCH 332/879] Added Morgan fingerprint example [skip ci] --- README.md | 1 + examples/morgan_fingerprints.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+) create mode 100644 examples/morgan_fingerprints.py diff --git a/README.md b/README.md index 670c6ea..cf60ef8 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ Or check out some examples: - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search.py) with Transformers - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing +- [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/morgan_fingerprints.py) with RDKit - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise - [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM diff --git a/examples/morgan_fingerprints.py b/examples/morgan_fingerprints.py new file mode 100644 index 0000000..afb56ec --- /dev/null +++ b/examples/morgan_fingerprints.py @@ -0,0 +1,32 @@ +# good resource +# https://www.rdkit.org/docs/GettingStartedInPython.html#morgan-fingerprints-circular-fingerprints + +from pgvector.psycopg import register_vector, Bit +import psycopg +from rdkit import Chem +from rdkit.Chem import AllChem + + +def generate_fingerprint(molecule): + fpgen = AllChem.GetMorganGenerator() + return fpgen.GetFingerprintAsNumPy(Chem.MolFromSmiles(molecule)) + + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS molecules') +conn.execute('CREATE TABLE molecules (id text PRIMARY KEY, fingerprint bit(2048))') + +molecules = ['Cc1ccccc1', 'Cc1ncccc1', 'c1ccccn1'] +for molecule in molecules: + fingerprint = generate_fingerprint(molecule) + conn.execute('INSERT INTO molecules (id, fingerprint) VALUES (%s, %s)', (molecule, Bit(fingerprint))) + +query_molecule = 'c1ccco1' +query_fingerprint = generate_fingerprint(query_molecule) +result = conn.execute('SELECT id, fingerprint <%%> %s AS distance FROM molecules ORDER BY distance LIMIT 5', (Bit(query_fingerprint),)).fetchall() +for row in result: + print(row) From 3bba751e2a7b729210a9da3235e26148347d6fff Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 27 Jun 2024 20:12:09 -0700 Subject: [PATCH 333/879] Updated image hash example [skip ci] --- examples/hash_image_search.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/examples/hash_image_search.py b/examples/hash_image_search.py index d30027d..33fef0e 100644 --- a/examples/hash_image_search.py +++ b/examples/hash_image_search.py @@ -1,5 +1,6 @@ from datasets import load_dataset import matplotlib.pyplot as plt +from pgvector.psycopg import register_vector, Bit import psycopg from imagehash import phash @@ -10,6 +11,9 @@ def hash_image(img): conn = psycopg.connect(dbname='pgvector_example', autocommit=True) +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + conn.execute('DROP TABLE IF EXISTS images') conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, hash bit(64))') @@ -23,13 +27,13 @@ def hash_image(img): cur = conn.cursor() with cur.copy('COPY images (hash) FROM STDIN') as copy: for image in images: - copy.write_row([image['hash']]) + copy.write_row([Bit(image['hash'])]) print('Querying hashes') results = [] for i in range(5): image = dataset['test'][i]['image'] - result = conn.execute('SELECT id FROM images ORDER BY bit_count(hash # %s) LIMIT 5', (hash_image(image),)).fetchall() + result = conn.execute('SELECT id FROM images ORDER BY hash <~> %s LIMIT 5', (hash_image(image),)).fetchall() nearest_images = [dataset['train'][row[0] - 1]['image'] for row in result] results.append([image] + nearest_images) From d36477ccc4796a5d2d4a7e91fd5114182c15787f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 28 Jun 2024 11:26:41 -0700 Subject: [PATCH 334/879] Added topic modeling example [skip ci] --- README.md | 1 + examples/requirements.txt | 1 + examples/topic_modeling.py | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 36 insertions(+) create mode 100644 examples/topic_modeling.py diff --git a/README.md b/README.md index cf60ef8..d2de992 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ Or check out some examples: - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing - [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/morgan_fingerprints.py) with RDKit +- [Topic modeling](https://github.com/pgvector/pgvector-python/blob/master/examples/topic_modeling.py) with Gensim - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit - [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise - [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM diff --git a/examples/requirements.txt b/examples/requirements.txt index 6832670..204e573 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -1,4 +1,5 @@ datasets +gensim imagehash implicit lightfm diff --git a/examples/topic_modeling.py b/examples/topic_modeling.py new file mode 100644 index 0000000..28d7fb2 --- /dev/null +++ b/examples/topic_modeling.py @@ -0,0 +1,34 @@ +from gensim.corpora.dictionary import Dictionary +from gensim.models import LdaModel +from gensim.utils import simple_preprocess +import numpy as np +from pgvector.psycopg import register_vector +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(20))') + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] + +tokens = [simple_preprocess(content) for content in input] +dictionary = Dictionary(tokens) +corpus = [dictionary.doc2bow(token) for token in tokens] +model = LdaModel(corpus, num_topics=20) + +for content, doc in zip(input, corpus): + embedding = np.array([v[1] for v in model[doc]]) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) + +document_id = 1 +neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() +for neighbor in neighbors: + print(neighbor[0]) From cbd8b9ee5bef83d7ed064c3389bd8d88a805923b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 28 Jun 2024 11:54:24 -0700 Subject: [PATCH 335/879] Fixed example [skip ci] --- examples/topic_modeling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/topic_modeling.py b/examples/topic_modeling.py index 28d7fb2..d4b8650 100644 --- a/examples/topic_modeling.py +++ b/examples/topic_modeling.py @@ -25,7 +25,7 @@ model = LdaModel(corpus, num_topics=20) for content, doc in zip(input, corpus): - embedding = np.array([v[1] for v in model[doc]]) + embedding = np.array([v[1] for v in model.get_document_topics(doc, minimum_probability=0)]) conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 From 88b808b7988b4372414ada22659f9709e0fd5ba6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 28 Jun 2024 12:02:20 -0700 Subject: [PATCH 336/879] Improved example [skip ci] --- examples/topic_modeling.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/topic_modeling.py b/examples/topic_modeling.py index d4b8650..6fa6b97 100644 --- a/examples/topic_modeling.py +++ b/examples/topic_modeling.py @@ -19,9 +19,9 @@ 'The bear is growling' ] -tokens = [simple_preprocess(content) for content in input] -dictionary = Dictionary(tokens) -corpus = [dictionary.doc2bow(token) for token in tokens] +docs = [simple_preprocess(content) for content in input] +dictionary = Dictionary(docs) +corpus = [dictionary.doc2bow(tokens) for tokens in docs] model = LdaModel(corpus, num_topics=20) for content, doc in zip(input, corpus): From 41692021c22fed6ebe62cbdad024f0e1e28c30d3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 28 Jun 2024 12:35:47 -0700 Subject: [PATCH 337/879] Improved naming [skip ci] --- examples/topic_modeling.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/topic_modeling.py b/examples/topic_modeling.py index 6fa6b97..875ac88 100644 --- a/examples/topic_modeling.py +++ b/examples/topic_modeling.py @@ -24,8 +24,8 @@ corpus = [dictionary.doc2bow(tokens) for tokens in docs] model = LdaModel(corpus, num_topics=20) -for content, doc in zip(input, corpus): - embedding = np.array([v[1] for v in model.get_document_topics(doc, minimum_probability=0)]) +for content, bow in zip(input, corpus): + embedding = np.array([v[1] for v in model.get_document_topics(bow, minimum_probability=0)]) conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 From 25074e04aa48fbe526762b0bca7dc3eb2d3dbf78 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 28 Jun 2024 12:39:50 -0700 Subject: [PATCH 338/879] Improved example [skip ci] --- examples/topic_modeling.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/topic_modeling.py b/examples/topic_modeling.py index 875ac88..cfbf18d 100644 --- a/examples/topic_modeling.py +++ b/examples/topic_modeling.py @@ -21,6 +21,7 @@ docs = [simple_preprocess(content) for content in input] dictionary = Dictionary(docs) +dictionary.filter_extremes(no_below=1) corpus = [dictionary.doc2bow(tokens) for tokens in docs] model = LdaModel(corpus, num_topics=20) From a55d2c26b9406009376fd6713d6a691d4ffe2893 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jul 2024 11:26:19 -0700 Subject: [PATCH 339/879] Added test for COPY TO with Psycopg 3 --- tests/test_psycopg.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 79ac190..44588a8 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -118,25 +118,34 @@ def test_sparsevec_text_format(self): assert res.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) - def test_text_copy(self): + def test_text_copy_from(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: copy.write_row([embedding, HalfVector(embedding), '101', SparseVector(embedding)]) - def test_binary_copy(self): + def test_binary_copy_from(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) - def test_binary_copy_set_types(self): + def test_binary_copy_from_set_types(self): embedding = np.array([1.5, 2, 3]) cur = conn.cursor() with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) + def test_binary_copy_to_set_types(self): + embedding = np.array([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) + cur = conn.cursor() + with cur.copy("COPY psycopg_items (embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + copy.set_types(['vector']) + for row in copy.rows(): + assert np.array_equal(row[0], embedding) + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) From ae23e0b01ca5a4111d57fc9d632dd8cfb2f114f2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jul 2024 11:33:23 -0700 Subject: [PATCH 340/879] Added more tests for COPY TO with Psycopg 3 --- tests/test_psycopg.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 44588a8..d08c7a4 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -137,6 +137,22 @@ def test_binary_copy_from_set_types(self): copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) + def test_text_copy_to(self): + embedding = np.array([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) + cur = conn.cursor() + with cur.copy("COPY psycopg_items (embedding) TO STDOUT") as copy: + for row in copy.rows(): + assert row[0] == "[1.5,2,3]" + + def test_binary_copy_to(self): + embedding = np.array([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) + cur = conn.cursor() + with cur.copy("COPY psycopg_items (embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + for row in copy.rows(): + assert Vector.from_binary(row[0]).to_list() == [1.5, 2, 3] + def test_binary_copy_to_set_types(self): embedding = np.array([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) From 29b4ee330a019428ca515662c31738f9b28f9fdd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jul 2024 11:37:29 -0700 Subject: [PATCH 341/879] Improved tests for COPY TO with Psycopg 3 [skip ci] --- tests/test_psycopg.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index d08c7a4..76d2840 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -139,28 +139,34 @@ def test_binary_copy_from_set_types(self): def test_text_copy_to(self): embedding = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) + half_embedding = HalfVector([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding) TO STDOUT") as copy: + with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT") as copy: for row in copy.rows(): assert row[0] == "[1.5,2,3]" + assert row[1] == "[1.5,2,3]" def test_binary_copy_to(self): embedding = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) + half_embedding = HalfVector([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: for row in copy.rows(): assert Vector.from_binary(row[0]).to_list() == [1.5, 2, 3] + assert HalfVector.from_binary(row[1]).to_list() == [1.5, 2, 3] def test_binary_copy_to_set_types(self): embedding = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s)', (embedding,)) + half_embedding = HalfVector([1.5, 2, 3]) + conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: - copy.set_types(['vector']) + with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + copy.set_types(['vector', 'halfvec']) for row in copy.rows(): assert np.array_equal(row[0], embedding) + assert row[1].to_list() == [1.5, 2, 3] @pytest.mark.asyncio async def test_async(self): From a07361bee7b248fe54c6b55263f6113c589e2ba3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 4 Jul 2024 08:52:02 -0700 Subject: [PATCH 342/879] Fixed NameError when vector type not found with Psycopg - fixes #77 --- pgvector/psycopg/vector.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 55e31a6..0f62ca9 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,3 +1,4 @@ +import psycopg from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from ..utils import Vector From 2e253ce10169595686b0ac9b0e5c609f9297cf47 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 10 Jul 2024 00:17:12 -0700 Subject: [PATCH 343/879] Fixed error with Psycopg 2 and pgvector < 0.7 - fixes #79 --- CHANGELOG.md | 4 ++++ pgvector/psycopg2/register.py | 24 ++++++++++-------------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 999a859..be87327 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.3.1 (unreleased) + +- Fixed error with Psycopg 2 and pgvector < 0.7 + ## 0.3.0 (2024-06-25) - Added support for `halfvec`, `bit`, and `sparsevec` types to Django diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 0ffd461..f299e86 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -7,20 +7,16 @@ def register_vector(conn_or_curs=None): cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs - try: - cur.execute('SELECT NULL::vector') - register_vector_info(cur.description[0][1]) - except psycopg2.errors.UndefinedObject: + cur.execute("SELECT typname, oid FROM pg_type WHERE typname IN ('vector', 'halfvec', 'sparsevec')") + type_info = dict(cur.fetchall()) + + if 'vector' not in type_info: raise psycopg2.ProgrammingError('vector type not found in the database') - try: - cur.execute('SELECT NULL::halfvec') - register_halfvec_info(cur.description[0][1]) - except psycopg2.errors.UndefinedObject: - pass + register_vector_info(type_info['vector']) + + if 'halfvec' in type_info: + register_halfvec_info(type_info['halfvec']) - try: - cur.execute('SELECT NULL::sparsevec') - register_sparsevec_info(cur.description[0][1]) - except psycopg2.errors.UndefinedObject: - pass + if 'sparsevec' in type_info: + register_sparsevec_info(type_info['sparsevec']) From 633cbd724380d445f47e405b801964c4b60fba6a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 10 Jul 2024 00:31:37 -0700 Subject: [PATCH 344/879] Updated changelog [skip ci] --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index be87327..00714d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.3.1 (unreleased) - Fixed error with Psycopg 2 and pgvector < 0.7 +- Fixed error message when `vector` type not found with Psycopg 3 ## 0.3.0 (2024-06-25) From 5c65394fd080d47e9813288e951e75511543f703 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 10 Jul 2024 12:21:05 -0700 Subject: [PATCH 345/879] Fixed error parsing zero sparse vectors --- CHANGELOG.md | 1 + pgvector/utils/sparsevec.py | 10 ++++++---- tests/test_sparse_vector.py | 4 ++++ 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 00714d8..972215d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.3.1 (unreleased) +- Fixed error parsing zero sparse vectors - Fixed error with Psycopg 2 and pgvector < 0.7 - Fixed error message when `vector` type not found with Psycopg 3 diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 1c81876..fd9ccff 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -95,10 +95,12 @@ def from_text(cls, value): elements, dim = value.split('/', 2) indices = [] values = [] - for e in elements[1:-1].split(','): - i, v = e.split(':', 2) - indices.append(int(i) - 1) - values.append(float(v)) + # split on empty string returns single element list + if len(elements) > 2: + for e in elements[1:-1].split(','): + i, v = e.split(':', 2) + indices.append(int(i) - 1) + values.append(float(v)) return cls._from_parts(int(dim), indices, values) @classmethod diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 643e841..06fe81a 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -63,3 +63,7 @@ def test_values(self): def test_to_coo(self): assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] + + def test_zero_vector_text(self): + vec = SparseVector({}, 3) + assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() From fcd7f613ca24fbfa50f2e8919dfe3382216519c5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 10 Jul 2024 19:07:14 -0700 Subject: [PATCH 346/879] Fixed backwards compatibility of type info query for Psycopg 2 --- pgvector/psycopg2/register.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index f299e86..e587721 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -7,7 +7,8 @@ def register_vector(conn_or_curs=None): cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs - cur.execute("SELECT typname, oid FROM pg_type WHERE typname IN ('vector', 'halfvec', 'sparsevec')") + # 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'))") type_info = dict(cur.fetchall()) if 'vector' not in type_info: From 13450b05a741d3258ce60cfd352e95dbe55892b8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 10 Jul 2024 20:33:25 -0700 Subject: [PATCH 347/879] Version bump to 0.3.1 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 972215d..1ffc98c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.1 (unreleased) +## 0.3.1 (2024-07-10) - Fixed error parsing zero sparse vectors - Fixed error with Psycopg 2 and pgvector < 0.7 diff --git a/setup.py b/setup.py index 37e0f68..cd7f909 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.3.0', + version='0.3.1', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From 182fe8491326d4e63b6bd5f376022a31fc6ecde5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 15 Jul 2024 17:04:13 -0700 Subject: [PATCH 348/879] Added Cohere example [skip ci] --- README.md | 1 + examples/cohere_embeddings.py | 34 ++++++++++++++++++++++++++++++++++ examples/requirements.txt | 1 + 3 files changed, 36 insertions(+) create mode 100644 examples/cohere_embeddings.py diff --git a/README.md b/README.md index d2de992..531973b 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ And follow the instructions for your database library: Or check out some examples: - [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai_embeddings.py) with OpenAI +- [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere_embeddings.py) with Cohere - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_embeddings.py) with SentenceTransformers - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) diff --git a/examples/cohere_embeddings.py b/examples/cohere_embeddings.py new file mode 100644 index 0000000..780352a --- /dev/null +++ b/examples/cohere_embeddings.py @@ -0,0 +1,34 @@ +import cohere +import numpy as np +from pgvector.psycopg import register_vector, Bit +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))') + + +def fetch_embeddings(input, input_type): + co = cohere.Client() + response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary']) + return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] + + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +embeddings = fetch_embeddings(input, 'search_document') +for content, embedding in zip(input, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Bit(embedding))) + +query = 'forest' +query_embedding = fetch_embeddings([query], 'search_query')[0] +result = conn.execute('SELECT content FROM documents ORDER BY embedding <~> %s LIMIT 5', (Bit(query_embedding),)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/requirements.txt b/examples/requirements.txt index 204e573..8b03fab 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -1,3 +1,4 @@ +cohere datasets gensim imagehash From 9c98a2d9730424e0705c97f4e684a65bbb7fedd5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 15 Jul 2024 17:15:39 -0700 Subject: [PATCH 349/879] Added todo [skip ci] --- pgvector/utils/bit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 8a6f2ac..fabac73 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -7,6 +7,7 @@ def __init__(self, value): if isinstance(value, str): self._value = self.from_text(value)._value else: + # TODO use np.unpackbits for uint8 array value = np.asarray(value, dtype=bool) if value.ndim != 1: From b87552b004bc0d665ebc2f93801cad9ce02c853b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 16 Jul 2024 17:45:33 -0700 Subject: [PATCH 350/879] Improved example [skip ci] --- examples/bulk_loading.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/bulk_loading.py b/examples/bulk_loading.py index 1df8fe7..ba058c4 100644 --- a/examples/bulk_loading.py +++ b/examples/bulk_loading.py @@ -38,7 +38,8 @@ print('\nSuccess!') # create any indexes *after* loading initial data (skipping for this example) -if False: +create_index = False +if create_index: print('Creating index') conn.execute("SET maintenance_work_mem = '8GB'") conn.execute('SET max_parallel_maintenance_workers = 7') From 4672a4b8a83b12c5abc5775839386928d94c6788 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Jul 2024 10:37:37 -0700 Subject: [PATCH 351/879] Fixed error with asyncpg and pgvector < 0.7 - fixes #83 --- CHANGELOG.md | 4 ++++ pgvector/asyncpg/register.py | 28 ++++++++++++++++------------ 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ffc98c..3b13d08 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.3.2 (unreleased) + +- Fixed error with asyncpg and pgvector < 0.7 + ## 0.3.1 (2024-07-10) - Fixed error parsing zero sparse vectors diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 7148ffa..ad75c2d 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -9,16 +9,20 @@ async def register_vector(conn): format='binary' ) - await conn.set_type_codec( - 'halfvec', - encoder=HalfVector._to_db_binary, - decoder=HalfVector._from_db_binary, - format='binary' - ) + try: + await conn.set_type_codec( + 'halfvec', + encoder=HalfVector._to_db_binary, + decoder=HalfVector._from_db_binary, + format='binary' + ) - await conn.set_type_codec( - 'sparsevec', - encoder=SparseVector._to_db_binary, - decoder=SparseVector._from_db_binary, - format='binary' - ) + await conn.set_type_codec( + 'sparsevec', + encoder=SparseVector._to_db_binary, + decoder=SparseVector._from_db_binary, + format='binary' + ) + except ValueError as e: + if not str(e).startswith('unknown type:'): + raise e From e77ee138d318b15c6276e0117a5905cc77e78527 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Jul 2024 13:15:18 -0700 Subject: [PATCH 352/879] Updated publish task [skip ci] --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 2199d93..c64f942 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ lint: publish: clean python3 setup.py bdist_wheel --universal ls dist - # twine upload dist/* + twine upload dist/* make clean clean: From 4498caaece8b72e4d4c2f67d400fbc6eb14a0ae4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Jul 2024 13:32:03 -0700 Subject: [PATCH 353/879] Added test for Bit constructor with uint8 --- pgvector/utils/bit.py | 7 ++++++- tests/test_bit.py | 6 ++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index fabac73..f2de6d7 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -7,7 +7,12 @@ def __init__(self, value): if isinstance(value, str): self._value = self.from_text(value)._value else: - # TODO use np.unpackbits for uint8 array + # TODO change in 0.4.0 + # if isinstance(value, np.ndarray) and value.dtype == np.uint8: + # value = np.unpackbits(value) + # else: + # value = np.asarray(value, dtype=bool) + value = np.asarray(value, dtype=bool) if value.ndim != 1: diff --git a/tests/test_bit.py b/tests/test_bit.py index c42bd05..32ab87b 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -13,6 +13,12 @@ def test_tuple(self): def test_str(self): assert Bit('101').to_list() == [True, False, True] + def test_ndarray_uint8(self): + arr = np.array([254, 7, 0], dtype=np.uint8) + # TODO change in 0.4.0 + # assert Bit(arr).to_text() == '111111100000011100000000' + assert Bit(arr).to_text() == '110' + def test_ndarray_same_object(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] From 50fbcabe46ac0a503a5fb74d19fad838f75ae5e0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Jul 2024 13:33:59 -0700 Subject: [PATCH 354/879] Added todo [skip ci] --- pgvector/utils/bit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index f2de6d7..51f7556 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -8,6 +8,7 @@ def __init__(self, value): self._value = self.from_text(value)._value else: # TODO change in 0.4.0 + # TODO raise if dtype not bool or uint8 # if isinstance(value, np.ndarray) and value.dtype == np.uint8: # value = np.unpackbits(value) # else: From 4f721eb96c4c26b85479536c3c6df1076ffd9dbe Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Jul 2024 13:52:12 -0700 Subject: [PATCH 355/879] Version bump to 0.3.2 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b13d08..e47f34f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.2 (unreleased) +## 0.3.2 (2024-07-17) - Fixed error with asyncpg and pgvector < 0.7 diff --git a/setup.py b/setup.py index cd7f909..db3df5c 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.3.1', + version='0.3.2', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From 2fc3ccdc1567b389f94c81be262ca282d87e5bd1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 2 Aug 2024 22:46:06 -0700 Subject: [PATCH 356/879] Updated readme [skip ci] --- README.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/README.md b/README.md index 531973b..b75cca0 100644 --- a/README.md +++ b/README.md @@ -530,3 +530,11 @@ pip install -r requirements.txt createdb pgvector_python_test pytest ``` + +To run an example: + +```sh +cd examples +pip install -r requirements.txt +python3 bulk_loading.py +``` From 93780f9f5952730e0c689b5bfc82d2dd93401e7d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 2 Aug 2024 22:49:42 -0700 Subject: [PATCH 357/879] Updated readme [skip ci] --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index b75cca0..2b8253b 100644 --- a/README.md +++ b/README.md @@ -536,5 +536,6 @@ To run an example: ```sh cd examples pip install -r requirements.txt +createdb pgvector_example python3 bulk_loading.py ``` From 3e32cef2688ced358d4d1f9c4bb9c8211d15998e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 12 Aug 2024 12:31:55 -0700 Subject: [PATCH 358/879] Added logging for Django tests [skip ci] --- tests/test_django.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 421966f..79837a3 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -7,6 +7,7 @@ from django.forms import ModelForm from math import sqrt import numpy as np +import os import pgvector.django from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVector, SparseVector from unittest import mock @@ -17,7 +18,25 @@ 'ENGINE': 'django.db.backends.postgresql', 'NAME': 'pgvector_python_test', } - } + }, + LOGGING={ + 'version': 1, + 'handlers': { + 'console': { + 'class': 'logging.StreamHandler' + } + }, + 'loggers': { + 'django.db.backends': { + 'handlers': ['console'], + 'level': 'DEBUG' + }, + 'django.db.backends.schema': { + 'level': 'WARNING' + } + } + }, + DEBUG=('VERBOSE' in os.environ) ) django.setup() From 04c39510910ec0b037e33817b6f54846fa1b4624 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 12 Aug 2024 12:32:46 -0700 Subject: [PATCH 359/879] Improved code [skip ci] --- tests/test_django.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 79837a3..06d86a4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -19,6 +19,7 @@ 'NAME': 'pgvector_python_test', } }, + DEBUG=('VERBOSE' in os.environ), LOGGING={ 'version': 1, 'handlers': { @@ -35,8 +36,7 @@ 'level': 'WARNING' } } - }, - DEBUG=('VERBOSE' in os.environ) + } ) django.setup() From 7fbb252788b64bf2cfa0d8bd4b59e7cd7a5a61c2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 12 Aug 2024 15:33:03 -0700 Subject: [PATCH 360/879] Added test for vector[] with Psycopg 3 [skip ci] --- tests/test_psycopg.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 76d2840..c4e1c22 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -7,7 +7,7 @@ conn.execute('CREATE EXTENSION IF NOT EXISTS vector') conn.execute('DROP TABLE IF EXISTS psycopg_items') -conn.execute('CREATE TABLE psycopg_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))') +conn.execute('CREATE TABLE psycopg_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector[])') register_vector(conn) @@ -168,6 +168,14 @@ def test_binary_copy_to_set_types(self): assert np.array_equal(row[0], embedding) assert row[1].to_list() == [1.5, 2, 3] + def test_vector_array(self): + embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + conn.execute('INSERT INTO psycopg_items (embeddings) VALUES (%s)', (embeddings,)) + + res = conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id').fetchone() + assert np.array_equal(res[0][0], embeddings[0]) + assert np.array_equal(res[0][1], embeddings[1]) + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) From 94c8c5494801ec3b254bb5bc7d40dde426a884d4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 12 Aug 2024 23:39:06 -0700 Subject: [PATCH 361/879] Added ColBERT example [skip ci] --- README.md | 1 + examples/colbert_exact.py | 48 +++++++++++++++++++++++++++++++++++++++ examples/requirements.txt | 1 + 3 files changed, 50 insertions(+) create mode 100644 examples/colbert_exact.py diff --git a/README.md b/README.md index 2b8253b..3d197ee 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ Or check out some examples: - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search.py) with Transformers +- [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert_exact.py) with ColBERT - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing - [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/morgan_fingerprints.py) with RDKit diff --git a/examples/colbert_exact.py b/examples/colbert_exact.py new file mode 100644 index 0000000..f953824 --- /dev/null +++ b/examples/colbert_exact.py @@ -0,0 +1,48 @@ +from colbert.infra import ColBERTConfig +from colbert.modeling.checkpoint import Checkpoint +import numpy as np +from pgvector.psycopg import register_vector +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embeddings vector(128)[])') +conn.execute(""" +CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ + WITH queries AS ( + SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) + ), + documents AS ( + SELECT unnest(document) AS document + ), + similarities AS ( + SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents + ), + max_similarities AS ( + SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number + ) + SELECT SUM(max_similarity) FROM max_similarities +$$ LANGUAGE SQL +""") + +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=ColBERTConfig()) + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +doc_embeddings = checkpoint.docFromText(input) +for content, embeddings in zip(input, doc_embeddings): + embeddings = [e.numpy() for e in embeddings if e.count_nonzero() > 0] + conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) + +query = 'puppy' +query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query], bsize=1)[0]] +result = conn.execute('SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() +for row in result: + print(row) diff --git a/examples/requirements.txt b/examples/requirements.txt index 8b03fab..45d8f17 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -1,4 +1,5 @@ cohere +colbert-ai datasets gensim imagehash From 02ee9161d96ca410b45a7374e202de7e73590127 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 13 Aug 2024 00:13:05 -0700 Subject: [PATCH 362/879] Improved example [skip ci] --- examples/colbert_exact.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/examples/colbert_exact.py b/examples/colbert_exact.py index f953824..dd6802d 100644 --- a/examples/colbert_exact.py +++ b/examples/colbert_exact.py @@ -1,6 +1,5 @@ from colbert.infra import ColBERTConfig from colbert.modeling.checkpoint import Checkpoint -import numpy as np from pgvector.psycopg import register_vector import psycopg @@ -29,7 +28,7 @@ $$ LANGUAGE SQL """) -checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=ColBERTConfig()) +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=ColBERTConfig(), verbose=0) input = [ 'The dog is barking', @@ -42,7 +41,7 @@ conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) query = 'puppy' -query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query], bsize=1)[0]] +query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query])[0]] result = conn.execute('SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() for row in result: print(row) From 1d98ca2b1eff47bd81ec5129859f2a4c750eed54 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 13 Aug 2024 00:14:46 -0700 Subject: [PATCH 363/879] Improved example [skip ci] --- examples/colbert_exact.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/colbert_exact.py b/examples/colbert_exact.py index dd6802d..672de4c 100644 --- a/examples/colbert_exact.py +++ b/examples/colbert_exact.py @@ -35,9 +35,9 @@ 'The cat is purring', 'The bear is growling' ] -doc_embeddings = checkpoint.docFromText(input) +doc_embeddings = checkpoint.docFromText(input, keep_dims=False) for content, embeddings in zip(input, doc_embeddings): - embeddings = [e.numpy() for e in embeddings if e.count_nonzero() > 0] + embeddings = [e.numpy() for e in embeddings] conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) query = 'puppy' From 780c1966be1ff4c205c026bd9dfa8255f7e4f5e5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 13 Aug 2024 12:29:07 -0700 Subject: [PATCH 364/879] Improved example [skip ci] --- examples/colbert_exact.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/colbert_exact.py b/examples/colbert_exact.py index 672de4c..06e4e76 100644 --- a/examples/colbert_exact.py +++ b/examples/colbert_exact.py @@ -28,7 +28,8 @@ $$ LANGUAGE SQL """) -checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=ColBERTConfig(), verbose=0) +config = ColBERTConfig(query_maxlen=5) +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) input = [ 'The dog is barking', From 472bdabde3f6a6b34a929063500ea4311f835b33 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 13 Aug 2024 12:33:17 -0700 Subject: [PATCH 365/879] Improved example [skip ci] --- examples/colbert_exact.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/colbert_exact.py b/examples/colbert_exact.py index 06e4e76..1c90b47 100644 --- a/examples/colbert_exact.py +++ b/examples/colbert_exact.py @@ -28,7 +28,7 @@ $$ LANGUAGE SQL """) -config = ColBERTConfig(query_maxlen=5) +config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) input = [ From a8a1bf106c079d330e1af3d9bbfb02df84fa6de6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 25 Aug 2024 12:38:22 -0700 Subject: [PATCH 366/879] Added tests for Psycopg 2 cursor factories --- tests/test_psycopg2.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 54da6a7..e6ba996 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,6 +1,7 @@ import numpy as np from pgvector.psycopg2 import register_vector, SparseVector import psycopg2 +from psycopg2.extras import DictCursor, NamedTupleCursor conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -53,3 +54,16 @@ def test_sparsevec(self): res = cur.fetchall() assert res[0][0].to_list() == [1.5, 2, 3] assert res[1][0] is None + + def test_cursor_factory(self): + for cursor_factory in [DictCursor, NamedTupleCursor]: + conn = psycopg2.connect(dbname='pgvector_python_test') + cur = conn.cursor(cursor_factory=cursor_factory) + register_vector(cur) + conn.close() + + def test_cursor_factory_connection(self): + for cursor_factory in [DictCursor, NamedTupleCursor]: + conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) + register_vector(conn) + conn.close() From d0a3c5a8c7fac6e03cd6f8bac9319e3e6d0746e8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 25 Aug 2024 12:43:49 -0700 Subject: [PATCH 367/879] Improved support for cursor factories with Psycopg 2 - closes #89 --- CHANGELOG.md | 4 ++++ pgvector/psycopg2/register.py | 4 +++- tests/test_psycopg2.py | 6 +++--- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e47f34f..6a856c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.3.3 (unreleased) + +- Improved support for cursor factories with Psycopg 2 + ## 0.3.2 (2024-07-17) - Fixed error with asyncpg and pgvector < 0.7 diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index e587721..60a7aa7 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -1,11 +1,13 @@ import psycopg2 +from psycopg2.extensions import cursor from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info def register_vector(conn_or_curs=None): - cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs + conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection + cur = conn.cursor(cursor_factory=cursor) # 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'))") diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index e6ba996..24e99d6 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,7 +1,7 @@ import numpy as np from pgvector.psycopg2 import register_vector, SparseVector import psycopg2 -from psycopg2.extras import DictCursor, NamedTupleCursor +from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -56,14 +56,14 @@ def test_sparsevec(self): assert res[1][0] is None def test_cursor_factory(self): - for cursor_factory in [DictCursor, NamedTupleCursor]: + for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) register_vector(cur) conn.close() def test_cursor_factory_connection(self): - for cursor_factory in [DictCursor, NamedTupleCursor]: + for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) register_vector(conn) conn.close() From 7e7a851d413d07cd7701ba7c411ed3093b531934 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 28 Aug 2024 13:35:47 -0700 Subject: [PATCH 368/879] Moved examples to separate directories [skip ci] --- README.md | 36 +++++++++---------- examples/{citus.py => citus/example.py} | 0 examples/citus/requirements.txt | 3 ++ .../example.py} | 0 examples/cohere/requirements.txt | 3 ++ .../{colbert_exact.py => colbert/exact.py} | 0 examples/colbert/requirements.txt | 3 ++ .../{topic_modeling.py => gensim/example.py} | 0 examples/gensim/requirements.txt | 5 +++ .../cross_encoder.py} | 0 examples/hybrid_search/requirements.txt | 3 ++ .../rrf.py} | 0 .../example.py} | 0 examples/image_search/requirements.txt | 6 ++++ .../example.py} | 2 +- examples/imagehash/requirements.txt | 5 +++ .../{implicit_recs.py => implicit/example.py} | 0 examples/implicit/requirements.txt | 4 +++ .../{lightfm_recs.py => lightfm/example.py} | 0 examples/lightfm/requirements.txt | 4 +++ .../{bulk_loading.py => loading/example.py} | 0 examples/loading/requirements.txt | 3 ++ .../example.py} | 0 examples/openai/requirements.txt | 3 ++ .../example.py} | 0 examples/rdkit/requirements.txt | 3 ++ examples/requirements.txt | 13 ------- .../example.py} | 0 .../sentence_transformers/requirements.txt | 3 ++ .../example.py} | 0 examples/sparse_search/requirements.txt | 5 +++ .../{surprise_recs.py => surprise/example.py} | 0 examples/surprise/requirements.txt | 4 +++ 33 files changed, 76 insertions(+), 32 deletions(-) rename examples/{citus.py => citus/example.py} (100%) create mode 100644 examples/citus/requirements.txt rename examples/{cohere_embeddings.py => cohere/example.py} (100%) create mode 100644 examples/cohere/requirements.txt rename examples/{colbert_exact.py => colbert/exact.py} (100%) create mode 100644 examples/colbert/requirements.txt rename examples/{topic_modeling.py => gensim/example.py} (100%) create mode 100644 examples/gensim/requirements.txt rename examples/{hybrid_search.py => hybrid_search/cross_encoder.py} (100%) create mode 100644 examples/hybrid_search/requirements.txt rename examples/{hybrid_search_rrf.py => hybrid_search/rrf.py} (100%) rename examples/{pytorch_image_search.py => image_search/example.py} (100%) create mode 100644 examples/image_search/requirements.txt rename examples/{hash_image_search.py => imagehash/example.py} (100%) create mode 100644 examples/imagehash/requirements.txt rename examples/{implicit_recs.py => implicit/example.py} (100%) create mode 100644 examples/implicit/requirements.txt rename examples/{lightfm_recs.py => lightfm/example.py} (100%) create mode 100644 examples/lightfm/requirements.txt rename examples/{bulk_loading.py => loading/example.py} (100%) create mode 100644 examples/loading/requirements.txt rename examples/{openai_embeddings.py => openai/example.py} (100%) create mode 100644 examples/openai/requirements.txt rename examples/{morgan_fingerprints.py => rdkit/example.py} (100%) create mode 100644 examples/rdkit/requirements.txt delete mode 100644 examples/requirements.txt rename examples/{sentence_embeddings.py => sentence_transformers/example.py} (100%) create mode 100644 examples/sentence_transformers/requirements.txt rename examples/{sparse_search.py => sparse_search/example.py} (100%) create mode 100644 examples/sparse_search/requirements.txt rename examples/{surprise_recs.py => surprise/example.py} (100%) create mode 100644 examples/surprise/requirements.txt diff --git a/README.md b/README.md index 3d197ee..f68fda4 100644 --- a/README.md +++ b/README.md @@ -26,22 +26,22 @@ And follow the instructions for your database library: Or check out some examples: -- [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai_embeddings.py) with OpenAI -- [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere_embeddings.py) with Cohere -- [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_embeddings.py) with SentenceTransformers -- [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) -- [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder) -- [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search.py) with Transformers -- [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert_exact.py) with ColBERT -- [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/pytorch_image_search.py) with PyTorch -- [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/hash_image_search.py) with perceptual hashing -- [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/morgan_fingerprints.py) with RDKit -- [Topic modeling](https://github.com/pgvector/pgvector-python/blob/master/examples/topic_modeling.py) with Gensim -- [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit_recs.py) with Implicit -- [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise_recs.py) with Surprise -- [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm_recs.py) with LightFM -- [Horizontal scaling](https://github.com/pgvector/pgvector-python/blob/master/examples/citus.py) with Citus -- [Bulk loading](https://github.com/pgvector/pgvector-python/blob/master/examples/bulk_loading.py) with `COPY` +- [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai/example.py) with OpenAI +- [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere/example.py) with Cohere +- [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_transformers/example.py) with SentenceTransformers +- [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) +- [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py) with SentenceTransformers (cross-encoder) +- [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search/example.py) with Transformers +- [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert/exact.py) with ColBERT +- [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/image_search/example.py) with PyTorch +- [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py) with perceptual hashing +- [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/rdkit/example.py) with RDKit +- [Topic modeling](https://github.com/pgvector/pgvector-python/blob/master/examples/gensim/example.py) with Gensim +- [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit/example.py) with Implicit +- [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise/example.py) with Surprise +- [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm/example.py) with LightFM +- [Horizontal scaling](https://github.com/pgvector/pgvector-python/blob/master/examples/citus/example.py) with Citus +- [Bulk loading](https://github.com/pgvector/pgvector-python/blob/master/examples/loading/example.py) with `COPY` ## Django @@ -535,8 +535,8 @@ pytest To run an example: ```sh -cd examples +cd examples/loading pip install -r requirements.txt createdb pgvector_example -python3 bulk_loading.py +python3 example.py ``` diff --git a/examples/citus.py b/examples/citus/example.py similarity index 100% rename from examples/citus.py rename to examples/citus/example.py diff --git a/examples/citus/requirements.txt b/examples/citus/requirements.txt new file mode 100644 index 0000000..1cf8ee9 --- /dev/null +++ b/examples/citus/requirements.txt @@ -0,0 +1,3 @@ +numpy +pgvector +psycopg[binary] diff --git a/examples/cohere_embeddings.py b/examples/cohere/example.py similarity index 100% rename from examples/cohere_embeddings.py rename to examples/cohere/example.py diff --git a/examples/cohere/requirements.txt b/examples/cohere/requirements.txt new file mode 100644 index 0000000..22fd056 --- /dev/null +++ b/examples/cohere/requirements.txt @@ -0,0 +1,3 @@ +cohere +pgvector +psycopg[binary] diff --git a/examples/colbert_exact.py b/examples/colbert/exact.py similarity index 100% rename from examples/colbert_exact.py rename to examples/colbert/exact.py diff --git a/examples/colbert/requirements.txt b/examples/colbert/requirements.txt new file mode 100644 index 0000000..4402ce8 --- /dev/null +++ b/examples/colbert/requirements.txt @@ -0,0 +1,3 @@ +colbert-ai +pgvector +psycopg[binary] diff --git a/examples/topic_modeling.py b/examples/gensim/example.py similarity index 100% rename from examples/topic_modeling.py rename to examples/gensim/example.py diff --git a/examples/gensim/requirements.txt b/examples/gensim/requirements.txt new file mode 100644 index 0000000..15411cd --- /dev/null +++ b/examples/gensim/requirements.txt @@ -0,0 +1,5 @@ +gensim +numpy +pgvector +psycopg[binary] +scipy<1.13 diff --git a/examples/hybrid_search.py b/examples/hybrid_search/cross_encoder.py similarity index 100% rename from examples/hybrid_search.py rename to examples/hybrid_search/cross_encoder.py diff --git a/examples/hybrid_search/requirements.txt b/examples/hybrid_search/requirements.txt new file mode 100644 index 0000000..237dcd1 --- /dev/null +++ b/examples/hybrid_search/requirements.txt @@ -0,0 +1,3 @@ +pgvector +psycopg[binary] +sentence-transformers diff --git a/examples/hybrid_search_rrf.py b/examples/hybrid_search/rrf.py similarity index 100% rename from examples/hybrid_search_rrf.py rename to examples/hybrid_search/rrf.py diff --git a/examples/pytorch_image_search.py b/examples/image_search/example.py similarity index 100% rename from examples/pytorch_image_search.py rename to examples/image_search/example.py diff --git a/examples/image_search/requirements.txt b/examples/image_search/requirements.txt new file mode 100644 index 0000000..3d82365 --- /dev/null +++ b/examples/image_search/requirements.txt @@ -0,0 +1,6 @@ +matplotlib +pgvector +psycopg[binary] +torch +torchvision +tqdm diff --git a/examples/hash_image_search.py b/examples/imagehash/example.py similarity index 100% rename from examples/hash_image_search.py rename to examples/imagehash/example.py index 33fef0e..f49af40 100644 --- a/examples/hash_image_search.py +++ b/examples/imagehash/example.py @@ -1,8 +1,8 @@ from datasets import load_dataset +from imagehash import phash import matplotlib.pyplot as plt from pgvector.psycopg import register_vector, Bit import psycopg -from imagehash import phash def hash_image(img): diff --git a/examples/imagehash/requirements.txt b/examples/imagehash/requirements.txt new file mode 100644 index 0000000..e3971e6 --- /dev/null +++ b/examples/imagehash/requirements.txt @@ -0,0 +1,5 @@ +datasets +imagehash +matplotlib +pgvector +psycopg[binary] diff --git a/examples/implicit_recs.py b/examples/implicit/example.py similarity index 100% rename from examples/implicit_recs.py rename to examples/implicit/example.py diff --git a/examples/implicit/requirements.txt b/examples/implicit/requirements.txt new file mode 100644 index 0000000..8f04b58 --- /dev/null +++ b/examples/implicit/requirements.txt @@ -0,0 +1,4 @@ +implicit +pgvector +psycopg[binary] +SQLAlchemy diff --git a/examples/lightfm_recs.py b/examples/lightfm/example.py similarity index 100% rename from examples/lightfm_recs.py rename to examples/lightfm/example.py diff --git a/examples/lightfm/requirements.txt b/examples/lightfm/requirements.txt new file mode 100644 index 0000000..cfa5f51 --- /dev/null +++ b/examples/lightfm/requirements.txt @@ -0,0 +1,4 @@ +lightfm +pgvector +psycopg[binary] +SQLAlchemy diff --git a/examples/bulk_loading.py b/examples/loading/example.py similarity index 100% rename from examples/bulk_loading.py rename to examples/loading/example.py diff --git a/examples/loading/requirements.txt b/examples/loading/requirements.txt new file mode 100644 index 0000000..1cf8ee9 --- /dev/null +++ b/examples/loading/requirements.txt @@ -0,0 +1,3 @@ +numpy +pgvector +psycopg[binary] diff --git a/examples/openai_embeddings.py b/examples/openai/example.py similarity index 100% rename from examples/openai_embeddings.py rename to examples/openai/example.py diff --git a/examples/openai/requirements.txt b/examples/openai/requirements.txt new file mode 100644 index 0000000..18587e2 --- /dev/null +++ b/examples/openai/requirements.txt @@ -0,0 +1,3 @@ +openai +pgvector +psycopg[binary] diff --git a/examples/morgan_fingerprints.py b/examples/rdkit/example.py similarity index 100% rename from examples/morgan_fingerprints.py rename to examples/rdkit/example.py diff --git a/examples/rdkit/requirements.txt b/examples/rdkit/requirements.txt new file mode 100644 index 0000000..85a3e4f --- /dev/null +++ b/examples/rdkit/requirements.txt @@ -0,0 +1,3 @@ +pgvector +psycopg[binary] +rdkit diff --git a/examples/requirements.txt b/examples/requirements.txt deleted file mode 100644 index 45d8f17..0000000 --- a/examples/requirements.txt +++ /dev/null @@ -1,13 +0,0 @@ -cohere -colbert-ai -datasets -gensim -imagehash -implicit -lightfm -matplotlib -openai -torch -torchvision -scikit-surprise -sentence-transformers diff --git a/examples/sentence_embeddings.py b/examples/sentence_transformers/example.py similarity index 100% rename from examples/sentence_embeddings.py rename to examples/sentence_transformers/example.py diff --git a/examples/sentence_transformers/requirements.txt b/examples/sentence_transformers/requirements.txt new file mode 100644 index 0000000..237dcd1 --- /dev/null +++ b/examples/sentence_transformers/requirements.txt @@ -0,0 +1,3 @@ +pgvector +psycopg[binary] +sentence-transformers diff --git a/examples/sparse_search.py b/examples/sparse_search/example.py similarity index 100% rename from examples/sparse_search.py rename to examples/sparse_search/example.py diff --git a/examples/sparse_search/requirements.txt b/examples/sparse_search/requirements.txt new file mode 100644 index 0000000..3de81c7 --- /dev/null +++ b/examples/sparse_search/requirements.txt @@ -0,0 +1,5 @@ +numpy +pgvector +psycopg[binary] +torch +transformers diff --git a/examples/surprise_recs.py b/examples/surprise/example.py similarity index 100% rename from examples/surprise_recs.py rename to examples/surprise/example.py diff --git a/examples/surprise/requirements.txt b/examples/surprise/requirements.txt new file mode 100644 index 0000000..cb2dca4 --- /dev/null +++ b/examples/surprise/requirements.txt @@ -0,0 +1,4 @@ +pgvector +psycopg[binary] +scikit-surprise +SQLAlchemy From c431eb658cb24a0545b5d106d83c3c8029b981fb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 9 Sep 2024 16:37:31 -0700 Subject: [PATCH 369/879] Version bump to 0.3.3 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a856c8..0bace8e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.3 (unreleased) +## 0.3.3 (2024-09-09) - Improved support for cursor factories with Psycopg 2 diff --git a/setup.py b/setup.py index db3df5c..ae636c6 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.3.2', + version='0.3.3', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From 3461181c17bb5814fbe80d9cd1f486ed543810e0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 20 Sep 2024 19:20:00 -0700 Subject: [PATCH 370/879] Added schema option for asyncpg - closes #90 --- CHANGELOG.md | 4 ++++ pgvector/asyncpg/register.py | 5 ++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0bace8e..bff6664 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.3.4 (unreleased) + +- Added `schema` option for asyncpg + ## 0.3.3 (2024-09-09) - Improved support for cursor factories with Psycopg 2 diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index ad75c2d..a388058 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -1,9 +1,10 @@ from ..utils import Vector, HalfVector, SparseVector -async def register_vector(conn): +async def register_vector(conn, schema='public'): await conn.set_type_codec( 'vector', + schema=schema, encoder=Vector._to_db_binary, decoder=Vector._from_db_binary, format='binary' @@ -12,6 +13,7 @@ async def register_vector(conn): try: await conn.set_type_codec( 'halfvec', + schema=schema, encoder=HalfVector._to_db_binary, decoder=HalfVector._from_db_binary, format='binary' @@ -19,6 +21,7 @@ async def register_vector(conn): await conn.set_type_codec( 'sparsevec', + schema=schema, encoder=SparseVector._to_db_binary, decoder=SparseVector._from_db_binary, format='binary' From e3f7bf029b6549bf10482a6dd511f5ecb5b33a01 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 22 Sep 2024 10:51:53 -0700 Subject: [PATCH 371/879] Updated example [skip ci] --- examples/loading/example.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/examples/loading/example.py b/examples/loading/example.py index ba058c4..0702129 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -31,10 +31,6 @@ copy.write_row([embedding]) - # flush data - while conn.pgconn.flush() == 1: - pass - print('\nSuccess!') # create any indexes *after* loading initial data (skipping for this example) From 6b648ee8a3d6ebfb8013024adfe59231881ce9b1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 26 Sep 2024 20:00:38 -0700 Subject: [PATCH 372/879] Version bump to 0.3.4 [skip ci] --- CHANGELOG.md | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bff6664..51ea394 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.4 (unreleased) +## 0.3.4 (2024-09-26) - Added `schema` option for asyncpg diff --git a/setup.py b/setup.py index ae636c6..0cebd22 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='pgvector', - version='0.3.3', + version='0.3.4', description='pgvector support for Python', long_description=long_description, long_description_content_type='text/markdown', From f0db3a412797c88246077a1a66b35b5fdecfe0da Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 1 Oct 2024 20:30:58 -0700 Subject: [PATCH 373/879] Switched to pyproject.toml --- Makefile | 2 +- pyproject.toml | 23 +++++++++++++++++++++++ pytest.ini | 2 -- setup.py | 30 ------------------------------ 4 files changed, 24 insertions(+), 33 deletions(-) create mode 100644 pyproject.toml delete mode 100644 pytest.ini delete mode 100644 setup.py diff --git a/Makefile b/Makefile index c64f942..e790e44 100644 --- a/Makefile +++ b/Makefile @@ -2,7 +2,7 @@ lint: pycodestyle . --ignore=E501 publish: clean - python3 setup.py bdist_wheel --universal + python3 -m build ls dist twine upload dist/* make clean diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..226c0a4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,23 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + +[project] +name = "pgvector" +version = "0.3.4" +description = "pgvector support for Python" +readme = "README.md" +authors = [ + {name = "Andrew Kane", email = "andrew@ankane.org"} +] +license = {text = "MIT"} +requires-python = ">= 3.8" +dependencies = [ + "numpy" +] + +[project.urls] +Homepage = "https://github.com/pgvector/pgvector-python" + +[tool.pytest.ini_options] +asyncio_mode = "auto" diff --git a/pytest.ini b/pytest.ini deleted file mode 100644 index 2f4c80e..0000000 --- a/pytest.ini +++ /dev/null @@ -1,2 +0,0 @@ -[pytest] -asyncio_mode = auto diff --git a/setup.py b/setup.py deleted file mode 100644 index 0cebd22..0000000 --- a/setup.py +++ /dev/null @@ -1,30 +0,0 @@ -from setuptools import setup - -with open('README.md', 'r', encoding='utf-8') as fh: - long_description = fh.read() - -setup( - name='pgvector', - version='0.3.4', - description='pgvector support for Python', - long_description=long_description, - long_description_content_type='text/markdown', - url='https://github.com/pgvector/pgvector-python', - author='Andrew Kane', - author_email='andrew@ankane.org', - license='MIT', - packages=[ - 'pgvector.asyncpg', - 'pgvector.django', - 'pgvector.peewee', - 'pgvector.psycopg', - 'pgvector.psycopg2', - 'pgvector.sqlalchemy', - 'pgvector.utils' - ], - python_requires='>=3.8', - install_requires=[ - 'numpy' - ], - zip_safe=False -) From 4483f03047d9c4095a0de75d16c402fad5e00dfc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 1 Oct 2024 20:33:50 -0700 Subject: [PATCH 374/879] Removed top-level __init__ to be consistent with previous packaging --- pgvector/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 pgvector/__init__.py diff --git a/pgvector/__init__.py b/pgvector/__init__.py deleted file mode 100644 index e69de29..0000000 From e6e43f563b4e9361d3afa553ef1f8169941fdaa7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 00:17:17 -0700 Subject: [PATCH 375/879] Improved Makefile [skip ci] --- Makefile | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/Makefile b/Makefile index e790e44..f0831c4 100644 --- a/Makefile +++ b/Makefile @@ -1,11 +1,13 @@ +.PHONY: lint build publish clean + lint: pycodestyle . --ignore=E501 -publish: clean +build: python3 -m build - ls dist + +publish: clean build twine upload dist/* - make clean clean: - rm -rf .pytest_cache build dist pgvector.egg-info + rm -rf .pytest_cache dist pgvector.egg-info From 9018d36200cc04ba931766acb068623cd8cb77d8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 10:16:51 -0700 Subject: [PATCH 376/879] Fixed deprecation warning with tests [skip ci] --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 226c0a4..25fcacc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,3 +21,4 @@ Homepage = "https://github.com/pgvector/pgvector-python" [tool.pytest.ini_options] asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" From b88ebed0089bed592ef8759c24c1b18bc772e5d2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 10:42:50 -0700 Subject: [PATCH 377/879] Added globally option for Psycopg 2 --- CHANGELOG.md | 4 ++++ pgvector/psycopg2/halfvec.py | 4 ++-- pgvector/psycopg2/register.py | 10 ++++++---- pgvector/psycopg2/sparsevec.py | 4 ++-- pgvector/psycopg2/vector.py | 4 ++-- 5 files changed, 16 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 51ea394..e6fd85c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.3.5 (unreleased) + +- Added `globally` option for Psycopg 2 + ## 0.3.4 (2024-09-26) - Added `schema` option for asyncpg diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index e3c0cdf..0fd66c1 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -14,7 +14,7 @@ def cast_halfvec(value, cur): return HalfVector._from_db(value) -def register_halfvec_info(oid): +def register_halfvec_info(oid, scope): halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) - register_type(halfvec) + register_type(halfvec, scope) register_adapter(HalfVector, HalfvecAdapter) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 60a7aa7..ae457f8 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -5,9 +5,11 @@ from .vector import register_vector_info -def register_vector(conn_or_curs=None): +# TODO make globally False by default in 0.4.0 +def register_vector(conn_or_curs=None, globally=True): 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'))") @@ -16,10 +18,10 @@ def register_vector(conn_or_curs=None): if 'vector' not in type_info: raise psycopg2.ProgrammingError('vector type not found in the database') - register_vector_info(type_info['vector']) + register_vector_info(type_info['vector'], scope) if 'halfvec' in type_info: - register_halfvec_info(type_info['halfvec']) + register_halfvec_info(type_info['halfvec'], scope) if 'sparsevec' in type_info: - register_sparsevec_info(type_info['sparsevec']) + register_sparsevec_info(type_info['sparsevec'], scope) diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 7cdf38c..03c7f7c 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -14,7 +14,7 @@ def cast_sparsevec(value, cur): return SparseVector._from_db(value) -def register_sparsevec_info(oid): +def register_sparsevec_info(oid, scope): sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) - register_type(sparsevec) + register_type(sparsevec, scope) register_adapter(SparseVector, SparsevecAdapter) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index c895f86..28c2f00 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -15,7 +15,7 @@ def cast_vector(value, cur): return Vector._from_db(value) -def register_vector_info(oid): +def register_vector_info(oid, scope): vector = new_type((oid,), 'VECTOR', cast_vector) - register_type(vector) + register_type(vector, scope) register_adapter(np.ndarray, VectorAdapter) From 70ae49a5e2e7e6c393f3b37041049091cf3007d5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 10:48:09 -0700 Subject: [PATCH 378/879] Updated tests to use globally=False --- tests/test_peewee.py | 12 ++++++++---- tests/test_psycopg2.py | 6 +++--- tests/test_sqlalchemy.py | 6 ++++-- tests/test_sqlmodel.py | 6 ++++-- 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 0882890..72502db 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -169,7 +169,8 @@ def test_vector_avg(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding)).scalar() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding)).scalar() @@ -177,7 +178,8 @@ def test_vector_sum(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding)).scalar() - assert np.array_equal(sum, np.array([5, 7, 9])) + # does not type cast + assert sum == '[5,7,9]' def test_halfvec_avg(self): avg = Item.select(fn.avg(Item.half_embedding)).scalar() @@ -185,7 +187,8 @@ def test_halfvec_avg(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.half_embedding)).scalar() - assert avg.to_list() == [2.5, 3.5, 4.5] + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_halfvec_sum(self): sum = Item.select(fn.sum(Item.half_embedding)).scalar() @@ -193,7 +196,8 @@ def test_halfvec_sum(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.half_embedding)).scalar() - assert sum.to_list() == [5, 7, 9] + # does not type cast + assert sum == '[5,7,9]' def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 24e99d6..9e5100c 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -11,7 +11,7 @@ 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))') -register_vector(cur) +register_vector(cur, globally=False) class TestPsycopg2: @@ -59,11 +59,11 @@ def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) - register_vector(cur) + register_vector(cur, globally=False) conn.close() def test_cursor_factory_connection(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) - register_vector(conn) + register_vector(conn, globally=False) conn.close() diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index edce3dc..d7d0756 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -337,7 +337,8 @@ def test_avg(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = session.query(func.avg(Item.embedding)).first()[0] - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_avg_orm(self): with Session(engine) as session: @@ -346,7 +347,8 @@ def test_avg_orm(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = session.scalars(select(func.avg(Item.embedding))).first() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_sum(self): with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 5685ce6..e1716e1 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -203,7 +203,8 @@ def test_vector_avg(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = session.exec(select(func.avg(Item.embedding))).first() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_vector_sum(self): with Session(engine) as session: @@ -221,7 +222,8 @@ def test_halfvec_avg(self): 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() - assert avg.to_list() == [2.5, 3.5, 4.5] + # does not type cast + assert avg == '[2.5,3.5,4.5]' def test_halfvec_sum(self): with Session(engine) as session: From d0aa5d3769ae05babdb03c74d6aabc03f7f62330 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 11:32:10 -0700 Subject: [PATCH 379/879] Fixed indentation [skip ci] --- tests/test_sqlmodel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index e1716e1..1211d8c 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -203,7 +203,7 @@ def test_vector_avg(self): 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 + # does not type cast assert avg == '[2.5,3.5,4.5]' def test_vector_sum(self): From 8879d95c1a75434cc413c238dec550e3c8167ec4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 11:46:37 -0700 Subject: [PATCH 380/879] Added note about register_adapter [skip ci] --- pgvector/psycopg2/register.py | 1 + tests/test_sqlalchemy.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index ae457f8..95d03a4 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -6,6 +6,7 @@ # 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): conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection cur = conn.cursor(cursor_factory=cursor) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index d7d0756..9419ee5 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -403,8 +403,9 @@ def test_insert(self): def test_insert_bulk(self): session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) - def test_insert_text(self): - session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) + # register_vector in psycopg2 tests change this behavior + # def test_insert_text(self): + # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) def test_automap(self): metadata = MetaData() From 89ec21dac00fb4e4836931e8eead4311d4bf5114 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 11:47:39 -0700 Subject: [PATCH 381/879] Updated tests to support SQLAlchemy 1 --- tests/test_sqlalchemy.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 9419ee5..77f5ab3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -4,10 +4,17 @@ from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer from sqlalchemy.exc import StatementError from sqlalchemy.ext.automap import automap_base -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import declarative_base, Session from sqlalchemy.sql import func +try: + from sqlalchemy.orm import mapped_column + from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + sqlalchemy_version = 2 +except ImportError: + mapped_column = Column + sqlalchemy_version = 1 + engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) @@ -418,6 +425,7 @@ def test_automap(self): assert item.embedding.tolist() == [1, 2, 3] @pytest.mark.asyncio + @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_async(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) From f65c3614c07184f73e0af854e92f89a928f8c3e6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 12:40:50 -0700 Subject: [PATCH 382/879] Fixed type casting for Peewee aggregations [skip ci] --- README.md | 2 +- tests/test_peewee.py | 28 ++++++++++++---------------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index f68fda4..4ed1b7a 100644 --- a/README.md +++ b/README.md @@ -496,7 +496,7 @@ Average vectors ```python from peewee import fn -Item.select(fn.avg(Item.embedding)).scalar() +Item.select(fn.avg(Item.embedding).coerce(True)).scalar() ``` Also supports `sum` diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 72502db..8785b60 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -164,40 +164,36 @@ def test_where(self): assert [v.id for v in items] == [1] def test_vector_avg(self): - avg = Item.select(fn.avg(Item.embedding)).scalar() + avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() assert avg is None Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) - avg = Item.select(fn.avg(Item.embedding)).scalar() - # does not type cast - assert avg == '[2.5,3.5,4.5]' + avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() + assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) def test_vector_sum(self): - sum = Item.select(fn.sum(Item.embedding)).scalar() + sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() assert sum is None Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) - sum = Item.select(fn.sum(Item.embedding)).scalar() - # does not type cast - assert sum == '[5,7,9]' + sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() + assert np.array_equal(sum, np.array([5, 7, 9])) def test_halfvec_avg(self): - avg = Item.select(fn.avg(Item.half_embedding)).scalar() + avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() assert avg is None Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) - avg = Item.select(fn.avg(Item.half_embedding)).scalar() - # does not type cast - assert avg == '[2.5,3.5,4.5]' + avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() + assert avg.to_list() == [2.5, 3.5, 4.5] def test_halfvec_sum(self): - sum = Item.select(fn.sum(Item.half_embedding)).scalar() + sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() assert sum is None Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) - sum = Item.select(fn.sum(Item.half_embedding)).scalar() - # does not type cast - assert sum == '[5,7,9]' + sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() + assert sum.to_list() == [5, 7, 9] def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) From e90260ae9a04ec776eba435b9bbc21cec008b570 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 13:00:15 -0700 Subject: [PATCH 383/879] Added avg function with type casting to SQLAlchemy - #44 Co-authored-by: lucasgadams --- CHANGELOG.md | 1 + pgvector/sqlalchemy/__init__.py | 5 ++++- pgvector/sqlalchemy/functions.py | 8 +++++++ tests/test_sqlalchemy.py | 36 +++++++++++++++----------------- tests/test_sqlmodel.py | 36 +++++++++++++++----------------- 5 files changed, 47 insertions(+), 39 deletions(-) create mode 100644 pgvector/sqlalchemy/functions.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e6fd85c..38b1830 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 67b1d16..382f53f 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/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' ] diff --git a/pgvector/sqlalchemy/functions.py b/pgvector/sqlalchemy/functions.py new file mode 100644 index 0000000..8f7ecc8 --- /dev/null +++ b/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' diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 77f5ab3..1c84108 100644 --- a/tests/test_sqlalchemy.py +++ b/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]) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 1211d8c..f6f26d6 100644 --- a/tests/test_sqlmodel.py +++ b/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]) From 6bc66f828943c4d93e8c89f9b18ec5738b251d40 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 2 Oct 2024 13:12:25 -0700 Subject: [PATCH 384/879] Wait to add avg and sum to __all__ [skip ci] --- pgvector/sqlalchemy/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 382f53f..6377c2c 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -14,6 +14,7 @@ 'SPARSEVEC', 'HalfVector', 'SparseVector', - 'avg', - 'sum' + # TODO add avg and sum in 0.4.0 + # 'avg', + # 'sum' ] From 107b81efd0144564483fa59ee3e7acdf814cb6f3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 5 Oct 2024 10:43:36 -0700 Subject: [PATCH 385/879] Removed unused code [skip ci] --- examples/sparse_search/example.py | 1 - pgvector/peewee/bit.py | 2 +- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/sparsevec.py | 2 +- pgvector/peewee/vector.py | 2 +- pgvector/psycopg/bit.py | 2 +- pgvector/psycopg/register.py | 1 - tests/test_asyncpg.py | 1 - tests/test_half_vector.py | 2 +- tests/test_sqlmodel.py | 1 - tests/test_vector.py | 2 +- 11 files changed, 7 insertions(+), 11 deletions(-) diff --git a/examples/sparse_search/example.py b/examples/sparse_search/example.py index 6ce33e8..fa6074e 100644 --- a/examples/sparse_search/example.py +++ b/examples/sparse_search/example.py @@ -2,7 +2,6 @@ # https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/ # https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1 -import numpy as np from pgvector.psycopg import register_vector, SparseVector import psycopg import torch diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index 8a3bfbc..ee5f12f 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -1,4 +1,4 @@ -from peewee import Expression, Field, Value +from peewee import Expression, Field class FixedBitField(Field): diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index bed7d1f..deaa14d 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,4 +1,4 @@ -from peewee import Expression, Field, Value +from peewee import Expression, Field from ..utils import HalfVector diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 8bba5cf..67f7d1b 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -1,4 +1,4 @@ -from peewee import Expression, Field, Value +from peewee import Expression, Field from ..utils import SparseVector diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index a9ebf87..22a87e5 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,4 +1,4 @@ -from peewee import Expression, Field, Value +from peewee import Expression, Field from ..utils import Vector diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index 605c6f0..f8eeb61 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,4 +1,4 @@ -from psycopg.adapt import Loader, Dumper +from psycopg.adapt import Dumper from psycopg.pq import Format from ..utils import Bit diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py index 7f54a31..b93fd3e 100644 --- a/pgvector/psycopg/register.py +++ b/pgvector/psycopg/register.py @@ -1,4 +1,3 @@ -import psycopg from psycopg.types import TypeInfo from .bit import register_bit_info from .halfvec import register_halfvec_info diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 829883e..e1b6250 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,4 +1,3 @@ -import asyncio import asyncpg import numpy as np from pgvector.asyncpg import register_vector, SparseVector diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 2a0d3a3..fdaa5f7 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -8,7 +8,7 @@ def test_list(self): assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] def test_list_str(self): - with pytest.raises(ValueError, match='could not convert string to float') as error: + with pytest.raises(ValueError, match='could not convert string to float'): HalfVector([1, 'two', 3]) def test_tuple(self): diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index f6f26d6..4cb0e9b 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -3,7 +3,6 @@ import pytest from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError -from sqlalchemy.sql import func from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text from typing import Any, Optional diff --git a/tests/test_vector.py b/tests/test_vector.py index 90d6a9b..1be2bc0 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -8,7 +8,7 @@ def test_list(self): assert Vector([1, 2, 3]).to_list() == [1, 2, 3] def test_list_str(self): - with pytest.raises(ValueError, match='could not convert string to float') as error: + with pytest.raises(ValueError, match='could not convert string to float'): Vector([1, 'two', 3]) def test_tuple(self): From 75a5d50a253c0cafbff8f6cb8d6ad9ad3b6c015b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 5 Oct 2024 11:03:42 -0700 Subject: [PATCH 386/879] Improved code [skip ci] --- pgvector/asyncpg/__init__.py | 1 + pgvector/sqlalchemy/__init__.py | 7 +++---- pgvector/sqlalchemy/functions.py | 6 ++++++ pgvector/utils/__init__.py | 15 +++++++++++---- 4 files changed, 21 insertions(+), 8 deletions(-) diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index f4e8754..543b882 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -3,6 +3,7 @@ __all__ = [ 'register_vector', + 'Vector', 'HalfVector', 'SparseVector' ] diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 6377c2c..4955eeb 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -8,13 +8,12 @@ __all__ = [ 'Vector', - 'VECTOR' + 'VECTOR', 'HALFVEC', 'BIT', 'SPARSEVEC', 'HalfVector', 'SparseVector', - # TODO add avg and sum in 0.4.0 - # 'avg', - # 'sum' + 'avg', + 'sum' ] diff --git a/pgvector/sqlalchemy/functions.py b/pgvector/sqlalchemy/functions.py index 8f7ecc8..72e3ca7 100644 --- a/pgvector/sqlalchemy/functions.py +++ b/pgvector/sqlalchemy/functions.py @@ -6,3 +6,9 @@ class avg(ReturnTypeFromArgs): inherit_cache = True package = 'pgvector' + + +__all__ = [ + 'avg', + 'sum' +] diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index 69e8cfb..3c01160 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1,4 +1,11 @@ -from .bit import * -from .halfvec import * -from .sparsevec import * -from .vector import * +from .bit import Bit +from .halfvec import HalfVector +from .sparsevec import SparseVector +from .vector import Vector + +__all__ = [ + 'Vector', + 'HalfVector', + 'Bit', + 'SparseVector' +] From 43b809fa519e9822406d90ad2630463c1ecea245 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 5 Oct 2024 18:19:25 -0700 Subject: [PATCH 387/879] Version bump to 0.3.5 [skip ci] --- CHANGELOG.md | 2 +- README.md | 4 ++-- pyproject.toml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 38b1830..b89f3a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.5 (unreleased) +## 0.3.5 (2024-10-05) - Added `avg` function with type casting to SQLAlchemy - Added `globally` option for Psycopg 2 diff --git a/README.md b/README.md index 4ed1b7a..613c387 100644 --- a/README.md +++ b/README.md @@ -183,7 +183,7 @@ session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) Average vectors ```python -from sqlalchemy.sql import func +from pgvector.sqlalchemy import avg session.scalars(select(func.avg(Item.embedding))).first() ``` @@ -265,7 +265,7 @@ session.exec(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) Average vectors ```python -from sqlalchemy.sql import func +from pgvector.sqlalchemy import avg session.exec(select(func.avg(Item.embedding))).first() ``` diff --git a/pyproject.toml b/pyproject.toml index 25fcacc..cceeab0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.3.4" +version = "0.3.5" description = "pgvector support for Python" readme = "README.md" authors = [ From 0852a1f9d2a3ae0a4d90c891d73c2a466d3361d9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 15:15:36 -0700 Subject: [PATCH 388/879] Added test for arrays with SQLAlchemy - #96 --- tests/test_sqlalchemy.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 1c84108..5f0c3ef 100644 --- a/tests/test_sqlalchemy.py +++ b/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 @@ def test_core(self): 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,14 @@ def test_automap(self): 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() + + # this fails if the driver does not cast arrays + # item = session.get(Item, 1) + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_async(self): From 82b7ab9922a9499c114035b7b8e0e8fc1c12a0a1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 15:35:21 -0700 Subject: [PATCH 389/879] Added arrays option for Psycopg 2 --- CHANGELOG.md | 4 ++++ pgvector/psycopg2/halfvec.py | 9 +++++++-- pgvector/psycopg2/register.py | 11 ++++++----- pgvector/psycopg2/sparsevec.py | 9 +++++++-- pgvector/psycopg2/vector.py | 9 +++++++-- tests/test_psycopg2.py | 33 ++++++++++++++++++++++++++++++--- 6 files changed, 61 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b89f3a0..efaa85c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.3.6 (unreleased) + +- Added `arrays` option for Psycopg 2 + ## 0.3.5 (2024-10-05) - Added `avg` function with type casting to SQLAlchemy diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 0fd66c1..b50e89b 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/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) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 95d03a4..7752852 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -7,22 +7,23 @@ # 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) diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 03c7f7c..a542807 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/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) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 28c2f00..9861f01 100644 --- a/pgvector/psycopg2/vector.py +++ b/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) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 9e5100c..c93fce4 100644 --- a/tests/test_psycopg2.py +++ b/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 @@ 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 @@ def test_sparsevec(self): 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') From 7887a34d5f1914592fcbd52be13b18b7ed43d39e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 15:51:16 -0700 Subject: [PATCH 390/879] Added test for arrays with asyncpg --- tests/test_asyncpg.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index e1b6250..7a68a9e 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -94,6 +94,24 @@ async def test_sparsevec(self): 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): From ba393fe2711142a49a8654304f9fb1051b1e2570 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 15:59:01 -0700 Subject: [PATCH 391/879] Added test for arrays with Django --- tests/test_django.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index 06d86a4..6990141 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,5 +1,6 @@ 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 @@ -46,6 +47,7 @@ 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) class Meta: app_label = 'django_app' @@ -82,6 +84,7 @@ 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)), ], ), migrations.AddIndex( @@ -433,3 +436,9 @@ def test_missing(self): 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() + + # this fails if the driver does not cast arrays + # item = Item.objects.get(pk=1) From c6c7d765fb84eaa29e7d7fd974d8aeaaa0734871 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 16:18:30 -0700 Subject: [PATCH 392/879] Added test for arrays with Peewee [skip ci] --- tests/test_peewee.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 8785b60..9666388 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -199,3 +199,25 @@ def test_get_or_create(self): 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])) From 2ff631e1e947954313d7d35426651c4788115057 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 16:25:51 -0700 Subject: [PATCH 393/879] Improved test for arrays with SQLAlchemy [skip ci] --- tests/test_sqlalchemy.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 5f0c3ef..e783ac0 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -425,12 +425,19 @@ def test_automap(self): assert item.embedding.tolist() == [1, 2, 3] def test_vector_array(self): + from pgvector.psycopg2 import register_vector + session = Session(engine) session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() - # this fails if the driver does not cast arrays - # item = session.get(Item, 1) + with engine.connect() as connection: + 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+') From 417f3ec7de0c8aaac7ce53ef8089ce7b0baa2483 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 16:26:52 -0700 Subject: [PATCH 394/879] Improved test [skip ci] --- tests/test_sqlalchemy.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index e783ac0..1ca0ea3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -425,13 +425,12 @@ def test_automap(self): assert item.embedding.tolist() == [1, 2, 3] def test_vector_array(self): - from pgvector.psycopg2 import register_vector - 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 From f2f30e1dd2691523a68ef347f847b38b44f5fa79 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 13 Oct 2024 16:30:50 -0700 Subject: [PATCH 395/879] Improved test for arrays with Django [skip ci] --- tests/test_django.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 6990141..bbeee11 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -440,5 +440,11 @@ def test_missing(self): def test_vector_array(self): Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]).save() - # this fails if the driver does not cast arrays - # item = Item.objects.get(pk=1) + 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] From 59a3efcbb00c1139a53fe0777e3400e15f4559f4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 16 Oct 2024 23:05:48 -0700 Subject: [PATCH 396/879] Updated readme [skip ci] --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 613c387..acd625d 100644 --- a/README.md +++ b/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` From 79d41119dabd2211601807263cee874c3a043d9f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 22 Oct 2024 13:21:05 -0700 Subject: [PATCH 397/879] Added tests for double and numeric arrays --- tests/test_django.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index bbeee11..5ab5f81 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -3,7 +3,8 @@ 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 @@ -48,6 +49,8 @@ class Item(models.Model): 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' @@ -85,6 +88,8 @@ class Migration(migrations.Migration): ('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( @@ -448,3 +453,23 @@ def test_vector_array(self): 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] From d24839b48f78059b2732f8dfb398fc1ffc8070cc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 26 Oct 2024 17:13:38 -0700 Subject: [PATCH 398/879] Version bump to 0.3.6 [skip ci] --- CHANGELOG.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index efaa85c..3a517d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.3.6 (unreleased) +## 0.3.6 (2024-10-26) - Added `arrays` option for Psycopg 2 diff --git a/pyproject.toml b/pyproject.toml index cceeab0..a6a6609 100644 --- a/pyproject.toml +++ b/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 = [ From 1037d7e4c05948b6b5bfc6f8d43e0e7730224f63 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 28 Oct 2024 19:37:18 -0700 Subject: [PATCH 399/879] Added test for half-precision indexing with SQLAlchemy - #98 --- tests/test_sqlalchemy.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 1ca0ea3..8a032ef 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -46,6 +46,15 @@ class Item(Base): ) index.create(engine) +half_precision_index = Index( + 'sqlalchemy_orm_half_precision_index', + func.cast(Item.embedding, HALFVEC(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'halfvec_l2_ops'} +) +half_precision_index.create(engine) + def create_items(): session = Session(engine) @@ -438,6 +447,12 @@ def test_vector_array(self): assert item.embeddings[0].tolist() == [1, 2, 3] assert item.embeddings[1].tolist() == [4, 5, 6] + def test_half_precision(self): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_async(self): From 32a8d04b06b0f3e77d639e9a9ed275a67fa1e36f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 28 Oct 2024 19:52:49 -0700 Subject: [PATCH 400/879] Added docs for half-precision indexing with SQLAlchemy - #98 [skip ci] --- README.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/README.md b/README.md index acd625d..bbf5973 100644 --- a/README.md +++ b/README.md @@ -214,6 +214,29 @@ index.create(engine) Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +#### Half-Precision Indexing + +Index vectors at half-precision + +```python +from pgvector.sqlalchemy import HALFVEC +from sqlalchemy.sql import func + +index = Index( + 'my_index', + func.cast(Item.embedding, HALFVEC(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) +``` + +Get the nearest neighbors + +```python +session.scalars(select(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])).limit(5)) +``` + ## SQLModel Enable the extension From 2c8fe09f824bc509ae692d2932fe1a0bc15b6923 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 28 Oct 2024 20:19:06 -0700 Subject: [PATCH 401/879] Fixed example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index bbf5973..10cca79 100644 --- a/README.md +++ b/README.md @@ -227,7 +227,7 @@ index = Index( func.cast(Item.embedding, HALFVEC(3)).label('embedding'), postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, - postgresql_ops={'embedding': 'vector_l2_ops'} + postgresql_ops={'embedding': 'halfvec_l2_ops'} ) ``` From 3d9ff72a270bbfedc2d579db7f2d03b4048dfbc0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 28 Oct 2024 20:21:06 -0700 Subject: [PATCH 402/879] Improved example [skip ci] --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 10cca79..917d1a7 100644 --- a/README.md +++ b/README.md @@ -234,7 +234,8 @@ index = Index( Get the nearest neighbors ```python -session.scalars(select(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])).limit(5)) +order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2]) +session.scalars(select(Item).order_by(order).limit(5)) ``` ## SQLModel From 25a30264599c9646b743efdfe1d28b99d6208f90 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 3 Nov 2024 19:30:28 -0800 Subject: [PATCH 403/879] Updated test [skip ci] --- tests/test_django.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 5ab5f81..92f3733 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -65,7 +65,7 @@ class Meta: name='hnsw_idx', fields=['embedding'], m=16, - ef_construction=100, + ef_construction=64, opclasses=['vector_l2_ops'] ) ] From 06a48c4699486b3dc2ab843104594d0fc4539038 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 09:12:53 -0800 Subject: [PATCH 404/879] Added pool example and tests for Psycopg 3 - closes #100 --- README.md | 9 +++++++++ requirements.txt | 2 +- tests/test_psycopg.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 917d1a7..0989ba7 100644 --- a/README.md +++ b/README.md @@ -338,6 +338,15 @@ from pgvector.psycopg import register_vector register_vector(conn) ``` +For [connection pools](https://www.psycopg.org/psycopg3/docs/advanced/pool.html), use + +```python +def configure(conn): + register_vector(conn) + +pool = ConnectionPool(configure=configure) +``` + For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use ```python diff --git a/requirements.txt b/requirements.txt index c1e11f3..0e30959 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,7 @@ asyncpg Django numpy peewee -psycopg[binary] +psycopg[binary,pool] psycopg2-binary pytest pytest-asyncio diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index c4e1c22..5802b2b 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,6 +1,7 @@ import numpy as np from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector, Vector import psycopg +from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) @@ -176,6 +177,18 @@ def test_vector_array(self): assert np.array_equal(res[0][0], embeddings[0]) assert np.array_equal(res[0][1], embeddings[1]) + def test_pool(self): + def configure(conn): + register_vector(conn) + + pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure) + + with pool.connection() as conn: + res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() + assert np.array_equal(res[0], np.array([1, 2, 3])) + + pool.close() + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) @@ -195,3 +208,19 @@ async def test_async(self): assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None + + @pytest.mark.asyncio + async def test_async_pool(self): + async def configure(conn): + await register_vector_async(conn) + + pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure) + await pool.open() + + async with pool.connection() as conn: + async with conn.cursor() as cur: + await cur.execute("SELECT '[1,2,3]'::vector") + res = await cur.fetchone() + assert np.array_equal(res[0], np.array([1, 2, 3])) + + await pool.close() From 49072f2e37ff97b07b422aa0a41c4d3bd312879f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 09:33:14 -0800 Subject: [PATCH 405/879] Updated readme [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 0989ba7..37a4737 100644 --- a/README.md +++ b/README.md @@ -344,7 +344,7 @@ For [connection pools](https://www.psycopg.org/psycopg3/docs/advanced/pool.html) def configure(conn): register_vector(conn) -pool = ConnectionPool(configure=configure) +pool = ConnectionPool(..., configure=configure) ``` For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use From 9e1c421c62a84f650cf8be73e0768bfab6715e7b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 11:18:41 -0800 Subject: [PATCH 406/879] Added docs and test for half-precision indexing with Django --- README.md | 25 +++++++++++++++++++++++++ tests/test_django.py | 25 ++++++++++++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 37a4737..938207a 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,31 @@ class Item(models.Model): Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +#### Half-Precision Indexing + +Index vectors at half-precision + +```python +from django.contrib.postgres.indexes import OpClass +from django.db.models.functions import Cast +from pgvector.django import HalfVectorField + +index = HnswIndex( + OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), + name='my_index', + m=16, + ef_construction=64 +) +``` + +Note: Add `'django.contrib.postgres'` to `INSTALLED_APPS` to use `OpClass` + +Get the nearest neighbors + +```python +Item.objects.order_by(L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2]))[:5] +``` + ## SQLAlchemy Enable the extension diff --git a/tests/test_django.py b/tests/test_django.py index 92f3733..353087e 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,6 +1,7 @@ import django from django.conf import settings from django.contrib.postgres.fields import ArrayField +from django.contrib.postgres.indexes import OpClass from django.core import serializers from django.db import connection, migrations, models from django.db.models import Avg, Sum, FloatField, DecimalField @@ -38,7 +39,12 @@ 'level': 'WARNING' } } - } + }, + # needed for OpClass + # https://docs.djangoproject.com/en/5.1/ref/contrib/postgres/indexes/#opclass-expressions + INSTALLED_APPS=[ + 'django.contrib.postgres' + ] ) django.setup() @@ -67,6 +73,12 @@ class Meta: m=16, ef_construction=64, opclasses=['vector_l2_ops'] + ), + HnswIndex( + OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), + name='hnsw_half_precision_idx', + m=16, + ef_construction=64 ) ] @@ -99,6 +111,10 @@ class Migration(migrations.Migration): migrations.AddIndex( model_name='item', index=pgvector.django.HnswIndex(fields=['embedding'], m=16, ef_construction=64, name='hnsw_idx', opclasses=['vector_l2_ops']), + ), + migrations.AddIndex( + model_name='item', + index=pgvector.django.HnswIndex(OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), m=16, ef_construction=64, name='hnsw_half_precision_idx'), ) ] @@ -473,3 +489,10 @@ def test_numeric_array(self): 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] + + def test_half_precision(self): + create_items() + distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [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)] From afcd67a4422dd390e07662bfe81a8fbfab571301 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 11:23:17 -0800 Subject: [PATCH 407/879] Updated readme [skip ci] --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 938207a..1000900 100644 --- a/README.md +++ b/README.md @@ -155,7 +155,8 @@ Note: Add `'django.contrib.postgres'` to `INSTALLED_APPS` to use `OpClass` Get the nearest neighbors ```python -Item.objects.order_by(L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2]))[:5] +distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2]) +Item.objects.order_by(distance)[:5] ``` ## SQLAlchemy From 441b26ec3dfbdfb6013ffdf18df083614d9fc5ff Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 11:30:19 -0800 Subject: [PATCH 408/879] Updated example [skip ci] --- README.md | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 1000900..b61059e 100644 --- a/README.md +++ b/README.md @@ -142,12 +142,16 @@ from django.contrib.postgres.indexes import OpClass from django.db.models.functions import Cast from pgvector.django import HalfVectorField -index = HnswIndex( - OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), - name='my_index', - m=16, - ef_construction=64 -) +class Item(models.Model): + class Meta: + indexes = [ + HnswIndex( + OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), + name='my_index', + m=16, + ef_construction=64 + ) + ] ``` Note: Add `'django.contrib.postgres'` to `INSTALLED_APPS` to use `OpClass` From 78e64594843f2eb833cca77d8f43c33971806963 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 11:31:05 -0800 Subject: [PATCH 409/879] Updated example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b61059e..94fed55 100644 --- a/README.md +++ b/README.md @@ -140,7 +140,7 @@ Index vectors at half-precision ```python from django.contrib.postgres.indexes import OpClass from django.db.models.functions import Cast -from pgvector.django import HalfVectorField +from pgvector.django import HnswIndex, HalfVectorField class Item(models.Model): class Meta: From 75e14d80c80975938c0f1c64f59901686f1cd24b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 12:02:54 -0800 Subject: [PATCH 410/879] Added pool test for Psycopg 2 --- tests/test_psycopg2.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index c93fce4..85f08aa 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -2,6 +2,7 @@ from pgvector.psycopg2 import register_vector, HalfVector, SparseVector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor +from psycopg2.pool import ThreadedConnectionPool conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -94,3 +95,21 @@ def test_cursor_factory_connection(self): conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) register_vector(conn, globally=False) conn.close() + + def test_pool(self): + pool = ThreadedConnectionPool(1, 3, dbname='pgvector_python_test') + + conn = pool.getconn() + try: + cur = conn.cursor() + + # use globally=True for apps + register_vector(cur, globally=False) + + cur.execute("SELECT '[1,2,3]'::vector") + res = cur.fetchone() + assert np.array_equal(res[0], np.array([1, 2, 3])) + finally: + pool.putconn(conn) + + pool.closeall() From 706cebcb4c10f5fc6288757744fcfe94cb461a0b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 12:04:39 -0800 Subject: [PATCH 411/879] Improved test --- tests/test_psycopg2.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 85f08aa..3730eb8 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -101,11 +101,14 @@ def test_pool(self): conn = pool.getconn() try: - cur = conn.cursor() - # use globally=True for apps - register_vector(cur, globally=False) + register_vector(conn, globally=False) + finally: + pool.putconn(conn) + conn = pool.getconn() + try: + cur = conn.cursor() cur.execute("SELECT '[1,2,3]'::vector") res = cur.fetchone() assert np.array_equal(res[0], np.array([1, 2, 3])) From 812a85e7ce40d42382d84244b25e2f44eddf2e94 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 12:04:59 -0800 Subject: [PATCH 412/879] Improved test [skip ci] --- tests/test_psycopg2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 3730eb8..3f52385 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -97,7 +97,7 @@ def test_cursor_factory_connection(self): conn.close() def test_pool(self): - pool = ThreadedConnectionPool(1, 3, dbname='pgvector_python_test') + pool = ThreadedConnectionPool(1, 1, dbname='pgvector_python_test') conn = pool.getconn() try: From 07a3b2b6eec65d332041dcec136ac9c75291bc2b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 5 Nov 2024 12:09:23 -0800 Subject: [PATCH 413/879] Updated comment [skip ci] --- tests/test_psycopg2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 3f52385..c3cd3cd 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -101,7 +101,7 @@ def test_pool(self): conn = pool.getconn() try: - # use globally=True for apps + # use globally=True for apps to ensure registered with all connections register_vector(conn, globally=False) finally: pool.putconn(conn) From ea32504ef8538c781fab1f579fcaec7b417b7163 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 11 Nov 2024 20:56:23 -0800 Subject: [PATCH 414/879] Updated pgvector on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f8bcaa3..04f1c21 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From 664b8ee8692a42236ff9b236ec2da635342b96c7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 18:28:01 -0800 Subject: [PATCH 415/879] Added test for halfvec arrays with SQLAlchemy - #101 --- tests/test_sqlalchemy.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 8a032ef..c9aa900 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -32,6 +32,7 @@ class Item(Base): binary_embedding = mapped_column(BIT(3)) sparse_embedding = mapped_column(SPARSEVEC(3)) embeddings = mapped_column(ARRAY(VECTOR(3))) + half_embeddings = mapped_column(ARRAY(HALFVEC(3))) Base.metadata.drop_all(engine) @@ -447,6 +448,20 @@ def test_vector_array(self): assert item.embeddings[0].tolist() == [1, 2, 3] assert item.embeddings[1].tolist() == [4, 5, 6] + def test_halfvec_array(self): + session = Session(engine) + session.add(Item(id=1, half_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.half_embeddings[0].to_list() == [1, 2, 3] + assert item.half_embeddings[1].to_list() == [4, 5, 6] + def test_half_precision(self): create_items() with Session(engine) as session: From 1c7e6a5fb3ea31512dacf71eaf4165eae9fa60e8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 19:04:21 -0800 Subject: [PATCH 416/879] Added docs for arrays with SQLAlchemy [skip ci] --- README.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/README.md b/README.md index 94fed55..44f1d93 100644 --- a/README.md +++ b/README.md @@ -268,6 +268,27 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2]) session.scalars(select(Item).order_by(order).limit(5)) ``` +#### Arrays + +Add an array column + +```python +from pgvector.sqlalchemy import Vector +from sqlalchemy import ARRAY + +class Item(Base): + embeddings = mapped_column(ARRAY(Vector(3))) +``` + +And register the types with the underlying driver + +```python +from pgvector.psycopg2 import register_vector + +with engine.connect() as connection: + register_vector(connection.connection.dbapi_connection, globally=True, arrays=True) +``` + ## SQLModel Enable the extension From 0a760663b1acd993c7caf364c8c087c50306a01f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 20:42:27 -0800 Subject: [PATCH 417/879] Use connection from session in example and tests --- README.md | 2 +- tests/test_sqlalchemy.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 44f1d93..757ade9 100644 --- a/README.md +++ b/README.md @@ -285,7 +285,7 @@ And register the types with the underlying driver ```python from pgvector.psycopg2 import register_vector -with engine.connect() as connection: +with session.connection() as connection: register_vector(connection.connection.dbapi_connection, globally=True, arrays=True) ``` diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index c9aa900..57cc12b 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -439,12 +439,12 @@ def test_vector_array(self): session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() - with engine.connect() as connection: + with session.connection() 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) + item = session.get(Item, 1) assert item.embeddings[0].tolist() == [1, 2, 3] assert item.embeddings[1].tolist() == [4, 5, 6] @@ -453,12 +453,12 @@ def test_halfvec_array(self): session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() - with engine.connect() as connection: + with session.connection() 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) + item = session.get(Item, 1) assert item.half_embeddings[0].to_list() == [1, 2, 3] assert item.half_embeddings[1].to_list() == [4, 5, 6] From 030def94b19329fa29c71f5273183f82c0550fd3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 20:55:15 -0800 Subject: [PATCH 418/879] Improved example and tests for arrays with SQLAlchemy - #101 [skip ci] --- README.md | 6 ++++-- tests/test_sqlalchemy.py | 39 ++++++++++++++++++++------------------- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 757ade9..991c51f 100644 --- a/README.md +++ b/README.md @@ -284,9 +284,11 @@ And register the types with the underlying driver ```python from pgvector.psycopg2 import register_vector +from sqlalchemy import engine -with session.connection() as connection: - register_vector(connection.connection.dbapi_connection, globally=True, arrays=True) +@event.listens_for(engine, "connect") +def connect(dbapi_connection, connection_record): + register_vector(dbapi_connection, arrays=True) ``` ## SQLModel diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 57cc12b..f8e4bb1 100644 --- a/tests/test_sqlalchemy.py +++ b/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, ARRAY +from sqlalchemy import create_engine, event, 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 @@ -20,6 +20,15 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() +array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') + + +@event.listens_for(array_engine, "connect") +def connect(dbapi_connection, connection_record): + from pgvector.psycopg2 import register_vector + register_vector(dbapi_connection, globally=False, arrays=True) + + Base = declarative_base() @@ -435,32 +444,24 @@ def test_automap(self): assert item.embedding.tolist() == [1, 2, 3] def test_vector_array(self): - session = Session(engine) + session = Session(array_engine) session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() - with session.connection() 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.get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + # this fails if the driver does not cast arrays + item = session.get(Item, 1) + assert item.embeddings[0].tolist() == [1, 2, 3] + assert item.embeddings[1].tolist() == [4, 5, 6] def test_halfvec_array(self): - session = Session(engine) + session = Session(array_engine) session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() - with session.connection() 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.get(Item, 1) - assert item.half_embeddings[0].to_list() == [1, 2, 3] - assert item.half_embeddings[1].to_list() == [4, 5, 6] + # this fails if the driver does not cast arrays + item = session.get(Item, 1) + assert item.half_embeddings[0].to_list() == [1, 2, 3] + assert item.half_embeddings[1].to_list() == [4, 5, 6] def test_half_precision(self): create_items() From d23844ef10dcd4297a9e2f3671ed8e851e0a2db1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 21:00:22 -0800 Subject: [PATCH 419/879] Fixed example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 991c51f..516f3aa 100644 --- a/README.md +++ b/README.md @@ -284,7 +284,7 @@ And register the types with the underlying driver ```python from pgvector.psycopg2 import register_vector -from sqlalchemy import engine +from sqlalchemy import event @event.listens_for(engine, "connect") def connect(dbapi_connection, connection_record): From 04aa5bca2ee60c73de91507e5eb7472a6cf6d7a6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 22:39:46 -0800 Subject: [PATCH 420/879] Added test for arrays with SQLAlchemy async - #101 --- tests/test_sqlalchemy.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f8e4bb1..77c03fc 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -483,3 +483,25 @@ async def test_async(self): assert avg.first() == '[2.5,3.5,4.5]' await engine.dispose() + + @pytest.mark.asyncio + @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') + async def test_async_vector_array(self): + engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async_session = async_sessionmaker(engine, expire_on_commit=False) + + @event.listens_for(engine.sync_engine, "connect") + def connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector_async + dbapi_connection.run_async(register_vector_async) + + async with async_session() as session: + async with session.begin(): + session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + + # this fails if the driver does not cast arrays + item = await session.get(Item, 1) + assert item.embeddings[0].tolist() == [1, 2, 3] + assert item.embeddings[1].tolist() == [4, 5, 6] + + await engine.dispose() From dbc44f4533e9edaa376dc4d4a18fea235c5e2187 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 12 Nov 2024 22:48:46 -0800 Subject: [PATCH 421/879] Added more examples for arrays with SQLAlchemy [skip ci] --- README.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/README.md b/README.md index 516f3aa..0dedce9 100644 --- a/README.md +++ b/README.md @@ -282,6 +282,30 @@ class Item(Base): And register the types with the underlying driver +For Psycopg 3, use + +```python +from pgvector.psycopg import register_vector +from sqlalchemy import event + +@event.listens_for(engine, "connect") +def connect(dbapi_connection, connection_record): + register_vector(dbapi_connection) +``` + +For [async connections](https://docs.sqlalchemy.org/en/20/orm/extensions/asyncio.html) with Psycopg 3, use + +```python +from pgvector.psycopg import register_vector_async +from sqlalchemy import event + +@event.listens_for(engine.sync_engine, "connect") +def connect(dbapi_connection, connection_record): + dbapi_connection.run_async(register_vector_async) +``` + +For Psycopg 2, use + ```python from pgvector.psycopg2 import register_vector from sqlalchemy import event From 368b363bbf9a48fe42bc114991c4e97ee140cdeb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 15:09:06 -0800 Subject: [PATCH 422/879] Added ColBERT example for binary embeddings [skip ci] --- examples/colbert/exact_binary.py | 53 ++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 examples/colbert/exact_binary.py diff --git a/examples/colbert/exact_binary.py b/examples/colbert/exact_binary.py new file mode 100644 index 0000000..8d398e2 --- /dev/null +++ b/examples/colbert/exact_binary.py @@ -0,0 +1,53 @@ +from colbert.infra import ColBERTConfig +from colbert.modeling.checkpoint import Checkpoint +from pgvector.psycopg import register_vector, Bit +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embeddings bit(128)[])') +conn.execute(""" +CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) RETURNS double precision AS $$ + WITH queries AS ( + SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) + ), + documents AS ( + SELECT unnest(document) AS document + ), + similarities AS ( + SELECT query_number, 1 - ((document <~> query) / bit_length(query)) AS similarity FROM queries CROSS JOIN documents + ), + max_similarities AS ( + SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number + ) + SELECT SUM(max_similarity) FROM max_similarities +$$ LANGUAGE SQL +""") + + +def binary_quantize(embeddings): + return [Bit(e.numpy()) for e in (embeddings > 0)] + + +config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +doc_embeddings = checkpoint.docFromText(input, keep_dims=False) +for content, embeddings in zip(input, doc_embeddings): + embeddings = binary_quantize(embeddings) + conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) + +query = 'puppy' +query_embeddings = binary_quantize(checkpoint.queryFromText([query])[0]) +result = conn.execute('SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() +for row in result: + print(row) From 267d7960156b6866c300229a10b79b89d670ea39 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 15:41:25 -0800 Subject: [PATCH 423/879] Added ColPali / ColQwen2 example [skip ci] --- README.md | 1 + examples/colpali/exact.py | 52 +++++++++++++++++++++++++++++++ examples/colpali/requirements.txt | 4 +++ 3 files changed, 57 insertions(+) create mode 100644 examples/colpali/exact.py create mode 100644 examples/colpali/requirements.txt diff --git a/README.md b/README.md index 0dedce9..224fe57 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ Or check out some examples: - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py) with SentenceTransformers (cross-encoder) - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search/example.py) with Transformers - [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert/exact.py) with ColBERT +- [Document retrieval](https://github.com/pgvector/pgvector-python/blob/master/examples/colpali/exact.py) with ColPali - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/image_search/example.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py) with perceptual hashing - [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/rdkit/example.py) with RDKit diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py new file mode 100644 index 0000000..408bc7f --- /dev/null +++ b/examples/colpali/exact.py @@ -0,0 +1,52 @@ +from colpali_engine.models import ColQwen2, ColQwen2Processor +from datasets import load_dataset +from pgvector.psycopg import register_vector +import psycopg +import torch + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, embeddings vector(128)[])') +conn.execute(""" +CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ + WITH queries AS ( + SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) + ), + documents AS ( + SELECT unnest(document) AS document + ), + similarities AS ( + SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents + ), + max_similarities AS ( + SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number + ) + SELECT SUM(max_similarity) FROM max_similarities +$$ LANGUAGE SQL +""") + + +device = 'mps' if torch.backends.mps.is_available() else 'cpu' +model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval() +processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0') + + +def generate_embeddings(processed): + with torch.no_grad(): + return model(**processed.to(model.device)).to(device='cpu', dtype=torch.float32) + + +input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image'] +for content in input: + embeddings = [e.numpy() for e in generate_embeddings(processor.process_images([content]))[0]] + conn.execute('INSERT INTO documents (embeddings) VALUES (%s)', (embeddings,)) + +query = 'dividend' +query_embeddings = [e.numpy() for e in generate_embeddings(processor.process_queries([query]))[0]] +result = conn.execute('SELECT id, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() +for row in result: + print(row) diff --git a/examples/colpali/requirements.txt b/examples/colpali/requirements.txt new file mode 100644 index 0000000..4cf770d --- /dev/null +++ b/examples/colpali/requirements.txt @@ -0,0 +1,4 @@ +colpali-engine +datasets +pgvector +psycopg[binary] From 7d8a4173d988b5e9debaba7b4d6320d61879e76c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 15:47:10 -0800 Subject: [PATCH 424/879] Updated ColPali example to use binary quantization [skip ci] --- examples/colbert/exact_binary.py | 53 -------------------------------- examples/colpali/exact.py | 16 ++++++---- 2 files changed, 10 insertions(+), 59 deletions(-) delete mode 100644 examples/colbert/exact_binary.py diff --git a/examples/colbert/exact_binary.py b/examples/colbert/exact_binary.py deleted file mode 100644 index 8d398e2..0000000 --- a/examples/colbert/exact_binary.py +++ /dev/null @@ -1,53 +0,0 @@ -from colbert.infra import ColBERTConfig -from colbert.modeling.checkpoint import Checkpoint -from pgvector.psycopg import register_vector, Bit -import psycopg - -conn = psycopg.connect(dbname='pgvector_example', autocommit=True) - -conn.execute('CREATE EXTENSION IF NOT EXISTS vector') -register_vector(conn) - -conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embeddings bit(128)[])') -conn.execute(""" -CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) RETURNS double precision AS $$ - WITH queries AS ( - SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) - ), - documents AS ( - SELECT unnest(document) AS document - ), - similarities AS ( - SELECT query_number, 1 - ((document <~> query) / bit_length(query)) AS similarity FROM queries CROSS JOIN documents - ), - max_similarities AS ( - SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number - ) - SELECT SUM(max_similarity) FROM max_similarities -$$ LANGUAGE SQL -""") - - -def binary_quantize(embeddings): - return [Bit(e.numpy()) for e in (embeddings > 0)] - - -config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) -checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) - -input = [ - 'The dog is barking', - 'The cat is purring', - 'The bear is growling' -] -doc_embeddings = checkpoint.docFromText(input, keep_dims=False) -for content, embeddings in zip(input, doc_embeddings): - embeddings = binary_quantize(embeddings) - conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) - -query = 'puppy' -query_embeddings = binary_quantize(checkpoint.queryFromText([query])[0]) -result = conn.execute('SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() -for row in result: - print(row) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index 408bc7f..9fffc5f 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -1,6 +1,6 @@ from colpali_engine.models import ColQwen2, ColQwen2Processor from datasets import load_dataset -from pgvector.psycopg import register_vector +from pgvector.psycopg import register_vector, Bit import psycopg import torch @@ -10,9 +10,9 @@ register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, embeddings vector(128)[])') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, embeddings bit(128)[])') conn.execute(""" -CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ +CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) RETURNS double precision AS $$ WITH queries AS ( SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) ), @@ -20,7 +20,7 @@ SELECT unnest(document) AS document ), similarities AS ( - SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents + SELECT query_number, 1 - ((document <~> query) / bit_length(query)) AS similarity FROM queries CROSS JOIN documents ), max_similarities AS ( SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number @@ -40,13 +40,17 @@ def generate_embeddings(processed): return model(**processed.to(model.device)).to(device='cpu', dtype=torch.float32) +def binary_quantize(embedding): + return Bit(embedding > 0) + + input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image'] for content in input: - embeddings = [e.numpy() for e in generate_embeddings(processor.process_images([content]))[0]] + embeddings = [binary_quantize(e.numpy()) for e in generate_embeddings(processor.process_images([content]))[0]] conn.execute('INSERT INTO documents (embeddings) VALUES (%s)', (embeddings,)) query = 'dividend' -query_embeddings = [e.numpy() for e in generate_embeddings(processor.process_queries([query]))[0]] +query_embeddings = [binary_quantize(e.numpy()) for e in generate_embeddings(processor.process_queries([query]))[0]] result = conn.execute('SELECT id, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() for row in result: print(row) From d73a412de5fcb6d225b4d90865f0c4e514a142d3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 15:54:52 -0800 Subject: [PATCH 425/879] Updated ColPali example to use get_torch_device [skip ci] --- examples/colpali/exact.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index 9fffc5f..6eac7a4 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -1,4 +1,5 @@ from colpali_engine.models import ColQwen2, ColQwen2Processor +from colpali_engine.utils.torch_utils import get_torch_device from datasets import load_dataset from pgvector.psycopg import register_vector, Bit import psycopg @@ -30,7 +31,7 @@ """) -device = 'mps' if torch.backends.mps.is_available() else 'cpu' +device = get_torch_device('auto') model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval() processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0') From 7b6a46a014144f05ba174a53510ed69fd113b100 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 15:55:26 -0800 Subject: [PATCH 426/879] Removed extra line [skip ci] --- examples/colpali/exact.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index 6eac7a4..06d1828 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -30,7 +30,6 @@ $$ LANGUAGE SQL """) - device = get_torch_device('auto') model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval() processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0') From 4998aa1daefe95eb7550bc92ed875d6193d73b57 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 16:19:48 -0800 Subject: [PATCH 427/879] Updated readme [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 224fe57..260d389 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Or check out some examples: - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py) with SentenceTransformers (cross-encoder) - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search/example.py) with Transformers - [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert/exact.py) with ColBERT -- [Document retrieval](https://github.com/pgvector/pgvector-python/blob/master/examples/colpali/exact.py) with ColPali +- [Visual document retrieval](https://github.com/pgvector/pgvector-python/blob/master/examples/colpali/exact.py) with ColPali - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/image_search/example.py) with PyTorch - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py) with perceptual hashing - [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/rdkit/example.py) with RDKit From af7b04f271612c3f663e9a508f9c44564272e3a8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 16:40:44 -0800 Subject: [PATCH 428/879] Updated example [skip ci] --- examples/colpali/exact.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index 06d1828..c6f1467 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -37,20 +37,20 @@ def generate_embeddings(processed): with torch.no_grad(): - return model(**processed.to(model.device)).to(device='cpu', dtype=torch.float32) + return model(**processed.to(model.device)) def binary_quantize(embedding): - return Bit(embedding > 0) + return Bit(embedding.gt(0).numpy(force=True)) input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image'] for content in input: - embeddings = [binary_quantize(e.numpy()) for e in generate_embeddings(processor.process_images([content]))[0]] + embeddings = [binary_quantize(e) for e in generate_embeddings(processor.process_images([content]))[0]] conn.execute('INSERT INTO documents (embeddings) VALUES (%s)', (embeddings,)) query = 'dividend' -query_embeddings = [binary_quantize(e.numpy()) for e in generate_embeddings(processor.process_queries([query]))[0]] +query_embeddings = [binary_quantize(e) for e in generate_embeddings(processor.process_queries([query]))[0]] result = conn.execute('SELECT id, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() for row in result: print(row) From ca637bff37674592f08b9f65c75249b0d709746e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 16 Nov 2024 16:44:27 -0800 Subject: [PATCH 429/879] Updated example [skip ci] --- examples/colpali/exact.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index c6f1467..80bb603 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -37,11 +37,11 @@ def generate_embeddings(processed): with torch.no_grad(): - return model(**processed.to(model.device)) + return model(**processed.to(model.device)).to(torch.float32).numpy(force=True) def binary_quantize(embedding): - return Bit(embedding.gt(0).numpy(force=True)) + return Bit(embedding > 0) input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image'] From 5c35a5399aa3101f35e09e941a4e7cce0218e1ef Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 30 Nov 2024 03:30:34 -0800 Subject: [PATCH 430/879] Added test for binary quantization with SQLAlchemy - #98 [skip ci] --- tests/test_sqlalchemy.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 77c03fc..0380c89 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -65,6 +65,14 @@ class Item(Base): ) half_precision_index.create(engine) +binary_quantize_index = Index( + 'sqlalchemy_orm_binary_quantize_index', + func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'bit_hamming_ops'} +) +binary_quantize_index.create(engine) def create_items(): session = Session(engine) @@ -469,6 +477,18 @@ def test_half_precision(self): items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] + def test_binary_quantize(self): + session = Session(engine) + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + + with Session(engine) as session: + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + items = session.query(Item).order_by(distance).all() + assert [v.id for v in items] == [2, 3, 1] + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_async(self): From 230fe853d58105df1951fbbbc1730469b341f056 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 19 Dec 2024 11:59:52 -0800 Subject: [PATCH 431/879] Fixed spacing [skip ci] --- tests/test_sqlalchemy.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0380c89..9ab706a 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -74,6 +74,7 @@ class Item(Base): ) binary_quantize_index.create(engine) + def create_items(): session = Session(engine) session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) From 057806e44f6943230699d99c742621baeb9023c3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 19 Dec 2024 12:00:52 -0800 Subject: [PATCH 432/879] Added test for bit type with SQLAlchemy and asyncpg - #110 --- tests/test_sqlalchemy.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 9ab706a..0b53252 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -526,3 +526,20 @@ def connect(dbapi_connection, connection_record): assert item.embeddings[1].tolist() == [4, 5, 6] await engine.dispose() + + @pytest.mark.asyncio + @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') + async def test_asyncpg_bit(self): + import asyncpg + + engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + embedding = asyncpg.BitString('101') + session.add(Item(id=1, binary_embedding=embedding)) + item = await session.get(Item, 1) + assert item.binary_embedding == embedding + + await engine.dispose() From 57b6a61149c1f009ae55cccc17a9b5900e335f72 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 19 Dec 2024 12:06:02 -0800 Subject: [PATCH 433/879] Improved asyncpg test [skip ci] --- tests/test_asyncpg.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 7a68a9e..48d1e32 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -59,10 +59,11 @@ async def test_bit(self): await register_vector(conn) - embedding = asyncpg.BitString.from_int(5, length=3) + embedding = asyncpg.BitString('101') await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + assert res[0]['embedding'].as_string() == '101' assert res[0]['embedding'].to_int() == 5 assert res[1]['embedding'] is None From 47ad76d88f72cf07ffa238e4ad2714b672346149 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 19:35:43 -0800 Subject: [PATCH 434/879] Improved SQLModel example --- README.md | 3 +-- tests/test_sqlmodel.py | 8 ++++---- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 260d389..ca399ea 100644 --- a/README.md +++ b/README.md @@ -328,10 +328,9 @@ Add a vector column ```python from pgvector.sqlalchemy import Vector -from sqlalchemy import Column class Item(SQLModel, table=True): - embedding: Any = Field(sa_column=Column(Vector(3))) + embedding: Any = Field(sa_type=Vector(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 4cb0e9b..8a1c86c 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -15,10 +15,10 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[Any] = Field(default=None, sa_column=Column(VECTOR(3))) - half_embedding: Optional[Any] = Field(default=None, sa_column=Column(HALFVEC(3))) - binary_embedding: Optional[Any] = Field(default=None, sa_column=Column(BIT(3))) - sparse_embedding: Optional[Any] = Field(default=None, sa_column=Column(SPARSEVEC(3))) + embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) + half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) + binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) + sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) SQLModel.metadata.drop_all(engine) From b3e8908d3b3d74eba016b1cdc5bc7b1df1ad92bf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 19:36:07 -0800 Subject: [PATCH 435/879] Removed unneeded code [skip ci] --- tests/test_django.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 353087e..2c53d82 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -86,9 +86,6 @@ class Meta: class Migration(migrations.Migration): initial = True - dependencies = [ - ] - operations = [ VectorExtension(), migrations.CreateModel( From edd9b4ba02160ef429c4e44455eb0bfe6c781092 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 19:37:05 -0800 Subject: [PATCH 436/879] Test with Python 3.13 on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 04f1c21..562ba94 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.12, 3.8] + python: [3.13, 3.8] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 From 1b9df46f9542f3262f6c93a1a858c1414a0ffdc5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 19:38:38 -0800 Subject: [PATCH 437/879] Improved test code [skip ci] --- tests/test_sqlmodel.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 8a1c86c..851afd8 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,9 +1,8 @@ import numpy as np from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum import pytest -from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError -from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text +from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text from typing import Any, Optional engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') From a3c611f3f141a00c42b311f387278bb4f3ee4bcf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 19:42:16 -0800 Subject: [PATCH 438/879] Updated examples [skip ci] --- README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index ca399ea..d7a7e6c 100644 --- a/README.md +++ b/README.md @@ -175,10 +175,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(Base): - embedding = mapped_column(Vector(3)) + embedding = mapped_column(VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -274,11 +274,11 @@ session.scalars(select(Item).order_by(order).limit(5)) Add an array column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(Vector(3))) + embeddings = mapped_column(ARRAY(VECTOR(3))) ``` And register the types with the underlying driver @@ -327,10 +327,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(SQLModel, table=True): - embedding: Any = Field(sa_type=Vector(3)) + embedding: Any = Field(sa_type=VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` From 2ba2a855164f6f0947f17b94201a46d5ad615e6c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 20:00:48 -0800 Subject: [PATCH 439/879] Improved example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d7a7e6c..794cf91 100644 --- a/README.md +++ b/README.md @@ -376,7 +376,7 @@ Also supports `sum` Add an approximate index ```python -from sqlalchemy import Index +from sqlmodel import Index index = Index( 'my_index', From c6d2ddd429c10316ef329dd07ab86fe192bc71a0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 22:09:23 -0800 Subject: [PATCH 440/879] Improved sparsevec tests [skip ci] --- pgvector/utils/sparsevec.py | 1 + tests/test_psycopg.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index fd9ccff..a370c5e 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -108,6 +108,7 @@ def from_binary(cls, value): dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) + # TODO convert indices and values to lists in 0.4.0 return cls._from_parts(int(dim), indices, values) @classmethod diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 5802b2b..6d4f34a 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -110,12 +110,19 @@ def test_sparsevec(self): def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] + assert res.dimensions() == 6 + # TODO convert indices and values to lists in 0.4.0 + assert res.indices() == (0, 2, 4) + assert res.values() == (1.5, 2, 3) assert res.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] + assert res.dimensions() == 6 + assert res.indices() == [0, 2, 4] + assert res.values() == [1.5, 2, 3] assert res.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) From 9d9f45b800f3731e213f7b06bf3374e177ad86d5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 22:11:22 -0800 Subject: [PATCH 441/879] Added todo [skip ci] --- pgvector/psycopg2/register.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 7752852..08a69a9 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -5,6 +5,7 @@ from .vector import register_vector_info +# TODO remove default value for conn_or_curs in 0.4.0 # TODO make globally False by default in 0.4.0 # note: register_adapter is always global # TODO make arrays True by defalt in 0.4.0 From 972b6739788f5a09ec270bed552182a052e994c5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 22:18:26 -0800 Subject: [PATCH 442/879] Updated license year [skip ci] --- LICENSE.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/LICENSE.txt b/LICENSE.txt index d205f4e..b612d6d 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2024 Andrew Kane +Copyright (c) 2021-2025 Andrew Kane Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal From 57b7d3ba12781871045a378221d90bc972a3d5c1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 23:26:02 -0800 Subject: [PATCH 443/879] Added test for vector type with SQLAlchemy and asyncpg - #114 --- tests/test_sqlalchemy.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0b53252..6fc0adf 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -527,6 +527,29 @@ def connect(dbapi_connection, connection_record): await engine.dispose() + @pytest.mark.asyncio + @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') + async def test_asyncpg_vector(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 = np.array([1, 2, 3]) + session.add(Item(id=1, embedding=embedding)) + item = await session.get(Item, 1) + assert np.array_equal(item.embedding, embedding) + + await engine.dispose() + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_asyncpg_bit(self): From bf9a0a469983eabb1b1b38c6ba2495e3c4c2b8ce Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 23:29:37 -0800 Subject: [PATCH 444/879] Added tests for halfvec and sparsevec types with SQLAlchemy and asyncpg [skip ci] --- tests/test_sqlalchemy.py | 46 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 6fc0adf..40068e9 100644 --- a/tests/test_sqlalchemy.py +++ b/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() From 257eb3b92c9f02e2ca266a15c6c8b93ebc94082a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 23:30:55 -0800 Subject: [PATCH 445/879] Simplified tests [skip ci] --- tests/test_sqlalchemy.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 40068e9..519a388 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -530,8 +530,6 @@ def connect(dbapi_connection, connection_record): @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_asyncpg_vector(self): - import asyncpg - engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) @@ -553,8 +551,6 @@ async def test_asyncpg_vector(self): @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) @@ -593,8 +589,6 @@ async def test_asyncpg_bit(self): @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) From 91f5d34c11f0064c83ca08b7e69055ce6ef03124 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 23:32:53 -0800 Subject: [PATCH 446/879] Added test for vector[] type with SQLAlchemy and asyncpg [skip ci] --- tests/test_sqlalchemy.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 519a388..7e8b888 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -606,3 +606,26 @@ async def test_asyncpg_sparsevec(self): assert item.sparse_embedding.to_list() == embedding await engine.dispose() + + @pytest.mark.asyncio + @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') + async def test_asyncpg_vector_array(self): + 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(): + session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + + # this fails if the driver does not cast arrays + item = await session.get(Item, 1) + assert item.embeddings[0].tolist() == [1, 2, 3] + assert item.embeddings[1].tolist() == [4, 5, 6] + + await engine.dispose() From f7eeb3a04554b9adf82a5073d08fc757c41604a3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 23:36:40 -0800 Subject: [PATCH 447/879] Improved test code [skip ci] --- tests/test_sqlalchemy.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 7e8b888..f3d045f 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -492,7 +492,7 @@ def test_binary_quantize(self): @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_async(self): + async def test_async_avg(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) @@ -622,8 +622,6 @@ async def test_asyncpg_vector_array(self): async with async_session() as session: async with session.begin(): session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - - # this fails if the driver does not cast arrays item = await session.get(Item, 1) assert item.embeddings[0].tolist() == [1, 2, 3] assert item.embeddings[1].tolist() == [4, 5, 6] From 2d2563d702ee319a33d17b27549bce035a6c7348 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 8 Feb 2025 23:39:55 -0800 Subject: [PATCH 448/879] Improved test names [skip ci] --- tests/test_sqlalchemy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f3d045f..fd46e74 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -492,7 +492,7 @@ def test_binary_quantize(self): @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_async_avg(self): + async def test_psycopg_async_avg(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) @@ -507,7 +507,7 @@ async def test_async_avg(self): @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_async_vector_array(self): + async def test_psycopg_async_vector_array(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) From d828239fb466e11a8fb02c7e35a052dbbce3e5b8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 01:00:50 -0800 Subject: [PATCH 449/879] Revert "Updated examples [skip ci]" This reverts commit a3c611f3f141a00c42b311f387278bb4f3ee4bcf. --- README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 794cf91..5a59c9d 100644 --- a/README.md +++ b/README.md @@ -175,10 +175,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import VECTOR +from pgvector.sqlalchemy import Vector class Item(Base): - embedding = mapped_column(VECTOR(3)) + embedding = mapped_column(Vector(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -274,11 +274,11 @@ session.scalars(select(Item).order_by(order).limit(5)) Add an array column ```python -from pgvector.sqlalchemy import VECTOR +from pgvector.sqlalchemy import Vector from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(VECTOR(3))) + embeddings = mapped_column(ARRAY(Vector(3))) ``` And register the types with the underlying driver @@ -327,10 +327,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import VECTOR +from pgvector.sqlalchemy import Vector class Item(SQLModel, table=True): - embedding: Any = Field(sa_type=VECTOR(3)) + embedding: Any = Field(sa_type=Vector(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` From 8a7040d2ee79ac8fc6313538ffbc38ebad3ac197 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 01:34:02 -0800 Subject: [PATCH 450/879] Removed unneeded code [skip ci] --- examples/citus/example.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/citus/example.py b/examples/citus/example.py index d448204..915c25f 100644 --- a/examples/citus/example.py +++ b/examples/citus/example.py @@ -40,9 +40,6 @@ for i in range(rows): copy.write_row([embeddings[i], categories[i]]) - while conn.pgconn.flush() == 1: - pass - print('Creating index in parallel') conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') From 00cd08e6c44077b99f378edbd007b2483ff406f7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 11:09:23 -0800 Subject: [PATCH 451/879] Improved tests --- tests/test_sqlalchemy.py | 139 ++++++++++++++++++++------------------- 1 file changed, 71 insertions(+), 68 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index fd46e74..405cd21 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -76,11 +76,11 @@ class Item(Base): def create_items(): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) - session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) - session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) - session.commit() + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) + session.commit() class TestSqlalchemy: @@ -129,11 +129,11 @@ def test_orm(self): item2 = Item(embedding=[4, 5, 6]) item3 = Item() - session = Session(engine) - session.add(item) - session.add(item2) - session.add(item3) - session.commit() + with Session(engine) as session: + session.add(item) + session.add(item2) + session.add(item3) + session.commit() stmt = select(Item) with Session(engine) as session: @@ -148,11 +148,11 @@ def test_orm(self): assert items[2].embedding is None def test_vector(self): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.embedding.tolist() == [1, 2, 3] def test_vector_l2_distance(self): create_items() @@ -203,11 +203,11 @@ def test_vector_l1_distance_orm(self): assert [v.id for v in items] == [1, 3, 2] def test_halfvec(self): - session = Session(engine) - session.add(Item(id=1, half_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding.to_list() == [1, 2, 3] def test_halfvec_l2_distance(self): create_items() @@ -258,11 +258,11 @@ def test_halfvec_l1_distance_orm(self): assert [v.id for v in items] == [1, 3, 2] def test_bit(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' + with Session(engine) as session: + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' def test_bit_hamming_distance(self): create_items() @@ -289,11 +289,11 @@ def test_bit_jaccard_distance_orm(self): assert [v.id for v in items] == [2, 3, 1] def test_sparsevec(self): - session = Session(engine) - session.add(Item(id=1, sparse_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() @@ -405,24 +405,24 @@ def test_sum_orm(self): def test_bad_dimensions(self): item = Item(embedding=[1, 2]) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + session.commit() def test_bad_ndim(self): item = Item(embedding=np.array([[1, 2, 3]])) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected ndim to be 1'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='expected ndim to be 1'): + session.commit() def test_bad_dtype(self): item = Item(embedding=np.array(['one', 'two', 'three'])) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='could not convert string to float'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() def test_inspect(self): columns = inspect(engine).get_columns('sqlalchemy_orm_item') @@ -433,14 +433,17 @@ def test_literal_binds(self): assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) def test_insert(self): - session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) + with Session(engine) as session: + session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) def test_insert_bulk(self): - session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) + with Session(engine) as session: + session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) # register_vector in psycopg2 tests change this behavior # def test_insert_text(self): - # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) + # with Session(engine) as session: + # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) def test_automap(self): metadata = MetaData() @@ -448,29 +451,30 @@ def test_automap(self): AutoBase = automap_base(metadata=metadata) AutoBase.prepare() AutoItem = AutoBase.classes.sqlalchemy_orm_item - session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) - item = session.query(AutoItem).first() - assert item.embedding.tolist() == [1, 2, 3] + with Session(engine) as session: + session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) + item = session.query(AutoItem).first() + assert item.embedding.tolist() == [1, 2, 3] def test_vector_array(self): - session = Session(array_engine) - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - session.commit() + with Session(array_engine) as session: + session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.commit() - # this fails if the driver does not cast arrays - item = session.get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + # this fails if the driver does not cast arrays + item = session.get(Item, 1) + assert item.embeddings[0].tolist() == [1, 2, 3] + assert item.embeddings[1].tolist() == [4, 5, 6] def test_halfvec_array(self): - session = Session(array_engine) - session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - session.commit() + with Session(array_engine) as session: + session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.commit() - # this fails if the driver does not cast arrays - item = session.get(Item, 1) - assert item.half_embeddings[0].to_list() == [1, 2, 3] - assert item.half_embeddings[1].to_list() == [4, 5, 6] + # this fails if the driver does not cast arrays + item = session.get(Item, 1) + assert item.half_embeddings[0].to_list() == [1, 2, 3] + assert item.half_embeddings[1].to_list() == [4, 5, 6] def test_half_precision(self): create_items() @@ -479,13 +483,12 @@ def test_half_precision(self): assert [v.id for v in items] == [1, 3, 2] def test_binary_quantize(self): - session = Session(engine) - session.add(Item(id=1, embedding=[-1, -2, -3])) - session.add(Item(id=2, embedding=[1, -2, 3])) - session.add(Item(id=3, embedding=[1, 2, 3])) - session.commit() - with Session(engine) as session: + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] From 7837e92d72eef265e075e8ea5aa305e159e41437 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 11:15:18 -0800 Subject: [PATCH 452/879] Added more tests for SQLAlchemy --- tests/test_sqlalchemy.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 405cd21..79b3c50 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -29,6 +29,8 @@ def connect(dbapi_connection, connection_record): register_vector(dbapi_connection, globally=False, arrays=True) +psycopg3_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + Base = declarative_base() @@ -493,6 +495,34 @@ def test_binary_quantize(self): items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] + def test_psycopg_vector(self): + with Session(psycopg3_engine) as session: + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.embedding.tolist() == [1, 2, 3] + + def test_psycopg_halfvec(self): + with Session(psycopg3_engine) as session: + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding.to_list() == [1, 2, 3] + + def test_psycopg_bit(self): + with Session(psycopg3_engine) as session: + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' + + def test_psycopg_sparsevec(self): + with Session(psycopg3_engine) as session: + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding.to_list() == [1, 2, 3] + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_psycopg_async_avg(self): From f08cec7a0522b19942a02df14f3f396f0773c912 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 11:46:42 -0800 Subject: [PATCH 453/879] Parameterize SQLAlchemy tests --- tests/test_sqlalchemy.py | 181 ++++++++++++++++++--------------------- 1 file changed, 84 insertions(+), 97 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 79b3c50..a4ac860 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -15,8 +15,15 @@ mapped_column = Column sqlalchemy_version = 1 -engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -with Session(engine) as session: +psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +engines = [psycopg2_engine] + +if sqlalchemy_version > 1: + psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + engines.append(psycopg_engine) + +setup_engine = engines[0] +with Session(setup_engine) as session: session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() @@ -29,8 +36,6 @@ def connect(dbapi_connection, connection_record): register_vector(dbapi_connection, globally=False, arrays=True) -psycopg3_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') - Base = declarative_base() @@ -46,8 +51,8 @@ class Item(Base): half_embeddings = mapped_column(ARRAY(HALFVEC(3))) -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) +Base.metadata.drop_all(setup_engine) +Base.metadata.create_all(setup_engine) index = Index( 'sqlalchemy_orm_index', @@ -56,7 +61,7 @@ class Item(Base): postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} ) -index.create(engine) +index.create(setup_engine) half_precision_index = Index( 'sqlalchemy_orm_half_precision_index', @@ -65,7 +70,7 @@ class Item(Base): postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'halfvec_l2_ops'} ) -half_precision_index.create(engine) +half_precision_index.create(setup_engine) binary_quantize_index = Index( 'sqlalchemy_orm_binary_quantize_index', @@ -74,24 +79,29 @@ class Item(Base): postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'bit_hamming_ops'} ) -binary_quantize_index.create(engine) +binary_quantize_index.create(setup_engine) def create_items(): - with Session(engine) as session: + with Session(setup_engine) as session: session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) session.commit() +def delete_items(): + with Session(setup_engine) as session: + session.query(Item).delete() + session.commit() + + +@pytest.mark.parametrize("engine", engines) class TestSqlalchemy: - def setup_method(self, test_method): - with Session(engine) as session: - session.query(Item).delete() - session.commit() + def setup_method(self): + delete_items() - def test_core(self): + def test_core(self, engine): metadata = MetaData() item_table = Table( @@ -126,7 +136,7 @@ def test_core(self): ) hnsw_index.create(engine) - def test_orm(self): + def test_orm(self, engine): item = Item(embedding=np.array([1.5, 2, 3])) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -140,236 +150,236 @@ def test_orm(self): stmt = select(Item) with Session(engine) as session: items = [v[0] for v in session.execute(stmt).all()] - assert items[0].id == 1 - assert items[1].id == 2 - assert items[2].id == 3 + assert items[0].id in [1, 4] + assert items[1].id in [2, 5] + assert items[2].id in [3, 6] assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) assert items[0].embedding.dtype == np.float32 assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) assert items[1].embedding.dtype == np.float32 assert items[2].embedding is None - def test_vector(self): + def test_vector(self, engine): with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) assert item.embedding.tolist() == [1, 2, 3] - def test_vector_l2_distance(self): + def test_vector_l2_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l2_distance_orm(self): + def test_vector_l2_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_vector_max_inner_product_orm(self): + def test_vector_max_inner_product_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_vector_cosine_distance_orm(self): + def test_vector_cosine_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l1_distance_orm(self): + def test_vector_l1_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self): + def test_halfvec(self, engine): with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) assert item.half_embedding.to_list() == [1, 2, 3] - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l2_distance_orm(self): + def test_halfvec_l2_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_max_inner_product_orm(self): + def test_halfvec_max_inner_product_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_cosine_distance_orm(self): + def test_halfvec_cosine_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l1_distance_orm(self): + def test_halfvec_l1_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self): + def test_bit(self, engine): with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() item = session.get(Item, 1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_hamming_distance_orm(self): + def test_bit_hamming_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance_orm(self): + def test_bit_jaccard_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self): + def test_sparsevec(self, engine): with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) assert item.sparse_embedding.to_list() == [1, 2, 3] - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l2_distance_orm(self): + def test_sparsevec_l2_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_max_inner_product_orm(self): + def test_sparsevec_max_inner_product_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_cosine_distance_orm(self): + def test_sparsevec_cosine_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l1_distance_orm(self): + def test_sparsevec_l1_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_filter(self): + def test_filter(self, engine): create_items() with Session(engine) as session: items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() assert [v.id for v in items] == [1] - def test_filter_orm(self): + def test_filter_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] - def test_select(self): + def test_select(self, engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() assert items[0] == 3 - def test_select_orm(self): + def test_select_orm(self, engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() assert items[0] == 3 - def test_avg(self): + def test_avg(self, engine): with Session(engine) as session: res = session.query(avg(Item.embedding)).first()[0] assert res is None @@ -378,7 +388,7 @@ def test_avg(self): 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): + def test_avg_orm(self, engine): with Session(engine) as session: res = session.scalars(select(avg(Item.embedding))).first() assert res is None @@ -387,7 +397,7 @@ def test_avg_orm(self): 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): + def test_sum(self, engine): with Session(engine) as session: res = session.query(sum(Item.embedding)).first()[0] assert res is None @@ -396,7 +406,7 @@ def test_sum(self): res = session.query(sum(Item.embedding)).first()[0] assert np.array_equal(res, np.array([5, 7, 9])) - def test_sum_orm(self): + def test_sum_orm(self, engine): with Session(engine) as session: res = session.scalars(select(sum(Item.embedding))).first() assert res is None @@ -405,40 +415,40 @@ def test_sum_orm(self): res = session.scalars(select(sum(Item.embedding))).first() assert np.array_equal(res, np.array([5, 7, 9])) - def test_bad_dimensions(self): + def test_bad_dimensions(self, engine): item = Item(embedding=[1, 2]) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() - def test_bad_ndim(self): + def test_bad_ndim(self, engine): item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() - def test_bad_dtype(self): + def test_bad_dtype(self, engine): item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='could not convert string to float'): session.commit() - def test_inspect(self): + def test_inspect(self, engine): columns = inspect(engine).get_columns('sqlalchemy_orm_item') assert isinstance(columns[1]['type'], VECTOR) - def test_literal_binds(self): + def test_literal_binds(self, engine): sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) - def test_insert(self): + def test_insert(self, engine): with Session(engine) as session: session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) - def test_insert_bulk(self): + def test_insert_bulk(self, engine): with Session(engine) as session: session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) @@ -447,7 +457,7 @@ def test_insert_bulk(self): # with Session(engine) as session: # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) - def test_automap(self): + def test_automap(self, engine): metadata = MetaData() metadata.reflect(engine, only=['sqlalchemy_orm_item']) AutoBase = automap_base(metadata=metadata) @@ -458,7 +468,7 @@ def test_automap(self): item = session.query(AutoItem).first() assert item.embedding.tolist() == [1, 2, 3] - def test_vector_array(self): + def test_vector_array(self, engine): with Session(array_engine) as session: session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() @@ -468,7 +478,7 @@ def test_vector_array(self): assert item.embeddings[0].tolist() == [1, 2, 3] assert item.embeddings[1].tolist() == [4, 5, 6] - def test_halfvec_array(self): + def test_halfvec_array(self, engine): with Session(array_engine) as session: session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() @@ -478,13 +488,13 @@ def test_halfvec_array(self): assert item.half_embeddings[0].to_list() == [1, 2, 3] assert item.half_embeddings[1].to_list() == [4, 5, 6] - def test_half_precision(self): + def test_half_precision(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_binary_quantize(self): + def test_binary_quantize(self, engine): with Session(engine) as session: session.add(Item(id=1, embedding=[-1, -2, -3])) session.add(Item(id=2, embedding=[1, -2, 3])) @@ -495,33 +505,10 @@ def test_binary_quantize(self): items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] - def test_psycopg_vector(self): - with Session(psycopg3_engine) as session: - session.add(Item(id=1, embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] - - def test_psycopg_halfvec(self): - with Session(psycopg3_engine) as session: - session.add(Item(id=1, half_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] - def test_psycopg_bit(self): - with Session(psycopg3_engine) as session: - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' - - def test_psycopg_sparsevec(self): - with Session(psycopg3_engine) as session: - session.add(Item(id=1, sparse_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] +class TestSqlalchemyAsync: + def setup_method(self): + delete_items() @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') From 5e381602a739ca5307f02c75ee57d219555f5ada Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 11:54:57 -0800 Subject: [PATCH 454/879] Improved array tests --- tests/test_sqlalchemy.py | 49 +++++++++++++++++++++++----------------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index a4ac860..b1f3e85 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -27,10 +27,11 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() -array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +psycopg2_array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +array_engines = [psycopg2_array_engine] -@event.listens_for(array_engine, "connect") +@event.listens_for(psycopg2_array_engine, "connect") def connect(dbapi_connection, connection_record): from pgvector.psycopg2 import register_vector register_vector(dbapi_connection, globally=False, arrays=True) @@ -468,8 +469,31 @@ def test_automap(self, engine): item = session.query(AutoItem).first() assert item.embedding.tolist() == [1, 2, 3] + def test_half_precision(self, engine): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_binary_quantize(self, engine): + with Session(engine) as session: + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + items = session.query(Item).order_by(distance).all() + assert [v.id for v in items] == [2, 3, 1] + + +@pytest.mark.parametrize("engine", array_engines) +class TestSqlalchemyArray: + def setup_method(self): + delete_items() + def test_vector_array(self, engine): - with Session(array_engine) as session: + with Session(engine) as session: session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() @@ -479,7 +503,7 @@ def test_vector_array(self, engine): assert item.embeddings[1].tolist() == [4, 5, 6] def test_halfvec_array(self, engine): - with Session(array_engine) as session: + with Session(engine) as session: session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) session.commit() @@ -488,23 +512,6 @@ def test_halfvec_array(self, engine): assert item.half_embeddings[0].to_list() == [1, 2, 3] assert item.half_embeddings[1].to_list() == [4, 5, 6] - def test_half_precision(self, engine): - create_items() - with Session(engine) as session: - items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() - assert [v.id for v in items] == [1, 3, 2] - - def test_binary_quantize(self, engine): - with Session(engine) as session: - session.add(Item(id=1, embedding=[-1, -2, -3])) - session.add(Item(id=2, embedding=[1, -2, 3])) - session.add(Item(id=3, embedding=[1, 2, 3])) - session.commit() - - distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) - items = session.query(Item).order_by(distance).all() - assert [v.id for v in items] == [2, 3, 1] - class TestSqlalchemyAsync: def setup_method(self): From f82e44f231e498c86839735de1658ef7b8cb11a1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:01:45 -0800 Subject: [PATCH 455/879] Added tests for SQLAlchemy with pg8000 --- requirements.txt | 1 + tests/test_sqlalchemy.py | 16 ++++++++++++---- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/requirements.txt b/requirements.txt index 0e30959..a13be06 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ asyncpg Django numpy peewee +pg8000 psycopg[binary,pool] psycopg2-binary pytest diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index b1f3e85..37e803d 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,4 +1,5 @@ import numpy as np +import os from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY @@ -16,7 +17,8 @@ sqlalchemy_version = 1 psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -engines = [psycopg2_engine] +pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ['USER']}@localhost/pgvector_python_test') +engines = [psycopg2_engine, pg8000_engine] if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') @@ -151,9 +153,9 @@ def test_orm(self, engine): stmt = select(Item) with Session(engine) as session: items = [v[0] for v in session.execute(stmt).all()] - assert items[0].id in [1, 4] - assert items[1].id in [2, 5] - assert items[2].id in [3, 6] + assert items[0].id in [1, 4, 7] + assert items[1].id in [2, 5, 8] + assert items[2].id in [3, 6, 9] assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) assert items[0].embedding.dtype == np.float32 assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) @@ -290,12 +292,18 @@ def test_bit_hamming_distance_orm(self, engine): assert [v.id for v in items] == [2, 3, 1] def test_bit_jaccard_distance(self, engine): + if engine == pg8000_engine: + return + create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] def test_bit_jaccard_distance_orm(self, engine): + if engine == pg8000_engine: + return + create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) From cfcc2ea7b8b942c47c378bf47a4490c5acb50ec7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:02:44 -0800 Subject: [PATCH 456/879] Updated style [skip ci] --- tests/test_sqlalchemy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 37e803d..aa2ad97 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -99,7 +99,7 @@ def delete_items(): session.commit() -@pytest.mark.parametrize("engine", engines) +@pytest.mark.parametrize('engine', engines) class TestSqlalchemy: def setup_method(self): delete_items() @@ -495,7 +495,7 @@ def test_binary_quantize(self, engine): assert [v.id for v in items] == [2, 3, 1] -@pytest.mark.parametrize("engine", array_engines) +@pytest.mark.parametrize('engine', array_engines) class TestSqlalchemyArray: def setup_method(self): delete_items() From 95403d5268e11ab6efef969f46f086e3f57e2b52 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:05:34 -0800 Subject: [PATCH 457/879] Added tests for arrays with SQLAlchemy and Psycopg 3 --- tests/test_sqlalchemy.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index aa2ad97..f4a6bce 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -34,11 +34,21 @@ @event.listens_for(psycopg2_array_engine, "connect") -def connect(dbapi_connection, connection_record): +def psycopg2_connect(dbapi_connection, connection_record): from pgvector.psycopg2 import register_vector register_vector(dbapi_connection, globally=False, arrays=True) +if sqlalchemy_version > 1: + psycopg_array_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + array_engines.append(psycopg_array_engine) + + @event.listens_for(psycopg_array_engine, "connect") + def psycopg_connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector + register_vector(dbapi_connection) + + Base = declarative_base() From c74e090f26a02fc920ef910265ac0e4f2eb7cbde Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:11:46 -0800 Subject: [PATCH 458/879] Fixed CI --- tests/test_sqlalchemy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f4a6bce..7dbc565 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -17,7 +17,7 @@ sqlalchemy_version = 1 psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ['USER']}@localhost/pgvector_python_test') +pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') engines = [psycopg2_engine, pg8000_engine] if sqlalchemy_version > 1: From b350d6a8d45d02ea954fad945d194896c50fbc1e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:13:46 -0800 Subject: [PATCH 459/879] Simplified test code [skip ci] --- tests/test_django.py | 2 +- tests/test_peewee.py | 2 +- tests/test_psycopg.py | 2 +- tests/test_psycopg2.py | 2 +- tests/test_sqlmodel.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 2c53d82..ea15771 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -158,7 +158,7 @@ class Meta: class TestDjango: - def setup_method(self, test_method): + def setup_method(self): Item.objects.all().delete() def test_vector(self): diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 9666388..e98a0ec 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -36,7 +36,7 @@ def create_items(): class TestPeewee: - def setup_method(self, test_method): + def setup_method(self): Item.truncate_table() def test_vector(self): diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 6d4f34a..90f80b6 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -14,7 +14,7 @@ class TestPsycopg: - def setup_method(self, test_method): + def setup_method(self): conn.execute('DELETE FROM psycopg_items') def test_vector(self): diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index c3cd3cd..d661f12 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -16,7 +16,7 @@ class TestPsycopg2: - def setup_method(self, test_method): + def setup_method(self): cur.execute('DELETE FROM psycopg2_items') def test_vector(self): diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 851afd8..e0330d2 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -42,7 +42,7 @@ def create_items(): class TestSqlmodel: - def setup_method(self, test_method): + def setup_method(self): with Session(engine) as session: session.exec(delete(Item)) session.commit() From 651df0844b3c6790414ec2e8ed75330ad80406af Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:15:52 -0800 Subject: [PATCH 460/879] Improved SQLModel tests [skip ci] --- tests/test_sqlmodel.py | 68 +++++++++++++++++++++--------------------- 1 file changed, 34 insertions(+), 34 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index e0330d2..373834f 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -34,11 +34,11 @@ class Item(SQLModel, table=True): def create_items(): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) - session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) - session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) - session.commit() + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) + session.commit() class TestSqlmodel: @@ -52,11 +52,11 @@ def test_orm(self): item2 = Item(embedding=[4, 5, 6]) item3 = Item() - session = Session(engine) - session.add(item) - session.add(item2) - session.add(item3) - session.commit() + with Session(engine) as session: + session.add(item) + session.add(item2) + session.add(item3) + session.commit() stmt = select(Item) with Session(engine) as session: @@ -71,11 +71,11 @@ def test_orm(self): assert items[2].embedding is None def test_vector(self): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.embedding.tolist() == [1, 2, 3] def test_vector_l2_distance(self): create_items() @@ -102,11 +102,11 @@ def test_vector_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] def test_halfvec(self): - session = Session(engine) - session.add(Item(id=1, half_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding.to_list() == [1, 2, 3] def test_halfvec_l2_distance(self): create_items() @@ -133,11 +133,11 @@ def test_halfvec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] def test_bit(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' + with Session(engine) as session: + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' def test_bit_hamming_distance(self): create_items() @@ -152,11 +152,11 @@ def test_bit_jaccard_distance(self): assert [v.id for v in items] == [2, 3, 1] def test_sparsevec(self): - session = Session(engine) - session.add(Item(id=1, sparse_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding.to_list() == [1, 2, 3] def test_sparsevec_l2_distance(self): create_items() @@ -232,7 +232,7 @@ def test_halfvec_sum(self): def test_bad_dimensions(self): item = Item(embedding=[1, 2]) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + session.commit() From 2883156b461f08fe32be81439d8e653ac1c41c5a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:42:19 -0800 Subject: [PATCH 461/879] Improved tests for async SQLAlchemy engines [skip ci] --- tests/test_sqlalchemy.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 7dbc565..6e1d496 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -19,11 +19,18 @@ psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') engines = [psycopg2_engine, pg8000_engine] +async_engines = [] if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') engines.append(psycopg_engine) + psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async_engines.append(psycopg_async_engine) + + asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + async_engines.append(asyncpg_engine) + setup_engine = engines[0] with Session(setup_engine) as session: session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) @@ -531,14 +538,13 @@ def test_halfvec_array(self, engine): assert item.half_embeddings[1].to_list() == [4, 5, 6] +@pytest.mark.parametrize('engine', async_engines) class TestSqlalchemyAsync: def setup_method(self): delete_items() @pytest.mark.asyncio - @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_psycopg_async_avg(self): - engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async def test_psycopg_async_avg(self, engine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -550,6 +556,11 @@ async def test_psycopg_async_avg(self): await engine.dispose() + +class TestSqlalchemyAsync2: + def setup_method(self): + delete_items() + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_psycopg_async_vector_array(self): From 86331f0ee6650adcdb655b5d092f1c24d3b0fa84 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:45:22 -0800 Subject: [PATCH 462/879] Improved tests for async SQLAlchemy engines [skip ci] --- tests/test_sqlalchemy.py | 104 +++++++++++++++------------------------ 1 file changed, 39 insertions(+), 65 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 6e1d496..689615b 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -544,122 +544,96 @@ def setup_method(self): delete_items() @pytest.mark.asyncio - async def test_psycopg_async_avg(self, engine): + async def test_vector(self, engine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - session.add(Item(embedding=[1, 2, 3])) - session.add(Item(embedding=[4, 5, 6])) - avg = await session.scalars(select(func.avg(Item.embedding))) - assert avg.first() == '[2.5,3.5,4.5]' + embedding = np.array([1, 2, 3]) + session.add(Item(id=1, embedding=embedding)) + item = await session.get(Item, 1) + assert np.array_equal(item.embedding, embedding) await engine.dispose() - -class TestSqlalchemyAsync2: - def setup_method(self): - delete_items() - @pytest.mark.asyncio - @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_psycopg_async_vector_array(self): - engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async def test_halfvec(self, engine): async_session = async_sessionmaker(engine, expire_on_commit=False) - @event.listens_for(engine.sync_engine, "connect") - def connect(dbapi_connection, connection_record): - from pgvector.psycopg import register_vector_async - dbapi_connection.run_async(register_vector_async) - async with async_session() as session: async with session.begin(): - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - - # this fails if the driver does not cast arrays + embedding = [1, 2, 3] + session.add(Item(id=1, half_embedding=embedding)) item = await session.get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + 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_vector(self): - engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') - async_session = async_sessionmaker(engine, expire_on_commit=False) + async def test_bit(self, engine): + import asyncpg - # 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_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - embedding = np.array([1, 2, 3]) - session.add(Item(id=1, embedding=embedding)) + embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' + session.add(Item(id=1, binary_embedding=embedding)) item = await session.get(Item, 1) - assert np.array_equal(item.embedding, embedding) + 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_halfvec(self): - engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + async def test_sparsevec(self, engine): 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)) + session.add(Item(id=1, sparse_embedding=embedding)) item = await session.get(Item, 1) - assert item.half_embedding.to_list() == embedding + assert item.sparse_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): - import asyncpg - - engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + async def test_avg(self, engine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - embedding = asyncpg.BitString('101') - session.add(Item(id=1, binary_embedding=embedding)) - item = await session.get(Item, 1) - assert item.binary_embedding == embedding + session.add(Item(embedding=[1, 2, 3])) + session.add(Item(embedding=[4, 5, 6])) + avg = await session.scalars(select(func.avg(Item.embedding))) + assert avg.first() == '[2.5,3.5,4.5]' await engine.dispose() + +class TestSqlalchemyAsync2: + def setup_method(self): + delete_items() + @pytest.mark.asyncio @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_asyncpg_sparsevec(self): - engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + async def test_psycopg_async_vector_array(self): + engine = create_async_engine('postgresql+psycopg://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) + @event.listens_for(engine.sync_engine, "connect") + def connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector_async + dbapi_connection.run_async(register_vector_async) async with async_session() as session: async with session.begin(): - embedding = [1, 2, 3] - session.add(Item(id=1, sparse_embedding=embedding)) + session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + + # this fails if the driver does not cast arrays item = await session.get(Item, 1) - assert item.sparse_embedding.to_list() == embedding + assert item.embeddings[0].tolist() == [1, 2, 3] + assert item.embeddings[1].tolist() == [4, 5, 6] await engine.dispose() From 224c18a47c2a2d652fe1f7267449e61c290b249f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 12:49:14 -0800 Subject: [PATCH 463/879] Simplified test code [skip ci] --- tests/test_sqlalchemy.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 689615b..07e29e7 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,3 +1,4 @@ +import asyncpg import numpy as np import os from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum @@ -571,8 +572,6 @@ async def test_halfvec(self, engine): @pytest.mark.asyncio async def test_bit(self, engine): - import asyncpg - async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -611,13 +610,13 @@ async def test_avg(self, engine): await engine.dispose() -class TestSqlalchemyAsync2: +@pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') +class TestSqlalchemyAsyncArray: def setup_method(self): delete_items() @pytest.mark.asyncio - @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_psycopg_async_vector_array(self): + async def test_psycopg_vector_array(self): engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) @@ -638,7 +637,6 @@ def connect(dbapi_connection, connection_record): await engine.dispose() @pytest.mark.asyncio - @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') async def test_asyncpg_vector_array(self): engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') async_session = async_sessionmaker(engine, expire_on_commit=False) From a1d89971cd10ebe0dc11969ec532011eaa8a9a78 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:02:05 -0800 Subject: [PATCH 464/879] Improved test code [skip ci] --- tests/test_sqlalchemy.py | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 07e29e7..8868df4 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -19,7 +19,17 @@ psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') +psycopg2_array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') + + +@event.listens_for(psycopg2_array_engine, "connect") +def psycopg2_connect(dbapi_connection, connection_record): + from pgvector.psycopg2 import register_vector + register_vector(dbapi_connection, globally=False, arrays=True) + + engines = [psycopg2_engine, pg8000_engine] +array_engines = [psycopg2_array_engine] async_engines = [] if sqlalchemy_version > 1: @@ -32,22 +42,6 @@ asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') async_engines.append(asyncpg_engine) -setup_engine = engines[0] -with Session(setup_engine) as session: - session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - session.commit() - -psycopg2_array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -array_engines = [psycopg2_array_engine] - - -@event.listens_for(psycopg2_array_engine, "connect") -def psycopg2_connect(dbapi_connection, connection_record): - from pgvector.psycopg2 import register_vector - register_vector(dbapi_connection, globally=False, arrays=True) - - -if sqlalchemy_version > 1: psycopg_array_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') array_engines.append(psycopg_array_engine) @@ -56,6 +50,10 @@ def psycopg_connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector register_vector(dbapi_connection) +setup_engine = engines[0] +with Session(setup_engine) as session: + session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) + session.commit() Base = declarative_base() From c792451a76fea51352ad0a5f952c97eaeaea70d7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:12:57 -0800 Subject: [PATCH 465/879] Test more engine configurations --- tests/test_sqlalchemy.py | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 8868df4..dffa07b 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -19,37 +19,39 @@ psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') -psycopg2_array_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -@event.listens_for(psycopg2_array_engine, "connect") +@event.listens_for(psycopg2_type_engine, "connect") def psycopg2_connect(dbapi_connection, connection_record): from pgvector.psycopg2 import register_vector register_vector(dbapi_connection, globally=False, arrays=True) -engines = [psycopg2_engine, pg8000_engine] -array_engines = [psycopg2_array_engine] +engines = [psycopg2_engine, pg8000_engine, psycopg2_type_engine] +array_engines = [psycopg2_type_engine] async_engines = [] if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') engines.append(psycopg_engine) + psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + + @event.listens_for(psycopg_type_engine, "connect") + def psycopg_connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector + register_vector(dbapi_connection) + + engines.append(psycopg_type_engine) + array_engines.append(psycopg_type_engine) + psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_engines.append(psycopg_async_engine) asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') async_engines.append(asyncpg_engine) - psycopg_array_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') - array_engines.append(psycopg_array_engine) - - @event.listens_for(psycopg_array_engine, "connect") - def psycopg_connect(dbapi_connection, connection_record): - from pgvector.psycopg import register_vector - register_vector(dbapi_connection) - setup_engine = engines[0] with Session(setup_engine) as session: session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) @@ -169,9 +171,10 @@ def test_orm(self, engine): stmt = select(Item) with Session(engine) as session: items = [v[0] for v in session.execute(stmt).all()] - assert items[0].id in [1, 4, 7] - assert items[1].id in [2, 5, 8] - assert items[2].id in [3, 6, 9] + # TODO improve + assert items[0].id % 3 == 1 + assert items[1].id % 3 == 2 + assert items[2].id % 3 == 0 assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) assert items[0].embedding.dtype == np.float32 assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) From 88873e54365ca6086a1c960e2ced19ee98ea2bb2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:21:36 -0800 Subject: [PATCH 466/879] Improved tests for async SQLAlchemy engines [skip ci] --- tests/test_sqlalchemy.py | 58 +++++++++++++++++++--------------------- 1 file changed, 28 insertions(+), 30 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index dffa07b..a245ffc 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -31,6 +31,7 @@ def psycopg2_connect(dbapi_connection, connection_record): engines = [psycopg2_engine, pg8000_engine, psycopg2_type_engine] array_engines = [psycopg2_type_engine] async_engines = [] +async_array_engines = [] if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') @@ -46,11 +47,32 @@ def psycopg_connect(dbapi_connection, connection_record): engines.append(psycopg_type_engine) array_engines.append(psycopg_type_engine) + psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + + @event.listens_for(psycopg_async_type_engine.sync_engine, "connect") + def connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector_async + dbapi_connection.run_async(register_vector_async) + + async_engines.append(psycopg_async_type_engine) + async_array_engines.append(psycopg_async_type_engine) + psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') async_engines.append(psycopg_async_engine) asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') async_engines.append(asyncpg_engine) + async_array_engines.append(asyncpg_engine) + + asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + + @event.listens_for(asyncpg_type_engine.sync_engine, "connect") + def connect(dbapi_connection, connection_record): + from pgvector.asyncpg import register_vector + dbapi_connection.run_async(register_vector) + + # TODO do not throw error when types are registered + # async_array_engines.append(asyncpg_type_engine) setup_engine = engines[0] with Session(setup_engine) as session: @@ -599,6 +621,10 @@ async def test_sparsevec(self, engine): @pytest.mark.asyncio async def test_avg(self, engine): + # TODO do not throw error when types are registered + if engine == psycopg_async_type_engine: + return + async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -611,43 +637,15 @@ async def test_avg(self, engine): await engine.dispose() -@pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') +@pytest.mark.parametrize('engine', async_array_engines) class TestSqlalchemyAsyncArray: def setup_method(self): delete_items() @pytest.mark.asyncio - async def test_psycopg_vector_array(self): - engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async def test_vector_array(self, engine): async_session = async_sessionmaker(engine, expire_on_commit=False) - @event.listens_for(engine.sync_engine, "connect") - def connect(dbapi_connection, connection_record): - from pgvector.psycopg import register_vector_async - dbapi_connection.run_async(register_vector_async) - - async with async_session() as session: - async with session.begin(): - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - - # this fails if the driver does not cast arrays - item = await session.get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] - - await engine.dispose() - - @pytest.mark.asyncio - async def test_asyncpg_vector_array(self): - 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(): session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) From e78a8d5f866c2577644a64eafcaf9939c4b9ab8c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:28:23 -0800 Subject: [PATCH 467/879] Improved test code [skip ci] --- tests/test_sqlalchemy.py | 34 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index a245ffc..4b26922 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -18,7 +18,6 @@ sqlalchemy_version = 1 psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -28,15 +27,10 @@ def psycopg2_connect(dbapi_connection, connection_record): register_vector(dbapi_connection, globally=False, arrays=True) -engines = [psycopg2_engine, pg8000_engine, psycopg2_type_engine] -array_engines = [psycopg2_type_engine] -async_engines = [] -async_array_engines = [] +pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') - engines.append(psycopg_engine) - psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') @event.listens_for(psycopg_type_engine, "connect") @@ -44,9 +38,7 @@ def psycopg_connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector register_vector(dbapi_connection) - engines.append(psycopg_type_engine) - array_engines.append(psycopg_type_engine) - + psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') @event.listens_for(psycopg_async_type_engine.sync_engine, "connect") @@ -54,16 +46,7 @@ def connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector_async dbapi_connection.run_async(register_vector_async) - async_engines.append(psycopg_async_type_engine) - async_array_engines.append(psycopg_async_type_engine) - - psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') - async_engines.append(psycopg_async_engine) - asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') - async_engines.append(asyncpg_engine) - async_array_engines.append(asyncpg_engine) - asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') @event.listens_for(asyncpg_type_engine.sync_engine, "connect") @@ -71,8 +54,17 @@ def connect(dbapi_connection, connection_record): from pgvector.asyncpg import register_vector dbapi_connection.run_async(register_vector) - # TODO do not throw error when types are registered - # async_array_engines.append(asyncpg_type_engine) +engines = [psycopg2_engine, psycopg2_type_engine, pg8000_engine] +array_engines = [psycopg2_type_engine] +async_engines = [] +async_array_engines = [] + +if sqlalchemy_version > 1: + engines += [psycopg_engine, psycopg_type_engine] + array_engines += [psycopg_type_engine] + async_engines += [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine] + # TODO add asyncpg_type_engine + async_array_engines += [psycopg_async_type_engine, asyncpg_engine] setup_engine = engines[0] with Session(setup_engine) as session: From a2699639d7fd468ea68442d72227d5099ad8a64b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:33:09 -0800 Subject: [PATCH 468/879] Updated todo [skip ci] --- tests/test_sqlalchemy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 4b26922..09df9b9 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -62,8 +62,8 @@ def connect(dbapi_connection, connection_record): if sqlalchemy_version > 1: engines += [psycopg_engine, psycopg_type_engine] array_engines += [psycopg_type_engine] + # TODO support asyncpg_type_engine async_engines += [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine] - # TODO add asyncpg_type_engine async_array_engines += [psycopg_async_type_engine, asyncpg_engine] setup_engine = engines[0] From 7cd310b5cf986fab1da536a94e9f9c74379e46b9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:36:12 -0800 Subject: [PATCH 469/879] Improved test [skip ci] --- tests/test_sqlalchemy.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 09df9b9..563e3a3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -613,10 +613,6 @@ async def test_sparsevec(self, engine): @pytest.mark.asyncio async def test_avg(self, engine): - # TODO do not throw error when types are registered - if engine == psycopg_async_type_engine: - return - async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -624,7 +620,10 @@ async def test_avg(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) avg = await session.scalars(select(func.avg(Item.embedding))) - assert avg.first() == '[2.5,3.5,4.5]' + if engine == psycopg_async_type_engine: + assert avg.first().tolist() == [2.5, 3.5, 4.5] + else: + assert avg.first() == '[2.5,3.5,4.5]' await engine.dispose() From cae30a1d1b0b23620abb9c9ff4c7084ca5bac1ee Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 13:37:07 -0800 Subject: [PATCH 470/879] Improved test [skip ci] --- tests/test_sqlalchemy.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 563e3a3..aa5d81f 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -619,11 +619,8 @@ async def test_avg(self, engine): async with session.begin(): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - avg = await session.scalars(select(func.avg(Item.embedding))) - if engine == psycopg_async_type_engine: - assert avg.first().tolist() == [2.5, 3.5, 4.5] - else: - assert avg.first() == '[2.5,3.5,4.5]' + res = await session.scalars(select(avg(Item.embedding))) + assert res.first().tolist() == [2.5, 3.5, 4.5] await engine.dispose() From 3de7832d164b82e929e08d928501b081c93e3a5a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:27:35 -0800 Subject: [PATCH 471/879] Dropped support for Python < 3.9 --- .github/workflows/build.yml | 2 +- CHANGELOG.md | 4 ++++ pyproject.toml | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 562ba94..dc53dfe 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.13, 3.8] + python: [3.13, 3.9] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a517d8..1788ff4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.4.0 (unreleased) + +- Dropped support for Python < 3.9 + ## 0.3.6 (2024-10-26) - Added `arrays` option for Psycopg 2 diff --git a/pyproject.toml b/pyproject.toml index a6a6609..0f291f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ authors = [ {name = "Andrew Kane", email = "andrew@ankane.org"} ] license = {text = "MIT"} -requires-python = ">= 3.8" +requires-python = ">= 3.9" dependencies = [ "numpy" ] From 37b148f459863ee6f6c448fb93eecef569e7eb40 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:29:04 -0800 Subject: [PATCH 472/879] Removed default value [skip ci] --- pgvector/psycopg2/register.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 08a69a9..2be292f 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -5,11 +5,10 @@ from .vector import register_vector_info -# TODO remove default value for conn_or_curs in 0.4.0 # TODO make globally False by default in 0.4.0 # note: register_adapter is always global # TODO make arrays True by defalt in 0.4.0 -def register_vector(conn_or_curs=None, globally=True, arrays=False): +def register_vector(conn_or_curs, 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 From 8a621a3ae96a85320475180b2120cb6d92c095a4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:33:27 -0800 Subject: [PATCH 473/879] Changed default values of globally and arrays for register_type with Psycopg 2 [skip ci] --- pgvector/psycopg2/register.py | 4 +--- tests/test_psycopg2.py | 8 ++++---- tests/test_sqlalchemy.py | 2 +- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 2be292f..1bc9d44 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -5,10 +5,8 @@ from .vector import register_vector_info -# TODO make globally False by default in 0.4.0 # note: register_adapter is always global -# TODO make arrays True by defalt in 0.4.0 -def register_vector(conn_or_curs, globally=True, arrays=False): +def register_vector(conn_or_curs, globally=False, arrays=True): 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 diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index d661f12..85aa0e8 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -12,7 +12,7 @@ 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), embeddings vector[], half_embeddings halfvec[], sparse_embeddings sparsevec[])') -register_vector(cur, globally=False, arrays=True) +register_vector(cur) class TestPsycopg2: @@ -87,13 +87,13 @@ def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) - register_vector(cur, globally=False) + register_vector(cur) conn.close() def test_cursor_factory_connection(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) - register_vector(conn, globally=False) + register_vector(conn) conn.close() def test_pool(self): @@ -102,7 +102,7 @@ def test_pool(self): conn = pool.getconn() try: # use globally=True for apps to ensure registered with all connections - register_vector(conn, globally=False) + register_vector(conn) finally: pool.putconn(conn) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index aa5d81f..067a153 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -24,7 +24,7 @@ @event.listens_for(psycopg2_type_engine, "connect") def psycopg2_connect(dbapi_connection, connection_record): from pgvector.psycopg2 import register_vector - register_vector(dbapi_connection, globally=False, arrays=True) + register_vector(dbapi_connection) pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') From 32b09c0272545322b90d38139ee625a9a7809a71 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:36:23 -0800 Subject: [PATCH 474/879] Fixed indices and values methods returning tuple instead of list in certain cases [skip ci] --- pgvector/utils/sparsevec.py | 3 +-- tests/test_psycopg.py | 5 ++--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index a370c5e..0398106 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -108,8 +108,7 @@ def from_binary(cls, value): dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - # TODO convert indices and values to lists in 0.4.0 - return cls._from_parts(int(dim), indices, values) + return cls._from_parts(int(dim), list(indices), list(values)) @classmethod def _from_parts(cls, dim, indices, values): diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 90f80b6..cf5f09a 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -111,9 +111,8 @@ def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] assert res.dimensions() == 6 - # TODO convert indices and values to lists in 0.4.0 - assert res.indices() == (0, 2, 4) - assert res.values() == (1.5, 2, 3) + assert res.indices() == [0, 2, 4] + assert res.values() == [1.5, 2, 3] assert res.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) From c10799c3dec3dea699fc4590d3c3baa688023b23 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:38:50 -0800 Subject: [PATCH 475/879] Added support for Vector class to Psycopg 2 [skip ci] --- pgvector/psycopg2/__init__.py | 3 ++- pgvector/psycopg2/vector.py | 1 + tests/test_psycopg2.py | 12 +++++++++++- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 7c95295..f109203 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,8 +1,9 @@ from .register import register_vector -from ..utils import HalfVector, SparseVector +from ..utils import HalfVector, SparseVector, Vector __all__ = [ 'register_vector', + 'Vector', 'HalfVector', 'SparseVector' ] diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 9861f01..5bd00bb 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -24,3 +24,4 @@ def register_vector_info(oid, array_oid, scope): register_type(vectorarray, scope) register_adapter(np.ndarray, VectorAdapter) + register_adapter(Vector, VectorAdapter) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 85aa0e8..f927d86 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.psycopg2 import register_vector, HalfVector, SparseVector +from pgvector.psycopg2 import register_vector, Vector, HalfVector, SparseVector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor from psycopg2.pool import ThreadedConnectionPool @@ -29,6 +29,16 @@ def test_vector(self): assert res[0][0].dtype == np.float32 assert res[1][0] is None + def test_vector_class(self): + embedding = Vector([1.5, 2, 3]) + cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert np.array_equal(res[0][0], embedding.to_numpy()) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None + def test_halfvec(self): embedding = [1.5, 2, 3] cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) From e138e55d187c1b03f9b827c6849aabb1f9697eaf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:50:52 -0800 Subject: [PATCH 476/879] Fixed equality for types [skip ci] --- pgvector/django/functions.py | 4 ++++ pgvector/utils/bit.py | 5 +++++ pgvector/utils/halfvec.py | 5 +++++ pgvector/utils/sparsevec.py | 5 +++++ pgvector/utils/vector.py | 5 +++++ tests/test_bit.py | 4 ++++ tests/test_half_vector.py | 4 ++++ tests/test_sparse_vector.py | 5 +++++ tests/test_vector.py | 4 ++++ 9 files changed, 41 insertions(+) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index da9fbf8..6c14c3d 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -13,6 +13,10 @@ def __init__(self, expression, vector, **extra): vector = Value(SparseVector._to_db(vector)) else: vector = Value(Vector._to_db(vector)) + + # prevent error with unhashable types + self._constructor_args = ((expression, vector), extra) + super().__init__(expression, vector, **extra) diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py index 51f7556..227edc1 100644 --- a/pgvector/utils/bit.py +++ b/pgvector/utils/bit.py @@ -24,6 +24,11 @@ def __init__(self, value): def __repr__(self): return f'Bit({self.to_text()})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return np.array_equal(self.to_numpy(), other.to_numpy()) + return False + def to_list(self): return self._value.tolist() diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py index e1e5051..f335f2f 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/utils/halfvec.py @@ -16,6 +16,11 @@ def __init__(self, value): def __repr__(self): return f'HalfVector({self.to_list()})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return np.array_equal(self.to_numpy(), other.to_numpy()) + return False + def dimensions(self): return len(self._value) diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py index 0398106..8df2dfd 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/utils/sparsevec.py @@ -26,6 +26,11 @@ def __repr__(self): elements = dict(zip(self._indices, self._values)) return f'SparseVector({elements}, {self._dim})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values() + return False + def dimensions(self): return self._dim diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py index 3fa2f35..ebbcafd 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/utils/vector.py @@ -16,6 +16,11 @@ def __init__(self, value): def __repr__(self): return f'Vector({self.to_list()})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return np.array_equal(self.to_numpy(), other.to_numpy()) + return False + def dimensions(self): return len(self._value) diff --git a/tests/test_bit.py b/tests/test_bit.py index 32ab87b..a7e0093 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -37,3 +37,7 @@ def test_ndim_zero(self): def test_repr(self): assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' + + def test_equality(self): + assert Bit([True, False, True]) == Bit([True, False, True]) + assert Bit([True, False, True]) != Bit([True, False, False]) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index fdaa5f7..77a7869 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -38,5 +38,9 @@ def test_repr(self): assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' + def test_equality(self): + assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) + assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) + def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 06fe81a..24d8c20 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -52,6 +52,11 @@ def test_repr(self): assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' + def test_equality(self): + assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) + assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) + assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) + def test_dimensions(self): assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 diff --git a/tests/test_vector.py b/tests/test_vector.py index 1be2bc0..fe14dea 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -38,5 +38,9 @@ def test_repr(self): assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' + def test_equality(self): + assert Vector([1, 2, 3]) == Vector([1, 2, 3]) + assert Vector([1, 2, 3]) != Vector([1, 2, 4]) + def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 From 838ea0c73b1669c94de274eccd82f58d83ea55b4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 20:57:03 -0800 Subject: [PATCH 477/879] Moved classes to pgvector module [skip ci] --- pgvector/__init__.py | 11 +++++++++++ pgvector/{utils => }/bit.py | 0 pgvector/{utils => }/halfvec.py | 0 pgvector/{utils => }/sparsevec.py | 0 pgvector/utils/__init__.py | 5 +---- pgvector/{utils => }/vector.py | 0 6 files changed, 12 insertions(+), 4 deletions(-) create mode 100644 pgvector/__init__.py rename pgvector/{utils => }/bit.py (100%) rename pgvector/{utils => }/halfvec.py (100%) rename pgvector/{utils => }/sparsevec.py (100%) rename pgvector/{utils => }/vector.py (100%) diff --git a/pgvector/__init__.py b/pgvector/__init__.py new file mode 100644 index 0000000..3c01160 --- /dev/null +++ b/pgvector/__init__.py @@ -0,0 +1,11 @@ +from .bit import Bit +from .halfvec import HalfVector +from .sparsevec import SparseVector +from .vector import Vector + +__all__ = [ + 'Vector', + 'HalfVector', + 'Bit', + 'SparseVector' +] diff --git a/pgvector/utils/bit.py b/pgvector/bit.py similarity index 100% rename from pgvector/utils/bit.py rename to pgvector/bit.py diff --git a/pgvector/utils/halfvec.py b/pgvector/halfvec.py similarity index 100% rename from pgvector/utils/halfvec.py rename to pgvector/halfvec.py diff --git a/pgvector/utils/sparsevec.py b/pgvector/sparsevec.py similarity index 100% rename from pgvector/utils/sparsevec.py rename to pgvector/sparsevec.py diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index 3c01160..1dcc240 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1,7 +1,4 @@ -from .bit import Bit -from .halfvec import HalfVector -from .sparsevec import SparseVector -from .vector import Vector +from .. import Bit, HalfVector, SparseVector, Vector __all__ = [ 'Vector', diff --git a/pgvector/utils/vector.py b/pgvector/vector.py similarity index 100% rename from pgvector/utils/vector.py rename to pgvector/vector.py From 0ac00b4e3d39ea1ddefd8573588f7de2e60d112f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 21:00:08 -0800 Subject: [PATCH 478/879] Improved imports for tests [skip ci] --- tests/test_asyncpg.py | 3 ++- tests/test_bit.py | 2 +- tests/test_django.py | 3 ++- tests/test_half_vector.py | 2 +- tests/test_peewee.py | 3 ++- tests/test_psycopg.py | 3 ++- tests/test_psycopg2.py | 3 ++- tests/test_sparse_vector.py | 2 +- tests/test_sqlalchemy.py | 3 ++- tests/test_sqlmodel.py | 3 ++- tests/test_vector.py | 2 +- 11 files changed, 18 insertions(+), 11 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 48d1e32..3c36048 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,6 +1,7 @@ import asyncpg import numpy as np -from pgvector.asyncpg import register_vector, SparseVector +from pgvector import SparseVector +from pgvector.asyncpg import register_vector import pytest diff --git a/tests/test_bit.py b/tests/test_bit.py index a7e0093..e0dcfe6 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.utils import Bit +from pgvector import Bit import pytest diff --git a/tests/test_django.py b/tests/test_django.py index ea15771..65082a3 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -12,7 +12,8 @@ import numpy as np import os import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVector, SparseVector +from pgvector import HalfVector, SparseVector +from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock settings.configure( diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 77a7869..6a94c2e 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.utils import HalfVector +from pgvector import HalfVector import pytest diff --git a/tests/test_peewee.py b/tests/test_peewee.py index e98a0ec..670d880 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,8 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField, SparseVector +from pgvector import SparseVector +from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField db = PostgresqlDatabase('pgvector_python_test') diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index cf5f09a..6a9d0b7 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,5 +1,6 @@ import numpy as np -from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector, Vector +from pgvector import Bit, HalfVector, SparseVector, Vector +from pgvector.psycopg import register_vector, register_vector_async import psycopg from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index f927d86..1994c87 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,5 +1,6 @@ import numpy as np -from pgvector.psycopg2 import register_vector, Vector, HalfVector, SparseVector +from pgvector import HalfVector, SparseVector, Vector +from pgvector.psycopg2 import register_vector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor from psycopg2.pool import ThreadedConnectionPool diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 24d8c20..b5e7fe8 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.utils import SparseVector +from pgvector import SparseVector import pytest from scipy.sparse import coo_array diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 067a153..052edd7 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,7 +1,8 @@ import asyncpg import numpy as np import os -from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum +from pgvector import SparseVector +from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY from sqlalchemy.exc import StatementError diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 373834f..b0e8ccd 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,6 @@ import numpy as np -from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum +from pgvector import SparseVector +from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy.exc import StatementError from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text diff --git a/tests/test_vector.py b/tests/test_vector.py index fe14dea..406637f 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.utils import Vector +from pgvector import Vector import pytest From 435e31654831d303342a1100a8dd32b6c1fe42a6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 21:03:28 -0800 Subject: [PATCH 479/879] Improved imports --- pgvector/asyncpg/__init__.py | 4 +++- pgvector/asyncpg/register.py | 2 +- pgvector/django/__init__.py | 4 +++- pgvector/django/functions.py | 2 +- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- pgvector/peewee/__init__.py | 4 +++- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/sparsevec.py | 2 +- pgvector/peewee/vector.py | 2 +- pgvector/psycopg/__init__.py | 4 +++- pgvector/psycopg/bit.py | 2 +- pgvector/psycopg/halfvec.py | 2 +- pgvector/psycopg/sparsevec.py | 2 +- pgvector/psycopg/vector.py | 2 +- pgvector/psycopg2/__init__.py | 4 +++- pgvector/psycopg2/halfvec.py | 2 +- pgvector/psycopg2/sparsevec.py | 2 +- pgvector/psycopg2/vector.py | 2 +- pgvector/sqlalchemy/__init__.py | 4 +++- pgvector/sqlalchemy/halfvec.py | 2 +- pgvector/sqlalchemy/sparsevec.py | 2 +- pgvector/sqlalchemy/vector.py | 2 +- 24 files changed, 36 insertions(+), 24 deletions(-) diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 543b882..c6a3b4e 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,5 +1,7 @@ from .register import register_vector -from ..utils import Vector, HalfVector, SparseVector + +# TODO remove +from .. import Vector, HalfVector, SparseVector __all__ = [ 'register_vector', diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index a388058..63726f3 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -1,4 +1,4 @@ -from ..utils import Vector, HalfVector, SparseVector +from .. import Vector, HalfVector, SparseVector async def register_vector(conn, schema='public'): diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 09978a9..43c64a3 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -5,7 +5,9 @@ from .indexes import IvfflatIndex, HnswIndex from .sparsevec import SparseVectorField from .vector import VectorField -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'VectorExtension', diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 6c14c3d..9df4fdb 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,5 +1,5 @@ from django.db.models import FloatField, Func, Value -from ..utils import Vector, HalfVector, SparseVector +from .. import Vector, HalfVector, SparseVector class DistanceBase(Func): diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 6b59a7f..3aeb90f 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -1,6 +1,6 @@ from django import forms from django.db.models import Field -from ..utils import HalfVector +from .. import HalfVector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index d0d2d07..580f27c 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -1,6 +1,6 @@ from django import forms from django.db.models import Field -from ..utils import SparseVector +from .. import SparseVector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index a89d540..861cfde 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -1,7 +1,7 @@ from django import forms from django.db.models import Field import numpy as np -from ..utils import Vector +from .. import Vector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 945e0dc..df21200 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -2,7 +2,9 @@ from .halfvec import HalfVectorField from .sparsevec import SparseVectorField from .vector import VectorField -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'VectorField', diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index deaa14d..0901fd2 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,5 +1,5 @@ from peewee import Expression, Field -from ..utils import HalfVector +from .. import HalfVector class HalfVectorField(Field): diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 67f7d1b..86dea73 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -1,5 +1,5 @@ from peewee import Expression, Field -from ..utils import SparseVector +from .. import SparseVector class SparseVectorField(Field): diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 22a87e5..83f9997 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,5 +1,5 @@ from peewee import Expression, Field -from ..utils import Vector +from .. import Vector class VectorField(Field): diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 9007c37..980af84 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,5 +1,7 @@ from .register import register_vector, register_vector_async -from ..utils import Bit, HalfVector, SparseVector, Vector + +# TODO remove +from .. import Bit, HalfVector, SparseVector, Vector __all__ = [ 'register_vector', diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index f8eeb61..cffe8fb 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,6 +1,6 @@ from psycopg.adapt import Dumper from psycopg.pq import Format -from ..utils import Bit +from .. import Bit class BitDumper(Dumper): diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 351d2cb..b3a0060 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,6 +1,6 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import HalfVector +from .. import HalfVector class HalfVectorDumper(Dumper): diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 435fd06..384a0e1 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,6 +1,6 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import SparseVector +from .. import SparseVector class SparseVectorDumper(Dumper): diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 0f62ca9..db9e826 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,7 +1,7 @@ import psycopg from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import Vector +from .. import Vector class VectorDumper(Dumper): diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index f109203..b40c673 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,5 +1,7 @@ from .register import register_vector -from ..utils import HalfVector, SparseVector, Vector + +# TODO remove +from .. import HalfVector, SparseVector, Vector __all__ = [ 'register_vector', diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index b50e89b..0a4c736 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,5 +1,5 @@ from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type -from ..utils import HalfVector +from .. import HalfVector class HalfvecAdapter: diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index a542807..148eff2 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,5 +1,5 @@ from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type -from ..utils import SparseVector +from .. import SparseVector class SparsevecAdapter: diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 5bd00bb..562de18 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,6 +1,6 @@ import numpy as np from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type -from ..utils import Vector +from .. import Vector class VectorAdapter: diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 4955eeb..52adf88 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -4,7 +4,9 @@ from .sparsevec import SPARSEVEC from .vector import VECTOR from .vector import VECTOR as Vector -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'Vector', diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 639f77b..10688b5 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import HalfVector +from .. import HalfVector class HALFVEC(UserDefinedType): diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 370f5d1..0058679 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import SparseVector +from .. import SparseVector class SPARSEVEC(UserDefinedType): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index f57a045..5a1e11f 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import Vector +from .. import Vector class VECTOR(UserDefinedType): From 08e29e1acdcdf03965f7ffb4e1e552688df51785 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 21:04:40 -0800 Subject: [PATCH 480/879] Added todo [skip ci] --- pgvector/utils/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index 1dcc240..8cdb5d6 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1,3 +1,4 @@ +# TODO remove from .. import Bit, HalfVector, SparseVector, Vector __all__ = [ From 1c0ff62b65718899915cd51466c63b9b60c3787f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 21:50:36 -0800 Subject: [PATCH 481/879] Updated changelog [skip ci] --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1788ff4..df60740 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.4.0 (unreleased) +- Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes - Dropped support for Python < 3.9 ## 0.3.6 (2024-10-26) From f618edb43395795451d3079b1def7f6c8cbb76ba Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 21:55:07 -0800 Subject: [PATCH 482/879] Updated changelog [skip ci] --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index df60740..24aebf9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.4.0 (unreleased) - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes +- Fixed `indices` and `values` methods of `SparseVector` returning tuple instead of list in some cases - Dropped support for Python < 3.9 ## 0.3.6 (2024-10-26) From 537f3ba72519314b2bc5b65f7d625594a496345e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 22:01:16 -0800 Subject: [PATCH 483/879] Updated changelog [skip ci] --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 24aebf9..42e9bff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.4.0 (unreleased) +- Added top-level `pgvector` package - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes - Fixed `indices` and `values` methods of `SparseVector` returning tuple instead of list in some cases - Dropped support for Python < 3.9 From 571bf4287cfe8a2371f477250c02efa3f62f67a1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 9 Feb 2025 22:03:52 -0800 Subject: [PATCH 484/879] Updated changelog [skip ci] --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 42e9bff..f53a2ce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ ## 0.4.0 (unreleased) - Added top-level `pgvector` package +- Changed `globally` option to default to `False` for Psycopg 2 +- Changed `arrays` option to default to `True` for Psycopg 2 - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes - Fixed `indices` and `values` methods of `SparseVector` returning tuple instead of list in some cases - Dropped support for Python < 3.9 From 1676e3ead391493375ff6958a5b80c78080cf01e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 17:13:55 -0800 Subject: [PATCH 485/879] Test SQLAlchemy 1 on CI --- .github/workflows/build.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index dc53dfe..4d4e8ed 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,3 +24,6 @@ jobs: make sudo make install - run: pytest + + - run: pip install "SQLAlchemy<2" -U + - run: pytest tests/test_sqlalchemy.py From ac9fd532f77c1497df250e519238f7f5d627f645 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 17:49:10 -0800 Subject: [PATCH 486/879] Improved Bit constructor for uint8 NumPy arrays --- pgvector/bit.py | 11 ++++------- tests/test_bit.py | 4 +--- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 227edc1..36da723 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -7,14 +7,11 @@ def __init__(self, value): if isinstance(value, str): self._value = self.from_text(value)._value else: - # TODO change in 0.4.0 # TODO raise if dtype not bool or uint8 - # if isinstance(value, np.ndarray) and value.dtype == np.uint8: - # value = np.unpackbits(value) - # else: - # value = np.asarray(value, dtype=bool) - - value = np.asarray(value, dtype=bool) + if isinstance(value, np.ndarray) and value.dtype == np.uint8: + value = np.unpackbits(value) + else: + value = np.asarray(value, dtype=bool) if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/tests/test_bit.py b/tests/test_bit.py index e0dcfe6..1d771ca 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -15,9 +15,7 @@ def test_str(self): def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) - # TODO change in 0.4.0 - # assert Bit(arr).to_text() == '111111100000011100000000' - assert Bit(arr).to_text() == '110' + assert Bit(arr).to_text() == '111111100000011100000000' def test_ndarray_same_object(self): arr = np.array([True, False, True]) From 1b25460e6184bb744b9c71c9c5b95852bdf7c63f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:02:23 -0800 Subject: [PATCH 487/879] Raise error for unexpected dtype for Bit constructor [skip ci] --- pgvector/bit.py | 8 +++++--- tests/test_bit.py | 6 ++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 36da723..4be7385 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -7,9 +7,11 @@ def __init__(self, value): if isinstance(value, str): self._value = self.from_text(value)._value else: - # TODO raise if dtype not bool or uint8 - if isinstance(value, np.ndarray) and value.dtype == np.uint8: - value = np.unpackbits(value) + if isinstance(value, np.ndarray): + if value.dtype == np.uint8: + value = np.unpackbits(value).astype(bool) + elif value.dtype != np.bool: + raise ValueError('expected dtype to be bool or uint8') else: value = np.asarray(value, dtype=bool) diff --git a/tests/test_bit.py b/tests/test_bit.py index 1d771ca..5e1bff2 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -17,6 +17,12 @@ def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) assert Bit(arr).to_text() == '111111100000011100000000' + def test_ndarray_uint16(self): + arr = np.array([254, 7, 0], dtype=np.uint16) + with pytest.raises(ValueError) as error: + Bit(arr) + assert str(error.value) == 'expected dtype to be bool or uint8' + def test_ndarray_same_object(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] From 8b927161f7856415436159d0b2c804280261a759 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:11:57 -0800 Subject: [PATCH 488/879] Improved asyncpg tests [skip ci] --- tests/test_asyncpg.py | 45 ++++++++++++++++++++++++++----------------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 3c36048..982ea8d 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,6 +1,6 @@ import asyncpg import numpy as np -from pgvector import SparseVector +from pgvector import HalfVector, SparseVector, Vector from pgvector.asyncpg import register_vector import pytest @@ -15,13 +15,15 @@ async def test_vector(self): await register_vector(conn) - embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding = Vector([1.5, 2, 3]) + embedding2 = np.array([4.5, 5, 6]) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding) + assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) assert res[0]['embedding'].dtype == np.float32 - assert res[1]['embedding'] is None + assert np.array_equal(res[1]['embedding'], embedding2) + assert res[2]['embedding'] is None # ensures binary format is correct text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") @@ -38,12 +40,14 @@ async def test_halfvec(self): await register_vector(conn) - embedding = [1.5, 2, 3] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding = HalfVector([1.5, 2, 3]) + embedding2 = [4.5, 5, 6] + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['embedding'].to_list() == [1.5, 2, 3] - assert res[1]['embedding'] is None + assert res[0]['embedding'] == embedding + assert res[1]['embedding'] == HalfVector(embedding2) + assert res[2]['embedding'] is None # ensures binary format is correct text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") @@ -87,7 +91,7 @@ async def test_sparsevec(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['embedding'].to_list() == [1.5, 2, 3] + assert res[0]['embedding'] == embedding assert res[1]['embedding'] is None # ensures binary format is correct @@ -105,12 +109,15 @@ async def test_vector_array(self): 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]) + embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] + embeddings2 = [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[]), (ARRAY[$3, $4]::vector[])", embeddings[0], embeddings[1], embeddings2[0], embeddings2[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]) + assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy()) + assert np.array_equal(res[0]['embeddings'][1], embeddings[1].to_numpy()) + assert np.array_equal(res[1]['embeddings'][0], embeddings2[0]) + assert np.array_equal(res[1]['embeddings'][1], embeddings2[1]) await conn.close() @@ -126,10 +133,12 @@ async def init(conn): await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') - embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding = Vector([1.5, 2, 3]) + embedding2 = np.array([1.5, 2, 3]) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding) + assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) assert res[0]['embedding'].dtype == np.float32 - assert res[1]['embedding'] is None + assert np.array_equal(res[1]['embedding'], embedding2) + assert res[2]['embedding'] is None From 9f825f2e8360a4f6ec8af0341584817e5191008c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:14:09 -0800 Subject: [PATCH 489/879] Improved asyncpg test [skip ci] --- tests/test_asyncpg.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 982ea8d..34d66a1 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -110,8 +110,10 @@ async def test_vector_array(self): await register_vector(conn) embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] + await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) + embeddings2 = [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[]), (ARRAY[$3, $4]::vector[])", embeddings[0], embeddings[1], embeddings2[0], embeddings2[1]) + await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy()) From bb3b32ccf9718c3675767de3e226d3638c1f82ea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:20:20 -0800 Subject: [PATCH 490/879] Improved tests [skip ci] --- tests/test_peewee.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 670d880..d7028c3 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector import SparseVector +from pgvector import HalfVector, SparseVector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField db = PostgresqlDatabase('pgvector_python_test') @@ -77,7 +77,7 @@ def test_vector_l1_distance(self): def test_halfvec(self): Item.create(id=1, half_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.half_embedding.to_list() == [1, 2, 3] + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): create_items() @@ -129,7 +129,7 @@ def test_bit_jaccard_distance(self): def test_sparsevec(self): Item.create(id=1, sparse_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): create_items() @@ -186,7 +186,7 @@ def test_halfvec_avg(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() - assert avg.to_list() == [2.5, 3.5, 4.5] + assert avg == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self): sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() @@ -194,7 +194,7 @@ def test_halfvec_sum(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() - assert sum.to_list() == [5, 7, 9] + assert sum == HalfVector([5, 7, 9]) def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) From c7cd058ea3145fd7cdcb45f712c0f4450ddbe16e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:24:06 -0800 Subject: [PATCH 491/879] Improved tests [skip ci] --- tests/test_django.py | 12 ++++++------ tests/test_psycopg.py | 8 ++++++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 65082a3..f187ad4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -199,7 +199,7 @@ def test_vector_l1_distance(self): def test_halfvec(self): Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert item.half_embedding.to_list() == [1, 2, 3] + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): create_items() @@ -251,7 +251,7 @@ def test_bit_jaccard_distance(self): def test_sparsevec(self): Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() item = Item.objects.get(pk=1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): create_items() @@ -309,7 +309,7 @@ def test_halfvec_avg(self): Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] - assert avg.to_list() == [2.5, 3.5, 4.5] + assert avg == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self): sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] @@ -317,7 +317,7 @@ def test_halfvec_sum(self): Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] - assert sum.to_list() == [5, 7, 9] + assert sum == HalfVector([5, 7, 9]) def test_serialization(self): create_items() @@ -375,7 +375,7 @@ def test_halfvec_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert [4, 5, 6] == Item.objects.get(pk=1).half_embedding.to_list() + assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) def test_halfvec_form_save_missing(self): Item(id=1).save() @@ -432,7 +432,7 @@ def test_sparsevec_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list() + assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) def test_sparesevec_form_save_missing(self): Item(id=1).save() diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 6a9d0b7..e2a40b2 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -69,17 +69,19 @@ def test_halfvec(self): conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0] - assert res.to_list() == [1.5, 2, 3] + assert res == HalfVector([1.5, 2, 3]) def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] + assert res == HalfVector([1.5, 2, 3]) assert res.to_list() == [1.5, 2, 3] assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] + assert res == HalfVector([1.5, 2, 3]) assert res.to_list() == [1.5, 2, 3] assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) @@ -106,11 +108,12 @@ def test_sparsevec(self): conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] - assert res.to_list() == [1.5, 2, 3] + assert res == SparseVector([1.5, 2, 3]) def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] + assert res == embedding assert res.dimensions() == 6 assert res.indices() == [0, 2, 4] assert res.values() == [1.5, 2, 3] @@ -120,6 +123,7 @@ def test_sparsevec_binary_format(self): def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] + assert res == embedding assert res.dimensions() == 6 assert res.indices() == [0, 2, 4] assert res.values() == [1.5, 2, 3] From 8441b463ccc8738a951dd8fd2c9ac8b8b292c774 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:26:05 -0800 Subject: [PATCH 492/879] Improved tests [skip ci] --- tests/test_psycopg2.py | 12 ++++++------ tests/test_sparse_vector.py | 1 + 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 1994c87..71e0015 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -46,7 +46,7 @@ def test_halfvec(self): cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0].to_list() == [1.5, 2, 3] + assert res[0][0] == HalfVector([1.5, 2, 3]) assert res[1][0] is None def test_bit(self): @@ -64,7 +64,7 @@ def test_sparsevec(self): cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0].to_list() == [1.5, 2, 3] + assert res[0][0] == SparseVector([1.5, 2, 3]) assert res[1][0] is None def test_vector_array(self): @@ -82,8 +82,8 @@ def test_halfvec_array(self): 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] + assert res[0][0] == HalfVector([1.5, 2, 3]) + assert res[0][1] == HalfVector([4.5, 5, 6]) def test_sparsevec_array(self): embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] @@ -91,8 +91,8 @@ def test_sparsevec_array(self): 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] + assert res[0][0] == SparseVector([1.5, 2, 3]) + assert res[0][1] == SparseVector([4.5, 5, 6]) def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index b5e7fe8..29c3ea7 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -56,6 +56,7 @@ def test_equality(self): assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) + assert SparseVector({}, 1) != SparseVector({}, 2) def test_dimensions(self): assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 From 6b8857a3146cf581bebcf32eb81a37135aa2fc15 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:28:20 -0800 Subject: [PATCH 493/879] Improved tests [skip ci] --- tests/test_sqlalchemy.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 052edd7..d791bd6 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,7 +1,7 @@ import asyncpg import numpy as np import os -from pgvector import SparseVector +from pgvector import HalfVector, SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY @@ -256,7 +256,7 @@ def test_halfvec(self, engine): session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self, engine): create_items() @@ -348,7 +348,7 @@ def test_sparsevec(self, engine): session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self, engine): create_items() @@ -551,8 +551,8 @@ def test_halfvec_array(self, engine): # this fails if the driver does not cast arrays item = session.get(Item, 1) - assert item.half_embeddings[0].to_list() == [1, 2, 3] - assert item.half_embeddings[1].to_list() == [4, 5, 6] + assert item.half_embeddings[0] == HalfVector([1, 2, 3]) + assert item.half_embeddings[1] == HalfVector([4, 5, 6]) @pytest.mark.parametrize('engine', async_engines) @@ -582,7 +582,7 @@ async def test_halfvec(self, engine): 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 + assert item.half_embedding == HalfVector(embedding) await engine.dispose() @@ -608,7 +608,7 @@ async def test_sparsevec(self, engine): 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 + assert item.sparse_embedding == SparseVector(embedding) await engine.dispose() From 022dd061b4ebdfb7e39b23abe782bc2d89ec9e98 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:29:57 -0800 Subject: [PATCH 494/879] Improved tests [skip ci] --- tests/test_sqlalchemy.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index d791bd6..4b1e516 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,7 +1,7 @@ import asyncpg import numpy as np import os -from pgvector import HalfVector, SparseVector +from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY @@ -637,9 +637,14 @@ async def test_vector_array(self, engine): async with async_session() as session: async with session.begin(): - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) item = await session.get(Item, 1) assert item.embeddings[0].tolist() == [1, 2, 3] assert item.embeddings[1].tolist() == [4, 5, 6] + session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + item = await session.get(Item, 2) + assert item.embeddings[0].tolist() == [1, 2, 3] + assert item.embeddings[1].tolist() == [4, 5, 6] + await engine.dispose() From bb02ee2742714cb4b566b95deb71a82539977dd4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:48:17 -0800 Subject: [PATCH 495/879] Improved tests [skip ci] --- tests/test_sqlmodel.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index b0e8ccd..8a472b1 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector import SparseVector +from pgvector import HalfVector, SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy.exc import StatementError @@ -107,7 +107,7 @@ def test_halfvec(self): session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): create_items() @@ -157,7 +157,7 @@ def test_sparsevec(self): session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): create_items() @@ -220,7 +220,7 @@ def test_halfvec_avg(self): session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(avg(Item.half_embedding))).first() - assert res.to_list() == [2.5, 3.5, 4.5] + assert res == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self): with Session(engine) as session: @@ -229,7 +229,7 @@ def test_halfvec_sum(self): session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(sum(Item.half_embedding))).first() - assert res.to_list() == [5, 7, 9] + assert res == HalfVector([5, 7, 9]) def test_bad_dimensions(self): item = Item(embedding=[1, 2]) From 340caa58195fc5e7b99eed8ab7fe4e4e912fd73c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:53:17 -0800 Subject: [PATCH 496/879] Improved tests [skip ci] --- tests/test_psycopg.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index e2a40b2..0859be7 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -75,6 +75,7 @@ def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] assert res == HalfVector([1.5, 2, 3]) + # TODO move assert res.to_list() == [1.5, 2, 3] assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) @@ -82,6 +83,7 @@ def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] assert res == HalfVector([1.5, 2, 3]) + # TODO move assert res.to_list() == [1.5, 2, 3] assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) @@ -114,6 +116,7 @@ def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] assert res == embedding + # TODO move assert res.dimensions() == 6 assert res.indices() == [0, 2, 4] assert res.values() == [1.5, 2, 3] @@ -124,6 +127,7 @@ def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] assert res == embedding + # TODO move assert res.dimensions() == 6 assert res.indices() == [0, 2, 4] assert res.values() == [1.5, 2, 3] @@ -166,8 +170,8 @@ def test_binary_copy_to(self): cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: for row in copy.rows(): - assert Vector.from_binary(row[0]).to_list() == [1.5, 2, 3] - assert HalfVector.from_binary(row[1]).to_list() == [1.5, 2, 3] + assert np.array_equal(Vector.from_binary(row[0]).to_numpy(), embedding) + assert HalfVector.from_binary(row[1]) == half_embedding def test_binary_copy_to_set_types(self): embedding = np.array([1.5, 2, 3]) @@ -178,7 +182,7 @@ def test_binary_copy_to_set_types(self): copy.set_types(['vector', 'halfvec']) for row in copy.rows(): assert np.array_equal(row[0], embedding) - assert row[1].to_list() == [1.5, 2, 3] + assert row[1] == half_embedding def test_vector_array(self): embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] From e6edb2a68f1a93df94c725c6d5ba29654694feab Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 18:56:52 -0800 Subject: [PATCH 497/879] Improved tests [skip ci] --- tests/test_half_vector.py | 5 +++++ tests/test_psycopg.py | 9 --------- tests/test_sparse_vector.py | 8 ++++++++ 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 6a94c2e..a17699a 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -44,3 +44,8 @@ def test_equality(self): def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 + + def test_from_text(self): + vec = HalfVector.from_text('[1.5,2,3]') + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 0859be7..24ab321 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -83,9 +83,6 @@ def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] assert res == HalfVector([1.5, 2, 3]) - # TODO move - assert res.to_list() == [1.5, 2, 3] - assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) def test_bit(self): embedding = Bit([True, False, True]) @@ -127,12 +124,6 @@ def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] assert res == embedding - # TODO move - assert res.dimensions() == 6 - assert res.indices() == [0, 2, 4] - assert res.values() == [1.5, 2, 3] - assert res.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_text_copy_from(self): embedding = np.array([1.5, 2, 3]) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 29c3ea7..fb01b5e 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -73,3 +73,11 @@ def test_to_coo(self): def test_zero_vector_text(self): vec = SparseVector({}, 3) assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() + + def test_from_text(self): + vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') + assert vec.dimensions() == 6 + assert vec.indices() == [0, 2, 4] + assert vec.values() == [1.5, 2, 3] + assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] + assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) From b57a2e9ed35c68eaa22afe27ce93401b0190adc2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 19:06:11 -0800 Subject: [PATCH 498/879] Improved tests [skip ci] --- tests/test_half_vector.py | 8 ++++++++ tests/test_psycopg.py | 9 --------- tests/test_sparse_vector.py | 11 +++++++++++ tests/test_vector.py | 13 +++++++++++++ 4 files changed, 32 insertions(+), 9 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index a17699a..9c0b041 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,6 +1,7 @@ import numpy as np from pgvector import HalfVector import pytest +from struct import pack class TestHalfVector: @@ -49,3 +50,10 @@ def test_from_text(self): vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + + def test_from_binary(self): + data = pack('>HH3e', 3, 0, *[1.5, 2, 3]) + vec = HalfVector.from_binary(data) + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + assert vec.to_binary() == data diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 24ab321..f61b4e3 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -75,9 +75,6 @@ def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] assert res == HalfVector([1.5, 2, 3]) - # TODO move - assert res.to_list() == [1.5, 2, 3] - assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) @@ -113,12 +110,6 @@ def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] assert res == embedding - # TODO move - assert res.dimensions() == 6 - assert res.indices() == [0, 2, 4] - assert res.values() == [1.5, 2, 3] - assert res.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index fb01b5e..fb51db9 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -2,6 +2,7 @@ from pgvector import SparseVector import pytest from scipy.sparse import coo_array +from struct import pack class TestSparseVector: @@ -81,3 +82,13 @@ def test_from_text(self): assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) + + def test_from_binary(self): + data = pack('>iii3i3f', 6, 3, 0, *[0, 2, 4], *[1.5, 2, 3]) + vec = SparseVector.from_binary(data) + assert vec.dimensions() == 6 + assert vec.indices() == [0, 2, 4] + assert vec.values() == [1.5, 2, 3] + assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] + assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) + assert vec.to_binary() == data diff --git a/tests/test_vector.py b/tests/test_vector.py index 406637f..094dd34 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,6 +1,7 @@ import numpy as np from pgvector import Vector import pytest +from struct import pack class TestVector: @@ -44,3 +45,15 @@ def test_equality(self): def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 + + def test_from_text(self): + vec = Vector.from_text('[1.5,2,3]') + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + + def test_from_binary(self): + data = pack('>HH3f', 3, 0, *[1.5, 2, 3]) + vec = Vector.from_binary(data) + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + assert vec.to_binary() == data From b6ccb3043fb1dac552b4dcdf6ecb947434d3b234 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 19:06:51 -0800 Subject: [PATCH 499/879] Improved tests [skip ci] --- tests/test_half_vector.py | 2 +- tests/test_sparse_vector.py | 2 +- tests/test_vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 9c0b041..756adc2 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -52,7 +52,7 @@ def test_from_text(self): assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) def test_from_binary(self): - data = pack('>HH3e', 3, 0, *[1.5, 2, 3]) + data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index fb51db9..cf5b016 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -84,7 +84,7 @@ def test_from_text(self): assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) def test_from_binary(self): - data = pack('>iii3i3f', 6, 3, 0, *[0, 2, 4], *[1.5, 2, 3]) + data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) vec = SparseVector.from_binary(data) assert vec.dimensions() == 6 assert vec.indices() == [0, 2, 4] diff --git a/tests/test_vector.py b/tests/test_vector.py index 094dd34..c367a7a 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -52,7 +52,7 @@ def test_from_text(self): assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) def test_from_binary(self): - data = pack('>HH3f', 3, 0, *[1.5, 2, 3]) + data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) From e566d4c9b4968b232c2348e9e608d06ee90b6253 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 19:20:13 -0800 Subject: [PATCH 500/879] Improved tests [skip ci] --- tests/test_django.py | 12 ++++++------ tests/test_half_vector.py | 4 ++-- tests/test_peewee.py | 10 +++++----- tests/test_psycopg.py | 12 ++++++------ tests/test_psycopg2.py | 2 +- tests/test_sparse_vector.py | 8 ++++---- tests/test_sqlalchemy.py | 18 +++++++++--------- tests/test_sqlmodel.py | 2 +- tests/test_vector.py | 4 ++-- 9 files changed, 36 insertions(+), 36 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index f187ad4..7a8a6eb 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -165,7 +165,7 @@ def setup_method(self): def test_vector(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert np.array_equal(item.embedding, [1, 2, 3]) assert item.embedding.dtype == np.float32 def test_vector_l2_distance(self): @@ -293,7 +293,7 @@ def test_vector_avg(self): Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + assert np.array_equal(avg, [2.5, 3.5, 4.5]) def test_vector_sum(self): sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] @@ -301,7 +301,7 @@ def test_vector_sum(self): Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] - assert np.array_equal(sum, np.array([5, 7, 9])) + assert np.array_equal(sum, [5, 7, 9]) def test_halfvec_avg(self): avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] @@ -347,7 +347,7 @@ def test_vector_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() + assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6]) def test_vector_form_save_missing(self): Item(id=1).save() @@ -465,8 +465,8 @@ def test_vector_array(self): # 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] + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) def test_double_array(self): Item(id=1, double_embedding=[1, 1, 1]).save() diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 756adc2..78b4977 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -49,11 +49,11 @@ def test_dimensions(self): def test_from_text(self): vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) def test_from_binary(self): data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data diff --git a/tests/test_peewee.py b/tests/test_peewee.py index d7028c3..64fc009 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -43,7 +43,7 @@ def setup_method(self): def test_vector(self): Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert np.array_equal(item.embedding, [1, 2, 3]) assert item.embedding.dtype == np.float32 def test_vector_l2_distance(self): @@ -170,7 +170,7 @@ def test_vector_avg(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + assert np.array_equal(avg, [2.5, 3.5, 4.5]) def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() @@ -178,7 +178,7 @@ def test_vector_sum(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() - assert np.array_equal(sum, np.array([5, 7, 9])) + assert np.array_equal(sum, [5, 7, 9]) def test_halfvec_avg(self): avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() @@ -220,5 +220,5 @@ class Meta: # 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])) + # assert np.array_equal(item.embeddings[0], [1, 2, 3]) + # assert np.array_equal(item.embeddings[1], [4, 5, 6]) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index f61b4e3..698b34f 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -46,23 +46,23 @@ def test_vector_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, np.array([3, 2, 1.5])) + assert np.array_equal(res, [3, 2, 1.5]) def test_vector_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, np.array([3, 2, 1.5])) + assert np.array_equal(res, [3, 2, 1.5]) def test_vector_class_binary_format(self): embedding = Vector([1.5, 2, 3]) res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] - assert np.array_equal(res, np.array([1.5, 2, 3])) + assert np.array_equal(res, [1.5, 2, 3]) def test_vector_class_text_format(self): embedding = Vector([1.5, 2, 3]) res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, np.array([1.5, 2, 3])) + assert np.array_equal(res, [1.5, 2, 3]) def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) @@ -182,7 +182,7 @@ def configure(conn): with pool.connection() as conn: res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() - assert np.array_equal(res[0], np.array([1, 2, 3])) + assert np.array_equal(res[0], [1, 2, 3]) pool.close() @@ -218,6 +218,6 @@ async def configure(conn): async with conn.cursor() as cur: await cur.execute("SELECT '[1,2,3]'::vector") res = await cur.fetchone() - assert np.array_equal(res[0], np.array([1, 2, 3])) + assert np.array_equal(res[0], [1, 2, 3]) await pool.close() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 71e0015..8f56ef5 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -122,7 +122,7 @@ def test_pool(self): cur = conn.cursor() cur.execute("SELECT '[1,2,3]'::vector") res = cur.fetchone() - assert np.array_equal(res[0], np.array([1, 2, 3])) + assert np.array_equal(res[0], [1, 2, 3]) finally: pool.putconn(conn) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index cf5b016..dff03dd 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -9,7 +9,7 @@ class TestSparseVector: def test_list(self): vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] - assert vec.to_numpy().tolist() == [1, 0, 2, 0, 3, 0] + assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] def test_list_dimensions(self): @@ -69,7 +69,7 @@ def test_values(self): assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] def test_to_coo(self): - assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] + assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) def test_zero_vector_text(self): vec = SparseVector({}, 3) @@ -81,7 +81,7 @@ def test_from_text(self): assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) def test_from_binary(self): data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) @@ -90,5 +90,5 @@ def test_from_binary(self): assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_binary() == data diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 4b1e516..41c309f 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -201,7 +201,7 @@ def test_vector(self, engine): session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] + assert np.array_equal(item.embedding, [1, 2, 3]) def test_vector_l2_distance(self, engine): create_items() @@ -509,7 +509,7 @@ def test_automap(self, engine): with Session(engine) as session: session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) item = session.query(AutoItem).first() - assert item.embedding.tolist() == [1, 2, 3] + assert np.array_equal(item.embedding, [1, 2, 3]) def test_half_precision(self, engine): create_items() @@ -541,8 +541,8 @@ def test_vector_array(self, engine): # this fails if the driver does not cast arrays item = session.get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) def test_halfvec_array(self, engine): with Session(engine) as session: @@ -621,7 +621,7 @@ async def test_avg(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = await session.scalars(select(avg(Item.embedding))) - assert res.first().tolist() == [2.5, 3.5, 4.5] + assert np.array_equal(res.first(), [2.5, 3.5, 4.5]) await engine.dispose() @@ -639,12 +639,12 @@ async def test_vector_array(self, engine): async with session.begin(): session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) item = await session.get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) item = await session.get(Item, 2) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 8a472b1..f4994f4 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -76,7 +76,7 @@ def test_vector(self): session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] + assert np.array_equal(item.embedding, np.array([1, 2, 3])) def test_vector_l2_distance(self): create_items() diff --git a/tests/test_vector.py b/tests/test_vector.py index c367a7a..e5a16fe 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -49,11 +49,11 @@ def test_dimensions(self): def test_from_text(self): vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) def test_from_binary(self): data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), np.array([1.5, 2, 3])) + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data From 057eff226bdb992ebdd952628bf3d54996d9437d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 19:38:31 -0800 Subject: [PATCH 501/879] Improved tests [skip ci] --- tests/test_psycopg2.py | 6 ++---- tests/test_sqlalchemy.py | 3 +-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 8f56ef5..3e5c8c3 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -82,8 +82,7 @@ def test_halfvec_array(self): cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0][0] == HalfVector([1.5, 2, 3]) - assert res[0][1] == HalfVector([4.5, 5, 6]) + assert res[0] == [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] def test_sparsevec_array(self): embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] @@ -91,8 +90,7 @@ def test_sparsevec_array(self): cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0][0] == SparseVector([1.5, 2, 3]) - assert res[0][1] == SparseVector([4.5, 5, 6]) + assert res[0] == [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 41c309f..0d8d1ca 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -551,8 +551,7 @@ def test_halfvec_array(self, engine): # this fails if the driver does not cast arrays item = session.get(Item, 1) - assert item.half_embeddings[0] == HalfVector([1, 2, 3]) - assert item.half_embeddings[1] == HalfVector([4, 5, 6]) + assert item.half_embeddings == [HalfVector([1, 2, 3]), HalfVector([4, 5, 6])] @pytest.mark.parametrize('engine', async_engines) From 8443ff519ac39a9f0b9b2c7233b33accbe6f63ae Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 10 Feb 2025 20:05:01 -0800 Subject: [PATCH 502/879] Added missing dependency for example [skip ci] --- examples/implicit/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/implicit/requirements.txt b/examples/implicit/requirements.txt index 8f04b58..424abbd 100644 --- a/examples/implicit/requirements.txt +++ b/examples/implicit/requirements.txt @@ -1,3 +1,4 @@ +h5py implicit pgvector psycopg[binary] From 2496340bc5e91a0b5cad2462f276c7b488f2e36a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 14 Feb 2025 16:45:35 -0800 Subject: [PATCH 503/879] Added support for pg8000 --- CHANGELOG.md | 1 + README.md | 48 ++++++++++++++++++++++++++++- pgvector/pg8000/__init__.py | 5 ++++ pgvector/pg8000/register.py | 23 ++++++++++++++ tests/test_pg8000.py | 60 +++++++++++++++++++++++++++++++++++++ 5 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 pgvector/pg8000/__init__.py create mode 100644 pgvector/pg8000/register.py create mode 100644 tests/test_pg8000.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f53a2ce..ebc165a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.4.0 (unreleased) - Added top-level `pgvector` package +- Added support for pg8000 - Changed `globally` option to default to `False` for Psycopg 2 - Changed `arrays` option to default to `True` for Psycopg 2 - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes diff --git a/README.md b/README.md index 5a59c9d..7f980bd 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) @@ -22,6 +22,7 @@ And follow the instructions for your database library: - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) +- [pg8000](#pg8000) [unreleased] - [Peewee](#peewee) Or check out some examples: @@ -562,6 +563,51 @@ await conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +## pg8000 + +Enable the extension + +```python +conn.run('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the vector type with your connection + +```python +from pgvector.pg8000 import register_vector + +register_vector(conn) +``` + +Create a table + +```python +conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +``` + +Insert a vector + +```python +embedding = np.array([1, 2, 3]) +conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding) +``` + +Get the nearest neighbors to a vector + +```python +conn.run('SELECT * FROM items ORDER BY embedding <-> :embedding LIMIT 5', embedding=embedding) +``` + +Add an approximate index + +```python +conn.run('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') +# or +conn.run('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## Peewee Add a vector column diff --git a/pgvector/pg8000/__init__.py b/pgvector/pg8000/__init__.py new file mode 100644 index 0000000..b3b4440 --- /dev/null +++ b/pgvector/pg8000/__init__.py @@ -0,0 +1,5 @@ +from .register import register_vector + +__all__ = [ + 'register_vector' +] diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py new file mode 100644 index 0000000..15ee219 --- /dev/null +++ b/pgvector/pg8000/register.py @@ -0,0 +1,23 @@ +import numpy as np +from .. import Vector, HalfVector, SparseVector + + +def register_vector(conn): + # use to_regtype to get first matching type in search path + res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") + type_info = dict(res) + + if 'vector' not in type_info: + raise RuntimeError('vector type not found in the database') + + conn.register_out_adapter(Vector, Vector._to_db) + conn.register_out_adapter(np.ndarray, Vector._to_db) + conn.register_in_adapter(type_info['vector'], Vector._from_db) + + if 'halfvec' in type_info: + conn.register_out_adapter(HalfVector, HalfVector._to_db) + conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db) + + if 'sparsevec' in type_info: + conn.register_out_adapter(SparseVector, SparseVector._to_db) + conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py new file mode 100644 index 0000000..86c0fb1 --- /dev/null +++ b/tests/test_pg8000.py @@ -0,0 +1,60 @@ +import numpy as np +import os +from pgvector import HalfVector, SparseVector, Vector +from pgvector.pg8000 import register_vector +from pg8000.native import Connection + +conn = Connection(os.environ["USER"], database='pgvector_python_test') + +conn.run('CREATE EXTENSION IF NOT EXISTS vector') +conn.run('DROP TABLE IF EXISTS pg8000_items') +conn.run('CREATE TABLE pg8000_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(conn) + + +class TestPg8000: + def setup_method(self): + conn.run('DELETE FROM pg8000_items') + + def test_vector(self): + embedding = np.array([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') + assert np.array_equal(res[0][0], embedding) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None + + def test_vector_class(self): + embedding = Vector([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') + assert np.array_equal(res[0][0], embedding.to_numpy()) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None + + def test_halfvec(self): + embedding = HalfVector([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') + assert res[0][0] == embedding + assert res[1][0] is None + + def test_bit(self): + embedding = '101' + conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') + assert res[0][0] == '101' + assert res[1][0] is None + + def test_sparsevec(self): + embedding = SparseVector([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') + assert res[0][0] == embedding + assert res[1][0] is None From df1766b7f9ed6320958c04caf7f1b832d5320e4b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 14 Feb 2025 16:59:42 -0800 Subject: [PATCH 504/879] Simplified test code [skip ci] --- tests/test_pg8000.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 86c0fb1..4d3e474 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -8,7 +8,7 @@ conn.run('CREATE EXTENSION IF NOT EXISTS vector') conn.run('DROP TABLE IF EXISTS pg8000_items') -conn.run('CREATE TABLE pg8000_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[])') +conn.run('CREATE TABLE pg8000_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))') register_vector(conn) From 70ff5d4765bb156a45d806d3cd171b3a38f03fca Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 14 Feb 2025 17:00:58 -0800 Subject: [PATCH 505/879] Improved tests [skip ci] --- tests/test_psycopg2.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 3e5c8c3..7f4932d 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -49,6 +49,15 @@ def test_halfvec(self): assert res[0][0] == HalfVector([1.5, 2, 3]) assert res[1][0] is None + def test_halfvec_class(self): + embedding = HalfVector([1.5, 2, 3]) + cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][0] == embedding + assert res[1][0] is None + def test_bit(self): embedding = '101' cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) From ac9e398f511ca65f11f62f4296e94f2106367936 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 16 Feb 2025 13:15:19 -0800 Subject: [PATCH 506/879] Improved examples [skip ci] --- examples/cohere/example.py | 6 +++--- examples/openai/example.py | 25 ++++++++++++++--------- examples/sentence_transformers/example.py | 14 ++++++------- examples/sparse_search/example.py | 6 +++--- 4 files changed, 28 insertions(+), 23 deletions(-) diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 780352a..393d1e0 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -12,7 +12,7 @@ conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))') -def fetch_embeddings(input, input_type): +def embed(input, input_type): co = cohere.Client() response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary']) return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] @@ -23,12 +23,12 @@ def fetch_embeddings(input, input_type): 'The cat is purring', 'The bear is growling' ] -embeddings = fetch_embeddings(input, 'search_document') +embeddings = embed(input, 'search_document') for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Bit(embedding))) query = 'forest' -query_embedding = fetch_embeddings([query], 'search_query')[0] +query_embedding = embed([query], 'search_query')[0] result = conn.execute('SELECT content FROM documents ORDER BY embedding <~> %s LIMIT 5', (Bit(query_embedding),)).fetchall() for row in result: print(row[0]) diff --git a/examples/openai/example.py b/examples/openai/example.py index ebed3d0..b9a078c 100644 --- a/examples/openai/example.py +++ b/examples/openai/example.py @@ -1,3 +1,4 @@ +import numpy as np from openai import OpenAI from pgvector.psycopg import register_vector import psycopg @@ -10,20 +11,24 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))') + +def embed(input): + client = OpenAI() + response = client.embeddings.create(input=input, model='text-embedding-3-small') + return [v.embedding for v in response.data] + + input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] - -client = OpenAI() -response = client.embeddings.create(input=input, model='text-embedding-3-small') -embeddings = [v.embedding for v in response.data] - +embeddings = embed(input) for content, embedding in zip(input, embeddings): - conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, np.array(embedding))) -document_id = 1 -neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() -for neighbor in neighbors: - print(neighbor[0]) +query = 'forest' +query_embedding = embed([query])[0] +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (np.array(query_embedding),)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/sentence_transformers/example.py b/examples/sentence_transformers/example.py index d4e7f96..3a7dca5 100644 --- a/examples/sentence_transformers/example.py +++ b/examples/sentence_transformers/example.py @@ -10,19 +10,19 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') +model = SentenceTransformer('all-MiniLM-L6-v2') + input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] - -model = SentenceTransformer('all-MiniLM-L6-v2') embeddings = model.encode(input) - for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) -document_id = 1 -neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() -for neighbor in neighbors: - print(neighbor[0]) +query = 'forest' +query_embedding = model.encode(query) +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (query_embedding,)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/sparse_search/example.py b/examples/sparse_search/example.py index fa6074e..2b5daea 100644 --- a/examples/sparse_search/example.py +++ b/examples/sparse_search/example.py @@ -20,7 +20,7 @@ special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()] -def fetch_embeddings(input): +def embed(input): feature = tokenizer( input, padding=True, @@ -42,12 +42,12 @@ def fetch_embeddings(input): 'The cat is purring', 'The bear is growling' ] -embeddings = fetch_embeddings(input) +embeddings = embed(input) for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector(embedding))) query = 'forest' -query_embedding = fetch_embeddings([query])[0] +query_embedding = embed([query])[0] result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector(query_embedding),)).fetchall() for row in result: print(row[0]) From 1443c3c3ca11b9efadb07612758c2ba62fb4ec65 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 16 Feb 2025 13:50:50 -0800 Subject: [PATCH 507/879] Added halfvec example for OpenAI [skip ci] --- examples/openai/halfvec.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 examples/openai/halfvec.py diff --git a/examples/openai/halfvec.py b/examples/openai/halfvec.py new file mode 100644 index 0000000..185c785 --- /dev/null +++ b/examples/openai/halfvec.py @@ -0,0 +1,34 @@ +from openai import OpenAI +from pgvector.psycopg import register_vector, HalfVector +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding halfvec(3072))') +conn.execute('CREATE INDEX ON documents USING hnsw (embedding halfvec_cosine_ops)') + + +def embed(input): + client = OpenAI() + response = client.embeddings.create(input=input, model='text-embedding-3-large') + return [v.embedding for v in response.data] + + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +embeddings = embed(input) +for content, embedding in zip(input, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, HalfVector(embedding))) + +query = 'forest' +query_embedding = embed([query])[0] +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (HalfVector(query_embedding),)).fetchall() +for row in result: + print(row[0]) From 12146d74db24514831138b43ec69273e289cde1a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 16 Feb 2025 18:34:41 -0800 Subject: [PATCH 508/879] Improved example [skip ci] --- examples/sentence_transformers/example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/sentence_transformers/example.py b/examples/sentence_transformers/example.py index 3a7dca5..50997d9 100644 --- a/examples/sentence_transformers/example.py +++ b/examples/sentence_transformers/example.py @@ -10,7 +10,7 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') -model = SentenceTransformer('all-MiniLM-L6-v2') +model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') input = [ 'The dog is barking', From 78466224ec95a38441240753f090625056b87b1e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 19 Feb 2025 15:54:25 -0800 Subject: [PATCH 509/879] Added reference section to readme [skip ci] --- README.md | 93 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/README.md b/README.md index 7f980bd..299753e 100644 --- a/README.md +++ b/README.md @@ -665,6 +665,99 @@ Item.add_index('embedding vector_l2_ops', using='hnsw') Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +## Reference + +### Half Vectors + +Create a half vector from a list + +```python +vec = HalfVector([1, 2, 3]) +``` + +Or a NumPy array + +```python +vec = HalfVector(np.array([1, 2, 3])) +``` + +Get a list + +```python +lst = vec.to_list() +``` + +Get a NumPy array + +```python +arr = vec.to_numpy() +``` + +### Sparse Vectors + +Create a sparse vector from a list + +```python +vec = SparseVector([1, 0, 2, 0, 3, 0]) +``` + +Or a NumPy array + +```python +vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) +``` + +Or a SciPy sparse array + +```python +arr = coo_array(([1, 2, 3], ([0, 2, 4],)), shape=(6,)) +vec = SparseVector(arr) +``` + +Or a dictionary of non-zero elements + +```python +vec = SparseVector({0: 1, 2: 2, 4: 3}, 6) +``` + +Note: Indices start at 0 + +Get the number of dimensions + +```python +dim = vec.dimensions() +``` + +Get the indices of non-zero elements + +```python +indices = vec.indices() +``` + +Get the values of non-zero elements + +```python +values = vec.values() +``` + +Get a list + +```python +lst = vec.to_list() +``` + +Get a NumPy array + +```python +arr = vec.to_numpy() +``` + +Get a SciPy sparse array + +```python +arr = vec.to_coo() +``` + ## History View the [changelog](https://github.com/pgvector/pgvector-python/blob/master/CHANGELOG.md) From ac1a543ab33a09efa2758f0179cea6a89257b601 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 14:07:29 -0700 Subject: [PATCH 510/879] Improved validation for Bit constructor --- pgvector/bit.py | 4 +++- tests/test_bit.py | 5 +++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 4be7385..9a890a1 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -13,7 +13,9 @@ def __init__(self, value): elif value.dtype != np.bool: raise ValueError('expected dtype to be bool or uint8') else: - value = np.asarray(value, dtype=bool) + value = np.asarray(value) + if value.dtype != np.bool: + raise ValueError('expected dtype to be bool') if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/tests/test_bit.py b/tests/test_bit.py index 5e1bff2..0c661d0 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -7,6 +7,11 @@ class TestBit: def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] + def test_list_int(self): + with pytest.raises(ValueError) as error: + Bit([254, 7, 0]) + assert str(error.value) == 'expected dtype to be bool' + def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] From 900cbb38370eebfeebdd519482cfd1a30cf6e937 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 14:14:46 -0700 Subject: [PATCH 511/879] Improved error message --- pgvector/bit.py | 2 +- tests/test_bit.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 9a890a1..a8feb55 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -15,7 +15,7 @@ def __init__(self, value): else: value = np.asarray(value) if value.dtype != np.bool: - raise ValueError('expected dtype to be bool') + raise ValueError('expected all elements to be boolean') if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/tests/test_bit.py b/tests/test_bit.py index 0c661d0..ae27359 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -7,10 +7,15 @@ class TestBit: def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] + def test_list_none(self): + with pytest.raises(ValueError) as error: + Bit([True, None, True]) + assert str(error.value) == 'expected all elements to be boolean' + def test_list_int(self): with pytest.raises(ValueError) as error: Bit([254, 7, 0]) - assert str(error.value) == 'expected dtype to be bool' + assert str(error.value) == 'expected all elements to be boolean' def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] From 534ec18683d4c5e3058ba14d7810d0d5df7d8c55 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 14:34:18 -0700 Subject: [PATCH 512/879] Added support for bytes to Bit constructor --- CHANGELOG.md | 1 + pgvector/bit.py | 2 ++ tests/test_bit.py | 4 ++++ 3 files changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ebc165a..89e955a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ - Added top-level `pgvector` package - Added support for pg8000 +- Added support for `bytes` to `Bit` constructor - Changed `globally` option to default to `False` for Psycopg 2 - Changed `arrays` option to default to `True` for Psycopg 2 - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes diff --git a/pgvector/bit.py b/pgvector/bit.py index a8feb55..8766f65 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -6,6 +6,8 @@ class Bit: def __init__(self, value): if isinstance(value, str): self._value = self.from_text(value)._value + elif isinstance(value, bytes): + self._value = np.unpackbits(np.frombuffer(value, dtype=np.uint8)).astype(bool) else: if isinstance(value, np.ndarray): if value.dtype == np.uint8: diff --git a/tests/test_bit.py b/tests/test_bit.py index ae27359..571205f 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -23,6 +23,10 @@ def test_tuple(self): def test_str(self): assert Bit('101').to_list() == [True, False, True] + def test_bytes(self): + assert Bit(b'\xff\x00').to_list() == [True, True, True, True, True, True, True, True, False, False, False, False, False, False, False, False] + assert Bit(b'\xfe\x07').to_list() == [True, True, True, True, True, True, True, False, False, False, False, False, False, True, True, True] + def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) assert Bit(arr).to_text() == '111111100000011100000000' From 2d1b754773f8c4f41970b3f61b93b20460961f98 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 14:54:26 -0700 Subject: [PATCH 513/879] Restored backwards compatibility of Bit constructor --- pgvector/bit.py | 15 ++++++--------- tests/test_bit.py | 18 ++++++++---------- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 8766f65..935f0f0 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,5 +1,6 @@ import numpy as np from struct import pack, unpack_from +from warnings import warn class Bit: @@ -9,15 +10,11 @@ def __init__(self, value): elif isinstance(value, bytes): self._value = np.unpackbits(np.frombuffer(value, dtype=np.uint8)).astype(bool) else: - if isinstance(value, np.ndarray): - if value.dtype == np.uint8: - value = np.unpackbits(value).astype(bool) - elif value.dtype != np.bool: - raise ValueError('expected dtype to be bool or uint8') - else: - value = np.asarray(value) - if value.dtype != np.bool: - raise ValueError('expected all elements to be boolean') + value = np.asarray(value) + + if value.dtype != np.bool: + warn('expected elements to be boolean', stacklevel=2) + value = value.astype(bool) if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/tests/test_bit.py b/tests/test_bit.py index 571205f..a13f476 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -8,14 +8,12 @@ def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] def test_list_none(self): - with pytest.raises(ValueError) as error: - Bit([True, None, True]) - assert str(error.value) == 'expected all elements to be boolean' + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit([True, None, True]).to_text() == '101' def test_list_int(self): - with pytest.raises(ValueError) as error: - Bit([254, 7, 0]) - assert str(error.value) == 'expected all elements to be boolean' + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit([254, 7, 0]).to_text() == '110' def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] @@ -29,13 +27,13 @@ def test_bytes(self): def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) - assert Bit(arr).to_text() == '111111100000011100000000' + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit(arr).to_text() == '110' def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) - with pytest.raises(ValueError) as error: - Bit(arr) - assert str(error.value) == 'expected dtype to be bool or uint8' + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit(arr).to_text() == '110' def test_ndarray_same_object(self): arr = np.array([True, False, True]) From 2ce3f43e6693fec29e92fa84f7d46fefb96f98f0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 15:35:09 -0700 Subject: [PATCH 514/879] Improved internal representation of Bit class --- pgvector/bit.py | 47 +++++++++++++++++++++++++++-------------------- tests/test_bit.py | 10 +++++----- 2 files changed, 32 insertions(+), 25 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 935f0f0..72b8052 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -5,51 +5,58 @@ class Bit: def __init__(self, value): - if isinstance(value, str): - self._value = self.from_text(value)._value - elif isinstance(value, bytes): - self._value = np.unpackbits(np.frombuffer(value, dtype=np.uint8)).astype(bool) + if isinstance(value, bytes): + self._len = 8 * len(value) + self._data = value else: - value = np.asarray(value) + if isinstance(value, str): + value = [v != '0' for v in value] + else: + value = np.asarray(value) - if value.dtype != np.bool: - warn('expected elements to be boolean', stacklevel=2) - value = value.astype(bool) + if value.dtype != np.bool: + warn('expected elements to be boolean', stacklevel=2) + value = value.astype(bool) - if value.ndim != 1: - raise ValueError('expected ndim to be 1') + if value.ndim != 1: + raise ValueError('expected ndim to be 1') - self._value = value + self._len = len(value) + self._data = np.packbits(value).tobytes() def __repr__(self): return f'Bit({self.to_text()})' def __eq__(self, other): if isinstance(other, self.__class__): - return np.array_equal(self.to_numpy(), other.to_numpy()) + return self._len == other._len and self._data == other._data return False def to_list(self): - return self._value.tolist() + return self.to_numpy().tolist() def to_numpy(self): - return self._value + return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) def to_text(self): - return ''.join(self._value.astype(np.uint8).astype(str)) + return ''.join(format(v, '08b') for v in self._data)[:self._len] def to_binary(self): - return pack('>i', len(self._value)) + np.packbits(self._value).tobytes() + return pack('>i', self._len) + self._data @classmethod def from_text(cls, value): - return cls(np.asarray([v != '0' for v in value], dtype=bool)) + return cls(str(value)) @classmethod def from_binary(cls, value): - count = unpack_from('>i', value)[0] - buf = np.frombuffer(value, dtype=np.uint8, offset=4) - return cls(np.unpackbits(buf, count=count).astype(bool)) + if not isinstance(value, bytes): + raise ValueError('expected bytes') + + bit = cls.__new__(cls) + bit._len = unpack_from('>i', value)[0] + bit._data = value[4:] + return bit @classmethod def _to_db(cls, value): diff --git a/tests/test_bit.py b/tests/test_bit.py index a13f476..cf1275e 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -25,6 +25,11 @@ def test_bytes(self): assert Bit(b'\xff\x00').to_list() == [True, True, True, True, True, True, True, True, False, False, False, False, False, False, False, False] assert Bit(b'\xfe\x07').to_list() == [True, True, True, True, True, True, True, False, False, False, False, False, False, True, True, True] + def test_ndarray(self): + arr = np.array([True, False, True]) + assert Bit(arr).to_list() == [True, False, True] + assert np.array_equal(Bit(arr).to_numpy(), arr) + def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.warns(UserWarning, match='expected elements to be boolean'): @@ -35,11 +40,6 @@ def test_ndarray_uint16(self): with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' - def test_ndarray_same_object(self): - arr = np.array([True, False, True]) - assert Bit(arr).to_list() == [True, False, True] - assert Bit(arr).to_numpy() is arr - def test_ndim_two(self): with pytest.raises(ValueError) as error: Bit([[True, False], [True, False]]) From c2c17c2ab6365e55677bde47d1d13c63b4e87642 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 16:02:46 -0700 Subject: [PATCH 515/879] Removed warning for result of np.unpackbits --- pgvector/bit.py | 4 +++- tests/test_bit.py | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 72b8052..edfaec6 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -15,7 +15,9 @@ def __init__(self, value): value = np.asarray(value) if value.dtype != np.bool: - warn('expected elements to be boolean', stacklevel=2) + # allow result of np.unpackbits + if value.dtype != np.uint8 or np.any(value > 1): + warn('expected elements to be boolean', stacklevel=2) value = value.astype(bool) if value.ndim != 1: diff --git a/tests/test_bit.py b/tests/test_bit.py index cf1275e..ef049c7 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -30,6 +30,10 @@ def test_ndarray(self): assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) + def test_ndarray_unpackbits(self): + arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) + assert Bit(arr).to_text() == '111111100000011100000000' + def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.warns(UserWarning, match='expected elements to be boolean'): From 50fac76f7959a155444e46d9e11be42403b09b26 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 16:04:10 -0700 Subject: [PATCH 516/879] Improved test --- tests/test_bit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index ef049c7..5a71642 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -22,8 +22,8 @@ def test_str(self): assert Bit('101').to_list() == [True, False, True] def test_bytes(self): - assert Bit(b'\xff\x00').to_list() == [True, True, True, True, True, True, True, True, False, False, False, False, False, False, False, False] - assert Bit(b'\xfe\x07').to_list() == [True, True, True, True, True, True, True, False, False, False, False, False, False, True, True, True] + assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' + assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' def test_ndarray(self): arr = np.array([True, False, True]) From 92bb02a531fc012369ee20f065028aec230d5dcf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 16:05:17 -0700 Subject: [PATCH 517/879] Updated comment [skip ci] --- pgvector/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index edfaec6..26a9d8d 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -15,7 +15,7 @@ def __init__(self, value): value = np.asarray(value) if value.dtype != np.bool: - # allow result of np.unpackbits + # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): warn('expected elements to be boolean', stacklevel=2) value = value.astype(bool) From 4e22f9b26545f1b871cfba0fde21812ebc88ca84 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 16:16:01 -0700 Subject: [PATCH 518/879] Updated warning message --- pgvector/bit.py | 2 +- tests/test_bit.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 26a9d8d..e82b325 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -17,7 +17,7 @@ def __init__(self, value): if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): - warn('expected elements to be boolean', stacklevel=2) + warn('elements should be boolean', stacklevel=2) value = value.astype(bool) if value.ndim != 1: diff --git a/tests/test_bit.py b/tests/test_bit.py index 5a71642..e920228 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -8,11 +8,11 @@ def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] def test_list_none(self): - with pytest.warns(UserWarning, match='expected elements to be boolean'): + with pytest.warns(UserWarning, match='elements should be boolean'): assert Bit([True, None, True]).to_text() == '101' def test_list_int(self): - with pytest.warns(UserWarning, match='expected elements to be boolean'): + with pytest.warns(UserWarning, match='elements should be boolean'): assert Bit([254, 7, 0]).to_text() == '110' def test_tuple(self): @@ -36,12 +36,12 @@ def test_ndarray_unpackbits(self): def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) - with pytest.warns(UserWarning, match='expected elements to be boolean'): + with pytest.warns(UserWarning, match='elements should be boolean'): assert Bit(arr).to_text() == '110' def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) - with pytest.warns(UserWarning, match='expected elements to be boolean'): + with pytest.warns(UserWarning, match='elements should be boolean'): assert Bit(arr).to_text() == '110' def test_ndim_two(self): From 7a2dd806e79ad82960cc1a89159ca61f9a12a373 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 16:20:17 -0700 Subject: [PATCH 519/879] Revert "Updated warning message" This reverts commit 4e22f9b26545f1b871cfba0fde21812ebc88ca84. --- pgvector/bit.py | 2 +- tests/test_bit.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index e82b325..26a9d8d 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -17,7 +17,7 @@ def __init__(self, value): if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): - warn('elements should be boolean', stacklevel=2) + warn('expected elements to be boolean', stacklevel=2) value = value.astype(bool) if value.ndim != 1: diff --git a/tests/test_bit.py b/tests/test_bit.py index e920228..5a71642 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -8,11 +8,11 @@ def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] def test_list_none(self): - with pytest.warns(UserWarning, match='elements should be boolean'): + with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit([True, None, True]).to_text() == '101' def test_list_int(self): - with pytest.warns(UserWarning, match='elements should be boolean'): + with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit([254, 7, 0]).to_text() == '110' def test_tuple(self): @@ -36,12 +36,12 @@ def test_ndarray_unpackbits(self): def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) - with pytest.warns(UserWarning, match='elements should be boolean'): + with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) - with pytest.warns(UserWarning, match='elements should be boolean'): + with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' def test_ndim_two(self): From 6bb6df8cce6d5b03e1a8a9b683ae37faaf12db7a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 16:35:04 -0700 Subject: [PATCH 520/879] Removed unreleased import --- pgvector/psycopg2/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index b40c673..33e5124 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,11 +1,10 @@ from .register import register_vector # TODO remove -from .. import HalfVector, SparseVector, Vector +from .. import HalfVector, SparseVector __all__ = [ 'register_vector', - 'Vector', 'HalfVector', 'SparseVector' ] From a8f2a5f8428ae10d79be53c0367fc007eca4ab78 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 15 Mar 2025 17:53:02 -0700 Subject: [PATCH 521/879] Version bump to 0.4.0 [skip ci] --- CHANGELOG.md | 2 +- README.md | 2 +- pyproject.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 89e955a..d0e2730 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.4.0 (unreleased) +## 0.4.0 (2025-03-15) - Added top-level `pgvector` package - Added support for pg8000 diff --git a/README.md b/README.md index 299753e..b6bc055 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ And follow the instructions for your database library: - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) -- [pg8000](#pg8000) [unreleased] +- [pg8000](#pg8000) - [Peewee](#peewee) Or check out some examples: diff --git a/pyproject.toml b/pyproject.toml index 0f291f5..b889f4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.3.6" +version = "0.4.0" description = "pgvector support for Python" readme = "README.md" authors = [ From e19df465f0745aef4240f5388b5ca765137397be Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 17 Mar 2025 15:58:59 -0700 Subject: [PATCH 522/879] Added basic RAG example [skip ci] --- .gitignore | 1 + README.md | 1 + examples/rag/example.py | 65 +++++++++++++++++++++++++++++++++++ examples/rag/requirements.txt | 3 ++ 4 files changed, 70 insertions(+) create mode 100644 examples/rag/example.py create mode 100644 examples/rag/requirements.txt diff --git a/.gitignore b/.gitignore index f7ff659..c55ff44 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ venv/ *.pyc __pycache__ .pytest_cache/ +examples/rag/README.md diff --git a/README.md b/README.md index b6bc055..24d9bb9 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ And follow the instructions for your database library: Or check out some examples: +- [Retrieval-augmented generation](https://github.com/pgvector/pgvector-python/blob/master/examples/rag/example.py) with Ollama - [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai/example.py) with OpenAI - [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere/example.py) with Cohere - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_transformers/example.py) with SentenceTransformers diff --git a/examples/rag/example.py b/examples/rag/example.py new file mode 100644 index 0000000..4d5d307 --- /dev/null +++ b/examples/rag/example.py @@ -0,0 +1,65 @@ +# Run: +# ollama pull llama3.2 +# ollama pull nomic-embed-text +# ollama serve + +import numpy as np +import ollama +from pathlib import Path +from pgvector.psycopg import register_vector +import psycopg +import urllib.request + +query = 'What index types are supported?' +load_data = True + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +if load_data: + # get data + url = 'https://raw.githubusercontent.com/pgvector/pgvector/refs/heads/master/README.md' + dest = Path(__file__).parent / 'README.md' + if not dest.exists(): + urllib.request.urlretrieve(url, dest) + + with open(dest, encoding='utf-8') as f: + doc = f.read() + + # generate chunks + # TODO improve chunking + # TODO remove markdown + chunks = doc.split('\n## ') + + # embed chunks + # nomic-embed-text has task instruction prefix + input = ['search_document: ' + chunk for chunk in chunks] + embeddings = ollama.embed(model='nomic-embed-text', input=input).embeddings + + # create table + conn.execute('DROP TABLE IF EXISTS chunks') + conn.execute('CREATE TABLE chunks (id bigserial PRIMARY KEY, content text, embedding vector(768))') + + # store chunks + cur = conn.cursor() + with cur.copy('COPY chunks (content, embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: + copy.set_types(['text', 'vector']) + + for content, embedding in zip(chunks, embeddings): + copy.write_row([content, embedding]) + +# embed query +# nomic-embed-text has task instruction prefix +input = 'search_query: ' + query +embedding = ollama.embed(model='nomic-embed-text', input=input).embeddings[0] + +# retrieve chunks +result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (np.array(embedding),)).fetchall() +context = '\n\n'.join([row[0] for row in result]) + +# get answer +# TODO improve prompt +prompt = f'Answer this question: {query}\n\n{context}' +response = ollama.generate(model='llama3.2', prompt=prompt).response +print(response) diff --git a/examples/rag/requirements.txt b/examples/rag/requirements.txt new file mode 100644 index 0000000..4eb5864 --- /dev/null +++ b/examples/rag/requirements.txt @@ -0,0 +1,3 @@ +ollama +pgvector +psycopg[binary] From 1901b9cc8ab1eaf3a7415e3424509381a3399ccc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 24 Mar 2025 01:20:18 -0700 Subject: [PATCH 523/879] Improved test [skip ci] --- tests/test_sqlalchemy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0d8d1ca..5aec977 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -43,7 +43,7 @@ def psycopg_connect(dbapi_connection, connection_record): psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') @event.listens_for(psycopg_async_type_engine.sync_engine, "connect") - def connect(dbapi_connection, connection_record): + def psycopg_async_connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector_async dbapi_connection.run_async(register_vector_async) @@ -51,7 +51,7 @@ def connect(dbapi_connection, connection_record): asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') @event.listens_for(asyncpg_type_engine.sync_engine, "connect") - def connect(dbapi_connection, connection_record): + def asyncpg_connect(dbapi_connection, connection_record): from pgvector.asyncpg import register_vector dbapi_connection.run_async(register_vector) From eb654016181b69e9ed06871c39d8df329614cb66 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 8 Apr 2025 22:42:07 -0700 Subject: [PATCH 524/879] Added ColBERT example for approximate search - #123 [skip ci] --- examples/colbert/approximate.py | 75 +++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 examples/colbert/approximate.py diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py new file mode 100644 index 0000000..0508d0f --- /dev/null +++ b/examples/colbert/approximate.py @@ -0,0 +1,75 @@ +# approach from section 3.6 in https://arxiv.org/abs/2004.12832 + +from colbert.infra import ColBERTConfig +from colbert.modeling.checkpoint import Checkpoint +from pgvector.psycopg import register_vector +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('DROP TABLE IF EXISTS document_embeddings') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text)') +conn.execute('CREATE TABLE document_embeddings (id bigserial PRIMARY KEY, document_id bigint, embedding vector(128))') +conn.execute(""" +CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ + WITH queries AS ( + SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) + ), + documents AS ( + SELECT unnest(document) AS document + ), + similarities AS ( + SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents + ), + max_similarities AS ( + SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number + ) + SELECT SUM(max_similarity) FROM max_similarities +$$ LANGUAGE SQL +""") + +config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +doc_embeddings = checkpoint.docFromText(input, keep_dims=False) +for content, embeddings in zip(input, doc_embeddings): + with conn.transaction(): + result = conn.execute('INSERT INTO documents (content) VALUES (%s) RETURNING id', (content,)).fetchone() + params = [] + for embedding in embeddings: + params.extend([result[0], embedding.numpy()]) + values = ', '.join(['(%s, %s)' for _ in embeddings]) + conn.execute(f'INSERT INTO document_embeddings (document_id, embedding) VALUES {values}', params) + +conn.execute('CREATE INDEX ON document_embeddings (document_id)') +conn.execute('CREATE INDEX ON document_embeddings USING hnsw (embedding vector_cosine_ops)') + +query = 'puppy' +query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query])[0]] +approximate_stage = ' UNION ALL '.join(['(SELECT document_id FROM document_embeddings ORDER BY embedding <=> %s LIMIT 5)' for _ in query_embeddings]) +sql = f""" +WITH approximate_stage AS ( + {approximate_stage} +), +embeddings AS ( + SELECT document_id, array_agg(embedding) AS embeddings FROM document_embeddings + WHERE document_id IN (SELECT DISTINCT document_id FROM approximate_stage) + GROUP BY document_id +) +SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents +INNER JOIN embeddings ON embeddings.document_id = documents.id +ORDER BY max_sim DESC LIMIT 10 +""" +params = [v for v in query_embeddings] + [query_embeddings] +result = conn.execute(sql, params).fetchall() +for row in result: + print(row) From 8718cdde9f91490b39a06293ec48d8f26193334b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 8 Apr 2025 22:51:47 -0700 Subject: [PATCH 525/879] Updated comment [skip ci] --- examples/colbert/approximate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py index 0508d0f..fc1d396 100644 --- a/examples/colbert/approximate.py +++ b/examples/colbert/approximate.py @@ -1,4 +1,4 @@ -# approach from section 3.6 in https://arxiv.org/abs/2004.12832 +# based on section 3.6 of https://arxiv.org/abs/2004.12832 from colbert.infra import ColBERTConfig from colbert.modeling.checkpoint import Checkpoint From 123f74343b03a7910b8b66de4fc33127f4696430 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 8 Apr 2025 22:53:10 -0700 Subject: [PATCH 526/879] Improved example [skip ci] --- examples/colbert/approximate.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py index fc1d396..290e66d 100644 --- a/examples/colbert/approximate.py +++ b/examples/colbert/approximate.py @@ -12,8 +12,10 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('DROP TABLE IF EXISTS document_embeddings') + conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text)') conn.execute('CREATE TABLE document_embeddings (id bigserial PRIMARY KEY, document_id bigint, embedding vector(128))') + conn.execute(""" CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ WITH queries AS ( @@ -69,7 +71,7 @@ INNER JOIN embeddings ON embeddings.document_id = documents.id ORDER BY max_sim DESC LIMIT 10 """ -params = [v for v in query_embeddings] + [query_embeddings] +params = query_embeddings + [query_embeddings] result = conn.execute(sql, params).fetchall() for row in result: print(row) From bef31a81ced1517f33c5fd960e7ba10f2fd5d8e2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 8 Apr 2025 23:02:35 -0700 Subject: [PATCH 527/879] Improved ColBERT examples [skip ci] --- examples/colbert/approximate.py | 4 ++++ examples/colbert/exact.py | 4 ++++ examples/colbert/requirements.txt | 1 + 3 files changed, 9 insertions(+) diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py index 290e66d..623f913 100644 --- a/examples/colbert/approximate.py +++ b/examples/colbert/approximate.py @@ -4,6 +4,10 @@ from colbert.modeling.checkpoint import Checkpoint from pgvector.psycopg import register_vector import psycopg +import warnings + +# ignore warnings from colbert +warnings.filterwarnings('ignore') conn = psycopg.connect(dbname='pgvector_example', autocommit=True) diff --git a/examples/colbert/exact.py b/examples/colbert/exact.py index 1c90b47..ceed2e3 100644 --- a/examples/colbert/exact.py +++ b/examples/colbert/exact.py @@ -2,6 +2,10 @@ from colbert.modeling.checkpoint import Checkpoint from pgvector.psycopg import register_vector import psycopg +import warnings + +# ignore warnings from colbert +warnings.filterwarnings('ignore') conn = psycopg.connect(dbname='pgvector_example', autocommit=True) diff --git a/examples/colbert/requirements.txt b/examples/colbert/requirements.txt index 4402ce8..54b2cb9 100644 --- a/examples/colbert/requirements.txt +++ b/examples/colbert/requirements.txt @@ -1,3 +1,4 @@ colbert-ai pgvector psycopg[binary] +transformers==4.49.0 From 208b11a893c6e5a672481847251bc13a72c84165 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 8 Apr 2025 23:08:09 -0700 Subject: [PATCH 528/879] Improved examples[skip ci] --- examples/colbert/approximate.py | 6 +++--- examples/colbert/exact.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py index 623f913..14f1ce0 100644 --- a/examples/colbert/approximate.py +++ b/examples/colbert/approximate.py @@ -6,9 +6,6 @@ import psycopg import warnings -# ignore warnings from colbert -warnings.filterwarnings('ignore') - conn = psycopg.connect(dbname='pgvector_example', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') @@ -38,6 +35,9 @@ $$ LANGUAGE SQL """) +# ignore warnings from colbert +warnings.filterwarnings('ignore') + config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) diff --git a/examples/colbert/exact.py b/examples/colbert/exact.py index ceed2e3..c1ca236 100644 --- a/examples/colbert/exact.py +++ b/examples/colbert/exact.py @@ -4,9 +4,6 @@ import psycopg import warnings -# ignore warnings from colbert -warnings.filterwarnings('ignore') - conn = psycopg.connect(dbname='pgvector_example', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') @@ -32,6 +29,9 @@ $$ LANGUAGE SQL """) +# ignore warnings from colbert +warnings.filterwarnings('ignore') + config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) From 6ff9b8997e75632936230829bd557281c49e1891 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 8 Apr 2025 23:13:23 -0700 Subject: [PATCH 529/879] Updated ColBERT examples [skip ci] --- examples/colbert/approximate.py | 3 +-- examples/colbert/exact.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py index 14f1ce0..41f88b2 100644 --- a/examples/colbert/approximate.py +++ b/examples/colbert/approximate.py @@ -35,8 +35,7 @@ $$ LANGUAGE SQL """) -# ignore warnings from colbert -warnings.filterwarnings('ignore') +warnings.filterwarnings('ignore') # ignore warnings from colbert config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) diff --git a/examples/colbert/exact.py b/examples/colbert/exact.py index c1ca236..e6a2936 100644 --- a/examples/colbert/exact.py +++ b/examples/colbert/exact.py @@ -29,8 +29,7 @@ $$ LANGUAGE SQL """) -# ignore warnings from colbert -warnings.filterwarnings('ignore') +warnings.filterwarnings('ignore') # ignore warnings from colbert config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) From 3f9e9a20b9f08033e7dc4e61ff4c43b34951d2ec Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 15 Apr 2025 10:01:51 -0700 Subject: [PATCH 530/879] Updated Cohere example [skip ci] --- examples/cohere/example.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 393d1e0..5ef4eec 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -9,12 +9,12 @@ register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))') def embed(input, input_type): - co = cohere.Client() - response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary']) + co = cohere.ClientV2() + response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] From 713590a798190b34f4c43c4b097dbd61455113c3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 21 Apr 2025 03:09:36 -0700 Subject: [PATCH 531/879] Fixed SparseVector constructor for SciPy sparse matrices - fixes #127 --- CHANGELOG.md | 4 ++++ pgvector/sparsevec.py | 2 +- tests/test_sparse_vector.py | 14 +++++++++++++- 3 files changed, 18 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d0e2730..1bbd73c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.4.1 (unreleased) + +- Fixed `SparseVector` constructor for SciPy sparse matrices + ## 0.4.0 (2025-03-15) - Added top-level `pgvector` package diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 8df2dfd..895fbd0 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -85,7 +85,7 @@ def _from_sparse(self, value): if hasattr(value, 'coords'): # scipy 1.13+ - self._indices = value.coords[0].tolist() + self._indices = value.coords[-1].tolist() else: self._indices = value.col.tolist() self._values = value.data.tolist() diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index dff03dd..933cfff 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,7 +1,7 @@ import numpy as np from pgvector import SparseVector import pytest -from scipy.sparse import coo_array +from scipy.sparse import coo_array, csr_array, csr_matrix from struct import pack @@ -49,6 +49,18 @@ def test_dok_array(self): assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + def test_csr_array(self): + arr = csr_array(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(arr) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + + def test_csr_matrix(self): + mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(mat) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + def test_repr(self): assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' From 76afd8ec3013ac58bb6cc60a1b5b705f157ea18b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 21 Apr 2025 03:15:41 -0700 Subject: [PATCH 532/879] Added test for coo_matrix --- tests/test_sparse_vector.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 933cfff..0cf0a72 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,7 +1,7 @@ import numpy as np from pgvector import SparseVector import pytest -from scipy.sparse import coo_array, csr_array, csr_matrix +from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix from struct import pack @@ -43,6 +43,12 @@ def test_coo_array_dimensions(self): SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) assert str(error.value) == 'extra argument' + def test_coo_matrix(self): + mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(mat) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + def test_dok_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) From 809287f92847e1c609a9c395891da76f674379ea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 21 Apr 2025 03:20:20 -0700 Subject: [PATCH 533/879] Fixed CI --- tests/test_sparse_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 0cf0a72..d580f32 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -56,7 +56,7 @@ def test_dok_array(self): assert vec.indices() == [0, 2, 4] def test_csr_array(self): - arr = csr_array(np.array([1, 0, 2, 0, 3, 0])) + arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] From f9d2073df5cce39f0691ead6f9e030516baac7f8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 26 Apr 2025 11:56:00 -0700 Subject: [PATCH 534/879] Version bump to 0.4.1 [skip ci] --- CHANGELOG.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bbd73c..0ed80e3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.4.1 (unreleased) +## 0.4.1 (2025-04-26) - Fixed `SparseVector` constructor for SciPy sparse matrices diff --git a/pyproject.toml b/pyproject.toml index b889f4b..9395f9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.4.0" +version = "0.4.1" description = "pgvector support for Python" readme = "README.md" authors = [ From 7793bb069942fbcc2e77cf7349c59ffc28d8b6e0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 21 May 2025 18:16:18 -0700 Subject: [PATCH 535/879] Improved example [skip ci] --- examples/loading/example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/loading/example.py b/examples/loading/example.py index 0702129..7f3dce8 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -25,12 +25,12 @@ copy.set_types(['vector']) for i, embedding in enumerate(embeddings): + copy.write_row([embedding]) + # show progress if i % 10000 == 0: print('.', end='', flush=True) - copy.write_row([embedding]) - print('\nSuccess!') # create any indexes *after* loading initial data (skipping for this example) From 91088aacfadad37c9b8ea533b1e2b16b08d12ac4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 8 Jun 2025 16:28:24 -0700 Subject: [PATCH 536/879] Updated readme [skip ci] --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 24d9bb9..7c302b1 100644 --- a/README.md +++ b/README.md @@ -409,7 +409,7 @@ Enable the extension conn.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.psycopg import register_vector @@ -472,7 +472,7 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection or cursor +Register the types with your connection or cursor ```python from pgvector.psycopg2 import register_vector @@ -518,7 +518,7 @@ Enable the extension await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.asyncpg import register_vector @@ -572,7 +572,7 @@ Enable the extension conn.run('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.pg8000 import register_vector From ee3e71ca2c07a12a8332a3877c0ce14adc9a5da8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 16 Jun 2025 15:36:53 -0700 Subject: [PATCH 537/879] Updated format for license identifier --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9395f9e..0cfa183 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ readme = "README.md" authors = [ {name = "Andrew Kane", email = "andrew@ankane.org"} ] -license = {text = "MIT"} +license = "MIT" requires-python = ">= 3.9" dependencies = [ "numpy" From 33dee606229489c9ffb0cb5a1cd72bd4705ac618 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Sep 2025 16:35:08 -0700 Subject: [PATCH 538/879] Added support for str objects for bit type with SQLAlchemy - #137 Co-authored-by: Giacomo rua --- CHANGELOG.md | 4 ++++ pgvector/sqlalchemy/bit.py | 12 ++++++++++++ tests/test_sqlalchemy.py | 5 +++++ 3 files changed, 21 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ed80e3..f219b22 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.4.2 (unreleased) + +- Added support for `str` objects for `bit` type with SQLAlchemy + ## 0.4.1 (2025-04-26) - Fixed `SparseVector` constructor for SciPy sparse matrices diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 0f83f3c..1ea85c3 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -14,6 +14,18 @@ def get_col_spec(self, **kw): return 'BIT' return 'BIT(%d)' % self.length + def bind_processor(self, dialect): + if dialect.__class__.__name__ == 'PGDialect_asyncpg': + import asyncpg + + def process(value): + if isinstance(value, str): + return asyncpg.BitString(value) + return value + return process + else: + return super().bind_processor(dialect) + class comparator_factory(UserDefinedType.Comparator): def hamming_distance(self, other): return self.op('<~>', return_type=Float)(other) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 5aec977..cd7bad8 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -596,6 +596,11 @@ async def test_bit(self, engine): item = await session.get(Item, 1) assert item.binary_embedding == embedding + if engine == asyncpg_engine: + session.add(Item(id=2, binary_embedding='101')) + item = await session.get(Item, 2) + assert item.binary_embedding == embedding + await engine.dispose() @pytest.mark.asyncio From dc9a8f959995f009649fd230139ca41193e0a801 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 3 Sep 2025 12:04:21 -0700 Subject: [PATCH 539/879] Added test for binary quantization with re-ranking --- tests/test_sqlalchemy.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index cd7bad8..702eee1 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -528,6 +528,22 @@ def test_binary_quantize(self, engine): items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] + def test_binary_quantize_reranking(self, engine): + # recreate index (could also vacuum table) + binary_quantize_index.drop(setup_engine) + binary_quantize_index.create(setup_engine) + + with Session(engine) as session: + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + subquery = session.query(Item).order_by(distance).limit(20).subquery() + items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all() + assert [v.id for v in items] == [2, 3, 1] + @pytest.mark.parametrize('engine', array_engines) class TestSqlalchemyArray: From caf1a2e0dd7a1ba2ad0ca9f09b50516dcfffcdeb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 3 Sep 2025 12:09:18 -0700 Subject: [PATCH 540/879] Added docs for binary quantization with SQLAlchemy [skip ci] --- README.md | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/README.md b/README.md index 7c302b1..bfec8bb 100644 --- a/README.md +++ b/README.md @@ -271,6 +271,38 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2]) session.scalars(select(Item).order_by(order).limit(5)) ``` +#### Binary Quantization + +Use expression indexing for binary quantization + +```python +from pgvector.sqlalchemy import BIT +from sqlalchemy.sql import func + +index = Index( + 'my_index', + func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'bit_hamming_ops'} +) +``` + +Get the nearest neighbors by Hamming distance + +```python +order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) +session.scalars(select(Item).order_by(order).limit(5)) +``` + +Re-rank by the original vectors for better recall + +```python +order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) +subquery = session.query(Item).order_by(order).limit(20).subquery() +session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5)) +``` + #### Arrays Add an array column From c820a53bfb46196551de3c3f59f81b192d890574 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 3 Sep 2025 12:11:38 -0700 Subject: [PATCH 541/879] Simplified examples [skip ci] --- README.md | 2 -- tests/test_sqlalchemy.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/README.md b/README.md index bfec8bb..7cff86c 100644 --- a/README.md +++ b/README.md @@ -259,7 +259,6 @@ index = Index( 'my_index', func.cast(Item.embedding, HALFVEC(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'halfvec_l2_ops'} ) ``` @@ -283,7 +282,6 @@ index = Index( 'my_index', func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'bit_hamming_ops'} ) ``` diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 702eee1..c59c12e 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -103,7 +103,6 @@ class Item(Base): 'sqlalchemy_orm_half_precision_index', func.cast(Item.embedding, HALFVEC(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'halfvec_l2_ops'} ) half_precision_index.create(setup_engine) @@ -112,7 +111,6 @@ class Item(Base): 'sqlalchemy_orm_binary_quantize_index', func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'bit_hamming_ops'} ) binary_quantize_index.create(setup_engine) From 1a72b7571adf3325174b383aca85bfb3a5b925fa Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Sep 2025 00:12:48 -0700 Subject: [PATCH 542/879] Updated pgvector on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 4d4e8ed..d943ea0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.1 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From e211ba4029f204734f0c001fbb90f6a594d561ae Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 9 Oct 2025 23:19:50 -0700 Subject: [PATCH 543/879] Test with Python 3.14 on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d943ea0..52ab712 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.13, 3.9] + python: [3.14, 3.9] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 From e2986daf2b1533cc2c849f7e39350e31d57ac325 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 4 Dec 2025 16:47:23 -0800 Subject: [PATCH 544/879] Added support for Django 6 --- CHANGELOG.md | 1 + pgvector/django/extensions.py | 9 +++++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f219b22..62da0bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.4.2 (unreleased) +- Added support for Django 6 - Added support for `str` objects for `bit` type with SQLAlchemy ## 0.4.1 (2025-04-26) diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py index 0573f72..1d04739 100644 --- a/pgvector/django/extensions.py +++ b/pgvector/django/extensions.py @@ -1,6 +1,11 @@ +from django import VERSION from django.contrib.postgres.operations import CreateExtension class VectorExtension(CreateExtension): - def __init__(self): - self.name = 'vector' + if VERSION[0] >= 6: + def __init__(self, hints=None): + super().__init__('vector', hints=hints) + else: + def __init__(self): + self.name = 'vector' From 674f5ba3410c873d49f50fa9725b95d9db50c674 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 4 Dec 2025 16:50:18 -0800 Subject: [PATCH 545/879] Updated checkout action [skip ci] --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 52ab712..34f15d5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -8,7 +8,7 @@ jobs: matrix: python: [3.14, 3.9] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python }} From 2968f258f9486531bd1340cbda4ff8fcaf06cdc1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 4 Dec 2025 17:06:01 -0800 Subject: [PATCH 546/879] Version bump to 0.4.2 [skip ci] --- CHANGELOG.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 62da0bb..745335f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.4.2 (unreleased) +## 0.4.2 (2025-12-04) - Added support for Django 6 - Added support for `str` objects for `bit` type with SQLAlchemy diff --git a/pyproject.toml b/pyproject.toml index 0cfa183..6f91e04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.4.1" +version = "0.4.2" description = "pgvector support for Python" readme = "README.md" authors = [ From 05387da3c5ce0dc9f1d6ef238dcae118aa8176ea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Dec 2025 16:31:46 -0800 Subject: [PATCH 547/879] Updated examples [skip ci] --- README.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 7cff86c..a208ae0 100644 --- a/README.md +++ b/README.md @@ -177,10 +177,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(Base): - embedding = mapped_column(Vector(3)) + embedding = mapped_column(VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -306,11 +306,11 @@ session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([ Add an array column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(Vector(3))) + embeddings = mapped_column(ARRAY(VECTOR(3))) ``` And register the types with the underlying driver @@ -359,10 +359,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(SQLModel, table=True): - embedding: Any = Field(sa_type=Vector(3)) + embedding: Any = Field(sa_type=VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` From 6d8db07f74fd95b3673fd8149f3f805a15788f48 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Dec 2025 16:34:08 -0800 Subject: [PATCH 548/879] Updated readme [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index a208ae0..36d6c06 100644 --- a/README.md +++ b/README.md @@ -345,7 +345,7 @@ from sqlalchemy import event @event.listens_for(engine, "connect") def connect(dbapi_connection, connection_record): - register_vector(dbapi_connection, arrays=True) + register_vector(dbapi_connection) ``` ## SQLModel From b34f1c994e843dd7468d600b0f0ff5dbb949ec61 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 6 Dec 2025 16:38:13 -0800 Subject: [PATCH 549/879] Updated examples [skip ci] --- examples/implicit/example.py | 6 +++--- examples/lightfm/example.py | 6 +++--- examples/surprise/example.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/examples/implicit/example.py b/examples/implicit/example.py index f70eb8c..2cbf7c6 100644 --- a/examples/implicit/example.py +++ b/examples/implicit/example.py @@ -1,6 +1,6 @@ import implicit from implicit.datasets.movielens import get_movielens -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session @@ -16,7 +16,7 @@ class User(Base): __tablename__ = 'user' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) class Item(Base): @@ -24,7 +24,7 @@ class Item(Base): id = mapped_column(Integer, primary_key=True) title = mapped_column(String) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) diff --git a/examples/lightfm/example.py b/examples/lightfm/example.py index fcb9027..65031c4 100644 --- a/examples/lightfm/example.py +++ b/examples/lightfm/example.py @@ -1,6 +1,6 @@ from lightfm import LightFM from lightfm.datasets import fetch_movielens -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Float, Integer, String from sqlalchemy.orm import declarative_base, mapped_column, Session @@ -16,7 +16,7 @@ class User(Base): __tablename__ = 'user' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) class Item(Base): @@ -24,7 +24,7 @@ class Item(Base): id = mapped_column(Integer, primary_key=True) title = mapped_column(String) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) bias = mapped_column(Float) diff --git a/examples/surprise/example.py b/examples/surprise/example.py index bd7d18d..e413bcf 100644 --- a/examples/surprise/example.py +++ b/examples/surprise/example.py @@ -1,4 +1,4 @@ -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer from sqlalchemy.orm import declarative_base, mapped_column, Session from surprise import Dataset, SVD @@ -15,14 +15,14 @@ class User(Base): __tablename__ = 'user' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + factors = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) From e1dda975bdb2635f273cc3e6cc1b9c01780cec00 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Dec 2025 12:20:29 -0800 Subject: [PATCH 550/879] Switched to getuser for tests --- tests/test_pg8000.py | 4 ++-- tests/test_sqlalchemy.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 4d3e474..61fbc4c 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -1,10 +1,10 @@ +from getpass import getuser import numpy as np -import os from pgvector import HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection -conn = Connection(os.environ["USER"], database='pgvector_python_test') +conn = Connection(getuser(), database='pgvector_python_test') conn.run('CREATE EXTENSION IF NOT EXISTS vector') conn.run('DROP TABLE IF EXISTS pg8000_items') diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index c59c12e..4e870cc 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,6 +1,6 @@ import asyncpg +from getpass import getuser import numpy as np -import os from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest @@ -28,7 +28,7 @@ def psycopg2_connect(dbapi_connection, connection_record): register_vector(dbapi_connection) -pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') +pg8000_engine = create_engine(f'postgresql+pg8000://{getuser()}@localhost/pgvector_python_test') if sqlalchemy_version > 1: psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') From 5392e2cd3200574829610efedef1678dafbaa4d7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 17 Dec 2025 13:13:06 -0800 Subject: [PATCH 551/879] Switched to dependency groups --- .github/workflows/build.yml | 2 +- README.md | 2 +- pyproject.toml | 15 +++++++++++++++ requirements.txt | 12 ------------ 4 files changed, 17 insertions(+), 14 deletions(-) delete mode 100644 requirements.txt diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 34f15d5..7d0225b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -12,7 +12,7 @@ jobs: - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python }} - - run: pip install -r requirements.txt + - run: pip install --group dev - uses: ankane/setup-postgres@v1 with: database: pgvector_python_test diff --git a/README.md b/README.md index 36d6c06..7671c00 100644 --- a/README.md +++ b/README.md @@ -807,7 +807,7 @@ To get started with development: ```sh git clone https://github.com/pgvector/pgvector-python.git cd pgvector-python -pip install -r requirements.txt +pip install --group dev createdb pgvector_python_test pytest ``` diff --git a/pyproject.toml b/pyproject.toml index 6f91e04..5716d05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,21 @@ dependencies = [ [project.urls] Homepage = "https://github.com/pgvector/pgvector-python" +[dependency-groups] +dev = [ + "asyncpg", + "Django", + "peewee", + "pg8000", + "psycopg[binary,pool]", + "psycopg2-binary", + "pytest", + "pytest-asyncio", + "scipy", + "SQLAlchemy[asyncio]>=2", + "sqlmodel>=0.0.12" +] + [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index a13be06..0000000 --- a/requirements.txt +++ /dev/null @@ -1,12 +0,0 @@ -asyncpg -Django -numpy -peewee -pg8000 -psycopg[binary,pool] -psycopg2-binary -pytest -pytest-asyncio -scipy -SQLAlchemy[asyncio]>=2 -sqlmodel>=0.0.12 From 6c1fa981ad624f7ded69e3a60df8fa303466e19c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 18 Dec 2025 21:52:46 -0800 Subject: [PATCH 552/879] Updated example to pyproject.toml [skip ci] --- README.md | 2 +- examples/citus/pyproject.toml | 11 +++++++++++ examples/citus/requirements.txt | 3 --- examples/cohere/pyproject.toml | 11 +++++++++++ examples/cohere/requirements.txt | 3 --- examples/colbert/pyproject.toml | 12 ++++++++++++ examples/colbert/requirements.txt | 4 ---- examples/colpali/pyproject.toml | 12 ++++++++++++ examples/colpali/requirements.txt | 4 ---- examples/gensim/pyproject.toml | 13 +++++++++++++ examples/gensim/requirements.txt | 5 ----- examples/hybrid_search/pyproject.toml | 11 +++++++++++ examples/hybrid_search/requirements.txt | 3 --- examples/image_search/pyproject.toml | 14 ++++++++++++++ examples/image_search/requirements.txt | 6 ------ examples/imagehash/pyproject.toml | 13 +++++++++++++ examples/imagehash/requirements.txt | 5 ----- examples/implicit/pyproject.toml | 13 +++++++++++++ examples/implicit/requirements.txt | 5 ----- examples/lightfm/pyproject.toml | 12 ++++++++++++ examples/lightfm/requirements.txt | 4 ---- examples/loading/pyproject.toml | 11 +++++++++++ examples/loading/requirements.txt | 3 --- examples/openai/pyproject.toml | 11 +++++++++++ examples/openai/requirements.txt | 3 --- examples/rag/pyproject.toml | 11 +++++++++++ examples/rag/requirements.txt | 3 --- examples/rdkit/pyproject.toml | 11 +++++++++++ examples/rdkit/requirements.txt | 3 --- examples/sentence_transformers/pyproject.toml | 11 +++++++++++ examples/sentence_transformers/requirements.txt | 3 --- examples/sparse_search/pyproject.toml | 13 +++++++++++++ examples/sparse_search/requirements.txt | 5 ----- examples/surprise/pyproject.toml | 12 ++++++++++++ examples/surprise/requirements.txt | 4 ---- 35 files changed, 203 insertions(+), 67 deletions(-) create mode 100644 examples/citus/pyproject.toml delete mode 100644 examples/citus/requirements.txt create mode 100644 examples/cohere/pyproject.toml delete mode 100644 examples/cohere/requirements.txt create mode 100644 examples/colbert/pyproject.toml delete mode 100644 examples/colbert/requirements.txt create mode 100644 examples/colpali/pyproject.toml delete mode 100644 examples/colpali/requirements.txt create mode 100644 examples/gensim/pyproject.toml delete mode 100644 examples/gensim/requirements.txt create mode 100644 examples/hybrid_search/pyproject.toml delete mode 100644 examples/hybrid_search/requirements.txt create mode 100644 examples/image_search/pyproject.toml delete mode 100644 examples/image_search/requirements.txt create mode 100644 examples/imagehash/pyproject.toml delete mode 100644 examples/imagehash/requirements.txt create mode 100644 examples/implicit/pyproject.toml delete mode 100644 examples/implicit/requirements.txt create mode 100644 examples/lightfm/pyproject.toml delete mode 100644 examples/lightfm/requirements.txt create mode 100644 examples/loading/pyproject.toml delete mode 100644 examples/loading/requirements.txt create mode 100644 examples/openai/pyproject.toml delete mode 100644 examples/openai/requirements.txt create mode 100644 examples/rag/pyproject.toml delete mode 100644 examples/rag/requirements.txt create mode 100644 examples/rdkit/pyproject.toml delete mode 100644 examples/rdkit/requirements.txt create mode 100644 examples/sentence_transformers/pyproject.toml delete mode 100644 examples/sentence_transformers/requirements.txt create mode 100644 examples/sparse_search/pyproject.toml delete mode 100644 examples/sparse_search/requirements.txt create mode 100644 examples/surprise/pyproject.toml delete mode 100644 examples/surprise/requirements.txt diff --git a/README.md b/README.md index 7671c00..95d5fbe 100644 --- a/README.md +++ b/README.md @@ -816,7 +816,7 @@ To run an example: ```sh cd examples/loading -pip install -r requirements.txt +pip install --group dev createdb pgvector_example python3 example.py ``` diff --git a/examples/citus/pyproject.toml b/examples/citus/pyproject.toml new file mode 100644 index 0000000..ee40a36 --- /dev/null +++ b/examples/citus/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/citus/requirements.txt b/examples/citus/requirements.txt deleted file mode 100644 index 1cf8ee9..0000000 --- a/examples/citus/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -numpy -pgvector -psycopg[binary] diff --git a/examples/cohere/pyproject.toml b/examples/cohere/pyproject.toml new file mode 100644 index 0000000..f0c88b7 --- /dev/null +++ b/examples/cohere/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "cohere", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/cohere/requirements.txt b/examples/cohere/requirements.txt deleted file mode 100644 index 22fd056..0000000 --- a/examples/cohere/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -cohere -pgvector -psycopg[binary] diff --git a/examples/colbert/pyproject.toml b/examples/colbert/pyproject.toml new file mode 100644 index 0000000..face4d2 --- /dev/null +++ b/examples/colbert/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "colbert-ai", + "pgvector", + "psycopg[binary]", + "transformers==4.49.0" +] diff --git a/examples/colbert/requirements.txt b/examples/colbert/requirements.txt deleted file mode 100644 index 54b2cb9..0000000 --- a/examples/colbert/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -colbert-ai -pgvector -psycopg[binary] -transformers==4.49.0 diff --git a/examples/colpali/pyproject.toml b/examples/colpali/pyproject.toml new file mode 100644 index 0000000..23fb23f --- /dev/null +++ b/examples/colpali/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "colpali-engine", + "datasets", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/colpali/requirements.txt b/examples/colpali/requirements.txt deleted file mode 100644 index 4cf770d..0000000 --- a/examples/colpali/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -colpali-engine -datasets -pgvector -psycopg[binary] diff --git a/examples/gensim/pyproject.toml b/examples/gensim/pyproject.toml new file mode 100644 index 0000000..7a33423 --- /dev/null +++ b/examples/gensim/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "gensim", + "numpy", + "pgvector", + "psycopg[binary]", + "scipy<1.13" +] diff --git a/examples/gensim/requirements.txt b/examples/gensim/requirements.txt deleted file mode 100644 index 15411cd..0000000 --- a/examples/gensim/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -gensim -numpy -pgvector -psycopg[binary] -scipy<1.13 diff --git a/examples/hybrid_search/pyproject.toml b/examples/hybrid_search/pyproject.toml new file mode 100644 index 0000000..b5a904a --- /dev/null +++ b/examples/hybrid_search/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/hybrid_search/requirements.txt b/examples/hybrid_search/requirements.txt deleted file mode 100644 index 237dcd1..0000000 --- a/examples/hybrid_search/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -sentence-transformers diff --git a/examples/image_search/pyproject.toml b/examples/image_search/pyproject.toml new file mode 100644 index 0000000..7644382 --- /dev/null +++ b/examples/image_search/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "matplotlib", + "pgvector", + "psycopg[binary]", + "torch", + "torchvision", + "tqdm" +] diff --git a/examples/image_search/requirements.txt b/examples/image_search/requirements.txt deleted file mode 100644 index 3d82365..0000000 --- a/examples/image_search/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -matplotlib -pgvector -psycopg[binary] -torch -torchvision -tqdm diff --git a/examples/imagehash/pyproject.toml b/examples/imagehash/pyproject.toml new file mode 100644 index 0000000..cf06c2b --- /dev/null +++ b/examples/imagehash/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "datasets", + "imagehash", + "matplotlib", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/imagehash/requirements.txt b/examples/imagehash/requirements.txt deleted file mode 100644 index e3971e6..0000000 --- a/examples/imagehash/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -datasets -imagehash -matplotlib -pgvector -psycopg[binary] diff --git a/examples/implicit/pyproject.toml b/examples/implicit/pyproject.toml new file mode 100644 index 0000000..c03b187 --- /dev/null +++ b/examples/implicit/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "h5py", + "implicit", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/implicit/requirements.txt b/examples/implicit/requirements.txt deleted file mode 100644 index 424abbd..0000000 --- a/examples/implicit/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -h5py -implicit -pgvector -psycopg[binary] -SQLAlchemy diff --git a/examples/lightfm/pyproject.toml b/examples/lightfm/pyproject.toml new file mode 100644 index 0000000..c202058 --- /dev/null +++ b/examples/lightfm/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "lightfm", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/lightfm/requirements.txt b/examples/lightfm/requirements.txt deleted file mode 100644 index cfa5f51..0000000 --- a/examples/lightfm/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -lightfm -pgvector -psycopg[binary] -SQLAlchemy diff --git a/examples/loading/pyproject.toml b/examples/loading/pyproject.toml new file mode 100644 index 0000000..ee40a36 --- /dev/null +++ b/examples/loading/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/loading/requirements.txt b/examples/loading/requirements.txt deleted file mode 100644 index 1cf8ee9..0000000 --- a/examples/loading/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -numpy -pgvector -psycopg[binary] diff --git a/examples/openai/pyproject.toml b/examples/openai/pyproject.toml new file mode 100644 index 0000000..3e6661a --- /dev/null +++ b/examples/openai/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "openai", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/openai/requirements.txt b/examples/openai/requirements.txt deleted file mode 100644 index 18587e2..0000000 --- a/examples/openai/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -openai -pgvector -psycopg[binary] diff --git a/examples/rag/pyproject.toml b/examples/rag/pyproject.toml new file mode 100644 index 0000000..fa0dcfd --- /dev/null +++ b/examples/rag/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "ollama", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/rag/requirements.txt b/examples/rag/requirements.txt deleted file mode 100644 index 4eb5864..0000000 --- a/examples/rag/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -ollama -pgvector -psycopg[binary] diff --git a/examples/rdkit/pyproject.toml b/examples/rdkit/pyproject.toml new file mode 100644 index 0000000..f8c035a --- /dev/null +++ b/examples/rdkit/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "rdkit" +] diff --git a/examples/rdkit/requirements.txt b/examples/rdkit/requirements.txt deleted file mode 100644 index 85a3e4f..0000000 --- a/examples/rdkit/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -rdkit diff --git a/examples/sentence_transformers/pyproject.toml b/examples/sentence_transformers/pyproject.toml new file mode 100644 index 0000000..b5a904a --- /dev/null +++ b/examples/sentence_transformers/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/sentence_transformers/requirements.txt b/examples/sentence_transformers/requirements.txt deleted file mode 100644 index 237dcd1..0000000 --- a/examples/sentence_transformers/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -sentence-transformers diff --git a/examples/sparse_search/pyproject.toml b/examples/sparse_search/pyproject.toml new file mode 100644 index 0000000..7927c34 --- /dev/null +++ b/examples/sparse_search/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]", + "torch", + "transformers" +] diff --git a/examples/sparse_search/requirements.txt b/examples/sparse_search/requirements.txt deleted file mode 100644 index 3de81c7..0000000 --- a/examples/sparse_search/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -numpy -pgvector -psycopg[binary] -torch -transformers diff --git a/examples/surprise/pyproject.toml b/examples/surprise/pyproject.toml new file mode 100644 index 0000000..94c6f13 --- /dev/null +++ b/examples/surprise/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.9" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "scikit-surprise", + "SQLAlchemy" +] diff --git a/examples/surprise/requirements.txt b/examples/surprise/requirements.txt deleted file mode 100644 index cb2dca4..0000000 --- a/examples/surprise/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -pgvector -psycopg[binary] -scikit-surprise -SQLAlchemy From 016478953c094d908096e775c4366309601a2a38 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 19 Dec 2025 16:35:14 -0800 Subject: [PATCH 553/879] Updated gitignore [skip ci] --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index c55ff44..5556c9f 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ venv/ *.pyc __pycache__ .pytest_cache/ +*.lock examples/rag/README.md From a3520fef143f01a93bc2e70b00070080d2996c29 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 7 Jan 2026 15:26:48 -0800 Subject: [PATCH 554/879] Updated CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7d0225b..7f46dd6 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -9,7 +9,7 @@ jobs: python: [3.14, 3.9] steps: - uses: actions/checkout@v5 - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python }} - run: pip install --group dev From 8bd3c61fc0954fc7d0e1060ce81a91315ca15010 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 22 Jan 2026 10:20:14 -0800 Subject: [PATCH 555/879] Added tests for subqueries with SQLAlchemy - closes #147 --- tests/test_sqlalchemy.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 4e870cc..7558942 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -396,6 +396,20 @@ def test_sparsevec_l1_distance_orm(self, engine): items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] + def test_subquery(self, engine): + create_items() + with Session(engine) as session: + subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() + items = session.query(Item).order_by(Item.embedding.l2_distance(subquery)).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_subquery_orm(self, engine): + create_items() + with Session(engine) as session: + subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() + items = session.scalars(select(Item).order_by(Item.embedding.l2_distance(subquery))) + assert [v.id for v in items] == [1, 3, 2] + def test_filter(self, engine): create_items() with Session(engine) as session: From 5b06584aaf1588b04441e622e28bd9b0aa09924a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 28 Feb 2026 18:44:10 -0800 Subject: [PATCH 556/879] Updated pgvector on CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7f46dd6..85db5b6 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.8.1 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From f606ee2ae8d188600f85a1d20c01524f6c3e8828 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 18 May 2026 20:15:56 -0700 Subject: [PATCH 557/879] Updated checkout action --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 85db5b6..fb4badc 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -8,7 +8,7 @@ jobs: matrix: python: [3.14, 3.9] steps: - - uses: actions/checkout@v5 + - uses: actions/checkout@v6 - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python }} From e835b44219a3faf0de160e76a6d866358f53e182 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 11:52:20 -0700 Subject: [PATCH 558/879] Dropped support for Python < 3.10 [skip ci] --- .github/workflows/build.yml | 2 +- CHANGELOG.md | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index fb4badc..6c7cf3b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.14, 3.9] + python: [3.14, 3.10] steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 diff --git a/CHANGELOG.md b/CHANGELOG.md index 745335f..85f9390 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.5.0 (unreleased) + +- Dropped support for Python < 3.10 + ## 0.4.2 (2025-12-04) - Added support for Django 6 From 7b5663fa3fec88cd7c7a1775a63b297451d39ae5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 12:07:28 -0700 Subject: [PATCH 559/879] Updated lint [skip ci] --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index f0831c4..7462142 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ .PHONY: lint build publish clean lint: - pycodestyle . --ignore=E501 + pycodestyle . --ignore=E501 --exclude=.venv build: python3 -m build From fa5fd61f34bde8b8de99852549d0094b44bf3f40 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 12:08:12 -0700 Subject: [PATCH 560/879] Dropped support for SQLAlchemy < 2 --- .github/workflows/build.yml | 3 -- CHANGELOG.md | 1 + tests/test_sqlalchemy.py | 82 +++++++++++++++++-------------------- 3 files changed, 38 insertions(+), 48 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 6c7cf3b..5bac05c 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,6 +24,3 @@ jobs: make sudo make install - run: pytest - - - run: pip install "SQLAlchemy<2" -U - - run: pytest tests/test_sqlalchemy.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 85f9390..b5923fa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.5.0 (unreleased) - Dropped support for Python < 3.10 +- Dropped support for SQLAlchemy < 2 ## 0.4.2 (2025-12-04) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 7558942..8a260d8 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -6,18 +6,11 @@ import pytest from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY from sqlalchemy.exc import StatementError +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.ext.automap import automap_base -from sqlalchemy.orm import declarative_base, Session +from sqlalchemy.orm import declarative_base, mapped_column, Session from sqlalchemy.sql import func -try: - from sqlalchemy.orm import mapped_column - from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine - sqlalchemy_version = 2 -except ImportError: - mapped_column = Column - sqlalchemy_version = 1 - psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -30,42 +23,41 @@ def psycopg2_connect(dbapi_connection, connection_record): pg8000_engine = create_engine(f'postgresql+pg8000://{getuser()}@localhost/pgvector_python_test') -if sqlalchemy_version > 1: - psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') - psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') - - @event.listens_for(psycopg_type_engine, "connect") - def psycopg_connect(dbapi_connection, connection_record): - from pgvector.psycopg import register_vector - register_vector(dbapi_connection) - - psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') - psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') - - @event.listens_for(psycopg_async_type_engine.sync_engine, "connect") - def psycopg_async_connect(dbapi_connection, connection_record): - from pgvector.psycopg import register_vector_async - dbapi_connection.run_async(register_vector_async) - - asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') - asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') - - @event.listens_for(asyncpg_type_engine.sync_engine, "connect") - def asyncpg_connect(dbapi_connection, connection_record): - from pgvector.asyncpg import register_vector - dbapi_connection.run_async(register_vector) - -engines = [psycopg2_engine, psycopg2_type_engine, pg8000_engine] -array_engines = [psycopg2_type_engine] -async_engines = [] -async_array_engines = [] - -if sqlalchemy_version > 1: - engines += [psycopg_engine, psycopg_type_engine] - array_engines += [psycopg_type_engine] - # TODO support asyncpg_type_engine - async_engines += [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine] - async_array_engines += [psycopg_async_type_engine, asyncpg_engine] +psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') +psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + + +@event.listens_for(psycopg_type_engine, "connect") +def psycopg_connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector + register_vector(dbapi_connection) + + +psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') +psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + + +@event.listens_for(psycopg_async_type_engine.sync_engine, "connect") +def psycopg_async_connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector_async + dbapi_connection.run_async(register_vector_async) + + +asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') +asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + + +@event.listens_for(asyncpg_type_engine.sync_engine, "connect") +def asyncpg_connect(dbapi_connection, connection_record): + from pgvector.asyncpg import register_vector + dbapi_connection.run_async(register_vector) + + +engines = [psycopg_engine, psycopg_type_engine, psycopg2_engine, psycopg2_type_engine, pg8000_engine] +array_engines = [psycopg_type_engine, psycopg2_type_engine] +# TODO support asyncpg_type_engine +async_engines = [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine] +async_array_engines = [psycopg_async_type_engine, asyncpg_engine] setup_engine = engines[0] with Session(setup_engine) as session: From 751e10572f22d8cf8b85464eee65837ee94f6feb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 12:11:36 -0700 Subject: [PATCH 561/879] Fixed CI --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 5bac05c..1ed395f 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.14, 3.10] + python: [3.14, "3.10"] steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 From 68e1a80f08cf0e88604e656da9a5aefa8aa4d90b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 12:30:11 -0700 Subject: [PATCH 562/879] Improved typechecking for psycopg tests [skip ci] --- tests/test_psycopg.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 698b34f..fb6aa4c 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -29,73 +29,73 @@ def test_vector(self): def test_vector_binary_format(self): embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert np.array_equal(res, embedding) def test_vector_text_format(self): embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert np.array_equal(res, embedding) def test_vector_binary_format_correct(self): embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector::text', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] assert res == '[1.5,2,3]' def test_vector_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous - res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert np.array_equal(res, [3, 2, 1.5]) def test_vector_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous - res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] assert np.array_equal(res, [3, 2, 1.5]) def test_vector_class_binary_format(self): embedding = Vector([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert np.array_equal(res, [1.5, 2, 3]) def test_vector_class_text_format(self): embedding = Vector([1.5, 2, 3]) - res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert np.array_equal(res, [1.5, 2, 3]) def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + res = next(conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id'))[0] assert res == HalfVector([1.5, 2, 3]) def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) - res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::halfvec', (embedding,), binary=True))[0] assert res == HalfVector([1.5, 2, 3]) def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) - res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::halfvec', (embedding,)))[0] assert res == HalfVector([1.5, 2, 3]) def test_bit(self): embedding = Bit([True, False, True]) conn.execute('INSERT INTO psycopg_items (binary_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + res = next(conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id'))[0] assert res == '101' def test_bit_binary_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) - res = conn.execute('SELECT %b::bit(9)', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::bit(9)', (embedding,), binary=True))[0] assert repr(Bit.from_binary(res)) == 'Bit(010100001)' def test_bit_text_format(self): embedding = Bit([False, True, False, True, False, False, False, False, True]) - res = conn.execute('SELECT %t::bit(9)', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::bit(9)', (embedding,)))[0] assert res == '010100001' assert repr(Bit(res)) == 'Bit(010100001)' @@ -103,17 +103,17 @@ def test_sparsevec(self): embedding = SparseVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + res = next(conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id'))[0] assert res == SparseVector([1.5, 2, 3]) def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) - res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::sparsevec', (embedding,), binary=True))[0] assert res == embedding def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) - res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::sparsevec', (embedding,)))[0] assert res == embedding def test_text_copy_from(self): @@ -170,7 +170,7 @@ def test_vector_array(self): embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] conn.execute('INSERT INTO psycopg_items (embeddings) VALUES (%s)', (embeddings,)) - res = conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id').fetchone() + res = next(conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id')) assert np.array_equal(res[0][0], embeddings[0]) assert np.array_equal(res[0][1], embeddings[1]) From 87d09131c95c1af4a9c81fd0cf687b13b535ee04 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 12:52:23 -0700 Subject: [PATCH 563/879] Improved typechecking for tests [skip ci] --- tests/test_asyncpg.py | 2 +- tests/test_sqlalchemy.py | 26 +++++++++++++------------- tests/test_sqlmodel.py | 8 ++++---- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 34d66a1..cabfc5a 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -64,7 +64,7 @@ async def test_bit(self): await register_vector(conn) - embedding = asyncpg.BitString('101') + embedding = asyncpg.BitString('101') # type: ignore await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 8a260d8..95b3da7 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -190,7 +190,7 @@ def test_vector(self, engine): with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert np.array_equal(item.embedding, [1, 2, 3]) def test_vector_l2_distance(self, engine): @@ -245,7 +245,7 @@ def test_halfvec(self, engine): with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self, engine): @@ -300,7 +300,7 @@ def test_bit(self, engine): with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.binary_embedding == '101' def test_bit_hamming_distance(self, engine): @@ -337,7 +337,7 @@ def test_sparsevec(self, engine): with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self, engine): @@ -560,7 +560,7 @@ def test_vector_array(self, engine): session.commit() # this fails if the driver does not cast arrays - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert np.array_equal(item.embeddings[0], [1, 2, 3]) assert np.array_equal(item.embeddings[1], [4, 5, 6]) @@ -570,7 +570,7 @@ def test_halfvec_array(self, engine): session.commit() # this fails if the driver does not cast arrays - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.half_embeddings == [HalfVector([1, 2, 3]), HalfVector([4, 5, 6])] @@ -587,7 +587,7 @@ async def test_vector(self, engine): async with session.begin(): embedding = np.array([1, 2, 3]) session.add(Item(id=1, embedding=embedding)) - item = await session.get(Item, 1) + item = await session.get_one(Item, 1) assert np.array_equal(item.embedding, embedding) await engine.dispose() @@ -600,7 +600,7 @@ async def test_halfvec(self, engine): async with session.begin(): embedding = [1, 2, 3] session.add(Item(id=1, half_embedding=embedding)) - item = await session.get(Item, 1) + item = await session.get_one(Item, 1) assert item.half_embedding == HalfVector(embedding) await engine.dispose() @@ -613,12 +613,12 @@ async def test_bit(self, engine): async with session.begin(): embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' session.add(Item(id=1, binary_embedding=embedding)) - item = await session.get(Item, 1) + item = await session.get_one(Item, 1) assert item.binary_embedding == embedding if engine == asyncpg_engine: session.add(Item(id=2, binary_embedding='101')) - item = await session.get(Item, 2) + item = await session.get_one(Item, 2) assert item.binary_embedding == embedding await engine.dispose() @@ -631,7 +631,7 @@ async def test_sparsevec(self, engine): async with session.begin(): embedding = [1, 2, 3] session.add(Item(id=1, sparse_embedding=embedding)) - item = await session.get(Item, 1) + item = await session.get_one(Item, 1) assert item.sparse_embedding == SparseVector(embedding) await engine.dispose() @@ -662,12 +662,12 @@ async def test_vector_array(self, engine): async with async_session() as session: async with session.begin(): session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) - item = await session.get(Item, 1) + item = await session.get_one(Item, 1) assert np.array_equal(item.embeddings[0], [1, 2, 3]) assert np.array_equal(item.embeddings[1], [4, 5, 6]) session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - item = await session.get(Item, 2) + item = await session.get_one(Item, 2) assert np.array_equal(item.embeddings[0], [1, 2, 3]) assert np.array_equal(item.embeddings[1], [4, 5, 6]) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index f4994f4..616beaa 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -75,7 +75,7 @@ def test_vector(self): with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert np.array_equal(item.embedding, np.array([1, 2, 3])) def test_vector_l2_distance(self): @@ -106,7 +106,7 @@ def test_halfvec(self): with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): @@ -137,7 +137,7 @@ def test_bit(self): with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.binary_embedding == '101' def test_bit_hamming_distance(self): @@ -156,7 +156,7 @@ def test_sparsevec(self): with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): From f083137e40524f7f370b3fe88088335081d316e7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 12:59:06 -0700 Subject: [PATCH 564/879] Improved typechecking for tests [skip ci] --- tests/test_sqlalchemy.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 95b3da7..f57edc4 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -428,11 +428,11 @@ def test_select_orm(self, engine): def test_avg(self, engine): with Session(engine) as session: - res = session.query(avg(Item.embedding)).first()[0] + res = session.query(avg(Item.embedding)).one()[0] assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.query(avg(Item.embedding)).first()[0] + res = session.query(avg(Item.embedding)).one()[0] assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) def test_avg_orm(self, engine): @@ -441,16 +441,16 @@ def test_avg_orm(self, engine): assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.scalars(select(avg(Item.embedding))).first() + res = session.scalars(select(avg(Item.embedding))).one() assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) def test_sum(self, engine): with Session(engine) as session: - res = session.query(sum(Item.embedding)).first()[0] + res = session.query(sum(Item.embedding)).one()[0] assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.query(sum(Item.embedding)).first()[0] + res = session.query(sum(Item.embedding)).one()[0] assert np.array_equal(res, np.array([5, 7, 9])) def test_sum_orm(self, engine): @@ -459,7 +459,7 @@ def test_sum_orm(self, engine): assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.scalars(select(sum(Item.embedding))).first() + res = session.scalars(select(sum(Item.embedding))).one() assert np.array_equal(res, np.array([5, 7, 9])) def test_bad_dimensions(self, engine): @@ -611,7 +611,7 @@ async def test_bit(self, engine): async with async_session() as session: async with session.begin(): - embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' + embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' # type: ignore session.add(Item(id=1, binary_embedding=embedding)) item = await session.get_one(Item, 1) assert item.binary_embedding == embedding @@ -645,7 +645,7 @@ async def test_avg(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = await session.scalars(select(avg(Item.embedding))) - assert np.array_equal(res.first(), [2.5, 3.5, 4.5]) + assert np.array_equal(res.one(), [2.5, 3.5, 4.5]) await engine.dispose() From 6fe7c083a78749b3e79323d8bbedc2d0ec5d84d4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 13:02:24 -0700 Subject: [PATCH 565/879] Improved typechecking for tests [skip ci] --- tests/test_sqlmodel.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 616beaa..5f74148 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -8,17 +8,17 @@ engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: - session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) + session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) # type: ignore class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) - half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) - binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) - sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) + embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) # type: ignore + half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) # type: ignore + binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) # type: ignore + sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore SQLModel.metadata.drop_all(engine) @@ -202,7 +202,7 @@ def test_vector_avg(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(avg(Item.embedding))).first() - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) + assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) # type: ignore def test_vector_sum(self): with Session(engine) as session: @@ -211,7 +211,7 @@ def test_vector_sum(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(sum(Item.embedding))).first() - assert np.array_equal(res, np.array([5, 7, 9])) + assert np.array_equal(res, np.array([5, 7, 9])) # type: ignore def test_halfvec_avg(self): with Session(engine) as session: From 93fabd2e5167e8076ed77f153f6c9b69e0caf6c0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 13:09:14 -0700 Subject: [PATCH 566/879] Fixed typechecking errors [skip ci] --- pgvector/django/bit.py | 2 +- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- pgvector/sqlalchemy/bit.py | 2 +- pgvector/sqlalchemy/halfvec.py | 4 ++-- pgvector/sqlalchemy/sparsevec.py | 4 ++-- pgvector/sqlalchemy/vector.py | 4 ++-- 8 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 2cc847a..29a31e1 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -21,7 +21,7 @@ def db_type(self, connection): return 'bit' return 'bit(%d)' % self.length - def formfield(self, **kwargs): + def formfield(self, **kwargs): # type: ignore return super().formfield(form_class=BitFormField, **kwargs) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 3aeb90f..36101ca 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -40,7 +40,7 @@ def get_prep_value(self, value): def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs): + def formfield(self, **kwargs): # type: ignore return super().formfield(form_class=HalfVectorFormField, **kwargs) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 580f27c..7f4db77 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -35,7 +35,7 @@ def get_prep_value(self, value): def value_to_string(self, obj): return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs): + def formfield(self, **kwargs): # type: ignore return super().formfield(form_class=SparseVectorFormField, **kwargs) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 861cfde..dde96f8 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -48,7 +48,7 @@ def run_validators(self, value): value = value.tolist() super().run_validators(value) - def formfield(self, **kwargs): + def formfield(self, **kwargs): # type: ignore return super().formfield(form_class=VectorFormField, **kwargs) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 1ea85c3..bbf822e 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -35,4 +35,4 @@ def jaccard_distance(self, other): # for reflection -ischema_names['bit'] = BIT +ischema_names['bit'] = BIT # type: ignore diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 10688b5..f51d1a6 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -25,7 +25,7 @@ def literal_processor(self, dialect): string_literal_processor = self._string._cached_literal_processor(dialect) def process(value): - return string_literal_processor(HalfVector._to_db(value, self.dim)) + return string_literal_processor(HalfVector._to_db(value, self.dim)) # type: ignore return process def result_processor(self, dialect, coltype): @@ -48,4 +48,4 @@ def l1_distance(self, other): # for reflection -ischema_names['halfvec'] = HALFVEC +ischema_names['halfvec'] = HALFVEC # type: ignore diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 0058679..6828bfd 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -25,7 +25,7 @@ def literal_processor(self, dialect): string_literal_processor = self._string._cached_literal_processor(dialect) def process(value): - return string_literal_processor(SparseVector._to_db(value, self.dim)) + return string_literal_processor(SparseVector._to_db(value, self.dim)) # type: ignore return process def result_processor(self, dialect, coltype): @@ -48,4 +48,4 @@ def l1_distance(self, other): # for reflection -ischema_names['sparsevec'] = SPARSEVEC +ischema_names['sparsevec'] = SPARSEVEC # type: ignore diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 5a1e11f..a9e6f22 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -25,7 +25,7 @@ def literal_processor(self, dialect): string_literal_processor = self._string._cached_literal_processor(dialect) def process(value): - return string_literal_processor(Vector._to_db(value, self.dim)) + return string_literal_processor(Vector._to_db(value, self.dim)) # type: ignore return process def result_processor(self, dialect, coltype): @@ -48,4 +48,4 @@ def l1_distance(self, other): # for reflection -ischema_names['vector'] = VECTOR +ischema_names['vector'] = VECTOR # type: ignore From 563ebdfa836e9b9e07b23b8046a298862447e12a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 13:15:54 -0700 Subject: [PATCH 567/879] Improved typechecking for tests [skip ci] --- Makefile | 3 +++ tests/test_django.py | 2 ++ tests/test_sqlmodel.py | 44 +++++++++++++++++++++--------------------- 3 files changed, 27 insertions(+), 22 deletions(-) diff --git a/Makefile b/Makefile index 7462142..ccf995a 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,9 @@ lint: pycodestyle . --ignore=E501 --exclude=.venv +check: + ty check pgvector tests + build: python3 -m build diff --git a/tests/test_django.py b/tests/test_django.py index 7a8a6eb..9e44401 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,3 +1,5 @@ +# type: ignore + import django from django.conf import settings from django.contrib.postgres.fields import ArrayField diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 5f74148..d67c735 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -26,7 +26,7 @@ class Item(SQLModel, table=True): index = Index( 'sqlmodel_index', - Item.embedding, + Item.embedding, # type: ignore postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} @@ -65,10 +65,10 @@ def test_orm(self): assert items[0].id == 1 assert items[1].id == 2 assert items[2].id == 3 - assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) - assert items[0].embedding.dtype == np.float32 - assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) - assert items[1].embedding.dtype == np.float32 + assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) # type: ignore + assert items[0].embedding.dtype == np.float32 # type: ignore + assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) # type: ignore + assert items[1].embedding.dtype == np.float32 # type: ignore assert items[2].embedding is None def test_vector(self): @@ -76,30 +76,30 @@ def test_vector(self): session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert np.array_equal(item.embedding, np.array([1, 2, 3])) # type: ignore def test_vector_l2_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_vector_max_inner_product(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_vector_cosine_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] def test_vector_l1_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_halfvec(self): @@ -112,25 +112,25 @@ def test_halfvec(self): def test_halfvec_l2_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_halfvec_max_inner_product(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_halfvec_cosine_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] def test_halfvec_l1_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_bit(self): @@ -143,13 +143,13 @@ def test_bit(self): def test_bit_hamming_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) + items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_bit_jaccard_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) + items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_sparsevec(self): @@ -162,37 +162,37 @@ def test_sparsevec(self): def test_sparsevec_l2_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_sparsevec_max_inner_product(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_sparsevec_cosine_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] def test_sparsevec_l1_distance(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_filter(self): create_items() with Session(engine) as session: - items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) + items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) # type: ignore assert [v.id for v in items] == [1] def test_select(self): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() + items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() # type: ignore assert items[0] == 3 def test_vector_avg(self): From 6a2df2a08f8e6a76fae05b0a4fb8466d723ddc48 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 31 May 2026 16:11:51 -0700 Subject: [PATCH 568/879] Dropped support for Python < 3.10, part 2 [skip ci] --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5716d05..3fb3ab5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ authors = [ {name = "Andrew Kane", email = "andrew@ankane.org"} ] license = "MIT" -requires-python = ">= 3.9" +requires-python = ">= 3.10" dependencies = [ "numpy" ] From c1233e2c05bc2f2e0de9042e5d02f81e0c8514b2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 18:50:55 -0700 Subject: [PATCH 569/879] Fixed link [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 95d5fbe..24f72bb 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://codeberg.org/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) From afa127fcd84639d7a0f5f02bf46b47a533e47d6c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 19:42:51 -0700 Subject: [PATCH 570/879] Added type hints [skip ci] --- CHANGELOG.md | 1 + pgvector/asyncpg/register.py | 3 ++- pgvector/bit.py | 28 +++++++++++---------- pgvector/halfvec.py | 34 +++++++++++++------------ pgvector/pg8000/register.py | 3 ++- pgvector/psycopg/bit.py | 11 +++++--- pgvector/psycopg/halfvec.py | 20 +++++++++------ pgvector/psycopg/register.py | 8 ++++-- pgvector/psycopg/sparsevec.py | 20 +++++++++------ pgvector/psycopg/vector.py | 21 ++++++++++------ pgvector/psycopg2/halfvec.py | 11 ++++---- pgvector/psycopg2/register.py | 4 +-- pgvector/psycopg2/sparsevec.py | 11 ++++---- pgvector/psycopg2/vector.py | 11 ++++---- pgvector/py.typed | 0 pgvector/sparsevec.py | 46 ++++++++++++++++++---------------- pgvector/vector.py | 30 +++++++++++----------- 17 files changed, 152 insertions(+), 110 deletions(-) create mode 100644 pgvector/py.typed diff --git a/CHANGELOG.md b/CHANGELOG.md index b5923fa..9ff616b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.5.0 (unreleased) +- Added type hints - Dropped support for Python < 3.10 - Dropped support for SQLAlchemy < 2 diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 63726f3..6718850 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -1,7 +1,8 @@ +from asyncpg import Connection from .. import Vector, HalfVector, SparseVector -async def register_vector(conn, schema='public'): +async def register_vector(conn: Connection, schema: str = 'public') -> None: await conn.set_type_codec( 'vector', schema=schema, diff --git a/pgvector/bit.py b/pgvector/bit.py index 26a9d8d..cb1bec4 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,10 +1,12 @@ +from __future__ import annotations import numpy as np from struct import pack, unpack_from +from typing import Any from warnings import warn class Bit: - def __init__(self, value): + def __init__(self, value: Any) -> None: if isinstance(value, bytes): self._len = 8 * len(value) self._data = value @@ -26,32 +28,32 @@ def __init__(self, value): self._len = len(value) self._data = np.packbits(value).tobytes() - def __repr__(self): + def __repr__(self) -> str: return f'Bit({self.to_text()})' - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: if isinstance(other, self.__class__): return self._len == other._len and self._data == other._data return False - def to_list(self): + def to_list(self) -> list[bool]: return self.to_numpy().tolist() - def to_numpy(self): + def to_numpy(self) -> np.ndarray: return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) - def to_text(self): + def to_text(self) -> str: return ''.join(format(v, '08b') for v in self._data)[:self._len] - def to_binary(self): + def to_binary(self) -> bytes: return pack('>i', self._len) + self._data @classmethod - def from_text(cls, value): + def from_text(cls, value: str) -> Bit: return cls(str(value)) @classmethod - def from_binary(cls, value): + def from_binary(cls, value: bytes) -> Bit: if not isinstance(value, bytes): raise ValueError('expected bytes') @@ -61,15 +63,15 @@ def from_binary(cls, value): return bit @classmethod - def _to_db(cls, value): - if not isinstance(value, cls): + def _to_db(cls, value: Bit) -> str: + if not isinstance(value, Bit): raise ValueError('expected bit') return value.to_text() @classmethod - def _to_db_binary(cls, value): - if not isinstance(value, cls): + def _to_db_binary(cls, value: Bit) -> bytes: + if not isinstance(value, Bit): raise ValueError('expected bit') return value.to_binary() diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index f335f2f..755b9d6 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,9 +1,11 @@ +from __future__ import annotations import numpy as np from struct import pack, unpack_from +from typing import Any class HalfVector: - def __init__(self, value): + def __init__(self, value: Any) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f2': value = np.asarray(value, dtype='>f2') @@ -13,40 +15,40 @@ def __init__(self, value): self._value = value - def __repr__(self): + def __repr__(self) -> str: return f'HalfVector({self.to_list()})' - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: if isinstance(other, self.__class__): return np.array_equal(self.to_numpy(), other.to_numpy()) return False - def dimensions(self): + def dimensions(self) -> int: return len(self._value) - def to_list(self): + def to_list(self) -> list[float]: return self._value.tolist() - def to_numpy(self): + def to_numpy(self) -> np.ndarray: return self._value - def to_text(self): + def to_text(self) -> str: return '[' + ','.join([str(float(v)) for v in self._value]) + ']' - def to_binary(self): + def to_binary(self) -> bytes: return pack('>HH', self.dimensions(), 0) + self._value.tobytes() @classmethod - def from_text(cls, value): + def from_text(cls, value: str) -> HalfVector: return cls([float(v) for v in value[1:-1].split(',')]) @classmethod - def from_binary(cls, value): + def from_binary(cls, value: bytes) -> HalfVector: dim, unused = unpack_from('>HH', value) return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) @classmethod - def _to_db(cls, value, dim=None): + def _to_db(cls, value: Any, dim: int | None = None) -> str | None: if value is None: return value @@ -59,7 +61,7 @@ def _to_db(cls, value, dim=None): return value.to_text() @classmethod - def _to_db_binary(cls, value): + def _to_db_binary(cls, value: Any) -> bytes | None: if value is None: return value @@ -69,15 +71,15 @@ def _to_db_binary(cls, value): return value.to_binary() @classmethod - def _from_db(cls, value): - if value is None or isinstance(value, cls): + def _from_db(cls, value: str | HalfVector | None) -> HalfVector | None: + if value is None or isinstance(value, HalfVector): return value return cls.from_text(value) @classmethod - def _from_db_binary(cls, value): - if value is None or isinstance(value, cls): + def _from_db_binary(cls, value: bytes | HalfVector | None) -> HalfVector | None: + if value is None or isinstance(value, HalfVector): return value return cls.from_binary(value) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index 15ee219..476495b 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -1,8 +1,9 @@ import numpy as np +from pg8000.native import Connection from .. import Vector, HalfVector, SparseVector -def register_vector(conn): +def register_vector(conn: Connection) -> None: # use to_regtype to get first matching type in search path res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") type_info = dict(res) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index cffe8fb..8cfacbf 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,13 +1,18 @@ +from psycopg import BaseConnection +from psycopg.types import TypeInfo from psycopg.adapt import Dumper from psycopg.pq import Format +from typing import Any, TypeAlias from .. import Bit +Buffer: TypeAlias = bytes | bytearray | memoryview + class BitDumper(Dumper): format = Format.TEXT - def dump(self, obj): + def dump(self, obj: Bit) -> Buffer | None: return Bit._to_db(obj).encode('utf8') @@ -15,11 +20,11 @@ class BitBinaryDumper(BitDumper): format = Format.BINARY - def dump(self, obj): + def dump(self, obj: Bit) -> Buffer | None: return Bit._to_db_binary(obj) -def register_bit_info(context, info): +def register_bit_info(context: BaseConnection[Any], info: TypeInfo) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index b3a0060..5ded56d 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,21 +1,27 @@ +from psycopg import BaseConnection from psycopg.adapt import Loader, Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any, TypeAlias from .. import HalfVector +Buffer: TypeAlias = bytes | bytearray | memoryview + class HalfVectorDumper(Dumper): format = Format.TEXT - def dump(self, obj): - return HalfVector._to_db(obj).encode('utf8') + def dump(self, obj: HalfVector) -> Buffer | None: + value = HalfVector._to_db(obj) + return value if value is None else value.encode('utf8') class HalfVectorBinaryDumper(HalfVectorDumper): format = Format.BINARY - def dump(self, obj): + def dump(self, obj: HalfVector) -> Buffer | None: return HalfVector._to_db_binary(obj) @@ -23,7 +29,7 @@ class HalfVectorLoader(Loader): format = Format.TEXT - def load(self, data): + def load(self, data: Buffer) -> HalfVector | None: if isinstance(data, memoryview): data = bytes(data) return HalfVector._from_db(data.decode('utf8')) @@ -33,13 +39,13 @@ class HalfVectorBinaryLoader(HalfVectorLoader): format = Format.BINARY - def load(self, data): - if isinstance(data, memoryview): + def load(self, data: Buffer) -> HalfVector | None: + if isinstance(data, (bytearray, memoryview)): data = bytes(data) return HalfVector._from_db_binary(data) -def register_halfvec_info(context, info): +def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py index b93fd3e..0256b94 100644 --- a/pgvector/psycopg/register.py +++ b/pgvector/psycopg/register.py @@ -1,15 +1,18 @@ +from psycopg import AsyncConnection, Connection from psycopg.types import TypeInfo +from typing import Any from .bit import register_bit_info from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -def register_vector(context): +def register_vector(context: Connection[Any]) -> None: info = TypeInfo.fetch(context, 'vector') register_vector_info(context, info) info = TypeInfo.fetch(context, 'bit') + assert info is not None register_bit_info(context, info) info = TypeInfo.fetch(context, 'halfvec') @@ -21,11 +24,12 @@ def register_vector(context): register_sparsevec_info(context, info) -async def register_vector_async(context): +async def register_vector_async(context: AsyncConnection[Any]) -> None: info = await TypeInfo.fetch(context, 'vector') register_vector_info(context, info) info = await TypeInfo.fetch(context, 'bit') + assert info is not None register_bit_info(context, info) info = await TypeInfo.fetch(context, 'halfvec') diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 384a0e1..b209b8f 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,21 +1,27 @@ +from psycopg import BaseConnection from psycopg.adapt import Loader, Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any, TypeAlias from .. import SparseVector +Buffer: TypeAlias = bytes | bytearray | memoryview + class SparseVectorDumper(Dumper): format = Format.TEXT - def dump(self, obj): - return SparseVector._to_db(obj).encode('utf8') + def dump(self, obj: SparseVector) -> Buffer | None: + value = SparseVector._to_db(obj) + return value if value is None else value.encode('utf8') class SparseVectorBinaryDumper(SparseVectorDumper): format = Format.BINARY - def dump(self, obj): + def dump(self, obj: SparseVector) -> Buffer | None: return SparseVector._to_db_binary(obj) @@ -23,7 +29,7 @@ class SparseVectorLoader(Loader): format = Format.TEXT - def load(self, data): + def load(self, data: Buffer) -> SparseVector | None: if isinstance(data, memoryview): data = bytes(data) return SparseVector._from_db(data.decode('utf8')) @@ -33,13 +39,13 @@ class SparseVectorBinaryLoader(SparseVectorLoader): format = Format.BINARY - def load(self, data): - if isinstance(data, memoryview): + def load(self, data: Buffer) -> SparseVector | None: + if isinstance(data, (bytearray, memoryview)): data = bytes(data) return SparseVector._from_db_binary(data) -def register_sparsevec_info(context, info): +def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index db9e826..1492017 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,22 +1,29 @@ +import numpy as np import psycopg +from psycopg import BaseConnection from psycopg.adapt import Loader, Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any, TypeAlias from .. import Vector +Buffer: TypeAlias = bytes | bytearray | memoryview + class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj): - return Vector._to_db(obj).encode('utf8') + def dump(self, obj: Vector) -> Buffer | None: + value = Vector._to_db(obj) + return value if value is None else value.encode('utf8') class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj): + def dump(self, obj: Vector) -> Buffer | None: return Vector._to_db_binary(obj) @@ -24,7 +31,7 @@ class VectorLoader(Loader): format = Format.TEXT - def load(self, data): + def load(self, data: Buffer) -> np.ndarray | None: if isinstance(data, memoryview): data = bytes(data) return Vector._from_db(data.decode('utf8')) @@ -34,13 +41,13 @@ class VectorBinaryLoader(VectorLoader): format = Format.BINARY - def load(self, data): - if isinstance(data, memoryview): + def load(self, data: Buffer) -> np.ndarray | None: + if isinstance(data, (bytearray, memoryview)): data = bytes(data) return Vector._from_db_binary(data) -def register_vector_info(context, info): +def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) -> None: if info is None: raise psycopg.ProgrammingError('vector type not found in the database') info.register(context) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 0a4c736..5b72f5c 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,20 +1,21 @@ -from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import HalfVector class HalfvecAdapter: - def __init__(self, value): + def __init__(self, value: Any) -> None: self._value = value - def getquoted(self): + def getquoted(self) -> Any: return adapt(HalfVector._to_db(self._value)).getquoted() -def cast_halfvec(value, cur): +def cast_halfvec(value: str | None, cur: cursor) -> HalfVector | None: return HalfVector._from_db(value) -def register_halfvec_info(oid, array_oid, scope): +def register_halfvec_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) register_type(halfvec, scope) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 1bc9d44..22beb8c 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -1,12 +1,12 @@ import psycopg2 -from psycopg2.extensions import cursor +from psycopg2.extensions import connection, cursor from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info # note: register_adapter is always global -def register_vector(conn_or_curs, globally=False, arrays=True): +def register_vector(conn_or_curs: connection | cursor, globally: bool = False, arrays: bool = True) -> None: 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 diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 148eff2..073e45b 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,20 +1,21 @@ -from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import SparseVector class SparsevecAdapter: - def __init__(self, value): + def __init__(self, value: Any) -> None: self._value = value - def getquoted(self): + def getquoted(self) -> Any: return adapt(SparseVector._to_db(self._value)).getquoted() -def cast_sparsevec(value, cur): +def cast_sparsevec(value: str | None, cur: cursor) -> SparseVector | None: return SparseVector._from_db(value) -def register_sparsevec_info(oid, array_oid, scope): +def register_sparsevec_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) register_type(sparsevec, scope) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 562de18..0e2d1e3 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,21 +1,22 @@ import numpy as np -from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import Vector class VectorAdapter: - def __init__(self, value): + def __init__(self, value: Any) -> None: self._value = value - def getquoted(self): + def getquoted(self) -> Any: return adapt(Vector._to_db(self._value)).getquoted() -def cast_vector(value, cur): +def cast_vector(value: str | None, cur: cursor) -> np.ndarray | None: return Vector._from_db(value) -def register_vector_info(oid, array_oid, scope): +def register_vector_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: vector = new_type((oid,), 'VECTOR', cast_vector) register_type(vector, scope) diff --git a/pgvector/py.typed b/pgvector/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 895fbd0..388116d 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,11 +1,13 @@ +from __future__ import annotations import numpy as np from struct import pack, unpack_from +from typing import Any NO_DEFAULT = object() class SparseVector: - def __init__(self, value, dimensions=NO_DEFAULT, /): + def __init__(self, value: Any, dimensions: int | Any = NO_DEFAULT, /) -> None: if value.__class__.__module__.startswith('scipy.sparse.'): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') @@ -22,22 +24,22 @@ def __init__(self, value, dimensions=NO_DEFAULT, /): self._from_dense(value) - def __repr__(self): + def __repr__(self) -> str: elements = dict(zip(self._indices, self._values)) return f'SparseVector({elements}, {self._dim})' - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: if isinstance(other, self.__class__): return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values() return False - def dimensions(self): + def dimensions(self) -> int: return self._dim - def indices(self): + def indices(self) -> list[int]: return self._indices - def values(self): + def values(self) -> list[float]: return self._values def to_coo(self): @@ -46,26 +48,26 @@ def to_coo(self): coords = ([0] * len(self._indices), self._indices) return coo_array((self._values, coords), shape=(1, self._dim)) - def to_list(self): + def to_list(self) -> list[float]: vec = [0.0] * self._dim for i, v in zip(self._indices, self._values): vec[i] = v return vec - def to_numpy(self): + def to_numpy(self) -> np.ndarray: vec = np.repeat(0.0, self._dim).astype(np.float32) for i, v in zip(self._indices, self._values): vec[i] = v return vec - def to_text(self): + def to_text(self) -> str: return '{' + ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + '}/' + str(int(self._dim)) - def to_binary(self): + def to_binary(self) -> bytes: nnz = len(self._indices) return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) - def _from_dict(self, d, dim): + def _from_dict(self, d: dict[int, float], dim: int) -> None: elements = [(i, v) for i, v in d.items() if v != 0] elements.sort() @@ -73,7 +75,7 @@ def _from_dict(self, d, dim): self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] - def _from_sparse(self, value): + def _from_sparse(self, value: Any) -> None: value = value.tocoo() if value.ndim == 1: @@ -90,13 +92,13 @@ def _from_sparse(self, value): self._indices = value.col.tolist() self._values = value.data.tolist() - def _from_dense(self, value): + def _from_dense(self, value: list[float]) -> None: self._dim = len(value) self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] @classmethod - def from_text(cls, value): + def from_text(cls, value: str) -> SparseVector: elements, dim = value.split('/', 2) indices = [] values = [] @@ -109,14 +111,14 @@ def from_text(cls, value): return cls._from_parts(int(dim), indices, values) @classmethod - def from_binary(cls, value): + def from_binary(cls, value: bytes) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) return cls._from_parts(int(dim), list(indices), list(values)) @classmethod - def _from_parts(cls, dim, indices, values): + def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> SparseVector: vec = cls.__new__(cls) vec._dim = dim vec._indices = indices @@ -124,7 +126,7 @@ def _from_parts(cls, dim, indices, values): return vec @classmethod - def _to_db(cls, value, dim=None): + def _to_db(cls, value: Any, dim: int | None = None) -> str | None: if value is None: return value @@ -137,7 +139,7 @@ def _to_db(cls, value, dim=None): return value.to_text() @classmethod - def _to_db_binary(cls, value): + def _to_db_binary(cls, value: Any) -> bytes | None: if value is None: return value @@ -147,15 +149,15 @@ def _to_db_binary(cls, value): return value.to_binary() @classmethod - def _from_db(cls, value): - if value is None or isinstance(value, cls): + def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None: + if value is None or isinstance(value, SparseVector): return value return cls.from_text(value) @classmethod - def _from_db_binary(cls, value): - if value is None or isinstance(value, cls): + def _from_db_binary(cls, value: bytes | SparseVector | None) -> SparseVector | None: + if value is None or isinstance(value, SparseVector): return value return cls.from_binary(value) diff --git a/pgvector/vector.py b/pgvector/vector.py index ebbcafd..f5ba742 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,9 +1,11 @@ +from __future__ import annotations import numpy as np from struct import pack, unpack_from +from typing import Any class Vector: - def __init__(self, value): + def __init__(self, value: Any) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f4': value = np.asarray(value, dtype='>f4') @@ -13,40 +15,40 @@ def __init__(self, value): self._value = value - def __repr__(self): + def __repr__(self) -> str: return f'Vector({self.to_list()})' - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: if isinstance(other, self.__class__): return np.array_equal(self.to_numpy(), other.to_numpy()) return False - def dimensions(self): + def dimensions(self) -> int: return len(self._value) - def to_list(self): + def to_list(self) -> list[float]: return self._value.tolist() - def to_numpy(self): + def to_numpy(self) -> np.ndarray: return self._value - def to_text(self): + def to_text(self) -> str: return '[' + ','.join([str(float(v)) for v in self._value]) + ']' - def to_binary(self): + def to_binary(self) -> bytes: return pack('>HH', self.dimensions(), 0) + self._value.tobytes() @classmethod - def from_text(cls, value): + def from_text(cls, value: str) -> Vector: return cls([float(v) for v in value[1:-1].split(',')]) @classmethod - def from_binary(cls, value): + def from_binary(cls, value: bytes) -> Vector: dim, unused = unpack_from('>HH', value) return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) @classmethod - def _to_db(cls, value, dim=None): + def _to_db(cls, value: Any, dim: int | None = None) -> str | None: if value is None: return value @@ -59,7 +61,7 @@ def _to_db(cls, value, dim=None): return value.to_text() @classmethod - def _to_db_binary(cls, value): + def _to_db_binary(cls, value: Any) -> bytes | None: if value is None: return value @@ -69,14 +71,14 @@ def _to_db_binary(cls, value): return value.to_binary() @classmethod - def _from_db(cls, value): + def _from_db(cls, value: str | np.ndarray | None) -> np.ndarray | None: if value is None or isinstance(value, np.ndarray): return value return cls.from_text(value).to_numpy().astype(np.float32) @classmethod - def _from_db_binary(cls, value): + def _from_db_binary(cls, value: bytes | np.ndarray | None) -> np.ndarray | None: if value is None or isinstance(value, np.ndarray): return value From d95fa2ef366039ac4b5e3badd679b074b9f1085a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 20:05:25 -0700 Subject: [PATCH 571/879] Added more type hints [skip ci] --- pgvector/sqlalchemy/bit.py | 12 +++++++----- pgvector/sqlalchemy/halfvec.py | 20 +++++++++++--------- pgvector/sqlalchemy/sparsevec.py | 20 +++++++++++--------- pgvector/sqlalchemy/vector.py | 20 +++++++++++--------- 4 files changed, 40 insertions(+), 32 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index bbf822e..c0bb6ba 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -1,20 +1,22 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float +from sqlalchemy import Dialect +from typing import Any class BIT(UserDefinedType): cache_ok = True - def __init__(self, length=None): + def __init__(self, length: int | None = None) -> None: super(UserDefinedType, self).__init__() self.length = length - def get_col_spec(self, **kw): + def get_col_spec(self, **kw) -> str: if self.length is None: return 'BIT' return 'BIT(%d)' % self.length - def bind_processor(self, dialect): + def bind_processor(self, dialect: Dialect) -> Any: if dialect.__class__.__name__ == 'PGDialect_asyncpg': import asyncpg @@ -27,10 +29,10 @@ def process(value): return super().bind_processor(dialect) class comparator_factory(UserDefinedType.Comparator): - def hamming_distance(self, other): + def hamming_distance(self, other: Any) -> Any: return self.op('<~>', return_type=Float)(other) - def jaccard_distance(self, other): + def jaccard_distance(self, other: Any) -> Any: return self.op('<%>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index f51d1a6..3612cfc 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,5 +1,7 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy import Dialect +from typing import Any from .. import HalfVector @@ -7,43 +9,43 @@ class HALFVEC(UserDefinedType): cache_ok = True _string = String() - def __init__(self, dim=None): + def __init__(self, dim: int | None = None) -> None: super(UserDefinedType, self).__init__() self.dim = dim - def get_col_spec(self, **kw): + def get_col_spec(self, **kw) -> str: if self.dim is None: return 'HALFVEC' return 'HALFVEC(%d)' % self.dim - def bind_processor(self, dialect): + def bind_processor(self, dialect: Dialect): def process(value): return HalfVector._to_db(value, self.dim) return process - def literal_processor(self, dialect): + def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) def process(value): return string_literal_processor(HalfVector._to_db(value, self.dim)) # type: ignore return process - def result_processor(self, dialect, coltype): + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value): return HalfVector._from_db(value) return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other): + def l2_distance(self, other: Any) -> Any: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other): + def max_inner_product(self, other: Any) -> Any: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other): + def cosine_distance(self, other: Any) -> Any: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other): + def l1_distance(self, other: Any) -> Any: return self.op('<+>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 6828bfd..e409bf1 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,5 +1,7 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy import Dialect +from typing import Any from .. import SparseVector @@ -7,43 +9,43 @@ class SPARSEVEC(UserDefinedType): cache_ok = True _string = String() - def __init__(self, dim=None): + def __init__(self, dim: int | None = None) -> None: super(UserDefinedType, self).__init__() self.dim = dim - def get_col_spec(self, **kw): + def get_col_spec(self, **kw) -> str: if self.dim is None: return 'SPARSEVEC' return 'SPARSEVEC(%d)' % self.dim - def bind_processor(self, dialect): + def bind_processor(self, dialect: Dialect) -> Any: def process(value): return SparseVector._to_db(value, self.dim) return process - def literal_processor(self, dialect): + def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) def process(value): return string_literal_processor(SparseVector._to_db(value, self.dim)) # type: ignore return process - def result_processor(self, dialect, coltype): + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value): return SparseVector._from_db(value) return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other): + def l2_distance(self, other: Any) -> Any: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other): + def max_inner_product(self, other: Any) -> Any: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other): + def cosine_distance(self, other: Any) -> Any: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other): + def l1_distance(self, other: Any) -> Any: return self.op('<+>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index a9e6f22..187e039 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,5 +1,7 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy import Dialect +from typing import Any from .. import Vector @@ -7,43 +9,43 @@ class VECTOR(UserDefinedType): cache_ok = True _string = String() - def __init__(self, dim=None): + def __init__(self, dim: int | None = None) -> None: super(UserDefinedType, self).__init__() self.dim = dim - def get_col_spec(self, **kw): + def get_col_spec(self, **kw) -> str: if self.dim is None: return 'VECTOR' return 'VECTOR(%d)' % self.dim - def bind_processor(self, dialect): + def bind_processor(self, dialect: Dialect) -> Any: def process(value): return Vector._to_db(value, self.dim) return process - def literal_processor(self, dialect): + def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) def process(value): return string_literal_processor(Vector._to_db(value, self.dim)) # type: ignore return process - def result_processor(self, dialect, coltype): + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value): return Vector._from_db(value) return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other): + def l2_distance(self, other: Any) -> Any: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other): + def max_inner_product(self, other: Any) -> Any: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other): + def cosine_distance(self, other: Any) -> Any: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other): + def l1_distance(self, other: Any) -> Any: return self.op('<+>', return_type=Float)(other) From 0673ca77c13482844807b7c03670898c32f23f2c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 20:08:14 -0700 Subject: [PATCH 572/879] Improved type hints [skip ci] --- pgvector/sqlalchemy/bit.py | 6 +++--- pgvector/sqlalchemy/halfvec.py | 10 +++++----- pgvector/sqlalchemy/sparsevec.py | 10 +++++----- pgvector/sqlalchemy/vector.py | 10 +++++----- 4 files changed, 18 insertions(+), 18 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index c0bb6ba..f38c994 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float -from sqlalchemy import Dialect +from sqlalchemy import Dialect, Operators from typing import Any @@ -29,10 +29,10 @@ def process(value): return super().bind_processor(dialect) class comparator_factory(UserDefinedType.Comparator): - def hamming_distance(self, other: Any) -> Any: + def hamming_distance(self, other: Any) -> Operators: return self.op('<~>', return_type=Float)(other) - def jaccard_distance(self, other: Any) -> Any: + def jaccard_distance(self, other: Any) -> Operators: return self.op('<%>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 3612cfc..f964c3a 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from sqlalchemy import Dialect +from sqlalchemy import Dialect, Operators from typing import Any from .. import HalfVector @@ -36,16 +36,16 @@ def process(value): return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other: Any) -> Any: + def l2_distance(self, other: Any) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: Any) -> Any: + def max_inner_product(self, other: Any) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: Any) -> Any: + def cosine_distance(self, other: Any) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: Any) -> Any: + def l1_distance(self, other: Any) -> Operators: return self.op('<+>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index e409bf1..17964c7 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from sqlalchemy import Dialect +from sqlalchemy import Dialect, Operators from typing import Any from .. import SparseVector @@ -36,16 +36,16 @@ def process(value): return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other: Any) -> Any: + def l2_distance(self, other: Any) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: Any) -> Any: + def max_inner_product(self, other: Any) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: Any) -> Any: + def cosine_distance(self, other: Any) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: Any) -> Any: + def l1_distance(self, other: Any) -> Operators: return self.op('<+>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 187e039..e1e301e 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from sqlalchemy import Dialect +from sqlalchemy import Dialect, Operators from typing import Any from .. import Vector @@ -36,16 +36,16 @@ def process(value): return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other: Any) -> Any: + def l2_distance(self, other: Any) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: Any) -> Any: + def max_inner_product(self, other: Any) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: Any) -> Any: + def cosine_distance(self, other: Any) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: Any) -> Any: + def l1_distance(self, other: Any) -> Operators: return self.op('<+>', return_type=Float)(other) From f4498f287984dbcb81c6a55909dee51df6cadd2b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 20:23:30 -0700 Subject: [PATCH 573/879] Improved type hints [skip ci] --- pgvector/peewee/bit.py | 13 +++++++------ pgvector/peewee/halfvec.py | 21 +++++++++++---------- pgvector/peewee/sparsevec.py | 21 +++++++++++---------- pgvector/peewee/vector.py | 22 ++++++++++++---------- pgvector/sqlalchemy/halfvec.py | 2 +- 5 files changed, 42 insertions(+), 37 deletions(-) diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index ee5f12f..a2bb797 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -1,21 +1,22 @@ from peewee import Expression, Field +from typing import Any class FixedBitField(Field): field_type = 'bit' - def __init__(self, max_length=None, *args, **kwargs): + def __init__(self, max_length: int | None = None, *args, **kwargs) -> None: self.max_length = max_length super(FixedBitField, self).__init__(*args, **kwargs) - def get_modifiers(self): - return self.max_length and [self.max_length] or None + def get_modifiers(self) -> list[int] | None: + return [self.max_length] if self.max_length else None - def _distance(self, op, vector): + def _distance(self, op: str, vector: Any) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def hamming_distance(self, vector): + def hamming_distance(self, vector: Any) -> Expression: return self._distance('<~>', vector) - def jaccard_distance(self, vector): + def jaccard_distance(self, vector: Any) -> Expression: return self._distance('<%%>', vector) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 0901fd2..234da40 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,34 +1,35 @@ from peewee import Expression, Field +from typing import Any from .. import HalfVector class HalfVectorField(Field): field_type = 'halfvec' - def __init__(self, dimensions=None, *args, **kwargs): + def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None: self.dimensions = dimensions super(HalfVectorField, self).__init__(*args, **kwargs) - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None - def db_value(self, value): + def db_value(self, value: Any) -> str | None: return HalfVector._to_db(value) - def python_value(self, value): + def python_value(self, value: Any) -> HalfVector | None: return HalfVector._from_db(value) - def _distance(self, op, vector): + def _distance(self, op: str, vector: Any) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector): + def l2_distance(self, vector: Any) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector): + def max_inner_product(self, vector: Any) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector): + def cosine_distance(self, vector: Any) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector): + def l1_distance(self, vector: Any) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 86dea73..816ff68 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -1,34 +1,35 @@ from peewee import Expression, Field +from typing import Any from .. import SparseVector class SparseVectorField(Field): field_type = 'sparsevec' - def __init__(self, dimensions=None, *args, **kwargs): + def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None: self.dimensions = dimensions super(SparseVectorField, self).__init__(*args, **kwargs) - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None - def db_value(self, value): + def db_value(self, value: Any) -> str | None: return SparseVector._to_db(value) - def python_value(self, value): + def python_value(self, value: Any) -> SparseVector | None: return SparseVector._from_db(value) - def _distance(self, op, vector): + def _distance(self, op: str, vector: Any) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector): + def l2_distance(self, vector: Any) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector): + def max_inner_product(self, vector: Any) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector): + def cosine_distance(self, vector: Any) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector): + def l1_distance(self, vector: Any) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 83f9997..6cdbef6 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,34 +1,36 @@ +import numpy as np from peewee import Expression, Field +from typing import Any from .. import Vector class VectorField(Field): field_type = 'vector' - def __init__(self, dimensions=None, *args, **kwargs): + def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None: self.dimensions = dimensions super(VectorField, self).__init__(*args, **kwargs) - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None - def db_value(self, value): + def db_value(self, value: Any) -> str | None: return Vector._to_db(value) - def python_value(self, value): + def python_value(self, value: Any) -> np.ndarray | None: return Vector._from_db(value) - def _distance(self, op, vector): + def _distance(self, op: str, vector: Any) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector): + def l2_distance(self, vector: Any) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector): + def max_inner_product(self, vector: Any) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector): + def cosine_distance(self, vector: Any) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector): + def l1_distance(self, vector: Any) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index f964c3a..703861a 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -18,7 +18,7 @@ def get_col_spec(self, **kw) -> str: return 'HALFVEC' return 'HALFVEC(%d)' % self.dim - def bind_processor(self, dialect: Dialect): + def bind_processor(self, dialect: Dialect) -> Any: def process(value): return HalfVector._to_db(value, self.dim) return process From 90fd6da9c54e22575ba94e5a9e9b96b7ca7a1c04 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 20:25:13 -0700 Subject: [PATCH 574/879] Fixed error with mypy [skip ci] --- pgvector/django/extensions.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py index 1d04739..9974e99 100644 --- a/pgvector/django/extensions.py +++ b/pgvector/django/extensions.py @@ -3,9 +3,8 @@ class VectorExtension(CreateExtension): - if VERSION[0] >= 6: - def __init__(self, hints=None): + def __init__(self, hints=None): + if VERSION[0] >= 6: super().__init__('vector', hints=hints) - else: - def __init__(self): + else: self.name = 'vector' From 9c2d44add965e726d6ee4fc6b3993b2d1fe2f3d7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 21:00:00 -0700 Subject: [PATCH 575/879] Added more type hints [skip ci] --- pgvector/django/bit.py | 11 ++++++----- pgvector/django/extensions.py | 3 ++- pgvector/django/functions.py | 5 +++-- pgvector/django/halfvec.py | 19 ++++++++++--------- pgvector/django/indexes.py | 13 +++++++------ pgvector/django/sparsevec.py | 21 +++++++++++---------- pgvector/django/vector.py | 27 ++++++++++++++------------- 7 files changed, 53 insertions(+), 46 deletions(-) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 29a31e1..cfc98ec 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -1,32 +1,33 @@ from django import forms from django.db.models import Field +from typing import Any # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ class BitField(Field): description = 'Bit string' - def __init__(self, *args, length=None, **kwargs): + def __init__(self, *args: Any, length: int | None = None, **kwargs: Any) -> None: self.length = length super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple: name, path, args, kwargs = super().deconstruct() if self.length is not None: kwargs['length'] = self.length return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.length is None: return 'bit' return 'bit(%d)' % self.length - def formfield(self, **kwargs): # type: ignore + def formfield(self, **kwargs: Any): # type: ignore return super().formfield(form_class=BitFormField, **kwargs) class BitFormField(forms.CharField): - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py index 9974e99..debac5d 100644 --- a/pgvector/django/extensions.py +++ b/pgvector/django/extensions.py @@ -1,9 +1,10 @@ from django import VERSION from django.contrib.postgres.operations import CreateExtension +from typing import Any class VectorExtension(CreateExtension): - def __init__(self, hints=None): + def __init__(self, hints: Any = None) -> None: if VERSION[0] >= 6: super().__init__('vector', hints=hints) else: diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 9df4fdb..af76a67 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,11 +1,12 @@ from django.db.models import FloatField, Func, Value from .. import Vector, HalfVector, SparseVector +from typing import Any class DistanceBase(Func): output_field = FloatField() - def __init__(self, expression, vector, **extra): + def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): if isinstance(vector, HalfVector): vector = Value(HalfVector._to_db(vector)) @@ -23,7 +24,7 @@ def __init__(self, expression, vector, **extra): class BitDistanceBase(Func): output_field = FloatField() - def __init__(self, expression, vector, **extra): + def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): vector = Value(vector) super().__init__(expression, vector, **extra) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 36101ca..973cc0d 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -1,5 +1,6 @@ from django import forms from django.db.models import Field +from typing import Any from .. import HalfVector @@ -8,25 +9,25 @@ class HalfVectorField(Field): description = 'Half vector' empty_strings_allowed = False - def __init__(self, *args, dimensions=None, **kwargs): + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None: self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.dimensions is None: return 'halfvec' return 'halfvec(%d)' % self.dimensions - def from_db_value(self, value, expression, connection): + def from_db_value(self, value: Any, expression: Any, connection: Any) -> HalfVector | None: return HalfVector._from_db(value) - def to_python(self, value): + def to_python(self, value: Any) -> HalfVector | None: if value is None or isinstance(value, HalfVector): return value elif isinstance(value, str): @@ -34,10 +35,10 @@ def to_python(self, value): else: return HalfVector(value) - def get_prep_value(self, value): + def get_prep_value(self, value: Any) -> str | None: return HalfVector._to_db(value) - def value_to_string(self, obj): + def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) def formfield(self, **kwargs): # type: ignore @@ -45,7 +46,7 @@ def formfield(self, **kwargs): # type: ignore class HalfVectorWidget(forms.TextInput): - def format_value(self, value): + def format_value(self, value: Any) -> str | None: if isinstance(value, HalfVector): value = value.to_list() return super().format_value(value) @@ -54,7 +55,7 @@ def format_value(self, value): class HalfVectorFormField(forms.CharField): widget = HalfVectorWidget - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/django/indexes.py b/pgvector/django/indexes.py index 5bec0eb..d18330f 100644 --- a/pgvector/django/indexes.py +++ b/pgvector/django/indexes.py @@ -1,20 +1,21 @@ from django.contrib.postgres.indexes import PostgresIndex +from typing import Any class IvfflatIndex(PostgresIndex): suffix = 'ivfflat' - def __init__(self, *expressions, lists=None, **kwargs): + def __init__(self, *expressions: Any, lists: int | None = None, **kwargs: Any) -> None: self.lists = lists super().__init__(*expressions, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple: path, args, kwargs = super().deconstruct() if self.lists is not None: kwargs['lists'] = self.lists return path, args, kwargs - def get_with_params(self): + def get_with_params(self) -> list[str]: with_params = [] if self.lists is not None: with_params.append('lists = %d' % self.lists) @@ -24,12 +25,12 @@ def get_with_params(self): class HnswIndex(PostgresIndex): suffix = 'hnsw' - def __init__(self, *expressions, m=None, ef_construction=None, **kwargs): + def __init__(self, *expressions: Any, m: int | None = None, ef_construction: int | None = None, **kwargs: Any) -> None: self.m = m self.ef_construction = ef_construction super().__init__(*expressions, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple: path, args, kwargs = super().deconstruct() if self.m is not None: kwargs['m'] = self.m @@ -37,7 +38,7 @@ def deconstruct(self): kwargs['ef_construction'] = self.ef_construction return path, args, kwargs - def get_with_params(self): + def get_with_params(self) -> list[str]: with_params = [] if self.m is not None: with_params.append('m = %d' % self.m) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 7f4db77..c264d7e 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -1,5 +1,6 @@ from django import forms from django.db.models import Field +from typing import Any from .. import SparseVector @@ -8,39 +9,39 @@ class SparseVectorField(Field): description = 'Sparse vector' empty_strings_allowed = False - def __init__(self, *args, dimensions=None, **kwargs): + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any): self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.dimensions is None: return 'sparsevec' return 'sparsevec(%d)' % self.dimensions - def from_db_value(self, value, expression, connection): + def from_db_value(self, value: Any, expression: Any, connection: Any) -> SparseVector | None: return SparseVector._from_db(value) - def to_python(self, value): + def to_python(self, value: Any) -> SparseVector | None: return SparseVector._from_db(value) - def get_prep_value(self, value): + def get_prep_value(self, value: Any) -> str | None: return SparseVector._to_db(value) - def value_to_string(self, obj): + def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs): # type: ignore + def formfield(self, **kwargs: Any): # type: ignore return super().formfield(form_class=SparseVectorFormField, **kwargs) class SparseVectorWidget(forms.TextInput): - def format_value(self, value): + def format_value(self, value: Any) -> str | None: if isinstance(value, SparseVector): value = value.to_text() return super().format_value(value) @@ -49,7 +50,7 @@ def format_value(self, value): class SparseVectorFormField(forms.CharField): widget = SparseVectorWidget - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index dde96f8..08c04c0 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -1,6 +1,7 @@ from django import forms from django.db.models import Field import numpy as np +from typing import Any from .. import Vector @@ -9,51 +10,51 @@ class VectorField(Field): description = 'Vector' empty_strings_allowed = False - def __init__(self, *args, dimensions=None, **kwargs): + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None: self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.dimensions is None: return 'vector' return 'vector(%d)' % self.dimensions - def from_db_value(self, value, expression, connection): + def from_db_value(self, value: Any, expression: Any, connection: Any) -> np.ndarray | None: return Vector._from_db(value) - def to_python(self, value): + def to_python(self, value: Any) -> np.ndarray | None: if isinstance(value, list): return np.array(value, dtype=np.float32) return Vector._from_db(value) - def get_prep_value(self, value): + def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) - def value_to_string(self, obj): + def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def validate(self, value, model_instance): + def validate(self, value: Any, model_instance: Any) -> None: if isinstance(value, np.ndarray): value = value.tolist() super().validate(value, model_instance) - def run_validators(self, value): + def run_validators(self, value: Any) -> None: if isinstance(value, np.ndarray): value = value.tolist() super().run_validators(value) - def formfield(self, **kwargs): # type: ignore + def formfield(self, **kwargs: Any): # type: ignore return super().formfield(form_class=VectorFormField, **kwargs) class VectorWidget(forms.TextInput): - def format_value(self, value): + def format_value(self, value: Any) -> str | None: if isinstance(value, np.ndarray): value = value.tolist() return super().format_value(value) @@ -62,12 +63,12 @@ def format_value(self, value): class VectorFormField(forms.CharField): widget = VectorWidget - def has_changed(self, initial, data): + def has_changed(self, initial: Any, data: Any) -> bool: if isinstance(initial, np.ndarray): initial = initial.tolist() return super().has_changed(initial, data) - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) From c9998e6a019a5dda32d81793c8c05dd1bc2ba668 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 21:07:26 -0700 Subject: [PATCH 576/879] Improved type hints [skip ci] --- pgvector/django/bit.py | 2 +- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- pgvector/peewee/bit.py | 2 +- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/sparsevec.py | 2 +- pgvector/peewee/vector.py | 2 +- pgvector/sqlalchemy/bit.py | 2 +- pgvector/sqlalchemy/halfvec.py | 2 +- pgvector/sqlalchemy/sparsevec.py | 2 +- pgvector/sqlalchemy/vector.py | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index cfc98ec..167f550 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -22,7 +22,7 @@ def db_type(self, connection: Any) -> str: return 'bit' return 'bit(%d)' % self.length - def formfield(self, **kwargs: Any): # type: ignore + def formfield(self, **kwargs: Any) -> forms.Field: # type: ignore return super().formfield(form_class=BitFormField, **kwargs) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 973cc0d..d445f75 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -41,7 +41,7 @@ def get_prep_value(self, value: Any) -> str | None: def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs): # type: ignore + def formfield(self, **kwargs) -> forms.Field: # type: ignore return super().formfield(form_class=HalfVectorFormField, **kwargs) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index c264d7e..3b2bb52 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -36,7 +36,7 @@ def get_prep_value(self, value: Any) -> str | None: def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs: Any): # type: ignore + def formfield(self, **kwargs: Any) -> forms.Field: # type: ignore return super().formfield(form_class=SparseVectorFormField, **kwargs) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 08c04c0..64d1f5e 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -49,7 +49,7 @@ def run_validators(self, value: Any) -> None: value = value.tolist() super().run_validators(value) - def formfield(self, **kwargs: Any): # type: ignore + def formfield(self, **kwargs: Any) -> forms.Field: # type: ignore return super().formfield(form_class=VectorFormField, **kwargs) diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index a2bb797..3128bcc 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -5,7 +5,7 @@ class FixedBitField(Field): field_type = 'bit' - def __init__(self, max_length: int | None = None, *args, **kwargs) -> None: + def __init__(self, max_length: int | None = None, *args: Any, **kwargs: Any) -> None: self.max_length = max_length super(FixedBitField, self).__init__(*args, **kwargs) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 234da40..992d8e7 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -6,7 +6,7 @@ class HalfVectorField(Field): field_type = 'halfvec' - def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None: + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions super(HalfVectorField, self).__init__(*args, **kwargs) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 816ff68..dfb403e 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -6,7 +6,7 @@ class SparseVectorField(Field): field_type = 'sparsevec' - def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None: + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions super(SparseVectorField, self).__init__(*args, **kwargs) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 6cdbef6..43b9f4b 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -7,7 +7,7 @@ class VectorField(Field): field_type = 'vector' - def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None: + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions super(VectorField, self).__init__(*args, **kwargs) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index f38c994..21caa31 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -11,7 +11,7 @@ def __init__(self, length: int | None = None) -> None: super(UserDefinedType, self).__init__() self.length = length - def get_col_spec(self, **kw) -> str: + def get_col_spec(self, **kw: Any) -> str: if self.length is None: return 'BIT' return 'BIT(%d)' % self.length diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 703861a..4e422a2 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -13,7 +13,7 @@ def __init__(self, dim: int | None = None) -> None: super(UserDefinedType, self).__init__() self.dim = dim - def get_col_spec(self, **kw) -> str: + def get_col_spec(self, **kw: Any) -> str: if self.dim is None: return 'HALFVEC' return 'HALFVEC(%d)' % self.dim diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 17964c7..d4e5cf0 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -13,7 +13,7 @@ def __init__(self, dim: int | None = None) -> None: super(UserDefinedType, self).__init__() self.dim = dim - def get_col_spec(self, **kw) -> str: + def get_col_spec(self, **kw: Any) -> str: if self.dim is None: return 'SPARSEVEC' return 'SPARSEVEC(%d)' % self.dim diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index e1e301e..251e1d5 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -13,7 +13,7 @@ def __init__(self, dim: int | None = None) -> None: super(UserDefinedType, self).__init__() self.dim = dim - def get_col_spec(self, **kw) -> str: + def get_col_spec(self, **kw: Any) -> str: if self.dim is None: return 'VECTOR' return 'VECTOR(%d)' % self.dim From 12f3dc514d7a8caee684b495a49cac989c444cfe Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 21:13:22 -0700 Subject: [PATCH 577/879] Improved type hints [skip ci] --- pgvector/django/bit.py | 8 ++++++-- pgvector/django/halfvec.py | 8 ++++++-- pgvector/django/sparsevec.py | 8 ++++++-- pgvector/django/vector.py | 8 ++++++-- 4 files changed, 24 insertions(+), 8 deletions(-) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 167f550..8086591 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -22,8 +22,12 @@ def db_type(self, connection: Any) -> str: return 'bit' return 'bit(%d)' % self.length - def formfield(self, **kwargs: Any) -> forms.Field: # type: ignore - return super().formfield(form_class=BitFormField, **kwargs) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=BitFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class BitFormField(forms.CharField): diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index d445f75..dcca99f 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -41,8 +41,12 @@ def get_prep_value(self, value: Any) -> str | None: def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs) -> forms.Field: # type: ignore - return super().formfield(form_class=HalfVectorFormField, **kwargs) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=HalfVectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class HalfVectorWidget(forms.TextInput): diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 3b2bb52..531d53c 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -36,8 +36,12 @@ def get_prep_value(self, value: Any) -> str | None: def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def formfield(self, **kwargs: Any) -> forms.Field: # type: ignore - return super().formfield(form_class=SparseVectorFormField, **kwargs) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=SparseVectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class SparseVectorWidget(forms.TextInput): diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 64d1f5e..7e34963 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -49,8 +49,12 @@ def run_validators(self, value: Any) -> None: value = value.tolist() super().run_validators(value) - def formfield(self, **kwargs: Any) -> forms.Field: # type: ignore - return super().formfield(form_class=VectorFormField, **kwargs) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=VectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class VectorWidget(forms.TextInput): From 4c2b0723c9c4f0b3ffdb33d35248b45a07e73009 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 21:23:52 -0700 Subject: [PATCH 578/879] Improved type hints [skip ci] --- pgvector/sparsevec.py | 2 +- pgvector/sqlalchemy/bit.py | 2 +- pgvector/sqlalchemy/halfvec.py | 6 +++--- pgvector/sqlalchemy/sparsevec.py | 6 +++--- pgvector/sqlalchemy/vector.py | 7 ++++--- 5 files changed, 12 insertions(+), 11 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 388116d..6a00981 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -42,7 +42,7 @@ def indices(self) -> list[int]: def values(self) -> list[float]: return self._values - def to_coo(self): + def to_coo(self) -> Any: from scipy.sparse import coo_array coords = ([0] * len(self._indices), self._indices) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 21caa31..6be2ab9 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -20,7 +20,7 @@ def bind_processor(self, dialect: Dialect) -> Any: if dialect.__class__.__name__ == 'PGDialect_asyncpg': import asyncpg - def process(value): + def process(value: Any) -> Any: if isinstance(value, str): return asyncpg.BitString(value) return value diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 4e422a2..8c050ec 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -19,19 +19,19 @@ def get_col_spec(self, **kw: Any) -> str: return 'HALFVEC(%d)' % self.dim def bind_processor(self, dialect: Dialect) -> Any: - def process(value): + def process(value: Any) -> str | None: return HalfVector._to_db(value, self.dim) return process def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) - def process(value): + def process(value: Any) -> Any: return string_literal_processor(HalfVector._to_db(value, self.dim)) # type: ignore return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: - def process(value): + def process(value: Any) -> HalfVector | None: return HalfVector._from_db(value) return process diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index d4e5cf0..a809c2a 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -19,19 +19,19 @@ def get_col_spec(self, **kw: Any) -> str: return 'SPARSEVEC(%d)' % self.dim def bind_processor(self, dialect: Dialect) -> Any: - def process(value): + def process(value: Any) -> str | None: return SparseVector._to_db(value, self.dim) return process def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) - def process(value): + def process(value: Any) -> Any: return string_literal_processor(SparseVector._to_db(value, self.dim)) # type: ignore return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: - def process(value): + def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 251e1d5..76f9497 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,3 +1,4 @@ +import numpy as np from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String from sqlalchemy import Dialect, Operators @@ -19,19 +20,19 @@ def get_col_spec(self, **kw: Any) -> str: return 'VECTOR(%d)' % self.dim def bind_processor(self, dialect: Dialect) -> Any: - def process(value): + def process(value: Any) -> str | None: return Vector._to_db(value, self.dim) return process def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) - def process(value): + def process(value: Any) -> Any: return string_literal_processor(Vector._to_db(value, self.dim)) # type: ignore return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: - def process(value): + def process(value: Any) -> np.ndarray | None: return Vector._from_db(value) return process From 03bbac707d87e6d53d4759190a2d5acdca1f42f9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 21:44:58 -0700 Subject: [PATCH 579/879] Improved type hints [skip ci] --- pgvector/django/bit.py | 2 +- pgvector/django/halfvec.py | 2 +- pgvector/django/indexes.py | 4 ++-- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 8086591..642770b 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -11,7 +11,7 @@ def __init__(self, *args: Any, length: int | None = None, **kwargs: Any) -> None self.length = length super().__init__(*args, **kwargs) - def deconstruct(self) -> tuple: + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.length is not None: kwargs['length'] = self.length diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index dcca99f..a2ec481 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -13,7 +13,7 @@ def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self) -> tuple: + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions diff --git a/pgvector/django/indexes.py b/pgvector/django/indexes.py index d18330f..e095351 100644 --- a/pgvector/django/indexes.py +++ b/pgvector/django/indexes.py @@ -9,7 +9,7 @@ def __init__(self, *expressions: Any, lists: int | None = None, **kwargs: Any) - self.lists = lists super().__init__(*expressions, **kwargs) - def deconstruct(self) -> tuple: + def deconstruct(self) -> tuple[Any, Any, Any]: path, args, kwargs = super().deconstruct() if self.lists is not None: kwargs['lists'] = self.lists @@ -30,7 +30,7 @@ def __init__(self, *expressions: Any, m: int | None = None, ef_construction: int self.ef_construction = ef_construction super().__init__(*expressions, **kwargs) - def deconstruct(self) -> tuple: + def deconstruct(self) -> tuple[Any, Any, Any]: path, args, kwargs = super().deconstruct() if self.m is not None: kwargs['m'] = self.m diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 531d53c..53a2c37 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -13,7 +13,7 @@ def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any): self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self) -> tuple: + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 7e34963..7a12f78 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -14,7 +14,7 @@ def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self) -> tuple: + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions From b4b684d89c4714350491832cf3533d76476bcf33 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 22:09:50 -0700 Subject: [PATCH 580/879] Improved type hints [skip ci] --- pgvector/bit.py | 2 +- pgvector/halfvec.py | 6 +++--- pgvector/sparsevec.py | 6 +++--- pgvector/vector.py | 6 +++--- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index cb1bec4..ac99a14 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -31,7 +31,7 @@ def __init__(self, value: Any) -> None: def __repr__(self) -> str: return f'Bit({self.to_text()})' - def __eq__(self, other: Any) -> bool: + def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): return self._len == other._len and self._data == other._data return False diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 755b9d6..62be2bb 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -18,7 +18,7 @@ def __init__(self, value: Any) -> None: def __repr__(self) -> str: return f'HalfVector({self.to_list()})' - def __eq__(self, other: Any) -> bool: + def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): return np.array_equal(self.to_numpy(), other.to_numpy()) return False @@ -48,7 +48,7 @@ def from_binary(cls, value: bytes) -> HalfVector: return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) @classmethod - def _to_db(cls, value: Any, dim: int | None = None) -> str | None: + def _to_db(cls, value: object, dim: int | None = None) -> str | None: if value is None: return value @@ -61,7 +61,7 @@ def _to_db(cls, value: Any, dim: int | None = None) -> str | None: return value.to_text() @classmethod - def _to_db_binary(cls, value: Any) -> bytes | None: + def _to_db_binary(cls, value: object) -> bytes | None: if value is None: return value diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 6a00981..d0704ae 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -28,7 +28,7 @@ def __repr__(self) -> str: elements = dict(zip(self._indices, self._values)) return f'SparseVector({elements}, {self._dim})' - def __eq__(self, other: Any) -> bool: + def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values() return False @@ -126,7 +126,7 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars return vec @classmethod - def _to_db(cls, value: Any, dim: int | None = None) -> str | None: + def _to_db(cls, value: object, dim: int | None = None) -> str | None: if value is None: return value @@ -139,7 +139,7 @@ def _to_db(cls, value: Any, dim: int | None = None) -> str | None: return value.to_text() @classmethod - def _to_db_binary(cls, value: Any) -> bytes | None: + def _to_db_binary(cls, value: object) -> bytes | None: if value is None: return value diff --git a/pgvector/vector.py b/pgvector/vector.py index f5ba742..b90222e 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -18,7 +18,7 @@ def __init__(self, value: Any) -> None: def __repr__(self) -> str: return f'Vector({self.to_list()})' - def __eq__(self, other: Any) -> bool: + def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): return np.array_equal(self.to_numpy(), other.to_numpy()) return False @@ -48,7 +48,7 @@ def from_binary(cls, value: bytes) -> Vector: return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) @classmethod - def _to_db(cls, value: Any, dim: int | None = None) -> str | None: + def _to_db(cls, value: object, dim: int | None = None) -> str | None: if value is None: return value @@ -61,7 +61,7 @@ def _to_db(cls, value: Any, dim: int | None = None) -> str | None: return value.to_text() @classmethod - def _to_db_binary(cls, value: Any) -> bytes | None: + def _to_db_binary(cls, value: object) -> bytes | None: if value is None: return value From 8598dc08db2269de49adb15cf440270ddb24cae2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 22:36:51 -0700 Subject: [PATCH 581/879] Improved type hints [skip ci] --- pgvector/peewee/bit.py | 6 +++--- pgvector/peewee/halfvec.py | 10 +++++----- pgvector/peewee/sparsevec.py | 10 +++++----- pgvector/peewee/vector.py | 10 +++++----- pgvector/sqlalchemy/bit.py | 4 ++-- pgvector/sqlalchemy/halfvec.py | 8 ++++---- pgvector/sqlalchemy/sparsevec.py | 8 ++++---- pgvector/sqlalchemy/vector.py | 8 ++++---- 8 files changed, 32 insertions(+), 32 deletions(-) diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index 3128bcc..14c4555 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -12,11 +12,11 @@ def __init__(self, max_length: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.max_length] if self.max_length else None - def _distance(self, op: str, vector: Any) -> Expression: + def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def hamming_distance(self, vector: Any) -> Expression: + def hamming_distance(self, vector: object) -> Expression: return self._distance('<~>', vector) - def jaccard_distance(self, vector: Any) -> Expression: + def jaccard_distance(self, vector: object) -> Expression: return self._distance('<%%>', vector) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 992d8e7..6d55d8f 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -19,17 +19,17 @@ def db_value(self, value: Any) -> str | None: def python_value(self, value: Any) -> HalfVector | None: return HalfVector._from_db(value) - def _distance(self, op: str, vector: Any) -> Expression: + def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector: Any) -> Expression: + def l2_distance(self, vector: object) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector: Any) -> Expression: + def max_inner_product(self, vector: object) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector: Any) -> Expression: + def cosine_distance(self, vector: object) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector: Any) -> Expression: + def l1_distance(self, vector: object) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index dfb403e..7e66dd3 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -19,17 +19,17 @@ def db_value(self, value: Any) -> str | None: def python_value(self, value: Any) -> SparseVector | None: return SparseVector._from_db(value) - def _distance(self, op: str, vector: Any) -> Expression: + def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector: Any) -> Expression: + def l2_distance(self, vector: object) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector: Any) -> Expression: + def max_inner_product(self, vector: object) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector: Any) -> Expression: + def cosine_distance(self, vector: object) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector: Any) -> Expression: + def l1_distance(self, vector: object) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 43b9f4b..f61aa3d 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -20,17 +20,17 @@ def db_value(self, value: Any) -> str | None: def python_value(self, value: Any) -> np.ndarray | None: return Vector._from_db(value) - def _distance(self, op: str, vector: Any) -> Expression: + def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector: Any) -> Expression: + def l2_distance(self, vector: object) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector: Any) -> Expression: + def max_inner_product(self, vector: object) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector: Any) -> Expression: + def cosine_distance(self, vector: object) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector: Any) -> Expression: + def l1_distance(self, vector: object) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 6be2ab9..35b6753 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -29,10 +29,10 @@ def process(value: Any) -> Any: return super().bind_processor(dialect) class comparator_factory(UserDefinedType.Comparator): - def hamming_distance(self, other: Any) -> Operators: + def hamming_distance(self, other: object) -> Operators: return self.op('<~>', return_type=Float)(other) - def jaccard_distance(self, other: Any) -> Operators: + def jaccard_distance(self, other: object) -> Operators: return self.op('<%>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 8c050ec..af13317 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -36,16 +36,16 @@ def process(value: Any) -> HalfVector | None: return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other: Any) -> Operators: + def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: Any) -> Operators: + def max_inner_product(self, other: object) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: Any) -> Operators: + def cosine_distance(self, other: object) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: Any) -> Operators: + def l1_distance(self, other: object) -> Operators: return self.op('<+>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index a809c2a..235a779 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -36,16 +36,16 @@ def process(value: Any) -> SparseVector | None: return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other: Any) -> Operators: + def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: Any) -> Operators: + def max_inner_product(self, other: object) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: Any) -> Operators: + def cosine_distance(self, other: object) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: Any) -> Operators: + def l1_distance(self, other: object) -> Operators: return self.op('<+>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 76f9497..cf01c2f 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -37,16 +37,16 @@ def process(value: Any) -> np.ndarray | None: return process class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other: Any) -> Operators: + def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: Any) -> Operators: + def max_inner_product(self, other: object) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: Any) -> Operators: + def cosine_distance(self, other: object) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: Any) -> Operators: + def l1_distance(self, other: object) -> Operators: return self.op('<+>', return_type=Float)(other) From 3a067c64ec428a989fa294972a50507d0d75b270 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 22:38:09 -0700 Subject: [PATCH 582/879] Improved type hints [skip ci] --- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/sparsevec.py | 2 +- pgvector/peewee/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 6d55d8f..161fb92 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None - def db_value(self, value: Any) -> str | None: + def db_value(self, value: object) -> str | None: return HalfVector._to_db(value) def python_value(self, value: Any) -> HalfVector | None: diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 7e66dd3..24b37b9 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None - def db_value(self, value: Any) -> str | None: + def db_value(self, value: object) -> str | None: return SparseVector._to_db(value) def python_value(self, value: Any) -> SparseVector | None: diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index f61aa3d..aa9509b 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -14,7 +14,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None - def db_value(self, value: Any) -> str | None: + def db_value(self, value: object) -> str | None: return Vector._to_db(value) def python_value(self, value: Any) -> np.ndarray | None: From 03d113511f2fefbfe178cc977325b12398244de6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 1 Jun 2026 22:42:45 -0700 Subject: [PATCH 583/879] Improved type hints [skip ci] --- pgvector/psycopg2/halfvec.py | 5 ++--- pgvector/psycopg2/sparsevec.py | 5 ++--- pgvector/psycopg2/vector.py | 5 ++--- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 5b72f5c..1313938 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,13 +1,12 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type -from typing import Any from .. import HalfVector class HalfvecAdapter: - def __init__(self, value: Any) -> None: + def __init__(self, value: object) -> None: self._value = value - def getquoted(self) -> Any: + def getquoted(self) -> bytes: return adapt(HalfVector._to_db(self._value)).getquoted() diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 073e45b..34a4810 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,13 +1,12 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type -from typing import Any from .. import SparseVector class SparsevecAdapter: - def __init__(self, value: Any) -> None: + def __init__(self, value: object) -> None: self._value = value - def getquoted(self) -> Any: + def getquoted(self) -> bytes: return adapt(SparseVector._to_db(self._value)).getquoted() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 0e2d1e3..2a2c617 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,14 +1,13 @@ import numpy as np from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type -from typing import Any from .. import Vector class VectorAdapter: - def __init__(self, value: Any) -> None: + def __init__(self, value: object) -> None: self._value = value - def getquoted(self) -> Any: + def getquoted(self) -> bytes: return adapt(Vector._to_db(self._value)).getquoted() From bc02e895bd99913a2b007c77679bcc35a6195fdc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 09:42:57 -0700 Subject: [PATCH 584/879] Moved code [skip ci] --- pgvector/psycopg/bit.py | 3 ++- pgvector/psycopg/register.py | 2 -- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index 8cfacbf..a3f3512 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -24,7 +24,8 @@ def dump(self, obj: Bit) -> Buffer | None: return Bit._to_db_binary(obj) -def register_bit_info(context: BaseConnection[Any], info: TypeInfo) -> None: +def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None) -> None: + assert info is not None info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py index 0256b94..15c2fc2 100644 --- a/pgvector/psycopg/register.py +++ b/pgvector/psycopg/register.py @@ -12,7 +12,6 @@ def register_vector(context: Connection[Any]) -> None: register_vector_info(context, info) info = TypeInfo.fetch(context, 'bit') - assert info is not None register_bit_info(context, info) info = TypeInfo.fetch(context, 'halfvec') @@ -29,7 +28,6 @@ async def register_vector_async(context: AsyncConnection[Any]) -> None: register_vector_info(context, info) info = await TypeInfo.fetch(context, 'bit') - assert info is not None register_bit_info(context, info) info = await TypeInfo.fetch(context, 'halfvec') From b1f9c8d8e0fe90509b35d4d4c2207e82e174a300 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:05:02 -0700 Subject: [PATCH 585/879] Improved type hints [skip ci] --- pgvector/halfvec.py | 8 +++++--- pgvector/sparsevec.py | 2 +- pgvector/vector.py | 8 +++++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 62be2bb..7b24037 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,19 +1,21 @@ from __future__ import annotations import numpy as np from struct import pack, unpack_from -from typing import Any class HalfVector: - def __init__(self, value: Any) -> None: + def __init__(self, value: object) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f2': value = np.asarray(value, dtype='>f2') + # for mypy + assert isinstance(value, np.ndarray) + if value.ndim != 1: raise ValueError('expected ndim to be 1') - self._value = value + self._value = np.atleast_1d(value) def __repr__(self) -> str: return f'HalfVector({self.to_list()})' diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index d0704ae..c8da4d2 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -7,7 +7,7 @@ class SparseVector: - def __init__(self, value: Any, dimensions: int | Any = NO_DEFAULT, /) -> None: + def __init__(self, value: dict[int, float] | list[float] | Any, dimensions: int | Any = NO_DEFAULT, /) -> None: if value.__class__.__module__.startswith('scipy.sparse.'): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') diff --git a/pgvector/vector.py b/pgvector/vector.py index b90222e..272da3a 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,19 +1,21 @@ from __future__ import annotations import numpy as np from struct import pack, unpack_from -from typing import Any class Vector: - def __init__(self, value: Any) -> None: + def __init__(self, value: object) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f4': value = np.asarray(value, dtype='>f4') + # for mypy + assert isinstance(value, np.ndarray) + if value.ndim != 1: raise ValueError('expected ndim to be 1') - self._value = value + self._value = np.atleast_1d(value) def __repr__(self) -> str: return f'Vector({self.to_list()})' From 09a98372b02f718f1f9f910b13bfb4a5362e02bd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:14:14 -0700 Subject: [PATCH 586/879] Added overloads [skip ci] --- pgvector/sparsevec.py | 14 +++++++++++++- tests/test_sparse_vector.py | 4 ++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index c8da4d2..3685b31 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,12 +1,24 @@ from __future__ import annotations import numpy as np from struct import pack, unpack_from -from typing import Any +from typing import Any, overload NO_DEFAULT = object() class SparseVector: + @overload + def __init__(self, value: dict[int, float], dimensions: int, /) -> None: + ... + + @overload + def __init__(self, value: list[float], /) -> None: + ... + + @overload + def __init__(self, value: Any, /) -> None: + ... + def __init__(self, value: dict[int, float] | list[float] | Any, dimensions: int | Any = NO_DEFAULT, /) -> None: if value.__class__.__module__.startswith('scipy.sparse.'): if dimensions is not NO_DEFAULT: diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index d580f32..8753e43 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -14,7 +14,7 @@ def test_list(self): def test_list_dimensions(self): with pytest.raises(ValueError) as error: - SparseVector([1, 0, 2, 0, 3, 0], 6) + SparseVector([1, 0, 2, 0, 3, 0], 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' def test_ndarray(self): @@ -40,7 +40,7 @@ def test_coo_array(self): def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: - SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) + SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' def test_coo_matrix(self): From 66c4515858aee43e3893cc78a91df609e55362da Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:19:23 -0700 Subject: [PATCH 587/879] Improved type hints [skip ci] --- pgvector/bit.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index ac99a14..f95a18f 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,12 +1,11 @@ from __future__ import annotations import numpy as np from struct import pack, unpack_from -from typing import Any from warnings import warn class Bit: - def __init__(self, value: Any) -> None: + def __init__(self, value: object) -> None: if isinstance(value, bytes): self._len = 8 * len(value) self._data = value @@ -16,12 +15,18 @@ def __init__(self, value: Any) -> None: else: value = np.asarray(value) + # for mypy + assert isinstance(value, np.ndarray) + if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): warn('expected elements to be boolean', stacklevel=2) value = value.astype(bool) + # for mypy + assert isinstance(value, np.ndarray) + if value.ndim != 1: raise ValueError('expected ndim to be 1') From 110bd425df368a6c8a412bb70c6d32967408aae4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:43:52 -0700 Subject: [PATCH 588/879] Improved type hints [skip ci] --- pgvector/halfvec.py | 6 +++--- pgvector/vector.py | 6 +++--- tests/test_half_vector.py | 6 +++--- tests/test_sqlalchemy.py | 2 +- tests/test_vector.py | 6 +++--- 5 files changed, 13 insertions(+), 13 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 7b24037..2d9507c 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -4,7 +4,7 @@ class HalfVector: - def __init__(self, value: object) -> None: + def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f2': value = np.asarray(value, dtype='>f2') @@ -55,7 +55,7 @@ def _to_db(cls, value: object, dim: int | None = None) -> str | None: return value if not isinstance(value, cls): - value = cls(value) + value = cls(value) # ty: ignore[invalid-argument-type] if dim is not None and value.dimensions() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) @@ -68,7 +68,7 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value if not isinstance(value, cls): - value = cls(value) + value = cls(value) # ty: ignore[invalid-argument-type] return value.to_binary() diff --git a/pgvector/vector.py b/pgvector/vector.py index 272da3a..2af8416 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -4,7 +4,7 @@ class Vector: - def __init__(self, value: object) -> None: + def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f4': value = np.asarray(value, dtype='>f4') @@ -55,7 +55,7 @@ def _to_db(cls, value: object, dim: int | None = None) -> str | None: return value if not isinstance(value, cls): - value = cls(value) + value = cls(value) # ty: ignore[invalid-argument-type] if dim is not None and value.dimensions() != dim: raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) @@ -68,7 +68,7 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value if not isinstance(value, cls): - value = cls(value) + value = cls(value) # ty: ignore[invalid-argument-type] return value.to_binary() diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 78b4977..f7fe347 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -10,7 +10,7 @@ def test_list(self): def test_list_str(self): with pytest.raises(ValueError, match='could not convert string to float'): - HalfVector([1, 'two', 3]) + HalfVector([1, 'two', 3]) # ty: ignore[invalid-argument-type] def test_tuple(self): assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3] @@ -27,12 +27,12 @@ def test_ndarray_same_object(self): def test_ndim_two(self): with pytest.raises(ValueError) as error: - HalfVector([[1, 2], [3, 4]]) + HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected ndim to be 1' def test_ndim_zero(self): with pytest.raises(ValueError) as error: - HalfVector(1) + HalfVector(1) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected ndim to be 1' def test_repr(self): diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f57edc4..e1acdc6 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -601,7 +601,7 @@ async def test_halfvec(self, engine): embedding = [1, 2, 3] session.add(Item(id=1, half_embedding=embedding)) item = await session.get_one(Item, 1) - assert item.half_embedding == HalfVector(embedding) + assert item.half_embedding == HalfVector(embedding) # ty: ignore[invalid-argument-type] await engine.dispose() diff --git a/tests/test_vector.py b/tests/test_vector.py index e5a16fe..789efea 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -10,7 +10,7 @@ def test_list(self): def test_list_str(self): with pytest.raises(ValueError, match='could not convert string to float'): - Vector([1, 'two', 3]) + Vector([1, 'two', 3]) # ty: ignore[invalid-argument-type] def test_tuple(self): assert Vector((1, 2, 3)).to_list() == [1, 2, 3] @@ -27,12 +27,12 @@ def test_ndarray_same_object(self): def test_ndim_two(self): with pytest.raises(ValueError) as error: - Vector([[1, 2], [3, 4]]) + Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected ndim to be 1' def test_ndim_zero(self): with pytest.raises(ValueError) as error: - Vector(1) + Vector(1) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected ndim to be 1' def test_repr(self): From 03dd0e0e76cc4aff49c8dc8c7fa95e448db21382 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:48:25 -0700 Subject: [PATCH 589/879] Fixed error with mypy [skip ci] --- tests/test_sqlalchemy.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index e1acdc6..dbb6f12 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -8,7 +8,7 @@ from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.ext.automap import automap_base -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Session from sqlalchemy.sql import func psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -64,7 +64,9 @@ def asyncpg_connect(dbapi_connection, connection_record): session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class Item(Base): From fbea9d7ed6eb007b404e85628047740473b11ecb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:52:15 -0700 Subject: [PATCH 590/879] Improved type hints [skip ci] --- pgvector/bit.py | 2 +- tests/test_bit.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index f95a18f..08ab910 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -5,7 +5,7 @@ class Bit: - def __init__(self, value: object) -> None: + def __init__(self, value: bytes | str | list[bool] | tuple[bool, ...] | np.ndarray) -> None: if isinstance(value, bytes): self._len = 8 * len(value) self._data = value diff --git a/tests/test_bit.py b/tests/test_bit.py index 5a71642..0e29eb7 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -9,11 +9,11 @@ def test_list(self): def test_list_none(self): with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit([True, None, True]).to_text() == '101' + assert Bit([True, None, True]).to_text() == '101' # ty: ignore[invalid-argument-type] def test_list_int(self): with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit([254, 7, 0]).to_text() == '110' + assert Bit([254, 7, 0]).to_text() == '110' # ty: ignore[invalid-argument-type] def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] @@ -46,12 +46,12 @@ def test_ndarray_uint16(self): def test_ndim_two(self): with pytest.raises(ValueError) as error: - Bit([[True, False], [True, False]]) + Bit([[True, False], [True, False]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected ndim to be 1' def test_ndim_zero(self): with pytest.raises(ValueError) as error: - Bit(True) + Bit(True) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected ndim to be 1' def test_repr(self): From 0178236b34cad23346c36810fc6c9d0133aa311f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 10:53:46 -0700 Subject: [PATCH 591/879] Added comments [skip ci] --- pgvector/halfvec.py | 1 + pgvector/vector.py | 1 + 2 files changed, 2 insertions(+) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 2d9507c..638de94 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -15,6 +15,7 @@ def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None: if value.ndim != 1: raise ValueError('expected ndim to be 1') + # atleast_1d for ty self._value = np.atleast_1d(value) def __repr__(self) -> str: diff --git a/pgvector/vector.py b/pgvector/vector.py index 2af8416..a101a6f 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -15,6 +15,7 @@ def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None: if value.ndim != 1: raise ValueError('expected ndim to be 1') + # atleast_1d for ty self._value = np.atleast_1d(value) def __repr__(self) -> str: From e86d7c40938fdb4ef85f54ea28ba9bbd8baed03b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 11:20:23 -0700 Subject: [PATCH 592/879] Improved type hints [skip ci] --- pgvector/bit.py | 2 +- pgvector/halfvec.py | 3 ++- pgvector/vector.py | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 08ab910..70a5df3 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -44,7 +44,7 @@ def __eq__(self, other: object) -> bool: def to_list(self) -> list[bool]: return self.to_numpy().tolist() - def to_numpy(self) -> np.ndarray: + def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]: return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) def to_text(self) -> str: diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 638de94..3e1745d 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -32,7 +32,8 @@ def dimensions(self) -> int: def to_list(self) -> list[float]: return self._value.tolist() - def to_numpy(self) -> np.ndarray: + def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: + # TODO return native endian return self._value def to_text(self) -> str: diff --git a/pgvector/vector.py b/pgvector/vector.py index a101a6f..67d1232 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -32,7 +32,8 @@ def dimensions(self) -> int: def to_list(self) -> list[float]: return self._value.tolist() - def to_numpy(self) -> np.ndarray: + def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: + # TODO return native endian return self._value def to_text(self) -> str: From 1913c8449c000f1023fde3fbd8b86c8d95dc9ef0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 11:25:53 -0700 Subject: [PATCH 593/879] Improved type hints [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 3685b31..205b0b1 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -66,7 +66,7 @@ def to_list(self) -> list[float]: vec[i] = v return vec - def to_numpy(self) -> np.ndarray: + def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: vec = np.repeat(0.0, self._dim).astype(np.float32) for i, v in zip(self._indices, self._values): vec[i] = v From 6effd62635b8c2928e3bd2ee0b17b960059a7ba7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 11:40:30 -0700 Subject: [PATCH 594/879] Dropped support for Python < 3.10 for examples [skip ci] --- examples/citus/pyproject.toml | 2 +- examples/cohere/pyproject.toml | 2 +- examples/colbert/pyproject.toml | 2 +- examples/colpali/pyproject.toml | 2 +- examples/gensim/pyproject.toml | 2 +- examples/hybrid_search/pyproject.toml | 2 +- examples/image_search/pyproject.toml | 2 +- examples/imagehash/pyproject.toml | 2 +- examples/implicit/pyproject.toml | 2 +- examples/lightfm/pyproject.toml | 2 +- examples/loading/pyproject.toml | 2 +- examples/openai/pyproject.toml | 2 +- examples/rag/pyproject.toml | 2 +- examples/rdkit/pyproject.toml | 2 +- examples/sentence_transformers/pyproject.toml | 2 +- examples/sparse_search/pyproject.toml | 2 +- examples/surprise/pyproject.toml | 2 +- 17 files changed, 17 insertions(+), 17 deletions(-) diff --git a/examples/citus/pyproject.toml b/examples/citus/pyproject.toml index ee40a36..3fb37d4 100644 --- a/examples/citus/pyproject.toml +++ b/examples/citus/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/cohere/pyproject.toml b/examples/cohere/pyproject.toml index f0c88b7..1228b3d 100644 --- a/examples/cohere/pyproject.toml +++ b/examples/cohere/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/colbert/pyproject.toml b/examples/colbert/pyproject.toml index face4d2..7cf1e3f 100644 --- a/examples/colbert/pyproject.toml +++ b/examples/colbert/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/colpali/pyproject.toml b/examples/colpali/pyproject.toml index 23fb23f..20c26a7 100644 --- a/examples/colpali/pyproject.toml +++ b/examples/colpali/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/gensim/pyproject.toml b/examples/gensim/pyproject.toml index 7a33423..4d2865a 100644 --- a/examples/gensim/pyproject.toml +++ b/examples/gensim/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/hybrid_search/pyproject.toml b/examples/hybrid_search/pyproject.toml index b5a904a..130f287 100644 --- a/examples/hybrid_search/pyproject.toml +++ b/examples/hybrid_search/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/image_search/pyproject.toml b/examples/image_search/pyproject.toml index 7644382..8ada0a7 100644 --- a/examples/image_search/pyproject.toml +++ b/examples/image_search/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/imagehash/pyproject.toml b/examples/imagehash/pyproject.toml index cf06c2b..4f0e2b4 100644 --- a/examples/imagehash/pyproject.toml +++ b/examples/imagehash/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/implicit/pyproject.toml b/examples/implicit/pyproject.toml index c03b187..ec305ee 100644 --- a/examples/implicit/pyproject.toml +++ b/examples/implicit/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/lightfm/pyproject.toml b/examples/lightfm/pyproject.toml index c202058..047e3f1 100644 --- a/examples/lightfm/pyproject.toml +++ b/examples/lightfm/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/loading/pyproject.toml b/examples/loading/pyproject.toml index ee40a36..3fb37d4 100644 --- a/examples/loading/pyproject.toml +++ b/examples/loading/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/openai/pyproject.toml b/examples/openai/pyproject.toml index 3e6661a..5665e02 100644 --- a/examples/openai/pyproject.toml +++ b/examples/openai/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/rag/pyproject.toml b/examples/rag/pyproject.toml index fa0dcfd..2e74790 100644 --- a/examples/rag/pyproject.toml +++ b/examples/rag/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/rdkit/pyproject.toml b/examples/rdkit/pyproject.toml index f8c035a..b717809 100644 --- a/examples/rdkit/pyproject.toml +++ b/examples/rdkit/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/sentence_transformers/pyproject.toml b/examples/sentence_transformers/pyproject.toml index b5a904a..130f287 100644 --- a/examples/sentence_transformers/pyproject.toml +++ b/examples/sentence_transformers/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/sparse_search/pyproject.toml b/examples/sparse_search/pyproject.toml index 7927c34..d9adf93 100644 --- a/examples/sparse_search/pyproject.toml +++ b/examples/sparse_search/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ diff --git a/examples/surprise/pyproject.toml b/examples/surprise/pyproject.toml index 94c6f13..bd0cff2 100644 --- a/examples/surprise/pyproject.toml +++ b/examples/surprise/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "example" version = "0.1.0" -requires-python = ">= 3.9" +requires-python = ">= 3.10" [dependency-groups] dev = [ From 35c1abf3e4cb06b29b996336b99d5787edef3ca1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 12:06:09 -0700 Subject: [PATCH 595/879] Removed re-exported classes --- CHANGELOG.md | 1 + pgvector/asyncpg/__init__.py | 8 +------- pgvector/django/__init__.py | 7 +------ pgvector/peewee/__init__.py | 7 +------ pgvector/psycopg/__init__.py | 7 ------- pgvector/psycopg2/__init__.py | 7 +------ pgvector/sqlalchemy/__init__.py | 5 ----- 7 files changed, 5 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ff616b..4ba2b90 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.5.0 (unreleased) - Added type hints +- Removed re-exported classes (use top-level `pgvector` package instead) - Dropped support for Python < 3.10 - Dropped support for SQLAlchemy < 2 diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index c6a3b4e..b3b4440 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,11 +1,5 @@ from .register import register_vector -# TODO remove -from .. import Vector, HalfVector, SparseVector - __all__ = [ - 'register_vector', - 'Vector', - 'HalfVector', - 'SparseVector' + 'register_vector' ] diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 43c64a3..3ed0f83 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -6,9 +6,6 @@ from .sparsevec import SparseVectorField from .vector import VectorField -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ 'VectorExtension', 'VectorField', @@ -22,7 +19,5 @@ 'CosineDistance', 'L1Distance', 'HammingDistance', - 'JaccardDistance', - 'HalfVector', - 'SparseVector' + 'JaccardDistance' ] diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index df21200..f557517 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -3,14 +3,9 @@ from .sparsevec import SparseVectorField from .vector import VectorField -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ 'VectorField', 'HalfVectorField', 'FixedBitField', - 'SparseVectorField', - 'HalfVector', - 'SparseVector' + 'SparseVectorField' ] diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 980af84..8922eda 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,13 +1,6 @@ from .register import register_vector, register_vector_async -# TODO remove -from .. import Bit, HalfVector, SparseVector, Vector - __all__ = [ 'register_vector', 'register_vector_async', - 'Vector', - 'HalfVector', - 'Bit', - 'SparseVector' ] diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 33e5124..b3b4440 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,10 +1,5 @@ from .register import register_vector -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ - 'register_vector', - 'HalfVector', - 'SparseVector' + 'register_vector' ] diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 52adf88..593876f 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -5,17 +5,12 @@ from .vector import VECTOR from .vector import VECTOR as Vector -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ 'Vector', 'VECTOR', 'HALFVEC', 'BIT', 'SPARSEVEC', - 'HalfVector', - 'SparseVector', 'avg', 'sum' ] From b028150fe62dc2065206dab1d1bfea2ed331a134 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 12:07:05 -0700 Subject: [PATCH 596/879] Removed utils package [skip ci] --- CHANGELOG.md | 1 + pgvector/utils/__init__.py | 9 --------- 2 files changed, 1 insertion(+), 9 deletions(-) delete mode 100644 pgvector/utils/__init__.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ba2b90..c535b96 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.5.0 (unreleased) - Added type hints +- Removed `utils` package (use top-level `pgvector` package instead) - Removed re-exported classes (use top-level `pgvector` package instead) - Dropped support for Python < 3.10 - Dropped support for SQLAlchemy < 2 diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py deleted file mode 100644 index 8cdb5d6..0000000 --- a/pgvector/utils/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# TODO remove -from .. import Bit, HalfVector, SparseVector, Vector - -__all__ = [ - 'Vector', - 'HalfVector', - 'Bit', - 'SparseVector' -] From c4c46ab203fd6077ede85704d7c6e5802a2aee2f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 12:10:50 -0700 Subject: [PATCH 597/879] Removed support for tuples from Vector, HalfVector, and Bit constructors [skip ci] --- pgvector/bit.py | 2 +- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- tests/test_bit.py | 3 --- tests/test_half_vector.py | 3 --- tests/test_vector.py | 3 --- 6 files changed, 3 insertions(+), 12 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 70a5df3..9a65593 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -5,7 +5,7 @@ class Bit: - def __init__(self, value: bytes | str | list[bool] | tuple[bool, ...] | np.ndarray) -> None: + def __init__(self, value: bytes | str | list[bool] | np.ndarray) -> None: if isinstance(value, bytes): self._len = 8 * len(value) self._data = value diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 3e1745d..c52a33c 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -4,7 +4,7 @@ class HalfVector: - def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None: + def __init__(self, value: list[float] | np.ndarray) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f2': value = np.asarray(value, dtype='>f2') diff --git a/pgvector/vector.py b/pgvector/vector.py index 67d1232..51b4939 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -4,7 +4,7 @@ class Vector: - def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None: + def __init__(self, value: list[float] | np.ndarray) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f4': value = np.asarray(value, dtype='>f4') diff --git a/tests/test_bit.py b/tests/test_bit.py index 0e29eb7..8663f36 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -15,9 +15,6 @@ def test_list_int(self): with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit([254, 7, 0]).to_text() == '110' # ty: ignore[invalid-argument-type] - def test_tuple(self): - assert Bit((True, False, True)).to_list() == [True, False, True] - def test_str(self): assert Bit('101').to_list() == [True, False, True] diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index f7fe347..b619cdf 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -12,9 +12,6 @@ def test_list_str(self): with pytest.raises(ValueError, match='could not convert string to float'): HalfVector([1, 'two', 3]) # ty: ignore[invalid-argument-type] - def test_tuple(self): - assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3] - def test_ndarray(self): arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] diff --git a/tests/test_vector.py b/tests/test_vector.py index 789efea..90550b0 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -12,9 +12,6 @@ def test_list_str(self): with pytest.raises(ValueError, match='could not convert string to float'): Vector([1, 'two', 3]) # ty: ignore[invalid-argument-type] - def test_tuple(self): - assert Vector((1, 2, 3)).to_list() == [1, 2, 3] - def test_ndarray(self): arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] From ae0f549aa7f58af6a96a50aa98e4345789ac9401 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 14:20:26 -0700 Subject: [PATCH 598/879] Improved type hints [skip ci] --- pgvector/bit.py | 2 +- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 9a65593..cf4fdb0 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -5,7 +5,7 @@ class Bit: - def __init__(self, value: bytes | str | list[bool] | np.ndarray) -> None: + def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None: if isinstance(value, bytes): self._len = 8 * len(value) self._data = value diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index c52a33c..b7a1e77 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -4,7 +4,7 @@ class HalfVector: - def __init__(self, value: list[float] | np.ndarray) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f2': value = np.asarray(value, dtype='>f2') diff --git a/pgvector/vector.py b/pgvector/vector.py index 51b4939..96bcd17 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -4,7 +4,7 @@ class Vector: - def __init__(self, value: list[float] | np.ndarray) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: # asarray still copies if same dtype if not isinstance(value, np.ndarray) or value.dtype != '>f4': value = np.asarray(value, dtype='>f4') From 033923bf9c7dd8b697a778c14f204c1da5142b85 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 14:47:10 -0700 Subject: [PATCH 599/879] Improved internal storage format --- pgvector/halfvec.py | 50 +++++++++++++++++++++++---------------- pgvector/vector.py | 50 +++++++++++++++++++++++---------------- tests/test_half_vector.py | 12 ++++------ tests/test_vector.py | 12 ++++------ 4 files changed, 66 insertions(+), 58 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index b7a1e77..2b3554e 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,46 +1,52 @@ from __future__ import annotations import numpy as np -from struct import pack, unpack_from +import struct class HalfVector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: - # asarray still copies if same dtype - if not isinstance(value, np.ndarray) or value.dtype != '>f2': - value = np.asarray(value, dtype='>f2') - - # for mypy - assert isinstance(value, np.ndarray) - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - # atleast_1d for ty - self._value = np.atleast_1d(value) + if isinstance(value, list): + dim = len(value) + try: + self._value = struct.pack(f'>HH{dim}e', dim, 0, *value) + except struct.error as e: + raise ValueError('expected list[float]') + elif isinstance(value, np.ndarray): + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + # asarray still copies if same dtype + if value.dtype != '>f2': + value = np.asarray(value, dtype='>f2') + + self._value = struct.pack('>HH', value.shape[0], 0) + value.tobytes() + else: + raise ValueError('expected list or ndarray') def __repr__(self) -> str: return f'HalfVector({self.to_list()})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return np.array_equal(self.to_numpy(), other.to_numpy()) + return self.to_binary() == other.to_binary() return False def dimensions(self) -> int: - return len(self._value) + dim, = struct.unpack_from('>H', self._value) + return dim def to_list(self) -> list[float]: - return self._value.tolist() + return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:])) def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: # TODO return native endian - return self._value + return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4) def to_text(self) -> str: - return '[' + ','.join([str(float(v)) for v in self._value]) + ']' + return '[' + ','.join([str(v) for v in self.to_list()]) + ']' def to_binary(self) -> bytes: - return pack('>HH', self.dimensions(), 0) + self._value.tobytes() + return self._value @classmethod def from_text(cls, value: str) -> HalfVector: @@ -48,8 +54,10 @@ def from_text(cls, value: str) -> HalfVector: @classmethod def from_binary(cls, value: bytes) -> HalfVector: - dim, unused = unpack_from('>HH', value) - return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) + # TODO check dimensions/length and unused + vec = cls.__new__(cls) + vec._value = value + return vec @classmethod def _to_db(cls, value: object, dim: int | None = None) -> str | None: diff --git a/pgvector/vector.py b/pgvector/vector.py index 96bcd17..6a0ed89 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,46 +1,52 @@ from __future__ import annotations import numpy as np -from struct import pack, unpack_from +import struct class Vector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: - # asarray still copies if same dtype - if not isinstance(value, np.ndarray) or value.dtype != '>f4': - value = np.asarray(value, dtype='>f4') - - # for mypy - assert isinstance(value, np.ndarray) - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - # atleast_1d for ty - self._value = np.atleast_1d(value) + if isinstance(value, list): + dim = len(value) + try: + self._value = struct.pack(f'>HH{dim}f', dim, 0, *value) + except struct.error as e: + raise ValueError('expected list[float]') + elif isinstance(value, np.ndarray): + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + # asarray still copies if same dtype + if value.dtype != '>f4': + value = np.asarray(value, dtype='>f4') + + self._value = struct.pack('>HH', value.shape[0], 0) + value.tobytes() + else: + raise ValueError('expected list or ndarray') def __repr__(self) -> str: return f'Vector({self.to_list()})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return np.array_equal(self.to_numpy(), other.to_numpy()) + return self.to_binary() == other.to_binary() return False def dimensions(self) -> int: - return len(self._value) + dim, = struct.unpack_from('>H', self._value) + return dim def to_list(self) -> list[float]: - return self._value.tolist() + return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:])) def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: # TODO return native endian - return self._value + return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4) def to_text(self) -> str: - return '[' + ','.join([str(float(v)) for v in self._value]) + ']' + return '[' + ','.join([str(v) for v in self.to_list()]) + ']' def to_binary(self) -> bytes: - return pack('>HH', self.dimensions(), 0) + self._value.tobytes() + return self._value @classmethod def from_text(cls, value: str) -> Vector: @@ -48,8 +54,10 @@ def from_text(cls, value: str) -> Vector: @classmethod def from_binary(cls, value: bytes) -> Vector: - dim, unused = unpack_from('>HH', value) - return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) + # TODO check dimensions/length and unused + vec = cls.__new__(cls) + vec._value = value + return vec @classmethod def _to_db(cls, value: object, dim: int | None = None) -> str | None: diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index b619cdf..fbc1142 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -9,28 +9,24 @@ def test_list(self): assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] def test_list_str(self): - with pytest.raises(ValueError, match='could not convert string to float'): + with pytest.raises(ValueError) as error: HalfVector([1, 'two', 3]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[float]' def test_ndarray(self): arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr - def test_ndarray_same_object(self): - arr = np.array([1, 2, 3], dtype='>f2') - assert HalfVector(arr).to_list() == [1, 2, 3] - assert HalfVector(arr).to_numpy() is arr - def test_ndim_two(self): with pytest.raises(ValueError) as error: HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected ndim to be 1' + assert str(error.value) == 'expected list[float]' def test_ndim_zero(self): with pytest.raises(ValueError) as error: HalfVector(1) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected ndim to be 1' + assert str(error.value) == 'expected list or ndarray' def test_repr(self): assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' diff --git a/tests/test_vector.py b/tests/test_vector.py index 90550b0..cfb7878 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -9,28 +9,24 @@ def test_list(self): assert Vector([1, 2, 3]).to_list() == [1, 2, 3] def test_list_str(self): - with pytest.raises(ValueError, match='could not convert string to float'): + with pytest.raises(ValueError) as error: Vector([1, 'two', 3]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[float]' def test_ndarray(self): arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr - def test_ndarray_same_object(self): - arr = np.array([1, 2, 3], dtype='>f4') - assert Vector(arr).to_list() == [1, 2, 3] - assert Vector(arr).to_numpy() is arr - def test_ndim_two(self): with pytest.raises(ValueError) as error: Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected ndim to be 1' + assert str(error.value) == 'expected list[float]' def test_ndim_zero(self): with pytest.raises(ValueError) as error: Vector(1) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected ndim to be 1' + assert str(error.value) == 'expected list or ndarray' def test_repr(self): assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' From d3e376e18e25dd6d730fb1db1adbb0b463784f69 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 14:59:08 -0700 Subject: [PATCH 600/879] Improved internal storage format for Bit [skip ci] --- pgvector/bit.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index cf4fdb0..8e92653 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -7,8 +7,8 @@ class Bit: def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None: if isinstance(value, bytes): - self._len = 8 * len(value) - self._data = value + _len = 8 * len(value) + _data = value else: if isinstance(value, str): value = [v != '0' for v in value] @@ -30,28 +30,34 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d if value.ndim != 1: raise ValueError('expected ndim to be 1') - self._len = len(value) - self._data = np.packbits(value).tobytes() + _len = len(value) + _data = np.packbits(value).tobytes() + + self._value = pack('>i', _len) + _data def __repr__(self) -> str: return f'Bit({self.to_text()})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return self._len == other._len and self._data == other._data + return self.to_binary() == other.to_binary() return False + def _len(self): + _len, = unpack_from('>i', self._value) + return _len + def to_list(self) -> list[bool]: return self.to_numpy().tolist() def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]: - return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) + return np.unpackbits(np.frombuffer(self._value[4:], dtype=np.uint8), count=self._len()).astype(bool) def to_text(self) -> str: - return ''.join(format(v, '08b') for v in self._data)[:self._len] + return ''.join(format(v, '08b') for v in self._value[4:])[:self._len()] def to_binary(self) -> bytes: - return pack('>i', self._len) + self._data + return self._value @classmethod def from_text(cls, value: str) -> Bit: @@ -63,8 +69,7 @@ def from_binary(cls, value: bytes) -> Bit: raise ValueError('expected bytes') bit = cls.__new__(cls) - bit._len = unpack_from('>i', value)[0] - bit._data = value[4:] + bit._value = value return bit @classmethod From f74d85ccca54683189ef13bc47dacd0ea9ade75b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 15:04:04 -0700 Subject: [PATCH 601/879] Added tests for to_numpy [skip ci] --- tests/test_half_vector.py | 6 ++++++ tests/test_vector.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index fbc1142..5f218ea 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -39,6 +39,12 @@ def test_equality(self): def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 + def test_to_numpy_readonly(self): + arr = HalfVector([1, 2, 3]).to_numpy() + with pytest.raises(ValueError) as error: + arr[0] = 4 + assert str(error.value) == 'assignment destination is read-only' + def test_from_text(self): vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] diff --git a/tests/test_vector.py b/tests/test_vector.py index cfb7878..f67d48e 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -39,6 +39,12 @@ def test_equality(self): def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 + def test_to_numpy_readonly(self): + arr = Vector([1, 2, 3]).to_numpy() + with pytest.raises(ValueError) as error: + arr[0] = 4 + assert str(error.value) == 'assignment destination is read-only' + def test_from_text(self): vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] From 42767e56850f23e1e4bd57a28d74e021e6044ea5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 15:06:39 -0700 Subject: [PATCH 602/879] Removed todos [skip ci] --- pgvector/halfvec.py | 1 - pgvector/vector.py | 1 - 2 files changed, 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 2b3554e..0a302e5 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -39,7 +39,6 @@ def to_list(self) -> list[float]: return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:])) def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: - # TODO return native endian return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4) def to_text(self) -> str: diff --git a/pgvector/vector.py b/pgvector/vector.py index 6a0ed89..58d7fe2 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -39,7 +39,6 @@ def to_list(self) -> list[float]: return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:])) def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: - # TODO return native endian return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4) def to_text(self) -> str: From d5dfd138bf841baf6474e5d69e672f582957740f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 15:07:32 -0700 Subject: [PATCH 603/879] Removed unneeded code [skip ci] --- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 0a302e5..85118ef 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -9,7 +9,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float dim = len(value) try: self._value = struct.pack(f'>HH{dim}e', dim, 0, *value) - except struct.error as e: + except struct.error: raise ValueError('expected list[float]') elif isinstance(value, np.ndarray): if value.ndim != 1: diff --git a/pgvector/vector.py b/pgvector/vector.py index 58d7fe2..1245442 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -9,7 +9,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float dim = len(value) try: self._value = struct.pack(f'>HH{dim}f', dim, 0, *value) - except struct.error as e: + except struct.error: raise ValueError('expected list[float]') elif isinstance(value, np.ndarray): if value.ndim != 1: From f58842c91c5e3332d507587f14ebc566b7619d5c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 15:44:45 -0700 Subject: [PATCH 604/879] Improved tests [skip ci] --- tests/test_half_vector.py | 15 +++++++++------ tests/test_vector.py | 15 +++++++++------ 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 5f218ea..a5c303c 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -8,22 +8,25 @@ class TestHalfVector: def test_list(self): assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] + def test_list_empty(self): + assert HalfVector([]).to_list() == [] + def test_list_str(self): with pytest.raises(ValueError) as error: HalfVector([1, 'two', 3]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' + def test_list_list(self): + with pytest.raises(ValueError) as error: + HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[float]' + def test_ndarray(self): arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr - def test_ndim_two(self): - with pytest.raises(ValueError) as error: - HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected list[float]' - - def test_ndim_zero(self): + def test_int(self): with pytest.raises(ValueError) as error: HalfVector(1) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list or ndarray' diff --git a/tests/test_vector.py b/tests/test_vector.py index f67d48e..337c7a7 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -8,22 +8,25 @@ class TestVector: def test_list(self): assert Vector([1, 2, 3]).to_list() == [1, 2, 3] + def test_list_empty(self): + assert Vector([]).to_list() == [] + def test_list_str(self): with pytest.raises(ValueError) as error: Vector([1, 'two', 3]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' + def test_list_list(self): + with pytest.raises(ValueError) as error: + Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[float]' + def test_ndarray(self): arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr - def test_ndim_two(self): - with pytest.raises(ValueError) as error: - Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected list[float]' - - def test_ndim_zero(self): + def test_int(self): with pytest.raises(ValueError) as error: Vector(1) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list or ndarray' From 60990c27136f6f017d5ed140e21259fe3665acc8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 16:09:02 -0700 Subject: [PATCH 605/879] Changed vector type to return Vector object instead of NumPy array - closes #99 --- CHANGELOG.md | 1 + pgvector/django/vector.py | 33 +++++++++------------------- pgvector/peewee/vector.py | 3 +-- pgvector/psycopg/vector.py | 4 ++-- pgvector/psycopg2/vector.py | 2 +- pgvector/sqlalchemy/vector.py | 3 +-- pgvector/vector.py | 12 +++++----- tests/test_asyncpg.py | 16 +++++--------- tests/test_django.py | 16 ++++++-------- tests/test_peewee.py | 9 ++++---- tests/test_pg8000.py | 6 ++--- tests/test_psycopg.py | 38 +++++++++++++++----------------- tests/test_psycopg2.py | 22 +++++++++---------- tests/test_sqlalchemy.py | 41 +++++++++++++++-------------------- tests/test_sqlmodel.py | 14 +++++------- 15 files changed, 93 insertions(+), 127 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c535b96..760f1e4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.5.0 (unreleased) - Added type hints +- Changed `vector` type to return `Vector` object instead of NumPy array - Removed `utils` package (use top-level `pgvector` package instead) - Removed re-exported classes (use top-level `pgvector` package instead) - Dropped support for Python < 3.10 diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 7a12f78..87d8989 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -1,6 +1,5 @@ from django import forms from django.db.models import Field -import numpy as np from typing import Any from .. import Vector @@ -25,13 +24,16 @@ def db_type(self, connection: Any) -> str: return 'vector' return 'vector(%d)' % self.dimensions - def from_db_value(self, value: Any, expression: Any, connection: Any) -> np.ndarray | None: + def from_db_value(self, value: Any, expression: Any, connection: Any) -> Vector | None: return Vector._from_db(value) - def to_python(self, value: Any) -> np.ndarray | None: - if isinstance(value, list): - return np.array(value, dtype=np.float32) - return Vector._from_db(value) + def to_python(self, value: Any) -> Vector | None: + if value is None or isinstance(value, Vector): + return value + elif isinstance(value, str): + return Vector._from_db(value) + else: + return Vector(value) def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) @@ -39,16 +41,6 @@ def get_prep_value(self, value: Any) -> str | None: def value_to_string(self, obj: Any) -> str | None: return self.get_prep_value(self.value_from_object(obj)) - def validate(self, value: Any, model_instance: Any) -> None: - if isinstance(value, np.ndarray): - value = value.tolist() - super().validate(value, model_instance) - - def run_validators(self, value: Any) -> None: - if isinstance(value, np.ndarray): - value = value.tolist() - super().run_validators(value) - def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: return super().formfield( form_class=VectorFormField if form_class is None else form_class, @@ -59,19 +51,14 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw class VectorWidget(forms.TextInput): def format_value(self, value: Any) -> str | None: - if isinstance(value, np.ndarray): - value = value.tolist() + if isinstance(value, Vector): + value = value.to_list() return super().format_value(value) class VectorFormField(forms.CharField): widget = VectorWidget - def has_changed(self, initial: Any, data: Any) -> bool: - if isinstance(initial, np.ndarray): - initial = initial.tolist() - return super().has_changed(initial, data) - def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index aa9509b..9652071 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,4 +1,3 @@ -import numpy as np from peewee import Expression, Field from typing import Any from .. import Vector @@ -17,7 +16,7 @@ def get_modifiers(self) -> list[int] | None: def db_value(self, value: object) -> str | None: return Vector._to_db(value) - def python_value(self, value: Any) -> np.ndarray | None: + def python_value(self, value: Any) -> Vector | None: return Vector._from_db(value) def _distance(self, op: str, vector: object) -> Expression: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 1492017..dcdc770 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -31,7 +31,7 @@ class VectorLoader(Loader): format = Format.TEXT - def load(self, data: Buffer) -> np.ndarray | None: + def load(self, data: Buffer) -> Vector | None: if isinstance(data, memoryview): data = bytes(data) return Vector._from_db(data.decode('utf8')) @@ -41,7 +41,7 @@ class VectorBinaryLoader(VectorLoader): format = Format.BINARY - def load(self, data: Buffer) -> np.ndarray | None: + def load(self, data: Buffer) -> Vector | None: if isinstance(data, (bytearray, memoryview)): data = bytes(data) return Vector._from_db_binary(data) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 2a2c617..ff5ad99 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -11,7 +11,7 @@ def getquoted(self) -> bytes: return adapt(Vector._to_db(self._value)).getquoted() -def cast_vector(value: str | None, cur: cursor) -> np.ndarray | None: +def cast_vector(value: str | None, cur: cursor) -> Vector | None: return Vector._from_db(value) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index cf01c2f..14d66f3 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,4 +1,3 @@ -import numpy as np from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String from sqlalchemy import Dialect, Operators @@ -32,7 +31,7 @@ def process(value: Any) -> Any: return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: - def process(value: Any) -> np.ndarray | None: + def process(value: Any) -> Vector | None: return Vector._from_db(value) return process diff --git a/pgvector/vector.py b/pgvector/vector.py index 1245442..525eace 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -82,15 +82,15 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value.to_binary() @classmethod - def _from_db(cls, value: str | np.ndarray | None) -> np.ndarray | None: - if value is None or isinstance(value, np.ndarray): + def _from_db(cls, value: str | Vector | None) -> Vector | None: + if value is None or isinstance(value, Vector): return value - return cls.from_text(value).to_numpy().astype(np.float32) + return cls.from_text(value) @classmethod - def _from_db_binary(cls, value: bytes | np.ndarray | None) -> np.ndarray | None: - if value is None or isinstance(value, np.ndarray): + def _from_db_binary(cls, value: bytes | Vector | None) -> Vector | None: + if value is None or isinstance(value, Vector): return value - return cls.from_binary(value).to_numpy().astype(np.float32) + return cls.from_binary(value) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index cabfc5a..45c8c5d 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -20,9 +20,8 @@ async def test_vector(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) - assert res[0]['embedding'].dtype == np.float32 - assert np.array_equal(res[1]['embedding'], embedding2) + assert res[0]['embedding'] == embedding + assert res[1]['embedding'] == Vector(embedding2) assert res[2]['embedding'] is None # ensures binary format is correct @@ -116,10 +115,8 @@ async def test_vector_array(self): await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy()) - assert np.array_equal(res[0]['embeddings'][1], embeddings[1].to_numpy()) - assert np.array_equal(res[1]['embeddings'][0], embeddings2[0]) - assert np.array_equal(res[1]['embeddings'][1], embeddings2[1]) + assert res[0]['embeddings'] == embeddings + assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] await conn.close() @@ -140,7 +137,6 @@ async def init(conn): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) - assert res[0]['embedding'].dtype == np.float32 - assert np.array_equal(res[1]['embedding'], embedding2) + assert res[0]['embedding'] == embedding + assert res[1]['embedding'] == Vector(embedding2) assert res[2]['embedding'] is None diff --git a/tests/test_django.py b/tests/test_django.py index 9e44401..fb37629 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -14,7 +14,7 @@ import numpy as np import os import pgvector.django -from pgvector import HalfVector, SparseVector +from pgvector import HalfVector, SparseVector, Vector from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock @@ -167,12 +167,11 @@ def setup_method(self): def test_vector(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert np.array_equal(item.embedding, [1, 2, 3]) - assert item.embedding.dtype == np.float32 + assert item.embedding == Vector([1, 2, 3]) def test_vector_l2_distance(self): create_items() - distance = L2Distance('embedding', [1, 1, 1]) + distance = L2Distance('embedding', Vector([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)] @@ -295,7 +294,7 @@ def test_vector_avg(self): Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] - assert np.array_equal(avg, [2.5, 3.5, 4.5]) + assert avg == Vector([2.5, 3.5, 4.5]) def test_vector_sum(self): sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] @@ -303,7 +302,7 @@ def test_vector_sum(self): Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] - assert np.array_equal(sum, [5, 7, 9]) + assert sum == Vector([5, 7, 9]) def test_halfvec_avg(self): avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] @@ -349,7 +348,7 @@ def test_vector_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6]) + assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6]) def test_vector_form_save_missing(self): Item(id=1).save() @@ -467,8 +466,7 @@ def test_vector_array(self): # this fails if the driver does not cast arrays item = Item.objects.get(pk=1) - assert np.array_equal(item.embeddings[0], [1, 2, 3]) - assert np.array_equal(item.embeddings[1], [4, 5, 6]) + assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] def test_double_array(self): Item(id=1, double_embedding=[1, 1, 1]).save() diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 64fc009..826608a 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,7 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector import HalfVector, SparseVector +from pgvector import HalfVector, SparseVector, Vector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField db = PostgresqlDatabase('pgvector_python_test') @@ -43,8 +43,7 @@ def setup_method(self): def test_vector(self): Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert np.array_equal(item.embedding, [1, 2, 3]) - assert item.embedding.dtype == np.float32 + assert item.embedding == Vector([1, 2, 3]) def test_vector_l2_distance(self): create_items() @@ -170,7 +169,7 @@ def test_vector_avg(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() - assert np.array_equal(avg, [2.5, 3.5, 4.5]) + assert avg == Vector([2.5, 3.5, 4.5]) def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() @@ -178,7 +177,7 @@ def test_vector_sum(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() - assert np.array_equal(sum, [5, 7, 9]) + assert sum == Vector([5, 7, 9]) def test_halfvec_avg(self): avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 61fbc4c..2e4e884 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -22,8 +22,7 @@ def test_vector(self): conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert np.array_equal(res[0][0], embedding) - assert res[0][0].dtype == np.float32 + assert res[0][0] == Vector([1.5, 2, 3]) assert res[1][0] is None def test_vector_class(self): @@ -31,8 +30,7 @@ def test_vector_class(self): conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert np.array_equal(res[0][0], embedding.to_numpy()) - assert res[0][0].dtype == np.float32 + assert res[0][0] == embedding assert res[1][0] is None def test_halfvec(self): diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index fb6aa4c..71eff2b 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -23,19 +23,18 @@ def test_vector(self): conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() - assert np.array_equal(res[0][0], embedding) - assert res[0][0].dtype == np.float32 + assert res[0][0] == Vector(embedding) assert res[1][0] is None def test_vector_binary_format(self): embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] - assert np.array_equal(res, embedding) + assert res == Vector(embedding) def test_vector_text_format(self): embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] - assert np.array_equal(res, embedding) + assert res == Vector(embedding) def test_vector_binary_format_correct(self): embedding = np.array([1.5, 2, 3]) @@ -46,23 +45,23 @@ def test_vector_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] - assert np.array_equal(res, [3, 2, 1.5]) + assert res == Vector([3, 2, 1.5]) def test_vector_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] - assert np.array_equal(res, [3, 2, 1.5]) + assert res == Vector([3, 2, 1.5]) def test_vector_class_binary_format(self): embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] - assert np.array_equal(res, [1.5, 2, 3]) + assert res == embedding def test_vector_class_text_format(self): embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] - assert np.array_equal(res, [1.5, 2, 3]) + assert res == embedding def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) @@ -146,33 +145,33 @@ def test_text_copy_to(self): assert row[1] == "[1.5,2,3]" def test_binary_copy_to(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: for row in copy.rows(): - assert np.array_equal(Vector.from_binary(row[0]).to_numpy(), embedding) + assert Vector.from_binary(row[0]) == embedding assert HalfVector.from_binary(row[1]) == half_embedding def test_binary_copy_to_set_types(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: copy.set_types(['vector', 'halfvec']) for row in copy.rows(): - assert np.array_equal(row[0], embedding) + assert row[0] == embedding assert row[1] == half_embedding def test_vector_array(self): - embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] conn.execute('INSERT INTO psycopg_items (embeddings) VALUES (%s)', (embeddings,)) res = next(conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id')) - assert np.array_equal(res[0][0], embeddings[0]) - assert np.array_equal(res[0][1], embeddings[1]) + assert res[0][0] == embeddings[0] + assert res[0][1] == embeddings[1] def test_pool(self): def configure(conn): @@ -182,7 +181,7 @@ def configure(conn): with pool.connection() as conn: res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() - assert np.array_equal(res[0], [1, 2, 3]) + assert res[0] == Vector([1, 2, 3]) pool.close() @@ -196,14 +195,13 @@ async def test_async(self): await register_vector_async(conn) - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) await conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) async with conn.cursor() as cur: await cur.execute('SELECT * FROM psycopg_items ORDER BY id') res = await cur.fetchall() - assert np.array_equal(res[0][1], embedding) - assert res[0][1].dtype == np.float32 + assert res[0][1] == embedding assert res[1][1] is None @pytest.mark.asyncio @@ -218,6 +216,6 @@ async def configure(conn): async with conn.cursor() as cur: await cur.execute("SELECT '[1,2,3]'::vector") res = await cur.fetchone() - assert np.array_equal(res[0], [1, 2, 3]) + assert res[0] == Vector([1, 2, 3]) await pool.close() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 7f4932d..7029962 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -26,8 +26,7 @@ def test_vector(self): cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][0], embedding) - assert res[0][0].dtype == np.float32 + assert res[0][0] == Vector(embedding) assert res[1][0] is None def test_vector_class(self): @@ -36,17 +35,16 @@ def test_vector_class(self): cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][0], embedding.to_numpy()) - assert res[0][0].dtype == np.float32 + assert res[0][0] == embedding assert res[1][0] is None def test_halfvec(self): - embedding = [1.5, 2, 3] + embedding = HalfVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == HalfVector([1.5, 2, 3]) + assert res[0][0] == embedding assert res[1][0] is None def test_halfvec_class(self): @@ -64,7 +62,7 @@ def test_bit(self): cur.execute('SELECT binary_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == '101' + assert res[0][0] == embedding assert res[1][0] is None def test_sparsevec(self): @@ -73,17 +71,17 @@ def test_sparsevec(self): cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == SparseVector([1.5, 2, 3]) + assert res[0][0] == embedding assert res[1][0] is None def test_vector_array(self): - embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + embeddings = [Vector([1.5, 2, 3]), Vector([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]) + assert res[0][0] == embeddings[0] + assert res[0][1] == embeddings[1] def test_halfvec_array(self): embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] @@ -129,7 +127,7 @@ def test_pool(self): cur = conn.cursor() cur.execute("SELECT '[1,2,3]'::vector") res = cur.fetchone() - assert np.array_equal(res[0], [1, 2, 3]) + assert res[0] == Vector([1, 2, 3]) finally: pool.putconn(conn) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index dbb6f12..f4ec787 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -182,10 +182,8 @@ def test_orm(self, engine): assert items[0].id % 3 == 1 assert items[1].id % 3 == 2 assert items[2].id % 3 == 0 - assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) - assert items[0].embedding.dtype == np.float32 - assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) - assert items[1].embedding.dtype == np.float32 + assert items[0].embedding == Vector([1.5, 2, 3]) + assert items[1].embedding == Vector([4, 5, 6]) assert items[2].embedding is None def test_vector(self, engine): @@ -193,7 +191,7 @@ def test_vector(self, engine): session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert np.array_equal(item.embedding, [1, 2, 3]) + assert item.embedding == Vector([1, 2, 3]) def test_vector_l2_distance(self, engine): create_items() @@ -435,7 +433,7 @@ def test_avg(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.query(avg(Item.embedding)).one()[0] - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) + assert res == Vector([2.5, 3.5, 4.5]) def test_avg_orm(self, engine): with Session(engine) as session: @@ -444,7 +442,7 @@ def test_avg_orm(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.scalars(select(avg(Item.embedding))).one() - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) + assert res == Vector([2.5, 3.5, 4.5]) def test_sum(self, engine): with Session(engine) as session: @@ -453,7 +451,7 @@ def test_sum(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.query(sum(Item.embedding)).one()[0] - assert np.array_equal(res, np.array([5, 7, 9])) + assert res == Vector([5, 7, 9]) def test_sum_orm(self, engine): with Session(engine) as session: @@ -462,7 +460,7 @@ def test_sum_orm(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.scalars(select(sum(Item.embedding))).one() - assert np.array_equal(res, np.array([5, 7, 9])) + assert res == Vector([5, 7, 9]) def test_bad_dimensions(self, engine): item = Item(embedding=[1, 2]) @@ -515,7 +513,7 @@ def test_automap(self, engine): with Session(engine) as session: session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) item = session.query(AutoItem).first() - assert np.array_equal(item.embedding, [1, 2, 3]) + assert item.embedding == Vector([1, 2, 3]) def test_half_precision(self, engine): create_items() @@ -563,8 +561,7 @@ def test_vector_array(self, engine): # this fails if the driver does not cast arrays item = session.get_one(Item, 1) - assert np.array_equal(item.embeddings[0], [1, 2, 3]) - assert np.array_equal(item.embeddings[1], [4, 5, 6]) + assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] def test_halfvec_array(self, engine): with Session(engine) as session: @@ -587,10 +584,10 @@ async def test_vector(self, engine): async with async_session() as session: async with session.begin(): - embedding = np.array([1, 2, 3]) + embedding = Vector([1, 2, 3]) session.add(Item(id=1, embedding=embedding)) item = await session.get_one(Item, 1) - assert np.array_equal(item.embedding, embedding) + assert item.embedding == embedding await engine.dispose() @@ -600,10 +597,10 @@ async def test_halfvec(self, engine): async with async_session() as session: async with session.begin(): - embedding = [1, 2, 3] + embedding = HalfVector([1, 2, 3]) session.add(Item(id=1, half_embedding=embedding)) item = await session.get_one(Item, 1) - assert item.half_embedding == HalfVector(embedding) # ty: ignore[invalid-argument-type] + assert item.half_embedding == embedding await engine.dispose() @@ -631,10 +628,10 @@ async def test_sparsevec(self, engine): async with async_session() as session: async with session.begin(): - embedding = [1, 2, 3] + embedding = SparseVector([1, 2, 3]) session.add(Item(id=1, sparse_embedding=embedding)) item = await session.get_one(Item, 1) - assert item.sparse_embedding == SparseVector(embedding) + assert item.sparse_embedding == embedding await engine.dispose() @@ -647,7 +644,7 @@ async def test_avg(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = await session.scalars(select(avg(Item.embedding))) - assert np.array_equal(res.one(), [2.5, 3.5, 4.5]) + assert res.first() == Vector([2.5, 3.5, 4.5]) await engine.dispose() @@ -665,12 +662,10 @@ async def test_vector_array(self, engine): async with session.begin(): session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) item = await session.get_one(Item, 1) - assert np.array_equal(item.embeddings[0], [1, 2, 3]) - assert np.array_equal(item.embeddings[1], [4, 5, 6]) + assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) item = await session.get_one(Item, 2) - assert np.array_equal(item.embeddings[0], [1, 2, 3]) - assert np.array_equal(item.embeddings[1], [4, 5, 6]) + assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index d67c735..4bb9d97 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector import HalfVector, SparseVector +from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy.exc import StatementError @@ -65,10 +65,8 @@ def test_orm(self): assert items[0].id == 1 assert items[1].id == 2 assert items[2].id == 3 - assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) # type: ignore - assert items[0].embedding.dtype == np.float32 # type: ignore - assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) # type: ignore - assert items[1].embedding.dtype == np.float32 # type: ignore + assert items[0].embedding == Vector([1.5, 2, 3]) + assert items[1].embedding == Vector([4, 5, 6]) assert items[2].embedding is None def test_vector(self): @@ -76,7 +74,7 @@ def test_vector(self): session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) # type: ignore + assert item.embedding == Vector([1, 2, 3]) def test_vector_l2_distance(self): create_items() @@ -202,7 +200,7 @@ def test_vector_avg(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(avg(Item.embedding))).first() - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) # type: ignore + assert res == Vector([2.5, 3.5, 4.5]) def test_vector_sum(self): with Session(engine) as session: @@ -211,7 +209,7 @@ def test_vector_sum(self): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(sum(Item.embedding))).first() - assert np.array_equal(res, np.array([5, 7, 9])) # type: ignore + assert res == Vector([5, 7, 9]) def test_halfvec_avg(self): with Session(engine) as session: From 2155f99e7c1b747763461902e257e4fdf710f0e3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 16:26:40 -0700 Subject: [PATCH 606/879] Improved naming [skip ci] --- pgvector/bit.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 8e92653..75a51e1 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -7,8 +7,8 @@ class Bit: def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None: if isinstance(value, bytes): - _len = 8 * len(value) - _data = value + length = 8 * len(value) + data = value else: if isinstance(value, str): value = [v != '0' for v in value] @@ -30,10 +30,10 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d if value.ndim != 1: raise ValueError('expected ndim to be 1') - _len = len(value) - _data = np.packbits(value).tobytes() + length = len(value) + data = np.packbits(value).tobytes() - self._value = pack('>i', _len) + _data + self._value = pack('>i', length) + data def __repr__(self) -> str: return f'Bit({self.to_text()})' @@ -43,18 +43,18 @@ def __eq__(self, other: object) -> bool: return self.to_binary() == other.to_binary() return False - def _len(self): - _len, = unpack_from('>i', self._value) - return _len + def _length(self): + length, = unpack_from('>i', self._value) + return length def to_list(self) -> list[bool]: return self.to_numpy().tolist() def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]: - return np.unpackbits(np.frombuffer(self._value[4:], dtype=np.uint8), count=self._len()).astype(bool) + return np.unpackbits(np.frombuffer(self._value[4:], dtype=np.uint8), count=self._length()).astype(bool) def to_text(self) -> str: - return ''.join(format(v, '08b') for v in self._value[4:])[:self._len()] + return ''.join(format(v, '08b') for v in self._value[4:])[:self._length()] def to_binary(self) -> bytes: return self._value From 057740bc1b24d9835116841485d2abc4ec0e377a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 17:24:02 -0700 Subject: [PATCH 607/879] Reduced dependency on NumPy [skip ci] --- pgvector/pg8000/register.py | 7 +++++-- pgvector/psycopg/vector.py | 1 - pgvector/psycopg2/vector.py | 7 +++++-- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index 476495b..e26b083 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -1,4 +1,3 @@ -import numpy as np from pg8000.native import Connection from .. import Vector, HalfVector, SparseVector @@ -12,7 +11,11 @@ def register_vector(conn: Connection) -> None: raise RuntimeError('vector type not found in the database') conn.register_out_adapter(Vector, Vector._to_db) - conn.register_out_adapter(np.ndarray, Vector._to_db) + try: + import numpy as np + conn.register_out_adapter(np.ndarray, Vector._to_db) + except ImportError: + pass conn.register_in_adapter(type_info['vector'], Vector._from_db) if 'halfvec' in type_info: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index dcdc770..cf1694c 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,4 +1,3 @@ -import numpy as np import psycopg from psycopg import BaseConnection from psycopg.adapt import Loader, Dumper diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index ff5ad99..9ad5902 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,4 +1,3 @@ -import numpy as np from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type from .. import Vector @@ -23,5 +22,9 @@ def register_vector_info(oid: int, array_oid: int | None, scope: connection | cu vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector) register_type(vectorarray, scope) - register_adapter(np.ndarray, VectorAdapter) register_adapter(Vector, VectorAdapter) + try: + import numpy as np + register_adapter(np.ndarray, VectorAdapter) + except ImportError: + pass From d13c5dd2d9e638cefc8c0154cb7d83582ac180d3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 18:08:59 -0700 Subject: [PATCH 608/879] Added test for Bit [skip ci] --- tests/test_bit.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index 8663f36..d523619 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -18,6 +18,10 @@ def test_list_int(self): def test_str(self): assert Bit('101').to_list() == [True, False, True] + def test_str_two(self): + # TODO raise + assert Bit('201').to_list() == [True, False, True] + def test_bytes(self): assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' From f770ccffe6f2c4e3fe47fb50a51f1d613c03db75 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 18:19:42 -0700 Subject: [PATCH 609/879] Reduced dependency on NumPy [skip ci] --- pgvector/bit.py | 14 +++++++++++--- tests/test_bit.py | 5 +++-- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 75a51e1..62418cf 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -11,7 +11,15 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d data = value else: if isinstance(value, str): - value = [v != '0' for v in value] + length = len(value) + + if length % 8 != 0: + value += '0'*(8 - (length % 8)) + + try: + data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + except ValueError: + raise ValueError('expected bit string') else: value = np.asarray(value) @@ -30,8 +38,8 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d if value.ndim != 1: raise ValueError('expected ndim to be 1') - length = len(value) - data = np.packbits(value).tobytes() + length = len(value) + data = np.packbits(value).tobytes() self._value = pack('>i', length) + data diff --git a/tests/test_bit.py b/tests/test_bit.py index d523619..ed7d532 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -19,8 +19,9 @@ def test_str(self): assert Bit('101').to_list() == [True, False, True] def test_str_two(self): - # TODO raise - assert Bit('201').to_list() == [True, False, True] + with pytest.raises(ValueError) as error: + Bit('201') + assert str(error.value) == 'expected bit string' def test_bytes(self): assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' From e065b87f8c6de7dc15e57409d434f4d93af074c8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 18:20:51 -0700 Subject: [PATCH 610/879] Fixed lint [skip ci] --- pgvector/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 62418cf..6a980d5 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -14,7 +14,7 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d length = len(value) if length % 8 != 0: - value += '0'*(8 - (length % 8)) + value += '0' * (8 - (length % 8)) try: data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') From 5bc91283168a15126b29e0944993992765978901 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 18:33:02 -0700 Subject: [PATCH 611/879] Improved errors for Bit and reduced dependency on NumPy [skip ci] --- pgvector/bit.py | 19 +++++++++++++------ tests/test_bit.py | 14 ++++++++------ 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 6a980d5..0b0243a 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -10,6 +10,16 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d length = 8 * len(value) data = value else: + if isinstance(value, list): + def bit_value(v: bool) -> str: + if v is True: + return '1' + if v is False: + return '0' + raise ValueError('expected list[bool]') + + value = ''.join([bit_value(v) for v in value]) + if isinstance(value, str): length = len(value) @@ -20,12 +30,7 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') except ValueError: raise ValueError('expected bit string') - else: - value = np.asarray(value) - - # for mypy - assert isinstance(value, np.ndarray) - + elif isinstance(value, np.ndarray): if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): @@ -40,6 +45,8 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d length = len(value) data = np.packbits(value).tobytes() + else: + raise ValueError('expected bytes, str, list, or ndarray') self._value = pack('>i', length) + data diff --git a/tests/test_bit.py b/tests/test_bit.py index ed7d532..4eaf103 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -8,12 +8,14 @@ def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] def test_list_none(self): - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit([True, None, True]).to_text() == '101' # ty: ignore[invalid-argument-type] + with pytest.raises(ValueError) as error: + Bit([True, None, True]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[bool]' def test_list_int(self): - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit([254, 7, 0]).to_text() == '110' # ty: ignore[invalid-argument-type] + with pytest.raises(ValueError) as error: + Bit([254, 7, 0]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[bool]' def test_str(self): assert Bit('101').to_list() == [True, False, True] @@ -49,12 +51,12 @@ def test_ndarray_uint16(self): def test_ndim_two(self): with pytest.raises(ValueError) as error: Bit([[True, False], [True, False]]) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected ndim to be 1' + assert str(error.value) == 'expected list[bool]' def test_ndim_zero(self): with pytest.raises(ValueError) as error: Bit(True) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected ndim to be 1' + assert str(error.value) == 'expected bytes, str, list, or ndarray' def test_repr(self): assert repr(Bit([True, False, True])) == 'Bit(101)' From 44d3e2eb454d0b6ec1aaa1da4edd3137a4229d03 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 19:11:18 -0700 Subject: [PATCH 612/879] Removed dependency on NumPy - closes #146 --- CHANGELOG.md | 1 + pgvector/bit.py | 11 +++++-- pgvector/halfvec.py | 8 +++-- pgvector/sparsevec.py | 6 +++- pgvector/vector.py | 8 +++-- pyproject.toml | 6 ++-- tests/test_asyncpg.py | 7 ++-- tests/test_bit.py | 10 +++++- tests/test_django.py | 8 +++-- tests/test_half_vector.py | 14 ++++++-- tests/test_peewee.py | 3 +- tests/test_pg8000.py | 18 +++++++---- tests/test_psycopg.py | 64 +++++++++++++++++++++---------------- tests/test_psycopg2.py | 27 +++++++--------- tests/test_sparse_vector.py | 29 ++++++++++++++--- tests/test_sqlalchemy.py | 27 ++++++++++------ tests/test_sqlmodel.py | 6 +++- tests/test_vector.py | 14 ++++++-- 18 files changed, 177 insertions(+), 90 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 760f1e4..44b66b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Changed `vector` type to return `Vector` object instead of NumPy array - Removed `utils` package (use top-level `pgvector` package instead) - Removed re-exported classes (use top-level `pgvector` package instead) +- Removed dependency on NumPy - Dropped support for Python < 3.10 - Dropped support for SQLAlchemy < 2 diff --git a/pgvector/bit.py b/pgvector/bit.py index 0b0243a..c8925f0 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,8 +1,12 @@ from __future__ import annotations -import numpy as np from struct import pack, unpack_from from warnings import warn +try: + import numpy as np +except ImportError: + np = None + class Bit: def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None: @@ -30,7 +34,7 @@ def bit_value(v: bool) -> str: data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') except ValueError: raise ValueError('expected bit string') - elif isinstance(value, np.ndarray): + elif np is not None and isinstance(value, np.ndarray): if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): @@ -63,7 +67,8 @@ def _length(self): return length def to_list(self) -> list[bool]: - return self.to_numpy().tolist() + # TODO improve + return [v != '0' for v in self.to_text()] def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]: return np.unpackbits(np.frombuffer(self._value[4:], dtype=np.uint8), count=self._length()).astype(bool) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 85118ef..f317709 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,7 +1,11 @@ from __future__ import annotations -import numpy as np import struct +try: + import numpy as np +except ImportError: + np = None + class HalfVector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: @@ -11,7 +15,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float self._value = struct.pack(f'>HH{dim}e', dim, 0, *value) except struct.error: raise ValueError('expected list[float]') - elif isinstance(value, np.ndarray): + elif np is not None and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 205b0b1..3a17b0f 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,8 +1,12 @@ from __future__ import annotations -import numpy as np from struct import pack, unpack_from from typing import Any, overload +try: + import numpy as np +except ImportError: + np = None + NO_DEFAULT = object() diff --git a/pgvector/vector.py b/pgvector/vector.py index 525eace..e6c6e60 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,7 +1,11 @@ from __future__ import annotations -import numpy as np import struct +try: + import numpy as np +except ImportError: + np = None + class Vector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: @@ -11,7 +15,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float self._value = struct.pack(f'>HH{dim}f', dim, 0, *value) except struct.error: raise ValueError('expected list[float]') - elif isinstance(value, np.ndarray): + elif np is not None and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/pyproject.toml b/pyproject.toml index 3fb3ab5..5a7d99e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,6 @@ authors = [ license = "MIT" requires-python = ">= 3.10" dependencies = [ - "numpy" ] [project.urls] @@ -29,10 +28,13 @@ dev = [ "psycopg2-binary", "pytest", "pytest-asyncio", - "scipy", "SQLAlchemy[asyncio]>=2", "sqlmodel>=0.0.12" ] +dev-optional = [ + "numpy", + "scipy" +] [tool.pytest.ini_options] asyncio_mode = "auto" diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 45c8c5d..b5a1484 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,5 +1,4 @@ import asyncpg -import numpy as np from pgvector import HalfVector, SparseVector, Vector from pgvector.asyncpg import register_vector import pytest @@ -16,7 +15,7 @@ async def test_vector(self): await register_vector(conn) embedding = Vector([1.5, 2, 3]) - embedding2 = np.array([4.5, 5, 6]) + embedding2 = [4.5, 5, 6] await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") @@ -111,7 +110,7 @@ async def test_vector_array(self): embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) - embeddings2 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]] await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") @@ -133,7 +132,7 @@ async def init(conn): await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') embedding = Vector([1.5, 2, 3]) - embedding2 = np.array([1.5, 2, 3]) + embedding2 = [1.5, 2, 3] await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") diff --git a/tests/test_bit.py b/tests/test_bit.py index 4eaf103..2f5293d 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,7 +1,11 @@ -import numpy as np from pgvector import Bit import pytest +try: + import numpy as np +except ImportError: + np = None + class TestBit: def test_list(self): @@ -29,20 +33,24 @@ def test_bytes(self): assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray_unpackbits(self): arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.warns(UserWarning, match='expected elements to be boolean'): diff --git a/tests/test_django.py b/tests/test_django.py index fb37629..37265ba 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -11,13 +11,17 @@ from django.db.migrations.loader import MigrationLoader from django.forms import ModelForm from math import sqrt -import numpy as np import os import pgvector.django from pgvector import HalfVector, SparseVector, Vector from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock +try: + import numpy as np +except ImportError: + np = None + settings.configure( DATABASES={ 'default': { @@ -458,7 +462,7 @@ def test_missing(self): 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() + Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]).save() with connection.cursor() as cursor: from pgvector.psycopg import register_vector diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index a5c303c..0080971 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,8 +1,12 @@ -import numpy as np from pgvector import HalfVector import pytest from struct import pack +try: + import numpy as np +except ImportError: + np = None + class TestHalfVector: def test_list(self): @@ -21,6 +25,7 @@ def test_list_list(self): HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray(self): arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] @@ -42,6 +47,7 @@ def test_equality(self): def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 + @pytest.mark.skipif(np is None, reason='NumPy required') def test_to_numpy_readonly(self): arr = HalfVector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: @@ -51,11 +57,13 @@ def test_to_numpy_readonly(self): def test_from_text(self): vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) def test_from_binary(self): data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 826608a..095a5db 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,5 +1,4 @@ from math import sqrt -import numpy as np from peewee import Model, PostgresqlDatabase, fn from pgvector import HalfVector, SparseVector, Vector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField @@ -197,7 +196,7 @@ def test_halfvec_sum(self): def test_get_or_create(self): 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=Vector([4, 5, 6])) Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) def test_vector_array(self): diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 2e4e884..8281ba5 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -1,8 +1,13 @@ from getpass import getuser -import numpy as np from pgvector import HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection +import pytest + +try: + import numpy as np +except ImportError: + np = None conn = Connection(getuser(), database='pgvector_python_test') @@ -18,19 +23,20 @@ def setup_method(self): conn.run('DELETE FROM pg8000_items') def test_vector(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == Vector([1.5, 2, 3]) + assert res[0][0] == embedding assert res[1][0] is None - def test_vector_class(self): - embedding = Vector([1.5, 2, 3]) + @pytest.mark.skipif(np is None, reason='NumPy required') + def test_vector_numpy(self): + embedding = np.array([1.5, 2, 3]) conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == embedding + assert res[0][0] == Vector([1.5, 2, 3]) assert res[1][0] is None def test_halfvec(self): diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 71eff2b..7ca6257 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,10 +1,14 @@ -import numpy as np from pgvector import Bit, HalfVector, SparseVector, Vector from pgvector.psycopg import register_vector, register_vector_async import psycopg from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest +try: + import numpy as np +except ImportError: + np = None + conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') @@ -19,49 +23,53 @@ def setup_method(self): conn.execute('DELETE FROM psycopg_items') def test_vector(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() - assert res[0][0] == Vector(embedding) + assert res[0][0] == embedding assert res[1][0] is None def test_vector_binary_format(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] - assert res == Vector(embedding) + assert res == embedding def test_vector_text_format(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] - assert res == Vector(embedding) + assert res == embedding def test_vector_binary_format_correct(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] assert res == '[1.5,2,3]' - def test_vector_text_format_non_contiguous(self): - embedding = np.flipud(np.array([1.5, 2, 3])) - assert not embedding.data.contiguous + @pytest.mark.skipif(np is None, reason='NumPy required') + def test_vector_numpy_binary_format(self): + embedding = np.array([1.5, 2, 3]) + res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] + assert res == Vector(embedding) + + @pytest.mark.skipif(np is None, reason='NumPy required') + def test_vector_numpy_text_format(self): + embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] - assert res == Vector([3, 2, 1.5]) + assert res == Vector(embedding) - def test_vector_binary_format_non_contiguous(self): + @pytest.mark.skipif(np is None, reason='NumPy required') + def test_vector_numpy_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] assert res == Vector([3, 2, 1.5]) - def test_vector_class_binary_format(self): - embedding = Vector([1.5, 2, 3]) - res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] - assert res == embedding - - def test_vector_class_text_format(self): - embedding = Vector([1.5, 2, 3]) + @pytest.mark.skipif(np is None, reason='NumPy required') + def test_vector_numpy_text_format_non_contiguous(self): + embedding = np.flipud(np.array([1.5, 2, 3])) + assert not embedding.data.contiguous res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] - assert res == embedding + assert res == Vector([3, 2, 1.5]) def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) @@ -116,26 +124,26 @@ def test_sparsevec_text_format(self): assert res == embedding def test_text_copy_from(self): - embedding = np.array([1.5, 2, 3]) + embedding = [1.5, 2, 3] cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: - copy.write_row([embedding, HalfVector(embedding), '101', SparseVector(embedding)]) + copy.write_row([Vector(embedding), HalfVector(embedding), '101', SparseVector(embedding)]) def test_binary_copy_from(self): - embedding = np.array([1.5, 2, 3]) + embedding = [1.5, 2, 3] cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) + copy.write_row([Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) def test_binary_copy_from_set_types(self): - embedding = np.array([1.5, 2, 3]) + embedding = [1.5, 2, 3] cur = conn.cursor() with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) - copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) + copy.write_row([1, Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) def test_text_copy_to(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 7029962..df5a308 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,9 +1,14 @@ -import numpy as np from pgvector import HalfVector, SparseVector, Vector from pgvector.psycopg2 import register_vector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor from psycopg2.pool import ThreadedConnectionPool +import pytest + +try: + import numpy as np +except ImportError: + np = None conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -21,21 +26,22 @@ def setup_method(self): cur.execute('DELETE FROM psycopg2_items') def test_vector(self): - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == Vector(embedding) + assert res[0][0] == embedding assert res[1][0] is None - def test_vector_class(self): - embedding = Vector([1.5, 2, 3]) + @pytest.mark.skipif(np is None, reason='NumPy required') + def test_vector_numpy(self): + embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == embedding + assert res[0][0] == Vector(embedding) assert res[1][0] is None def test_halfvec(self): @@ -47,15 +53,6 @@ def test_halfvec(self): assert res[0][0] == embedding assert res[1][0] is None - def test_halfvec_class(self): - embedding = HalfVector([1.5, 2, 3]) - cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) - - cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') - res = cur.fetchall() - assert res[0][0] == embedding - assert res[1][0] is None - def test_bit(self): embedding = '101' cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 8753e43..e785643 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,15 +1,24 @@ -import numpy as np from pgvector import SparseVector import pytest -from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix from struct import pack +try: + import numpy as np +except ImportError: + np = None + +try: + from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix +except ImportError: + coo_array = None + class TestSparseVector: def test_list(self): vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] def test_list_dimensions(self): @@ -17,6 +26,7 @@ def test_list_dimensions(self): SparseVector([1, 0, 2, 0, 3, 0], 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray(self): vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) assert vec.to_list() == [1, 0, 2, 0, 3, 0] @@ -32,35 +42,41 @@ def test_dict_no_dimensions(self): SparseVector({0: 1, 2: 2, 4: 3}) assert str(error.value) == 'missing dimensions' + @pytest.mark.skipif(coo_array is None, reason='SciPy required') def test_coo_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + @pytest.mark.skipif(coo_array is None, reason='SciPy required') def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' + @pytest.mark.skipif(coo_array is None, reason='SciPy required') def test_coo_matrix(self): mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + @pytest.mark.skipif(coo_array is None, reason='SciPy required') def test_dok_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + @pytest.mark.skipif(coo_array is None, reason='SciPy required') def test_csr_array(self): arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + @pytest.mark.skipif(coo_array is None, reason='SciPy required') def test_csr_matrix(self): mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) @@ -86,6 +102,7 @@ def test_indices(self): def test_values(self): assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] + @pytest.mark.skipif(np is None or coo_array is None, reason='NumPy and SciPy required') def test_to_coo(self): assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) @@ -99,7 +116,8 @@ def test_from_text(self): assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) def test_from_binary(self): data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) @@ -108,5 +126,6 @@ def test_from_binary(self): assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_binary() == data diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f4ec787..5a9bec3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,6 +1,5 @@ import asyncpg from getpass import getuser -import numpy as np from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest @@ -11,6 +10,11 @@ from sqlalchemy.orm import mapped_column, DeclarativeBase, Session from sqlalchemy.sql import func +try: + import numpy as np +except ImportError: + np = None + psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -165,7 +169,7 @@ def test_core(self, engine): hnsw_index.create(engine) def test_orm(self, engine): - item = Item(embedding=np.array([1.5, 2, 3])) + item = Item(embedding=Vector([1.5, 2, 3])) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -469,6 +473,7 @@ def test_bad_dimensions(self, engine): with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() + @pytest.mark.skipif(np is None, reason='NumPy required') def test_bad_ndim(self, engine): item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: @@ -476,6 +481,7 @@ def test_bad_ndim(self, engine): with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() + @pytest.mark.skipif(np is None, reason='NumPy required') def test_bad_dtype(self, engine): item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: @@ -493,11 +499,11 @@ def test_literal_binds(self, engine): def test_insert(self, engine): with Session(engine) as session: - session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) + session.execute(insert(Item).values(embedding=Vector([1, 2, 3]))) def test_insert_bulk(self, engine): with Session(engine) as session: - session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) + session.execute(insert(Item), [{'embedding': Vector([1, 2, 3])}]) # register_vector in psycopg2 tests change this behavior # def test_insert_text(self): @@ -511,7 +517,7 @@ def test_automap(self, engine): AutoBase.prepare() AutoItem = AutoBase.classes.sqlalchemy_orm_item with Session(engine) as session: - session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) + session.execute(insert(AutoItem), [{'embedding': Vector([1, 2, 3])}]) item = session.query(AutoItem).first() assert item.embedding == Vector([1, 2, 3]) @@ -556,7 +562,7 @@ def setup_method(self): def test_vector_array(self, engine): with Session(engine) as session: - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) session.commit() # this fails if the driver does not cast arrays @@ -565,7 +571,7 @@ def test_vector_array(self, engine): def test_halfvec_array(self, engine): with Session(engine) as session: - session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.add(Item(id=1, half_embeddings=[HalfVector([1, 2, 3]), HalfVector([4, 5, 6])])) session.commit() # this fails if the driver does not cast arrays @@ -664,8 +670,9 @@ async def test_vector_array(self, engine): item = await session.get_one(Item, 1) assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] - session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - item = await session.get_one(Item, 2) - assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] + if np is not None: + session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + item = await session.get_one(Item, 2) + assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 4bb9d97..cc916bf 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,4 +1,3 @@ -import numpy as np from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest @@ -6,6 +5,11 @@ from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text from typing import Any, Optional +try: + import numpy as np +except ImportError: + np = None + engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) # type: ignore diff --git a/tests/test_vector.py b/tests/test_vector.py index 337c7a7..ab6fb1a 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,8 +1,12 @@ -import numpy as np from pgvector import Vector import pytest from struct import pack +try: + import numpy as np +except ImportError: + np = None + class TestVector: def test_list(self): @@ -21,6 +25,7 @@ def test_list_list(self): Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' + @pytest.mark.skipif(np is None, reason='NumPy required') def test_ndarray(self): arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] @@ -42,6 +47,7 @@ def test_equality(self): def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 + @pytest.mark.skipif(np is None, reason='NumPy required') def test_to_numpy_readonly(self): arr = Vector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: @@ -51,11 +57,13 @@ def test_to_numpy_readonly(self): def test_from_text(self): vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) def test_from_binary(self): data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data From 77d878bbb2eae107f644adb62b2ca376903b5dea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 19:11:45 -0700 Subject: [PATCH 613/879] Updated changelog [skip ci] --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 44b66b3..6d89fdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.5.0 (unreleased) +## 1.0.0 (unreleased) - Added type hints - Changed `vector` type to return `Vector` object instead of NumPy array From 2bf00f8c326507335a9faf3a7a058b74e7b48255 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 19:20:47 -0700 Subject: [PATCH 614/879] Updated style [skip ci] --- pgvector/pg8000/register.py | 3 ++- pgvector/psycopg2/vector.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index e26b083..5d44b7c 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -11,12 +11,13 @@ def register_vector(conn: Connection) -> None: raise RuntimeError('vector type not found in the database') conn.register_out_adapter(Vector, Vector._to_db) + conn.register_in_adapter(type_info['vector'], Vector._from_db) + try: import numpy as np conn.register_out_adapter(np.ndarray, Vector._to_db) except ImportError: pass - conn.register_in_adapter(type_info['vector'], Vector._from_db) if 'halfvec' in type_info: conn.register_out_adapter(HalfVector, HalfVector._to_db) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 9ad5902..f8832a6 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -23,6 +23,7 @@ def register_vector_info(oid: int, array_oid: int | None, scope: connection | cu register_type(vectorarray, scope) register_adapter(Vector, VectorAdapter) + try: import numpy as np register_adapter(np.ndarray, VectorAdapter) From f610d5ec1debba32ad662c28665d3bf841e2441e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 20:20:31 -0700 Subject: [PATCH 615/879] Fixed typechecking [skip ci] --- pgvector/bit.py | 5 +++-- pgvector/halfvec.py | 5 +++-- pgvector/sparsevec.py | 2 +- pgvector/vector.py | 5 +++-- tests/test_bit.py | 11 ++++++----- tests/test_django.py | 3 ++- tests/test_half_vector.py | 7 ++++--- tests/test_pg8000.py | 5 +++-- tests/test_psycopg.py | 11 ++++++----- tests/test_psycopg2.py | 5 +++-- tests/test_sparse_vector.py | 22 ++++++++++++---------- tests/test_sqlalchemy.py | 7 ++++--- tests/test_sqlmodel.py | 3 ++- tests/test_vector.py | 7 ++++--- 14 files changed, 56 insertions(+), 42 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index c8925f0..6961a40 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -4,8 +4,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False class Bit: @@ -34,7 +35,7 @@ def bit_value(v: bool) -> str: data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') except ValueError: raise ValueError('expected bit string') - elif np is not None and isinstance(value, np.ndarray): + elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index f317709..deb6f76 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -3,8 +3,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False class HalfVector: @@ -15,7 +16,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float self._value = struct.pack(f'>HH{dim}e', dim, 0, *value) except struct.error: raise ValueError('expected list[float]') - elif np is not None and isinstance(value, np.ndarray): + elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 3a17b0f..22ee39d 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -5,7 +5,7 @@ try: import numpy as np except ImportError: - np = None + pass NO_DEFAULT = object() diff --git a/pgvector/vector.py b/pgvector/vector.py index e6c6e60..1c1d678 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -3,8 +3,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False class Vector: @@ -15,7 +16,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float self._value = struct.pack(f'>HH{dim}f', dim, 0, *value) except struct.error: raise ValueError('expected list[float]') - elif np is not None and isinstance(value, np.ndarray): + elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') diff --git a/tests/test_bit.py b/tests/test_bit.py index 2f5293d..5b33d92 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -3,8 +3,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False class TestBit: @@ -33,24 +34,24 @@ def test_bytes(self): assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_unpackbits(self): arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.warns(UserWarning, match='expected elements to be boolean'): diff --git a/tests/test_django.py b/tests/test_django.py index 37265ba..c64d030 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -19,8 +19,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False settings.configure( DATABASES={ diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 0080971..1a9176f 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -4,8 +4,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False class TestHalfVector: @@ -25,7 +26,7 @@ def test_list_list(self): HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] @@ -47,7 +48,7 @@ def test_equality(self): def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_to_numpy_readonly(self): arr = HalfVector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 8281ba5..c36cccd 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -6,8 +6,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False conn = Connection(getuser(), database='pgvector_python_test') @@ -30,7 +31,7 @@ def test_vector(self): assert res[0][0] == embedding assert res[1][0] is None - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy(self): embedding = np.array([1.5, 2, 3]) conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 7ca6257..19c3421 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -6,8 +6,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) @@ -45,26 +46,26 @@ def test_vector_binary_format_correct(self): res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] assert res == '[1.5,2,3]' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_binary_format(self): embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert res == Vector(embedding) - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_text_format(self): embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert res == Vector(embedding) - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] assert res == Vector([3, 2, 1.5]) - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index df5a308..d0d7c86 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -7,8 +7,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -34,7 +35,7 @@ def test_vector(self): assert res[0][0] == embedding assert res[1][0] is None - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy(self): embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index e785643..55cded6 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -4,13 +4,15 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False try: from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix + SCIPY_AVAILABLE = True except ImportError: - coo_array = None + SCIPY_AVAILABLE = False class TestSparseVector: @@ -26,7 +28,7 @@ def test_list_dimensions(self): SparseVector([1, 0, 2, 0, 3, 0], 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) assert vec.to_list() == [1, 0, 2, 0, 3, 0] @@ -42,41 +44,41 @@ def test_dict_no_dimensions(self): SparseVector({0: 1, 2: 2, 4: 3}) assert str(error.value) == 'missing dimensions' - @pytest.mark.skipif(coo_array is None, reason='SciPy required') + @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(coo_array is None, reason='SciPy required') + @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' - @pytest.mark.skipif(coo_array is None, reason='SciPy required') + @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') def test_coo_matrix(self): mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(coo_array is None, reason='SciPy required') + @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') def test_dok_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(coo_array is None, reason='SciPy required') + @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') def test_csr_array(self): arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(coo_array is None, reason='SciPy required') + @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') def test_csr_matrix(self): mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) @@ -102,7 +104,7 @@ def test_indices(self): def test_values(self): assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] - @pytest.mark.skipif(np is None or coo_array is None, reason='NumPy and SciPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE or SCIPY_AVAILABLE, reason='NumPy and SciPy required') def test_to_coo(self): assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 5a9bec3..a6dc382 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -12,8 +12,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -473,7 +474,7 @@ def test_bad_dimensions(self, engine): with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_bad_ndim(self, engine): item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: @@ -481,7 +482,7 @@ def test_bad_ndim(self, engine): with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_bad_dtype(self, engine): item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index cc916bf..4dcaedd 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -7,8 +7,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: diff --git a/tests/test_vector.py b/tests/test_vector.py index ab6fb1a..9446ca8 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -4,8 +4,9 @@ try: import numpy as np + NUMPY_AVAILABLE = True except ImportError: - np = None + NUMPY_AVAILABLE = False class TestVector: @@ -25,7 +26,7 @@ def test_list_list(self): Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] @@ -47,7 +48,7 @@ def test_equality(self): def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 - @pytest.mark.skipif(np is None, reason='NumPy required') + @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') def test_to_numpy_readonly(self): arr = Vector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: From b1e3669eff86b1e1fd8ecb2c4bb756db21fd6619 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 20:29:34 -0700 Subject: [PATCH 616/879] Added checks to from_binary [skip ci] --- pgvector/bit.py | 5 +++++ pgvector/halfvec.py | 9 ++++++++- pgvector/vector.py | 9 ++++++++- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 6961a40..e2069af 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -89,6 +89,11 @@ def from_binary(cls, value: bytes) -> Bit: if not isinstance(value, bytes): raise ValueError('expected bytes') + length, = unpack_from('>i', value) + + if len(value) != 4 + (length + 7) // 8: + raise ValueError('invalid length') + bit = cls.__new__(cls) bit._value = value return bit diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index deb6f76..d89cdb6 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -58,7 +58,14 @@ def from_text(cls, value: str) -> HalfVector: @classmethod def from_binary(cls, value: bytes) -> HalfVector: - # TODO check dimensions/length and unused + dim, unused = struct.unpack_from('>HH', value) + + if len(value) != 4 + 2 * dim: + raise ValueError('invalid length') + + if unused != 0: + raise ValueError('expected unused to be 0') + vec = cls.__new__(cls) vec._value = value return vec diff --git a/pgvector/vector.py b/pgvector/vector.py index 1c1d678..bd5a2f5 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -58,7 +58,14 @@ def from_text(cls, value: str) -> Vector: @classmethod def from_binary(cls, value: bytes) -> Vector: - # TODO check dimensions/length and unused + dim, unused = struct.unpack_from('>HH', value) + + if len(value) != 4 + 4 * dim: + raise ValueError('invalid length') + + if unused != 0: + raise ValueError('expected unused to be 0') + vec = cls.__new__(cls) vec._value = value return vec From 90ed89c66f1e6ee298f928f2219f971725835e09 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 20:33:38 -0700 Subject: [PATCH 617/879] Added checks to from_binary for SparseVector [skip ci] --- pgvector/sparsevec.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 22ee39d..09b796c 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -129,6 +129,13 @@ def from_text(cls, value: str) -> SparseVector: @classmethod def from_binary(cls, value: bytes) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) + + if len(value) != 12 + 8 * nnz: + raise ValueError('invalid length') + + if unused != 0: + raise ValueError('expected unused to be 0') + indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) return cls._from_parts(int(dim), list(indices), list(values)) From 05fe2cb8866e7eb32a58559eee5aeca8680e6c53 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 20:43:46 -0700 Subject: [PATCH 618/879] Removed unneeded check [skip ci] --- pgvector/bit.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index e2069af..eafd31a 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -42,9 +42,6 @@ def bit_value(v: bool) -> str: warn('expected elements to be boolean', stacklevel=2) value = value.astype(bool) - # for mypy - assert isinstance(value, np.ndarray) - if value.ndim != 1: raise ValueError('expected ndim to be 1') From a9a3ea15af65de6349b06dae40239138f4136aa4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:05:26 -0700 Subject: [PATCH 619/879] Fixed tests --- tests/test_bit.py | 8 ++++---- tests/test_half_vector.py | 8 ++++---- tests/test_pg8000.py | 2 +- tests/test_psycopg.py | 8 ++++---- tests/test_psycopg2.py | 2 +- tests/test_sparse_vector.py | 22 +++++++++++----------- tests/test_sqlalchemy.py | 6 +++--- tests/test_vector.py | 8 ++++---- 8 files changed, 32 insertions(+), 32 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 5b33d92..f87d339 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -34,24 +34,24 @@ def test_bytes(self): assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_unpackbits(self): arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.warns(UserWarning, match='expected elements to be boolean'): diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 1a9176f..d859bd8 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -26,7 +26,7 @@ def test_list_list(self): HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] @@ -48,7 +48,7 @@ def test_equality(self): def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_to_numpy_readonly(self): arr = HalfVector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: @@ -58,13 +58,13 @@ def test_to_numpy_readonly(self): def test_from_text(self): vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) def test_from_binary(self): data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index c36cccd..057dbc9 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -31,7 +31,7 @@ def test_vector(self): assert res[0][0] == embedding assert res[1][0] is None - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy(self): embedding = np.array([1.5, 2, 3]) conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 19c3421..ca04c97 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -46,26 +46,26 @@ def test_vector_binary_format_correct(self): res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] assert res == '[1.5,2,3]' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_binary_format(self): embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert res == Vector(embedding) - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_text_format(self): embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert res == Vector(embedding) - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] assert res == Vector([3, 2, 1.5]) - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index d0d7c86..d12eae0 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -35,7 +35,7 @@ def test_vector(self): assert res[0][0] == embedding assert res[1][0] is None - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy(self): embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 55cded6..4da2f3f 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -19,7 +19,7 @@ class TestSparseVector: def test_list(self): vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] @@ -28,7 +28,7 @@ def test_list_dimensions(self): SparseVector([1, 0, 2, 0, 3, 0], 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) assert vec.to_list() == [1, 0, 2, 0, 3, 0] @@ -44,41 +44,41 @@ def test_dict_no_dimensions(self): SparseVector({0: 1, 2: 2, 4: 3}) assert str(error.value) == 'missing dimensions' - @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') + @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') + @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # ty: ignore[invalid-argument-type] assert str(error.value) == 'extra argument' - @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') + @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_matrix(self): mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') + @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_dok_array(self): arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') + @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_csr_array(self): arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(SCIPY_AVAILABLE, reason='SciPy required') + @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_csr_matrix(self): mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) @@ -104,7 +104,7 @@ def test_indices(self): def test_values(self): assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] - @pytest.mark.skipif(NUMPY_AVAILABLE or SCIPY_AVAILABLE, reason='NumPy and SciPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE or not SCIPY_AVAILABLE, reason='NumPy and SciPy required') def test_to_coo(self): assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) @@ -118,7 +118,7 @@ def test_from_text(self): assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) def test_from_binary(self): @@ -128,6 +128,6 @@ def test_from_binary(self): assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_binary() == data diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index a6dc382..35b8be2 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -474,7 +474,7 @@ def test_bad_dimensions(self, engine): with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_bad_ndim(self, engine): item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: @@ -482,7 +482,7 @@ def test_bad_ndim(self, engine): with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_bad_dtype(self, engine): item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: @@ -671,7 +671,7 @@ async def test_vector_array(self, engine): item = await session.get_one(Item, 1) assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] - if np is not None: + if NUMPY_AVAILABLE: session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) item = await session.get_one(Item, 2) assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] diff --git a/tests/test_vector.py b/tests/test_vector.py index 9446ca8..a8b8e28 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -26,7 +26,7 @@ def test_list_list(self): Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[float]' - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self): arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] @@ -48,7 +48,7 @@ def test_equality(self): def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 - @pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required') + @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_to_numpy_readonly(self): arr = Vector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: @@ -58,13 +58,13 @@ def test_to_numpy_readonly(self): def test_from_text(self): vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) def test_from_binary(self): data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - if np is not None: + if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data From 0e2acf939ea646c078b6236449a18529d8d63ae7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:07:58 -0700 Subject: [PATCH 620/879] Improved tests [skip ci] --- tests/test_bit.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index f87d339..3b31c62 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -22,6 +22,11 @@ def test_list_int(self): Bit([254, 7, 0]) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected list[bool]' + def test_list_list(self): + with pytest.raises(ValueError) as error: + Bit([[True, False], [True, False]]) # ty: ignore[invalid-argument-type] + assert str(error.value) == 'expected list[bool]' + def test_str(self): assert Bit('101').to_list() == [True, False, True] @@ -57,12 +62,7 @@ def test_ndarray_uint16(self): with pytest.warns(UserWarning, match='expected elements to be boolean'): assert Bit(arr).to_text() == '110' - def test_ndim_two(self): - with pytest.raises(ValueError) as error: - Bit([[True, False], [True, False]]) # ty: ignore[invalid-argument-type] - assert str(error.value) == 'expected list[bool]' - - def test_ndim_zero(self): + def test_bool(self): with pytest.raises(ValueError) as error: Bit(True) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected bytes, str, list, or ndarray' From a885b1c4aca5606d770619b3af0ee44d7d555017 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:09:57 -0700 Subject: [PATCH 621/879] Changed warning to error [skip ci] --- pgvector/bit.py | 2 +- tests/test_bit.py | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index eafd31a..ff991c8 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -39,7 +39,7 @@ def bit_value(v: bool) -> str: if value.dtype != np.bool: # skip warning for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): - warn('expected elements to be boolean', stacklevel=2) + raise ValueError('expected elements to be boolean') value = value.astype(bool) if value.ndim != 1: diff --git a/tests/test_bit.py b/tests/test_bit.py index 3b31c62..e96af23 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -53,14 +53,16 @@ def test_ndarray_unpackbits(self): @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint8(self): arr = np.array([254, 7, 0], dtype=np.uint8) - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit(arr).to_text() == '110' + with pytest.raises(ValueError) as error: + Bit(arr) + assert str(error.value) == 'expected elements to be boolean' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit(arr).to_text() == '110' + with pytest.raises(ValueError) as error: + Bit(arr) + assert str(error.value) == 'expected elements to be boolean' def test_bool(self): with pytest.raises(ValueError) as error: From a27a8b0b962792a498edd15185b64942c05530e4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:10:48 -0700 Subject: [PATCH 622/879] Improved code [skip ci] --- pgvector/bit.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index ff991c8..ba116cb 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,6 +1,5 @@ from __future__ import annotations from struct import pack, unpack_from -from warnings import warn try: import numpy as np @@ -37,7 +36,7 @@ def bit_value(v: bool) -> str: raise ValueError('expected bit string') elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.dtype != np.bool: - # skip warning for result of np.unpackbits + # skip error for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): raise ValueError('expected elements to be boolean') value = value.astype(bool) From f89f83bd9ddfdb60537dc707fef97faebd412093 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:49:55 -0700 Subject: [PATCH 623/879] Updated license year [skip ci] --- LICENSE.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/LICENSE.txt b/LICENSE.txt index b612d6d..17e5210 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2025 Andrew Kane +Copyright (c) 2021-2026 Andrew Kane Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal From af6f9fc26b47767e26f88e8b18437a5c3e2ada5b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:51:28 -0700 Subject: [PATCH 624/879] Added optional dependency tests to CI --- .github/workflows/build.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 1ed395f..c1d335e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,3 +24,6 @@ jobs: make sudo make install - run: pytest + + - run: pip install --group dev-optional + - run: pytest From 536c3fbe968699c387b30c6904ad86f3f9073382 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 21:59:31 -0700 Subject: [PATCH 625/879] DRYed tests [skip ci] --- tests/test_asyncpg.py | 35 +++++++++++------------------------ 1 file changed, 11 insertions(+), 24 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index b5a1484..01e6a7e 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -5,14 +5,17 @@ class TestAsyncpg: - @pytest.mark.asyncio - async def test_vector(self): + async def setup_connection(self): conn = await asyncpg.connect(database='pgvector_python_test') await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + await register_vector(conn) await conn.execute('DROP TABLE IF EXISTS asyncpg_items') - await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') + return conn - await register_vector(conn) + @pytest.mark.asyncio + async def test_vector(self): + conn = await self.setup_connection(); + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') embedding = Vector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] @@ -31,13 +34,9 @@ async def test_vector(self): @pytest.mark.asyncio async def test_halfvec(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') + conn = await self.setup_connection(); await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') - await register_vector(conn) - embedding = HalfVector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) @@ -55,13 +54,9 @@ async def test_halfvec(self): @pytest.mark.asyncio async def test_bit(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') + conn = await self.setup_connection(); await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') - await register_vector(conn) - embedding = asyncpg.BitString('101') # type: ignore await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) @@ -78,13 +73,9 @@ async def test_bit(self): @pytest.mark.asyncio async def test_sparsevec(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') + conn = await self.setup_connection(); await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') - await register_vector(conn) - embedding = SparseVector([1.5, 2, 3]) await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) @@ -100,13 +91,9 @@ async def test_sparsevec(self): @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') + conn = await self.setup_connection(); await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') - await register_vector(conn) - embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) From d31d279784a16026c8787730a9d375f191a6161b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:06:20 -0700 Subject: [PATCH 626/879] Improved test [skip ci] --- tests/test_asyncpg.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 01e6a7e..5beb8a9 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -3,6 +3,12 @@ from pgvector.asyncpg import register_vector import pytest +try: + import numpy as np + NUMPY_AVAILABLE = True +except ImportError: + NUMPY_AVAILABLE = False + class TestAsyncpg: async def setup_connection(self): @@ -19,12 +25,14 @@ async def test_vector(self): embedding = Vector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) + embedding3 = np.array([7.5, 8, 9]) if NUMPY_AVAILABLE else [7.5, 8, 9] + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), (NULL)", embedding, embedding2, embedding3) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embedding'] == embedding assert res[1]['embedding'] == Vector(embedding2) - assert res[2]['embedding'] is None + assert res[2]['embedding'] == Vector(embedding3) + assert res[3]['embedding'] is None # ensures binary format is correct text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") From 061de102b9a2538544c9d9ab2cdc1100d20cd937 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:09:50 -0700 Subject: [PATCH 627/879] Improved tests [skip ci] --- tests/test_asyncpg.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 5beb8a9..4901616 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -26,7 +26,8 @@ async def test_vector(self): embedding = Vector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] embedding3 = np.array([7.5, 8, 9]) if NUMPY_AVAILABLE else [7.5, 8, 9] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), (NULL)", embedding, embedding2, embedding3) + embedding4 = None + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), ($4)", embedding, embedding2, embedding3, embedding4) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embedding'] == embedding @@ -47,7 +48,8 @@ async def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) + embedding3 = None + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)", embedding, embedding2, embedding3) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embedding'] == embedding @@ -66,7 +68,8 @@ async def test_bit(self): await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') embedding = asyncpg.BitString('101') # type: ignore - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding2 = None + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embedding'].as_string() == '101' @@ -85,7 +88,8 @@ async def test_sparsevec(self): await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') embedding = SparseVector([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding2 = None + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embedding'] == embedding @@ -108,9 +112,15 @@ async def test_vector_array(self): embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]] await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) + if NUMPY_AVAILABLE: + embeddings3 = [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[])", embeddings3[0], embeddings3[1]) + res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embeddings'] == embeddings assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] + if NUMPY_AVAILABLE: + assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] await conn.close() @@ -128,7 +138,8 @@ async def init(conn): embedding = Vector([1.5, 2, 3]) embedding2 = [1.5, 2, 3] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) + embedding3 = None + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)", embedding, embedding2, embedding3) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") assert res[0]['embedding'] == embedding From 57f6d942b71eda3e87ebcb523eecc72c89d55c91 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:13:35 -0700 Subject: [PATCH 628/879] Improved tests [skip ci] --- tests/test_pg8000.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 057dbc9..0c29cf6 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -25,7 +25,8 @@ def setup_method(self): def test_vector(self): embedding = Vector([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') assert res[0][0] == embedding @@ -42,7 +43,8 @@ def test_vector_numpy(self): def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') assert res[0][0] == embedding @@ -50,7 +52,8 @@ def test_halfvec(self): def test_bit(self): embedding = '101' - conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') assert res[0][0] == '101' @@ -58,7 +61,8 @@ def test_bit(self): def test_sparsevec(self): embedding = SparseVector([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') assert res[0][0] == embedding From 061ae69816ca77a25647a4400b7340b6f4834933 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:18:08 -0700 Subject: [PATCH 629/879] Improved tests [skip ci] --- tests/test_asyncpg.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 4901616..153b98c 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -15,12 +15,12 @@ async def setup_connection(self): conn = await asyncpg.connect(database='pgvector_python_test') await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') await register_vector(conn) - await conn.execute('DROP TABLE IF EXISTS asyncpg_items') return conn @pytest.mark.asyncio async def test_vector(self): conn = await self.setup_connection(); + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') embedding = Vector([1.5, 2, 3]) @@ -44,6 +44,7 @@ async def test_vector(self): @pytest.mark.asyncio async def test_halfvec(self): conn = await self.setup_connection(); + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') embedding = HalfVector([1.5, 2, 3]) @@ -65,6 +66,7 @@ async def test_halfvec(self): @pytest.mark.asyncio async def test_bit(self): conn = await self.setup_connection(); + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') embedding = asyncpg.BitString('101') # type: ignore @@ -85,6 +87,7 @@ async def test_bit(self): @pytest.mark.asyncio async def test_sparsevec(self): conn = await self.setup_connection(); + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') embedding = SparseVector([1.5, 2, 3]) @@ -104,6 +107,7 @@ async def test_sparsevec(self): @pytest.mark.asyncio async def test_vector_array(self): conn = await self.setup_connection(); + await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] From 605482b5adfa84ff91ff7e141310fbcc0f8fd5b1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:27:18 -0700 Subject: [PATCH 630/879] Added Bit test [skip ci] --- tests/test_bit.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index e96af23..ff708d0 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,5 +1,6 @@ from pgvector import Bit import pytest +import random try: import numpy as np @@ -69,6 +70,10 @@ def test_bool(self): Bit(True) # ty: ignore[invalid-argument-type] assert str(error.value) == 'expected bytes, str, list, or ndarray' + def test_random(self): + value = ''.join(random.choices(['0', '1'], k=random.randint(1024, 2048))) + assert Bit(value).to_text() == value + def test_repr(self): assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' From a920f0414be5ab52d5fb6e9fb0706e5878c9c9b3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:30:16 -0700 Subject: [PATCH 631/879] Improved tests [skip ci] --- tests/test_sparse_vector.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 4da2f3f..9a18f98 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -23,6 +23,9 @@ def test_list(self): assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] + def test_list_empty(self): + assert SparseVector([]).to_list() == [] + def test_list_dimensions(self): with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) # ty: ignore[invalid-argument-type] @@ -39,6 +42,9 @@ def test_dict(self): assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + def test_dict_empty(self): + assert SparseVector({}, 0).to_list() == [] + def test_dict_no_dimensions(self): with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) From 8a3dadb1b3f2ae260a0211e1cd3245bfa10a5cf9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:33:12 -0700 Subject: [PATCH 632/879] Simplified tests [skip ci] --- tests/test_pg8000.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 0c29cf6..75f9bee 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -25,21 +25,14 @@ def setup_method(self): def test_vector(self): embedding = Vector([1.5, 2, 3]) - embedding2 = None - conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) + embedding2 = np.array([4.5, 5, 6]) if NUMPY_AVAILABLE else [4.5, 5, 6] + embedding3 = None + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2), (:embedding3)', embedding=embedding, embedding2=embedding2, embedding3=embedding3) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') assert res[0][0] == embedding - assert res[1][0] is None - - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_vector_numpy(self): - embedding = np.array([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) - - res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == Vector([1.5, 2, 3]) - assert res[1][0] is None + assert res[1][0] == Vector(embedding2) + assert res[2][0] is None def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) From 41305ded9abd202dc80fcbc5b57f54e386e50b2b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:34:49 -0700 Subject: [PATCH 633/879] Improved test [skip ci] --- tests/test_psycopg.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index ca04c97..80ab538 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -25,7 +25,8 @@ def setup_method(self): def test_vector(self): embedding = Vector([1.5, 2, 3]) - conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) + embedding2 = None + conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (%s)', (embedding, embedding2)) res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() assert res[0][0] == embedding From 706ec697775252619465bfb523ae8d3d10661837 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:37:05 -0700 Subject: [PATCH 634/879] Improved test [skip ci] --- tests/test_psycopg2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index d12eae0..466fef1 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -28,7 +28,8 @@ def setup_method(self): def test_vector(self): embedding = Vector([1.5, 2, 3]) - cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) + embedding2 = None + cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (%s)', (embedding, embedding2)) cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() From 058ba28d4408f09d2f7184fa76d0cfe9f0be950f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:40:43 -0700 Subject: [PATCH 635/879] Fixed test [skip ci] --- tests/test_pg8000.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 75f9bee..718d0f7 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -25,13 +25,13 @@ def setup_method(self): def test_vector(self): embedding = Vector([1.5, 2, 3]) - embedding2 = np.array([4.5, 5, 6]) if NUMPY_AVAILABLE else [4.5, 5, 6] + embedding2 = np.array([4.5, 5, 6]) if NUMPY_AVAILABLE else Vector([4.5, 5, 6]) embedding3 = None conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2), (:embedding3)', embedding=embedding, embedding2=embedding2, embedding3=embedding3) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') assert res[0][0] == embedding - assert res[1][0] == Vector(embedding2) + assert res[1][0] == Vector([4.5, 5, 6]) assert res[2][0] is None def test_halfvec(self): From e67e8462323dafb5e2b016c02ccd3c3b9a0d6a25 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:52:09 -0700 Subject: [PATCH 636/879] Improved code [skip ci] --- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index d89cdb6..87fafdc 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -47,7 +47,7 @@ def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4) def to_text(self) -> str: - return '[' + ','.join([str(v) for v in self.to_list()]) + ']' + return f'[{','.join([str(v) for v in self.to_list()])}]' def to_binary(self) -> bytes: return self._value diff --git a/pgvector/vector.py b/pgvector/vector.py index bd5a2f5..3049e3c 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -47,7 +47,7 @@ def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4) def to_text(self) -> str: - return '[' + ','.join([str(v) for v in self.to_list()]) + ']' + return f'[{','.join([str(v) for v in self.to_list()])}]' def to_binary(self) -> bytes: return self._value From 6deaa5f5ccbd563bf86c4fa38ab40b636631ddc7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 22:59:07 -0700 Subject: [PATCH 637/879] Improved Bit constructor code [skip ci] --- pgvector/bit.py | 49 ++++++++++++++++++++++++------------------------- 1 file changed, 24 insertions(+), 25 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index ba116cb..1052cee 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -13,7 +13,7 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d if isinstance(value, bytes): length = 8 * len(value) data = value - else: + elif isinstance(value, (list, str)): if isinstance(value, list): def bit_value(v: bool) -> str: if v is True: @@ -24,30 +24,29 @@ def bit_value(v: bool) -> str: value = ''.join([bit_value(v) for v in value]) - if isinstance(value, str): - length = len(value) - - if length % 8 != 0: - value += '0' * (8 - (length % 8)) - - try: - data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') - except ValueError: - raise ValueError('expected bit string') - elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): - if value.dtype != np.bool: - # skip error for result of np.unpackbits - if value.dtype != np.uint8 or np.any(value > 1): - raise ValueError('expected elements to be boolean') - value = value.astype(bool) - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - length = len(value) - data = np.packbits(value).tobytes() - else: - raise ValueError('expected bytes, str, list, or ndarray') + length = len(value) + + if length % 8 != 0: + value += '0' * (8 - (length % 8)) + + try: + data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + except ValueError: + raise ValueError('expected bit string') + elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): + if value.dtype != np.bool: + # skip error for result of np.unpackbits + if value.dtype != np.uint8 or np.any(value > 1): + raise ValueError('expected elements to be boolean') + value = value.astype(bool) + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + length = len(value) + data = np.packbits(value).tobytes() + else: + raise ValueError('expected bytes, str, list, or ndarray') self._value = pack('>i', length) + data From 634215cfc921e0f778d57c18f073f1fa258a3daf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 23:00:43 -0700 Subject: [PATCH 638/879] Fixed quoting for Python < 3.12 [skip ci] --- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 87fafdc..4728065 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -47,7 +47,7 @@ def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4) def to_text(self) -> str: - return f'[{','.join([str(v) for v in self.to_list()])}]' + return f'[{",".join([str(v) for v in self.to_list()])}]' def to_binary(self) -> bytes: return self._value diff --git a/pgvector/vector.py b/pgvector/vector.py index 3049e3c..257ad8f 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -47,7 +47,7 @@ def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4) def to_text(self) -> str: - return f'[{','.join([str(v) for v in self.to_list()])}]' + return f'[{",".join([str(v) for v in self.to_list()])}]' def to_binary(self) -> bytes: return self._value From d96381375dd84e0f9ff907b1745ff8b0627ce4c2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 23:21:13 -0700 Subject: [PATCH 639/879] Updated style [skip ci] --- tests/test_asyncpg.py | 36 ++++++++++++++++++------------------ tests/test_django.py | 2 +- tests/test_psycopg.py | 16 ++++++++-------- tests/test_sqlalchemy.py | 8 ++++---- 4 files changed, 31 insertions(+), 31 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 153b98c..2b96de0 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -27,16 +27,16 @@ async def test_vector(self): embedding2 = [4.5, 5, 6] embedding3 = np.array([7.5, 8, 9]) if NUMPY_AVAILABLE else [7.5, 8, 9] embedding4 = None - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), ($4)", embedding, embedding2, embedding3, embedding4) + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), ($4)', embedding, embedding2, embedding3, embedding4) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'] == embedding assert res[1]['embedding'] == Vector(embedding2) assert res[2]['embedding'] == Vector(embedding3) assert res[3]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @@ -50,15 +50,15 @@ async def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] embedding3 = None - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)", embedding, embedding2, embedding3) + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)', embedding, embedding2, embedding3) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'] == embedding assert res[1]['embedding'] == HalfVector(embedding2) assert res[2]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @@ -71,15 +71,15 @@ async def test_bit(self): embedding = asyncpg.BitString('101') # type: ignore embedding2 = None - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)", embedding, embedding2) + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)', embedding, embedding2) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'].as_string() == '101' assert res[0]['embedding'].to_int() == 5 assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '101' await conn.close() @@ -92,14 +92,14 @@ async def test_sparsevec(self): embedding = SparseVector([1.5, 2, 3]) embedding2 = None - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)", embedding, embedding2) + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)', embedding, embedding2) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'] == embedding assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '{1:1.5,2:2,3:3}/3' await conn.close() @@ -111,16 +111,16 @@ async def test_vector_array(self): await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] - await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES ($1)', embeddings) embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]] - await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings2[0], embeddings2[1]) if NUMPY_AVAILABLE: embeddings3 = [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[])", embeddings3[0], embeddings3[1]) + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings3[0], embeddings3[1]) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embeddings'] == embeddings assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] if NUMPY_AVAILABLE: @@ -143,9 +143,9 @@ async def init(conn): embedding = Vector([1.5, 2, 3]) embedding2 = [1.5, 2, 3] embedding3 = None - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)", embedding, embedding2, embedding3) + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)', embedding, embedding2, embedding3) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'] == embedding assert res[1]['embedding'] == Vector(embedding2) assert res[2]['embedding'] is None diff --git a/tests/test_django.py b/tests/test_django.py index c64d030..6bcf67a 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -131,7 +131,7 @@ class Migration(migrations.Migration): sql_statements = loader.collect_sql([(migration, False)]) with connection.cursor() as cursor: - cursor.execute("DROP TABLE IF EXISTS django_app_item") + cursor.execute('DROP TABLE IF EXISTS django_app_item') cursor.execute('\n'.join(sql_statements)) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 80ab538..13dc7c8 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -128,19 +128,19 @@ def test_sparsevec_text_format(self): def test_text_copy_from(self): embedding = [1.5, 2, 3] cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN') as copy: copy.write_row([Vector(embedding), HalfVector(embedding), '101', SparseVector(embedding)]) def test_binary_copy_from(self): embedding = [1.5, 2, 3] cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: copy.write_row([Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) def test_binary_copy_from_set_types(self): embedding = [1.5, 2, 3] cur = conn.cursor() - with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) copy.write_row([1, Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) @@ -149,17 +149,17 @@ def test_text_copy_to(self): half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding) TO STDOUT') as copy: for row in copy.rows(): - assert row[0] == "[1.5,2,3]" - assert row[1] == "[1.5,2,3]" + assert row[0] == '[1.5,2,3]' + assert row[1] == '[1.5,2,3]' def test_binary_copy_to(self): embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)') as copy: for row in copy.rows(): assert Vector.from_binary(row[0]) == embedding assert HalfVector.from_binary(row[1]) == half_embedding @@ -169,7 +169,7 @@ def test_binary_copy_to_set_types(self): half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)') as copy: copy.set_types(['vector', 'halfvec']) for row in copy.rows(): assert row[0] == embedding diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 35b8be2..3d2c1d3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -20,7 +20,7 @@ psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -@event.listens_for(psycopg2_type_engine, "connect") +@event.listens_for(psycopg2_type_engine, 'connect') def psycopg2_connect(dbapi_connection, connection_record): from pgvector.psycopg2 import register_vector register_vector(dbapi_connection) @@ -32,7 +32,7 @@ def psycopg2_connect(dbapi_connection, connection_record): psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') -@event.listens_for(psycopg_type_engine, "connect") +@event.listens_for(psycopg_type_engine, 'connect') def psycopg_connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector register_vector(dbapi_connection) @@ -42,7 +42,7 @@ def psycopg_connect(dbapi_connection, connection_record): psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') -@event.listens_for(psycopg_async_type_engine.sync_engine, "connect") +@event.listens_for(psycopg_async_type_engine.sync_engine, 'connect') def psycopg_async_connect(dbapi_connection, connection_record): from pgvector.psycopg import register_vector_async dbapi_connection.run_async(register_vector_async) @@ -52,7 +52,7 @@ def psycopg_async_connect(dbapi_connection, connection_record): asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') -@event.listens_for(asyncpg_type_engine.sync_engine, "connect") +@event.listens_for(asyncpg_type_engine.sync_engine, 'connect') def asyncpg_connect(dbapi_connection, connection_record): from pgvector.asyncpg import register_vector dbapi_connection.run_async(register_vector) From 6cfc910fd60fa90dfcee995fc86a535e7cc2aca1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 23:45:23 -0700 Subject: [PATCH 640/879] Updated style [skip ci] --- tests/test_asyncpg.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 2b96de0..3df3d39 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -19,7 +19,7 @@ async def setup_connection(self): @pytest.mark.asyncio async def test_vector(self): - conn = await self.setup_connection(); + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') @@ -43,7 +43,7 @@ async def test_vector(self): @pytest.mark.asyncio async def test_halfvec(self): - conn = await self.setup_connection(); + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') @@ -65,7 +65,7 @@ async def test_halfvec(self): @pytest.mark.asyncio async def test_bit(self): - conn = await self.setup_connection(); + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') @@ -86,7 +86,7 @@ async def test_bit(self): @pytest.mark.asyncio async def test_sparsevec(self): - conn = await self.setup_connection(); + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') @@ -106,7 +106,7 @@ async def test_sparsevec(self): @pytest.mark.asyncio async def test_vector_array(self): - conn = await self.setup_connection(); + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') From 965ca1774271263ccdc586417196b6d49d9cb4ae Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 23:45:37 -0700 Subject: [PATCH 641/879] Updated style [skip ci] --- pgvector/psycopg/bit.py | 2 -- pgvector/psycopg/halfvec.py | 4 ---- pgvector/psycopg/sparsevec.py | 4 ---- pgvector/psycopg/vector.py | 4 ---- 4 files changed, 14 deletions(-) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index a3f3512..f0f4021 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -9,7 +9,6 @@ class BitDumper(Dumper): - format = Format.TEXT def dump(self, obj: Bit) -> Buffer | None: @@ -17,7 +16,6 @@ def dump(self, obj: Bit) -> Buffer | None: class BitBinaryDumper(BitDumper): - format = Format.BINARY def dump(self, obj: Bit) -> Buffer | None: diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 5ded56d..6eaf501 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -9,7 +9,6 @@ class HalfVectorDumper(Dumper): - format = Format.TEXT def dump(self, obj: HalfVector) -> Buffer | None: @@ -18,7 +17,6 @@ def dump(self, obj: HalfVector) -> Buffer | None: class HalfVectorBinaryDumper(HalfVectorDumper): - format = Format.BINARY def dump(self, obj: HalfVector) -> Buffer | None: @@ -26,7 +24,6 @@ def dump(self, obj: HalfVector) -> Buffer | None: class HalfVectorLoader(Loader): - format = Format.TEXT def load(self, data: Buffer) -> HalfVector | None: @@ -36,7 +33,6 @@ def load(self, data: Buffer) -> HalfVector | None: class HalfVectorBinaryLoader(HalfVectorLoader): - format = Format.BINARY def load(self, data: Buffer) -> HalfVector | None: diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index b209b8f..16eaf5a 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -9,7 +9,6 @@ class SparseVectorDumper(Dumper): - format = Format.TEXT def dump(self, obj: SparseVector) -> Buffer | None: @@ -18,7 +17,6 @@ def dump(self, obj: SparseVector) -> Buffer | None: class SparseVectorBinaryDumper(SparseVectorDumper): - format = Format.BINARY def dump(self, obj: SparseVector) -> Buffer | None: @@ -26,7 +24,6 @@ def dump(self, obj: SparseVector) -> Buffer | None: class SparseVectorLoader(Loader): - format = Format.TEXT def load(self, data: Buffer) -> SparseVector | None: @@ -36,7 +33,6 @@ def load(self, data: Buffer) -> SparseVector | None: class SparseVectorBinaryLoader(SparseVectorLoader): - format = Format.BINARY def load(self, data: Buffer) -> SparseVector | None: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index cf1694c..f04cabd 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -10,7 +10,6 @@ class VectorDumper(Dumper): - format = Format.TEXT def dump(self, obj: Vector) -> Buffer | None: @@ -19,7 +18,6 @@ def dump(self, obj: Vector) -> Buffer | None: class VectorBinaryDumper(VectorDumper): - format = Format.BINARY def dump(self, obj: Vector) -> Buffer | None: @@ -27,7 +25,6 @@ def dump(self, obj: Vector) -> Buffer | None: class VectorLoader(Loader): - format = Format.TEXT def load(self, data: Buffer) -> Vector | None: @@ -37,7 +34,6 @@ def load(self, data: Buffer) -> Vector | None: class VectorBinaryLoader(VectorLoader): - format = Format.BINARY def load(self, data: Buffer) -> Vector | None: From 203dad334abde1f770ad411455a994da00ed84ac Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 2 Jun 2026 23:53:51 -0700 Subject: [PATCH 642/879] Improved type hints [skip ci] --- pgvector/psycopg/vector.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index f04cabd..7daf273 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -3,16 +3,19 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import Any, TypeAlias +from typing import Any, TypeAlias, TYPE_CHECKING from .. import Vector Buffer: TypeAlias = bytes | bytearray | memoryview +if TYPE_CHECKING: + import numpy as np + class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: Vector) -> Buffer | None: + def dump(self, obj: Vector | np.ndarray) -> Buffer | None: value = Vector._to_db(obj) return value if value is None else value.encode('utf8') @@ -20,7 +23,7 @@ def dump(self, obj: Vector) -> Buffer | None: class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: Vector) -> Buffer | None: + def dump(self, obj: Vector | np.ndarray) -> Buffer | None: return Vector._to_db_binary(obj) From 07ccccefbb747e1280b052f2dc18bce6b922e206 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 16:52:39 -0700 Subject: [PATCH 643/879] Removed dimensions check from SQLAlchemy for consistency with built-in types --- pgvector/halfvec.py | 5 +---- pgvector/sparsevec.py | 5 +---- pgvector/sqlalchemy/halfvec.py | 4 ++-- pgvector/sqlalchemy/sparsevec.py | 4 ++-- pgvector/sqlalchemy/vector.py | 4 ++-- pgvector/vector.py | 5 +---- 6 files changed, 9 insertions(+), 18 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 4728065..a9610d5 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -71,16 +71,13 @@ def from_binary(cls, value: bytes) -> HalfVector: return vec @classmethod - def _to_db(cls, value: object, dim: int | None = None) -> str | None: + def _to_db(cls, value: object) -> str | None: if value is None: return value if not isinstance(value, cls): value = cls(value) # ty: ignore[invalid-argument-type] - if dim is not None and value.dimensions() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) - return value.to_text() @classmethod diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 09b796c..0d76563 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -149,16 +149,13 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars return vec @classmethod - def _to_db(cls, value: object, dim: int | None = None) -> str | None: + def _to_db(cls, value: object) -> str | None: if value is None: return value if not isinstance(value, cls): value = cls(value) - if dim is not None and value.dimensions() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) - return value.to_text() @classmethod diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index af13317..49d5df0 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -20,14 +20,14 @@ def get_col_spec(self, **kw: Any) -> str: def bind_processor(self, dialect: Dialect) -> Any: def process(value: Any) -> str | None: - return HalfVector._to_db(value, self.dim) + return HalfVector._to_db(value) return process def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) def process(value: Any) -> Any: - return string_literal_processor(HalfVector._to_db(value, self.dim)) # type: ignore + return string_literal_processor(HalfVector._to_db(value)) # type: ignore return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 235a779..2dfea01 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -20,14 +20,14 @@ def get_col_spec(self, **kw: Any) -> str: def bind_processor(self, dialect: Dialect) -> Any: def process(value: Any) -> str | None: - return SparseVector._to_db(value, self.dim) + return SparseVector._to_db(value) return process def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) def process(value: Any) -> Any: - return string_literal_processor(SparseVector._to_db(value, self.dim)) # type: ignore + return string_literal_processor(SparseVector._to_db(value)) # type: ignore return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 14d66f3..20ad30f 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -20,14 +20,14 @@ def get_col_spec(self, **kw: Any) -> str: def bind_processor(self, dialect: Dialect) -> Any: def process(value: Any) -> str | None: - return Vector._to_db(value, self.dim) + return Vector._to_db(value) return process def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) def process(value: Any) -> Any: - return string_literal_processor(Vector._to_db(value, self.dim)) # type: ignore + return string_literal_processor(Vector._to_db(value)) # type: ignore return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: diff --git a/pgvector/vector.py b/pgvector/vector.py index 257ad8f..b48f223 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -71,16 +71,13 @@ def from_binary(cls, value: bytes) -> Vector: return vec @classmethod - def _to_db(cls, value: object, dim: int | None = None) -> str | None: + def _to_db(cls, value: object) -> str | None: if value is None: return value if not isinstance(value, cls): value = cls(value) # ty: ignore[invalid-argument-type] - if dim is not None and value.dimensions() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) - return value.to_text() @classmethod From 135cd53d057b9b5081476d170dded801f1fbda7c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 17:16:07 -0700 Subject: [PATCH 644/879] Simplified code --- pgvector/asyncpg/register.py | 6 +++--- pgvector/django/halfvec.py | 2 +- pgvector/django/vector.py | 2 +- pgvector/halfvec.py | 7 ------- pgvector/pg8000/register.py | 6 +++--- pgvector/psycopg/bit.py | 4 ++-- pgvector/psycopg/halfvec.py | 9 ++++----- pgvector/psycopg/sparsevec.py | 9 ++++----- pgvector/psycopg/vector.py | 13 ++++++++----- pgvector/sparsevec.py | 7 ------- pgvector/vector.py | 7 ------- 11 files changed, 26 insertions(+), 46 deletions(-) diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 6718850..67983da 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -7,7 +7,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None: 'vector', schema=schema, encoder=Vector._to_db_binary, - decoder=Vector._from_db_binary, + decoder=Vector.from_binary, format='binary' ) @@ -16,7 +16,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None: 'halfvec', schema=schema, encoder=HalfVector._to_db_binary, - decoder=HalfVector._from_db_binary, + decoder=HalfVector.from_binary, format='binary' ) @@ -24,7 +24,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None: 'sparsevec', schema=schema, encoder=SparseVector._to_db_binary, - decoder=SparseVector._from_db_binary, + decoder=SparseVector.from_binary, format='binary' ) except ValueError as e: diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index a2ec481..a3901b4 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -31,7 +31,7 @@ def to_python(self, value: Any) -> HalfVector | None: if value is None or isinstance(value, HalfVector): return value elif isinstance(value, str): - return HalfVector._from_db(value) + return HalfVector.from_text(value) else: return HalfVector(value) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 87d8989..3aafc11 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -31,7 +31,7 @@ def to_python(self, value: Any) -> Vector | None: if value is None or isinstance(value, Vector): return value elif isinstance(value, str): - return Vector._from_db(value) + return Vector.from_text(value) else: return Vector(value) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index a9610d5..860803c 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -96,10 +96,3 @@ def _from_db(cls, value: str | HalfVector | None) -> HalfVector | None: return value return cls.from_text(value) - - @classmethod - def _from_db_binary(cls, value: bytes | HalfVector | None) -> HalfVector | None: - if value is None or isinstance(value, HalfVector): - return value - - return cls.from_binary(value) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index 5d44b7c..6e84b7f 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -11,7 +11,7 @@ def register_vector(conn: Connection) -> None: raise RuntimeError('vector type not found in the database') conn.register_out_adapter(Vector, Vector._to_db) - conn.register_in_adapter(type_info['vector'], Vector._from_db) + conn.register_in_adapter(type_info['vector'], Vector.from_text) try: import numpy as np @@ -21,8 +21,8 @@ def register_vector(conn: Connection) -> None: if 'halfvec' in type_info: conn.register_out_adapter(HalfVector, HalfVector._to_db) - conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db) + conn.register_in_adapter(type_info['halfvec'], HalfVector.from_text) if 'sparsevec' in type_info: conn.register_out_adapter(SparseVector, SparseVector._to_db) - conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db) + conn.register_in_adapter(type_info['sparsevec'], SparseVector.from_text) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index f0f4021..a6ac6b2 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -12,14 +12,14 @@ class BitDumper(Dumper): format = Format.TEXT def dump(self, obj: Bit) -> Buffer | None: - return Bit._to_db(obj).encode('utf8') + return obj.to_text().encode('utf8') class BitBinaryDumper(BitDumper): format = Format.BINARY def dump(self, obj: Bit) -> Buffer | None: - return Bit._to_db_binary(obj) + return obj.to_binary() def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None) -> None: diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 6eaf501..8c9a8c8 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -12,15 +12,14 @@ class HalfVectorDumper(Dumper): format = Format.TEXT def dump(self, obj: HalfVector) -> Buffer | None: - value = HalfVector._to_db(obj) - return value if value is None else value.encode('utf8') + return obj.to_text().encode('utf8') class HalfVectorBinaryDumper(HalfVectorDumper): format = Format.BINARY def dump(self, obj: HalfVector) -> Buffer | None: - return HalfVector._to_db_binary(obj) + return obj.to_binary() class HalfVectorLoader(Loader): @@ -29,7 +28,7 @@ class HalfVectorLoader(Loader): def load(self, data: Buffer) -> HalfVector | None: if isinstance(data, memoryview): data = bytes(data) - return HalfVector._from_db(data.decode('utf8')) + return HalfVector.from_text(data.decode('utf8')) class HalfVectorBinaryLoader(HalfVectorLoader): @@ -38,7 +37,7 @@ class HalfVectorBinaryLoader(HalfVectorLoader): def load(self, data: Buffer) -> HalfVector | None: if isinstance(data, (bytearray, memoryview)): data = bytes(data) - return HalfVector._from_db_binary(data) + return HalfVector.from_binary(data) def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo) -> None: diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 16eaf5a..2a98e6b 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -12,15 +12,14 @@ class SparseVectorDumper(Dumper): format = Format.TEXT def dump(self, obj: SparseVector) -> Buffer | None: - value = SparseVector._to_db(obj) - return value if value is None else value.encode('utf8') + return obj.to_text().encode('utf8') class SparseVectorBinaryDumper(SparseVectorDumper): format = Format.BINARY def dump(self, obj: SparseVector) -> Buffer | None: - return SparseVector._to_db_binary(obj) + return obj.to_binary() class SparseVectorLoader(Loader): @@ -29,7 +28,7 @@ class SparseVectorLoader(Loader): def load(self, data: Buffer) -> SparseVector | None: if isinstance(data, memoryview): data = bytes(data) - return SparseVector._from_db(data.decode('utf8')) + return SparseVector.from_text(data.decode('utf8')) class SparseVectorBinaryLoader(SparseVectorLoader): @@ -38,7 +37,7 @@ class SparseVectorBinaryLoader(SparseVectorLoader): def load(self, data: Buffer) -> SparseVector | None: if isinstance(data, (bytearray, memoryview)): data = bytes(data) - return SparseVector._from_db_binary(data) + return SparseVector.from_binary(data) def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo) -> None: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 7daf273..a858ab6 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -16,15 +16,18 @@ class VectorDumper(Dumper): format = Format.TEXT def dump(self, obj: Vector | np.ndarray) -> Buffer | None: - value = Vector._to_db(obj) - return value if value is None else value.encode('utf8') + if not isinstance(obj, Vector): + obj = Vector(obj) + return obj.to_text().encode('utf8') class VectorBinaryDumper(VectorDumper): format = Format.BINARY def dump(self, obj: Vector | np.ndarray) -> Buffer | None: - return Vector._to_db_binary(obj) + if not isinstance(obj, Vector): + obj = Vector(obj) + return obj.to_binary() class VectorLoader(Loader): @@ -33,7 +36,7 @@ class VectorLoader(Loader): def load(self, data: Buffer) -> Vector | None: if isinstance(data, memoryview): data = bytes(data) - return Vector._from_db(data.decode('utf8')) + return Vector.from_text(data.decode('utf8')) class VectorBinaryLoader(VectorLoader): @@ -42,7 +45,7 @@ class VectorBinaryLoader(VectorLoader): def load(self, data: Buffer) -> Vector | None: if isinstance(data, (bytearray, memoryview)): data = bytes(data) - return Vector._from_db_binary(data) + return Vector.from_binary(data) def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) -> None: diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 0d76563..3fd9b53 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -174,10 +174,3 @@ def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None: return value return cls.from_text(value) - - @classmethod - def _from_db_binary(cls, value: bytes | SparseVector | None) -> SparseVector | None: - if value is None or isinstance(value, SparseVector): - return value - - return cls.from_binary(value) diff --git a/pgvector/vector.py b/pgvector/vector.py index b48f223..c682fe7 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -96,10 +96,3 @@ def _from_db(cls, value: str | Vector | None) -> Vector | None: return value return cls.from_text(value) - - @classmethod - def _from_db_binary(cls, value: bytes | Vector | None) -> Vector | None: - if value is None or isinstance(value, Vector): - return value - - return cls.from_binary(value) From 15f1ac7c688a0ed1ab157284b08f04368b692c9e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 17:22:15 -0700 Subject: [PATCH 645/879] Simplified Psycopg 2 code --- pgvector/psycopg2/halfvec.py | 4 +++- pgvector/psycopg2/sparsevec.py | 4 +++- pgvector/psycopg2/vector.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 1313938..cf26576 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -11,7 +11,9 @@ def getquoted(self) -> bytes: def cast_halfvec(value: str | None, cur: cursor) -> HalfVector | None: - return HalfVector._from_db(value) + if value is None: + return None + return HalfVector.from_text(value) def register_halfvec_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 34a4810..9e8c819 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -11,7 +11,9 @@ def getquoted(self) -> bytes: def cast_sparsevec(value: str | None, cur: cursor) -> SparseVector | None: - return SparseVector._from_db(value) + if value is None: + return None + return SparseVector.from_text(value) def register_sparsevec_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index f8832a6..7a3e3c4 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -11,7 +11,9 @@ def getquoted(self) -> bytes: def cast_vector(value: str | None, cur: cursor) -> Vector | None: - return Vector._from_db(value) + if value is None: + return None + return Vector.from_text(value) def register_vector_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: From 1868e6d86c649a84c392eaf400a741b1c63a8375 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 17:27:44 -0700 Subject: [PATCH 646/879] Simplified code --- pgvector/psycopg2/halfvec.py | 4 ++-- pgvector/psycopg2/sparsevec.py | 4 ++-- pgvector/psycopg2/vector.py | 10 ++++++++-- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index cf26576..8c95f7e 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -3,11 +3,11 @@ class HalfvecAdapter: - def __init__(self, value: object) -> None: + def __init__(self, value: HalfVector) -> None: self._value = value def getquoted(self) -> bytes: - return adapt(HalfVector._to_db(self._value)).getquoted() + return adapt(self._value.to_text()).getquoted() def cast_halfvec(value: str | None, cur: cursor) -> HalfVector | None: diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 9e8c819..6fb7f2f 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -3,11 +3,11 @@ class SparsevecAdapter: - def __init__(self, value: object) -> None: + def __init__(self, value: SparseVector) -> None: self._value = value def getquoted(self) -> bytes: - return adapt(SparseVector._to_db(self._value)).getquoted() + return adapt(self._value.to_text()).getquoted() def cast_sparsevec(value: str | None, cur: cursor) -> SparseVector | None: diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 7a3e3c4..b193eb6 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,13 +1,19 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import TYPE_CHECKING from .. import Vector +if TYPE_CHECKING: + import numpy as np + class VectorAdapter: - def __init__(self, value: object) -> None: + def __init__(self, value: Vector | np.ndarray) -> None: + if not isinstance(value, Vector): + value = Vector(value) self._value = value def getquoted(self) -> bytes: - return adapt(Vector._to_db(self._value)).getquoted() + return adapt(self._value.to_text()).getquoted() def cast_vector(value: str | None, cur: cursor) -> Vector | None: From 12ffb43dbf87253a9341eef1eff640615dc091cd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 18:05:50 -0700 Subject: [PATCH 647/879] Fixed CI --- pgvector/psycopg/vector.py | 13 +++++++++---- pgvector/psycopg2/vector.py | 11 +++++------ 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index a858ab6..b9608b3 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -3,13 +3,16 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import Any, TypeAlias, TYPE_CHECKING +from typing import Any, TypeAlias from .. import Vector Buffer: TypeAlias = bytes | bytearray | memoryview -if TYPE_CHECKING: +try: import numpy as np + NUMPY_AVAILABLE = True +except ImportError: + NUMPY_AVAILABLE = False class VectorDumper(Dumper): @@ -58,9 +61,11 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) -> binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) adapters = context.adapters - adapters.register_dumper('numpy.ndarray', text_dumper) - adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_dumper(Vector, text_dumper) adapters.register_dumper(Vector, binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) + + if NUMPY_AVAILABLE: + adapters.register_dumper(np.ndarray, text_dumper) + adapters.register_dumper(np.ndarray, binary_dumper) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index b193eb6..254ffcb 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,9 +1,11 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type -from typing import TYPE_CHECKING from .. import Vector -if TYPE_CHECKING: +try: import numpy as np + NUMPY_AVAILABLE = True +except ImportError: + NUMPY_AVAILABLE = False class VectorAdapter: @@ -32,8 +34,5 @@ def register_vector_info(oid: int, array_oid: int | None, scope: connection | cu register_adapter(Vector, VectorAdapter) - try: - import numpy as np + if NUMPY_AVAILABLE: register_adapter(np.ndarray, VectorAdapter) - except ImportError: - pass From 075a40703c26e66ccd0b338785e2a36cc38d023f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 18:13:28 -0700 Subject: [PATCH 648/879] Fixed CI --- pgvector/psycopg/vector.py | 4 ++-- pgvector/psycopg2/vector.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index b9608b3..b8b3611 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -18,7 +18,7 @@ class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: Vector | np.ndarray) -> Buffer | None: + def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_text().encode('utf8') @@ -27,7 +27,7 @@ def dump(self, obj: Vector | np.ndarray) -> Buffer | None: class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: Vector | np.ndarray) -> Buffer | None: + def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_binary() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 254ffcb..4571fc5 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -9,7 +9,7 @@ class VectorAdapter: - def __init__(self, value: Vector | np.ndarray) -> None: + def __init__(self, value: 'Vector | np.ndarray') -> None: if not isinstance(value, Vector): value = Vector(value) self._value = value From f6358361ff9b7889b1efa9cfcaa0b8560807603f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 21:07:25 -0700 Subject: [PATCH 649/879] Improved example [skip ci] --- examples/surprise/example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/surprise/example.py b/examples/surprise/example.py index e413bcf..40760e2 100644 --- a/examples/surprise/example.py +++ b/examples/surprise/example.py @@ -40,10 +40,10 @@ class Item(Base): session.execute(insert(User), users) session.execute(insert(Item), items) -user = session.get(User, 1) +user = session.get_one(User, 1) items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) print('user-based recs:', [item.id for item in items]) -item = session.get(Item, 50) +item = session.get_one(Item, 50) items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.id for item in items]) From cfd141ac419ff0826bad32809e8d5eda2d5b692b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 21:08:40 -0700 Subject: [PATCH 650/879] Improved example [skip ci] --- examples/surprise/example.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/surprise/example.py b/examples/surprise/example.py index 40760e2..6a40784 100644 --- a/examples/surprise/example.py +++ b/examples/surprise/example.py @@ -1,6 +1,6 @@ from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Session from surprise import Dataset, SVD engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -8,7 +8,9 @@ conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class User(Base): From 711e021f04932607787aa021c08fe96a6154a5f3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 9 Jun 2026 23:30:58 -0700 Subject: [PATCH 651/879] Updated examples [skip ci] --- examples/implicit/example.py | 10 ++++++---- examples/lightfm/example.py | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/examples/implicit/example.py b/examples/implicit/example.py index 2cbf7c6..ebb1ff1 100644 --- a/examples/implicit/example.py +++ b/examples/implicit/example.py @@ -2,14 +2,16 @@ from implicit.datasets.movielens import get_movielens from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer, String -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class User(Base): @@ -41,10 +43,10 @@ class Item(Base): session.execute(insert(User), users) session.execute(insert(Item), items) -user = session.get(User, 1) +user = session.get_one(User, 1) items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) print('user-based recs:', [item.title for item in items]) -item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() +item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).one() items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) diff --git a/examples/lightfm/example.py b/examples/lightfm/example.py index 65031c4..c869e89 100644 --- a/examples/lightfm/example.py +++ b/examples/lightfm/example.py @@ -2,14 +2,16 @@ from lightfm.datasets import fetch_movielens from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Float, Integer, String -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class User(Base): @@ -45,12 +47,12 @@ class Item(Base): session.execute(insert(User), users) session.execute(insert(Item), items) -user = session.get(User, 1) +user = session.get_one(User, 1) # subtract item bias for negative inner product items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors) - Item.bias).limit(5)) print('user-based recs:', [item.title for item in items]) # broken due to https://github.com/lyst/lightfm/issues/682 -item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() +item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).one() items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) From 0a45b840cec148681762334d4c59ebed8aba0d26 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:31:42 -0700 Subject: [PATCH 652/879] Improved example [skip ci] --- examples/sparse_search/example.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/sparse_search/example.py b/examples/sparse_search/example.py index 2b5daea..e3fec01 100644 --- a/examples/sparse_search/example.py +++ b/examples/sparse_search/example.py @@ -2,7 +2,8 @@ # https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/ # https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1 -from pgvector.psycopg import register_vector, SparseVector +from pgvector import SparseVector +from pgvector.psycopg import register_vector import psycopg import torch from transformers import AutoModelForMaskedLM, AutoTokenizer From f5cc7d090aece6d92abef4be8f636adf68248d11 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:33:22 -0700 Subject: [PATCH 653/879] Improved example [skip ci] --- examples/rdkit/example.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/rdkit/example.py b/examples/rdkit/example.py index afb56ec..dec2993 100644 --- a/examples/rdkit/example.py +++ b/examples/rdkit/example.py @@ -1,7 +1,8 @@ # good resource # https://www.rdkit.org/docs/GettingStartedInPython.html#morgan-fingerprints-circular-fingerprints -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg from rdkit import Chem from rdkit.Chem import AllChem From de5ae8e88f10b48aa22944f57fc4c2e72fa15441 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:37:27 -0700 Subject: [PATCH 654/879] Improved example [skip ci] --- examples/rag/example.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/rag/example.py b/examples/rag/example.py index 4d5d307..ddda8c3 100644 --- a/examples/rag/example.py +++ b/examples/rag/example.py @@ -1,11 +1,11 @@ # Run: +# ollama serve # ollama pull llama3.2 # ollama pull nomic-embed-text -# ollama serve -import numpy as np import ollama from pathlib import Path +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg import urllib.request @@ -55,7 +55,7 @@ embedding = ollama.embed(model='nomic-embed-text', input=input).embeddings[0] # retrieve chunks -result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (np.array(embedding),)).fetchall() +result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (Vector(embedding),)).fetchall() context = '\n\n'.join([row[0] for row in result]) # get answer From f9863c30227b8f2deb0bcd0794de77a063850a69 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:38:57 -0700 Subject: [PATCH 655/879] Improved example [skip ci] --- examples/openai/example.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/openai/example.py b/examples/openai/example.py index b9a078c..ba1fa7a 100644 --- a/examples/openai/example.py +++ b/examples/openai/example.py @@ -1,5 +1,5 @@ -import numpy as np from openai import OpenAI +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -25,10 +25,10 @@ def embed(input): ] embeddings = embed(input) for content, embedding in zip(input, embeddings): - conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, np.array(embedding))) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Vector(embedding))) query = 'forest' query_embedding = embed([query])[0] -result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (np.array(query_embedding),)).fetchall() +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (Vector(query_embedding),)).fetchall() for row in result: print(row[0]) From ba06214225b1cac40d3603f7d2793640115327fb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:45:55 -0700 Subject: [PATCH 656/879] Improved example [skip ci] --- examples/imagehash/example.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/imagehash/example.py b/examples/imagehash/example.py index f49af40..595dcb5 100644 --- a/examples/imagehash/example.py +++ b/examples/imagehash/example.py @@ -1,7 +1,8 @@ from datasets import load_dataset from imagehash import phash import matplotlib.pyplot as plt -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg @@ -18,7 +19,7 @@ def hash_image(img): conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, hash bit(64))') print('Loading dataset') -dataset = load_dataset('mnist') +dataset = load_dataset('ylecun/mnist') print('Generating hashes') images = [{'hash': hash_image(row['image'])} for row in dataset['train']] From 31cbefae92080e23dd4bee5e67a8ca6b4e559407 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:46:54 -0700 Subject: [PATCH 657/879] Improved example [skip ci] --- examples/colpali/exact.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index 80bb603..313c8bf 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -1,7 +1,8 @@ from colpali_engine.models import ColQwen2, ColQwen2Processor from colpali_engine.utils.torch_utils import get_torch_device from datasets import load_dataset -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg import torch From 34b39bc4ffc2b238f2273854404f8bca5dd6a9a0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:47:54 -0700 Subject: [PATCH 658/879] Improved example [skip ci] --- examples/cohere/example.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 5ef4eec..50a030a 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -1,6 +1,7 @@ import cohere import numpy as np -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg conn = psycopg.connect(dbname='pgvector_example', autocommit=True) From 6c6fb37721876ea6521946407b98cc75fd7f1754 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:53:18 -0700 Subject: [PATCH 659/879] Improved example [skip ci] --- examples/cohere/example.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 50a030a..5cf8d69 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -1,5 +1,4 @@ import cohere -import numpy as np from pgvector import Bit from pgvector.psycopg import register_vector import psycopg @@ -16,7 +15,7 @@ def embed(input, input_type): co = cohere.ClientV2() response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) - return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] + return [bytes(embedding) for embedding in response.embeddings.ubinary] input = [ From b0309bff46a2155c8ac85367c458b11ebdc54a69 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 16:55:26 -0700 Subject: [PATCH 660/879] Improved example [skip ci] --- examples/gensim/example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/gensim/example.py b/examples/gensim/example.py index cfbf18d..2fc63c8 100644 --- a/examples/gensim/example.py +++ b/examples/gensim/example.py @@ -1,7 +1,7 @@ from gensim.corpora.dictionary import Dictionary from gensim.models import LdaModel from gensim.utils import simple_preprocess -import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -26,7 +26,7 @@ model = LdaModel(corpus, num_topics=20) for content, bow in zip(input, corpus): - embedding = np.array([v[1] for v in model.get_document_topics(bow, minimum_probability=0)]) + embedding = Vector([v[1] for v in model.get_document_topics(bow, minimum_probability=0)]) conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 From 13c57d2a7be54a011cef58c5344eb932f85a897e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 25 Jun 2026 17:35:48 -0700 Subject: [PATCH 661/879] Improved example [skip ci] --- examples/rag/example.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/rag/example.py b/examples/rag/example.py index ddda8c3..9d5a784 100644 --- a/examples/rag/example.py +++ b/examples/rag/example.py @@ -51,8 +51,7 @@ # embed query # nomic-embed-text has task instruction prefix -input = 'search_query: ' + query -embedding = ollama.embed(model='nomic-embed-text', input=input).embeddings[0] +embedding = ollama.embed(model='nomic-embed-text', input=f'search_query: {query}').embeddings[0] # retrieve chunks result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (Vector(embedding),)).fetchall() From 80e7570dd396c60a0b59560505bd0caf4f947275 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 17:53:54 -0700 Subject: [PATCH 662/879] Added more tests for Django --- tests/test_django.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_django.py b/tests/test_django.py index 6bcf67a..ac2b626 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -181,6 +181,19 @@ def test_vector_l2_distance(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] + def test_vector_l2_distance_list(self): + create_items() + distance = L2Distance('embedding', [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)] + + def test_vector_l2_distance_none(self): + create_items() + distance = L2Distance('embedding', None) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.distance for v in items] == [None, None, None] + def test_vector_max_inner_product(self): create_items() distance = MaxInnerProduct('embedding', [1, 1, 1]) From 8b9f9cafc28d766508a54f5e64bc1e0b4fb40e31 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 17:59:51 -0700 Subject: [PATCH 663/879] Improved code --- pgvector/django/functions.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index af76a67..6c58f34 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -8,12 +8,10 @@ class DistanceBase(Func): def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): - if isinstance(vector, HalfVector): - vector = Value(HalfVector._to_db(vector)) - elif isinstance(vector, SparseVector): - vector = Value(SparseVector._to_db(vector)) - else: - vector = Value(Vector._to_db(vector)) + if isinstance(vector, (Vector, HalfVector, SparseVector)): + vector = Value(vector.to_text()) + elif vector is not None: + vector = Value(Vector(vector).to_text()) # prevent error with unhashable types self._constructor_args = ((expression, vector), extra) From d90ad1b112669572ef41afd73391480ac621004e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 18:24:58 -0700 Subject: [PATCH 664/879] Updated super style --- pgvector/peewee/bit.py | 2 +- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/sparsevec.py | 2 +- pgvector/peewee/vector.py | 2 +- pgvector/sqlalchemy/bit.py | 2 +- pgvector/sqlalchemy/halfvec.py | 2 +- pgvector/sqlalchemy/sparsevec.py | 2 +- pgvector/sqlalchemy/vector.py | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index 14c4555..d38c08e 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -7,7 +7,7 @@ class FixedBitField(Field): def __init__(self, max_length: int | None = None, *args: Any, **kwargs: Any) -> None: self.max_length = max_length - super(FixedBitField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) def get_modifiers(self) -> list[int] | None: return [self.max_length] if self.max_length else None diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 161fb92..cb76976 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -8,7 +8,7 @@ class HalfVectorField(Field): def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions - super(HalfVectorField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 24b37b9..20943a5 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -8,7 +8,7 @@ class SparseVectorField(Field): def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions - super(SparseVectorField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 9652071..18fb75a 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -8,7 +8,7 @@ class VectorField(Field): def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions - super(VectorField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 35b6753..59d221a 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -8,7 +8,7 @@ class BIT(UserDefinedType): cache_ok = True def __init__(self, length: int | None = None) -> None: - super(UserDefinedType, self).__init__() + super().__init__() self.length = length def get_col_spec(self, **kw: Any) -> str: diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 49d5df0..2945bc6 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -10,7 +10,7 @@ class HALFVEC(UserDefinedType): _string = String() def __init__(self, dim: int | None = None) -> None: - super(UserDefinedType, self).__init__() + super().__init__() self.dim = dim def get_col_spec(self, **kw: Any) -> str: diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 2dfea01..1c3938e 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -10,7 +10,7 @@ class SPARSEVEC(UserDefinedType): _string = String() def __init__(self, dim: int | None = None) -> None: - super(UserDefinedType, self).__init__() + super().__init__() self.dim = dim def get_col_spec(self, **kw: Any) -> str: diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 20ad30f..5ac359c 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -10,7 +10,7 @@ class VECTOR(UserDefinedType): _string = String() def __init__(self, dim: int | None = None) -> None: - super(UserDefinedType, self).__init__() + super().__init__() self.dim = dim def get_col_spec(self, **kw: Any) -> str: From 410b02c2c927452528b1b3551cda60cfbb867388 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 20:37:27 -0700 Subject: [PATCH 665/879] Improved type comments [skip ci] --- pgvector/halfvec.py | 4 ++-- pgvector/vector.py | 4 ++-- tests/test_bit.py | 8 ++++---- tests/test_half_vector.py | 6 +++--- tests/test_sparse_vector.py | 4 ++-- tests/test_vector.py | 6 +++--- 6 files changed, 16 insertions(+), 16 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 860803c..9831e27 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -76,7 +76,7 @@ def _to_db(cls, value: object) -> str | None: return value if not isinstance(value, cls): - value = cls(value) # ty: ignore[invalid-argument-type] + value = cls(value) # type: ignore return value.to_text() @@ -86,7 +86,7 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value if not isinstance(value, cls): - value = cls(value) # ty: ignore[invalid-argument-type] + value = cls(value) # type: ignore return value.to_binary() diff --git a/pgvector/vector.py b/pgvector/vector.py index c682fe7..721a2ec 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -76,7 +76,7 @@ def _to_db(cls, value: object) -> str | None: return value if not isinstance(value, cls): - value = cls(value) # ty: ignore[invalid-argument-type] + value = cls(value) # type: ignore return value.to_text() @@ -86,7 +86,7 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value if not isinstance(value, cls): - value = cls(value) # ty: ignore[invalid-argument-type] + value = cls(value) # type: ignore return value.to_binary() diff --git a/tests/test_bit.py b/tests/test_bit.py index ff708d0..5b82f7c 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -15,17 +15,17 @@ def test_list(self): def test_list_none(self): with pytest.raises(ValueError) as error: - Bit([True, None, True]) # ty: ignore[invalid-argument-type] + Bit([True, None, True]) # type: ignore assert str(error.value) == 'expected list[bool]' def test_list_int(self): with pytest.raises(ValueError) as error: - Bit([254, 7, 0]) # ty: ignore[invalid-argument-type] + Bit([254, 7, 0]) # type: ignore assert str(error.value) == 'expected list[bool]' def test_list_list(self): with pytest.raises(ValueError) as error: - Bit([[True, False], [True, False]]) # ty: ignore[invalid-argument-type] + Bit([[True, False], [True, False]]) # type: ignore assert str(error.value) == 'expected list[bool]' def test_str(self): @@ -67,7 +67,7 @@ def test_ndarray_uint16(self): def test_bool(self): with pytest.raises(ValueError) as error: - Bit(True) # ty: ignore[invalid-argument-type] + Bit(True) # type: ignore assert str(error.value) == 'expected bytes, str, list, or ndarray' def test_random(self): diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index d859bd8..362f1b3 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -18,12 +18,12 @@ def test_list_empty(self): def test_list_str(self): with pytest.raises(ValueError) as error: - HalfVector([1, 'two', 3]) # ty: ignore[invalid-argument-type] + HalfVector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' def test_list_list(self): with pytest.raises(ValueError) as error: - HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] + HalfVector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') @@ -34,7 +34,7 @@ def test_ndarray(self): def test_int(self): with pytest.raises(ValueError) as error: - HalfVector(1) # ty: ignore[invalid-argument-type] + HalfVector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' def test_repr(self): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 9a18f98..490472f 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -28,7 +28,7 @@ def test_list_empty(self): def test_list_dimensions(self): with pytest.raises(ValueError) as error: - SparseVector([1, 0, 2, 0, 3, 0], 6) # ty: ignore[invalid-argument-type] + SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore assert str(error.value) == 'extra argument' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') @@ -60,7 +60,7 @@ def test_coo_array(self): @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: - SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # ty: ignore[invalid-argument-type] + SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore assert str(error.value) == 'extra argument' @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') diff --git a/tests/test_vector.py b/tests/test_vector.py index a8b8e28..3392df6 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -18,12 +18,12 @@ def test_list_empty(self): def test_list_str(self): with pytest.raises(ValueError) as error: - Vector([1, 'two', 3]) # ty: ignore[invalid-argument-type] + Vector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' def test_list_list(self): with pytest.raises(ValueError) as error: - Vector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type] + Vector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') @@ -34,7 +34,7 @@ def test_ndarray(self): def test_int(self): with pytest.raises(ValueError) as error: - Vector(1) # ty: ignore[invalid-argument-type] + Vector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' def test_repr(self): From b367d86815ea8abd02cad831bba1c9a5427ec47c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 20:56:52 -0700 Subject: [PATCH 666/879] Fixed typechecking [skip ci] --- tests/test_peewee.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 095a5db..3c0cd9c 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -2,6 +2,7 @@ from peewee import Model, PostgresqlDatabase, fn from pgvector import HalfVector, SparseVector, Vector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField +from typing import Any db = PostgresqlDatabase('pgvector_python_test') @@ -12,6 +13,7 @@ class Meta: class Item(BaseModel): + id: Any # for typing embedding = VectorField(dimensions=3, null=True) half_embedding = HalfVectorField(dimensions=3, null=True) binary_embedding = FixedBitField(max_length=3, null=True) From 72b3a16cf8c7e55730aeb0abfa6f4f2703bed789 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:15:14 -0700 Subject: [PATCH 667/879] Improved type hints [skip ci] --- Makefile | 2 +- pgvector/psycopg2/register.py | 2 +- pgvector/sqlalchemy/bit.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index ccf995a..cd5b0ac 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ lint: pycodestyle . --ignore=E501 --exclude=.venv check: - ty check pgvector tests + ty check pgvector tests --ignore unused-type-ignore-comment build: python3 -m build diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 22beb8c..f0004ce 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -7,7 +7,7 @@ # note: register_adapter is always global def register_vector(conn_or_curs: connection | cursor, globally: bool = False, arrays: bool = True) -> None: - conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection + conn = conn_or_curs if isinstance(conn_or_curs, connection) else conn_or_curs.connection cur = conn.cursor(cursor_factory=cursor) scope = None if globally else conn_or_curs diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 59d221a..eb6769a 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -22,7 +22,7 @@ def bind_processor(self, dialect: Dialect) -> Any: def process(value: Any) -> Any: if isinstance(value, str): - return asyncpg.BitString(value) + return asyncpg.BitString(value) # type: ignore return value return process else: From c636fbfdf944c0330241a5f8400b35dd2d6acb11 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:23:44 -0700 Subject: [PATCH 668/879] Improved type hints [skip ci] --- tests/test_bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 5b82f7c..0161d72 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -62,7 +62,7 @@ def test_ndarray_uint8(self): def test_ndarray_uint16(self): arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.raises(ValueError) as error: - Bit(arr) + Bit(arr) # type: ignore assert str(error.value) == 'expected elements to be boolean' def test_bool(self): From 6ef248444f4b294fdaefc15c529199a16a7602b9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:29:31 -0700 Subject: [PATCH 669/879] Improved type hints [skip ci] --- tests/test_sqlalchemy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 3d2c1d3..b8c8c89 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -535,7 +535,7 @@ def test_binary_quantize(self, engine): session.add(Item(id=3, embedding=[1, 2, 3])) session.commit() - distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) # type: ignore items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] @@ -550,7 +550,7 @@ def test_binary_quantize_reranking(self, engine): session.add(Item(id=3, embedding=[1, 2, 3])) session.commit() - distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) # type: ignore subquery = session.query(Item).order_by(distance).limit(20).subquery() items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all() assert [v.id for v in items] == [2, 3, 1] From e9af82e67b566a1d72b255545f37c806e71122dd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:32:50 -0700 Subject: [PATCH 670/879] Fixed typechecking --- pgvector/sqlalchemy/bit.py | 4 ++-- pgvector/sqlalchemy/halfvec.py | 4 ++-- pgvector/sqlalchemy/sparsevec.py | 4 ++-- pgvector/sqlalchemy/vector.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index eb6769a..cd81dc9 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -1,5 +1,5 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float +from sqlalchemy.types import UserDefinedType, TypeEngine, Float from sqlalchemy import Dialect, Operators from typing import Any @@ -28,7 +28,7 @@ def process(value: Any) -> Any: else: return super().bind_processor(dialect) - class comparator_factory(UserDefinedType.Comparator): + class comparator_factory(TypeEngine.Comparator): def hamming_distance(self, other: object) -> Operators: return self.op('<~>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 2945bc6..0abf8bb 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,5 +1,5 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String from sqlalchemy import Dialect, Operators from typing import Any from .. import HalfVector @@ -35,7 +35,7 @@ def process(value: Any) -> HalfVector | None: return HalfVector._from_db(value) return process - class comparator_factory(UserDefinedType.Comparator): + class comparator_factory(TypeEngine.Comparator): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 1c3938e..7782b09 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,5 +1,5 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String from sqlalchemy import Dialect, Operators from typing import Any from .. import SparseVector @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class comparator_factory(UserDefinedType.Comparator): + class comparator_factory(TypeEngine.Comparator): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 5ac359c..9b38823 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,5 +1,5 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String from sqlalchemy import Dialect, Operators from typing import Any from .. import Vector @@ -35,7 +35,7 @@ def process(value: Any) -> Vector | None: return Vector._from_db(value) return process - class comparator_factory(UserDefinedType.Comparator): + class comparator_factory(TypeEngine.Comparator): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) From d0511913d9127cb127326ef241792533653ba788 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:46:36 -0700 Subject: [PATCH 671/879] Fixed typechecking [skip ci] --- pgvector/django/halfvec.py | 5 +++-- pgvector/django/sparsevec.py | 5 +++-- pgvector/django/vector.py | 5 +++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index a3901b4..6059f87 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -38,8 +38,9 @@ def to_python(self, value: Any) -> HalfVector | None: def get_prep_value(self, value: Any) -> str | None: return HalfVector._to_db(value) - def value_to_string(self, obj: Any) -> str | None: - return self.get_prep_value(self.value_from_object(obj)) + def value_to_string(self, obj: Any) -> str: + value = self.get_prep_value(self.value_from_object(obj)) + return '' if value is None else value def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: return super().formfield( diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 53a2c37..bec664f 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -33,8 +33,9 @@ def to_python(self, value: Any) -> SparseVector | None: def get_prep_value(self, value: Any) -> str | None: return SparseVector._to_db(value) - def value_to_string(self, obj: Any) -> str | None: - return self.get_prep_value(self.value_from_object(obj)) + def value_to_string(self, obj: Any) -> str: + value = self.get_prep_value(self.value_from_object(obj)) + return '' if value is None else value def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: return super().formfield( diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 3aafc11..89b20d6 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -38,8 +38,9 @@ def to_python(self, value: Any) -> Vector | None: def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) - def value_to_string(self, obj: Any) -> str | None: - return self.get_prep_value(self.value_from_object(obj)) + def value_to_string(self, obj: Any) -> str: + value = self.get_prep_value(self.value_from_object(obj)) + return '' if value is None else value def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: return super().formfield( From 788fbaa21a460d1ad727aefe2702a6ab6fe4854d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:51:38 -0700 Subject: [PATCH 672/879] Improved type hints [skip ci] --- pgvector/django/functions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 6c58f34..73b35b7 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -4,7 +4,7 @@ class DistanceBase(Func): - output_field = FloatField() + output_field = FloatField() # type: ignore def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): @@ -20,7 +20,7 @@ def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: class BitDistanceBase(Func): - output_field = FloatField() + output_field = FloatField() # type: ignore def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): From f96bb59c71d65886f99518b7efb5abc7757f833c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:51:52 -0700 Subject: [PATCH 673/879] Fixed typechecking [skip ci] --- pgvector/psycopg2/halfvec.py | 4 +++- pgvector/psycopg2/sparsevec.py | 4 +++- pgvector/psycopg2/vector.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 8c95f7e..9fa311b 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -10,9 +10,11 @@ def getquoted(self) -> bytes: return adapt(self._value.to_text()).getquoted() -def cast_halfvec(value: str | None, cur: cursor) -> HalfVector | None: +def cast_halfvec(value: str | bytes | None, cur: cursor) -> HalfVector | None: if value is None: return None + if isinstance(value, bytes): + raise ValueError('expected str') return HalfVector.from_text(value) diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 6fb7f2f..b733200 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -10,9 +10,11 @@ def getquoted(self) -> bytes: return adapt(self._value.to_text()).getquoted() -def cast_sparsevec(value: str | None, cur: cursor) -> SparseVector | None: +def cast_sparsevec(value: str | bytes | None, cur: cursor) -> SparseVector | None: if value is None: return None + if isinstance(value, bytes): + raise ValueError('expected str') return SparseVector.from_text(value) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 4571fc5..0eca9db 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -18,9 +18,11 @@ def getquoted(self) -> bytes: return adapt(self._value.to_text()).getquoted() -def cast_vector(value: str | None, cur: cursor) -> Vector | None: +def cast_vector(value: str | bytes | None, cur: cursor) -> Vector | None: if value is None: return None + if isinstance(value, bytes): + raise ValueError('expected str') return Vector.from_text(value) From c34c3b946ea05ebfcfed2d20bbc353bcc92c25a6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:55:27 -0700 Subject: [PATCH 674/879] Fixed typechecking [skip ci] --- pgvector/pg8000/register.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index 6e84b7f..c9a0138 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -1,10 +1,11 @@ from pg8000.native import Connection +from typing import Any from .. import Vector, HalfVector, SparseVector def register_vector(conn: Connection) -> None: # use to_regtype to get first matching type in search path - res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") + res: Any = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") type_info = dict(res) if 'vector' not in type_info: From c8aaf62779af84ec91e7a570cd3b5cc4aa8c6062 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 22:57:58 -0700 Subject: [PATCH 675/879] Improved typing [skip ci] --- pgvector/pg8000/register.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index c9a0138..520fd10 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -1,11 +1,11 @@ from pg8000.native import Connection -from typing import Any +from typing import cast from .. import Vector, HalfVector, SparseVector def register_vector(conn: Connection) -> None: # use to_regtype to get first matching type in search path - res: Any = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") + res = cast(list[tuple[str, int]], conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")) type_info = dict(res) if 'vector' not in type_info: From 5f80f4e3bc0b68bdf1e3b430fd7910b03a1b2e56 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 23:04:24 -0700 Subject: [PATCH 676/879] Fixed typechecking [skip ci] --- pgvector/sqlalchemy/bit.py | 4 +++- pgvector/sqlalchemy/halfvec.py | 4 +++- pgvector/sqlalchemy/sparsevec.py | 4 +++- pgvector/sqlalchemy/vector.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index cd81dc9..e65012e 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -28,13 +28,15 @@ def process(value: Any) -> Any: else: return super().bind_processor(dialect) - class comparator_factory(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator): def hamming_distance(self, other: object) -> Operators: return self.op('<~>', return_type=Float)(other) def jaccard_distance(self, other: object) -> Operators: return self.op('<%>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['bit'] = BIT # type: ignore diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 0abf8bb..c54ee85 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -35,7 +35,7 @@ def process(value: Any) -> HalfVector | None: return HalfVector._from_db(value) return process - class comparator_factory(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,6 +48,8 @@ def cosine_distance(self, other: object) -> Operators: def l1_distance(self, other: object) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['halfvec'] = HALFVEC # type: ignore diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 7782b09..9aeaf08 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class comparator_factory(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,6 +48,8 @@ def cosine_distance(self, other: object) -> Operators: def l1_distance(self, other: object) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['sparsevec'] = SPARSEVEC # type: ignore diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 9b38823..0c85f48 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -35,7 +35,7 @@ def process(value: Any) -> Vector | None: return Vector._from_db(value) return process - class comparator_factory(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,6 +48,8 @@ def cosine_distance(self, other: object) -> Operators: def l1_distance(self, other: object) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['vector'] = VECTOR # type: ignore From 96e1b052c4ab81fc2f393c46e141bc3dadda4d4d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 23:10:49 -0700 Subject: [PATCH 677/879] Fixed typechecking [skip ci] --- tests/test_psycopg.py | 4 ++-- tests/test_psycopg2.py | 9 ++++----- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 13dc7c8..4c8df2e 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -191,7 +191,7 @@ def configure(conn): with pool.connection() as conn: res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() - assert res[0] == Vector([1, 2, 3]) + assert res == (Vector([1, 2, 3]),) pool.close() @@ -226,6 +226,6 @@ async def configure(conn): async with conn.cursor() as cur: await cur.execute("SELECT '[1,2,3]'::vector") res = await cur.fetchone() - assert res[0] == Vector([1, 2, 3]) + assert res == (Vector([1, 2, 3]),) await pool.close() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 466fef1..edb6e44 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -79,8 +79,7 @@ def test_vector_array(self): cur.execute('SELECT embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0][0] == embeddings[0] - assert res[0][1] == embeddings[1] + assert res == (embeddings,) def test_halfvec_array(self): embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] @@ -88,7 +87,7 @@ def test_halfvec_array(self): cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0] == [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] + assert res == ([HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])],) def test_sparsevec_array(self): embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] @@ -96,7 +95,7 @@ def test_sparsevec_array(self): cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0] == [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] + assert res == ([SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])],) def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: @@ -126,7 +125,7 @@ def test_pool(self): cur = conn.cursor() cur.execute("SELECT '[1,2,3]'::vector") res = cur.fetchone() - assert res[0] == Vector([1, 2, 3]) + assert res == (Vector([1, 2, 3]),) finally: pool.putconn(conn) From 6f26aefbf73aa32b9996981715d635f2982a34b6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 23:13:40 -0700 Subject: [PATCH 678/879] Fixed typechecking [skip ci] --- tests/test_pg8000.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 718d0f7..0c4fb18 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -30,9 +30,7 @@ def test_vector(self): conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2), (:embedding3)', embedding=embedding, embedding2=embedding2, embedding3=embedding3) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == embedding - assert res[1][0] == Vector([4.5, 5, 6]) - assert res[2][0] is None + assert res == [[embedding], [Vector([4.5, 5, 6])], [None]] def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) @@ -40,8 +38,7 @@ def test_halfvec(self): conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == embedding - assert res[1][0] is None + assert res == [[embedding], [None]] def test_bit(self): embedding = '101' @@ -49,8 +46,7 @@ def test_bit(self): conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == '101' - assert res[1][0] is None + assert res == [['101'], [None]] def test_sparsevec(self): embedding = SparseVector([1.5, 2, 3]) @@ -58,5 +54,4 @@ def test_sparsevec(self): conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == embedding - assert res[1][0] is None + assert res == [[embedding], [None]] From 892180d4fbd7a0943e83597601015e9f3a283a58 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 29 Jun 2026 23:16:38 -0700 Subject: [PATCH 679/879] Fixed typechecking [skip ci] --- tests/test_sqlalchemy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index b8c8c89..92fdf3b 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -423,13 +423,13 @@ def test_select(self, engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() - assert items[0] == 3 + assert items == (3,) def test_select_orm(self, engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() - assert items[0] == 3 + assert items == [3] def test_avg(self, engine): with Session(engine) as session: @@ -520,7 +520,7 @@ def test_automap(self, engine): with Session(engine) as session: session.execute(insert(AutoItem), [{'embedding': Vector([1, 2, 3])}]) item = session.query(AutoItem).first() - assert item.embedding == Vector([1, 2, 3]) + assert item is not None and item.embedding == Vector([1, 2, 3]) def test_half_precision(self, engine): create_items() From 096d2e21f9ca458406d8e8f68c34cac478aa273e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 02:55:48 -0700 Subject: [PATCH 680/879] Fixed typechecking [skip ci] --- tests/test_peewee.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 3c0cd9c..2c7d6a2 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -19,7 +19,7 @@ class Item(BaseModel): binary_embedding = FixedBitField(max_length=3, null=True) sparse_embedding = SparseVectorField(dimensions=3, null=True) - class Meta: + class Meta: # type: ignore table_name = 'peewee_item' @@ -161,7 +161,7 @@ def test_sparsevec_l1_distance(self): def test_where(self): create_items() - items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) + items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) # type: ignore assert [v.id for v in items] == [1] def test_vector_avg(self): @@ -207,9 +207,9 @@ def test_vector_array(self): ext_db = PostgresqlExtDatabase('pgvector_python_test') class ExtItem(BaseModel): - embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False) + embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False) # type: ignore - class Meta: + class Meta: # type: ignore database = ext_db table_name = 'peewee_ext_item' From d4560b8640344fa89e004ca5f1e5050e5bc325cf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 02:57:56 -0700 Subject: [PATCH 681/879] Fixed typechecking [skip ci] --- tests/test_asyncpg.py | 2 +- tests/test_sqlmodel.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 3df3d39..a925f63 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -124,7 +124,7 @@ async def test_vector_array(self): assert res[0]['embeddings'] == embeddings assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] if NUMPY_AVAILABLE: - assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] + assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] # type: ignore await conn.close() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 4dcaedd..7df08ff 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -17,7 +17,7 @@ class Item(SQLModel, table=True): - __tablename__ = 'sqlmodel_item' + __tablename__ = 'sqlmodel_item' # type: ignore id: Optional[int] = Field(default=None, primary_key=True) embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) # type: ignore From 6d9e25812665edb320c7658105a567f9e0a6e212 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 03:06:16 -0700 Subject: [PATCH 682/879] Added pyright and pyrefly config to pyproject.toml [skip ci] --- pyproject.toml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 5a7d99e..9c90300 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,3 +39,14 @@ dev-optional = [ [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" + +[tool.pyright] +venv = ".venv" +venvPath = "." +reportPossiblyUnboundVariable = "none" + +[tool.pyrefly] +project-includes = [ + "pgvector/**/*.py", + "tests/**/*.py" +] From ff46e05bec73f84118ce0642bfeb5e1a92aaa9eb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 03:16:33 -0700 Subject: [PATCH 683/879] Added mypy config to pyproject.toml [skip ci] --- pyproject.toml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 9c90300..bc8271d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,14 @@ dev-optional = [ asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" +[tool.mypy] +python_executable = ".venv/bin/python" +files = [ + "pgvector", + "tests" +] +ignore_missing_imports = true + [tool.pyright] venv = ".venv" venvPath = "." From 0eb8905387a99a88564cff33d800b5eded630b88 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 03:18:41 -0700 Subject: [PATCH 684/879] Updated typing config [skip ci] --- pyproject.toml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bc8271d..eeefa5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,17 +41,21 @@ asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" [tool.mypy] -python_executable = ".venv/bin/python" files = [ "pgvector", "tests" ] ignore_missing_imports = true +python_executable = ".venv/bin/python" [tool.pyright] +include = [ + "pgvector", + "tests" +] +reportPossiblyUnboundVariable = "none" venv = ".venv" venvPath = "." -reportPossiblyUnboundVariable = "none" [tool.pyrefly] project-includes = [ From 19ef9e7b7a2932fecedcb5dd7becd3e57f4cbf3c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 03:21:43 -0700 Subject: [PATCH 685/879] Fixed typechecking [skip ci] --- tests/test_psycopg2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index edb6e44..76bae07 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -106,7 +106,7 @@ def test_cursor_factory(self): def test_cursor_factory_connection(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: - conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) + conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) # type: ignore register_vector(conn) conn.close() From 217a7e3e03402f0acf5aae94a0adc9360660b590 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 03:29:28 -0700 Subject: [PATCH 686/879] Improved typing [skip ci] --- pgvector/bit.py | 2 +- pgvector/sqlalchemy/functions.py | 3 ++- pyproject.toml | 2 ++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 1052cee..e45336e 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -58,7 +58,7 @@ def __eq__(self, other: object) -> bool: return self.to_binary() == other.to_binary() return False - def _length(self): + def _length(self) -> int: length, = unpack_from('>i', self._value) return length diff --git a/pgvector/sqlalchemy/functions.py b/pgvector/sqlalchemy/functions.py index 72e3ca7..ca4cae9 100644 --- a/pgvector/sqlalchemy/functions.py +++ b/pgvector/sqlalchemy/functions.py @@ -1,9 +1,10 @@ # https://docs.sqlalchemy.org/en/20/core/functions.html # include sum for a consistent API from sqlalchemy.sql.functions import ReturnTypeFromArgs, sum +from typing import Any -class avg(ReturnTypeFromArgs): +class avg(ReturnTypeFromArgs[Any]): inherit_cache = True package = 'pgvector' diff --git a/pyproject.toml b/pyproject.toml index eeefa5d..968945c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,8 @@ files = [ "pgvector", "tests" ] +strict = true +disable_error_code = ["unused-ignore"] ignore_missing_imports = true python_executable = ".venv/bin/python" From d2a5fb0c7c0be945164cf58fe1ef468ec390aec3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 03:35:28 -0700 Subject: [PATCH 687/879] Improved typing [skip ci] --- pgvector/sqlalchemy/bit.py | 4 ++-- pgvector/sqlalchemy/halfvec.py | 4 ++-- pgvector/sqlalchemy/sparsevec.py | 4 ++-- pgvector/sqlalchemy/vector.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index e65012e..2679ee3 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -4,7 +4,7 @@ from typing import Any -class BIT(UserDefinedType): +class BIT(UserDefinedType[Any]): cache_ok = True def __init__(self, length: int | None = None) -> None: @@ -28,7 +28,7 @@ def process(value: Any) -> Any: else: return super().bind_processor(dialect) - class Comparator(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator[Any]): def hamming_distance(self, other: object) -> Operators: return self.op('<~>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index c54ee85..5c16bfd 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -5,7 +5,7 @@ from .. import HalfVector -class HALFVEC(UserDefinedType): +class HALFVEC(UserDefinedType[Any]): cache_ok = True _string = String() @@ -35,7 +35,7 @@ def process(value: Any) -> HalfVector | None: return HalfVector._from_db(value) return process - class Comparator(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator[Any]): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 9aeaf08..8baff48 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -5,7 +5,7 @@ from .. import SparseVector -class SPARSEVEC(UserDefinedType): +class SPARSEVEC(UserDefinedType[Any]): cache_ok = True _string = String() @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class Comparator(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator[Any]): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 0c85f48..b049758 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -5,7 +5,7 @@ from .. import Vector -class VECTOR(UserDefinedType): +class VECTOR(UserDefinedType[Any]): cache_ok = True _string = String() @@ -35,7 +35,7 @@ def process(value: Any) -> Vector | None: return Vector._from_db(value) return process - class Comparator(TypeEngine.Comparator): + class Comparator(TypeEngine.Comparator[Any]): def l2_distance(self, other: object) -> Operators: return self.op('<->', return_type=Float)(other) From 23d287a6ee5bbdde92e32f660edc4cb8c29591d6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 04:33:19 -0700 Subject: [PATCH 688/879] Improved typing [skip ci] --- pgvector/bit.py | 3 ++- pgvector/halfvec.py | 3 ++- pgvector/vector.py | 3 ++- pyproject.toml | 3 ++- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index e45336e..b5c48e6 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,5 +1,6 @@ from __future__ import annotations from struct import pack, unpack_from +from typing import cast try: import numpy as np @@ -60,7 +61,7 @@ def __eq__(self, other: object) -> bool: def _length(self) -> int: length, = unpack_from('>i', self._value) - return length + return cast(int, length) def to_list(self) -> list[bool]: # TODO improve diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 9831e27..7a14d12 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,5 +1,6 @@ from __future__ import annotations import struct +from typing import cast try: import numpy as np @@ -38,7 +39,7 @@ def __eq__(self, other: object) -> bool: def dimensions(self) -> int: dim, = struct.unpack_from('>H', self._value) - return dim + return cast(int, dim) def to_list(self) -> list[float]: return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:])) diff --git a/pgvector/vector.py b/pgvector/vector.py index 721a2ec..1841d93 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,5 +1,6 @@ from __future__ import annotations import struct +from typing import cast try: import numpy as np @@ -38,7 +39,7 @@ def __eq__(self, other: object) -> bool: def dimensions(self) -> int: dim, = struct.unpack_from('>H', self._value) - return dim + return cast(int, dim) def to_list(self) -> list[float]: return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:])) diff --git a/pyproject.toml b/pyproject.toml index 968945c..a8be6cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,8 +46,9 @@ files = [ "tests" ] strict = true -disable_error_code = ["unused-ignore"] +disable_error_code = ["unused-ignore", "no-untyped-def", "no-untyped-call"] ignore_missing_imports = true +disallow_subclassing_any = false python_executable = ".venv/bin/python" [tool.pyright] From 4a99efcd2e8e7bba6ab1e531801257ad691a6344 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 04:40:09 -0700 Subject: [PATCH 689/879] Improved typing [skip ci] --- tests/test_psycopg.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 4c8df2e..1ac91ce 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,8 +1,10 @@ from pgvector import Bit, HalfVector, SparseVector, Vector from pgvector.psycopg import register_vector, register_vector_async import psycopg +from psycopg import Connection, AsyncConnection from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest +from typing import Any try: import numpy as np @@ -184,7 +186,7 @@ def test_vector_array(self): assert res[0][1] == embeddings[1] def test_pool(self): - def configure(conn): + def configure(conn: Connection[Any]): register_vector(conn) pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure) @@ -216,7 +218,7 @@ async def test_async(self): @pytest.mark.asyncio async def test_async_pool(self): - async def configure(conn): + async def configure(conn: AsyncConnection[Any]): await register_vector_async(conn) pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure) From a4ef932c79f1ac27b555b59a85e710b65f68bbfb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:06:25 -0700 Subject: [PATCH 690/879] Improved typing [skip ci] --- pyproject.toml | 1 + tests/test_sqlalchemy.py | 130 +++++++++++++++++++-------------------- 2 files changed, 66 insertions(+), 65 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a8be6cb..2691854 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ include = [ "pgvector", "tests" ] +typeCheckingMode = "strict" reportPossiblyUnboundVariable = "none" venv = ".venv" venvPath = "." diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 92fdf3b..053b967 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -3,9 +3,9 @@ from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest -from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY +from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine from sqlalchemy.exc import StatementError -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncEngine from sqlalchemy.ext.automap import automap_base from sqlalchemy.orm import mapped_column, DeclarativeBase, Session from sqlalchemy.sql import func @@ -134,7 +134,7 @@ class TestSqlalchemy: def setup_method(self): delete_items() - def test_core(self, engine): + def test_core(self, engine: Engine): metadata = MetaData() item_table = Table( @@ -169,7 +169,7 @@ def test_core(self, engine): ) hnsw_index.create(engine) - def test_orm(self, engine): + def test_orm(self, engine: Engine): item = Item(embedding=Vector([1.5, 2, 3])) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -191,136 +191,136 @@ def test_orm(self, engine): assert items[1].embedding == Vector([4, 5, 6]) assert items[2].embedding is None - def test_vector(self, engine): + def test_vector(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.embedding == Vector([1, 2, 3]) - def test_vector_l2_distance(self, engine): + def test_vector_l2_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l2_distance_orm(self, engine): + def test_vector_l2_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self, engine): + def test_vector_max_inner_product(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_vector_max_inner_product_orm(self, engine): + def test_vector_max_inner_product_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self, engine): + def test_vector_cosine_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_vector_cosine_distance_orm(self, engine): + def test_vector_cosine_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self, engine): + def test_vector_l1_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l1_distance_orm(self, engine): + def test_vector_l1_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self, engine): + def test_halfvec(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.half_embedding == HalfVector([1, 2, 3]) - def test_halfvec_l2_distance(self, engine): + def test_halfvec_l2_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l2_distance_orm(self, engine): + def test_halfvec_l2_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self, engine): + def test_halfvec_max_inner_product(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_max_inner_product_orm(self, engine): + def test_halfvec_max_inner_product_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self, engine): + def test_halfvec_cosine_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_cosine_distance_orm(self, engine): + def test_halfvec_cosine_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self, engine): + def test_halfvec_l1_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l1_distance_orm(self, engine): + def test_halfvec_l1_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self, engine): + def test_bit(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() item = session.get_one(Item, 1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self, engine): + def test_bit_hamming_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_hamming_distance_orm(self, engine): + def test_bit_hamming_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self, engine): + def test_bit_jaccard_distance(self, engine: Engine): if engine == pg8000_engine: return @@ -329,7 +329,7 @@ def test_bit_jaccard_distance(self, engine): items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance_orm(self, engine): + def test_bit_jaccard_distance_orm(self, engine: Engine): if engine == pg8000_engine: return @@ -338,100 +338,100 @@ def test_bit_jaccard_distance_orm(self, engine): items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self, engine): + def test_sparsevec(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self, engine): + def test_sparsevec_l2_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l2_distance_orm(self, engine): + def test_sparsevec_l2_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self, engine): + def test_sparsevec_max_inner_product(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_max_inner_product_orm(self, engine): + def test_sparsevec_max_inner_product_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self, engine): + def test_sparsevec_cosine_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_cosine_distance_orm(self, engine): + def test_sparsevec_cosine_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self, engine): + def test_sparsevec_l1_distance(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l1_distance_orm(self, engine): + def test_sparsevec_l1_distance_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_subquery(self, engine): + def test_subquery(self, engine: Engine): create_items() with Session(engine) as session: subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() items = session.query(Item).order_by(Item.embedding.l2_distance(subquery)).all() assert [v.id for v in items] == [1, 3, 2] - def test_subquery_orm(self, engine): + def test_subquery_orm(self, engine: Engine): create_items() with Session(engine) as session: subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() items = session.scalars(select(Item).order_by(Item.embedding.l2_distance(subquery))) assert [v.id for v in items] == [1, 3, 2] - def test_filter(self, engine): + def test_filter(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() assert [v.id for v in items] == [1] - def test_filter_orm(self, engine): + def test_filter_orm(self, engine: Engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] - def test_select(self, engine): + def test_select(self, engine: Engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() assert items == (3,) - def test_select_orm(self, engine): + def test_select_orm(self, engine: Engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() assert items == [3] - def test_avg(self, engine): + def test_avg(self, engine: Engine): with Session(engine) as session: res = session.query(avg(Item.embedding)).one()[0] assert res is None @@ -440,7 +440,7 @@ def test_avg(self, engine): res = session.query(avg(Item.embedding)).one()[0] assert res == Vector([2.5, 3.5, 4.5]) - def test_avg_orm(self, engine): + def test_avg_orm(self, engine: Engine): with Session(engine) as session: res = session.scalars(select(avg(Item.embedding))).first() assert res is None @@ -449,7 +449,7 @@ def test_avg_orm(self, engine): res = session.scalars(select(avg(Item.embedding))).one() assert res == Vector([2.5, 3.5, 4.5]) - def test_sum(self, engine): + def test_sum(self, engine: Engine): with Session(engine) as session: res = session.query(sum(Item.embedding)).one()[0] assert res is None @@ -458,7 +458,7 @@ def test_sum(self, engine): res = session.query(sum(Item.embedding)).one()[0] assert res == Vector([5, 7, 9]) - def test_sum_orm(self, engine): + def test_sum_orm(self, engine: Engine): with Session(engine) as session: res = session.scalars(select(sum(Item.embedding))).first() assert res is None @@ -467,7 +467,7 @@ def test_sum_orm(self, engine): res = session.scalars(select(sum(Item.embedding))).one() assert res == Vector([5, 7, 9]) - def test_bad_dimensions(self, engine): + def test_bad_dimensions(self, engine: Engine): item = Item(embedding=[1, 2]) with Session(engine) as session: session.add(item) @@ -475,7 +475,7 @@ def test_bad_dimensions(self, engine): session.commit() @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_bad_ndim(self, engine): + def test_bad_ndim(self, engine: Engine): item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: session.add(item) @@ -483,26 +483,26 @@ def test_bad_ndim(self, engine): session.commit() @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_bad_dtype(self, engine): + def test_bad_dtype(self, engine: Engine): item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='could not convert string to float'): session.commit() - def test_inspect(self, engine): + def test_inspect(self, engine: Engine): columns = inspect(engine).get_columns('sqlalchemy_orm_item') assert isinstance(columns[1]['type'], VECTOR) - def test_literal_binds(self, engine): + def test_literal_binds(self, engine: Engine): sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) - def test_insert(self, engine): + def test_insert(self, engine: Engine): with Session(engine) as session: session.execute(insert(Item).values(embedding=Vector([1, 2, 3]))) - def test_insert_bulk(self, engine): + def test_insert_bulk(self, engine: Engine): with Session(engine) as session: session.execute(insert(Item), [{'embedding': Vector([1, 2, 3])}]) @@ -511,7 +511,7 @@ def test_insert_bulk(self, engine): # with Session(engine) as session: # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) - def test_automap(self, engine): + def test_automap(self, engine: Engine): metadata = MetaData() metadata.reflect(engine, only=['sqlalchemy_orm_item']) AutoBase = automap_base(metadata=metadata) @@ -522,13 +522,13 @@ def test_automap(self, engine): item = session.query(AutoItem).first() assert item is not None and item.embedding == Vector([1, 2, 3]) - def test_half_precision(self, engine): + def test_half_precision(self, engine: Engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_binary_quantize(self, engine): + def test_binary_quantize(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, embedding=[-1, -2, -3])) session.add(Item(id=2, embedding=[1, -2, 3])) @@ -539,7 +539,7 @@ def test_binary_quantize(self, engine): items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] - def test_binary_quantize_reranking(self, engine): + def test_binary_quantize_reranking(self, engine: Engine): # recreate index (could also vacuum table) binary_quantize_index.drop(setup_engine) binary_quantize_index.create(setup_engine) @@ -561,7 +561,7 @@ class TestSqlalchemyArray: def setup_method(self): delete_items() - def test_vector_array(self, engine): + def test_vector_array(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) session.commit() @@ -570,7 +570,7 @@ def test_vector_array(self, engine): item = session.get_one(Item, 1) assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] - def test_halfvec_array(self, engine): + def test_halfvec_array(self, engine: Engine): with Session(engine) as session: session.add(Item(id=1, half_embeddings=[HalfVector([1, 2, 3]), HalfVector([4, 5, 6])])) session.commit() @@ -586,7 +586,7 @@ def setup_method(self): delete_items() @pytest.mark.asyncio - async def test_vector(self, engine): + async def test_vector(self, engine: AsyncEngine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -599,7 +599,7 @@ async def test_vector(self, engine): await engine.dispose() @pytest.mark.asyncio - async def test_halfvec(self, engine): + async def test_halfvec(self, engine: AsyncEngine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -612,7 +612,7 @@ async def test_halfvec(self, engine): await engine.dispose() @pytest.mark.asyncio - async def test_bit(self, engine): + async def test_bit(self, engine: AsyncEngine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -630,7 +630,7 @@ async def test_bit(self, engine): await engine.dispose() @pytest.mark.asyncio - async def test_sparsevec(self, engine): + async def test_sparsevec(self, engine: AsyncEngine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -643,7 +643,7 @@ async def test_sparsevec(self, engine): await engine.dispose() @pytest.mark.asyncio - async def test_avg(self, engine): + async def test_avg(self, engine: AsyncEngine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -662,7 +662,7 @@ def setup_method(self): delete_items() @pytest.mark.asyncio - async def test_vector_array(self, engine): + async def test_vector_array(self, engine: AsyncEngine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: From e0632b8ad5e0ed4b975e78194796accf7ae24b1a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:10:54 -0700 Subject: [PATCH 691/879] Improved typing [skip ci] --- pyproject.toml | 2 +- tests/test_asyncpg.py | 3 ++- tests/test_peewee.py | 2 +- tests/test_sqlalchemy.py | 4 ++-- tests/test_sqlmodel.py | 2 +- 5 files changed, 7 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2691854..b4d4699 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,7 @@ files = [ "tests" ] strict = true -disable_error_code = ["unused-ignore", "no-untyped-def", "no-untyped-call"] +disable_error_code = ["unused-ignore", "no-untyped-def"] ignore_missing_imports = true disallow_subclassing_any = false python_executable = ".venv/bin/python" diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index a925f63..017710b 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,4 +1,5 @@ import asyncpg +from asyncpg import Connection from pgvector import HalfVector, SparseVector, Vector from pgvector.asyncpg import register_vector import pytest @@ -11,7 +12,7 @@ class TestAsyncpg: - async def setup_connection(self): + async def setup_connection(self) -> Connection: conn = await asyncpg.connect(database='pgvector_python_test') await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') await register_vector(conn) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 2c7d6a2..e8e204f 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -31,7 +31,7 @@ class Meta: # type: ignore db.create_tables([Item]) -def create_items(): +def create_items() -> None: Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])) Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])) Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 053b967..5e1163a 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -115,7 +115,7 @@ class Item(Base): binary_quantize_index.create(setup_engine) -def create_items(): +def create_items() -> None: with Session(setup_engine) as session: session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) @@ -123,7 +123,7 @@ def create_items(): session.commit() -def delete_items(): +def delete_items() -> None: with Session(setup_engine) as session: session.query(Item).delete() session.commit() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 7df08ff..4cdc140 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -39,7 +39,7 @@ class Item(SQLModel, table=True): index.create(engine) -def create_items(): +def create_items() -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) From e73bac28c6f962631d2f3934c12285fa592d4f95 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:24:06 -0700 Subject: [PATCH 692/879] Added return type annotations to tests [skip ci] --- pyproject.toml | 2 +- tests/test_asyncpg.py | 14 ++-- tests/test_bit.py | 30 ++++---- tests/test_half_vector.py | 24 +++---- tests/test_peewee.py | 52 +++++++------- tests/test_pg8000.py | 10 +-- tests/test_psycopg.py | 60 ++++++++-------- tests/test_psycopg2.py | 24 +++---- tests/test_sparse_vector.py | 44 ++++++------ tests/test_sqlalchemy.py | 140 ++++++++++++++++++------------------ tests/test_sqlmodel.py | 54 +++++++------- tests/test_vector.py | 24 +++---- 12 files changed, 239 insertions(+), 239 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b4d4699..faf1aa1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,7 @@ files = [ "tests" ] strict = true -disable_error_code = ["unused-ignore", "no-untyped-def"] +disable_error_code = ["unused-ignore"] ignore_missing_imports = true disallow_subclassing_any = false python_executable = ".venv/bin/python" diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 017710b..b8b42bf 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -19,7 +19,7 @@ async def setup_connection(self) -> Connection: return conn @pytest.mark.asyncio - async def test_vector(self): + async def test_vector(self) -> None: conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') @@ -43,7 +43,7 @@ async def test_vector(self): await conn.close() @pytest.mark.asyncio - async def test_halfvec(self): + async def test_halfvec(self) -> None: conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') @@ -65,7 +65,7 @@ async def test_halfvec(self): await conn.close() @pytest.mark.asyncio - async def test_bit(self): + async def test_bit(self) -> None: conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') @@ -86,7 +86,7 @@ async def test_bit(self): await conn.close() @pytest.mark.asyncio - async def test_sparsevec(self): + async def test_sparsevec(self) -> None: conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') @@ -106,7 +106,7 @@ async def test_sparsevec(self): await conn.close() @pytest.mark.asyncio - async def test_vector_array(self): + async def test_vector_array(self) -> None: conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') @@ -130,8 +130,8 @@ async def test_vector_array(self): await conn.close() @pytest.mark.asyncio - async def test_pool(self): - async def init(conn): + async def test_pool(self) -> None: + async def init(conn) -> None: await register_vector(conn) pool = await asyncpg.create_pool(database='pgvector_python_test', init=init) diff --git a/tests/test_bit.py b/tests/test_bit.py index 0161d72..5e1da9c 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -10,74 +10,74 @@ class TestBit: - def test_list(self): + def test_list(self) -> None: assert Bit([True, False, True]).to_list() == [True, False, True] - def test_list_none(self): + def test_list_none(self) -> None: with pytest.raises(ValueError) as error: Bit([True, None, True]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_list_int(self): + def test_list_int(self) -> None: with pytest.raises(ValueError) as error: Bit([254, 7, 0]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_list_list(self): + def test_list_list(self) -> None: with pytest.raises(ValueError) as error: Bit([[True, False], [True, False]]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_str(self): + def test_str(self) -> None: assert Bit('101').to_list() == [True, False, True] - def test_str_two(self): + def test_str_two(self) -> None: with pytest.raises(ValueError) as error: Bit('201') assert str(error.value) == 'expected bit string' - def test_bytes(self): + def test_bytes(self) -> None: assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray(self): + def test_ndarray(self) -> None: arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray_unpackbits(self): + def test_ndarray_unpackbits(self) -> None: arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray_uint8(self): + def test_ndarray_uint8(self) -> None: arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.raises(ValueError) as error: Bit(arr) assert str(error.value) == 'expected elements to be boolean' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray_uint16(self): + def test_ndarray_uint16(self) -> None: arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.raises(ValueError) as error: Bit(arr) # type: ignore assert str(error.value) == 'expected elements to be boolean' - def test_bool(self): + def test_bool(self) -> None: with pytest.raises(ValueError) as error: Bit(True) # type: ignore assert str(error.value) == 'expected bytes, str, list, or ndarray' - def test_random(self): + def test_random(self) -> None: value = ''.join(random.choices(['0', '1'], k=random.randint(1024, 2048))) assert Bit(value).to_text() == value - def test_repr(self): + def test_repr(self) -> None: assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' - def test_equality(self): + def test_equality(self) -> None: assert Bit([True, False, True]) == Bit([True, False, True]) assert Bit([True, False, True]) != Bit([True, False, False]) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 362f1b3..6bd819c 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -10,58 +10,58 @@ class TestHalfVector: - def test_list(self): + def test_list(self) -> None: assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] - def test_list_empty(self): + def test_list_empty(self) -> None: assert HalfVector([]).to_list() == [] - def test_list_str(self): + def test_list_str(self) -> None: with pytest.raises(ValueError) as error: HalfVector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_list_list(self): + def test_list_list(self) -> None: with pytest.raises(ValueError) as error: HalfVector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray(self): + def test_ndarray(self) -> None: arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr - def test_int(self): + def test_int(self) -> None: with pytest.raises(ValueError) as error: HalfVector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' - def test_repr(self): + def test_repr(self) -> None: assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' - def test_equality(self): + def test_equality(self) -> None: assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) - def test_dimensions(self): + def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_to_numpy_readonly(self): + def test_to_numpy_readonly(self) -> None: arr = HalfVector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: arr[0] = 4 assert str(error.value) == 'assignment destination is read-only' - def test_from_text(self): + def test_from_text(self) -> None: vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - def test_from_binary(self): + def test_from_binary(self) -> None: data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] diff --git a/tests/test_peewee.py b/tests/test_peewee.py index e8e204f..f1a4758 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -38,133 +38,133 @@ def create_items() -> None: class TestPeewee: - def setup_method(self): + def setup_method(self) -> None: Item.truncate_table() - def test_vector(self): + def test_vector(self) -> None: Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) assert item.embedding == Vector([1, 2, 3]) - def test_vector_l2_distance(self): + def test_vector_l2_distance(self) -> None: create_items() distance = Item.embedding.l2_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self) -> None: create_items() distance = Item.embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self) -> None: create_items() distance = Item.embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self) -> None: create_items() distance = Item.embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_halfvec(self): + def test_halfvec(self) -> None: Item.create(id=1, half_embedding=[1, 2, 3]) item = Item.get_by_id(1) assert item.half_embedding == HalfVector([1, 2, 3]) - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self) -> None: create_items() distance = Item.half_embedding.l2_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self) -> None: create_items() distance = Item.half_embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self) -> None: create_items() distance = Item.half_embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self) -> None: create_items() distance = Item.half_embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_bit(self): + def test_bit(self) -> None: Item.create(id=1, binary_embedding='101') item = Item.get_by_id(1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self) -> None: create_items() distance = Item.binary_embedding.hamming_distance('101') items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self) -> None: create_items() distance = Item.binary_embedding.jaccard_distance('101') items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] # assert [v.distance for v in items] == [0, 1/3, 1] - def test_sparsevec(self): + def test_sparsevec(self) -> None: Item.create(id=1, sparse_embedding=[1, 2, 3]) item = Item.get_by_id(1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() distance = Item.sparse_embedding.l2_distance(SparseVector([1, 1, 1])) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() distance = Item.sparse_embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() distance = Item.sparse_embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self) -> None: create_items() distance = Item.sparse_embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_where(self): + def test_where(self) -> None: create_items() items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) # type: ignore assert [v.id for v in items] == [1] - def test_vector_avg(self): + def test_vector_avg(self) -> None: avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() assert avg is None Item.create(embedding=[1, 2, 3]) @@ -172,7 +172,7 @@ def test_vector_avg(self): avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() assert avg == Vector([2.5, 3.5, 4.5]) - def test_vector_sum(self): + def test_vector_sum(self) -> None: sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() assert sum is None Item.create(embedding=[1, 2, 3]) @@ -180,7 +180,7 @@ def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() assert sum == Vector([5, 7, 9]) - def test_halfvec_avg(self): + def test_halfvec_avg(self) -> None: avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() assert avg is None Item.create(half_embedding=[1, 2, 3]) @@ -188,7 +188,7 @@ def test_halfvec_avg(self): avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() assert avg == HalfVector([2.5, 3.5, 4.5]) - def test_halfvec_sum(self): + def test_halfvec_sum(self) -> None: sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() assert sum is None Item.create(half_embedding=[1, 2, 3]) @@ -196,12 +196,12 @@ def test_halfvec_sum(self): sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() assert sum == HalfVector([5, 7, 9]) - def test_get_or_create(self): + def test_get_or_create(self) -> None: Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) Item.get_or_create(embedding=Vector([4, 5, 6])) Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) - def test_vector_array(self): + def test_vector_array(self) -> None: from playhouse.postgres_ext import PostgresqlExtDatabase, ArrayField ext_db = PostgresqlExtDatabase('pgvector_python_test') diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 0c4fb18..c7ad466 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -20,10 +20,10 @@ class TestPg8000: - def setup_method(self): + def setup_method(self) -> None: conn.run('DELETE FROM pg8000_items') - def test_vector(self): + def test_vector(self) -> None: embedding = Vector([1.5, 2, 3]) embedding2 = np.array([4.5, 5, 6]) if NUMPY_AVAILABLE else Vector([4.5, 5, 6]) embedding3 = None @@ -32,7 +32,7 @@ def test_vector(self): res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') assert res == [[embedding], [Vector([4.5, 5, 6])], [None]] - def test_halfvec(self): + def test_halfvec(self) -> None: embedding = HalfVector([1.5, 2, 3]) embedding2 = None conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) @@ -40,7 +40,7 @@ def test_halfvec(self): res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') assert res == [[embedding], [None]] - def test_bit(self): + def test_bit(self) -> None: embedding = '101' embedding2 = None conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) @@ -48,7 +48,7 @@ def test_bit(self): res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') assert res == [['101'], [None]] - def test_sparsevec(self): + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) embedding2 = None conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 1ac91ce..6e8820d 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -22,10 +22,10 @@ class TestPsycopg: - def setup_method(self): + def setup_method(self) -> None: conn.execute('DELETE FROM psycopg_items') - def test_vector(self): + def test_vector(self) -> None: embedding = Vector([1.5, 2, 3]) embedding2 = None conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (%s)', (embedding, embedding2)) @@ -34,119 +34,119 @@ def test_vector(self): assert res[0][0] == embedding assert res[1][0] is None - def test_vector_binary_format(self): + def test_vector_binary_format(self) -> None: embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert res == embedding - def test_vector_text_format(self): + def test_vector_text_format(self) -> None: embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert res == embedding - def test_vector_binary_format_correct(self): + def test_vector_binary_format_correct(self) -> None: embedding = Vector([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] assert res == '[1.5,2,3]' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_vector_numpy_binary_format(self): + def test_vector_numpy_binary_format(self) -> None: embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert res == Vector(embedding) @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_vector_numpy_text_format(self): + def test_vector_numpy_text_format(self) -> None: embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert res == Vector(embedding) @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_vector_numpy_binary_format_non_contiguous(self): + def test_vector_numpy_binary_format_non_contiguous(self) -> None: embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] assert res == Vector([3, 2, 1.5]) @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_vector_numpy_text_format_non_contiguous(self): + def test_vector_numpy_text_format_non_contiguous(self) -> None: embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert res == Vector([3, 2, 1.5]) - def test_halfvec(self): + def test_halfvec(self) -> None: embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) res = next(conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id'))[0] assert res == HalfVector([1.5, 2, 3]) - def test_halfvec_binary_format(self): + def test_halfvec_binary_format(self) -> None: embedding = HalfVector([1.5, 2, 3]) res = next(conn.execute('SELECT %b::halfvec', (embedding,), binary=True))[0] assert res == HalfVector([1.5, 2, 3]) - def test_halfvec_text_format(self): + def test_halfvec_text_format(self) -> None: embedding = HalfVector([1.5, 2, 3]) res = next(conn.execute('SELECT %t::halfvec', (embedding,)))[0] assert res == HalfVector([1.5, 2, 3]) - def test_bit(self): + def test_bit(self) -> None: embedding = Bit([True, False, True]) conn.execute('INSERT INTO psycopg_items (binary_embedding) VALUES (%s)', (embedding,)) res = next(conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id'))[0] assert res == '101' - def test_bit_binary_format(self): + def test_bit_binary_format(self) -> None: embedding = Bit([False, True, False, True, False, False, False, False, True]) res = next(conn.execute('SELECT %b::bit(9)', (embedding,), binary=True))[0] assert repr(Bit.from_binary(res)) == 'Bit(010100001)' - def test_bit_text_format(self): + def test_bit_text_format(self) -> None: embedding = Bit([False, True, False, True, False, False, False, False, True]) res = next(conn.execute('SELECT %t::bit(9)', (embedding,)))[0] assert res == '010100001' assert repr(Bit(res)) == 'Bit(010100001)' - def test_sparsevec(self): + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) res = next(conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id'))[0] assert res == SparseVector([1.5, 2, 3]) - def test_sparsevec_binary_format(self): + def test_sparsevec_binary_format(self) -> None: embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = next(conn.execute('SELECT %b::sparsevec', (embedding,), binary=True))[0] assert res == embedding - def test_sparsevec_text_format(self): + def test_sparsevec_text_format(self) -> None: embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = next(conn.execute('SELECT %t::sparsevec', (embedding,)))[0] assert res == embedding - def test_text_copy_from(self): + def test_text_copy_from(self) -> None: embedding = [1.5, 2, 3] cur = conn.cursor() with cur.copy('COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN') as copy: copy.write_row([Vector(embedding), HalfVector(embedding), '101', SparseVector(embedding)]) - def test_binary_copy_from(self): + def test_binary_copy_from(self) -> None: embedding = [1.5, 2, 3] cur = conn.cursor() with cur.copy('COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: copy.write_row([Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) - def test_binary_copy_from_set_types(self): + def test_binary_copy_from_set_types(self) -> None: embedding = [1.5, 2, 3] cur = conn.cursor() with cur.copy('COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) copy.write_row([1, Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) - def test_text_copy_to(self): + def test_text_copy_to(self) -> None: embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) @@ -156,7 +156,7 @@ def test_text_copy_to(self): assert row[0] == '[1.5,2,3]' assert row[1] == '[1.5,2,3]' - def test_binary_copy_to(self): + def test_binary_copy_to(self) -> None: embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) @@ -166,7 +166,7 @@ def test_binary_copy_to(self): assert Vector.from_binary(row[0]) == embedding assert HalfVector.from_binary(row[1]) == half_embedding - def test_binary_copy_to_set_types(self): + def test_binary_copy_to_set_types(self) -> None: embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) @@ -177,7 +177,7 @@ def test_binary_copy_to_set_types(self): assert row[0] == embedding assert row[1] == half_embedding - def test_vector_array(self): + def test_vector_array(self) -> None: embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] conn.execute('INSERT INTO psycopg_items (embeddings) VALUES (%s)', (embeddings,)) @@ -185,8 +185,8 @@ def test_vector_array(self): assert res[0][0] == embeddings[0] assert res[0][1] == embeddings[1] - def test_pool(self): - def configure(conn: Connection[Any]): + def test_pool(self) -> None: + def configure(conn: Connection[Any]) -> None: register_vector(conn) pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure) @@ -198,7 +198,7 @@ def configure(conn: Connection[Any]): pool.close() @pytest.mark.asyncio - async def test_async(self): + async def test_async(self) -> None: conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') @@ -217,8 +217,8 @@ async def test_async(self): assert res[1][1] is None @pytest.mark.asyncio - async def test_async_pool(self): - async def configure(conn: AsyncConnection[Any]): + async def test_async_pool(self) -> None: + async def configure(conn: AsyncConnection[Any]) -> None: await register_vector_async(conn) pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure) diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 76bae07..9b42d35 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -23,10 +23,10 @@ class TestPsycopg2: - def setup_method(self): + def setup_method(self) -> None: cur.execute('DELETE FROM psycopg2_items') - def test_vector(self): + def test_vector(self) -> None: embedding = Vector([1.5, 2, 3]) embedding2 = None cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (%s)', (embedding, embedding2)) @@ -37,7 +37,7 @@ def test_vector(self): assert res[1][0] is None @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_vector_numpy(self): + def test_vector_numpy(self) -> None: embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) @@ -46,7 +46,7 @@ def test_vector_numpy(self): assert res[0][0] == Vector(embedding) assert res[1][0] is None - def test_halfvec(self): + def test_halfvec(self) -> None: embedding = HalfVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) @@ -55,7 +55,7 @@ def test_halfvec(self): assert res[0][0] == embedding assert res[1][0] is None - def test_bit(self): + def test_bit(self) -> None: embedding = '101' cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) @@ -64,7 +64,7 @@ def test_bit(self): assert res[0][0] == embedding assert res[1][0] is None - def test_sparsevec(self): + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) @@ -73,7 +73,7 @@ def test_sparsevec(self): assert res[0][0] == embedding assert res[1][0] is None - def test_vector_array(self): + def test_vector_array(self) -> None: embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] cur.execute('INSERT INTO psycopg2_items (embeddings) VALUES (%s::vector[])', (embeddings,)) @@ -81,7 +81,7 @@ def test_vector_array(self): res = cur.fetchone() assert res == (embeddings,) - def test_halfvec_array(self): + def test_halfvec_array(self) -> None: embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] cur.execute('INSERT INTO psycopg2_items (half_embeddings) VALUES (%s::halfvec[])', (embeddings,)) @@ -89,7 +89,7 @@ def test_halfvec_array(self): res = cur.fetchone() assert res == ([HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])],) - def test_sparsevec_array(self): + def test_sparsevec_array(self) -> None: embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,)) @@ -97,20 +97,20 @@ def test_sparsevec_array(self): res = cur.fetchone() assert res == ([SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])],) - def test_cursor_factory(self): + def test_cursor_factory(self) -> None: for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) register_vector(cur) conn.close() - def test_cursor_factory_connection(self): + def test_cursor_factory_connection(self) -> None: for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) # type: ignore register_vector(conn) conn.close() - def test_pool(self): + def test_pool(self) -> None: pool = ThreadedConnectionPool(1, 1, dbname='pgvector_python_test') conn = pool.getconn() diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 490472f..7720644 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -16,109 +16,109 @@ class TestSparseVector: - def test_list(self): + def test_list(self) -> None: vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] - def test_list_empty(self): + def test_list_empty(self) -> None: assert SparseVector([]).to_list() == [] - def test_list_dimensions(self): + def test_list_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore assert str(error.value) == 'extra argument' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray(self): + def test_ndarray(self) -> None: vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict(self): + def test_dict(self) -> None: vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict_empty(self): + def test_dict_empty(self) -> None: assert SparseVector({}, 0).to_list() == [] - def test_dict_no_dimensions(self): + def test_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) assert str(error.value) == 'missing dimensions' @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') - def test_coo_array(self): + def test_coo_array(self) -> None: arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') - def test_coo_array_dimensions(self): + def test_coo_array_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore assert str(error.value) == 'extra argument' @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') - def test_coo_matrix(self): + def test_coo_matrix(self) -> None: mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') - def test_dok_array(self): + def test_dok_array(self) -> None: arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') - def test_csr_array(self): + def test_csr_array(self) -> None: arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') - def test_csr_matrix(self): + def test_csr_matrix(self) -> None: mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_repr(self): + def test_repr(self) -> None: assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' - def test_equality(self): + def test_equality(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert SparseVector({}, 1) != SparseVector({}, 2) - def test_dimensions(self): + def test_dimensions(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 - def test_indices(self): + def test_indices(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] - def test_values(self): + def test_values(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] @pytest.mark.skipif(not NUMPY_AVAILABLE or not SCIPY_AVAILABLE, reason='NumPy and SciPy required') - def test_to_coo(self): + def test_to_coo(self) -> None: assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) - def test_zero_vector_text(self): + def test_zero_vector_text(self) -> None: vec = SparseVector({}, 3) assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() - def test_from_text(self): + def test_from_text(self) -> None: vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') assert vec.dimensions() == 6 assert vec.indices() == [0, 2, 4] @@ -127,7 +127,7 @@ def test_from_text(self): if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) - def test_from_binary(self): + def test_from_binary(self) -> None: data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) vec = SparseVector.from_binary(data) assert vec.dimensions() == 6 diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 5e1163a..08621b9 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -33,7 +33,7 @@ def psycopg2_connect(dbapi_connection, connection_record): @event.listens_for(psycopg_type_engine, 'connect') -def psycopg_connect(dbapi_connection, connection_record): +def psycopg_connect(dbapi_connection, connection_record) -> None: from pgvector.psycopg import register_vector register_vector(dbapi_connection) @@ -43,7 +43,7 @@ def psycopg_connect(dbapi_connection, connection_record): @event.listens_for(psycopg_async_type_engine.sync_engine, 'connect') -def psycopg_async_connect(dbapi_connection, connection_record): +def psycopg_async_connect(dbapi_connection, connection_record) -> None: from pgvector.psycopg import register_vector_async dbapi_connection.run_async(register_vector_async) @@ -53,7 +53,7 @@ def psycopg_async_connect(dbapi_connection, connection_record): @event.listens_for(asyncpg_type_engine.sync_engine, 'connect') -def asyncpg_connect(dbapi_connection, connection_record): +def asyncpg_connect(dbapi_connection, connection_record) -> None: from pgvector.asyncpg import register_vector dbapi_connection.run_async(register_vector) @@ -131,10 +131,10 @@ def delete_items() -> None: @pytest.mark.parametrize('engine', engines) class TestSqlalchemy: - def setup_method(self): + def setup_method(self) -> None: delete_items() - def test_core(self, engine: Engine): + def test_core(self, engine: Engine) -> None: metadata = MetaData() item_table = Table( @@ -169,7 +169,7 @@ def test_core(self, engine: Engine): ) hnsw_index.create(engine) - def test_orm(self, engine: Engine): + def test_orm(self, engine: Engine) -> None: item = Item(embedding=Vector([1.5, 2, 3])) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -191,136 +191,136 @@ def test_orm(self, engine: Engine): assert items[1].embedding == Vector([4, 5, 6]) assert items[2].embedding is None - def test_vector(self, engine: Engine): + def test_vector(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.embedding == Vector([1, 2, 3]) - def test_vector_l2_distance(self, engine: Engine): + def test_vector_l2_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l2_distance_orm(self, engine: Engine): + def test_vector_l2_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self, engine: Engine): + def test_vector_max_inner_product(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_vector_max_inner_product_orm(self, engine: Engine): + def test_vector_max_inner_product_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self, engine: Engine): + def test_vector_cosine_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_vector_cosine_distance_orm(self, engine: Engine): + def test_vector_cosine_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self, engine: Engine): + def test_vector_l1_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l1_distance_orm(self, engine: Engine): + def test_vector_l1_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self, engine: Engine): + def test_halfvec(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.half_embedding == HalfVector([1, 2, 3]) - def test_halfvec_l2_distance(self, engine: Engine): + def test_halfvec_l2_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l2_distance_orm(self, engine: Engine): + def test_halfvec_l2_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self, engine: Engine): + def test_halfvec_max_inner_product(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_max_inner_product_orm(self, engine: Engine): + def test_halfvec_max_inner_product_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self, engine: Engine): + def test_halfvec_cosine_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_cosine_distance_orm(self, engine: Engine): + def test_halfvec_cosine_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self, engine: Engine): + def test_halfvec_l1_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l1_distance_orm(self, engine: Engine): + def test_halfvec_l1_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self, engine: Engine): + def test_bit(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() item = session.get_one(Item, 1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self, engine: Engine): + def test_bit_hamming_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_hamming_distance_orm(self, engine: Engine): + def test_bit_hamming_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self, engine: Engine): + def test_bit_jaccard_distance(self, engine: Engine) -> None: if engine == pg8000_engine: return @@ -329,7 +329,7 @@ def test_bit_jaccard_distance(self, engine: Engine): items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance_orm(self, engine: Engine): + def test_bit_jaccard_distance_orm(self, engine: Engine) -> None: if engine == pg8000_engine: return @@ -338,100 +338,100 @@ def test_bit_jaccard_distance_orm(self, engine: Engine): items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self, engine: Engine): + def test_sparsevec(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self, engine: Engine): + def test_sparsevec_l2_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l2_distance_orm(self, engine: Engine): + def test_sparsevec_l2_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self, engine: Engine): + def test_sparsevec_max_inner_product(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_max_inner_product_orm(self, engine: Engine): + def test_sparsevec_max_inner_product_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self, engine: Engine): + def test_sparsevec_cosine_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_cosine_distance_orm(self, engine: Engine): + def test_sparsevec_cosine_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self, engine: Engine): + def test_sparsevec_l1_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l1_distance_orm(self, engine: Engine): + def test_sparsevec_l1_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_subquery(self, engine: Engine): + def test_subquery(self, engine: Engine) -> None: create_items() with Session(engine) as session: subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() items = session.query(Item).order_by(Item.embedding.l2_distance(subquery)).all() assert [v.id for v in items] == [1, 3, 2] - def test_subquery_orm(self, engine: Engine): + def test_subquery_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() items = session.scalars(select(Item).order_by(Item.embedding.l2_distance(subquery))) assert [v.id for v in items] == [1, 3, 2] - def test_filter(self, engine: Engine): + def test_filter(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() assert [v.id for v in items] == [1] - def test_filter_orm(self, engine: Engine): + def test_filter_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] - def test_select(self, engine: Engine): + def test_select(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() assert items == (3,) - def test_select_orm(self, engine: Engine): + def test_select_orm(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() assert items == [3] - def test_avg(self, engine: Engine): + def test_avg(self, engine: Engine) -> None: with Session(engine) as session: res = session.query(avg(Item.embedding)).one()[0] assert res is None @@ -440,7 +440,7 @@ def test_avg(self, engine: Engine): res = session.query(avg(Item.embedding)).one()[0] assert res == Vector([2.5, 3.5, 4.5]) - def test_avg_orm(self, engine: Engine): + def test_avg_orm(self, engine: Engine) -> None: with Session(engine) as session: res = session.scalars(select(avg(Item.embedding))).first() assert res is None @@ -449,7 +449,7 @@ def test_avg_orm(self, engine: Engine): res = session.scalars(select(avg(Item.embedding))).one() assert res == Vector([2.5, 3.5, 4.5]) - def test_sum(self, engine: Engine): + def test_sum(self, engine: Engine) -> None: with Session(engine) as session: res = session.query(sum(Item.embedding)).one()[0] assert res is None @@ -458,7 +458,7 @@ def test_sum(self, engine: Engine): res = session.query(sum(Item.embedding)).one()[0] assert res == Vector([5, 7, 9]) - def test_sum_orm(self, engine: Engine): + def test_sum_orm(self, engine: Engine) -> None: with Session(engine) as session: res = session.scalars(select(sum(Item.embedding))).first() assert res is None @@ -467,7 +467,7 @@ def test_sum_orm(self, engine: Engine): res = session.scalars(select(sum(Item.embedding))).one() assert res == Vector([5, 7, 9]) - def test_bad_dimensions(self, engine: Engine): + def test_bad_dimensions(self, engine: Engine) -> None: item = Item(embedding=[1, 2]) with Session(engine) as session: session.add(item) @@ -475,7 +475,7 @@ def test_bad_dimensions(self, engine: Engine): session.commit() @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_bad_ndim(self, engine: Engine): + def test_bad_ndim(self, engine: Engine) -> None: item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: session.add(item) @@ -483,26 +483,26 @@ def test_bad_ndim(self, engine: Engine): session.commit() @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_bad_dtype(self, engine: Engine): + def test_bad_dtype(self, engine: Engine) -> None: item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='could not convert string to float'): session.commit() - def test_inspect(self, engine: Engine): + def test_inspect(self, engine: Engine) -> None: columns = inspect(engine).get_columns('sqlalchemy_orm_item') assert isinstance(columns[1]['type'], VECTOR) - def test_literal_binds(self, engine: Engine): + def test_literal_binds(self, engine: Engine) -> None: sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) - def test_insert(self, engine: Engine): + def test_insert(self, engine: Engine) -> None: with Session(engine) as session: session.execute(insert(Item).values(embedding=Vector([1, 2, 3]))) - def test_insert_bulk(self, engine: Engine): + def test_insert_bulk(self, engine: Engine) -> None: with Session(engine) as session: session.execute(insert(Item), [{'embedding': Vector([1, 2, 3])}]) @@ -511,7 +511,7 @@ def test_insert_bulk(self, engine: Engine): # with Session(engine) as session: # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) - def test_automap(self, engine: Engine): + def test_automap(self, engine: Engine) -> None: metadata = MetaData() metadata.reflect(engine, only=['sqlalchemy_orm_item']) AutoBase = automap_base(metadata=metadata) @@ -522,13 +522,13 @@ def test_automap(self, engine: Engine): item = session.query(AutoItem).first() assert item is not None and item.embedding == Vector([1, 2, 3]) - def test_half_precision(self, engine: Engine): + def test_half_precision(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_binary_quantize(self, engine: Engine): + def test_binary_quantize(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[-1, -2, -3])) session.add(Item(id=2, embedding=[1, -2, 3])) @@ -539,7 +539,7 @@ def test_binary_quantize(self, engine: Engine): items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] - def test_binary_quantize_reranking(self, engine: Engine): + def test_binary_quantize_reranking(self, engine: Engine) -> None: # recreate index (could also vacuum table) binary_quantize_index.drop(setup_engine) binary_quantize_index.create(setup_engine) @@ -558,10 +558,10 @@ def test_binary_quantize_reranking(self, engine: Engine): @pytest.mark.parametrize('engine', array_engines) class TestSqlalchemyArray: - def setup_method(self): + def setup_method(self) -> None: delete_items() - def test_vector_array(self, engine: Engine): + def test_vector_array(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) session.commit() @@ -570,7 +570,7 @@ def test_vector_array(self, engine: Engine): item = session.get_one(Item, 1) assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] - def test_halfvec_array(self, engine: Engine): + def test_halfvec_array(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, half_embeddings=[HalfVector([1, 2, 3]), HalfVector([4, 5, 6])])) session.commit() @@ -582,11 +582,11 @@ def test_halfvec_array(self, engine: Engine): @pytest.mark.parametrize('engine', async_engines) class TestSqlalchemyAsync: - def setup_method(self): + def setup_method(self) -> None: delete_items() @pytest.mark.asyncio - async def test_vector(self, engine: AsyncEngine): + async def test_vector(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -599,7 +599,7 @@ async def test_vector(self, engine: AsyncEngine): await engine.dispose() @pytest.mark.asyncio - async def test_halfvec(self, engine: AsyncEngine): + async def test_halfvec(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -612,7 +612,7 @@ async def test_halfvec(self, engine: AsyncEngine): await engine.dispose() @pytest.mark.asyncio - async def test_bit(self, engine: AsyncEngine): + async def test_bit(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -630,7 +630,7 @@ async def test_bit(self, engine: AsyncEngine): await engine.dispose() @pytest.mark.asyncio - async def test_sparsevec(self, engine: AsyncEngine): + async def test_sparsevec(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -643,7 +643,7 @@ async def test_sparsevec(self, engine: AsyncEngine): await engine.dispose() @pytest.mark.asyncio - async def test_avg(self, engine: AsyncEngine): + async def test_avg(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -658,11 +658,11 @@ async def test_avg(self, engine: AsyncEngine): @pytest.mark.parametrize('engine', async_array_engines) class TestSqlalchemyAsyncArray: - def setup_method(self): + def setup_method(self) -> None: delete_items() @pytest.mark.asyncio - async def test_vector_array(self, engine: AsyncEngine): + async def test_vector_array(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 4cdc140..0d98f9e 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -48,12 +48,12 @@ def create_items() -> None: class TestSqlmodel: - def setup_method(self): + def setup_method(self) -> None: with Session(engine) as session: session.exec(delete(Item)) session.commit() - def test_orm(self): + def test_orm(self) -> None: item = Item(embedding=[1.5, 2, 3]) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -74,131 +74,131 @@ def test_orm(self): assert items[1].embedding == Vector([4, 5, 6]) assert items[2].embedding is None - def test_vector(self): + def test_vector(self) -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.embedding == Vector([1, 2, 3]) - def test_vector_l2_distance(self): + def test_vector_l2_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self): + def test_halfvec(self) -> None: with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.half_embedding == HalfVector([1, 2, 3]) - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_bit(self): + def test_bit(self) -> None: with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() item = session.get_one(Item, 1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self): + def test_sparsevec(self) -> None: with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_filter(self): + def test_filter(self) -> None: create_items() with Session(engine) as session: items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) # type: ignore assert [v.id for v in items] == [1] - def test_select(self): + def test_select(self) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() # type: ignore assert items[0] == 3 - def test_vector_avg(self): + def test_vector_avg(self) -> None: with Session(engine) as session: res = session.exec(select(avg(Item.embedding))).first() assert res is None @@ -207,7 +207,7 @@ def test_vector_avg(self): res = session.exec(select(avg(Item.embedding))).first() assert res == Vector([2.5, 3.5, 4.5]) - def test_vector_sum(self): + def test_vector_sum(self) -> None: with Session(engine) as session: res = session.exec(select(sum(Item.embedding))).first() assert res is None @@ -216,7 +216,7 @@ def test_vector_sum(self): res = session.exec(select(sum(Item.embedding))).first() assert res == Vector([5, 7, 9]) - def test_halfvec_avg(self): + def test_halfvec_avg(self) -> None: with Session(engine) as session: res = session.exec(select(avg(Item.half_embedding))).first() assert res is None @@ -225,7 +225,7 @@ def test_halfvec_avg(self): res = session.exec(select(avg(Item.half_embedding))).first() assert res == HalfVector([2.5, 3.5, 4.5]) - def test_halfvec_sum(self): + def test_halfvec_sum(self) -> None: with Session(engine) as session: res = session.exec(select(sum(Item.half_embedding))).first() assert res is None @@ -234,7 +234,7 @@ def test_halfvec_sum(self): res = session.exec(select(sum(Item.half_embedding))).first() assert res == HalfVector([5, 7, 9]) - def test_bad_dimensions(self): + def test_bad_dimensions(self) -> None: item = Item(embedding=[1, 2]) with Session(engine) as session: session.add(item) diff --git a/tests/test_vector.py b/tests/test_vector.py index 3392df6..db23648 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -10,58 +10,58 @@ class TestVector: - def test_list(self): + def test_list(self) -> None: assert Vector([1, 2, 3]).to_list() == [1, 2, 3] - def test_list_empty(self): + def test_list_empty(self) -> None: assert Vector([]).to_list() == [] - def test_list_str(self): + def test_list_str(self) -> None: with pytest.raises(ValueError) as error: Vector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_list_list(self): + def test_list_list(self) -> None: with pytest.raises(ValueError) as error: Vector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_ndarray(self): + def test_ndarray(self) -> None: arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr - def test_int(self): + def test_int(self) -> None: with pytest.raises(ValueError) as error: Vector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' - def test_repr(self): + def test_repr(self) -> None: assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' - def test_equality(self): + def test_equality(self) -> None: assert Vector([1, 2, 3]) == Vector([1, 2, 3]) assert Vector([1, 2, 3]) != Vector([1, 2, 4]) - def test_dimensions(self): + def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_to_numpy_readonly(self): + def test_to_numpy_readonly(self) -> None: arr = Vector([1, 2, 3]).to_numpy() with pytest.raises(ValueError) as error: arr[0] = 4 assert str(error.value) == 'assignment destination is read-only' - def test_from_text(self): + def test_from_text(self) -> None: vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - def test_from_binary(self): + def test_from_binary(self) -> None: data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] From 22958488b6cf866c91219deab736611e07485413 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:26:13 -0700 Subject: [PATCH 693/879] Improved typing [skip ci] --- tests/test_asyncpg.py | 2 +- tests/test_sqlalchemy.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index b8b42bf..8ff9f3b 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -131,7 +131,7 @@ async def test_vector_array(self) -> None: @pytest.mark.asyncio async def test_pool(self) -> None: - async def init(conn) -> None: + async def init(conn: Connection) -> None: await register_vector(conn) pool = await asyncpg.create_pool(database='pgvector_python_test', init=init) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 08621b9..c376fa8 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -9,6 +9,7 @@ from sqlalchemy.ext.automap import automap_base from sqlalchemy.orm import mapped_column, DeclarativeBase, Session from sqlalchemy.sql import func +from typing import Any try: import numpy as np @@ -21,7 +22,7 @@ @event.listens_for(psycopg2_type_engine, 'connect') -def psycopg2_connect(dbapi_connection, connection_record): +def psycopg2_connect(dbapi_connection: Any, connection_record: Any) -> None: from pgvector.psycopg2 import register_vector register_vector(dbapi_connection) @@ -33,7 +34,7 @@ def psycopg2_connect(dbapi_connection, connection_record): @event.listens_for(psycopg_type_engine, 'connect') -def psycopg_connect(dbapi_connection, connection_record) -> None: +def psycopg_connect(dbapi_connection: Any, connection_record: Any) -> None: from pgvector.psycopg import register_vector register_vector(dbapi_connection) @@ -43,7 +44,7 @@ def psycopg_connect(dbapi_connection, connection_record) -> None: @event.listens_for(psycopg_async_type_engine.sync_engine, 'connect') -def psycopg_async_connect(dbapi_connection, connection_record) -> None: +def psycopg_async_connect(dbapi_connection: Any, connection_record: Any) -> None: from pgvector.psycopg import register_vector_async dbapi_connection.run_async(register_vector_async) @@ -53,7 +54,7 @@ def psycopg_async_connect(dbapi_connection, connection_record) -> None: @event.listens_for(asyncpg_type_engine.sync_engine, 'connect') -def asyncpg_connect(dbapi_connection, connection_record) -> None: +def asyncpg_connect(dbapi_connection: Any, connection_record: Any) -> None: from pgvector.asyncpg import register_vector dbapi_connection.run_async(register_vector) From 910e82bdebf7d1f3dd1384e8882b08332c25af27 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:29:25 -0700 Subject: [PATCH 694/879] Moved ty config [skip ci] --- Makefile | 3 --- pyproject.toml | 9 +++++++++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index cd5b0ac..7462142 100644 --- a/Makefile +++ b/Makefile @@ -3,9 +3,6 @@ lint: pycodestyle . --ignore=E501 --exclude=.venv -check: - ty check pgvector tests --ignore unused-type-ignore-comment - build: python3 -m build diff --git a/pyproject.toml b/pyproject.toml index faf1aa1..d005b48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,15 @@ ignore_missing_imports = true disallow_subclassing_any = false python_executable = ".venv/bin/python" +[tool.ty.src] +include = [ + "pgvector", + "tests" +] + +[tool.ty.rules] +unused-type-ignore-comment = "ignore" + [tool.pyright] include = [ "pgvector", From d58bb3e44f867418d060b3b074b52fcee9ea9603 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:41:56 -0700 Subject: [PATCH 695/879] Improved pyright config [skip ci] --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index d005b48..9a5362a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,10 @@ include = [ "tests" ] typeCheckingMode = "strict" +reportMissingTypeStubs = "none" reportPossiblyUnboundVariable = "none" +reportUnknownMemberType = "none" +reportUnknownVariableType = "none" venv = ".venv" venvPath = "." From 01950d1a276369a37bb0e252e31a7f453188954f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:43:20 -0700 Subject: [PATCH 696/879] Improved pyrefly config [skip ci] --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 9a5362a..57e520e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,3 +78,4 @@ project-includes = [ "pgvector/**/*.py", "tests/**/*.py" ] +preset = "strict" From f6f713ccf55d8ff132d191bdef6b73047999afe0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:45:29 -0700 Subject: [PATCH 697/879] Improved pyrefly config [skip ci] --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 57e520e..cf9b3fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,3 +79,6 @@ project-includes = [ "tests/**/*.py" ] preset = "strict" + +[tool.pyrefly.errors] +missing-override-decorator = false From 6a1e7bf9c7b74d35f762d349b8d818abd42e2a32 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:47:10 -0700 Subject: [PATCH 698/879] Removed unused imports --- tests/test_pg8000.py | 1 - tests/test_sqlmodel.py | 6 ------ 2 files changed, 7 deletions(-) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index c7ad466..ffb90ac 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -2,7 +2,6 @@ from pgvector import HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection -import pytest try: import numpy as np diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 0d98f9e..ee0070f 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -5,12 +5,6 @@ from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text from typing import Any, Optional -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False - engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) # type: ignore From 087a2dfc09651b4457f20cbdc5e42ffd22a6af85 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:48:45 -0700 Subject: [PATCH 699/879] Improved pyright config [skip ci] --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index cf9b3fb..b385fe0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,10 +66,13 @@ include = [ "tests" ] typeCheckingMode = "strict" +reportConstantRedefinition = "none" reportMissingTypeStubs = "none" reportPossiblyUnboundVariable = "none" +reportUnknownArgumentType = "none" reportUnknownMemberType = "none" reportUnknownVariableType = "none" +reportUnnecessaryIsInstance = "none" venv = ".venv" venvPath = "." From 56283d41dcca9106447c85af7ba847f5f1260f82 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 17:54:26 -0700 Subject: [PATCH 700/879] Fixed typechecking [skip ci] --- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- pgvector/psycopg2/halfvec.py | 3 ++- pgvector/psycopg2/sparsevec.py | 3 ++- pgvector/psycopg2/vector.py | 3 ++- 6 files changed, 9 insertions(+), 6 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 6059f87..f3d3719 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -51,7 +51,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw class HalfVectorWidget(forms.TextInput): - def format_value(self, value: Any) -> str | None: + def format_value(self, value: Any) -> Any: if isinstance(value, HalfVector): value = value.to_list() return super().format_value(value) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index bec664f..cd289a4 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -46,7 +46,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw class SparseVectorWidget(forms.TextInput): - def format_value(self, value: Any) -> str | None: + def format_value(self, value: Any) -> Any: if isinstance(value, SparseVector): value = value.to_text() return super().format_value(value) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 89b20d6..e5ab444 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -51,7 +51,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw class VectorWidget(forms.TextInput): - def format_value(self, value: Any) -> str | None: + def format_value(self, value: Any) -> Any: if isinstance(value, Vector): value = value.to_list() return super().format_value(value) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 9fa311b..4616095 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,4 +1,5 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import HalfVector @@ -6,7 +7,7 @@ class HalfvecAdapter: def __init__(self, value: HalfVector) -> None: self._value = value - def getquoted(self) -> bytes: + def getquoted(self) -> Any: return adapt(self._value.to_text()).getquoted() diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index b733200..9b475ff 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,4 +1,5 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import SparseVector @@ -6,7 +7,7 @@ class SparsevecAdapter: def __init__(self, value: SparseVector) -> None: self._value = value - def getquoted(self) -> bytes: + def getquoted(self) -> Any: return adapt(self._value.to_text()).getquoted() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 0eca9db..0d11240 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,4 +1,5 @@ from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import Vector try: @@ -14,7 +15,7 @@ def __init__(self, value: 'Vector | np.ndarray') -> None: value = Vector(value) self._value = value - def getquoted(self) -> bytes: + def getquoted(self) -> Any: return adapt(self._value.to_text()).getquoted() From f6fb996fd61db25b866a293da0a500a63d4d6b2d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 30 Jun 2026 18:18:18 -0700 Subject: [PATCH 701/879] Improved code [skip ci] --- pgvector/bit.py | 4 ++-- pgvector/halfvec.py | 4 ++-- pgvector/vector.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index b5c48e6..c9bb50d 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -60,8 +60,8 @@ def __eq__(self, other: object) -> bool: return False def _length(self) -> int: - length, = unpack_from('>i', self._value) - return cast(int, length) + length, = cast(tuple[int], unpack_from('>i', self._value)) + return length def to_list(self) -> list[bool]: # TODO improve diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 7a14d12..95e5254 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -38,8 +38,8 @@ def __eq__(self, other: object) -> bool: return False def dimensions(self) -> int: - dim, = struct.unpack_from('>H', self._value) - return cast(int, dim) + dim, = cast(tuple[int], struct.unpack_from('>H', self._value)) + return dim def to_list(self) -> list[float]: return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:])) diff --git a/pgvector/vector.py b/pgvector/vector.py index 1841d93..a4b6041 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -38,8 +38,8 @@ def __eq__(self, other: object) -> bool: return False def dimensions(self) -> int: - dim, = struct.unpack_from('>H', self._value) - return cast(int, dim) + dim, = cast(tuple[int], struct.unpack_from('>H', self._value)) + return dim def to_list(self) -> list[float]: return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:])) From ddf5c183ac82d7da87cf87e86ae1b9d9a024ba91 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 22:45:15 -0700 Subject: [PATCH 702/879] Improved pyright config [skip ci] --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index b385fe0..3dc4c11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ typeCheckingMode = "strict" reportConstantRedefinition = "none" reportMissingTypeStubs = "none" reportPossiblyUnboundVariable = "none" +reportPrivateUsage = "none" reportUnknownArgumentType = "none" reportUnknownMemberType = "none" reportUnknownVariableType = "none" From 4a10b3561d5006038a8b7f26e28ebe0fc9efe50c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 22:46:34 -0700 Subject: [PATCH 703/879] Improved typing [skip ci] --- tests/test_sqlalchemy.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index c376fa8..9b18160 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -7,7 +7,7 @@ from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncEngine from sqlalchemy.ext.automap import automap_base -from sqlalchemy.orm import mapped_column, DeclarativeBase, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session from sqlalchemy.sql import func from typing import Any @@ -78,13 +78,13 @@ class Base(DeclarativeBase): class Item(Base): __tablename__ = 'sqlalchemy_orm_item' - id = mapped_column(Integer, primary_key=True) - embedding = mapped_column(VECTOR(3)) - half_embedding = mapped_column(HALFVEC(3)) - binary_embedding = mapped_column(BIT(3)) - sparse_embedding = mapped_column(SPARSEVEC(3)) - embeddings = mapped_column(ARRAY(VECTOR(3))) - half_embeddings = mapped_column(ARRAY(HALFVEC(3))) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + embedding: Mapped[Vector] = mapped_column(VECTOR(3)) + half_embedding: Mapped[HalfVector] = mapped_column(HALFVEC(3)) + binary_embedding: Mapped[Any] = mapped_column(BIT(3)) + sparse_embedding: Mapped[SparseVector] = mapped_column(SPARSEVEC(3)) + embeddings: Mapped[list[Vector]] = mapped_column(ARRAY(VECTOR(3))) + half_embeddings: Mapped[list[HalfVector]] = mapped_column(ARRAY(HALFVEC(3))) Base.metadata.drop_all(setup_engine) From b5c9e3900ceb9739cdad4f7824dc24649f51a4bb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 22:48:34 -0700 Subject: [PATCH 704/879] Improved typing [skip ci] --- tests/test_sqlmodel.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index ee0070f..af31504 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -3,7 +3,7 @@ import pytest from sqlalchemy.exc import StatementError from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text -from typing import Any, Optional +from typing import Any engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: @@ -13,11 +13,11 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' # type: ignore - id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) # type: ignore - half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) # type: ignore - binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) # type: ignore - sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore + id: int | None = Field(default=None, primary_key=True) + embedding: Any = Field(default=None, sa_type=VECTOR(3)) # type: ignore + half_embedding: Any = Field(default=None, sa_type=HALFVEC(3)) # type: ignore + binary_embedding: Any = Field(default=None, sa_type=BIT(3)) # type: ignore + sparse_embedding: Any = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore SQLModel.metadata.drop_all(engine) From 2096c58e8825961d0444a47fbdb0e36e2bee5f7f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 22:51:57 -0700 Subject: [PATCH 705/879] Improved typing [skip ci] --- tests/test_sqlmodel.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index af31504..d5a123c 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -3,7 +3,6 @@ import pytest from sqlalchemy.exc import StatementError from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text -from typing import Any engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: @@ -14,10 +13,10 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' # type: ignore id: int | None = Field(default=None, primary_key=True) - embedding: Any = Field(default=None, sa_type=VECTOR(3)) # type: ignore - half_embedding: Any = Field(default=None, sa_type=HALFVEC(3)) # type: ignore - binary_embedding: Any = Field(default=None, sa_type=BIT(3)) # type: ignore - sparse_embedding: Any = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore + embedding: Vector | list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore + half_embedding: HalfVector | list[float] | None = Field(default=None, sa_type=HALFVEC(3)) # type: ignore + binary_embedding: str | None = Field(default=None, sa_type=BIT(3)) # type: ignore + sparse_embedding: SparseVector | list[float] | None = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore SQLModel.metadata.drop_all(engine) From 662050838828939331874eb432769eb8cd95b159 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 23:07:10 -0700 Subject: [PATCH 706/879] Improved typing [skip ci] --- tests/test_django.py | 104 +++++++++++++++++++++---------------------- 1 file changed, 51 insertions(+), 53 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index ac2b626..294f5bb 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,5 +1,3 @@ -# type: ignore - import django from django.conf import settings from django.contrib.postgres.fields import ArrayField @@ -135,7 +133,7 @@ class Migration(migrations.Migration): cursor.execute('\n'.join(sql_statements)) -def create_items(): +def create_items() -> None: Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])).save() Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])).save() Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])).save() @@ -166,147 +164,147 @@ class Meta: class TestDjango: - def setup_method(self): + def setup_method(self) -> None: Item.objects.all().delete() - def test_vector(self): + def test_vector(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) assert item.embedding == Vector([1, 2, 3]) - def test_vector_l2_distance(self): + def test_vector_l2_distance(self) -> None: create_items() distance = L2Distance('embedding', Vector([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)] - def test_vector_l2_distance_list(self): + def test_vector_l2_distance_list(self) -> None: create_items() distance = L2Distance('embedding', [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)] - def test_vector_l2_distance_none(self): + def test_vector_l2_distance_none(self) -> None: create_items() distance = L2Distance('embedding', None) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.distance for v in items] == [None, None, None] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self) -> None: create_items() distance = CosineDistance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self) -> None: create_items() distance = L1Distance('embedding', [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, 3] - def test_halfvec(self): + def test_halfvec(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) assert item.half_embedding == HalfVector([1, 2, 3]) - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self) -> None: create_items() distance = L2Distance('half_embedding', HalfVector([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)] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self) -> None: create_items() distance = CosineDistance('half_embedding', HalfVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self) -> None: create_items() distance = L1Distance('half_embedding', HalfVector([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, 3] - def test_bit(self): + def test_bit(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self) -> None: create_items() distance = HammingDistance('binary_embedding', '101') items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self) -> None: create_items() distance = JaccardDistance('binary_embedding', '101') items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] # assert [v.distance for v in items] == [0, 1/3, 1] - def test_sparsevec(self): + def test_sparsevec(self) -> None: Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() item = Item.objects.get(pk=1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() distance = L2Distance('sparse_embedding', SparseVector([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)] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() distance = CosineDistance('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self) -> None: create_items() distance = L1Distance('sparse_embedding', SparseVector([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, 3] - def test_filter(self): + def test_filter(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) items = Item.objects.alias(distance=distance).filter(distance__lt=1) assert [v.id for v in items] == [1] - def test_vector_avg(self): + def test_vector_avg(self) -> None: avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert avg is None Item(embedding=[1, 2, 3]).save() @@ -314,7 +312,7 @@ def test_vector_avg(self): avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert avg == Vector([2.5, 3.5, 4.5]) - def test_vector_sum(self): + def test_vector_sum(self) -> None: sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert sum is None Item(embedding=[1, 2, 3]).save() @@ -322,7 +320,7 @@ def test_vector_sum(self): sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert sum == Vector([5, 7, 9]) - def test_halfvec_avg(self): + def test_halfvec_avg(self) -> None: avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] assert avg is None Item(half_embedding=[1, 2, 3]).save() @@ -330,7 +328,7 @@ def test_halfvec_avg(self): avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] assert avg == HalfVector([2.5, 3.5, 4.5]) - def test_halfvec_sum(self): + def test_halfvec_sum(self) -> None: sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] assert sum is None Item(half_embedding=[1, 2, 3]).save() @@ -338,7 +336,7 @@ def test_halfvec_sum(self): sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] assert sum == HalfVector([5, 7, 9]) - def test_serialization(self): + def test_serialization(self) -> None: create_items() items = Item.objects.all() for format in ['json', 'xml']: @@ -348,18 +346,18 @@ def test_serialization(self): for obj in serializers.deserialize(format, data): obj.save() - def test_vector_form(self): + def test_vector_form(self) -> None: form = VectorForm(data={'embedding': '[1, 2, 3]'}) assert form.is_valid() assert 'value="[1, 2, 3]"' in form.as_div() - def test_vector_form_instance(self): + def test_vector_form_instance(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item) assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() - def test_vector_form_save(self): + def test_vector_form_save(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item, data={'embedding': '[4, 5, 6]'}) @@ -368,7 +366,7 @@ def test_vector_form_save(self): assert form.save() assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6]) - def test_vector_form_save_missing(self): + def test_vector_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item, data={'embedding': ''}) @@ -376,18 +374,18 @@ def test_vector_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).embedding is None - def test_halfvec_form(self): + def test_halfvec_form(self) -> None: form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) assert form.is_valid() assert 'value="[1, 2, 3]"' in form.as_div() - def test_halfvec_form_instance(self): + def test_halfvec_form_instance(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item) assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() - def test_halfvec_form_save(self): + def test_halfvec_form_save(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': '[4, 5, 6]'}) @@ -396,7 +394,7 @@ def test_halfvec_form_save(self): assert form.save() assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) - def test_halfvec_form_save_missing(self): + def test_halfvec_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': ''}) @@ -404,18 +402,18 @@ def test_halfvec_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).half_embedding is None - def test_bit_form(self): + def test_bit_form(self) -> None: form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() assert 'value="101"' in form.as_div() - def test_bit_form_instance(self): + def test_bit_form_instance(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) form = BitForm(instance=item) assert 'value="101"' in form.as_div() - def test_bit_form_save(self): + def test_bit_form_save(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) form = BitForm(instance=item, data={'binary_embedding': '010'}) @@ -424,7 +422,7 @@ def test_bit_form_save(self): assert form.save() assert '010' == Item.objects.get(pk=1).binary_embedding - def test_bit_form_save_missing(self): + def test_bit_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = BitForm(instance=item, data={'binary_embedding': ''}) @@ -432,19 +430,19 @@ def test_bit_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).binary_embedding is None - def test_sparsevec_form(self): + def test_sparsevec_form(self) -> None: form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) assert form.is_valid() assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() - def test_sparsevec_form_instance(self): + def test_sparsevec_form_instance(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item) # TODO improve assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() - def test_sparsevec_form_save(self): + def test_sparsevec_form_save(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item, data={'sparse_embedding': '{1:4,2:5,3:6}/3'}) @@ -453,7 +451,7 @@ def test_sparsevec_form_save(self): assert form.save() assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) - def test_sparesevec_form_save_missing(self): + def test_sparesevec_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item, data={'sparse_embedding': ''}) @@ -461,21 +459,21 @@ def test_sparesevec_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).sparse_embedding is None - def test_clean(self): + def test_clean(self) -> None: item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector([1, 2, 3])) item.full_clean() - def test_get_or_create(self): + def test_get_or_create(self) -> None: Item.objects.get_or_create(embedding=[1, 2, 3]) - def test_missing(self): + def test_missing(self) -> None: Item().save() assert Item.objects.first().embedding is None 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): + def test_vector_array(self) -> None: Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]).save() with connection.cursor() as cursor: @@ -486,7 +484,7 @@ def test_vector_array(self): item = Item.objects.get(pk=1) assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] - def test_double_array(self): + def test_double_array(self) -> None: 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() @@ -496,7 +494,7 @@ def test_double_array(self): assert [v.distance for v in items] == [0, 1, sqrt(3)] assert items[1].double_embedding == [1, 1, 2] - def test_numeric_array(self): + def test_numeric_array(self) -> None: 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() @@ -506,7 +504,7 @@ def test_numeric_array(self): assert [v.distance for v in items] == [0, 1, sqrt(3)] assert items[1].numeric_embedding == [1, 1, 2] - def test_half_precision(self): + def test_half_precision(self) -> None: create_items() distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) From 1121306beefcacd4e6073a310a2fa970a77f7da7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 23:19:08 -0700 Subject: [PATCH 707/879] Removed unused import [skip ci] --- tests/test_django.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 294f5bb..7747ec4 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -15,12 +15,6 @@ from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False - settings.configure( DATABASES={ 'default': { From 0e8ae4e3d63fe3b72e4c4bcd61faef860bcd8c8b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 23:19:27 -0700 Subject: [PATCH 708/879] Fixed tests --- tests/test_sqlalchemy.py | 12 ++++++------ tests/test_sqlmodel.py | 2 ++ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 9b18160..6edcda7 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -79,12 +79,12 @@ class Item(Base): __tablename__ = 'sqlalchemy_orm_item' id: Mapped[int] = mapped_column(Integer, primary_key=True) - embedding: Mapped[Vector] = mapped_column(VECTOR(3)) - half_embedding: Mapped[HalfVector] = mapped_column(HALFVEC(3)) - binary_embedding: Mapped[Any] = mapped_column(BIT(3)) - sparse_embedding: Mapped[SparseVector] = mapped_column(SPARSEVEC(3)) - embeddings: Mapped[list[Vector]] = mapped_column(ARRAY(VECTOR(3))) - half_embeddings: Mapped[list[HalfVector]] = mapped_column(ARRAY(HALFVEC(3))) + embedding: Mapped[Vector | None] = mapped_column(VECTOR(3)) + half_embedding: Mapped[HalfVector | None] = mapped_column(HALFVEC(3)) + binary_embedding: Mapped[Any | None] = mapped_column(BIT(3)) + sparse_embedding: Mapped[SparseVector | None] = mapped_column(SPARSEVEC(3)) + embeddings: Mapped[list[Vector] | None] = mapped_column(ARRAY(VECTOR(3))) + half_embeddings: Mapped[list[HalfVector] | None] = mapped_column(ARRAY(HALFVEC(3))) Base.metadata.drop_all(setup_engine) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index d5a123c..b6f4510 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,5 +1,6 @@ from pgvector import HalfVector, SparseVector, Vector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum +from pydantic import ConfigDict import pytest from sqlalchemy.exc import StatementError from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text @@ -11,6 +12,7 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' # type: ignore + model_config = ConfigDict(arbitrary_types_allowed=True) # type: ignore id: int | None = Field(default=None, primary_key=True) embedding: Vector | list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore From e755615b6463cae104be5fa4ec89c0ed436159e3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 1 Jul 2026 23:29:06 -0700 Subject: [PATCH 709/879] Fixed typechecking [skip ci] --- tests/test_django.py | 136 +++++++++++++++++++++---------------------- 1 file changed, 68 insertions(+), 68 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 7747ec4..613d404 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -159,180 +159,180 @@ class Meta: class TestDjango: def setup_method(self) -> None: - Item.objects.all().delete() + Item.objects.all().delete() # type: ignore def test_vector(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore assert item.embedding == Vector([1, 2, 3]) def test_vector_l2_distance(self) -> None: create_items() distance = L2Distance('embedding', Vector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] def test_vector_l2_distance_list(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] def test_vector_l2_distance_none(self) -> None: create_items() distance = L2Distance('embedding', None) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.distance for v in items] == [None, None, None] def test_vector_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_vector_cosine_distance(self) -> None: create_items() distance = CosineDistance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_vector_l1_distance(self) -> None: create_items() distance = L1Distance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] def test_halfvec(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self) -> None: create_items() distance = L2Distance('half_embedding', HalfVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] def test_halfvec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_halfvec_cosine_distance(self) -> None: create_items() distance = CosineDistance('half_embedding', HalfVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_halfvec_l1_distance(self) -> None: create_items() distance = L1Distance('half_embedding', HalfVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] def test_bit(self) -> None: Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore assert item.binary_embedding == '101' def test_bit_hamming_distance(self) -> None: create_items() distance = HammingDistance('binary_embedding', '101') - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] def test_bit_jaccard_distance(self) -> None: create_items() distance = JaccardDistance('binary_embedding', '101') - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [2, 3, 1] # assert [v.distance for v in items] == [0, 1/3, 1] def test_sparsevec(self) -> None: Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self) -> None: create_items() distance = L2Distance('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] def test_sparsevec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_sparsevec_cosine_distance(self) -> None: create_items() distance = CosineDistance('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_sparsevec_l1_distance(self) -> None: create_items() distance = L1Distance('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] def test_filter(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) - items = Item.objects.alias(distance=distance).filter(distance__lt=1) + items = Item.objects.alias(distance=distance).filter(distance__lt=1) # type: ignore assert [v.id for v in items] == [1] def test_vector_avg(self) -> None: - avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] + avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] # type: ignore assert avg is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() - avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] + avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] # type: ignore assert avg == Vector([2.5, 3.5, 4.5]) def test_vector_sum(self) -> None: - sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] + sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] # type: ignore assert sum is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() - sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] + sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] # type: ignore assert sum == Vector([5, 7, 9]) def test_halfvec_avg(self) -> None: - avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] + avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] # type: ignore assert avg is None Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() - avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] + avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] # type: ignore assert avg == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self) -> None: - sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] # type: ignore assert sum is None Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() - sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] # type: ignore assert sum == HalfVector([5, 7, 9]) def test_serialization(self) -> None: create_items() - items = Item.objects.all() + items = Item.objects.all() # type: ignore for format in ['json', 'xml']: data = serializers.serialize(format, items) with mock.patch('django.core.serializers.python.apps.get_model') as get_model: @@ -343,129 +343,129 @@ def test_serialization(self) -> None: def test_vector_form(self) -> None: form = VectorForm(data={'embedding': '[1, 2, 3]'}) assert form.is_valid() - assert 'value="[1, 2, 3]"' in form.as_div() + assert 'value="[1, 2, 3]"' in form.as_div() # type: ignore def test_vector_form_instance(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = VectorForm(instance=item) - assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() # type: ignore def test_vector_form_save(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = VectorForm(instance=item, data={'embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6]) + assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6]) # type: ignore def test_vector_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = VectorForm(instance=item, data={'embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).embedding is None + assert Item.objects.get(pk=1).embedding is None # type: ignore def test_halfvec_form(self) -> None: form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) assert form.is_valid() - assert 'value="[1, 2, 3]"' in form.as_div() + assert 'value="[1, 2, 3]"' in form.as_div() # type: ignore def test_halfvec_form_instance(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = HalfVectorForm(instance=item) - assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() # type: ignore def test_halfvec_form_save(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = HalfVectorForm(instance=item, data={'half_embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) + assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) # type: ignore def test_halfvec_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = HalfVectorForm(instance=item, data={'half_embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).half_embedding is None + assert Item.objects.get(pk=1).half_embedding is None # type: ignore def test_bit_form(self) -> None: form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() - assert 'value="101"' in form.as_div() + assert 'value="101"' in form.as_div() # type: ignore def test_bit_form_instance(self) -> None: Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = BitForm(instance=item) - assert 'value="101"' in form.as_div() + assert 'value="101"' in form.as_div() # type: ignore def test_bit_form_save(self) -> None: Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = BitForm(instance=item, data={'binary_embedding': '010'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert '010' == Item.objects.get(pk=1).binary_embedding + assert '010' == Item.objects.get(pk=1).binary_embedding # type: ignore def test_bit_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = BitForm(instance=item, data={'binary_embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).binary_embedding is None + assert Item.objects.get(pk=1).binary_embedding is None # type: ignore def test_sparsevec_form(self) -> None: form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) assert form.is_valid() - assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() + assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() # type: ignore def test_sparsevec_form_instance(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = SparseVectorForm(instance=item) # TODO improve - assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() + assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() # type: ignore def test_sparsevec_form_save(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = SparseVectorForm(instance=item, data={'sparse_embedding': '{1:4,2:5,3:6}/3'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) + assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) # type: ignore def test_sparesevec_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore form = SparseVectorForm(instance=item, data={'sparse_embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).sparse_embedding is None + assert Item.objects.get(pk=1).sparse_embedding is None # type: ignore def test_clean(self) -> None: item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector([1, 2, 3])) item.full_clean() def test_get_or_create(self) -> None: - Item.objects.get_or_create(embedding=[1, 2, 3]) + Item.objects.get_or_create(embedding=[1, 2, 3]) # type: ignore def test_missing(self) -> None: Item().save() - assert Item.objects.first().embedding is None - assert Item.objects.first().half_embedding is None - assert Item.objects.first().binary_embedding is None - assert Item.objects.first().sparse_embedding is None + assert Item.objects.first().embedding is None # type: ignore + assert Item.objects.first().half_embedding is None # type: ignore + assert Item.objects.first().binary_embedding is None # type: ignore + assert Item.objects.first().sparse_embedding is None # type: ignore def test_vector_array(self) -> None: Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]).save() @@ -475,7 +475,7 @@ def test_vector_array(self) -> None: register_vector(cursor.connection) # this fails if the driver does not cast arrays - item = Item.objects.get(pk=1) + item = Item.objects.get(pk=1) # type: ignore assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] def test_double_array(self) -> None: @@ -483,7 +483,7 @@ def test_double_array(self) -> None: 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) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore 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] @@ -493,7 +493,7 @@ def test_numeric_array(self) -> None: 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) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore 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] @@ -501,6 +501,6 @@ def test_numeric_array(self) -> None: def test_half_precision(self) -> None: create_items() distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) + items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] From 8c6412d1cd770f346325a8a4ffd498a201a5ee50 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 00:20:07 -0700 Subject: [PATCH 710/879] Improved typing [skip ci] --- pgvector/sparsevec.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 3fd9b53..22a8bff 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -116,8 +116,8 @@ def _from_dense(self, value: list[float]) -> None: @classmethod def from_text(cls, value: str) -> SparseVector: elements, dim = value.split('/', 2) - indices = [] - values = [] + indices: list[int] = [] + values: list[float] = [] # split on empty string returns single element list if len(elements) > 2: for e in elements[1:-1].split(','): From 27a16c28c2e092993f0d4e7d325056ee6f5fb777 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 11:38:34 -0700 Subject: [PATCH 711/879] Updated changelog [skip ci] --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d89fdb..37bd613 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## 1.0.0 (unreleased) -- Added type hints +- Added experimental support for type hints - Changed `vector` type to return `Vector` object instead of NumPy array - Removed `utils` package (use top-level `pgvector` package instead) - Removed re-exported classes (use top-level `pgvector` package instead) From dc7e3e2a8418c673b973ed9b472417fb7962263a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 12:05:21 -0700 Subject: [PATCH 712/879] Switched to TypeDecorator for bit type for SQLAlchemy --- pgvector/sqlalchemy/bit.py | 40 +++++++++++--------------------------- 1 file changed, 11 insertions(+), 29 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 2679ee3..a86ce1a 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -1,42 +1,24 @@ -from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, TypeEngine, Float -from sqlalchemy import Dialect, Operators +from sqlalchemy.dialects.postgresql.base import PGBit +from sqlalchemy.types import TypeDecorator, Float +from sqlalchemy import Operators from typing import Any -class BIT(UserDefinedType[Any]): +class BIT(TypeDecorator[Any]): + impl = PGBit cache_ok = True - def __init__(self, length: int | None = None) -> None: - super().__init__() - self.length = length - - def get_col_spec(self, **kw: Any) -> str: - if self.length is None: - return 'BIT' - return 'BIT(%d)' % self.length - - def bind_processor(self, dialect: Dialect) -> Any: - if dialect.__class__.__name__ == 'PGDialect_asyncpg': + def process_bind_param(self, value: Any, dialect: Any) -> Any: + if dialect.__class__.__name__ == 'PGDialect_asyncpg' and isinstance(value, str): import asyncpg + return asyncpg.BitString(value) # type: ignore + return value - def process(value: Any) -> Any: - if isinstance(value, str): - return asyncpg.BitString(value) # type: ignore - return value - return process - else: - return super().bind_processor(dialect) - - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeDecorator.Comparator[Any]): def hamming_distance(self, other: object) -> Operators: return self.op('<~>', return_type=Float)(other) def jaccard_distance(self, other: object) -> Operators: return self.op('<%>', return_type=Float)(other) - comparator_factory = Comparator - - -# for reflection -ischema_names['bit'] = BIT # type: ignore + comparator_factory = Comparator # type: ignore From 81b118cbf324f0ace221f1cd5aa85ecaf8cd29d2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 13:06:47 -0700 Subject: [PATCH 713/879] Changed vector and halfvec types to return list for Django, SQLAlchemy, SQLModel, and Peewee --- CHANGELOG.md | 3 +- pgvector/django/halfvec.py | 27 ++++++---------- pgvector/django/vector.py | 27 ++++++---------- pgvector/peewee/halfvec.py | 5 +-- pgvector/peewee/vector.py | 5 +-- pgvector/sqlalchemy/halfvec.py | 5 +-- pgvector/sqlalchemy/vector.py | 5 +-- tests/test_django.py | 32 +++++++++---------- tests/test_peewee.py | 14 ++++----- tests/test_sqlalchemy.py | 56 +++++++++++++++++----------------- tests/test_sqlmodel.py | 22 ++++++------- 11 files changed, 96 insertions(+), 105 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 37bd613..42ff5d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,8 @@ ## 1.0.0 (unreleased) - Added experimental support for type hints -- Changed `vector` type to return `Vector` object instead of NumPy array +- Changed `vector` and `halfvec` types to return list for Django, SQLAlchemy, SQLModel, and Peewee +- Changed `vector` type to return `Vector` object for Psycopg 3, Psycopg 2, asyncpg, and pg8000 - Removed `utils` package (use top-level `pgvector` package instead) - Removed re-exported classes (use top-level `pgvector` package instead) - Removed dependency on NumPy diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index f3d3719..ed158ab 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -24,16 +24,18 @@ def db_type(self, connection: Any) -> str: return 'halfvec' return 'halfvec(%d)' % self.dimensions - def from_db_value(self, value: Any, expression: Any, connection: Any) -> HalfVector | None: - return HalfVector._from_db(value) + def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: + value = HalfVector._from_db(value) + return None if value is None else value.to_list() - def to_python(self, value: Any) -> HalfVector | None: - if value is None or isinstance(value, HalfVector): + def to_python(self, value: Any) -> list[float] | None: + if value is None or isinstance(value, list): return value - elif isinstance(value, str): - return HalfVector.from_text(value) - else: - return HalfVector(value) + if isinstance(value, str): + value = HalfVector.from_text(value) + if not isinstance(value, HalfVector): + value = HalfVector(value) + return value.to_list() def get_prep_value(self, value: Any) -> str | None: return HalfVector._to_db(value) @@ -50,16 +52,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw ) -class HalfVectorWidget(forms.TextInput): - def format_value(self, value: Any) -> Any: - if isinstance(value, HalfVector): - value = value.to_list() - return super().format_value(value) - - class HalfVectorFormField(forms.CharField): - widget = HalfVectorWidget - def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index e5ab444..3529c8f 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -24,16 +24,18 @@ def db_type(self, connection: Any) -> str: return 'vector' return 'vector(%d)' % self.dimensions - def from_db_value(self, value: Any, expression: Any, connection: Any) -> Vector | None: - return Vector._from_db(value) + def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: + value = Vector._from_db(value) + return None if value is None else value.to_list() - def to_python(self, value: Any) -> Vector | None: - if value is None or isinstance(value, Vector): + def to_python(self, value: Any) -> list[float] | None: + if value is None or isinstance(value, list): return value - elif isinstance(value, str): - return Vector.from_text(value) - else: - return Vector(value) + if isinstance(value, str): + value = Vector.from_text(value) + if not isinstance(value, Vector): + value = Vector(value) + return value.to_list() def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) @@ -50,16 +52,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw ) -class VectorWidget(forms.TextInput): - def format_value(self, value: Any) -> Any: - if isinstance(value, Vector): - value = value.to_list() - return super().format_value(value) - - class VectorFormField(forms.CharField): - widget = VectorWidget - def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index cb76976..d718ed4 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -16,8 +16,9 @@ def get_modifiers(self) -> list[int] | None: def db_value(self, value: object) -> str | None: return HalfVector._to_db(value) - def python_value(self, value: Any) -> HalfVector | None: - return HalfVector._from_db(value) + def python_value(self, value: Any) -> list[float] | None: + value = HalfVector._from_db(value) + return None if value is None else value.to_list() def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 18fb75a..6e5f11b 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -16,8 +16,9 @@ def get_modifiers(self) -> list[int] | None: def db_value(self, value: object) -> str | None: return Vector._to_db(value) - def python_value(self, value: Any) -> Vector | None: - return Vector._from_db(value) + def python_value(self, value: Any) -> list[float] | None: + value = Vector._from_db(value) + return None if value is None else value.to_list() def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 5c16bfd..a5b4444 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -31,8 +31,9 @@ def process(value: Any) -> Any: return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: - def process(value: Any) -> HalfVector | None: - return HalfVector._from_db(value) + def process(value: Any) -> list[float] | None: + value = HalfVector._from_db(value) + return None if value is None else value.to_list() return process class Comparator(TypeEngine.Comparator[Any]): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index b049758..385b1ee 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -31,8 +31,9 @@ def process(value: Any) -> Any: return process def result_processor(self, dialect: Dialect, coltype: Any) -> Any: - def process(value: Any) -> Vector | None: - return Vector._from_db(value) + def process(value: Any) -> list[float] | None: + value = Vector._from_db(value) + return None if value is None else value.to_list() return process class Comparator(TypeEngine.Comparator[Any]): diff --git a/tests/test_django.py b/tests/test_django.py index 613d404..03529ae 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -11,7 +11,7 @@ from math import sqrt import os import pgvector.django -from pgvector import HalfVector, SparseVector, Vector +from pgvector import SparseVector from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock @@ -164,11 +164,11 @@ def setup_method(self) -> None: def test_vector(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) # type: ignore - assert item.embedding == Vector([1, 2, 3]) + assert item.embedding == [1, 2, 3] def test_vector_l2_distance(self) -> None: create_items() - distance = L2Distance('embedding', Vector([1, 1, 1])) + distance = L2Distance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] @@ -210,32 +210,32 @@ def test_vector_l1_distance(self) -> None: def test_halfvec(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) # type: ignore - assert item.half_embedding == HalfVector([1, 2, 3]) + assert item.half_embedding == [1, 2, 3] def test_halfvec_l2_distance(self) -> None: create_items() - distance = L2Distance('half_embedding', HalfVector([1, 1, 1])) + distance = L2Distance('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] def test_halfvec_max_inner_product(self) -> None: create_items() - distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1])) + distance = MaxInnerProduct('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_halfvec_cosine_distance(self) -> None: create_items() - distance = CosineDistance('half_embedding', HalfVector([1, 1, 1])) + distance = CosineDistance('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_halfvec_l1_distance(self) -> None: create_items() - distance = L1Distance('half_embedding', HalfVector([1, 1, 1])) + distance = L1Distance('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] @@ -304,7 +304,7 @@ def test_vector_avg(self) -> None: Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] # type: ignore - assert avg == Vector([2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] def test_vector_sum(self) -> None: sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] # type: ignore @@ -312,7 +312,7 @@ def test_vector_sum(self) -> None: Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] # type: ignore - assert sum == Vector([5, 7, 9]) + assert sum == [5, 7, 9] def test_halfvec_avg(self) -> None: avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] # type: ignore @@ -320,7 +320,7 @@ def test_halfvec_avg(self) -> None: Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] # type: ignore - assert avg == HalfVector([2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] def test_halfvec_sum(self) -> None: sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] # type: ignore @@ -328,7 +328,7 @@ def test_halfvec_sum(self) -> None: Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] # type: ignore - assert sum == HalfVector([5, 7, 9]) + assert sum == [5, 7, 9] def test_serialization(self) -> None: create_items() @@ -358,7 +358,7 @@ def test_vector_form_save(self) -> None: assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6]) # type: ignore + assert Item.objects.get(pk=1).embedding == [4, 5, 6] # type: ignore def test_vector_form_save_missing(self) -> None: Item(id=1).save() @@ -386,7 +386,7 @@ def test_halfvec_form_save(self) -> None: assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) # type: ignore + assert Item.objects.get(pk=1).half_embedding == [4, 5, 6] # type: ignore def test_halfvec_form_save_missing(self) -> None: Item(id=1).save() @@ -468,7 +468,7 @@ def test_missing(self) -> None: assert Item.objects.first().sparse_embedding is None # type: ignore def test_vector_array(self) -> None: - Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]).save() + Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]]).save() with connection.cursor() as cursor: from pgvector.psycopg import register_vector @@ -476,7 +476,7 @@ def test_vector_array(self) -> None: # this fails if the driver does not cast arrays item = Item.objects.get(pk=1) # type: ignore - assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] def test_double_array(self) -> None: Item(id=1, double_embedding=[1, 1, 1]).save() diff --git a/tests/test_peewee.py b/tests/test_peewee.py index f1a4758..79358ed 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,6 +1,6 @@ from math import sqrt from peewee import Model, PostgresqlDatabase, fn -from pgvector import HalfVector, SparseVector, Vector +from pgvector import SparseVector, Vector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField from typing import Any @@ -44,7 +44,7 @@ def setup_method(self) -> None: def test_vector(self) -> None: Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.embedding == Vector([1, 2, 3]) + assert item.embedding == [1, 2, 3] def test_vector_l2_distance(self) -> None: create_items() @@ -77,7 +77,7 @@ def test_vector_l1_distance(self) -> None: def test_halfvec(self) -> None: Item.create(id=1, half_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.half_embedding == HalfVector([1, 2, 3]) + assert item.half_embedding == [1, 2, 3] def test_halfvec_l2_distance(self) -> None: create_items() @@ -170,7 +170,7 @@ def test_vector_avg(self) -> None: Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() - assert avg == Vector([2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] def test_vector_sum(self) -> None: sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() @@ -178,7 +178,7 @@ def test_vector_sum(self) -> None: Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() - assert sum == Vector([5, 7, 9]) + assert sum == [5, 7, 9] def test_halfvec_avg(self) -> None: avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() @@ -186,7 +186,7 @@ def test_halfvec_avg(self) -> None: Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() - assert avg == HalfVector([2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] def test_halfvec_sum(self) -> None: sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() @@ -194,7 +194,7 @@ def test_halfvec_sum(self) -> None: Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() - assert sum == HalfVector([5, 7, 9]) + assert sum == [5, 7, 9] def test_get_or_create(self) -> None: Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 6edcda7..9fc653d 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,6 +1,6 @@ import asyncpg from getpass import getuser -from pgvector import HalfVector, SparseVector, Vector +from pgvector import SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine @@ -79,12 +79,12 @@ class Item(Base): __tablename__ = 'sqlalchemy_orm_item' id: Mapped[int] = mapped_column(Integer, primary_key=True) - embedding: Mapped[Vector | None] = mapped_column(VECTOR(3)) - half_embedding: Mapped[HalfVector | None] = mapped_column(HALFVEC(3)) + embedding: Mapped[list[float] | None] = mapped_column(VECTOR(3)) + half_embedding: Mapped[list[float] | None] = mapped_column(HALFVEC(3)) binary_embedding: Mapped[Any | None] = mapped_column(BIT(3)) sparse_embedding: Mapped[SparseVector | None] = mapped_column(SPARSEVEC(3)) - embeddings: Mapped[list[Vector] | None] = mapped_column(ARRAY(VECTOR(3))) - half_embeddings: Mapped[list[HalfVector] | None] = mapped_column(ARRAY(HALFVEC(3))) + embeddings: Mapped[list[list[float]] | None] = mapped_column(ARRAY(VECTOR(3), dimensions=1)) + half_embeddings: Mapped[list[list[float]] | None] = mapped_column(ARRAY(HALFVEC(3), dimensions=1)) Base.metadata.drop_all(setup_engine) @@ -171,7 +171,7 @@ def test_core(self, engine: Engine) -> None: hnsw_index.create(engine) def test_orm(self, engine: Engine) -> None: - item = Item(embedding=Vector([1.5, 2, 3])) + item = Item(embedding=[1.5, 2, 3]) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -188,8 +188,8 @@ def test_orm(self, engine: Engine) -> None: assert items[0].id % 3 == 1 assert items[1].id % 3 == 2 assert items[2].id % 3 == 0 - assert items[0].embedding == Vector([1.5, 2, 3]) - assert items[1].embedding == Vector([4, 5, 6]) + assert items[0].embedding == [1.5, 2, 3] + assert items[1].embedding == [4, 5, 6] assert items[2].embedding is None def test_vector(self, engine: Engine) -> None: @@ -197,7 +197,7 @@ def test_vector(self, engine: Engine) -> None: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert item.embedding == Vector([1, 2, 3]) + assert item.embedding == [1, 2, 3] def test_vector_l2_distance(self, engine: Engine) -> None: create_items() @@ -252,7 +252,7 @@ def test_halfvec(self, engine: Engine) -> None: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert item.half_embedding == HalfVector([1, 2, 3]) + assert item.half_embedding == [1, 2, 3] def test_halfvec_l2_distance(self, engine: Engine) -> None: create_items() @@ -439,7 +439,7 @@ def test_avg(self, engine: Engine) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.query(avg(Item.embedding)).one()[0] - assert res == Vector([2.5, 3.5, 4.5]) + assert res == [2.5, 3.5, 4.5] def test_avg_orm(self, engine: Engine) -> None: with Session(engine) as session: @@ -448,7 +448,7 @@ def test_avg_orm(self, engine: Engine) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.scalars(select(avg(Item.embedding))).one() - assert res == Vector([2.5, 3.5, 4.5]) + assert res == [2.5, 3.5, 4.5] def test_sum(self, engine: Engine) -> None: with Session(engine) as session: @@ -457,7 +457,7 @@ def test_sum(self, engine: Engine) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.query(sum(Item.embedding)).one()[0] - assert res == Vector([5, 7, 9]) + assert res == [5, 7, 9] def test_sum_orm(self, engine: Engine) -> None: with Session(engine) as session: @@ -466,7 +466,7 @@ def test_sum_orm(self, engine: Engine) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.scalars(select(sum(Item.embedding))).one() - assert res == Vector([5, 7, 9]) + assert res == [5, 7, 9] def test_bad_dimensions(self, engine: Engine) -> None: item = Item(embedding=[1, 2]) @@ -501,11 +501,11 @@ def test_literal_binds(self, engine: Engine) -> None: def test_insert(self, engine: Engine) -> None: with Session(engine) as session: - session.execute(insert(Item).values(embedding=Vector([1, 2, 3]))) + session.execute(insert(Item).values(embedding=[1, 2, 3])) def test_insert_bulk(self, engine: Engine) -> None: with Session(engine) as session: - session.execute(insert(Item), [{'embedding': Vector([1, 2, 3])}]) + session.execute(insert(Item), [{'embedding': [1, 2, 3]}]) # register_vector in psycopg2 tests change this behavior # def test_insert_text(self): @@ -519,9 +519,9 @@ def test_automap(self, engine: Engine) -> None: AutoBase.prepare() AutoItem = AutoBase.classes.sqlalchemy_orm_item with Session(engine) as session: - session.execute(insert(AutoItem), [{'embedding': Vector([1, 2, 3])}]) + session.execute(insert(AutoItem), [{'embedding': [1, 2, 3]}]) item = session.query(AutoItem).first() - assert item is not None and item.embedding == Vector([1, 2, 3]) + assert item is not None and item.embedding == [1, 2, 3] def test_half_precision(self, engine: Engine) -> None: create_items() @@ -564,21 +564,21 @@ def setup_method(self) -> None: def test_vector_array(self, engine: Engine) -> None: with Session(engine) as session: - session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) + session.add(Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]])) session.commit() # this fails if the driver does not cast arrays item = session.get_one(Item, 1) - assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] def test_halfvec_array(self, engine: Engine) -> None: with Session(engine) as session: - session.add(Item(id=1, half_embeddings=[HalfVector([1, 2, 3]), HalfVector([4, 5, 6])])) + session.add(Item(id=1, half_embeddings=[[1, 2, 3], [4, 5, 6]])) session.commit() # this fails if the driver does not cast arrays item = session.get_one(Item, 1) - assert item.half_embeddings == [HalfVector([1, 2, 3]), HalfVector([4, 5, 6])] + assert item.half_embeddings == [[1, 2, 3], [4, 5, 6]] @pytest.mark.parametrize('engine', async_engines) @@ -592,7 +592,7 @@ async def test_vector(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): - embedding = Vector([1, 2, 3]) + embedding = [1, 2, 3] session.add(Item(id=1, embedding=embedding)) item = await session.get_one(Item, 1) assert item.embedding == embedding @@ -605,7 +605,7 @@ async def test_halfvec(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): - embedding = HalfVector([1, 2, 3]) + embedding = [1, 2, 3] session.add(Item(id=1, half_embedding=embedding)) item = await session.get_one(Item, 1) assert item.half_embedding == embedding @@ -652,7 +652,7 @@ async def test_avg(self, engine: AsyncEngine) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = await session.scalars(select(avg(Item.embedding))) - assert res.first() == Vector([2.5, 3.5, 4.5]) + assert res.first() == [2.5, 3.5, 4.5] await engine.dispose() @@ -668,13 +668,13 @@ async def test_vector_array(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): - session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])])) + session.add(Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]])) item = await session.get_one(Item, 1) - assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] if NUMPY_AVAILABLE: session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) item = await session.get_one(Item, 2) - assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])] + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index b6f4510..bac3aee 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,4 +1,4 @@ -from pgvector import HalfVector, SparseVector, Vector +from pgvector import SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum from pydantic import ConfigDict import pytest @@ -15,8 +15,8 @@ class Item(SQLModel, table=True): model_config = ConfigDict(arbitrary_types_allowed=True) # type: ignore id: int | None = Field(default=None, primary_key=True) - embedding: Vector | list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore - half_embedding: HalfVector | list[float] | None = Field(default=None, sa_type=HALFVEC(3)) # type: ignore + embedding: list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore + half_embedding: list[float] | None = Field(default=None, sa_type=HALFVEC(3)) # type: ignore binary_embedding: str | None = Field(default=None, sa_type=BIT(3)) # type: ignore sparse_embedding: SparseVector | list[float] | None = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore @@ -65,8 +65,8 @@ def test_orm(self) -> None: assert items[0].id == 1 assert items[1].id == 2 assert items[2].id == 3 - assert items[0].embedding == Vector([1.5, 2, 3]) - assert items[1].embedding == Vector([4, 5, 6]) + assert items[0].embedding == [1.5, 2, 3] + assert items[1].embedding == [4, 5, 6] assert items[2].embedding is None def test_vector(self) -> None: @@ -74,7 +74,7 @@ def test_vector(self) -> None: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert item.embedding == Vector([1, 2, 3]) + assert item.embedding == [1, 2, 3] def test_vector_l2_distance(self) -> None: create_items() @@ -105,7 +105,7 @@ def test_halfvec(self) -> None: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() item = session.get_one(Item, 1) - assert item.half_embedding == HalfVector([1, 2, 3]) + assert item.half_embedding == [1, 2, 3] def test_halfvec_l2_distance(self) -> None: create_items() @@ -200,7 +200,7 @@ def test_vector_avg(self) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(avg(Item.embedding))).first() - assert res == Vector([2.5, 3.5, 4.5]) + assert res == [2.5, 3.5, 4.5] def test_vector_sum(self) -> None: with Session(engine) as session: @@ -209,7 +209,7 @@ def test_vector_sum(self) -> None: session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(sum(Item.embedding))).first() - assert res == Vector([5, 7, 9]) + assert res == [5, 7, 9] def test_halfvec_avg(self) -> None: with Session(engine) as session: @@ -218,7 +218,7 @@ def test_halfvec_avg(self) -> None: session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(avg(Item.half_embedding))).first() - assert res == HalfVector([2.5, 3.5, 4.5]) + assert res == [2.5, 3.5, 4.5] def test_halfvec_sum(self) -> None: with Session(engine) as session: @@ -227,7 +227,7 @@ def test_halfvec_sum(self) -> None: session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(sum(Item.half_embedding))).first() - assert res == HalfVector([5, 7, 9]) + assert res == [5, 7, 9] def test_bad_dimensions(self) -> None: item = Item(embedding=[1, 2]) From dc604e46a46bfea2e6a8afb3291e9136e2d32939 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 13:21:01 -0700 Subject: [PATCH 714/879] Added fast paths for higher-level libraries [skip ci] --- pgvector/halfvec.py | 4 ++++ pgvector/vector.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 95e5254..2b77873 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -76,6 +76,10 @@ def _to_db(cls, value: object) -> str | None: if value is None: return value + # fast path for high-level libraries + if isinstance(value, list): + return f'[{",".join([str(float(v)) for v in value])}]' + if not isinstance(value, cls): value = cls(value) # type: ignore diff --git a/pgvector/vector.py b/pgvector/vector.py index a4b6041..d06dea1 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -76,6 +76,10 @@ def _to_db(cls, value: object) -> str | None: if value is None: return value + # fast path for high-level libraries + if isinstance(value, list): + return f'[{",".join([str(float(v)) for v in value])}]' + if not isinstance(value, cls): value = cls(value) # type: ignore From c39ddd9c0f2f274c2dcafb8073fe228042732e82 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 13:26:43 -0700 Subject: [PATCH 715/879] Fixed typechecking [skip ci] --- pgvector/django/halfvec.py | 4 ++-- pgvector/django/vector.py | 4 ++-- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index ed158ab..87c76f0 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -33,9 +33,9 @@ def to_python(self, value: Any) -> list[float] | None: return value if isinstance(value, str): value = HalfVector.from_text(value) - if not isinstance(value, HalfVector): + elif not isinstance(value, HalfVector): value = HalfVector(value) - return value.to_list() + return value.to_list() # type: ignore def get_prep_value(self, value: Any) -> str | None: return HalfVector._to_db(value) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 3529c8f..dbe6a1f 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -33,9 +33,9 @@ def to_python(self, value: Any) -> list[float] | None: return value if isinstance(value, str): value = Vector.from_text(value) - if not isinstance(value, Vector): + elif not isinstance(value, Vector): value = Vector(value) - return value.to_list() + return value.to_list() # type: ignore def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 2b77873..16f67a3 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -78,7 +78,7 @@ def _to_db(cls, value: object) -> str | None: # fast path for high-level libraries if isinstance(value, list): - return f'[{",".join([str(float(v)) for v in value])}]' + return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore if not isinstance(value, cls): value = cls(value) # type: ignore diff --git a/pgvector/vector.py b/pgvector/vector.py index d06dea1..f3df071 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -78,7 +78,7 @@ def _to_db(cls, value: object) -> str | None: # fast path for high-level libraries if isinstance(value, list): - return f'[{",".join([str(float(v)) for v in value])}]' + return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore if not isinstance(value, cls): value = cls(value) # type: ignore From 26deadd7a6160308e6876e61a943a83b9264751a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 13:46:33 -0700 Subject: [PATCH 716/879] Improved code [skip ci] --- pgvector/django/halfvec.py | 3 +-- pgvector/django/vector.py | 3 +-- pgvector/halfvec.py | 8 ++++---- pgvector/peewee/halfvec.py | 3 +-- pgvector/peewee/vector.py | 3 +-- pgvector/sqlalchemy/halfvec.py | 3 +-- pgvector/sqlalchemy/vector.py | 3 +-- pgvector/vector.py | 8 ++++---- 8 files changed, 14 insertions(+), 20 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 87c76f0..d501dbb 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -25,8 +25,7 @@ def db_type(self, connection: Any) -> str: return 'halfvec(%d)' % self.dimensions def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: - value = HalfVector._from_db(value) - return None if value is None else value.to_list() + return None if value is None else HalfVector._from_db(value) def to_python(self, value: Any) -> list[float] | None: if value is None or isinstance(value, list): diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index dbe6a1f..cfee0ca 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -25,8 +25,7 @@ def db_type(self, connection: Any) -> str: return 'vector(%d)' % self.dimensions def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: - value = Vector._from_db(value) - return None if value is None else value.to_list() + return None if value is None else Vector._from_db(value) def to_python(self, value: Any) -> list[float] | None: if value is None or isinstance(value, list): diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 16f67a3..e38eccf 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -96,8 +96,8 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value.to_binary() @classmethod - def _from_db(cls, value: str | HalfVector | None) -> HalfVector | None: - if value is None or isinstance(value, HalfVector): - return value + def _from_db(cls, value: str | HalfVector) -> list[float]: + if isinstance(value, str): + value = cls.from_text(value) - return cls.from_text(value) + return value.to_list() diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index d718ed4..5fc2b0b 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -17,8 +17,7 @@ def db_value(self, value: object) -> str | None: return HalfVector._to_db(value) def python_value(self, value: Any) -> list[float] | None: - value = HalfVector._from_db(value) - return None if value is None else value.to_list() + return None if value is None else HalfVector._from_db(value) def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 6e5f11b..5419b22 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -17,8 +17,7 @@ def db_value(self, value: object) -> str | None: return Vector._to_db(value) def python_value(self, value: Any) -> list[float] | None: - value = Vector._from_db(value) - return None if value is None else value.to_list() + return None if value is None else Vector._from_db(value) def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index a5b4444..a4529f6 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -32,8 +32,7 @@ def process(value: Any) -> Any: def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value: Any) -> list[float] | None: - value = HalfVector._from_db(value) - return None if value is None else value.to_list() + return None if value is None else HalfVector._from_db(value) return process class Comparator(TypeEngine.Comparator[Any]): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 385b1ee..e6e79b5 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -32,8 +32,7 @@ def process(value: Any) -> Any: def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value: Any) -> list[float] | None: - value = Vector._from_db(value) - return None if value is None else value.to_list() + return None if value is None else Vector._from_db(value) return process class Comparator(TypeEngine.Comparator[Any]): diff --git a/pgvector/vector.py b/pgvector/vector.py index f3df071..d767f27 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -96,8 +96,8 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value.to_binary() @classmethod - def _from_db(cls, value: str | Vector | None) -> Vector | None: - if value is None or isinstance(value, Vector): - return value + def _from_db(cls, value: str | Vector) -> list[float]: + if isinstance(value, str): + value = cls.from_text(value) - return cls.from_text(value) + return value.to_list() From b07fdb17787117d4de2f029aa5e4599ecfc37b64 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 13:50:00 -0700 Subject: [PATCH 717/879] Improved code [skip ci] --- pgvector/django/halfvec.py | 2 +- pgvector/django/vector.py | 2 +- pgvector/halfvec.py | 5 ++++- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/vector.py | 2 +- pgvector/sqlalchemy/halfvec.py | 2 +- pgvector/sqlalchemy/vector.py | 2 +- pgvector/vector.py | 5 ++++- 8 files changed, 14 insertions(+), 8 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index d501dbb..98cc013 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -25,7 +25,7 @@ def db_type(self, connection: Any) -> str: return 'halfvec(%d)' % self.dimensions def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: - return None if value is None else HalfVector._from_db(value) + return HalfVector._from_db(value) def to_python(self, value: Any) -> list[float] | None: if value is None or isinstance(value, list): diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index cfee0ca..c8d37fb 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -25,7 +25,7 @@ def db_type(self, connection: Any) -> str: return 'vector(%d)' % self.dimensions def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: - return None if value is None else Vector._from_db(value) + return Vector._from_db(value) def to_python(self, value: Any) -> list[float] | None: if value is None or isinstance(value, list): diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index e38eccf..8a4dea0 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -96,7 +96,10 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value.to_binary() @classmethod - def _from_db(cls, value: str | HalfVector) -> list[float]: + def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: + if value is None: + return value + if isinstance(value, str): value = cls.from_text(value) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 5fc2b0b..72ef759 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -17,7 +17,7 @@ def db_value(self, value: object) -> str | None: return HalfVector._to_db(value) def python_value(self, value: Any) -> list[float] | None: - return None if value is None else HalfVector._from_db(value) + return HalfVector._from_db(value) def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 5419b22..d6864bc 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -17,7 +17,7 @@ def db_value(self, value: object) -> str | None: return Vector._to_db(value) def python_value(self, value: Any) -> list[float] | None: - return None if value is None else Vector._from_db(value) + return Vector._from_db(value) def _distance(self, op: str, vector: object) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index a4529f6..f0759e5 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -32,7 +32,7 @@ def process(value: Any) -> Any: def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value: Any) -> list[float] | None: - return None if value is None else HalfVector._from_db(value) + return HalfVector._from_db(value) return process class Comparator(TypeEngine.Comparator[Any]): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index e6e79b5..199afd7 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -32,7 +32,7 @@ def process(value: Any) -> Any: def result_processor(self, dialect: Dialect, coltype: Any) -> Any: def process(value: Any) -> list[float] | None: - return None if value is None else Vector._from_db(value) + return Vector._from_db(value) return process class Comparator(TypeEngine.Comparator[Any]): diff --git a/pgvector/vector.py b/pgvector/vector.py index d767f27..fc1a6b6 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -96,7 +96,10 @@ def _to_db_binary(cls, value: object) -> bytes | None: return value.to_binary() @classmethod - def _from_db(cls, value: str | Vector) -> list[float]: + def _from_db(cls, value: str | Vector | None) -> list[float] | None: + if value is None: + return value + if isinstance(value, str): value = cls.from_text(value) From 0b8fbd7b297803fb4f087d9ca3111ae072b456d1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 14:12:06 -0700 Subject: [PATCH 718/879] Switched to list for internal storage format for Vector and HalfVector [skip ci] --- pgvector/halfvec.py | 34 +++++++++++++++------------------- pgvector/vector.py | 34 +++++++++++++++------------------- tests/test_half_vector.py | 12 ++++-------- tests/test_vector.py | 12 ++++-------- 4 files changed, 38 insertions(+), 54 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 8a4dea0..bef3acd 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,6 +1,5 @@ from __future__ import annotations import struct -from typing import cast try: import numpy as np @@ -12,50 +11,47 @@ class HalfVector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: if isinstance(value, list): - dim = len(value) try: - self._value = struct.pack(f'>HH{dim}e', dim, 0, *value) - except struct.error: + self._value = [float(v) for v in value] + except (TypeError, ValueError): raise ValueError('expected list[float]') elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') - # asarray still copies if same dtype - if value.dtype != '>f2': - value = np.asarray(value, dtype='>f2') - - self._value = struct.pack('>HH', value.shape[0], 0) + value.tobytes() + self._value = [float(v) for v in value] else: raise ValueError('expected list or ndarray') def __repr__(self) -> str: - return f'HalfVector({self.to_list()})' + return f'HalfVector({self._value})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return self.to_binary() == other.to_binary() + return self._value == other._value return False def dimensions(self) -> int: - dim, = cast(tuple[int], struct.unpack_from('>H', self._value)) - return dim + return len(self._value) def to_list(self) -> list[float]: - return list(struct.unpack_from(f'>{self.dimensions()}e', self._value[4:])) + return self._value def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: - return np.frombuffer(self._value, dtype='>f2', count=self.dimensions(), offset=4) + return np.array(self._value, dtype=np.float16) def to_text(self) -> str: - return f'[{",".join([str(v) for v in self.to_list()])}]' + return f'[{",".join([str(v) for v in self._value])}]' def to_binary(self) -> bytes: - return self._value + dim = len(self._value) + return struct.pack(f'>HH{dim}e', dim, 0, *self._value) @classmethod def from_text(cls, value: str) -> HalfVector: - return cls([float(v) for v in value[1:-1].split(',')]) + vec = cls.__new__(cls) + vec._value = [float(v) for v in value[1:-1].split(',')] + return vec @classmethod def from_binary(cls, value: bytes) -> HalfVector: @@ -68,7 +64,7 @@ def from_binary(cls, value: bytes) -> HalfVector: raise ValueError('expected unused to be 0') vec = cls.__new__(cls) - vec._value = value + vec._value = list(struct.unpack_from(f'>{dim}e', value[4:])) return vec @classmethod diff --git a/pgvector/vector.py b/pgvector/vector.py index fc1a6b6..0550eb3 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,6 +1,5 @@ from __future__ import annotations import struct -from typing import cast try: import numpy as np @@ -12,50 +11,47 @@ class Vector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: if isinstance(value, list): - dim = len(value) try: - self._value = struct.pack(f'>HH{dim}f', dim, 0, *value) - except struct.error: + self._value = [float(v) for v in value] + except (TypeError, ValueError): raise ValueError('expected list[float]') elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') - # asarray still copies if same dtype - if value.dtype != '>f4': - value = np.asarray(value, dtype='>f4') - - self._value = struct.pack('>HH', value.shape[0], 0) + value.tobytes() + self._value = [float(v) for v in value] else: raise ValueError('expected list or ndarray') def __repr__(self) -> str: - return f'Vector({self.to_list()})' + return f'Vector({self._value})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return self.to_binary() == other.to_binary() + return self._value == other._value return False def dimensions(self) -> int: - dim, = cast(tuple[int], struct.unpack_from('>H', self._value)) - return dim + return len(self._value) def to_list(self) -> list[float]: - return list(struct.unpack_from(f'>{self.dimensions()}f', self._value[4:])) + return self._value def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: - return np.frombuffer(self._value, dtype='>f4', count=self.dimensions(), offset=4) + return np.array(self._value, dtype=np.float32) def to_text(self) -> str: - return f'[{",".join([str(v) for v in self.to_list()])}]' + return f'[{",".join([str(v) for v in self._value])}]' def to_binary(self) -> bytes: - return self._value + dim = len(self._value) + return struct.pack(f'>HH{dim}f', dim, 0, *self._value) @classmethod def from_text(cls, value: str) -> Vector: - return cls([float(v) for v in value[1:-1].split(',')]) + vec = cls.__new__(cls) + vec._value = [float(v) for v in value[1:-1].split(',')] + return vec @classmethod def from_binary(cls, value: bytes) -> Vector: @@ -68,7 +64,7 @@ def from_binary(cls, value: bytes) -> Vector: raise ValueError('expected unused to be 0') vec = cls.__new__(cls) - vec._value = value + vec._value = list(struct.unpack_from(f'>{dim}f', value[4:])) return vec @classmethod diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 6bd819c..b9526d4 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -11,7 +11,9 @@ class TestHalfVector: def test_list(self) -> None: - assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] + arr = [1.0, 2.0, 3.0] + assert HalfVector(arr).to_list() == arr + assert HalfVector(arr).to_list() is not arr def test_list_empty(self) -> None: assert HalfVector([]).to_list() == [] @@ -31,6 +33,7 @@ def test_ndarray(self) -> None: arr = np.array([1, 2, 3]) assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr + assert HalfVector(arr).to_numpy().dtype == np.float16 def test_int(self) -> None: with pytest.raises(ValueError) as error: @@ -48,13 +51,6 @@ def test_equality(self) -> None: def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_to_numpy_readonly(self) -> None: - arr = HalfVector([1, 2, 3]).to_numpy() - with pytest.raises(ValueError) as error: - arr[0] = 4 - assert str(error.value) == 'assignment destination is read-only' - def test_from_text(self) -> None: vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] diff --git a/tests/test_vector.py b/tests/test_vector.py index db23648..499d9a2 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -11,7 +11,9 @@ class TestVector: def test_list(self) -> None: - assert Vector([1, 2, 3]).to_list() == [1, 2, 3] + arr = [1.0, 2.0, 3.0] + assert Vector(arr).to_list() == arr + assert Vector(arr).to_list() is not arr def test_list_empty(self) -> None: assert Vector([]).to_list() == [] @@ -31,6 +33,7 @@ def test_ndarray(self) -> None: arr = np.array([1, 2, 3]) assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr + assert Vector(arr).to_numpy().dtype == np.float32 def test_int(self) -> None: with pytest.raises(ValueError) as error: @@ -48,13 +51,6 @@ def test_equality(self) -> None: def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') - def test_to_numpy_readonly(self) -> None: - arr = Vector([1, 2, 3]).to_numpy() - with pytest.raises(ValueError) as error: - arr[0] = 4 - assert str(error.value) == 'assignment destination is read-only' - def test_from_text(self) -> None: vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] From e771e695600de3a2770d8d58746dbc622b6ff975 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 14:34:22 -0700 Subject: [PATCH 719/879] Restored fast path for NumPy --- pgvector/psycopg/vector.py | 4 +--- pgvector/vector.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index b8b3611..d70a07c 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -28,9 +28,7 @@ class VectorBinaryDumper(VectorDumper): format = Format.BINARY def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: - if not isinstance(obj, Vector): - obj = Vector(obj) - return obj.to_binary() + return Vector._to_db_binary(obj) class VectorLoader(Loader): diff --git a/pgvector/vector.py b/pgvector/vector.py index 0550eb3..e08d3a0 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -86,6 +86,16 @@ def _to_db_binary(cls, value: object) -> bytes | None: if value is None: return value + # fast path for NumPy + if NUMPY_AVAILABLE and isinstance(value, np.ndarray): + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + if value.dtype != '>f4': + value = np.asarray(value, dtype='>f4') + + return struct.pack('>HH', len(value), 0) + value.tobytes() + if not isinstance(value, cls): value = cls(value) # type: ignore From 81c717b1024454cd7fc1e92efcd5ece31952a19f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 15:23:19 -0700 Subject: [PATCH 720/879] Switched to array for internal storage format for Vector and HalfVector [skip ci] --- pgvector/halfvec.py | 37 ++++++++++++++++++++++++------------- pgvector/vector.py | 43 +++++++++++++++++++++---------------------- 2 files changed, 45 insertions(+), 35 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index bef3acd..d11e795 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,5 +1,7 @@ from __future__ import annotations +import array import struct +import sys try: import numpy as np @@ -11,20 +13,24 @@ class HalfVector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: if isinstance(value, list): + dim = len(value) try: - self._value = [float(v) for v in value] - except (TypeError, ValueError): + self._value = array.array('H', struct.pack(f'{dim}e', *value)) + except struct.error: raise ValueError('expected list[float]') elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') - self._value = [float(v) for v in value] + if value.dtype != np.float16: + value = np.asarray(value, dtype=np.float16) + + self._value = array.array('H', value.tobytes()) else: raise ValueError('expected list or ndarray') def __repr__(self) -> str: - return f'HalfVector({self._value})' + return f'HalfVector({self.to_list()})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): @@ -35,23 +41,26 @@ def dimensions(self) -> int: return len(self._value) def to_list(self) -> list[float]: - return self._value + dim = len(self._value) + return list(struct.unpack(f'{dim}e', self._value.tobytes())) def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: - return np.array(self._value, dtype=np.float16) + return np.frombuffer(self._value, dtype=np.float16) def to_text(self) -> str: - return f'[{",".join([str(v) for v in self._value])}]' + return f'[{",".join([str(v) for v in self.to_list()])}]' def to_binary(self) -> bytes: - dim = len(self._value) - return struct.pack(f'>HH{dim}e', dim, 0, *self._value) + if sys.byteorder == 'big': + value = self._value + else: + value = array.array('H', self._value) + value.byteswap() + return struct.pack(f'>HH', len(value), 0) + value.tobytes() @classmethod def from_text(cls, value: str) -> HalfVector: - vec = cls.__new__(cls) - vec._value = [float(v) for v in value[1:-1].split(',')] - return vec + return cls([float(v) for v in value[1:-1].split(',')]) @classmethod def from_binary(cls, value: bytes) -> HalfVector: @@ -64,7 +73,9 @@ def from_binary(cls, value: bytes) -> HalfVector: raise ValueError('expected unused to be 0') vec = cls.__new__(cls) - vec._value = list(struct.unpack_from(f'>{dim}e', value[4:])) + vec._value = array.array('H', value[4:]) + if sys.byteorder != 'big': + vec._value.byteswap() return vec @classmethod diff --git a/pgvector/vector.py b/pgvector/vector.py index e08d3a0..60b8157 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,5 +1,7 @@ from __future__ import annotations +import array import struct +import sys try: import numpy as np @@ -12,19 +14,22 @@ class Vector: def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: if isinstance(value, list): try: - self._value = [float(v) for v in value] - except (TypeError, ValueError): + self._value = array.array('f', value) + except TypeError: raise ValueError('expected list[float]') elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): if value.ndim != 1: raise ValueError('expected ndim to be 1') - self._value = [float(v) for v in value] + if value.dtype != np.float32: + value = np.asarray(value, dtype=np.float32) + + self._value = array.array('f', value.tobytes()) else: raise ValueError('expected list or ndarray') def __repr__(self) -> str: - return f'Vector({self._value})' + return f'Vector({self.to_list()})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): @@ -35,23 +40,25 @@ def dimensions(self) -> int: return len(self._value) def to_list(self) -> list[float]: - return self._value + return self._value.tolist() def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: - return np.array(self._value, dtype=np.float32) + return np.frombuffer(self._value, dtype=np.float32) def to_text(self) -> str: return f'[{",".join([str(v) for v in self._value])}]' def to_binary(self) -> bytes: - dim = len(self._value) - return struct.pack(f'>HH{dim}f', dim, 0, *self._value) + if sys.byteorder == 'big': + value = self._value + else: + value = array.array('f', self._value) + value.byteswap() + return struct.pack(f'>HH', len(value), 0) + value.tobytes() @classmethod def from_text(cls, value: str) -> Vector: - vec = cls.__new__(cls) - vec._value = [float(v) for v in value[1:-1].split(',')] - return vec + return cls([float(v) for v in value[1:-1].split(',')]) @classmethod def from_binary(cls, value: bytes) -> Vector: @@ -64,7 +71,9 @@ def from_binary(cls, value: bytes) -> Vector: raise ValueError('expected unused to be 0') vec = cls.__new__(cls) - vec._value = list(struct.unpack_from(f'>{dim}f', value[4:])) + vec._value = array.array('f', value[4:]) + if sys.byteorder != 'big': + vec._value.byteswap() return vec @classmethod @@ -86,16 +95,6 @@ def _to_db_binary(cls, value: object) -> bytes | None: if value is None: return value - # fast path for NumPy - if NUMPY_AVAILABLE and isinstance(value, np.ndarray): - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - if value.dtype != '>f4': - value = np.asarray(value, dtype='>f4') - - return struct.pack('>HH', len(value), 0) + value.tobytes() - if not isinstance(value, cls): value = cls(value) # type: ignore From 88bef310ba264071873fc820e14e4855db5c1036 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 15:39:58 -0700 Subject: [PATCH 721/879] Added tests for non-contiguous arrays [skip ci] --- examples/loading/example.py | 3 ++- examples/loading/pyproject.toml | 3 +++ tests/test_half_vector.py | 3 +++ tests/test_vector.py | 3 +++ 4 files changed, 11 insertions(+), 1 deletion(-) diff --git a/examples/loading/example.py b/examples/loading/example.py index 7f3dce8..4c0e579 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -1,4 +1,5 @@ import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -25,7 +26,7 @@ copy.set_types(['vector']) for i, embedding in enumerate(embeddings): - copy.write_row([embedding]) + copy.write_row([Vector(embedding)]) # show progress if i % 10000 == 0: diff --git a/examples/loading/pyproject.toml b/examples/loading/pyproject.toml index 3fb37d4..7247ec0 100644 --- a/examples/loading/pyproject.toml +++ b/examples/loading/pyproject.toml @@ -9,3 +9,6 @@ dev = [ "pgvector", "psycopg[binary]" ] + +[tool.uv.sources] +pgvector = { path = "../..", editable = true } diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index b9526d4..3a9ae99 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -34,6 +34,9 @@ def test_ndarray(self) -> None: assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr assert HalfVector(arr).to_numpy().dtype == np.float16 + # non-contiguous + assert HalfVector(np.flip(arr)).to_list() == [3, 2, 1] + assert HalfVector(np.flip(arr)).to_binary() == HalfVector([3, 2, 1]).to_binary() def test_int(self) -> None: with pytest.raises(ValueError) as error: diff --git a/tests/test_vector.py b/tests/test_vector.py index 499d9a2..75443a9 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -34,6 +34,9 @@ def test_ndarray(self) -> None: assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr assert Vector(arr).to_numpy().dtype == np.float32 + # non-contiguous + assert Vector(np.flip(arr)).to_list() == [3, 2, 1] + assert Vector(np.flip(arr)).to_binary() == Vector([3, 2, 1]).to_binary() def test_int(self) -> None: with pytest.raises(ValueError) as error: From 0b4a61c6ebc0909433a892e258458df4c427e527 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 15:45:20 -0700 Subject: [PATCH 722/879] Added tests for endianness [skip ci] --- tests/test_half_vector.py | 5 ++++- tests/test_vector.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 3a9ae99..7235797 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -30,13 +30,16 @@ def test_list_list(self) -> None: @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self) -> None: - arr = np.array([1, 2, 3]) + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr assert HalfVector(arr).to_numpy().dtype == np.float16 # non-contiguous assert HalfVector(np.flip(arr)).to_list() == [3, 2, 1] assert HalfVector(np.flip(arr)).to_binary() == HalfVector([3, 2, 1]).to_binary() + # big endian + assert HalfVector(arr.astype('>f2')).to_list() == [1, 2, 3] + assert HalfVector(arr.astype('>f2')).to_binary() == HalfVector([1, 2, 3]).to_binary() def test_int(self) -> None: with pytest.raises(ValueError) as error: diff --git a/tests/test_vector.py b/tests/test_vector.py index 75443a9..f4f35e8 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -30,13 +30,16 @@ def test_list_list(self) -> None: @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self) -> None: - arr = np.array([1, 2, 3]) + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr assert Vector(arr).to_numpy().dtype == np.float32 # non-contiguous assert Vector(np.flip(arr)).to_list() == [3, 2, 1] assert Vector(np.flip(arr)).to_binary() == Vector([3, 2, 1]).to_binary() + # big endian + assert Vector(arr.astype('>f4')).to_list() == [1, 2, 3] + assert Vector(arr.astype('>f4')).to_binary() == Vector([1, 2, 3]).to_binary() def test_int(self) -> None: with pytest.raises(ValueError) as error: From fdd559aeebb7335ab3e19fe21493e750025e1955 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 15:49:16 -0700 Subject: [PATCH 723/879] Simplified code [skip ci] --- pgvector/django/halfvec.py | 8 ++------ pgvector/django/vector.py | 8 ++------ 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 98cc013..88e0512 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -28,13 +28,9 @@ def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[fl return HalfVector._from_db(value) def to_python(self, value: Any) -> list[float] | None: - if value is None or isinstance(value, list): + if isinstance(value, list): return value - if isinstance(value, str): - value = HalfVector.from_text(value) - elif not isinstance(value, HalfVector): - value = HalfVector(value) - return value.to_list() # type: ignore + return HalfVector._from_db(value) def get_prep_value(self, value: Any) -> str | None: return HalfVector._to_db(value) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index c8d37fb..331999d 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -28,13 +28,9 @@ def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[fl return Vector._from_db(value) def to_python(self, value: Any) -> list[float] | None: - if value is None or isinstance(value, list): + if isinstance(value, list): return value - if isinstance(value, str): - value = Vector.from_text(value) - elif not isinstance(value, Vector): - value = Vector(value) - return value.to_list() # type: ignore + return Vector._from_db(value) def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) From 326cff5e9bf23f7e3248ae404aeecf470c117e91 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:04:37 -0700 Subject: [PATCH 724/879] Improved code [skip ci] --- pgvector/sparsevec.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 22a8bff..b08a064 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -136,9 +136,9 @@ def from_binary(cls, value: bytes) -> SparseVector: if unused != 0: raise ValueError('expected unused to be 0') - indices = unpack_from(f'>{nnz}i', value, 12) - values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return cls._from_parts(int(dim), list(indices), list(values)) + indices = list(unpack_from(f'>{nnz}i', value, 12)) + values = list(unpack_from(f'>{nnz}f', value, 12 + nnz * 4)) + return cls._from_parts(dim, indices, values) @classmethod def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> SparseVector: From 6c1238cc43ed862ed2eba229f02c00d0685bf9a3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:17:13 -0700 Subject: [PATCH 725/879] Added comments [skip ci] --- pgvector/halfvec.py | 1 + pgvector/vector.py | 1 + 2 files changed, 2 insertions(+) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index d11e795..60fcac9 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -25,6 +25,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float if value.dtype != np.float16: value = np.asarray(value, dtype=np.float16) + # tobytes() important for performance self._value = array.array('H', value.tobytes()) else: raise ValueError('expected list or ndarray') diff --git a/pgvector/vector.py b/pgvector/vector.py index 60b8157..6a9d505 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -24,6 +24,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float if value.dtype != np.float32: value = np.asarray(value, dtype=np.float32) + # tobytes() important for performance self._value = array.array('f', value.tobytes()) else: raise ValueError('expected list or ndarray') From 33293bfd7e15f3f53eed0aa92a9f7509cd63a4be Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:20:29 -0700 Subject: [PATCH 726/879] Added more tests for NumPy arrays [skip ci] --- tests/test_half_vector.py | 4 ++++ tests/test_vector.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 7235797..29143d1 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -41,6 +41,10 @@ def test_ndarray(self) -> None: assert HalfVector(arr.astype('>f2')).to_list() == [1, 2, 3] assert HalfVector(arr.astype('>f2')).to_binary() == HalfVector([1, 2, 3]).to_binary() + with pytest.raises(ValueError) as error: + HalfVector(np.array(['one', 'two', 'three'])) + assert 'could not convert string to float' in str(error.value) + def test_int(self) -> None: with pytest.raises(ValueError) as error: HalfVector(1) # type: ignore diff --git a/tests/test_vector.py b/tests/test_vector.py index f4f35e8..c9a2184 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -41,6 +41,10 @@ def test_ndarray(self) -> None: assert Vector(arr.astype('>f4')).to_list() == [1, 2, 3] assert Vector(arr.astype('>f4')).to_binary() == Vector([1, 2, 3]).to_binary() + with pytest.raises(ValueError) as error: + Vector(np.array(['one', 'two', 'three'])) + assert 'could not convert string to float' in str(error.value) + def test_int(self) -> None: with pytest.raises(ValueError) as error: Vector(1) # type: ignore From 433f167601ece1c8d174b0f0f2195203d48ae430 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:42:05 -0700 Subject: [PATCH 727/879] Improved typing [skip ci] --- tests/test_sqlalchemy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 9fc653d..475ce3e 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -81,7 +81,7 @@ class Item(Base): id: Mapped[int] = mapped_column(Integer, primary_key=True) embedding: Mapped[list[float] | None] = mapped_column(VECTOR(3)) half_embedding: Mapped[list[float] | None] = mapped_column(HALFVEC(3)) - binary_embedding: Mapped[Any | None] = mapped_column(BIT(3)) + binary_embedding: Mapped[str | None] = mapped_column(BIT(3)) sparse_embedding: Mapped[SparseVector | None] = mapped_column(SPARSEVEC(3)) embeddings: Mapped[list[list[float]] | None] = mapped_column(ARRAY(VECTOR(3), dimensions=1)) half_embeddings: Mapped[list[list[float]] | None] = mapped_column(ARRAY(HALFVEC(3), dimensions=1)) From f2ecf3c29d0fd636a1821f10816dd683d89bad8c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:57:36 -0700 Subject: [PATCH 728/879] Restored example [skip ci] --- examples/loading/example.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/loading/example.py b/examples/loading/example.py index 4c0e579..7f3dce8 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -1,5 +1,4 @@ import numpy as np -from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -26,7 +25,7 @@ copy.set_types(['vector']) for i, embedding in enumerate(embeddings): - copy.write_row([Vector(embedding)]) + copy.write_row([embedding]) # show progress if i % 10000 == 0: From b4287d5c99588bc3a34dd7c0f1bb46a36dcd6ff1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:58:33 -0700 Subject: [PATCH 729/879] Restored VectorBinaryDumper [skip ci] --- pgvector/psycopg/vector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index d70a07c..b8b3611 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -28,7 +28,9 @@ class VectorBinaryDumper(VectorDumper): format = Format.BINARY def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: - return Vector._to_db_binary(obj) + if not isinstance(obj, Vector): + obj = Vector(obj) + return obj.to_binary() class VectorLoader(Loader): From dd28e2836454d562887cdaaff07dd5e7acf9dfe2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 16:59:23 -0700 Subject: [PATCH 730/879] Removed unused method [skip ci] --- pgvector/bit.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index c9bb50d..e1af56c 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -100,10 +100,3 @@ def _to_db(cls, value: Bit) -> str: raise ValueError('expected bit') return value.to_text() - - @classmethod - def _to_db_binary(cls, value: Bit) -> bytes: - if not isinstance(value, Bit): - raise ValueError('expected bit') - - return value.to_binary() From 04b2fd5a4c31cc4d7f7eb44a4fe968698c4fe9d6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:00:40 -0700 Subject: [PATCH 731/879] Removed unused method [skip ci] --- pgvector/bit.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index e1af56c..6afa41c 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -93,10 +93,3 @@ def from_binary(cls, value: bytes) -> Bit: bit = cls.__new__(cls) bit._value = value return bit - - @classmethod - def _to_db(cls, value: Bit) -> str: - if not isinstance(value, Bit): - raise ValueError('expected bit') - - return value.to_text() From 4dba27b9072a70a8fcd354b0522b95ab2e3a4d39 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:18:31 -0700 Subject: [PATCH 732/879] Reduced allocations [skip ci] --- pgvector/halfvec.py | 8 ++++++-- pgvector/vector.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 60fcac9..0e2736e 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -61,7 +61,7 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> HalfVector: - return cls([float(v) for v in value[1:-1].split(',')]) + return cls(cls._list_from_text(value)) @classmethod def from_binary(cls, value: bytes) -> HalfVector: @@ -109,6 +109,10 @@ def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: return value if isinstance(value, str): - value = cls.from_text(value) + return cls._list_from_text(value) return value.to_list() + + @classmethod + def _list_from_text(cls, value: str) -> list[float]: + return [float(v) for v in value[1:-1].split(',')] diff --git a/pgvector/vector.py b/pgvector/vector.py index 6a9d505..d7c0667 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -59,7 +59,7 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> Vector: - return cls([float(v) for v in value[1:-1].split(',')]) + return cls(cls._list_from_text(value)) @classmethod def from_binary(cls, value: bytes) -> Vector: @@ -107,6 +107,10 @@ def _from_db(cls, value: str | Vector | None) -> list[float] | None: return value if isinstance(value, str): - value = cls.from_text(value) + return cls._list_from_text(value) return value.to_list() + + @classmethod + def _list_from_text(cls, value: str) -> list[float]: + return [float(v) for v in value[1:-1].split(',')] From c399b606fdf8918507bce296c18a62231acc2e5f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:31:30 -0700 Subject: [PATCH 733/879] Improved code [skip ci] --- pgvector/halfvec.py | 7 ++++--- pgvector/vector.py | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 0e2736e..d5c014f 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -108,10 +108,11 @@ def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: if value is None: return value - if isinstance(value, str): - return cls._list_from_text(value) + # Handle conversion from low-level driver + if isinstance(value, HalfVector): + return value.to_list() - return value.to_list() + return cls._list_from_text(value) @classmethod def _list_from_text(cls, value: str) -> list[float]: diff --git a/pgvector/vector.py b/pgvector/vector.py index d7c0667..9a0fb30 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -106,10 +106,11 @@ def _from_db(cls, value: str | Vector | None) -> list[float] | None: if value is None: return value - if isinstance(value, str): - return cls._list_from_text(value) + # Handle conversion from low-level driver + if isinstance(value, Vector): + return value.to_list() - return value.to_list() + return cls._list_from_text(value) @classmethod def _list_from_text(cls, value: str) -> list[float]: From 9ad82cb185cf1843e309c418b3494a1f6f93d4be Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:33:53 -0700 Subject: [PATCH 734/879] Added comment [skip ci] --- pgvector/sparsevec.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index b08a064..60372d8 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -170,6 +170,7 @@ def _to_db_binary(cls, value: object) -> bytes | None: @classmethod def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None: + # Handle conversion from low-level driver if value is None or isinstance(value, SparseVector): return value From 3c1dd2d8897ec8579476d314ebc0647b619485dd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:36:38 -0700 Subject: [PATCH 735/879] Removed comments [skip ci] --- pgvector/halfvec.py | 1 - pgvector/sparsevec.py | 1 - pgvector/vector.py | 1 - 3 files changed, 3 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index d5c014f..8ad52c7 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -108,7 +108,6 @@ def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: if value is None: return value - # Handle conversion from low-level driver if isinstance(value, HalfVector): return value.to_list() diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 60372d8..b08a064 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -170,7 +170,6 @@ def _to_db_binary(cls, value: object) -> bytes | None: @classmethod def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None: - # Handle conversion from low-level driver if value is None or isinstance(value, SparseVector): return value diff --git a/pgvector/vector.py b/pgvector/vector.py index 9a0fb30..213cf43 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -106,7 +106,6 @@ def _from_db(cls, value: str | Vector | None) -> list[float] | None: if value is None: return value - # Handle conversion from low-level driver if isinstance(value, Vector): return value.to_list() From f7795913040bfe068d0978245eabe289c14509a4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:46:54 -0700 Subject: [PATCH 736/879] Improved register_vector for pg8000 [skip ci] --- pgvector/pg8000/register.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index 520fd10..adf7912 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -11,19 +11,19 @@ def register_vector(conn: Connection) -> None: if 'vector' not in type_info: raise RuntimeError('vector type not found in the database') - conn.register_out_adapter(Vector, Vector._to_db) + conn.register_out_adapter(Vector, lambda v: v.to_text()) conn.register_in_adapter(type_info['vector'], Vector.from_text) try: import numpy as np - conn.register_out_adapter(np.ndarray, Vector._to_db) + conn.register_out_adapter(np.ndarray, lambda v: Vector(v).to_text()) except ImportError: pass if 'halfvec' in type_info: - conn.register_out_adapter(HalfVector, HalfVector._to_db) + conn.register_out_adapter(HalfVector, lambda v: v.to_text()) conn.register_in_adapter(type_info['halfvec'], HalfVector.from_text) if 'sparsevec' in type_info: - conn.register_out_adapter(SparseVector, SparseVector._to_db) + conn.register_out_adapter(SparseVector, lambda v: v.to_text()) conn.register_in_adapter(type_info['sparsevec'], SparseVector.from_text) From 2c95c2ae6aafdb4994f10d3556d9f66673f5de68 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:51:52 -0700 Subject: [PATCH 737/879] Updated pyright config [skip ci] --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 3dc4c11..a342fa7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,7 @@ reportMissingTypeStubs = "none" reportPossiblyUnboundVariable = "none" reportPrivateUsage = "none" reportUnknownArgumentType = "none" +reportUnknownLambdaType = "none" reportUnknownMemberType = "none" reportUnknownVariableType = "none" reportUnnecessaryIsInstance = "none" From c912f50d39f1a2fdcc5e65ec7a240457243bb850 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 17:56:12 -0700 Subject: [PATCH 738/879] Simplified _to_db_binary [skip ci] --- pgvector/halfvec.py | 5 +---- pgvector/sparsevec.py | 5 +---- pgvector/vector.py | 5 +---- 3 files changed, 3 insertions(+), 12 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 8ad52c7..bfd678b 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -94,10 +94,7 @@ def _to_db(cls, value: object) -> str | None: return value.to_text() @classmethod - def _to_db_binary(cls, value: object) -> bytes | None: - if value is None: - return value - + def _to_db_binary(cls, value: object) -> bytes: if not isinstance(value, cls): value = cls(value) # type: ignore diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index b08a064..d4efe13 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -159,10 +159,7 @@ def _to_db(cls, value: object) -> str | None: return value.to_text() @classmethod - def _to_db_binary(cls, value: object) -> bytes | None: - if value is None: - return value - + def _to_db_binary(cls, value: object) -> bytes: if not isinstance(value, cls): value = cls(value) diff --git a/pgvector/vector.py b/pgvector/vector.py index 213cf43..00d9e82 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -92,10 +92,7 @@ def _to_db(cls, value: object) -> str | None: return value.to_text() @classmethod - def _to_db_binary(cls, value: object) -> bytes | None: - if value is None: - return value - + def _to_db_binary(cls, value: object) -> bytes: if not isinstance(value, cls): value = cls(value) # type: ignore From 8f2ec857651ff595d2675034305d26fa50f0e9a6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:00:18 -0700 Subject: [PATCH 739/879] Removed _to_db_binary methods [skip ci] --- pgvector/asyncpg/register.py | 6 +++--- pgvector/halfvec.py | 7 ------- pgvector/sparsevec.py | 7 ------- pgvector/vector.py | 7 ------- 4 files changed, 3 insertions(+), 24 deletions(-) diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 67983da..effa056 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -6,7 +6,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None: await conn.set_type_codec( 'vector', schema=schema, - encoder=Vector._to_db_binary, + encoder=lambda v: (v if isinstance(v, Vector) else Vector(v)).to_binary(), decoder=Vector.from_binary, format='binary' ) @@ -15,7 +15,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None: await conn.set_type_codec( 'halfvec', schema=schema, - encoder=HalfVector._to_db_binary, + encoder=lambda v: (v if isinstance(v, HalfVector) else HalfVector(v)).to_binary(), decoder=HalfVector.from_binary, format='binary' ) @@ -23,7 +23,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None: await conn.set_type_codec( 'sparsevec', schema=schema, - encoder=SparseVector._to_db_binary, + encoder=lambda v: (v if isinstance(v, SparseVector) else SparseVector(v)).to_binary(), decoder=SparseVector.from_binary, format='binary' ) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index bfd678b..4da0756 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -93,13 +93,6 @@ def _to_db(cls, value: object) -> str | None: return value.to_text() - @classmethod - def _to_db_binary(cls, value: object) -> bytes: - if not isinstance(value, cls): - value = cls(value) # type: ignore - - return value.to_binary() - @classmethod def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: if value is None: diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index d4efe13..0e737af 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -158,13 +158,6 @@ def _to_db(cls, value: object) -> str | None: return value.to_text() - @classmethod - def _to_db_binary(cls, value: object) -> bytes: - if not isinstance(value, cls): - value = cls(value) - - return value.to_binary() - @classmethod def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None: if value is None or isinstance(value, SparseVector): diff --git a/pgvector/vector.py b/pgvector/vector.py index 00d9e82..bfd6590 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -91,13 +91,6 @@ def _to_db(cls, value: object) -> str | None: return value.to_text() - @classmethod - def _to_db_binary(cls, value: object) -> bytes: - if not isinstance(value, cls): - value = cls(value) # type: ignore - - return value.to_binary() - @classmethod def _from_db(cls, value: str | Vector | None) -> list[float] | None: if value is None: From eeb1b7e58aada462da44eb0e64eba2a38fda84ce Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:04:46 -0700 Subject: [PATCH 740/879] Updated naming [skip ci] --- pgvector/halfvec.py | 6 +++--- pgvector/vector.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 4da0756..4e28047 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -61,7 +61,7 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> HalfVector: - return cls(cls._list_from_text(value)) + return cls(cls._from_text(value)) @classmethod def from_binary(cls, value: bytes) -> HalfVector: @@ -101,8 +101,8 @@ def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: if isinstance(value, HalfVector): return value.to_list() - return cls._list_from_text(value) + return cls._from_text(value) @classmethod - def _list_from_text(cls, value: str) -> list[float]: + def _from_text(cls, value: str) -> list[float]: return [float(v) for v in value[1:-1].split(',')] diff --git a/pgvector/vector.py b/pgvector/vector.py index bfd6590..782222a 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -59,7 +59,7 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> Vector: - return cls(cls._list_from_text(value)) + return cls(cls._from_text(value)) @classmethod def from_binary(cls, value: bytes) -> Vector: @@ -99,8 +99,8 @@ def _from_db(cls, value: str | Vector | None) -> list[float] | None: if isinstance(value, Vector): return value.to_list() - return cls._list_from_text(value) + return cls._from_text(value) @classmethod - def _list_from_text(cls, value: str) -> list[float]: + def _from_text(cls, value: str) -> list[float]: return [float(v) for v in value[1:-1].split(',')] From b5987c05a9840a6a09176d1c7c7421da107efd10 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:07:24 -0700 Subject: [PATCH 741/879] Moved _from_text methods [skip ci] --- pgvector/halfvec.py | 8 ++++---- pgvector/vector.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 4e28047..91c3ab1 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -79,6 +79,10 @@ def from_binary(cls, value: bytes) -> HalfVector: vec._value.byteswap() return vec + @classmethod + def _from_text(cls, value: str) -> list[float]: + return [float(v) for v in value[1:-1].split(',')] + @classmethod def _to_db(cls, value: object) -> str | None: if value is None: @@ -102,7 +106,3 @@ def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: return value.to_list() return cls._from_text(value) - - @classmethod - def _from_text(cls, value: str) -> list[float]: - return [float(v) for v in value[1:-1].split(',')] diff --git a/pgvector/vector.py b/pgvector/vector.py index 782222a..a6e83e6 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -77,6 +77,10 @@ def from_binary(cls, value: bytes) -> Vector: vec._value.byteswap() return vec + @classmethod + def _from_text(cls, value: str) -> list[float]: + return [float(v) for v in value[1:-1].split(',')] + @classmethod def _to_db(cls, value: object) -> str | None: if value is None: @@ -100,7 +104,3 @@ def _from_db(cls, value: str | Vector | None) -> list[float] | None: return value.to_list() return cls._from_text(value) - - @classmethod - def _from_text(cls, value: str) -> list[float]: - return [float(v) for v in value[1:-1].split(',')] From 239ad25466b2077ebcd3efe35489594bce02302f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:25:20 -0700 Subject: [PATCH 742/879] Updated comments [skip ci] --- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 91c3ab1..fe874fd 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -88,7 +88,7 @@ def _to_db(cls, value: object) -> str | None: if value is None: return value - # fast path for high-level libraries + # fewer allocations for lists if isinstance(value, list): return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore diff --git a/pgvector/vector.py b/pgvector/vector.py index a6e83e6..16846e5 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -86,7 +86,7 @@ def _to_db(cls, value: object) -> str | None: if value is None: return value - # fast path for high-level libraries + # fewer allocations for lists if isinstance(value, list): return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore From f6052533b997726d9050be3449bfb219125fbada Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:31:38 -0700 Subject: [PATCH 743/879] Added test for bad type [skip ci] --- tests/test_sqlalchemy.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 475ce3e..0f5cbf3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -468,6 +468,13 @@ def test_sum_orm(self, engine: Engine) -> None: res = session.scalars(select(sum(Item.embedding))).one() assert res == [5, 7, 9] + def test_bad_type(self, engine: Engine) -> None: + item = Item(embedding=[1, 'two']) + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() + def test_bad_dimensions(self, engine: Engine) -> None: item = Item(embedding=[1, 2]) with Session(engine) as session: From 858e6b1cdfc2beb7f4a0161f67e919ecb29fe358 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:33:12 -0700 Subject: [PATCH 744/879] Added more tests for bad type --- tests/test_sqlalchemy.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0f5cbf3..f0be2a2 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -468,13 +468,27 @@ def test_sum_orm(self, engine: Engine) -> None: res = session.scalars(select(sum(Item.embedding))).one() assert res == [5, 7, 9] - def test_bad_type(self, engine: Engine) -> None: + def test_vector_bad_type(self, engine: Engine) -> None: item = Item(embedding=[1, 'two']) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='could not convert string to float'): session.commit() + def test_halfvec_bad_type(self, engine: Engine) -> None: + item = Item(half_embedding=[1, 'two']) + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() + + def test_sparsevec_bad_type(self, engine: Engine) -> None: + item = Item(sparse_embedding=[1, 'two']) + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() + def test_bad_dimensions(self, engine: Engine) -> None: item = Item(embedding=[1, 2]) with Session(engine) as session: From cfc040468f54f8e7ec1a095aa08523e85a90167f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:35:07 -0700 Subject: [PATCH 745/879] Added test for bad type [skip ci] --- tests/test_sparse_vector.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 7720644..ed419c5 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -26,6 +26,10 @@ def test_list(self) -> None: def test_list_empty(self) -> None: assert SparseVector([]).to_list() == [] + def test_list_str(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float') as error: + SparseVector([1, 'two', 3]) # type: ignore + def test_list_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore From e6372ce23d4e727f11b382d9893d7b654ee57dc9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:49:06 -0700 Subject: [PATCH 746/879] Split length and data for Bit [skip ci] --- pgvector/bit.py | 32 ++++++++++++++------------------ 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 6afa41c..ddacf71 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,6 +1,5 @@ from __future__ import annotations from struct import pack, unpack_from -from typing import cast try: import numpy as np @@ -12,8 +11,8 @@ class Bit: def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None: if isinstance(value, bytes): - length = 8 * len(value) - data = value + self._length = 8 * len(value) + self._data = value elif isinstance(value, (list, str)): if isinstance(value, list): def bit_value(v: bool) -> str: @@ -30,8 +29,9 @@ def bit_value(v: bool) -> str: if length % 8 != 0: value += '0' * (8 - (length % 8)) + self._length = length try: - data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') except ValueError: raise ValueError('expected bit string') elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): @@ -44,37 +44,31 @@ def bit_value(v: bool) -> str: if value.ndim != 1: raise ValueError('expected ndim to be 1') - length = len(value) - data = np.packbits(value).tobytes() + self._length = len(value) + self._data = np.packbits(value).tobytes() else: raise ValueError('expected bytes, str, list, or ndarray') - self._value = pack('>i', length) + data - def __repr__(self) -> str: return f'Bit({self.to_text()})' def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return self.to_binary() == other.to_binary() + return self._length == other._length and self._data == other._data return False - def _length(self) -> int: - length, = cast(tuple[int], unpack_from('>i', self._value)) - return length - def to_list(self) -> list[bool]: # TODO improve return [v != '0' for v in self.to_text()] def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]: - return np.unpackbits(np.frombuffer(self._value[4:], dtype=np.uint8), count=self._length()).astype(bool) + return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool) def to_text(self) -> str: - return ''.join(format(v, '08b') for v in self._value[4:])[:self._length()] + return ''.join(format(v, '08b') for v in self._data)[:self._length] def to_binary(self) -> bytes: - return self._value + return pack('>i', self._length) + self._data @classmethod def from_text(cls, value: str) -> Bit: @@ -86,10 +80,12 @@ def from_binary(cls, value: bytes) -> Bit: raise ValueError('expected bytes') length, = unpack_from('>i', value) + data = value[4:] - if len(value) != 4 + (length + 7) // 8: + if len(data) != (length + 7) // 8: raise ValueError('invalid length') bit = cls.__new__(cls) - bit._value = value + bit._length = length + bit._data = data return bit From 7574a9d618a1cdb69179cb6e700964224f841874 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:49:36 -0700 Subject: [PATCH 747/879] Removed unused variable [skip ci] --- tests/test_sparse_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index ed419c5..74220bb 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -27,7 +27,7 @@ def test_list_empty(self) -> None: assert SparseVector([]).to_list() == [] def test_list_str(self) -> None: - with pytest.raises(ValueError, match='could not convert string to float') as error: + with pytest.raises(ValueError, match='could not convert string to float'): SparseVector([1, 'two', 3]) # type: ignore def test_list_dimensions(self) -> None: From 94c386d545a34abae5f2316bd112476d93556b08 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:50:52 -0700 Subject: [PATCH 748/879] Updated style [skip ci] --- pgvector/bit.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index ddacf71..38dffee 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -25,7 +25,6 @@ def bit_value(v: bool) -> str: value = ''.join([bit_value(v) for v in value]) length = len(value) - if length % 8 != 0: value += '0' * (8 - (length % 8)) From 7482ce1833e8937ead02097243d29baa60872453 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 18:57:14 -0700 Subject: [PATCH 749/879] Updated style [skip ci] --- pgvector/psycopg/vector.py | 4 ++-- pgvector/psycopg2/vector.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index b8b3611..b9608b3 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -18,7 +18,7 @@ class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: + def dump(self, obj: Vector | np.ndarray) -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_text().encode('utf8') @@ -27,7 +27,7 @@ def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: + def dump(self, obj: Vector | np.ndarray) -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_binary() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 0d11240..4cfcfc7 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -10,7 +10,7 @@ class VectorAdapter: - def __init__(self, value: 'Vector | np.ndarray') -> None: + def __init__(self, value: Vector | np.ndarray) -> None: if not isinstance(value, Vector): value = Vector(value) self._value = value From 54dfcdc0eaa73e3ab95e1bc4edfd1d6a9d166a35 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:00:17 -0700 Subject: [PATCH 750/879] Improved readability [skip ci] --- pgvector/halfvec.py | 11 +++++++---- pgvector/vector.py | 11 +++++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index fe874fd..876bf99 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -66,17 +66,20 @@ def from_text(cls, value: str) -> HalfVector: @classmethod def from_binary(cls, value: bytes) -> HalfVector: dim, unused = struct.unpack_from('>HH', value) + data = value[4:] - if len(value) != 4 + 2 * dim: + if len(data) != 2 * dim: raise ValueError('invalid length') if unused != 0: raise ValueError('expected unused to be 0') - vec = cls.__new__(cls) - vec._value = array.array('H', value[4:]) + arr = array.array('H', data) if sys.byteorder != 'big': - vec._value.byteswap() + arr.byteswap() + + vec = cls.__new__(cls) + vec._value = arr return vec @classmethod diff --git a/pgvector/vector.py b/pgvector/vector.py index 16846e5..922608d 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -64,17 +64,20 @@ def from_text(cls, value: str) -> Vector: @classmethod def from_binary(cls, value: bytes) -> Vector: dim, unused = struct.unpack_from('>HH', value) + data = value[4:] - if len(value) != 4 + 4 * dim: + if len(data) != 4 * dim: raise ValueError('invalid length') if unused != 0: raise ValueError('expected unused to be 0') - vec = cls.__new__(cls) - vec._value = array.array('f', value[4:]) + arr = array.array('f', data) if sys.byteorder != 'big': - vec._value.byteswap() + arr.byteswap() + + vec = cls.__new__(cls) + vec._value = arr return vec @classmethod From 730ab85ae858a634ef418c11d5129e1bbc6a5541 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:31:18 -0700 Subject: [PATCH 751/879] Reduced copying --- pgvector/bit.py | 7 ++----- pgvector/halfvec.py | 7 ++++--- pgvector/psycopg/halfvec.py | 2 -- pgvector/psycopg/sparsevec.py | 2 -- pgvector/psycopg/vector.py | 2 -- pgvector/sparsevec.py | 2 +- pgvector/vector.py | 7 ++++--- 7 files changed, 11 insertions(+), 18 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 38dffee..e93d11d 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -74,12 +74,9 @@ def from_text(cls, value: str) -> Bit: return cls(str(value)) @classmethod - def from_binary(cls, value: bytes) -> Bit: - if not isinstance(value, bytes): - raise ValueError('expected bytes') - + def from_binary(cls, value: bytes | bytearray | memoryview) -> Bit: length, = unpack_from('>i', value) - data = value[4:] + data = memoryview(value)[4:].tobytes() if len(data) != (length + 7) // 8: raise ValueError('invalid length') diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 876bf99..320c606 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -64,9 +64,9 @@ def from_text(cls, value: str) -> HalfVector: return cls(cls._from_text(value)) @classmethod - def from_binary(cls, value: bytes) -> HalfVector: + def from_binary(cls, value: bytes | bytearray | memoryview) -> HalfVector: dim, unused = struct.unpack_from('>HH', value) - data = value[4:] + data = memoryview(value)[4:] if len(data) != 2 * dim: raise ValueError('invalid length') @@ -74,7 +74,8 @@ def from_binary(cls, value: bytes) -> HalfVector: if unused != 0: raise ValueError('expected unused to be 0') - arr = array.array('H', data) + arr = array.array('H') + arr.frombytes(data) if sys.byteorder != 'big': arr.byteswap() diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 8c9a8c8..f16d569 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -35,8 +35,6 @@ class HalfVectorBinaryLoader(HalfVectorLoader): format = Format.BINARY def load(self, data: Buffer) -> HalfVector | None: - if isinstance(data, (bytearray, memoryview)): - data = bytes(data) return HalfVector.from_binary(data) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 2a98e6b..6cf8e00 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -35,8 +35,6 @@ class SparseVectorBinaryLoader(SparseVectorLoader): format = Format.BINARY def load(self, data: Buffer) -> SparseVector | None: - if isinstance(data, (bytearray, memoryview)): - data = bytes(data) return SparseVector.from_binary(data) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index b9608b3..ddc1163 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -46,8 +46,6 @@ class VectorBinaryLoader(VectorLoader): format = Format.BINARY def load(self, data: Buffer) -> Vector | None: - if isinstance(data, (bytearray, memoryview)): - data = bytes(data) return Vector.from_binary(data) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 0e737af..fdea411 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -127,7 +127,7 @@ def from_text(cls, value: str) -> SparseVector: return cls._from_parts(int(dim), indices, values) @classmethod - def from_binary(cls, value: bytes) -> SparseVector: + def from_binary(cls, value: bytes | bytearray | memoryview) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) if len(value) != 12 + 8 * nnz: diff --git a/pgvector/vector.py b/pgvector/vector.py index 922608d..dc5b9a1 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -62,9 +62,9 @@ def from_text(cls, value: str) -> Vector: return cls(cls._from_text(value)) @classmethod - def from_binary(cls, value: bytes) -> Vector: + def from_binary(cls, value: bytes | bytearray | memoryview) -> Vector: dim, unused = struct.unpack_from('>HH', value) - data = value[4:] + data = memoryview(value)[4:] if len(data) != 4 * dim: raise ValueError('invalid length') @@ -72,7 +72,8 @@ def from_binary(cls, value: bytes) -> Vector: if unused != 0: raise ValueError('expected unused to be 0') - arr = array.array('f', data) + arr = array.array('f') + arr.frombytes(data) if sys.byteorder != 'big': arr.byteswap() From e7a9a005c837cefea93b6944106ad8db9d203fde Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:35:28 -0700 Subject: [PATCH 752/879] Switched to tobytes [skip ci] --- pgvector/psycopg/halfvec.py | 2 +- pgvector/psycopg/sparsevec.py | 2 +- pgvector/psycopg/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index f16d569..b09134c 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -27,7 +27,7 @@ class HalfVectorLoader(Loader): def load(self, data: Buffer) -> HalfVector | None: if isinstance(data, memoryview): - data = bytes(data) + data = data.tobytes() return HalfVector.from_text(data.decode('utf8')) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 6cf8e00..f1e275c 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -27,7 +27,7 @@ class SparseVectorLoader(Loader): def load(self, data: Buffer) -> SparseVector | None: if isinstance(data, memoryview): - data = bytes(data) + data = data.tobytes() return SparseVector.from_text(data.decode('utf8')) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index ddc1163..2801ea9 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -38,7 +38,7 @@ class VectorLoader(Loader): def load(self, data: Buffer) -> Vector | None: if isinstance(data, memoryview): - data = bytes(data) + data = data.tobytes() return Vector.from_text(data.decode('utf8')) From 7f7780b56c6a0cc7df103b113756834d1434d9bb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:36:03 -0700 Subject: [PATCH 753/879] Removed unneeded cast [skip ci] --- pgvector/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index e93d11d..d7d0dfd 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -71,7 +71,7 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> Bit: - return cls(str(value)) + return cls(value) @classmethod def from_binary(cls, value: bytes | bytearray | memoryview) -> Bit: From a3b3b19b66aa7bf62d52fa37fe83c0b9be51e2fa Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:36:59 -0700 Subject: [PATCH 754/879] Restored cast [skip ci] --- pgvector/bit.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index d7d0dfd..70ac4fb 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -71,7 +71,8 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> Bit: - return cls(value) + # cast to ensure always use str constructor + return cls(str(value)) @classmethod def from_binary(cls, value: bytes | bytearray | memoryview) -> Bit: From e3e4308731d3aab82f9e91f60072fd970874ef2b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:37:14 -0700 Subject: [PATCH 755/879] Updated comment [skip ci] --- pgvector/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 70ac4fb..8af8e1c 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -71,7 +71,7 @@ def to_binary(self) -> bytes: @classmethod def from_text(cls, value: str) -> Bit: - # cast to ensure always use str constructor + # cast to ensure always uses str constructor return cls(str(value)) @classmethod From 191bee7ec4b639691d4c93c3b00b8a2923c7c6c2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 20:51:18 -0700 Subject: [PATCH 756/879] Reduced copying [skip ci] --- pgvector/halfvec.py | 8 +++----- pgvector/vector.py | 8 +++----- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 320c606..3e987c6 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -22,11 +22,9 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float if value.ndim != 1: raise ValueError('expected ndim to be 1') - if value.dtype != np.float16: - value = np.asarray(value, dtype=np.float16) - - # tobytes() important for performance - self._value = array.array('H', value.tobytes()) + arr = array.array('H') + arr.frombytes(value.astype(np.float16, order='C', copy=False).data.cast('B')) + self._value = arr else: raise ValueError('expected list or ndarray') diff --git a/pgvector/vector.py b/pgvector/vector.py index dc5b9a1..b9d131a 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -21,11 +21,9 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float if value.ndim != 1: raise ValueError('expected ndim to be 1') - if value.dtype != np.float32: - value = np.asarray(value, dtype=np.float32) - - # tobytes() important for performance - self._value = array.array('f', value.tobytes()) + arr = array.array('f') + arr.frombytes(value.astype(np.float32, order='C', copy=False).data.cast('B')) + self._value = arr else: raise ValueError('expected list or ndarray') From 120276fdb7b75b2272fd7e1388b20489e27032b8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:07:39 -0700 Subject: [PATCH 757/879] Removed tobytes --- pgvector/halfvec.py | 4 ++-- pgvector/vector.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 3e987c6..a2d660a 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -41,7 +41,7 @@ def dimensions(self) -> int: def to_list(self) -> list[float]: dim = len(self._value) - return list(struct.unpack(f'{dim}e', self._value.tobytes())) + return list(struct.unpack(f'{dim}e', self._value)) def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: return np.frombuffer(self._value, dtype=np.float16) @@ -55,7 +55,7 @@ def to_binary(self) -> bytes: else: value = array.array('H', self._value) value.byteswap() - return struct.pack(f'>HH', len(value), 0) + value.tobytes() + return struct.pack(f'>HH', len(value), 0) + memoryview(value) @classmethod def from_text(cls, value: str) -> HalfVector: diff --git a/pgvector/vector.py b/pgvector/vector.py index b9d131a..5d8fc32 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -53,7 +53,7 @@ def to_binary(self) -> bytes: else: value = array.array('f', self._value) value.byteswap() - return struct.pack(f'>HH', len(value), 0) + value.tobytes() + return struct.pack(f'>HH', len(value), 0) + memoryview(value) @classmethod def from_text(cls, value: str) -> Vector: From 7002b0f9d21ac36b1d992a5d49c77be70b5e2637 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:09:18 -0700 Subject: [PATCH 758/879] Revert "Updated style [skip ci]" This reverts commit 7482ce1833e8937ead02097243d29baa60872453. --- pgvector/psycopg/vector.py | 4 ++-- pgvector/psycopg2/vector.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 2801ea9..8df258f 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -18,7 +18,7 @@ class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: Vector | np.ndarray) -> Buffer | None: + def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_text().encode('utf8') @@ -27,7 +27,7 @@ def dump(self, obj: Vector | np.ndarray) -> Buffer | None: class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: Vector | np.ndarray) -> Buffer | None: + def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_binary() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 4cfcfc7..0d11240 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -10,7 +10,7 @@ class VectorAdapter: - def __init__(self, value: Vector | np.ndarray) -> None: + def __init__(self, value: 'Vector | np.ndarray') -> None: if not isinstance(value, Vector): value = Vector(value) self._value = value From 31c4a944f5c42c2d0ae5f60451be27ad170cc404 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:21:08 -0700 Subject: [PATCH 759/879] Added tests for from_text and from_binary for Bit --- tests/test_bit.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index 5e1da9c..958e07b 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,6 +1,7 @@ from pgvector import Bit import pytest import random +from struct import pack try: import numpy as np @@ -81,3 +82,17 @@ def test_repr(self) -> None: def test_equality(self) -> None: assert Bit([True, False, True]) == Bit([True, False, True]) assert Bit([True, False, True]) != Bit([True, False, False]) + + def test_from_text(self) -> None: + vec = Bit.from_text('101') + assert vec.to_list() == [True, False, True] + if NUMPY_AVAILABLE: + assert np.array_equal(vec.to_numpy(), [True, False, True]) + + def test_from_binary(self) -> None: + data = pack('>iB', 3, 5 << 5) + vec = Bit.from_binary(data) + assert vec.to_list() == [True, False, True] + if NUMPY_AVAILABLE: + assert np.array_equal(vec.to_numpy(), [True, False, True]) + assert vec.to_binary() == data From d5550772e1a92f3dfeb37cde9c149f7270ca1587 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:22:41 -0700 Subject: [PATCH 760/879] Updated CI --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c1d335e..2d307ce 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -8,7 +8,7 @@ jobs: matrix: python: [3.14, "3.10"] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python }} @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.4 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From c8e17bbe328d327506a08c938c2af6d587099b98 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:52:41 -0700 Subject: [PATCH 761/879] Added tests for to_text methods [skip ci] --- tests/test_bit.py | 1 + tests/test_half_vector.py | 1 + tests/test_sparse_vector.py | 1 + tests/test_vector.py | 1 + 4 files changed, 4 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index 958e07b..3502746 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -88,6 +88,7 @@ def test_from_text(self) -> None: assert vec.to_list() == [True, False, True] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [True, False, True]) + assert vec.to_text() == '101' def test_from_binary(self) -> None: data = pack('>iB', 3, 5 << 5) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 29143d1..d9f5a96 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -66,6 +66,7 @@ def test_from_text(self) -> None: assert vec.to_list() == [1.5, 2, 3] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + assert vec.to_text() == '[1.5,2.0,3.0]' def test_from_binary(self) -> None: data = pack('>HH3e', 3, 0, 1.5, 2, 3) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 74220bb..28a9ea7 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -130,6 +130,7 @@ def test_from_text(self) -> None: assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' def test_from_binary(self) -> None: data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) diff --git a/tests/test_vector.py b/tests/test_vector.py index c9a2184..1c22057 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -66,6 +66,7 @@ def test_from_text(self) -> None: assert vec.to_list() == [1.5, 2, 3] if NUMPY_AVAILABLE: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + assert vec.to_text() == '[1.5,2.0,3.0]' def test_from_binary(self) -> None: data = pack('>HH3f', 3, 0, 1.5, 2, 3) From 53f5b391eca13a4dab7f07b63c5e0367858647e2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:57:21 -0700 Subject: [PATCH 762/879] Switched to memoryview for consistency [skip ci] --- pgvector/sparsevec.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index fdea411..0ba80c1 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -129,15 +129,16 @@ def from_text(cls, value: str) -> SparseVector: @classmethod def from_binary(cls, value: bytes | bytearray | memoryview) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) + data = memoryview(value)[12:] - if len(value) != 12 + 8 * nnz: + if len(data) != 8 * nnz: raise ValueError('invalid length') if unused != 0: raise ValueError('expected unused to be 0') - indices = list(unpack_from(f'>{nnz}i', value, 12)) - values = list(unpack_from(f'>{nnz}f', value, 12 + nnz * 4)) + indices = list(unpack_from(f'>{nnz}i', data)) + values = list(unpack_from(f'>{nnz}f', data[nnz * 4:])) return cls._from_parts(dim, indices, values) @classmethod From 076b8d8b70cc451ad066b84d4d5d81904ddc9813 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 21:58:17 -0700 Subject: [PATCH 763/879] Reduced allocations [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 0ba80c1..ad38854 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -138,7 +138,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> SparseVector: raise ValueError('expected unused to be 0') indices = list(unpack_from(f'>{nnz}i', data)) - values = list(unpack_from(f'>{nnz}f', data[nnz * 4:])) + values = list(unpack_from(f'>{nnz}f', data, nnz * 4)) return cls._from_parts(dim, indices, values) @classmethod From 53688306f0f654470a9446e6588e0444b3a41f03 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 22:00:35 -0700 Subject: [PATCH 764/879] Revert "Switched to memoryview for consistency [skip ci]" This reverts commit 53f5b391eca13a4dab7f07b63c5e0367858647e2. --- pgvector/sparsevec.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index ad38854..fdea411 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -129,16 +129,15 @@ def from_text(cls, value: str) -> SparseVector: @classmethod def from_binary(cls, value: bytes | bytearray | memoryview) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) - data = memoryview(value)[12:] - if len(data) != 8 * nnz: + if len(value) != 12 + 8 * nnz: raise ValueError('invalid length') if unused != 0: raise ValueError('expected unused to be 0') - indices = list(unpack_from(f'>{nnz}i', data)) - values = list(unpack_from(f'>{nnz}f', data, nnz * 4)) + indices = list(unpack_from(f'>{nnz}i', value, 12)) + values = list(unpack_from(f'>{nnz}f', value, 12 + nnz * 4)) return cls._from_parts(dim, indices, values) @classmethod From 4b5f8551d7c0398070eebe8da052f252ce2bdb76 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 22:38:08 -0700 Subject: [PATCH 765/879] Improved typing [skip ci] --- pgvector/halfvec.py | 6 +++--- pgvector/peewee/halfvec.py | 2 +- pgvector/peewee/sparsevec.py | 2 +- pgvector/peewee/vector.py | 2 +- pgvector/sparsevec.py | 4 ++-- pgvector/vector.py | 6 +++--- 6 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index a2d660a..85e73fc 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -86,7 +86,7 @@ def _from_text(cls, value: str) -> list[float]: return [float(v) for v in value[1:-1].split(',')] @classmethod - def _to_db(cls, value: object) -> str | None: + def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | HalfVector | None) -> str | None: if value is None: return value @@ -94,8 +94,8 @@ def _to_db(cls, value: object) -> str | None: if isinstance(value, list): return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore - if not isinstance(value, cls): - value = cls(value) # type: ignore + if not isinstance(value, HalfVector): + value = cls(value) return value.to_text() diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 72ef759..988737f 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None - def db_value(self, value: object) -> str | None: + def db_value(self, value: Any) -> str | None: return HalfVector._to_db(value) def python_value(self, value: Any) -> list[float] | None: diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 20943a5..f3684f0 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None - def db_value(self, value: object) -> str | None: + def db_value(self, value: Any) -> str | None: return SparseVector._to_db(value) def python_value(self, value: Any) -> SparseVector | None: diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index d6864bc..1ac52f9 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.dimensions] if self.dimensions else None - def db_value(self, value: object) -> str | None: + def db_value(self, value: Any) -> str | None: return Vector._to_db(value) def python_value(self, value: Any) -> list[float] | None: diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index fdea411..73571f1 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -149,11 +149,11 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars return vec @classmethod - def _to_db(cls, value: object) -> str | None: + def _to_db(cls, value: list[float] | Any | SparseVector | None) -> str | None: if value is None: return value - if not isinstance(value, cls): + if not isinstance(value, SparseVector): value = cls(value) return value.to_text() diff --git a/pgvector/vector.py b/pgvector/vector.py index 5d8fc32..f3bbd1c 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -84,7 +84,7 @@ def _from_text(cls, value: str) -> list[float]: return [float(v) for v in value[1:-1].split(',')] @classmethod - def _to_db(cls, value: object) -> str | None: + def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | Vector | None) -> str | None: if value is None: return value @@ -92,8 +92,8 @@ def _to_db(cls, value: object) -> str | None: if isinstance(value, list): return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore - if not isinstance(value, cls): - value = cls(value) # type: ignore + if not isinstance(value, Vector): + value = cls(value) return value.to_text() From 4bcf31c01b65ac01ab75ef29e0572fdf5180fc58 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 22:52:21 -0700 Subject: [PATCH 766/879] Improved typing [skip ci] --- pgvector/sparsevec.py | 22 ++++++++++++++-------- tests/test_sparse_vector.py | 2 +- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 73571f1..73077fa 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -7,6 +7,12 @@ except ImportError: pass +try: + from scipy.sparse import sparray, spmatrix + SCIPY_AVAILABLE = True +except ImportError: + SCIPY_AVAILABLE = False + NO_DEFAULT = object() @@ -16,15 +22,15 @@ def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... @overload - def __init__(self, value: list[float], /) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]], /) -> None: ... @overload - def __init__(self, value: Any, /) -> None: + def __init__(self, value: sparray | spmatrix, /) -> None: ... - def __init__(self, value: dict[int, float] | list[float] | Any, dimensions: int | Any = NO_DEFAULT, /) -> None: - if value.__class__.__module__.startswith('scipy.sparse.'): + def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: + if SCIPY_AVAILABLE and isinstance(value, (sparray, spmatrix)): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') @@ -33,12 +39,12 @@ def __init__(self, value: dict[int, float] | list[float] | Any, dimensions: int if dimensions is NO_DEFAULT: raise ValueError('missing dimensions') - self._from_dict(value, dimensions) + self._from_dict(value, dimensions) # type: ignore else: if dimensions is not NO_DEFAULT: raise ValueError('extra argument') - self._from_dense(value) + self._from_dense(value) # type: ignore def __repr__(self) -> str: elements = dict(zip(self._indices, self._values)) @@ -108,7 +114,7 @@ def _from_sparse(self, value: Any) -> None: self._indices = value.col.tolist() self._values = value.data.tolist() - def _from_dense(self, value: list[float]) -> None: + def _from_dense(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: self._dim = len(value) self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] @@ -149,7 +155,7 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars return vec @classmethod - def _to_db(cls, value: list[float] | Any | SparseVector | None) -> str | None: + def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix | SparseVector | None) -> str | None: if value is None: return value diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 28a9ea7..89bffa3 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -51,7 +51,7 @@ def test_dict_empty(self) -> None: def test_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: - SparseVector({0: 1, 2: 2, 4: 3}) + SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore assert str(error.value) == 'missing dimensions' @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') From 3e7d56bf09c73d95852c996afba3d569cb9f3d73 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 23:07:01 -0700 Subject: [PATCH 767/879] Improved typing --- pgvector/sparsevec.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 73077fa..698d754 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -8,7 +8,7 @@ pass try: - from scipy.sparse import sparray, spmatrix + from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix SCIPY_AVAILABLE = True except ImportError: SCIPY_AVAILABLE = False @@ -34,7 +34,7 @@ def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int] if dimensions is not NO_DEFAULT: raise ValueError('extra argument') - self._from_sparse(value) + self._from_sparse(value) # type: ignore elif isinstance(value, dict): if dimensions is NO_DEFAULT: raise ValueError('missing dimensions') @@ -97,13 +97,13 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] - def _from_sparse(self, value: Any) -> None: - value = value.tocoo() + def _from_sparse(self, arr: sparray | spmatrix) -> None: + value: coo_array | coo_matrix = arr.tocoo() # type: ignore if value.ndim == 1: - self._dim = value.shape[0] - elif value.ndim == 2 and value.shape[0] == 1: - self._dim = value.shape[1] + self._dim = value.shape[0] # type: ignore + elif value.ndim == 2 and value.shape[0] == 1: # type: ignore + self._dim = value.shape[1] # type: ignore else: raise ValueError('expected ndim to be 1') From ab79b728c0c5bfb6c2ac4460f2bb839d5e005de2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 23:08:35 -0700 Subject: [PATCH 768/879] Improved typing --- pgvector/sparsevec.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 698d754..bf93f63 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -64,9 +64,7 @@ def indices(self) -> list[int]: def values(self) -> list[float]: return self._values - def to_coo(self) -> Any: - from scipy.sparse import coo_array - + def to_coo(self) -> coo_array: coords = ([0] * len(self._indices), self._indices) return coo_array((self._values, coords), shape=(1, self._dim)) From c892874d5c1bf00b29fc81db20a23d463b481ef4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 23:10:49 -0700 Subject: [PATCH 769/879] Improved typing [skip ci] --- pgvector/sparsevec.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index bf93f63..4a8eaa9 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -22,11 +22,7 @@ def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... @overload - def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]], /) -> None: - ... - - @overload - def __init__(self, value: sparray | spmatrix, /) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix, /) -> None: ... def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: From c1d00b0696afd0f87cc5635aa5d30204a0f822b3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 2 Jul 2026 23:29:30 -0700 Subject: [PATCH 770/879] Fixed typechecking for Python 3.10 [skip ci] --- pgvector/bit.py | 6 +++--- pgvector/django/extensions.py | 2 +- pgvector/halfvec.py | 4 ++-- pgvector/psycopg/vector.py | 4 ++-- pgvector/psycopg2/vector.py | 2 +- pgvector/sparsevec.py | 6 +++--- pgvector/vector.py | 4 ++-- 7 files changed, 14 insertions(+), 14 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 8af8e1c..5be31e9 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -9,7 +9,7 @@ class Bit: - def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None: + def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]]) -> None: if isinstance(value, bytes): self._length = 8 * len(value) self._data = value @@ -44,7 +44,7 @@ def bit_value(v: bool) -> str: raise ValueError('expected ndim to be 1') self._length = len(value) - self._data = np.packbits(value).tobytes() + self._data = np.packbits(value).tobytes() # type: ignore else: raise ValueError('expected bytes, str, list, or ndarray') @@ -60,7 +60,7 @@ def to_list(self) -> list[bool]: # TODO improve return [v != '0' for v in self.to_text()] - def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]: + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]: return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool) def to_text(self) -> str: diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py index debac5d..d75884f 100644 --- a/pgvector/django/extensions.py +++ b/pgvector/django/extensions.py @@ -6,6 +6,6 @@ class VectorExtension(CreateExtension): def __init__(self, hints: Any = None) -> None: if VERSION[0] >= 6: - super().__init__('vector', hints=hints) + super().__init__('vector', hints=hints) # type: ignore else: self.name = 'vector' diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 85e73fc..e2b2bc0 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -11,7 +11,7 @@ class HalfVector: - def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None: if isinstance(value, list): dim = len(value) try: @@ -43,7 +43,7 @@ def to_list(self) -> list[float]: dim = len(self._value) return list(struct.unpack(f'{dim}e', self._value)) - def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]: + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float16]]: return np.frombuffer(self._value, dtype=np.float16) def to_text(self) -> str: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 8df258f..a2951a1 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -18,7 +18,7 @@ class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: + def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_text().encode('utf8') @@ -27,7 +27,7 @@ def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None: + def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_binary() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 0d11240..b96ad86 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -10,7 +10,7 @@ class VectorAdapter: - def __init__(self, value: 'Vector | np.ndarray') -> None: + def __init__(self, value: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> None: if not isinstance(value, Vector): value = Vector(value) self._value = value diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 4a8eaa9..f5b88dc 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -22,10 +22,10 @@ def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... @overload - def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix, /) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, /) -> None: ... - def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: + def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: if SCIPY_AVAILABLE and isinstance(value, (sparray, spmatrix)): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') @@ -70,7 +70,7 @@ def to_list(self) -> list[float]: vec[i] = v return vec - def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: vec = np.repeat(0.0, self._dim).astype(np.float32) for i, v in zip(self._indices, self._values): vec[i] = v diff --git a/pgvector/vector.py b/pgvector/vector.py index f3bbd1c..867dc5e 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -11,7 +11,7 @@ class Vector: - def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: + def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None: if isinstance(value, list): try: self._value = array.array('f', value) @@ -41,7 +41,7 @@ def dimensions(self) -> int: def to_list(self) -> list[float]: return self._value.tolist() - def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]: + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: return np.frombuffer(self._value, dtype=np.float32) def to_text(self) -> str: From 24a95cb47acca2a35e62905f9b8b061e629af5ea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 00:23:00 -0700 Subject: [PATCH 771/879] Updated style [skip ci] --- pgvector/psycopg/vector.py | 5 +++-- pgvector/psycopg2/vector.py | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index a2951a1..8f3ddcc 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,3 +1,4 @@ +from __future__ import annotations import psycopg from psycopg import BaseConnection from psycopg.adapt import Loader, Dumper @@ -18,7 +19,7 @@ class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> Buffer | None: + def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_text().encode('utf8') @@ -27,7 +28,7 @@ def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]] class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> Buffer | None: + def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_binary() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index b96ad86..633b506 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,3 +1,4 @@ +from __future__ import annotations from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type from typing import Any from .. import Vector @@ -10,7 +11,7 @@ class VectorAdapter: - def __init__(self, value: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> None: + def __init__(self, value: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None: if not isinstance(value, Vector): value = Vector(value) self._value = value From ad4a7e89716b430342ac985149acbd9fc38c7f70 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 01:20:11 -0700 Subject: [PATCH 772/879] Removed unnecessary importing --- pgvector/_utils.py | 13 +++++++++++++ pgvector/bit.py | 13 ++++++++----- pgvector/halfvec.py | 12 +++++++----- pgvector/psycopg/vector.py | 13 ++++--------- pgvector/psycopg2/vector.py | 12 ++++++------ pgvector/sparsevec.py | 19 +++++++++---------- pgvector/vector.py | 12 +++++++----- 7 files changed, 54 insertions(+), 40 deletions(-) create mode 100644 pgvector/_utils.py diff --git a/pgvector/_utils.py b/pgvector/_utils.py new file mode 100644 index 0000000..bfef843 --- /dev/null +++ b/pgvector/_utils.py @@ -0,0 +1,13 @@ +import sys + + +def is_ndarray(value: object) -> bool: + if (numpy := sys.modules.get('numpy')): + return isinstance(value, numpy.ndarray) + return False + + +def is_sparse_array(value: object) -> bool: + if (sparse := sys.modules.get('scipy.sparse')): + return isinstance(value, (sparse.sparray, sparse.spmatrix)) + return False diff --git a/pgvector/bit.py b/pgvector/bit.py index 5be31e9..2cd7af1 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,11 +1,10 @@ from __future__ import annotations from struct import pack, unpack_from +from typing import TYPE_CHECKING +from ._utils import is_ndarray -try: +if TYPE_CHECKING: import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False class Bit: @@ -33,7 +32,9 @@ def bit_value(v: bool) -> str: self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') except ValueError: raise ValueError('expected bit string') - elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): + elif is_ndarray(value): + import numpy as np + if value.dtype != np.bool: # skip error for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): @@ -61,6 +62,8 @@ def to_list(self) -> list[bool]: return [v != '0' for v in self.to_text()] def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]: + import numpy as np + return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool) def to_text(self) -> str: diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index e2b2bc0..ec809bb 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -2,12 +2,11 @@ import array import struct import sys +from typing import TYPE_CHECKING +from ._utils import is_ndarray -try: +if TYPE_CHECKING: import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False class HalfVector: @@ -18,7 +17,9 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np. self._value = array.array('H', struct.pack(f'{dim}e', *value)) except struct.error: raise ValueError('expected list[float]') - elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): + elif is_ndarray(value): + import numpy as np + if value.ndim != 1: raise ValueError('expected ndim to be 1') @@ -44,6 +45,7 @@ def to_list(self) -> list[float]: return list(struct.unpack(f'{dim}e', self._value)) def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float16]]: + import numpy as np return np.frombuffer(self._value, dtype=np.float16) def to_text(self) -> str: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 8f3ddcc..63ca6ce 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -4,16 +4,13 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import Any, TypeAlias +from typing import TYPE_CHECKING, Any, TypeAlias from .. import Vector Buffer: TypeAlias = bytes | bytearray | memoryview -try: +if TYPE_CHECKING: import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False class VectorDumper(Dumper): @@ -62,9 +59,7 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) -> adapters = context.adapters adapters.register_dumper(Vector, text_dumper) adapters.register_dumper(Vector, binary_dumper) + adapters.register_dumper('numpy.ndarray', text_dumper) + adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) - - if NUMPY_AVAILABLE: - adapters.register_dumper(np.ndarray, text_dumper) - adapters.register_dumper(np.ndarray, binary_dumper) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 633b506..95fbcb9 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,13 +1,10 @@ from __future__ import annotations from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type -from typing import Any +from typing import TYPE_CHECKING, Any from .. import Vector -try: +if TYPE_CHECKING: import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False class VectorAdapter: @@ -38,5 +35,8 @@ def register_vector_info(oid: int, array_oid: int | None, scope: connection | cu register_adapter(Vector, VectorAdapter) - if NUMPY_AVAILABLE: + try: + import numpy as np register_adapter(np.ndarray, VectorAdapter) + except ImportError: + pass diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index f5b88dc..af84d2c 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,17 +1,12 @@ from __future__ import annotations from struct import pack, unpack_from -from typing import Any, overload +from typing import TYPE_CHECKING, Any, overload +from ._utils import is_sparse_array -try: +if TYPE_CHECKING: import numpy as np -except ImportError: - pass - -try: from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix - SCIPY_AVAILABLE = True -except ImportError: - SCIPY_AVAILABLE = False + NO_DEFAULT = object() @@ -26,7 +21,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np. ... def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: - if SCIPY_AVAILABLE and isinstance(value, (sparray, spmatrix)): + if is_sparse_array(value): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') @@ -61,6 +56,8 @@ def values(self) -> list[float]: return self._values def to_coo(self) -> coo_array: + from scipy.sparse import coo_array + coords = ([0] * len(self._indices), self._indices) return coo_array((self._values, coords), shape=(1, self._dim)) @@ -71,6 +68,8 @@ def to_list(self) -> list[float]: return vec def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: + import numpy as np + vec = np.repeat(0.0, self._dim).astype(np.float32) for i, v in zip(self._indices, self._values): vec[i] = v diff --git a/pgvector/vector.py b/pgvector/vector.py index 867dc5e..7cfdeb0 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -2,12 +2,11 @@ import array import struct import sys +from typing import TYPE_CHECKING +from ._utils import is_ndarray -try: +if TYPE_CHECKING: import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False class Vector: @@ -17,7 +16,9 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np. self._value = array.array('f', value) except TypeError: raise ValueError('expected list[float]') - elif NUMPY_AVAILABLE and isinstance(value, np.ndarray): + elif is_ndarray(value): + import numpy as np + if value.ndim != 1: raise ValueError('expected ndim to be 1') @@ -42,6 +43,7 @@ def to_list(self) -> list[float]: return self._value.tolist() def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: + import numpy as np return np.frombuffer(self._value, dtype=np.float32) def to_text(self) -> str: From c223ef4f8b95f7590ca2a407faf7d99bb8d46e34 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 01:51:45 -0700 Subject: [PATCH 773/879] Improved optional test imports --- tests/conftest.py | 8 +++++ tests/test_asyncpg.py | 13 +++----- tests/test_bit.py | 27 +++++++++------- tests/test_half_vector.py | 15 ++++----- tests/test_pg8000.py | 9 ++---- tests/test_psycopg.py | 23 ++++++++------ tests/test_psycopg2.py | 11 +++---- tests/test_sparse_vector.py | 62 +++++++++++++++++++++---------------- tests/test_sqlalchemy.py | 17 +++++----- tests/test_vector.py | 15 ++++----- 10 files changed, 102 insertions(+), 98 deletions(-) create mode 100644 tests/conftest.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e2878ab --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +try: + import numpy +except ImportError: + numpy = None # type: ignore + +__all__ = [ + 'numpy' +] diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 8ff9f3b..f1a94d8 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -3,12 +3,7 @@ from pgvector import HalfVector, SparseVector, Vector from pgvector.asyncpg import register_vector import pytest - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np class TestAsyncpg: @@ -26,7 +21,7 @@ async def test_vector(self) -> None: embedding = Vector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] - embedding3 = np.array([7.5, 8, 9]) if NUMPY_AVAILABLE else [7.5, 8, 9] + embedding3 = np.array([7.5, 8, 9]) if np is not None else [7.5, 8, 9] embedding4 = None await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), ($4)', embedding, embedding2, embedding3, embedding4) @@ -117,14 +112,14 @@ async def test_vector_array(self) -> None: embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]] await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings2[0], embeddings2[1]) - if NUMPY_AVAILABLE: + if np is not None: embeddings3 = [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[])', embeddings3[0], embeddings3[1]) res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embeddings'] == embeddings assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] - if NUMPY_AVAILABLE: + if np is not None: assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] # type: ignore await conn.close() diff --git a/tests/test_bit.py b/tests/test_bit.py index 3502746..ed46e09 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -2,12 +2,7 @@ import pytest import random from struct import pack - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np class TestBit: @@ -41,26 +36,34 @@ def test_bytes(self) -> None: assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_unpackbits(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint8(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([254, 7, 0], dtype=np.uint8) with pytest.raises(ValueError) as error: Bit(arr) assert str(error.value) == 'expected elements to be boolean' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray_uint16(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.raises(ValueError) as error: Bit(arr) # type: ignore @@ -86,7 +89,7 @@ def test_equality(self) -> None: def test_from_text(self) -> None: vec = Bit.from_text('101') assert vec.to_list() == [True, False, True] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [True, False, True]) assert vec.to_text() == '101' @@ -94,6 +97,6 @@ def test_from_binary(self) -> None: data = pack('>iB', 3, 5 << 5) vec = Bit.from_binary(data) assert vec.to_list() == [True, False, True] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [True, False, True]) assert vec.to_binary() == data diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index d9f5a96..4b033e3 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,12 +1,7 @@ from pgvector import HalfVector import pytest from struct import pack - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np class TestHalfVector: @@ -28,8 +23,10 @@ def test_list_list(self) -> None: HalfVector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr @@ -64,7 +61,7 @@ def test_dimensions(self) -> None: def test_from_text(self) -> None: vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_text() == '[1.5,2.0,3.0]' @@ -72,6 +69,6 @@ def test_from_binary(self) -> None: data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index ffb90ac..8c8e02c 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -2,12 +2,7 @@ from pgvector import HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np conn = Connection(getuser(), database='pgvector_python_test') @@ -24,7 +19,7 @@ def setup_method(self) -> None: def test_vector(self) -> None: embedding = Vector([1.5, 2, 3]) - embedding2 = np.array([4.5, 5, 6]) if NUMPY_AVAILABLE else Vector([4.5, 5, 6]) + embedding2 = np.array([4.5, 5, 6]) if np is not None else Vector([4.5, 5, 6]) embedding3 = None conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2), (:embedding3)', embedding=embedding, embedding2=embedding2, embedding3=embedding3) diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 6e8820d..04863e9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -5,12 +5,7 @@ from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest from typing import Any - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) @@ -49,27 +44,35 @@ def test_vector_binary_format_correct(self) -> None: res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] assert res == '[1.5,2,3]' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_binary_format(self) -> None: + if np is None: + pytest.skip('NumPy required') + embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] assert res == Vector(embedding) - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_text_format(self) -> None: + if np is None: + pytest.skip('NumPy required') + embedding = np.array([1.5, 2, 3]) res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] assert res == Vector(embedding) - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_binary_format_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') + embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] assert res == Vector([3, 2, 1.5]) - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy_text_format_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') + embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 9b42d35..badb7c2 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -4,12 +4,7 @@ from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor from psycopg2.pool import ThreadedConnectionPool import pytest - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -36,8 +31,10 @@ def test_vector(self) -> None: assert res[0][0] == embedding assert res[1][0] is None - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_vector_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 89bffa3..a4ada33 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,25 +1,19 @@ from pgvector import SparseVector import pytest from struct import pack +from .conftest import numpy as np try: - import numpy as np - NUMPY_AVAILABLE = True + from scipy import sparse except ImportError: - NUMPY_AVAILABLE = False - -try: - from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix - SCIPY_AVAILABLE = True -except ImportError: - SCIPY_AVAILABLE = False + sparse = None # type: ignore class TestSparseVector: def test_list(self) -> None: vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] @@ -35,8 +29,10 @@ def test_list_dimensions(self) -> None: SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore assert str(error.value) == 'extra argument' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') + vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] @@ -54,43 +50,55 @@ def test_dict_no_dimensions(self) -> None: SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore assert str(error.value) == 'missing dimensions' - @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array(self) -> None: - arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_array_dimensions(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + with pytest.raises(ValueError) as error: - SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore + SparseVector(sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore assert str(error.value) == 'extra argument' - @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_coo_matrix(self) -> None: - mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + mat = sparse.coo_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_dok_array(self) -> None: - arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_csr_array(self) -> None: - arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - @pytest.mark.skipif(not SCIPY_AVAILABLE, reason='SciPy required') def test_csr_matrix(self) -> None: - mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + mat = sparse.csr_matrix(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(mat) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] @@ -114,8 +122,10 @@ def test_indices(self) -> None: def test_values(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] - @pytest.mark.skipif(not NUMPY_AVAILABLE or not SCIPY_AVAILABLE, reason='NumPy and SciPy required') def test_to_coo(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) def test_zero_vector_text(self) -> None: @@ -128,7 +138,7 @@ def test_from_text(self) -> None: assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' @@ -139,6 +149,6 @@ def test_from_binary(self) -> None: assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_binary() == data diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index f0be2a2..bf12083 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -10,12 +10,7 @@ from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session from sqlalchemy.sql import func from typing import Any - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -496,16 +491,20 @@ def test_bad_dimensions(self, engine: Engine) -> None: with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_bad_ndim(self, engine: Engine) -> None: + if np is None: + pytest.skip('NumPy required') + item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_bad_dtype(self, engine: Engine) -> None: + if np is None: + pytest.skip('NumPy required') + item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: session.add(item) @@ -693,7 +692,7 @@ async def test_vector_array(self, engine: AsyncEngine) -> None: item = await session.get_one(Item, 1) assert item.embeddings == [[1, 2, 3], [4, 5, 6]] - if NUMPY_AVAILABLE: + if np is not None: session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) item = await session.get_one(Item, 2) assert item.embeddings == [[1, 2, 3], [4, 5, 6]] diff --git a/tests/test_vector.py b/tests/test_vector.py index 1c22057..314ddcd 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,12 +1,7 @@ from pgvector import Vector import pytest from struct import pack - -try: - import numpy as np - NUMPY_AVAILABLE = True -except ImportError: - NUMPY_AVAILABLE = False +from .conftest import numpy as np class TestVector: @@ -28,8 +23,10 @@ def test_list_list(self) -> None: Vector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' - @pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required') def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr @@ -64,7 +61,7 @@ def test_dimensions(self) -> None: def test_from_text(self) -> None: vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_text() == '[1.5,2.0,3.0]' @@ -72,6 +69,6 @@ def test_from_binary(self) -> None: data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - if NUMPY_AVAILABLE: + if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data From 643f8d1e4d444d887a19cea7e5e109b4ddfb10a8 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 02:09:38 -0700 Subject: [PATCH 774/879] Added ndarray type alias --- pgvector/_utils.py | 6 ++++++ pgvector/halfvec.py | 5 +++-- pgvector/psycopg/vector.py | 6 +++--- pgvector/psycopg2/vector.py | 4 ++-- pgvector/sparsevec.py | 9 +++++---- pgvector/vector.py | 5 +++-- 6 files changed, 22 insertions(+), 13 deletions(-) diff --git a/pgvector/_utils.py b/pgvector/_utils.py index bfef843..1995168 100644 --- a/pgvector/_utils.py +++ b/pgvector/_utils.py @@ -1,4 +1,10 @@ import sys +from typing import TYPE_CHECKING, TypeAlias + +if TYPE_CHECKING: + import numpy as np + + ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]] def is_ndarray(value: object) -> bool: diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index ec809bb..ae0aec6 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -7,10 +7,11 @@ if TYPE_CHECKING: import numpy as np + from ._utils import ndarray class HalfVector: - def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None: + def __init__(self, value: list[float] | ndarray) -> None: if isinstance(value, list): dim = len(value) try: @@ -88,7 +89,7 @@ def _from_text(cls, value: str) -> list[float]: return [float(v) for v in value[1:-1].split(',')] @classmethod - def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | HalfVector | None) -> str | None: + def _to_db(cls, value: list[float] | ndarray | HalfVector | None) -> str | None: if value is None: return value diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 63ca6ce..8163a9c 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -10,13 +10,13 @@ Buffer: TypeAlias = bytes | bytearray | memoryview if TYPE_CHECKING: - import numpy as np + from .._utils import ndarray class VectorDumper(Dumper): format = Format.TEXT - def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> Buffer | None: + def dump(self, obj: Vector | ndarray) -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_text().encode('utf8') @@ -25,7 +25,7 @@ def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) class VectorBinaryDumper(VectorDumper): format = Format.BINARY - def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> Buffer | None: + def dump(self, obj: Vector | ndarray) -> Buffer | None: if not isinstance(obj, Vector): obj = Vector(obj) return obj.to_binary() diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 95fbcb9..4f8176e 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -4,11 +4,11 @@ from .. import Vector if TYPE_CHECKING: - import numpy as np + from .._utils import ndarray class VectorAdapter: - def __init__(self, value: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None: + def __init__(self, value: Vector | ndarray) -> None: if not isinstance(value, Vector): value = Vector(value) self._value = value diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index af84d2c..e70aa20 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -6,6 +6,7 @@ if TYPE_CHECKING: import numpy as np from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix + from ._utils import ndarray NO_DEFAULT = object() @@ -17,10 +18,10 @@ def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... @overload - def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, /) -> None: + def __init__(self, value: list[float] | ndarray | sparray | spmatrix, /) -> None: ... - def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: + def __init__(self, value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: if is_sparse_array(value): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') @@ -107,7 +108,7 @@ def _from_sparse(self, arr: sparray | spmatrix) -> None: self._indices = value.col.tolist() self._values = value.data.tolist() - def _from_dense(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None: + def _from_dense(self, value: list[float] | ndarray) -> None: self._dim = len(value) self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] @@ -148,7 +149,7 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars return vec @classmethod - def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix | SparseVector | None) -> str | None: + def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector | None) -> str | None: if value is None: return value diff --git a/pgvector/vector.py b/pgvector/vector.py index 7cfdeb0..8792897 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -7,10 +7,11 @@ if TYPE_CHECKING: import numpy as np + from ._utils import ndarray class Vector: - def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None: + def __init__(self, value: list[float] | ndarray) -> None: if isinstance(value, list): try: self._value = array.array('f', value) @@ -86,7 +87,7 @@ def _from_text(cls, value: str) -> list[float]: return [float(v) for v in value[1:-1].split(',')] @classmethod - def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | Vector | None) -> str | None: + def _to_db(cls, value: list[float] | ndarray | Vector | None) -> str | None: if value is None: return value From f3b7d35a3733de5d339b835da6ac5863f8ced0cb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 02:37:47 -0700 Subject: [PATCH 775/879] Switched to import for Buffer --- pgvector/psycopg/bit.py | 7 +++---- pgvector/psycopg/halfvec.py | 5 ++--- pgvector/psycopg/sparsevec.py | 5 ++--- pgvector/psycopg/vector.py | 5 ++--- 4 files changed, 9 insertions(+), 13 deletions(-) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index a6ac6b2..06fe8de 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,12 +1,11 @@ from psycopg import BaseConnection -from psycopg.types import TypeInfo +from psycopg.abc import Buffer from psycopg.adapt import Dumper from psycopg.pq import Format -from typing import Any, TypeAlias +from psycopg.types import TypeInfo +from typing import Any from .. import Bit -Buffer: TypeAlias = bytes | bytearray | memoryview - class BitDumper(Dumper): format = Format.TEXT diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index b09134c..a7f5157 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,12 +1,11 @@ from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import Any, TypeAlias +from typing import Any from .. import HalfVector -Buffer: TypeAlias = bytes | bytearray | memoryview - class HalfVectorDumper(Dumper): format = Format.TEXT diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index f1e275c..4202a7b 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,12 +1,11 @@ from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import Any, TypeAlias +from typing import Any from .. import SparseVector -Buffer: TypeAlias = bytes | bytearray | memoryview - class SparseVectorDumper(Dumper): format = Format.TEXT diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 8163a9c..ebb9e75 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,14 +1,13 @@ from __future__ import annotations import psycopg from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import TYPE_CHECKING, Any, TypeAlias +from typing import TYPE_CHECKING, Any from .. import Vector -Buffer: TypeAlias = bytes | bytearray | memoryview - if TYPE_CHECKING: from .._utils import ndarray From 1b62e145674bdad212d7878019343b0299a547a9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 02:41:18 -0700 Subject: [PATCH 776/879] Improved pyright config [skip ci] --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a342fa7..aef991f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,6 @@ include = [ "tests" ] typeCheckingMode = "strict" -reportConstantRedefinition = "none" reportMissingTypeStubs = "none" reportPossiblyUnboundVariable = "none" reportPrivateUsage = "none" From 178088362c205c57e7d677501d4bf9ff17dc5269 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 03:05:37 -0700 Subject: [PATCH 777/879] Improved ndarray imports --- pgvector/_utils.py | 4 ++++ pgvector/halfvec.py | 3 +-- pgvector/psycopg/vector.py | 6 ++---- pgvector/psycopg2/vector.py | 6 ++---- pgvector/sparsevec.py | 3 +-- pgvector/vector.py | 3 +-- 6 files changed, 11 insertions(+), 14 deletions(-) diff --git a/pgvector/_utils.py b/pgvector/_utils.py index 1995168..2c7f7cd 100644 --- a/pgvector/_utils.py +++ b/pgvector/_utils.py @@ -5,6 +5,10 @@ import numpy as np ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]] +else: + # any value works since not type checking + # TODO use Never when Python 3.10 no longer supported + ndarray = None def is_ndarray(value: object) -> bool: diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index ae0aec6..f542781 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -3,11 +3,10 @@ import struct import sys from typing import TYPE_CHECKING -from ._utils import is_ndarray +from ._utils import is_ndarray, ndarray if TYPE_CHECKING: import numpy as np - from ._utils import ndarray class HalfVector: diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index ebb9e75..37580c8 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -5,11 +5,9 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format from psycopg.types import TypeInfo -from typing import TYPE_CHECKING, Any +from typing import Any from .. import Vector - -if TYPE_CHECKING: - from .._utils import ndarray +from .._utils import ndarray class VectorDumper(Dumper): diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 4f8176e..73baf51 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,10 +1,8 @@ from __future__ import annotations from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type -from typing import TYPE_CHECKING, Any +from typing import Any from .. import Vector - -if TYPE_CHECKING: - from .._utils import ndarray +from .._utils import ndarray class VectorAdapter: diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index e70aa20..712f626 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,12 +1,11 @@ from __future__ import annotations from struct import pack, unpack_from from typing import TYPE_CHECKING, Any, overload -from ._utils import is_sparse_array +from ._utils import is_sparse_array, ndarray if TYPE_CHECKING: import numpy as np from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix - from ._utils import ndarray NO_DEFAULT = object() diff --git a/pgvector/vector.py b/pgvector/vector.py index 8792897..0eb258e 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -3,11 +3,10 @@ import struct import sys from typing import TYPE_CHECKING -from ._utils import is_ndarray +from ._utils import is_ndarray, ndarray if TYPE_CHECKING: import numpy as np - from ._utils import ndarray class Vector: From 72fc947da115681c07c71637aa99bf43e09df59c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 03:54:30 -0700 Subject: [PATCH 778/879] Updated style [skip ci] --- pgvector/halfvec.py | 2 +- pgvector/vector.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index f542781..21725fb 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -57,7 +57,7 @@ def to_binary(self) -> bytes: else: value = array.array('H', self._value) value.byteswap() - return struct.pack(f'>HH', len(value), 0) + memoryview(value) + return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod def from_text(cls, value: str) -> HalfVector: diff --git a/pgvector/vector.py b/pgvector/vector.py index 0eb258e..0965a05 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -55,7 +55,7 @@ def to_binary(self) -> bytes: else: value = array.array('f', self._value) value.byteswap() - return struct.pack(f'>HH', len(value), 0) + memoryview(value) + return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod def from_text(cls, value: str) -> Vector: From cfe759ae0a936d67961a9bce0fd1bc6abb48d9c7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:11:39 -0700 Subject: [PATCH 779/879] Added mypy and ty to CI --- .github/workflows/build.yml | 4 ++++ pyproject.toml | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 2d307ce..c215202 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -27,3 +27,7 @@ jobs: - run: pip install --group dev-optional - run: pytest + + - run: pip install --group typing + - run: mypy + - run: ty check diff --git a/pyproject.toml b/pyproject.toml index aef991f..ce98f65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,11 @@ dev-optional = [ "numpy", "scipy" ] +typing = [ + "mypy", + "ty", + {include-group = "dev-optional"} +] [tool.pytest.ini_options] asyncio_mode = "auto" From ed6967b06ad4fb83e923c0dff38d3e121d2b8730 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:15:23 -0700 Subject: [PATCH 780/879] Fixed CI --- pyproject.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ce98f65..549b7a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,10 +51,9 @@ files = [ "tests" ] strict = true -disable_error_code = ["unused-ignore"] +disable_error_code = ["unused-ignore", "untyped-decorator"] ignore_missing_imports = true disallow_subclassing_any = false -python_executable = ".venv/bin/python" [tool.ty.src] include = [ From 392467211297df4c97a0bc4cc55b805d41846464 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:25:37 -0700 Subject: [PATCH 781/879] Fixed CI --- tests/test_peewee.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 79358ed..2599a77 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -18,6 +18,7 @@ class Item(BaseModel): half_embedding = HalfVectorField(dimensions=3, null=True) binary_embedding = FixedBitField(max_length=3, null=True) sparse_embedding = SparseVectorField(dimensions=3, null=True) + distance: float # for typing class Meta: # type: ignore table_name = 'peewee_item' @@ -26,7 +27,7 @@ class Meta: # type: ignore Item.add_index('embedding vector_l2_ops', using='hnsw') db.connect() -db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') +db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') # type: ignore db.drop_tables([Item]) db.create_tables([Item]) @@ -197,9 +198,9 @@ def test_halfvec_sum(self) -> None: assert sum == [5, 7, 9] def test_get_or_create(self) -> None: - Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) - Item.get_or_create(embedding=Vector([4, 5, 6])) - Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) + Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) # type: ignore + Item.get_or_create(embedding=Vector([4, 5, 6])) # type: ignore + Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) # type: ignore def test_vector_array(self) -> None: from playhouse.postgres_ext import PostgresqlExtDatabase, ArrayField From 006756c506adfa1b51f0dc1e4c16135b6d242319 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:36:44 -0700 Subject: [PATCH 782/879] Updated style [skip ci] --- tests/test_peewee.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 2599a77..294389d 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -18,7 +18,7 @@ class Item(BaseModel): half_embedding = HalfVectorField(dimensions=3, null=True) binary_embedding = FixedBitField(max_length=3, null=True) sparse_embedding = SparseVectorField(dimensions=3, null=True) - distance: float # for typing + distance: float # for typing class Meta: # type: ignore table_name = 'peewee_item' From 66eef363bb5940bb86893d55ab7df9dd3be47937 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:39:10 -0700 Subject: [PATCH 783/879] Improved typing for NO_DEFAULT --- pgvector/sparsevec.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 712f626..27477d0 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,6 +1,6 @@ from __future__ import annotations from struct import pack, unpack_from -from typing import TYPE_CHECKING, Any, overload +from typing import TYPE_CHECKING, Final, overload from ._utils import is_sparse_array, ndarray if TYPE_CHECKING: @@ -8,7 +8,11 @@ from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix -NO_DEFAULT = object() +class Sentinel: + pass + + +NO_DEFAULT: Final[Sentinel] = Sentinel() class SparseVector: @@ -20,7 +24,7 @@ def __init__(self, value: dict[int, float], dimensions: int, /) -> None: def __init__(self, value: list[float] | ndarray | sparray | spmatrix, /) -> None: ... - def __init__(self, value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None: + def __init__(self, value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, dimensions: int | Sentinel = NO_DEFAULT, /) -> None: if is_sparse_array(value): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') From cf0f4e9882af07a14119b92fecdaaba6a5da3375 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:54:27 -0700 Subject: [PATCH 784/879] Improved to_text method for SparseVector [skip ci] --- pgvector/sparsevec.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 27477d0..f4fe2df 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -80,7 +80,8 @@ def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: return vec def to_text(self) -> str: - return '{' + ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + '}/' + str(int(self._dim)) + elements = ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + return f'{{{elements}}}/{int(self._dim)}' def to_binary(self) -> bytes: nnz = len(self._indices) From a6465d1059e27f75bd82164e5ee26f4567b7b665 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 04:56:58 -0700 Subject: [PATCH 785/879] Improved to_numpy method for SparseVector [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index f4fe2df..064bca7 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -74,7 +74,7 @@ def to_list(self) -> list[float]: def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: import numpy as np - vec = np.repeat(0.0, self._dim).astype(np.float32) + vec = np.zeros(self._dim, dtype=np.float32) for i, v in zip(self._indices, self._values): vec[i] = v return vec From 9409d6483710b9236e4875701c007e2b37ec6ed9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 05:18:03 -0700 Subject: [PATCH 786/879] Improved typing [skip ci] --- pgvector/sparsevec.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 064bca7..3d7fa48 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,6 +1,6 @@ from __future__ import annotations from struct import pack, unpack_from -from typing import TYPE_CHECKING, Final, overload +from typing import TYPE_CHECKING, Final, cast, overload from ._utils import is_sparse_array, ndarray if TYPE_CHECKING: @@ -98,10 +98,11 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: def _from_sparse(self, arr: sparray | spmatrix) -> None: value: coo_array | coo_matrix = arr.tocoo() # type: ignore - if value.ndim == 1: - self._dim = value.shape[0] # type: ignore - elif value.ndim == 2 and value.shape[0] == 1: # type: ignore - self._dim = value.shape[1] # type: ignore + shape = cast(tuple[int], value.shape) + if len(shape) == 1: + self._dim = shape[0] + elif len(shape) == 2 and shape[0] == 1: + self._dim = shape[1] else: raise ValueError('expected ndim to be 1') From 24765f6ac27cd68f0ef6e5066ad620354a2f7824 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 05:33:28 -0700 Subject: [PATCH 787/879] Added typing to attributes --- pgvector/bit.py | 3 +++ pgvector/halfvec.py | 12 +++++++----- pgvector/sparsevec.py | 4 ++++ pgvector/vector.py | 12 +++++++----- 4 files changed, 21 insertions(+), 10 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 2cd7af1..6bba8e9 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -8,6 +8,9 @@ class Bit: + _length: int + _data: bytes + def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]]) -> None: if isinstance(value, bytes): self._length = 8 * len(value) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 21725fb..539b957 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,5 +1,5 @@ from __future__ import annotations -import array +from array import array import struct import sys from typing import TYPE_CHECKING @@ -10,11 +10,13 @@ class HalfVector: + _value: array[int] + def __init__(self, value: list[float] | ndarray) -> None: if isinstance(value, list): dim = len(value) try: - self._value = array.array('H', struct.pack(f'{dim}e', *value)) + self._value = array('H', struct.pack(f'{dim}e', *value)) except struct.error: raise ValueError('expected list[float]') elif is_ndarray(value): @@ -23,7 +25,7 @@ def __init__(self, value: list[float] | ndarray) -> None: if value.ndim != 1: raise ValueError('expected ndim to be 1') - arr = array.array('H') + arr = array('H') arr.frombytes(value.astype(np.float16, order='C', copy=False).data.cast('B')) self._value = arr else: @@ -55,7 +57,7 @@ def to_binary(self) -> bytes: if sys.byteorder == 'big': value = self._value else: - value = array.array('H', self._value) + value = array('H', self._value) value.byteswap() return struct.pack('>HH', len(value), 0) + memoryview(value) @@ -74,7 +76,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> HalfVector: if unused != 0: raise ValueError('expected unused to be 0') - arr = array.array('H') + arr = array('H') arr.frombytes(data) if sys.byteorder != 'big': arr.byteswap() diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 3d7fa48..fc0f09c 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -16,6 +16,10 @@ class Sentinel: class SparseVector: + _dim: int + _indices: list[int] + _values: list[float] + @overload def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... diff --git a/pgvector/vector.py b/pgvector/vector.py index 0965a05..a8f6645 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,5 +1,5 @@ from __future__ import annotations -import array +from array import array import struct import sys from typing import TYPE_CHECKING @@ -10,10 +10,12 @@ class Vector: + _value: array[float] + def __init__(self, value: list[float] | ndarray) -> None: if isinstance(value, list): try: - self._value = array.array('f', value) + self._value = array('f', value) except TypeError: raise ValueError('expected list[float]') elif is_ndarray(value): @@ -22,7 +24,7 @@ def __init__(self, value: list[float] | ndarray) -> None: if value.ndim != 1: raise ValueError('expected ndim to be 1') - arr = array.array('f') + arr = array('f') arr.frombytes(value.astype(np.float32, order='C', copy=False).data.cast('B')) self._value = arr else: @@ -53,7 +55,7 @@ def to_binary(self) -> bytes: if sys.byteorder == 'big': value = self._value else: - value = array.array('f', self._value) + value = array('f', self._value) value.byteswap() return struct.pack('>HH', len(value), 0) + memoryview(value) @@ -72,7 +74,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> Vector: if unused != 0: raise ValueError('expected unused to be 0') - arr = array.array('f') + arr = array('f') arr.frombytes(data) if sys.byteorder != 'big': arr.byteswap() From 515f8611979fa2d847e7966d4963abdf079bbd7a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 05:38:59 -0700 Subject: [PATCH 788/879] Added comment [skip ci] --- pgvector/halfvec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 539b957..9ca60f7 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -10,7 +10,7 @@ class HalfVector: - _value: array[int] + _value: array[int] # uses uint16 since no float16 def __init__(self, value: list[float] | ndarray) -> None: if isinstance(value, list): From 2ae28ba366d4cdae37c6370e74f91ff5b7cb7131 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 05:53:22 -0700 Subject: [PATCH 789/879] Improved performance of Bit constructor with list argument [skip ci] --- pgvector/bit.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 6bba8e9..180df43 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -17,15 +17,12 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], self._data = value elif isinstance(value, (list, str)): if isinstance(value, list): - def bit_value(v: bool) -> str: - if v is True: - return '1' - if v is False: - return '0' + bits = {True: '1', False: '0'} + try: + value = ''.join([bits[v] for v in value]) + except (KeyError, TypeError): raise ValueError('expected list[bool]') - value = ''.join([bit_value(v) for v in value]) - length = len(value) if length % 8 != 0: value += '0' * (8 - (length % 8)) From 2cff2f861c8872de7951c167d22ab2d524dd82d6 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 11:26:53 -0700 Subject: [PATCH 790/879] Reduced copying in _from_sparse method [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index fc0f09c..44ebfd1 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -100,7 +100,7 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: self._values = [float(v[1]) for v in elements] def _from_sparse(self, arr: sparray | spmatrix) -> None: - value: coo_array | coo_matrix = arr.tocoo() # type: ignore + value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore shape = cast(tuple[int], value.shape) if len(shape) == 1: From 05a7c0ce105160440b680bd03547235e7480cd82 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 12:52:15 -0700 Subject: [PATCH 791/879] Improved __eq__ for SparseVector [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 44ebfd1..b03404a 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -51,7 +51,7 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): - return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values() + return self._dim == other._dim and self._indices == other._indices and self._values == other._values return False def dimensions(self) -> int: From 4808799df338b877b4f1e573a897263e0d4a6c30 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 14:28:57 -0700 Subject: [PATCH 792/879] Updated style [skip ci] --- pgvector/sparsevec.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index b03404a..ca2c7aa 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -21,14 +21,17 @@ class SparseVector: _values: list[float] @overload - def __init__(self, value: dict[int, float], dimensions: int, /) -> None: - ... + def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... @overload - def __init__(self, value: list[float] | ndarray | sparray | spmatrix, /) -> None: - ... - - def __init__(self, value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, dimensions: int | Sentinel = NO_DEFAULT, /) -> None: + def __init__(self, value: list[float] | ndarray | sparray | spmatrix, /) -> None: ... + + def __init__( + self, + value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, + dimensions: int | Sentinel = NO_DEFAULT, + / + ) -> None: if is_sparse_array(value): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') From 892657ea5d7675944864f181585012c2163bd120 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 14:38:55 -0700 Subject: [PATCH 793/879] Switched to positional-only parameters --- pgvector/_utils.py | 4 ++-- pgvector/bit.py | 8 ++++---- pgvector/halfvec.py | 14 +++++++------- pgvector/sparsevec.py | 16 ++++++++-------- pgvector/vector.py | 14 +++++++------- 5 files changed, 28 insertions(+), 28 deletions(-) diff --git a/pgvector/_utils.py b/pgvector/_utils.py index 2c7f7cd..4cafc32 100644 --- a/pgvector/_utils.py +++ b/pgvector/_utils.py @@ -11,13 +11,13 @@ ndarray = None -def is_ndarray(value: object) -> bool: +def is_ndarray(value: object, /) -> bool: if (numpy := sys.modules.get('numpy')): return isinstance(value, numpy.ndarray) return False -def is_sparse_array(value: object) -> bool: +def is_sparse_array(value: object, /) -> bool: if (sparse := sys.modules.get('scipy.sparse')): return isinstance(value, (sparse.sparray, sparse.spmatrix)) return False diff --git a/pgvector/bit.py b/pgvector/bit.py index 180df43..7d6ba61 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -11,7 +11,7 @@ class Bit: _length: int _data: bytes - def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]]) -> None: + def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]], /) -> None: if isinstance(value, bytes): self._length = 8 * len(value) self._data = value @@ -52,7 +52,7 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], def __repr__(self) -> str: return f'Bit({self.to_text()})' - def __eq__(self, other: object) -> bool: + def __eq__(self, other: object, /) -> bool: if isinstance(other, self.__class__): return self._length == other._length and self._data == other._data return False @@ -73,12 +73,12 @@ def to_binary(self) -> bytes: return pack('>i', self._length) + self._data @classmethod - def from_text(cls, value: str) -> Bit: + def from_text(cls, value: str, /) -> Bit: # cast to ensure always uses str constructor return cls(str(value)) @classmethod - def from_binary(cls, value: bytes | bytearray | memoryview) -> Bit: + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Bit: length, = unpack_from('>i', value) data = memoryview(value)[4:].tobytes() diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 9ca60f7..90eaf3a 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -12,7 +12,7 @@ class HalfVector: _value: array[int] # uses uint16 since no float16 - def __init__(self, value: list[float] | ndarray) -> None: + def __init__(self, value: list[float] | ndarray, /) -> None: if isinstance(value, list): dim = len(value) try: @@ -34,7 +34,7 @@ def __init__(self, value: list[float] | ndarray) -> None: def __repr__(self) -> str: return f'HalfVector({self.to_list()})' - def __eq__(self, other: object) -> bool: + def __eq__(self, other: object, /) -> bool: if isinstance(other, self.__class__): return self._value == other._value return False @@ -62,11 +62,11 @@ def to_binary(self) -> bytes: return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod - def from_text(cls, value: str) -> HalfVector: + def from_text(cls, value: str, /) -> HalfVector: return cls(cls._from_text(value)) @classmethod - def from_binary(cls, value: bytes | bytearray | memoryview) -> HalfVector: + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: dim, unused = struct.unpack_from('>HH', value) data = memoryview(value)[4:] @@ -86,11 +86,11 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> HalfVector: return vec @classmethod - def _from_text(cls, value: str) -> list[float]: + def _from_text(cls, value: str, /) -> list[float]: return [float(v) for v in value[1:-1].split(',')] @classmethod - def _to_db(cls, value: list[float] | ndarray | HalfVector | None) -> str | None: + def _to_db(cls, value: list[float] | ndarray | HalfVector | None, /) -> str | None: if value is None: return value @@ -104,7 +104,7 @@ def _to_db(cls, value: list[float] | ndarray | HalfVector | None) -> str | None: return value.to_text() @classmethod - def _from_db(cls, value: str | HalfVector | None) -> list[float] | None: + def _from_db(cls, value: str | HalfVector | None, /) -> list[float] | None: if value is None: return value diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index ca2c7aa..04435f3 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -52,7 +52,7 @@ def __repr__(self) -> str: elements = dict(zip(self._indices, self._values)) return f'SparseVector({elements}, {self._dim})' - def __eq__(self, other: object) -> bool: + def __eq__(self, other: object, /) -> bool: if isinstance(other, self.__class__): return self._dim == other._dim and self._indices == other._indices and self._values == other._values return False @@ -102,7 +102,7 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] - def _from_sparse(self, arr: sparray | spmatrix) -> None: + def _from_sparse(self, arr: sparray | spmatrix, /) -> None: value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore shape = cast(tuple[int], value.shape) @@ -120,13 +120,13 @@ def _from_sparse(self, arr: sparray | spmatrix) -> None: self._indices = value.col.tolist() self._values = value.data.tolist() - def _from_dense(self, value: list[float] | ndarray) -> None: + def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] @classmethod - def from_text(cls, value: str) -> SparseVector: + def from_text(cls, value: str, /) -> SparseVector: elements, dim = value.split('/', 2) indices: list[int] = [] values: list[float] = [] @@ -139,7 +139,7 @@ def from_text(cls, value: str) -> SparseVector: return cls._from_parts(int(dim), indices, values) @classmethod - def from_binary(cls, value: bytes | bytearray | memoryview) -> SparseVector: + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) if len(value) != 12 + 8 * nnz: @@ -153,7 +153,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> SparseVector: return cls._from_parts(dim, indices, values) @classmethod - def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> SparseVector: + def _from_parts(cls, dim: int, indices: list[int], values: list[float], /) -> SparseVector: vec = cls.__new__(cls) vec._dim = dim vec._indices = indices @@ -161,7 +161,7 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars return vec @classmethod - def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector | None) -> str | None: + def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector | None, /) -> str | None: if value is None: return value @@ -171,7 +171,7 @@ def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector return value.to_text() @classmethod - def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None: + def _from_db(cls, value: str | SparseVector | None, /) -> SparseVector | None: if value is None or isinstance(value, SparseVector): return value diff --git a/pgvector/vector.py b/pgvector/vector.py index a8f6645..53fee7e 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -12,7 +12,7 @@ class Vector: _value: array[float] - def __init__(self, value: list[float] | ndarray) -> None: + def __init__(self, value: list[float] | ndarray, /) -> None: if isinstance(value, list): try: self._value = array('f', value) @@ -33,7 +33,7 @@ def __init__(self, value: list[float] | ndarray) -> None: def __repr__(self) -> str: return f'Vector({self.to_list()})' - def __eq__(self, other: object) -> bool: + def __eq__(self, other: object, /) -> bool: if isinstance(other, self.__class__): return self._value == other._value return False @@ -60,11 +60,11 @@ def to_binary(self) -> bytes: return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod - def from_text(cls, value: str) -> Vector: + def from_text(cls, value: str, /) -> Vector: return cls(cls._from_text(value)) @classmethod - def from_binary(cls, value: bytes | bytearray | memoryview) -> Vector: + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: dim, unused = struct.unpack_from('>HH', value) data = memoryview(value)[4:] @@ -84,11 +84,11 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> Vector: return vec @classmethod - def _from_text(cls, value: str) -> list[float]: + def _from_text(cls, value: str, /) -> list[float]: return [float(v) for v in value[1:-1].split(',')] @classmethod - def _to_db(cls, value: list[float] | ndarray | Vector | None) -> str | None: + def _to_db(cls, value: list[float] | ndarray | Vector | None, /) -> str | None: if value is None: return value @@ -102,7 +102,7 @@ def _to_db(cls, value: list[float] | ndarray | Vector | None) -> str | None: return value.to_text() @classmethod - def _from_db(cls, value: str | Vector | None) -> list[float] | None: + def _from_db(cls, value: str | Vector | None, /) -> list[float] | None: if value is None: return value From 1e57b6f6b072fbebe413e190e01c791c6fa442cb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 14:39:55 -0700 Subject: [PATCH 794/879] Updated style [skip ci] --- pgvector/bit.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 7d6ba61..64315aa 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -11,7 +11,11 @@ class Bit: _length: int _data: bytes - def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]], /) -> None: + def __init__( + self, + value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]], + / + ) -> None: if isinstance(value, bytes): self._length = 8 * len(value) self._data = value From 7fdf69d36562fda0b84bd6f3fafd1d8027e5fb6c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 14:47:33 -0700 Subject: [PATCH 795/879] Updated __eq__ methods to return NotImplemented --- pgvector/bit.py | 6 +++--- pgvector/halfvec.py | 6 +++--- pgvector/sparsevec.py | 6 +++--- pgvector/vector.py | 6 +++--- tests/test_bit.py | 1 + tests/test_half_vector.py | 1 + tests/test_sparse_vector.py | 1 + tests/test_vector.py | 1 + 8 files changed, 16 insertions(+), 12 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 64315aa..2b04912 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -57,9 +57,9 @@ def __repr__(self) -> str: return f'Bit({self.to_text()})' def __eq__(self, other: object, /) -> bool: - if isinstance(other, self.__class__): - return self._length == other._length and self._data == other._data - return False + if not isinstance(other, self.__class__): + return NotImplemented + return self._length == other._length and self._data == other._data def to_list(self) -> list[bool]: # TODO improve diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 90eaf3a..1a2c8c7 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -35,9 +35,9 @@ def __repr__(self) -> str: return f'HalfVector({self.to_list()})' def __eq__(self, other: object, /) -> bool: - if isinstance(other, self.__class__): - return self._value == other._value - return False + if not isinstance(other, self.__class__): + return NotImplemented + return self._value == other._value def dimensions(self) -> int: return len(self._value) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 04435f3..bd32c1f 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -53,9 +53,9 @@ def __repr__(self) -> str: return f'SparseVector({elements}, {self._dim})' def __eq__(self, other: object, /) -> bool: - if isinstance(other, self.__class__): - return self._dim == other._dim and self._indices == other._indices and self._values == other._values - return False + if not isinstance(other, self.__class__): + return NotImplemented + return self._dim == other._dim and self._indices == other._indices and self._values == other._values def dimensions(self) -> int: return self._dim diff --git a/pgvector/vector.py b/pgvector/vector.py index 53fee7e..d357113 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -34,9 +34,9 @@ def __repr__(self) -> str: return f'Vector({self.to_list()})' def __eq__(self, other: object, /) -> bool: - if isinstance(other, self.__class__): - return self._value == other._value - return False + if not isinstance(other, self.__class__): + return NotImplemented + return self._value == other._value def dimensions(self) -> int: return len(self._value) diff --git a/tests/test_bit.py b/tests/test_bit.py index ed46e09..ed13374 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -85,6 +85,7 @@ def test_repr(self) -> None: def test_equality(self) -> None: assert Bit([True, False, True]) == Bit([True, False, True]) assert Bit([True, False, True]) != Bit([True, False, False]) + assert Bit([True, False, True]) != 1 def test_from_text(self) -> None: vec = Bit.from_text('101') diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 4b033e3..26c5ac7 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -54,6 +54,7 @@ def test_repr(self) -> None: def test_equality(self) -> None: assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) + assert HalfVector([1, 2, 3]) != 1 def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index a4ada33..7d5f5d1 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -112,6 +112,7 @@ def test_equality(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert SparseVector({}, 1) != SparseVector({}, 2) + assert SparseVector([1, 0, 2, 0, 3, 0]) != 1 def test_dimensions(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 diff --git a/tests/test_vector.py b/tests/test_vector.py index 314ddcd..a08370e 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -54,6 +54,7 @@ def test_repr(self) -> None: def test_equality(self) -> None: assert Vector([1, 2, 3]) == Vector([1, 2, 3]) assert Vector([1, 2, 3]) != Vector([1, 2, 4]) + assert Vector([1, 2, 3]) != 1 def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 From 1654c9dbec76962638c3d08585234b97c7fc4889 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:27:36 -0700 Subject: [PATCH 796/879] Switched to positional-only parameters in more places [skip ci] --- pgvector/asyncpg/register.py | 2 +- pgvector/peewee/bit.py | 6 +++--- pgvector/peewee/halfvec.py | 10 +++++----- pgvector/peewee/sparsevec.py | 10 +++++----- pgvector/peewee/vector.py | 10 +++++----- pgvector/pg8000/register.py | 2 +- pgvector/psycopg/register.py | 4 ++-- pgvector/psycopg2/register.py | 2 +- pgvector/sqlalchemy/bit.py | 4 ++-- pgvector/sqlalchemy/halfvec.py | 8 ++++---- pgvector/sqlalchemy/sparsevec.py | 8 ++++---- pgvector/sqlalchemy/vector.py | 8 ++++---- 12 files changed, 37 insertions(+), 37 deletions(-) diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index effa056..6e01813 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -2,7 +2,7 @@ from .. import Vector, HalfVector, SparseVector -async def register_vector(conn: Connection, schema: str = 'public') -> None: +async def register_vector(conn: Connection, /, schema: str = 'public') -> None: await conn.set_type_codec( 'vector', schema=schema, diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index d38c08e..fcc7d8f 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -12,11 +12,11 @@ def __init__(self, max_length: int | None = None, *args: Any, **kwargs: Any) -> def get_modifiers(self) -> list[int] | None: return [self.max_length] if self.max_length else None - def _distance(self, op: str, vector: object) -> Expression: + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def hamming_distance(self, vector: object) -> Expression: + def hamming_distance(self, vector: object, /) -> Expression: return self._distance('<~>', vector) - def jaccard_distance(self, vector: object) -> Expression: + def jaccard_distance(self, vector: object, /) -> Expression: return self._distance('<%%>', vector) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 988737f..5e1e92c 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -19,17 +19,17 @@ def db_value(self, value: Any) -> str | None: def python_value(self, value: Any) -> list[float] | None: return HalfVector._from_db(value) - def _distance(self, op: str, vector: object) -> Expression: + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector: object) -> Expression: + def l2_distance(self, vector: object, /) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector: object) -> Expression: + def max_inner_product(self, vector: object, /) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector: object) -> Expression: + def cosine_distance(self, vector: object, /) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector: object) -> Expression: + def l1_distance(self, vector: object, /) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index f3684f0..dd83644 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -19,17 +19,17 @@ def db_value(self, value: Any) -> str | None: def python_value(self, value: Any) -> SparseVector | None: return SparseVector._from_db(value) - def _distance(self, op: str, vector: object) -> Expression: + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector: object) -> Expression: + def l2_distance(self, vector: object, /) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector: object) -> Expression: + def max_inner_product(self, vector: object, /) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector: object) -> Expression: + def cosine_distance(self, vector: object, /) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector: object) -> Expression: + def l1_distance(self, vector: object, /) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 1ac52f9..a28b6b5 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -19,17 +19,17 @@ def db_value(self, value: Any) -> str | None: def python_value(self, value: Any) -> list[float] | None: return Vector._from_db(value) - def _distance(self, op: str, vector: object) -> Expression: + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector: object) -> Expression: + def l2_distance(self, vector: object, /) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector: object) -> Expression: + def max_inner_product(self, vector: object, /) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector: object) -> Expression: + def cosine_distance(self, vector: object, /) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector: object) -> Expression: + def l1_distance(self, vector: object, /) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index adf7912..f3978d7 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -3,7 +3,7 @@ from .. import Vector, HalfVector, SparseVector -def register_vector(conn: Connection) -> None: +def register_vector(conn: Connection, /) -> None: # use to_regtype to get first matching type in search path res = cast(list[tuple[str, int]], conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")) type_info = dict(res) diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py index 15c2fc2..ed95042 100644 --- a/pgvector/psycopg/register.py +++ b/pgvector/psycopg/register.py @@ -7,7 +7,7 @@ from .vector import register_vector_info -def register_vector(context: Connection[Any]) -> None: +def register_vector(context: Connection[Any], /) -> None: info = TypeInfo.fetch(context, 'vector') register_vector_info(context, info) @@ -23,7 +23,7 @@ def register_vector(context: Connection[Any]) -> None: register_sparsevec_info(context, info) -async def register_vector_async(context: AsyncConnection[Any]) -> None: +async def register_vector_async(context: AsyncConnection[Any], /) -> None: info = await TypeInfo.fetch(context, 'vector') register_vector_info(context, info) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index f0004ce..95d1e56 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -6,7 +6,7 @@ # note: register_adapter is always global -def register_vector(conn_or_curs: connection | cursor, globally: bool = False, arrays: bool = True) -> None: +def register_vector(conn_or_curs: connection | cursor, /, globally: bool = False, arrays: bool = True) -> None: conn = conn_or_curs if isinstance(conn_or_curs, connection) else conn_or_curs.connection cur = conn.cursor(cursor_factory=cursor) scope = None if globally else conn_or_curs diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index a86ce1a..7c861ef 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -15,10 +15,10 @@ def process_bind_param(self, value: Any, dialect: Any) -> Any: return value class Comparator(TypeDecorator.Comparator[Any]): - def hamming_distance(self, other: object) -> Operators: + def hamming_distance(self, other: object, /) -> Operators: return self.op('<~>', return_type=Float)(other) - def jaccard_distance(self, other: object) -> Operators: + def jaccard_distance(self, other: object, /) -> Operators: return self.op('<%>', return_type=Float)(other) comparator_factory = Comparator # type: ignore diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index f0759e5..9d717c3 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -36,16 +36,16 @@ def process(value: Any) -> list[float] | None: return process class Comparator(TypeEngine.Comparator[Any]): - def l2_distance(self, other: object) -> Operators: + def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: object) -> Operators: + def max_inner_product(self, other: object, /) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: object) -> Operators: + def cosine_distance(self, other: object, /) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: object) -> Operators: + def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) comparator_factory = Comparator diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 8baff48..9bd1478 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -36,16 +36,16 @@ def process(value: Any) -> SparseVector | None: return process class Comparator(TypeEngine.Comparator[Any]): - def l2_distance(self, other: object) -> Operators: + def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: object) -> Operators: + def max_inner_product(self, other: object, /) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: object) -> Operators: + def cosine_distance(self, other: object, /) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: object) -> Operators: + def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) comparator_factory = Comparator diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 199afd7..83acb09 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -36,16 +36,16 @@ def process(value: Any) -> list[float] | None: return process class Comparator(TypeEngine.Comparator[Any]): - def l2_distance(self, other: object) -> Operators: + def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other: object) -> Operators: + def max_inner_product(self, other: object, /) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other: object) -> Operators: + def cosine_distance(self, other: object, /) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other: object) -> Operators: + def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) comparator_factory = Comparator From 841c6cc3802c3d458e067157ea3217f07e771360 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:36:03 -0700 Subject: [PATCH 797/879] Switched to positional-only parameters in more places [skip ci] --- pgvector/django/functions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 73b35b7..352c631 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -6,7 +6,7 @@ class DistanceBase(Func): output_field = FloatField() # type: ignore - def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: + def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): if isinstance(vector, (Vector, HalfVector, SparseVector)): vector = Value(vector.to_text()) @@ -22,7 +22,7 @@ def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: class BitDistanceBase(Func): output_field = FloatField() # type: ignore - def __init__(self, expression: Any, vector: Any, **extra: Any) -> None: + def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): vector = Value(vector) super().__init__(expression, vector, **extra) From 58ea2e872cf075b5fd49454a8e26040c163343bf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:38:09 -0700 Subject: [PATCH 798/879] Switched to keyword-only parameters [skip ci] --- pgvector/asyncpg/register.py | 2 +- pgvector/psycopg2/register.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 6e01813..0a2715f 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -2,7 +2,7 @@ from .. import Vector, HalfVector, SparseVector -async def register_vector(conn: Connection, /, schema: str = 'public') -> None: +async def register_vector(conn: Connection, /, *, schema: str = 'public') -> None: await conn.set_type_codec( 'vector', schema=schema, diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 95d1e56..2a5af49 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -6,7 +6,7 @@ # note: register_adapter is always global -def register_vector(conn_or_curs: connection | cursor, /, globally: bool = False, arrays: bool = True) -> None: +def register_vector(conn_or_curs: connection | cursor, /, *, globally: bool = False, arrays: bool = True) -> None: conn = conn_or_curs if isinstance(conn_or_curs, connection) else conn_or_curs.connection cur = conn.cursor(cursor_factory=cursor) scope = None if globally else conn_or_curs From a2cbbd8ec0ccc7e8fe33e56aa057bf70f27ea14b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:42:35 -0700 Subject: [PATCH 799/879] Switched to positional-only parameters in more places [skip ci] --- pgvector/psycopg/bit.py | 2 +- pgvector/psycopg/halfvec.py | 2 +- pgvector/psycopg/sparsevec.py | 2 +- pgvector/psycopg/vector.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index 06fe8de..f899497 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -21,7 +21,7 @@ def dump(self, obj: Bit) -> Buffer | None: return obj.to_binary() -def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None) -> None: +def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> None: assert info is not None info.register(context) diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index a7f5157..9ed0453 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -37,7 +37,7 @@ def load(self, data: Buffer) -> HalfVector | None: return HalfVector.from_binary(data) -def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo) -> None: +def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 4202a7b..57fd19a 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -37,7 +37,7 @@ def load(self, data: Buffer) -> SparseVector | None: return SparseVector.from_binary(data) -def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo) -> None: +def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 37580c8..7ef16b0 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -44,7 +44,7 @@ def load(self, data: Buffer) -> Vector | None: return Vector.from_binary(data) -def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) -> None: +def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> None: if info is None: raise psycopg.ProgrammingError('vector type not found in the database') info.register(context) From eedd9fec4291de72ac72ea69e038a6de51834631 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:43:53 -0700 Subject: [PATCH 800/879] Switched to positional-only parameters in more places [skip ci] --- pgvector/psycopg2/halfvec.py | 2 +- pgvector/psycopg2/sparsevec.py | 2 +- pgvector/psycopg2/vector.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 4616095..e1be330 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -19,7 +19,7 @@ def cast_halfvec(value: str | bytes | None, cur: cursor) -> HalfVector | None: return HalfVector.from_text(value) -def register_halfvec_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: +def register_halfvec_info(oid: int, array_oid: int | None, scope: connection | cursor | None, /) -> None: halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) register_type(halfvec, scope) diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 9b475ff..47dbc66 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -19,7 +19,7 @@ def cast_sparsevec(value: str | bytes | None, cur: cursor) -> SparseVector | Non return SparseVector.from_text(value) -def register_sparsevec_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: +def register_sparsevec_info(oid: int, array_oid: int | None, scope: connection | cursor | None, /) -> None: sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) register_type(sparsevec, scope) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 73baf51..7d37648 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -23,7 +23,7 @@ def cast_vector(value: str | bytes | None, cur: cursor) -> Vector | None: return Vector.from_text(value) -def register_vector_info(oid: int, array_oid: int | None, scope: connection | cursor | None) -> None: +def register_vector_info(oid: int, array_oid: int | None, scope: connection | cursor | None, /) -> None: vector = new_type((oid,), 'VECTOR', cast_vector) register_type(vector, scope) From 277e5f08ead2650ef39b42acd3216f1880f7b91e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:56:40 -0700 Subject: [PATCH 801/879] Updated urls [skip ci] --- pgvector/django/bit.py | 2 +- pgvector/django/halfvec.py | 2 +- pgvector/django/sparsevec.py | 2 +- pgvector/django/vector.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 642770b..eab0534 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -3,7 +3,7 @@ from typing import Any -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class BitField(Field): description = 'Bit string' diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 88e0512..1ab4a47 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -4,7 +4,7 @@ from .. import HalfVector -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class HalfVectorField(Field): description = 'Half vector' empty_strings_allowed = False diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index cd289a4..a10c133 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -4,7 +4,7 @@ from .. import SparseVector -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class SparseVectorField(Field): description = 'Sparse vector' empty_strings_allowed = False diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 331999d..42970c6 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -4,7 +4,7 @@ from .. import Vector -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class VectorField(Field): description = 'Vector' empty_strings_allowed = False From 0c02a20682710d727ac6983581e4d3bca74957eb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:57:38 -0700 Subject: [PATCH 802/879] Removed version from sqlmodel [skip ci] --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 549b7a9..4a9763b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dev = [ "pytest", "pytest-asyncio", "SQLAlchemy[asyncio]>=2", - "sqlmodel>=0.0.12" + "sqlmodel" ] dev-optional = [ "numpy", From 98bdd96b6608cf6d7c895a0f19ee8d660ac92449 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 15:59:25 -0700 Subject: [PATCH 803/879] Updated changelog [skip ci] --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 42ff5d7..f7598df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 1.0.0 (unreleased) +## 0.5.0 (unreleased) - Added experimental support for type hints - Changed `vector` and `halfvec` types to return list for Django, SQLAlchemy, SQLModel, and Peewee From c8a5832c0dae49688e69551e199dc92a0735f4b2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:09:57 -0700 Subject: [PATCH 804/879] Improved typing [skip ci] --- tests/test_sqlmodel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index bac3aee..583810c 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -3,7 +3,7 @@ from pydantic import ConfigDict import pytest from sqlalchemy.exc import StatementError -from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text +from sqlmodel import Field, Index, Session, SQLModel, col, create_engine, delete, select, text engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: @@ -26,7 +26,7 @@ class Item(SQLModel, table=True): index = Index( 'sqlmodel_index', - Item.embedding, # type: ignore + col(Item.embedding), # col() for typing postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} From 7d171b99067ed1937af1658e2081238d75865cb1 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:17:24 -0700 Subject: [PATCH 805/879] Improved typing for SQLModel tests [skip ci] --- tests/test_sqlmodel.py | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 583810c..b5f21d5 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -15,6 +15,9 @@ class Item(SQLModel, table=True): model_config = ConfigDict(arbitrary_types_allowed=True) # type: ignore id: int | None = Field(default=None, primary_key=True) + # typing issues + # https://github.com/fastapi/sqlmodel/discussions/1228 + # https://github.com/fastapi/sqlmodel/pull/1345 embedding: list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore half_embedding: list[float] | None = Field(default=None, sa_type=HALFVEC(3)) # type: ignore binary_embedding: str | None = Field(default=None, sa_type=BIT(3)) # type: ignore @@ -79,25 +82,25 @@ def test_vector(self) -> None: def test_vector_l2_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.embedding).l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_vector_max_inner_product(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.embedding).max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_vector_cosine_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.embedding).cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] def test_vector_l1_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.embedding).l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_halfvec(self) -> None: @@ -110,25 +113,25 @@ def test_halfvec(self) -> None: def test_halfvec_l2_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.half_embedding).l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_halfvec_max_inner_product(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.half_embedding).max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_halfvec_cosine_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.half_embedding).cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] def test_halfvec_l1_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.half_embedding).l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_bit(self) -> None: @@ -141,13 +144,13 @@ def test_bit(self) -> None: def test_bit_hamming_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.binary_embedding).hamming_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_bit_jaccard_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.binary_embedding).jaccard_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_sparsevec(self) -> None: @@ -160,37 +163,37 @@ def test_sparsevec(self) -> None: def test_sparsevec_l2_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_sparsevec_max_inner_product(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] def test_sparsevec_cosine_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] def test_sparsevec_l1_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) # type: ignore + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] def test_filter(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) # type: ignore + items = session.exec(select(Item).filter(col(Item.embedding).l2_distance([1, 1, 1]) < 1)) # type: ignore assert [v.id for v in items] == [1] def test_select(self) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() # type: ignore + items = session.exec(select(col(Item.embedding).l2_distance([1, 1, 1]))).all() # type: ignore assert items[0] == 3 def test_vector_avg(self) -> None: From 86d4c7f108a300c70d973361b213f33753bda115 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:21:54 -0700 Subject: [PATCH 806/879] Improved typing [skip ci] --- tests/test_asyncpg.py | 9 ++++----- tests/test_sqlalchemy.py | 1 + 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index f1a94d8..1b07eff 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -65,6 +65,7 @@ async def test_bit(self) -> None: await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') + # https://github.com/MagicStack/py-pgproto/pull/32 embedding = asyncpg.BitString('101') # type: ignore embedding2 = None await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)', embedding, embedding2) @@ -112,15 +113,13 @@ async def test_vector_array(self) -> None: embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]] await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings2[0], embeddings2[1]) - if np is not None: - embeddings3 = [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[])', embeddings3[0], embeddings3[1]) + embeddings3 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] if np is not None else [[1.5, 2, 3], [4.5, 5, 6]] + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings3[0], embeddings3[1]) res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embeddings'] == embeddings assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] - if np is not None: - assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] # type: ignore + assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] await conn.close() diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index bf12083..7080d19 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -638,6 +638,7 @@ async def test_bit(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): + # https://github.com/MagicStack/py-pgproto/pull/32 embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' # type: ignore session.add(Item(id=1, binary_embedding=embedding)) item = await session.get_one(Item, 1) From e83dd7106eacdc70bcdd43fdf4c8acbb0e240610 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:35:44 -0700 Subject: [PATCH 807/879] Improved typing [skip ci] --- tests/test_sqlalchemy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 7080d19..e8f221f 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -3,7 +3,7 @@ from pgvector import SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest -from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine +from sqlalchemy import create_engine, event, insert, inspect, literal, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncEngine from sqlalchemy.ext.automap import automap_base @@ -556,7 +556,7 @@ def test_binary_quantize(self, engine: Engine) -> None: session.add(Item(id=3, embedding=[1, 2, 3])) session.commit() - distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) # type: ignore + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast(literal([3, -1, 2], VECTOR), VECTOR(3)))) items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] @@ -571,7 +571,7 @@ def test_binary_quantize_reranking(self, engine: Engine) -> None: session.add(Item(id=3, embedding=[1, 2, 3])) session.commit() - distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) # type: ignore + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast(literal([3, -1, 2], VECTOR), VECTOR(3)))) subquery = session.query(Item).order_by(distance).limit(20).subquery() items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all() assert [v.id for v in items] == [2, 3, 1] From 935118b356e766827cea9d03caf8b641d03941d2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:45:58 -0700 Subject: [PATCH 808/879] Improved typing for Django tests --- tests/test_django.py | 56 +++++++++++++++++++++++--------------------- 1 file changed, 29 insertions(+), 27 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 03529ae..5bb51ae 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -58,6 +58,8 @@ class Item(models.Model): double_embedding = ArrayField(FloatField(), null=True, blank=True) numeric_embedding = ArrayField(DecimalField(max_digits=20, decimal_places=10), null=True, blank=True) + objects: models.Manager # for typing + class Meta: app_label = 'django_app' indexes = [ @@ -159,51 +161,51 @@ class Meta: class TestDjango: def setup_method(self) -> None: - Item.objects.all().delete() # type: ignore + Item.objects.all().delete() def test_vector(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) assert item.embedding == [1, 2, 3] def test_vector_l2_distance(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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)] def test_vector_l2_distance_list(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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)] def test_vector_l2_distance_none(self) -> None: create_items() distance = L2Distance('embedding', None) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.distance for v in items] == [None, None, None] def test_vector_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_vector_cosine_distance(self) -> None: create_items() distance = CosineDistance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_vector_l1_distance(self) -> None: create_items() distance = L1Distance('embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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, 3] @@ -215,28 +217,28 @@ def test_halfvec(self) -> None: def test_halfvec_l2_distance(self) -> None: create_items() distance = L2Distance('half_embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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)] def test_halfvec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('half_embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_halfvec_cosine_distance(self) -> None: create_items() distance = CosineDistance('half_embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_halfvec_l1_distance(self) -> None: create_items() distance = L1Distance('half_embedding', [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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, 3] @@ -248,14 +250,14 @@ def test_bit(self) -> None: def test_bit_hamming_distance(self) -> None: create_items() distance = HammingDistance('binary_embedding', '101') - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] def test_bit_jaccard_distance(self) -> None: create_items() distance = JaccardDistance('binary_embedding', '101') - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] # assert [v.distance for v in items] == [0, 1/3, 1] @@ -267,67 +269,67 @@ def test_sparsevec(self) -> None: def test_sparsevec_l2_distance(self) -> None: create_items() distance = L2Distance('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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)] def test_sparsevec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] def test_sparsevec_cosine_distance(self) -> None: create_items() distance = CosineDistance('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] def test_sparsevec_l1_distance(self) -> None: create_items() distance = L1Distance('sparse_embedding', SparseVector([1, 1, 1])) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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, 3] def test_filter(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) - items = Item.objects.alias(distance=distance).filter(distance__lt=1) # type: ignore + items = Item.objects.alias(distance=distance).filter(distance__lt=1) assert [v.id for v in items] == [1] def test_vector_avg(self) -> None: - avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] # type: ignore + avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert avg is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() - avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] # type: ignore + avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert avg == [2.5, 3.5, 4.5] def test_vector_sum(self) -> None: - sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] # type: ignore + sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert sum is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() - sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] # type: ignore + sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert sum == [5, 7, 9] def test_halfvec_avg(self) -> None: - avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] # type: ignore + avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] assert avg is None Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() - avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] # type: ignore + avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] assert avg == [2.5, 3.5, 4.5] def test_halfvec_sum(self) -> None: - sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] # type: ignore + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] assert sum is None Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() - sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] # type: ignore + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] assert sum == [5, 7, 9] def test_serialization(self) -> None: From 95cd7181510138ea4c11e24a19e6f94658368559 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:48:14 -0700 Subject: [PATCH 809/879] Improved typing for Django tests [skip ci] --- tests/test_django.py | 66 ++++++++++++++++++++++---------------------- 1 file changed, 33 insertions(+), 33 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 5bb51ae..3e1a366 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -211,7 +211,7 @@ def test_vector_l1_distance(self) -> None: def test_halfvec(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) assert item.half_embedding == [1, 2, 3] def test_halfvec_l2_distance(self) -> None: @@ -244,7 +244,7 @@ def test_halfvec_l1_distance(self) -> None: def test_bit(self) -> None: Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) assert item.binary_embedding == '101' def test_bit_hamming_distance(self) -> None: @@ -263,7 +263,7 @@ def test_bit_jaccard_distance(self) -> None: def test_sparsevec(self) -> None: Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self) -> None: @@ -334,7 +334,7 @@ def test_halfvec_sum(self) -> None: def test_serialization(self) -> None: create_items() - items = Item.objects.all() # type: ignore + items = Item.objects.all() for format in ['json', 'xml']: data = serializers.serialize(format, items) with mock.patch('django.core.serializers.python.apps.get_model') as get_model: @@ -349,26 +349,26 @@ def test_vector_form(self) -> None: def test_vector_form_instance(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = VectorForm(instance=item) assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() # type: ignore def test_vector_form_save(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = VectorForm(instance=item, data={'embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).embedding == [4, 5, 6] # type: ignore + assert Item.objects.get(pk=1).embedding == [4, 5, 6] def test_vector_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = VectorForm(instance=item, data={'embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).embedding is None # type: ignore + assert Item.objects.get(pk=1).embedding is None def test_halfvec_form(self) -> None: form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) @@ -377,26 +377,26 @@ def test_halfvec_form(self) -> None: def test_halfvec_form_instance(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item) assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() # type: ignore def test_halfvec_form_save(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).half_embedding == [4, 5, 6] # type: ignore + assert Item.objects.get(pk=1).half_embedding == [4, 5, 6] def test_halfvec_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).half_embedding is None # type: ignore + assert Item.objects.get(pk=1).half_embedding is None def test_bit_form(self) -> None: form = BitForm(data={'binary_embedding': '101'}) @@ -405,26 +405,26 @@ def test_bit_form(self) -> None: def test_bit_form_instance(self) -> None: Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = BitForm(instance=item) assert 'value="101"' in form.as_div() # type: ignore def test_bit_form_save(self) -> None: Item(id=1, binary_embedding='101').save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = BitForm(instance=item, data={'binary_embedding': '010'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert '010' == Item.objects.get(pk=1).binary_embedding # type: ignore + assert '010' == Item.objects.get(pk=1).binary_embedding def test_bit_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = BitForm(instance=item, data={'binary_embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).binary_embedding is None # type: ignore + assert Item.objects.get(pk=1).binary_embedding is None def test_sparsevec_form(self) -> None: form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) @@ -433,41 +433,41 @@ def test_sparsevec_form(self) -> None: def test_sparsevec_form_instance(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item) # TODO improve assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() # type: ignore def test_sparsevec_form_save(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item, data={'sparse_embedding': '{1:4,2:5,3:6}/3'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) # type: ignore + assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) def test_sparesevec_form_save_missing(self) -> None: Item(id=1).save() - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item, data={'sparse_embedding': ''}) assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).sparse_embedding is None # type: ignore + assert Item.objects.get(pk=1).sparse_embedding is None def test_clean(self) -> None: item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector([1, 2, 3])) item.full_clean() def test_get_or_create(self) -> None: - Item.objects.get_or_create(embedding=[1, 2, 3]) # type: ignore + Item.objects.get_or_create(embedding=[1, 2, 3]) def test_missing(self) -> None: Item().save() - assert Item.objects.first().embedding is None # type: ignore - assert Item.objects.first().half_embedding is None # type: ignore - assert Item.objects.first().binary_embedding is None # type: ignore - assert Item.objects.first().sparse_embedding is None # type: ignore + assert Item.objects.first().embedding is None + 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) -> None: Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]]).save() @@ -477,7 +477,7 @@ def test_vector_array(self) -> None: register_vector(cursor.connection) # this fails if the driver does not cast arrays - item = Item.objects.get(pk=1) # type: ignore + item = Item.objects.get(pk=1) assert item.embeddings == [[1, 2, 3], [4, 5, 6]] def test_double_array(self) -> None: @@ -485,7 +485,7 @@ def test_double_array(self) -> None: 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) # type: ignore + 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] @@ -495,7 +495,7 @@ def test_numeric_array(self) -> None: 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) # type: ignore + 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] @@ -503,6 +503,6 @@ def test_numeric_array(self) -> None: def test_half_precision(self) -> None: create_items() distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [1, 1, 1]) - items = Item.objects.annotate(distance=distance).order_by(distance) # type: ignore + 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)] From b7bc8ca979928143c7b6d6cdf755c0a6f426c223 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:49:47 -0700 Subject: [PATCH 810/879] Improved typing for Django tests [skip ci] --- tests/test_django.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_django.py b/tests/test_django.py index 3e1a366..6572867 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -345,13 +345,13 @@ def test_serialization(self) -> None: def test_vector_form(self) -> None: form = VectorForm(data={'embedding': '[1, 2, 3]'}) assert form.is_valid() - assert 'value="[1, 2, 3]"' in form.as_div() # type: ignore + assert 'value="[1, 2, 3]"' in str(form.as_div()) def test_vector_form_instance(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item) - assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() # type: ignore + assert 'value="[1.0, 2.0, 3.0]"' in str(form.as_div()) def test_vector_form_save(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() @@ -373,13 +373,13 @@ def test_vector_form_save_missing(self) -> None: def test_halfvec_form(self) -> None: form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) assert form.is_valid() - assert 'value="[1, 2, 3]"' in form.as_div() # type: ignore + assert 'value="[1, 2, 3]"' in str(form.as_div()) def test_halfvec_form_instance(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item) - assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() # type: ignore + assert 'value="[1.0, 2.0, 3.0]"' in str(form.as_div()) def test_halfvec_form_save(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() @@ -401,13 +401,13 @@ def test_halfvec_form_save_missing(self) -> None: def test_bit_form(self) -> None: form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() - assert 'value="101"' in form.as_div() # type: ignore + assert 'value="101"' in str(form.as_div()) def test_bit_form_instance(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) form = BitForm(instance=item) - assert 'value="101"' in form.as_div() # type: ignore + assert 'value="101"' in str(form.as_div()) def test_bit_form_save(self) -> None: Item(id=1, binary_embedding='101').save() @@ -429,14 +429,14 @@ def test_bit_form_save_missing(self) -> None: def test_sparsevec_form(self) -> None: form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) assert form.is_valid() - assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() # type: ignore + assert 'value="{1:1,2:2,3:3}/3"' in str(form.as_div()) def test_sparsevec_form_instance(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item) # TODO improve - assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() # type: ignore + assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in str(form.as_div()) def test_sparsevec_form_save(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() From 9d171fe77ccf58bbc42a91ec60a4184079610c2a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 16:58:07 -0700 Subject: [PATCH 811/879] Added comment [skip ci] --- tests/test_sqlmodel.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index b5f21d5..3eb09c8 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -7,6 +7,9 @@ engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: + # typing issue + # https://github.com/fastapi/sqlmodel/issues/909 + # https://github.com/fastapi/sqlmodel/pull/1657 session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) # type: ignore From 7a89b94040553e0df22fc16c5c60275f10b0803c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:06:29 -0700 Subject: [PATCH 812/879] Improved typing for Peewee tests --- tests/test_peewee.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 294389d..59910ad 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -2,7 +2,6 @@ from peewee import Model, PostgresqlDatabase, fn from pgvector import SparseVector, Vector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField -from typing import Any db = PostgresqlDatabase('pgvector_python_test') @@ -13,7 +12,7 @@ class Meta: class Item(BaseModel): - id: Any # for typing + id: int # for typing embedding = VectorField(dimensions=3, null=True) half_embedding = HalfVectorField(dimensions=3, null=True) binary_embedding = FixedBitField(max_length=3, null=True) From 1d8e9fc7877f257c25fdb25a118fcef6ddf58f4e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:12:29 -0700 Subject: [PATCH 813/879] Fixed type hint [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index bd32c1f..4205abf 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -105,7 +105,7 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: def _from_sparse(self, arr: sparray | spmatrix, /) -> None: value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore - shape = cast(tuple[int], value.shape) + shape = cast(tuple[int, ...], value.shape) if len(shape) == 1: self._dim = shape[0] elif len(shape) == 2 and shape[0] == 1: From dcbe3e7e1b6dd51f457a833035d7f06f7da6cd04 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:21:42 -0700 Subject: [PATCH 814/879] Improved pyright config [skip ci] --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4a9763b..29fedcd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,6 @@ include = [ ] typeCheckingMode = "strict" reportMissingTypeStubs = "none" -reportPossiblyUnboundVariable = "none" reportPrivateUsage = "none" reportUnknownArgumentType = "none" reportUnknownLambdaType = "none" From e766f92d33290f9bca28e3cb582dda58d9cb3e3e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:23:14 -0700 Subject: [PATCH 815/879] Improved pyright config [skip ci] --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 29fedcd..91f1257 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,7 +76,6 @@ reportUnknownArgumentType = "none" reportUnknownLambdaType = "none" reportUnknownMemberType = "none" reportUnknownVariableType = "none" -reportUnnecessaryIsInstance = "none" venv = ".venv" venvPath = "." From 0a5356c795498682ff21db329144b60c7ec971af Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:24:04 -0700 Subject: [PATCH 816/879] Improved pyright config [skip ci] --- pyproject.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 91f1257..bfb7361 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,8 +76,6 @@ reportUnknownArgumentType = "none" reportUnknownLambdaType = "none" reportUnknownMemberType = "none" reportUnknownVariableType = "none" -venv = ".venv" -venvPath = "." [tool.pyrefly] project-includes = [ From 3cb5fee5598bcdb57211ce1088f7115b1cc6286b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:46:08 -0700 Subject: [PATCH 817/879] Updated lint task [skip ci] --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 7462142..7e5aab0 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ .PHONY: lint build publish clean lint: - pycodestyle . --ignore=E501 --exclude=.venv + pycodestyle . --ignore=E501,E704,W504 --exclude=.venv build: python3 -m build From 663eff8d395d1d6c19b8d9df567e4e67dd492f6c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 17:51:47 -0700 Subject: [PATCH 818/879] Updated comments [skip ci] --- tests/test_asyncpg.py | 1 + tests/test_sqlalchemy.py | 1 + tests/test_sqlmodel.py | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 1b07eff..38e5fd0 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -65,6 +65,7 @@ async def test_bit(self) -> None: await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') + # typing issue # https://github.com/MagicStack/py-pgproto/pull/32 embedding = asyncpg.BitString('101') # type: ignore embedding2 = None diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index e8f221f..2df5e34 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -638,6 +638,7 @@ async def test_bit(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): + # typing issue # https://github.com/MagicStack/py-pgproto/pull/32 embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' # type: ignore session.add(Item(id=1, binary_embedding=embedding)) diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 3eb09c8..853384e 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -18,7 +18,7 @@ class Item(SQLModel, table=True): model_config = ConfigDict(arbitrary_types_allowed=True) # type: ignore id: int | None = Field(default=None, primary_key=True) - # typing issues + # typing issue # https://github.com/fastapi/sqlmodel/discussions/1228 # https://github.com/fastapi/sqlmodel/pull/1345 embedding: list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore From 8862287b4bed87154fd2931787a3de4ce6669669 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 3 Jul 2026 20:03:53 -0700 Subject: [PATCH 819/879] Updated dependency groups --- .github/workflows/build.yml | 2 +- pyproject.toml | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c215202..952be69 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -25,7 +25,7 @@ jobs: sudo make install - run: pytest - - run: pip install --group dev-optional + - run: pip install --group numpy --group scipy - run: pytest - run: pip install --group typing diff --git a/pyproject.toml b/pyproject.toml index bfb7361..9c7cd15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,14 +31,13 @@ dev = [ "SQLAlchemy[asyncio]>=2", "sqlmodel" ] -dev-optional = [ - "numpy", - "scipy" -] +numpy = ["numpy"] +scipy = ["scipy"] typing = [ "mypy", "ty", - {include-group = "dev-optional"} + "numpy", + "scipy" ] [tool.pytest.ini_options] From f43a82827a7d8209e7cdee6cbb50946c163c8bd9 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:20:41 -0700 Subject: [PATCH 820/879] Ensure values are always floats for SparseVector --- pgvector/sparsevec.py | 2 +- tests/test_sparse_vector.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 4205abf..3631f78 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -118,7 +118,7 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: self._indices = value.coords[-1].tolist() else: self._indices = value.col.tolist() - self._values = value.data.tolist() + self._values = [float(v) for v in value.data] def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 7d5f5d1..1924f98 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -58,6 +58,7 @@ def test_coo_array(self) -> None: vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + assert isinstance(vec.values()[0], float) def test_coo_array_dimensions(self) -> None: if np is None or sparse is None: From bddb17662c633d7416b2c5686c2108a904a43925 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:25:22 -0700 Subject: [PATCH 821/879] Updated readme [skip ci] --- README.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 24f72bb..ee6c5f5 100644 --- a/README.md +++ b/README.md @@ -473,7 +473,9 @@ conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3)) Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) ``` @@ -519,7 +521,9 @@ cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))' Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) ``` @@ -574,7 +578,9 @@ await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vect Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding) ``` @@ -619,7 +625,9 @@ conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding) ``` From eb33d03364d44e45242ae440ae3ef322f8d69a89 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:28:10 -0700 Subject: [PATCH 822/879] Updated readme [skip ci] --- README.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/README.md b/README.md index ee6c5f5..33db4bd 100644 --- a/README.md +++ b/README.md @@ -706,6 +706,32 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc ## Reference +### Vectors + +Create a vector from a list + +```python +vec = Vector([1, 2, 3]) +``` + +Or a NumPy array + +```python +vec = Vector(np.array([1, 2, 3])) +``` + +Get a list + +```python +lst = vec.to_list() +``` + +Get a NumPy array + +```python +arr = vec.to_numpy() +``` + ### Half Vectors Create a half vector from a list From 28191fc8b9a5702cd07b2854dbf89b1e0d263100 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:36:58 -0700 Subject: [PATCH 823/879] Removed todo [skip ci] --- tests/test_django.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_django.py b/tests/test_django.py index 6572867..e5f0c6f 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -435,7 +435,6 @@ def test_sparsevec_form_instance(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item) - # TODO improve assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in str(form.as_div()) def test_sparsevec_form_save(self) -> None: From 0c18e0aa5de9d9a4ff1bd1ad58c8839bb5ce814c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:39:29 -0700 Subject: [PATCH 824/879] Addressed todo --- tests/test_sqlalchemy.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 2df5e34..3dd7bac 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -176,13 +176,9 @@ def test_orm(self, engine: Engine) -> None: session.add(item3) session.commit() - stmt = select(Item) + stmt = select(Item).order_by(Item.id) with Session(engine) as session: items = [v[0] for v in session.execute(stmt).all()] - # TODO improve - assert items[0].id % 3 == 1 - assert items[1].id % 3 == 2 - assert items[2].id % 3 == 0 assert items[0].embedding == [1.5, 2, 3] assert items[1].embedding == [4, 5, 6] assert items[2].embedding is None From dc6877060099dcd2485e23c9facb04d3e6df3faf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:47:03 -0700 Subject: [PATCH 825/879] Updated readme [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 33db4bd..12b76c1 100644 --- a/README.md +++ b/README.md @@ -310,7 +310,7 @@ from pgvector.sqlalchemy import VECTOR from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(VECTOR(3))) + embeddings = mapped_column(ARRAY(VECTOR(3), dimensions=1)) ``` And register the types with the underlying driver From f6e113e1ad10968473e630d001bc473b6c9a62ed Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:49:02 -0700 Subject: [PATCH 826/879] Updated example [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 12b76c1..86a72c7 100644 --- a/README.md +++ b/README.md @@ -362,7 +362,7 @@ Add a vector column from pgvector.sqlalchemy import VECTOR class Item(SQLModel, table=True): - embedding: Any = Field(sa_type=VECTOR(3)) + embedding: list[float] = Field(sa_type=VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` From 9f0d6b97640486eacf858143430b40d11634a845 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:53:03 -0700 Subject: [PATCH 827/879] Updated examples [skip ci] --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 86a72c7..26dccfb 100644 --- a/README.md +++ b/README.md @@ -180,7 +180,7 @@ Add a vector column from pgvector.sqlalchemy import VECTOR class Item(Base): - embedding = mapped_column(VECTOR(3)) + embedding: Mapped[list[float]] = mapped_column(VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -310,7 +310,7 @@ from pgvector.sqlalchemy import VECTOR from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(VECTOR(3), dimensions=1)) + embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1)) ``` And register the types with the underlying driver From 780cc73f723fd2300cd16ae7362c9ab00453ecfc Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 13:57:56 -0700 Subject: [PATCH 828/879] Updated example [skip ci] --- examples/loading/example.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/loading/example.py b/examples/loading/example.py index 7f3dce8..4c0e579 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -1,4 +1,5 @@ import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -25,7 +26,7 @@ copy.set_types(['vector']) for i, embedding in enumerate(embeddings): - copy.write_row([embedding]) + copy.write_row([Vector(embedding)]) # show progress if i % 10000 == 0: From efa144252a96be79c695923d2717b0555bea8fcd Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 14:03:59 -0700 Subject: [PATCH 829/879] Updated examples [skip ci] --- examples/lightfm/example.py | 14 +++++++------- examples/openai/halfvec.py | 3 ++- examples/surprise/example.py | 10 +++++----- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/examples/lightfm/example.py b/examples/lightfm/example.py index c869e89..728adc1 100644 --- a/examples/lightfm/example.py +++ b/examples/lightfm/example.py @@ -2,7 +2,7 @@ from lightfm.datasets import fetch_movielens from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Float, Integer, String -from sqlalchemy.orm import mapped_column, DeclarativeBase, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: @@ -17,17 +17,17 @@ class Base(DeclarativeBase): class User(Base): __tablename__ = 'user' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(VECTOR(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' - id = mapped_column(Integer, primary_key=True) - title = mapped_column(String) - factors = mapped_column(VECTOR(20)) - bias = mapped_column(Float) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + title: Mapped[str] = mapped_column(String) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) + bias: Mapped[float] = mapped_column(Float) Base.metadata.drop_all(engine) diff --git a/examples/openai/halfvec.py b/examples/openai/halfvec.py index 185c785..b2234f6 100644 --- a/examples/openai/halfvec.py +++ b/examples/openai/halfvec.py @@ -1,5 +1,6 @@ from openai import OpenAI -from pgvector.psycopg import register_vector, HalfVector +from pgvector import HalfVector +from pgvector.psycopg import register_vector import psycopg conn = psycopg.connect(dbname='pgvector_example', autocommit=True) diff --git a/examples/surprise/example.py b/examples/surprise/example.py index 6a40784..699ea31 100644 --- a/examples/surprise/example.py +++ b/examples/surprise/example.py @@ -1,6 +1,6 @@ from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer -from sqlalchemy.orm import mapped_column, DeclarativeBase, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session from surprise import Dataset, SVD engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -16,15 +16,15 @@ class Base(DeclarativeBase): class User(Base): __tablename__ = 'user' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(VECTOR(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(VECTOR(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) From 45b054c53fb174ca12ecb08157f7875bd4795356 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 14:08:52 -0700 Subject: [PATCH 830/879] Updated example [skip ci] --- examples/implicit/example.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/implicit/example.py b/examples/implicit/example.py index ebb1ff1..a4bd0f4 100644 --- a/examples/implicit/example.py +++ b/examples/implicit/example.py @@ -2,7 +2,7 @@ from implicit.datasets.movielens import get_movielens from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer, String -from sqlalchemy.orm import mapped_column, DeclarativeBase, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: @@ -17,16 +17,16 @@ class Base(DeclarativeBase): class User(Base): __tablename__ = 'user' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(VECTOR(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' - id = mapped_column(Integer, primary_key=True) - title = mapped_column(String) - factors = mapped_column(VECTOR(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + title: Mapped[str] = mapped_column(String) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) From a810f27ca93c185e81e98b05d87622d736e89568 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 5 Jul 2026 14:09:49 -0700 Subject: [PATCH 831/879] Updated example [skip ci] --- examples/citus/example.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/citus/example.py b/examples/citus/example.py index 915c25f..d1af5b9 100644 --- a/examples/citus/example.py +++ b/examples/citus/example.py @@ -1,4 +1,5 @@ import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -38,7 +39,7 @@ copy.set_types(['vector', 'bigint']) for i in range(rows): - copy.write_row([embeddings[i], categories[i]]) + copy.write_row([Vector(embeddings[i]), categories[i]]) print('Creating index in parallel') conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') From 60739dfd6cb9d674f32afa4184d43e6aff9dfbcf Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 6 Jul 2026 11:25:13 -0700 Subject: [PATCH 832/879] Version bump to 0.5.0 [skip ci] --- CHANGELOG.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f7598df..e509532 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.5.0 (unreleased) +## 0.5.0 (2026-07-06) - Added experimental support for type hints - Changed `vector` and `halfvec` types to return list for Django, SQLAlchemy, SQLModel, and Peewee diff --git a/pyproject.toml b/pyproject.toml index 9c7cd15..ad130b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.4.2" +version = "0.5.0" description = "pgvector support for Python" readme = "README.md" authors = [ From 86f1c40c1033ed87acdd6ce199540ba2d0ee0423 Mon Sep 17 00:00:00 2001 From: Jiahui Miao <130162274+HUIIIM@users.noreply.github.com> Date: Fri, 2 Oct 2026 00:31:58 -0400 Subject: [PATCH 833/879] Fix dtype check for numpy boolean array (#158) --- pgvector/bit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 2b04912..dfcf10b 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -39,7 +39,7 @@ def __init__( elif is_ndarray(value): import numpy as np - if value.dtype != np.bool: + if value.dtype != np.bool_: # skip error for result of np.unpackbits if value.dtype != np.uint8 or np.any(value > 1): raise ValueError('expected elements to be boolean') From e07659f531e2761d6a4dfa65cbe2ae0f5389f1da Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 1 Oct 2026 21:35:44 -0700 Subject: [PATCH 834/879] Updating typing for NumPy 1.24 to 1.26 - #158 [skip ci] --- pgvector/bit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index dfcf10b..5340d55 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -13,7 +13,7 @@ class Bit: def __init__( self, - value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]], + value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool_ | np.uint8]], / ) -> None: if isinstance(value, bytes): @@ -65,7 +65,7 @@ def to_list(self) -> list[bool]: # TODO improve return [v != '0' for v in self.to_text()] - def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]: + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool_]]: import numpy as np return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool) From 99a6776a2a8458ca2f8a9fd804fde205fff405ae Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 1 Oct 2026 21:45:16 -0700 Subject: [PATCH 835/879] Updated changelog [skip ci] --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e509532..e68bcae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.5.1 (unreleased) + +- Fixed error with `Bit` constructor for NumPy 1.24-1.26 + ## 0.5.0 (2026-07-06) - Added experimental support for type hints From fd74f7842cc28ef5030ba7c9627e1ce66148234e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 3 Oct 2026 13:34:17 -0700 Subject: [PATCH 836/879] Fixed error with unordered SciPy sparse arrays - resolves #160 Co-authored-by: Christopher Krah --- CHANGELOG.md | 1 + pgvector/sparsevec.py | 10 +++++++--- tests/test_sparse_vector.py | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e68bcae..20229cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.5.1 (unreleased) +- Fixed error with unordered SciPy sparse arrays - Fixed error with `Bit` constructor for NumPy 1.24-1.26 ## 0.5.0 (2026-07-06) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 3631f78..5287ef2 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -115,10 +115,14 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: if hasattr(value, 'coords'): # scipy 1.13+ - self._indices = value.coords[-1].tolist() + indices = value.coords[-1].tolist() else: - self._indices = value.col.tolist() - self._values = [float(v) for v in value.data] + indices = value.col.tolist() + + elements = sorted(zip(indices, value.data)) + + self._indices = [v[0] for v in elements] + self._values = [float(v[1]) for v in elements] def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 1924f98..cef3c62 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -54,7 +54,7 @@ def test_coo_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') - arr = sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])) + arr = sparse.coo_array(([3, 1, 2], ([4, 0, 2],)), shape=(6,)) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] From d3c5d4b780730472a69420ba0b6607bd58d13e45 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:04:45 -0700 Subject: [PATCH 837/879] Improved performance by checking has_canonical_format [skip ci] --- pgvector/sparsevec.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 5287ef2..a7c08dc 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -119,10 +119,14 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: else: indices = value.col.tolist() - elements = sorted(zip(indices, value.data)) - - self._indices = [v[0] for v in elements] - self._values = [float(v[1]) for v in elements] + # spipy 1.12+ + if hasattr(value, 'has_canonical_format') and value.has_canonical_format: + self._indices = indices + self._values = [float(v) for v in value.data] + else: + elements = sorted(zip(indices, value.data)) + self._indices = [v[0] for v in elements] + self._values = [float(v[1]) for v in elements] def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) From cd18c3f1074e08a9bbe5d72c98f5c838e369caab Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:13:46 -0700 Subject: [PATCH 838/879] Fixed error with SciPy sparse arrays with zeros --- pgvector/sparsevec.py | 16 ++++++++-------- tests/test_sparse_vector.py | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index a7c08dc..3df7c67 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -119,14 +119,14 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: else: indices = value.col.tolist() - # spipy 1.12+ - if hasattr(value, 'has_canonical_format') and value.has_canonical_format: - self._indices = indices - self._values = [float(v) for v in value.data] - else: - elements = sorted(zip(indices, value.data)) - self._indices = [v[0] for v in elements] - self._values = [float(v[1]) for v in elements] + elements = [v for v in zip(indices, value.data) if v[1] != 0] + + # has_canonical_format added in spipy 1.12+ + if not hasattr(value, 'has_canonical_format') or not value.has_canonical_format: + elements.sort() + + self._indices = [v[0] for v in elements] + self._values = [float(v[1]) for v in elements] def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index cef3c62..b663696 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -54,7 +54,7 @@ def test_coo_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') - arr = sparse.coo_array(([3, 1, 2], ([4, 0, 2],)), shape=(6,)) + arr = sparse.coo_array(([2, 3, 1, 0], ([2, 4, 0, 3],)), shape=(6,)) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] From 63f5a483777209e30420e02da13abe46cbc4f2af Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:18:23 -0700 Subject: [PATCH 839/879] Fixed typo [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 3df7c67..7a20d40 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -121,7 +121,7 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: elements = [v for v in zip(indices, value.data) if v[1] != 0] - # has_canonical_format added in spipy 1.12+ + # has_canonical_format added in scipy 1.12+ if not hasattr(value, 'has_canonical_format') or not value.has_canonical_format: elements.sort() From c451fb017d479d9ce9af7e08c4ffdd06f7210bc4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:19:46 -0700 Subject: [PATCH 840/879] Improved code consistency [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 7a20d40..5844ad7 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -119,7 +119,7 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: else: indices = value.col.tolist() - elements = [v for v in zip(indices, value.data) if v[1] != 0] + elements = [(i, v) for i, v in zip(indices, value.data) if v != 0] # has_canonical_format added in scipy 1.12+ if not hasattr(value, 'has_canonical_format') or not value.has_canonical_format: From f4dde1961e81db596660cdc619eb7609667bb9b3 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:40:14 -0700 Subject: [PATCH 841/879] Fixed handling of duplicates for SciPy sparse arrays - fixes #162 --- pgvector/sparsevec.py | 5 +++++ tests/test_sparse_vector.py | 10 ++++++++++ 2 files changed, 15 insertions(+) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 5844ad7..0d3434c 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -105,6 +105,11 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: def _from_sparse(self, arr: sparray | spmatrix, /) -> None: value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore + # has_canonical_format added in scipy 1.12+ + if not hasattr(value, 'has_canonical_format') or not value.has_canonical_format: + value = value.copy() + value.sum_duplicates() + shape = cast(tuple[int, ...], value.shape) if len(shape) == 1: self._dim = shape[0] diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index b663696..262e23a 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -60,6 +60,16 @@ def test_coo_array(self) -> None: assert vec.indices() == [0, 2, 4] assert isinstance(vec.values()[0], float) + def test_coo_array_duplicates(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.coo_array(([1, 2], ([1, 1],)), shape=(3,)) + assert arr.todense().tolist() == [0, 3, 0] + vec = SparseVector(arr) + assert vec.to_list() == [0, 3, 0] + assert vec.indices() == [1] + def test_coo_array_dimensions(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') From f3a05b42a9161f6f30fec9e31e161a21a2e2f3d2 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:53:23 -0700 Subject: [PATCH 842/879] Updated changelog [skip ci] --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 20229cb..b74cdf5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ ## 0.5.1 (unreleased) -- Fixed error with unordered SciPy sparse arrays +- Fixed error with non-canonical SciPy sparse arrays - Fixed error with `Bit` constructor for NumPy 1.24-1.26 ## 0.5.0 (2026-07-06) From 90056379c41574db1c0df367f9e02e79644ed995 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 11:55:44 -0700 Subject: [PATCH 843/879] Updated changelog [skip ci] --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b74cdf5..cd577c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.5.1 (unreleased) - Fixed error with non-canonical SciPy sparse arrays +- Fixed error with SciPy sparse arrays with explicit zeros - Fixed error with `Bit` constructor for NumPy 1.24-1.26 ## 0.5.0 (2026-07-06) From 474522bafad1385c4bdda439700f39226b309ec0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sun, 4 Oct 2026 12:34:28 -0700 Subject: [PATCH 844/879] Improved validation for bit strings - fixes #161 --- pgvector/bit.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 5340d55..a157b30 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -26,16 +26,15 @@ def __init__( value = ''.join([bits[v] for v in value]) except (KeyError, TypeError): raise ValueError('expected list[bool]') + elif not set(value).issubset({'0', '1'}): + raise ValueError('expected bit string') length = len(value) if length % 8 != 0: value += '0' * (8 - (length % 8)) self._length = length - try: - self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') - except ValueError: - raise ValueError('expected bit string') + self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') elif is_ndarray(value): import numpy as np From 5275164edb3a2847d4a5aec1af5ae90f7dc47b53 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 8 Oct 2026 18:42:26 -0700 Subject: [PATCH 845/879] Version bump to 0.5.1 [skip ci] --- CHANGELOG.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cd577c8..666ca29 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,4 @@ -## 0.5.1 (unreleased) +## 0.5.1 (2026-10-08) - Fixed error with non-canonical SciPy sparse arrays - Fixed error with SciPy sparse arrays with explicit zeros diff --git a/pyproject.toml b/pyproject.toml index ad130b5..3ff09d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.5.0" +version = "0.5.1" description = "pgvector support for Python" readme = "README.md" authors = [ From 62d9388d2572c21e6abc0de6023c62cf58b40251 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 10 Oct 2026 14:23:22 -0700 Subject: [PATCH 846/879] Fixed tests for SQLAlchemy 2.1 [skip ci] --- tests/test_sqlalchemy.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 3dd7bac..eef30b3 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -12,6 +12,12 @@ from typing import Any from .conftest import numpy as np +try: + from sqlalchemy.dialects.postgresql import BitString + sqlalchemy_version = 2.1 +except ImportError: + sqlalchemy_version = 2 + psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -420,7 +426,7 @@ def test_select(self, engine: Engine) -> None: def test_select_orm(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() + items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() # type: ignore assert items == [3] def test_avg(self, engine: Engine) -> None: @@ -634,9 +640,15 @@ async def test_bit(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): - # typing issue - # https://github.com/MagicStack/py-pgproto/pull/32 - embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' # type: ignore + embedding: Any + if sqlalchemy_version >= 2.1: + embedding = BitString('101') + elif engine == asyncpg_engine: + # typing issue + # https://github.com/MagicStack/py-pgproto/pull/32 + embedding = asyncpg.BitString('101') # type: ignore + else: + embedding = '101' session.add(Item(id=1, binary_embedding=embedding)) item = await session.get_one(Item, 1) assert item.binary_embedding == embedding From 0c33b6dd187e25fdd6048781999b09a2f11fa28b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 10 Oct 2026 14:23:43 -0700 Subject: [PATCH 847/879] Updated CI --- .github/workflows/build.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 952be69..5a01635 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,10 +6,10 @@ jobs: strategy: fail-fast: false matrix: - python: [3.14, "3.10"] + python: [3.15, "3.10"] steps: - uses: actions/checkout@v7 - - uses: actions/setup-python@v6 + - uses: actions/setup-python@v7 with: python-version: ${{ matrix.python }} - run: pip install --group dev @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.8.4 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.7 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install From ffdcea9ebe686fdc5a17e8bb2c16c0805975dfc0 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 10 Oct 2026 14:37:13 -0700 Subject: [PATCH 848/879] Fixed CI --- tests/test_sqlalchemy.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index eef30b3..dc58aeb 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -4,6 +4,7 @@ from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, literal, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine +from sqlalchemy.dialects import postgresql from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncEngine from sqlalchemy.ext.automap import automap_base @@ -12,11 +13,7 @@ from typing import Any from .conftest import numpy as np -try: - from sqlalchemy.dialects.postgresql import BitString - sqlalchemy_version = 2.1 -except ImportError: - sqlalchemy_version = 2 +sqlalchemy_version = 2.1 if hasattr(postgresql, 'BitString') else 2 psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -642,7 +639,7 @@ async def test_bit(self, engine: AsyncEngine) -> None: async with session.begin(): embedding: Any if sqlalchemy_version >= 2.1: - embedding = BitString('101') + embedding = postgresql.BitString('101') # type: ignore elif engine == asyncpg_engine: # typing issue # https://github.com/MagicStack/py-pgproto/pull/32 From 57b8e24ed659cba1bcd646b3abddea5a3855d63e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 10 Oct 2026 14:47:59 -0700 Subject: [PATCH 849/879] Fixed typing --- pgvector/sqlalchemy/bit.py | 4 +--- pgvector/sqlalchemy/halfvec.py | 4 +--- pgvector/sqlalchemy/sparsevec.py | 4 +--- pgvector/sqlalchemy/vector.py | 4 +--- 4 files changed, 4 insertions(+), 12 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 7c861ef..56b5dcc 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -14,11 +14,9 @@ def process_bind_param(self, value: Any, dialect: Any) -> Any: return asyncpg.BitString(value) # type: ignore return value - class Comparator(TypeDecorator.Comparator[Any]): + class comparator_factory(TypeDecorator.Comparator[Any]): def hamming_distance(self, other: object, /) -> Operators: return self.op('<~>', return_type=Float)(other) def jaccard_distance(self, other: object, /) -> Operators: return self.op('<%>', return_type=Float)(other) - - comparator_factory = Comparator # type: ignore diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 9d717c3..e726e0f 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return HalfVector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class comparator_factory(TypeEngine.Comparator[Any]): def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,8 +48,6 @@ def cosine_distance(self, other: object, /) -> Operators: def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) - comparator_factory = Comparator - # for reflection ischema_names['halfvec'] = HALFVEC # type: ignore diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 9bd1478..2c94104 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class comparator_factory(TypeEngine.Comparator[Any]): def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,8 +48,6 @@ def cosine_distance(self, other: object, /) -> Operators: def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) - comparator_factory = Comparator - # for reflection ischema_names['sparsevec'] = SPARSEVEC # type: ignore diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 83acb09..b86c91d 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return Vector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class comparator_factory(TypeEngine.Comparator[Any]): def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,8 +48,6 @@ def cosine_distance(self, other: object, /) -> Operators: def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) - comparator_factory = Comparator - # for reflection ischema_names['vector'] = VECTOR # type: ignore From 40ef30fa67373ac613c5764e240a2c8806555b94 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 10 Oct 2026 15:05:10 -0700 Subject: [PATCH 850/879] Revert "Fixed typing" This reverts commit 57b8e24ed659cba1bcd646b3abddea5a3855d63e. --- pgvector/sqlalchemy/bit.py | 4 +++- pgvector/sqlalchemy/halfvec.py | 4 +++- pgvector/sqlalchemy/sparsevec.py | 4 +++- pgvector/sqlalchemy/vector.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 56b5dcc..7c861ef 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -14,9 +14,11 @@ def process_bind_param(self, value: Any, dialect: Any) -> Any: return asyncpg.BitString(value) # type: ignore return value - class comparator_factory(TypeDecorator.Comparator[Any]): + class Comparator(TypeDecorator.Comparator[Any]): def hamming_distance(self, other: object, /) -> Operators: return self.op('<~>', return_type=Float)(other) def jaccard_distance(self, other: object, /) -> Operators: return self.op('<%>', return_type=Float)(other) + + comparator_factory = Comparator # type: ignore diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index e726e0f..9d717c3 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return HalfVector._from_db(value) return process - class comparator_factory(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,6 +48,8 @@ def cosine_distance(self, other: object, /) -> Operators: def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['halfvec'] = HALFVEC # type: ignore diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 2c94104..9bd1478 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class comparator_factory(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,6 +48,8 @@ def cosine_distance(self, other: object, /) -> Operators: def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['sparsevec'] = SPARSEVEC # type: ignore diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index b86c91d..83acb09 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return Vector._from_db(value) return process - class comparator_factory(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) @@ -48,6 +48,8 @@ def cosine_distance(self, other: object, /) -> Operators: def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection ischema_names['vector'] = VECTOR # type: ignore From 1cbee1ac5cee097d1ccb7ff811c18c003653ae75 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Sat, 10 Oct 2026 15:13:57 -0700 Subject: [PATCH 851/879] Fixed CI --- pgvector/sqlalchemy/bit.py | 2 +- pgvector/sqlalchemy/halfvec.py | 2 +- pgvector/sqlalchemy/sparsevec.py | 2 +- pgvector/sqlalchemy/vector.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 7c861ef..a5af3b0 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -14,7 +14,7 @@ def process_bind_param(self, value: Any, dialect: Any) -> Any: return asyncpg.BitString(value) # type: ignore return value - class Comparator(TypeDecorator.Comparator[Any]): + class Comparator(TypeDecorator.Comparator[Any]): # type: ignore def hamming_distance(self, other: object, /) -> Operators: return self.op('<~>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 9d717c3..d82c5ac 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return HalfVector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): # type: ignore def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 9bd1478..5cc910d 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): # type: ignore def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 83acb09..4e7e4b3 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return Vector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): # type: ignore def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) From 8cef6e4e867b8d07971da1a3d6170fb6e75e3c40 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 13 Oct 2026 13:11:57 -0700 Subject: [PATCH 852/879] Fixed equality for HalfVector with signed zeros - resolves #164 Co-authored-by: Zhuo Zhao <193034689+chaduo@users.noreply.github.com> --- CHANGELOG.md | 4 ++++ pgvector/halfvec.py | 2 +- tests/test_half_vector.py | 1 + tests/test_vector.py | 1 + 4 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 666ca29..9561ff5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.5.2 (unreleased) + +- Fixed equality for `HalfVector` with signed zeros + ## 0.5.1 (2026-10-08) - Fixed error with non-canonical SciPy sparse arrays diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 1a2c8c7..fa35194 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -37,7 +37,7 @@ def __repr__(self) -> str: def __eq__(self, other: object, /) -> bool: if not isinstance(other, self.__class__): return NotImplemented - return self._value == other._value + return self.to_list() == other.to_list() def dimensions(self) -> int: return len(self._value) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 26c5ac7..56bda26 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -55,6 +55,7 @@ def test_equality(self) -> None: assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) assert HalfVector([1, 2, 3]) != 1 + assert HalfVector([0.0]) == HalfVector([-0.0]) def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 diff --git a/tests/test_vector.py b/tests/test_vector.py index a08370e..c5b75d0 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -55,6 +55,7 @@ def test_equality(self) -> None: assert Vector([1, 2, 3]) == Vector([1, 2, 3]) assert Vector([1, 2, 3]) != Vector([1, 2, 4]) assert Vector([1, 2, 3]) != 1 + assert Vector([0.0]) == Vector([-0.0]) def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 From 2915a599ac1a22876fe5d05fdb55bc40feadaa36 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 14 Oct 2026 14:16:26 -0700 Subject: [PATCH 853/879] Fixed error with Bit constructor and empty strings - fixes #165 and resolves #166 Co-authored-by: Zaakir Hussain --- CHANGELOG.md | 1 + pgvector/bit.py | 5 ++++- tests/test_bit.py | 3 +++ 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9561ff5..e1c1f48 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.5.2 (unreleased) +- Fixed error with `Bit` constructor and empty strings - Fixed equality for `HalfVector` with signed zeros ## 0.5.1 (2026-10-08) diff --git a/pgvector/bit.py b/pgvector/bit.py index a157b30..098ebb2 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -34,7 +34,10 @@ def __init__( value += '0' * (8 - (length % 8)) self._length = length - self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + if length > 0: + self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + else: + self._data = b'' elif is_ndarray(value): import numpy as np diff --git a/tests/test_bit.py b/tests/test_bit.py index ed13374..2714ac2 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -27,6 +27,9 @@ def test_list_list(self) -> None: def test_str(self) -> None: assert Bit('101').to_list() == [True, False, True] + def test_str_empty(self) -> None: + assert Bit('').to_list() == [] + def test_str_two(self) -> None: with pytest.raises(ValueError) as error: Bit('201') From 7329d8a3973cfe9390296f12b471adcd6070e2da Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 14 Oct 2026 14:22:51 -0700 Subject: [PATCH 854/879] Added more tests for zero length bit [skip ci] --- tests/test_bit.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index 2714ac2..6487348 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -9,6 +9,9 @@ class TestBit: def test_list(self) -> None: assert Bit([True, False, True]).to_list() == [True, False, True] + def test_list_empty(self) -> None: + assert Bit([]).to_list() == [] + def test_list_none(self) -> None: with pytest.raises(ValueError) as error: Bit([True, None, True]) # type: ignore @@ -39,6 +42,9 @@ def test_bytes(self) -> None: assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' + def test_bytes_empty(self) -> None: + assert Bit(b'').to_text() == '' + def test_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') From 4003693d0f398cf87068bcac404a0d082291953a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 16 Oct 2026 13:44:16 -0700 Subject: [PATCH 855/879] Prepared validation for SparseVector indexes - resolves #167 [skip ci] Co-authored-by: AryanRG --- pgvector/sparsevec.py | 8 ++++++++ tests/test_sparse_vector.py | 10 ++++++++++ 2 files changed, 18 insertions(+) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 0d3434c..bf8f3b4 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -101,6 +101,8 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: self._dim = int(dim) self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] + # TODO uncomment in 0.6.0 + # self._validate() def _from_sparse(self, arr: sparray | spmatrix, /) -> None: value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore @@ -138,6 +140,10 @@ def _from_dense(self, value: list[float] | ndarray, /) -> None: self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] + def _validate(self): + if any(i < 0 or i >= self._dim for i in self._indices): + raise ValueError('sparsevec index out of bounds') + @classmethod def from_text(cls, value: str, /) -> SparseVector: elements, dim = value.split('/', 2) @@ -171,6 +177,8 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float], /) -> Sp vec._dim = dim vec._indices = indices vec._values = values + # TODO uncomment in 0.6.0 + # vec._validate() return vec @classmethod diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 262e23a..78ca985 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -45,6 +45,16 @@ def test_dict(self) -> None: def test_dict_empty(self) -> None: assert SparseVector({}, 0).to_list() == [] + # TODO uncomment in 0.6.0 + # def test_dict_negative_index(self) -> None: + # with pytest.raises(ValueError, match='sparsevec index out of bounds'): + # SparseVector({-1: 1}, 3) + + # TODO uncomment in 0.6.0 + # def test_dict_higher_index(self) -> None: + # with pytest.raises(ValueError, match='sparsevec index out of bounds'): + # SparseVector({3: 1}, 3) + def test_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore From 6ce28526a13590884e87f9613bbc058860e9c708 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 16 Oct 2026 13:46:40 -0700 Subject: [PATCH 856/879] Prepared validation for SparseVector dimensions [skip ci] --- pgvector/sparsevec.py | 3 +++ tests/test_sparse_vector.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index bf8f3b4..335a9ad 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -141,6 +141,9 @@ def _from_dense(self, value: list[float] | ndarray, /) -> None: self._values = [float(value[i]) for i in self._indices] def _validate(self): + if self._dim < 0: + raise ValueError('sparsevec cannot have negative dimensions') + if any(i < 0 or i >= self._dim for i in self._indices): raise ValueError('sparsevec index out of bounds') diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 78ca985..9d64c3d 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -55,6 +55,11 @@ def test_dict_empty(self) -> None: # with pytest.raises(ValueError, match='sparsevec index out of bounds'): # SparseVector({3: 1}, 3) + # TODO uncomment in 0.6.0 + # def test_dict_negative_dimensions(self) -> None: + # with pytest.raises(ValueError, match='sparsevec cannot have negative dimensions'): + # SparseVector({}, -1) + def test_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore From 8be9673245133041bb696c1071b153a20f88585b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 16 Oct 2026 22:16:24 -0700 Subject: [PATCH 857/879] Added typing [skip ci] --- pgvector/sparsevec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 335a9ad..6a477a9 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -140,7 +140,7 @@ def _from_dense(self, value: list[float] | ndarray, /) -> None: self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] - def _validate(self): + def _validate(self) -> None: if self._dim < 0: raise ValueError('sparsevec cannot have negative dimensions') From 063995bcfaeb58aa86b3ec6306987f09cce18f1b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 16 Oct 2026 22:24:54 -0700 Subject: [PATCH 858/879] Prepared validation for SparseVector order [skip ci] --- pgvector/sparsevec.py | 6 ++++++ tests/test_sparse_vector.py | 10 ++++++++++ 2 files changed, 16 insertions(+) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 6a477a9..c3ac7e5 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -147,6 +147,11 @@ def _validate(self) -> None: if any(i < 0 or i >= self._dim for i in self._indices): raise ValueError('sparsevec index out of bounds') + def _validate_sorted(self) -> None: + for i in range(len(self._indices) - 1): + if self._indices[i] >= self._indices[i + 1]: + raise ValueError('sparsevec indices must be in ascending order') + @classmethod def from_text(cls, value: str, /) -> SparseVector: elements, dim = value.split('/', 2) @@ -182,6 +187,7 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float], /) -> Sp vec._values = values # TODO uncomment in 0.6.0 # vec._validate() + # vec._validate_sorted() return vec @classmethod diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 9d64c3d..2b8fd1f 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -169,6 +169,16 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' + # TODO uncomment in 0.6.0 + # def test_from_text_unsorted(self) -> None: + # with pytest.raises(ValueError, match='sparsevec indices must be in ascending order'): + # SparseVector.from_text('{2:1,1:1}/3') + + # TODO uncomment in 0.6.0 + # def test_from_text_duplicate(self) -> None: + # with pytest.raises(ValueError, match='sparsevec indices must be in ascending order'): + # SparseVector.from_text('{2:1,2:1}/3') + def test_from_binary(self) -> None: data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) vec = SparseVector.from_binary(data) From b5d9b7084024a26d4d7b0dda92280a23e7f6d518 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Fri, 16 Oct 2026 22:33:07 -0700 Subject: [PATCH 859/879] Improved _from_dense method --- pgvector/sparsevec.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index c3ac7e5..76e1356 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -137,8 +137,13 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) - self._indices = [i for i, v in enumerate(value) if v != 0] - self._values = [float(value[i]) for i in self._indices] + self._indices = [] + self._values = [] + for i, v in enumerate(value): + v = float(v) + if v != 0: + self._indices.append(i) + self._values.append(v) def _validate(self) -> None: if self._dim < 0: From 97e9fefe23fdc787e3bea96f06b1d2686f8b6d5d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 21 Oct 2026 13:01:51 -0700 Subject: [PATCH 860/879] Fixed memory leak with register_vector and register_vector_async for Psycopg 3 - fixes #168 --- CHANGELOG.md | 1 + pgvector/psycopg/bit.py | 15 +++++++++++++-- pgvector/psycopg/halfvec.py | 15 +++++++++++++-- pgvector/psycopg/sparsevec.py | 15 +++++++++++++-- pgvector/psycopg/vector.py | 15 +++++++++++++-- 5 files changed, 53 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e1c1f48..f033cf1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,6 @@ ## 0.5.2 (unreleased) +- Fixed memory leak with `register_vector` and `register_vector_async` for Psycopg 3 - Fixed error with `Bit` constructor and empty strings - Fixed equality for `HalfVector` with signed zeros diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index f899497..d00c019 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,3 +1,4 @@ +from functools import cache from psycopg import BaseConnection from psycopg.abc import Buffer from psycopg.adapt import Dumper @@ -26,9 +27,19 @@ def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (BitDumper,), {'oid': info.oid}) - binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid}) + text_dumper = bit_text_dumper(info.oid) + binary_dumper = bit_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(Bit, text_dumper) adapters.register_dumper(Bit, binary_dumper) + + +@cache +def bit_text_dumper(oid: int): + return type('', (BitDumper,), {'oid': oid}) + + +@cache +def bit_binary_dumper(oid: int): + return type('', (BitBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 9ed0453..4ca14bb 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,3 +1,4 @@ +from functools import cache from psycopg import BaseConnection from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper @@ -41,11 +42,21 @@ def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> No info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid}) - binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid}) + text_dumper = halfvec_text_dumper(info.oid) + binary_dumper = halfvec_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(HalfVector, text_dumper) adapters.register_dumper(HalfVector, binary_dumper) adapters.register_loader(info.oid, HalfVectorLoader) adapters.register_loader(info.oid, HalfVectorBinaryLoader) + + +@cache +def halfvec_text_dumper(oid: int): + return type('', (HalfVectorDumper,), {'oid': oid}) + + +@cache +def halfvec_binary_dumper(oid: int): + return type('', (HalfVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 57fd19a..2796c94 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,3 +1,4 @@ +from functools import cache from psycopg import BaseConnection from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper @@ -41,11 +42,21 @@ def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid}) - binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid}) + text_dumper = sparsevec_text_dumper(info.oid) + binary_dumper = sparsevec_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(SparseVector, text_dumper) adapters.register_dumper(SparseVector, binary_dumper) adapters.register_loader(info.oid, SparseVectorLoader) adapters.register_loader(info.oid, SparseVectorBinaryLoader) + + +@cache +def sparsevec_text_dumper(oid: int): + return type('', (SparseVectorDumper,), {'oid': oid}) + + +@cache +def sparsevec_binary_dumper(oid: int): + return type('', (SparseVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 7ef16b0..16df1c9 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,4 +1,5 @@ from __future__ import annotations +from functools import cache import psycopg from psycopg import BaseConnection from psycopg.abc import Buffer @@ -50,8 +51,8 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (VectorDumper,), {'oid': info.oid}) - binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) + text_dumper = vector_text_dumper(info.oid) + binary_dumper = vector_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(Vector, text_dumper) @@ -60,3 +61,13 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) + + +@cache +def vector_text_dumper(oid: int): + return type('', (VectorDumper,), {'oid': oid}) + + +@cache +def vector_binary_dumper(oid: int): + return type('', (VectorBinaryDumper,), {'oid': oid}) From e0f3ac37b8f05216cd02a3e4cdd502bcbdb0539d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 21 Oct 2026 13:23:29 -0700 Subject: [PATCH 861/879] Added missing typing --- pgvector/psycopg/bit.py | 4 ++-- pgvector/psycopg/halfvec.py | 4 ++-- pgvector/psycopg/sparsevec.py | 4 ++-- pgvector/psycopg/vector.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index d00c019..4cb502f 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -36,10 +36,10 @@ def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> @cache -def bit_text_dumper(oid: int): +def bit_text_dumper(oid: int) -> type[BitDumper]: return type('', (BitDumper,), {'oid': oid}) @cache -def bit_binary_dumper(oid: int): +def bit_binary_dumper(oid: int) -> type[BitBinaryDumper]: return type('', (BitBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 4ca14bb..05e7994 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -53,10 +53,10 @@ def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> No @cache -def halfvec_text_dumper(oid: int): +def halfvec_text_dumper(oid: int) -> type[HalfVectorDumper]: return type('', (HalfVectorDumper,), {'oid': oid}) @cache -def halfvec_binary_dumper(oid: int): +def halfvec_binary_dumper(oid: int) -> type[HalfVectorBinaryDumper]: return type('', (HalfVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 2796c94..0c953bf 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -53,10 +53,10 @@ def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> @cache -def sparsevec_text_dumper(oid: int): +def sparsevec_text_dumper(oid: int) -> type[SparseVectorDumper]: return type('', (SparseVectorDumper,), {'oid': oid}) @cache -def sparsevec_binary_dumper(oid: int): +def sparsevec_binary_dumper(oid: int) -> type[SparseVectorBinaryDumper]: return type('', (SparseVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 16df1c9..f4b66e4 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -64,10 +64,10 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) @cache -def vector_text_dumper(oid: int): +def vector_text_dumper(oid: int) -> type[VectorDumper]: return type('', (VectorDumper,), {'oid': oid}) @cache -def vector_binary_dumper(oid: int): +def vector_binary_dumper(oid: int) -> type[VectorBinaryDumper]: return type('', (VectorBinaryDumper,), {'oid': oid}) From a75513769f8cc5a696bab2949a81294004600a4d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 11:50:26 -0700 Subject: [PATCH 862/879] Fixed Bit.from_binary with negative length - resolves #173 Co-authored-by: beyondsundial <67931863+beyondsundial@users.noreply.github.com> --- pgvector/bit.py | 2 +- tests/test_bit.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/pgvector/bit.py b/pgvector/bit.py index 098ebb2..a9370e4 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -88,7 +88,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Bit: length, = unpack_from('>i', value) data = memoryview(value)[4:].tobytes() - if len(data) != (length + 7) // 8: + if length < 0 or len(data) != (length + 7) // 8: raise ValueError('invalid length') bit = cls.__new__(cls) diff --git a/tests/test_bit.py b/tests/test_bit.py index 6487348..caeb33f 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -110,3 +110,8 @@ def test_from_binary(self) -> None: if np is not None: assert np.array_equal(vec.to_numpy(), [True, False, True]) assert vec.to_binary() == data + + def test_from_binary_negative_length(self) -> None: + with pytest.raises(ValueError) as error: + Bit.from_binary(pack('>i', -1)) + assert str(error.value) == 'invalid length' From d702a70e05b1a4396266e37caf86766d0a66d7d4 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 17:22:06 -0700 Subject: [PATCH 863/879] Improved tests for from_binary methods --- tests/test_bit.py | 18 +++++++++++++----- tests/test_half_vector.py | 17 +++++++++++++++-- tests/test_sparse_vector.py | 17 +++++++++++++++-- tests/test_vector.py | 17 +++++++++++++++-- 4 files changed, 58 insertions(+), 11 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index caeb33f..2125fbf 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,7 +1,7 @@ from pgvector import Bit import pytest import random -from struct import pack +import struct from .conftest import numpy as np @@ -104,14 +104,22 @@ def test_from_text(self) -> None: assert vec.to_text() == '101' def test_from_binary(self) -> None: - data = pack('>iB', 3, 5 << 5) + data = struct.pack('>iB', 3, 5 << 5) vec = Bit.from_binary(data) assert vec.to_list() == [True, False, True] if np is not None: assert np.array_equal(vec.to_numpy(), [True, False, True]) assert vec.to_binary() == data + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + Bit.from_binary(b'') + def test_from_binary_negative_length(self) -> None: - with pytest.raises(ValueError) as error: - Bit.from_binary(pack('>i', -1)) - assert str(error.value) == 'invalid length' + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>i', -1)) + + def test_from_binary_invalid_length(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>i', 1)) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 56bda26..6932c48 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,6 +1,6 @@ from pgvector import HalfVector import pytest -from struct import pack +import struct from .conftest import numpy as np @@ -68,9 +68,22 @@ def test_from_text(self) -> None: assert vec.to_text() == '[1.5,2.0,3.0]' def test_from_binary(self) -> None: - data = pack('>HH3e', 3, 0, 1.5, 2, 3) + data = struct.pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + HalfVector.from_binary(b'') + + def test_from_binary_negative_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hh', -1, 0)) + + def test_from_binary_non_zero_unused(self) -> None: + with pytest.raises(ValueError, match='expected unused to be 0'): + HalfVector.from_binary(struct.pack('>hh', 0, 1)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 2b8fd1f..4609e2e 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,6 +1,6 @@ from pgvector import SparseVector import pytest -from struct import pack +import struct from .conftest import numpy as np try: @@ -180,7 +180,7 @@ def test_from_text(self) -> None: # SparseVector.from_text('{2:1,2:1}/3') def test_from_binary(self) -> None: - data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) + data = struct.pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) vec = SparseVector.from_binary(data) assert vec.dimensions() == 6 assert vec.indices() == [0, 2, 4] @@ -189,3 +189,16 @@ def test_from_binary(self) -> None: if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_binary() == data + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): + SparseVector.from_binary(b'') + + def test_from_binary_negative_nnz(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii', 0, -1, 0)) + + def test_from_binary_non_zero_unused(self) -> None: + with pytest.raises(ValueError, match='expected unused to be 0'): + SparseVector.from_binary(struct.pack('>iii', 0, 0, 1)) diff --git a/tests/test_vector.py b/tests/test_vector.py index c5b75d0..23f1e0d 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,6 +1,6 @@ from pgvector import Vector import pytest -from struct import pack +import struct from .conftest import numpy as np @@ -68,9 +68,22 @@ def test_from_text(self) -> None: assert vec.to_text() == '[1.5,2.0,3.0]' def test_from_binary(self) -> None: - data = pack('>HH3f', 3, 0, 1.5, 2, 3) + data = struct.pack('>hh3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] if np is not None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + Vector.from_binary(b'') + + def test_from_binary_negative_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hh', -1, 0)) + + def test_from_binary_non_zero_unused(self) -> None: + with pytest.raises(ValueError, match='expected unused to be 0'): + Vector.from_binary(struct.pack('>hh', 0, 1)) From d2976a0f3f859d114c3385fbbb758dd125feefaa Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 17:26:35 -0700 Subject: [PATCH 864/879] Improved format characters [skip ci] --- pgvector/halfvec.py | 3 ++- pgvector/vector.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index fa35194..8658b09 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -59,6 +59,7 @@ def to_binary(self) -> bytes: else: value = array('H', self._value) value.byteswap() + # TODO change to >hh in 0.6.0 return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod @@ -67,7 +68,7 @@ def from_text(cls, value: str, /) -> HalfVector: @classmethod def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: - dim, unused = struct.unpack_from('>HH', value) + dim, unused = struct.unpack_from('>hh', value) data = memoryview(value)[4:] if len(data) != 2 * dim: diff --git a/pgvector/vector.py b/pgvector/vector.py index d357113..10f70fa 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -57,6 +57,7 @@ def to_binary(self) -> bytes: else: value = array('f', self._value) value.byteswap() + # TODO change to >hh in 0.6.0 return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod @@ -65,7 +66,7 @@ def from_text(cls, value: str, /) -> Vector: @classmethod def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: - dim, unused = struct.unpack_from('>HH', value) + dim, unused = struct.unpack_from('>hh', value) data = memoryview(value)[4:] if len(data) != 4 * dim: From 47a6300e764157f01af582956c6f93eb0e15bbac Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 17:35:34 -0700 Subject: [PATCH 865/879] Improved tests for from_binary methods [skip ci] --- tests/test_bit.py | 8 ++++---- tests/test_half_vector.py | 4 ++++ tests/test_sparse_vector.py | 4 ++++ tests/test_vector.py | 4 ++++ 4 files changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 2125fbf..6823c46 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -116,10 +116,10 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Bit.from_binary(b'') - def test_from_binary_negative_length(self) -> None: - with pytest.raises(ValueError, match='invalid length'): - Bit.from_binary(struct.pack('>i', -1)) - def test_from_binary_invalid_length(self) -> None: with pytest.raises(ValueError, match='invalid length'): Bit.from_binary(struct.pack('>i', 1)) + + def test_from_binary_negative_length(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>i', -1)) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 6932c48..6d9146f 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -80,6 +80,10 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): HalfVector.from_binary(b'') + def test_from_binary_invalid_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hh', 1, 0)) + def test_from_binary_negative_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): HalfVector.from_binary(struct.pack('>hh', -1, 0)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 4609e2e..b05ccee 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -195,6 +195,10 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): SparseVector.from_binary(b'') + def test_from_binary_invalid_nnz(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii', 1, 1, 0)) + def test_from_binary_negative_nnz(self) -> None: with pytest.raises(ValueError, match='invalid length'): SparseVector.from_binary(struct.pack('>iii', 0, -1, 0)) diff --git a/tests/test_vector.py b/tests/test_vector.py index 23f1e0d..00e497e 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -80,6 +80,10 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Vector.from_binary(b'') + def test_from_binary_invalid_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hh', 1, 0)) + def test_from_binary_negative_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): Vector.from_binary(struct.pack('>hh', -1, 0)) From 845ee44437086c869e1d7a181dcd924b91b513de Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 17:55:25 -0700 Subject: [PATCH 866/879] Improved tests for from_binary methods [skip ci] --- tests/test_bit.py | 8 ++++++-- tests/test_half_vector.py | 8 ++++++-- tests/test_sparse_vector.py | 8 ++++++-- tests/test_vector.py | 8 ++++++-- 4 files changed, 24 insertions(+), 8 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 6823c46..a7df304 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -116,9 +116,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Bit.from_binary(b'') - def test_from_binary_invalid_length(self) -> None: + def test_from_binary_low_length(self) -> None: with pytest.raises(ValueError, match='invalid length'): - Bit.from_binary(struct.pack('>i', 1)) + Bit.from_binary(struct.pack('>iBB', 8, 0, 0)) + + def test_from_binary_high_length(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>iBB', 17, 0, 0)) def test_from_binary_negative_length(self) -> None: with pytest.raises(ValueError, match='invalid length'): diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 6d9146f..032b923 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -80,9 +80,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): HalfVector.from_binary(b'') - def test_from_binary_invalid_dimensions(self) -> None: + def test_from_binary_low_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): - HalfVector.from_binary(struct.pack('>hh', 1, 0)) + HalfVector.from_binary(struct.pack('>hhee', 1, 0, 1, 2)) + + def test_from_binary_high_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hhee', 3, 0, 1, 2)) def test_from_binary_negative_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index b05ccee..e1be6b7 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -195,9 +195,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): SparseVector.from_binary(b'') - def test_from_binary_invalid_nnz(self) -> None: + def test_from_binary_low_nnz(self) -> None: with pytest.raises(ValueError, match='invalid length'): - SparseVector.from_binary(struct.pack('>iii', 1, 1, 0)) + SparseVector.from_binary(struct.pack('>iii2i2f', 3, 1, 0, 0, 1, 1, 2)) + + def test_from_binary_high_nnz(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii2i2f', 3, 3, 0, 0, 1, 1, 2)) def test_from_binary_negative_nnz(self) -> None: with pytest.raises(ValueError, match='invalid length'): diff --git a/tests/test_vector.py b/tests/test_vector.py index 00e497e..544b66a 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -80,9 +80,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Vector.from_binary(b'') - def test_from_binary_invalid_dimensions(self) -> None: + def test_from_binary_low_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): - Vector.from_binary(struct.pack('>hh', 1, 0)) + Vector.from_binary(struct.pack('>hhff', 1, 0, 1, 2)) + + def test_from_binary_high_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hhff', 3, 0, 1, 2)) def test_from_binary_negative_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): From 6c8276d7fd2d8b9779499e6577ef8c8614da6dc5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:05:47 -0700 Subject: [PATCH 867/879] Improved tests for from_text methods [skip ci] --- tests/test_bit.py | 3 +++ tests/test_half_vector.py | 13 +++++++++++++ tests/test_sparse_vector.py | 3 +++ tests/test_vector.py | 13 +++++++++++++ 4 files changed, 32 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index a7df304..5f60f77 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -103,6 +103,9 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [True, False, True]) assert vec.to_text() == '101' + def test_from_text_zero_length(self) -> None: + assert Bit.from_text('') == Bit([]) + def test_from_binary(self) -> None: data = struct.pack('>iB', 3, 5 << 5) vec = Bit.from_binary(data) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 032b923..74adf97 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -67,6 +67,19 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_text() == '[1.5,2.0,3.0]' + def test_from_text_zero_dimensions(self) -> None: + # TODO no error + with pytest.raises(ValueError, match='could not convert string to float'): + HalfVector.from_text('[]') + + def test_from_text_open_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert HalfVector.from_text('{1]') + + def test_from_text_close_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert HalfVector.from_text('[1}') + def test_from_binary(self) -> None: data = struct.pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index e1be6b7..ae0a8d0 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -169,6 +169,9 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' + def test_from_text_zero_dimensions(self) -> None: + assert SparseVector.from_text('{}/0') == SparseVector([]) + # TODO uncomment in 0.6.0 # def test_from_text_unsorted(self) -> None: # with pytest.raises(ValueError, match='sparsevec indices must be in ascending order'): diff --git a/tests/test_vector.py b/tests/test_vector.py index 544b66a..be08bc3 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -67,6 +67,19 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_text() == '[1.5,2.0,3.0]' + def test_from_text_zero_dimensions(self) -> None: + # TODO no error + with pytest.raises(ValueError, match='could not convert string to float'): + Vector.from_text('[]') + + def test_from_text_open_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert Vector.from_text('{1]') + + def test_from_text_close_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert Vector.from_text('[1}') + def test_from_binary(self) -> None: data = struct.pack('>hh3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) From 90ac9314fcdc0b55fa9986e141da0bf4fddf4a5b Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:07:50 -0700 Subject: [PATCH 868/879] Improved tests for from_text methods [skip ci] --- tests/test_bit.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_bit.py b/tests/test_bit.py index 5f60f77..d0ab71a 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -106,6 +106,10 @@ def test_from_text(self) -> None: def test_from_text_zero_length(self) -> None: assert Bit.from_text('') == Bit([]) + def test_from_text_invalid_character(self) -> None: + with pytest.raises(ValueError, match='expected bit string'): + Bit.from_text('2') + def test_from_binary(self) -> None: data = struct.pack('>iB', 3, 5 << 5) vec = Bit.from_binary(data) From ba94991860251f8604bb16563f920642d74c40c7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:13:19 -0700 Subject: [PATCH 869/879] Added support for empty vectors to from_text methods [skip ci] --- pgvector/halfvec.py | 2 ++ pgvector/vector.py | 2 ++ tests/test_half_vector.py | 4 +--- tests/test_vector.py | 4 +--- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 8658b09..2b3b325 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -88,6 +88,8 @@ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: @classmethod def _from_text(cls, value: str, /) -> list[float]: + if len(value) == 2: + return [] return [float(v) for v in value[1:-1].split(',')] @classmethod diff --git a/pgvector/vector.py b/pgvector/vector.py index 10f70fa..d86cba0 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -86,6 +86,8 @@ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: @classmethod def _from_text(cls, value: str, /) -> list[float]: + if len(value) == 2: + return [] return [float(v) for v in value[1:-1].split(',')] @classmethod diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 74adf97..be5251b 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -68,9 +68,7 @@ def test_from_text(self) -> None: assert vec.to_text() == '[1.5,2.0,3.0]' def test_from_text_zero_dimensions(self) -> None: - # TODO no error - with pytest.raises(ValueError, match='could not convert string to float'): - HalfVector.from_text('[]') + assert HalfVector.from_text('[]') == HalfVector([]) def test_from_text_open_bracket(self) -> None: # TODO raise ValueError in 0.6.0 diff --git a/tests/test_vector.py b/tests/test_vector.py index be08bc3..80e6502 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -68,9 +68,7 @@ def test_from_text(self) -> None: assert vec.to_text() == '[1.5,2.0,3.0]' def test_from_text_zero_dimensions(self) -> None: - # TODO no error - with pytest.raises(ValueError, match='could not convert string to float'): - Vector.from_text('[]') + assert Vector.from_text('[]') == Vector([]) def test_from_text_open_bracket(self) -> None: # TODO raise ValueError in 0.6.0 From 2a35b9b2991e6a01b706ae926a3c4351e44dca12 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:16:02 -0700 Subject: [PATCH 870/879] Improved tests for from_text methods --- tests/test_half_vector.py | 4 ++++ tests/test_sparse_vector.py | 4 ++++ tests/test_vector.py | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index be5251b..688bf0f 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -67,6 +67,10 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_text() == '[1.5,2.0,3.0]' + def test_from_text_zero_length(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + HalfVector.from_text('') + def test_from_text_zero_dimensions(self) -> None: assert HalfVector.from_text('[]') == HalfVector([]) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index ae0a8d0..73a68f5 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -169,6 +169,10 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' + def test_from_text_zero_length(self) -> None: + with pytest.raises(ValueError, match='not enough values to unpack'): + SparseVector.from_text('') + def test_from_text_zero_dimensions(self) -> None: assert SparseVector.from_text('{}/0') == SparseVector([]) diff --git a/tests/test_vector.py b/tests/test_vector.py index 80e6502..681c925 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -67,6 +67,10 @@ def test_from_text(self) -> None: assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_text() == '[1.5,2.0,3.0]' + def test_from_text_zero_length(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + Vector.from_text('') + def test_from_text_zero_dimensions(self) -> None: assert Vector.from_text('[]') == Vector([]) From 67913de4887f43deb1fbf3d187cb6e698861da20 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:21:11 -0700 Subject: [PATCH 871/879] Simplified tests [skip ci] --- tests/test_half_vector.py | 13 ++----------- tests/test_sparse_vector.py | 18 ++---------------- tests/test_vector.py | 13 ++----------- 3 files changed, 6 insertions(+), 38 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 688bf0f..954594b 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -61,11 +61,7 @@ def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 def test_from_text(self) -> None: - vec = HalfVector.from_text('[1.5,2,3]') - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_text() == '[1.5,2.0,3.0]' + assert HalfVector.from_text('[1.5,2,3]') == HalfVector([1.5, 2, 3]) def test_from_text_zero_length(self) -> None: with pytest.raises(ValueError, match='could not convert string to float'): @@ -83,12 +79,7 @@ def test_from_text_close_bracket(self) -> None: assert HalfVector.from_text('[1}') def test_from_binary(self) -> None: - data = struct.pack('>HH3e', 3, 0, 1.5, 2, 3) - vec = HalfVector.from_binary(data) - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_binary() == data + assert HalfVector.from_binary(struct.pack('>hh3e', 3, 0, 1.5, 2, 3)) == HalfVector([1.5, 2, 3]) def test_from_binary_zero_length(self) -> None: # TODO change to ValueError('invalid length') in 0.6.0 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 73a68f5..86ae333 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -160,14 +160,7 @@ def test_zero_vector_text(self) -> None: assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() def test_from_text(self) -> None: - vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') - assert vec.dimensions() == 6 - assert vec.indices() == [0, 2, 4] - assert vec.values() == [1.5, 2, 3] - assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) - assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' + assert SparseVector.from_text('{1:1.5,3:2,5:3}/6') == SparseVector([1.5, 0, 2, 0, 3, 0]) def test_from_text_zero_length(self) -> None: with pytest.raises(ValueError, match='not enough values to unpack'): @@ -188,14 +181,7 @@ def test_from_text_zero_dimensions(self) -> None: def test_from_binary(self) -> None: data = struct.pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) - vec = SparseVector.from_binary(data) - assert vec.dimensions() == 6 - assert vec.indices() == [0, 2, 4] - assert vec.values() == [1.5, 2, 3] - assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) - assert vec.to_binary() == data + assert SparseVector.from_binary(data) == SparseVector([1.5, 0, 2, 0, 3, 0]) def test_from_binary_zero_length(self) -> None: # TODO change to ValueError('invalid length') in 0.6.0 diff --git a/tests/test_vector.py b/tests/test_vector.py index 681c925..d4b1138 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -61,11 +61,7 @@ def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 def test_from_text(self) -> None: - vec = Vector.from_text('[1.5,2,3]') - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_text() == '[1.5,2.0,3.0]' + assert Vector.from_text('[1.5,2,3]') == Vector([1.5, 2, 3]) def test_from_text_zero_length(self) -> None: with pytest.raises(ValueError, match='could not convert string to float'): @@ -83,12 +79,7 @@ def test_from_text_close_bracket(self) -> None: assert Vector.from_text('[1}') def test_from_binary(self) -> None: - data = struct.pack('>hh3f', 3, 0, 1.5, 2, 3) - vec = Vector.from_binary(data) - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_binary() == data + assert Vector.from_binary(struct.pack('>hh3f', 3, 0, 1.5, 2, 3)) == Vector([1.5, 2, 3]) def test_from_binary_zero_length(self) -> None: # TODO change to ValueError('invalid length') in 0.6.0 From a594893a604aebc433c000f421b9183ad39529e7 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:33:25 -0700 Subject: [PATCH 872/879] Improved tests for from_text methods [skip ci] --- tests/test_half_vector.py | 4 ++++ tests/test_sparse_vector.py | 8 ++++++++ tests/test_vector.py | 4 ++++ 3 files changed, 16 insertions(+) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 954594b..6f3b5af 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -78,6 +78,10 @@ def test_from_text_close_bracket(self) -> None: # TODO raise ValueError in 0.6.0 assert HalfVector.from_text('[1}') + def test_from_text_non_number(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + assert HalfVector.from_text('[a]') + def test_from_binary(self) -> None: assert HalfVector.from_binary(struct.pack('>hh3e', 3, 0, 1.5, 2, 3)) == HalfVector([1.5, 2, 3]) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 86ae333..7139595 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -169,6 +169,14 @@ def test_from_text_zero_length(self) -> None: def test_from_text_zero_dimensions(self) -> None: assert SparseVector.from_text('{}/0') == SparseVector([]) + def test_from_text_non_number_index(self) -> None: + with pytest.raises(ValueError, match='invalid literal for int'): + assert SparseVector.from_text('{a:1}/1') + + def test_from_text_non_number_value(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + assert SparseVector.from_text('{1:a}/1') + # TODO uncomment in 0.6.0 # def test_from_text_unsorted(self) -> None: # with pytest.raises(ValueError, match='sparsevec indices must be in ascending order'): diff --git a/tests/test_vector.py b/tests/test_vector.py index d4b1138..075e862 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -78,6 +78,10 @@ def test_from_text_close_bracket(self) -> None: # TODO raise ValueError in 0.6.0 assert Vector.from_text('[1}') + def test_from_text_non_number(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + assert Vector.from_text('[a]') + def test_from_binary(self) -> None: assert Vector.from_binary(struct.pack('>hh3f', 3, 0, 1.5, 2, 3)) == Vector([1.5, 2, 3]) From 2109b244e6c96e8d173f97d4186f2147aa28cc4d Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:35:44 -0700 Subject: [PATCH 873/879] Simplified tests [skip ci] --- tests/test_bit.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index d0ab71a..4e018d0 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -97,11 +97,7 @@ def test_equality(self) -> None: assert Bit([True, False, True]) != 1 def test_from_text(self) -> None: - vec = Bit.from_text('101') - assert vec.to_list() == [True, False, True] - if np is not None: - assert np.array_equal(vec.to_numpy(), [True, False, True]) - assert vec.to_text() == '101' + assert Bit.from_text('101') == Bit('101') def test_from_text_zero_length(self) -> None: assert Bit.from_text('') == Bit([]) @@ -111,12 +107,7 @@ def test_from_text_invalid_character(self) -> None: Bit.from_text('2') def test_from_binary(self) -> None: - data = struct.pack('>iB', 3, 5 << 5) - vec = Bit.from_binary(data) - assert vec.to_list() == [True, False, True] - if np is not None: - assert np.array_equal(vec.to_numpy(), [True, False, True]) - assert vec.to_binary() == data + assert Bit.from_binary(struct.pack('>iB', 3, 5 << 5)) == Bit('101') def test_from_binary_zero_length(self) -> None: # TODO change to ValueError('invalid length') in 0.6.0 From bb914602af73d564b438a8a3d5e00d1bd63e3e62 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 18:43:21 -0700 Subject: [PATCH 874/879] Improved test names [skip ci] --- tests/test_bit.py | 34 +++++++++++++++++----------------- tests/test_half_vector.py | 14 +++++++------- tests/test_sparse_vector.py | 32 ++++++++++++++++---------------- tests/test_vector.py | 14 +++++++------- 4 files changed, 47 insertions(+), 47 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 4e018d0..76c9d57 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -6,46 +6,46 @@ class TestBit: - def test_list(self) -> None: + def test_init_list(self) -> None: assert Bit([True, False, True]).to_list() == [True, False, True] - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert Bit([]).to_list() == [] - def test_list_none(self) -> None: + def test_init_list_none(self) -> None: with pytest.raises(ValueError) as error: Bit([True, None, True]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_list_int(self) -> None: + def test_init_list_int(self) -> None: with pytest.raises(ValueError) as error: Bit([254, 7, 0]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_list_list(self) -> None: + def test_init_list_list(self) -> None: with pytest.raises(ValueError) as error: Bit([[True, False], [True, False]]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_str(self) -> None: + def test_init_str(self) -> None: assert Bit('101').to_list() == [True, False, True] - def test_str_empty(self) -> None: + def test_init_str_empty(self) -> None: assert Bit('').to_list() == [] - def test_str_two(self) -> None: + def test_init_str_two(self) -> None: with pytest.raises(ValueError) as error: Bit('201') assert str(error.value) == 'expected bit string' - def test_bytes(self) -> None: + def test_init_bytes(self) -> None: assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - def test_bytes_empty(self) -> None: + def test_init_bytes_empty(self) -> None: assert Bit(b'').to_text() == '' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -53,14 +53,14 @@ def test_ndarray(self) -> None: assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) - def test_ndarray_unpackbits(self) -> None: + def test_init_ndarray_unpackbits(self) -> None: if np is None: pytest.skip('NumPy required') arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' - def test_ndarray_uint8(self) -> None: + def test_init_ndarray_uint8(self) -> None: if np is None: pytest.skip('NumPy required') @@ -69,7 +69,7 @@ def test_ndarray_uint8(self) -> None: Bit(arr) assert str(error.value) == 'expected elements to be boolean' - def test_ndarray_uint16(self) -> None: + def test_init_ndarray_uint16(self) -> None: if np is None: pytest.skip('NumPy required') @@ -78,12 +78,12 @@ def test_ndarray_uint16(self) -> None: Bit(arr) # type: ignore assert str(error.value) == 'expected elements to be boolean' - def test_bool(self) -> None: + def test_init_bool(self) -> None: with pytest.raises(ValueError) as error: Bit(True) # type: ignore assert str(error.value) == 'expected bytes, str, list, or ndarray' - def test_random(self) -> None: + def test_init_random(self) -> None: value = ''.join(random.choices(['0', '1'], k=random.randint(1024, 2048))) assert Bit(value).to_text() == value @@ -91,7 +91,7 @@ def test_repr(self) -> None: assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' - def test_equality(self) -> None: + def test_eq(self) -> None: assert Bit([True, False, True]) == Bit([True, False, True]) assert Bit([True, False, True]) != Bit([True, False, False]) assert Bit([True, False, True]) != 1 diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 6f3b5af..39354ee 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -5,25 +5,25 @@ class TestHalfVector: - def test_list(self) -> None: + def test_init_list(self) -> None: arr = [1.0, 2.0, 3.0] assert HalfVector(arr).to_list() == arr assert HalfVector(arr).to_list() is not arr - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert HalfVector([]).to_list() == [] - def test_list_str(self) -> None: + def test_init_list_str(self) -> None: with pytest.raises(ValueError) as error: HalfVector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_list_list(self) -> None: + def test_init_list_list(self) -> None: with pytest.raises(ValueError) as error: HalfVector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -42,7 +42,7 @@ def test_ndarray(self) -> None: HalfVector(np.array(['one', 'two', 'three'])) assert 'could not convert string to float' in str(error.value) - def test_int(self) -> None: + def test_init_int(self) -> None: with pytest.raises(ValueError) as error: HalfVector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' @@ -51,7 +51,7 @@ def test_repr(self) -> None: assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' - def test_equality(self) -> None: + def test_eq(self) -> None: assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) assert HalfVector([1, 2, 3]) != 1 diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 7139595..a644260 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -10,26 +10,26 @@ class TestSparseVector: - def test_list(self) -> None: + def test_init_list(self) -> None: vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] if np is not None: assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert SparseVector([]).to_list() == [] - def test_list_str(self) -> None: + def test_init_list_str(self) -> None: with pytest.raises(ValueError, match='could not convert string to float'): SparseVector([1, 'two', 3]) # type: ignore - def test_list_dimensions(self) -> None: + def test_init_list_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore assert str(error.value) == 'extra argument' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -37,12 +37,12 @@ def test_ndarray(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict(self) -> None: + def test_init_dict(self) -> None: vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict_empty(self) -> None: + def test_init_dict_empty(self) -> None: assert SparseVector({}, 0).to_list() == [] # TODO uncomment in 0.6.0 @@ -60,12 +60,12 @@ def test_dict_empty(self) -> None: # with pytest.raises(ValueError, match='sparsevec cannot have negative dimensions'): # SparseVector({}, -1) - def test_dict_no_dimensions(self) -> None: + def test_init_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore assert str(error.value) == 'missing dimensions' - def test_coo_array(self) -> None: + def test_init_coo_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -75,7 +75,7 @@ def test_coo_array(self) -> None: assert vec.indices() == [0, 2, 4] assert isinstance(vec.values()[0], float) - def test_coo_array_duplicates(self) -> None: + def test_init_coo_array_duplicates(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -85,7 +85,7 @@ def test_coo_array_duplicates(self) -> None: assert vec.to_list() == [0, 3, 0] assert vec.indices() == [1] - def test_coo_array_dimensions(self) -> None: + def test_init_coo_array_dimensions(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -93,7 +93,7 @@ def test_coo_array_dimensions(self) -> None: SparseVector(sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore assert str(error.value) == 'extra argument' - def test_coo_matrix(self) -> None: + def test_init_coo_matrix(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -102,7 +102,7 @@ def test_coo_matrix(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dok_array(self) -> None: + def test_init_dok_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -111,7 +111,7 @@ def test_dok_array(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_csr_array(self) -> None: + def test_init_csr_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -120,7 +120,7 @@ def test_csr_array(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_csr_matrix(self) -> None: + def test_init_csr_matrix(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -133,7 +133,7 @@ def test_repr(self) -> None: assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' - def test_equality(self) -> None: + def test_eq(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) diff --git a/tests/test_vector.py b/tests/test_vector.py index 075e862..d1fc6e5 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -5,25 +5,25 @@ class TestVector: - def test_list(self) -> None: + def test_init_list(self) -> None: arr = [1.0, 2.0, 3.0] assert Vector(arr).to_list() == arr assert Vector(arr).to_list() is not arr - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert Vector([]).to_list() == [] - def test_list_str(self) -> None: + def test_init_list_str(self) -> None: with pytest.raises(ValueError) as error: Vector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_list_list(self) -> None: + def test_init_list_list(self) -> None: with pytest.raises(ValueError) as error: Vector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -42,7 +42,7 @@ def test_ndarray(self) -> None: Vector(np.array(['one', 'two', 'three'])) assert 'could not convert string to float' in str(error.value) - def test_int(self) -> None: + def test_init_int(self) -> None: with pytest.raises(ValueError) as error: Vector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' @@ -51,7 +51,7 @@ def test_repr(self) -> None: assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' - def test_equality(self) -> None: + def test_eq(self) -> None: assert Vector([1, 2, 3]) == Vector([1, 2, 3]) assert Vector([1, 2, 3]) != Vector([1, 2, 4]) assert Vector([1, 2, 3]) != 1 From 312ce70b78f2589379d016321e83b87eaddff9a5 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 22:37:51 -0700 Subject: [PATCH 875/879] Added more tests --- tests/test_bit.py | 17 +++++++++++++++++ tests/test_half_vector.py | 17 +++++++++++++++++ tests/test_sparse_vector.py | 19 ++++++++++++++++--- tests/test_vector.py | 17 +++++++++++++++++ 4 files changed, 67 insertions(+), 3 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 76c9d57..34a48e3 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -96,6 +96,23 @@ def test_eq(self) -> None: assert Bit([True, False, True]) != Bit([True, False, False]) assert Bit([True, False, True]) != 1 + def test_to_list(self) -> None: + assert Bit('101').to_list() == [True, False, True] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = Bit('101').to_numpy() + assert np.array_equal(arr, np.array([True, False, True])) + assert arr.dtype == np.bool_ + + def test_to_text(self) -> None: + assert Bit('101').to_text() == '101' + + def test_to_binary(self) -> None: + assert Bit('101').to_binary() == struct.pack('>iB', 3, 5 << 5) + def test_from_text(self) -> None: assert Bit.from_text('101') == Bit('101') diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 39354ee..d9d8580 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -60,6 +60,23 @@ def test_eq(self) -> None: def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 + def test_to_list(self) -> None: + assert HalfVector([1.5, 2, 3]).to_list() == [1.5, 2, 3] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = HalfVector([1.5, 2, 3]).to_numpy() + assert np.array_equal(arr, np.array([1.5, 2, 3])) + assert arr.dtype == np.float16 + + def test_to_text(self) -> None: + assert HalfVector([1.5, 2, 3]).to_text() == '[1.5,2.0,3.0]' + + def test_to_binary(self) -> None: + assert HalfVector([1.5, 2, 3]).to_binary() == struct.pack('>hh3e', 3, 0, 1.5, 2, 3) + def test_from_text(self) -> None: assert HalfVector.from_text('[1.5,2,3]') == HalfVector([1.5, 2, 3]) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index a644260..6666baf 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -155,9 +155,22 @@ def test_to_coo(self) -> None: assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) - def test_zero_vector_text(self) -> None: - vec = SparseVector({}, 3) - assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() + def test_to_list(self) -> None: + assert SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_list() == [1.5, 0, 2, 0, 3, 0] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_numpy() + assert np.array_equal(arr, np.array([1.5, 0, 2, 0, 3, 0])) + assert arr.dtype == np.float32 + + def test_to_text(self) -> None: + assert SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_text() == '{1:1.5,3:2.0,5:3.0}/6' + + def test_to_binary(self) -> None: + assert SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_binary() == struct.pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) def test_from_text(self) -> None: assert SparseVector.from_text('{1:1.5,3:2,5:3}/6') == SparseVector([1.5, 0, 2, 0, 3, 0]) diff --git a/tests/test_vector.py b/tests/test_vector.py index d1fc6e5..621d0fe 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -60,6 +60,23 @@ def test_eq(self) -> None: def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 + def test_to_list(self) -> None: + assert Vector([1.5, 2, 3]).to_list() == [1.5, 2, 3] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = Vector([1.5, 2, 3]).to_numpy() + assert np.array_equal(arr, np.array([1.5, 2, 3])) + assert arr.dtype == np.float32 + + def test_to_text(self) -> None: + assert Vector([1.5, 2, 3]).to_text() == '[1.5,2.0,3.0]' + + def test_to_binary(self) -> None: + assert Vector([1.5, 2, 3]).to_binary() == struct.pack('>hh3f', 3, 0, 1.5, 2, 3) + def test_from_text(self) -> None: assert Vector.from_text('[1.5,2,3]') == Vector([1.5, 2, 3]) From d4ca080473d1cadd6e318fd9aeda864d02d3612e Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 22:45:54 -0700 Subject: [PATCH 876/879] Split tests [skip ci] --- tests/test_half_vector.py | 18 ++++++++++++++++-- tests/test_vector.py | 18 ++++++++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index d9d8580..c99c0b2 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -31,13 +31,27 @@ def test_init_ndarray(self) -> None: assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr assert HalfVector(arr).to_numpy().dtype == np.float16 - # non-contiguous + + def test_init_ndarray_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(np.flip(arr)).to_list() == [3, 2, 1] assert HalfVector(np.flip(arr)).to_binary() == HalfVector([3, 2, 1]).to_binary() - # big endian + + def test_init_ndarray_big_endian(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(arr.astype('>f2')).to_list() == [1, 2, 3] assert HalfVector(arr.astype('>f2')).to_binary() == HalfVector([1, 2, 3]).to_binary() + def test_init_ndarray_str(self) -> None: + if np is None: + pytest.skip('NumPy required') + with pytest.raises(ValueError) as error: HalfVector(np.array(['one', 'two', 'three'])) assert 'could not convert string to float' in str(error.value) diff --git a/tests/test_vector.py b/tests/test_vector.py index 621d0fe..b1eda33 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -31,13 +31,27 @@ def test_init_ndarray(self) -> None: assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr assert Vector(arr).to_numpy().dtype == np.float32 - # non-contiguous + + def test_init_ndarray_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(np.flip(arr)).to_list() == [3, 2, 1] assert Vector(np.flip(arr)).to_binary() == Vector([3, 2, 1]).to_binary() - # big endian + + def test_init_ndarray_big_endian(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(arr.astype('>f4')).to_list() == [1, 2, 3] assert Vector(arr.astype('>f4')).to_binary() == Vector([1, 2, 3]).to_binary() + def test_init_ndarray_str(self) -> None: + if np is None: + pytest.skip('NumPy required') + with pytest.raises(ValueError) as error: Vector(np.array(['one', 'two', 'three'])) assert 'could not convert string to float' in str(error.value) From c314b69f678eaa959f45958ef0547f4ec8672824 Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 23:00:51 -0700 Subject: [PATCH 877/879] Improved test names [skip ci] --- tests/test_bit.py | 4 ++-- tests/test_half_vector.py | 4 ++-- tests/test_sparse_vector.py | 4 ++-- tests/test_vector.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index 34a48e3..bffbd5b 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -131,11 +131,11 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Bit.from_binary(b'') - def test_from_binary_low_length(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): Bit.from_binary(struct.pack('>iBB', 8, 0, 0)) - def test_from_binary_high_length(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): Bit.from_binary(struct.pack('>iBB', 17, 0, 0)) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index c99c0b2..f0648d5 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -121,11 +121,11 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): HalfVector.from_binary(b'') - def test_from_binary_low_dimensions(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): HalfVector.from_binary(struct.pack('>hhee', 1, 0, 1, 2)) - def test_from_binary_high_dimensions(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): HalfVector.from_binary(struct.pack('>hhee', 3, 0, 1, 2)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 6666baf..502e820 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -209,11 +209,11 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): SparseVector.from_binary(b'') - def test_from_binary_low_nnz(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): SparseVector.from_binary(struct.pack('>iii2i2f', 3, 1, 0, 0, 1, 1, 2)) - def test_from_binary_high_nnz(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): SparseVector.from_binary(struct.pack('>iii2i2f', 3, 3, 0, 0, 1, 1, 2)) diff --git a/tests/test_vector.py b/tests/test_vector.py index b1eda33..883e3a6 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -121,11 +121,11 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Vector.from_binary(b'') - def test_from_binary_low_dimensions(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): Vector.from_binary(struct.pack('>hhff', 1, 0, 1, 2)) - def test_from_binary_high_dimensions(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): Vector.from_binary(struct.pack('>hhff', 3, 0, 1, 2)) From 062b477cde8f722aeffc9f58f7ad1f91751f848c Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Tue, 27 Oct 2026 23:27:30 -0700 Subject: [PATCH 878/879] Improved tests --- tests/test_bit.py | 8 ++++---- tests/test_half_vector.py | 8 ++++---- tests/test_sparse_vector.py | 8 ++++---- tests/test_vector.py | 8 ++++---- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/test_bit.py b/tests/test_bit.py index bffbd5b..9179cd3 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -131,13 +131,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Bit.from_binary(b'') - def test_from_binary_extra_data(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - Bit.from_binary(struct.pack('>iBB', 8, 0, 0)) + Bit.from_binary(struct.pack('>iB', 9, 0)) - def test_from_binary_missing_data(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - Bit.from_binary(struct.pack('>iBB', 17, 0, 0)) + Bit.from_binary(struct.pack('>iBBB', 16, 0, 0, 0)) def test_from_binary_negative_length(self) -> None: with pytest.raises(ValueError, match='invalid length'): diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index f0648d5..a907c94 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -121,13 +121,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): HalfVector.from_binary(b'') - def test_from_binary_extra_data(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - HalfVector.from_binary(struct.pack('>hhee', 1, 0, 1, 2)) + HalfVector.from_binary(struct.pack('>hhe', 2, 0, 1)) - def test_from_binary_missing_data(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - HalfVector.from_binary(struct.pack('>hhee', 3, 0, 1, 2)) + HalfVector.from_binary(struct.pack('>hheee', 2, 0, 1, 2, 3)) def test_from_binary_negative_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 502e820..e3bb8cc 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -209,13 +209,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): SparseVector.from_binary(b'') - def test_from_binary_extra_data(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - SparseVector.from_binary(struct.pack('>iii2i2f', 3, 1, 0, 0, 1, 1, 2)) + SparseVector.from_binary(struct.pack('>iii2if', 2, 2, 0, 0, 1, 1)) - def test_from_binary_missing_data(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - SparseVector.from_binary(struct.pack('>iii2i2f', 3, 3, 0, 0, 1, 1, 2)) + SparseVector.from_binary(struct.pack('>iii2ifff', 2, 2, 0, 0, 1, 1, 2, 3)) def test_from_binary_negative_nnz(self) -> None: with pytest.raises(ValueError, match='invalid length'): diff --git a/tests/test_vector.py b/tests/test_vector.py index 883e3a6..13a9cc7 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -121,13 +121,13 @@ def test_from_binary_zero_length(self) -> None: with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): Vector.from_binary(b'') - def test_from_binary_extra_data(self) -> None: + def test_from_binary_missing_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - Vector.from_binary(struct.pack('>hhff', 1, 0, 1, 2)) + Vector.from_binary(struct.pack('>hhf', 2, 0, 1)) - def test_from_binary_missing_data(self) -> None: + def test_from_binary_extra_data(self) -> None: with pytest.raises(ValueError, match='invalid length'): - Vector.from_binary(struct.pack('>hhff', 3, 0, 1, 2)) + Vector.from_binary(struct.pack('>hhfff', 2, 0, 1, 2, 3)) def test_from_binary_negative_dimensions(self) -> None: with pytest.raises(ValueError, match='invalid length'): From d2ace87ca497517b73ffa2c88a50896f65a2b17a Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Wed, 4 Nov 2026 18:53:08 -0800 Subject: [PATCH 879/879] Updated Cohere example [skip ci] --- examples/cohere/example.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 5cf8d69..9b32073 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -9,12 +9,12 @@ register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(2048))') def embed(input, input_type): co = cohere.ClientV2() - response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) + response = co.embed(texts=input, model='embed-v5.0-pro', input_type=input_type, embedding_types=['ubinary']) return [bytes(embedding) for embedding in response.embeddings.ubinary]