pgvector.rs的Python绑定
项目描述
PGVecto.rs支持Python
PGVecto.rs Python库,支持Django、SQLAlchemy和Psycopg 3。
向量 | 稀疏向量 | 半精度向量 | 二进制向量 | |
---|---|---|---|---|
SQLAlchemy | ✅插入 | ✅插入 | ✅插入 | ✅插入 |
Psycopg3 | ✅插入 ✅复制 | ✅插入 ✅复制 | ✅插入 ✅复制 | ✅插入 ✅复制 |
Django | ✅插入 | ✅插入 | ✅插入 | ✅插入 |
用法
从PyPI安装
pip install pgvecto_rs
并使用它与您的数据库库
或作为独立的SDK
需求
要初始化pgvecto.rs实例,您可以运行我们的官方镜像 快速开始
您可以从 发布页面 获取最新的标签。例如,它可能是
docker run \
--name pgvecto-rs-demo \
-e POSTGRES_PASSWORD=mysecretpassword \
-p 5432:5432 \
-d tensorchord/pgvecto-rs:pg16-v0.3.0
SQLAlchemy
安装依赖项
pip install "pgvecto_rs[sqlalchemy]"
初始化连接
from sqlalchemy import create_engine
from sqlalchemy.orm import Session
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"
engine = create_engine(URL)
with Session(engine) as session:
pass
启用扩展
from sqlalchemy import text
session.execute(text('CREATE EXTENSION IF NOT EXISTS vectors'))
创建模型
from pgvecto_rs.sqlalchemy import Vector
class Item(Base):
embedding = mapped_column(Vector(3))
所有支持的类型都在此表中显示
原生类型 | SQLAlchemy的类型 | 对应于pgvector-python |
---|---|---|
vector | VECTOR | VECTOR |
svector | SVECTOR | SPARSEVEC |
vecf16 | VECF16 | HALFVEC |
bvector | BVECTOR | BIT |
插入向量
from sqlalchemy import insert
stmt = insert(Item).values(embedding=[1, 2, 3])
session.execute(stmt)
session.commit()
添加近似索引
from sqlalchemy import Index
from pgvecto_rs.types import IndexOption, Hnsw, Ivf
index = Index(
"emb_idx_1",
Item.embedding,
postgresql_using="vectors",
postgresql_with={
"options": f"$${IndexOption(index=Ivf(), threads=1).dumps()}$$"
},
postgresql_ops={"embedding": "vector_l2_ops"},
)
# or
index = Index(
"emb_idx_2",
Item.embedding,
postgresql_using="vectors",
postgresql_with={
"options": f"$${IndexOption(index=Hnsw()).dumps()}$$"
},
postgresql_ops={"embedding": "vector_l2_ops"},
)
# Apply changes
index.create(session.bind)
获取向量的最近邻
from sqlalchemy import select
session.scalars(select(Item.embedding).order_by(Item.embedding.l2_distance(target.embedding)))
也支持 max_inner_product
、cosine_distance
和 jaccard_distance(for BVECTOR)
获取一定距离内的项
session.scalars(select(Item).filter(Item.embedding.l2_distance([3, 1, 2]) < 5))
请参阅 examples/sqlalchemy_example.py 和 tests/test_sqlalchemy.py 以获取更多示例
Psycopg3
安装依赖项
pip install "pgvecto_rs[psycopg3]"
初始化连接
import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"
with psycopg.connect(URL) as conn:
pass
启用扩展并注册向量类型
from pgvecto_rs.psycopg import register_vector
conn.execute('CREATE EXTENSION IF NOT EXISTS vectors')
register_vector(conn)
# or asynchronously
# await register_vector_async(conn)
创建一个表
conn.execute('CREATE TABLE items (embedding vector(3))')
将向量插入或复制到表中
conn.execute('INSERT INTO items (embedding) VALUES (%s)', ([1, 2, 3],))
# or faster, copy it
with conn.cursor() as cursor, cursor.copy(
"COPY items (embedding) FROM STDIN (FORMAT BINARY)"
) as copy:
copy.write_row([np.array([1, 2, 3])])
添加近似索引
from pgvecto_rs.types import IndexOption, Hnsw, Ivf
conn.execute(
"CREATE INDEX emb_idx_1 ON items USING \
vectors (embedding vector_l2_ops) WITH (options=$${}$$);".format(
IndexOption(index=Hnsw(), threads=1).dumps()
),
)
# or
conn.execute(
"CREATE INDEX emb_idx_2 ON items USING \
vectors (embedding vector_l2_ops) WITH (options=$${}$$);".format(
IndexOption(index=Ivf()).dumps()
),
)
# Apply all changes
conn.commit()
获取向量的最近邻
conn.execute('SELECT * FROM items ORDER BY embedding <-> %s LIMIT 5', (embedding,)).fetchall()
获取距离
conn.execute('SELECT embedding <-> %s FROM items \
ORDER BY embedding <-> %s', (embedding, embedding)).fetchall()
获取一定距离内的项
conn.execute('SELECT * FROM items WHERE embedding <-> %s < 1.0 \
ORDER BY embedding <-> %s', (embedding, embedding)).fetchall()
请参阅 examples/psycopg_example.py 和 tests/test_psycopg.py 以获取更多示例
Django
安装依赖项
pip install "pgvecto_rs[django]"
创建一个迁移以启用扩展
from pgvecto_rs.django import VectorExtension
class Migration(migrations.Migration):
operations = [
VectorExtension()
]
将向量字段添加到您的模型中
from pgvecto_rs.django import VectorField
class Document(models.Model):
embedding = VectorField(dimensions=3)
所有支持的类型都在此表中显示
原生类型 | Django 类型 | 对应于pgvector-python |
---|---|---|
vector | VectorField | VectorField |
svector | SparseVectorField | SparseVectorField |
vecf16 | Float16VectorField | HalfVectorField |
bvector | BinaryVectorField | BitField |
插入向量
Item(embedding=[1, 2, 3]).save()
添加近似索引
from django.db import models
from pgvecto_rs.django import HnswIndex, IvfIndex
from pgvecto_rs.types import IndexOption, Hnsw
class Item(models.Model):
class Meta:
indexes = [
HnswIndex(
name="emb_idx_1",
fields=["embedding"],
opclasses=["vector_l2_ops"],
m=16,
ef_construction=100,
threads=1,
)
# or
IvfIndex(
name="emb_idx_2",
fields=["embedding"],
nlist=3,
opclasses=["vector_l2_ops"],
),
]
获取向量的最近邻
from pgvecto_rs.django import L2Distance
Item.objects.order_by(L2Distance('embedding', [3, 1, 2]))[:5]
也支持 MaxInnerProduct
、CosineDistance
和 JaccardDistance(for BinaryVectorField)
获取距离
Item.objects.annotate(distance=L2Distance('embedding', [3, 1, 2]))
获取一定距离内的项
Item.objects.alias(distance=L2Distance('embedding', [3, 1, 2])).filter(distance__lt=5)
请参阅 examples/django_example.py 和 tests/test_django.py 以获取更多示例。
SDK
我们的 SDK 设计为直接使用 pgvecto.rs。您可以使用 pgvecto.rs 的强大功能进行相似性搜索或带过滤器的检索,而无需编写任何 SQL 代码。
安装依赖项
pip install "pgvecto_rs[sdk]"
一个最小示例
from pgvecto_rs.sdk import PGVectoRs, Record
# Create a client
client = PGVectoRs(
db_url="postgresql+psycopg://postgres:mysecretpassword@localhost:5432/postgres",
table_name="example",
dimension=3,
)
try:
# Add some records
client.add_records(
[
Record.from_text("hello 1", [1, 2, 3]),
Record.from_text("hello 2", [1, 2, 4]),
]
)
# Search with default operator (sqrt_euclid).
# The results is sorted by distance
for rec, dis in client.search([1, 2, 5]):
print(rec.text)
print(dis)
finally:
# Clean up (i.e. drop the table)
client.drop()
输出
hello 2
1.0
hello 1
4.0
请参阅 examples/sdk_example.py 和 tests/test_sdk.py 以获取更多示例。
开发
此包由 PDM 管理。
设置事项
pdm venv create
pdm use # select the venv inside the project path
pdm sync -d -G :all --no-isolation
# lock requirement
# need pdm >=2.17: https://pdm-project.org/latest/usage/lock-targets/#separate-lock-files-or-merge-into-one
pdm lock -d -G :all --python=">=3.9"
pdm lock -d -G :all --python="<3.9" --append
# install package to local
# `--no-isolation` is required for scipy
pdm install -d --no-isolation
运行 lint
pdm run format
pdm run fix
pdm run check
在当前环境中运行测试
pdm run test
测试
Tox 用于在本地上测试包。
在所有环境中运行测试
tox run
致谢
我们想对 pgvector-python 存储库的创建者和贡献者表示我们的感谢,他们的宝贵代码和架构对开发此存储库产生了很大影响。
项目详情
下载文件
下载适用于您的平台的文件。如果您不确定选择哪个,请了解有关 安装包 的更多信息。