Added Citus example [skip ci]

Author: ankane

README.md

@@ -37,6 +37,7 @@ Or check out some examples:
 
 - [Embeddings](src/test/java/com/pgvector/OpenAITest.java) with OpenAI
 - [Binary embeddings](src/test/java/com/pgvector/CohereTest.java) with Cohere
+- [Horizontal scaling](src/test/java/com/pgvector/CitusTest.java) with Citus
 - [Bulk loading](src/test/java/com/pgvector/LoadingTest.java) with `COPY`
 
 ## JDBC (Java)

src/test/java/com/pgvector/CitusTest.java

@@ -0,0 +1,96 @@
+package com.pgvector;
+
+import java.io.UnsupportedEncodingException;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Random;
+import com.pgvector.PGvector;
+import org.postgresql.PGConnection;
+import org.postgresql.copy.CopyIn;
+import org.postgresql.copy.CopyManager;
+import org.postgresql.core.BaseConnection;
+import org.junit.jupiter.api.Test;
+
+public class CitusTest {
+    @Test
+    void example() throws SQLException {
+        if (System.getenv("TEST_CITUS") == null) {
+            return;
+        }
+
+        // generate data
+        int rows = 1000000;
+        int dimensions = 128;
+        ArrayList<float[]> embeddings = new ArrayList<>(rows);
+        ArrayList<Integer> categories = new ArrayList<>(rows);
+        Random rnd = new Random();
+        for (int i = 0; i < rows; i++) {
+            float[] embedding = new float[dimensions];
+            for (int j = 0; j < dimensions; j++) {
+                embedding[j] = (float) Math.random();
+            }
+            embeddings.add(embedding);
+            categories.add(rnd.nextInt(100));
+        }
+
+        // enable extensions
+        Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_citus");
+        Statement setupStmt = conn.createStatement();
+        setupStmt.executeUpdate("CREATE EXTENSION IF NOT EXISTS citus");
+        setupStmt.executeUpdate("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
+        setupStmt.executeUpdate("ALTER DATABASE pgvector_citus SET maintenance_work_mem = '512MB'");
+        setupStmt.executeUpdate("ALTER DATABASE pgvector_citus SET hnsw.ef_search = 20");
+        conn.close();
+
+        // reconnect for updated GUC variables to take effect
+        conn = DriverManager.getConnection("jdbc:postgresql://localhost:5432/pgvector_citus");
+        PGvector.addVectorType(conn);
+
+        System.out.println("Creating distributed table");
+        setupStmt = conn.createStatement();
+        setupStmt.executeUpdate("DROP TABLE IF EXISTS items");
+        setupStmt.executeUpdate(String.format("CREATE TABLE items (id bigserial, embedding vector(%d), category_id bigint, PRIMARY KEY (id, category_id))", dimensions));
+        setupStmt.executeUpdate("SET citus.shard_count = 4");
+        setupStmt.executeQuery("SELECT create_distributed_table('items', 'category_id')");
+
+        System.out.println("Loading data in parallel");
+        CopyManager copyManager = new CopyManager((BaseConnection) conn);
+        // TODO use binary format
+        CopyIn copyIn = copyManager.copyIn("COPY items (embedding, category_id) FROM STDIN");
+        for (int i = 0; i < rows; i++) {
+            PGvector embedding = new PGvector(embeddings.get(i));
+            byte[] bytes = String.format("%s\t%d\n", embedding.getValue(), categories.get(i)).getBytes();
+            copyIn.writeToCopy(bytes, 0, bytes.length);
+        }
+        copyIn.endCopy();
+
+        System.out.println("Creating index in parallel");
+        Statement createIndexStmt = conn.createStatement();
+        createIndexStmt.executeUpdate("CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)");
+
+        System.out.println("Running distributed queries");
+        for (int i = 0; i < 10; i++) {
+            PreparedStatement queryStmt = conn.prepareStatement("SELECT id FROM items ORDER BY embedding <-> ? LIMIT 10");
+            queryStmt.setObject(1, new PGvector(embeddings.get(rnd.nextInt(rows))));
+            ResultSet rs = queryStmt.executeQuery();
+            ArrayList<Long> ids = new ArrayList<>();
+            while (rs.next()) {
+                ids.add(rs.getLong("id"));
+            }
+            System.out.println(ids.toString());
+        }
+
+        conn.close();
+    }
+}