Skip to content

Commit

Permalink
Add match macro to perform KNN query
Browse files Browse the repository at this point in the history
  • Loading branch information
joelpaulkoch committed Nov 29, 2024
1 parent 522865b commit 27600e0
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 1 deletion.
9 changes: 9 additions & 0 deletions lib/sqlite_vec/ecto/query.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 4 additions & 1 deletion notebooks/usage_with_ecto.livemd
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
)
)
```
24 changes: 24 additions & 0 deletions test/sqlite_vec/bit_ecto_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
22 changes: 22 additions & 0 deletions test/sqlite_vec/float32_ecto_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -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])

Expand Down
20 changes: 20 additions & 0 deletions test/sqlite_vec/int8_ecto_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 27600e0

Please sign in to comment.