pgvector support for Python
Supports SQLAlchemy and Psycopg 2
Run:
pip install pgvectorAnd follow the instructions for your database library:
Use the vector type
from pgvector.sqlalchemy import Vector
# Core
item_table = Table(
'item',
metadata,
Column('factors', Vector(3))
)
# ORM
class Item(Base):
factors = Column(Vector(3))Add an index
index = Index(
'my_index',
item_table.c.factors, # or Item.factors for ORM
postgresql_using='ivfflat',
postgresql_with={'lists': 100}
)
index.create(engine)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
Register the vector type with your connection or cursor
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()View the changelog
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/ankane/pgvector-python.git
cd pgvector-python
pip install -r requirements.txt
pytest