Skip to content

Commit

Permalink
Added support for creating a SparseVector from a map
Browse files Browse the repository at this point in the history
  • Loading branch information
ankane committed Jun 20, 2024
1 parent e22334d commit a619f7d
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
13 changes: 13 additions & 0 deletions lib/pgvector/sparse_vector.ex
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,19 @@ defmodule Pgvector.SparseVector do
end
end

@doc """
Creates a new sparse vector from a map
"""
def new(map, dimensions) when is_map(map) do
{indices, values} =
map
|> Enum.sort_by(fn {k, _} -> k end)
|> Enum.filter(fn {_, v} -> v != 0 end)
|> Enum.unzip()

new(dimensions, indices, values)
end

defp new(dim, indices, values) do
nnz = indices |> length()
indices = for v <- indices, into: "", do: <<v::signed-32>>
Expand Down
5 changes: 5 additions & 0 deletions test/sparse_vector_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@ defmodule SparseVectorTest do
assert tensor == tensor |> Pgvector.SparseVector.new() |> Pgvector.to_tensor()
end

test "map" do
map = %{0 => 1.0, 2 => 2.0, 4 => 3.0}
assert [1.0, 0.0, 2.0, 0.0, 3.0, 0.0] == map |> Pgvector.SparseVector.new(6) |> Pgvector.to_list()
end

test "inspect" do
vector = Pgvector.SparseVector.new([1, 2, 3])
assert "Pgvector.SparseVector.new([1.0, 2.0, 3.0])" == inspect(vector)
Expand Down

0 comments on commit a619f7d

Please sign in to comment.