GitHub - maparent/pgvector-python: pgvector support for Python

pgvector support for Python

Great for online recommendations 🎉

Supports Django, SQLAlchemy, Psycopg 3, Psycopg 2, and asyncpg

Installation

Run:

And follow the instructions for your database library:

Or check out some examples:

Image search with PyTorch
Implicit feedback recommendations with Implicit
Explicit feedback recommendations with Surprise
Recommendations with LightFM

Django

Create the extension

from pgvector.django import VectorExtension

class Migration(migrations.Migration):
    operations = [
        VectorExtension()
    ]

Add a vector field

from pgvector.django import VectorField

class Item(models.Model):
    factors = VectorField(dimensions=3)

Insert a vector

item = Item(factors=[1, 2, 3])
item.save()

Get the nearest neighbors to a vector

from pgvector.django import L2Distance

Item.objects.order_by(L2Distance('factors', [3, 1, 2]))[:5]

Also supports MaxInnerProduct and CosineDistance

Add an approximate index

from pgvector.django import IvfflatIndex

class Item(models.Model):
    class Meta:
        indexes = [
            IvfflatIndex(
                name='my_index',
                fields=['factors'],
                lists=100,
                opclasses=['vector_l2_ops']
            )
        ]

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

SQLAlchemy

Add a vector column

from pgvector.sqlalchemy import Vector

class Item(Base):
    factors = Column(Vector(3))

Insert a vector

item = Item(factors=[1, 2, 3])
session.add(item)
session.commit()

Get the nearest neighbors to a vector

session.query(Item).order_by(Item.factors.l2_distance([3, 1, 2])).limit(5).all()

Also supports max_inner_product and cosine_distance

Add an approximate index

index = Index('my_index', Item.factors,
    postgresql_using='ivfflat',
    postgresql_with={'lists': 100},
    postgresql_ops={'factors': 'vector_l2_ops'}
)
index.create(engine)

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

Psycopg 3

from pgvector.psycopg import register_vector

register_vector(conn)

Insert a vector

factors = np.array([1, 2, 3])
conn.execute('INSERT INTO item (factors) VALUES (%s)', (factors,))

Get the nearest neighbors to a vector

conn.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,)).fetchall()

Psycopg 2

from pgvector.psycopg2 import register_vector

register_vector(conn)

Insert a vector

factors = np.array([1, 2, 3])
cur.execute('INSERT INTO item (factors) VALUES (%s)', (factors,))

Get the nearest neighbors to a vector

cur.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,))
cur.fetchall()

asyncpg

from pgvector.asyncpg import register_vector

await register_vector(conn)

Insert a vector

factors = np.array([1, 2, 3])
await conn.execute('INSERT INTO item (factors) VALUES ($1)', factors)

Get the nearest neighbors to a vector

await conn.fetch('SELECT * FROM item ORDER BY factors <-> $1 LIMIT 5', factors)

History

View the changelog

Contributing

Everyone is encouraged to help improve this project. Here are a few ways you can help:

Report bugs
Fix bugs and submit pull requests
Write, clarify, or fix documentation
Suggest or add new features

To get started with development:

git clone https://github.com/pgvector/pgvector-python.git
cd pgvector-python
pip install -r requirements.txt
createdb pgvector_python_test
pytest