Skip to content

Commit

Permalink
docs: add PortalCorp example for pgvector
Browse files Browse the repository at this point in the history
  • Loading branch information
brianheineman committed Aug 4, 2024
1 parent afac84f commit 148b5f3
Show file tree
Hide file tree
Showing 6 changed files with 192 additions and 20 deletions.
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"] }
130 changes: 130 additions & 0 deletions examples/portal_corp_extension/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
#![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 portal corp 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?;
// configure_extension(&pool).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(target_os = "linux")]
#[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

0 comments on commit 148b5f3

Please sign in to comment.