diff --git a/CHANGELOG.md b/CHANGELOG.md index 10bb506..4a1a6f2 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/lib/pgvector/ecto/query.ex b/lib/pgvector/ecto/query.ex index e70213d..c236a7b 100644 --- a/lib/pgvector/ecto/query.ex +++ b/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 diff --git a/test/ecto_test.exs b/test/ecto_test.exs index 371d9ad..1e6138e 100644 --- a/test/ecto_test.exs +++ b/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]