defmodule HNSWLib.Index do
@moduledoc """
Documentation for `HNSWLib.Index`.
"""
defstruct [:space, :dim, :reference]
alias __MODULE__, as: T
alias HNSWLib.Helper
@doc """
Construct a new Index
##### Positional Parameters
- *space*: `:cosine` | `:ip` | `:l2`.
An atom that indicates the vector space. Valid values are
- `:cosine`, cosine space
- `:ip`, inner product space
- `:l2`, L2 space
- *dim*: `non_neg_integer()`.
Number of dimensions for each vector.
- *max_elements*: `non_neg_integer()`.
Number of maximum elements.
##### Keyword Paramters
- *m*: `non_neg_integer()`.
`M` is tightly connected with internal dimensionality of the data
strongly affects the memory consumption
- *ef_construction*: `non_neg_integer()`.
controls index search speed/build speed tradeoff
- *random_seed*: `non_neg_integer()`.
- *allow_replace_deleted*: `boolean()`.
"""
@spec new(:cosine | :ip | :l2, non_neg_integer(), non_neg_integer(), [
{:m, non_neg_integer()},
{:ef_construction, non_neg_integer()},
{:random_seed, non_neg_integer()},
{:allow_replace_deleted, boolean()}
]) :: {:ok, %T{}} | {:error, String.t()}
def new(space, dim, max_elements, opts \\ [])
when (space == :l2 or space == :ip or space == :cosine) and is_integer(dim) and dim >= 0 and
is_integer(max_elements) and max_elements >= 0 do
m = Helper.get_keyword!(opts, :m, :non_neg_integer, 16)
ef_construction = Helper.get_keyword!(opts, :ef_construction, :non_neg_integer, 200)
random_seed = Helper.get_keyword!(opts, :random_seed, :non_neg_integer, 100)
allow_replace_deleted = Helper.get_keyword!(opts, :allow_replace_deleted, :boolean, false)
with {:ok, ref} <-
HNSWLib.Nif.index_new(
space,
dim,
max_elements,
m,
ef_construction,
random_seed,
allow_replace_deleted
) do
{:ok,
%T{
space: space,
dim: dim,
reference: ref
}}
else
{:error, reason} ->
{:error, reason}
end
end
@doc """
Query the index with a single vector or a list of vectors.
##### Positional Parameters
- *query*: `Nx.Tensor.t() | binary() | [binary()]`.
A vector or a list of vectors to query.
If *query* is a list of vectors, the vectors must be of the same dimension.
##### Keyword Paramters
- *k*: `pos_integer()`.
Number of nearest neighbors to return.
- *num_threads*: `integer()`.
Number of threads to use.
"""
@spec knn_query(%T{}, Nx.Tensor.t() | binary() | [binary()], [
{:k, pos_integer()},
{:num_threads, integer()}
# {:filter, function()}
]) :: {:ok, Nx.Tensor.t(), Nx.Tensor.t()} | {:error, String.t()}
def knn_query(self, query, opts \\ [])
def knn_query(self = %T{}, query, opts) when is_binary(query) do
k = Helper.get_keyword!(opts, :k, :pos_integer, 1)
num_threads = Helper.get_keyword!(opts, :num_threads, :integer, -1)
Helper.might_be_float_data!(query)
features = trunc(byte_size(query) / Helper.float_size())
Helper.ensure_vector_dimension!(self, features, true)
_do_knn_query(self, query, k, num_threads, nil, 1, features)
end
def knn_query(self = %T{}, query, opts) when is_list(query) do
k = Helper.get_keyword!(opts, :k, :pos_integer, 1)
num_threads = Helper.get_keyword!(opts, :num_threads, :integer, -1)
{rows, features} = Helper.list_of_binary(query)
Helper.ensure_vector_dimension!(self, features, true)
_do_knn_query(self, IO.iodata_to_binary(query), k, num_threads, nil, rows, features)
end
def knn_query(self = %T{}, query = %Nx.Tensor{}, opts) do
k = Helper.get_keyword!(opts, :k, :pos_integer, 1)
num_threads = Helper.get_keyword!(opts, :num_threads, :integer, -1)
{f32_data, rows, features} = Helper.verify_data_tensor!(self, query)
_do_knn_query(self, f32_data, k, num_threads, nil, rows, features)
end
defp _do_knn_query(self = %T{}, query, k, num_threads, filter, rows, features) do
case HNSWLib.Nif.index_knn_query(
self.reference,
query,
k,
num_threads,
filter,
rows,
features
) do
{:ok, labels, dists, rows, k, label_bits, dist_bits} ->
labels = Nx.reshape(Nx.from_binary(labels, :"u#{label_bits}"), {rows, k})
dists = Nx.reshape(Nx.from_binary(dists, :"f#{dist_bits}"), {rows, k})
{:ok, labels, dists}
{:error, reason} ->
{:error, reason}
end
end
@doc """
Get a list of existing IDs in the index.
"""
@spec get_ids_list(%T{}) :: {:ok, [integer()]} | {:error, String.t()}
def get_ids_list(self = %T{}) do
HNSWLib.Nif.index_get_ids_list(self.reference)
end
@doc """
Get the ef parameter.
"""
@spec get_ef(%T{}) :: {:ok, non_neg_integer()} | {:error, String.t()}
def get_ef(self = %T{}) do
HNSWLib.Nif.index_get_ef(self.reference)
end
@doc """
Set the ef parameter.
"""
@spec set_ef(%T{}, non_neg_integer()) :: :ok | {:error, String.t()}
def set_ef(self = %T{}, new_ef) when is_integer(new_ef) and new_ef >= 0 do
HNSWLib.Nif.index_set_ef(self.reference, new_ef)
end
@doc """
Get the number of threads to use.
"""
@spec get_num_threads(%T{}) :: {:ok, integer()} | {:error, String.t()}
def get_num_threads(self = %T{}) do
HNSWLib.Nif.index_get_num_threads(self.reference)
end
@doc """
Set the number of threads to use.
"""
@spec set_num_threads(%T{}, integer()) :: {:ok, integer()} | {:error, String.t()}
def set_num_threads(self = %T{}, new_num_threads) do
HNSWLib.Nif.index_set_num_threads(self.reference, new_num_threads)
end
@doc """
Save current index to disk.
##### Positional Parameters
- *path*: `Path.t()`.
Path to save the index to.
"""
@spec save_index(%T{}, Path.t()) :: {:ok, integer()} | {:error, String.t()}
def save_index(self = %T{}, path) when is_binary(path) do
HNSWLib.Nif.index_save_index(self.reference, path)
end
@doc """
Load index from disk.
##### Positional Parameters
- *space*: `:cosine` | `:ip` | `:l2`.
An atom that indicates the vector space. Valid values are
- `:cosine`, cosine space
- `:ip`, inner product space
- `:l2`, L2 space
- *dim*: `non_neg_integer()`.
Number of dimensions for each vector.
- *path*: `Path.t()`.
Path to load the index from.
##### Keyword Parameters
- *max_elements*: `non_neg_integer()`.
Maximum number of elements to load from the index.
If set to 0, all elements will be loaded.
Default: 0.
- *allow_replace_deleted*: `boolean()`.
"""
@spec load_index(:cosine | :ip | :l2, non_neg_integer(), Path.t(), [
{:max_elements, non_neg_integer()},
{:allow_replace_deleted, boolean()}
]) :: {:ok, %T{}} | {:error, String.t()}
def load_index(space, dim, path, opts \\ [])
when (space == :l2 or space == :ip or space == :cosine) and is_integer(dim) and dim >= 0 and
is_binary(path) and is_list(opts) do
max_elements = Helper.get_keyword!(opts, :max_elements, :non_neg_integer, 0)
allow_replace_deleted = Helper.get_keyword!(opts, :allow_replace_deleted, :boolean, false)
with {:ok, ref} <-
HNSWLib.Nif.index_load_index(space, dim, path, max_elements, allow_replace_deleted) do
{:ok,
%T{
space: space,
dim: dim,
reference: ref
}}
else
{:error, reason} ->
{:error, reason}
end
end
@doc """
Mark a label as deleted.
##### Positional Parameters
- *label*: `non_neg_integer()`.
Label to mark as deleted.
"""
@spec mark_deleted(%T{}, non_neg_integer()) :: :ok | {:error, String.t()}
def mark_deleted(self = %T{}, label) when is_integer(label) and label >= 0 do
HNSWLib.Nif.index_mark_deleted(self.reference, label)
end
@doc """
Unmark a label as deleted.
##### Positional Parameters
- *label*: `non_neg_integer()`.
Label to unmark as deleted.
"""
@spec unmark_deleted(%T{}, non_neg_integer()) :: :ok | {:error, String.t()}
def unmark_deleted(self = %T{}, label) when is_integer(label) and label >= 0 do
HNSWLib.Nif.index_unmark_deleted(self.reference, label)
end
@doc """
Add items to the index.
##### Positional Parameters
- *data*: `Nx.Tensor.t()`.
Data to add to the index.
##### Keyword Parameters
- *ids*: `Nx.Tensor.t() | [non_neg_integer()] | nil`.
IDs to assign to the data.
If `nil`, IDs will be assigned sequentially starting from 0.
Defaults to `nil`.
- *num_threads*: `integer()`.
Number of threads to use.
If set to `-1`, the number of threads will be automatically determined.
Defaults to `-1`.
- *replace_deleted*: `boolean()`.
Whether to replace deleted items.
Defaults to `false`.
"""
@spec add_items(%T{}, Nx.Tensor.t(), [
{:ids, Nx.Tensor.t() | [non_neg_integer()] | nil},
{:num_threads, integer()},
{:replace_deleted, false}
]) :: :ok | {:error, String.t()}
def add_items(self, data, opts \\ [])
def add_items(self = %T{}, data = %Nx.Tensor{}, opts) when is_list(opts) do
num_threads = Helper.get_keyword!(opts, :num_threads, :integer, -1)
replace_deleted = Helper.get_keyword!(opts, :replace_deleted, :boolean, false)
ids = Helper.normalize_ids!(opts[:ids])
{f32_data, rows, features} = Helper.verify_data_tensor!(self, data)
HNSWLib.Nif.index_add_items(
self.reference,
f32_data,
ids,
num_threads,
replace_deleted,
rows,
features
)
end
@doc """
Retrieve items from the index using IDs.
##### Positional Parameters
- *ids*: `Nx.Tensor.t() | [non_neg_integer()]`.
IDs to retrieve.
"""
@spec get_items(%T{}, Nx.Tensor.t() | [integer()]) :: {:ok, [binary()]} | {:error, String.t()}
def get_items(self = %T{}, ids) do
ids = Helper.normalize_ids!(ids)
HNSWLib.Nif.index_get_items(self.reference, ids)
end
@doc """
Resize the index.
##### Positional Parameters
- *new_size*: `non_neg_integer()`.
New size of the index.
"""
@spec resize_index(%T{}, non_neg_integer()) :: :ok | {:error, String.t()}
def resize_index(self = %T{}, new_size) when is_integer(new_size) and new_size >= 0 do
HNSWLib.Nif.index_resize_index(self.reference, new_size)
end
@doc """
Get the maximum number of elements the index can hold.
"""
@spec get_max_elements(%T{}) :: {:ok, integer()} | {:error, String.t()}
def get_max_elements(self = %T{}) do
HNSWLib.Nif.index_get_max_elements(self.reference)
end
@doc """
Get the current number of elements in the index.
"""
@spec get_current_count(%T{}) :: {:ok, integer()} | {:error, String.t()}
def get_current_count(self = %T{}) do
HNSWLib.Nif.index_get_current_count(self.reference)
end
@doc """
Get the ef_construction parameter.
"""
@spec get_ef_construction(%T{}) :: {:ok, integer()} | {:error, String.t()}
def get_ef_construction(self = %T{}) do
HNSWLib.Nif.index_get_ef_construction(self.reference)
end
@doc """
Get the M parameter.
"""
@spec get_m(%T{}) :: {:ok, integer()} | {:error, String.t()}
def get_m(self = %T{}) do
HNSWLib.Nif.index_get_m(self.reference)
end
end
