PGVecto.rs支持Python

PGVecto.rs Python库，支持Django、SQLAlchemy和Psycopg 3。

	向量	稀疏向量	半精度向量	二进制向量
SQLAlchemy	✅插入	✅插入	✅插入	✅插入
Psycopg3	✅插入 ✅复制	✅插入 ✅复制	✅插入 ✅复制	✅插入 ✅复制
Django	✅插入	✅插入	✅插入	✅插入

用法

从PyPI安装

pip install pgvecto_rs

并使用它与您的数据库库

• SQLAlchemy
• Psycopg3
• Django

或作为独立的SDK

• 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 存储库的创建者和贡献者表示我们的感谢，他们的宝贵代码和架构对开发此存储库产生了很大影响。