Added binary_quantize and subvector functions for Ecto

Author: ankane

CHANGELOG.md

@@ -4,6 +4,7 @@
 - Added support for `bit` type to Ecto
 - Added `Pgvector.extensions/0` function
 - Added `l1_distance`, `hamming_distance`, and `jaccard_distance` functions for Ecto
+- Added `binary_quantize` and `subvector` functions for Ecto
 - Dropped support for Elixir < 1.12
 
 ## 0.2.1 (2023-09-25)

lib/pgvector/ecto/query.ex

@@ -57,5 +57,23 @@ if Code.ensure_loaded?(Ecto) do
         fragment("(? <%> ?)", unquote(column), unquote(value))
       end
     end
+
+    @doc """
+    Returns the binary quantization
+    """
+    defmacro binary_quantize(value) do
+      quote do
+        fragment("binary_quantize(?)", unquote(value))
+      end
+    end
+
+    @doc """
+    Returns a subvector
+    """
+    defmacro subvector(value, start, count) do
+      quote do
+        fragment("subvector(?, ?, ?)", unquote(value), unquote(start), unquote(count))
+      end
+    end
   end
 end

test/ecto_test.exs

@@ -56,6 +56,16 @@ defmodule EctoTest do
     assert Enum.map(items, fn v -> v.id end) == [1, 3, 2]
   end
 
+  test "vector binary_quantize" do
+    item = Repo.one(first(from i in Item, select: binary_quantize(i.embedding)))
+    assert item == <<1::1, 1::1, 1::1>>
+  end
+
+  test "vector subvector" do
+    result = Repo.one(last(from i in Item, select: subvector(i.embedding, 2, 2)))
+    assert result |> Pgvector.to_list() == [1, 2]
+  end
+
   test "halfvec l2 distance" do
     items = Repo.all(from i in Item, order_by: l2_distance(i.half_embedding, ^Pgvector.HalfVector.new([1, 1, 1])), limit: 5)
     assert Enum.map(items, fn v -> v.id end) == [1, 3, 2]
@@ -82,6 +92,16 @@ defmodule EctoTest do
     assert Enum.map(items, fn v -> v.id end) == [1, 3, 2]
   end
 
+  test "halfvec binary_quantize" do
+    item = Repo.one(first(from i in Item, select: binary_quantize(i.half_embedding)))
+    assert item == <<1::1, 1::1, 1::1>>
+  end
+
+  test "halfvec subvector" do
+    result = Repo.one(last(from i in Item, select: subvector(i.half_embedding, 2, 2)))
+    assert result |> Pgvector.to_list() == [1, 2]
+  end
+
   test "bit hamming distance" do
     items = Repo.all(from i in Item, order_by: hamming_distance(i.binary_embedding, ^<<1::1, 0::1, 1::1>>), limit: 5)
     assert Enum.map(items, fn v -> v.id end) == [2, 3, 1]