├── .github └── workflows │ └── build.yml ├── .gitignore ├── CHANGELOG.md ├── LICENSE.txt ├── Makefile ├── README.md ├── examples ├── citus │ ├── example.py │ └── requirements.txt ├── cohere │ ├── example.py │ └── requirements.txt ├── colbert │ ├── approximate.py │ ├── exact.py │ └── requirements.txt ├── colpali │ ├── exact.py │ └── requirements.txt ├── gensim │ ├── example.py │ └── requirements.txt ├── hybrid_search │ ├── cross_encoder.py │ ├── requirements.txt │ └── rrf.py ├── image_search │ ├── example.py │ └── requirements.txt ├── imagehash │ ├── example.py │ └── requirements.txt ├── implicit │ ├── example.py │ └── requirements.txt ├── lightfm │ ├── example.py │ └── requirements.txt ├── loading │ ├── example.py │ └── requirements.txt ├── openai │ ├── example.py │ ├── halfvec.py │ └── requirements.txt ├── rag │ ├── example.py │ └── requirements.txt ├── rdkit │ ├── example.py │ └── requirements.txt ├── sentence_transformers │ ├── example.py │ └── requirements.txt ├── sparse_search │ ├── example.py │ └── requirements.txt └── surprise │ ├── example.py │ └── requirements.txt ├── pgvector ├── __init__.py ├── asyncpg │ ├── __init__.py │ └── register.py ├── bit.py ├── django │ ├── __init__.py │ ├── bit.py │ ├── extensions.py │ ├── functions.py │ ├── halfvec.py │ ├── indexes.py │ ├── sparsevec.py │ └── vector.py ├── halfvec.py ├── peewee │ ├── __init__.py │ ├── bit.py │ ├── halfvec.py │ ├── sparsevec.py │ └── vector.py ├── pg8000 │ ├── __init__.py │ └── register.py ├── psycopg │ ├── __init__.py │ ├── bit.py │ ├── halfvec.py │ ├── register.py │ ├── sparsevec.py │ └── vector.py ├── psycopg2 │ ├── __init__.py │ ├── halfvec.py │ ├── register.py │ ├── sparsevec.py │ └── vector.py ├── sparsevec.py ├── sqlalchemy │ ├── __init__.py │ ├── bit.py │ ├── functions.py │ ├── halfvec.py │ ├── sparsevec.py │ └── vector.py ├── utils │ └── __init__.py └── vector.py ├── pyproject.toml ├── requirements.txt └── tests ├── __init__.py ├── test_asyncpg.py ├── test_bit.py ├── test_django.py ├── test_half_vector.py ├── test_peewee.py ├── test_pg8000.py ├── test_psycopg.py ├── test_psycopg2.py ├── test_sparse_vector.py ├── test_sqlalchemy.py ├── test_sqlmodel.py └── test_vector.py /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | on: [push, pull_request] 3 | jobs: 4 | build: 5 | runs-on: ubuntu-latest 6 | strategy: 7 | fail-fast: false 8 | matrix: 9 | python: [3.13, 3.9] 10 | steps: 11 | - uses: actions/checkout@v4 12 | - uses: actions/setup-python@v5 13 | with: 14 | python-version: ${{ matrix.python }} 15 | - run: pip install -r requirements.txt 16 | - uses: ankane/setup-postgres@v1 17 | with: 18 | database: pgvector_python_test 19 | dev-files: true 20 | - run: | 21 | cd /tmp 22 | git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git 23 | cd pgvector 24 | make 25 | sudo make install 26 | - run: pytest 27 | 28 | - run: pip install "SQLAlchemy<2" -U 29 | - run: pytest tests/test_sqlalchemy.py 30 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | dist/ 3 | *.egg-info/ 4 | venv/ 5 | .cache/ 6 | *.pyc 7 | __pycache__ 8 | .pytest_cache/ 9 | examples/rag/README.md 10 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## 0.4.1 (2025-04-26) 2 | 3 | - Fixed `SparseVector` constructor for SciPy sparse matrices 4 | 5 | ## 0.4.0 (2025-03-15) 6 | 7 | - Added top-level `pgvector` package 8 | - Added support for pg8000 9 | - Added support for `bytes` to `Bit` constructor 10 | - Changed `globally` option to default to `False` for Psycopg 2 11 | - Changed `arrays` option to default to `True` for Psycopg 2 12 | - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes 13 | - Fixed `indices` and `values` methods of `SparseVector` returning tuple instead of list in some cases 14 | - Dropped support for Python < 3.9 15 | 16 | ## 0.3.6 (2024-10-26) 17 | 18 | - Added `arrays` option for Psycopg 2 19 | 20 | ## 0.3.5 (2024-10-05) 21 | 22 | - Added `avg` function with type casting to SQLAlchemy 23 | - Added `globally` option for Psycopg 2 24 | 25 | ## 0.3.4 (2024-09-26) 26 | 27 | - Added `schema` option for asyncpg 28 | 29 | ## 0.3.3 (2024-09-09) 30 | 31 | - Improved support for cursor factories with Psycopg 2 32 | 33 | ## 0.3.2 (2024-07-17) 34 | 35 | - Fixed error with asyncpg and pgvector < 0.7 36 | 37 | ## 0.3.1 (2024-07-10) 38 | 39 | - Fixed error parsing zero sparse vectors 40 | - Fixed error with Psycopg 2 and pgvector < 0.7 41 | - Fixed error message when `vector` type not found with Psycopg 3 42 | 43 | ## 0.3.0 (2024-06-25) 44 | 45 | - Added support for `halfvec`, `bit`, and `sparsevec` types to Django 46 | - Added support for `halfvec`, `bit`, and `sparsevec` types to SQLAlchemy and SQLModel 47 | - Added support for `halfvec` and `sparsevec` types to Psycopg 3 48 | - Added support for `halfvec` and `sparsevec` types to Psycopg 2 49 | - Added support for `halfvec` and `sparsevec` types to asyncpg 50 | - Added support for `halfvec`, `bit`, and `sparsevec` types to Peewee 51 | - Added `L1Distance`, `HammingDistance`, and `JaccardDistance` for Django 52 | - Added `l1_distance`, `hamming_distance`, and `jaccard_distance` for SQLAlchemy and SQLModel 53 | - Added `l1_distance`, `hamming_distance`, and `jaccard_distance` for Peewee 54 | 55 | ## 0.2.5 (2024-02-07) 56 | 57 | - Added literal binds support for SQLAlchemy 58 | 59 | ## 0.2.4 (2023-11-24) 60 | 61 | - Improved reflection with SQLAlchemy 62 | 63 | ## 0.2.3 (2023-09-25) 64 | 65 | - Fixed null values with Django 66 | - Fixed `full_clean` with Django 67 | 68 | ## 0.2.2 (2023-09-08) 69 | 70 | - Added support for Peewee 71 | - Added `HnswIndex` for Django 72 | 73 | ## 0.2.1 (2023-07-31) 74 | 75 | - Fixed form issues with Django 76 | 77 | ## 0.2.0 (2023-07-23) 78 | 79 | - Fixed form validation with Django 80 | - Dropped support for Python < 3.8 81 | 82 | ## 0.1.8 (2023-05-20) 83 | 84 | - Fixed serialization with Django 85 | 86 | ## 0.1.7 (2023-05-11) 87 | 88 | - Added `register_vector_async` for Psycopg 3 89 | - Fixed `set_types` for Psycopg 3 90 | 91 | ## 0.1.6 (2022-05-22) 92 | 93 | - Fixed return type for distance operators with SQLAlchemy 94 | 95 | ## 0.1.5 (2022-01-14) 96 | 97 | - Fixed `operator does not exist` error with Django 98 | - Fixed warning with SQLAlchemy 1.4.28+ 99 | 100 | ## 0.1.4 (2021-10-12) 101 | 102 | - Updated Psycopg 3 integration for 3.0 release (no longer experimental) 103 | 104 | ## 0.1.3 (2021-06-22) 105 | 106 | - Added support for asyncpg 107 | - Added experimental support for Psycopg 3 108 | 109 | ## 0.1.2 (2021-06-13) 110 | 111 | - Added Django support 112 | 113 | ## 0.1.1 (2021-06-12) 114 | 115 | - Added `l2_distance`, `max_inner_product`, and `cosine_distance` for SQLAlchemy 116 | 117 | ## 0.1.0 (2021-06-11) 118 | 119 | - First release 120 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2021-2025 Andrew Kane 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | THE SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: lint build publish clean 2 | 3 | lint: 4 | pycodestyle . --ignore=E501 5 | 6 | build: 7 | python3 -m build 8 | 9 | publish: clean build 10 | twine upload dist/* 11 | 12 | clean: 13 | rm -rf .pytest_cache dist pgvector.egg-info 14 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pgvector-python 2 | 3 | [pgvector](https://github.com/pgvector/pgvector) support for Python 4 | 5 | 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) 6 | 7 | [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions) 8 | 9 | ## Installation 10 | 11 | Run: 12 | 13 | ```sh 14 | pip install pgvector 15 | ``` 16 | 17 | And follow the instructions for your database library: 18 | 19 | - [Django](#django) 20 | - [SQLAlchemy](#sqlalchemy) 21 | - [SQLModel](#sqlmodel) 22 | - [Psycopg 3](#psycopg-3) 23 | - [Psycopg 2](#psycopg-2) 24 | - [asyncpg](#asyncpg) 25 | - [pg8000](#pg8000) 26 | - [Peewee](#peewee) 27 | 28 | Or check out some examples: 29 | 30 | - [Retrieval-augmented generation](https://github.com/pgvector/pgvector-python/blob/master/examples/rag/example.py) with Ollama 31 | - [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai/example.py) with OpenAI 32 | - [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere/example.py) with Cohere 33 | - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_transformers/example.py) with SentenceTransformers 34 | - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/rrf.py) with SentenceTransformers (Reciprocal Rank Fusion) 35 | - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search/cross_encoder.py) with SentenceTransformers (cross-encoder) 36 | - [Sparse search](https://github.com/pgvector/pgvector-python/blob/master/examples/sparse_search/example.py) with Transformers 37 | - [Late interaction search](https://github.com/pgvector/pgvector-python/blob/master/examples/colbert/exact.py) with ColBERT 38 | - [Visual document retrieval](https://github.com/pgvector/pgvector-python/blob/master/examples/colpali/exact.py) with ColPali 39 | - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/image_search/example.py) with PyTorch 40 | - [Image search](https://github.com/pgvector/pgvector-python/blob/master/examples/imagehash/example.py) with perceptual hashing 41 | - [Morgan fingerprints](https://github.com/pgvector/pgvector-python/blob/master/examples/rdkit/example.py) with RDKit 42 | - [Topic modeling](https://github.com/pgvector/pgvector-python/blob/master/examples/gensim/example.py) with Gensim 43 | - [Implicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/implicit/example.py) with Implicit 44 | - [Explicit feedback recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/surprise/example.py) with Surprise 45 | - [Recommendations](https://github.com/pgvector/pgvector-python/blob/master/examples/lightfm/example.py) with LightFM 46 | - [Horizontal scaling](https://github.com/pgvector/pgvector-python/blob/master/examples/citus/example.py) with Citus 47 | - [Bulk loading](https://github.com/pgvector/pgvector-python/blob/master/examples/loading/example.py) with `COPY` 48 | 49 | ## Django 50 | 51 | Create a migration to enable the extension 52 | 53 | ```python 54 | from pgvector.django import VectorExtension 55 | 56 | class Migration(migrations.Migration): 57 | operations = [ 58 | VectorExtension() 59 | ] 60 | ``` 61 | 62 | Add a vector field to your model 63 | 64 | ```python 65 | from pgvector.django import VectorField 66 | 67 | class Item(models.Model): 68 | embedding = VectorField(dimensions=3) 69 | ``` 70 | 71 | Also supports `HalfVectorField`, `BitField`, and `SparseVectorField` 72 | 73 | Insert a vector 74 | 75 | ```python 76 | item = Item(embedding=[1, 2, 3]) 77 | item.save() 78 | ``` 79 | 80 | Get the nearest neighbors to a vector 81 | 82 | ```python 83 | from pgvector.django import L2Distance 84 | 85 | Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5] 86 | ``` 87 | 88 | Also supports `MaxInnerProduct`, `CosineDistance`, `L1Distance`, `HammingDistance`, and `JaccardDistance` 89 | 90 | Get the distance 91 | 92 | ```python 93 | Item.objects.annotate(distance=L2Distance('embedding', [3, 1, 2])) 94 | ``` 95 | 96 | Get items within a certain distance 97 | 98 | ```python 99 | Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance__lt=5) 100 | ``` 101 | 102 | Average vectors 103 | 104 | ```python 105 | from django.db.models import Avg 106 | 107 | Item.objects.aggregate(Avg('embedding')) 108 | ``` 109 | 110 | Also supports `Sum` 111 | 112 | Add an approximate index 113 | 114 | ```python 115 | from pgvector.django import HnswIndex, IvfflatIndex 116 | 117 | class Item(models.Model): 118 | class Meta: 119 | indexes = [ 120 | HnswIndex( 121 | name='my_index', 122 | fields=['embedding'], 123 | m=16, 124 | ef_construction=64, 125 | opclasses=['vector_l2_ops'] 126 | ), 127 | # or 128 | IvfflatIndex( 129 | name='my_index', 130 | fields=['embedding'], 131 | lists=100, 132 | opclasses=['vector_l2_ops'] 133 | ) 134 | ] 135 | ``` 136 | 137 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 138 | 139 | #### Half-Precision Indexing 140 | 141 | Index vectors at half-precision 142 | 143 | ```python 144 | from django.contrib.postgres.indexes import OpClass 145 | from django.db.models.functions import Cast 146 | from pgvector.django import HnswIndex, HalfVectorField 147 | 148 | class Item(models.Model): 149 | class Meta: 150 | indexes = [ 151 | HnswIndex( 152 | OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), 153 | name='my_index', 154 | m=16, 155 | ef_construction=64 156 | ) 157 | ] 158 | ``` 159 | 160 | Note: Add `'django.contrib.postgres'` to `INSTALLED_APPS` to use `OpClass` 161 | 162 | Get the nearest neighbors 163 | 164 | ```python 165 | distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2]) 166 | Item.objects.order_by(distance)[:5] 167 | ``` 168 | 169 | ## SQLAlchemy 170 | 171 | Enable the extension 172 | 173 | ```python 174 | session.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) 175 | ``` 176 | 177 | Add a vector column 178 | 179 | ```python 180 | from pgvector.sqlalchemy import Vector 181 | 182 | class Item(Base): 183 | embedding = mapped_column(Vector(3)) 184 | ``` 185 | 186 | Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` 187 | 188 | Insert a vector 189 | 190 | ```python 191 | item = Item(embedding=[1, 2, 3]) 192 | session.add(item) 193 | session.commit() 194 | ``` 195 | 196 | Get the nearest neighbors to a vector 197 | 198 | ```python 199 | session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) 200 | ``` 201 | 202 | Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` 203 | 204 | Get the distance 205 | 206 | ```python 207 | session.scalars(select(Item.embedding.l2_distance([3, 1, 2]))) 208 | ``` 209 | 210 | Get items within a certain distance 211 | 212 | ```python 213 | session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) 214 | ``` 215 | 216 | Average vectors 217 | 218 | ```python 219 | from pgvector.sqlalchemy import avg 220 | 221 | session.scalars(select(avg(Item.embedding))).first() 222 | ``` 223 | 224 | Also supports `sum` 225 | 226 | Add an approximate index 227 | 228 | ```python 229 | index = Index( 230 | 'my_index', 231 | Item.embedding, 232 | postgresql_using='hnsw', 233 | postgresql_with={'m': 16, 'ef_construction': 64}, 234 | postgresql_ops={'embedding': 'vector_l2_ops'} 235 | ) 236 | # or 237 | index = Index( 238 | 'my_index', 239 | Item.embedding, 240 | postgresql_using='ivfflat', 241 | postgresql_with={'lists': 100}, 242 | postgresql_ops={'embedding': 'vector_l2_ops'} 243 | ) 244 | 245 | index.create(engine) 246 | ``` 247 | 248 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 249 | 250 | #### Half-Precision Indexing 251 | 252 | Index vectors at half-precision 253 | 254 | ```python 255 | from pgvector.sqlalchemy import HALFVEC 256 | from sqlalchemy.sql import func 257 | 258 | index = Index( 259 | 'my_index', 260 | func.cast(Item.embedding, HALFVEC(3)).label('embedding'), 261 | postgresql_using='hnsw', 262 | postgresql_with={'m': 16, 'ef_construction': 64}, 263 | postgresql_ops={'embedding': 'halfvec_l2_ops'} 264 | ) 265 | ``` 266 | 267 | Get the nearest neighbors 268 | 269 | ```python 270 | order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2]) 271 | session.scalars(select(Item).order_by(order).limit(5)) 272 | ``` 273 | 274 | #### Arrays 275 | 276 | Add an array column 277 | 278 | ```python 279 | from pgvector.sqlalchemy import Vector 280 | from sqlalchemy import ARRAY 281 | 282 | class Item(Base): 283 | embeddings = mapped_column(ARRAY(Vector(3))) 284 | ``` 285 | 286 | And register the types with the underlying driver 287 | 288 | For Psycopg 3, use 289 | 290 | ```python 291 | from pgvector.psycopg import register_vector 292 | from sqlalchemy import event 293 | 294 | @event.listens_for(engine, "connect") 295 | def connect(dbapi_connection, connection_record): 296 | register_vector(dbapi_connection) 297 | ``` 298 | 299 | For [async connections](https://docs.sqlalchemy.org/en/20/orm/extensions/asyncio.html) with Psycopg 3, use 300 | 301 | ```python 302 | from pgvector.psycopg import register_vector_async 303 | from sqlalchemy import event 304 | 305 | @event.listens_for(engine.sync_engine, "connect") 306 | def connect(dbapi_connection, connection_record): 307 | dbapi_connection.run_async(register_vector_async) 308 | ``` 309 | 310 | For Psycopg 2, use 311 | 312 | ```python 313 | from pgvector.psycopg2 import register_vector 314 | from sqlalchemy import event 315 | 316 | @event.listens_for(engine, "connect") 317 | def connect(dbapi_connection, connection_record): 318 | register_vector(dbapi_connection, arrays=True) 319 | ``` 320 | 321 | ## SQLModel 322 | 323 | Enable the extension 324 | 325 | ```python 326 | session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) 327 | ``` 328 | 329 | Add a vector column 330 | 331 | ```python 332 | from pgvector.sqlalchemy import Vector 333 | 334 | class Item(SQLModel, table=True): 335 | embedding: Any = Field(sa_type=Vector(3)) 336 | ``` 337 | 338 | Also supports `HALFVEC`, `BIT`, and `SPARSEVEC` 339 | 340 | Insert a vector 341 | 342 | ```python 343 | item = Item(embedding=[1, 2, 3]) 344 | session.add(item) 345 | session.commit() 346 | ``` 347 | 348 | Get the nearest neighbors to a vector 349 | 350 | ```python 351 | session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5)) 352 | ``` 353 | 354 | Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` 355 | 356 | Get the distance 357 | 358 | ```python 359 | session.exec(select(Item.embedding.l2_distance([3, 1, 2]))) 360 | ``` 361 | 362 | Get items within a certain distance 363 | 364 | ```python 365 | session.exec(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5)) 366 | ``` 367 | 368 | Average vectors 369 | 370 | ```python 371 | from pgvector.sqlalchemy import avg 372 | 373 | session.exec(select(avg(Item.embedding))).first() 374 | ``` 375 | 376 | Also supports `sum` 377 | 378 | Add an approximate index 379 | 380 | ```python 381 | from sqlmodel import Index 382 | 383 | index = Index( 384 | 'my_index', 385 | Item.embedding, 386 | postgresql_using='hnsw', 387 | postgresql_with={'m': 16, 'ef_construction': 64}, 388 | postgresql_ops={'embedding': 'vector_l2_ops'} 389 | ) 390 | # or 391 | index = Index( 392 | 'my_index', 393 | Item.embedding, 394 | postgresql_using='ivfflat', 395 | postgresql_with={'lists': 100}, 396 | postgresql_ops={'embedding': 'vector_l2_ops'} 397 | ) 398 | 399 | index.create(engine) 400 | ``` 401 | 402 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 403 | 404 | ## Psycopg 3 405 | 406 | Enable the extension 407 | 408 | ```python 409 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 410 | ``` 411 | 412 | Register the vector type with your connection 413 | 414 | ```python 415 | from pgvector.psycopg import register_vector 416 | 417 | register_vector(conn) 418 | ``` 419 | 420 | For [connection pools](https://www.psycopg.org/psycopg3/docs/advanced/pool.html), use 421 | 422 | ```python 423 | def configure(conn): 424 | register_vector(conn) 425 | 426 | pool = ConnectionPool(..., configure=configure) 427 | ``` 428 | 429 | For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use 430 | 431 | ```python 432 | from pgvector.psycopg import register_vector_async 433 | 434 | await register_vector_async(conn) 435 | ``` 436 | 437 | Create a table 438 | 439 | ```python 440 | conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') 441 | ``` 442 | 443 | Insert a vector 444 | 445 | ```python 446 | embedding = np.array([1, 2, 3]) 447 | conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) 448 | ``` 449 | 450 | Get the nearest neighbors to a vector 451 | 452 | ```python 453 | conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall() 454 | ``` 455 | 456 | Add an approximate index 457 | 458 | ```python 459 | conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') 460 | # or 461 | conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') 462 | ``` 463 | 464 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 465 | 466 | ## Psycopg 2 467 | 468 | Enable the extension 469 | 470 | ```python 471 | cur = conn.cursor() 472 | cur.execute('CREATE EXTENSION IF NOT EXISTS vector') 473 | ``` 474 | 475 | Register the vector type with your connection or cursor 476 | 477 | ```python 478 | from pgvector.psycopg2 import register_vector 479 | 480 | register_vector(conn) 481 | ``` 482 | 483 | Create a table 484 | 485 | ```python 486 | cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') 487 | ``` 488 | 489 | Insert a vector 490 | 491 | ```python 492 | embedding = np.array([1, 2, 3]) 493 | cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,)) 494 | ``` 495 | 496 | Get the nearest neighbors to a vector 497 | 498 | ```python 499 | cur.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)) 500 | cur.fetchall() 501 | ``` 502 | 503 | Add an approximate index 504 | 505 | ```python 506 | cur.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') 507 | # or 508 | cur.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') 509 | ``` 510 | 511 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 512 | 513 | ## asyncpg 514 | 515 | Enable the extension 516 | 517 | ```python 518 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 519 | ``` 520 | 521 | Register the vector type with your connection 522 | 523 | ```python 524 | from pgvector.asyncpg import register_vector 525 | 526 | await register_vector(conn) 527 | ``` 528 | 529 | or your pool 530 | 531 | ```python 532 | async def init(conn): 533 | await register_vector(conn) 534 | 535 | pool = await asyncpg.create_pool(..., init=init) 536 | ``` 537 | 538 | Create a table 539 | 540 | ```python 541 | await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') 542 | ``` 543 | 544 | Insert a vector 545 | 546 | ```python 547 | embedding = np.array([1, 2, 3]) 548 | await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding) 549 | ``` 550 | 551 | Get the nearest neighbors to a vector 552 | 553 | ```python 554 | await conn.fetch('SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5', embedding) 555 | ``` 556 | 557 | Add an approximate index 558 | 559 | ```python 560 | await conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') 561 | # or 562 | await conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') 563 | ``` 564 | 565 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 566 | 567 | ## pg8000 568 | 569 | Enable the extension 570 | 571 | ```python 572 | conn.run('CREATE EXTENSION IF NOT EXISTS vector') 573 | ``` 574 | 575 | Register the vector type with your connection 576 | 577 | ```python 578 | from pgvector.pg8000 import register_vector 579 | 580 | register_vector(conn) 581 | ``` 582 | 583 | Create a table 584 | 585 | ```python 586 | conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))') 587 | ``` 588 | 589 | Insert a vector 590 | 591 | ```python 592 | embedding = np.array([1, 2, 3]) 593 | conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding) 594 | ``` 595 | 596 | Get the nearest neighbors to a vector 597 | 598 | ```python 599 | conn.run('SELECT * FROM items ORDER BY embedding <-> :embedding LIMIT 5', embedding=embedding) 600 | ``` 601 | 602 | Add an approximate index 603 | 604 | ```python 605 | conn.run('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') 606 | # or 607 | conn.run('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)') 608 | ``` 609 | 610 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 611 | 612 | ## Peewee 613 | 614 | Add a vector column 615 | 616 | ```python 617 | from pgvector.peewee import VectorField 618 | 619 | class Item(BaseModel): 620 | embedding = VectorField(dimensions=3) 621 | ``` 622 | 623 | Also supports `HalfVectorField`, `FixedBitField`, and `SparseVectorField` 624 | 625 | Insert a vector 626 | 627 | ```python 628 | item = Item.create(embedding=[1, 2, 3]) 629 | ``` 630 | 631 | Get the nearest neighbors to a vector 632 | 633 | ```python 634 | Item.select().order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5) 635 | ``` 636 | 637 | Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance` 638 | 639 | Get the distance 640 | 641 | ```python 642 | Item.select(Item.embedding.l2_distance([3, 1, 2]).alias('distance')) 643 | ``` 644 | 645 | Get items within a certain distance 646 | 647 | ```python 648 | Item.select().where(Item.embedding.l2_distance([3, 1, 2]) < 5) 649 | ``` 650 | 651 | Average vectors 652 | 653 | ```python 654 | from peewee import fn 655 | 656 | Item.select(fn.avg(Item.embedding).coerce(True)).scalar() 657 | ``` 658 | 659 | Also supports `sum` 660 | 661 | Add an approximate index 662 | 663 | ```python 664 | Item.add_index('embedding vector_l2_ops', using='hnsw') 665 | ``` 666 | 667 | Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance 668 | 669 | ## Reference 670 | 671 | ### Half Vectors 672 | 673 | Create a half vector from a list 674 | 675 | ```python 676 | vec = HalfVector([1, 2, 3]) 677 | ``` 678 | 679 | Or a NumPy array 680 | 681 | ```python 682 | vec = HalfVector(np.array([1, 2, 3])) 683 | ``` 684 | 685 | Get a list 686 | 687 | ```python 688 | lst = vec.to_list() 689 | ``` 690 | 691 | Get a NumPy array 692 | 693 | ```python 694 | arr = vec.to_numpy() 695 | ``` 696 | 697 | ### Sparse Vectors 698 | 699 | Create a sparse vector from a list 700 | 701 | ```python 702 | vec = SparseVector([1, 0, 2, 0, 3, 0]) 703 | ``` 704 | 705 | Or a NumPy array 706 | 707 | ```python 708 | vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) 709 | ``` 710 | 711 | Or a SciPy sparse array 712 | 713 | ```python 714 | arr = coo_array(([1, 2, 3], ([0, 2, 4],)), shape=(6,)) 715 | vec = SparseVector(arr) 716 | ``` 717 | 718 | Or a dictionary of non-zero elements 719 | 720 | ```python 721 | vec = SparseVector({0: 1, 2: 2, 4: 3}, 6) 722 | ``` 723 | 724 | Note: Indices start at 0 725 | 726 | Get the number of dimensions 727 | 728 | ```python 729 | dim = vec.dimensions() 730 | ``` 731 | 732 | Get the indices of non-zero elements 733 | 734 | ```python 735 | indices = vec.indices() 736 | ``` 737 | 738 | Get the values of non-zero elements 739 | 740 | ```python 741 | values = vec.values() 742 | ``` 743 | 744 | Get a list 745 | 746 | ```python 747 | lst = vec.to_list() 748 | ``` 749 | 750 | Get a NumPy array 751 | 752 | ```python 753 | arr = vec.to_numpy() 754 | ``` 755 | 756 | Get a SciPy sparse array 757 | 758 | ```python 759 | arr = vec.to_coo() 760 | ``` 761 | 762 | ## History 763 | 764 | View the [changelog](https://github.com/pgvector/pgvector-python/blob/master/CHANGELOG.md) 765 | 766 | ## Contributing 767 | 768 | Everyone is encouraged to help improve this project. Here are a few ways you can help: 769 | 770 | - [Report bugs](https://github.com/pgvector/pgvector-python/issues) 771 | - Fix bugs and [submit pull requests](https://github.com/pgvector/pgvector-python/pulls) 772 | - Write, clarify, or fix documentation 773 | - Suggest or add new features 774 | 775 | To get started with development: 776 | 777 | ```sh 778 | git clone https://github.com/pgvector/pgvector-python.git 779 | cd pgvector-python 780 | pip install -r requirements.txt 781 | createdb pgvector_python_test 782 | pytest 783 | ``` 784 | 785 | To run an example: 786 | 787 | ```sh 788 | cd examples/loading 789 | pip install -r requirements.txt 790 | createdb pgvector_example 791 | python3 example.py 792 | ``` 793 | -------------------------------------------------------------------------------- /examples/citus/example.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector.psycopg import register_vector 3 | import psycopg 4 | 5 | # generate random data 6 | rows = 100000 7 | dimensions = 128 8 | embeddings = np.random.rand(rows, dimensions) 9 | categories = np.random.randint(100, size=rows).tolist() 10 | queries = np.random.rand(10, dimensions) 11 | 12 | # enable extensions 13 | conn = psycopg.connect(dbname='pgvector_citus', autocommit=True) 14 | conn.execute('CREATE EXTENSION IF NOT EXISTS citus') 15 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 16 | 17 | # GUC variables set on the session do not propagate to Citus workers 18 | # https://github.com/citusdata/citus/issues/462 19 | # you can either: 20 | # 1. set them on the system, user, or database and reconnect 21 | # 2. set them for a transaction with SET LOCAL 22 | conn.execute("ALTER DATABASE pgvector_citus SET maintenance_work_mem = '512MB'") 23 | conn.execute('ALTER DATABASE pgvector_citus SET hnsw.ef_search = 20') 24 | conn.close() 25 | 26 | # reconnect for updated GUC variables to take effect 27 | conn = psycopg.connect(dbname='pgvector_citus', autocommit=True) 28 | register_vector(conn) 29 | 30 | print('Creating distributed table') 31 | conn.execute('DROP TABLE IF EXISTS items') 32 | conn.execute('CREATE TABLE items (id bigserial, embedding vector(%d), category_id bigint, PRIMARY KEY (id, category_id))' % dimensions) 33 | conn.execute('SET citus.shard_count = 4') 34 | conn.execute("SELECT create_distributed_table('items', 'category_id')") 35 | 36 | print('Loading data in parallel') 37 | with conn.cursor().copy('COPY items (embedding, category_id) FROM STDIN WITH (FORMAT BINARY)') as copy: 38 | copy.set_types(['vector', 'bigint']) 39 | 40 | for i in range(rows): 41 | copy.write_row([embeddings[i], categories[i]]) 42 | 43 | print('Creating index in parallel') 44 | conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)') 45 | 46 | print('Running distributed queries') 47 | for query in queries: 48 | items = conn.execute('SELECT id FROM items ORDER BY embedding <-> %s LIMIT 10', (query,)).fetchall() 49 | print([r[0] for r in items]) 50 | -------------------------------------------------------------------------------- /examples/citus/requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | pgvector 3 | psycopg[binary] 4 | -------------------------------------------------------------------------------- /examples/cohere/example.py: -------------------------------------------------------------------------------- 1 | import cohere 2 | import numpy as np 3 | from pgvector.psycopg import register_vector, Bit 4 | import psycopg 5 | 6 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 7 | 8 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 9 | register_vector(conn) 10 | 11 | conn.execute('DROP TABLE IF EXISTS documents') 12 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))') 13 | 14 | 15 | def embed(input, input_type): 16 | co = cohere.ClientV2() 17 | response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) 18 | return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary] 19 | 20 | 21 | input = [ 22 | 'The dog is barking', 23 | 'The cat is purring', 24 | 'The bear is growling' 25 | ] 26 | embeddings = embed(input, 'search_document') 27 | for content, embedding in zip(input, embeddings): 28 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Bit(embedding))) 29 | 30 | query = 'forest' 31 | query_embedding = embed([query], 'search_query')[0] 32 | result = conn.execute('SELECT content FROM documents ORDER BY embedding <~> %s LIMIT 5', (Bit(query_embedding),)).fetchall() 33 | for row in result: 34 | print(row[0]) 35 | -------------------------------------------------------------------------------- /examples/cohere/requirements.txt: -------------------------------------------------------------------------------- 1 | cohere 2 | pgvector 3 | psycopg[binary] 4 | -------------------------------------------------------------------------------- /examples/colbert/approximate.py: -------------------------------------------------------------------------------- 1 | # based on section 3.6 of https://arxiv.org/abs/2004.12832 2 | 3 | from colbert.infra import ColBERTConfig 4 | from colbert.modeling.checkpoint import Checkpoint 5 | from pgvector.psycopg import register_vector 6 | import psycopg 7 | import warnings 8 | 9 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 10 | 11 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 12 | register_vector(conn) 13 | 14 | conn.execute('DROP TABLE IF EXISTS documents') 15 | conn.execute('DROP TABLE IF EXISTS document_embeddings') 16 | 17 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text)') 18 | conn.execute('CREATE TABLE document_embeddings (id bigserial PRIMARY KEY, document_id bigint, embedding vector(128))') 19 | 20 | conn.execute(""" 21 | CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ 22 | WITH queries AS ( 23 | SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) 24 | ), 25 | documents AS ( 26 | SELECT unnest(document) AS document 27 | ), 28 | similarities AS ( 29 | SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents 30 | ), 31 | max_similarities AS ( 32 | SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number 33 | ) 34 | SELECT SUM(max_similarity) FROM max_similarities 35 | $$ LANGUAGE SQL 36 | """) 37 | 38 | warnings.filterwarnings('ignore') # ignore warnings from colbert 39 | 40 | config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) 41 | checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) 42 | 43 | input = [ 44 | 'The dog is barking', 45 | 'The cat is purring', 46 | 'The bear is growling' 47 | ] 48 | doc_embeddings = checkpoint.docFromText(input, keep_dims=False) 49 | for content, embeddings in zip(input, doc_embeddings): 50 | with conn.transaction(): 51 | result = conn.execute('INSERT INTO documents (content) VALUES (%s) RETURNING id', (content,)).fetchone() 52 | params = [] 53 | for embedding in embeddings: 54 | params.extend([result[0], embedding.numpy()]) 55 | values = ', '.join(['(%s, %s)' for _ in embeddings]) 56 | conn.execute(f'INSERT INTO document_embeddings (document_id, embedding) VALUES {values}', params) 57 | 58 | conn.execute('CREATE INDEX ON document_embeddings (document_id)') 59 | conn.execute('CREATE INDEX ON document_embeddings USING hnsw (embedding vector_cosine_ops)') 60 | 61 | query = 'puppy' 62 | query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query])[0]] 63 | approximate_stage = ' UNION ALL '.join(['(SELECT document_id FROM document_embeddings ORDER BY embedding <=> %s LIMIT 5)' for _ in query_embeddings]) 64 | sql = f""" 65 | WITH approximate_stage AS ( 66 | {approximate_stage} 67 | ), 68 | embeddings AS ( 69 | SELECT document_id, array_agg(embedding) AS embeddings FROM document_embeddings 70 | WHERE document_id IN (SELECT DISTINCT document_id FROM approximate_stage) 71 | GROUP BY document_id 72 | ) 73 | SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents 74 | INNER JOIN embeddings ON embeddings.document_id = documents.id 75 | ORDER BY max_sim DESC LIMIT 10 76 | """ 77 | params = query_embeddings + [query_embeddings] 78 | result = conn.execute(sql, params).fetchall() 79 | for row in result: 80 | print(row) 81 | -------------------------------------------------------------------------------- /examples/colbert/exact.py: -------------------------------------------------------------------------------- 1 | from colbert.infra import ColBERTConfig 2 | from colbert.modeling.checkpoint import Checkpoint 3 | from pgvector.psycopg import register_vector 4 | import psycopg 5 | import warnings 6 | 7 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 8 | 9 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 10 | register_vector(conn) 11 | 12 | conn.execute('DROP TABLE IF EXISTS documents') 13 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embeddings vector(128)[])') 14 | conn.execute(""" 15 | CREATE OR REPLACE FUNCTION max_sim(document vector[], query vector[]) RETURNS double precision AS $$ 16 | WITH queries AS ( 17 | SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) 18 | ), 19 | documents AS ( 20 | SELECT unnest(document) AS document 21 | ), 22 | similarities AS ( 23 | SELECT query_number, 1 - (document <=> query) AS similarity FROM queries CROSS JOIN documents 24 | ), 25 | max_similarities AS ( 26 | SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number 27 | ) 28 | SELECT SUM(max_similarity) FROM max_similarities 29 | $$ LANGUAGE SQL 30 | """) 31 | 32 | warnings.filterwarnings('ignore') # ignore warnings from colbert 33 | 34 | config = ColBERTConfig(doc_maxlen=220, query_maxlen=32) 35 | checkpoint = Checkpoint('colbert-ir/colbertv2.0', colbert_config=config, verbose=0) 36 | 37 | input = [ 38 | 'The dog is barking', 39 | 'The cat is purring', 40 | 'The bear is growling' 41 | ] 42 | doc_embeddings = checkpoint.docFromText(input, keep_dims=False) 43 | for content, embeddings in zip(input, doc_embeddings): 44 | embeddings = [e.numpy() for e in embeddings] 45 | conn.execute('INSERT INTO documents (content, embeddings) VALUES (%s, %s)', (content, embeddings)) 46 | 47 | query = 'puppy' 48 | query_embeddings = [e.numpy() for e in checkpoint.queryFromText([query])[0]] 49 | result = conn.execute('SELECT content, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() 50 | for row in result: 51 | print(row) 52 | -------------------------------------------------------------------------------- /examples/colbert/requirements.txt: -------------------------------------------------------------------------------- 1 | colbert-ai 2 | pgvector 3 | psycopg[binary] 4 | transformers==4.49.0 5 | -------------------------------------------------------------------------------- /examples/colpali/exact.py: -------------------------------------------------------------------------------- 1 | from colpali_engine.models import ColQwen2, ColQwen2Processor 2 | from colpali_engine.utils.torch_utils import get_torch_device 3 | from datasets import load_dataset 4 | from pgvector.psycopg import register_vector, Bit 5 | import psycopg 6 | import torch 7 | 8 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 9 | 10 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 11 | register_vector(conn) 12 | 13 | conn.execute('DROP TABLE IF EXISTS documents') 14 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, embeddings bit(128)[])') 15 | conn.execute(""" 16 | CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) RETURNS double precision AS $$ 17 | WITH queries AS ( 18 | SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) 19 | ), 20 | documents AS ( 21 | SELECT unnest(document) AS document 22 | ), 23 | similarities AS ( 24 | SELECT query_number, 1 - ((document <~> query) / bit_length(query)) AS similarity FROM queries CROSS JOIN documents 25 | ), 26 | max_similarities AS ( 27 | SELECT MAX(similarity) AS max_similarity FROM similarities GROUP BY query_number 28 | ) 29 | SELECT SUM(max_similarity) FROM max_similarities 30 | $$ LANGUAGE SQL 31 | """) 32 | 33 | device = get_torch_device('auto') 34 | model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval() 35 | processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0') 36 | 37 | 38 | def generate_embeddings(processed): 39 | with torch.no_grad(): 40 | return model(**processed.to(model.device)).to(torch.float32).numpy(force=True) 41 | 42 | 43 | def binary_quantize(embedding): 44 | return Bit(embedding > 0) 45 | 46 | 47 | input = load_dataset('vidore/docvqa_test_subsampled', split='test[:3]')['image'] 48 | for content in input: 49 | embeddings = [binary_quantize(e) for e in generate_embeddings(processor.process_images([content]))[0]] 50 | conn.execute('INSERT INTO documents (embeddings) VALUES (%s)', (embeddings,)) 51 | 52 | query = 'dividend' 53 | query_embeddings = [binary_quantize(e) for e in generate_embeddings(processor.process_queries([query]))[0]] 54 | result = conn.execute('SELECT id, max_sim(embeddings, %s) AS max_sim FROM documents ORDER BY max_sim DESC LIMIT 5', (query_embeddings,)).fetchall() 55 | for row in result: 56 | print(row) 57 | -------------------------------------------------------------------------------- /examples/colpali/requirements.txt: -------------------------------------------------------------------------------- 1 | colpali-engine 2 | datasets 3 | pgvector 4 | psycopg[binary] 5 | -------------------------------------------------------------------------------- /examples/gensim/example.py: -------------------------------------------------------------------------------- 1 | from gensim.corpora.dictionary import Dictionary 2 | from gensim.models import LdaModel 3 | from gensim.utils import simple_preprocess 4 | import numpy as np 5 | from pgvector.psycopg import register_vector 6 | import psycopg 7 | 8 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 9 | 10 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 11 | register_vector(conn) 12 | 13 | conn.execute('DROP TABLE IF EXISTS documents') 14 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(20))') 15 | 16 | input = [ 17 | 'The dog is barking', 18 | 'The cat is purring', 19 | 'The bear is growling' 20 | ] 21 | 22 | docs = [simple_preprocess(content) for content in input] 23 | dictionary = Dictionary(docs) 24 | dictionary.filter_extremes(no_below=1) 25 | corpus = [dictionary.doc2bow(tokens) for tokens in docs] 26 | model = LdaModel(corpus, num_topics=20) 27 | 28 | for content, bow in zip(input, corpus): 29 | embedding = np.array([v[1] for v in model.get_document_topics(bow, minimum_probability=0)]) 30 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) 31 | 32 | document_id = 1 33 | neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() 34 | for neighbor in neighbors: 35 | print(neighbor[0]) 36 | -------------------------------------------------------------------------------- /examples/gensim/requirements.txt: -------------------------------------------------------------------------------- 1 | gensim 2 | numpy 3 | pgvector 4 | psycopg[binary] 5 | scipy<1.13 6 | -------------------------------------------------------------------------------- /examples/hybrid_search/cross_encoder.py: -------------------------------------------------------------------------------- 1 | # good resources 2 | # https://qdrant.tech/articles/hybrid-search/ 3 | # https://www.sbert.net/examples/applications/semantic-search/README.html 4 | 5 | import asyncio 6 | import itertools 7 | from pgvector.psycopg import register_vector_async 8 | import psycopg 9 | from sentence_transformers import CrossEncoder, SentenceTransformer 10 | 11 | sentences = [ 12 | 'The dog is barking', 13 | 'The cat is purring', 14 | 'The bear is growling' 15 | ] 16 | query = 'growling bear' 17 | 18 | 19 | async def create_schema(conn): 20 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 21 | await register_vector_async(conn) 22 | 23 | await conn.execute('DROP TABLE IF EXISTS documents') 24 | await conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') 25 | await conn.execute("CREATE INDEX ON documents USING GIN (to_tsvector('english', content))") 26 | 27 | 28 | async def insert_data(conn): 29 | model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') 30 | embeddings = model.encode(sentences) 31 | 32 | sql = 'INSERT INTO documents (content, embedding) VALUES ' + ', '.join(['(%s, %s)' for _ in embeddings]) 33 | params = list(itertools.chain(*zip(sentences, embeddings))) 34 | await conn.execute(sql, params) 35 | 36 | 37 | async def semantic_search(conn, query): 38 | model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') 39 | embedding = model.encode(query) 40 | 41 | async with conn.cursor() as cur: 42 | await cur.execute('SELECT id, content FROM documents ORDER BY embedding <=> %s LIMIT 5', (embedding,)) 43 | return await cur.fetchall() 44 | 45 | 46 | async def keyword_search(conn, query): 47 | async with conn.cursor() as cur: 48 | await cur.execute("SELECT id, content FROM documents, plainto_tsquery('english', %s) query WHERE to_tsvector('english', content) @@ query ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC LIMIT 5", (query,)) 49 | return await cur.fetchall() 50 | 51 | 52 | def rerank(query, results): 53 | # deduplicate 54 | results = set(itertools.chain(*results)) 55 | 56 | # re-rank 57 | encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') 58 | scores = encoder.predict([(query, item[1]) for item in results]) 59 | return [v for _, v in sorted(zip(scores, results), reverse=True)] 60 | 61 | 62 | async def main(): 63 | conn = await psycopg.AsyncConnection.connect(dbname='pgvector_example', autocommit=True) 64 | await create_schema(conn) 65 | await insert_data(conn) 66 | 67 | # perform queries in parallel 68 | results = await asyncio.gather(semantic_search(conn, query), keyword_search(conn, query)) 69 | results = rerank(query, results) 70 | print(results) 71 | 72 | 73 | asyncio.run(main()) 74 | -------------------------------------------------------------------------------- /examples/hybrid_search/requirements.txt: -------------------------------------------------------------------------------- 1 | pgvector 2 | psycopg[binary] 3 | sentence-transformers 4 | -------------------------------------------------------------------------------- /examples/hybrid_search/rrf.py: -------------------------------------------------------------------------------- 1 | from pgvector.psycopg import register_vector 2 | import psycopg 3 | from sentence_transformers import SentenceTransformer 4 | 5 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 6 | 7 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 8 | register_vector(conn) 9 | 10 | conn.execute('DROP TABLE IF EXISTS documents') 11 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') 12 | conn.execute("CREATE INDEX ON documents USING GIN (to_tsvector('english', content))") 13 | 14 | sentences = [ 15 | 'The dog is barking', 16 | 'The cat is purring', 17 | 'The bear is growling' 18 | ] 19 | model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') 20 | embeddings = model.encode(sentences) 21 | for content, embedding in zip(sentences, embeddings): 22 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) 23 | 24 | sql = """ 25 | WITH semantic_search AS ( 26 | SELECT id, RANK () OVER (ORDER BY embedding <=> %(embedding)s) AS rank 27 | FROM documents 28 | ORDER BY embedding <=> %(embedding)s 29 | LIMIT 20 30 | ), 31 | keyword_search AS ( 32 | SELECT id, RANK () OVER (ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC) 33 | FROM documents, plainto_tsquery('english', %(query)s) query 34 | WHERE to_tsvector('english', content) @@ query 35 | ORDER BY ts_rank_cd(to_tsvector('english', content), query) DESC 36 | LIMIT 20 37 | ) 38 | SELECT 39 | COALESCE(semantic_search.id, keyword_search.id) AS id, 40 | COALESCE(1.0 / (%(k)s + semantic_search.rank), 0.0) + 41 | COALESCE(1.0 / (%(k)s + keyword_search.rank), 0.0) AS score 42 | FROM semantic_search 43 | FULL OUTER JOIN keyword_search ON semantic_search.id = keyword_search.id 44 | ORDER BY score DESC 45 | LIMIT 5 46 | """ 47 | query = 'growling bear' 48 | embedding = model.encode(query) 49 | k = 60 50 | results = conn.execute(sql, {'query': query, 'embedding': embedding, 'k': k}).fetchall() 51 | for row in results: 52 | print('document:', row[0], 'RRF score:', row[1]) 53 | -------------------------------------------------------------------------------- /examples/image_search/example.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | from pgvector.psycopg import register_vector 3 | import psycopg 4 | import tempfile 5 | import torch 6 | import torchvision 7 | from tqdm import tqdm 8 | 9 | seed = True 10 | 11 | # establish connection 12 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 13 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 14 | register_vector(conn) 15 | 16 | # load images 17 | transform = torchvision.transforms.Compose([ 18 | torchvision.transforms.ToTensor(), 19 | torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) 20 | ]) 21 | dataset = torchvision.datasets.CIFAR10(root=tempfile.gettempdir(), train=True, download=True, transform=transform) 22 | dataloader = torch.utils.data.DataLoader(dataset, batch_size=1000) 23 | 24 | # load pretrained model 25 | device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') 26 | model = torchvision.models.resnet18(weights='DEFAULT') 27 | model.fc = torch.nn.Identity() 28 | model.to(device) 29 | model.eval() 30 | 31 | 32 | def generate_embeddings(inputs): 33 | return model(inputs.to(device)).detach().cpu().numpy() 34 | 35 | 36 | # generate and store embeddings 37 | if seed: 38 | conn.execute('DROP TABLE IF EXISTS images') 39 | conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, embedding vector(512))') 40 | 41 | print('Generating embeddings') 42 | for data in tqdm(dataloader): 43 | embeddings = generate_embeddings(data[0]) 44 | 45 | sql = 'INSERT INTO images (embedding) VALUES ' + ','.join(['(%s)' for _ in embeddings]) 46 | params = [embedding for embedding in embeddings] 47 | conn.execute(sql, params) 48 | 49 | # load 5 random unseen images 50 | queryset = torchvision.datasets.CIFAR10(root=tempfile.gettempdir(), train=False, download=True, transform=transform) 51 | queryloader = torch.utils.data.DataLoader(queryset, batch_size=5, shuffle=True) 52 | images = next(iter(queryloader))[0] 53 | 54 | # generate and query embeddings 55 | results = [] 56 | embeddings = generate_embeddings(images) 57 | for image, embedding in zip(images, embeddings): 58 | result = conn.execute('SELECT id FROM images ORDER BY embedding <=> %s LIMIT 5', (embedding,)).fetchall() 59 | nearest_images = [dataset[row[0] - 1][0] for row in result] 60 | results.append([image] + nearest_images) 61 | 62 | # show images 63 | fig, axs = plt.subplots(len(results), len(results[0])) 64 | for i, result in enumerate(results): 65 | for j, image in enumerate(result): 66 | ax = axs[i, j] 67 | ax.imshow((image / 2 + 0.5).permute(1, 2, 0).numpy()) 68 | ax.set_axis_off() 69 | plt.show(block=True) 70 | -------------------------------------------------------------------------------- /examples/image_search/requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | pgvector 3 | psycopg[binary] 4 | torch 5 | torchvision 6 | tqdm 7 | -------------------------------------------------------------------------------- /examples/imagehash/example.py: -------------------------------------------------------------------------------- 1 | from datasets import load_dataset 2 | from imagehash import phash 3 | import matplotlib.pyplot as plt 4 | from pgvector.psycopg import register_vector, Bit 5 | import psycopg 6 | 7 | 8 | def hash_image(img): 9 | return ''.join(['1' if v else '0' for v in phash(img).hash.flatten()]) 10 | 11 | 12 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 13 | 14 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 15 | register_vector(conn) 16 | 17 | conn.execute('DROP TABLE IF EXISTS images') 18 | conn.execute('CREATE TABLE images (id bigserial PRIMARY KEY, hash bit(64))') 19 | 20 | print('Loading dataset') 21 | dataset = load_dataset('mnist') 22 | 23 | print('Generating hashes') 24 | images = [{'hash': hash_image(row['image'])} for row in dataset['train']] 25 | 26 | print('Storing hashes') 27 | cur = conn.cursor() 28 | with cur.copy('COPY images (hash) FROM STDIN') as copy: 29 | for image in images: 30 | copy.write_row([Bit(image['hash'])]) 31 | 32 | print('Querying hashes') 33 | results = [] 34 | for i in range(5): 35 | image = dataset['test'][i]['image'] 36 | result = conn.execute('SELECT id FROM images ORDER BY hash <~> %s LIMIT 5', (hash_image(image),)).fetchall() 37 | nearest_images = [dataset['train'][row[0] - 1]['image'] for row in result] 38 | results.append([image] + nearest_images) 39 | 40 | print('Showing results (first column is query image)') 41 | fig, axs = plt.subplots(len(results), len(results[0])) 42 | for i, result in enumerate(results): 43 | for j, image in enumerate(result): 44 | ax = axs[i, j] 45 | ax.imshow(image) 46 | ax.set_axis_off() 47 | plt.show(block=True) 48 | -------------------------------------------------------------------------------- /examples/imagehash/requirements.txt: -------------------------------------------------------------------------------- 1 | datasets 2 | imagehash 3 | matplotlib 4 | pgvector 5 | psycopg[binary] 6 | -------------------------------------------------------------------------------- /examples/implicit/example.py: -------------------------------------------------------------------------------- 1 | import implicit 2 | from implicit.datasets.movielens import get_movielens 3 | from pgvector.sqlalchemy import Vector 4 | from sqlalchemy import create_engine, insert, select, text, Integer, String 5 | from sqlalchemy.orm import declarative_base, mapped_column, Session 6 | 7 | engine = create_engine('postgresql+psycopg://localhost/pgvector_example') 8 | with engine.connect() as conn: 9 | conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) 10 | conn.commit() 11 | 12 | Base = declarative_base() 13 | 14 | 15 | class User(Base): 16 | __tablename__ = 'user' 17 | 18 | id = mapped_column(Integer, primary_key=True) 19 | factors = mapped_column(Vector(20)) 20 | 21 | 22 | class Item(Base): 23 | __tablename__ = 'item' 24 | 25 | id = mapped_column(Integer, primary_key=True) 26 | title = mapped_column(String) 27 | factors = mapped_column(Vector(20)) 28 | 29 | 30 | Base.metadata.drop_all(engine) 31 | Base.metadata.create_all(engine) 32 | 33 | titles, ratings = get_movielens('100k') 34 | model = implicit.als.AlternatingLeastSquares(factors=20) 35 | model.fit(ratings) 36 | 37 | users = [dict(factors=factors) for i, factors in enumerate(model.user_factors)] 38 | items = [dict(title=titles[i], factors=factors) for i, factors in enumerate(model.item_factors)] 39 | 40 | session = Session(engine) 41 | session.execute(insert(User), users) 42 | session.execute(insert(Item), items) 43 | 44 | user = session.get(User, 1) 45 | items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) 46 | print('user-based recs:', [item.title for item in items]) 47 | 48 | item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() 49 | items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) 50 | print('item-based recs:', [item.title for item in items]) 51 | -------------------------------------------------------------------------------- /examples/implicit/requirements.txt: -------------------------------------------------------------------------------- 1 | h5py 2 | implicit 3 | pgvector 4 | psycopg[binary] 5 | SQLAlchemy 6 | -------------------------------------------------------------------------------- /examples/lightfm/example.py: -------------------------------------------------------------------------------- 1 | from lightfm import LightFM 2 | from lightfm.datasets import fetch_movielens 3 | from pgvector.sqlalchemy import Vector 4 | from sqlalchemy import create_engine, insert, select, text, Float, Integer, String 5 | from sqlalchemy.orm import declarative_base, mapped_column, Session 6 | 7 | engine = create_engine('postgresql+psycopg://localhost/pgvector_example') 8 | with engine.connect() as conn: 9 | conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) 10 | conn.commit() 11 | 12 | Base = declarative_base() 13 | 14 | 15 | class User(Base): 16 | __tablename__ = 'user' 17 | 18 | id = mapped_column(Integer, primary_key=True) 19 | factors = mapped_column(Vector(20)) 20 | 21 | 22 | class Item(Base): 23 | __tablename__ = 'item' 24 | 25 | id = mapped_column(Integer, primary_key=True) 26 | title = mapped_column(String) 27 | factors = mapped_column(Vector(20)) 28 | bias = mapped_column(Float) 29 | 30 | 31 | Base.metadata.drop_all(engine) 32 | Base.metadata.create_all(engine) 33 | 34 | data = fetch_movielens(min_rating=5.0) 35 | model = LightFM(loss='warp', no_components=20) 36 | model.fit(data['train'], epochs=30) 37 | 38 | user_biases, user_factors = model.get_user_representations() 39 | item_biases, item_factors = model.get_item_representations() 40 | 41 | users = [dict(factors=factors) for i, factors in enumerate(user_factors)] 42 | items = [dict(title=data['item_labels'][i], factors=factors, bias=item_biases[i].item()) for i, factors in enumerate(item_factors)] 43 | 44 | session = Session(engine) 45 | session.execute(insert(User), users) 46 | session.execute(insert(Item), items) 47 | 48 | user = session.get(User, 1) 49 | # subtract item bias for negative inner product 50 | items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors) - Item.bias).limit(5)) 51 | print('user-based recs:', [item.title for item in items]) 52 | 53 | # broken due to https://github.com/lyst/lightfm/issues/682 54 | item = session.scalars(select(Item).filter(Item.title == 'Star Wars (1977)')).first() 55 | items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) 56 | print('item-based recs:', [item.title for item in items]) 57 | -------------------------------------------------------------------------------- /examples/lightfm/requirements.txt: -------------------------------------------------------------------------------- 1 | lightfm 2 | pgvector 3 | psycopg[binary] 4 | SQLAlchemy 5 | -------------------------------------------------------------------------------- /examples/loading/example.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector.psycopg import register_vector 3 | import psycopg 4 | 5 | # generate random data 6 | rows = 1000000 7 | dimensions = 128 8 | embeddings = np.random.rand(rows, dimensions) 9 | 10 | # enable extension 11 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 12 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 13 | register_vector(conn) 14 | 15 | # create table 16 | conn.execute('DROP TABLE IF EXISTS items') 17 | conn.execute(f'CREATE TABLE items (id bigserial, embedding vector({dimensions}))') 18 | 19 | # load data 20 | print(f'Loading {len(embeddings)} rows') 21 | cur = conn.cursor() 22 | with cur.copy('COPY items (embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: 23 | # use set_types for binary copy 24 | # https://www.psycopg.org/psycopg3/docs/basic/copy.html#binary-copy 25 | copy.set_types(['vector']) 26 | 27 | for i, embedding in enumerate(embeddings): 28 | copy.write_row([embedding]) 29 | 30 | # show progress 31 | if i % 10000 == 0: 32 | print('.', end='', flush=True) 33 | 34 | print('\nSuccess!') 35 | 36 | # create any indexes *after* loading initial data (skipping for this example) 37 | create_index = False 38 | if create_index: 39 | print('Creating index') 40 | conn.execute("SET maintenance_work_mem = '8GB'") 41 | conn.execute('SET max_parallel_maintenance_workers = 7') 42 | conn.execute('CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)') 43 | 44 | # update planner statistics for good measure 45 | conn.execute('ANALYZE items') 46 | -------------------------------------------------------------------------------- /examples/loading/requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | pgvector 3 | psycopg[binary] 4 | -------------------------------------------------------------------------------- /examples/openai/example.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from openai import OpenAI 3 | from pgvector.psycopg import register_vector 4 | import psycopg 5 | 6 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 7 | 8 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 9 | register_vector(conn) 10 | 11 | conn.execute('DROP TABLE IF EXISTS documents') 12 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))') 13 | 14 | 15 | def embed(input): 16 | client = OpenAI() 17 | response = client.embeddings.create(input=input, model='text-embedding-3-small') 18 | return [v.embedding for v in response.data] 19 | 20 | 21 | input = [ 22 | 'The dog is barking', 23 | 'The cat is purring', 24 | 'The bear is growling' 25 | ] 26 | embeddings = embed(input) 27 | for content, embedding in zip(input, embeddings): 28 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, np.array(embedding))) 29 | 30 | query = 'forest' 31 | query_embedding = embed([query])[0] 32 | result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (np.array(query_embedding),)).fetchall() 33 | for row in result: 34 | print(row[0]) 35 | -------------------------------------------------------------------------------- /examples/openai/halfvec.py: -------------------------------------------------------------------------------- 1 | from openai import OpenAI 2 | from pgvector.psycopg import register_vector, HalfVector 3 | import psycopg 4 | 5 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 6 | 7 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 8 | register_vector(conn) 9 | 10 | conn.execute('DROP TABLE IF EXISTS documents') 11 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding halfvec(3072))') 12 | conn.execute('CREATE INDEX ON documents USING hnsw (embedding halfvec_cosine_ops)') 13 | 14 | 15 | def embed(input): 16 | client = OpenAI() 17 | response = client.embeddings.create(input=input, model='text-embedding-3-large') 18 | return [v.embedding for v in response.data] 19 | 20 | 21 | input = [ 22 | 'The dog is barking', 23 | 'The cat is purring', 24 | 'The bear is growling' 25 | ] 26 | embeddings = embed(input) 27 | for content, embedding in zip(input, embeddings): 28 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, HalfVector(embedding))) 29 | 30 | query = 'forest' 31 | query_embedding = embed([query])[0] 32 | result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (HalfVector(query_embedding),)).fetchall() 33 | for row in result: 34 | print(row[0]) 35 | -------------------------------------------------------------------------------- /examples/openai/requirements.txt: -------------------------------------------------------------------------------- 1 | openai 2 | pgvector 3 | psycopg[binary] 4 | -------------------------------------------------------------------------------- /examples/rag/example.py: -------------------------------------------------------------------------------- 1 | # Run: 2 | # ollama pull llama3.2 3 | # ollama pull nomic-embed-text 4 | # ollama serve 5 | 6 | import numpy as np 7 | import ollama 8 | from pathlib import Path 9 | from pgvector.psycopg import register_vector 10 | import psycopg 11 | import urllib.request 12 | 13 | query = 'What index types are supported?' 14 | load_data = True 15 | 16 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 17 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 18 | register_vector(conn) 19 | 20 | if load_data: 21 | # get data 22 | url = 'https://raw.githubusercontent.com/pgvector/pgvector/refs/heads/master/README.md' 23 | dest = Path(__file__).parent / 'README.md' 24 | if not dest.exists(): 25 | urllib.request.urlretrieve(url, dest) 26 | 27 | with open(dest, encoding='utf-8') as f: 28 | doc = f.read() 29 | 30 | # generate chunks 31 | # TODO improve chunking 32 | # TODO remove markdown 33 | chunks = doc.split('\n## ') 34 | 35 | # embed chunks 36 | # nomic-embed-text has task instruction prefix 37 | input = ['search_document: ' + chunk for chunk in chunks] 38 | embeddings = ollama.embed(model='nomic-embed-text', input=input).embeddings 39 | 40 | # create table 41 | conn.execute('DROP TABLE IF EXISTS chunks') 42 | conn.execute('CREATE TABLE chunks (id bigserial PRIMARY KEY, content text, embedding vector(768))') 43 | 44 | # store chunks 45 | cur = conn.cursor() 46 | with cur.copy('COPY chunks (content, embedding) FROM STDIN WITH (FORMAT BINARY)') as copy: 47 | copy.set_types(['text', 'vector']) 48 | 49 | for content, embedding in zip(chunks, embeddings): 50 | copy.write_row([content, embedding]) 51 | 52 | # embed query 53 | # nomic-embed-text has task instruction prefix 54 | input = 'search_query: ' + query 55 | embedding = ollama.embed(model='nomic-embed-text', input=input).embeddings[0] 56 | 57 | # retrieve chunks 58 | result = conn.execute('SELECT content FROM chunks ORDER BY embedding <=> %s LIMIT 5', (np.array(embedding),)).fetchall() 59 | context = '\n\n'.join([row[0] for row in result]) 60 | 61 | # get answer 62 | # TODO improve prompt 63 | prompt = f'Answer this question: {query}\n\n{context}' 64 | response = ollama.generate(model='llama3.2', prompt=prompt).response 65 | print(response) 66 | -------------------------------------------------------------------------------- /examples/rag/requirements.txt: -------------------------------------------------------------------------------- 1 | ollama 2 | pgvector 3 | psycopg[binary] 4 | -------------------------------------------------------------------------------- /examples/rdkit/example.py: -------------------------------------------------------------------------------- 1 | # good resource 2 | # https://www.rdkit.org/docs/GettingStartedInPython.html#morgan-fingerprints-circular-fingerprints 3 | 4 | from pgvector.psycopg import register_vector, Bit 5 | import psycopg 6 | from rdkit import Chem 7 | from rdkit.Chem import AllChem 8 | 9 | 10 | def generate_fingerprint(molecule): 11 | fpgen = AllChem.GetMorganGenerator() 12 | return fpgen.GetFingerprintAsNumPy(Chem.MolFromSmiles(molecule)) 13 | 14 | 15 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 16 | 17 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 18 | register_vector(conn) 19 | 20 | conn.execute('DROP TABLE IF EXISTS molecules') 21 | conn.execute('CREATE TABLE molecules (id text PRIMARY KEY, fingerprint bit(2048))') 22 | 23 | molecules = ['Cc1ccccc1', 'Cc1ncccc1', 'c1ccccn1'] 24 | for molecule in molecules: 25 | fingerprint = generate_fingerprint(molecule) 26 | conn.execute('INSERT INTO molecules (id, fingerprint) VALUES (%s, %s)', (molecule, Bit(fingerprint))) 27 | 28 | query_molecule = 'c1ccco1' 29 | query_fingerprint = generate_fingerprint(query_molecule) 30 | result = conn.execute('SELECT id, fingerprint <%%> %s AS distance FROM molecules ORDER BY distance LIMIT 5', (Bit(query_fingerprint),)).fetchall() 31 | for row in result: 32 | print(row) 33 | -------------------------------------------------------------------------------- /examples/rdkit/requirements.txt: -------------------------------------------------------------------------------- 1 | pgvector 2 | psycopg[binary] 3 | rdkit 4 | -------------------------------------------------------------------------------- /examples/sentence_transformers/example.py: -------------------------------------------------------------------------------- 1 | from pgvector.psycopg import register_vector 2 | import psycopg 3 | from sentence_transformers import SentenceTransformer 4 | 5 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 6 | 7 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 8 | register_vector(conn) 9 | 10 | conn.execute('DROP TABLE IF EXISTS documents') 11 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') 12 | 13 | model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') 14 | 15 | input = [ 16 | 'The dog is barking', 17 | 'The cat is purring', 18 | 'The bear is growling' 19 | ] 20 | embeddings = model.encode(input) 21 | for content, embedding in zip(input, embeddings): 22 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) 23 | 24 | query = 'forest' 25 | query_embedding = model.encode(query) 26 | result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (query_embedding,)).fetchall() 27 | for row in result: 28 | print(row[0]) 29 | -------------------------------------------------------------------------------- /examples/sentence_transformers/requirements.txt: -------------------------------------------------------------------------------- 1 | pgvector 2 | psycopg[binary] 3 | sentence-transformers 4 | -------------------------------------------------------------------------------- /examples/sparse_search/example.py: -------------------------------------------------------------------------------- 1 | # good resources 2 | # https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/ 3 | # https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1 4 | 5 | from pgvector.psycopg import register_vector, SparseVector 6 | import psycopg 7 | import torch 8 | from transformers import AutoModelForMaskedLM, AutoTokenizer 9 | 10 | conn = psycopg.connect(dbname='pgvector_example', autocommit=True) 11 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 12 | register_vector(conn) 13 | 14 | conn.execute('DROP TABLE IF EXISTS documents') 15 | conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding sparsevec(30522))') 16 | 17 | model_id = 'opensearch-project/opensearch-neural-sparse-encoding-v1' 18 | model = AutoModelForMaskedLM.from_pretrained(model_id) 19 | tokenizer = AutoTokenizer.from_pretrained(model_id) 20 | special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()] 21 | 22 | 23 | def embed(input): 24 | feature = tokenizer( 25 | input, 26 | padding=True, 27 | truncation=True, 28 | return_tensors='pt', 29 | return_token_type_ids=False 30 | ) 31 | output = model(**feature)[0] 32 | 33 | values, _ = torch.max(output * feature['attention_mask'].unsqueeze(-1), dim=1) 34 | values = torch.log(1 + torch.relu(values)) 35 | values[:, special_token_ids] = 0 36 | return values.detach().cpu().numpy() 37 | 38 | 39 | # note: works much better with longer content 40 | input = [ 41 | 'The dog is barking', 42 | 'The cat is purring', 43 | 'The bear is growling' 44 | ] 45 | embeddings = embed(input) 46 | for content, embedding in zip(input, embeddings): 47 | conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector(embedding))) 48 | 49 | query = 'forest' 50 | query_embedding = embed([query])[0] 51 | result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector(query_embedding),)).fetchall() 52 | for row in result: 53 | print(row[0]) 54 | -------------------------------------------------------------------------------- /examples/sparse_search/requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | pgvector 3 | psycopg[binary] 4 | torch 5 | transformers 6 | -------------------------------------------------------------------------------- /examples/surprise/example.py: -------------------------------------------------------------------------------- 1 | from pgvector.sqlalchemy import Vector 2 | from sqlalchemy import create_engine, insert, select, text, Integer 3 | from sqlalchemy.orm import declarative_base, mapped_column, Session 4 | from surprise import Dataset, SVD 5 | 6 | engine = create_engine('postgresql+psycopg://localhost/pgvector_example') 7 | with engine.connect() as conn: 8 | conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector')) 9 | conn.commit() 10 | 11 | Base = declarative_base() 12 | 13 | 14 | class User(Base): 15 | __tablename__ = 'user' 16 | 17 | id = mapped_column(Integer, primary_key=True) 18 | factors = mapped_column(Vector(20)) 19 | 20 | 21 | class Item(Base): 22 | __tablename__ = 'item' 23 | 24 | id = mapped_column(Integer, primary_key=True) 25 | factors = mapped_column(Vector(20)) 26 | 27 | 28 | Base.metadata.drop_all(engine) 29 | Base.metadata.create_all(engine) 30 | 31 | data = Dataset.load_builtin('ml-100k') 32 | trainset = data.build_full_trainset() 33 | algo = SVD(n_factors=20, biased=False) 34 | algo.fit(trainset) 35 | 36 | users = [dict(id=trainset.to_raw_uid(i), factors=algo.pu[i]) for i in trainset.all_users()] 37 | items = [dict(id=trainset.to_raw_iid(i), factors=algo.qi[i]) for i in trainset.all_items()] 38 | 39 | session = Session(engine) 40 | session.execute(insert(User), users) 41 | session.execute(insert(Item), items) 42 | 43 | user = session.get(User, 1) 44 | items = session.scalars(select(Item).order_by(Item.factors.max_inner_product(user.factors)).limit(5)) 45 | print('user-based recs:', [item.id for item in items]) 46 | 47 | item = session.get(Item, 50) 48 | items = session.scalars(select(Item).filter(Item.id != item.id).order_by(Item.factors.cosine_distance(item.factors)).limit(5)) 49 | print('item-based recs:', [item.id for item in items]) 50 | -------------------------------------------------------------------------------- /examples/surprise/requirements.txt: -------------------------------------------------------------------------------- 1 | pgvector 2 | psycopg[binary] 3 | scikit-surprise 4 | SQLAlchemy 5 | -------------------------------------------------------------------------------- /pgvector/__init__.py: -------------------------------------------------------------------------------- 1 | from .bit import Bit 2 | from .halfvec import HalfVector 3 | from .sparsevec import SparseVector 4 | from .vector import Vector 5 | 6 | __all__ = [ 7 | 'Vector', 8 | 'HalfVector', 9 | 'Bit', 10 | 'SparseVector' 11 | ] 12 | -------------------------------------------------------------------------------- /pgvector/asyncpg/__init__.py: -------------------------------------------------------------------------------- 1 | from .register import register_vector 2 | 3 | # TODO remove 4 | from .. import Vector, HalfVector, SparseVector 5 | 6 | __all__ = [ 7 | 'register_vector', 8 | 'Vector', 9 | 'HalfVector', 10 | 'SparseVector' 11 | ] 12 | -------------------------------------------------------------------------------- /pgvector/asyncpg/register.py: -------------------------------------------------------------------------------- 1 | from .. import Vector, HalfVector, SparseVector 2 | 3 | 4 | async def register_vector(conn, schema='public'): 5 | await conn.set_type_codec( 6 | 'vector', 7 | schema=schema, 8 | encoder=Vector._to_db_binary, 9 | decoder=Vector._from_db_binary, 10 | format='binary' 11 | ) 12 | 13 | try: 14 | await conn.set_type_codec( 15 | 'halfvec', 16 | schema=schema, 17 | encoder=HalfVector._to_db_binary, 18 | decoder=HalfVector._from_db_binary, 19 | format='binary' 20 | ) 21 | 22 | await conn.set_type_codec( 23 | 'sparsevec', 24 | schema=schema, 25 | encoder=SparseVector._to_db_binary, 26 | decoder=SparseVector._from_db_binary, 27 | format='binary' 28 | ) 29 | except ValueError as e: 30 | if not str(e).startswith('unknown type:'): 31 | raise e 32 | -------------------------------------------------------------------------------- /pgvector/bit.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from struct import pack, unpack_from 3 | from warnings import warn 4 | 5 | 6 | class Bit: 7 | def __init__(self, value): 8 | if isinstance(value, bytes): 9 | self._len = 8 * len(value) 10 | self._data = value 11 | else: 12 | if isinstance(value, str): 13 | value = [v != '0' for v in value] 14 | else: 15 | value = np.asarray(value) 16 | 17 | if value.dtype != np.bool: 18 | # skip warning for result of np.unpackbits 19 | if value.dtype != np.uint8 or np.any(value > 1): 20 | warn('expected elements to be boolean', stacklevel=2) 21 | value = value.astype(bool) 22 | 23 | if value.ndim != 1: 24 | raise ValueError('expected ndim to be 1') 25 | 26 | self._len = len(value) 27 | self._data = np.packbits(value).tobytes() 28 | 29 | def __repr__(self): 30 | return f'Bit({self.to_text()})' 31 | 32 | def __eq__(self, other): 33 | if isinstance(other, self.__class__): 34 | return self._len == other._len and self._data == other._data 35 | return False 36 | 37 | def to_list(self): 38 | return self.to_numpy().tolist() 39 | 40 | def to_numpy(self): 41 | return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool) 42 | 43 | def to_text(self): 44 | return ''.join(format(v, '08b') for v in self._data)[:self._len] 45 | 46 | def to_binary(self): 47 | return pack('>i', self._len) + self._data 48 | 49 | @classmethod 50 | def from_text(cls, value): 51 | return cls(str(value)) 52 | 53 | @classmethod 54 | def from_binary(cls, value): 55 | if not isinstance(value, bytes): 56 | raise ValueError('expected bytes') 57 | 58 | bit = cls.__new__(cls) 59 | bit._len = unpack_from('>i', value)[0] 60 | bit._data = value[4:] 61 | return bit 62 | 63 | @classmethod 64 | def _to_db(cls, value): 65 | if not isinstance(value, cls): 66 | raise ValueError('expected bit') 67 | 68 | return value.to_text() 69 | 70 | @classmethod 71 | def _to_db_binary(cls, value): 72 | if not isinstance(value, cls): 73 | raise ValueError('expected bit') 74 | 75 | return value.to_binary() 76 | -------------------------------------------------------------------------------- /pgvector/django/__init__.py: -------------------------------------------------------------------------------- 1 | from .bit import BitField 2 | from .extensions import VectorExtension 3 | from .functions import L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance 4 | from .halfvec import HalfVectorField 5 | from .indexes import IvfflatIndex, HnswIndex 6 | from .sparsevec import SparseVectorField 7 | from .vector import VectorField 8 | 9 | # TODO remove 10 | from .. import HalfVector, SparseVector 11 | 12 | __all__ = [ 13 | 'VectorExtension', 14 | 'VectorField', 15 | 'HalfVectorField', 16 | 'BitField', 17 | 'SparseVectorField', 18 | 'IvfflatIndex', 19 | 'HnswIndex', 20 | 'L2Distance', 21 | 'MaxInnerProduct', 22 | 'CosineDistance', 23 | 'L1Distance', 24 | 'HammingDistance', 25 | 'JaccardDistance', 26 | 'HalfVector', 27 | 'SparseVector' 28 | ] 29 | -------------------------------------------------------------------------------- /pgvector/django/bit.py: -------------------------------------------------------------------------------- 1 | from django import forms 2 | from django.db.models import Field 3 | 4 | 5 | # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ 6 | class BitField(Field): 7 | description = 'Bit string' 8 | 9 | def __init__(self, *args, length=None, **kwargs): 10 | self.length = length 11 | super().__init__(*args, **kwargs) 12 | 13 | def deconstruct(self): 14 | name, path, args, kwargs = super().deconstruct() 15 | if self.length is not None: 16 | kwargs['length'] = self.length 17 | return name, path, args, kwargs 18 | 19 | def db_type(self, connection): 20 | if self.length is None: 21 | return 'bit' 22 | return 'bit(%d)' % self.length 23 | 24 | def formfield(self, **kwargs): 25 | return super().formfield(form_class=BitFormField, **kwargs) 26 | 27 | 28 | class BitFormField(forms.CharField): 29 | def to_python(self, value): 30 | if isinstance(value, str) and value == '': 31 | return None 32 | return super().to_python(value) 33 | -------------------------------------------------------------------------------- /pgvector/django/extensions.py: -------------------------------------------------------------------------------- 1 | from django.contrib.postgres.operations import CreateExtension 2 | 3 | 4 | class VectorExtension(CreateExtension): 5 | def __init__(self): 6 | self.name = 'vector' 7 | -------------------------------------------------------------------------------- /pgvector/django/functions.py: -------------------------------------------------------------------------------- 1 | from django.db.models import FloatField, Func, Value 2 | from .. import Vector, HalfVector, SparseVector 3 | 4 | 5 | class DistanceBase(Func): 6 | output_field = FloatField() 7 | 8 | def __init__(self, expression, vector, **extra): 9 | if not hasattr(vector, 'resolve_expression'): 10 | if isinstance(vector, HalfVector): 11 | vector = Value(HalfVector._to_db(vector)) 12 | elif isinstance(vector, SparseVector): 13 | vector = Value(SparseVector._to_db(vector)) 14 | else: 15 | vector = Value(Vector._to_db(vector)) 16 | 17 | # prevent error with unhashable types 18 | self._constructor_args = ((expression, vector), extra) 19 | 20 | super().__init__(expression, vector, **extra) 21 | 22 | 23 | class BitDistanceBase(Func): 24 | output_field = FloatField() 25 | 26 | def __init__(self, expression, vector, **extra): 27 | if not hasattr(vector, 'resolve_expression'): 28 | vector = Value(vector) 29 | super().__init__(expression, vector, **extra) 30 | 31 | 32 | class L2Distance(DistanceBase): 33 | function = '' 34 | arg_joiner = ' <-> ' 35 | 36 | 37 | class MaxInnerProduct(DistanceBase): 38 | function = '' 39 | arg_joiner = ' <#> ' 40 | 41 | 42 | class CosineDistance(DistanceBase): 43 | function = '' 44 | arg_joiner = ' <=> ' 45 | 46 | 47 | class L1Distance(DistanceBase): 48 | function = '' 49 | arg_joiner = ' <+> ' 50 | 51 | 52 | class HammingDistance(BitDistanceBase): 53 | function = '' 54 | arg_joiner = ' <~> ' 55 | 56 | 57 | class JaccardDistance(BitDistanceBase): 58 | function = '' 59 | arg_joiner = ' <%%> ' 60 | -------------------------------------------------------------------------------- /pgvector/django/halfvec.py: -------------------------------------------------------------------------------- 1 | from django import forms 2 | from django.db.models import Field 3 | from .. import HalfVector 4 | 5 | 6 | # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ 7 | class HalfVectorField(Field): 8 | description = 'Half vector' 9 | empty_strings_allowed = False 10 | 11 | def __init__(self, *args, dimensions=None, **kwargs): 12 | self.dimensions = dimensions 13 | super().__init__(*args, **kwargs) 14 | 15 | def deconstruct(self): 16 | name, path, args, kwargs = super().deconstruct() 17 | if self.dimensions is not None: 18 | kwargs['dimensions'] = self.dimensions 19 | return name, path, args, kwargs 20 | 21 | def db_type(self, connection): 22 | if self.dimensions is None: 23 | return 'halfvec' 24 | return 'halfvec(%d)' % self.dimensions 25 | 26 | def from_db_value(self, value, expression, connection): 27 | return HalfVector._from_db(value) 28 | 29 | def to_python(self, value): 30 | if value is None or isinstance(value, HalfVector): 31 | return value 32 | elif isinstance(value, str): 33 | return HalfVector._from_db(value) 34 | else: 35 | return HalfVector(value) 36 | 37 | def get_prep_value(self, value): 38 | return HalfVector._to_db(value) 39 | 40 | def value_to_string(self, obj): 41 | return self.get_prep_value(self.value_from_object(obj)) 42 | 43 | def formfield(self, **kwargs): 44 | return super().formfield(form_class=HalfVectorFormField, **kwargs) 45 | 46 | 47 | class HalfVectorWidget(forms.TextInput): 48 | def format_value(self, value): 49 | if isinstance(value, HalfVector): 50 | value = value.to_list() 51 | return super().format_value(value) 52 | 53 | 54 | class HalfVectorFormField(forms.CharField): 55 | widget = HalfVectorWidget 56 | 57 | def to_python(self, value): 58 | if isinstance(value, str) and value == '': 59 | return None 60 | return super().to_python(value) 61 | -------------------------------------------------------------------------------- /pgvector/django/indexes.py: -------------------------------------------------------------------------------- 1 | from django.contrib.postgres.indexes import PostgresIndex 2 | 3 | 4 | class IvfflatIndex(PostgresIndex): 5 | suffix = 'ivfflat' 6 | 7 | def __init__(self, *expressions, lists=None, **kwargs): 8 | self.lists = lists 9 | super().__init__(*expressions, **kwargs) 10 | 11 | def deconstruct(self): 12 | path, args, kwargs = super().deconstruct() 13 | if self.lists is not None: 14 | kwargs['lists'] = self.lists 15 | return path, args, kwargs 16 | 17 | def get_with_params(self): 18 | with_params = [] 19 | if self.lists is not None: 20 | with_params.append('lists = %d' % self.lists) 21 | return with_params 22 | 23 | 24 | class HnswIndex(PostgresIndex): 25 | suffix = 'hnsw' 26 | 27 | def __init__(self, *expressions, m=None, ef_construction=None, **kwargs): 28 | self.m = m 29 | self.ef_construction = ef_construction 30 | super().__init__(*expressions, **kwargs) 31 | 32 | def deconstruct(self): 33 | path, args, kwargs = super().deconstruct() 34 | if self.m is not None: 35 | kwargs['m'] = self.m 36 | if self.ef_construction is not None: 37 | kwargs['ef_construction'] = self.ef_construction 38 | return path, args, kwargs 39 | 40 | def get_with_params(self): 41 | with_params = [] 42 | if self.m is not None: 43 | with_params.append('m = %d' % self.m) 44 | if self.ef_construction is not None: 45 | with_params.append('ef_construction = %d' % self.ef_construction) 46 | return with_params 47 | -------------------------------------------------------------------------------- /pgvector/django/sparsevec.py: -------------------------------------------------------------------------------- 1 | from django import forms 2 | from django.db.models import Field 3 | from .. import SparseVector 4 | 5 | 6 | # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ 7 | class SparseVectorField(Field): 8 | description = 'Sparse vector' 9 | empty_strings_allowed = False 10 | 11 | def __init__(self, *args, dimensions=None, **kwargs): 12 | self.dimensions = dimensions 13 | super().__init__(*args, **kwargs) 14 | 15 | def deconstruct(self): 16 | name, path, args, kwargs = super().deconstruct() 17 | if self.dimensions is not None: 18 | kwargs['dimensions'] = self.dimensions 19 | return name, path, args, kwargs 20 | 21 | def db_type(self, connection): 22 | if self.dimensions is None: 23 | return 'sparsevec' 24 | return 'sparsevec(%d)' % self.dimensions 25 | 26 | def from_db_value(self, value, expression, connection): 27 | return SparseVector._from_db(value) 28 | 29 | def to_python(self, value): 30 | return SparseVector._from_db(value) 31 | 32 | def get_prep_value(self, value): 33 | return SparseVector._to_db(value) 34 | 35 | def value_to_string(self, obj): 36 | return self.get_prep_value(self.value_from_object(obj)) 37 | 38 | def formfield(self, **kwargs): 39 | return super().formfield(form_class=SparseVectorFormField, **kwargs) 40 | 41 | 42 | class SparseVectorWidget(forms.TextInput): 43 | def format_value(self, value): 44 | if isinstance(value, SparseVector): 45 | value = value.to_text() 46 | return super().format_value(value) 47 | 48 | 49 | class SparseVectorFormField(forms.CharField): 50 | widget = SparseVectorWidget 51 | 52 | def to_python(self, value): 53 | if isinstance(value, str) and value == '': 54 | return None 55 | return super().to_python(value) 56 | -------------------------------------------------------------------------------- /pgvector/django/vector.py: -------------------------------------------------------------------------------- 1 | from django import forms 2 | from django.db.models import Field 3 | import numpy as np 4 | from .. import Vector 5 | 6 | 7 | # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/ 8 | class VectorField(Field): 9 | description = 'Vector' 10 | empty_strings_allowed = False 11 | 12 | def __init__(self, *args, dimensions=None, **kwargs): 13 | self.dimensions = dimensions 14 | super().__init__(*args, **kwargs) 15 | 16 | def deconstruct(self): 17 | name, path, args, kwargs = super().deconstruct() 18 | if self.dimensions is not None: 19 | kwargs['dimensions'] = self.dimensions 20 | return name, path, args, kwargs 21 | 22 | def db_type(self, connection): 23 | if self.dimensions is None: 24 | return 'vector' 25 | return 'vector(%d)' % self.dimensions 26 | 27 | def from_db_value(self, value, expression, connection): 28 | return Vector._from_db(value) 29 | 30 | def to_python(self, value): 31 | if isinstance(value, list): 32 | return np.array(value, dtype=np.float32) 33 | return Vector._from_db(value) 34 | 35 | def get_prep_value(self, value): 36 | return Vector._to_db(value) 37 | 38 | def value_to_string(self, obj): 39 | return self.get_prep_value(self.value_from_object(obj)) 40 | 41 | def validate(self, value, model_instance): 42 | if isinstance(value, np.ndarray): 43 | value = value.tolist() 44 | super().validate(value, model_instance) 45 | 46 | def run_validators(self, value): 47 | if isinstance(value, np.ndarray): 48 | value = value.tolist() 49 | super().run_validators(value) 50 | 51 | def formfield(self, **kwargs): 52 | return super().formfield(form_class=VectorFormField, **kwargs) 53 | 54 | 55 | class VectorWidget(forms.TextInput): 56 | def format_value(self, value): 57 | if isinstance(value, np.ndarray): 58 | value = value.tolist() 59 | return super().format_value(value) 60 | 61 | 62 | class VectorFormField(forms.CharField): 63 | widget = VectorWidget 64 | 65 | def has_changed(self, initial, data): 66 | if isinstance(initial, np.ndarray): 67 | initial = initial.tolist() 68 | return super().has_changed(initial, data) 69 | 70 | def to_python(self, value): 71 | if isinstance(value, str) and value == '': 72 | return None 73 | return super().to_python(value) 74 | -------------------------------------------------------------------------------- /pgvector/halfvec.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from struct import pack, unpack_from 3 | 4 | 5 | class HalfVector: 6 | def __init__(self, value): 7 | # asarray still copies if same dtype 8 | if not isinstance(value, np.ndarray) or value.dtype != '>f2': 9 | value = np.asarray(value, dtype='>f2') 10 | 11 | if value.ndim != 1: 12 | raise ValueError('expected ndim to be 1') 13 | 14 | self._value = value 15 | 16 | def __repr__(self): 17 | return f'HalfVector({self.to_list()})' 18 | 19 | def __eq__(self, other): 20 | if isinstance(other, self.__class__): 21 | return np.array_equal(self.to_numpy(), other.to_numpy()) 22 | return False 23 | 24 | def dimensions(self): 25 | return len(self._value) 26 | 27 | def to_list(self): 28 | return self._value.tolist() 29 | 30 | def to_numpy(self): 31 | return self._value 32 | 33 | def to_text(self): 34 | return '[' + ','.join([str(float(v)) for v in self._value]) + ']' 35 | 36 | def to_binary(self): 37 | return pack('>HH', self.dimensions(), 0) + self._value.tobytes() 38 | 39 | @classmethod 40 | def from_text(cls, value): 41 | return cls([float(v) for v in value[1:-1].split(',')]) 42 | 43 | @classmethod 44 | def from_binary(cls, value): 45 | dim, unused = unpack_from('>HH', value) 46 | return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4)) 47 | 48 | @classmethod 49 | def _to_db(cls, value, dim=None): 50 | if value is None: 51 | return value 52 | 53 | if not isinstance(value, cls): 54 | value = cls(value) 55 | 56 | if dim is not None and value.dimensions() != dim: 57 | raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) 58 | 59 | return value.to_text() 60 | 61 | @classmethod 62 | def _to_db_binary(cls, value): 63 | if value is None: 64 | return value 65 | 66 | if not isinstance(value, cls): 67 | value = cls(value) 68 | 69 | return value.to_binary() 70 | 71 | @classmethod 72 | def _from_db(cls, value): 73 | if value is None or isinstance(value, cls): 74 | return value 75 | 76 | return cls.from_text(value) 77 | 78 | @classmethod 79 | def _from_db_binary(cls, value): 80 | if value is None or isinstance(value, cls): 81 | return value 82 | 83 | return cls.from_binary(value) 84 | -------------------------------------------------------------------------------- /pgvector/peewee/__init__.py: -------------------------------------------------------------------------------- 1 | from .bit import FixedBitField 2 | from .halfvec import HalfVectorField 3 | from .sparsevec import SparseVectorField 4 | from .vector import VectorField 5 | 6 | # TODO remove 7 | from .. import HalfVector, SparseVector 8 | 9 | __all__ = [ 10 | 'VectorField', 11 | 'HalfVectorField', 12 | 'FixedBitField', 13 | 'SparseVectorField', 14 | 'HalfVector', 15 | 'SparseVector' 16 | ] 17 | -------------------------------------------------------------------------------- /pgvector/peewee/bit.py: -------------------------------------------------------------------------------- 1 | from peewee import Expression, Field 2 | 3 | 4 | class FixedBitField(Field): 5 | field_type = 'bit' 6 | 7 | def __init__(self, max_length=None, *args, **kwargs): 8 | self.max_length = max_length 9 | super(FixedBitField, self).__init__(*args, **kwargs) 10 | 11 | def get_modifiers(self): 12 | return self.max_length and [self.max_length] or None 13 | 14 | def _distance(self, op, vector): 15 | return Expression(lhs=self, op=op, rhs=self.to_value(vector)) 16 | 17 | def hamming_distance(self, vector): 18 | return self._distance('<~>', vector) 19 | 20 | def jaccard_distance(self, vector): 21 | return self._distance('<%%>', vector) 22 | -------------------------------------------------------------------------------- /pgvector/peewee/halfvec.py: -------------------------------------------------------------------------------- 1 | from peewee import Expression, Field 2 | from .. import HalfVector 3 | 4 | 5 | class HalfVectorField(Field): 6 | field_type = 'halfvec' 7 | 8 | def __init__(self, dimensions=None, *args, **kwargs): 9 | self.dimensions = dimensions 10 | super(HalfVectorField, self).__init__(*args, **kwargs) 11 | 12 | def get_modifiers(self): 13 | return self.dimensions and [self.dimensions] or None 14 | 15 | def db_value(self, value): 16 | return HalfVector._to_db(value) 17 | 18 | def python_value(self, value): 19 | return HalfVector._from_db(value) 20 | 21 | def _distance(self, op, vector): 22 | return Expression(lhs=self, op=op, rhs=self.to_value(vector)) 23 | 24 | def l2_distance(self, vector): 25 | return self._distance('<->', vector) 26 | 27 | def max_inner_product(self, vector): 28 | return self._distance('<#>', vector) 29 | 30 | def cosine_distance(self, vector): 31 | return self._distance('<=>', vector) 32 | 33 | def l1_distance(self, vector): 34 | return self._distance('<+>', vector) 35 | -------------------------------------------------------------------------------- /pgvector/peewee/sparsevec.py: -------------------------------------------------------------------------------- 1 | from peewee import Expression, Field 2 | from .. import SparseVector 3 | 4 | 5 | class SparseVectorField(Field): 6 | field_type = 'sparsevec' 7 | 8 | def __init__(self, dimensions=None, *args, **kwargs): 9 | self.dimensions = dimensions 10 | super(SparseVectorField, self).__init__(*args, **kwargs) 11 | 12 | def get_modifiers(self): 13 | return self.dimensions and [self.dimensions] or None 14 | 15 | def db_value(self, value): 16 | return SparseVector._to_db(value) 17 | 18 | def python_value(self, value): 19 | return SparseVector._from_db(value) 20 | 21 | def _distance(self, op, vector): 22 | return Expression(lhs=self, op=op, rhs=self.to_value(vector)) 23 | 24 | def l2_distance(self, vector): 25 | return self._distance('<->', vector) 26 | 27 | def max_inner_product(self, vector): 28 | return self._distance('<#>', vector) 29 | 30 | def cosine_distance(self, vector): 31 | return self._distance('<=>', vector) 32 | 33 | def l1_distance(self, vector): 34 | return self._distance('<+>', vector) 35 | -------------------------------------------------------------------------------- /pgvector/peewee/vector.py: -------------------------------------------------------------------------------- 1 | from peewee import Expression, Field 2 | from .. import Vector 3 | 4 | 5 | class VectorField(Field): 6 | field_type = 'vector' 7 | 8 | def __init__(self, dimensions=None, *args, **kwargs): 9 | self.dimensions = dimensions 10 | super(VectorField, self).__init__(*args, **kwargs) 11 | 12 | def get_modifiers(self): 13 | return self.dimensions and [self.dimensions] or None 14 | 15 | def db_value(self, value): 16 | return Vector._to_db(value) 17 | 18 | def python_value(self, value): 19 | return Vector._from_db(value) 20 | 21 | def _distance(self, op, vector): 22 | return Expression(lhs=self, op=op, rhs=self.to_value(vector)) 23 | 24 | def l2_distance(self, vector): 25 | return self._distance('<->', vector) 26 | 27 | def max_inner_product(self, vector): 28 | return self._distance('<#>', vector) 29 | 30 | def cosine_distance(self, vector): 31 | return self._distance('<=>', vector) 32 | 33 | def l1_distance(self, vector): 34 | return self._distance('<+>', vector) 35 | -------------------------------------------------------------------------------- /pgvector/pg8000/__init__.py: -------------------------------------------------------------------------------- 1 | from .register import register_vector 2 | 3 | __all__ = [ 4 | 'register_vector' 5 | ] 6 | -------------------------------------------------------------------------------- /pgvector/pg8000/register.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .. import Vector, HalfVector, SparseVector 3 | 4 | 5 | def register_vector(conn): 6 | # use to_regtype to get first matching type in search path 7 | res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))") 8 | type_info = dict(res) 9 | 10 | if 'vector' not in type_info: 11 | raise RuntimeError('vector type not found in the database') 12 | 13 | conn.register_out_adapter(Vector, Vector._to_db) 14 | conn.register_out_adapter(np.ndarray, Vector._to_db) 15 | conn.register_in_adapter(type_info['vector'], Vector._from_db) 16 | 17 | if 'halfvec' in type_info: 18 | conn.register_out_adapter(HalfVector, HalfVector._to_db) 19 | conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db) 20 | 21 | if 'sparsevec' in type_info: 22 | conn.register_out_adapter(SparseVector, SparseVector._to_db) 23 | conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db) 24 | -------------------------------------------------------------------------------- /pgvector/psycopg/__init__.py: -------------------------------------------------------------------------------- 1 | from .register import register_vector, register_vector_async 2 | 3 | # TODO remove 4 | from .. import Bit, HalfVector, SparseVector, Vector 5 | 6 | __all__ = [ 7 | 'register_vector', 8 | 'register_vector_async', 9 | 'Vector', 10 | 'HalfVector', 11 | 'Bit', 12 | 'SparseVector' 13 | ] 14 | -------------------------------------------------------------------------------- /pgvector/psycopg/bit.py: -------------------------------------------------------------------------------- 1 | from psycopg.adapt import Dumper 2 | from psycopg.pq import Format 3 | from .. import Bit 4 | 5 | 6 | class BitDumper(Dumper): 7 | 8 | format = Format.TEXT 9 | 10 | def dump(self, obj): 11 | return Bit._to_db(obj).encode('utf8') 12 | 13 | 14 | class BitBinaryDumper(BitDumper): 15 | 16 | format = Format.BINARY 17 | 18 | def dump(self, obj): 19 | return Bit._to_db_binary(obj) 20 | 21 | 22 | def register_bit_info(context, info): 23 | info.register(context) 24 | 25 | # add oid to anonymous class for set_types 26 | text_dumper = type('', (BitDumper,), {'oid': info.oid}) 27 | binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid}) 28 | 29 | adapters = context.adapters 30 | adapters.register_dumper(Bit, text_dumper) 31 | adapters.register_dumper(Bit, binary_dumper) 32 | -------------------------------------------------------------------------------- /pgvector/psycopg/halfvec.py: -------------------------------------------------------------------------------- 1 | from psycopg.adapt import Loader, Dumper 2 | from psycopg.pq import Format 3 | from .. import HalfVector 4 | 5 | 6 | class HalfVectorDumper(Dumper): 7 | 8 | format = Format.TEXT 9 | 10 | def dump(self, obj): 11 | return HalfVector._to_db(obj).encode('utf8') 12 | 13 | 14 | class HalfVectorBinaryDumper(HalfVectorDumper): 15 | 16 | format = Format.BINARY 17 | 18 | def dump(self, obj): 19 | return HalfVector._to_db_binary(obj) 20 | 21 | 22 | class HalfVectorLoader(Loader): 23 | 24 | format = Format.TEXT 25 | 26 | def load(self, data): 27 | if isinstance(data, memoryview): 28 | data = bytes(data) 29 | return HalfVector._from_db(data.decode('utf8')) 30 | 31 | 32 | class HalfVectorBinaryLoader(HalfVectorLoader): 33 | 34 | format = Format.BINARY 35 | 36 | def load(self, data): 37 | if isinstance(data, memoryview): 38 | data = bytes(data) 39 | return HalfVector._from_db_binary(data) 40 | 41 | 42 | def register_halfvec_info(context, info): 43 | info.register(context) 44 | 45 | # add oid to anonymous class for set_types 46 | text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid}) 47 | binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid}) 48 | 49 | adapters = context.adapters 50 | adapters.register_dumper(HalfVector, text_dumper) 51 | adapters.register_dumper(HalfVector, binary_dumper) 52 | adapters.register_loader(info.oid, HalfVectorLoader) 53 | adapters.register_loader(info.oid, HalfVectorBinaryLoader) 54 | -------------------------------------------------------------------------------- /pgvector/psycopg/register.py: -------------------------------------------------------------------------------- 1 | from psycopg.types import TypeInfo 2 | from .bit import register_bit_info 3 | from .halfvec import register_halfvec_info 4 | from .sparsevec import register_sparsevec_info 5 | from .vector import register_vector_info 6 | 7 | 8 | def register_vector(context): 9 | info = TypeInfo.fetch(context, 'vector') 10 | register_vector_info(context, info) 11 | 12 | info = TypeInfo.fetch(context, 'bit') 13 | register_bit_info(context, info) 14 | 15 | info = TypeInfo.fetch(context, 'halfvec') 16 | if info is not None: 17 | register_halfvec_info(context, info) 18 | 19 | info = TypeInfo.fetch(context, 'sparsevec') 20 | if info is not None: 21 | register_sparsevec_info(context, info) 22 | 23 | 24 | async def register_vector_async(context): 25 | info = await TypeInfo.fetch(context, 'vector') 26 | register_vector_info(context, info) 27 | 28 | info = await TypeInfo.fetch(context, 'bit') 29 | register_bit_info(context, info) 30 | 31 | info = await TypeInfo.fetch(context, 'halfvec') 32 | if info is not None: 33 | register_halfvec_info(context, info) 34 | 35 | info = await TypeInfo.fetch(context, 'sparsevec') 36 | if info is not None: 37 | register_sparsevec_info(context, info) 38 | -------------------------------------------------------------------------------- /pgvector/psycopg/sparsevec.py: -------------------------------------------------------------------------------- 1 | from psycopg.adapt import Loader, Dumper 2 | from psycopg.pq import Format 3 | from .. import SparseVector 4 | 5 | 6 | class SparseVectorDumper(Dumper): 7 | 8 | format = Format.TEXT 9 | 10 | def dump(self, obj): 11 | return SparseVector._to_db(obj).encode('utf8') 12 | 13 | 14 | class SparseVectorBinaryDumper(SparseVectorDumper): 15 | 16 | format = Format.BINARY 17 | 18 | def dump(self, obj): 19 | return SparseVector._to_db_binary(obj) 20 | 21 | 22 | class SparseVectorLoader(Loader): 23 | 24 | format = Format.TEXT 25 | 26 | def load(self, data): 27 | if isinstance(data, memoryview): 28 | data = bytes(data) 29 | return SparseVector._from_db(data.decode('utf8')) 30 | 31 | 32 | class SparseVectorBinaryLoader(SparseVectorLoader): 33 | 34 | format = Format.BINARY 35 | 36 | def load(self, data): 37 | if isinstance(data, memoryview): 38 | data = bytes(data) 39 | return SparseVector._from_db_binary(data) 40 | 41 | 42 | def register_sparsevec_info(context, info): 43 | info.register(context) 44 | 45 | # add oid to anonymous class for set_types 46 | text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid}) 47 | binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid}) 48 | 49 | adapters = context.adapters 50 | adapters.register_dumper(SparseVector, text_dumper) 51 | adapters.register_dumper(SparseVector, binary_dumper) 52 | adapters.register_loader(info.oid, SparseVectorLoader) 53 | adapters.register_loader(info.oid, SparseVectorBinaryLoader) 54 | -------------------------------------------------------------------------------- /pgvector/psycopg/vector.py: -------------------------------------------------------------------------------- 1 | import psycopg 2 | from psycopg.adapt import Loader, Dumper 3 | from psycopg.pq import Format 4 | from .. import Vector 5 | 6 | 7 | class VectorDumper(Dumper): 8 | 9 | format = Format.TEXT 10 | 11 | def dump(self, obj): 12 | return Vector._to_db(obj).encode('utf8') 13 | 14 | 15 | class VectorBinaryDumper(VectorDumper): 16 | 17 | format = Format.BINARY 18 | 19 | def dump(self, obj): 20 | return Vector._to_db_binary(obj) 21 | 22 | 23 | class VectorLoader(Loader): 24 | 25 | format = Format.TEXT 26 | 27 | def load(self, data): 28 | if isinstance(data, memoryview): 29 | data = bytes(data) 30 | return Vector._from_db(data.decode('utf8')) 31 | 32 | 33 | class VectorBinaryLoader(VectorLoader): 34 | 35 | format = Format.BINARY 36 | 37 | def load(self, data): 38 | if isinstance(data, memoryview): 39 | data = bytes(data) 40 | return Vector._from_db_binary(data) 41 | 42 | 43 | def register_vector_info(context, info): 44 | if info is None: 45 | raise psycopg.ProgrammingError('vector type not found in the database') 46 | info.register(context) 47 | 48 | # add oid to anonymous class for set_types 49 | text_dumper = type('', (VectorDumper,), {'oid': info.oid}) 50 | binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) 51 | 52 | adapters = context.adapters 53 | adapters.register_dumper('numpy.ndarray', text_dumper) 54 | adapters.register_dumper('numpy.ndarray', binary_dumper) 55 | adapters.register_dumper(Vector, text_dumper) 56 | adapters.register_dumper(Vector, binary_dumper) 57 | adapters.register_loader(info.oid, VectorLoader) 58 | adapters.register_loader(info.oid, VectorBinaryLoader) 59 | -------------------------------------------------------------------------------- /pgvector/psycopg2/__init__.py: -------------------------------------------------------------------------------- 1 | from .register import register_vector 2 | 3 | # TODO remove 4 | from .. import HalfVector, SparseVector 5 | 6 | __all__ = [ 7 | 'register_vector', 8 | 'HalfVector', 9 | 'SparseVector' 10 | ] 11 | -------------------------------------------------------------------------------- /pgvector/psycopg2/halfvec.py: -------------------------------------------------------------------------------- 1 | from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type 2 | from .. import HalfVector 3 | 4 | 5 | class HalfvecAdapter: 6 | def __init__(self, value): 7 | self._value = value 8 | 9 | def getquoted(self): 10 | return adapt(HalfVector._to_db(self._value)).getquoted() 11 | 12 | 13 | def cast_halfvec(value, cur): 14 | return HalfVector._from_db(value) 15 | 16 | 17 | def register_halfvec_info(oid, array_oid, scope): 18 | halfvec = new_type((oid,), 'HALFVEC', cast_halfvec) 19 | register_type(halfvec, scope) 20 | 21 | if array_oid is not None: 22 | halfvecarray = new_array_type((array_oid,), 'HALFVECARRAY', halfvec) 23 | register_type(halfvecarray, scope) 24 | 25 | register_adapter(HalfVector, HalfvecAdapter) 26 | -------------------------------------------------------------------------------- /pgvector/psycopg2/register.py: -------------------------------------------------------------------------------- 1 | import psycopg2 2 | from psycopg2.extensions import cursor 3 | from .halfvec import register_halfvec_info 4 | from .sparsevec import register_sparsevec_info 5 | from .vector import register_vector_info 6 | 7 | 8 | # note: register_adapter is always global 9 | def register_vector(conn_or_curs, globally=False, arrays=True): 10 | conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection 11 | cur = conn.cursor(cursor_factory=cursor) 12 | scope = None if globally else conn_or_curs 13 | 14 | # use to_regtype to get first matching type in search path 15 | cur.execute("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('_vector'), to_regtype('halfvec'), to_regtype('_halfvec'), to_regtype('sparsevec'), to_regtype('_sparsevec'))") 16 | type_info = dict(cur.fetchall()) 17 | 18 | if 'vector' not in type_info: 19 | raise psycopg2.ProgrammingError('vector type not found in the database') 20 | 21 | register_vector_info(type_info['vector'], type_info['_vector'] if arrays else None, scope) 22 | 23 | if 'halfvec' in type_info: 24 | register_halfvec_info(type_info['halfvec'], type_info['_halfvec'] if arrays else None, scope) 25 | 26 | if 'sparsevec' in type_info: 27 | register_sparsevec_info(type_info['sparsevec'], type_info['_sparsevec'] if arrays else None, scope) 28 | -------------------------------------------------------------------------------- /pgvector/psycopg2/sparsevec.py: -------------------------------------------------------------------------------- 1 | from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type 2 | from .. import SparseVector 3 | 4 | 5 | class SparsevecAdapter: 6 | def __init__(self, value): 7 | self._value = value 8 | 9 | def getquoted(self): 10 | return adapt(SparseVector._to_db(self._value)).getquoted() 11 | 12 | 13 | def cast_sparsevec(value, cur): 14 | return SparseVector._from_db(value) 15 | 16 | 17 | def register_sparsevec_info(oid, array_oid, scope): 18 | sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec) 19 | register_type(sparsevec, scope) 20 | 21 | if array_oid is not None: 22 | sparsevecarray = new_array_type((array_oid,), 'SPARSEVECARRAY', sparsevec) 23 | register_type(sparsevecarray, scope) 24 | 25 | register_adapter(SparseVector, SparsevecAdapter) 26 | -------------------------------------------------------------------------------- /pgvector/psycopg2/vector.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type 3 | from .. import Vector 4 | 5 | 6 | class VectorAdapter: 7 | def __init__(self, value): 8 | self._value = value 9 | 10 | def getquoted(self): 11 | return adapt(Vector._to_db(self._value)).getquoted() 12 | 13 | 14 | def cast_vector(value, cur): 15 | return Vector._from_db(value) 16 | 17 | 18 | def register_vector_info(oid, array_oid, scope): 19 | vector = new_type((oid,), 'VECTOR', cast_vector) 20 | register_type(vector, scope) 21 | 22 | if array_oid is not None: 23 | vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector) 24 | register_type(vectorarray, scope) 25 | 26 | register_adapter(np.ndarray, VectorAdapter) 27 | register_adapter(Vector, VectorAdapter) 28 | -------------------------------------------------------------------------------- /pgvector/sparsevec.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from struct import pack, unpack_from 3 | 4 | NO_DEFAULT = object() 5 | 6 | 7 | class SparseVector: 8 | def __init__(self, value, dimensions=NO_DEFAULT, /): 9 | if value.__class__.__module__.startswith('scipy.sparse.'): 10 | if dimensions is not NO_DEFAULT: 11 | raise ValueError('extra argument') 12 | 13 | self._from_sparse(value) 14 | elif isinstance(value, dict): 15 | if dimensions is NO_DEFAULT: 16 | raise ValueError('missing dimensions') 17 | 18 | self._from_dict(value, dimensions) 19 | else: 20 | if dimensions is not NO_DEFAULT: 21 | raise ValueError('extra argument') 22 | 23 | self._from_dense(value) 24 | 25 | def __repr__(self): 26 | elements = dict(zip(self._indices, self._values)) 27 | return f'SparseVector({elements}, {self._dim})' 28 | 29 | def __eq__(self, other): 30 | if isinstance(other, self.__class__): 31 | return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values() 32 | return False 33 | 34 | def dimensions(self): 35 | return self._dim 36 | 37 | def indices(self): 38 | return self._indices 39 | 40 | def values(self): 41 | return self._values 42 | 43 | def to_coo(self): 44 | from scipy.sparse import coo_array 45 | 46 | coords = ([0] * len(self._indices), self._indices) 47 | return coo_array((self._values, coords), shape=(1, self._dim)) 48 | 49 | def to_list(self): 50 | vec = [0.0] * self._dim 51 | for i, v in zip(self._indices, self._values): 52 | vec[i] = v 53 | return vec 54 | 55 | def to_numpy(self): 56 | vec = np.repeat(0.0, self._dim).astype(np.float32) 57 | for i, v in zip(self._indices, self._values): 58 | vec[i] = v 59 | return vec 60 | 61 | def to_text(self): 62 | return '{' + ','.join([f'{int(i) + 1}:{float(v)}' for i, v in zip(self._indices, self._values)]) + '}/' + str(int(self._dim)) 63 | 64 | def to_binary(self): 65 | nnz = len(self._indices) 66 | return pack(f'>iii{nnz}i{nnz}f', self._dim, nnz, 0, *self._indices, *self._values) 67 | 68 | def _from_dict(self, d, dim): 69 | elements = [(i, v) for i, v in d.items() if v != 0] 70 | elements.sort() 71 | 72 | self._dim = int(dim) 73 | self._indices = [int(v[0]) for v in elements] 74 | self._values = [float(v[1]) for v in elements] 75 | 76 | def _from_sparse(self, value): 77 | value = value.tocoo() 78 | 79 | if value.ndim == 1: 80 | self._dim = value.shape[0] 81 | elif value.ndim == 2 and value.shape[0] == 1: 82 | self._dim = value.shape[1] 83 | else: 84 | raise ValueError('expected ndim to be 1') 85 | 86 | if hasattr(value, 'coords'): 87 | # scipy 1.13+ 88 | self._indices = value.coords[-1].tolist() 89 | else: 90 | self._indices = value.col.tolist() 91 | self._values = value.data.tolist() 92 | 93 | def _from_dense(self, value): 94 | self._dim = len(value) 95 | self._indices = [i for i, v in enumerate(value) if v != 0] 96 | self._values = [float(value[i]) for i in self._indices] 97 | 98 | @classmethod 99 | def from_text(cls, value): 100 | elements, dim = value.split('/', 2) 101 | indices = [] 102 | values = [] 103 | # split on empty string returns single element list 104 | if len(elements) > 2: 105 | for e in elements[1:-1].split(','): 106 | i, v = e.split(':', 2) 107 | indices.append(int(i) - 1) 108 | values.append(float(v)) 109 | return cls._from_parts(int(dim), indices, values) 110 | 111 | @classmethod 112 | def from_binary(cls, value): 113 | dim, nnz, unused = unpack_from('>iii', value) 114 | indices = unpack_from(f'>{nnz}i', value, 12) 115 | values = unpack_from(f'>{nnz}f', value, 12 + nnz * 4) 116 | return cls._from_parts(int(dim), list(indices), list(values)) 117 | 118 | @classmethod 119 | def _from_parts(cls, dim, indices, values): 120 | vec = cls.__new__(cls) 121 | vec._dim = dim 122 | vec._indices = indices 123 | vec._values = values 124 | return vec 125 | 126 | @classmethod 127 | def _to_db(cls, value, dim=None): 128 | if value is None: 129 | return value 130 | 131 | if not isinstance(value, cls): 132 | value = cls(value) 133 | 134 | if dim is not None and value.dimensions() != dim: 135 | raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) 136 | 137 | return value.to_text() 138 | 139 | @classmethod 140 | def _to_db_binary(cls, value): 141 | if value is None: 142 | return value 143 | 144 | if not isinstance(value, cls): 145 | value = cls(value) 146 | 147 | return value.to_binary() 148 | 149 | @classmethod 150 | def _from_db(cls, value): 151 | if value is None or isinstance(value, cls): 152 | return value 153 | 154 | return cls.from_text(value) 155 | 156 | @classmethod 157 | def _from_db_binary(cls, value): 158 | if value is None or isinstance(value, cls): 159 | return value 160 | 161 | return cls.from_binary(value) 162 | -------------------------------------------------------------------------------- /pgvector/sqlalchemy/__init__.py: -------------------------------------------------------------------------------- 1 | from .bit import BIT 2 | from .functions import avg, sum 3 | from .halfvec import HALFVEC 4 | from .sparsevec import SPARSEVEC 5 | from .vector import VECTOR 6 | from .vector import VECTOR as Vector 7 | 8 | # TODO remove 9 | from .. import HalfVector, SparseVector 10 | 11 | __all__ = [ 12 | 'Vector', 13 | 'VECTOR', 14 | 'HALFVEC', 15 | 'BIT', 16 | 'SPARSEVEC', 17 | 'HalfVector', 18 | 'SparseVector', 19 | 'avg', 20 | 'sum' 21 | ] 22 | -------------------------------------------------------------------------------- /pgvector/sqlalchemy/bit.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy.dialects.postgresql.base import ischema_names 2 | from sqlalchemy.types import UserDefinedType, Float 3 | 4 | 5 | class BIT(UserDefinedType): 6 | cache_ok = True 7 | 8 | def __init__(self, length=None): 9 | super(UserDefinedType, self).__init__() 10 | self.length = length 11 | 12 | def get_col_spec(self, **kw): 13 | if self.length is None: 14 | return 'BIT' 15 | return 'BIT(%d)' % self.length 16 | 17 | class comparator_factory(UserDefinedType.Comparator): 18 | def hamming_distance(self, other): 19 | return self.op('<~>', return_type=Float)(other) 20 | 21 | def jaccard_distance(self, other): 22 | return self.op('<%>', return_type=Float)(other) 23 | 24 | 25 | # for reflection 26 | ischema_names['bit'] = BIT 27 | -------------------------------------------------------------------------------- /pgvector/sqlalchemy/functions.py: -------------------------------------------------------------------------------- 1 | # https://docs.sqlalchemy.org/en/20/core/functions.html 2 | # include sum for a consistent API 3 | from sqlalchemy.sql.functions import ReturnTypeFromArgs, sum 4 | 5 | 6 | class avg(ReturnTypeFromArgs): 7 | inherit_cache = True 8 | package = 'pgvector' 9 | 10 | 11 | __all__ = [ 12 | 'avg', 13 | 'sum' 14 | ] 15 | -------------------------------------------------------------------------------- /pgvector/sqlalchemy/halfvec.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy.dialects.postgresql.base import ischema_names 2 | from sqlalchemy.types import UserDefinedType, Float, String 3 | from .. import HalfVector 4 | 5 | 6 | class HALFVEC(UserDefinedType): 7 | cache_ok = True 8 | _string = String() 9 | 10 | def __init__(self, dim=None): 11 | super(UserDefinedType, self).__init__() 12 | self.dim = dim 13 | 14 | def get_col_spec(self, **kw): 15 | if self.dim is None: 16 | return 'HALFVEC' 17 | return 'HALFVEC(%d)' % self.dim 18 | 19 | def bind_processor(self, dialect): 20 | def process(value): 21 | return HalfVector._to_db(value, self.dim) 22 | return process 23 | 24 | def literal_processor(self, dialect): 25 | string_literal_processor = self._string._cached_literal_processor(dialect) 26 | 27 | def process(value): 28 | return string_literal_processor(HalfVector._to_db(value, self.dim)) 29 | return process 30 | 31 | def result_processor(self, dialect, coltype): 32 | def process(value): 33 | return HalfVector._from_db(value) 34 | return process 35 | 36 | class comparator_factory(UserDefinedType.Comparator): 37 | def l2_distance(self, other): 38 | return self.op('<->', return_type=Float)(other) 39 | 40 | def max_inner_product(self, other): 41 | return self.op('<#>', return_type=Float)(other) 42 | 43 | def cosine_distance(self, other): 44 | return self.op('<=>', return_type=Float)(other) 45 | 46 | def l1_distance(self, other): 47 | return self.op('<+>', return_type=Float)(other) 48 | 49 | 50 | # for reflection 51 | ischema_names['halfvec'] = HALFVEC 52 | -------------------------------------------------------------------------------- /pgvector/sqlalchemy/sparsevec.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy.dialects.postgresql.base import ischema_names 2 | from sqlalchemy.types import UserDefinedType, Float, String 3 | from .. import SparseVector 4 | 5 | 6 | class SPARSEVEC(UserDefinedType): 7 | cache_ok = True 8 | _string = String() 9 | 10 | def __init__(self, dim=None): 11 | super(UserDefinedType, self).__init__() 12 | self.dim = dim 13 | 14 | def get_col_spec(self, **kw): 15 | if self.dim is None: 16 | return 'SPARSEVEC' 17 | return 'SPARSEVEC(%d)' % self.dim 18 | 19 | def bind_processor(self, dialect): 20 | def process(value): 21 | return SparseVector._to_db(value, self.dim) 22 | return process 23 | 24 | def literal_processor(self, dialect): 25 | string_literal_processor = self._string._cached_literal_processor(dialect) 26 | 27 | def process(value): 28 | return string_literal_processor(SparseVector._to_db(value, self.dim)) 29 | return process 30 | 31 | def result_processor(self, dialect, coltype): 32 | def process(value): 33 | return SparseVector._from_db(value) 34 | return process 35 | 36 | class comparator_factory(UserDefinedType.Comparator): 37 | def l2_distance(self, other): 38 | return self.op('<->', return_type=Float)(other) 39 | 40 | def max_inner_product(self, other): 41 | return self.op('<#>', return_type=Float)(other) 42 | 43 | def cosine_distance(self, other): 44 | return self.op('<=>', return_type=Float)(other) 45 | 46 | def l1_distance(self, other): 47 | return self.op('<+>', return_type=Float)(other) 48 | 49 | 50 | # for reflection 51 | ischema_names['sparsevec'] = SPARSEVEC 52 | -------------------------------------------------------------------------------- /pgvector/sqlalchemy/vector.py: -------------------------------------------------------------------------------- 1 | from sqlalchemy.dialects.postgresql.base import ischema_names 2 | from sqlalchemy.types import UserDefinedType, Float, String 3 | from .. import Vector 4 | 5 | 6 | class VECTOR(UserDefinedType): 7 | cache_ok = True 8 | _string = String() 9 | 10 | def __init__(self, dim=None): 11 | super(UserDefinedType, self).__init__() 12 | self.dim = dim 13 | 14 | def get_col_spec(self, **kw): 15 | if self.dim is None: 16 | return 'VECTOR' 17 | return 'VECTOR(%d)' % self.dim 18 | 19 | def bind_processor(self, dialect): 20 | def process(value): 21 | return Vector._to_db(value, self.dim) 22 | return process 23 | 24 | def literal_processor(self, dialect): 25 | string_literal_processor = self._string._cached_literal_processor(dialect) 26 | 27 | def process(value): 28 | return string_literal_processor(Vector._to_db(value, self.dim)) 29 | return process 30 | 31 | def result_processor(self, dialect, coltype): 32 | def process(value): 33 | return Vector._from_db(value) 34 | return process 35 | 36 | class comparator_factory(UserDefinedType.Comparator): 37 | def l2_distance(self, other): 38 | return self.op('<->', return_type=Float)(other) 39 | 40 | def max_inner_product(self, other): 41 | return self.op('<#>', return_type=Float)(other) 42 | 43 | def cosine_distance(self, other): 44 | return self.op('<=>', return_type=Float)(other) 45 | 46 | def l1_distance(self, other): 47 | return self.op('<+>', return_type=Float)(other) 48 | 49 | 50 | # for reflection 51 | ischema_names['vector'] = VECTOR 52 | -------------------------------------------------------------------------------- /pgvector/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # TODO remove 2 | from .. import Bit, HalfVector, SparseVector, Vector 3 | 4 | __all__ = [ 5 | 'Vector', 6 | 'HalfVector', 7 | 'Bit', 8 | 'SparseVector' 9 | ] 10 | -------------------------------------------------------------------------------- /pgvector/vector.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from struct import pack, unpack_from 3 | 4 | 5 | class Vector: 6 | def __init__(self, value): 7 | # asarray still copies if same dtype 8 | if not isinstance(value, np.ndarray) or value.dtype != '>f4': 9 | value = np.asarray(value, dtype='>f4') 10 | 11 | if value.ndim != 1: 12 | raise ValueError('expected ndim to be 1') 13 | 14 | self._value = value 15 | 16 | def __repr__(self): 17 | return f'Vector({self.to_list()})' 18 | 19 | def __eq__(self, other): 20 | if isinstance(other, self.__class__): 21 | return np.array_equal(self.to_numpy(), other.to_numpy()) 22 | return False 23 | 24 | def dimensions(self): 25 | return len(self._value) 26 | 27 | def to_list(self): 28 | return self._value.tolist() 29 | 30 | def to_numpy(self): 31 | return self._value 32 | 33 | def to_text(self): 34 | return '[' + ','.join([str(float(v)) for v in self._value]) + ']' 35 | 36 | def to_binary(self): 37 | return pack('>HH', self.dimensions(), 0) + self._value.tobytes() 38 | 39 | @classmethod 40 | def from_text(cls, value): 41 | return cls([float(v) for v in value[1:-1].split(',')]) 42 | 43 | @classmethod 44 | def from_binary(cls, value): 45 | dim, unused = unpack_from('>HH', value) 46 | return cls(np.frombuffer(value, dtype='>f4', count=dim, offset=4)) 47 | 48 | @classmethod 49 | def _to_db(cls, value, dim=None): 50 | if value is None: 51 | return value 52 | 53 | if not isinstance(value, cls): 54 | value = cls(value) 55 | 56 | if dim is not None and value.dimensions() != dim: 57 | raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions())) 58 | 59 | return value.to_text() 60 | 61 | @classmethod 62 | def _to_db_binary(cls, value): 63 | if value is None: 64 | return value 65 | 66 | if not isinstance(value, cls): 67 | value = cls(value) 68 | 69 | return value.to_binary() 70 | 71 | @classmethod 72 | def _from_db(cls, value): 73 | if value is None or isinstance(value, np.ndarray): 74 | return value 75 | 76 | return cls.from_text(value).to_numpy().astype(np.float32) 77 | 78 | @classmethod 79 | def _from_db_binary(cls, value): 80 | if value is None or isinstance(value, np.ndarray): 81 | return value 82 | 83 | return cls.from_binary(value).to_numpy().astype(np.float32) 84 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "pgvector" 7 | version = "0.4.1" 8 | description = "pgvector support for Python" 9 | readme = "README.md" 10 | authors = [ 11 | {name = "Andrew Kane", email = "andrew@ankane.org"} 12 | ] 13 | license = {text = "MIT"} 14 | requires-python = ">= 3.9" 15 | dependencies = [ 16 | "numpy" 17 | ] 18 | 19 | [project.urls] 20 | Homepage = "https://github.com/pgvector/pgvector-python" 21 | 22 | [tool.pytest.ini_options] 23 | asyncio_mode = "auto" 24 | asyncio_default_fixture_loop_scope = "function" 25 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | asyncpg 2 | Django 3 | numpy 4 | peewee 5 | pg8000 6 | psycopg[binary,pool] 7 | psycopg2-binary 8 | pytest 9 | pytest-asyncio 10 | scipy 11 | SQLAlchemy[asyncio]>=2 12 | sqlmodel>=0.0.12 13 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pgvector/pgvector-python/7793bb069942fbcc2e77cf7349c59ffc28d8b6e0/tests/__init__.py -------------------------------------------------------------------------------- /tests/test_asyncpg.py: -------------------------------------------------------------------------------- 1 | import asyncpg 2 | import numpy as np 3 | from pgvector import HalfVector, SparseVector, Vector 4 | from pgvector.asyncpg import register_vector 5 | import pytest 6 | 7 | 8 | class TestAsyncpg: 9 | @pytest.mark.asyncio 10 | async def test_vector(self): 11 | conn = await asyncpg.connect(database='pgvector_python_test') 12 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 13 | await conn.execute('DROP TABLE IF EXISTS asyncpg_items') 14 | await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') 15 | 16 | await register_vector(conn) 17 | 18 | embedding = Vector([1.5, 2, 3]) 19 | embedding2 = np.array([4.5, 5, 6]) 20 | await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) 21 | 22 | res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") 23 | assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) 24 | assert res[0]['embedding'].dtype == np.float32 25 | assert np.array_equal(res[1]['embedding'], embedding2) 26 | assert res[2]['embedding'] is None 27 | 28 | # ensures binary format is correct 29 | text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") 30 | assert text_res[0]['embedding'] == '[1.5,2,3]' 31 | 32 | await conn.close() 33 | 34 | @pytest.mark.asyncio 35 | async def test_halfvec(self): 36 | conn = await asyncpg.connect(database='pgvector_python_test') 37 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 38 | await conn.execute('DROP TABLE IF EXISTS asyncpg_items') 39 | await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding halfvec(3))') 40 | 41 | await register_vector(conn) 42 | 43 | embedding = HalfVector([1.5, 2, 3]) 44 | embedding2 = [4.5, 5, 6] 45 | await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) 46 | 47 | res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") 48 | assert res[0]['embedding'] == embedding 49 | assert res[1]['embedding'] == HalfVector(embedding2) 50 | assert res[2]['embedding'] is None 51 | 52 | # ensures binary format is correct 53 | text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") 54 | assert text_res[0]['embedding'] == '[1.5,2,3]' 55 | 56 | await conn.close() 57 | 58 | @pytest.mark.asyncio 59 | async def test_bit(self): 60 | conn = await asyncpg.connect(database='pgvector_python_test') 61 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 62 | await conn.execute('DROP TABLE IF EXISTS asyncpg_items') 63 | await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding bit(3))') 64 | 65 | await register_vector(conn) 66 | 67 | embedding = asyncpg.BitString('101') 68 | await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) 69 | 70 | res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") 71 | assert res[0]['embedding'].as_string() == '101' 72 | assert res[0]['embedding'].to_int() == 5 73 | assert res[1]['embedding'] is None 74 | 75 | # ensures binary format is correct 76 | text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") 77 | assert text_res[0]['embedding'] == '101' 78 | 79 | await conn.close() 80 | 81 | @pytest.mark.asyncio 82 | async def test_sparsevec(self): 83 | conn = await asyncpg.connect(database='pgvector_python_test') 84 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 85 | await conn.execute('DROP TABLE IF EXISTS asyncpg_items') 86 | await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding sparsevec(3))') 87 | 88 | await register_vector(conn) 89 | 90 | embedding = SparseVector([1.5, 2, 3]) 91 | await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding) 92 | 93 | res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") 94 | assert res[0]['embedding'] == embedding 95 | assert res[1]['embedding'] is None 96 | 97 | # ensures binary format is correct 98 | text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1") 99 | assert text_res[0]['embedding'] == '{1:1.5,2:2,3:3}/3' 100 | 101 | await conn.close() 102 | 103 | @pytest.mark.asyncio 104 | async def test_vector_array(self): 105 | conn = await asyncpg.connect(database='pgvector_python_test') 106 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 107 | await conn.execute('DROP TABLE IF EXISTS asyncpg_items') 108 | await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embeddings vector[])') 109 | 110 | await register_vector(conn) 111 | 112 | embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])] 113 | await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings) 114 | 115 | embeddings2 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] 116 | await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1]) 117 | 118 | res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") 119 | assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy()) 120 | assert np.array_equal(res[0]['embeddings'][1], embeddings[1].to_numpy()) 121 | assert np.array_equal(res[1]['embeddings'][0], embeddings2[0]) 122 | assert np.array_equal(res[1]['embeddings'][1], embeddings2[1]) 123 | 124 | await conn.close() 125 | 126 | @pytest.mark.asyncio 127 | async def test_pool(self): 128 | async def init(conn): 129 | await register_vector(conn) 130 | 131 | pool = await asyncpg.create_pool(database='pgvector_python_test', init=init) 132 | 133 | async with pool.acquire() as conn: 134 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 135 | await conn.execute('DROP TABLE IF EXISTS asyncpg_items') 136 | await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))') 137 | 138 | embedding = Vector([1.5, 2, 3]) 139 | embedding2 = np.array([1.5, 2, 3]) 140 | await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2) 141 | 142 | res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id") 143 | assert np.array_equal(res[0]['embedding'], embedding.to_numpy()) 144 | assert res[0]['embedding'].dtype == np.float32 145 | assert np.array_equal(res[1]['embedding'], embedding2) 146 | assert res[2]['embedding'] is None 147 | -------------------------------------------------------------------------------- /tests/test_bit.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import Bit 3 | import pytest 4 | 5 | 6 | class TestBit: 7 | def test_list(self): 8 | assert Bit([True, False, True]).to_list() == [True, False, True] 9 | 10 | def test_list_none(self): 11 | with pytest.warns(UserWarning, match='expected elements to be boolean'): 12 | assert Bit([True, None, True]).to_text() == '101' 13 | 14 | def test_list_int(self): 15 | with pytest.warns(UserWarning, match='expected elements to be boolean'): 16 | assert Bit([254, 7, 0]).to_text() == '110' 17 | 18 | def test_tuple(self): 19 | assert Bit((True, False, True)).to_list() == [True, False, True] 20 | 21 | def test_str(self): 22 | assert Bit('101').to_list() == [True, False, True] 23 | 24 | def test_bytes(self): 25 | assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' 26 | assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' 27 | 28 | def test_ndarray(self): 29 | arr = np.array([True, False, True]) 30 | assert Bit(arr).to_list() == [True, False, True] 31 | assert np.array_equal(Bit(arr).to_numpy(), arr) 32 | 33 | def test_ndarray_unpackbits(self): 34 | arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) 35 | assert Bit(arr).to_text() == '111111100000011100000000' 36 | 37 | def test_ndarray_uint8(self): 38 | arr = np.array([254, 7, 0], dtype=np.uint8) 39 | with pytest.warns(UserWarning, match='expected elements to be boolean'): 40 | assert Bit(arr).to_text() == '110' 41 | 42 | def test_ndarray_uint16(self): 43 | arr = np.array([254, 7, 0], dtype=np.uint16) 44 | with pytest.warns(UserWarning, match='expected elements to be boolean'): 45 | assert Bit(arr).to_text() == '110' 46 | 47 | def test_ndim_two(self): 48 | with pytest.raises(ValueError) as error: 49 | Bit([[True, False], [True, False]]) 50 | assert str(error.value) == 'expected ndim to be 1' 51 | 52 | def test_ndim_zero(self): 53 | with pytest.raises(ValueError) as error: 54 | Bit(True) 55 | assert str(error.value) == 'expected ndim to be 1' 56 | 57 | def test_repr(self): 58 | assert repr(Bit([True, False, True])) == 'Bit(101)' 59 | assert str(Bit([True, False, True])) == 'Bit(101)' 60 | 61 | def test_equality(self): 62 | assert Bit([True, False, True]) == Bit([True, False, True]) 63 | assert Bit([True, False, True]) != Bit([True, False, False]) 64 | -------------------------------------------------------------------------------- /tests/test_django.py: -------------------------------------------------------------------------------- 1 | import django 2 | from django.conf import settings 3 | from django.contrib.postgres.fields import ArrayField 4 | from django.contrib.postgres.indexes import OpClass 5 | from django.core import serializers 6 | from django.db import connection, migrations, models 7 | from django.db.models import Avg, Sum, FloatField, DecimalField 8 | from django.db.models.functions import Cast 9 | from django.db.migrations.loader import MigrationLoader 10 | from django.forms import ModelForm 11 | from math import sqrt 12 | import numpy as np 13 | import os 14 | import pgvector.django 15 | from pgvector import HalfVector, SparseVector 16 | from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance 17 | from unittest import mock 18 | 19 | settings.configure( 20 | DATABASES={ 21 | 'default': { 22 | 'ENGINE': 'django.db.backends.postgresql', 23 | 'NAME': 'pgvector_python_test', 24 | } 25 | }, 26 | DEBUG=('VERBOSE' in os.environ), 27 | LOGGING={ 28 | 'version': 1, 29 | 'handlers': { 30 | 'console': { 31 | 'class': 'logging.StreamHandler' 32 | } 33 | }, 34 | 'loggers': { 35 | 'django.db.backends': { 36 | 'handlers': ['console'], 37 | 'level': 'DEBUG' 38 | }, 39 | 'django.db.backends.schema': { 40 | 'level': 'WARNING' 41 | } 42 | } 43 | }, 44 | # needed for OpClass 45 | # https://docs.djangoproject.com/en/5.1/ref/contrib/postgres/indexes/#opclass-expressions 46 | INSTALLED_APPS=[ 47 | 'django.contrib.postgres' 48 | ] 49 | ) 50 | django.setup() 51 | 52 | 53 | class Item(models.Model): 54 | embedding = VectorField(dimensions=3, null=True, blank=True) 55 | half_embedding = HalfVectorField(dimensions=3, null=True, blank=True) 56 | binary_embedding = BitField(length=3, null=True, blank=True) 57 | sparse_embedding = SparseVectorField(dimensions=3, null=True, blank=True) 58 | embeddings = ArrayField(VectorField(dimensions=3), null=True, blank=True) 59 | double_embedding = ArrayField(FloatField(), null=True, blank=True) 60 | numeric_embedding = ArrayField(DecimalField(max_digits=20, decimal_places=10), null=True, blank=True) 61 | 62 | class Meta: 63 | app_label = 'django_app' 64 | indexes = [ 65 | IvfflatIndex( 66 | name='ivfflat_idx', 67 | fields=['embedding'], 68 | lists=100, 69 | opclasses=['vector_l2_ops'] 70 | ), 71 | HnswIndex( 72 | name='hnsw_idx', 73 | fields=['embedding'], 74 | m=16, 75 | ef_construction=64, 76 | opclasses=['vector_l2_ops'] 77 | ), 78 | HnswIndex( 79 | OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), 80 | name='hnsw_half_precision_idx', 81 | m=16, 82 | ef_construction=64 83 | ) 84 | ] 85 | 86 | 87 | class Migration(migrations.Migration): 88 | initial = True 89 | 90 | operations = [ 91 | VectorExtension(), 92 | migrations.CreateModel( 93 | name='Item', 94 | fields=[ 95 | ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), 96 | ('embedding', pgvector.django.VectorField(dimensions=3, null=True, blank=True)), 97 | ('half_embedding', pgvector.django.HalfVectorField(dimensions=3, null=True, blank=True)), 98 | ('binary_embedding', pgvector.django.BitField(length=3, null=True, blank=True)), 99 | ('sparse_embedding', pgvector.django.SparseVectorField(dimensions=3, null=True, blank=True)), 100 | ('embeddings', ArrayField(pgvector.django.VectorField(dimensions=3), null=True, blank=True)), 101 | ('double_embedding', ArrayField(FloatField(), null=True, blank=True)), 102 | ('numeric_embedding', ArrayField(DecimalField(max_digits=20, decimal_places=10), null=True, blank=True)), 103 | ], 104 | ), 105 | migrations.AddIndex( 106 | model_name='item', 107 | index=pgvector.django.IvfflatIndex(fields=['embedding'], lists=1, name='ivfflat_idx', opclasses=['vector_l2_ops']), 108 | ), 109 | migrations.AddIndex( 110 | model_name='item', 111 | index=pgvector.django.HnswIndex(fields=['embedding'], m=16, ef_construction=64, name='hnsw_idx', opclasses=['vector_l2_ops']), 112 | ), 113 | migrations.AddIndex( 114 | model_name='item', 115 | index=pgvector.django.HnswIndex(OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'), m=16, ef_construction=64, name='hnsw_half_precision_idx'), 116 | ) 117 | ] 118 | 119 | 120 | # probably a better way to do this 121 | migration = Migration('initial', 'django_app') 122 | loader = MigrationLoader(connection, replace_migrations=False) 123 | loader.graph.add_node(('django_app', migration.name), migration) 124 | sql_statements = loader.collect_sql([(migration, False)]) 125 | 126 | with connection.cursor() as cursor: 127 | cursor.execute("DROP TABLE IF EXISTS django_app_item") 128 | cursor.execute('\n'.join(sql_statements)) 129 | 130 | 131 | def create_items(): 132 | Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])).save() 133 | Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])).save() 134 | Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])).save() 135 | 136 | 137 | class VectorForm(ModelForm): 138 | class Meta: 139 | model = Item 140 | fields = ['embedding'] 141 | 142 | 143 | class HalfVectorForm(ModelForm): 144 | class Meta: 145 | model = Item 146 | fields = ['half_embedding'] 147 | 148 | 149 | class BitForm(ModelForm): 150 | class Meta: 151 | model = Item 152 | fields = ['binary_embedding'] 153 | 154 | 155 | class SparseVectorForm(ModelForm): 156 | class Meta: 157 | model = Item 158 | fields = ['sparse_embedding'] 159 | 160 | 161 | class TestDjango: 162 | def setup_method(self): 163 | Item.objects.all().delete() 164 | 165 | def test_vector(self): 166 | Item(id=1, embedding=[1, 2, 3]).save() 167 | item = Item.objects.get(pk=1) 168 | assert np.array_equal(item.embedding, [1, 2, 3]) 169 | assert item.embedding.dtype == np.float32 170 | 171 | def test_vector_l2_distance(self): 172 | create_items() 173 | distance = L2Distance('embedding', [1, 1, 1]) 174 | items = Item.objects.annotate(distance=distance).order_by(distance) 175 | assert [v.id for v in items] == [1, 3, 2] 176 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 177 | 178 | def test_vector_max_inner_product(self): 179 | create_items() 180 | distance = MaxInnerProduct('embedding', [1, 1, 1]) 181 | items = Item.objects.annotate(distance=distance).order_by(distance) 182 | assert [v.id for v in items] == [2, 3, 1] 183 | assert [v.distance for v in items] == [-6, -4, -3] 184 | 185 | def test_vector_cosine_distance(self): 186 | create_items() 187 | distance = CosineDistance('embedding', [1, 1, 1]) 188 | items = Item.objects.annotate(distance=distance).order_by(distance) 189 | assert [v.id for v in items] == [1, 2, 3] 190 | assert [v.distance for v in items] == [0, 0, 0.05719095841793653] 191 | 192 | def test_vector_l1_distance(self): 193 | create_items() 194 | distance = L1Distance('embedding', [1, 1, 1]) 195 | items = Item.objects.annotate(distance=distance).order_by(distance) 196 | assert [v.id for v in items] == [1, 3, 2] 197 | assert [v.distance for v in items] == [0, 1, 3] 198 | 199 | def test_halfvec(self): 200 | Item(id=1, half_embedding=[1, 2, 3]).save() 201 | item = Item.objects.get(pk=1) 202 | assert item.half_embedding == HalfVector([1, 2, 3]) 203 | 204 | def test_halfvec_l2_distance(self): 205 | create_items() 206 | distance = L2Distance('half_embedding', HalfVector([1, 1, 1])) 207 | items = Item.objects.annotate(distance=distance).order_by(distance) 208 | assert [v.id for v in items] == [1, 3, 2] 209 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 210 | 211 | def test_halfvec_max_inner_product(self): 212 | create_items() 213 | distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1])) 214 | items = Item.objects.annotate(distance=distance).order_by(distance) 215 | assert [v.id for v in items] == [2, 3, 1] 216 | assert [v.distance for v in items] == [-6, -4, -3] 217 | 218 | def test_halfvec_cosine_distance(self): 219 | create_items() 220 | distance = CosineDistance('half_embedding', HalfVector([1, 1, 1])) 221 | items = Item.objects.annotate(distance=distance).order_by(distance) 222 | assert [v.id for v in items] == [1, 2, 3] 223 | assert [v.distance for v in items] == [0, 0, 0.05719095841793653] 224 | 225 | def test_halfvec_l1_distance(self): 226 | create_items() 227 | distance = L1Distance('half_embedding', HalfVector([1, 1, 1])) 228 | items = Item.objects.annotate(distance=distance).order_by(distance) 229 | assert [v.id for v in items] == [1, 3, 2] 230 | assert [v.distance for v in items] == [0, 1, 3] 231 | 232 | def test_bit(self): 233 | Item(id=1, binary_embedding='101').save() 234 | item = Item.objects.get(pk=1) 235 | assert item.binary_embedding == '101' 236 | 237 | def test_bit_hamming_distance(self): 238 | create_items() 239 | distance = HammingDistance('binary_embedding', '101') 240 | items = Item.objects.annotate(distance=distance).order_by(distance) 241 | assert [v.id for v in items] == [2, 3, 1] 242 | assert [v.distance for v in items] == [0, 1, 2] 243 | 244 | def test_bit_jaccard_distance(self): 245 | create_items() 246 | distance = JaccardDistance('binary_embedding', '101') 247 | items = Item.objects.annotate(distance=distance).order_by(distance) 248 | assert [v.id for v in items] == [2, 3, 1] 249 | # assert [v.distance for v in items] == [0, 1/3, 1] 250 | 251 | def test_sparsevec(self): 252 | Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save() 253 | item = Item.objects.get(pk=1) 254 | assert item.sparse_embedding == SparseVector([1, 2, 3]) 255 | 256 | def test_sparsevec_l2_distance(self): 257 | create_items() 258 | distance = L2Distance('sparse_embedding', SparseVector([1, 1, 1])) 259 | items = Item.objects.annotate(distance=distance).order_by(distance) 260 | assert [v.id for v in items] == [1, 3, 2] 261 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 262 | 263 | def test_sparsevec_max_inner_product(self): 264 | create_items() 265 | distance = MaxInnerProduct('sparse_embedding', SparseVector([1, 1, 1])) 266 | items = Item.objects.annotate(distance=distance).order_by(distance) 267 | assert [v.id for v in items] == [2, 3, 1] 268 | assert [v.distance for v in items] == [-6, -4, -3] 269 | 270 | def test_sparsevec_cosine_distance(self): 271 | create_items() 272 | distance = CosineDistance('sparse_embedding', SparseVector([1, 1, 1])) 273 | items = Item.objects.annotate(distance=distance).order_by(distance) 274 | assert [v.id for v in items] == [1, 2, 3] 275 | assert [v.distance for v in items] == [0, 0, 0.05719095841793653] 276 | 277 | def test_sparsevec_l1_distance(self): 278 | create_items() 279 | distance = L1Distance('sparse_embedding', SparseVector([1, 1, 1])) 280 | items = Item.objects.annotate(distance=distance).order_by(distance) 281 | assert [v.id for v in items] == [1, 3, 2] 282 | assert [v.distance for v in items] == [0, 1, 3] 283 | 284 | def test_filter(self): 285 | create_items() 286 | distance = L2Distance('embedding', [1, 1, 1]) 287 | items = Item.objects.alias(distance=distance).filter(distance__lt=1) 288 | assert [v.id for v in items] == [1] 289 | 290 | def test_vector_avg(self): 291 | avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] 292 | assert avg is None 293 | Item(embedding=[1, 2, 3]).save() 294 | Item(embedding=[4, 5, 6]).save() 295 | avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg'] 296 | assert np.array_equal(avg, [2.5, 3.5, 4.5]) 297 | 298 | def test_vector_sum(self): 299 | sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] 300 | assert sum is None 301 | Item(embedding=[1, 2, 3]).save() 302 | Item(embedding=[4, 5, 6]).save() 303 | sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum'] 304 | assert np.array_equal(sum, [5, 7, 9]) 305 | 306 | def test_halfvec_avg(self): 307 | avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] 308 | assert avg is None 309 | Item(half_embedding=[1, 2, 3]).save() 310 | Item(half_embedding=[4, 5, 6]).save() 311 | avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg'] 312 | assert avg == HalfVector([2.5, 3.5, 4.5]) 313 | 314 | def test_halfvec_sum(self): 315 | sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] 316 | assert sum is None 317 | Item(half_embedding=[1, 2, 3]).save() 318 | Item(half_embedding=[4, 5, 6]).save() 319 | sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum'] 320 | assert sum == HalfVector([5, 7, 9]) 321 | 322 | def test_serialization(self): 323 | create_items() 324 | items = Item.objects.all() 325 | for format in ['json', 'xml']: 326 | data = serializers.serialize(format, items) 327 | with mock.patch('django.core.serializers.python.apps.get_model') as get_model: 328 | get_model.return_value = Item 329 | for obj in serializers.deserialize(format, data): 330 | obj.save() 331 | 332 | def test_vector_form(self): 333 | form = VectorForm(data={'embedding': '[1, 2, 3]'}) 334 | assert form.is_valid() 335 | assert 'value="[1, 2, 3]"' in form.as_div() 336 | 337 | def test_vector_form_instance(self): 338 | Item(id=1, embedding=[1, 2, 3]).save() 339 | item = Item.objects.get(pk=1) 340 | form = VectorForm(instance=item) 341 | assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() 342 | 343 | def test_vector_form_save(self): 344 | Item(id=1, embedding=[1, 2, 3]).save() 345 | item = Item.objects.get(pk=1) 346 | form = VectorForm(instance=item, data={'embedding': '[4, 5, 6]'}) 347 | assert form.has_changed() 348 | assert form.is_valid() 349 | assert form.save() 350 | assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6]) 351 | 352 | def test_vector_form_save_missing(self): 353 | Item(id=1).save() 354 | item = Item.objects.get(pk=1) 355 | form = VectorForm(instance=item, data={'embedding': ''}) 356 | assert form.is_valid() 357 | assert form.save() 358 | assert Item.objects.get(pk=1).embedding is None 359 | 360 | def test_halfvec_form(self): 361 | form = HalfVectorForm(data={'half_embedding': '[1, 2, 3]'}) 362 | assert form.is_valid() 363 | assert 'value="[1, 2, 3]"' in form.as_div() 364 | 365 | def test_halfvec_form_instance(self): 366 | Item(id=1, half_embedding=[1, 2, 3]).save() 367 | item = Item.objects.get(pk=1) 368 | form = HalfVectorForm(instance=item) 369 | assert 'value="[1.0, 2.0, 3.0]"' in form.as_div() 370 | 371 | def test_halfvec_form_save(self): 372 | Item(id=1, half_embedding=[1, 2, 3]).save() 373 | item = Item.objects.get(pk=1) 374 | form = HalfVectorForm(instance=item, data={'half_embedding': '[4, 5, 6]'}) 375 | assert form.has_changed() 376 | assert form.is_valid() 377 | assert form.save() 378 | assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6]) 379 | 380 | def test_halfvec_form_save_missing(self): 381 | Item(id=1).save() 382 | item = Item.objects.get(pk=1) 383 | form = HalfVectorForm(instance=item, data={'half_embedding': ''}) 384 | assert form.is_valid() 385 | assert form.save() 386 | assert Item.objects.get(pk=1).half_embedding is None 387 | 388 | def test_bit_form(self): 389 | form = BitForm(data={'binary_embedding': '101'}) 390 | assert form.is_valid() 391 | assert 'value="101"' in form.as_div() 392 | 393 | def test_bit_form_instance(self): 394 | Item(id=1, binary_embedding='101').save() 395 | item = Item.objects.get(pk=1) 396 | form = BitForm(instance=item) 397 | assert 'value="101"' in form.as_div() 398 | 399 | def test_bit_form_save(self): 400 | Item(id=1, binary_embedding='101').save() 401 | item = Item.objects.get(pk=1) 402 | form = BitForm(instance=item, data={'binary_embedding': '010'}) 403 | assert form.has_changed() 404 | assert form.is_valid() 405 | assert form.save() 406 | assert '010' == Item.objects.get(pk=1).binary_embedding 407 | 408 | def test_bit_form_save_missing(self): 409 | Item(id=1).save() 410 | item = Item.objects.get(pk=1) 411 | form = BitForm(instance=item, data={'binary_embedding': ''}) 412 | assert form.is_valid() 413 | assert form.save() 414 | assert Item.objects.get(pk=1).binary_embedding is None 415 | 416 | def test_sparsevec_form(self): 417 | form = SparseVectorForm(data={'sparse_embedding': '{1:1,2:2,3:3}/3'}) 418 | assert form.is_valid() 419 | assert 'value="{1:1,2:2,3:3}/3"' in form.as_div() 420 | 421 | def test_sparsevec_form_instance(self): 422 | Item(id=1, sparse_embedding=[1, 2, 3]).save() 423 | item = Item.objects.get(pk=1) 424 | form = SparseVectorForm(instance=item) 425 | # TODO improve 426 | assert 'value="{1:1.0,2:2.0,3:3.0}/3"' in form.as_div() 427 | 428 | def test_sparsevec_form_save(self): 429 | Item(id=1, sparse_embedding=[1, 2, 3]).save() 430 | item = Item.objects.get(pk=1) 431 | form = SparseVectorForm(instance=item, data={'sparse_embedding': '{1:4,2:5,3:6}/3'}) 432 | assert form.has_changed() 433 | assert form.is_valid() 434 | assert form.save() 435 | assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6]) 436 | 437 | def test_sparesevec_form_save_missing(self): 438 | Item(id=1).save() 439 | item = Item.objects.get(pk=1) 440 | form = SparseVectorForm(instance=item, data={'sparse_embedding': ''}) 441 | assert form.is_valid() 442 | assert form.save() 443 | assert Item.objects.get(pk=1).sparse_embedding is None 444 | 445 | def test_clean(self): 446 | item = Item(id=1, embedding=[1, 2, 3], half_embedding=[1, 2, 3], binary_embedding='101', sparse_embedding=SparseVector([1, 2, 3])) 447 | item.full_clean() 448 | 449 | def test_get_or_create(self): 450 | Item.objects.get_or_create(embedding=[1, 2, 3]) 451 | 452 | def test_missing(self): 453 | Item().save() 454 | assert Item.objects.first().embedding is None 455 | assert Item.objects.first().half_embedding is None 456 | assert Item.objects.first().binary_embedding is None 457 | assert Item.objects.first().sparse_embedding is None 458 | 459 | def test_vector_array(self): 460 | Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]).save() 461 | 462 | with connection.cursor() as cursor: 463 | from pgvector.psycopg import register_vector 464 | register_vector(cursor.connection) 465 | 466 | # this fails if the driver does not cast arrays 467 | item = Item.objects.get(pk=1) 468 | assert np.array_equal(item.embeddings[0], [1, 2, 3]) 469 | assert np.array_equal(item.embeddings[1], [4, 5, 6]) 470 | 471 | def test_double_array(self): 472 | Item(id=1, double_embedding=[1, 1, 1]).save() 473 | Item(id=2, double_embedding=[2, 2, 2]).save() 474 | Item(id=3, double_embedding=[1, 1, 2]).save() 475 | distance = L2Distance(Cast('double_embedding', VectorField()), [1, 1, 1]) 476 | items = Item.objects.annotate(distance=distance).order_by(distance) 477 | assert [v.id for v in items] == [1, 3, 2] 478 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 479 | assert items[1].double_embedding == [1, 1, 2] 480 | 481 | def test_numeric_array(self): 482 | Item(id=1, numeric_embedding=[1, 1, 1]).save() 483 | Item(id=2, numeric_embedding=[2, 2, 2]).save() 484 | Item(id=3, numeric_embedding=[1, 1, 2]).save() 485 | distance = L2Distance(Cast('numeric_embedding', VectorField()), [1, 1, 1]) 486 | items = Item.objects.annotate(distance=distance).order_by(distance) 487 | assert [v.id for v in items] == [1, 3, 2] 488 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 489 | assert items[1].numeric_embedding == [1, 1, 2] 490 | 491 | def test_half_precision(self): 492 | create_items() 493 | distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [1, 1, 1]) 494 | items = Item.objects.annotate(distance=distance).order_by(distance) 495 | assert [v.id for v in items] == [1, 3, 2] 496 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 497 | -------------------------------------------------------------------------------- /tests/test_half_vector.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import HalfVector 3 | import pytest 4 | from struct import pack 5 | 6 | 7 | class TestHalfVector: 8 | def test_list(self): 9 | assert HalfVector([1, 2, 3]).to_list() == [1, 2, 3] 10 | 11 | def test_list_str(self): 12 | with pytest.raises(ValueError, match='could not convert string to float'): 13 | HalfVector([1, 'two', 3]) 14 | 15 | def test_tuple(self): 16 | assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3] 17 | 18 | def test_ndarray(self): 19 | arr = np.array([1, 2, 3]) 20 | assert HalfVector(arr).to_list() == [1, 2, 3] 21 | assert HalfVector(arr).to_numpy() is not arr 22 | 23 | def test_ndarray_same_object(self): 24 | arr = np.array([1, 2, 3], dtype='>f2') 25 | assert HalfVector(arr).to_list() == [1, 2, 3] 26 | assert HalfVector(arr).to_numpy() is arr 27 | 28 | def test_ndim_two(self): 29 | with pytest.raises(ValueError) as error: 30 | HalfVector([[1, 2], [3, 4]]) 31 | assert str(error.value) == 'expected ndim to be 1' 32 | 33 | def test_ndim_zero(self): 34 | with pytest.raises(ValueError) as error: 35 | HalfVector(1) 36 | assert str(error.value) == 'expected ndim to be 1' 37 | 38 | def test_repr(self): 39 | assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' 40 | assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' 41 | 42 | def test_equality(self): 43 | assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) 44 | assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) 45 | 46 | def test_dimensions(self): 47 | assert HalfVector([1, 2, 3]).dimensions() == 3 48 | 49 | def test_from_text(self): 50 | vec = HalfVector.from_text('[1.5,2,3]') 51 | assert vec.to_list() == [1.5, 2, 3] 52 | assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) 53 | 54 | def test_from_binary(self): 55 | data = pack('>HH3e', 3, 0, 1.5, 2, 3) 56 | vec = HalfVector.from_binary(data) 57 | assert vec.to_list() == [1.5, 2, 3] 58 | assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) 59 | assert vec.to_binary() == data 60 | -------------------------------------------------------------------------------- /tests/test_peewee.py: -------------------------------------------------------------------------------- 1 | from math import sqrt 2 | import numpy as np 3 | from peewee import Model, PostgresqlDatabase, fn 4 | from pgvector import HalfVector, SparseVector 5 | from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField 6 | 7 | db = PostgresqlDatabase('pgvector_python_test') 8 | 9 | 10 | class BaseModel(Model): 11 | class Meta: 12 | database = db 13 | 14 | 15 | class Item(BaseModel): 16 | embedding = VectorField(dimensions=3, null=True) 17 | half_embedding = HalfVectorField(dimensions=3, null=True) 18 | binary_embedding = FixedBitField(max_length=3, null=True) 19 | sparse_embedding = SparseVectorField(dimensions=3, null=True) 20 | 21 | class Meta: 22 | table_name = 'peewee_item' 23 | 24 | 25 | Item.add_index('embedding vector_l2_ops', using='hnsw') 26 | 27 | db.connect() 28 | db.execute_sql('CREATE EXTENSION IF NOT EXISTS vector') 29 | db.drop_tables([Item]) 30 | db.create_tables([Item]) 31 | 32 | 33 | def create_items(): 34 | Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])) 35 | Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])) 36 | Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])) 37 | 38 | 39 | class TestPeewee: 40 | def setup_method(self): 41 | Item.truncate_table() 42 | 43 | def test_vector(self): 44 | Item.create(id=1, embedding=[1, 2, 3]) 45 | item = Item.get_by_id(1) 46 | assert np.array_equal(item.embedding, [1, 2, 3]) 47 | assert item.embedding.dtype == np.float32 48 | 49 | def test_vector_l2_distance(self): 50 | create_items() 51 | distance = Item.embedding.l2_distance([1, 1, 1]) 52 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 53 | assert [v.id for v in items] == [1, 3, 2] 54 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 55 | 56 | def test_vector_max_inner_product(self): 57 | create_items() 58 | distance = Item.embedding.max_inner_product([1, 1, 1]) 59 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 60 | assert [v.id for v in items] == [2, 3, 1] 61 | assert [v.distance for v in items] == [-6, -4, -3] 62 | 63 | def test_vector_cosine_distance(self): 64 | create_items() 65 | distance = Item.embedding.cosine_distance([1, 1, 1]) 66 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 67 | assert [v.id for v in items] == [1, 2, 3] 68 | assert [v.distance for v in items] == [0, 0, 0.05719095841793653] 69 | 70 | def test_vector_l1_distance(self): 71 | create_items() 72 | distance = Item.embedding.l1_distance([1, 1, 1]) 73 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 74 | assert [v.id for v in items] == [1, 3, 2] 75 | assert [v.distance for v in items] == [0, 1, 3] 76 | 77 | def test_halfvec(self): 78 | Item.create(id=1, half_embedding=[1, 2, 3]) 79 | item = Item.get_by_id(1) 80 | assert item.half_embedding == HalfVector([1, 2, 3]) 81 | 82 | def test_halfvec_l2_distance(self): 83 | create_items() 84 | distance = Item.half_embedding.l2_distance([1, 1, 1]) 85 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 86 | assert [v.id for v in items] == [1, 3, 2] 87 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 88 | 89 | def test_halfvec_max_inner_product(self): 90 | create_items() 91 | distance = Item.half_embedding.max_inner_product([1, 1, 1]) 92 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 93 | assert [v.id for v in items] == [2, 3, 1] 94 | assert [v.distance for v in items] == [-6, -4, -3] 95 | 96 | def test_halfvec_cosine_distance(self): 97 | create_items() 98 | distance = Item.half_embedding.cosine_distance([1, 1, 1]) 99 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 100 | assert [v.id for v in items] == [1, 2, 3] 101 | assert [v.distance for v in items] == [0, 0, 0.05719095841793653] 102 | 103 | def test_halfvec_l1_distance(self): 104 | create_items() 105 | distance = Item.half_embedding.l1_distance([1, 1, 1]) 106 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 107 | assert [v.id for v in items] == [1, 3, 2] 108 | assert [v.distance for v in items] == [0, 1, 3] 109 | 110 | def test_bit(self): 111 | Item.create(id=1, binary_embedding='101') 112 | item = Item.get_by_id(1) 113 | assert item.binary_embedding == '101' 114 | 115 | def test_bit_hamming_distance(self): 116 | create_items() 117 | distance = Item.binary_embedding.hamming_distance('101') 118 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 119 | assert [v.id for v in items] == [2, 3, 1] 120 | assert [v.distance for v in items] == [0, 1, 2] 121 | 122 | def test_bit_jaccard_distance(self): 123 | create_items() 124 | distance = Item.binary_embedding.jaccard_distance('101') 125 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 126 | assert [v.id for v in items] == [2, 3, 1] 127 | # assert [v.distance for v in items] == [0, 1/3, 1] 128 | 129 | def test_sparsevec(self): 130 | Item.create(id=1, sparse_embedding=[1, 2, 3]) 131 | item = Item.get_by_id(1) 132 | assert item.sparse_embedding == SparseVector([1, 2, 3]) 133 | 134 | def test_sparsevec_l2_distance(self): 135 | create_items() 136 | distance = Item.sparse_embedding.l2_distance(SparseVector([1, 1, 1])) 137 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 138 | assert [v.id for v in items] == [1, 3, 2] 139 | assert [v.distance for v in items] == [0, 1, sqrt(3)] 140 | 141 | def test_sparsevec_max_inner_product(self): 142 | create_items() 143 | distance = Item.sparse_embedding.max_inner_product([1, 1, 1]) 144 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 145 | assert [v.id for v in items] == [2, 3, 1] 146 | assert [v.distance for v in items] == [-6, -4, -3] 147 | 148 | def test_sparsevec_cosine_distance(self): 149 | create_items() 150 | distance = Item.sparse_embedding.cosine_distance([1, 1, 1]) 151 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 152 | assert [v.id for v in items] == [1, 2, 3] 153 | assert [v.distance for v in items] == [0, 0, 0.05719095841793653] 154 | 155 | def test_sparsevec_l1_distance(self): 156 | create_items() 157 | distance = Item.sparse_embedding.l1_distance([1, 1, 1]) 158 | items = Item.select(Item.id, distance.alias('distance')).order_by(distance).limit(5) 159 | assert [v.id for v in items] == [1, 3, 2] 160 | assert [v.distance for v in items] == [0, 1, 3] 161 | 162 | def test_where(self): 163 | create_items() 164 | items = Item.select().where(Item.embedding.l2_distance([1, 1, 1]) < 1) 165 | assert [v.id for v in items] == [1] 166 | 167 | def test_vector_avg(self): 168 | avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() 169 | assert avg is None 170 | Item.create(embedding=[1, 2, 3]) 171 | Item.create(embedding=[4, 5, 6]) 172 | avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar() 173 | assert np.array_equal(avg, [2.5, 3.5, 4.5]) 174 | 175 | def test_vector_sum(self): 176 | sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() 177 | assert sum is None 178 | Item.create(embedding=[1, 2, 3]) 179 | Item.create(embedding=[4, 5, 6]) 180 | sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar() 181 | assert np.array_equal(sum, [5, 7, 9]) 182 | 183 | def test_halfvec_avg(self): 184 | avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() 185 | assert avg is None 186 | Item.create(half_embedding=[1, 2, 3]) 187 | Item.create(half_embedding=[4, 5, 6]) 188 | avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar() 189 | assert avg == HalfVector([2.5, 3.5, 4.5]) 190 | 191 | def test_halfvec_sum(self): 192 | sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() 193 | assert sum is None 194 | Item.create(half_embedding=[1, 2, 3]) 195 | Item.create(half_embedding=[4, 5, 6]) 196 | sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar() 197 | assert sum == HalfVector([5, 7, 9]) 198 | 199 | def test_get_or_create(self): 200 | Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]}) 201 | Item.get_or_create(embedding=np.array([4, 5, 6])) 202 | Item.get_or_create(embedding=Item.embedding.to_value([7, 8, 9])) 203 | 204 | def test_vector_array(self): 205 | from playhouse.postgres_ext import PostgresqlExtDatabase, ArrayField 206 | 207 | ext_db = PostgresqlExtDatabase('pgvector_python_test') 208 | 209 | class ExtItem(BaseModel): 210 | embeddings = ArrayField(VectorField, field_kwargs={'dimensions': 3}, index=False) 211 | 212 | class Meta: 213 | database = ext_db 214 | table_name = 'peewee_ext_item' 215 | 216 | ext_db.connect() 217 | ext_db.drop_tables([ExtItem]) 218 | ext_db.create_tables([ExtItem]) 219 | 220 | # fails with column "embeddings" is of type vector[] but expression is of type text[] 221 | # ExtItem.create(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]) 222 | # item = ExtItem.get_by_id(1) 223 | # assert np.array_equal(item.embeddings[0], [1, 2, 3]) 224 | # assert np.array_equal(item.embeddings[1], [4, 5, 6]) 225 | -------------------------------------------------------------------------------- /tests/test_pg8000.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | from pgvector import HalfVector, SparseVector, Vector 4 | from pgvector.pg8000 import register_vector 5 | from pg8000.native import Connection 6 | 7 | conn = Connection(os.environ["USER"], database='pgvector_python_test') 8 | 9 | conn.run('CREATE EXTENSION IF NOT EXISTS vector') 10 | conn.run('DROP TABLE IF EXISTS pg8000_items') 11 | conn.run('CREATE TABLE pg8000_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))') 12 | 13 | register_vector(conn) 14 | 15 | 16 | class TestPg8000: 17 | def setup_method(self): 18 | conn.run('DELETE FROM pg8000_items') 19 | 20 | def test_vector(self): 21 | embedding = np.array([1.5, 2, 3]) 22 | conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) 23 | 24 | res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') 25 | assert np.array_equal(res[0][0], embedding) 26 | assert res[0][0].dtype == np.float32 27 | assert res[1][0] is None 28 | 29 | def test_vector_class(self): 30 | embedding = Vector([1.5, 2, 3]) 31 | conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding) 32 | 33 | res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id') 34 | assert np.array_equal(res[0][0], embedding.to_numpy()) 35 | assert res[0][0].dtype == np.float32 36 | assert res[1][0] is None 37 | 38 | def test_halfvec(self): 39 | embedding = HalfVector([1.5, 2, 3]) 40 | conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (NULL)', embedding=embedding) 41 | 42 | res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id') 43 | assert res[0][0] == embedding 44 | assert res[1][0] is None 45 | 46 | def test_bit(self): 47 | embedding = '101' 48 | conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (NULL)', embedding=embedding) 49 | 50 | res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') 51 | assert res[0][0] == '101' 52 | assert res[1][0] is None 53 | 54 | def test_sparsevec(self): 55 | embedding = SparseVector([1.5, 2, 3]) 56 | conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (NULL)', embedding=embedding) 57 | 58 | res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id') 59 | assert res[0][0] == embedding 60 | assert res[1][0] is None 61 | -------------------------------------------------------------------------------- /tests/test_psycopg.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import Bit, HalfVector, SparseVector, Vector 3 | from pgvector.psycopg import register_vector, register_vector_async 4 | import psycopg 5 | from psycopg_pool import ConnectionPool, AsyncConnectionPool 6 | import pytest 7 | 8 | conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True) 9 | 10 | conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 11 | conn.execute('DROP TABLE IF EXISTS psycopg_items') 12 | conn.execute('CREATE TABLE psycopg_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector[])') 13 | 14 | register_vector(conn) 15 | 16 | 17 | class TestPsycopg: 18 | def setup_method(self): 19 | conn.execute('DELETE FROM psycopg_items') 20 | 21 | def test_vector(self): 22 | embedding = np.array([1.5, 2, 3]) 23 | conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) 24 | 25 | res = conn.execute('SELECT embedding FROM psycopg_items ORDER BY id').fetchall() 26 | assert np.array_equal(res[0][0], embedding) 27 | assert res[0][0].dtype == np.float32 28 | assert res[1][0] is None 29 | 30 | def test_vector_binary_format(self): 31 | embedding = np.array([1.5, 2, 3]) 32 | res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] 33 | assert np.array_equal(res, embedding) 34 | 35 | def test_vector_text_format(self): 36 | embedding = np.array([1.5, 2, 3]) 37 | res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] 38 | assert np.array_equal(res, embedding) 39 | 40 | def test_vector_binary_format_correct(self): 41 | embedding = np.array([1.5, 2, 3]) 42 | res = conn.execute('SELECT %b::vector::text', (embedding,)).fetchone()[0] 43 | assert res == '[1.5,2,3]' 44 | 45 | def test_vector_text_format_non_contiguous(self): 46 | embedding = np.flipud(np.array([1.5, 2, 3])) 47 | assert not embedding.data.contiguous 48 | res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] 49 | assert np.array_equal(res, [3, 2, 1.5]) 50 | 51 | def test_vector_binary_format_non_contiguous(self): 52 | embedding = np.flipud(np.array([1.5, 2, 3])) 53 | assert not embedding.data.contiguous 54 | res = conn.execute('SELECT %b::vector', (embedding,)).fetchone()[0] 55 | assert np.array_equal(res, [3, 2, 1.5]) 56 | 57 | def test_vector_class_binary_format(self): 58 | embedding = Vector([1.5, 2, 3]) 59 | res = conn.execute('SELECT %b::vector', (embedding,), binary=True).fetchone()[0] 60 | assert np.array_equal(res, [1.5, 2, 3]) 61 | 62 | def test_vector_class_text_format(self): 63 | embedding = Vector([1.5, 2, 3]) 64 | res = conn.execute('SELECT %t::vector', (embedding,)).fetchone()[0] 65 | assert np.array_equal(res, [1.5, 2, 3]) 66 | 67 | def test_halfvec(self): 68 | embedding = HalfVector([1.5, 2, 3]) 69 | conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,)) 70 | 71 | res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0] 72 | assert res == HalfVector([1.5, 2, 3]) 73 | 74 | def test_halfvec_binary_format(self): 75 | embedding = HalfVector([1.5, 2, 3]) 76 | res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0] 77 | assert res == HalfVector([1.5, 2, 3]) 78 | 79 | def test_halfvec_text_format(self): 80 | embedding = HalfVector([1.5, 2, 3]) 81 | res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0] 82 | assert res == HalfVector([1.5, 2, 3]) 83 | 84 | def test_bit(self): 85 | embedding = Bit([True, False, True]) 86 | conn.execute('INSERT INTO psycopg_items (binary_embedding) VALUES (%s)', (embedding,)) 87 | 88 | res = conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id').fetchone()[0] 89 | assert res == '101' 90 | 91 | def test_bit_binary_format(self): 92 | embedding = Bit([False, True, False, True, False, False, False, False, True]) 93 | res = conn.execute('SELECT %b::bit(9)', (embedding,), binary=True).fetchone()[0] 94 | assert repr(Bit.from_binary(res)) == 'Bit(010100001)' 95 | 96 | def test_bit_text_format(self): 97 | embedding = Bit([False, True, False, True, False, False, False, False, True]) 98 | res = conn.execute('SELECT %t::bit(9)', (embedding,)).fetchone()[0] 99 | assert res == '010100001' 100 | assert repr(Bit(res)) == 'Bit(010100001)' 101 | 102 | def test_sparsevec(self): 103 | embedding = SparseVector([1.5, 2, 3]) 104 | conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,)) 105 | 106 | res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0] 107 | assert res == SparseVector([1.5, 2, 3]) 108 | 109 | def test_sparsevec_binary_format(self): 110 | embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) 111 | res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0] 112 | assert res == embedding 113 | 114 | def test_sparsevec_text_format(self): 115 | embedding = SparseVector([1.5, 0, 2, 0, 3, 0]) 116 | res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0] 117 | assert res == embedding 118 | 119 | def test_text_copy_from(self): 120 | embedding = np.array([1.5, 2, 3]) 121 | cur = conn.cursor() 122 | with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN") as copy: 123 | copy.write_row([embedding, HalfVector(embedding), '101', SparseVector(embedding)]) 124 | 125 | def test_binary_copy_from(self): 126 | embedding = np.array([1.5, 2, 3]) 127 | cur = conn.cursor() 128 | with cur.copy("COPY psycopg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: 129 | copy.write_row([embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) 130 | 131 | def test_binary_copy_from_set_types(self): 132 | embedding = np.array([1.5, 2, 3]) 133 | cur = conn.cursor() 134 | with cur.copy("COPY psycopg_items (id, embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)") as copy: 135 | copy.set_types(['int8', 'vector', 'halfvec', 'bit', 'sparsevec']) 136 | copy.write_row([1, embedding, HalfVector(embedding), Bit('101'), SparseVector(embedding)]) 137 | 138 | def test_text_copy_to(self): 139 | embedding = np.array([1.5, 2, 3]) 140 | half_embedding = HalfVector([1.5, 2, 3]) 141 | conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) 142 | cur = conn.cursor() 143 | with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT") as copy: 144 | for row in copy.rows(): 145 | assert row[0] == "[1.5,2,3]" 146 | assert row[1] == "[1.5,2,3]" 147 | 148 | def test_binary_copy_to(self): 149 | embedding = np.array([1.5, 2, 3]) 150 | half_embedding = HalfVector([1.5, 2, 3]) 151 | conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) 152 | cur = conn.cursor() 153 | with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: 154 | for row in copy.rows(): 155 | assert np.array_equal(Vector.from_binary(row[0]).to_numpy(), embedding) 156 | assert HalfVector.from_binary(row[1]) == half_embedding 157 | 158 | def test_binary_copy_to_set_types(self): 159 | embedding = np.array([1.5, 2, 3]) 160 | half_embedding = HalfVector([1.5, 2, 3]) 161 | conn.execute('INSERT INTO psycopg_items (embedding, half_embedding) VALUES (%s, %s)', (embedding, half_embedding)) 162 | cur = conn.cursor() 163 | with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy: 164 | copy.set_types(['vector', 'halfvec']) 165 | for row in copy.rows(): 166 | assert np.array_equal(row[0], embedding) 167 | assert row[1] == half_embedding 168 | 169 | def test_vector_array(self): 170 | embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] 171 | conn.execute('INSERT INTO psycopg_items (embeddings) VALUES (%s)', (embeddings,)) 172 | 173 | res = conn.execute('SELECT embeddings FROM psycopg_items ORDER BY id').fetchone() 174 | assert np.array_equal(res[0][0], embeddings[0]) 175 | assert np.array_equal(res[0][1], embeddings[1]) 176 | 177 | def test_pool(self): 178 | def configure(conn): 179 | register_vector(conn) 180 | 181 | pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure) 182 | 183 | with pool.connection() as conn: 184 | res = conn.execute("SELECT '[1,2,3]'::vector").fetchone() 185 | assert np.array_equal(res[0], [1, 2, 3]) 186 | 187 | pool.close() 188 | 189 | @pytest.mark.asyncio 190 | async def test_async(self): 191 | conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True) 192 | 193 | await conn.execute('CREATE EXTENSION IF NOT EXISTS vector') 194 | await conn.execute('DROP TABLE IF EXISTS psycopg_items') 195 | await conn.execute('CREATE TABLE psycopg_items (id bigserial PRIMARY KEY, embedding vector(3))') 196 | 197 | await register_vector_async(conn) 198 | 199 | embedding = np.array([1.5, 2, 3]) 200 | await conn.execute('INSERT INTO psycopg_items (embedding) VALUES (%s), (NULL)', (embedding,)) 201 | 202 | async with conn.cursor() as cur: 203 | await cur.execute('SELECT * FROM psycopg_items ORDER BY id') 204 | res = await cur.fetchall() 205 | assert np.array_equal(res[0][1], embedding) 206 | assert res[0][1].dtype == np.float32 207 | assert res[1][1] is None 208 | 209 | @pytest.mark.asyncio 210 | async def test_async_pool(self): 211 | async def configure(conn): 212 | await register_vector_async(conn) 213 | 214 | pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure) 215 | await pool.open() 216 | 217 | async with pool.connection() as conn: 218 | async with conn.cursor() as cur: 219 | await cur.execute("SELECT '[1,2,3]'::vector") 220 | res = await cur.fetchone() 221 | assert np.array_equal(res[0], [1, 2, 3]) 222 | 223 | await pool.close() 224 | -------------------------------------------------------------------------------- /tests/test_psycopg2.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import HalfVector, SparseVector, Vector 3 | from pgvector.psycopg2 import register_vector 4 | import psycopg2 5 | from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor 6 | from psycopg2.pool import ThreadedConnectionPool 7 | 8 | conn = psycopg2.connect(dbname='pgvector_python_test') 9 | conn.autocommit = True 10 | 11 | cur = conn.cursor() 12 | cur.execute('CREATE EXTENSION IF NOT EXISTS vector') 13 | cur.execute('DROP TABLE IF EXISTS psycopg2_items') 14 | cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector[], half_embeddings halfvec[], sparse_embeddings sparsevec[])') 15 | 16 | register_vector(cur) 17 | 18 | 19 | class TestPsycopg2: 20 | def setup_method(self): 21 | cur.execute('DELETE FROM psycopg2_items') 22 | 23 | def test_vector(self): 24 | embedding = np.array([1.5, 2, 3]) 25 | cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) 26 | 27 | cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') 28 | res = cur.fetchall() 29 | assert np.array_equal(res[0][0], embedding) 30 | assert res[0][0].dtype == np.float32 31 | assert res[1][0] is None 32 | 33 | def test_vector_class(self): 34 | embedding = Vector([1.5, 2, 3]) 35 | cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,)) 36 | 37 | cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id') 38 | res = cur.fetchall() 39 | assert np.array_equal(res[0][0], embedding.to_numpy()) 40 | assert res[0][0].dtype == np.float32 41 | assert res[1][0] is None 42 | 43 | def test_halfvec(self): 44 | embedding = [1.5, 2, 3] 45 | cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) 46 | 47 | cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') 48 | res = cur.fetchall() 49 | assert res[0][0] == HalfVector([1.5, 2, 3]) 50 | assert res[1][0] is None 51 | 52 | def test_halfvec_class(self): 53 | embedding = HalfVector([1.5, 2, 3]) 54 | cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,)) 55 | 56 | cur.execute('SELECT half_embedding FROM psycopg2_items ORDER BY id') 57 | res = cur.fetchall() 58 | assert res[0][0] == embedding 59 | assert res[1][0] is None 60 | 61 | def test_bit(self): 62 | embedding = '101' 63 | cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) 64 | 65 | cur.execute('SELECT binary_embedding FROM psycopg2_items ORDER BY id') 66 | res = cur.fetchall() 67 | assert res[0][0] == '101' 68 | assert res[1][0] is None 69 | 70 | def test_sparsevec(self): 71 | embedding = SparseVector([1.5, 2, 3]) 72 | cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,)) 73 | 74 | cur.execute('SELECT sparse_embedding FROM psycopg2_items ORDER BY id') 75 | res = cur.fetchall() 76 | assert res[0][0] == SparseVector([1.5, 2, 3]) 77 | assert res[1][0] is None 78 | 79 | def test_vector_array(self): 80 | embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])] 81 | cur.execute('INSERT INTO psycopg2_items (embeddings) VALUES (%s::vector[])', (embeddings,)) 82 | 83 | cur.execute('SELECT embeddings FROM psycopg2_items ORDER BY id') 84 | res = cur.fetchone() 85 | assert np.array_equal(res[0][0], embeddings[0]) 86 | assert np.array_equal(res[0][1], embeddings[1]) 87 | 88 | def test_halfvec_array(self): 89 | embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] 90 | cur.execute('INSERT INTO psycopg2_items (half_embeddings) VALUES (%s::halfvec[])', (embeddings,)) 91 | 92 | cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id') 93 | res = cur.fetchone() 94 | assert res[0] == [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])] 95 | 96 | def test_sparsevec_array(self): 97 | embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] 98 | cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,)) 99 | 100 | cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id') 101 | res = cur.fetchone() 102 | assert res[0] == [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])] 103 | 104 | def test_cursor_factory(self): 105 | for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: 106 | conn = psycopg2.connect(dbname='pgvector_python_test') 107 | cur = conn.cursor(cursor_factory=cursor_factory) 108 | register_vector(cur) 109 | conn.close() 110 | 111 | def test_cursor_factory_connection(self): 112 | for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]: 113 | conn = psycopg2.connect(dbname='pgvector_python_test', cursor_factory=cursor_factory) 114 | register_vector(conn) 115 | conn.close() 116 | 117 | def test_pool(self): 118 | pool = ThreadedConnectionPool(1, 1, dbname='pgvector_python_test') 119 | 120 | conn = pool.getconn() 121 | try: 122 | # use globally=True for apps to ensure registered with all connections 123 | register_vector(conn) 124 | finally: 125 | pool.putconn(conn) 126 | 127 | conn = pool.getconn() 128 | try: 129 | cur = conn.cursor() 130 | cur.execute("SELECT '[1,2,3]'::vector") 131 | res = cur.fetchone() 132 | assert np.array_equal(res[0], [1, 2, 3]) 133 | finally: 134 | pool.putconn(conn) 135 | 136 | pool.closeall() 137 | -------------------------------------------------------------------------------- /tests/test_sparse_vector.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import SparseVector 3 | import pytest 4 | from scipy.sparse import coo_array, coo_matrix, csr_array, csr_matrix 5 | from struct import pack 6 | 7 | 8 | class TestSparseVector: 9 | def test_list(self): 10 | vec = SparseVector([1, 0, 2, 0, 3, 0]) 11 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 12 | assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) 13 | assert vec.indices() == [0, 2, 4] 14 | 15 | def test_list_dimensions(self): 16 | with pytest.raises(ValueError) as error: 17 | SparseVector([1, 0, 2, 0, 3, 0], 6) 18 | assert str(error.value) == 'extra argument' 19 | 20 | def test_ndarray(self): 21 | vec = SparseVector(np.array([1, 0, 2, 0, 3, 0])) 22 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 23 | assert vec.indices() == [0, 2, 4] 24 | 25 | def test_dict(self): 26 | vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) 27 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 28 | assert vec.indices() == [0, 2, 4] 29 | 30 | def test_dict_no_dimensions(self): 31 | with pytest.raises(ValueError) as error: 32 | SparseVector({0: 1, 2: 2, 4: 3}) 33 | assert str(error.value) == 'missing dimensions' 34 | 35 | def test_coo_array(self): 36 | arr = coo_array(np.array([1, 0, 2, 0, 3, 0])) 37 | vec = SparseVector(arr) 38 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 39 | assert vec.indices() == [0, 2, 4] 40 | 41 | def test_coo_array_dimensions(self): 42 | with pytest.raises(ValueError) as error: 43 | SparseVector(coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) 44 | assert str(error.value) == 'extra argument' 45 | 46 | def test_coo_matrix(self): 47 | mat = coo_matrix(np.array([1, 0, 2, 0, 3, 0])) 48 | vec = SparseVector(mat) 49 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 50 | assert vec.indices() == [0, 2, 4] 51 | 52 | def test_dok_array(self): 53 | arr = coo_array(np.array([1, 0, 2, 0, 3, 0])).todok() 54 | vec = SparseVector(arr) 55 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 56 | assert vec.indices() == [0, 2, 4] 57 | 58 | def test_csr_array(self): 59 | arr = csr_array(np.array([[1, 0, 2, 0, 3, 0]])) 60 | vec = SparseVector(arr) 61 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 62 | assert vec.indices() == [0, 2, 4] 63 | 64 | def test_csr_matrix(self): 65 | mat = csr_matrix(np.array([1, 0, 2, 0, 3, 0])) 66 | vec = SparseVector(mat) 67 | assert vec.to_list() == [1, 0, 2, 0, 3, 0] 68 | assert vec.indices() == [0, 2, 4] 69 | 70 | def test_repr(self): 71 | assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' 72 | assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' 73 | 74 | def test_equality(self): 75 | assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) 76 | assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) 77 | assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) 78 | assert SparseVector({}, 1) != SparseVector({}, 2) 79 | 80 | def test_dimensions(self): 81 | assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6 82 | 83 | def test_indices(self): 84 | assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4] 85 | 86 | def test_values(self): 87 | assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3] 88 | 89 | def test_to_coo(self): 90 | assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) 91 | 92 | def test_zero_vector_text(self): 93 | vec = SparseVector({}, 3) 94 | assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() 95 | 96 | def test_from_text(self): 97 | vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') 98 | assert vec.dimensions() == 6 99 | assert vec.indices() == [0, 2, 4] 100 | assert vec.values() == [1.5, 2, 3] 101 | assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] 102 | assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) 103 | 104 | def test_from_binary(self): 105 | data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) 106 | vec = SparseVector.from_binary(data) 107 | assert vec.dimensions() == 6 108 | assert vec.indices() == [0, 2, 4] 109 | assert vec.values() == [1.5, 2, 3] 110 | assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] 111 | assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) 112 | assert vec.to_binary() == data 113 | -------------------------------------------------------------------------------- /tests/test_sqlmodel.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import HalfVector, SparseVector 3 | from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum 4 | import pytest 5 | from sqlalchemy.exc import StatementError 6 | from sqlmodel import Field, Index, Session, SQLModel, create_engine, delete, select, text 7 | from typing import Any, Optional 8 | 9 | engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') 10 | with Session(engine) as session: 11 | session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) 12 | 13 | 14 | class Item(SQLModel, table=True): 15 | __tablename__ = 'sqlmodel_item' 16 | 17 | id: Optional[int] = Field(default=None, primary_key=True) 18 | embedding: Optional[Any] = Field(default=None, sa_type=VECTOR(3)) 19 | half_embedding: Optional[Any] = Field(default=None, sa_type=HALFVEC(3)) 20 | binary_embedding: Optional[Any] = Field(default=None, sa_type=BIT(3)) 21 | sparse_embedding: Optional[Any] = Field(default=None, sa_type=SPARSEVEC(3)) 22 | 23 | 24 | SQLModel.metadata.drop_all(engine) 25 | SQLModel.metadata.create_all(engine) 26 | 27 | index = Index( 28 | 'sqlmodel_index', 29 | Item.embedding, 30 | postgresql_using='hnsw', 31 | postgresql_with={'m': 16, 'ef_construction': 64}, 32 | postgresql_ops={'embedding': 'vector_l2_ops'} 33 | ) 34 | index.create(engine) 35 | 36 | 37 | def create_items(): 38 | with Session(engine) as session: 39 | session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))) 40 | session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))) 41 | session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))) 42 | session.commit() 43 | 44 | 45 | class TestSqlmodel: 46 | def setup_method(self): 47 | with Session(engine) as session: 48 | session.exec(delete(Item)) 49 | session.commit() 50 | 51 | def test_orm(self): 52 | item = Item(embedding=[1.5, 2, 3]) 53 | item2 = Item(embedding=[4, 5, 6]) 54 | item3 = Item() 55 | 56 | with Session(engine) as session: 57 | session.add(item) 58 | session.add(item2) 59 | session.add(item3) 60 | session.commit() 61 | 62 | stmt = select(Item) 63 | with Session(engine) as session: 64 | items = session.exec(stmt).all() 65 | assert items[0].id == 1 66 | assert items[1].id == 2 67 | assert items[2].id == 3 68 | assert np.array_equal(items[0].embedding, np.array([1.5, 2, 3])) 69 | assert items[0].embedding.dtype == np.float32 70 | assert np.array_equal(items[1].embedding, np.array([4, 5, 6])) 71 | assert items[1].embedding.dtype == np.float32 72 | assert items[2].embedding is None 73 | 74 | def test_vector(self): 75 | with Session(engine) as session: 76 | session.add(Item(id=1, embedding=[1, 2, 3])) 77 | session.commit() 78 | item = session.get(Item, 1) 79 | assert np.array_equal(item.embedding, np.array([1, 2, 3])) 80 | 81 | def test_vector_l2_distance(self): 82 | create_items() 83 | with Session(engine) as session: 84 | items = session.exec(select(Item).order_by(Item.embedding.l2_distance([1, 1, 1]))) 85 | assert [v.id for v in items] == [1, 3, 2] 86 | 87 | def test_vector_max_inner_product(self): 88 | create_items() 89 | with Session(engine) as session: 90 | items = session.exec(select(Item).order_by(Item.embedding.max_inner_product([1, 1, 1]))) 91 | assert [v.id for v in items] == [2, 3, 1] 92 | 93 | def test_vector_cosine_distance(self): 94 | create_items() 95 | with Session(engine) as session: 96 | items = session.exec(select(Item).order_by(Item.embedding.cosine_distance([1, 1, 1]))) 97 | assert [v.id for v in items] == [1, 2, 3] 98 | 99 | def test_vector_l1_distance(self): 100 | create_items() 101 | with Session(engine) as session: 102 | items = session.exec(select(Item).order_by(Item.embedding.l1_distance([1, 1, 1]))) 103 | assert [v.id for v in items] == [1, 3, 2] 104 | 105 | def test_halfvec(self): 106 | with Session(engine) as session: 107 | session.add(Item(id=1, half_embedding=[1, 2, 3])) 108 | session.commit() 109 | item = session.get(Item, 1) 110 | assert item.half_embedding == HalfVector([1, 2, 3]) 111 | 112 | def test_halfvec_l2_distance(self): 113 | create_items() 114 | with Session(engine) as session: 115 | items = session.exec(select(Item).order_by(Item.half_embedding.l2_distance([1, 1, 1]))) 116 | assert [v.id for v in items] == [1, 3, 2] 117 | 118 | def test_halfvec_max_inner_product(self): 119 | create_items() 120 | with Session(engine) as session: 121 | items = session.exec(select(Item).order_by(Item.half_embedding.max_inner_product([1, 1, 1]))) 122 | assert [v.id for v in items] == [2, 3, 1] 123 | 124 | def test_halfvec_cosine_distance(self): 125 | create_items() 126 | with Session(engine) as session: 127 | items = session.exec(select(Item).order_by(Item.half_embedding.cosine_distance([1, 1, 1]))) 128 | assert [v.id for v in items] == [1, 2, 3] 129 | 130 | def test_halfvec_l1_distance(self): 131 | create_items() 132 | with Session(engine) as session: 133 | items = session.exec(select(Item).order_by(Item.half_embedding.l1_distance([1, 1, 1]))) 134 | assert [v.id for v in items] == [1, 3, 2] 135 | 136 | def test_bit(self): 137 | with Session(engine) as session: 138 | session.add(Item(id=1, binary_embedding='101')) 139 | session.commit() 140 | item = session.get(Item, 1) 141 | assert item.binary_embedding == '101' 142 | 143 | def test_bit_hamming_distance(self): 144 | create_items() 145 | with Session(engine) as session: 146 | items = session.exec(select(Item).order_by(Item.binary_embedding.hamming_distance('101'))) 147 | assert [v.id for v in items] == [2, 3, 1] 148 | 149 | def test_bit_jaccard_distance(self): 150 | create_items() 151 | with Session(engine) as session: 152 | items = session.exec(select(Item).order_by(Item.binary_embedding.jaccard_distance('101'))) 153 | assert [v.id for v in items] == [2, 3, 1] 154 | 155 | def test_sparsevec(self): 156 | with Session(engine) as session: 157 | session.add(Item(id=1, sparse_embedding=[1, 2, 3])) 158 | session.commit() 159 | item = session.get(Item, 1) 160 | assert item.sparse_embedding == SparseVector([1, 2, 3]) 161 | 162 | def test_sparsevec_l2_distance(self): 163 | create_items() 164 | with Session(engine) as session: 165 | items = session.exec(select(Item).order_by(Item.sparse_embedding.l2_distance([1, 1, 1]))) 166 | assert [v.id for v in items] == [1, 3, 2] 167 | 168 | def test_sparsevec_max_inner_product(self): 169 | create_items() 170 | with Session(engine) as session: 171 | items = session.exec(select(Item).order_by(Item.sparse_embedding.max_inner_product([1, 1, 1]))) 172 | assert [v.id for v in items] == [2, 3, 1] 173 | 174 | def test_sparsevec_cosine_distance(self): 175 | create_items() 176 | with Session(engine) as session: 177 | items = session.exec(select(Item).order_by(Item.sparse_embedding.cosine_distance([1, 1, 1]))) 178 | assert [v.id for v in items] == [1, 2, 3] 179 | 180 | def test_sparsevec_l1_distance(self): 181 | create_items() 182 | with Session(engine) as session: 183 | items = session.exec(select(Item).order_by(Item.sparse_embedding.l1_distance([1, 1, 1]))) 184 | assert [v.id for v in items] == [1, 3, 2] 185 | 186 | def test_filter(self): 187 | create_items() 188 | with Session(engine) as session: 189 | items = session.exec(select(Item).filter(Item.embedding.l2_distance([1, 1, 1]) < 1)) 190 | assert [v.id for v in items] == [1] 191 | 192 | def test_select(self): 193 | with Session(engine) as session: 194 | session.add(Item(embedding=[2, 3, 3])) 195 | items = session.exec(select(Item.embedding.l2_distance([1, 1, 1]))).all() 196 | assert items[0] == 3 197 | 198 | def test_vector_avg(self): 199 | with Session(engine) as session: 200 | res = session.exec(select(avg(Item.embedding))).first() 201 | assert res is None 202 | session.add(Item(embedding=[1, 2, 3])) 203 | session.add(Item(embedding=[4, 5, 6])) 204 | res = session.exec(select(avg(Item.embedding))).first() 205 | assert np.array_equal(res, np.array([2.5, 3.5, 4.5])) 206 | 207 | def test_vector_sum(self): 208 | with Session(engine) as session: 209 | res = session.exec(select(sum(Item.embedding))).first() 210 | assert res is None 211 | session.add(Item(embedding=[1, 2, 3])) 212 | session.add(Item(embedding=[4, 5, 6])) 213 | res = session.exec(select(sum(Item.embedding))).first() 214 | assert np.array_equal(res, np.array([5, 7, 9])) 215 | 216 | def test_halfvec_avg(self): 217 | with Session(engine) as session: 218 | res = session.exec(select(avg(Item.half_embedding))).first() 219 | assert res is None 220 | session.add(Item(half_embedding=[1, 2, 3])) 221 | session.add(Item(half_embedding=[4, 5, 6])) 222 | res = session.exec(select(avg(Item.half_embedding))).first() 223 | assert res == HalfVector([2.5, 3.5, 4.5]) 224 | 225 | def test_halfvec_sum(self): 226 | with Session(engine) as session: 227 | res = session.exec(select(sum(Item.half_embedding))).first() 228 | assert res is None 229 | session.add(Item(half_embedding=[1, 2, 3])) 230 | session.add(Item(half_embedding=[4, 5, 6])) 231 | res = session.exec(select(sum(Item.half_embedding))).first() 232 | assert res == HalfVector([5, 7, 9]) 233 | 234 | def test_bad_dimensions(self): 235 | item = Item(embedding=[1, 2]) 236 | with Session(engine) as session: 237 | session.add(item) 238 | with pytest.raises(StatementError, match='expected 3 dimensions, not 2'): 239 | session.commit() 240 | -------------------------------------------------------------------------------- /tests/test_vector.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from pgvector import Vector 3 | import pytest 4 | from struct import pack 5 | 6 | 7 | class TestVector: 8 | def test_list(self): 9 | assert Vector([1, 2, 3]).to_list() == [1, 2, 3] 10 | 11 | def test_list_str(self): 12 | with pytest.raises(ValueError, match='could not convert string to float'): 13 | Vector([1, 'two', 3]) 14 | 15 | def test_tuple(self): 16 | assert Vector((1, 2, 3)).to_list() == [1, 2, 3] 17 | 18 | def test_ndarray(self): 19 | arr = np.array([1, 2, 3]) 20 | assert Vector(arr).to_list() == [1, 2, 3] 21 | assert Vector(arr).to_numpy() is not arr 22 | 23 | def test_ndarray_same_object(self): 24 | arr = np.array([1, 2, 3], dtype='>f4') 25 | assert Vector(arr).to_list() == [1, 2, 3] 26 | assert Vector(arr).to_numpy() is arr 27 | 28 | def test_ndim_two(self): 29 | with pytest.raises(ValueError) as error: 30 | Vector([[1, 2], [3, 4]]) 31 | assert str(error.value) == 'expected ndim to be 1' 32 | 33 | def test_ndim_zero(self): 34 | with pytest.raises(ValueError) as error: 35 | Vector(1) 36 | assert str(error.value) == 'expected ndim to be 1' 37 | 38 | def test_repr(self): 39 | assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' 40 | assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' 41 | 42 | def test_equality(self): 43 | assert Vector([1, 2, 3]) == Vector([1, 2, 3]) 44 | assert Vector([1, 2, 3]) != Vector([1, 2, 4]) 45 | 46 | def test_dimensions(self): 47 | assert Vector([1, 2, 3]).dimensions() == 3 48 | 49 | def test_from_text(self): 50 | vec = Vector.from_text('[1.5,2,3]') 51 | assert vec.to_list() == [1.5, 2, 3] 52 | assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) 53 | 54 | def test_from_binary(self): 55 | data = pack('>HH3f', 3, 0, 1.5, 2, 3) 56 | vec = Vector.from_binary(data) 57 | assert vec.to_list() == [1.5, 2, 3] 58 | assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) 59 | assert vec.to_binary() == data 60 | --------------------------------------------------------------------------------