usearch 0.21.0 → 0.22.1

package/python/README.md

@@ -0,0 +1,283 @@
+# USearch for Python
+
+## Installation
+
+```sh
+pip install usearch
+```
+
+## Quickstart
+
+```python
+import numpy as np
+from usearch.index import Index
+
+index = Index(
+    ndim=3, # Define the number of dimensions in input vectors
+    metric='cos', # Choose 'l2sq', 'haversine' or other metric, default = 'ip'
+    dtype='f32', # Quantize to 'f16' or 'f8' if needed, default = 'f32'
+    connectivity=16, # How frequent should the connections in the graph be, optional
+    expansion_add=128, # Control the recall of indexing, optional
+    expansion_search=64, # Control the quality of search, optional
+)
+
+vector = np.array([0.2, 0.6, 0.4])
+index.add(42, vector)
+matches, distances, count = index.search(vector, 10)
+
+assert len(index) == 1
+assert count == 1
+assert matches[0] == 42
+assert distances[0] <= 0.001
+assert np.allclose(index[42], vector)
+```
+
+Python bindings are implemented with [`pybind/pybind11`](https://github.com/pybind/pybind11).
+Assuming the presence of Global Interpreter Lock in Python, we spawn threads in the C++ layer on large insertions.
+
+## Serialization
+
+```py
+index.save('index.usearch')
+index.load('index.usearch') # Copy the whole index into memory
+index.view('index.usearch') # View from disk without loading in memory
+```
+
+If you don't know anything about the index except its path, there are two more endpoints to know:
+
+```py
+Index.metadata('index.usearch') -> IndexMetadata
+Index.restore('index.usearch', view=False) -> Index
+```
+
+## Batch Operations
+
+Adding or querying a batch of entries is identical to adding a single vector.
+The difference would be in the shape of the tensors.
+
+```py
+n = 100
+labels = np.arange(n)
+vectors = np.random.uniform(0, 0.3, (n, index.ndim)).astype(np.float32)
+
+index.add(labels, vectors, threads=..., copy=...)
+matches, distances, counts = index.search(vectors, 10, threads=...)
+
+assert matches.shape[0] == vectors.shape[0]
+assert counts[0] <= 10
+```
+
+You can also override the default `threads` and `copy` arguments in bulk workloads.
+The first controls the number of threads spawned for the task.
+The second controls whether the vector itself will be persisted inside the index.
+If you can preserve the lifetime of the vector somewhere else, you can avoid the copy.
+
+## User-Defined Metrics and JIT in Python
+
+### [Numba][numba]
+
+Assuming the language boundary exists between Python user code and C++ implementation, there are more efficient solutions than passing a Python callable to the engine.
+Luckily, with the help of [Numba][numba], we can JIT compile a function with a matching signature and pass it down to the engine.
+
+```py
+from numba import cfunc, types, carray
+
+ndim = 256
+signature = types.float32(
+    types.CPointer(types.float32),
+    types.CPointer(types.float32))
+
+@cfunc(signature)
+def inner_product(a, b):
+    a_array = carray(a, ndim)
+    b_array = carray(b, ndim)
+    c = 0.0
+    for i in range(ndim):
+        c += a_array[i] * b_array[i]
+    return 1 - c
+
+index = Index(ndim=ndim, metric=CompiledMetric(
+    pointer=inner_product.address,
+    kind=MetricKind.IP,
+    signature=MetricSignature.ArrayArray,
+))
+```
+
+Alternatively, you can avoid pre-defining the number of dimensions, and pass it separately:
+
+```py
+signature = types.float32(
+    types.CPointer(types.float32),
+    types.CPointer(types.float32),
+    types.uint64)
+
+@cfunc(signature)
+def inner_product(a, b, ndim):
+    a_array = carray(a, ndim)
+    b_array = carray(b, ndim)
+    c = 0.0
+    for i in range(ndim):
+        c += a_array[i] * b_array[i]
+    return 1 - c
+
+index = Index(ndim=ndim, metric=CompiledMetric(
+    pointer=inner_product.address,
+    kind=MetricKind.IP,
+    signature=MetricSignature.ArrayArraySize,
+))
+```
+
+```sh
+pip install numba
+```
+
+### [Cppyy][cppyy]
+
+Similarly, you can use Cppyy with Cling to JIT-compile native C or C++ code and pass it to USearch, which may be a good idea, if you want to explicitly request loop-unrolling or other low-level optimizations!
+
+```py
+import cppyy
+import cppyy.ll
+
+cppyy.cppdef("""
+float inner_product(float *a, float *b) {
+    float result = 0;
+#pragma unroll
+    for (size_t i = 0; i != ndim; ++i)
+        result += a[i] * b[i];
+    return 1 - result;
+}
+""".replace("ndim", str(ndim)))
+
+function = cppyy.gbl.inner_product
+index = Index(ndim=ndim, metric=CompiledMetric(
+    pointer=cppyy.ll.addressof(function),
+    kind=MetricKind.IP,
+    signature=MetricSignature.ArrayArraySize,
+))
+```
+
+```sh
+conda install -c conda-forge cppyy
+```
+
+### [PeachPy][peachpy]
+
+We have covered JIT-ing Python with Numba and C++ with Cppyy and Cling.
+How about writing Assembly directly?
+That is also possible.
+Below is an example of constructing the "Inner Product" distance for 8-dimensional `f32` vectors for x86 using [PeachPy][peachpy].
+
+```py
+from peachpy import (
+    Argument,
+    ptr,
+    float_,
+    const_float_,
+)
+from peachpy.x86_64 import (
+    abi,
+    Function,
+    uarch,
+    isa,
+    GeneralPurposeRegister64,
+    LOAD,
+    YMMRegister,
+    VSUBPS,
+    VADDPS,
+    VHADDPS,
+    VMOVUPS,
+    VFMADD231PS,
+    VPERM2F128,
+    VXORPS,
+    RETURN,
+)
+
+a = Argument(ptr(const_float_), name="a")
+b = Argument(ptr(const_float_), name="b")
+
+with Function(
+    "inner_product", (a, b), float_, target=uarch.default + isa.avx2
+) as asm_function:
+    
+    # Request two 64-bit general-purpose registers for addresses
+    reg_a, reg_b = GeneralPurposeRegister64(), GeneralPurposeRegister64()
+    LOAD.ARGUMENT(reg_a, a)
+    LOAD.ARGUMENT(reg_b, b)
+
+    # Load the vectors
+    ymm_a = YMMRegister()
+    ymm_b = YMMRegister()
+    VMOVUPS(ymm_a, [reg_a])
+    VMOVUPS(ymm_b, [reg_b])
+
+    # Prepare the accumulator
+    ymm_c = YMMRegister()
+    ymm_one = YMMRegister()
+    VXORPS(ymm_c, ymm_c, ymm_c)
+    VXORPS(ymm_one, ymm_one, ymm_one)
+
+    # Accumulate A and B product into C
+    VFMADD231PS(ymm_c, ymm_a, ymm_b)
+
+    # Reduce the contents of a YMM register
+    ymm_c_permuted = YMMRegister()
+    VPERM2F128(ymm_c_permuted, ymm_c, ymm_c, 1)
+    VADDPS(ymm_c, ymm_c, ymm_c_permuted)
+    VHADDPS(ymm_c, ymm_c, ymm_c)
+    VHADDPS(ymm_c, ymm_c, ymm_c)
+
+    # Negate the values, to go from "similarity" to "distance"
+    VSUBPS(ymm_c, ymm_one, ymm_c)
+
+    # A common convention is to return floats in XMM registers
+    RETURN(ymm_c.as_xmm)
+
+python_function = asm_function.finalize(abi.detect()).encode().load()
+metric = CompiledMetric(
+    pointer=python_function.loader.code_address,
+    kind=MetricKind.IP,
+    signature=MetricSignature.ArrayArray,
+)
+index = Index(ndim=ndim, metric=metric)
+```
+
+[numba]: https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
+[cppyy]: https://cppyy.readthedocs.io/en/latest/
+[peachpy]: https://github.com/Maratyszcza/PeachPy
+
+## Tooling
+
+To work with `bbin`, `fbin`, `ibin`, `hbin` matrix files USearch provides `load_matrix` and `save_matrix`.
+Such files are standard in k-ANN tasks and represent a binary object with all the scalars, prepended by two 32-bit integers - the number of rows and columns in the matrix.
+
+```py
+from usearch.index import Index
+from usearch.io import load_matrix, save_matrix
+
+vectors = load_matrix('deep1B.fbin')
+index = Index(ndim=vectors.shape[1])
+index.add(labels, vectors)
+```
+
+One may often want to evaluate the quality of the constructed index before running in production.
+The trivial way is to measure `recall@1` on the entries already present in the index.
+
+```py
+from usearch.eval import recall_members
+
+assert recall_members(index, exact=True) == 1
+print(recall_members(index, exact=False))
+```
+
+In case you have some ground-truth data for more than one entry, you compare search results against expected values:
+
+```py
+from usearch.eval import relevance, dcg, ndcg, random_vectors
+
+vectors = random_vectors(index=index)
+matches_approximate = index.search(vectors)
+matches_exact = index.search(vectors, exact=True)
+relevance_scores = relevance(matches_exact, matches_approximate)
+print(dcg(relevance_scores), ndcg(relevance_scores))
+```

package/python/lib.cpp

@@ -15,6 +15,7 @@
 #include <limits> // `std::numeric_limits`
 #include <thread> // `std::thread`
 
+#define _CRT_SECURE_NO_WARNINGS
 #define PY_SSIZE_T_CLEAN
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -45,6 +46,7 @@ using distance_t = punned_distance_t;
 using dense_index_t = punned_small_t;
 using dense_add_result_t = typename dense_index_t::add_result_t;
 using dense_search_result_t = typename dense_index_t::search_result_t;
+using dense_labeling_result_t = typename dense_index_t::labeling_result_t;
 
 struct dense_index_py_t : public dense_index_t {
     using native_t = dense_index_t;
@@ -181,8 +183,8 @@ static void add_one_to_index(dense_index_py_t& index, label_t label, py::buffer
     if (vector_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
         throw std::invalid_argument("The number of vector dimensions doesn't match!");
 
-    if (index.size() + 1 >= index.capacity())
-        index.reserve(ceil2(index.size() + 1));
+    if (!index.reserve(ceil2(index.size() + 1)))
+        throw std::invalid_argument("Out of memory!");
 
     add_config_t config;
     config.store_vector = copy;
@@ -224,9 +226,6 @@ static void add_many_to_index(                                      //
     dense_index_py_t& index, py::buffer labels, py::buffer vectors, //
     bool copy, std::size_t threads) {
 
-    if (index.limits().threads_add < threads)
-        throw std::invalid_argument("Can't use that many threads!");
-
     py::buffer_info labels_info = labels.request();
     py::buffer_info vectors_info = vectors.request();
 
@@ -248,8 +247,10 @@ static void add_many_to_index(                                      //
     if (labels_count != vectors_count)
         throw std::invalid_argument("Number of labels and vectors must match!");
 
-    if (index.size() + vectors_count >= index.capacity())
-        index.reserve(ceil2(index.size() + vectors_count));
+    if (!threads)
+        threads = std::thread::hardware_concurrency();
+    if (!index.reserve(index_limits_t(ceil2(index.size() + vectors_count), threads)))
+        throw std::invalid_argument("Out of memory!");
 
     switch (numpy_string_to_kind(vectors_info.format)) {
     case scalar_kind_t::b1x8_k: add_typed_to_index<b1x8_t>(index, labels_info, vectors_info, copy, threads); break;
@@ -317,6 +318,11 @@ static void search_typed(                                   //
     Py_ssize_t vectors_count = vectors_info.shape[0];
     char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
 
+    if (!threads)
+        threads = std::thread::hardware_concurrency();
+    if (!index.reserve(index_limits_t(index.size(), threads)))
+        throw std::invalid_argument("Out of memory!");
+
     executor_default_t{threads}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
         search_config_t config;
         config.thread = thread_idx;
@@ -450,16 +456,17 @@ template <typename index_at> py::object get_member(index_at const& index, label_
         throw std::invalid_argument("Incompatible scalars in the query matrix!");
 }
 
-template <typename index_at> py::array_t<label_t> get_labels(index_at const& index, std::size_t limit) {
+template <typename index_at>
+py::array_t<label_t> get_labels(index_at const& index, std::size_t offset, std::size_t limit) {
     limit = std::min(index.size(), limit);
     py::array_t<label_t> result_py(static_cast<Py_ssize_t>(limit));
     auto result_py1d = result_py.template mutable_unchecked<1>();
-    index.export_labels(&result_py1d(0), limit);
+    index.export_labels(&result_py1d(0), offset, limit);
     return result_py;
 }
 
 template <typename index_at> py::array_t<label_t> get_all_labels(index_at const& index) {
-    return get_labels(index, index.size());
+    return get_labels(index, 0, index.size());
 }
 
 template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
@@ -666,6 +673,48 @@ PYBIND11_MODULE(compiled, m) {
         py::arg("threads") = 0           //
     );
 
+    i.def(
+        "rename",
+        [](dense_index_py_t& index, label_t from, label_t to) -> bool {
+            dense_labeling_result_t result = index.rename(from, to);
+            result.error.raise();
+            return result.completed;
+        },
+        py::arg("from"), py::arg("to"));
+
+    i.def(
+        "remove",
+        [](dense_index_py_t& index, label_t label, bool compact, std::size_t threads) -> bool {
+            dense_labeling_result_t result = index.remove(label);
+            result.error.raise();
+
+            if (!threads)
+                threads = std::thread::hardware_concurrency();
+            if (!index.reserve(index_limits_t(index.size(), threads)))
+                throw std::invalid_argument("Out of memory!");
+
+            index.compact(executor_default_t{threads});
+            return result.completed;
+        },
+        py::arg("label"), py::arg("compact"), py::arg("threads"));
+
+    i.def(
+        "remove",
+        [](dense_index_py_t& index, std::vector<label_t> const& labels, bool compact,
+           std::size_t threads) -> std::size_t {
+            dense_labeling_result_t result = index.remove(labels.begin(), labels.end());
+            result.error.raise();
+
+            if (!threads)
+                threads = std::thread::hardware_concurrency();
+            if (!index.reserve(index_limits_t(index.size(), threads)))
+                throw std::invalid_argument("Out of memory!");
+
+            index.compact(executor_default_t{threads});
+            return result.completed;
+        },
+        py::arg("label"), py::arg("compact"), py::arg("threads"));
+
     i.def("__len__", &dense_index_py_t::size);
     i.def_property_readonly("size", &dense_index_py_t::size);
     i.def_property_readonly("ndim", &dense_index_py_t::dimensions);
@@ -681,7 +730,8 @@ PYBIND11_MODULE(compiled, m) {
     i.def_property("expansion_search", &dense_index_py_t::expansion_search, &dense_index_py_t::change_expansion_search);
 
     i.def_property_readonly("labels", &get_all_labels<dense_index_py_t>);
-    i.def("get_labels", &get_labels<dense_index_py_t>, py::arg("limit") = std::numeric_limits<std::size_t>::max());
+    i.def("get_labels", &get_labels<dense_index_py_t>, py::arg("offset") = 0,
+          py::arg("limit") = std::numeric_limits<std::size_t>::max());
     i.def("__contains__", &dense_index_py_t::contains);
     i.def("__getitem__", &get_member<dense_index_py_t>, py::arg("label"), py::arg("dtype") = scalar_kind_t::f32_k);
 

package/python/scripts/join.py

@@ -5,7 +5,11 @@ wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/image
 wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/texts.fbin -P datasets/cc_3M/
 wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/uform_english_images.fbin -P datasets/cc_3M/
 wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/uform_english_texts.fbin -P datasets/cc_3M/
-rm -rf datasets/cc_3M/*.usearch
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clipbigg_images.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clipbigg_texts.fbin -P datasets/cc_3M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-arxiv-2m/resolve/main/e5base_abstract.fbin -P datasets/arxiv_2M/
+wget -nc https://huggingface.co/datasets/unum-cloud/ann-arxiv-2m/resolve/main/e5base_title.fbin -P datasets/arxiv_2M/
+rm -rf datasets/cc_3M/*.usearch datasets/arxiv_2M/*.usearch
 python python/scripts/join.py
 """
 from numpy import dot
@@ -45,13 +49,13 @@ metric = CompiledMetric(
 a = Index(
     a_mat.shape[1],
     metric=metric,
-    path=f"datasets/{a_name}.f32.usearch",
+    path=f"datasets/{a_name}-{max_elements}.f32.usearch",
     dtype="f32",
 )
 b = Index(
     b_mat.shape[1],
     metric=metric,
-    path=f"datasets/{b_name}.f32.usearch",
+    path=f"datasets/{b_name}-{max_elements}.f32.usearch",
     dtype="f32",
 )