pgvector
pgvector support for Python
Description
pgvector-python
pgvector support for Python
Supports Django, SQLAlchemy, SQLModel, Psycopg 3, Psycopg 2, asyncpg, pg8000, and Peewee
Installation
Run:
pip install pgvector
And follow the instructions for your database library:
Or check out some examples:
- Retrieval-augmented generation with Ollama
- Embeddings with OpenAI
- Binary embeddings with Cohere
- Sentence embeddings with SentenceTransformers
- Hybrid search with SentenceTransformers (Reciprocal Rank Fusion)
- Hybrid search with SentenceTransformers (cross-encoder)
- Sparse search with Transformers
- Late interaction search with ColBERT
- Visual document retrieval with ColPali
- Image search with PyTorch
- Image search with perceptual hashing
- Morgan fingerprints with RDKit
- Topic modeling with Gensim
- Implicit feedback recommendations with Implicit
- Explicit feedback recommendations with Surprise
- Recommendations with LightFM
- Horizontal scaling with Citus
- Bulk loading with
COPY
Django
Create a migration to enable the extension
from pgvector.django import VectorExtension
class Migration(migrations.Migration):
operations = [
VectorExtension()
]
Add a vector field to your model
from pgvector.django import VectorField
class Item(models.Model):
embedding = VectorField(dimensions=3)
Also supports HalfVectorField, BitField, and SparseVectorField
Insert a vector
item = Item(embedding=[1, 2, 3])
item.save()
Get the nearest neighbors to a vector
from pgvector.django import L2Distance
Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5]
Also supports MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, and JaccardDistance
Get the distance
Item.objects.annotate(distance=L2Distance('embedding', [3, 1, 2]))
Get items within a certain distance
Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance__lt=5)
Average vectors
from django.db.models import Avg
Item.objects.aggregate(Avg('embedding'))
Also supports Sum
Add an approximate index
from pgvector.django import HnswIndex, IvfflatIndex
class Item(models.Model):
class Meta:
indexes = [
HnswIndex(
name='my_index',
fields=['embedding'],
m=16,
ef_construction=64,
opclasses=['vector_l2_ops']
),
# or
IvfflatIndex(
name='my_index',
fields=['embedding'],
lists=100,
opclasses=['vector_l2_ops']
)
]
Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance
Half-Precision Indexing
Index vectors at half-precision
from django.contrib.postgres.indexes import OpClass
from django.db.models.functions import Cast
from pgvector.django import HnswIndex, HalfVectorField
class Item(models.Model):
class Meta:
indexes = [
HnswIndex(
OpClass(Cast('embedding', HalfVectorField(dimensions=3)), name='halfvec_l2_ops'),
name='my_index',
m=16,
ef_construction=64
)
]
Note: Add 'django.contrib.postgres' to INSTALLED_APPS to use OpClass
Get the nearest neighbors
distance = L2Distance(Cast('embedding', HalfVectorField(dimensions=3)), [3, 1, 2])
Item.objects.order_by(distance)[:5]
SQLAlchemy
Enable the extension
session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
Add a vector column
from pgvector.sqlalchemy import Vector
class Item(Base):
embedding = mapped_column(Vector(3))
Also supports HALFVEC, BIT, and SPARSEVEC
Insert a vector
item = Item(embedding=[1, 2, 3])
session.add(item)
session.commit()
Get the nearest neighbors to a vector
session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5))
Also supports max_inner_product, cosine_distance, l1_distance, hamming_distance, and jaccard_distance
Get the distance
session.scalars(select(Item.embedding.l2_distance([3, 1, 2])))
Get items within a certain distance
session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5))
Average vectors
from pgvector.sqlalchemy import avg
session.scalars(select(avg(Item.embedding))).first()
Also supports sum
Add an approximate index
index = Index(
'my_index',
Item.embedding,
postgresql_using='hnsw',
postgresql_with={'m': 16, 'ef_construction': 64},
postgresql_ops={'embedding': 'vector_l2_ops'}
)
# or
index = Index(
'my_index',
Item.embedding,
postgresql_using='ivfflat',
postgresql_with={'lists': 100},
postgresql_ops={'embedding': 'vector_l2_ops'}
)
index.create(engine)
Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance
Half-Precision Indexing
Index vectors at half-precision
from pgvector.sqlalchemy import HALFVEC
from sqlalchemy.sql import func
index = Index(
'my_index',
func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
postgresql_using='hnsw',
postgresql_ops={'embedding': 'halfvec_l2_ops'}
)
Get the nearest neighbors
order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
session.scalars(select(Item).order_by(order).limit(5))
Binary Quantization
Use expression indexing for binary quantization
from pgvector.sqlalchemy import BIT
from sqlalchemy.sql import func
index = Index(
'my_index',
func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
postgresql_using='hnsw',
postgresql_ops={'embedding': 'bit_hamming_ops'}
)
Get the nearest neighbors by Hamming distance
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
session.scalars(select(Item).order_by(order).limit(5))
Re-rank by the original vectors for better recall
order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
subquery = session.query(Item).order_by(order).limit(20).subquery()
session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
Arrays
Add an array column
from pgvector.sqlalchemy import Vector
from sqlalchemy import ARRAY
class Item(Base):
embeddings = mapped_column(ARRAY(Vector(3)))
And register the types with the underlying driver
For Psycopg 3, use
from pgvector.psycopg import register_vector
from sqlalchemy import event
@event.listens_for(engine, "connect")
def connect(dbapi_connection, connection_record):
register_vector(dbapi_connection)
For async connections with Psycopg 3, use
from pgvector.psycopg import register_vector_async
from sqlalchemy import event
@event.listens_for(engine.sync_engine, "connect")
def connect(dbapi_connection, connection_record):
dbapi_connection.run_async(register_vector_async)
For Psycopg 2, use
from pgvector.psycopg2 import register_vector
from sqlalchemy import event
@event.listens_for(engine, "connect")
def connect(dbapi_connection, connection_record):
register_vector(dbapi_connection, arrays=True)
SQLModel
Enable the extension
session.exec(text('CREATE EXTENSI