Add match macro to perform KNN query

Author: joelpaulkoch

lib/sqlite_vec/ecto/query.ex

@@ -31,6 +31,15 @@ if Code.ensure_loaded?(Ecto) do
       end
     end
 
+    @doc """
+    Performs a K-nearest-neighbors (KNN) query. You must specify a LIMIT or 'k = ?' constraint.
+    """
+    defmacro match(left, right) do
+      quote do
+        fragment("? match ?", unquote(left), unquote(right))
+      end
+    end
+
     @doc """
     Calculates the L2 euclidian distance between vectors a and b. Only valid for float32 or int8 vectors.
 

notebooks/usage_with_ecto.livemd

@@ -160,6 +160,8 @@ MyApp.Repo.insert(%VirtualEmbedding{
 })
 ```
 
+You can perform a K-nearest-neighbors query using `match` and `limit`.
+
 ```elixir
 import Ecto.Query
 import SqliteVec.Ecto.Query
@@ -168,7 +170,8 @@ v = SqliteVec.Float32.new([2, 2])
 
 MyApp.Repo.all(
   from(i in VirtualEmbedding,
-    where: fragment("? match ? and k=?", i.embedding, vec_f32(v), 3)
+    where: match(i.embedding, vec_f32(v)),
+    limit: 3
   )
 )
 ```

test/sqlite_vec/bit_ecto_test.exs

@@ -39,6 +39,30 @@ defmodule BitEctoTest do
     )
   end
 
+  test "match performs a KNN query" do
+    items =
+      Repo.all(
+        from(i in BitItem,
+          where:
+            match(
+              i.embedding,
+              vec_bit(SqliteVec.Bit.new([0, 0, 0, 0, 0, 0, 0, 1]))
+            ),
+          limit: 3
+        )
+      )
+
+    assert Enum.map(items, fn v ->
+             v.id
+           end) == [2, 3, 1]
+
+    assert Enum.map(items, fn v -> v.embedding |> SqliteVec.Bit.to_list() end) == [
+             [0, 0, 0, 0, 0, 0, 0, 0],
+             [0, 0, 0, 0, 1, 0, 1, 0],
+             [1, 1, 1, 1, 1, 1, 1, 1]
+           ]
+  end
+
   test "vector hamming distance" do
     items =
       Repo.all(

test/sqlite_vec/float32_ecto_test.exs

@@ -45,6 +45,28 @@ defmodule EctoTest do
     })
   end
 
+  test "match performs a KNN query" do
+    v = SqliteVec.Float32.new([2, 2])
+
+    items =
+      Repo.all(
+        from(i in Float32Item,
+          where: match(i.embedding, vec_f32(v)),
+          limit: 3
+        )
+      )
+
+    assert Enum.map(items, fn v ->
+             v.id
+           end) == [1, 3, 2]
+
+    assert Enum.map(items, fn v -> v.embedding |> SqliteVec.Float32.to_list() end) == [
+             [1.0, 2.0],
+             [3.0, 4.0],
+             [52.0, 43.0]
+           ]
+  end
+
   test "vector l2 distance" do
     v = SqliteVec.Float32.new([2, 2])
 

test/sqlite_vec/int8_ecto_test.exs

@@ -71,6 +71,26 @@ defmodule Int8EctoTest do
            ]
   end
 
+  test "match performs a KNN query" do
+    items =
+      Repo.all(
+        from(i in Int8Item,
+          where: match(i.embedding, vec_int8(SqliteVec.Int8.new([2, 2]))),
+          limit: 3
+        )
+      )
+
+    assert Enum.map(items, fn v ->
+             v.id
+           end) == [1, 3, 2]
+
+    assert Enum.map(items, fn v -> v.embedding |> SqliteVec.Int8.to_list() end) == [
+             [1, 2],
+             [3, 4],
+             [52, 43]
+           ]
+  end
+
   test "vector cosine distance" do
     items =
       Repo.all(