diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 4d4e8ed..952be69 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,24 +6,28 @@ jobs: strategy: fail-fast: false matrix: - python: [3.13, 3.9] + python: [3.14, "3.10"] steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v7 + - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python }} - - run: pip install -r requirements.txt + - run: pip install --group dev - uses: ankane/setup-postgres@v1 with: database: pgvector_python_test dev-files: true - run: | cd /tmp - git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.4 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 + - run: pip install --group numpy --group scipy + - run: pytest + + - run: pip install --group typing + - run: mypy + - run: ty check diff --git a/.gitignore b/.gitignore index c55ff44..5556c9f 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ venv/ *.pyc __pycache__ .pytest_cache/ +*.lock examples/rag/README.md diff --git a/CHANGELOG.md b/CHANGELOG.md index d0e2730..e509532 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,23 @@ +## 0.5.0 (2026-07-06) + +- Added experimental support for type hints +- Changed `vector` and `halfvec` types to return list for Django, SQLAlchemy, SQLModel, and Peewee +- Changed `vector` type to return `Vector` object for Psycopg 3, Psycopg 2, asyncpg, and pg8000 +- Removed `utils` package (use top-level `pgvector` package instead) +- Removed re-exported classes (use top-level `pgvector` package instead) +- Removed dependency on NumPy +- Dropped support for Python < 3.10 +- Dropped support for SQLAlchemy < 2 + +## 0.4.2 (2025-12-04) + +- Added support for Django 6 +- Added support for `str` objects for `bit` type with SQLAlchemy + +## 0.4.1 (2025-04-26) + +- Fixed `SparseVector` constructor for SciPy sparse matrices + ## 0.4.0 (2025-03-15) - Added top-level `pgvector` package diff --git a/LICENSE.txt b/LICENSE.txt index b612d6d..17e5210 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2021-2025 Andrew Kane +Copyright (c) 2021-2026 Andrew Kane Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/Makefile b/Makefile index f0831c4..7e5aab0 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ .PHONY: lint build publish clean lint: - pycodestyle . --ignore=E501 + pycodestyle . --ignore=E501,E704,W504 --exclude=.venv build: python3 -m build diff --git a/README.md b/README.md index 24d9bb9..26dccfb 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [pgvector](https://github.com/pgvector/pgvector) support for Python -Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee) +Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://codeberg.org/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee) [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) @@ -177,10 +177,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(Base): - embedding = mapped_column(Vector(3)) + embedding: Mapped[list[float]] = mapped_column(VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -259,7 +259,6 @@ index = Index( 'my_index', func.cast(Item.embedding, HALFVEC(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'halfvec_l2_ops'} ) ``` @@ -271,16 +270,47 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2]) session.scalars(select(Item).order_by(order).limit(5)) ``` +#### Binary Quantization + +Use expression indexing for binary quantization + +```python +from pgvector.sqlalchemy import BIT +from sqlalchemy.sql import func + +index = Index( + 'my_index', + func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), + postgresql_using='hnsw', + postgresql_ops={'embedding': 'bit_hamming_ops'} +) +``` + +Get the nearest neighbors by Hamming distance + +```python +order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) +session.scalars(select(Item).order_by(order).limit(5)) +``` + +Re-rank by the original vectors for better recall + +```python +order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) +subquery = session.query(Item).order_by(order).limit(20).subquery() +session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5)) +``` + #### Arrays Add an array column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import ARRAY class Item(Base): - embeddings = mapped_column(ARRAY(Vector(3))) + embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1)) ``` And register the types with the underlying driver @@ -315,7 +345,7 @@ from sqlalchemy import event @event.listens_for(engine, "connect") def connect(dbapi_connection, connection_record): - register_vector(dbapi_connection, arrays=True) + register_vector(dbapi_connection) ``` ## SQLModel @@ -329,10 +359,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) Add a vector column ```python -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR class Item(SQLModel, table=True): - embedding: Any = Field(sa_type=Vector(3)) + embedding: list[float] = Field(sa_type=VECTOR(3)) ``` Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` @@ -409,7 +439,7 @@ Enable the extension conn.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.psycopg import register_vector @@ -443,7 +473,9 @@ conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3)) Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) ``` @@ -472,7 +504,7 @@ cur = conn.cursor() cur.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection or cursor +Register the types with your connection or cursor ```python from pgvector.psycopg2 import register_vector @@ -489,7 +521,9 @@ cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))' Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) ``` @@ -518,7 +552,7 @@ Enable the extension await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.asyncpg import register_vector @@ -544,7 +578,9 @@ await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vect Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding) ``` @@ -572,7 +608,7 @@ Enable the extension conn.run('CREATE EXTENSION IF NOT EXISTS vector') ``` -Register the vector type with your connection +Register the types with your connection ```python from pgvector.pg8000 import register_vector @@ -589,7 +625,9 @@ conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') Insert a vector ```python -embedding = np.array([1, 2, 3]) +from pgvector import Vector + +embedding = Vector([1, 2, 3]) conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding) ``` @@ -668,6 +706,32 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc ## Reference +### Vectors + +Create a vector from a list + +```python +vec = Vector([1, 2, 3]) +``` + +Or a NumPy array + +```python +vec = Vector(np.array([1, 2, 3])) +``` + +Get a list + +```python +lst = vec.to_list() +``` + +Get a NumPy array + +```python +arr = vec.to_numpy() +``` + ### Half Vectors Create a half vector from a list @@ -777,7 +841,7 @@ To get started with development: ```sh git clone https://github.com/pgvector/pgvector-python.git cd pgvector-python -pip install -r requirements.txt +pip install --group dev createdb pgvector_python_test pytest ``` @@ -786,7 +850,7 @@ To run an example: ```sh cd examples/loading -pip install -r requirements.txt +pip install --group dev createdb pgvector_example python3 example.py ``` diff --git a/examples/citus/example.py b/examples/citus/example.py index 915c25f..d1af5b9 100644 --- a/examples/citus/example.py +++ b/examples/citus/example.py @@ -1,4 +1,5 @@ import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -38,7 +39,7 @@ copy.set_types(['vector', 'bigint']) for i in range(rows): - copy.write_row([embeddings[i], categories[i]]) + copy.write_row([Vector(embeddings[i]), categories[i]]) print('Creating index in parallel') conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') diff --git a/examples/citus/pyproject.toml b/examples/citus/pyproject.toml new file mode 100644 index 0000000..3fb37d4 --- /dev/null +++ b/examples/citus/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/citus/requirements.txt b/examples/citus/requirements.txt deleted file mode 100644 index 1cf8ee9..0000000 --- a/examples/citus/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -numpy -pgvector -psycopg[binary] diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 393d1e0..5cf8d69 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -1,6 +1,6 @@ import cohere -import numpy as np -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg conn = psycopg.connect(dbname='pgvector_example', autocommit=True) @@ -9,13 +9,13 @@ register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))') def embed(input, input_type): - co = cohere.Client() - response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary']) - return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] + co = cohere.ClientV2() + response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) + return [bytes(embedding) for embedding in response.embeddings.ubinary] input = [ diff --git a/examples/cohere/pyproject.toml b/examples/cohere/pyproject.toml new file mode 100644 index 0000000..1228b3d --- /dev/null +++ b/examples/cohere/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "cohere", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/cohere/requirements.txt b/examples/cohere/requirements.txt deleted file mode 100644 index 22fd056..0000000 --- a/examples/cohere/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -cohere -pgvector -psycopg[binary] diff --git a/examples/colbert/approximate.py b/examples/colbert/approximate.py new file mode 100644 index 0000000..41f88b2 --- /dev/null +++ b/examples/colbert/approximate.py @@ -0,0 +1,80 @@ +# based on section 3.6 of https://arxiv.org/abs/2004.12832 + +from colbert.infra import ColBERTConfig +from colbert.modeling.checkpoint import Checkpoint +from pgvector.psycopg import register_vector +import psycopg +import warnings + +conn = psycopg.connect(dbname='pgvector_example', autocommit=True) + +conn.execute('CREATE EXTENSION IF NOT EXISTS vector') +register_vector(conn) + +conn.execute('DROP TABLE IF EXISTS documents') +conn.execute('DROP TABLE IF EXISTS document_embeddings') + +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text)') +conn.execute('CREATE TABLE document_embeddings (id bigserial PRIMARY KEY, document_id bigint, embedding vector(128))') + +conn.execute(""" +CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ + WITH queries AS ( + SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) + ), + documents AS ( + SELECT unnest(document) AS document + ), + similarities AS ( + SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents + ), + max_similarities AS ( + SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number + ) + SELECT SUM(max_similarity) FROM max_similarities +$$ LANGUAGE SQL +""") + +warnings.filterwarnings('ignore') # ignore warnings from colbert + +config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) +checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) + +input = [ + 'The dog is barking', + 'The cat is purring', + 'The bear is growling' +] +doc_embeddings = checkpoint.docFromText(input, keep_dims=False) +for content, embeddings in zip(input, doc_embeddings): + with conn.transaction(): + result = conn.execute('INSERT INTO documents (content) VALUES (%s) RETURNING id', (content,)).fetchone() + params = [] + for embedding in embeddings: + params.extend([result[0], embedding.numpy()]) + values = ', '.join(['(%s, %s)' for _ in embeddings]) + conn.execute(f'INSERT INTO document_embeddings (document_id, embedding) VALUES {values}', params) + +conn.execute('CREATE INDEX ON document_embeddings (document_id)') +conn.execute('CREATE INDEX ON document_embeddings USING hnsw (embedding vector_cosine_ops)') + +query = 'puppy' +query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query])[0]] +approximate_stage = ' UNION ALL '.join(['(SELECT document_id FROM document_embeddings ORDER BY embedding <=> %s LIMIT 5)' for _ in query_embeddings]) +sql = f""" +WITH approximate_stage AS ( + {approximate_stage} +), +embeddings AS ( + SELECT document_id, array_agg(embedding) AS embeddings FROM document_embeddings + WHERE document_id IN (SELECT DISTINCT document_id FROM approximate_stage) + GROUP BY document_id +) +SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents +INNER JOIN embeddings ON embeddings.document_id = documents.id +ORDER BY max_sim DESC LIMIT 10 +""" +params = query_embeddings + [query_embeddings] +result = conn.execute(sql, params).fetchall() +for row in result: + print(row) diff --git a/examples/colbert/exact.py b/examples/colbert/exact.py index 1c90b47..e6a2936 100644 --- a/examples/colbert/exact.py +++ b/examples/colbert/exact.py @@ -2,6 +2,7 @@ from colbert.modeling.checkpoint import Checkpoint from pgvector.psycopg import register_vector import psycopg +import warnings conn = psycopg.connect(dbname='pgvector_example', autocommit=True) @@ -28,6 +29,8 @@ $$ LANGUAGE SQL """) +warnings.filterwarnings('ignore') # ignore warnings from colbert + config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) diff --git a/examples/colbert/pyproject.toml b/examples/colbert/pyproject.toml new file mode 100644 index 0000000..7cf1e3f --- /dev/null +++ b/examples/colbert/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "colbert-ai", + "pgvector", + "psycopg[binary]", + "transformers==4.49.0" +] diff --git a/examples/colbert/requirements.txt b/examples/colbert/requirements.txt deleted file mode 100644 index 4402ce8..0000000 --- a/examples/colbert/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -colbert-ai -pgvector -psycopg[binary] diff --git a/examples/colpali/exact.py b/examples/colpali/exact.py index 80bb603..313c8bf 100644 --- a/examples/colpali/exact.py +++ b/examples/colpali/exact.py @@ -1,7 +1,8 @@ from colpali_engine.models import ColQwen2, ColQwen2Processor from colpali_engine.utils.torch_utils import get_torch_device from datasets import load_dataset -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg import torch diff --git a/examples/colpali/pyproject.toml b/examples/colpali/pyproject.toml new file mode 100644 index 0000000..20c26a7 --- /dev/null +++ b/examples/colpali/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "colpali-engine", + "datasets", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/colpali/requirements.txt b/examples/colpali/requirements.txt deleted file mode 100644 index 4cf770d..0000000 --- a/examples/colpali/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -colpali-engine -datasets -pgvector -psycopg[binary] diff --git a/examples/gensim/example.py b/examples/gensim/example.py index cfbf18d..2fc63c8 100644 --- a/examples/gensim/example.py +++ b/examples/gensim/example.py @@ -1,7 +1,7 @@ from gensim.corpora.dictionary import Dictionary from gensim.models import LdaModel from gensim.utils import simple_preprocess -import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -26,7 +26,7 @@ model = LdaModel(corpus, num_topics=20) for content, bow in zip(input, corpus): - embedding = np.array([v[1] for v in model.get_document_topics(bow, minimum_probability=0)]) + embedding = Vector([v[1] for v in model.get_document_topics(bow, minimum_probability=0)]) conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 diff --git a/examples/gensim/pyproject.toml b/examples/gensim/pyproject.toml new file mode 100644 index 0000000..4d2865a --- /dev/null +++ b/examples/gensim/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "gensim", + "numpy", + "pgvector", + "psycopg[binary]", + "scipy<1.13" +] diff --git a/examples/gensim/requirements.txt b/examples/gensim/requirements.txt deleted file mode 100644 index 15411cd..0000000 --- a/examples/gensim/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -gensim -numpy -pgvector -psycopg[binary] -scipy<1.13 diff --git a/examples/hybrid_search/pyproject.toml b/examples/hybrid_search/pyproject.toml new file mode 100644 index 0000000..130f287 --- /dev/null +++ b/examples/hybrid_search/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/hybrid_search/requirements.txt b/examples/hybrid_search/requirements.txt deleted file mode 100644 index 237dcd1..0000000 --- a/examples/hybrid_search/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -sentence-transformers diff --git a/examples/image_search/pyproject.toml b/examples/image_search/pyproject.toml new file mode 100644 index 0000000..8ada0a7 --- /dev/null +++ b/examples/image_search/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "matplotlib", + "pgvector", + "psycopg[binary]", + "torch", + "torchvision", + "tqdm" +] diff --git a/examples/image_search/requirements.txt b/examples/image_search/requirements.txt deleted file mode 100644 index 3d82365..0000000 --- a/examples/image_search/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -matplotlib -pgvector -psycopg[binary] -torch -torchvision -tqdm diff --git a/examples/imagehash/example.py b/examples/imagehash/example.py index f49af40..595dcb5 100644 --- a/examples/imagehash/example.py +++ b/examples/imagehash/example.py @@ -1,7 +1,8 @@ from datasets import load_dataset from imagehash import phash import matplotlib.pyplot as plt -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg @@ -18,7 +19,7 @@ def hash_image(img): conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, hash bit(64))') print('Loading dataset') -dataset = load_dataset('mnist') +dataset = load_dataset('ylecun/mnist') print('Generating hashes') images = [{'hash': hash_image(row['image'])} for row in dataset['train']] diff --git a/examples/imagehash/pyproject.toml b/examples/imagehash/pyproject.toml new file mode 100644 index 0000000..4f0e2b4 --- /dev/null +++ b/examples/imagehash/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "datasets", + "imagehash", + "matplotlib", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/imagehash/requirements.txt b/examples/imagehash/requirements.txt deleted file mode 100644 index e3971e6..0000000 --- a/examples/imagehash/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -datasets -imagehash -matplotlib -pgvector -psycopg[binary] diff --git a/examples/implicit/example.py b/examples/implicit/example.py index f70eb8c..a4bd0f4 100644 --- a/examples/implicit/example.py +++ b/examples/implicit/example.py @@ -1,30 +1,32 @@ import implicit from implicit.datasets.movielens import get_movielens -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer, String -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class User(Base): __tablename__ = 'user' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' - id = mapped_column(Integer, primary_key=True) - title = mapped_column(String) - factors = mapped_column(Vector(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + title: Mapped[str] = mapped_column(String) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) @@ -41,10 +43,10 @@ class Item(Base): session.execute(insert(User), users) session.execute(insert(Item), items) -user = session.get(User, 1) +user = session.get_one(User, 1) items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) print('user-based recs:', [item.title for item in items]) -item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() +item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).one() items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) diff --git a/examples/implicit/pyproject.toml b/examples/implicit/pyproject.toml new file mode 100644 index 0000000..ec305ee --- /dev/null +++ b/examples/implicit/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "h5py", + "implicit", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/implicit/requirements.txt b/examples/implicit/requirements.txt deleted file mode 100644 index 424abbd..0000000 --- a/examples/implicit/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -h5py -implicit -pgvector -psycopg[binary] -SQLAlchemy diff --git a/examples/lightfm/example.py b/examples/lightfm/example.py index fcb9027..728adc1 100644 --- a/examples/lightfm/example.py +++ b/examples/lightfm/example.py @@ -1,31 +1,33 @@ from lightfm import LightFM from lightfm.datasets import fetch_movielens -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Float, Integer, String -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session engine = create_engine('postgresql+psycopg://localhost/pgvector_example') with engine.connect() as conn: conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class User(Base): __tablename__ = 'user' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' - id = mapped_column(Integer, primary_key=True) - title = mapped_column(String) - factors = mapped_column(Vector(20)) - bias = mapped_column(Float) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + title: Mapped[str] = mapped_column(String) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) + bias: Mapped[float] = mapped_column(Float) Base.metadata.drop_all(engine) @@ -45,12 +47,12 @@ class Item(Base): session.execute(insert(User), users) session.execute(insert(Item), items) -user = session.get(User, 1) +user = session.get_one(User, 1) # subtract item bias for negative inner product items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors) - Item.bias).limit(5)) print('user-based recs:', [item.title for item in items]) # broken due to https://github.com/lyst/lightfm/issues/682 -item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() +item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).one() items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.title for item in items]) diff --git a/examples/lightfm/pyproject.toml b/examples/lightfm/pyproject.toml new file mode 100644 index 0000000..047e3f1 --- /dev/null +++ b/examples/lightfm/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "lightfm", + "pgvector", + "psycopg[binary]", + "SQLAlchemy" +] diff --git a/examples/lightfm/requirements.txt b/examples/lightfm/requirements.txt deleted file mode 100644 index cfa5f51..0000000 --- a/examples/lightfm/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -lightfm -pgvector -psycopg[binary] -SQLAlchemy diff --git a/examples/loading/example.py b/examples/loading/example.py index 0702129..4c0e579 100644 --- a/examples/loading/example.py +++ b/examples/loading/example.py @@ -1,4 +1,5 @@ import numpy as np +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -25,12 +26,12 @@ copy.set_types(['vector']) for i, embedding in enumerate(embeddings): + copy.write_row([Vector(embedding)]) + # show progress if i % 10000 == 0: print('.', end='', flush=True) - copy.write_row([embedding]) - print('\nSuccess!') # create any indexes *after* loading initial data (skipping for this example) diff --git a/examples/loading/pyproject.toml b/examples/loading/pyproject.toml new file mode 100644 index 0000000..7247ec0 --- /dev/null +++ b/examples/loading/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]" +] + +[tool.uv.sources] +pgvector = { path = "../..", editable = true } diff --git a/examples/loading/requirements.txt b/examples/loading/requirements.txt deleted file mode 100644 index 1cf8ee9..0000000 --- a/examples/loading/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -numpy -pgvector -psycopg[binary] diff --git a/examples/openai/example.py b/examples/openai/example.py index b9a078c..ba1fa7a 100644 --- a/examples/openai/example.py +++ b/examples/openai/example.py @@ -1,5 +1,5 @@ -import numpy as np from openai import OpenAI +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg @@ -25,10 +25,10 @@ def embed(input): ] embeddings = embed(input) for content, embedding in zip(input, embeddings): - conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, np.array(embedding))) + conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Vector(embedding))) query = 'forest' query_embedding = embed([query])[0] -result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (np.array(query_embedding),)).fetchall() +result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (Vector(query_embedding),)).fetchall() for row in result: print(row[0]) diff --git a/examples/openai/halfvec.py b/examples/openai/halfvec.py index 185c785..b2234f6 100644 --- a/examples/openai/halfvec.py +++ b/examples/openai/halfvec.py @@ -1,5 +1,6 @@ from openai import OpenAI -from pgvector.psycopg import register_vector, HalfVector +from pgvector import HalfVector +from pgvector.psycopg import register_vector import psycopg conn = psycopg.connect(dbname='pgvector_example', autocommit=True) diff --git a/examples/openai/pyproject.toml b/examples/openai/pyproject.toml new file mode 100644 index 0000000..5665e02 --- /dev/null +++ b/examples/openai/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "openai", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/openai/requirements.txt b/examples/openai/requirements.txt deleted file mode 100644 index 18587e2..0000000 --- a/examples/openai/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -openai -pgvector -psycopg[binary] diff --git a/examples/rag/example.py b/examples/rag/example.py index 4d5d307..9d5a784 100644 --- a/examples/rag/example.py +++ b/examples/rag/example.py @@ -1,11 +1,11 @@ # Run: +# ollama serve # ollama pull llama3.2 # ollama pull nomic-embed-text -# ollama serve -import numpy as np import ollama from pathlib import Path +from pgvector import Vector from pgvector.psycopg import register_vector import psycopg import urllib.request @@ -51,11 +51,10 @@ # embed query # nomic-embed-text has task instruction prefix -input = 'search_query: ' + query -embedding = ollama.embed(model='nomic-embed-text', input=input).embeddings[0] +embedding = ollama.embed(model='nomic-embed-text', input=f'search_query: {query}').embeddings[0] # retrieve chunks -result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (np.array(embedding),)).fetchall() +result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (Vector(embedding),)).fetchall() context = '\n\n'.join([row[0] for row in result]) # get answer diff --git a/examples/rag/pyproject.toml b/examples/rag/pyproject.toml new file mode 100644 index 0000000..2e74790 --- /dev/null +++ b/examples/rag/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "ollama", + "pgvector", + "psycopg[binary]" +] diff --git a/examples/rag/requirements.txt b/examples/rag/requirements.txt deleted file mode 100644 index 4eb5864..0000000 --- a/examples/rag/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -ollama -pgvector -psycopg[binary] diff --git a/examples/rdkit/example.py b/examples/rdkit/example.py index afb56ec..dec2993 100644 --- a/examples/rdkit/example.py +++ b/examples/rdkit/example.py @@ -1,7 +1,8 @@ # good resource # https://www.rdkit.org/docs/GettingStartedInPython.html#morgan-fingerprints-circular-fingerprints -from pgvector.psycopg import register_vector, Bit +from pgvector import Bit +from pgvector.psycopg import register_vector import psycopg from rdkit import Chem from rdkit.Chem import AllChem diff --git a/examples/rdkit/pyproject.toml b/examples/rdkit/pyproject.toml new file mode 100644 index 0000000..b717809 --- /dev/null +++ b/examples/rdkit/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "rdkit" +] diff --git a/examples/rdkit/requirements.txt b/examples/rdkit/requirements.txt deleted file mode 100644 index 85a3e4f..0000000 --- a/examples/rdkit/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -rdkit diff --git a/examples/sentence_transformers/pyproject.toml b/examples/sentence_transformers/pyproject.toml new file mode 100644 index 0000000..130f287 --- /dev/null +++ b/examples/sentence_transformers/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "sentence-transformers" +] diff --git a/examples/sentence_transformers/requirements.txt b/examples/sentence_transformers/requirements.txt deleted file mode 100644 index 237dcd1..0000000 --- a/examples/sentence_transformers/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -pgvector -psycopg[binary] -sentence-transformers diff --git a/examples/sparse_search/example.py b/examples/sparse_search/example.py index 2b5daea..e3fec01 100644 --- a/examples/sparse_search/example.py +++ b/examples/sparse_search/example.py @@ -2,7 +2,8 @@ # https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/ # https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1 -from pgvector.psycopg import register_vector, SparseVector +from pgvector import SparseVector +from pgvector.psycopg import register_vector import psycopg import torch from transformers import AutoModelForMaskedLM, AutoTokenizer diff --git a/examples/sparse_search/pyproject.toml b/examples/sparse_search/pyproject.toml new file mode 100644 index 0000000..d9adf93 --- /dev/null +++ b/examples/sparse_search/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "numpy", + "pgvector", + "psycopg[binary]", + "torch", + "transformers" +] diff --git a/examples/sparse_search/requirements.txt b/examples/sparse_search/requirements.txt deleted file mode 100644 index 3de81c7..0000000 --- a/examples/sparse_search/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -numpy -pgvector -psycopg[binary] -torch -transformers diff --git a/examples/surprise/example.py b/examples/surprise/example.py index bd7d18d..699ea31 100644 --- a/examples/surprise/example.py +++ b/examples/surprise/example.py @@ -1,6 +1,6 @@ -from pgvector.sqlalchemy import Vector +from pgvector.sqlalchemy import VECTOR from sqlalchemy import create_engine, insert, select, text, Integer -from sqlalchemy.orm import declarative_base, mapped_column, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session from surprise import Dataset, SVD engine = create_engine('postgresql+psycopg://localhost/pgvector_example') @@ -8,21 +8,23 @@ conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) conn.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class User(Base): __tablename__ = 'user' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) class Item(Base): __tablename__ = 'item' - id = mapped_column(Integer, primary_key=True) - factors = mapped_column(Vector(20)) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + factors: Mapped[list[float]] = mapped_column(VECTOR(20)) Base.metadata.drop_all(engine) @@ -40,10 +42,10 @@ class Item(Base): session.execute(insert(User), users) session.execute(insert(Item), items) -user = session.get(User, 1) +user = session.get_one(User, 1) items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) print('user-based recs:', [item.id for item in items]) -item = session.get(Item, 50) +item = session.get_one(Item, 50) items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) print('item-based recs:', [item.id for item in items]) diff --git a/examples/surprise/pyproject.toml b/examples/surprise/pyproject.toml new file mode 100644 index 0000000..bd0cff2 --- /dev/null +++ b/examples/surprise/pyproject.toml @@ -0,0 +1,12 @@ +[project] +name = "example" +version = "0.1.0" +requires-python = ">= 3.10" + +[dependency-groups] +dev = [ + "pgvector", + "psycopg[binary]", + "scikit-surprise", + "SQLAlchemy" +] diff --git a/examples/surprise/requirements.txt b/examples/surprise/requirements.txt deleted file mode 100644 index cb2dca4..0000000 --- a/examples/surprise/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -pgvector -psycopg[binary] -scikit-surprise -SQLAlchemy diff --git a/pgvector/_utils.py b/pgvector/_utils.py new file mode 100644 index 0000000..4cafc32 --- /dev/null +++ b/pgvector/_utils.py @@ -0,0 +1,23 @@ +import sys +from typing import TYPE_CHECKING, TypeAlias + +if TYPE_CHECKING: + import numpy as np + + ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]] +else: + # any value works since not type checking + # TODO use Never when Python 3.10 no longer supported + ndarray = None + + +def is_ndarray(value: object, /) -> bool: + if (numpy := sys.modules.get('numpy')): + return isinstance(value, numpy.ndarray) + return False + + +def is_sparse_array(value: object, /) -> bool: + if (sparse := sys.modules.get('scipy.sparse')): + return isinstance(value, (sparse.sparray, sparse.spmatrix)) + return False diff --git a/pgvector/asyncpg/__init__.py b/pgvector/asyncpg/__init__.py index c6a3b4e..b3b4440 100644 --- a/pgvector/asyncpg/__init__.py +++ b/pgvector/asyncpg/__init__.py @@ -1,11 +1,5 @@ from .register import register_vector -# TODO remove -from .. import Vector, HalfVector, SparseVector - __all__ = [ - 'register_vector', - 'Vector', - 'HalfVector', - 'SparseVector' + 'register_vector' ] diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py index 63726f3..0a2715f 100644 --- a/pgvector/asyncpg/register.py +++ b/pgvector/asyncpg/register.py @@ -1,12 +1,13 @@ +from asyncpg import Connection from .. import Vector, HalfVector, SparseVector -async def register_vector(conn, schema='public'): +async def register_vector(conn: Connection, /, *, schema: str = 'public') -> None: await conn.set_type_codec( 'vector', schema=schema, - encoder=Vector._to_db_binary, - decoder=Vector._from_db_binary, + encoder=lambda v: (v if isinstance(v, Vector) else Vector(v)).to_binary(), + decoder=Vector.from_binary, format='binary' ) @@ -14,16 +15,16 @@ async def register_vector(conn, schema='public'): await conn.set_type_codec( 'halfvec', schema=schema, - encoder=HalfVector._to_db_binary, - decoder=HalfVector._from_db_binary, + encoder=lambda v: (v if isinstance(v, HalfVector) else HalfVector(v)).to_binary(), + decoder=HalfVector.from_binary, format='binary' ) await conn.set_type_codec( 'sparsevec', schema=schema, - encoder=SparseVector._to_db_binary, - decoder=SparseVector._from_db_binary, + encoder=lambda v: (v if isinstance(v, SparseVector) else SparseVector(v)).to_binary(), + decoder=SparseVector.from_binary, format='binary' ) except ValueError as e: diff --git a/pgvector/bit.py b/pgvector/bit.py index 26a9d8d..2b04912 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -1,75 +1,95 @@ -import numpy as np +from __future__ import annotations from struct import pack, unpack_from -from warnings import warn +from typing import TYPE_CHECKING +from ._utils import is_ndarray + +if TYPE_CHECKING: + import numpy as np class Bit: - def __init__(self, value): + _length: int + _data: bytes + + def __init__( + self, + value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]], + / + ) -> None: if isinstance(value, bytes): - self._len = 8 * len(value) + self._length = 8 * len(value) self._data = value + elif isinstance(value, (list, str)): + if isinstance(value, list): + bits = {True: '1', False: '0'} + try: + value = ''.join([bits[v] for v in value]) + except (KeyError, TypeError): + raise ValueError('expected list[bool]') + + length = len(value) + if length % 8 != 0: + value += '0' * (8 - (length % 8)) + + self._length = length + try: + self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + except ValueError: + raise ValueError('expected bit string') + elif is_ndarray(value): + import numpy as np + + if value.dtype != np.bool: + # skip error for result of np.unpackbits + if value.dtype != np.uint8 or np.any(value > 1): + raise ValueError('expected elements to be boolean') + value = value.astype(bool) + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + self._length = len(value) + self._data = np.packbits(value).tobytes() # type: ignore else: - 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') + raise ValueError('expected bytes, str, list, or ndarray') - self._len = len(value) - self._data = np.packbits(value).tobytes() - - def __repr__(self): + def __repr__(self) -> str: return f'Bit({self.to_text()})' - def __eq__(self, other): - if isinstance(other, self.__class__): - return self._len == other._len and self._data == other._data - return False + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, self.__class__): + return NotImplemented + return self._length == other._length and self._data == other._data + + def to_list(self) -> list[bool]: + # TODO improve + return [v != '0' for v in self.to_text()] - def to_list(self): - return self.to_numpy().tolist() + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]: + import numpy as np - def to_numpy(self): - return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) + return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool) - def to_text(self): - return ''.join(format(v, '08b') for v in self._data)[:self._len] + def to_text(self) -> str: + return ''.join(format(v, '08b') for v in self._data)[:self._length] - def to_binary(self): - return pack('>i', self._len) + self._data + def to_binary(self) -> bytes: + return pack('>i', self._length) + self._data @classmethod - def from_text(cls, value): + def from_text(cls, value: str, /) -> Bit: + # cast to ensure always uses str constructor return cls(str(value)) @classmethod - def from_binary(cls, value): - if not isinstance(value, bytes): - raise ValueError('expected bytes') + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Bit: + length, = unpack_from('>i', value) + data = memoryview(value)[4:].tobytes() + + if len(data) != (length + 7) // 8: + raise ValueError('invalid length') bit = cls.__new__(cls) - bit._len = unpack_from('>i', value)[0] - bit._data = value[4:] + bit._length = length + bit._data = data 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 43c64a3..3ed0f83 100644 --- a/pgvector/django/__init__.py +++ b/pgvector/django/__init__.py @@ -6,9 +6,6 @@ from .sparsevec import SparseVectorField from .vector import VectorField -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ 'VectorExtension', 'VectorField', @@ -22,7 +19,5 @@ 'CosineDistance', 'L1Distance', 'HammingDistance', - 'JaccardDistance', - 'HalfVector', - 'SparseVector' + 'JaccardDistance' ] diff --git a/pgvector/django/bit.py b/pgvector/django/bit.py index 2cc847a..eab0534 100644 --- a/pgvector/django/bit.py +++ b/pgvector/django/bit.py @@ -1,32 +1,37 @@ from django import forms from django.db.models import Field +from typing import Any -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class BitField(Field): description = 'Bit string' - def __init__(self, *args, length=None, **kwargs): + def __init__(self, *args: Any, length: int | None = None, **kwargs: Any) -> None: self.length = length super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.length is not None: kwargs['length'] = self.length return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.length is None: return 'bit' return 'bit(%d)' % self.length - def formfield(self, **kwargs): - return super().formfield(form_class=BitFormField, **kwargs) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=BitFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class BitFormField(forms.CharField): - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/django/extensions.py b/pgvector/django/extensions.py index 0573f72..d75884f 100644 --- a/pgvector/django/extensions.py +++ b/pgvector/django/extensions.py @@ -1,6 +1,11 @@ +from django import VERSION from django.contrib.postgres.operations import CreateExtension +from typing import Any class VectorExtension(CreateExtension): - def __init__(self): - self.name = 'vector' + def __init__(self, hints: Any = None) -> None: + if VERSION[0] >= 6: + super().__init__('vector', hints=hints) # type: ignore + else: + self.name = 'vector' diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py index 9df4fdb..352c631 100644 --- a/pgvector/django/functions.py +++ b/pgvector/django/functions.py @@ -1,18 +1,17 @@ from django.db.models import FloatField, Func, Value from .. import Vector, HalfVector, SparseVector +from typing import Any class DistanceBase(Func): - output_field = FloatField() + output_field = FloatField() # type: ignore - def __init__(self, expression, vector, **extra): + def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): - if isinstance(vector, HalfVector): - vector = Value(HalfVector._to_db(vector)) - elif isinstance(vector, SparseVector): - vector = Value(SparseVector._to_db(vector)) - else: - vector = Value(Vector._to_db(vector)) + if isinstance(vector, (Vector, HalfVector, SparseVector)): + vector = Value(vector.to_text()) + elif vector is not None: + vector = Value(Vector(vector).to_text()) # prevent error with unhashable types self._constructor_args = ((expression, vector), extra) @@ -21,9 +20,9 @@ def __init__(self, expression, vector, **extra): class BitDistanceBase(Func): - output_field = FloatField() + output_field = FloatField() # type: ignore - def __init__(self, expression, vector, **extra): + def __init__(self, expression: Any, vector: Any, /, **extra: Any) -> None: if not hasattr(vector, 'resolve_expression'): vector = Value(vector) super().__init__(expression, vector, **extra) diff --git a/pgvector/django/halfvec.py b/pgvector/django/halfvec.py index 3aeb90f..1ab4a47 100644 --- a/pgvector/django/halfvec.py +++ b/pgvector/django/halfvec.py @@ -1,60 +1,54 @@ from django import forms from django.db.models import Field +from typing import Any from .. import HalfVector -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class HalfVectorField(Field): description = 'Half vector' empty_strings_allowed = False - def __init__(self, *args, dimensions=None, **kwargs): + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None: self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.dimensions is None: return 'halfvec' return 'halfvec(%d)' % self.dimensions - def from_db_value(self, value, expression, connection): + def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: return HalfVector._from_db(value) - def to_python(self, value): - if value is None or isinstance(value, HalfVector): + def to_python(self, value: Any) -> list[float] | None: + if isinstance(value, list): return value - elif isinstance(value, str): - return HalfVector._from_db(value) - else: - return HalfVector(value) + return HalfVector._from_db(value) - def get_prep_value(self, value): + def get_prep_value(self, value: Any) -> str | None: return HalfVector._to_db(value) - def value_to_string(self, obj): - return self.get_prep_value(self.value_from_object(obj)) - - def formfield(self, **kwargs): - return super().formfield(form_class=HalfVectorFormField, **kwargs) + def value_to_string(self, obj: Any) -> str: + value = self.get_prep_value(self.value_from_object(obj)) + return '' if value is None else value - -class HalfVectorWidget(forms.TextInput): - def format_value(self, value): - if isinstance(value, HalfVector): - value = value.to_list() - return super().format_value(value) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=HalfVectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class HalfVectorFormField(forms.CharField): - widget = HalfVectorWidget - - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/django/indexes.py b/pgvector/django/indexes.py index 5bec0eb..e095351 100644 --- a/pgvector/django/indexes.py +++ b/pgvector/django/indexes.py @@ -1,20 +1,21 @@ from django.contrib.postgres.indexes import PostgresIndex +from typing import Any class IvfflatIndex(PostgresIndex): suffix = 'ivfflat' - def __init__(self, *expressions, lists=None, **kwargs): + def __init__(self, *expressions: Any, lists: int | None = None, **kwargs: Any) -> None: self.lists = lists super().__init__(*expressions, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple[Any, Any, Any]: path, args, kwargs = super().deconstruct() if self.lists is not None: kwargs['lists'] = self.lists return path, args, kwargs - def get_with_params(self): + def get_with_params(self) -> list[str]: with_params = [] if self.lists is not None: with_params.append('lists = %d' % self.lists) @@ -24,12 +25,12 @@ def get_with_params(self): class HnswIndex(PostgresIndex): suffix = 'hnsw' - def __init__(self, *expressions, m=None, ef_construction=None, **kwargs): + def __init__(self, *expressions: Any, m: int | None = None, ef_construction: int | None = None, **kwargs: Any) -> None: self.m = m self.ef_construction = ef_construction super().__init__(*expressions, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple[Any, Any, Any]: path, args, kwargs = super().deconstruct() if self.m is not None: kwargs['m'] = self.m @@ -37,7 +38,7 @@ def deconstruct(self): kwargs['ef_construction'] = self.ef_construction return path, args, kwargs - def get_with_params(self): + def get_with_params(self) -> list[str]: with_params = [] if self.m is not None: with_params.append('m = %d' % self.m) diff --git a/pgvector/django/sparsevec.py b/pgvector/django/sparsevec.py index 580f27c..a10c133 100644 --- a/pgvector/django/sparsevec.py +++ b/pgvector/django/sparsevec.py @@ -1,46 +1,52 @@ from django import forms from django.db.models import Field +from typing import Any from .. import SparseVector -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class SparseVectorField(Field): description = 'Sparse vector' empty_strings_allowed = False - def __init__(self, *args, dimensions=None, **kwargs): + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any): self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.dimensions is None: return 'sparsevec' return 'sparsevec(%d)' % self.dimensions - def from_db_value(self, value, expression, connection): + def from_db_value(self, value: Any, expression: Any, connection: Any) -> SparseVector | None: return SparseVector._from_db(value) - def to_python(self, value): + def to_python(self, value: Any) -> SparseVector | None: return SparseVector._from_db(value) - def get_prep_value(self, value): + def get_prep_value(self, value: Any) -> str | None: return SparseVector._to_db(value) - def value_to_string(self, obj): - return self.get_prep_value(self.value_from_object(obj)) + def value_to_string(self, obj: Any) -> str: + value = self.get_prep_value(self.value_from_object(obj)) + return '' if value is None else value - def formfield(self, **kwargs): - return super().formfield(form_class=SparseVectorFormField, **kwargs) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=SparseVectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class SparseVectorWidget(forms.TextInput): - def format_value(self, value): + def format_value(self, value: Any) -> Any: if isinstance(value, SparseVector): value = value.to_text() return super().format_value(value) @@ -49,7 +55,7 @@ def format_value(self, value): class SparseVectorFormField(forms.CharField): widget = SparseVectorWidget - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/django/vector.py b/pgvector/django/vector.py index 861cfde..42970c6 100644 --- a/pgvector/django/vector.py +++ b/pgvector/django/vector.py @@ -1,73 +1,54 @@ from django import forms from django.db.models import Field -import numpy as np +from typing import Any from .. import Vector -# https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ +# https://docs.djangoproject.com/en/6.0/howto/custom-model-fields/ class VectorField(Field): description = 'Vector' empty_strings_allowed = False - def __init__(self, *args, dimensions=None, **kwargs): + def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None: self.dimensions = dimensions super().__init__(*args, **kwargs) - def deconstruct(self): + def deconstruct(self) -> tuple[Any, Any, Any, Any]: name, path, args, kwargs = super().deconstruct() if self.dimensions is not None: kwargs['dimensions'] = self.dimensions return name, path, args, kwargs - def db_type(self, connection): + def db_type(self, connection: Any) -> str: if self.dimensions is None: return 'vector' return 'vector(%d)' % self.dimensions - def from_db_value(self, value, expression, connection): + def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None: return Vector._from_db(value) - def to_python(self, value): + def to_python(self, value: Any) -> list[float] | None: if isinstance(value, list): - return np.array(value, dtype=np.float32) + return value return Vector._from_db(value) - def get_prep_value(self, value): + def get_prep_value(self, value: Any) -> str | None: return Vector._to_db(value) - def value_to_string(self, obj): - return self.get_prep_value(self.value_from_object(obj)) + def value_to_string(self, obj: Any) -> str: + value = self.get_prep_value(self.value_from_object(obj)) + return '' if value is None else value - def validate(self, value, model_instance): - if isinstance(value, np.ndarray): - value = value.tolist() - super().validate(value, model_instance) - - def run_validators(self, value): - if isinstance(value, np.ndarray): - value = value.tolist() - super().run_validators(value) - - def formfield(self, **kwargs): - return super().formfield(form_class=VectorFormField, **kwargs) - - -class VectorWidget(forms.TextInput): - def format_value(self, value): - if isinstance(value, np.ndarray): - value = value.tolist() - return super().format_value(value) + def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field: + return super().formfield( + form_class=VectorFormField if form_class is None else form_class, + choices_form_class=choices_form_class, + **kwargs + ) class VectorFormField(forms.CharField): - widget = VectorWidget - - def has_changed(self, initial, data): - if isinstance(initial, np.ndarray): - initial = initial.tolist() - return super().has_changed(initial, data) - - def to_python(self, value): + def to_python(self, value: Any) -> Any: if isinstance(value, str) and value == '': return None return super().to_python(value) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index f335f2f..1a2c8c7 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -1,83 +1,114 @@ -import numpy as np -from struct import pack, unpack_from +from __future__ import annotations +from array import array +import struct +import sys +from typing import TYPE_CHECKING +from ._utils import is_ndarray, ndarray +if TYPE_CHECKING: + import numpy as np -class HalfVector: - def __init__(self, value): - # asarray still copies if same dtype - if not isinstance(value, np.ndarray) or value.dtype != '>f2': - value = np.asarray(value, dtype='>f2') - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - self._value = value - def __repr__(self): +class HalfVector: + _value: array[int] # uses uint16 since no float16 + + def __init__(self, value: list[float] | ndarray, /) -> None: + if isinstance(value, list): + dim = len(value) + try: + self._value = array('H', struct.pack(f'{dim}e', *value)) + except struct.error: + raise ValueError('expected list[float]') + elif is_ndarray(value): + import numpy as np + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + arr = array('H') + arr.frombytes(value.astype(np.float16, order='C', copy=False).data.cast('B')) + self._value = arr + else: + raise ValueError('expected list or ndarray') + + def __repr__(self) -> str: return f'HalfVector({self.to_list()})' - def __eq__(self, other): - if isinstance(other, self.__class__): - return np.array_equal(self.to_numpy(), other.to_numpy()) - return False + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, self.__class__): + return NotImplemented + return self._value == other._value - def dimensions(self): + def dimensions(self) -> int: return len(self._value) - def to_list(self): - return self._value.tolist() + def to_list(self) -> list[float]: + dim = len(self._value) + return list(struct.unpack(f'{dim}e', self._value)) - def to_numpy(self): - return self._value + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float16]]: + import numpy as np + return np.frombuffer(self._value, dtype=np.float16) - def to_text(self): - return '[' + ','.join([str(float(v)) for v in self._value]) + ']' + def to_text(self) -> str: + return f'[{",".join([str(v) for v in self.to_list()])}]' - def to_binary(self): - return pack('>HH', self.dimensions(), 0) + self._value.tobytes() + def to_binary(self) -> bytes: + if sys.byteorder == 'big': + value = self._value + else: + value = array('H', self._value) + value.byteswap() + return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod - def from_text(cls, value): - return cls([float(v) for v in value[1:-1].split(',')]) + def from_text(cls, value: str, /) -> HalfVector: + return cls(cls._from_text(value)) @classmethod - def from_binary(cls, value): - dim, unused = unpack_from('>HH', value) - return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: + dim, unused = struct.unpack_from('>HH', value) + data = memoryview(value)[4:] - @classmethod - def _to_db(cls, value, dim=None): - if value is None: - return value + if len(data) != 2 * dim: + raise ValueError('invalid length') - if not isinstance(value, cls): - value = cls(value) + if unused != 0: + raise ValueError('expected unused to be 0') - if dim is not None and value.dimensions() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) + arr = array('H') + arr.frombytes(data) + if sys.byteorder != 'big': + arr.byteswap() - return value.to_text() + vec = cls.__new__(cls) + vec._value = arr + return vec @classmethod - def _to_db_binary(cls, value): + def _from_text(cls, value: str, /) -> list[float]: + return [float(v) for v in value[1:-1].split(',')] + + @classmethod + def _to_db(cls, value: list[float] | ndarray | HalfVector | None, /) -> str | None: if value is None: return value - if not isinstance(value, cls): + # fewer allocations for lists + if isinstance(value, list): + return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore + + if not isinstance(value, HalfVector): value = cls(value) - return value.to_binary() + return value.to_text() @classmethod - def _from_db(cls, value): - if value is None or isinstance(value, cls): + def _from_db(cls, value: str | HalfVector | None, /) -> list[float] | None: + if value is None: return value - return cls.from_text(value) - - @classmethod - def _from_db_binary(cls, value): - if value is None or isinstance(value, cls): - return value + if isinstance(value, HalfVector): + return value.to_list() - return cls.from_binary(value) + return cls._from_text(value) diff --git a/pgvector/peewee/__init__.py b/pgvector/peewee/__init__.py index df21200..f557517 100644 --- a/pgvector/peewee/__init__.py +++ b/pgvector/peewee/__init__.py @@ -3,14 +3,9 @@ from .sparsevec import SparseVectorField from .vector import VectorField -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ 'VectorField', 'HalfVectorField', 'FixedBitField', - 'SparseVectorField', - 'HalfVector', - 'SparseVector' + 'SparseVectorField' ] diff --git a/pgvector/peewee/bit.py b/pgvector/peewee/bit.py index ee5f12f..fcc7d8f 100644 --- a/pgvector/peewee/bit.py +++ b/pgvector/peewee/bit.py @@ -1,21 +1,22 @@ from peewee import Expression, Field +from typing import Any class FixedBitField(Field): field_type = 'bit' - def __init__(self, max_length=None, *args, **kwargs): + def __init__(self, max_length: int | None = None, *args: Any, **kwargs: Any) -> None: self.max_length = max_length - super(FixedBitField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) - def get_modifiers(self): - return self.max_length and [self.max_length] or None + def get_modifiers(self) -> list[int] | None: + return [self.max_length] if self.max_length else None - def _distance(self, op, vector): + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def hamming_distance(self, vector): + def hamming_distance(self, vector: object, /) -> Expression: return self._distance('<~>', vector) - def jaccard_distance(self, vector): + def jaccard_distance(self, vector: object, /) -> Expression: return self._distance('<%%>', vector) diff --git a/pgvector/peewee/halfvec.py b/pgvector/peewee/halfvec.py index 0901fd2..5e1e92c 100644 --- a/pgvector/peewee/halfvec.py +++ b/pgvector/peewee/halfvec.py @@ -1,34 +1,35 @@ from peewee import Expression, Field +from typing import Any from .. import HalfVector class HalfVectorField(Field): field_type = 'halfvec' - def __init__(self, dimensions=None, *args, **kwargs): + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions - super(HalfVectorField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None - def db_value(self, value): + def db_value(self, value: Any) -> str | None: return HalfVector._to_db(value) - def python_value(self, value): + def python_value(self, value: Any) -> list[float] | None: return HalfVector._from_db(value) - def _distance(self, op, vector): + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector): + def l2_distance(self, vector: object, /) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector): + def max_inner_product(self, vector: object, /) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector): + def cosine_distance(self, vector: object, /) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector): + def l1_distance(self, vector: object, /) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/sparsevec.py b/pgvector/peewee/sparsevec.py index 86dea73..dd83644 100644 --- a/pgvector/peewee/sparsevec.py +++ b/pgvector/peewee/sparsevec.py @@ -1,34 +1,35 @@ from peewee import Expression, Field +from typing import Any from .. import SparseVector class SparseVectorField(Field): field_type = 'sparsevec' - def __init__(self, dimensions=None, *args, **kwargs): + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions - super(SparseVectorField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None - def db_value(self, value): + def db_value(self, value: Any) -> str | None: return SparseVector._to_db(value) - def python_value(self, value): + def python_value(self, value: Any) -> SparseVector | None: return SparseVector._from_db(value) - def _distance(self, op, vector): + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector): + def l2_distance(self, vector: object, /) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector): + def max_inner_product(self, vector: object, /) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector): + def cosine_distance(self, vector: object, /) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector): + def l1_distance(self, vector: object, /) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/peewee/vector.py b/pgvector/peewee/vector.py index 83f9997..a28b6b5 100644 --- a/pgvector/peewee/vector.py +++ b/pgvector/peewee/vector.py @@ -1,34 +1,35 @@ from peewee import Expression, Field +from typing import Any from .. import Vector class VectorField(Field): field_type = 'vector' - def __init__(self, dimensions=None, *args, **kwargs): + def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) -> None: self.dimensions = dimensions - super(VectorField, self).__init__(*args, **kwargs) + super().__init__(*args, **kwargs) - def get_modifiers(self): - return self.dimensions and [self.dimensions] or None + def get_modifiers(self) -> list[int] | None: + return [self.dimensions] if self.dimensions else None - def db_value(self, value): + def db_value(self, value: Any) -> str | None: return Vector._to_db(value) - def python_value(self, value): + def python_value(self, value: Any) -> list[float] | None: return Vector._from_db(value) - def _distance(self, op, vector): + def _distance(self, op: str, vector: object, /) -> Expression: return Expression(lhs=self, op=op, rhs=self.to_value(vector)) - def l2_distance(self, vector): + def l2_distance(self, vector: object, /) -> Expression: return self._distance('<->', vector) - def max_inner_product(self, vector): + def max_inner_product(self, vector: object, /) -> Expression: return self._distance('<#>', vector) - def cosine_distance(self, vector): + def cosine_distance(self, vector: object, /) -> Expression: return self._distance('<=>', vector) - def l1_distance(self, vector): + def l1_distance(self, vector: object, /) -> Expression: return self._distance('<+>', vector) diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index 15ee219..f3978d7 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -1,23 +1,29 @@ -import numpy as np +from pg8000.native import Connection +from typing import cast from .. import Vector, HalfVector, SparseVector -def register_vector(conn): +def register_vector(conn: Connection, /) -> None: # use to_regtype to get first matching type in search path - res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") + res = cast(list[tuple[str, int]], conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")) type_info = dict(res) if 'vector' not in type_info: raise RuntimeError('vector type not found in the database') - conn.register_out_adapter(Vector, Vector._to_db) - conn.register_out_adapter(np.ndarray, Vector._to_db) - conn.register_in_adapter(type_info['vector'], Vector._from_db) + conn.register_out_adapter(Vector, lambda v: v.to_text()) + conn.register_in_adapter(type_info['vector'], Vector.from_text) + + try: + import numpy as np + conn.register_out_adapter(np.ndarray, lambda v: Vector(v).to_text()) + except ImportError: + pass if 'halfvec' in type_info: - conn.register_out_adapter(HalfVector, HalfVector._to_db) - conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db) + conn.register_out_adapter(HalfVector, lambda v: v.to_text()) + conn.register_in_adapter(type_info['halfvec'], HalfVector.from_text) if 'sparsevec' in type_info: - conn.register_out_adapter(SparseVector, SparseVector._to_db) - conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db) + conn.register_out_adapter(SparseVector, lambda v: v.to_text()) + conn.register_in_adapter(type_info['sparsevec'], SparseVector.from_text) diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py index 980af84..8922eda 100644 --- a/pgvector/psycopg/__init__.py +++ b/pgvector/psycopg/__init__.py @@ -1,13 +1,6 @@ from .register import register_vector, register_vector_async -# TODO remove -from .. import Bit, HalfVector, SparseVector, Vector - __all__ = [ 'register_vector', 'register_vector_async', - 'Vector', - 'HalfVector', - 'Bit', - 'SparseVector' ] diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index cffe8fb..f899497 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,25 +1,28 @@ +from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any from .. import Bit class BitDumper(Dumper): - format = Format.TEXT - def dump(self, obj): - return Bit._to_db(obj).encode('utf8') + def dump(self, obj: Bit) -> Buffer | None: + return obj.to_text().encode('utf8') class BitBinaryDumper(BitDumper): - format = Format.BINARY - def dump(self, obj): - return Bit._to_db_binary(obj) + def dump(self, obj: Bit) -> Buffer | None: + return obj.to_binary() -def register_bit_info(context, info): +def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> None: + assert info is not None info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index b3a0060..9ed0453 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,45 +1,43 @@ +from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any from .. import HalfVector class HalfVectorDumper(Dumper): - format = Format.TEXT - def dump(self, obj): - return HalfVector._to_db(obj).encode('utf8') + def dump(self, obj: HalfVector) -> Buffer | None: + return obj.to_text().encode('utf8') class HalfVectorBinaryDumper(HalfVectorDumper): - format = Format.BINARY - def dump(self, obj): - return HalfVector._to_db_binary(obj) + def dump(self, obj: HalfVector) -> Buffer | None: + return obj.to_binary() class HalfVectorLoader(Loader): - format = Format.TEXT - def load(self, data): + def load(self, data: Buffer) -> HalfVector | None: if isinstance(data, memoryview): - data = bytes(data) - return HalfVector._from_db(data.decode('utf8')) + data = data.tobytes() + return HalfVector.from_text(data.decode('utf8')) class HalfVectorBinaryLoader(HalfVectorLoader): - format = Format.BINARY - def load(self, data): - if isinstance(data, memoryview): - data = bytes(data) - return HalfVector._from_db_binary(data) + def load(self, data: Buffer) -> HalfVector | None: + return HalfVector.from_binary(data) -def register_halfvec_info(context, info): +def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/register.py b/pgvector/psycopg/register.py index b93fd3e..ed95042 100644 --- a/pgvector/psycopg/register.py +++ b/pgvector/psycopg/register.py @@ -1,11 +1,13 @@ +from psycopg import AsyncConnection, Connection from psycopg.types import TypeInfo +from typing import Any from .bit import register_bit_info from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info -def register_vector(context): +def register_vector(context: Connection[Any], /) -> None: info = TypeInfo.fetch(context, 'vector') register_vector_info(context, info) @@ -21,7 +23,7 @@ def register_vector(context): register_sparsevec_info(context, info) -async def register_vector_async(context): +async def register_vector_async(context: AsyncConnection[Any], /) -> None: info = await TypeInfo.fetch(context, 'vector') register_vector_info(context, info) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 384a0e1..57fd19a 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,45 +1,43 @@ +from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any from .. import SparseVector class SparseVectorDumper(Dumper): - format = Format.TEXT - def dump(self, obj): - return SparseVector._to_db(obj).encode('utf8') + def dump(self, obj: SparseVector) -> Buffer | None: + return obj.to_text().encode('utf8') class SparseVectorBinaryDumper(SparseVectorDumper): - format = Format.BINARY - def dump(self, obj): - return SparseVector._to_db_binary(obj) + def dump(self, obj: SparseVector) -> Buffer | None: + return obj.to_binary() class SparseVectorLoader(Loader): - format = Format.TEXT - def load(self, data): + def load(self, data: Buffer) -> SparseVector | None: if isinstance(data, memoryview): - data = bytes(data) - return SparseVector._from_db(data.decode('utf8')) + data = data.tobytes() + return SparseVector.from_text(data.decode('utf8')) class SparseVectorBinaryLoader(SparseVectorLoader): - format = Format.BINARY - def load(self, data): - if isinstance(data, memoryview): - data = bytes(data) - return SparseVector._from_db_binary(data) + def load(self, data: Buffer) -> SparseVector | None: + return SparseVector.from_binary(data) -def register_sparsevec_info(context, info): +def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None: info.register(context) # add oid to anonymous class for set_types diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index db9e826..7ef16b0 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,46 +1,50 @@ +from __future__ import annotations import psycopg +from psycopg import BaseConnection +from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper from psycopg.pq import Format +from psycopg.types import TypeInfo +from typing import Any from .. import Vector +from .._utils import ndarray class VectorDumper(Dumper): - format = Format.TEXT - def dump(self, obj): - return Vector._to_db(obj).encode('utf8') + def dump(self, obj: Vector | ndarray) -> Buffer | None: + if not isinstance(obj, Vector): + obj = Vector(obj) + return obj.to_text().encode('utf8') class VectorBinaryDumper(VectorDumper): - format = Format.BINARY - def dump(self, obj): - return Vector._to_db_binary(obj) + def dump(self, obj: Vector | ndarray) -> Buffer | None: + if not isinstance(obj, Vector): + obj = Vector(obj) + return obj.to_binary() class VectorLoader(Loader): - format = Format.TEXT - def load(self, data): + def load(self, data: Buffer) -> Vector | None: if isinstance(data, memoryview): - data = bytes(data) - return Vector._from_db(data.decode('utf8')) + data = data.tobytes() + return Vector.from_text(data.decode('utf8')) class VectorBinaryLoader(VectorLoader): - format = Format.BINARY - def load(self, data): - if isinstance(data, memoryview): - data = bytes(data) - return Vector._from_db_binary(data) + def load(self, data: Buffer) -> Vector | None: + return Vector.from_binary(data) -def register_vector_info(context, info): +def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> None: if info is None: raise psycopg.ProgrammingError('vector type not found in the database') info.register(context) @@ -50,9 +54,9 @@ def register_vector_info(context, info): binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) adapters = context.adapters - adapters.register_dumper('numpy.ndarray', text_dumper) - adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_dumper(Vector, text_dumper) adapters.register_dumper(Vector, binary_dumper) + adapters.register_dumper('numpy.ndarray', text_dumper) + adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py index 33e5124..b3b4440 100644 --- a/pgvector/psycopg2/__init__.py +++ b/pgvector/psycopg2/__init__.py @@ -1,10 +1,5 @@ from .register import register_vector -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ - 'register_vector', - 'HalfVector', - 'SparseVector' + 'register_vector' ] diff --git a/pgvector/psycopg2/halfvec.py b/pgvector/psycopg2/halfvec.py index 0a4c736..e1be330 100644 --- a/pgvector/psycopg2/halfvec.py +++ b/pgvector/psycopg2/halfvec.py @@ -1,20 +1,25 @@ -from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import HalfVector class HalfvecAdapter: - def __init__(self, value): + def __init__(self, value: HalfVector) -> None: self._value = value - def getquoted(self): - return adapt(HalfVector._to_db(self._value)).getquoted() + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() -def cast_halfvec(value, cur): - return HalfVector._from_db(value) +def cast_halfvec(value: str | bytes | None, cur: cursor) -> HalfVector | None: + if value is None: + return None + if isinstance(value, bytes): + raise ValueError('expected str') + return HalfVector.from_text(value) -def register_halfvec_info(oid, array_oid, scope): +def register_halfvec_info(oid: int, array_oid: int | None, scope: connection | cursor | None, /) -> None: halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) register_type(halfvec, scope) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 1bc9d44..2a5af49 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -1,13 +1,13 @@ import psycopg2 -from psycopg2.extensions import cursor +from psycopg2.extensions import connection, cursor from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info # note: register_adapter is always global -def register_vector(conn_or_curs, globally=False, arrays=True): - conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection +def register_vector(conn_or_curs: connection | cursor, /, *, globally: bool = False, arrays: bool = True) -> None: + conn = conn_or_curs if isinstance(conn_or_curs, connection) else conn_or_curs.connection cur = conn.cursor(cursor_factory=cursor) scope = None if globally else conn_or_curs diff --git a/pgvector/psycopg2/sparsevec.py b/pgvector/psycopg2/sparsevec.py index 148eff2..47dbc66 100644 --- a/pgvector/psycopg2/sparsevec.py +++ b/pgvector/psycopg2/sparsevec.py @@ -1,20 +1,25 @@ -from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import SparseVector class SparsevecAdapter: - def __init__(self, value): + def __init__(self, value: SparseVector) -> None: self._value = value - def getquoted(self): - return adapt(SparseVector._to_db(self._value)).getquoted() + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() -def cast_sparsevec(value, cur): - return SparseVector._from_db(value) +def cast_sparsevec(value: str | bytes | None, cur: cursor) -> SparseVector | None: + if value is None: + return None + if isinstance(value, bytes): + raise ValueError('expected str') + return SparseVector.from_text(value) -def register_sparsevec_info(oid, array_oid, scope): +def register_sparsevec_info(oid: int, array_oid: int | None, scope: connection | cursor | None, /) -> None: sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) register_type(sparsevec, scope) diff --git a/pgvector/psycopg2/vector.py b/pgvector/psycopg2/vector.py index 562de18..7d37648 100644 --- a/pgvector/psycopg2/vector.py +++ b/pgvector/psycopg2/vector.py @@ -1,21 +1,29 @@ -import numpy as np -from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type +from __future__ import annotations +from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type +from typing import Any from .. import Vector +from .._utils import ndarray class VectorAdapter: - def __init__(self, value): + def __init__(self, value: Vector | ndarray) -> None: + if not isinstance(value, Vector): + value = Vector(value) self._value = value - def getquoted(self): - return adapt(Vector._to_db(self._value)).getquoted() + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() -def cast_vector(value, cur): - return Vector._from_db(value) +def cast_vector(value: str | bytes | None, cur: cursor) -> Vector | None: + if value is None: + return None + if isinstance(value, bytes): + raise ValueError('expected str') + return Vector.from_text(value) -def register_vector_info(oid, array_oid, scope): +def register_vector_info(oid: int, array_oid: int | None, scope: connection | cursor | None, /) -> None: vector = new_type((oid,), 'VECTOR', cast_vector) register_type(vector, scope) @@ -23,5 +31,10 @@ def register_vector_info(oid, array_oid, scope): vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector) register_type(vectorarray, scope) - register_adapter(np.ndarray, VectorAdapter) register_adapter(Vector, VectorAdapter) + + try: + import numpy as np + register_adapter(np.ndarray, VectorAdapter) + except ImportError: + pass diff --git a/pgvector/py.typed b/pgvector/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 8df2dfd..3631f78 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -1,71 +1,100 @@ -import numpy as np +from __future__ import annotations from struct import pack, unpack_from +from typing import TYPE_CHECKING, Final, cast, overload +from ._utils import is_sparse_array, ndarray -NO_DEFAULT = object() +if TYPE_CHECKING: + import numpy as np + from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix + + +class Sentinel: + pass + + +NO_DEFAULT: Final[Sentinel] = Sentinel() class SparseVector: - def __init__(self, value, dimensions=NO_DEFAULT, /): - if value.__class__.__module__.startswith('scipy.sparse.'): + _dim: int + _indices: list[int] + _values: list[float] + + @overload + def __init__(self, value: dict[int, float], dimensions: int, /) -> None: ... + + @overload + def __init__(self, value: list[float] | ndarray | sparray | spmatrix, /) -> None: ... + + def __init__( + self, + value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, + dimensions: int | Sentinel = NO_DEFAULT, + / + ) -> None: + if is_sparse_array(value): if dimensions is not NO_DEFAULT: raise ValueError('extra argument') - self._from_sparse(value) + self._from_sparse(value) # type: ignore elif isinstance(value, dict): if dimensions is NO_DEFAULT: raise ValueError('missing dimensions') - self._from_dict(value, dimensions) + self._from_dict(value, dimensions) # type: ignore else: if dimensions is not NO_DEFAULT: raise ValueError('extra argument') - self._from_dense(value) + self._from_dense(value) # type: ignore - def __repr__(self): + def __repr__(self) -> str: 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 __eq__(self, other: object, /) -> bool: + if not isinstance(other, self.__class__): + return NotImplemented + return self._dim == other._dim and self._indices == other._indices and self._values == other._values - def dimensions(self): + def dimensions(self) -> int: return self._dim - def indices(self): + def indices(self) -> list[int]: return self._indices - def values(self): + def values(self) -> list[float]: return self._values - def to_coo(self): + def to_coo(self) -> coo_array: from scipy.sparse import coo_array coords = ([0] * len(self._indices), self._indices) return coo_array((self._values, coords), shape=(1, self._dim)) - def to_list(self): + def to_list(self) -> list[float]: vec = [0.0] * self._dim for i, v in zip(self._indices, self._values): vec[i] = v return vec - def to_numpy(self): - vec = np.repeat(0.0, self._dim).astype(np.float32) + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: + import numpy as np + + vec = np.zeros(self._dim, dtype=np.float32) for i, v in zip(self._indices, self._values): vec[i] = v return vec - def to_text(self): - return '{' + ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + '}/' + str(int(self._dim)) + def to_text(self) -> str: + elements = ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + return f'{{{elements}}}/{int(self._dim)}' - def to_binary(self): + def to_binary(self) -> bytes: nnz = len(self._indices) return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) - def _from_dict(self, d, dim): + def _from_dict(self, d: dict[int, float], dim: int) -> None: elements = [(i, v) for i, v in d.items() if v != 0] elements.sort() @@ -73,33 +102,34 @@ def _from_dict(self, d, dim): self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] - def _from_sparse(self, value): - value = value.tocoo() + def _from_sparse(self, arr: sparray | spmatrix, /) -> None: + value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore - if value.ndim == 1: - self._dim = value.shape[0] - elif value.ndim == 2 and value.shape[0] == 1: - self._dim = value.shape[1] + shape = cast(tuple[int, ...], value.shape) + if len(shape) == 1: + self._dim = shape[0] + elif len(shape) == 2 and shape[0] == 1: + self._dim = shape[1] else: raise ValueError('expected ndim to be 1') if hasattr(value, 'coords'): # scipy 1.13+ - self._indices = value.coords[0].tolist() + self._indices = value.coords[-1].tolist() else: self._indices = value.col.tolist() - self._values = value.data.tolist() + self._values = [float(v) for v in value.data] - def _from_dense(self, value): + def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) self._indices = [i for i, v in enumerate(value) if v != 0] self._values = [float(value[i]) for i in self._indices] @classmethod - def from_text(cls, value): + def from_text(cls, value: str, /) -> SparseVector: elements, dim = value.split('/', 2) - indices = [] - values = [] + indices: list[int] = [] + values: list[float] = [] # split on empty string returns single element list if len(elements) > 2: for e in elements[1:-1].split(','): @@ -109,14 +139,21 @@ def from_text(cls, value): return cls._from_parts(int(dim), indices, values) @classmethod - def from_binary(cls, value): + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> SparseVector: dim, nnz, unused = unpack_from('>iii', value) - indices = unpack_from(f'>{nnz}i', value, 12) - values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) - return cls._from_parts(int(dim), list(indices), list(values)) + + if len(value) != 12 + 8 * nnz: + raise ValueError('invalid length') + + if unused != 0: + raise ValueError('expected unused to be 0') + + indices = list(unpack_from(f'>{nnz}i', value, 12)) + values = list(unpack_from(f'>{nnz}f', value, 12 + nnz * 4)) + return cls._from_parts(dim, indices, values) @classmethod - def _from_parts(cls, dim, indices, values): + def _from_parts(cls, dim: int, indices: list[int], values: list[float], /) -> SparseVector: vec = cls.__new__(cls) vec._dim = dim vec._indices = indices @@ -124,38 +161,18 @@ def _from_parts(cls, dim, indices, values): return vec @classmethod - def _to_db(cls, value, dim=None): + def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector | None, /) -> str | None: if value is None: return value - if not isinstance(value, cls): + if not isinstance(value, SparseVector): value = cls(value) - if dim is not None and value.dimensions() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) - return value.to_text() @classmethod - def _to_db_binary(cls, value): - if value is None: - return value - - if not isinstance(value, cls): - value = cls(value) - - return value.to_binary() - - @classmethod - def _from_db(cls, value): - if value is None or isinstance(value, cls): + def _from_db(cls, value: str | SparseVector | None, /) -> SparseVector | None: + if value is None or isinstance(value, SparseVector): return value return cls.from_text(value) - - @classmethod - def _from_db_binary(cls, value): - if value is None or isinstance(value, cls): - return value - - return cls.from_binary(value) diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py index 52adf88..593876f 100644 --- a/pgvector/sqlalchemy/__init__.py +++ b/pgvector/sqlalchemy/__init__.py @@ -5,17 +5,12 @@ from .vector import VECTOR from .vector import VECTOR as Vector -# TODO remove -from .. import HalfVector, SparseVector - __all__ = [ 'Vector', 'VECTOR', 'HALFVEC', 'BIT', 'SPARSEVEC', - 'HalfVector', - 'SparseVector', 'avg', 'sum' ] diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 0f83f3c..7c861ef 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -1,26 +1,24 @@ -from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float +from sqlalchemy.dialects.postgresql.base import PGBit +from sqlalchemy.types import TypeDecorator, Float +from sqlalchemy import Operators +from typing import Any -class BIT(UserDefinedType): +class BIT(TypeDecorator[Any]): + impl = PGBit cache_ok = True - def __init__(self, length=None): - super(UserDefinedType, self).__init__() - self.length = length + def process_bind_param(self, value: Any, dialect: Any) -> Any: + if dialect.__class__.__name__ == 'PGDialect_asyncpg' and isinstance(value, str): + import asyncpg + return asyncpg.BitString(value) # type: ignore + return value - def get_col_spec(self, **kw): - if self.length is None: - return 'BIT' - return 'BIT(%d)' % self.length - - class comparator_factory(UserDefinedType.Comparator): - def hamming_distance(self, other): + class Comparator(TypeDecorator.Comparator[Any]): + def hamming_distance(self, other: object, /) -> Operators: return self.op('<~>', return_type=Float)(other) - def jaccard_distance(self, other): + def jaccard_distance(self, other: object, /) -> Operators: return self.op('<%>', return_type=Float)(other) - -# for reflection -ischema_names['bit'] = BIT + comparator_factory = Comparator # type: ignore diff --git a/pgvector/sqlalchemy/functions.py b/pgvector/sqlalchemy/functions.py index 72e3ca7..ca4cae9 100644 --- a/pgvector/sqlalchemy/functions.py +++ b/pgvector/sqlalchemy/functions.py @@ -1,9 +1,10 @@ # https://docs.sqlalchemy.org/en/20/core/functions.html # include sum for a consistent API from sqlalchemy.sql.functions import ReturnTypeFromArgs, sum +from typing import Any -class avg(ReturnTypeFromArgs): +class avg(ReturnTypeFromArgs[Any]): inherit_cache = True package = 'pgvector' diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 10688b5..9d717c3 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -1,51 +1,55 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String +from sqlalchemy import Dialect, Operators +from typing import Any from .. import HalfVector -class HALFVEC(UserDefinedType): +class HALFVEC(UserDefinedType[Any]): cache_ok = True _string = String() - def __init__(self, dim=None): - super(UserDefinedType, self).__init__() + def __init__(self, dim: int | None = None) -> None: + super().__init__() self.dim = dim - def get_col_spec(self, **kw): + def get_col_spec(self, **kw: Any) -> str: if self.dim is None: return 'HALFVEC' return 'HALFVEC(%d)' % self.dim - def bind_processor(self, dialect): - def process(value): - return HalfVector._to_db(value, self.dim) + def bind_processor(self, dialect: Dialect) -> Any: + def process(value: Any) -> str | None: + return HalfVector._to_db(value) return process - def literal_processor(self, dialect): + def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) - def process(value): - return string_literal_processor(HalfVector._to_db(value, self.dim)) + def process(value: Any) -> Any: + return string_literal_processor(HalfVector._to_db(value)) # type: ignore return process - def result_processor(self, dialect, coltype): - def process(value): + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: + def process(value: Any) -> list[float] | None: return HalfVector._from_db(value) return process - class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other): + class Comparator(TypeEngine.Comparator[Any]): + def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other): + def max_inner_product(self, other: object, /) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other): + def cosine_distance(self, other: object, /) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other): + def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection -ischema_names['halfvec'] = HALFVEC +ischema_names['halfvec'] = HALFVEC # type: ignore diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 0058679..9bd1478 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -1,51 +1,55 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String +from sqlalchemy import Dialect, Operators +from typing import Any from .. import SparseVector -class SPARSEVEC(UserDefinedType): +class SPARSEVEC(UserDefinedType[Any]): cache_ok = True _string = String() - def __init__(self, dim=None): - super(UserDefinedType, self).__init__() + def __init__(self, dim: int | None = None) -> None: + super().__init__() self.dim = dim - def get_col_spec(self, **kw): + def get_col_spec(self, **kw: Any) -> str: if self.dim is None: return 'SPARSEVEC' return 'SPARSEVEC(%d)' % self.dim - def bind_processor(self, dialect): - def process(value): - return SparseVector._to_db(value, self.dim) + def bind_processor(self, dialect: Dialect) -> Any: + def process(value: Any) -> str | None: + return SparseVector._to_db(value) return process - def literal_processor(self, dialect): + def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) - def process(value): - return string_literal_processor(SparseVector._to_db(value, self.dim)) + def process(value: Any) -> Any: + return string_literal_processor(SparseVector._to_db(value)) # type: ignore return process - def result_processor(self, dialect, coltype): - def process(value): + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: + def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other): + class Comparator(TypeEngine.Comparator[Any]): + def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other): + def max_inner_product(self, other: object, /) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other): + def cosine_distance(self, other: object, /) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other): + def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection -ischema_names['sparsevec'] = SPARSEVEC +ischema_names['sparsevec'] = SPARSEVEC # type: ignore diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 5a1e11f..83acb09 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -1,51 +1,55 @@ from sqlalchemy.dialects.postgresql.base import ischema_names -from sqlalchemy.types import UserDefinedType, Float, String +from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String +from sqlalchemy import Dialect, Operators +from typing import Any from .. import Vector -class VECTOR(UserDefinedType): +class VECTOR(UserDefinedType[Any]): cache_ok = True _string = String() - def __init__(self, dim=None): - super(UserDefinedType, self).__init__() + def __init__(self, dim: int | None = None) -> None: + super().__init__() self.dim = dim - def get_col_spec(self, **kw): + def get_col_spec(self, **kw: Any) -> str: if self.dim is None: return 'VECTOR' return 'VECTOR(%d)' % self.dim - def bind_processor(self, dialect): - def process(value): - return Vector._to_db(value, self.dim) + def bind_processor(self, dialect: Dialect) -> Any: + def process(value: Any) -> str | None: + return Vector._to_db(value) return process - def literal_processor(self, dialect): + def literal_processor(self, dialect: Dialect) -> Any: string_literal_processor = self._string._cached_literal_processor(dialect) - def process(value): - return string_literal_processor(Vector._to_db(value, self.dim)) + def process(value: Any) -> Any: + return string_literal_processor(Vector._to_db(value)) # type: ignore return process - def result_processor(self, dialect, coltype): - def process(value): + def result_processor(self, dialect: Dialect, coltype: Any) -> Any: + def process(value: Any) -> list[float] | None: return Vector._from_db(value) return process - class comparator_factory(UserDefinedType.Comparator): - def l2_distance(self, other): + class Comparator(TypeEngine.Comparator[Any]): + def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) - def max_inner_product(self, other): + def max_inner_product(self, other: object, /) -> Operators: return self.op('<#>', return_type=Float)(other) - def cosine_distance(self, other): + def cosine_distance(self, other: object, /) -> Operators: return self.op('<=>', return_type=Float)(other) - def l1_distance(self, other): + def l1_distance(self, other: object, /) -> Operators: return self.op('<+>', return_type=Float)(other) + comparator_factory = Comparator + # for reflection -ischema_names['vector'] = VECTOR +ischema_names['vector'] = VECTOR # type: ignore diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py deleted file mode 100644 index 8cdb5d6..0000000 --- a/pgvector/utils/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# TODO remove -from .. import Bit, HalfVector, SparseVector, Vector - -__all__ = [ - 'Vector', - 'HalfVector', - 'Bit', - 'SparseVector' -] diff --git a/pgvector/vector.py b/pgvector/vector.py index ebbcafd..d357113 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -1,83 +1,112 @@ -import numpy as np -from struct import pack, unpack_from +from __future__ import annotations +from array import array +import struct +import sys +from typing import TYPE_CHECKING +from ._utils import is_ndarray, ndarray +if TYPE_CHECKING: + import numpy as np -class Vector: - def __init__(self, value): - # asarray still copies if same dtype - if not isinstance(value, np.ndarray) or value.dtype != '>f4': - value = np.asarray(value, dtype='>f4') - - if value.ndim != 1: - raise ValueError('expected ndim to be 1') - - self._value = value - def __repr__(self): +class Vector: + _value: array[float] + + def __init__(self, value: list[float] | ndarray, /) -> None: + if isinstance(value, list): + try: + self._value = array('f', value) + except TypeError: + raise ValueError('expected list[float]') + elif is_ndarray(value): + import numpy as np + + if value.ndim != 1: + raise ValueError('expected ndim to be 1') + + arr = array('f') + arr.frombytes(value.astype(np.float32, order='C', copy=False).data.cast('B')) + self._value = arr + else: + raise ValueError('expected list or ndarray') + + def __repr__(self) -> str: return f'Vector({self.to_list()})' - def __eq__(self, other): - if isinstance(other, self.__class__): - return np.array_equal(self.to_numpy(), other.to_numpy()) - return False + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, self.__class__): + return NotImplemented + return self._value == other._value - def dimensions(self): + def dimensions(self) -> int: return len(self._value) - def to_list(self): + def to_list(self) -> list[float]: return self._value.tolist() - def to_numpy(self): - return self._value + def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]: + import numpy as np + return np.frombuffer(self._value, dtype=np.float32) - def to_text(self): - return '[' + ','.join([str(float(v)) for v in self._value]) + ']' + def to_text(self) -> str: + return f'[{",".join([str(v) for v in self._value])}]' - def to_binary(self): - return pack('>HH', self.dimensions(), 0) + self._value.tobytes() + def to_binary(self) -> bytes: + if sys.byteorder == 'big': + value = self._value + else: + value = array('f', self._value) + value.byteswap() + return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod - def from_text(cls, value): - return cls([float(v) for v in value[1:-1].split(',')]) + def from_text(cls, value: str, /) -> Vector: + return cls(cls._from_text(value)) @classmethod - def from_binary(cls, value): - dim, unused = unpack_from('>HH', value) - return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) + def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: + dim, unused = struct.unpack_from('>HH', value) + data = memoryview(value)[4:] - @classmethod - def _to_db(cls, value, dim=None): - if value is None: - return value + if len(data) != 4 * dim: + raise ValueError('invalid length') - if not isinstance(value, cls): - value = cls(value) + if unused != 0: + raise ValueError('expected unused to be 0') - if dim is not None and value.dimensions() != dim: - raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) + arr = array('f') + arr.frombytes(data) + if sys.byteorder != 'big': + arr.byteswap() - return value.to_text() + vec = cls.__new__(cls) + vec._value = arr + return vec @classmethod - def _to_db_binary(cls, value): + def _from_text(cls, value: str, /) -> list[float]: + return [float(v) for v in value[1:-1].split(',')] + + @classmethod + def _to_db(cls, value: list[float] | ndarray | Vector | None, /) -> str | None: if value is None: return value - if not isinstance(value, cls): + # fewer allocations for lists + if isinstance(value, list): + return f'[{",".join([str(float(v)) for v in value])}]' # type: ignore + + if not isinstance(value, Vector): value = cls(value) - return value.to_binary() + return value.to_text() @classmethod - def _from_db(cls, value): - if value is None or isinstance(value, np.ndarray): + def _from_db(cls, value: str | Vector | None, /) -> list[float] | None: + if value is None: return value - return cls.from_text(value).to_numpy().astype(np.float32) - - @classmethod - def _from_db_binary(cls, value): - if value is None or isinstance(value, np.ndarray): - return value + if isinstance(value, Vector): + return value.to_list() - return cls.from_binary(value).to_numpy().astype(np.float32) + return cls._from_text(value) diff --git a/pyproject.toml b/pyproject.toml index b889f4b..ad130b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,21 +4,84 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.4.0" +version = "0.5.0" description = "pgvector support for Python" readme = "README.md" authors = [ {name = "Andrew Kane", email = "andrew@ankane.org"} ] -license = {text = "MIT"} -requires-python = ">= 3.9" +license = "MIT" +requires-python = ">= 3.10" dependencies = [ - "numpy" ] [project.urls] Homepage = "https://github.com/pgvector/pgvector-python" +[dependency-groups] +dev = [ + "asyncpg", + "Django", + "peewee", + "pg8000", + "psycopg[binary,pool]", + "psycopg2-binary", + "pytest", + "pytest-asyncio", + "SQLAlchemy[asyncio]>=2", + "sqlmodel" +] +numpy = ["numpy"] +scipy = ["scipy"] +typing = [ + "mypy", + "ty", + "numpy", + "scipy" +] + [tool.pytest.ini_options] asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" + +[tool.mypy] +files = [ + "pgvector", + "tests" +] +strict = true +disable_error_code = ["unused-ignore", "untyped-decorator"] +ignore_missing_imports = true +disallow_subclassing_any = false + +[tool.ty.src] +include = [ + "pgvector", + "tests" +] + +[tool.ty.rules] +unused-type-ignore-comment = "ignore" + +[tool.pyright] +include = [ + "pgvector", + "tests" +] +typeCheckingMode = "strict" +reportMissingTypeStubs = "none" +reportPrivateUsage = "none" +reportUnknownArgumentType = "none" +reportUnknownLambdaType = "none" +reportUnknownMemberType = "none" +reportUnknownVariableType = "none" + +[tool.pyrefly] +project-includes = [ + "pgvector/**/*.py", + "tests/**/*.py" +] +preset = "strict" + +[tool.pyrefly.errors] +missing-override-decorator = false diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index a13be06..0000000 --- a/requirements.txt +++ /dev/null @@ -1,12 +0,0 @@ -asyncpg -Django -numpy -peewee -pg8000 -psycopg[binary,pool] -psycopg2-binary -pytest -pytest-asyncio -scipy -SQLAlchemy[asyncio]>=2 -sqlmodel>=0.0.12 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e2878ab --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +try: + import numpy +except ImportError: + numpy = None # type: ignore + +__all__ = [ + 'numpy' +] diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py index 34d66a1..38e5fd0 100644 --- a/tests/test_asyncpg.py +++ b/tests/test_asyncpg.py @@ -1,131 +1,132 @@ import asyncpg -import numpy as np +from asyncpg import Connection from pgvector import HalfVector, SparseVector, Vector from pgvector.asyncpg import register_vector import pytest +from .conftest import numpy as np class TestAsyncpg: - @pytest.mark.asyncio - async def test_vector(self): + async def setup_connection(self) -> Connection: conn = await asyncpg.connect(database='pgvector_python_test') await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + await register_vector(conn) + return conn + + @pytest.mark.asyncio + async def test_vector(self) -> None: + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') - await register_vector(conn) - 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) + embedding2 = [4.5, 5, 6] + embedding3 = np.array([7.5, 8, 9]) if np is not None else [7.5, 8, 9] + embedding4 = None + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), ($4)', embedding, embedding2, embedding3, embedding4) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) - assert res[0]['embedding'].dtype == np.float32 - assert np.array_equal(res[1]['embedding'], embedding2) - assert res[2]['embedding'] is None + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') + assert res[0]['embedding'] == embedding + assert res[1]['embedding'] == Vector(embedding2) + assert res[2]['embedding'] == Vector(embedding3) + assert res[3]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @pytest.mark.asyncio - async def test_halfvec(self): - conn = await asyncpg.connect(database='pgvector_python_test') - await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + async def test_halfvec(self) -> None: + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') - await register_vector(conn) - embedding = HalfVector([1.5, 2, 3]) embedding2 = [4.5, 5, 6] - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) + embedding3 = None + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)', embedding, embedding2, embedding3) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'] == embedding assert res[1]['embedding'] == HalfVector(embedding2) assert res[2]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '[1.5,2,3]' await conn.close() @pytest.mark.asyncio - async def test_bit(self): - conn = await asyncpg.connect(database='pgvector_python_test') - await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + async def test_bit(self) -> None: + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') - await register_vector(conn) - - embedding = asyncpg.BitString('101') - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + # typing issue + # https://github.com/MagicStack/py-pgproto/pull/32 + embedding = asyncpg.BitString('101') # type: ignore + embedding2 = None + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)', embedding, embedding2) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'].as_string() == '101' assert res[0]['embedding'].to_int() == 5 assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '101' await conn.close() @pytest.mark.asyncio - async def test_sparsevec(self): - conn = await asyncpg.connect(database='pgvector_python_test') - await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + async def test_sparsevec(self) -> None: + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') - await register_vector(conn) - embedding = SparseVector([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) + embedding2 = None + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2)', embedding, embedding2) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') assert res[0]['embedding'] == embedding assert res[1]['embedding'] is None # ensures binary format is correct - text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") + text_res = await conn.fetch('SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1') assert text_res[0]['embedding'] == '{1:1.5,2:2,3:3}/3' await conn.close() @pytest.mark.asyncio - async def test_vector_array(self): - conn = await asyncpg.connect(database='pgvector_python_test') - await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') + async def test_vector_array(self) -> None: + conn = await self.setup_connection() await conn.execute('DROP TABLE IF EXISTS asyncpg_items') await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') - await register_vector(conn) - embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] - await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES ($1)', embeddings) + + embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]] + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings2[0], embeddings2[1]) - 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]) + embeddings3 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] if np is not None else [[1.5, 2, 3], [4.5, 5, 6]] + await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings3[0], embeddings3[1]) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert 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]) + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') + assert res[0]['embeddings'] == embeddings + assert res[1]['embeddings'] == [Vector(e) for e in embeddings2] + assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] await conn.close() @pytest.mark.asyncio - async def test_pool(self): - async def init(conn): + async def test_pool(self) -> None: + async def init(conn: Connection) -> None: await register_vector(conn) pool = await asyncpg.create_pool(database='pgvector_python_test', init=init) @@ -136,11 +137,11 @@ async def init(conn): await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') embedding = Vector([1.5, 2, 3]) - embedding2 = np.array([1.5, 2, 3]) - await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) + embedding2 = [1.5, 2, 3] + embedding3 = None + await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3)', embedding, embedding2, embedding3) - res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") - assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) - assert res[0]['embedding'].dtype == np.float32 - assert np.array_equal(res[1]['embedding'], embedding2) + res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id') + assert res[0]['embedding'] == embedding + assert res[1]['embedding'] == Vector(embedding2) assert res[2]['embedding'] is None diff --git a/tests/test_bit.py b/tests/test_bit.py index 5a71642..ed13374 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,63 +1,103 @@ -import numpy as np from pgvector import Bit import pytest +import random +from struct import pack +from .conftest import numpy as np class TestBit: - def test_list(self): + def test_list(self) -> None: assert Bit([True, False, True]).to_list() == [True, False, True] - def test_list_none(self): - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit([True, None, True]).to_text() == '101' + def test_list_none(self) -> None: + with pytest.raises(ValueError) as error: + Bit([True, None, True]) # type: ignore + assert str(error.value) == 'expected list[bool]' - def test_list_int(self): - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit([254, 7, 0]).to_text() == '110' + def test_list_int(self) -> None: + with pytest.raises(ValueError) as error: + Bit([254, 7, 0]) # type: ignore + assert str(error.value) == 'expected list[bool]' - def test_tuple(self): - assert Bit((True, False, True)).to_list() == [True, False, True] + def test_list_list(self) -> None: + with pytest.raises(ValueError) as error: + Bit([[True, False], [True, False]]) # type: ignore + assert str(error.value) == 'expected list[bool]' - def test_str(self): + def test_str(self) -> None: assert Bit('101').to_list() == [True, False, True] - def test_bytes(self): + def test_str_two(self) -> None: + with pytest.raises(ValueError) as error: + Bit('201') + assert str(error.value) == 'expected bit string' + + def test_bytes(self) -> None: assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - def test_ndarray(self): + def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([True, False, True]) assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) - def test_ndarray_unpackbits(self): + def test_ndarray_unpackbits(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' - def test_ndarray_uint8(self): + def test_ndarray_uint8(self) -> None: + if np is None: + pytest.skip('NumPy required') + arr = np.array([254, 7, 0], dtype=np.uint8) - with pytest.warns(UserWarning, match='expected elements to be boolean'): - assert Bit(arr).to_text() == '110' + with pytest.raises(ValueError) as error: + Bit(arr) + assert str(error.value) == 'expected elements to be boolean' - def test_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_ndarray_uint16(self) -> None: + if np is None: + pytest.skip('NumPy required') - def test_ndim_two(self): + arr = np.array([254, 7, 0], dtype=np.uint16) with pytest.raises(ValueError) as error: - Bit([[True, False], [True, False]]) - assert str(error.value) == 'expected ndim to be 1' + Bit(arr) # type: ignore + assert str(error.value) == 'expected elements to be boolean' - def test_ndim_zero(self): + def test_bool(self) -> None: with pytest.raises(ValueError) as error: - Bit(True) - assert str(error.value) == 'expected ndim to be 1' + Bit(True) # type: ignore + assert str(error.value) == 'expected bytes, str, list, or ndarray' - def test_repr(self): + def test_random(self) -> None: + value = ''.join(random.choices(['0', '1'], k=random.randint(1024, 2048))) + assert Bit(value).to_text() == value + + def test_repr(self) -> None: assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' - def test_equality(self): + def test_equality(self) -> None: assert Bit([True, False, True]) == Bit([True, False, True]) assert Bit([True, False, True]) != Bit([True, False, False]) + assert Bit([True, False, True]) != 1 + + def test_from_text(self) -> None: + vec = Bit.from_text('101') + assert vec.to_list() == [True, False, True] + if np is not None: + assert np.array_equal(vec.to_numpy(), [True, False, True]) + assert vec.to_text() == '101' + + def test_from_binary(self) -> None: + data = pack('>iB', 3, 5 << 5) + vec = Bit.from_binary(data) + assert vec.to_list() == [True, False, True] + if np is not None: + assert np.array_equal(vec.to_numpy(), [True, False, True]) + assert vec.to_binary() == data diff --git a/tests/test_django.py b/tests/test_django.py index 7a8a6eb..e5f0c6f 100644 --- a/tests/test_django.py +++ b/tests/test_django.py @@ -9,10 +9,9 @@ from django.db.migrations.loader import MigrationLoader from django.forms import ModelForm from math import sqrt -import numpy as np import os import pgvector.django -from pgvector import HalfVector, SparseVector +from pgvector import SparseVector from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance from unittest import mock @@ -59,6 +58,8 @@ class Item(models.Model): double_embedding = ArrayField(FloatField(), null=True, blank=True) numeric_embedding = ArrayField(DecimalField(max_digits=20, decimal_places=10), null=True, blank=True) + objects: models.Manager # for typing + class Meta: app_label = 'django_app' indexes = [ @@ -124,11 +125,11 @@ class Migration(migrations.Migration): sql_statements = loader.collect_sql([(migration, False)]) with connection.cursor() as cursor: - cursor.execute("DROP TABLE IF EXISTS django_app_item") + cursor.execute('DROP TABLE IF EXISTS django_app_item') cursor.execute('\n'.join(sql_statements)) -def create_items(): +def create_items() -> None: Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])).save() Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])).save() Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])).save() @@ -159,167 +160,179 @@ class Meta: class TestDjango: - def setup_method(self): + def setup_method(self) -> None: Item.objects.all().delete() - def test_vector(self): + def test_vector(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert np.array_equal(item.embedding, [1, 2, 3]) - assert item.embedding.dtype == np.float32 + assert item.embedding == [1, 2, 3] + + def test_vector_l2_distance(self) -> None: + create_items() + distance = L2Distance('embedding', [1, 1, 1]) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.id for v in items] == [1, 3, 2] + assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_vector_l2_distance(self): + def test_vector_l2_distance_list(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_vector_max_inner_product(self): + def test_vector_l2_distance_none(self) -> None: + create_items() + distance = L2Distance('embedding', None) + items = Item.objects.annotate(distance=distance).order_by(distance) + assert [v.distance for v in items] == [None, None, None] + + def test_vector_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self) -> None: create_items() distance = CosineDistance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self) -> None: create_items() distance = L1Distance('embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_halfvec(self): + def test_halfvec(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) - assert item.half_embedding == HalfVector([1, 2, 3]) + assert item.half_embedding == [1, 2, 3] - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self) -> None: create_items() - distance = L2Distance('half_embedding', HalfVector([1, 1, 1])) + distance = L2Distance('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self) -> None: create_items() - distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1])) + distance = MaxInnerProduct('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self) -> None: create_items() - distance = CosineDistance('half_embedding', HalfVector([1, 1, 1])) + distance = CosineDistance('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self) -> None: create_items() - distance = L1Distance('half_embedding', HalfVector([1, 1, 1])) + distance = L1Distance('half_embedding', [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_bit(self): + def test_bit(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self) -> None: create_items() distance = HammingDistance('binary_embedding', '101') items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self) -> None: create_items() distance = JaccardDistance('binary_embedding', '101') items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] # assert [v.distance for v in items] == [0, 1/3, 1] - def test_sparsevec(self): + def test_sparsevec(self) -> None: Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() item = Item.objects.get(pk=1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() distance = L2Distance('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() distance = MaxInnerProduct('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() distance = CosineDistance('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self) -> None: create_items() distance = L1Distance('sparse_embedding', SparseVector([1, 1, 1])) items = Item.objects.annotate(distance=distance).order_by(distance) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_filter(self): + def test_filter(self) -> None: create_items() distance = L2Distance('embedding', [1, 1, 1]) items = Item.objects.alias(distance=distance).filter(distance__lt=1) assert [v.id for v in items] == [1] - def test_vector_avg(self): + def test_vector_avg(self) -> None: avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] assert avg is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] - assert np.array_equal(avg, [2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] - def test_vector_sum(self): + def test_vector_sum(self) -> None: sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] assert sum is None Item(embedding=[1, 2, 3]).save() Item(embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] - assert np.array_equal(sum, [5, 7, 9]) + assert sum == [5, 7, 9] - def test_halfvec_avg(self): + def test_halfvec_avg(self) -> None: avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] assert avg is None Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] - assert avg == HalfVector([2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] - def test_halfvec_sum(self): + def test_halfvec_sum(self) -> None: sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] assert sum is None Item(half_embedding=[1, 2, 3]).save() Item(half_embedding=[4, 5, 6]).save() sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] - assert sum == HalfVector([5, 7, 9]) + assert sum == [5, 7, 9] - def test_serialization(self): + def test_serialization(self) -> None: create_items() items = Item.objects.all() for format in ['json', 'xml']: @@ -329,27 +342,27 @@ def test_serialization(self): for obj in serializers.deserialize(format, data): obj.save() - def test_vector_form(self): + def test_vector_form(self) -> None: form = VectorForm(data={'embedding': '[1, 2, 3]'}) assert form.is_valid() - assert 'value="[1, 2, 3]"' in form.as_div() + assert 'value="[1, 2, 3]"' in str(form.as_div()) - def test_vector_form_instance(self): + def test_vector_form_instance(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item) - assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + assert 'value="[1.0, 2.0, 3.0]"' in str(form.as_div()) - def test_vector_form_save(self): + def test_vector_form_save(self) -> None: Item(id=1, embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item, data={'embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6]) + assert Item.objects.get(pk=1).embedding == [4, 5, 6] - def test_vector_form_save_missing(self): + def test_vector_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = VectorForm(instance=item, data={'embedding': ''}) @@ -357,27 +370,27 @@ def test_vector_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).embedding is None - def test_halfvec_form(self): + def test_halfvec_form(self) -> None: form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) assert form.is_valid() - assert 'value="[1, 2, 3]"' in form.as_div() + assert 'value="[1, 2, 3]"' in str(form.as_div()) - def test_halfvec_form_instance(self): + def test_halfvec_form_instance(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item) - assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() + assert 'value="[1.0, 2.0, 3.0]"' in str(form.as_div()) - def test_halfvec_form_save(self): + def test_halfvec_form_save(self) -> None: Item(id=1, half_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': '[4, 5, 6]'}) assert form.has_changed() assert form.is_valid() assert form.save() - assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) + assert Item.objects.get(pk=1).half_embedding == [4, 5, 6] - def test_halfvec_form_save_missing(self): + def test_halfvec_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = HalfVectorForm(instance=item, data={'half_embedding': ''}) @@ -385,18 +398,18 @@ def test_halfvec_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).half_embedding is None - def test_bit_form(self): + def test_bit_form(self) -> None: form = BitForm(data={'binary_embedding': '101'}) assert form.is_valid() - assert 'value="101"' in form.as_div() + assert 'value="101"' in str(form.as_div()) - def test_bit_form_instance(self): + def test_bit_form_instance(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) form = BitForm(instance=item) - assert 'value="101"' in form.as_div() + assert 'value="101"' in str(form.as_div()) - def test_bit_form_save(self): + def test_bit_form_save(self) -> None: Item(id=1, binary_embedding='101').save() item = Item.objects.get(pk=1) form = BitForm(instance=item, data={'binary_embedding': '010'}) @@ -405,7 +418,7 @@ def test_bit_form_save(self): assert form.save() assert '010' == Item.objects.get(pk=1).binary_embedding - def test_bit_form_save_missing(self): + def test_bit_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = BitForm(instance=item, data={'binary_embedding': ''}) @@ -413,19 +426,18 @@ def test_bit_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).binary_embedding is None - def test_sparsevec_form(self): + def test_sparsevec_form(self) -> None: form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) assert form.is_valid() - assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() + assert 'value="{1:1,2:2,3:3}/3"' in str(form.as_div()) - def test_sparsevec_form_instance(self): + def test_sparsevec_form_instance(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item) - # TODO improve - assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() + assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in str(form.as_div()) - def test_sparsevec_form_save(self): + def test_sparsevec_form_save(self) -> None: Item(id=1, sparse_embedding=[1, 2, 3]).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item, data={'sparse_embedding': '{1:4,2:5,3:6}/3'}) @@ -434,7 +446,7 @@ def test_sparsevec_form_save(self): assert form.save() assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) - def test_sparesevec_form_save_missing(self): + def test_sparesevec_form_save_missing(self) -> None: Item(id=1).save() item = Item.objects.get(pk=1) form = SparseVectorForm(instance=item, data={'sparse_embedding': ''}) @@ -442,22 +454,22 @@ def test_sparesevec_form_save_missing(self): assert form.save() assert Item.objects.get(pk=1).sparse_embedding is None - def test_clean(self): + def test_clean(self) -> None: item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector([1, 2, 3])) item.full_clean() - def test_get_or_create(self): + def test_get_or_create(self) -> None: Item.objects.get_or_create(embedding=[1, 2, 3]) - def test_missing(self): + def test_missing(self) -> None: Item().save() assert Item.objects.first().embedding is None assert Item.objects.first().half_embedding is None assert Item.objects.first().binary_embedding is None assert Item.objects.first().sparse_embedding is None - def test_vector_array(self): - Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]).save() + def test_vector_array(self) -> None: + Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]]).save() with connection.cursor() as cursor: from pgvector.psycopg import register_vector @@ -465,10 +477,9 @@ def test_vector_array(self): # this fails if the driver does not cast arrays item = Item.objects.get(pk=1) - assert np.array_equal(item.embeddings[0], [1, 2, 3]) - assert np.array_equal(item.embeddings[1], [4, 5, 6]) + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] - def test_double_array(self): + def test_double_array(self) -> None: Item(id=1, double_embedding=[1, 1, 1]).save() Item(id=2, double_embedding=[2, 2, 2]).save() Item(id=3, double_embedding=[1, 1, 2]).save() @@ -478,7 +489,7 @@ def test_double_array(self): assert [v.distance for v in items] == [0, 1, sqrt(3)] assert items[1].double_embedding == [1, 1, 2] - def test_numeric_array(self): + def test_numeric_array(self) -> None: Item(id=1, numeric_embedding=[1, 1, 1]).save() Item(id=2, numeric_embedding=[2, 2, 2]).save() Item(id=3, numeric_embedding=[1, 1, 2]).save() @@ -488,7 +499,7 @@ def test_numeric_array(self): assert [v.distance for v in items] == [0, 1, sqrt(3)] assert items[1].numeric_embedding == [1, 1, 2] - def test_half_precision(self): + def test_half_precision(self) -> None: create_items() distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [1, 1, 1]) items = Item.objects.annotate(distance=distance).order_by(distance) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 78b4977..26c5ac7 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,59 +1,75 @@ -import numpy as np from pgvector import HalfVector import pytest from struct import pack +from .conftest import numpy as np class TestHalfVector: - def test_list(self): - assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] + def test_list(self) -> None: + arr = [1.0, 2.0, 3.0] + assert HalfVector(arr).to_list() == arr + assert HalfVector(arr).to_list() is not arr - def test_list_str(self): - with pytest.raises(ValueError, match='could not convert string to float'): - HalfVector([1, 'two', 3]) + def test_list_empty(self) -> None: + assert HalfVector([]).to_list() == [] - def test_tuple(self): - assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3] + def test_list_str(self) -> None: + with pytest.raises(ValueError) as error: + HalfVector([1, 'two', 3]) # type: ignore + assert str(error.value) == 'expected list[float]' - def test_ndarray(self): - arr = np.array([1, 2, 3]) - assert HalfVector(arr).to_list() == [1, 2, 3] - assert HalfVector(arr).to_numpy() is not arr + def test_list_list(self) -> None: + with pytest.raises(ValueError) as error: + HalfVector([[1, 2], [3, 4]]) # type: ignore + assert str(error.value) == 'expected list[float]' + + def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') - def test_ndarray_same_object(self): - arr = np.array([1, 2, 3], dtype='>f2') + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(arr).to_list() == [1, 2, 3] - assert HalfVector(arr).to_numpy() is arr + assert HalfVector(arr).to_numpy() is not arr + assert HalfVector(arr).to_numpy().dtype == np.float16 + # non-contiguous + assert HalfVector(np.flip(arr)).to_list() == [3, 2, 1] + assert HalfVector(np.flip(arr)).to_binary() == HalfVector([3, 2, 1]).to_binary() + # big endian + assert HalfVector(arr.astype('>f2')).to_list() == [1, 2, 3] + assert HalfVector(arr.astype('>f2')).to_binary() == HalfVector([1, 2, 3]).to_binary() - def test_ndim_two(self): with pytest.raises(ValueError) as error: - HalfVector([[1, 2], [3, 4]]) - assert str(error.value) == 'expected ndim to be 1' + HalfVector(np.array(['one', 'two', 'three'])) + assert 'could not convert string to float' in str(error.value) - def test_ndim_zero(self): + def test_int(self) -> None: with pytest.raises(ValueError) as error: - HalfVector(1) - assert str(error.value) == 'expected ndim to be 1' + HalfVector(1) # type: ignore + assert str(error.value) == 'expected list or ndarray' - def test_repr(self): + def test_repr(self) -> None: assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' - def test_equality(self): + def test_equality(self) -> None: assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) + assert HalfVector([1, 2, 3]) != 1 - def test_dimensions(self): + def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 - def test_from_text(self): + def test_from_text(self) -> None: vec = HalfVector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + assert vec.to_text() == '[1.5,2.0,3.0]' - def test_from_binary(self): + def test_from_binary(self) -> None: data = pack('>HH3e', 3, 0, 1.5, 2, 3) vec = HalfVector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data diff --git a/tests/test_peewee.py b/tests/test_peewee.py index 64fc009..59910ad 100644 --- a/tests/test_peewee.py +++ b/tests/test_peewee.py @@ -1,7 +1,6 @@ from math import sqrt -import numpy as np from peewee import Model, PostgresqlDatabase, fn -from pgvector import HalfVector, SparseVector +from pgvector import SparseVector, Vector from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField db = PostgresqlDatabase('pgvector_python_test') @@ -13,203 +12,204 @@ class Meta: class Item(BaseModel): + id: int # for typing embedding = VectorField(dimensions=3, null=True) half_embedding = HalfVectorField(dimensions=3, null=True) binary_embedding = FixedBitField(max_length=3, null=True) sparse_embedding = SparseVectorField(dimensions=3, null=True) + distance: float # for typing - class Meta: + class Meta: # type: ignore table_name = 'peewee_item' Item.add_index('embedding vector_l2_ops', using='hnsw') db.connect() -db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') +db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') # type: ignore db.drop_tables([Item]) db.create_tables([Item]) -def create_items(): +def create_items() -> None: Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])) Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])) Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])) class TestPeewee: - def setup_method(self): + def setup_method(self) -> None: Item.truncate_table() - def test_vector(self): + def test_vector(self) -> None: Item.create(id=1, embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert np.array_equal(item.embedding, [1, 2, 3]) - assert item.embedding.dtype == np.float32 + assert item.embedding == [1, 2, 3] - def test_vector_l2_distance(self): + def test_vector_l2_distance(self) -> None: create_items() distance = Item.embedding.l2_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self) -> None: create_items() distance = Item.embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self) -> None: create_items() distance = Item.embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self) -> None: create_items() distance = Item.embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_halfvec(self): + def test_halfvec(self) -> None: Item.create(id=1, half_embedding=[1, 2, 3]) item = Item.get_by_id(1) - assert item.half_embedding == HalfVector([1, 2, 3]) + assert item.half_embedding == [1, 2, 3] - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self) -> None: create_items() distance = Item.half_embedding.l2_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self) -> None: create_items() distance = Item.half_embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self) -> None: create_items() distance = Item.half_embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self) -> None: create_items() distance = Item.half_embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_bit(self): + def test_bit(self) -> None: Item.create(id=1, binary_embedding='101') item = Item.get_by_id(1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self) -> None: create_items() distance = Item.binary_embedding.hamming_distance('101') items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [0, 1, 2] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self) -> None: create_items() distance = Item.binary_embedding.jaccard_distance('101') items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] # assert [v.distance for v in items] == [0, 1/3, 1] - def test_sparsevec(self): + def test_sparsevec(self) -> None: Item.create(id=1, sparse_embedding=[1, 2, 3]) item = Item.get_by_id(1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() distance = Item.sparse_embedding.l2_distance(SparseVector([1, 1, 1])) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, sqrt(3)] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() distance = Item.sparse_embedding.max_inner_product([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [2, 3, 1] assert [v.distance for v in items] == [-6, -4, -3] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() distance = Item.sparse_embedding.cosine_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 2, 3] assert [v.distance for v in items] == [0, 0, 0.05719095841793653] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self) -> None: create_items() distance = Item.sparse_embedding.l1_distance([1, 1, 1]) items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) assert [v.id for v in items] == [1, 3, 2] assert [v.distance for v in items] == [0, 1, 3] - def test_where(self): + def test_where(self) -> None: create_items() - items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) + items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) # type: ignore assert [v.id for v in items] == [1] - def test_vector_avg(self): + def test_vector_avg(self) -> None: avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() assert avg is None Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() - assert np.array_equal(avg, [2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] - def test_vector_sum(self): + def test_vector_sum(self) -> None: sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() assert sum is None Item.create(embedding=[1, 2, 3]) Item.create(embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() - assert np.array_equal(sum, [5, 7, 9]) + assert sum == [5, 7, 9] - def test_halfvec_avg(self): + def test_halfvec_avg(self) -> None: avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() assert avg is None Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() - assert avg == HalfVector([2.5, 3.5, 4.5]) + assert avg == [2.5, 3.5, 4.5] - def test_halfvec_sum(self): + def test_halfvec_sum(self) -> None: sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() assert sum is None Item.create(half_embedding=[1, 2, 3]) Item.create(half_embedding=[4, 5, 6]) sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() - assert sum == HalfVector([5, 7, 9]) + assert sum == [5, 7, 9] - def test_get_or_create(self): - Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) - Item.get_or_create(embedding=np.array([4, 5, 6])) - Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) + def test_get_or_create(self) -> None: + Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) # type: ignore + Item.get_or_create(embedding=Vector([4, 5, 6])) # type: ignore + Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) # type: ignore - def test_vector_array(self): + def test_vector_array(self) -> None: from playhouse.postgres_ext import PostgresqlExtDatabase, ArrayField ext_db = PostgresqlExtDatabase('pgvector_python_test') class ExtItem(BaseModel): - embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False) + embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False) # type: ignore - class Meta: + class Meta: # type: ignore database = ext_db table_name = 'peewee_ext_item' diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 4d3e474..8c8e02c 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -1,10 +1,10 @@ -import numpy as np -import os +from getpass import getuser from pgvector import HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection +from .conftest import numpy as np -conn = Connection(os.environ["USER"], database='pgvector_python_test') +conn = Connection(getuser(), database='pgvector_python_test') conn.run('CREATE EXTENSION IF NOT EXISTS vector') conn.run('DROP TABLE IF EXISTS pg8000_items') @@ -14,47 +14,38 @@ class TestPg8000: - def setup_method(self): + def setup_method(self) -> None: 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): + def test_vector(self) -> None: embedding = Vector([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = np.array([4.5, 5, 6]) if np is not None else Vector([4.5, 5, 6]) + embedding3 = None + conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2), (:embedding3)', embedding=embedding, embedding2=embedding2, embedding3=embedding3) res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') - assert np.array_equal(res[0][0], embedding.to_numpy()) - assert res[0][0].dtype == np.float32 - assert res[1][0] is None + assert res == [[embedding], [Vector([4.5, 5, 6])], [None]] - def test_halfvec(self): + def test_halfvec(self) -> None: embedding = HalfVector([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == embedding - assert res[1][0] is None + assert res == [[embedding], [None]] - def test_bit(self): + def test_bit(self) -> None: embedding = '101' - conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == '101' - assert res[1][0] is None + assert res == [['101'], [None]] - def test_sparsevec(self): + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) - conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (NULL)', embedding=embedding) + embedding2 = None + conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') - assert res[0][0] == embedding - assert res[1][0] is None + assert res == [[embedding], [None]] diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py index 698b34f..04863e9 100644 --- a/tests/test_psycopg.py +++ b/tests/test_psycopg.py @@ -1,9 +1,11 @@ -import numpy as np from pgvector import Bit, HalfVector, SparseVector, Vector from pgvector.psycopg import register_vector, register_vector_async import psycopg +from psycopg import Connection, AsyncConnection from psycopg_pool import ConnectionPool, AsyncConnectionPool import pytest +from typing import Any +from .conftest import numpy as np conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) @@ -15,179 +17,191 @@ class TestPsycopg: - def setup_method(self): + def setup_method(self) -> None: conn.execute('DELETE FROM psycopg_items') - def test_vector(self): - embedding = np.array([1.5, 2, 3]) - conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) + def test_vector(self) -> None: + embedding = Vector([1.5, 2, 3]) + embedding2 = None + conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (%s)', (embedding, embedding2)) res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() - assert np.array_equal(res[0][0], embedding) - assert res[0][0].dtype == np.float32 + assert res[0][0] == embedding assert res[1][0] is None - def test_vector_binary_format(self): - embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] - assert np.array_equal(res, embedding) + def test_vector_binary_format(self) -> None: + embedding = Vector([1.5, 2, 3]) + res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] + assert res == embedding + + def test_vector_text_format(self) -> None: + embedding = Vector([1.5, 2, 3]) + res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] + assert res == embedding + + def test_vector_binary_format_correct(self) -> None: + embedding = Vector([1.5, 2, 3]) + res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0] + assert res == '[1.5,2,3]' + + def test_vector_numpy_binary_format(self) -> None: + if np is None: + pytest.skip('NumPy required') - def test_vector_text_format(self): embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] - assert np.array_equal(res, embedding) + res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0] + assert res == Vector(embedding) + + def test_vector_numpy_text_format(self) -> None: + if np is None: + pytest.skip('NumPy required') - def test_vector_binary_format_correct(self): embedding = np.array([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector::text', (embedding,)).fetchone()[0] - assert res == '[1.5,2,3]' + res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] + assert res == Vector(embedding) - 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, [3, 2, 1.5]) + def test_vector_numpy_binary_format_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') - 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, [3, 2, 1.5]) + res = next(conn.execute('SELECT %b::vector', (embedding,)))[0] + assert res == Vector([3, 2, 1.5]) - def test_vector_class_binary_format(self): - embedding = Vector([1.5, 2, 3]) - res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] - assert np.array_equal(res, [1.5, 2, 3]) + def test_vector_numpy_text_format_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') - 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, [1.5, 2, 3]) + embedding = np.flipud(np.array([1.5, 2, 3])) + assert not embedding.data.contiguous + res = next(conn.execute('SELECT %t::vector', (embedding,)))[0] + assert res == Vector([3, 2, 1.5]) - def test_halfvec(self): + def test_halfvec(self) -> None: embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + res = next(conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id'))[0] assert res == HalfVector([1.5, 2, 3]) - def test_halfvec_binary_format(self): + def test_halfvec_binary_format(self) -> None: embedding = HalfVector([1.5, 2, 3]) - res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::halfvec', (embedding,), binary=True))[0] assert res == HalfVector([1.5, 2, 3]) - def test_halfvec_text_format(self): + def test_halfvec_text_format(self) -> None: embedding = HalfVector([1.5, 2, 3]) - res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::halfvec', (embedding,)))[0] assert res == HalfVector([1.5, 2, 3]) - def test_bit(self): + def test_bit(self) -> None: embedding = Bit([True, False, True]) conn.execute('INSERT INTO psycopg_items (binary_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + res = next(conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id'))[0] assert res == '101' - def test_bit_binary_format(self): + def test_bit_binary_format(self) -> None: embedding = Bit([False, True, False, True, False, False, False, False, True]) - res = conn.execute('SELECT %b::bit(9)', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::bit(9)', (embedding,), binary=True))[0] assert repr(Bit.from_binary(res)) == 'Bit(010100001)' - def test_bit_text_format(self): + def test_bit_text_format(self) -> None: embedding = Bit([False, True, False, True, False, False, False, False, True]) - res = conn.execute('SELECT %t::bit(9)', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::bit(9)', (embedding,)))[0] assert res == '010100001' assert repr(Bit(res)) == 'Bit(010100001)' - def test_sparsevec(self): + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) - res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] + res = next(conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id'))[0] assert res == SparseVector([1.5, 2, 3]) - def test_sparsevec_binary_format(self): + def test_sparsevec_binary_format(self) -> None: embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) - res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] + res = next(conn.execute('SELECT %b::sparsevec', (embedding,), binary=True))[0] assert res == embedding - def test_sparsevec_text_format(self): + def test_sparsevec_text_format(self) -> None: embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) - res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] + res = next(conn.execute('SELECT %t::sparsevec', (embedding,)))[0] assert res == embedding - def test_text_copy_from(self): - embedding = np.array([1.5, 2, 3]) + def test_text_copy_from(self) -> None: + embedding = [1.5, 2, 3] cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: - copy.write_row([embedding, HalfVector(embedding), '101', SparseVector(embedding)]) + with cur.copy('COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN') as copy: + copy.write_row([Vector(embedding), HalfVector(embedding), '101', SparseVector(embedding)]) - def test_binary_copy_from(self): - embedding = np.array([1.5, 2, 3]) + def test_binary_copy_from(self) -> None: + embedding = [1.5, 2, 3] cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: - copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) + with cur.copy('COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: + copy.write_row([Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) - def test_binary_copy_from_set_types(self): - embedding = np.array([1.5, 2, 3]) + def test_binary_copy_from_set_types(self) -> None: + embedding = [1.5, 2, 3] cur = conn.cursor() - with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) - copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) + copy.write_row([1, Vector(embedding), HalfVector(embedding), Bit('101'), SparseVector(embedding)]) - def test_text_copy_to(self): - embedding = np.array([1.5, 2, 3]) + def test_text_copy_to(self) -> None: + embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding) TO STDOUT') as copy: for row in copy.rows(): - assert row[0] == "[1.5,2,3]" - assert row[1] == "[1.5,2,3]" + assert row[0] == '[1.5,2,3]' + assert row[1] == '[1.5,2,3]' - def test_binary_copy_to(self): - embedding = np.array([1.5, 2, 3]) + def test_binary_copy_to(self) -> None: + embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)') as copy: for row in copy.rows(): - assert np.array_equal(Vector.from_binary(row[0]).to_numpy(), embedding) + assert Vector.from_binary(row[0]) == embedding assert HalfVector.from_binary(row[1]) == half_embedding - def test_binary_copy_to_set_types(self): - embedding = np.array([1.5, 2, 3]) + def test_binary_copy_to_set_types(self) -> None: + embedding = Vector([1.5, 2, 3]) half_embedding = HalfVector([1.5, 2, 3]) conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) cur = conn.cursor() - with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: + with cur.copy('COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)') as copy: copy.set_types(['vector', 'halfvec']) for row in copy.rows(): - assert np.array_equal(row[0], embedding) + assert row[0] == embedding assert row[1] == half_embedding - def test_vector_array(self): - embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + def test_vector_array(self) -> None: + embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] conn.execute('INSERT INTO psycopg_items (embeddings) VALUES (%s)', (embeddings,)) - res = conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id').fetchone() - assert np.array_equal(res[0][0], embeddings[0]) - assert np.array_equal(res[0][1], embeddings[1]) + res = next(conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id')) + assert res[0][0] == embeddings[0] + assert res[0][1] == embeddings[1] - def test_pool(self): - def configure(conn): + def test_pool(self) -> None: + def configure(conn: Connection[Any]) -> None: register_vector(conn) pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure) with pool.connection() as conn: res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() - assert np.array_equal(res[0], [1, 2, 3]) + assert res == (Vector([1, 2, 3]),) pool.close() @pytest.mark.asyncio - async def test_async(self): + async def test_async(self) -> None: conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') @@ -196,19 +210,18 @@ async def test_async(self): await register_vector_async(conn) - embedding = np.array([1.5, 2, 3]) + embedding = Vector([1.5, 2, 3]) await conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) async with conn.cursor() as cur: await cur.execute('SELECT * FROM psycopg_items ORDER BY id') res = await cur.fetchall() - assert np.array_equal(res[0][1], embedding) - assert res[0][1].dtype == np.float32 + assert res[0][1] == embedding assert res[1][1] is None @pytest.mark.asyncio - async def test_async_pool(self): - async def configure(conn): + async def test_async_pool(self) -> None: + async def configure(conn: AsyncConnection[Any]) -> None: await register_vector_async(conn) pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure) @@ -218,6 +231,6 @@ async def configure(conn): async with conn.cursor() as cur: await cur.execute("SELECT '[1,2,3]'::vector") res = await cur.fetchone() - assert np.array_equal(res[0], [1, 2, 3]) + assert res == (Vector([1, 2, 3]),) await pool.close() diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index 7f4932d..badb7c2 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,9 +1,10 @@ -import numpy as np from pgvector import HalfVector, SparseVector, Vector from pgvector.psycopg2 import register_vector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor from psycopg2.pool import ThreadedConnectionPool +import pytest +from .conftest import numpy as np conn = psycopg2.connect(dbname='pgvector_python_test') conn.autocommit = True @@ -17,39 +18,32 @@ class TestPsycopg2: - def setup_method(self): + def setup_method(self) -> None: cur.execute('DELETE FROM psycopg2_items') - def test_vector(self): - embedding = np.array([1.5, 2, 3]) - cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) + def test_vector(self) -> None: + embedding = Vector([1.5, 2, 3]) + embedding2 = None + cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (%s)', (embedding, embedding2)) cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][0], embedding) - assert res[0][0].dtype == np.float32 + assert res[0][0] == embedding assert res[1][0] is None - def test_vector_class(self): - embedding = Vector([1.5, 2, 3]) + def test_vector_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + embedding = np.array([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert np.array_equal(res[0][0], embedding.to_numpy()) - assert res[0][0].dtype == np.float32 + assert res[0][0] == Vector(embedding) 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] == HalfVector([1.5, 2, 3]) - assert res[1][0] is None - - def test_halfvec_class(self): + def test_halfvec(self) -> None: embedding = HalfVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) @@ -58,63 +52,62 @@ def test_halfvec_class(self): assert res[0][0] == embedding assert res[1][0] is None - def test_bit(self): + def test_bit(self) -> None: embedding = '101' cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT binary_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == '101' + assert res[0][0] == embedding assert res[1][0] is None - def test_sparsevec(self): + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') res = cur.fetchall() - assert res[0][0] == SparseVector([1.5, 2, 3]) + assert res[0][0] == embedding assert res[1][0] is None - def test_vector_array(self): - embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] + def test_vector_array(self) -> None: + embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] cur.execute('INSERT INTO psycopg2_items (embeddings) VALUES (%s::vector[])', (embeddings,)) cur.execute('SELECT embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert np.array_equal(res[0][0], embeddings[0]) - assert np.array_equal(res[0][1], embeddings[1]) + assert res == (embeddings,) - def test_halfvec_array(self): + def test_halfvec_array(self) -> None: embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] cur.execute('INSERT INTO psycopg2_items (half_embeddings) VALUES (%s::halfvec[])', (embeddings,)) cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0] == [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] + assert res == ([HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])],) - def test_sparsevec_array(self): + def test_sparsevec_array(self) -> None: embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,)) cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id') res = cur.fetchone() - assert res[0] == [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] + assert res == ([SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])],) - def test_cursor_factory(self): + def test_cursor_factory(self) -> None: for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: conn = psycopg2.connect(dbname='pgvector_python_test') cur = conn.cursor(cursor_factory=cursor_factory) register_vector(cur) conn.close() - def test_cursor_factory_connection(self): + def test_cursor_factory_connection(self) -> None: for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: - conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) + conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) # type: ignore register_vector(conn) conn.close() - def test_pool(self): + def test_pool(self) -> None: pool = ThreadedConnectionPool(1, 1, dbname='pgvector_python_test') conn = pool.getconn() @@ -129,7 +122,7 @@ def test_pool(self): cur = conn.cursor() cur.execute("SELECT '[1,2,3]'::vector") res = cur.fetchone() - assert np.array_equal(res[0], [1, 2, 3]) + assert res == (Vector([1, 2, 3]),) finally: pool.putconn(conn) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index dff03dd..1924f98 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,94 +1,156 @@ -import numpy as np from pgvector import SparseVector import pytest -from scipy.sparse import coo_array from struct import pack +from .conftest import numpy as np + +try: + from scipy import sparse +except ImportError: + sparse = None # type: ignore class TestSparseVector: - def test_list(self): + def test_list(self) -> None: vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] - def test_list_dimensions(self): + def test_list_empty(self) -> None: + assert SparseVector([]).to_list() == [] + + def test_list_str(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + SparseVector([1, 'two', 3]) # type: ignore + + def test_list_dimensions(self) -> None: with pytest.raises(ValueError) as error: - SparseVector([1, 0, 2, 0, 3, 0], 6) + SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore assert str(error.value) == 'extra argument' - def test_ndarray(self): + def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') + vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict(self): + def test_dict(self) -> None: vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict_no_dimensions(self): + def test_dict_empty(self) -> None: + assert SparseVector({}, 0).to_list() == [] + + def test_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: - SparseVector({0: 1, 2: 2, 4: 3}) + SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore assert str(error.value) == 'missing dimensions' - def test_coo_array(self): - arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) + def test_coo_array(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] + assert isinstance(vec.values()[0], float) + + def test_coo_array_dimensions(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') - def test_coo_array_dimensions(self): with pytest.raises(ValueError) as error: - SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) + SparseVector(sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore assert str(error.value) == 'extra argument' - def test_dok_array(self): - arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() + def test_coo_matrix(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + mat = sparse.coo_matrix(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(mat) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + + def test_dok_array(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() + vec = SparseVector(arr) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + + def test_csr_array(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + arr = sparse.csr_array(np.array([[1, 0, 2, 0, 3, 0]])) vec = SparseVector(arr) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_repr(self): + def test_csr_matrix(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + + mat = sparse.csr_matrix(np.array([1, 0, 2, 0, 3, 0])) + vec = SparseVector(mat) + assert vec.to_list() == [1, 0, 2, 0, 3, 0] + assert vec.indices() == [0, 2, 4] + + def test_repr(self) -> None: assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' - def test_equality(self): + def test_equality(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert SparseVector({}, 1) != SparseVector({}, 2) + assert SparseVector([1, 0, 2, 0, 3, 0]) != 1 - def test_dimensions(self): + def test_dimensions(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 - def test_indices(self): + def test_indices(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] - def test_values(self): + def test_values(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] - def test_to_coo(self): + def test_to_coo(self) -> None: + if np is None or sparse is None: + pytest.skip('NumPy and SciPy required') + assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) - def test_zero_vector_text(self): + def test_zero_vector_text(self) -> None: vec = SparseVector({}, 3) assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() - def test_from_text(self): + def test_from_text(self) -> None: vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') assert vec.dimensions() == 6 assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' - def test_from_binary(self): + def test_from_binary(self) -> None: data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) vec = SparseVector.from_binary(data) assert vec.dimensions() == 6 assert vec.indices() == [0, 2, 4] assert vec.values() == [1.5, 2, 3] assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) assert vec.to_binary() == data diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 0d8d1ca..3dd7bac 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -1,90 +1,85 @@ import asyncpg -import numpy as np -import os -from pgvector import HalfVector, SparseVector, Vector +from getpass import getuser +from pgvector import SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest -from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY +from sqlalchemy import create_engine, event, insert, inspect, literal, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine from sqlalchemy.exc import StatementError +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncEngine from sqlalchemy.ext.automap import automap_base -from sqlalchemy.orm import declarative_base, Session +from sqlalchemy.orm import mapped_column, DeclarativeBase, Mapped, Session from sqlalchemy.sql import func - -try: - from sqlalchemy.orm import mapped_column - from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine - sqlalchemy_version = 2 -except ImportError: - mapped_column = Column - sqlalchemy_version = 1 +from typing import Any +from .conftest import numpy as np psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') -@event.listens_for(psycopg2_type_engine, "connect") -def psycopg2_connect(dbapi_connection, connection_record): +@event.listens_for(psycopg2_type_engine, 'connect') +def psycopg2_connect(dbapi_connection: Any, connection_record: Any) -> None: from pgvector.psycopg2 import register_vector register_vector(dbapi_connection) -pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test') +pg8000_engine = create_engine(f'postgresql+pg8000://{getuser()}@localhost/pgvector_python_test') + +psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') +psycopg_type_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test') + + +@event.listens_for(psycopg_type_engine, 'connect') +def psycopg_connect(dbapi_connection: Any, connection_record: Any) -> None: + from pgvector.psycopg import register_vector + register_vector(dbapi_connection) + + +psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') +psycopg_async_type_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test') -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) +@event.listens_for(psycopg_async_type_engine.sync_engine, 'connect') +def psycopg_async_connect(dbapi_connection: Any, connection_record: Any) -> None: + from pgvector.psycopg import register_vector_async + dbapi_connection.run_async(register_vector_async) - 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') - 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) +@event.listens_for(asyncpg_type_engine.sync_engine, 'connect') +def asyncpg_connect(dbapi_connection: Any, connection_record: Any) -> None: + from pgvector.asyncpg import register_vector + dbapi_connection.run_async(register_vector) -engines = [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] +engines = [psycopg_engine, psycopg_type_engine, psycopg2_engine, psycopg2_type_engine, pg8000_engine] +array_engines = [psycopg_type_engine, psycopg2_type_engine] +# TODO support asyncpg_type_engine +async_engines = [psycopg_async_engine, psycopg_async_type_engine, asyncpg_engine] +async_array_engines = [psycopg_async_type_engine, asyncpg_engine] setup_engine = engines[0] with Session(setup_engine) as session: session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) session.commit() -Base = declarative_base() + +class Base(DeclarativeBase): + pass class Item(Base): __tablename__ = 'sqlalchemy_orm_item' - id = mapped_column(Integer, primary_key=True) - embedding = mapped_column(VECTOR(3)) - half_embedding = mapped_column(HALFVEC(3)) - binary_embedding = mapped_column(BIT(3)) - sparse_embedding = mapped_column(SPARSEVEC(3)) - embeddings = mapped_column(ARRAY(VECTOR(3))) - half_embeddings = mapped_column(ARRAY(HALFVEC(3))) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + embedding: Mapped[list[float] | None] = mapped_column(VECTOR(3)) + half_embedding: Mapped[list[float] | None] = mapped_column(HALFVEC(3)) + binary_embedding: Mapped[str | None] = mapped_column(BIT(3)) + sparse_embedding: Mapped[SparseVector | None] = mapped_column(SPARSEVEC(3)) + embeddings: Mapped[list[list[float]] | None] = mapped_column(ARRAY(VECTOR(3), dimensions=1)) + half_embeddings: Mapped[list[list[float]] | None] = mapped_column(ARRAY(HALFVEC(3), dimensions=1)) Base.metadata.drop_all(setup_engine) @@ -103,7 +98,6 @@ class Item(Base): 'sqlalchemy_orm_half_precision_index', func.cast(Item.embedding, HALFVEC(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'halfvec_l2_ops'} ) half_precision_index.create(setup_engine) @@ -112,13 +106,12 @@ class Item(Base): 'sqlalchemy_orm_binary_quantize_index', func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'), postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'bit_hamming_ops'} ) binary_quantize_index.create(setup_engine) -def create_items(): +def create_items() -> None: with Session(setup_engine) as session: session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) @@ -126,7 +119,7 @@ def create_items(): session.commit() -def delete_items(): +def delete_items() -> None: with Session(setup_engine) as session: session.query(Item).delete() session.commit() @@ -134,10 +127,10 @@ def delete_items(): @pytest.mark.parametrize('engine', engines) class TestSqlalchemy: - def setup_method(self): + def setup_method(self) -> None: delete_items() - def test_core(self, engine): + def test_core(self, engine: Engine) -> None: metadata = MetaData() item_table = Table( @@ -172,8 +165,8 @@ def test_core(self, engine): ) hnsw_index.create(engine) - def test_orm(self, engine): - item = Item(embedding=np.array([1.5, 2, 3])) + def test_orm(self, engine: Engine) -> None: + item = Item(embedding=[1.5, 2, 3]) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -183,149 +176,143 @@ def test_orm(self, engine): session.add(item3) session.commit() - stmt = select(Item) + stmt = select(Item).order_by(Item.id) with Session(engine) as session: items = [v[0] for v in session.execute(stmt).all()] - # TODO improve - assert items[0].id % 3 == 1 - assert items[1].id % 3 == 2 - assert items[2].id % 3 == 0 - assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) - assert items[0].embedding.dtype == np.float32 - assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) - assert items[1].embedding.dtype == np.float32 + assert items[0].embedding == [1.5, 2, 3] + assert items[1].embedding == [4, 5, 6] assert items[2].embedding is None - def test_vector(self, engine): + def test_vector(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) - assert np.array_equal(item.embedding, [1, 2, 3]) + item = session.get_one(Item, 1) + assert item.embedding == [1, 2, 3] - def test_vector_l2_distance(self, engine): + def test_vector_l2_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l2_distance_orm(self, engine): + def test_vector_l2_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self, engine): + def test_vector_max_inner_product(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_vector_max_inner_product_orm(self, engine): + def test_vector_max_inner_product_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self, engine): + def test_vector_cosine_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_vector_cosine_distance_orm(self, engine): + def test_vector_cosine_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self, engine): + def test_vector_l1_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_vector_l1_distance_orm(self, engine): + def test_vector_l1_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self, engine): + def test_halfvec(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) - assert item.half_embedding == HalfVector([1, 2, 3]) + item = session.get_one(Item, 1) + assert item.half_embedding == [1, 2, 3] - def test_halfvec_l2_distance(self, engine): + def test_halfvec_l2_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l2_distance_orm(self, engine): + def test_halfvec_l2_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self, engine): + def test_halfvec_max_inner_product(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_max_inner_product_orm(self, engine): + def test_halfvec_max_inner_product_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self, engine): + def test_halfvec_cosine_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_cosine_distance_orm(self, engine): + def test_halfvec_cosine_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self, engine): + def test_halfvec_l1_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_l1_distance_orm(self, engine): + def test_halfvec_l1_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_bit(self, engine): + def test_bit(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self, engine): + def test_bit_hamming_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.binary_embedding.hamming_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_hamming_distance_orm(self, engine): + def test_bit_hamming_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self, engine): + def test_bit_jaccard_distance(self, engine: Engine) -> None: if engine == pg8000_engine: return @@ -334,7 +321,7 @@ def test_bit_jaccard_distance(self, engine): items = session.query(Item).order_by(Item.binary_embedding.jaccard_distance('101')).all() assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance_orm(self, engine): + def test_bit_jaccard_distance_orm(self, engine: Engine) -> None: if engine == pg8000_engine: return @@ -343,276 +330,339 @@ def test_bit_jaccard_distance_orm(self, engine): items = session.scalars(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self, engine): + def test_sparsevec(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self, engine): + def test_sparsevec_l2_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l2_distance_orm(self, engine): + def test_sparsevec_l2_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self, engine): + def test_sparsevec_max_inner_product(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1])).all() assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_max_inner_product_orm(self, engine): + def test_sparsevec_max_inner_product_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self, engine): + def test_sparsevec_cosine_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_cosine_distance_orm(self, engine): + def test_sparsevec_cosine_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self, engine): + def test_sparsevec_l1_distance(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_l1_distance_orm(self, engine): + def test_sparsevec_l1_distance_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) assert [v.id for v in items] == [1, 3, 2] - def test_filter(self, engine): + def test_subquery(self, engine: Engine) -> None: + create_items() + with Session(engine) as session: + subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() + items = session.query(Item).order_by(Item.embedding.l2_distance(subquery)).all() + assert [v.id for v in items] == [1, 3, 2] + + def test_subquery_orm(self, engine: Engine) -> None: + create_items() + with Session(engine) as session: + subquery = select(Item.embedding).filter_by(id=1).scalar_subquery() + items = session.scalars(select(Item).order_by(Item.embedding.l2_distance(subquery))) + assert [v.id for v in items] == [1, 3, 2] + + def test_filter(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1).all() assert [v.id for v in items] == [1] - def test_filter_orm(self, engine): + def test_filter_orm(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.scalars(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) assert [v.id for v in items] == [1] - def test_select(self, engine): + def test_select(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.query(Item.embedding.l2_distance([1, 1, 1])).first() - assert items[0] == 3 + assert items == (3,) - def test_select_orm(self, engine): + def test_select_orm(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() - assert items[0] == 3 + assert items == [3] - def test_avg(self, engine): + def test_avg(self, engine: Engine) -> None: with Session(engine) as session: - res = session.query(avg(Item.embedding)).first()[0] + res = session.query(avg(Item.embedding)).one()[0] assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.query(avg(Item.embedding)).first()[0] - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) + res = session.query(avg(Item.embedding)).one()[0] + assert res == [2.5, 3.5, 4.5] - def test_avg_orm(self, engine): + def test_avg_orm(self, engine: Engine) -> None: with Session(engine) as session: res = session.scalars(select(avg(Item.embedding))).first() assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.scalars(select(avg(Item.embedding))).first() - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) + res = session.scalars(select(avg(Item.embedding))).one() + assert res == [2.5, 3.5, 4.5] - def test_sum(self, engine): + def test_sum(self, engine: Engine) -> None: with Session(engine) as session: - res = session.query(sum(Item.embedding)).first()[0] + res = session.query(sum(Item.embedding)).one()[0] assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.query(sum(Item.embedding)).first()[0] - assert np.array_equal(res, np.array([5, 7, 9])) + res = session.query(sum(Item.embedding)).one()[0] + assert res == [5, 7, 9] - def test_sum_orm(self, engine): + def test_sum_orm(self, engine: Engine) -> None: with Session(engine) as session: res = session.scalars(select(sum(Item.embedding))).first() assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) - res = session.scalars(select(sum(Item.embedding))).first() - assert np.array_equal(res, np.array([5, 7, 9])) + res = session.scalars(select(sum(Item.embedding))).one() + assert res == [5, 7, 9] + + def test_vector_bad_type(self, engine: Engine) -> None: + item = Item(embedding=[1, 'two']) + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() - def test_bad_dimensions(self, engine): + def test_halfvec_bad_type(self, engine: Engine) -> None: + item = Item(half_embedding=[1, 'two']) + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() + + def test_sparsevec_bad_type(self, engine: Engine) -> None: + item = Item(sparse_embedding=[1, 'two']) + with Session(engine) as session: + session.add(item) + with pytest.raises(StatementError, match='could not convert string to float'): + session.commit() + + def test_bad_dimensions(self, engine: Engine) -> None: item = Item(embedding=[1, 2]) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): session.commit() - def test_bad_ndim(self, engine): + def test_bad_ndim(self, engine: Engine) -> None: + if np is None: + pytest.skip('NumPy required') + item = Item(embedding=np.array([[1, 2, 3]])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='expected ndim to be 1'): session.commit() - def test_bad_dtype(self, engine): + def test_bad_dtype(self, engine: Engine) -> None: + if np is None: + pytest.skip('NumPy required') + item = Item(embedding=np.array(['one', 'two', 'three'])) with Session(engine) as session: session.add(item) with pytest.raises(StatementError, match='could not convert string to float'): session.commit() - def test_inspect(self, engine): + def test_inspect(self, engine: Engine) -> None: columns = inspect(engine).get_columns('sqlalchemy_orm_item') assert isinstance(columns[1]['type'], VECTOR) - def test_literal_binds(self, engine): + def test_literal_binds(self, engine: Engine) -> None: sql = select(Item).order_by(Item.embedding.l2_distance([1, 2, 3])).compile(engine, compile_kwargs={'literal_binds': True}) assert "embedding <-> '[1.0,2.0,3.0]'" in str(sql) - def test_insert(self, engine): + def test_insert(self, engine: Engine) -> None: with Session(engine) as session: - session.execute(insert(Item).values(embedding=np.array([1, 2, 3]))) + session.execute(insert(Item).values(embedding=[1, 2, 3])) - def test_insert_bulk(self, engine): + def test_insert_bulk(self, engine: Engine) -> None: with Session(engine) as session: - session.execute(insert(Item), [{'embedding': np.array([1, 2, 3])}]) + session.execute(insert(Item), [{'embedding': [1, 2, 3]}]) # register_vector in psycopg2 tests change this behavior # def test_insert_text(self): # with Session(engine) as session: # session.execute(text('INSERT INTO sqlalchemy_orm_item (embedding) VALUES (:embedding)'), {'embedding': np.array([1, 2, 3])}) - def test_automap(self, engine): + def test_automap(self, engine: Engine) -> None: metadata = MetaData() metadata.reflect(engine, only=['sqlalchemy_orm_item']) AutoBase = automap_base(metadata=metadata) AutoBase.prepare() AutoItem = AutoBase.classes.sqlalchemy_orm_item with Session(engine) as session: - session.execute(insert(AutoItem), [{'embedding': np.array([1, 2, 3])}]) + session.execute(insert(AutoItem), [{'embedding': [1, 2, 3]}]) item = session.query(AutoItem).first() - assert np.array_equal(item.embedding, [1, 2, 3]) + assert item is not None and item.embedding == [1, 2, 3] - def test_half_precision(self, engine): + def test_half_precision(self, engine: Engine) -> None: create_items() with Session(engine) as session: items = session.query(Item).order_by(func.cast(Item.embedding, HALFVEC(3)).l2_distance([1, 1, 1])).all() assert [v.id for v in items] == [1, 3, 2] - def test_binary_quantize(self, engine): + def test_binary_quantize(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[-1, -2, -3])) session.add(Item(id=2, embedding=[1, -2, 3])) session.add(Item(id=3, embedding=[1, 2, 3])) session.commit() - distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3)))) + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast(literal([3, -1, 2], VECTOR), VECTOR(3)))) items = session.query(Item).order_by(distance).all() assert [v.id for v in items] == [2, 3, 1] + def test_binary_quantize_reranking(self, engine: Engine) -> None: + # recreate index (could also vacuum table) + binary_quantize_index.drop(setup_engine) + binary_quantize_index.create(setup_engine) + + with Session(engine) as session: + session.add(Item(id=1, embedding=[-1, -2, -3])) + session.add(Item(id=2, embedding=[1, -2, 3])) + session.add(Item(id=3, embedding=[1, 2, 3])) + session.commit() + + distance = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast(literal([3, -1, 2], VECTOR), VECTOR(3)))) + subquery = session.query(Item).order_by(distance).limit(20).subquery() + items = session.query(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5).all() + assert [v.id for v in items] == [2, 3, 1] + @pytest.mark.parametrize('engine', array_engines) class TestSqlalchemyArray: - def setup_method(self): + def setup_method(self) -> None: delete_items() - def test_vector_array(self, engine): + def test_vector_array(self, engine: Engine) -> None: with Session(engine) as session: - session.add(Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.add(Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]])) session.commit() # this fails if the driver does not cast arrays - 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]) + item = session.get_one(Item, 1) + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] - def test_halfvec_array(self, engine): + def test_halfvec_array(self, engine: Engine) -> None: with Session(engine) as session: - session.add(Item(id=1, half_embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + session.add(Item(id=1, half_embeddings=[[1, 2, 3], [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])] + item = session.get_one(Item, 1) + assert item.half_embeddings == [[1, 2, 3], [4, 5, 6]] @pytest.mark.parametrize('engine', async_engines) class TestSqlalchemyAsync: - def setup_method(self): + def setup_method(self) -> None: delete_items() @pytest.mark.asyncio - async def test_vector(self, engine): + async def test_vector(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - embedding = np.array([1, 2, 3]) + embedding = [1, 2, 3] session.add(Item(id=1, embedding=embedding)) - item = await session.get(Item, 1) - assert np.array_equal(item.embedding, embedding) + item = await session.get_one(Item, 1) + assert item.embedding == embedding await engine.dispose() @pytest.mark.asyncio - async def test_halfvec(self, engine): + async def test_halfvec(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): embedding = [1, 2, 3] session.add(Item(id=1, half_embedding=embedding)) - item = await session.get(Item, 1) - assert item.half_embedding == HalfVector(embedding) + item = await session.get_one(Item, 1) + assert item.half_embedding == embedding await engine.dispose() @pytest.mark.asyncio - async def test_bit(self, engine): + async def test_bit(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' + # typing issue + # https://github.com/MagicStack/py-pgproto/pull/32 + embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' # type: ignore session.add(Item(id=1, binary_embedding=embedding)) - item = await session.get(Item, 1) + item = await session.get_one(Item, 1) assert item.binary_embedding == embedding + if engine == asyncpg_engine: + session.add(Item(id=2, binary_embedding='101')) + item = await session.get_one(Item, 2) + assert item.binary_embedding == embedding + await engine.dispose() @pytest.mark.asyncio - async def test_sparsevec(self, engine): + async def test_sparsevec(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - embedding = [1, 2, 3] + embedding = SparseVector([1, 2, 3]) session.add(Item(id=1, sparse_embedding=embedding)) - item = await session.get(Item, 1) - assert item.sparse_embedding == SparseVector(embedding) + item = await session.get_one(Item, 1) + assert item.sparse_embedding == embedding await engine.dispose() @pytest.mark.asyncio - async def test_avg(self, engine): + async def test_avg(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: @@ -620,30 +670,29 @@ async def test_avg(self, engine): session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = await session.scalars(select(avg(Item.embedding))) - assert np.array_equal(res.first(), [2.5, 3.5, 4.5]) + assert res.first() == [2.5, 3.5, 4.5] await engine.dispose() @pytest.mark.parametrize('engine', async_array_engines) class TestSqlalchemyAsyncArray: - def setup_method(self): + def setup_method(self) -> None: delete_items() @pytest.mark.asyncio - async def test_vector_array(self, engine): + async def test_vector_array(self, engine: AsyncEngine) -> None: async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: async with session.begin(): - session.add(Item(id=1, embeddings=[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]) + session.add(Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]])) + item = await session.get_one(Item, 1) + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] + + if np is not None: + session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])) + item = await session.get_one(Item, 2) + assert item.embeddings == [[1, 2, 3], [4, 5, 6]] await engine.dispose() diff --git a/tests/test_sqlmodel.py b/tests/test_sqlmodel.py index f4994f4..853384e 100644 --- a/tests/test_sqlmodel.py +++ b/tests/test_sqlmodel.py @@ -1,24 +1,30 @@ -import numpy as np -from pgvector import HalfVector, SparseVector +from pgvector import SparseVector from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum +from pydantic import ConfigDict import pytest from sqlalchemy.exc import StatementError -from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text -from typing import Any, Optional +from sqlmodel import Field, Index, Session, SQLModel, col, create_engine, delete, select, text engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') with Session(engine) as session: - session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) + # typing issue + # https://github.com/fastapi/sqlmodel/issues/909 + # https://github.com/fastapi/sqlmodel/pull/1657 + session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) # type: ignore class Item(SQLModel, table=True): - __tablename__ = 'sqlmodel_item' + __tablename__ = 'sqlmodel_item' # type: ignore + model_config = ConfigDict(arbitrary_types_allowed=True) # type: ignore - id: Optional[int] = Field(default=None, primary_key=True) - embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) - half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) - binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) - sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) + id: int | None = Field(default=None, primary_key=True) + # typing issue + # https://github.com/fastapi/sqlmodel/discussions/1228 + # https://github.com/fastapi/sqlmodel/pull/1345 + embedding: list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore + half_embedding: list[float] | None = Field(default=None, sa_type=HALFVEC(3)) # type: ignore + binary_embedding: str | None = Field(default=None, sa_type=BIT(3)) # type: ignore + sparse_embedding: SparseVector | list[float] | None = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore SQLModel.metadata.drop_all(engine) @@ -26,7 +32,7 @@ class Item(SQLModel, table=True): index = Index( 'sqlmodel_index', - Item.embedding, + col(Item.embedding), # col() for typing postgresql_using='hnsw', postgresql_with={'m': 16, 'ef_construction': 64}, postgresql_ops={'embedding': 'vector_l2_ops'} @@ -34,7 +40,7 @@ class Item(SQLModel, table=True): index.create(engine) -def create_items(): +def create_items() -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) @@ -43,12 +49,12 @@ def create_items(): class TestSqlmodel: - def setup_method(self): + def setup_method(self) -> None: with Session(engine) as session: session.exec(delete(Item)) session.commit() - def test_orm(self): + def test_orm(self) -> None: item = Item(embedding=[1.5, 2, 3]) item2 = Item(embedding=[4, 5, 6]) item3 = Item() @@ -65,173 +71,171 @@ def test_orm(self): assert items[0].id == 1 assert items[1].id == 2 assert items[2].id == 3 - assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) - assert items[0].embedding.dtype == np.float32 - assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) - assert items[1].embedding.dtype == np.float32 + assert items[0].embedding == [1.5, 2, 3] + assert items[1].embedding == [4, 5, 6] assert items[2].embedding is None - def test_vector(self): + def test_vector(self) -> None: with Session(engine) as session: session.add(Item(id=1, embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) - assert np.array_equal(item.embedding, np.array([1, 2, 3])) + item = session.get_one(Item, 1) + assert item.embedding == [1, 2, 3] - def test_vector_l2_distance(self): + def test_vector_l2_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.embedding).l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_vector_max_inner_product(self): + def test_vector_max_inner_product(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.embedding).max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_vector_cosine_distance(self): + def test_vector_cosine_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.embedding).cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] - def test_vector_l1_distance(self): + def test_vector_l1_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.embedding).l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_halfvec(self): + def test_halfvec(self) -> None: with Session(engine) as session: session.add(Item(id=1, half_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) - assert item.half_embedding == HalfVector([1, 2, 3]) + item = session.get_one(Item, 1) + assert item.half_embedding == [1, 2, 3] - def test_halfvec_l2_distance(self): + def test_halfvec_l2_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.half_embedding).l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_halfvec_max_inner_product(self): + def test_halfvec_max_inner_product(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.half_embedding).max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_halfvec_cosine_distance(self): + def test_halfvec_cosine_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.half_embedding).cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] - def test_halfvec_l1_distance(self): + def test_halfvec_l1_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.half_embedding).l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_bit(self): + def test_bit(self) -> None: with Session(engine) as session: session.add(Item(id=1, binary_embedding='101')) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.binary_embedding == '101' - def test_bit_hamming_distance(self): + def test_bit_hamming_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) + items = session.exec(select(Item).order_by(col(Item.binary_embedding).hamming_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_bit_jaccard_distance(self): + def test_bit_jaccard_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) + items = session.exec(select(Item).order_by(col(Item.binary_embedding).jaccard_distance('101'))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec(self): + def test_sparsevec(self) -> None: with Session(engine) as session: session.add(Item(id=1, sparse_embedding=[1, 2, 3])) session.commit() - item = session.get(Item, 1) + item = session.get_one(Item, 1) assert item.sparse_embedding == SparseVector([1, 2, 3]) - def test_sparsevec_l2_distance(self): + def test_sparsevec_l2_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).l2_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_sparsevec_max_inner_product(self): + def test_sparsevec_max_inner_product(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).max_inner_product([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [2, 3, 1] - def test_sparsevec_cosine_distance(self): + def test_sparsevec_cosine_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).cosine_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 2, 3] - def test_sparsevec_l1_distance(self): + def test_sparsevec_l1_distance(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) + items = session.exec(select(Item).order_by(col(Item.sparse_embedding).l1_distance([1, 1, 1]))) # type: ignore assert [v.id for v in items] == [1, 3, 2] - def test_filter(self): + def test_filter(self) -> None: create_items() with Session(engine) as session: - items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) + items = session.exec(select(Item).filter(col(Item.embedding).l2_distance([1, 1, 1]) < 1)) # type: ignore assert [v.id for v in items] == [1] - def test_select(self): + def test_select(self) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() + items = session.exec(select(col(Item.embedding).l2_distance([1, 1, 1]))).all() # type: ignore assert items[0] == 3 - def test_vector_avg(self): + def test_vector_avg(self) -> None: with Session(engine) as session: res = session.exec(select(avg(Item.embedding))).first() assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(avg(Item.embedding))).first() - assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) + assert res == [2.5, 3.5, 4.5] - def test_vector_sum(self): + def test_vector_sum(self) -> None: with Session(engine) as session: res = session.exec(select(sum(Item.embedding))).first() assert res is None session.add(Item(embedding=[1, 2, 3])) session.add(Item(embedding=[4, 5, 6])) res = session.exec(select(sum(Item.embedding))).first() - assert np.array_equal(res, np.array([5, 7, 9])) + assert res == [5, 7, 9] - def test_halfvec_avg(self): + def test_halfvec_avg(self) -> None: with Session(engine) as session: res = session.exec(select(avg(Item.half_embedding))).first() assert res is None session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(avg(Item.half_embedding))).first() - assert res == HalfVector([2.5, 3.5, 4.5]) + assert res == [2.5, 3.5, 4.5] - def test_halfvec_sum(self): + def test_halfvec_sum(self) -> None: with Session(engine) as session: res = session.exec(select(sum(Item.half_embedding))).first() assert res is None session.add(Item(half_embedding=[1, 2, 3])) session.add(Item(half_embedding=[4, 5, 6])) res = session.exec(select(sum(Item.half_embedding))).first() - assert res == HalfVector([5, 7, 9]) + assert res == [5, 7, 9] - def test_bad_dimensions(self): + def test_bad_dimensions(self) -> None: item = Item(embedding=[1, 2]) with Session(engine) as session: session.add(item) diff --git a/tests/test_vector.py b/tests/test_vector.py index e5a16fe..a08370e 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,59 +1,75 @@ -import numpy as np from pgvector import Vector import pytest from struct import pack +from .conftest import numpy as np class TestVector: - def test_list(self): - assert Vector([1, 2, 3]).to_list() == [1, 2, 3] + def test_list(self) -> None: + arr = [1.0, 2.0, 3.0] + assert Vector(arr).to_list() == arr + assert Vector(arr).to_list() is not arr - def test_list_str(self): - with pytest.raises(ValueError, match='could not convert string to float'): - Vector([1, 'two', 3]) + def test_list_empty(self) -> None: + assert Vector([]).to_list() == [] - def test_tuple(self): - assert Vector((1, 2, 3)).to_list() == [1, 2, 3] + def test_list_str(self) -> None: + with pytest.raises(ValueError) as error: + Vector([1, 'two', 3]) # type: ignore + assert str(error.value) == 'expected list[float]' - def test_ndarray(self): - arr = np.array([1, 2, 3]) - assert Vector(arr).to_list() == [1, 2, 3] - assert Vector(arr).to_numpy() is not arr + def test_list_list(self) -> None: + with pytest.raises(ValueError) as error: + Vector([[1, 2], [3, 4]]) # type: ignore + assert str(error.value) == 'expected list[float]' + + def test_ndarray(self) -> None: + if np is None: + pytest.skip('NumPy required') - def test_ndarray_same_object(self): - arr = np.array([1, 2, 3], dtype='>f4') + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(arr).to_list() == [1, 2, 3] - assert Vector(arr).to_numpy() is arr + assert Vector(arr).to_numpy() is not arr + assert Vector(arr).to_numpy().dtype == np.float32 + # non-contiguous + assert Vector(np.flip(arr)).to_list() == [3, 2, 1] + assert Vector(np.flip(arr)).to_binary() == Vector([3, 2, 1]).to_binary() + # big endian + assert Vector(arr.astype('>f4')).to_list() == [1, 2, 3] + assert Vector(arr.astype('>f4')).to_binary() == Vector([1, 2, 3]).to_binary() - def test_ndim_two(self): with pytest.raises(ValueError) as error: - Vector([[1, 2], [3, 4]]) - assert str(error.value) == 'expected ndim to be 1' + Vector(np.array(['one', 'two', 'three'])) + assert 'could not convert string to float' in str(error.value) - def test_ndim_zero(self): + def test_int(self) -> None: with pytest.raises(ValueError) as error: - Vector(1) - assert str(error.value) == 'expected ndim to be 1' + Vector(1) # type: ignore + assert str(error.value) == 'expected list or ndarray' - def test_repr(self): + def test_repr(self) -> None: assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' - def test_equality(self): + def test_equality(self) -> None: assert Vector([1, 2, 3]) == Vector([1, 2, 3]) assert Vector([1, 2, 3]) != Vector([1, 2, 4]) + assert Vector([1, 2, 3]) != 1 - def test_dimensions(self): + def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 - def test_from_text(self): + def test_from_text(self) -> None: vec = Vector.from_text('[1.5,2,3]') assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + assert vec.to_text() == '[1.5,2.0,3.0]' - def test_from_binary(self): + def test_from_binary(self) -> None: data = pack('>HH3f', 3, 0, 1.5, 2, 3) vec = Vector.from_binary(data) assert vec.to_list() == [1.5, 2, 3] - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) + if np is not None: + assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) assert vec.to_binary() == data