Crate pgvector

Source
Expand description

§pgvector-rust

pgvector support for Rust

Supports Rust-Postgres, SQLx, and Diesel

Build Status

§Getting Started

Follow the instructions for your database library:

Or check out some examples:

§Rust-Postgres

Add this line to your application’s Cargo.toml under [dependencies]:

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

Enable the extension

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

Create a table

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

Create a vector from a Vec<f32>

use pgvector::Vector;

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

Insert a vector

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

Get the nearest neighbor

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

Retrieve a vector

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

Use Option if the value could be NULL

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

§SQLx

Add this line to your application’s Cargo.toml under [dependencies]:

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

For SQLx < 0.8, use version = "0.3" and this readme.

Enable the extension

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

Create a table

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

Create a vector from a Vec<f32>

use pgvector::Vector;

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

Insert a vector

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

Get the nearest neighbors

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

Retrieve a vector

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

§Diesel

Add this line to your application’s Cargo.toml under [dependencies]:

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

And update your application’s diesel.toml under [print_schema]:

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

Create a migration

diesel migration generate create_vector_extension

with up.sql:

CREATE EXTENSION vector

and down.sql:

DROP EXTENSION vector

Run the migration

diesel migration run

You can now use the vector type in future migrations

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

For models, use:

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>,
}

Create a vector from a Vec<f32>

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

Insert a vector

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

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

Get the nearest neighbors

use pgvector::VectorExpressionMethods;

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

Also supports max_inner_product, cosine_distance, l1_distance, hamming_distance, and jaccard_distance

Get the distances

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

Add an approximate index in a migration

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)

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

§Serialization

Use the serde feature to enable serialization

§Reference

§Vectors

Create a vector

use pgvector::Vector;

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

Convert to a Vec<f32>

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

Get a slice

let slice = vec.as_slice();

§Half Vectors

Note: Use the halfvec feature to enable half vectors

Create a half vector from a Vec<f16>

use half::f16;
use pgvector::HalfVector;

let vec = HalfVector::from(vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]);

Or a f32 slice

let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);

Convert to a Vec<f16>

let f16_vec: Vec<f16> = vec.into();

Get a slice

let slice = vec.as_slice();

§Binary Vectors

Create a binary vector from a slice of bits

use pgvector::Bit;

let vec = Bit::new(&[true, false, true]);

Or a slice of bytes

let vec = Bit::from_bytes(&[0b00000000, 0b11111111]);

Get the number of bits

let len = vec.len();

Get a slice of bytes

let bytes = vec.as_bytes();

§Sparse Vectors

Create a sparse vector from a dense vector

use pgvector::SparseVector;

let vec = SparseVector::from_dense(vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0]);

Or a map of non-zero elements

let map = HashMap::from([(0, 1.0), (2, 2.0), (4, 3.0)]);
let vec = SparseVector::from_map(&map, 6);

Note: Indices start at 0

Get the number of dimensions

let dim = vec.dimensions();

Get the indices of non-zero elements

let indices = vec.indices();

Get the values of non-zero elements

let values = vec.values();

Get a dense vector

let f32_vec = vec.to_vec();

§History

View the changelog

§Contributing

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

To get started with development:

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

To run an example:

cd examples/loading
createdb pgvector_example
cargo run

Structs§

Bit
A bit string.
HalfVector
A half vector.
SparseVector
A sparse vector.
Vector
A vector.