repository·master·Indexed 23 days ago
https://github.com/pgvector/pgvector-pythonPython support for the pgvector extension in PostgreSQL, enabling vector similarity search. Version 0.5.0 provides integration for various database libraries including Django, SQLAlchemy, SQLModel, Psycopg (2 and 3), asyncpg, pg8000, and Peewee ORM. It includes support for dense vectors (Vector), half-precision vectors (HalfVector), binary vectors (Bit), and sparse vectors (SparseVector), along with distance functions and approximate indexing (HNSW and IVFFlat).
When using SQL directly (Psycopg, asyncpg, pg8000), you can add approximate indices using hnsw or ivfflat. You can specify different operator classes for different distance metrics:
vector_l2_ops: L2 distancevector_ip_ops: Inner productvector_cosine_ops: Cosine distanceTo use pgvector with asyncpg, enable the extension and register types using register_vector. For connection pools, use the init argument in create_pool.
# Enable the extension
await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
# Register the types
from pgvector.asyncpg import register_vector
await register_vector(conn)
# For connection pools
async def init(conn):
await register_vector(conn)
pool = await asyncpg.create_pool(..., init=init)
# Create a table
await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
# Insert a vector
from pgvector import Vector
embedding = Vector([1, 2, 3])
await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding)
# Get nearest neighbors
await conn.fetch('SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5', embedding)For pg8000, enable the extension using conn.run, register types with register_vector, and use named parameters for vector insertion.
# Enable the extension
conn.run('CREATE EXTENSION IF NOT EXISTS vector')
# Register the types
from pgvector.pg8000 import register_vector
register_vector(conn)
# Create a table
conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
# Insert a vector
from pgvector import Vector
embedding = Vector([1, 2, 3])
conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding)
# Get nearest neighbors
conn.run('SELECT * FROM items ORDER BY embedding <-> :embedding LIMIT 5', embedding=embedding)To use pgvector with Psycopg 3, enable the extension, register the vector types with your connection, and use the Vector class for insertions. For connection pools, use the configure argument. For asynchronous connections, use register_vector_async.
# Enable the extension
conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
# Register the types
from pgvector.psycopg import register_vector
register_vector(conn)
# For connection pools
def configure(conn):
register_vector(conn)
pool = ConnectionPool(..., configure=configure)
# For async connections
from pgvector.psycopg import register_vector_async
await register_vector_async(conn)
# Create a table
conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
# Insert a vector
from pgvector import Vector
embedding = Vector([1, 2, 3])
conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
# Get nearest neighbors
conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall()Install the core package using pip:
pip install pgvectorAfter installation, you must follow specific instructions for your database library (e.g., Django, SQLAlchemy, SQLModel, Psycopg, etc.) to enable vector support.
SQLModel support is built on top of the SQLAlchemy implementation.
session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))Use sa_type=VECTOR(n) within the Field definition.
from pgvector.sqlalchemy import VECTOR
from sqlmodel import SQLModel, Field
class Item(SQLModel, table=True):
embedding: list[float] = Field(sa_type=VECTOR(3))Use the same distance methods available in SQLAlchemy (l2_distance, cosine_distance, etc.) via session.exec(select(...)).
from pgvector.sqlalchemy import VECTOR
from sqlmodel import SQLModel, Field
class Item(SQLModel, table=True):
embedding: list[float] = Field(sa_type=VECTOR(3))
# Querying
session.exec(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5))For Psycopg 2, enable the extension via a cursor, register types with the connection or cursor, and use the Vector class for data operations.
# Enable the extension
cur = conn.cursor()
cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
# Register the types
from pgvector.psycopg2 import register_vector
register_vector(conn)
# Create a table
cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
# Insert a vector
from pgvector import Vector
embedding = Vector([1, 2, 3])
cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
# Get nearest neighbors
cur.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,))
cur.fetchall()To use pgvector with SQLAlchemy, follow these steps:
Execute the SQL command to create the extension:
session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))Use the VECTOR type to add vector columns. Supported types include VECTOR, HALFVEC, BIT, and SPARSEVEC.
from pgvector.sqlalchemy import VECTOR
class Item(Base):
embedding: Mapped[list[float]] = mapped_column(VECTOR(3)).l2_distance(), .max_inner_product(), .cosine_distance(), .l1_distance(), .hamming_distance(), or .jaccard_distance() on the column object.pgvector.sqlalchemy.avg or sum.Create Index objects specifying postgresql_using='hnsw' or postgresql_using='ivfflat'. Specify operators via postgresql_ops (e.g., vector_l2_ops).
func.cast(Item.embedding, HALFVEC(3)) for indexing and querying.func.binary_quantize with BIT type for Hamming distance searches.ARRAY(VECTOR(n)) for columns containing multiple vectors.from pgvector.sqlalchemy import VECTOR
class Item(Base):
embedding: Mapped[list[float]] = mapped_column(VECTOR(3))
# Querying nearest neighbors
session.scalars(select(Item).order_by(Item.embedding.l2_distance([3, 1, 2])).limit(5))To use pgvector with Django, follow these steps:
Create a migration to enable the vector extension:
from pgvector.django import VectorExtension
class Migration(migrations.Migration):
operations = [
VectorExtension()
]Use VectorField to add vector columns. Supported fields include VectorField, HalfVectorField, BitField, and SparseVectorField.
from pgvector.django import VectorField
class Item(models.Model):
embedding = VectorField(dimensions=3)L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, or JaccardDistance in order_by..annotate() to get distances or .alias().filter() to filter by distance.Avg and Sum on vector fields.Add approximate indexes using HnswIndex or IvfflatIndex. Use vector_l2_ops for L2 distance, vector_ip_ops for inner product, and vector_cosine_ops for cosine distance.
from pgvector.django import VectorField
class Item(models.Model):
embedding = VectorField(dimensions=3)
# Querying nearest neighbors
from pgvector.django import L2Distance
Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5]pgvector with Peewee by using VectorField. It supports various distance operators like l2_distance, max_inner_product, cosine_distance, l1_distance, hamming_distance, and jaccard_distance. You can also perform aggregations like avg or sum on vector fields.HalfVector class allows you to create half-precision vector objects from Python lists or NumPy arrays.SparseVector class supports creating sparse vectors from lists, NumPy arrays, SciPy sparse arrays, or dictionaries of non-zero elements. It provides methods to access dimensions, indices, and values, and to convert back to lists, NumPy arrays, or SciPy COO arrays.