diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e821057..5a01635 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -3,18 +3,31 @@ on: [push, pull_request] jobs: build: runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python: [3.15, "3.10"] steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - - run: pip install -r requirements.txt - - uses: ankane/setup-postgres@v1 - with: - database: pgvector_python_test - - 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 - cd pgvector - make - sudo make install - - run: pytest + - uses: actions/checkout@v7 + - uses: actions/setup-python@v7 + with: + python-version: ${{ matrix.python }} + - run: pip install --group dev + - uses: ankane/setup-postgres@v1 + with: + database: pgvector_python_test + dev-files: true + - run: | + cd /tmp + git clone --branch v0.8.7 https://github.com/pgvector/pgvector.git + cd pgvector + make + sudo make install + - run: pytest + + - run: pip install --group numpy --group scipy + - run: pytest + + - run: pip install --group typing + - run: mypy + - run: ty check diff --git a/.gitignore b/.gitignore index 1d3b727..5556c9f 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,6 @@ venv/ .cache/ *.pyc __pycache__ +.pytest_cache/ +*.lock +examples/rag/README.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 16617ca..f033cf1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,125 @@ +## 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 + +## 0.5.1 (2026-10-08) + +- 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) + +- Added experimental support for type hints +- 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 +- Dropped support for Python < 3.10 +- Dropped support for SQLAlchemy < 2 + +## 0.4.2 (2025-12-04) + +- Added support for Django 6 +- Added support for `str` objects for `bit` type with SQLAlchemy + +## 0.4.1 (2025-04-26) + +- Fixed `SparseVector` constructor for SciPy sparse matrices + +## 0.4.0 (2025-03-15) + +- 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 +- 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) + +- Added `arrays` option for Psycopg 2 + +## 0.3.5 (2024-10-05) + +- Added `avg` function with type casting to SQLAlchemy +- Added `globally` option for Psycopg 2 + +## 0.3.4 (2024-09-26) + +- Added `schema` option for asyncpg + +## 0.3.3 (2024-09-09) + +- Improved support for cursor factories with Psycopg 2 + +## 0.3.2 (2024-07-17) + +- Fixed error with asyncpg and pgvector < 0.7 + +## 0.3.1 (2024-07-10) + +- 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 + +## 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 +- 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`, `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`, `hamming_distance`, and `jaccard_distance` for Peewee + +## 0.2.5 (2024-02-07) + +- Added literal binds support for SQLAlchemy + +## 0.2.4 (2023-11-24) + +- Improved reflection with SQLAlchemy + +## 0.2.3 (2023-09-25) + +- Fixed null values with Django +- Fixed `full_clean` with Django + +## 0.2.2 (2023-09-08) + +- Added support for Peewee +- Added `HnswIndex` for Django + +## 0.2.1 (2023-07-31) + +- Fixed form issues with Django + +## 0.2.0 (2023-07-23) + +- Fixed form validation with Django +- Dropped support for Python < 3.8 + +## 0.1.8 (2023-05-20) + +- Fixed serialization with Django + +## 0.1.7 (2023-05-11) + +- Added `register_vector_async` for Psycopg 3 +- Fixed `set_types` for Psycopg 3 + +## 0.1.6 (2022-05-22) + +- Fixed return type for distance operators with SQLAlchemy + ## 0.1.5 (2022-01-14) - Fixed `operator does not exist` error with Django @@ -5,12 +127,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) diff --git a/LICENSE.txt b/LICENSE.txt index 5474bcb..17e5210 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2022 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 diff --git a/Makefile b/Makefile index 2199d93..7e5aab0 100644 --- a/Makefile +++ b/Makefile @@ -1,11 +1,13 @@ +.PHONY: lint build publish clean + lint: - pycodestyle . --ignore=E501 + pycodestyle . --ignore=E501,E704,W504 --exclude=.venv + +build: + python3 -m build -publish: clean - python3 setup.py bdist_wheel --universal - ls dist - # twine upload dist/* - make clean +publish: clean build + twine upload dist/* clean: - rm -rf .pytest_cache build dist pgvector.egg-info + rm -rf .pytest_cache dist pgvector.egg-info diff --git a/README.md b/README.md index 5e4406e..26dccfb 100644 --- a/README.md +++ b/README.md @@ -2,11 +2,9 @@ [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), [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) -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) +[![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) ## Installation @@ -20,20 +18,37 @@ And follow the instructions for your database library: - [Django](#django) - [SQLAlchemy](#sqlalchemy) +- [SQLModel](#sqlmodel) - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) +- [pg8000](#pg8000) +- [Peewee](#peewee) Or check out some examples: -- [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 +- [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 +- [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 +- [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 +- [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 -Create the extension +Create a migration to enable the extension ```python from pgvector.django import VectorExtension @@ -44,19 +59,21 @@ class Migration(migrations.Migration): ] ``` -Add a vector field +Add a vector field to your model ```python from pgvector.django import VectorField class Item(models.Model): - factors = VectorField(dimensions=3) + embedding = VectorField(dimensions=3) ``` +Also supports `HalfVectorField`, `BitField`, and `SparseVectorField` + Insert a vector ```python -item = Item(factors=[1, 2, 3]) +item = Item(embedding=[1, 2, 3]) item.save() ``` @@ -65,22 +82,52 @@ 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`, `CosineDistance`, `L1Distance`, `HammingDistance`, and `JaccardDistance` + +Get the distance + +```python +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) +``` + +Average vectors + +```python +from django.db.models import Avg + +Item.objects.aggregate(Avg('embedding')) ``` -Also supports `MaxInnerProduct` and `CosineDistance` +Also supports `Sum` Add an approximate index ```python -from pgvector.django import IvfflatIndex +from pgvector.django import HnswIndex, IvfflatIndex class Item(models.Model): class Meta: indexes = [ + HnswIndex( + name='my_index', + fields=['embedding'], + m=16, + ef_construction=64, + opclasses=['vector_l2_ops'] + ), + # or IvfflatIndex( name='my_index', - fields=['factors'], + fields=['embedding'], lists=100, opclasses=['vector_l2_ops'] ) @@ -89,21 +136,59 @@ 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 HnswIndex, HalfVectorField + +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` + +Get the nearest neighbors + +```python +distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2]) +Item.objects.order_by(distance)[:5] +``` + ## SQLAlchemy +Enable the extension + +```python +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): - factors = Column(Vector(3)) + embedding: Mapped[list[float]] = mapped_column(VECTOR(3)) ``` +Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` + Insert a vector ```python -item = Item(factors=[1, 2, 3]) +item = Item(embedding=[1, 2, 3]) session.add(item) session.commit() ``` @@ -111,19 +196,236 @@ 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.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) +``` + +Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` + +Get the distance + +```python +session.scalars(select(Item.embedding.l2_distance([3, 1, 2]))) +``` + +Get items within a certain distance + +```python +session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) +``` + +Average vectors + +```python +from pgvector.sqlalchemy import avg + +session.scalars(select(avg(Item.embedding))).first() ``` -Also supports `max_inner_product` and `cosine_distance` +Also supports `sum` Add an approximate index ```python -index = Index('my_index', Item.factors, +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={'factors': 'vector_l2_ops'} + postgresql_ops={'embedding': 'vector_l2_ops'} ) + +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_ops={'embedding': 'halfvec_l2_ops'} +) +``` + +Get the nearest neighbors + +```python +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_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 + +```python +from pgvector.sqlalchemy import VECTOR +from sqlalchemy import ARRAY + +class Item(Base): + embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1)) +``` + +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 + +@event.listens_for(engine, "connect") +def connect(dbapi_connection, connection_record): + register_vector(dbapi_connection) +``` + +## SQLModel + +Enable the extension + +```python +session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) +``` + +Add a vector column + +```python +from pgvector.sqlalchemy import VECTOR + +class Item(SQLModel, table=True): + embedding: list[float] = Field(sa_type=VECTOR(3)) +``` + +Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` + +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`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_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 pgvector.sqlalchemy import avg + +session.exec(select(avg(Item.embedding))).first() +``` + +Also supports `sum` + +Add an approximate index + +```python +from sqlmodel 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) ``` @@ -131,7 +433,13 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc ## Psycopg 3 -Register the vector type with your connection +Enable the extension + +```python +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the types with your connection ```python from pgvector.psycopg import register_vector @@ -139,22 +447,64 @@ 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 +from pgvector.psycopg import register_vector_async + +await register_vector_async(conn) +``` + +Create a table + +```python +conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +``` + Insert a vector ```python -factors = np.array([1, 2, 3]) -conn.execute('INSERT INTO item (factors) VALUES (%s)', (factors,)) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) +conn.execute('INSERT INTO items (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 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 -Register the vector type with your connection or cursor +Enable the extension + +```python +cur = conn.cursor() +cur.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the types with your connection or cursor ```python from pgvector.psycopg2 import register_vector @@ -162,23 +512,47 @@ from pgvector.psycopg2 import register_vector register_vector(conn) ``` +Create a table + +```python +cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +``` + Insert a vector ```python -factors = np.array([1, 2, 3]) -cur.execute('INSERT INTO item (factors) VALUES (%s)', (factors,)) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) +cur.execute('INSERT INTO items (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 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 -Register the vector type with your connection +Enable the extension + +```python +await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the types with your connection ```python from pgvector.asyncpg import register_vector @@ -186,17 +560,267 @@ 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=init) +``` + +Create a table + +```python +await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +``` + +Insert a vector + +```python +from pgvector import Vector + +embedding = Vector([1, 2, 3]) +await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding) +``` + +Get the nearest neighbors to a vector + +```python +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 + +## pg8000 + +Enable the extension + +```python +conn.run('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the types 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 -factors = np.array([1, 2, 3]) -await conn.execute('INSERT INTO item (factors) VALUES ($1)', factors) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) +conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding) ``` Get the nearest neighbors to a vector ```python -await conn.fetch('SELECT * FROM item ORDER BY factors <-> $1 LIMIT 5', factors) +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 + +```python +from pgvector.peewee import VectorField + +class Item(BaseModel): + embedding = VectorField(dimensions=3) +``` + +Also supports `HalfVectorField`, `FixedBitField`, and `SparseVectorField` + +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`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_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) +``` + +Average vectors + +```python +from peewee import fn + +Item.select(fn.avg(Item.embedding).coerce(True)).scalar() +``` + +Also supports `sum` + +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 + +## 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 + +```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 @@ -217,7 +841,16 @@ 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 ``` + +To run an example: + +```sh +cd examples/loading +pip install --group dev +createdb pgvector_example +python3 example.py +``` diff --git a/examples/citus/example.py b/examples/citus/example.py new file mode 100644 index 0000000..d1af5b9 --- /dev/null +++ b/examples/citus/example.py @@ -0,0 +1,50 @@ +import numpy as np +from pgvector import Vector +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([Vector(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]) diff --git a/examples/citus/pyproject.toml b/examples/citus/pyproject.toml new file mode 100644 index 0000000..3fb37d4 --- /dev/null +++ b/examples/citus/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/cohere/example.py b/examples/cohere/example.py new file mode 100644 index 0000000..9b32073 --- /dev/null +++ b/examples/cohere/example.py @@ -0,0 +1,34 @@ +import cohere +from pgvector import Bit +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 bit(2048))') + + +def embed(input, input_type): + co = cohere.ClientV2() + 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] + + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +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 = 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/cohere/pyproject.toml b/examples/cohere/pyproject.toml new file mode 100644 index 0000000..1228b3d --- /dev/null +++ b/examples/cohere/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "cohere", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py new file mode 100644 index 0000000..41f88b2 --- /dev/null +++ b/examples/colbert/approximate.py @@ -0,0 +1,80 @@ +# based on section 3.6 of 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 +import warnings + +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 +""") + +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) + +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 = query_embeddings + [query_embeddings] +result = conn.execute(sql, params).fetchall() +for row in result: + print(row) diff --git a/examples/colbert/exact.py b/examples/colbert/exact.py new file mode 100644 index 0000000..e6a2936 --- /dev/null +++ b/examples/colbert/exact.py @@ -0,0 +1,51 @@ +from colbert.infra import ColBERTConfig +from colbert.modeling.checkpoint import Checkpoint +from pgvector.psycopg import register_vector +import psycopg +import warnings + +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 +""") + +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) + +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 = [e.numpy() for e in embeddings] + conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) + +query = 'puppy' +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) diff --git a/examples/colbert/pyproject.toml b/examples/colbert/pyproject.toml new file mode 100644 index 0000000..7cf1e3f --- /dev/null +++ b/examples/colbert/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "colbert-ai", + "pgvector", + "psycopg[binary]", + "transformers==4.49.0" +] diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py new file mode 100644 index 0000000..313c8bf --- /dev/null +++ b/examples/colpali/exact.py @@ -0,0 +1,57 @@ +from colpali_engine.models import ColQwen2, ColQwen2Processor +from colpali_engine.utils.torch_utils import get_torch_device +from datasets import load_dataset +from pgvector import Bit +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 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 +""") + +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') + + +def generate_embeddings(processed): + with torch.no_grad(): + return model(**processed.to(model.device)).to(torch.float32).numpy(force=True) + + +def binary_quantize(embedding): + return Bit(embedding > 0) + + +input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image'] +for content in input: + 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) 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/pyproject.toml b/examples/colpali/pyproject.toml new file mode 100644 index 0000000..20c26a7 --- /dev/null +++ b/examples/colpali/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "colpali-engine", + "datasets", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/gensim/example.py b/examples/gensim/example.py new file mode 100644 index 0000000..2fc63c8 --- /dev/null +++ b/examples/gensim/example.py @@ -0,0 +1,35 @@ +from gensim.corpora.dictionary import Dictionary +from gensim.models import LdaModel +from gensim.utils import simple_preprocess +from pgvector import Vector +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' +] + +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) + +for content, bow in zip(input, corpus): + 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 +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]) diff --git a/examples/gensim/pyproject.toml b/examples/gensim/pyproject.toml new file mode 100644 index 0000000..4d2865a --- /dev/null +++ b/examples/gensim/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "gensim", + "numpy", + "pgvector", + "psycopg[binary]", + "scipy<1.13" +] diff --git a/examples/hybrid_search/cross_encoder.py b/examples/hybrid_search/cross_encoder.py new file mode 100644 index 0000000..c4e6eea --- /dev/null +++ b/examples/hybrid_search/cross_encoder.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 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 documents (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 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 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() + + +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()) diff --git a/examples/hybrid_search/pyproject.toml b/examples/hybrid_search/pyproject.toml new file mode 100644 index 0000000..130f287 --- /dev/null +++ b/examples/hybrid_search/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/hybrid_search/rrf.py b/examples/hybrid_search/rrf.py new file mode 100644 index 0000000..74d60bf --- /dev/null +++ b/examples/hybrid_search/rrf.py @@ -0,0 +1,52 @@ +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 5 +""" +query = 'growling bear' +embedding = model.encode(query) +k = 60 +results = conn.execute(sql, {'query': query, 'embedding': embedding, 'k': k}).fetchall() +for row in results: + print('document:', row[0], 'RRF score:', row[1]) diff --git a/examples/pytorch_image_search.py b/examples/image_search/example.py similarity index 54% rename from examples/pytorch_image_search.py rename to examples/image_search/example.py index 3551e16..b4cf131 100644 --- a/examples/pytorch_image_search.py +++ b/examples/image_search/example.py @@ -8,14 +8,11 @@ seed = True - # 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) - # load images transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), @@ -24,47 +21,49 @@ 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 -model = torchvision.models.resnet18(pretrained=True) +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 +# generate and store 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) - conn.execute('CREATE INDEX ON image USING ivfflat (embedding vector_cosine_ops)') - - -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() - - # 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 image 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) diff --git a/examples/image_search/pyproject.toml b/examples/image_search/pyproject.toml new file mode 100644 index 0000000..8ada0a7 --- /dev/null +++ b/examples/image_search/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "matplotlib", + "pgvector", + "psycopg[binary]", + "torch", + "torchvision", + "tqdm" +] diff --git a/examples/imagehash/example.py b/examples/imagehash/example.py new file mode 100644 index 0000000..595dcb5 --- /dev/null +++ b/examples/imagehash/example.py @@ -0,0 +1,48 @@ +from datasets import load_dataset +from imagehash import phash +import matplotlib.pyplot as plt +from pgvector import Bit +from pgvector.psycopg import register_vector +import psycopg + + +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('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))') + +print('Loading dataset') +dataset = load_dataset('ylecun/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([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 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/imagehash/pyproject.toml b/examples/imagehash/pyproject.toml new file mode 100644 index 0000000..4f0e2b4 --- /dev/null +++ b/examples/imagehash/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "datasets", + "imagehash", + "matplotlib", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/implicit/example.py b/examples/implicit/example.py new file mode 100644 index 0000000..a4bd0f4 --- /dev/null +++ b/examples/implicit/example.py @@ -0,0 +1,52 @@ +import implicit +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, Mapped, 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() + + +class Base(DeclarativeBase): + pass + + +class User(Base): + __tablename__ = 'user' + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) + + +class Item(Base): + __tablename__ = 'item' + + 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) +Base.metadata.create_all(engine) + +titles, ratings = get_movielens('100k') +model = implicit.als.AlternatingLeastSquares(factors=20) +model.fit(ratings) + +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.execute(insert(User), users) +session.execute(insert(Item), items) + +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)')).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/implicit/pyproject.toml b/examples/implicit/pyproject.toml new file mode 100644 index 0000000..ec305ee --- /dev/null +++ b/examples/implicit/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "h5py", + "implicit", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/implicit_recs.py b/examples/implicit_recs.py deleted file mode 100644 index 10a1fe0..0000000 --- a/examples/implicit_recs.py +++ /dev/null @@ -1,51 +0,0 @@ -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 - -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example', future=True) -with engine.connect() as conn: - conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - conn.commit() - -Base = declarative_base() - - -class User(Base): - __tablename__ = 'user' - - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) - - -class Item(Base): - __tablename__ = 'item' - - id = Column(Integer, primary_key=True) - title = Column(String) - factors = Column(Vector(20)) - - -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) - -titles, ratings = get_movielens('100k') -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].decode('utf-8'), 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() -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() -print('item-based recs:', [item.title for item in items]) diff --git a/examples/lightfm/example.py b/examples/lightfm/example.py new file mode 100644 index 0000000..728adc1 --- /dev/null +++ b/examples/lightfm/example.py @@ -0,0 +1,58 @@ +from lightfm import LightFM +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, Mapped, 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() + + +class Base(DeclarativeBase): + pass + + +class User(Base): + __tablename__ = 'user' + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) + + +class Item(Base): + __tablename__ = 'item' + + 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) +Base.metadata.create_all(engine) + +data = fetch_movielens(min_rating=5.0) +model = LightFM(loss='warp', no_components=20) +model.fit(data['train'], epochs=30) + +user_biases, user_factors = model.get_user_representations() +item_biases, item_factors = model.get_item_representations() + +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.execute(insert(User), users) +session.execute(insert(Item), items) + +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)')).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/pyproject.toml b/examples/lightfm/pyproject.toml new file mode 100644 index 0000000..047e3f1 --- /dev/null +++ b/examples/lightfm/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "lightfm", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/lightfm_recs.py b/examples/lightfm_recs.py deleted file mode 100644 index 9b9bc5f..0000000 --- a/examples/lightfm_recs.py +++ /dev/null @@ -1,56 +0,0 @@ -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 - -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example', future=True) -with engine.connect() as conn: - conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - conn.commit() - -Base = declarative_base() - - -class User(Base): - __tablename__ = 'user' - - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) - - -class Item(Base): - __tablename__ = 'item' - - id = Column(Integer, primary_key=True) - title = Column(String) - factors = Column(Vector(20)) - bias = Column(Float) - - -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) - -data = fetch_movielens(min_rating=5.0) -model = LightFM(loss='warp', no_components=20) -model.fit(data['train'], epochs=30) - -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)] - -session = Session(engine) -session.bulk_insert_mappings(User, users) -session.bulk_insert_mappings(Item, items) -session.commit() - -user = session.query(User).get(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() -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() -print('item-based recs:', [item.title for item in items]) diff --git a/examples/loading/example.py b/examples/loading/example.py new file mode 100644 index 0000000..4c0e579 --- /dev/null +++ b/examples/loading/example.py @@ -0,0 +1,46 @@ +import numpy as np +from pgvector import Vector +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 +print(f'Loading {len(embeddings)} rows') +cur = conn.cursor() +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): + copy.write_row([Vector(embedding)]) + + # show progress + if i % 10000 == 0: + print('.', end='', flush=True) + +print('\nSuccess!') + +# create any indexes *after* loading initial data (skipping for this example) +create_index = False +if create_index: + 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') diff --git a/examples/loading/pyproject.toml b/examples/loading/pyproject.toml new file mode 100644 index 0000000..7247ec0 --- /dev/null +++ b/examples/loading/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] + +[tool.uv.sources] +pgvector = { path = "../..", editable = true } diff --git a/examples/openai/example.py b/examples/openai/example.py new file mode 100644 index 0000000..ba1fa7a --- /dev/null +++ b/examples/openai/example.py @@ -0,0 +1,34 @@ +from openai import OpenAI +from pgvector import Vector +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(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' +] +embeddings = embed(input) +for content, embedding in zip(input, embeddings): + 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', (Vector(query_embedding),)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/openai/halfvec.py b/examples/openai/halfvec.py new file mode 100644 index 0000000..b2234f6 --- /dev/null +++ b/examples/openai/halfvec.py @@ -0,0 +1,35 @@ +from openai import OpenAI +from pgvector import HalfVector +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 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]) diff --git a/examples/openai/pyproject.toml b/examples/openai/pyproject.toml new file mode 100644 index 0000000..5665e02 --- /dev/null +++ b/examples/openai/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "openai", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/rag/example.py b/examples/rag/example.py new file mode 100644 index 0000000..9d5a784 --- /dev/null +++ b/examples/rag/example.py @@ -0,0 +1,64 @@ +# Run: +# ollama serve +# ollama pull llama3.2 +# ollama pull nomic-embed-text + +import ollama +from pathlib import Path +from pgvector import Vector +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 +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() +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/pyproject.toml b/examples/rag/pyproject.toml new file mode 100644 index 0000000..2e74790 --- /dev/null +++ b/examples/rag/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "ollama", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/rdkit/example.py b/examples/rdkit/example.py new file mode 100644 index 0000000..dec2993 --- /dev/null +++ b/examples/rdkit/example.py @@ -0,0 +1,33 @@ +# good resource +# https://www.rdkit.org/docs/GettingStartedInPython.html#morgan-fingerprints-circular-fingerprints + +from pgvector import Bit +from pgvector.psycopg import register_vector +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) diff --git a/examples/rdkit/pyproject.toml b/examples/rdkit/pyproject.toml new file mode 100644 index 0000000..b717809 --- /dev/null +++ b/examples/rdkit/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "rdkit" +] diff --git a/examples/sentence_transformers/example.py b/examples/sentence_transformers/example.py new file mode 100644 index 0000000..50997d9 --- /dev/null +++ b/examples/sentence_transformers/example.py @@ -0,0 +1,28 @@ +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))') + +model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +embeddings = model.encode(input) +for content, embedding in zip(input, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) + +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/sentence_transformers/pyproject.toml b/examples/sentence_transformers/pyproject.toml new file mode 100644 index 0000000..130f287 --- /dev/null +++ b/examples/sentence_transformers/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/sparse_search/example.py b/examples/sparse_search/example.py new file mode 100644 index 0000000..e3fec01 --- /dev/null +++ b/examples/sparse_search/example.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 + +from pgvector import SparseVector +from pgvector.psycopg import register_vector +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 embed(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 = 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 = 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]) diff --git a/examples/sparse_search/pyproject.toml b/examples/sparse_search/pyproject.toml new file mode 100644 index 0000000..d9adf93 --- /dev/null +++ b/examples/sparse_search/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]", + "torch", + "transformers" +] diff --git a/examples/surprise/example.py b/examples/surprise/example.py new file mode 100644 index 0000000..699ea31 --- /dev/null +++ b/examples/surprise/example.py @@ -0,0 +1,51 @@ +from pgvector.sqlalchemy import VECTOR +from sqlalchemy import create_engine, insert, select, text, Integer +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session +from surprise import Dataset, SVD + +engine = create_engine('postgresql+psycopg://localhost/pgvector_example') +with engine.connect() as conn: + conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) + conn.commit() + + +class Base(DeclarativeBase): + pass + + +class User(Base): + __tablename__ = 'user' + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) + + +class Item(Base): + __tablename__ = 'item' + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) + + +Base.metadata.drop_all(engine) +Base.metadata.create_all(engine) + +data = Dataset.load_builtin('ml-100k') +trainset = data.build_full_trainset() +algo = SVD(n_factors=20, biased=False) +algo.fit(trainset) + +users = [dict(id=trainset.to_raw_uid(i), factors=algo.pu[i]) for i in trainset.all_users()] +items = [dict(id=trainset.to_raw_iid(i), factors=algo.qi[i]) for i in trainset.all_items()] + +session = Session(engine) +session.execute(insert(User), users) +session.execute(insert(Item), items) + +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_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]) diff --git a/examples/surprise/pyproject.toml b/examples/surprise/pyproject.toml new file mode 100644 index 0000000..bd0cff2 --- /dev/null +++ b/examples/surprise/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "scikit-surprise", + "SQLAlchemy" +] diff --git a/examples/surprise_recs.py b/examples/surprise_recs.py deleted file mode 100644 index 1d957c3..0000000 --- a/examples/surprise_recs.py +++ /dev/null @@ -1,50 +0,0 @@ -from pgvector.sqlalchemy import Vector -from sqlalchemy import create_engine, text, Column, Integer -from sqlalchemy.orm import declarative_base, Session -from surprise import Dataset, SVD - -engine = create_engine('postgresql+psycopg2://localhost/pgvector_example', future=True) -with engine.connect() as conn: - conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - conn.commit() - -Base = declarative_base() - - -class User(Base): - __tablename__ = 'user' - - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) - - -class Item(Base): - __tablename__ = 'item' - - id = Column(Integer, primary_key=True) - factors = Column(Vector(20)) - - -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) - -data = Dataset.load_builtin('ml-100k') -trainset = data.build_full_trainset() -algo = SVD(n_factors=20, biased=False) -algo.fit(trainset) - -users = [dict(id=trainset.to_raw_uid(i), factors=algo.pu[i]) for i in trainset.all_users()] -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() - -user = session.query(User).get(1) -items = session.query(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5).all() -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() -print('item-based recs:', [item.id for item in items]) 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.py b/pgvector/_utils.py new file mode 100644 index 0000000..4cafc32 --- /dev/null +++ b/pgvector/_utils.py @@ -0,0 +1,23 @@ +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]] +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: + 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/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 5797327..b3b4440 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,12 +1,5 @@ -from ..utils import from_db, from_db_binary, to_db, to_db_binary +from .register import register_vector -__all__ = ['register_vector'] - - -async def register_vector(conn): - await conn.set_type_codec( - 'vector', - encoder=to_db_binary, - decoder=from_db_binary, - format='binary' - ) +__all__ = [ + 'register_vector' +] diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py new file mode 100644 index 0000000..0a2715f --- /dev/null +++ b/pgvector/asyncpg/register.py @@ -0,0 +1,32 @@ +from asyncpg import Connection +from .. import Vector, HalfVector, SparseVector + + +async def register_vector(conn: Connection, /, *, schema: str = 'public') -> None: + await conn.set_type_codec( + 'vector', + schema=schema, + encoder=lambda v: (v if isinstance(v, Vector) else Vector(v)).to_binary(), + decoder=Vector.from_binary, + format='binary' + ) + + try: + await conn.set_type_codec( + 'halfvec', + schema=schema, + encoder=lambda v: (v if isinstance(v, HalfVector) else HalfVector(v)).to_binary(), + decoder=HalfVector.from_binary, + format='binary' + ) + + await conn.set_type_codec( + 'sparsevec', + schema=schema, + encoder=lambda v: (v if isinstance(v, SparseVector) else SparseVector(v)).to_binary(), + decoder=SparseVector.from_binary, + format='binary' + ) + except ValueError as e: + if not str(e).startswith('unknown type:'): + raise e diff --git a/pgvector/bit.py b/pgvector/bit.py new file mode 100644 index 0000000..a9370e4 --- /dev/null +++ b/pgvector/bit.py @@ -0,0 +1,97 @@ +from __future__ import annotations +from struct import pack, unpack_from +from typing import TYPE_CHECKING +from ._utils import is_ndarray + +if TYPE_CHECKING: + import numpy as np + + +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) + self._data = value + elif isinstance(value, (list, str)): + if isinstance(value, list): + bits = {True: '1', False: '0'} + try: + 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 + 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 + + 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') + + self._length = len(value) + self._data = np.packbits(value).tobytes() # type: ignore + else: + raise ValueError('expected bytes, str, list, or ndarray') + + def __repr__(self) -> str: + return f'Bit({self.to_text()})' + + def __eq__(self, other: object, /) -> bool: + 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 + 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: + return ''.join(format(v, '08b') for v in self._data)[:self._length] + + def to_binary(self) -> bytes: + return pack('>i', self._length) + self._data + + @classmethod + 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: + length, = unpack_from('>i', value) + data = memoryview(value)[4:].tobytes() + + if length < 0 or len(data) != (length + 7) // 8: + raise ValueError('invalid length') + + bit = cls.__new__(cls) + bit._length = length + bit._data = data + return bit diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 02e77cd..3ed0f83 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -1,84 +1,23 @@ -from django.contrib.postgres.operations import CreateExtension -from django.contrib.postgres.indexes import PostgresIndex -from django.db.models import Field, FloatField, Func, Value -from ..utils import from_db, to_db - -__all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance'] - - -class VectorExtension(CreateExtension): - def __init__(self): - self.name = 'vector' - - -# https://docs.djangoproject.com/en/3.2/howto/custom-model-fields/ -class VectorField(Field): - description = 'Vector' - - 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 from_db(value) - - def to_python(self, value): - return from_db(value) - - def get_prep_value(self, value): - return to_db(value) - - -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 DistanceBase(Func): - output_field = FloatField() - - def __init__(self, expression, vector, **extra): - if not hasattr(vector, 'resolve_expression'): - vector = Value(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 = ' <=> ' +from .bit import BitField +from .extensions import VectorExtension +from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance +from .halfvec import HalfVectorField +from .indexes import IvfflatIndex, HnswIndex +from .sparsevec import SparseVectorField +from .vector import VectorField + +__all__ = [ + 'VectorExtension', + 'VectorField', + 'HalfVectorField', + 'BitField', + 'SparseVectorField', + 'IvfflatIndex', + 'HnswIndex', + 'L2Distance', + 'MaxInnerProduct', + 'CosineDistance', + 'L1Distance', + 'HammingDistance', + 'JaccardDistance' +] diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py new file mode 100644 index 0000000..eab0534 --- /dev/null +++ b/pgvector/django/bit.py @@ -0,0 +1,37 @@ +from django import forms +from django.db.models import Field +from typing import Any + + +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ +class BitField(Field): + description = 'Bit string' + + def __init__(self, *args: Any, length: int | None = None, **kwargs: Any) -> None: + self.length = length + super().__init__(*args, **kwargs) + + def deconstruct(self) -> tuple[Any, Any, Any, Any]: + 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: Any) -> str: + if self.length is None: + return 'bit' + return 'bit(%d)' % self.length + + 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): + 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 new file mode 100644 index 0000000..d75884f --- /dev/null +++ b/pgvector/django/extensions.py @@ -0,0 +1,11 @@ +from django import VERSION +from django.contrib.postgres.operations import CreateExtension +from typing import Any + + +class VectorExtension(CreateExtension): + def __init__(self, hints: Any = None) -> None: + if VERSION[0] >= 6: + super().__init__('vector', hints=hints) # type: ignore + else: + self.name = 'vector' diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py new file mode 100644 index 0000000..352c631 --- /dev/null +++ b/pgvector/django/functions.py @@ -0,0 +1,58 @@ +from django.db.models import FloatField, Func, Value +from .. import Vector, HalfVector, SparseVector +from typing import Any + + +class DistanceBase(Func): + output_field = FloatField() # type: ignore + + 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()) + elif vector is not None: + vector = Value(Vector(vector).to_text()) + + # prevent error with unhashable types + self._constructor_args = ((expression, vector), extra) + + super().__init__(expression, vector, **extra) + + +class BitDistanceBase(Func): + output_field = FloatField() # type: ignore + + def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None: + if not hasattr(vector, 'resolve_expression'): + vector = Value(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 = ' <+> ' + + +class HammingDistance(BitDistanceBase): + function = '' + arg_joiner = ' <~> ' + + +class JaccardDistance(BitDistanceBase): + function = '' + arg_joiner = ' <%%> ' diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py new file mode 100644 index 0000000..1ab4a47 --- /dev/null +++ b/pgvector/django/halfvec.py @@ -0,0 +1,54 @@ +from django import forms +from django.db.models import Field +from typing import Any +from .. import HalfVector + + +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ +class HalfVectorField(Field): + description = 'Half vector' + empty_strings_allowed = False + + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None: + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def deconstruct(self) -> tuple[Any, Any, Any, Any]: + 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: Any) -> str: + if self.dimensions is None: + return 'halfvec' + return 'halfvec(%d)' % self.dimensions + + def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: + return HalfVector._from_db(value) + + def to_python(self, value: Any) -> list[float] | None: + if isinstance(value, list): + return value + return HalfVector._from_db(value) + + def get_prep_value(self, value: Any) -> str | None: + return HalfVector._to_db(value) + + 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( + form_class=HalfVectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) + + +class HalfVectorFormField(forms.CharField): + 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 new file mode 100644 index 0000000..e095351 --- /dev/null +++ b/pgvector/django/indexes.py @@ -0,0 +1,47 @@ +from django.contrib.postgres.indexes import PostgresIndex +from typing import Any + + +class IvfflatIndex(PostgresIndex): + suffix = 'ivfflat' + + def __init__(self, *expressions: Any, lists: int | None = None, **kwargs: Any) -> None: + self.lists = lists + super().__init__(*expressions, **kwargs) + + def deconstruct(self) -> tuple[Any, Any, Any]: + path, args, kwargs = super().deconstruct() + if self.lists is not None: + kwargs['lists'] = self.lists + return path, args, kwargs + + def get_with_params(self) -> list[str]: + 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: 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) -> tuple[Any, Any, Any]: + 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) -> list[str]: + 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 diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py new file mode 100644 index 0000000..a10c133 --- /dev/null +++ b/pgvector/django/sparsevec.py @@ -0,0 +1,61 @@ +from django import forms +from django.db.models import Field +from typing import Any +from .. import SparseVector + + +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ +class SparseVectorField(Field): + description = 'Sparse vector' + empty_strings_allowed = False + + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any): + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def deconstruct(self) -> tuple[Any, Any, Any, Any]: + 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: Any) -> str: + if self.dimensions is None: + return 'sparsevec' + return 'sparsevec(%d)' % self.dimensions + + def from_db_value(self, value: Any, expression: Any, connection: Any) -> SparseVector | None: + return SparseVector._from_db(value) + + def to_python(self, value: Any) -> SparseVector | None: + return SparseVector._from_db(value) + + def get_prep_value(self, value: Any) -> str | None: + return SparseVector._to_db(value) + + 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( + form_class=SparseVectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) + + +class SparseVectorWidget(forms.TextInput): + def format_value(self, value: Any) -> Any: + if isinstance(value, SparseVector): + value = value.to_text() + return super().format_value(value) + + +class SparseVectorFormField(forms.CharField): + widget = SparseVectorWidget + + 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 new file mode 100644 index 0000000..42970c6 --- /dev/null +++ b/pgvector/django/vector.py @@ -0,0 +1,54 @@ +from django import forms +from django.db.models import Field +from typing import Any +from .. import Vector + + +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ +class VectorField(Field): + description = 'Vector' + empty_strings_allowed = False + + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None: + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def deconstruct(self) -> tuple[Any, Any, Any, Any]: + 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: Any) -> str: + if self.dimensions is None: + return 'vector' + return 'vector(%d)' % self.dimensions + + def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: + return Vector._from_db(value) + + def to_python(self, value: Any) -> list[float] | None: + if isinstance(value, list): + return value + return Vector._from_db(value) + + def get_prep_value(self, value: Any) -> str | None: + return Vector._to_db(value) + + 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( + form_class=VectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) + + +class VectorFormField(forms.CharField): + def to_python(self, value: Any) -> Any: + if isinstance(value, str) and value == '': + return None + return super().to_python(value) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py new file mode 100644 index 0000000..2b3b325 --- /dev/null +++ b/pgvector/halfvec.py @@ -0,0 +1,117 @@ +from __future__ import annotations +from array import array +import struct +import sys +from typing import TYPE_CHECKING +from ._utils import is_ndarray, ndarray + +if TYPE_CHECKING: + import numpy as np + + +class HalfVector: + _value: array[int] # uses uint16 since no float16 + + def __init__(self, value: list[float] | ndarray, /) -> None: + if isinstance(value, list): + dim = len(value) + try: + self._value = array('H', struct.pack(f'{dim}e', *value)) + except struct.error: + raise ValueError('expected list[float]') + elif is_ndarray(value): + import numpy as np + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + arr = 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') + + def __repr__(self) -> str: + return f'HalfVector({self.to_list()})' + + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, self.__class__): + return NotImplemented + return self.to_list() == other.to_list() + + def dimensions(self) -> int: + return len(self._value) + + 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]]: + import numpy as np + return np.frombuffer(self._value, dtype=np.float16) + + def to_text(self) -> str: + return f'[{",".join([str(v) for v in self.to_list()])}]' + + def to_binary(self) -> bytes: + if sys.byteorder == 'big': + value = self._value + 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 + def from_text(cls, value: str, /) -> HalfVector: + return cls(cls._from_text(value)) + + @classmethod + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: + dim, unused = struct.unpack_from('>hh', value) + data = memoryview(value)[4:] + + if len(data) != 2 * dim: + raise ValueError('invalid length') + + if unused != 0: + raise ValueError('expected unused to be 0') + + arr = array('H') + arr.frombytes(data) + if sys.byteorder != 'big': + arr.byteswap() + + vec = cls.__new__(cls) + vec._value = arr + return vec + + @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 + def _to_db(cls, value: list[float] | ndarray | HalfVector | None, /) -> str | None: + if value is None: + return value + + # fewer allocations for lists + if isinstance(value, list): + return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore + + if not isinstance(value, HalfVector): + value = cls(value) + + return value.to_text() + + @classmethod + def _from_db(cls, value: str | HalfVector | None, /) -> list[float] | None: + if value is None: + return value + + if isinstance(value, HalfVector): + return value.to_list() + + return cls._from_text(value) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py new file mode 100644 index 0000000..f557517 --- /dev/null +++ b/pgvector/peewee/__init__.py @@ -0,0 +1,11 @@ +from .bit import FixedBitField +from .halfvec import HalfVectorField +from .sparsevec import SparseVectorField +from .vector import VectorField + +__all__ = [ + 'VectorField', + 'HalfVectorField', + 'FixedBitField', + 'SparseVectorField' +] diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py new file mode 100644 index 0000000..fcc7d8f --- /dev/null +++ b/pgvector/peewee/bit.py @@ -0,0 +1,22 @@ +from peewee import Expression, Field +from typing import Any + + +class FixedBitField(Field): + field_type = 'bit' + + def __init__(self, max_length: int | None = None, *args: Any, **kwargs: Any) -> None: + self.max_length = max_length + super().__init__(*args, **kwargs) + + 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: + return Expression(lhs=self, op=op, rhs=self.to_value(vector)) + + def hamming_distance(self, vector: object, /) -> Expression: + return self._distance('<~>', vector) + + def jaccard_distance(self, vector: object, /) -> Expression: + return self._distance('<%%>', vector) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py new file mode 100644 index 0000000..5e1e92c --- /dev/null +++ b/pgvector/peewee/halfvec.py @@ -0,0 +1,35 @@ +from peewee import Expression, Field +from typing import Any +from .. import HalfVector + + +class HalfVectorField(Field): + field_type = 'halfvec' + + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None + + def db_value(self, value: Any) -> str | None: + return HalfVector._to_db(value) + + def python_value(self, value: Any) -> list[float] | None: + return HalfVector._from_db(value) + + 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: + return self._distance('<->', vector) + + def max_inner_product(self, vector: object, /) -> Expression: + return self._distance('<#>', vector) + + def cosine_distance(self, vector: object, /) -> Expression: + return self._distance('<=>', vector) + + def l1_distance(self, vector: object, /) -> Expression: + return self._distance('<+>', vector) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py new file mode 100644 index 0000000..dd83644 --- /dev/null +++ b/pgvector/peewee/sparsevec.py @@ -0,0 +1,35 @@ +from peewee import Expression, Field +from typing import Any +from .. import SparseVector + + +class SparseVectorField(Field): + field_type = 'sparsevec' + + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None + + def db_value(self, value: Any) -> str | None: + return SparseVector._to_db(value) + + def python_value(self, value: Any) -> SparseVector | None: + return SparseVector._from_db(value) + + 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: + return self._distance('<->', vector) + + def max_inner_product(self, vector: object, /) -> Expression: + return self._distance('<#>', vector) + + def cosine_distance(self, vector: object, /) -> Expression: + return self._distance('<=>', vector) + + def l1_distance(self, vector: object, /) -> Expression: + return self._distance('<+>', vector) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py new file mode 100644 index 0000000..a28b6b5 --- /dev/null +++ b/pgvector/peewee/vector.py @@ -0,0 +1,35 @@ +from peewee import Expression, Field +from typing import Any +from .. import Vector + + +class VectorField(Field): + field_type = 'vector' + + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: + self.dimensions = dimensions + super().__init__(*args, **kwargs) + + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None + + def db_value(self, value: Any) -> str | None: + return Vector._to_db(value) + + def python_value(self, value: Any) -> list[float] | None: + return Vector._from_db(value) + + 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: + return self._distance('<->', vector) + + def max_inner_product(self, vector: object, /) -> Expression: + return self._distance('<#>', vector) + + def cosine_distance(self, vector: object, /) -> Expression: + return self._distance('<=>', vector) + + def l1_distance(self, vector: object, /) -> Expression: + return self._distance('<+>', vector) 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..f3978d7 --- /dev/null +++ b/pgvector/pg8000/register.py @@ -0,0 +1,29 @@ +from pg8000.native import Connection +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 = 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: + raise RuntimeError('vector type not found in the database') + + 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, lambda v: Vector(v).to_text()) + except ImportError: + pass + + if 'halfvec' in type_info: + 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, lambda v: v.to_text()) + conn.register_in_adapter(type_info['sparsevec'], SparseVector.from_text) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index ecb8fe6..8922eda 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,56 +1,6 @@ -import psycopg -from psycopg.adapt import Loader, Dumper -from psycopg.pq import Format -from psycopg.types import TypeInfo -from ..utils import from_db, from_db_binary, to_db, to_db_binary +from .register import register_vector, register_vector_async -__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') - if info is None: - raise psycopg.ProgrammingError('vector type not found in the database') - info.register(context) - - adapters = context.adapters - adapters.register_dumper('numpy.ndarray', VectorDumper) - adapters.register_dumper('numpy.ndarray', VectorBinaryDumper) - adapters.register_loader(info.oid, VectorLoader) - adapters.register_loader(info.oid, VectorBinaryLoader) +__all__ = [ + 'register_vector', + 'register_vector_async', +] diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py new file mode 100644 index 0000000..4cb502f --- /dev/null +++ b/pgvector/psycopg/bit.py @@ -0,0 +1,45 @@ +from functools import cache +from psycopg import BaseConnection +from psycopg.abc import Buffer +from psycopg.adapt import Dumper +from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any +from .. import Bit + + +class BitDumper(Dumper): + format = Format.TEXT + + def dump(self, obj: Bit) -> Buffer | None: + return obj.to_text().encode('utf8') + + +class BitBinaryDumper(BitDumper): + format = Format.BINARY + + def dump(self, obj: Bit) -> Buffer | None: + return obj.to_binary() + + +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 + 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) -> type[BitDumper]: + return type('', (BitDumper,), {'oid': oid}) + + +@cache +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 new file mode 100644 index 0000000..05e7994 --- /dev/null +++ b/pgvector/psycopg/halfvec.py @@ -0,0 +1,62 @@ +from functools import cache +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 +from .. import HalfVector + + +class HalfVectorDumper(Dumper): + format = Format.TEXT + + def dump(self, obj: HalfVector) -> Buffer | None: + return obj.to_text().encode('utf8') + + +class HalfVectorBinaryDumper(HalfVectorDumper): + format = Format.BINARY + + def dump(self, obj: HalfVector) -> Buffer | None: + return obj.to_binary() + + +class HalfVectorLoader(Loader): + format = Format.TEXT + + def load(self, data: Buffer) -> HalfVector | None: + if isinstance(data, memoryview): + data = data.tobytes() + return HalfVector.from_text(data.decode('utf8')) + + +class HalfVectorBinaryLoader(HalfVectorLoader): + format = Format.BINARY + + def load(self, data: Buffer) -> HalfVector | None: + return HalfVector.from_binary(data) + + +def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None: + info.register(context) + + # add oid to anonymous class for set_types + 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) -> type[HalfVectorDumper]: + return type('', (HalfVectorDumper,), {'oid': oid}) + + +@cache +def halfvec_binary_dumper(oid: int) -> type[HalfVectorBinaryDumper]: + return type('', (HalfVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py new file mode 100644 index 0000000..ed95042 --- /dev/null +++ b/pgvector/psycopg/register.py @@ -0,0 +1,39 @@ +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: Connection[Any], /) -> None: + 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: AsyncConnection[Any], /) -> None: + 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/sparsevec.py b/pgvector/psycopg/sparsevec.py new file mode 100644 index 0000000..0c953bf --- /dev/null +++ b/pgvector/psycopg/sparsevec.py @@ -0,0 +1,62 @@ +from functools import cache +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 +from .. import SparseVector + + +class SparseVectorDumper(Dumper): + format = Format.TEXT + + def dump(self, obj: SparseVector) -> Buffer | None: + return obj.to_text().encode('utf8') + + +class SparseVectorBinaryDumper(SparseVectorDumper): + format = Format.BINARY + + def dump(self, obj: SparseVector) -> Buffer | None: + return obj.to_binary() + + +class SparseVectorLoader(Loader): + format = Format.TEXT + + def load(self, data: Buffer) -> SparseVector | None: + if isinstance(data, memoryview): + data = data.tobytes() + return SparseVector.from_text(data.decode('utf8')) + + +class SparseVectorBinaryLoader(SparseVectorLoader): + format = Format.BINARY + + def load(self, data: Buffer) -> SparseVector | None: + return SparseVector.from_binary(data) + + +def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None: + info.register(context) + + # add oid to anonymous class for set_types + 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) -> type[SparseVectorDumper]: + return type('', (SparseVectorDumper,), {'oid': oid}) + + +@cache +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 new file mode 100644 index 0000000..f4b66e4 --- /dev/null +++ b/pgvector/psycopg/vector.py @@ -0,0 +1,73 @@ +from __future__ import annotations +from functools import cache +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 Any +from .. import Vector +from .._utils import ndarray + + +class VectorDumper(Dumper): + format = Format.TEXT + + def dump(self, obj: Vector | ndarray) -> Buffer | None: + 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 | ndarray) -> Buffer | None: + if not isinstance(obj, Vector): + obj = Vector(obj) + return obj.to_binary() + + +class VectorLoader(Loader): + format = Format.TEXT + + def load(self, data: Buffer) -> Vector | None: + if isinstance(data, memoryview): + data = data.tobytes() + return Vector.from_text(data.decode('utf8')) + + +class VectorBinaryLoader(VectorLoader): + format = Format.BINARY + + def load(self, data: Buffer) -> Vector | None: + return Vector.from_binary(data) + + +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) + + # add oid to anonymous class for set_types + text_dumper = vector_text_dumper(info.oid) + binary_dumper = vector_binary_dumper(info.oid) + + 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) + + +@cache +def vector_text_dumper(oid: int) -> type[VectorDumper]: + return type('', (VectorDumper,), {'oid': oid}) + + +@cache +def vector_binary_dumper(oid: int) -> type[VectorBinaryDumper]: + return type('', (VectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 94cc081..b3b4440 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,32 +1,5 @@ -import numpy as np -import psycopg2 -from psycopg2.extensions import adapt, new_type, register_adapter, register_type -from ..utils import from_db, to_db +from .register import register_vector -__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 - - try: - cur.execute('SELECT NULL::vector') - oid = cur.description[0][1] - 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) +__all__ = [ + 'register_vector' +] diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py new file mode 100644 index 0000000..e1be330 --- /dev/null +++ b/pgvector/psycopg2/halfvec.py @@ -0,0 +1,30 @@ +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: HalfVector) -> None: + self._value = value + + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() + + +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) + + +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) + + 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 new file mode 100644 index 0000000..2a5af49 --- /dev/null +++ b/pgvector/psycopg2/register.py @@ -0,0 +1,27 @@ +import psycopg2 +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: 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 + + # 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('_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'], type_info['_vector'] if arrays else None, scope) + + if 'halfvec' in type_info: + 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'], type_info['_sparsevec'] if arrays else None, scope) diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py new file mode 100644 index 0000000..47dbc66 --- /dev/null +++ b/pgvector/psycopg2/sparsevec.py @@ -0,0 +1,30 @@ +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: SparseVector) -> None: + self._value = value + + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() + + +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) + + +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) + + 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 new file mode 100644 index 0000000..7d37648 --- /dev/null +++ b/pgvector/psycopg2/vector.py @@ -0,0 +1,40 @@ +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 +from .._utils import ndarray + + +class VectorAdapter: + def __init__(self, value: Vector | ndarray) -> None: + if not isinstance(value, Vector): + value = Vector(value) + self._value = value + + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() + + +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) + + +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) + + if array_oid is not None: + vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector) + register_type(vectorarray, scope) + + register_adapter(Vector, VectorAdapter) + + try: + import numpy as np + register_adapter(np.ndarray, VectorAdapter) + except ImportError: + pass 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 new file mode 100644 index 0000000..76e1356 --- /dev/null +++ b/pgvector/sparsevec.py @@ -0,0 +1,213 @@ +from __future__ import annotations +from struct import pack, unpack_from +from typing import TYPE_CHECKING, Final, cast, overload +from ._utils import is_sparse_array, ndarray + +if TYPE_CHECKING: + import numpy as np + from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix + + +class Sentinel: + pass + + +NO_DEFAULT: Final[Sentinel] = Sentinel() + + +class SparseVector: + _dim: int + _indices: list[int] + _values: list[float] + + @overload + 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: + if is_sparse_array(value): + if dimensions is not NO_DEFAULT: + raise ValueError('extra argument') + + self._from_sparse(value) # type: ignore + elif isinstance(value, dict): + if dimensions is NO_DEFAULT: + raise ValueError('missing dimensions') + + self._from_dict(value, dimensions) # type: ignore + else: + if dimensions is not NO_DEFAULT: + raise ValueError('extra argument') + + self._from_dense(value) # type: ignore + + def __repr__(self) -> str: + elements = dict(zip(self._indices, self._values)) + return f'SparseVector({elements}, {self._dim})' + + def __eq__(self, other: object, /) -> bool: + 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 + + def indices(self) -> list[int]: + return self._indices + + 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)) + + 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) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: + import numpy as np + + vec = np.zeros(self._dim, dtype=np.float32) + for i, v in zip(self._indices, self._values): + vec[i] = v + return vec + + def to_text(self) -> str: + 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) + return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) + + 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() + + 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 + + # 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] + elif len(shape) == 2 and shape[0] == 1: + self._dim = shape[1] + else: + raise ValueError('expected ndim to be 1') + + if hasattr(value, 'coords'): + # scipy 1.13+ + indices = value.coords[-1].tolist() + else: + indices = value.col.tolist() + + 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: + 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) + 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: + 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') + + 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) + 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(','): + i, v = e.split(':', 2) + indices.append(int(i) - 1) + values.append(float(v)) + return cls._from_parts(int(dim), indices, values) + + @classmethod + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> 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 = 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: + vec = cls.__new__(cls) + vec._dim = dim + vec._indices = indices + vec._values = values + # TODO uncomment in 0.6.0 + # vec._validate() + # vec._validate_sorted() + return vec + + @classmethod + def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector | None, /) -> str | None: + if value is None: + return value + + if not isinstance(value, SparseVector): + value = cls(value) + + return value.to_text() + + @classmethod + 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) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 6956691..593876f 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -1,37 +1,16 @@ -from sqlalchemy.types import UserDefinedType -from ..utils import from_db, to_db - -__all__ = ['Vector'] - - -class Vector(UserDefinedType): - cache_ok = True - - 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 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('<->')(other) - - def max_inner_product(self, other): - return self.op('<#>')(other) - - def cosine_distance(self, other): - return self.op('<=>')(other) +from .bit import BIT +from .functions import avg, sum +from .halfvec import HALFVEC +from .sparsevec import SPARSEVEC +from .vector import VECTOR +from .vector import VECTOR as Vector + +__all__ = [ + 'Vector', + 'VECTOR', + 'HALFVEC', + 'BIT', + 'SPARSEVEC', + 'avg', + 'sum' +] diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py new file mode 100644 index 0000000..a5af3b0 --- /dev/null +++ b/pgvector/sqlalchemy/bit.py @@ -0,0 +1,24 @@ +from sqlalchemy.dialects.postgresql.base import PGBit +from sqlalchemy.types import TypeDecorator, Float +from sqlalchemy import Operators +from typing import Any + + +class BIT(TypeDecorator[Any]): + impl = PGBit + cache_ok = True + + 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 + + class Comparator(TypeDecorator.Comparator[Any]): # type: ignore + 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/functions.py b/pgvector/sqlalchemy/functions.py new file mode 100644 index 0000000..ca4cae9 --- /dev/null +++ b/pgvector/sqlalchemy/functions.py @@ -0,0 +1,15 @@ +# 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[Any]): + inherit_cache = True + package = 'pgvector' + + +__all__ = [ + 'avg', + 'sum' +] diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py new file mode 100644 index 0000000..d82c5ac --- /dev/null +++ b/pgvector/sqlalchemy/halfvec.py @@ -0,0 +1,55 @@ +from sqlalchemy.dialects.postgresql.base import ischema_names +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String +from sqlalchemy import Dialect, Operators +from typing import Any +from .. import HalfVector + + +class HALFVEC(UserDefinedType[Any]): + cache_ok = True + _string = String() + + def __init__(self, dim: int | None = None) -> None: + super().__init__() + self.dim = dim + + def get_col_spec(self, **kw: Any) -> str: + if self.dim is None: + return 'HALFVEC' + return 'HALFVEC(%d)' % self.dim + + def bind_processor(self, dialect: Dialect) -> Any: + def process(value: Any) -> str | None: + 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)) # type: ignore + return process + + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: + def process(value: Any) -> list[float] | None: + return HalfVector._from_db(value) + return process + + class Comparator(TypeEngine.Comparator[Any]): # type: ignore + def l2_distance(self, other: object, /) -> Operators: + return self.op('<->', return_type=Float)(other) + + def max_inner_product(self, other: object, /) -> Operators: + return self.op('<#>', return_type=Float)(other) + + def cosine_distance(self, other: object, /) -> Operators: + return self.op('<=>', return_type=Float)(other) + + 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 new file mode 100644 index 0000000..5cc910d --- /dev/null +++ b/pgvector/sqlalchemy/sparsevec.py @@ -0,0 +1,55 @@ +from sqlalchemy.dialects.postgresql.base import ischema_names +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String +from sqlalchemy import Dialect, Operators +from typing import Any +from .. import SparseVector + + +class SPARSEVEC(UserDefinedType[Any]): + cache_ok = True + _string = String() + + def __init__(self, dim: int | None = None) -> None: + super().__init__() + self.dim = dim + + def get_col_spec(self, **kw: Any) -> str: + if self.dim is None: + return 'SPARSEVEC' + return 'SPARSEVEC(%d)' % self.dim + + def bind_processor(self, dialect: Dialect) -> Any: + def process(value: Any) -> str | None: + 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)) # type: ignore + return process + + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: + def process(value: Any) -> SparseVector | None: + return SparseVector._from_db(value) + return process + + class Comparator(TypeEngine.Comparator[Any]): # type: ignore + def l2_distance(self, other: object, /) -> Operators: + return self.op('<->', return_type=Float)(other) + + def max_inner_product(self, other: object, /) -> Operators: + return self.op('<#>', return_type=Float)(other) + + def cosine_distance(self, other: object, /) -> Operators: + return self.op('<=>', return_type=Float)(other) + + 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 new file mode 100644 index 0000000..4e7e4b3 --- /dev/null +++ b/pgvector/sqlalchemy/vector.py @@ -0,0 +1,55 @@ +from sqlalchemy.dialects.postgresql.base import ischema_names +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String +from sqlalchemy import Dialect, Operators +from typing import Any +from .. import Vector + + +class VECTOR(UserDefinedType[Any]): + cache_ok = True + _string = String() + + def __init__(self, dim: int | None = None) -> None: + super().__init__() + self.dim = dim + + def get_col_spec(self, **kw: Any) -> str: + if self.dim is None: + return 'VECTOR' + return 'VECTOR(%d)' % self.dim + + def bind_processor(self, dialect: Dialect) -> Any: + def process(value: Any) -> str | None: + 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)) # type: ignore + return process + + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: + def process(value: Any) -> list[float] | None: + return Vector._from_db(value) + return process + + class Comparator(TypeEngine.Comparator[Any]): # type: ignore + def l2_distance(self, other: object, /) -> Operators: + return self.op('<->', return_type=Float)(other) + + def max_inner_product(self, other: object, /) -> Operators: + return self.op('<#>', return_type=Float)(other) + + def cosine_distance(self, other: object, /) -> Operators: + return self.op('<=>', return_type=Float)(other) + + 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 diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py deleted file mode 100644 index 5640b60..0000000 --- a/pgvector/utils/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -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() diff --git a/pgvector/vector.py b/pgvector/vector.py new file mode 100644 index 0000000..d86cba0 --- /dev/null +++ b/pgvector/vector.py @@ -0,0 +1,115 @@ +from __future__ import annotations +from array import array +import struct +import sys +from typing import TYPE_CHECKING +from ._utils import is_ndarray, ndarray + +if TYPE_CHECKING: + import numpy as np + + +class Vector: + _value: array[float] + + def __init__(self, value: list[float] | ndarray, /) -> None: + if isinstance(value, list): + try: + self._value = array('f', value) + except TypeError: + raise ValueError('expected list[float]') + elif is_ndarray(value): + import numpy as np + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + arr = 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') + + def __repr__(self) -> str: + return f'Vector({self.to_list()})' + + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, self.__class__): + return NotImplemented + return self._value == other._value + + def dimensions(self) -> int: + return len(self._value) + + 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: + return f'[{",".join([str(v) for v in self._value])}]' + + def to_binary(self) -> bytes: + if sys.byteorder == 'big': + value = self._value + 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 + def from_text(cls, value: str, /) -> Vector: + return cls(cls._from_text(value)) + + @classmethod + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: + dim, unused = struct.unpack_from('>hh', value) + data = memoryview(value)[4:] + + if len(data) != 4 * dim: + raise ValueError('invalid length') + + if unused != 0: + raise ValueError('expected unused to be 0') + + arr = array('f') + arr.frombytes(data) + if sys.byteorder != 'big': + arr.byteswap() + + vec = cls.__new__(cls) + vec._value = arr + return vec + + @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 + def _to_db(cls, value: list[float] | ndarray | Vector | None, /) -> str | None: + if value is None: + return value + + # fewer allocations for lists + if isinstance(value, list): + return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore + + if not isinstance(value, Vector): + value = cls(value) + + return value.to_text() + + @classmethod + def _from_db(cls, value: str | Vector | None, /) -> list[float] | None: + if value is None: + return value + + if isinstance(value, Vector): + return value.to_list() + + return cls._from_text(value) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3ff09d9 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,87 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + +[project] +name = "pgvector" +version = "0.5.1" +description = "pgvector support for Python" +readme = "README.md" +authors = [ + {name = "Andrew Kane", email = "andrew@ankane.org"} +] +license = "MIT" +requires-python = ">= 3.10" +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", + "SQLAlchemy[asyncio]>=2", + "sqlmodel" +] +numpy = ["numpy"] +scipy = ["scipy"] +typing = [ + "mypy", + "ty", + "numpy", + "scipy" +] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" + +[tool.mypy] +files = [ + "pgvector", + "tests" +] +strict = true +disable_error_code = ["unused-ignore", "untyped-decorator"] +ignore_missing_imports = true +disallow_subclassing_any = false + +[tool.ty.src] +include = [ + "pgvector", + "tests" +] + +[tool.ty.rules] +unused-type-ignore-comment = "ignore" + +[tool.pyright] +include = [ + "pgvector", + "tests" +] +typeCheckingMode = "strict" +reportMissingTypeStubs = "none" +reportPrivateUsage = "none" +reportUnknownArgumentType = "none" +reportUnknownLambdaType = "none" +reportUnknownMemberType = "none" +reportUnknownVariableType = "none" + +[tool.pyrefly] +project-includes = [ + "pgvector/**/*.py", + "tests/**/*.py" +] +preset = "strict" + +[tool.pyrefly.errors] +missing-override-decorator = false 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/requirements.txt b/requirements.txt deleted file mode 100644 index 712f909..0000000 --- a/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ -asyncpg -Django -numpy -psycopg[binary] -psycopg2-binary -pytest -pytest-asyncio -SQLAlchemy diff --git a/setup.py b/setup.py deleted file mode 100644 index 0f4d7da..0000000 --- a/setup.py +++ /dev/null @@ -1,29 +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.1.5', - 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.psycopg', - 'pgvector.psycopg2', - 'pgvector.sqlalchemy', - 'pgvector.utils' - ], - python_requires='>=3.6', - install_requires=[ - 'numpy' - ], - zip_safe=False -) 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 12ab806..38e5fd0 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,32 +1,147 @@ -import asyncio import asyncpg -import numpy as np +from asyncpg import Connection +from pgvector import HalfVector, SparseVector, Vector from pgvector.asyncpg import register_vector import pytest +from .conftest import numpy as np class TestAsyncpg: - @pytest.mark.asyncio - async def test_works(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 conn.execute('DROP TABLE IF EXISTS item') - await conn.execute('CREATE TABLE item (id bigserial primary key, factors vector(3))') - await register_vector(conn) + return conn + + @pytest.mark.asyncio + 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))') + + embedding = Vector([1.5, 2, 3]) + embedding2 = [4.5, 5, 6] + 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) + + 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') + assert text_res[0]['embedding'] == '[1.5,2,3]' + + await conn.close() + + @pytest.mark.asyncio + 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))') + + 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) + + 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') + assert text_res[0]['embedding'] == '[1.5,2,3]' + + await conn.close() + + @pytest.mark.asyncio + 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))') + + # typing issue + # 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) - factors = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO item (factors) VALUES ($1), (NULL)", factors) + 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 - 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 + # ensures binary format is correct + 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() + + @pytest.mark.asyncio + 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))') + + embedding = SparseVector([1.5, 2, 3]) + 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 + 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 asyncpg_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_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[])') + + embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] + 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]) + + 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] + assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] + + await conn.close() + + @pytest.mark.asyncio + async def test_pool(self) -> None: + async def init(conn: Connection) -> None: + 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 asyncpg_items') + await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') + + 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) + + 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_bit.py b/tests/test_bit.py new file mode 100644 index 0000000..9179cd3 --- /dev/null +++ b/tests/test_bit.py @@ -0,0 +1,144 @@ +from pgvector import Bit +import pytest +import random +import struct +from .conftest import numpy as np + + +class TestBit: + def test_init_list(self) -> None: + assert Bit([True, False, True]).to_list() == [True, False, True] + + def test_init_list_empty(self) -> None: + assert Bit([]).to_list() == [] + + 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_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_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_init_str(self) -> None: + assert Bit('101').to_list() == [True, False, True] + + def test_init_str_empty(self) -> None: + assert Bit('').to_list() == [] + + def test_init_str_two(self) -> None: + with pytest.raises(ValueError) as error: + Bit('201') + assert str(error.value) == 'expected bit string' + + 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_init_bytes_empty(self) -> None: + assert Bit(b'').to_text() == '' + + def test_init_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) + + 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_init_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' + + def test_init_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 + assert str(error.value) == 'expected elements to be boolean' + + 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_init_random(self) -> None: + value = ''.join(random.choices(['0', '1'], k=random.randint(1024, 2048))) + assert Bit(value).to_text() == value + + def test_repr(self) -> None: + assert repr(Bit([True, False, True])) == 'Bit(101)' + assert str(Bit([True, False, True])) == 'Bit(101)' + + 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 + + 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') + + 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: + 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 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + Bit.from_binary(b'') + + def test_from_binary_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>iB', 9, 0)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + 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'): + Bit.from_binary(struct.pack('>i', -1)) diff --git a/tests/test_django.py b/tests/test_django.py index 95e25f2..e5f0c6f 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,10 +1,19 @@ 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 +from django.db.models.functions import Cast from django.db.migrations.loader import MigrationLoader -import numpy as np +from django.forms import ModelForm +from math import sqrt +import os import pgvector.django -from pgvector.django import VectorExtension, VectorField, IvfflatIndex, L2Distance, MaxInnerProduct, CosineDistance +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 settings.configure( DATABASES={ @@ -12,22 +21,66 @@ 'ENGINE': 'django.db.backends.postgresql', 'NAME': 'pgvector_python_test', } - } + }, + DEBUG=('VERBOSE' in os.environ), + LOGGING={ + 'version': 1, + 'handlers': { + 'console': { + 'class': 'logging.StreamHandler' + } + }, + 'loggers': { + 'django.db.backends': { + 'handlers': ['console'], + 'level': 'DEBUG' + }, + 'django.db.backends.schema': { + '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() class Item(models.Model): - factors = VectorField(dimensions=3) + embedding = VectorField(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 = 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) + + objects: models.Manager # for typing class Meta: - app_label = 'myapp' + app_label = 'django_app' indexes = [ IvfflatIndex( - name='my_index', - fields=['factors'], + name='ivfflat_idx', + fields=['embedding'], lists=100, opclasses=['vector_l2_ops'] + ), + HnswIndex( + name='hnsw_idx', + fields=['embedding'], + 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 ) ] @@ -35,70 +88,420 @@ class Meta: class Migration(migrations.Migration): initial = True - dependencies = [ - ] - operations = [ VectorExtension(), migrations.CreateModel( 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, 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.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( 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='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']), + ), + 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'), ) ] # 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: - cursor.execute("DROP TABLE IF EXISTS myapp_item") + cursor.execute('DROP TABLE IF EXISTS django_app_item') cursor.execute('\n'.join(sql_statements)) -def create_items(): - vectors = [ - [1, 1, 1], - [2, 2, 2], - [1, 1, 2] - ] - for i, v in enumerate(vectors): - item = Item(id=i + 1, factors=v) - item.save() +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() + + +class VectorForm(ModelForm): + class Meta: + model = Item + fields = ['embedding'] + + +class HalfVectorForm(ModelForm): + class Meta: + model = Item + fields = ['half_embedding'] + + +class BitForm(ModelForm): + class Meta: + model = Item + fields = ['binary_embedding'] + + +class SparseVectorForm(ModelForm): + class Meta: + model = Item + fields = ['sparse_embedding'] class TestDjango: - def setup_method(self, test_method): + def setup_method(self) -> None: Item.objects.all().delete() - def test_works(self): - item = Item(id=1, factors=[1, 2, 3]) - item.save() + def test_vector(self) -> None: + Item(id=1, embedding=[1, 2, 3]).save() + 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) + 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) + 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) + 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) + 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) + 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) + 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) + assert item.half_embedding == [1, 2, 3] + + 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) + 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) + 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) + 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) + 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) + 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) + 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) + 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) - assert item.id == 1 - assert np.array_equal(item.factors, np.array([1, 2, 3])) - assert item.factors.dtype == np.float32 + assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() - items = Item.objects.order_by(L2Distance('factors', [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_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() - items = Item.objects.order_by(MaxInnerProduct('factors', [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_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() - items = Item.objects.order_by(CosineDistance('factors', [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) -> 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) -> 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) -> None: + 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 avg == [2.5, 3.5, 4.5] + + def test_vector_sum(self) -> 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() + 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'] + 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 == [2.5, 3.5, 4.5] + + 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() + Item(half_embedding=[4, 5, 6]).save() + sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] + assert sum == [5, 7, 9] + + def test_serialization(self) -> None: + create_items() + 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: + get_model.return_value = Item + for obj in serializers.deserialize(format, data): + obj.save() + + def test_vector_form(self) -> None: + form = VectorForm(data={'embedding': '[1, 2, 3]'}) + assert form.is_valid() + 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 str(form.as_div()) + + 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]'}) + assert form.has_changed() + assert form.is_valid() + assert form.save() + 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) + form = VectorForm(instance=item, data={'embedding': ''}) + assert form.is_valid() + assert form.save() + assert Item.objects.get(pk=1).embedding is 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 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 str(form.as_div()) + + 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]'}) + assert form.has_changed() + assert form.is_valid() + assert form.save() + 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) + form = HalfVectorForm(instance=item, data={'half_embedding': ''}) + assert form.is_valid() + assert form.save() + assert Item.objects.get(pk=1).half_embedding is None + + def test_bit_form(self) -> None: + form = BitForm(data={'binary_embedding': '101'}) + assert form.is_valid() + 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 str(form.as_div()) + + 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'}) + assert form.has_changed() + assert form.is_valid() + assert form.save() + 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) + 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) -> 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 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) + 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() + 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]) + + 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': ''}) + assert form.is_valid() + assert form.save() + 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]) + + 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) -> None: + Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]]).save() + + with connection.cursor() as cursor: + from pgvector.psycopg import register_vector + register_vector(cursor.connection) + + # this fails if the driver does not cast arrays + item = Item.objects.get(pk=1) + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] + + 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() + 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) -> 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() + 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] + + 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) + 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_half_vector.py b/tests/test_half_vector.py new file mode 100644 index 0000000..a907c94 --- /dev/null +++ b/tests/test_half_vector.py @@ -0,0 +1,138 @@ +from pgvector import HalfVector +import pytest +import struct +from .conftest import numpy as np + + +class TestHalfVector: + 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_init_list_empty(self) -> None: + assert HalfVector([]).to_list() == [] + + 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_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_init_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 + assert HalfVector(arr).to_numpy().dtype == np.float16 + + 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() + + 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) + + def test_init_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) -> 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_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 + assert HalfVector([0.0]) == HalfVector([-0.0]) + + 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]) + + 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([]) + + 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_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]) + + 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_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hhe', 2, 0, 1)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + 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'): + 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_peewee.py b/tests/test_peewee.py new file mode 100644 index 0000000..59910ad --- /dev/null +++ b/tests/test_peewee.py @@ -0,0 +1,224 @@ +from math import sqrt +from peewee import Model, PostgresqlDatabase, fn +from pgvector import SparseVector, Vector +from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField + +db = PostgresqlDatabase('pgvector_python_test') + + +class BaseModel(Model): + class Meta: + database = db + + +class Item(BaseModel): + 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) + sparse_embedding = SparseVectorField(dimensions=3, null=True) + distance: float # for typing + + class Meta: # type: ignore + table_name = 'peewee_item' + + +Item.add_index('embedding vector_l2_ops', using='hnsw') + +db.connect() +db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') # type: ignore +db.drop_tables([Item]) +db.create_tables([Item]) + + +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])) + + +class TestPeewee: + def setup_method(self) -> None: + Item.truncate_table() + + def test_vector(self) -> None: + Item.create(id=1, embedding=[1, 2, 3]) + item = Item.get_by_id(1) + assert item.embedding == [1, 2, 3] + + 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) -> 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) -> 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) -> 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) -> None: + Item.create(id=1, half_embedding=[1, 2, 3]) + item = Item.get_by_id(1) + assert item.half_embedding == [1, 2, 3] + + 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> None: + 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).coerce(True)).scalar() + assert avg == [2.5, 3.5, 4.5] + + 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]) + Item.create(embedding=[4, 5, 6]) + sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() + assert sum == [5, 7, 9] + + 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]) + Item.create(half_embedding=[4, 5, 6]) + avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() + 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() + 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).coerce(True)).scalar() + assert sum == [5, 7, 9] + + def test_get_or_create(self) -> None: + 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 + + ext_db = PostgresqlExtDatabase('pgvector_python_test') + + class ExtItem(BaseModel): + embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False) # type: ignore + + class Meta: # type: ignore + 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], [1, 2, 3]) + # assert np.array_equal(item.embeddings[1], [4, 5, 6]) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py new file mode 100644 index 0000000..8c8e02c --- /dev/null +++ b/tests/test_pg8000.py @@ -0,0 +1,51 @@ +from getpass import getuser +from pgvector import HalfVector, SparseVector, Vector +from pgvector.pg8000 import register_vector +from pg8000.native import Connection +from .conftest import numpy as np + +conn = Connection(getuser(), 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))') + +register_vector(conn) + + +class TestPg8000: + def setup_method(self) -> None: + conn.run('DELETE FROM pg8000_items') + + def test_vector(self) -> None: + embedding = Vector([1.5, 2, 3]) + 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) + + res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') + assert res == [[embedding], [Vector([4.5, 5, 6])], [None]] + + 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) + + res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') + assert res == [[embedding], [None]] + + def test_bit(self) -> None: + embedding = '101' + 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 == [['101'], [None]] + + 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) + + res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') + assert res == [[embedding], [None]] diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 3fe1edc..04863e9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,53 +1,236 @@ -import numpy as np -from pgvector.psycopg import register_vector +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 +from .conftest import numpy as np -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') -conn.execute('CREATE TABLE item (id bigserial primary key, factors 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), embeddings vector[])') register_vector(conn) class TestPsycopg: - 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,)) - - res = conn.execute('SELECT * FROM item ORDER BY id').fetchall() - assert np.array_equal(res[0][1], factors) - 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) - - 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) - - def test_binary_format_correct(self): - factors = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector::text', (factors,)).fetchone()[0] + def setup_method(self) -> None: + conn.execute('DELETE FROM psycopg_items') + + 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)) + + res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() + assert res[0][0] == embedding + assert res[1][0] is None + + 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) -> 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) -> None: + embedding = Vector([1.5, 2, 3]) + res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[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] - 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] - assert np.array_equal(res, np.array([3, 2, 1.5])) + 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) + + 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) + + 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]) + + 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] + assert res == Vector([3, 2, 1.5]) + + 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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)) + cur = conn.cursor() + 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) -> 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)) + 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]) == embedding + assert HalfVector.from_binary(row[1]) == half_embedding + + 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)) + 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 row[0] == embedding + assert row[1] == half_embedding + + 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,)) + + res = next(conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id')) + assert res[0][0] == embeddings[0] + assert res[0][1] == embeddings[1] + + 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) + + with pool.connection() as conn: + res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() + assert res == (Vector([1, 2, 3]),) + + pool.close() + + @pytest.mark.asyncio + 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') + 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 = 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 res[0][1] == embedding + assert res[1][1] is None + + @pytest.mark.asyncio + 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) + 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 res == (Vector([1, 2, 3]),) + + await pool.close() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index df86706..badb7c2 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,28 +1,129 @@ -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 +from .conftest import numpy as np conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True 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('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) class TestPsycopg2: - def setup_method(self, test_method): - cur.execute('DELETE FROM item') + def setup_method(self) -> None: + cur.execute('DELETE FROM psycopg2_items') - def test_works(self): - factors = np.array([1.5, 2, 3]) - cur.execute('INSERT INTO item (factors) VALUES (%s), (NULL)', (factors,)) + 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)) - cur.execute('SELECT * FROM item ORDER BY id') + cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][1], factors) - assert res[0][1].dtype == np.float32 - assert res[1][1] is None + assert res[0][0] == embedding + assert res[1][0] is None + + 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,)) + + cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][0] == Vector(embedding) + assert res[1][0] is None + + def test_halfvec(self) -> None: + 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) -> None: + 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] == embedding + assert res[1][0] is None + + def test_sparsevec(self) -> None: + 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') + res = cur.fetchall() + assert res[0][0] == embedding + assert res[1][0] is None + + 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,)) + + cur.execute('SELECT embeddings FROM psycopg2_items ORDER BY id') + res = cur.fetchone() + assert res == (embeddings,) + + 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,)) + + cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id') + res = cur.fetchone() + assert res == ([HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])],) + + 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,)) + + cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id') + res = cur.fetchone() + assert res == ([SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])],) + + 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) -> 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) -> None: + pool = ThreadedConnectionPool(1, 1, dbname='pgvector_python_test') + + conn = pool.getconn() + try: + # use globally=True for apps to ensure registered with all connections + register_vector(conn) + finally: + pool.putconn(conn) + + conn = pool.getconn() + try: + cur = conn.cursor() + cur.execute("SELECT '[1,2,3]'::vector") + res = cur.fetchone() + assert res == (Vector([1, 2, 3]),) + finally: + pool.putconn(conn) + + pool.closeall() diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py new file mode 100644 index 0000000..e3bb8cc --- /dev/null +++ b/tests/test_sparse_vector.py @@ -0,0 +1,226 @@ +from pgvector import SparseVector +import pytest +import struct +from .conftest import numpy as np + +try: + from scipy import sparse +except ImportError: + sparse = None # type: ignore + + +class TestSparseVector: + 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_init_list_empty(self) -> None: + assert SparseVector([]).to_list() == [] + + 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_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_init_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] + + 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_init_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) + + # 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_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_init_coo_array(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + 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] + assert isinstance(vec.values()[0], float) + + def test_init_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_init_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(sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore + assert str(error.value) == 'extra argument' + + def test_init_coo_matrix(self) -> None: + 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] + + def test_init_dok_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])).todok() + vec = SparseVector(arr) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + + def test_init_csr_array(self) -> None: + 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] + + def test_init_csr_matrix(self) -> None: + 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] + + 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_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) + 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 + + def test_indices(self) -> None: + assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] + + def test_values(self) -> None: + assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] + + 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_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]) + + 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([]) + + 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'): + # 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 = struct.pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) + 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 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): + SparseVector.from_binary(b'') + + def test_from_binary_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii2if', 2, 2, 0, 0, 1, 1)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + 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'): + 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_sqlalchemy.py b/tests/test_sqlalchemy.py index 11d551a..dc58aeb 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,127 +1,707 @@ -import numpy as np -from pgvector.sqlalchemy import Vector +import asyncpg +from getpass import getuser +from pgvector import SparseVector +from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest -from sqlalchemy import create_engine, select, text, MetaData, Table, Column, Index, Integer +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.orm import declarative_base, Session +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, Mapped, Session +from sqlalchemy.sql import func +from typing import Any +from .conftest import numpy as np -engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test', future=True) -with engine.connect() as con: - con.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) - con.commit() +sqlalchemy_version = 2.1 if hasattr(postgresql, 'BitString') else 2 -Base = declarative_base() +psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -class Item(Base): - __tablename__ = 'orm_item' +@event.listens_for(psycopg2_type_engine, 'connect') +def psycopg2_connect(dbapi_connection: Any, connection_record: Any) -> None: + from pgvector.psycopg2 import register_vector + register_vector(dbapi_connection) + + +pg8000_engine = create_engine(f'postgresql+pg8000://{getuser()}@localhost/pgvector_python_test') + +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: Any, connection_record: Any) -> None: + 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') - id = Column(Integer, primary_key=True) - factors = Column(Vector(3)) +@event.listens_for(psycopg_async_type_engine.sync_engine, 'connect') +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) -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) +asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') +asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') -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, factors=v)) + +@event.listens_for(asyncpg_type_engine.sync_engine, 'connect') +def asyncpg_connect(dbapi_connection: Any, connection_record: Any) -> None: + 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: + session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() +class Base(DeclarativeBase): + pass + + +class Item(Base): + __tablename__ = 'sqlalchemy_orm_item' + + 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[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)) + + +Base.metadata.drop_all(setup_engine) +Base.metadata.create_all(setup_engine) + +index = Index( + 'sqlalchemy_orm_index', + Item.embedding, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) +index.create(setup_engine) + +half_precision_index = Index( + 'sqlalchemy_orm_half_precision_index', + func.cast(Item.embedding, HALFVEC(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_ops={'embedding': 'halfvec_l2_ops'} +) +half_precision_index.create(setup_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_ops={'embedding': 'bit_hamming_ops'} +) +binary_quantize_index.create(setup_engine) + + +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]))) + 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() -> None: + 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) -> None: + delete_items() - def test_core(self): + def test_core(self, engine: Engine) -> None: metadata = MetaData() item_table = Table( - 'core_item', + 'sqlalchemy_core_item', metadata, Column('id', Integer, primary_key=True), - Column('factors', Vector(3)) + Column('embedding', VECTOR(3)), + Column('half_embedding', HALFVEC(3)), + Column('binary_embedding', BIT(3)), + Column('sparse_embedding', SPARSEVEC(3)), + Column('embeddings', ARRAY(VECTOR(3))) ) metadata.drop_all(engine) metadata.create_all(engine) - index = Index( - 'my_core_index', - item_table.c.factors, + ivfflat_index = Index( + 'sqlalchemy_core_ivfflat_index', + 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) + ivfflat_index.create(engine) - def test_orm(self): - item = Item(factors=np.array([1.5, 2, 3])) - item2 = Item(factors=[4, 5, 6]) + hnsw_index = Index( + 'sqlalchemy_core_hnsw_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, engine: Engine) -> None: + 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() + with Session(engine) as session: + session.add(item) + session.add(item2) + 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()] - 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 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: + 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 == [1, 2, 3] - def test_l2_distance(self): + def test_vector_l2_distance(self, engine: Engine) -> None: 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): + 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) -> 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) -> None: create_items() with Session(engine) as session: - items = session.query(Item).order_by(Item.factors.max_inner_product([1, 1, 1])).all() + 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, 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) -> None: create_items() with Session(engine) as session: - items = session.query(Item).order_by(Item.factors.cosine_distance([1, 1, 1])).all() + 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_bad_dimensions(self): - item = Item(factors=[1, 2]) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + 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) -> 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) -> 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 == [1, 2, 3] + + 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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_bad_ndim(self): - item = Item(factors=np.array([[1, 2, 3]])) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected ndim to be 1'): + 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) -> 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) -> 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) -> None: + 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: Engine) -> None: + if engine == pg8000_engine: + return - def test_bad_dtype(self): - item = Item(factors=np.array(['one', 'two', 'three'])) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='dtype must be numeric'): + 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, 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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) -> 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() # type: ignore + assert items == [3] + + def test_avg(self, engine: Engine) -> None: + with Session(engine) as session: + 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)).one()[0] + assert res == [2.5, 3.5, 4.5] + + 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 + session.add(Item(embedding=[1, 2, 3])) + session.add(Item(embedding=[4, 5, 6])) + res = session.scalars(select(avg(Item.embedding))).one() + assert res == [2.5, 3.5, 4.5] + + def test_sum(self, engine: Engine) -> None: + with Session(engine) as session: + 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)).one()[0] + assert res == [5, 7, 9] + + 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 + session.add(Item(embedding=[1, 2, 3])) + session.add(Item(embedding=[4, 5, 6])) + res = session.scalars(select(sum(Item.embedding))).one() + assert res == [5, 7, 9] + + 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: + session.add(item) + with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + session.commit() + + 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() + + 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) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() + + 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) -> 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) -> None: + with Session(engine) as session: + 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': [1, 2, 3]}]) + + # register_vector in psycopg2 tests change this behavior + # def test_insert_text(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, engine: Engine) -> None: + metadata = MetaData() + metadata.reflect(engine, only=['sqlalchemy_orm_item']) + AutoBase = automap_base(metadata=metadata) + AutoBase.prepare() + AutoItem = AutoBase.classes.sqlalchemy_orm_item + with Session(engine) as session: + session.execute(insert(AutoItem), [{'embedding': [1, 2, 3]}]) + item = session.query(AutoItem).first() + assert item is not None and item.embedding == [1, 2, 3] + + 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) -> None: + 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(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] + + 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) + + 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(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] + + +@pytest.mark.parametrize('engine', array_engines) +class TestSqlalchemyArray: + def setup_method(self) -> None: + delete_items() + + def test_vector_array(self, engine: Engine) -> None: + with Session(engine) as session: + 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 == [[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=[[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 == [[1, 2, 3], [4, 5, 6]] + + +@pytest.mark.parametrize('engine', async_engines) +class TestSqlalchemyAsync: + def setup_method(self) -> None: + delete_items() + + @pytest.mark.asyncio + async def test_vector(self, engine: AsyncEngine) -> None: + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + embedding = [1, 2, 3] + session.add(Item(id=1, embedding=embedding)) + item = await session.get_one(Item, 1) + assert item.embedding == embedding + + await engine.dispose() + + @pytest.mark.asyncio + async def test_halfvec(self, engine: AsyncEngine) -> None: + async_session = async_sessionmaker(engine, expire_on_commit=False) + + 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_one(Item, 1) + assert item.half_embedding == embedding + + await engine.dispose() + + @pytest.mark.asyncio + async def test_bit(self, engine: AsyncEngine) -> None: + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + embedding: Any + if sqlalchemy_version >= 2.1: + embedding = postgresql.BitString('101') # type: ignore + 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 + + if engine == asyncpg_engine: + session.add(Item(id=2, binary_embedding='101')) + item = await session.get_one(Item, 2) + assert item.binary_embedding == embedding + + await engine.dispose() + + @pytest.mark.asyncio + async def test_sparsevec(self, engine: AsyncEngine) -> None: + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + embedding = SparseVector([1, 2, 3]) + session.add(Item(id=1, sparse_embedding=embedding)) + item = await session.get_one(Item, 1) + assert item.sparse_embedding == embedding + + await engine.dispose() + + @pytest.mark.asyncio + async def test_avg(self, engine: AsyncEngine) -> None: + 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])) + res = await session.scalars(select(avg(Item.embedding))) + assert res.first() == [2.5, 3.5, 4.5] + + await engine.dispose() + + +@pytest.mark.parametrize('engine', async_array_engines) +class TestSqlalchemyAsyncArray: + def setup_method(self) -> None: + delete_items() + + @pytest.mark.asyncio + async def test_vector_array(self, engine: AsyncEngine) -> None: + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + session.add(Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]])) + item = await session.get_one(Item, 1) + assert item.embeddings == [[1, 2, 3], [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 == [[1, 2, 3], [4, 5, 6]] + + await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py new file mode 100644 index 0000000..853384e --- /dev/null +++ b/tests/test_sqlmodel.py @@ -0,0 +1,243 @@ +from pgvector import SparseVector +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, col, create_engine, delete, select, text + +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 + + +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) + # 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 + 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 + + +SQLModel.metadata.drop_all(engine) +SQLModel.metadata.create_all(engine) + +index = Index( + 'sqlmodel_index', + col(Item.embedding), # col() for typing + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'vector_l2_ops'} +) +index.create(engine) + + +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]))) + 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: + def setup_method(self) -> None: + with Session(engine) as session: + session.exec(delete(Item)) + session.commit() + + def test_orm(self) -> None: + item = Item(embedding=[1.5, 2, 3]) + item2 = Item(embedding=[4, 5, 6]) + item3 = Item() + + with Session(engine) as session: + 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 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: + 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 == [1, 2, 3] + + def test_vector_l2_distance(self) -> None: + create_items() + with Session(engine) as session: + 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(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(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(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: + 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 == [1, 2, 3] + + def test_halfvec_l2_distance(self) -> None: + create_items() + with Session(engine) as session: + 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(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(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(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: + 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) -> None: + create_items() + with Session(engine) as session: + 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(col(Item.binary_embedding).jaccard_distance('101'))) # type: ignore + assert [v.id for v in items] == [2, 3, 1] + + 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) -> None: + create_items() + with Session(engine) as session: + 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(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(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(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(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(col(Item.embedding).l2_distance([1, 1, 1]))).all() # type: ignore + assert items[0] == 3 + + def test_vector_avg(self) -> None: + with Session(engine) as session: + 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])) + res = session.exec(select(avg(Item.embedding))).first() + assert res == [2.5, 3.5, 4.5] + + def test_vector_sum(self) -> None: + with Session(engine) as session: + 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])) + res = session.exec(select(sum(Item.embedding))).first() + assert res == [5, 7, 9] + + def test_halfvec_avg(self) -> None: + with Session(engine) as session: + 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])) + res = session.exec(select(avg(Item.half_embedding))).first() + assert res == [2.5, 3.5, 4.5] + + def test_halfvec_sum(self) -> None: + with Session(engine) as session: + 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])) + res = session.exec(select(sum(Item.half_embedding))).first() + assert res == [5, 7, 9] + + def test_bad_dimensions(self) -> None: + 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() diff --git a/tests/test_vector.py b/tests/test_vector.py new file mode 100644 index 0000000..13a9cc7 --- /dev/null +++ b/tests/test_vector.py @@ -0,0 +1,138 @@ +from pgvector import Vector +import pytest +import struct +from .conftest import numpy as np + + +class TestVector: + 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_init_list_empty(self) -> None: + assert Vector([]).to_list() == [] + + 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_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_init_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 + assert Vector(arr).to_numpy().dtype == np.float32 + + 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() + + 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) + + def test_init_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) -> 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_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 + assert Vector([0.0]) == Vector([-0.0]) + + 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]) + + 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([]) + + 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_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]) + + 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_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hhf', 2, 0, 1)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + 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'): + 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))