pgvector-rust

pgvector对Rust的支持

支持 Rust-Postgres, SQLx, 和 Diesel

入门指南

遵循您数据库库的说明

• Rust-Postgres
• SQLx
• Diesel

或者查看一些示例

• 与OpenAI的嵌入
• 与Cohere的二进制嵌入
• 与Candle的句子嵌入
• 与Candle的混合搜索（互斥秩融合）
• 与Disco的推荐
• 使用COPY进行批量加载

Rust-Postgres

将此行添加到您的应用的Cargo.toml中的[dependencies]

pgvector = { version = "0.4", features = ["postgres"] }

启用扩展

client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;

创建一个表

client.execute("CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))", &[])?;

从一个Vec<f32>创建一个向量

use pgvector::Vector;

let embedding = Vector::from(vec![1.0, 2.0, 3.0]);

插入一个向量

client.execute("INSERT INTO items (embedding) VALUES ($1)", &[&embedding])?;

获取最近邻

let row = client.query_one(
    "SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 1",
    &[&embedding],
)?;

检索一个向量

let row = client.query_one("SELECT embedding FROM items LIMIT 1", &[])?;
let embedding: Vector = row.get(0);

如果值可能是NULL，请使用Option

let embedding: Option<Vector> = row.get(0);

SQLx

将此行添加到您的应用的Cargo.toml中的[dependencies]

pgvector = { version = "0.4", features = ["sqlx"] }

对于SQLx < 0.8，请使用version = "0.3"和此README.

启用扩展

sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
    .execute(&pool)
    .await?;

创建一个表

sqlx::query("CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))")
    .execute(&pool)
    .await?;

从一个Vec<f32>创建一个向量

use pgvector::Vector;

let embedding = Vector::from(vec![1.0, 2.0, 3.0]);

插入一个向量

sqlx::query("INSERT INTO items (embedding) VALUES ($1)")
    .bind(embedding)
    .execute(&pool)
    .await?;

获取最近邻

let rows = sqlx::query("SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 1")
    .bind(embedding)
    .fetch_all(&pool)
    .await?;

检索一个向量

let row = sqlx::query("SELECT embedding FROM items LIMIT 1").fetch_one(&pool).await?;
let embedding: Vector = row.try_get("embedding")?;

Diesel

将此行添加到您的应用的Cargo.toml中的[dependencies]

pgvector = { version = "0.4", features = ["diesel"] }

并更新您的应用的diesel.toml中的[print_schema]

import_types = ["diesel::sql_types::*", "pgvector::sql_types::*"]
generate_missing_sql_type_definitions = false

创建迁移

diesel migration generate create_vector_extension

使用up.sql

CREATE EXTENSION vector

和down.sql

DROP EXTENSION vector

运行迁移

diesel migration run

现在可以在未来的迁移中使用vector类型

CREATE TABLE items (
  id SERIAL PRIMARY KEY,
  embedding VECTOR(3)
)

对于模型，使用

use pgvector::Vector;

#[derive(Queryable)]
#[diesel(table_name = items)]
pub struct Item {
    pub id: i32,
    pub embedding: Option<Vector>,
}

#[derive(Insertable)]
#[diesel(table_name = items)]
pub struct NewItem {
    pub embedding: Option<Vector>,
}

从一个Vec<f32>创建一个向量

let embedding = Vector::from(vec![1.0, 2.0, 3.0]);

插入一个向量

let new_item = NewItem {
    embedding: Some(embedding)
};

diesel::insert_into(items::table)
    .values(&new_item)
    .get_result::<Item>(&mut conn)?;

获取最近邻

use pgvector::VectorExpressionMethods;

let neighbors = items::table
    .order(items::embedding.l2_distance(embedding))
    .limit(5)
    .load::<Item>(&mut conn)?;

也支持max_inner_product、cosine_distance、l1_distance、hamming_distance和jaccard_distance

获取距离

let distances = items::table
    .select(items::embedding.l2_distance(embedding))
    .load::<Option<f64>>(&mut conn)?;

在迁移中添加近似索引

CREATE INDEX my_index ON items USING hnsw (embedding vector_l2_ops)
-- or
CREATE INDEX my_index ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)

使用vector_ip_ops进行内积，使用vector_cosine_ops进行余弦距离

序列化

使用serde功能启用序列化

半精度向量

使用halfvec功能启用半精度向量

参考

将向量转换为Vec<f32>

let f32_vec: Vec<f32> = vec.into();

获取切片

let slice = vec.as_slice();

历史记录

查看变更日志

贡献

鼓励每个人帮助改进这个项目。以下是一些你可以帮助的方式

• 报告错误
• 修复错误并提交拉取请求
• 编写、阐明或修复文档
• 建议或添加新功能

开始开发

git clone https://github.com/pgvector/pgvector-rust.git
cd pgvector-rust
createdb pgvector_rust_test
cargo test --all-features