Added Citus example [skip ci]

Author: ankane

README.md

@@ -21,6 +21,7 @@ Or check out some examples:
 - [Sentence embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/candle/src/main.rs) with Candle
 - [Hybrid search](https://github.com/pgvector/pgvector-rust/blob/master/examples/hybrid_search/src/main.rs) with Candle (Reciprocal Rank Fusion)
 - [Recommendations](https://github.com/pgvector/pgvector-rust/blob/master/examples/disco/src/main.rs) with Disco
+- [Horizontal scaling](https://github.com/pgvector/pgvector-rust/blob/master/examples/citus/src/main.rs) with Citus
 - [Bulk loading](https://github.com/pgvector/pgvector-rust/blob/master/examples/loading/src/main.rs) with `COPY`
 
 ## Rust-Postgres

examples/citus/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "example"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[dependencies]
+pgvector = { path = "../..", features = ["postgres"] }
+postgres = "0.19"
+rand = "0.8"
+
+[profile.dev]
+opt-level = 1

examples/citus/src/main.rs

@@ -0,0 +1,98 @@
+use pgvector::Vector;
+use postgres::binary_copy::BinaryCopyInWriter;
+use postgres::types::{Kind, Type};
+use postgres::{Client, NoTls};
+use rand::Rng;
+use std::error::Error;
+
+fn main() -> Result<(), Box<dyn Error>> {
+    // generate random data
+    let rows = 100000;
+    let dimensions = 128;
+    let mut rng = rand::thread_rng();
+    let embeddings: Vec<Vec<f32>> = (0..rows)
+        .map(|_| (0..dimensions).map(|_| rng.gen()).collect())
+        .collect();
+    let categories: Vec<i64> = (0..rows).map(|_| rng.gen_range(1..=100)).collect();
+    let queries: Vec<Vec<f32>> = (0..10)
+        .map(|_| (0..dimensions).map(|_| rng.gen()).collect())
+        .collect();
+
+    // enable extensions
+    let mut client = Client::configure()
+        .host("localhost")
+        .dbname("pgvector_citus")
+        .user(std::env::var("USER")?.as_str())
+        .connect(NoTls)?;
+    client.execute("CREATE EXTENSION IF NOT EXISTS citus", &[])?;
+    client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
+
+    // GUC variables set on the session do not propagate to Citus workers
+    // https://github.com/citusdata/citus/issues/462
+    // you can either:
+    // 1. set them on the system, user, or database and reconnect
+    // 2. set them for a transaction with SET LOCAL
+    client.execute(
+        "ALTER DATABASE pgvector_citus SET maintenance_work_mem = '512MB'",
+        &[],
+    )?;
+    client.execute("ALTER DATABASE pgvector_citus SET hnsw.ef_search = 20", &[])?;
+    client.close()?;
+
+    // reconnect for updated GUC variables to take effect
+    let mut client = Client::configure()
+        .host("localhost")
+        .dbname("pgvector_citus")
+        .user(std::env::var("USER")?.as_str())
+        .connect(NoTls)?;
+
+    println!("Creating distributed table");
+    client.execute("DROP TABLE IF EXISTS items", &[])?;
+    client.execute(
+        &format!("CREATE TABLE items (id bigserial, embedding vector({dimensions}), category_id bigint, PRIMARY KEY (id, category_id))"),
+        &[],
+    )?;
+    client.execute("SET citus.shard_count = 4", &[])?;
+    client.execute(
+        "SELECT create_distributed_table('items', 'category_id')",
+        &[],
+    )?;
+
+    println!("Loading data in parallel");
+    let vector_type = get_type(&mut client, "vector")?;
+    let writer =
+        client.copy_in("COPY items (embedding, category_id) FROM STDIN WITH (FORMAT BINARY)")?;
+    let mut writer = BinaryCopyInWriter::new(writer, &[vector_type, Type::INT8]);
+    for (embedding, category) in embeddings.into_iter().zip(categories) {
+        writer.write(&[&Vector::from(embedding), &category])?;
+    }
+    writer.finish()?;
+
+    println!("Creating index in parallel");
+    client.execute(
+        "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)",
+        &[],
+    )?;
+
+    println!("Running distributed queries");
+    for query in queries {
+        let rows = client.query(
+            "SELECT id FROM items ORDER BY embedding <-> $1 LIMIT 10",
+            &[&Vector::from(query)],
+        )?;
+        let ids: Vec<i64> = rows.into_iter().map(|row| row.get(0)).collect();
+        println!("{:?}", ids);
+    }
+
+    Ok(())
+}
+
+fn get_type(client: &mut Client, name: &str) -> Result<Type, Box<dyn Error>> {
+    let row = client.query_one("SELECT pg_type.oid, nspname AS schema FROM pg_type INNER JOIN pg_namespace ON pg_namespace.oid = pg_type.typnamespace WHERE typname = $1", &[&name])?;
+    Ok(Type::new(
+        name.into(),
+        row.get("oid"),
+        Kind::Simple,
+        row.get("schema"),
+    ))
+}