Added test for halfvec with Postgrex

Author: ankane

test/postgrex_test.exs

@@ -1,4 +1,4 @@
-Postgrex.Types.define(PostgrexApp.PostgrexTypes, [Pgvector.Extensions.Vector], [])
+Postgrex.Types.define(PostgrexApp.PostgrexTypes, Pgvector.extensions(), [])
 
 # needed if postgrex is optional
 # Application.ensure_all_started(:postgrex)
@@ -10,7 +10,7 @@ defmodule PostgrexTest do
     {:ok, pid} = Postgrex.start_link(database: "pgvector_elixir_test", types: PostgrexApp.PostgrexTypes)
     Postgrex.query!(pid, "CREATE EXTENSION IF NOT EXISTS vector", [])
     Postgrex.query!(pid, "DROP TABLE IF EXISTS postgrex_items", [])
-    Postgrex.query!(pid, "CREATE TABLE postgrex_items (id bigserial primary key, embedding vector(3), binary_embedding bit(3))", [])
+    Postgrex.query!(pid, "CREATE TABLE postgrex_items (id bigserial primary key, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3))", [])
     {:ok, pid: pid}
   end
 
@@ -19,7 +19,7 @@ defmodule PostgrexTest do
     context
   end
 
-  test "l2 distance", %{pid: pid} = _context do
+  test "vector l2 distance", %{pid: pid} = _context do
     embeddings = [Pgvector.new([1, 1, 1]), [2, 2, 2], Nx.tensor([1, 1, 2], type: :f32)]
     Postgrex.query!(pid, "INSERT INTO postgrex_items (embedding) VALUES ($1), ($2), ($3)", embeddings)
 
@@ -30,6 +30,17 @@ defmodule PostgrexTest do
     assert Enum.map(result.rows, fn v -> Enum.at(v, 1) |> Pgvector.to_list() end) == [[1.0, 1.0, 1.0], [1.0, 1.0, 2.0], [2.0, 2.0, 2.0]]
   end
 
+  test "halfvec l2 distance", %{pid: pid} = _context do
+    embeddings = [Pgvector.HalfVector.new([1, 1, 1]), [2, 2, 2], Nx.tensor([1, 1, 2], type: :f16)]
+    Postgrex.query!(pid, "INSERT INTO postgrex_items (half_embedding) VALUES ($1), ($2), ($3)", embeddings)
+
+    result = Postgrex.query!(pid, "SELECT id, half_embedding FROM postgrex_items ORDER BY half_embedding <-> $1 LIMIT 5", [[1, 1, 1]])
+
+    assert ["id", "half_embedding"] == result.columns
+    assert Enum.map(result.rows, fn v -> Enum.at(v, 0) end) == [1, 3, 2]
+    assert Enum.map(result.rows, fn v -> Enum.at(v, 1) |> Pgvector.to_list() end) == [[1.0, 1.0, 1.0], [1.0, 1.0, 2.0], [2.0, 2.0, 2.0]]
+  end
+
   test "create index", %{pid: pid} = _context do
     Postgrex.query!(pid, "CREATE INDEX my_index ON postgrex_items USING ivfflat (embedding vector_l2_ops) WITH (lists = 1)", [])
   end