Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 38 additions & 14 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ human_bytes = { version = "0.4.3", default-features = false }
indoc = "2.0.5"
md-5 = "0.10.6"
num-format = "0.4.4"
pgvector = "0.4.0"
quick-xml = "0.36.1"
rand = "0.8.5"
regex = "1.10.6"
Expand Down
17 changes: 17 additions & 0 deletions examples/portal_corp_extension/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
[package]
edition.workspace = true
name = "portal_corp_extension"
publish = false
license.workspace = true
version.workspace = true

[dependencies]
anyhow = { workspace = true }
indoc = { workspace = true }
pgvector = { workspace = true, features = ["sqlx"] }
postgresql_embedded = { path = "../../postgresql_embedded" }
postgresql_extensions = { path = "../../postgresql_extensions" }
sqlx = { workspace = true, features = ["runtime-tokio"] }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
tokio = { workspace = true, features = ["full"] }
128 changes: 128 additions & 0 deletions examples/portal_corp_extension/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
#![forbid(unsafe_code)]
#![deny(clippy::pedantic)]

use anyhow::Result;
use indoc::indoc;
use pgvector::Vector;
use sqlx::{PgPool, Row};
use tracing::info;

use postgresql_embedded::{PostgreSQL, Settings, VersionReq};

/// Example of how to install and configure the PortalCorp pgvector extension.
///
/// See: <https://github.com/pgvector/pgvector?tab=readme-ov-file#getting-started>
#[tokio::main]
async fn main() -> Result<()> {
tracing_subscriber::fmt().compact().init();

info!("Installing PostgreSQL");
let settings = Settings {
version: VersionReq::parse("=16.3.0")?,
..Default::default()
};
let mut postgresql = PostgreSQL::new(settings);
postgresql.setup().await?;

info!("Installing the vector extension from PortalCorp");
postgresql_extensions::install(
postgresql.settings(),
"portal-corp",
"pgvector_compiled",
&VersionReq::parse("=0.16.12")?,
)
.await?;

info!("Starting PostgreSQL");
postgresql.start().await?;

let database_name = "vector-example";
info!("Creating database {database_name}");
postgresql.create_database(database_name).await?;

info!("Configuring extension");
let settings = postgresql.settings();
let database_url = settings.url(database_name);
let pool = PgPool::connect(database_url.as_str()).await?;
pool.close().await;

info!("Restarting database");
postgresql.stop().await?;
postgresql.start().await?;

info!("Enabling extension");
let pool = PgPool::connect(database_url.as_str()).await?;
enable_extension(&pool).await?;

info!("Creating table");
create_table(&pool).await?;

info!("Creating data");
create_data(&pool).await?;

info!("Get the nearest neighbors by L2 distance");
execute_query(
&pool,
"SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5",
)
.await?;

info!("Stopping database");
postgresql.stop().await?;
Ok(())
}

async fn enable_extension(pool: &PgPool) -> Result<()> {
sqlx::query("DROP EXTENSION IF EXISTS vector")
.execute(pool)
.await?;
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(pool)
.await?;
Ok(())
}

async fn create_table(pool: &PgPool) -> Result<()> {
sqlx::query(indoc! {"
CREATE TABLE IF NOT EXISTS items (
id bigserial PRIMARY KEY,
embedding vector(3) NOT NULL
)
"})
.execute(pool)
.await?;
Ok(())
}

async fn create_data(pool: &PgPool) -> Result<()> {
sqlx::query(indoc! {"
INSERT INTO items (embedding)
VALUES
('[1,2,3]'),
('[4,5,6]')
"})
.execute(pool)
.await?;
Ok(())
}

async fn execute_query(pool: &PgPool, query: &str) -> Result<()> {
info!("Query: {query}");
let rows = sqlx::query(query).fetch_all(pool).await?;
for row in rows {
let id: i64 = row.try_get("id")?;
let embedding: Vector = row.try_get("embedding")?;
info!("ID: {id}, Embedding: {embedding:?}");
}
Ok(())
}

#[cfg(test)]
mod test {
use super::*;

#[test]
fn test_main() -> Result<()> {
main()
}
}
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
edition.workspace = true
name = "vector_extension"
name = "tensor_chord_extension"
publish = false
license.workspace = true
version.workspace = true
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use tracing::info;

use postgresql_embedded::{PostgreSQL, Settings, VersionReq};

/// Example of how to install and configure the vector extension.
/// Example of how to install and configure the TensorChord vector extension.
///
/// See: <https://github.com/tensorchord/pgvecto.rs/?tab=readme-ov-file#quick-start>
#[tokio::main]
Expand Down Expand Up @@ -63,21 +63,21 @@ async fn main() -> Result<()> {
info!("Squared Euclidean Distance");
execute_query(
&pool,
"SELECT '[1, 2, 3]'::vector <-> '[3, 2, 1]'::vector AS value;",
"SELECT '[1, 2, 3]'::vector <-> '[3, 2, 1]'::vector AS value",
)
.await?;

info!("Negative Dot Product");
execute_query(
&pool,
"SELECT '[1, 2, 3]'::vector <#> '[3, 2, 1]'::vector AS value;",
"SELECT '[1, 2, 3]'::vector <#> '[3, 2, 1]'::vector AS value",
)
.await?;

info!("Cosine Distance");
execute_query(
&pool,
"SELECT '[1, 2, 3]'::vector <=> '[3, 2, 1]'::vector AS value;",
"SELECT '[1, 2, 3]'::vector <=> '[3, 2, 1]'::vector AS value",
)
.await?;

Expand All @@ -100,7 +100,7 @@ async fn enable_extension(pool: &PgPool) -> Result<()> {
sqlx::query("DROP EXTENSION IF EXISTS vectors")
.execute(pool)
.await?;
sqlx::query("CREATE EXTENSION vectors")
sqlx::query("CREATE EXTENSION IF NOT EXISTS vectors")
.execute(pool)
.await?;
Ok(())
Expand Down