Merge pull request #117 from theseus-rs/add-portal-corp-example

docs: add PortalCorp example for pgvector

Author: brianheineman

Cargo.lock

@@ -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"

Cargo.toml

@@ -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"] }

examples/portal_corp_extension/Cargo.toml

@@ -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()
+    }
+}

examples/portal_corp_extension/src/main.rs

@@ -1,6 +1,6 @@
 [package]
 edition.workspace = true
-name = "vector_extension"
+name = "tensor_chord_extension"
 publish = false
 license.workspace = true
 version.workspace = true

examples/vector_extension/Cargo.toml renamed to examples/tensor_chord_extension/Cargo.toml

@@ -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]
@@ -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?;
 
@@ -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(())