Added Citus example [skip ci]

Author: ankane

README.md

@@ -26,6 +26,7 @@ Or check out some examples:
 - [Sentence embeddings](https://github.com/pgvector/pgvector-elixir/blob/master/examples/bumblebee/example.exs) with Bumblebee
 - [Hybrid search](https://github.com/pgvector/pgvector-elixir/blob/master/examples/hybrid_search/example.exs) with Bumblebee (Reciprocal Rank Fusion)
 - [Sparse search](https://github.com/pgvector/pgvector-elixir/blob/master/examples/sparse_search/example.exs) with Bumblebee
+- [Horizontal scaling](https://github.com/pgvector/pgvector-elixir/blob/master/examples/citus/example.exs) with Citus
 - [Bulk loading](https://github.com/pgvector/pgvector-elixir/blob/master/examples/loading/example.exs) with `COPY`
 
 ## Ecto

examples/citus/example.exs

@@ -0,0 +1,109 @@
+Postgrex.Types.define(Example.PostgrexTypes, Pgvector.extensions(), [])
+
+rows = 100_000
+dimensions = 128
+
+IO.puts("Generating random data")
+
+key = Nx.Random.key(42)
+{embeddings, new_key} = Nx.Random.uniform(key, shape: {rows, dimensions})
+{categories, new_key} = Nx.Random.randint(new_key, 1, 100, shape: {rows})
+{queries, _new_key} = Nx.Random.uniform(new_key, shape: {10, dimensions})
+
+# enable extensions
+{:ok, pid} = Postgrex.start_link(database: "pgvector_citus", types: Example.PostgrexTypes)
+Postgrex.query!(pid, "CREATE EXTENSION IF NOT EXISTS citus", [])
+Postgrex.query!(pid, "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
+Postgrex.query!(pid, "ALTER DATABASE pgvector_citus SET maintenance_work_mem = '512MB'", [])
+Postgrex.query!(pid, "ALTER DATABASE pgvector_citus SET hnsw.ef_search = 20", [])
+# TODO close connection
+
+# reconnect for updated GUC variables to take effect
+{:ok, pid} = Postgrex.start_link(database: "pgvector_citus", types: Example.PostgrexTypes)
+
+IO.puts("Creating distributed table")
+
+Postgrex.query!(pid, "DROP TABLE IF EXISTS items", [])
+
+Postgrex.query!(
+  pid,
+  "CREATE TABLE items (id bigserial, embedding vector(#{dimensions}), category_id bigint, PRIMARY KEY (id, category_id))",
+  []
+)
+
+Postgrex.query!(pid, "SET citus.shard_count = 4", [])
+Postgrex.query!(pid, "SELECT create_distributed_table('items', 'category_id')", [])
+
+defmodule Example do
+  # https://www.postgresql.org/docs/current/sql-copy.html
+  def copy(stream, rows) do
+    signature = "PGCOPY\n\xFF\r\n\0"
+
+    Enum.into(
+      [
+        <<signature::binary, 0::unsigned-32, 0::unsigned-32>>,
+        Enum.map(rows, &copy_row(&1)),
+        <<-1::unsigned-16>>
+      ],
+      stream
+    )
+  end
+
+  defp copy_row(row) do
+    count = row |> length()
+    data = Enum.join(Enum.map(row, &copy_value(&1)))
+    <<count::unsigned-16, data::binary>>
+  end
+
+  defp copy_value(value) when is_struct(value, Pgvector) do
+    data = value |> Pgvector.to_binary()
+    <<IO.iodata_length(data)::unsigned-32, data::binary>>
+  end
+
+  defp copy_value(value) when is_integer(value) do
+    <<8::unsigned-32, value::64>>
+  end
+end
+
+IO.puts("Loading data in parallel")
+
+Postgrex.transaction(
+  pid,
+  fn conn ->
+    stream =
+      Postgrex.stream(
+        conn,
+        "COPY items (embedding, category_id) FROM STDIN WITH (FORMAT BINARY)",
+        []
+      )
+
+    rows =
+      Enum.map(Enum.zip(embeddings |> Nx.to_list(), categories |> Nx.to_list()), fn {v, c} ->
+        [v |> Pgvector.new(), c]
+      end)
+
+    stream |> Example.copy(rows)
+  end,
+  timeout: 30000
+)
+
+IO.puts("Creating index in parallel")
+
+Postgrex.query!(pid, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)", [])
+
+IO.puts("Running distributed queries")
+
+for query <- Nx.to_list(queries) do
+  result =
+    Postgrex.query!(pid, "SELECT id FROM items ORDER BY embedding <-> $1 LIMIT 10", [
+      query |> Pgvector.new()
+    ])
+
+  IO.inspect(Enum.map(result.rows, fn v -> List.first(v) end))
+end

examples/citus/mix.exs

@@ -0,0 +1,19 @@
+defmodule Example.MixProject do
+  use Mix.Project
+
+  def project do
+    [
+      app: :example,
+      version: "0.1.0",
+      deps: deps()
+    ]
+  end
+
+  defp deps do
+    [
+      {:pgvector, path: "../.."},
+      {:postgrex, "~> 0.17"},
+      {:nx, "~> 0.5"}
+    ]
+  end
+end