diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f8bcaa3..4d4e8ed 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,7 +6,7 @@ jobs: strategy: fail-fast: false matrix: - python: [3.12, 3.8] + python: [3.13, 3.9] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -19,8 +19,11 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.7.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install - run: pytest + + - run: pip install "SQLAlchemy<2" -U + - run: pytest tests/test_sqlalchemy.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a517d8..d0e2730 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,14 @@ +## 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 diff --git a/LICENSE.txt b/LICENSE.txt index d205f4e..b612d6d 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2024 Andrew Kane +Copyright (c) 2021-2025 Andrew Kane Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index acd625d..b6bc055 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) @@ -22,6 +22,7 @@ And follow the instructions for your database library: - [Psycopg 3](#psycopg-3) - [Psycopg 2](#psycopg-2) - [asyncpg](#asyncpg) +- [pg8000](#pg8000) - [Peewee](#peewee) Or check out some examples: @@ -33,6 +34,7 @@ Or check out some examples: - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py) with SentenceTransformers (cross-encoder) - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search/example.py) with Transformers - [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert/exact.py) with ColBERT +- [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 @@ -133,6 +135,36 @@ 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 @@ -214,6 +246,77 @@ index.create(engine) Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +#### Half-Precision Indexing + +Index vectors at half-precision + +```python +from pgvector.sqlalchemy import HALFVEC +from sqlalchemy.sql import func + +index = Index( + 'my_index', + func.cast(Item.embedding, HALFVEC(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': '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)) +``` + +#### Arrays + +Add an array column + +```python +from pgvector.sqlalchemy import Vector +from sqlalchemy import ARRAY + +class Item(Base): + embeddings = mapped_column(ARRAY(Vector(3))) +``` + +And register the types with the underlying driver + +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, arrays=True) +``` + ## SQLModel Enable the extension @@ -226,10 +329,9 @@ Add a vector column ```python from pgvector.sqlalchemy import Vector -from sqlalchemy import Column class Item(SQLModel, table=True): - embedding: Any = Field(sa_column=Column(Vector(3))) + embedding: Any = Field(sa_type=Vector(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -275,7 +377,7 @@ Also supports `sum` Add an approximate index ```python -from sqlalchemy import Index +from sqlmodel import Index index = Index( 'my_index', @@ -314,6 +416,15 @@ from pgvector.psycopg import register_vector register_vector(conn) ``` +For [connection pools](https://www.psycopg.org/psycopg3/docs/advanced/pool.html), use + +```python +def configure(conn): + register_vector(conn) + +pool = ConnectionPool(..., configure=configure) +``` + For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use ```python @@ -452,6 +563,51 @@ await conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +## pg8000 + +Enable the extension + +```python +conn.run('CREATE EXTENSION IF NOT EXISTS vector') +``` + +Register the vector type with your connection + +```python +from pgvector.pg8000 import register_vector + +register_vector(conn) +``` + +Create a table + +```python +conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') +``` + +Insert a vector + +```python +embedding = np.array([1, 2, 3]) +conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding) +``` + +Get the nearest neighbors to a vector + +```python +conn.run('SELECT * FROM items ORDER BY embedding <-> :embedding LIMIT 5', embedding=embedding) +``` + +Add an approximate index + +```python +conn.run('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') +# or +conn.run('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') +``` + +Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance + ## Peewee Add a vector column @@ -509,6 +665,99 @@ Item.add_index('embedding vector_l2_ops', using='hnsw') Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance +## Reference + +### Half Vectors + +Create a half vector from a list + +```python +vec = HalfVector([1, 2, 3]) +``` + +Or a NumPy array + +```python +vec = HalfVector(np.array([1, 2, 3])) +``` + +Get a list + +```python +lst = vec.to_list() +``` + +Get a NumPy array + +```python +arr = vec.to_numpy() +``` + +### Sparse Vectors + +Create a sparse vector from a list + +```python +vec = SparseVector([1, 0, 2, 0, 3, 0]) +``` + +Or a NumPy array + +```python +vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) +``` + +Or a SciPy sparse array + +```python +arr = coo_array(([1, 2, 3], ([0, 2, 4],)), shape=(6,)) +vec = SparseVector(arr) +``` + +Or a dictionary of non-zero elements + +```python +vec = SparseVector({0: 1, 2: 2, 4: 3}, 6) +``` + +Note: Indices start at 0 + +Get the number of dimensions + +```python +dim = vec.dimensions() +``` + +Get the indices of non-zero elements + +```python +indices = vec.indices() +``` + +Get the values of non-zero elements + +```python +values = vec.values() +``` + +Get a list + +```python +lst = vec.to_list() +``` + +Get a NumPy array + +```python +arr = vec.to_numpy() +``` + +Get a SciPy sparse array + +```python +arr = vec.to_coo() +``` + ## History View the [changelog](https://github.com/pgvector/pgvector-python/blob/master/CHANGELOG.md) diff --git a/examples/citus/example.py b/examples/citus/example.py index d448204..915c25f 100644 --- a/examples/citus/example.py +++ b/examples/citus/example.py @@ -40,9 +40,6 @@ for i in range(rows): copy.write_row([embeddings[i], categories[i]]) - while conn.pgconn.flush() == 1: - pass - print('Creating index in parallel') conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 780352a..393d1e0 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -12,7 +12,7 @@ conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))') -def fetch_embeddings(input, input_type): +def embed(input, input_type): co = cohere.Client() response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary']) return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] @@ -23,12 +23,12 @@ def fetch_embeddings(input, input_type): 'The cat is purring', 'The bear is growling' ] -embeddings = fetch_embeddings(input, 'search_document') +embeddings = embed(input, 'search_document') for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Bit(embedding))) query = 'forest' -query_embedding = fetch_embeddings([query], 'search_query')[0] +query_embedding = embed([query], 'search_query')[0] result = conn.execute('SELECT content FROM documents ORDER BY embedding <~> %s LIMIT 5', (Bit(query_embedding),)).fetchall() for row in result: print(row[0]) diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py new file mode 100644 index 0000000..80bb603 --- /dev/null +++ b/examples/colpali/exact.py @@ -0,0 +1,56 @@ +from colpali_engine.models import ColQwen2, ColQwen2Processor +from colpali_engine.utils.torch_utils import get_torch_device +from datasets import load_dataset +from pgvector.psycopg import register_vector, Bit +import psycopg +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/requirements.txt b/examples/colpali/requirements.txt new file mode 100644 index 0000000..4cf770d --- /dev/null +++ b/examples/colpali/requirements.txt @@ -0,0 +1,4 @@ +colpali-engine +datasets +pgvector +psycopg[binary] diff --git a/examples/implicit/requirements.txt b/examples/implicit/requirements.txt index 8f04b58..424abbd 100644 --- a/examples/implicit/requirements.txt +++ b/examples/implicit/requirements.txt @@ -1,3 +1,4 @@ +h5py implicit pgvector psycopg[binary] diff --git a/examples/openai/example.py b/examples/openai/example.py index ebed3d0..b9a078c 100644 --- a/examples/openai/example.py +++ b/examples/openai/example.py @@ -1,3 +1,4 @@ +import numpy as np from openai import OpenAI from pgvector.psycopg import register_vector import psycopg @@ -10,20 +11,24 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))') + +def embed(input): + client = OpenAI() + response = client.embeddings.create(input=input, model='text-embedding-3-small') + return [v.embedding for v in response.data] + + input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] - -client = OpenAI() -response = client.embeddings.create(input=input, model='text-embedding-3-small') -embeddings = [v.embedding for v in response.data] - +embeddings = embed(input) for content, embedding in zip(input, embeddings): - conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, np.array(embedding))) -document_id = 1 -neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() -for neighbor in neighbors: - print(neighbor[0]) +query = 'forest' +query_embedding = embed([query])[0] +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (np.array(query_embedding),)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/openai/halfvec.py b/examples/openai/halfvec.py new file mode 100644 index 0000000..185c785 --- /dev/null +++ b/examples/openai/halfvec.py @@ -0,0 +1,34 @@ +from openai import OpenAI +from pgvector.psycopg import register_vector, HalfVector +import psycopg + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding halfvec(3072))') +conn.execute('CREATE INDEX ON documents USING hnsw (embedding halfvec_cosine_ops)') + + +def embed(input): + client = OpenAI() + response = client.embeddings.create(input=input, model='text-embedding-3-large') + return [v.embedding for v in response.data] + + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +embeddings = embed(input) +for content, embedding in zip(input, embeddings): + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, HalfVector(embedding))) + +query = 'forest' +query_embedding = embed([query])[0] +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (HalfVector(query_embedding),)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/sentence_transformers/example.py b/examples/sentence_transformers/example.py index d4e7f96..50997d9 100644 --- a/examples/sentence_transformers/example.py +++ b/examples/sentence_transformers/example.py @@ -10,19 +10,19 @@ conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') +model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') + input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] - -model = SentenceTransformer('all-MiniLM-L6-v2') embeddings = model.encode(input) - for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) -document_id = 1 -neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() -for neighbor in neighbors: - print(neighbor[0]) +query = 'forest' +query_embedding = model.encode(query) +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (query_embedding,)).fetchall() +for row in result: + print(row[0]) diff --git a/examples/sparse_search/example.py b/examples/sparse_search/example.py index fa6074e..2b5daea 100644 --- a/examples/sparse_search/example.py +++ b/examples/sparse_search/example.py @@ -20,7 +20,7 @@ special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()] -def fetch_embeddings(input): +def embed(input): feature = tokenizer( input, padding=True, @@ -42,12 +42,12 @@ def fetch_embeddings(input): 'The cat is purring', 'The bear is growling' ] -embeddings = fetch_embeddings(input) +embeddings = embed(input) for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector(embedding))) query = 'forest' -query_embedding = fetch_embeddings([query])[0] +query_embedding = embed([query])[0] result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector(query_embedding),)).fetchall() for row in result: print(row[0]) 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/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index 543b882..c6a3b4e 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,5 +1,7 @@ from .register import register_vector -from ..utils import Vector, HalfVector, SparseVector + +# TODO remove +from .. import Vector, HalfVector, SparseVector __all__ = [ 'register_vector', diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index a388058..63726f3 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -1,4 +1,4 @@ -from ..utils import Vector, HalfVector, SparseVector +from .. import Vector, HalfVector, SparseVector async def register_vector(conn, schema='public'): diff --git a/pgvector/bit.py b/pgvector/bit.py new file mode 100644 index 0000000..26a9d8d --- /dev/null +++ b/pgvector/bit.py @@ -0,0 +1,75 @@ +import numpy as np +from struct import pack, unpack_from +from warnings import warn + + +class Bit: + def __init__(self, value): + if isinstance(value, bytes): + self._len = 8 * len(value) + self._data = value + else: + if isinstance(value, str): + value = [v != '0' for v in value] + else: + value = np.asarray(value) + + if value.dtype != np.bool: + # skip warning for result of np.unpackbits + if value.dtype != np.uint8 or np.any(value > 1): + warn('expected elements to be boolean', stacklevel=2) + value = value.astype(bool) + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + self._len = len(value) + self._data = np.packbits(value).tobytes() + + def __repr__(self): + return f'Bit({self.to_text()})' + + def __eq__(self, other): + if isinstance(other, self.__class__): + return self._len == other._len and self._data == other._data + return False + + def to_list(self): + return self.to_numpy().tolist() + + def to_numpy(self): + return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) + + def to_text(self): + return ''.join(format(v, '08b') for v in self._data)[:self._len] + + def to_binary(self): + return pack('>i', self._len) + self._data + + @classmethod + def from_text(cls, value): + return cls(str(value)) + + @classmethod + def from_binary(cls, value): + if not isinstance(value, bytes): + raise ValueError('expected bytes') + + bit = cls.__new__(cls) + bit._len = unpack_from('>i', value)[0] + bit._data = value[4:] + return bit + + @classmethod + def _to_db(cls, value): + if not isinstance(value, cls): + raise ValueError('expected bit') + + return value.to_text() + + @classmethod + def _to_db_binary(cls, value): + if not isinstance(value, cls): + raise ValueError('expected bit') + + return value.to_binary() diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py index 09978a9..43c64a3 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -5,7 +5,9 @@ from .indexes import IvfflatIndex, HnswIndex from .sparsevec import SparseVectorField from .vector import VectorField -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'VectorExtension', diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index da9fbf8..9df4fdb 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,5 +1,5 @@ from django.db.models import FloatField, Func, Value -from ..utils import Vector, HalfVector, SparseVector +from .. import Vector, HalfVector, SparseVector class DistanceBase(Func): @@ -13,6 +13,10 @@ def __init__(self, expression, vector, **extra): vector = Value(SparseVector._to_db(vector)) else: vector = Value(Vector._to_db(vector)) + + # prevent error with unhashable types + self._constructor_args = ((expression, vector), extra) + super().__init__(expression, vector, **extra) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 6b59a7f..3aeb90f 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -1,6 +1,6 @@ from django import forms from django.db.models import Field -from ..utils import HalfVector +from .. import HalfVector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index d0d2d07..580f27c 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -1,6 +1,6 @@ from django import forms from django.db.models import Field -from ..utils import SparseVector +from .. import SparseVector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index a89d540..861cfde 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -1,7 +1,7 @@ from django import forms from django.db.models import Field import numpy as np -from ..utils import Vector +from .. import Vector # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ diff --git a/pgvector/utils/halfvec.py b/pgvector/halfvec.py similarity index 92% rename from pgvector/utils/halfvec.py rename to pgvector/halfvec.py index e1e5051..f335f2f 100644 --- a/pgvector/utils/halfvec.py +++ b/pgvector/halfvec.py @@ -16,6 +16,11 @@ def __init__(self, value): def __repr__(self): return f'HalfVector({self.to_list()})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return np.array_equal(self.to_numpy(), other.to_numpy()) + return False + def dimensions(self): return len(self._value) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index 945e0dc..df21200 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -2,7 +2,9 @@ from .halfvec import HalfVectorField from .sparsevec import SparseVectorField from .vector import VectorField -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'VectorField', diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index deaa14d..0901fd2 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,5 +1,5 @@ from peewee import Expression, Field -from ..utils import HalfVector +from .. import HalfVector class HalfVectorField(Field): diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 67f7d1b..86dea73 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -1,5 +1,5 @@ from peewee import Expression, Field -from ..utils import SparseVector +from .. import SparseVector class SparseVectorField(Field): diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 22a87e5..83f9997 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,5 +1,5 @@ from peewee import Expression, Field -from ..utils import Vector +from .. import Vector class VectorField(Field): diff --git a/pgvector/pg8000/__init__.py b/pgvector/pg8000/__init__.py new file mode 100644 index 0000000..b3b4440 --- /dev/null +++ b/pgvector/pg8000/__init__.py @@ -0,0 +1,5 @@ +from .register import register_vector + +__all__ = [ + 'register_vector' +] diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py new file mode 100644 index 0000000..15ee219 --- /dev/null +++ b/pgvector/pg8000/register.py @@ -0,0 +1,23 @@ +import numpy as np +from .. import Vector, HalfVector, SparseVector + + +def register_vector(conn): + # use to_regtype to get first matching type in search path + res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") + type_info = dict(res) + + if 'vector' not in type_info: + raise RuntimeError('vector type not found in the database') + + conn.register_out_adapter(Vector, Vector._to_db) + conn.register_out_adapter(np.ndarray, Vector._to_db) + conn.register_in_adapter(type_info['vector'], Vector._from_db) + + if 'halfvec' in type_info: + conn.register_out_adapter(HalfVector, HalfVector._to_db) + conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db) + + if 'sparsevec' in type_info: + conn.register_out_adapter(SparseVector, SparseVector._to_db) + conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 9007c37..980af84 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,5 +1,7 @@ from .register import register_vector, register_vector_async -from ..utils import Bit, HalfVector, SparseVector, Vector + +# TODO remove +from .. import Bit, HalfVector, SparseVector, Vector __all__ = [ 'register_vector', diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index f8eeb61..cffe8fb 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,6 +1,6 @@ from psycopg.adapt import Dumper from psycopg.pq import Format -from ..utils import Bit +from .. import Bit class BitDumper(Dumper): diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 351d2cb..b3a0060 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,6 +1,6 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import HalfVector +from .. import HalfVector class HalfVectorDumper(Dumper): diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 435fd06..384a0e1 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,6 +1,6 @@ from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import SparseVector +from .. import SparseVector class SparseVectorDumper(Dumper): diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 0f62ca9..db9e826 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,7 +1,7 @@ import psycopg from psycopg.adapt import Loader, Dumper from psycopg.pq import Format -from ..utils import Vector +from .. import Vector class VectorDumper(Dumper): diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 7c95295..33e5124 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,5 +1,7 @@ from .register import register_vector -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'register_vector', diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index b50e89b..0a4c736 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,5 +1,5 @@ from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type -from ..utils import HalfVector +from .. import HalfVector class HalfvecAdapter: diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 7752852..1bc9d44 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -5,10 +5,8 @@ from .vector import register_vector_info -# TODO make globally False by default in 0.4.0 # note: register_adapter is always global -# TODO make arrays True by defalt in 0.4.0 -def register_vector(conn_or_curs=None, globally=True, arrays=False): +def register_vector(conn_or_curs, globally=False, arrays=True): conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection cur = conn.cursor(cursor_factory=cursor) scope = None if globally else conn_or_curs diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index a542807..148eff2 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,5 +1,5 @@ from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type -from ..utils import SparseVector +from .. import SparseVector class SparsevecAdapter: diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 9861f01..562de18 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,6 +1,6 @@ import numpy as np from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type -from ..utils import Vector +from .. import Vector class VectorAdapter: @@ -24,3 +24,4 @@ def register_vector_info(oid, array_oid, scope): register_type(vectorarray, scope) register_adapter(np.ndarray, VectorAdapter) + register_adapter(Vector, VectorAdapter) diff --git a/pgvector/utils/sparsevec.py b/pgvector/sparsevec.py similarity index 93% rename from pgvector/utils/sparsevec.py rename to pgvector/sparsevec.py index fd9ccff..8df2dfd 100644 --- a/pgvector/utils/sparsevec.py +++ b/pgvector/sparsevec.py @@ -26,6 +26,11 @@ def __repr__(self): elements = dict(zip(self._indices, self._values)) return f'SparseVector({elements}, {self._dim})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values() + return False + def dimensions(self): return self._dim @@ -108,7 +113,7 @@ def from_binary(cls, value): dim, nnz, unused = unpack_from('>iii', value) indices = unpack_from(f'>{nnz}i', value, 12) values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return cls._from_parts(int(dim), indices, values) + return cls._from_parts(int(dim), list(indices), list(values)) @classmethod def _from_parts(cls, dim, indices, values): diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 4955eeb..52adf88 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -4,7 +4,9 @@ from .sparsevec import SPARSEVEC from .vector import VECTOR from .vector import VECTOR as Vector -from ..utils import HalfVector, SparseVector + +# TODO remove +from .. import HalfVector, SparseVector __all__ = [ 'Vector', diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 639f77b..10688b5 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import HalfVector +from .. import HalfVector class HALFVEC(UserDefinedType): diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 370f5d1..0058679 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import SparseVector +from .. import SparseVector class SPARSEVEC(UserDefinedType): diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index f57a045..5a1e11f 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,6 +1,6 @@ from sqlalchemy.dialects.postgresql.base import ischema_names from sqlalchemy.types import UserDefinedType, Float, String -from ..utils import Vector +from .. import Vector class VECTOR(UserDefinedType): diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py index 3c01160..8cdb5d6 100644 --- a/pgvector/utils/__init__.py +++ b/pgvector/utils/__init__.py @@ -1,7 +1,5 @@ -from .bit import Bit -from .halfvec import HalfVector -from .sparsevec import SparseVector -from .vector import Vector +# TODO remove +from .. import Bit, HalfVector, SparseVector, Vector __all__ = [ 'Vector', diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py deleted file mode 100644 index 51f7556..0000000 --- a/pgvector/utils/bit.py +++ /dev/null @@ -1,61 +0,0 @@ -import numpy as np -from struct import pack, unpack_from - - -class Bit: - def __init__(self, value): - if isinstance(value, str): - self._value = self.from_text(value)._value - else: - # TODO change in 0.4.0 - # TODO raise if dtype not bool or uint8 - # if isinstance(value, np.ndarray) and value.dtype == np.uint8: - # value = np.unpackbits(value) - # else: - # value = np.asarray(value, dtype=bool) - - value = np.asarray(value, dtype=bool) - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - self._value = value - - def __repr__(self): - return f'Bit({self.to_text()})' - - def to_list(self): - return self._value.tolist() - - def to_numpy(self): - return self._value - - def to_text(self): - return ''.join(self._value.astype(np.uint8).astype(str)) - - def to_binary(self): - return pack('>i', len(self._value)) + np.packbits(self._value).tobytes() - - @classmethod - def from_text(cls, value): - return cls(np.asarray([v != '0' for v in value], dtype=bool)) - - @classmethod - def from_binary(cls, value): - count = unpack_from('>i', value)[0] - buf = np.frombuffer(value, dtype=np.uint8, offset=4) - return cls(np.unpackbits(buf, count=count).astype(bool)) - - @classmethod - def _to_db(cls, value): - if not isinstance(value, cls): - raise ValueError('expected bit') - - return value.to_text() - - @classmethod - def _to_db_binary(cls, value): - if not isinstance(value, cls): - raise ValueError('expected bit') - - return value.to_binary() diff --git a/pgvector/utils/vector.py b/pgvector/vector.py similarity index 92% rename from pgvector/utils/vector.py rename to pgvector/vector.py index 3fa2f35..ebbcafd 100644 --- a/pgvector/utils/vector.py +++ b/pgvector/vector.py @@ -16,6 +16,11 @@ def __init__(self, value): def __repr__(self): return f'Vector({self.to_list()})' + def __eq__(self, other): + if isinstance(other, self.__class__): + return np.array_equal(self.to_numpy(), other.to_numpy()) + return False + def dimensions(self): return len(self._value) diff --git a/pyproject.toml b/pyproject.toml index a6a6609..b889f4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,14 +4,14 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.3.6" +version = "0.4.0" description = "pgvector support for Python" readme = "README.md" authors = [ {name = "Andrew Kane", email = "andrew@ankane.org"} ] license = {text = "MIT"} -requires-python = ">= 3.8" +requires-python = ">= 3.9" dependencies = [ "numpy" ] diff --git a/requirements.txt b/requirements.txt index c1e11f3..a13be06 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,8 @@ asyncpg Django numpy peewee -psycopg[binary] +pg8000 +psycopg[binary,pool] psycopg2-binary pytest pytest-asyncio diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 7a68a9e..34d66a1 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,6 +1,7 @@ import asyncpg import numpy as np -from pgvector.asyncpg import register_vector, SparseVector +from pgvector import HalfVector, SparseVector, Vector +from pgvector.asyncpg import register_vector import pytest @@ -14,13 +15,15 @@ async def test_vector(self): await register_vector(conn) - embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding = Vector([1.5, 2, 3]) + embedding2 = np.array([4.5, 5, 6]) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding) + assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) assert res[0]['embedding'].dtype == np.float32 - assert res[1]['embedding'] is None + assert np.array_equal(res[1]['embedding'], embedding2) + assert res[2]['embedding'] is None # ensures binary format is correct text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") @@ -37,12 +40,14 @@ async def test_halfvec(self): await register_vector(conn) - embedding = [1.5, 2, 3] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding = HalfVector([1.5, 2, 3]) + embedding2 = [4.5, 5, 6] + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['embedding'].to_list() == [1.5, 2, 3] - assert res[1]['embedding'] is None + assert res[0]['embedding'] == embedding + assert res[1]['embedding'] == HalfVector(embedding2) + assert res[2]['embedding'] is None # ensures binary format is correct text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") @@ -59,10 +64,11 @@ async def test_bit(self): await register_vector(conn) - embedding = asyncpg.BitString.from_int(5, length=3) + embedding = asyncpg.BitString('101') await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + assert res[0]['embedding'].as_string() == '101' assert res[0]['embedding'].to_int() == 5 assert res[1]['embedding'] is None @@ -85,7 +91,7 @@ async def test_sparsevec(self): await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert res[0]['embedding'].to_list() == [1.5, 2, 3] + assert res[0]['embedding'] == embedding assert res[1]['embedding'] is None # ensures binary format is correct @@ -103,12 +109,17 @@ async def test_vector_array(self): await register_vector(conn) - embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] - await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings[0], embeddings[1]) + embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] + await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) + + embeddings2 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embeddings'][0], embeddings[0]) - assert np.array_equal(res[0]['embeddings'][1], embeddings[1]) + assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy()) + assert np.array_equal(res[0]['embeddings'][1], embeddings[1].to_numpy()) + assert np.array_equal(res[1]['embeddings'][0], embeddings2[0]) + assert np.array_equal(res[1]['embeddings'][1], embeddings2[1]) await conn.close() @@ -124,10 +135,12 @@ async def init(conn): await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') - embedding = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding = Vector([1.5, 2, 3]) + embedding2 = np.array([1.5, 2, 3]) + await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding) + assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) assert res[0]['embedding'].dtype == np.float32 - assert res[1]['embedding'] is None + assert np.array_equal(res[1]['embedding'], embedding2) + assert res[2]['embedding'] is None diff --git a/tests/test_bit.py b/tests/test_bit.py index 32ab87b..5a71642 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,5 +1,5 @@ import numpy as np -from pgvector.utils import Bit +from pgvector import Bit import pytest @@ -7,22 +7,42 @@ class TestBit: def test_list(self): assert Bit([True, False, True]).to_list() == [True, False, True] + def test_list_none(self): + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit([True, None, True]).to_text() == '101' + + def test_list_int(self): + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit([254, 7, 0]).to_text() == '110' + def test_tuple(self): assert Bit((True, False, True)).to_list() == [True, False, True] def test_str(self): assert Bit('101').to_list() == [True, False, True] - def test_ndarray_uint8(self): - arr = np.array([254, 7, 0], dtype=np.uint8) - # TODO change in 0.4.0 - # assert Bit(arr).to_text() == '111111100000011100000000' - assert Bit(arr).to_text() == '110' + def test_bytes(self): + assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' + assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - def test_ndarray_same_object(self): + def test_ndarray(self): arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] - assert Bit(arr).to_numpy() is arr + assert np.array_equal(Bit(arr).to_numpy(), arr) + + def test_ndarray_unpackbits(self): + arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) + assert Bit(arr).to_text() == '111111100000011100000000' + + def test_ndarray_uint8(self): + arr = np.array([254, 7, 0], dtype=np.uint8) + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit(arr).to_text() == '110' + + def test_ndarray_uint16(self): + arr = np.array([254, 7, 0], dtype=np.uint16) + with pytest.warns(UserWarning, match='expected elements to be boolean'): + assert Bit(arr).to_text() == '110' def test_ndim_two(self): with pytest.raises(ValueError) as error: @@ -37,3 +57,7 @@ def test_ndim_zero(self): def test_repr(self): assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' + + def test_equality(self): + assert Bit([True, False, True]) == Bit([True, False, True]) + assert Bit([True, False, True]) != Bit([True, False, False]) diff --git a/tests/test_django.py b/tests/test_django.py index 5ab5f81..7a8a6eb 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -1,6 +1,7 @@ import django from django.conf import settings from django.contrib.postgres.fields import ArrayField +from django.contrib.postgres.indexes import OpClass from django.core import serializers from django.db import connection, migrations, models from django.db.models import Avg, Sum, FloatField, DecimalField @@ -11,7 +12,8 @@ import numpy as np import os import pgvector.django -from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance, HalfVector, SparseVector +from pgvector import HalfVector, SparseVector +from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock settings.configure( @@ -38,7 +40,12 @@ 'level': 'WARNING' } } - } + }, + # needed for OpClass + # https://docs.djangoproject.com/en/5.1/ref/contrib/postgres/indexes/#opclass-expressions + INSTALLED_APPS=[ + 'django.contrib.postgres' + ] ) django.setup() @@ -65,8 +72,14 @@ class Meta: name='hnsw_idx', fields=['embedding'], m=16, - ef_construction=100, + 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 ) ] @@ -74,9 +87,6 @@ class Meta: class Migration(migrations.Migration): initial = True - dependencies = [ - ] - operations = [ VectorExtension(), migrations.CreateModel( @@ -99,6 +109,10 @@ class Migration(migrations.Migration): migrations.AddIndex( model_name='item', index=pgvector.django.HnswIndex(fields=['embedding'], m=16, ef_construction=64, name='hnsw_idx', opclasses=['vector_l2_ops']), + ), + migrations.AddIndex( + model_name='item', + index=pgvector.django.HnswIndex(OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), m=16, ef_construction=64, name='hnsw_half_precision_idx'), ) ] @@ -145,13 +159,13 @@ class Meta: class TestDjango: - def setup_method(self, test_method): + def setup_method(self): Item.objects.all().delete() def test_vector(self): Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert np.array_equal(item.embedding, [1, 2, 3]) assert item.embedding.dtype == np.float32 def test_vector_l2_distance(self): @@ -185,7 +199,7 @@ def test_vector_l1_distance(self): def test_halfvec(self): Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert item.half_embedding.to_list() == [1, 2, 3] + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): create_items() @@ -237,7 +251,7 @@ def test_bit_jaccard_distance(self): def test_sparsevec(self): Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() item = Item.objects.get(pk=1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): create_items() @@ -279,7 +293,7 @@ def test_vector_avg(self): Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + assert np.array_equal(avg, [2.5, 3.5, 4.5]) def test_vector_sum(self): sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] @@ -287,7 +301,7 @@ def test_vector_sum(self): Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] - assert np.array_equal(sum, np.array([5, 7, 9])) + assert np.array_equal(sum, [5, 7, 9]) def test_halfvec_avg(self): avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] @@ -295,7 +309,7 @@ def test_halfvec_avg(self): Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] - assert avg.to_list() == [2.5, 3.5, 4.5] + assert avg == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self): sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] @@ -303,7 +317,7 @@ def test_halfvec_sum(self): Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] - assert sum.to_list() == [5, 7, 9] + assert sum == HalfVector([5, 7, 9]) def test_serialization(self): create_items() @@ -333,7 +347,7 @@ def test_vector_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist() + assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6]) def test_vector_form_save_missing(self): Item(id=1).save() @@ -361,7 +375,7 @@ def test_halfvec_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert [4, 5, 6] == Item.objects.get(pk=1).half_embedding.to_list() + assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) def test_halfvec_form_save_missing(self): Item(id=1).save() @@ -418,7 +432,7 @@ def test_sparsevec_form_save(self): assert form.has_changed() assert form.is_valid() assert form.save() - assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list() + assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) def test_sparesevec_form_save_missing(self): Item(id=1).save() @@ -451,8 +465,8 @@ def test_vector_array(self): # this fails if the driver does not cast arrays item = Item.objects.get(pk=1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) def test_double_array(self): Item(id=1, double_embedding=[1, 1, 1]).save() @@ -473,3 +487,10 @@ def test_numeric_array(self): assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] assert items[1].numeric_embedding == [1, 1, 2] + + def test_half_precision(self): + create_items() + distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [1, 1, 1]) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, sqrt(3)] diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index fdaa5f7..78b4977 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,6 +1,7 @@ import numpy as np -from pgvector.utils import HalfVector +from pgvector import HalfVector import pytest +from struct import pack class TestHalfVector: @@ -38,5 +39,21 @@ def test_repr(self): assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' + def test_equality(self): + assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) + assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) + def test_dimensions(self): assert HalfVector([1, 2, 3]).dimensions() == 3 + + def test_from_text(self): + vec = HalfVector.from_text('[1.5,2,3]') + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + + def test_from_binary(self): + data = pack('>HH3e', 3, 0, 1.5, 2, 3) + vec = HalfVector.from_binary(data) + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + assert vec.to_binary() == data diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 9666388..64fc009 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,8 @@ from math import sqrt import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField, SparseVector +from pgvector import HalfVector, SparseVector +from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField db = PostgresqlDatabase('pgvector_python_test') @@ -36,13 +37,13 @@ def create_items(): class TestPeewee: - def setup_method(self, test_method): + def setup_method(self): Item.truncate_table() def test_vector(self): Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) + assert np.array_equal(item.embedding, [1, 2, 3]) assert item.embedding.dtype == np.float32 def test_vector_l2_distance(self): @@ -76,7 +77,7 @@ def test_vector_l1_distance(self): def test_halfvec(self): Item.create(id=1, half_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.half_embedding.to_list() == [1, 2, 3] + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): create_items() @@ -128,7 +129,7 @@ def test_bit_jaccard_distance(self): def test_sparsevec(self): Item.create(id=1, sparse_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): create_items() @@ -169,7 +170,7 @@ def test_vector_avg(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() - assert np.array_equal(avg, np.array([2.5, 3.5, 4.5])) + assert np.array_equal(avg, [2.5, 3.5, 4.5]) def test_vector_sum(self): sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() @@ -177,7 +178,7 @@ def test_vector_sum(self): Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() - assert np.array_equal(sum, np.array([5, 7, 9])) + assert np.array_equal(sum, [5, 7, 9]) def test_halfvec_avg(self): avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() @@ -185,7 +186,7 @@ def test_halfvec_avg(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() - assert avg.to_list() == [2.5, 3.5, 4.5] + assert avg == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self): sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() @@ -193,7 +194,7 @@ def test_halfvec_sum(self): Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() - assert sum.to_list() == [5, 7, 9] + assert sum == HalfVector([5, 7, 9]) def test_get_or_create(self): Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) @@ -219,5 +220,5 @@ class Meta: # fails with column "embeddings" is of type vector[] but expression is of type text[] # ExtItem.create(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]) # item = ExtItem.get_by_id(1) - # assert np.array_equal(item.embeddings[0], np.array([1, 2, 3])) - # assert np.array_equal(item.embeddings[1], np.array([4, 5, 6])) + # assert np.array_equal(item.embeddings[0], [1, 2, 3]) + # assert np.array_equal(item.embeddings[1], [4, 5, 6]) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py new file mode 100644 index 0000000..4d3e474 --- /dev/null +++ b/tests/test_pg8000.py @@ -0,0 +1,60 @@ +import numpy as np +import os +from pgvector import HalfVector, SparseVector, Vector +from pgvector.pg8000 import register_vector +from pg8000.native import Connection + +conn = Connection(os.environ["USER"], database='pgvector_python_test') + +conn.run('CREATE EXTENSION IF NOT EXISTS vector') +conn.run('DROP TABLE IF EXISTS pg8000_items') +conn.run('CREATE TABLE pg8000_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))') + +register_vector(conn) + + +class TestPg8000: + def setup_method(self): + conn.run('DELETE FROM pg8000_items') + + def test_vector(self): + embedding = np.array([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') + assert np.array_equal(res[0][0], embedding) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None + + def test_vector_class(self): + embedding = Vector([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') + assert np.array_equal(res[0][0], embedding.to_numpy()) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None + + def test_halfvec(self): + embedding = HalfVector([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') + assert res[0][0] == embedding + assert res[1][0] is None + + def test_bit(self): + embedding = '101' + conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') + assert res[0][0] == '101' + assert res[1][0] is None + + def test_sparsevec(self): + embedding = SparseVector([1.5, 2, 3]) + conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + + res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') + assert res[0][0] == embedding + assert res[1][0] is None diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index c4e1c22..698b34f 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,6 +1,8 @@ import numpy as np -from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector, Vector +from pgvector import Bit, HalfVector, SparseVector, Vector +from pgvector.psycopg import register_vector, register_vector_async import psycopg +from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) @@ -13,7 +15,7 @@ class TestPsycopg: - def setup_method(self, test_method): + def setup_method(self): conn.execute('DELETE FROM psycopg_items') def test_vector(self): @@ -44,42 +46,40 @@ def test_vector_text_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, np.array([3, 2, 1.5])) + assert np.array_equal(res, [3, 2, 1.5]) def test_vector_binary_format_non_contiguous(self): embedding = np.flipud(np.array([1.5, 2, 3])) assert not embedding.data.contiguous res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, np.array([3, 2, 1.5])) + assert np.array_equal(res, [3, 2, 1.5]) def test_vector_class_binary_format(self): embedding = Vector([1.5, 2, 3]) res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] - assert np.array_equal(res, np.array([1.5, 2, 3])) + assert np.array_equal(res, [1.5, 2, 3]) def test_vector_class_text_format(self): embedding = Vector([1.5, 2, 3]) res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, np.array([1.5, 2, 3])) + assert np.array_equal(res, [1.5, 2, 3]) def test_halfvec(self): embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0] - assert res.to_list() == [1.5, 2, 3] + assert res == HalfVector([1.5, 2, 3]) def test_halfvec_binary_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] - assert res.to_list() == [1.5, 2, 3] - assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) + assert res == HalfVector([1.5, 2, 3]) def test_halfvec_text_format(self): embedding = HalfVector([1.5, 2, 3]) res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] - assert res.to_list() == [1.5, 2, 3] - assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3])) + assert res == HalfVector([1.5, 2, 3]) def test_bit(self): embedding = Bit([True, False, True]) @@ -104,19 +104,17 @@ def test_sparsevec(self): conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] - assert res.to_list() == [1.5, 2, 3] + assert res == SparseVector([1.5, 2, 3]) def test_sparsevec_binary_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] - assert res.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) + assert res == embedding def test_sparsevec_text_format(self): embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] - assert res.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0])) + assert res == embedding def test_text_copy_from(self): embedding = np.array([1.5, 2, 3]) @@ -154,8 +152,8 @@ def test_binary_copy_to(self): cur = conn.cursor() with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: for row in copy.rows(): - assert Vector.from_binary(row[0]).to_list() == [1.5, 2, 3] - assert HalfVector.from_binary(row[1]).to_list() == [1.5, 2, 3] + assert np.array_equal(Vector.from_binary(row[0]).to_numpy(), embedding) + assert HalfVector.from_binary(row[1]) == half_embedding def test_binary_copy_to_set_types(self): embedding = np.array([1.5, 2, 3]) @@ -166,7 +164,7 @@ def test_binary_copy_to_set_types(self): copy.set_types(['vector', 'halfvec']) for row in copy.rows(): assert np.array_equal(row[0], embedding) - assert row[1].to_list() == [1.5, 2, 3] + assert row[1] == half_embedding def test_vector_array(self): embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] @@ -176,6 +174,18 @@ def test_vector_array(self): assert np.array_equal(res[0][0], embeddings[0]) assert np.array_equal(res[0][1], embeddings[1]) + def test_pool(self): + def configure(conn): + register_vector(conn) + + pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure) + + with pool.connection() as conn: + res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() + assert np.array_equal(res[0], [1, 2, 3]) + + pool.close() + @pytest.mark.asyncio async def test_async(self): conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) @@ -195,3 +205,19 @@ async def test_async(self): assert np.array_equal(res[0][1], embedding) assert res[0][1].dtype == np.float32 assert res[1][1] is None + + @pytest.mark.asyncio + async def test_async_pool(self): + async def configure(conn): + await register_vector_async(conn) + + pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure) + await pool.open() + + async with pool.connection() as conn: + async with conn.cursor() as cur: + await cur.execute("SELECT '[1,2,3]'::vector") + res = await cur.fetchone() + assert np.array_equal(res[0], [1, 2, 3]) + + await pool.close() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index c93fce4..7f4932d 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,7 +1,9 @@ import numpy as np -from pgvector.psycopg2 import register_vector, HalfVector, SparseVector +from pgvector import HalfVector, SparseVector, Vector +from pgvector.psycopg2 import register_vector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor +from psycopg2.pool import ThreadedConnectionPool conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -11,11 +13,11 @@ cur.execute('DROP TABLE IF EXISTS psycopg2_items') cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector[], half_embeddings halfvec[], sparse_embeddings sparsevec[])') -register_vector(cur, globally=False, arrays=True) +register_vector(cur) class TestPsycopg2: - def setup_method(self, test_method): + def setup_method(self): cur.execute('DELETE FROM psycopg2_items') def test_vector(self): @@ -28,13 +30,32 @@ def test_vector(self): assert res[0][0].dtype == np.float32 assert res[1][0] is None + def test_vector_class(self): + embedding = Vector([1.5, 2, 3]) + cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert np.array_equal(res[0][0], embedding.to_numpy()) + assert res[0][0].dtype == np.float32 + assert res[1][0] is None + def test_halfvec(self): embedding = [1.5, 2, 3] cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0].to_list() == [1.5, 2, 3] + assert res[0][0] == HalfVector([1.5, 2, 3]) + assert res[1][0] is None + + def test_halfvec_class(self): + embedding = HalfVector([1.5, 2, 3]) + cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][0] == embedding assert res[1][0] is None def test_bit(self): @@ -52,7 +73,7 @@ def test_sparsevec(self): cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0].to_list() == [1.5, 2, 3] + assert res[0][0] == SparseVector([1.5, 2, 3]) assert res[1][0] is None def test_vector_array(self): @@ -70,8 +91,7 @@ def test_halfvec_array(self): cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0][0].to_list() == [1.5, 2, 3] - assert res[0][1].to_list() == [4.5, 5, 6] + assert res[0] == [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] def test_sparsevec_array(self): embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] @@ -79,18 +99,38 @@ def test_sparsevec_array(self): cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0][0].to_list() == [1.5, 2, 3] - assert res[0][1].to_list() == [4.5, 5, 6] + assert res[0] == [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] def test_cursor_factory(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) - register_vector(cur, globally=False) + register_vector(cur) conn.close() def test_cursor_factory_connection(self): for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) - register_vector(conn, globally=False) + register_vector(conn) conn.close() + + def test_pool(self): + 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 np.array_equal(res[0], [1, 2, 3]) + finally: + pool.putconn(conn) + + pool.closeall() diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 06fe81a..dff03dd 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,14 +1,15 @@ import numpy as np -from pgvector.utils import SparseVector +from pgvector import SparseVector import pytest from scipy.sparse import coo_array +from struct import pack class TestSparseVector: def test_list(self): vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] - assert vec.to_numpy().tolist() == [1, 0, 2, 0, 3, 0] + assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] def test_list_dimensions(self): @@ -52,6 +53,12 @@ def test_repr(self): assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' + def test_equality(self): + assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) + assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) + assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) + assert SparseVector({}, 1) != SparseVector({}, 2) + def test_dimensions(self): assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 @@ -62,8 +69,26 @@ def test_values(self): assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] def test_to_coo(self): - assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]] + assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) def test_zero_vector_text(self): vec = SparseVector({}, 3) assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() + + def test_from_text(self): + vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') + assert vec.dimensions() == 6 + assert vec.indices() == [0, 2, 4] + assert vec.values() == [1.5, 2, 3] + assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + + def test_from_binary(self): + data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) + vec = SparseVector.from_binary(data) + assert vec.dimensions() == 6 + assert vec.indices() == [0, 2, 4] + assert vec.values() == [1.5, 2, 3] + assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + assert vec.to_binary() == data diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 1ca0ea3..0d8d1ca 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,7 +1,10 @@ +import asyncpg import numpy as np -from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum +import os +from pgvector import HalfVector, SparseVector, Vector +from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest -from sqlalchemy import create_engine, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY +from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY from sqlalchemy.exc import StatementError from sqlalchemy.ext.automap import automap_base from sqlalchemy.orm import declarative_base, Session @@ -15,8 +18,57 @@ mapped_column = Column sqlalchemy_version = 1 -engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -with Session(engine) as session: +psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') +psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') + + +@event.listens_for(psycopg2_type_engine, "connect") +def psycopg2_connect(dbapi_connection, connection_record): + from pgvector.psycopg2 import register_vector + register_vector(dbapi_connection) + + +pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') + +if sqlalchemy_version > 1: + psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + + @event.listens_for(psycopg_type_engine, "connect") + def psycopg_connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector + register_vector(dbapi_connection) + + psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + + @event.listens_for(psycopg_async_type_engine.sync_engine, "connect") + def connect(dbapi_connection, connection_record): + from pgvector.psycopg import register_vector_async + dbapi_connection.run_async(register_vector_async) + + asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + asyncpg_type_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test') + + @event.listens_for(asyncpg_type_engine.sync_engine, "connect") + def connect(dbapi_connection, connection_record): + from pgvector.asyncpg import register_vector + dbapi_connection.run_async(register_vector) + +engines = [psycopg2_engine, psycopg2_type_engine, pg8000_engine] +array_engines = [psycopg2_type_engine] +async_engines = [] +async_array_engines = [] + +if sqlalchemy_version > 1: + engines += [psycopg_engine, psycopg_type_engine] + array_engines += [psycopg_type_engine] + # TODO support asyncpg_type_engine + async_engines += [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine] + async_array_engines += [psycopg_async_type_engine, asyncpg_engine] + +setup_engine = engines[0] +with Session(setup_engine) as session: session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() @@ -32,10 +84,11 @@ class Item(Base): binary_embedding = mapped_column(BIT(3)) sparse_embedding = mapped_column(SPARSEVEC(3)) embeddings = mapped_column(ARRAY(VECTOR(3))) + half_embeddings = mapped_column(ARRAY(HALFVEC(3))) -Base.metadata.drop_all(engine) -Base.metadata.create_all(engine) +Base.metadata.drop_all(setup_engine) +Base.metadata.create_all(setup_engine) index = Index( 'sqlalchemy_orm_index', @@ -44,24 +97,47 @@ class Item(Base): postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} ) -index.create(engine) +index.create(setup_engine) + +half_precision_index = Index( + 'sqlalchemy_orm_half_precision_index', + func.cast(Item.embedding, HALFVEC(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'halfvec_l2_ops'} +) +half_precision_index.create(setup_engine) + +binary_quantize_index = Index( + 'sqlalchemy_orm_binary_quantize_index', + func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 64}, + postgresql_ops={'embedding': 'bit_hamming_ops'} +) +binary_quantize_index.create(setup_engine) def create_items(): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) - session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) - session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) - session.commit() + with Session(setup_engine) as session: + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) + session.commit() + + +def delete_items(): + with Session(setup_engine) as session: + session.query(Item).delete() + session.commit() +@pytest.mark.parametrize('engine', engines) class TestSqlalchemy: - def setup_method(self, test_method): - with Session(engine) as session: - session.query(Item).delete() - session.commit() + def setup_method(self): + delete_items() - def test_core(self): + def test_core(self, engine): metadata = MetaData() item_table = Table( @@ -96,250 +172,257 @@ def test_core(self): ) hnsw_index.create(engine) - def test_orm(self): + def test_orm(self, engine): item = Item(embedding=np.array([1.5, 2, 3])) item2 = Item(embedding=[4, 5, 6]) item3 = Item() - session = Session(engine) - session.add(item) - session.add(item2) - session.add(item3) - session.commit() + with Session(engine) as session: + session.add(item) + session.add(item2) + session.add(item3) + session.commit() stmt = select(Item) with Session(engine) as session: 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 + # TODO improve + assert items[0].id % 3 == 1 + assert items[1].id % 3 == 2 + assert items[2].id % 3 == 0 assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) assert items[0].embedding.dtype == np.float32 assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) assert items[1].embedding.dtype == np.float32 assert items[2].embedding is None - def test_vector(self): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] + def test_vector(self, engine): + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert np.array_equal(item.embedding, [1, 2, 3]) - def test_vector_l2_distance(self): + def test_vector_l2_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l2_distance_orm(self): + def test_vector_l2_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_vector_max_inner_product_orm(self): + def test_vector_max_inner_product_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_vector_cosine_distance_orm(self): + def test_vector_cosine_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l1_distance_orm(self): + def test_vector_l1_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self): - session = Session(engine) - session.add(Item(id=1, half_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] + def test_halfvec(self, engine): + with Session(engine) as session: + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding == HalfVector([1, 2, 3]) - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l2_distance_orm(self): + def test_halfvec_l2_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_max_inner_product_orm(self): + def test_halfvec_max_inner_product_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_cosine_distance_orm(self): + def test_halfvec_cosine_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l1_distance_orm(self): + def test_halfvec_l1_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' + def test_bit(self, engine): + with Session(engine) as session: + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_hamming_distance_orm(self): + def test_bit_hamming_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self, engine): + 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): + def test_bit_jaccard_distance_orm(self, engine): + if engine == pg8000_engine: + return + create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self): - session = Session(engine) - session.add(Item(id=1, sparse_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + def test_sparsevec(self, engine): + with Session(engine) as session: + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l2_distance_orm(self): + def test_sparsevec_l2_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_max_inner_product_orm(self): + def test_sparsevec_max_inner_product_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_cosine_distance_orm(self): + def test_sparsevec_cosine_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self, engine): create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l1_distance_orm(self): + def test_sparsevec_l1_distance_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_filter(self): + def test_filter(self, engine): create_items() with Session(engine) as session: items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() assert [v.id for v in items] == [1] - def test_filter_orm(self): + def test_filter_orm(self, engine): create_items() with Session(engine) as session: items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] - def test_select(self): + def test_select(self, engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() assert items[0] == 3 - def test_select_orm(self): + def test_select_orm(self, engine): with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() assert items[0] == 3 - def test_avg(self): + def test_avg(self, engine): with Session(engine) as session: res = session.query(avg(Item.embedding)).first()[0] assert res is None @@ -348,7 +431,7 @@ def test_avg(self): res = session.query(avg(Item.embedding)).first()[0] assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) - def test_avg_orm(self): + def test_avg_orm(self, engine): with Session(engine) as session: res = session.scalars(select(avg(Item.embedding))).first() assert res is None @@ -357,7 +440,7 @@ def test_avg_orm(self): res = session.scalars(select(avg(Item.embedding))).first() assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) - def test_sum(self): + def test_sum(self, engine): with Session(engine) as session: res = session.query(sum(Item.embedding)).first()[0] assert res is None @@ -366,7 +449,7 @@ def test_sum(self): res = session.query(sum(Item.embedding)).first()[0] assert np.array_equal(res, np.array([5, 7, 9])) - def test_sum_orm(self): + def test_sum_orm(self, engine): with Session(engine) as session: res = session.scalars(select(sum(Item.embedding))).first() assert res is None @@ -375,80 +458,192 @@ def test_sum_orm(self): res = session.scalars(select(sum(Item.embedding))).first() assert np.array_equal(res, np.array([5, 7, 9])) - def test_bad_dimensions(self): + def test_bad_dimensions(self, engine): item = Item(embedding=[1, 2]) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + session.commit() - def test_bad_ndim(self): + def test_bad_ndim(self, engine): item = Item(embedding=np.array([[1, 2, 3]])) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected ndim to be 1'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='expected ndim to be 1'): + session.commit() - def test_bad_dtype(self): + def test_bad_dtype(self, engine): item = Item(embedding=np.array(['one', 'two', 'three'])) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='could not convert string to float'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() - def test_inspect(self): + def test_inspect(self, engine): columns = inspect(engine).get_columns('sqlalchemy_orm_item') assert isinstance(columns[1]['type'], VECTOR) - def test_literal_binds(self): + def test_literal_binds(self, engine): sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) - def test_insert(self): - session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) + def test_insert(self, engine): + with Session(engine) as session: + session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) - def test_insert_bulk(self): - session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) + def test_insert_bulk(self, engine): + with Session(engine) as session: + session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) # register_vector in psycopg2 tests change this behavior # def test_insert_text(self): - # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) + # with Session(engine) as session: + # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) - def test_automap(self): + def test_automap(self, engine): metadata = MetaData() metadata.reflect(engine, only=['sqlalchemy_orm_item']) AutoBase = automap_base(metadata=metadata) AutoBase.prepare() AutoItem = AutoBase.classes.sqlalchemy_orm_item - session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) - item = session.query(AutoItem).first() - assert item.embedding.tolist() == [1, 2, 3] + with Session(engine) as session: + session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) + item = session.query(AutoItem).first() + assert np.array_equal(item.embedding, [1, 2, 3]) - def test_vector_array(self): - session = Session(engine) - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) - session.commit() + def test_half_precision(self, engine): + create_items() + with Session(engine) as session: + items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() + assert [v.id for v in items] == [1, 3, 2] - with engine.connect() as connection: - from pgvector.psycopg2 import register_vector - register_vector(connection.connection.dbapi_connection, globally=False, arrays=True) + def test_binary_quantize(self, engine): + with Session(engine) as session: + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + items = session.query(Item).order_by(distance).all() + assert [v.id for v in items] == [2, 3, 1] + + +@pytest.mark.parametrize('engine', array_engines) +class TestSqlalchemyArray: + def setup_method(self): + delete_items() + + def test_vector_array(self, engine): + with Session(engine) as session: + session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.commit() # this fails if the driver does not cast arrays - item = Session(bind=connection).get(Item, 1) - assert item.embeddings[0].tolist() == [1, 2, 3] - assert item.embeddings[1].tolist() == [4, 5, 6] + item = session.get(Item, 1) + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) + + def test_halfvec_array(self, engine): + with Session(engine) as session: + session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.commit() + + # this fails if the driver does not cast arrays + item = session.get(Item, 1) + assert item.half_embeddings == [HalfVector([1, 2, 3]), HalfVector([4, 5, 6])] + + +@pytest.mark.parametrize('engine', async_engines) +class TestSqlalchemyAsync: + def setup_method(self): + delete_items() + + @pytest.mark.asyncio + async def test_vector(self, engine): + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + embedding = np.array([1, 2, 3]) + session.add(Item(id=1, embedding=embedding)) + item = await session.get(Item, 1) + assert np.array_equal(item.embedding, embedding) + + await engine.dispose() + + @pytest.mark.asyncio + async def test_halfvec(self, engine): + 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(Item, 1) + assert item.half_embedding == HalfVector(embedding) + + await engine.dispose() @pytest.mark.asyncio - @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+') - async def test_async(self): - engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') + async def test_bit(self, engine): + async_session = async_sessionmaker(engine, expire_on_commit=False) + + async with async_session() as session: + async with session.begin(): + embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' + session.add(Item(id=1, binary_embedding=embedding)) + item = await session.get(Item, 1) + assert item.binary_embedding == embedding + + await engine.dispose() + + @pytest.mark.asyncio + async def test_sparsevec(self, engine): + 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, sparse_embedding=embedding)) + item = await session.get(Item, 1) + assert item.sparse_embedding == SparseVector(embedding) + + await engine.dispose() + + @pytest.mark.asyncio + async def test_avg(self, engine): async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - avg = await session.scalars(select(func.avg(Item.embedding))) - assert avg.first() == '[2.5,3.5,4.5]' + res = await session.scalars(select(avg(Item.embedding))) + assert np.array_equal(res.first(), [2.5, 3.5, 4.5]) + + await engine.dispose() + + +@pytest.mark.parametrize('engine', async_array_engines) +class TestSqlalchemyAsyncArray: + def setup_method(self): + delete_items() + + @pytest.mark.asyncio + async def test_vector_array(self, engine): + 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=[Vector([1, 2, 3]), Vector([4, 5, 6])])) + item = await session.get(Item, 1) + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) + + session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + item = await session.get(Item, 2) + assert np.array_equal(item.embeddings[0], [1, 2, 3]) + assert np.array_equal(item.embeddings[1], [4, 5, 6]) await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index 4cb0e9b..f4994f4 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,9 +1,9 @@ import numpy as np -from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, SparseVector, avg, sum +from pgvector import HalfVector, SparseVector +from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest -from sqlalchemy import Column, Index from sqlalchemy.exc import StatementError -from sqlmodel import Field, Session, SQLModel, create_engine, delete, select, text +from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text from typing import Any, Optional engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -15,10 +15,10 @@ class Item(SQLModel, table=True): __tablename__ = 'sqlmodel_item' id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[Any] = Field(default=None, sa_column=Column(VECTOR(3))) - half_embedding: Optional[Any] = Field(default=None, sa_column=Column(HALFVEC(3))) - binary_embedding: Optional[Any] = Field(default=None, sa_column=Column(BIT(3))) - sparse_embedding: Optional[Any] = Field(default=None, sa_column=Column(SPARSEVEC(3))) + embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) + half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) + binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) + sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) SQLModel.metadata.drop_all(engine) @@ -35,15 +35,15 @@ class Item(SQLModel, table=True): def create_items(): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) - session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) - session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) - session.commit() + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) + session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) + session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) + session.commit() class TestSqlmodel: - def setup_method(self, test_method): + def setup_method(self): with Session(engine) as session: session.exec(delete(Item)) session.commit() @@ -53,11 +53,11 @@ def test_orm(self): item2 = Item(embedding=[4, 5, 6]) item3 = Item() - session = Session(engine) - session.add(item) - session.add(item2) - session.add(item3) - session.commit() + with Session(engine) as session: + session.add(item) + session.add(item2) + session.add(item3) + session.commit() stmt = select(Item) with Session(engine) as session: @@ -72,11 +72,11 @@ def test_orm(self): assert items[2].embedding is None def test_vector(self): - session = Session(engine) - session.add(Item(id=1, embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.embedding.tolist() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert np.array_equal(item.embedding, np.array([1, 2, 3])) def test_vector_l2_distance(self): create_items() @@ -103,11 +103,11 @@ def test_vector_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] def test_halfvec(self): - session = Session(engine) - session.add(Item(id=1, half_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.half_embedding.to_list() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, half_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.half_embedding == HalfVector([1, 2, 3]) def test_halfvec_l2_distance(self): create_items() @@ -134,11 +134,11 @@ def test_halfvec_l1_distance(self): assert [v.id for v in items] == [1, 3, 2] def test_bit(self): - session = Session(engine) - session.add(Item(id=1, binary_embedding='101')) - session.commit() - item = session.get(Item, 1) - assert item.binary_embedding == '101' + with Session(engine) as session: + session.add(Item(id=1, binary_embedding='101')) + session.commit() + item = session.get(Item, 1) + assert item.binary_embedding == '101' def test_bit_hamming_distance(self): create_items() @@ -153,11 +153,11 @@ def test_bit_jaccard_distance(self): assert [v.id for v in items] == [2, 3, 1] def test_sparsevec(self): - session = Session(engine) - session.add(Item(id=1, sparse_embedding=[1, 2, 3])) - session.commit() - item = session.get(Item, 1) - assert item.sparse_embedding.to_list() == [1, 2, 3] + with Session(engine) as session: + session.add(Item(id=1, sparse_embedding=[1, 2, 3])) + session.commit() + item = session.get(Item, 1) + assert item.sparse_embedding == SparseVector([1, 2, 3]) def test_sparsevec_l2_distance(self): create_items() @@ -220,7 +220,7 @@ def test_halfvec_avg(self): session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(avg(Item.half_embedding))).first() - assert res.to_list() == [2.5, 3.5, 4.5] + assert res == HalfVector([2.5, 3.5, 4.5]) def test_halfvec_sum(self): with Session(engine) as session: @@ -229,11 +229,11 @@ def test_halfvec_sum(self): session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(sum(Item.half_embedding))).first() - assert res.to_list() == [5, 7, 9] + assert res == HalfVector([5, 7, 9]) def test_bad_dimensions(self): item = Item(embedding=[1, 2]) - session = Session(engine) - session.add(item) - with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): - session.commit() + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): + session.commit() diff --git a/tests/test_vector.py b/tests/test_vector.py index 1be2bc0..e5a16fe 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,6 +1,7 @@ import numpy as np -from pgvector.utils import Vector +from pgvector import Vector import pytest +from struct import pack class TestVector: @@ -38,5 +39,21 @@ def test_repr(self): assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' + def test_equality(self): + assert Vector([1, 2, 3]) == Vector([1, 2, 3]) + assert Vector([1, 2, 3]) != Vector([1, 2, 4]) + def test_dimensions(self): assert Vector([1, 2, 3]).dimensions() == 3 + + def test_from_text(self): + vec = Vector.from_text('[1.5,2,3]') + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + + def test_from_binary(self): + data = pack('>HH3f', 3, 0, 1.5, 2, 3) + vec = Vector.from_binary(data) + assert vec.to_list() == [1.5, 2, 3] + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + assert vec.to_binary() == data