usearch 0.18.6 → 0.19.0

package/python/lib.cpp

@@ -13,8 +13,9 @@
  *  @copyright Copyright (c) 2023
  */
 #define PY_SSIZE_T_CLEAN
-#define NOMINMAX // Some of our dependencies call `std::max(x, y)`, which crashes Windows builds
-#include <thread>
+#define NOMINMAX  // Some of our dependencies call `std::max(x, y)`, which crashes Windows builds
+#include <limits> // `std::numeric_limits`
+#include <thread> // `std::thread`
 
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -24,12 +25,23 @@
 using namespace unum::usearch;
 using namespace unum;
 
+/**
+ *  @brief  The signature of the user-defined function.
+ *          Can be just two array pointers, precompiled for a specific array length,
+ *          or include one or two array sizes as 64-bit unsigned integers.
+ */
+enum class metric_signature_t {
+    array_array_k = 0,
+    array_array_size_k,
+    array_size_array_size_k,
+};
+
 namespace py = pybind11;
 using py_shape_t = py::array::ShapeContainer;
 
 using label_t = std::uint32_t;
 using distance_t = punned_distance_t;
-using metric_t = punned_metric_t;
+using metric_t = index_punned_dense_metric_t;
 using id_t = std::uint32_t;
 using big_id_t = std::uint64_t;
 
@@ -50,7 +62,7 @@ struct punned_index_py_t : public punned_index_t {
 
 using set_member_t = std::uint32_t;
 using set_view_t = span_gt<set_member_t const>;
-using sparse_index_t = index_gt<jaccard_gt<set_member_t>, label_t, id_t, set_member_t>;
+using sparse_index_t = index_gt<jaccard_gt<set_member_t>, label_t, id_t>;
 
 struct sparse_index_py_t : public sparse_index_t {
     using native_t = sparse_index_t;
@@ -63,32 +75,58 @@ struct sparse_index_py_t : public sparse_index_t {
     sparse_index_py_t(native_t&& base) : native_t(std::move(base)) {}
 };
 
-template <typename scalar_at> metric_t udf(std::size_t metric_uintptr) {
-    return [metric_uintptr](punned_vector_view_t a, punned_vector_view_t b) -> distance_t {
-        using metric_raw_t = punned_distance_t (*)(scalar_at const*, scalar_at const*);
-        metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
-        return metric_ptr((scalar_at const*)a.data(), (scalar_at const*)b.data());
-    };
+template <typename scalar_at>
+metric_t udf(metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr) {
+    //
+    metric_t result;
+    result.kind_ = kind;
+    switch (signature) {
+    case metric_signature_t::array_array_k:
+        result.func_ = [metric_uintptr](punned_vector_view_t a, punned_vector_view_t b) -> distance_t {
+            using metric_raw_t = punned_distance_t (*)(scalar_at const*, scalar_at const*);
+            metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
+            return metric_ptr((scalar_at const*)a.data(), (scalar_at const*)b.data());
+        };
+        break;
+    case metric_signature_t::array_array_size_k:
+        result.func_ = [metric_uintptr](punned_vector_view_t a, punned_vector_view_t b) -> distance_t {
+            using metric_raw_t = punned_distance_t (*)(scalar_at const*, scalar_at const*, size_t);
+            metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
+            return metric_ptr((scalar_at const*)a.data(), (scalar_at const*)b.data(), a.size() / sizeof(scalar_at));
+        };
+        break;
+    case metric_signature_t::array_size_array_size_k:
+        result.func_ = [metric_uintptr](punned_vector_view_t a, punned_vector_view_t b) -> distance_t {
+            using metric_raw_t = punned_distance_t (*)(scalar_at const*, size_t, scalar_at const*, size_t);
+            metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
+            return metric_ptr(                                            //
+                (scalar_at const*)a.data(), a.size() / sizeof(scalar_at), //
+                (scalar_at const*)b.data(), b.size() / sizeof(scalar_at));
+        };
+        break;
+    }
+    return result;
 }
 
-metric_t udf(std::size_t metric_uintptr, scalar_kind_t accuracy) {
+metric_t udf(metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, scalar_kind_t accuracy) {
     switch (accuracy) {
-    case scalar_kind_t::f8_k: return udf<f8_bits_t>(metric_uintptr);
-    case scalar_kind_t::f16_k: return udf<f16_t>(metric_uintptr);
-    case scalar_kind_t::f32_k: return udf<f32_t>(metric_uintptr);
-    case scalar_kind_t::f64_k: return udf<f64_t>(metric_uintptr);
+    case scalar_kind_t::f8_k: return udf<f8_bits_t>(kind, signature, metric_uintptr);
+    case scalar_kind_t::f16_k: return udf<f16_t>(kind, signature, metric_uintptr);
+    case scalar_kind_t::f32_k: return udf<f32_t>(kind, signature, metric_uintptr);
+    case scalar_kind_t::f64_k: return udf<f64_t>(kind, signature, metric_uintptr);
     default: return {};
     }
 }
 
-static punned_index_py_t make_index( //
-    std::size_t dimensions,          //
-    scalar_kind_t scalar_kind,       //
-    metric_kind_t metric_kind,       //
-    std::size_t connectivity,        //
-    std::size_t expansion_add,       //
-    std::size_t expansion_search,    //
-    std::size_t metric_uintptr,      //
+static punned_index_py_t make_index(     //
+    std::size_t dimensions,              //
+    scalar_kind_t scalar_kind,           //
+    metric_kind_t metric_kind,           //
+    std::size_t connectivity,            //
+    std::size_t expansion_add,           //
+    std::size_t expansion_search,        //
+    metric_signature_t metric_signature, //
+    std::uintptr_t metric_uintptr,       //
     bool tune) {
 
     index_config_t config;
@@ -99,7 +137,8 @@ static punned_index_py_t make_index( //
 
     if (metric_uintptr)
         return punned_index_t::make( //
-            dimensions, udf(metric_uintptr, scalar_kind), config, scalar_kind, expansion_add, expansion_search);
+            dimensions, udf(metric_kind, metric_signature, metric_uintptr, scalar_kind), config, scalar_kind,
+            expansion_add, expansion_search);
     else
         return punned_index_t::make(dimensions, metric_kind, config, scalar_kind, expansion_add, expansion_search);
 }
@@ -417,14 +456,18 @@ 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 result_length = index.size();
-    py::array_t<label_t> result_py(static_cast<Py_ssize_t>(result_length));
+template <typename index_at> py::array_t<label_t> get_labels(index_at const& index, 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), result_length);
+    index.export_labels(&result_py1d(0), 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());
+}
+
 template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
     if (begin == end)
         return false;
@@ -539,6 +582,11 @@ PYBIND11_MODULE(compiled, m) {
     m.attr("USES_SIMSIMD") = py::int_(USEARCH_USE_SIMSIMD);
     m.attr("USES_NATIVE_F16") = py::int_(USEARCH_USE_NATIVE_F16);
 
+    py::enum_<metric_signature_t>(m, "MetricSignature")
+        .value("ArrayArray", metric_signature_t::array_array_k)
+        .value("ArrayArraySize", metric_signature_t::array_array_size_k)
+        .value("ArraySizeArraySize", metric_signature_t::array_size_array_size_k);
+
     py::enum_<metric_kind_t>(m, "MetricKind")
         .value("Unknown", metric_kind_t::unknown_k)
         .value("IP", metric_kind_t::ip_k)
@@ -558,18 +606,44 @@ PYBIND11_MODULE(compiled, m) {
         .value("F8", scalar_kind_t::f8_k)
         .value("B1", scalar_kind_t::b1x8_k);
 
+    auto h = py::class_<file_head_result_t>(m, "IndexMetadata");
+    h.def(py::init([](std::string const& path) {
+        file_head_result_t h = index_metadata(path.c_str());
+        h.error.raise();
+        return h;
+    }));
+    h.def_property_readonly("version", [](file_head_result_t const& h) {
+        return                                      //
+            std::to_string(h.version_major) + "." + //
+            std::to_string(h.version_minor) + "." + //
+            std::to_string(h.version_patch);
+    });
+    h.def_readonly("metric", &file_head_result_t::metric);
+    h.def_readonly("connectivity", &file_head_result_t::connectivity);
+    h.def_readonly("max_level", &file_head_result_t::max_level);
+    h.def_readonly("vector_alignment", &file_head_result_t::vector_alignment);
+    h.def_readonly("bytes_per_label", &file_head_result_t::bytes_per_label);
+    h.def_readonly("bytes_per_id", &file_head_result_t::bytes_per_id);
+    h.def_readonly("scalar_kind", &file_head_result_t::scalar_kind);
+    h.def_readonly("size", &file_head_result_t::size);
+    h.def_readonly("entry_idx", &file_head_result_t::entry_idx);
+    h.def_readonly("bytes_for_graphs", &file_head_result_t::bytes_for_graphs);
+    h.def_readonly("bytes_for_vectors", &file_head_result_t::bytes_for_vectors);
+    h.def_readonly("bytes_checksum", &file_head_result_t::bytes_checksum);
+
     auto i = py::class_<punned_index_py_t>(m, "Index");
 
-    i.def(py::init(&make_index),                                    //
-          py::kw_only(),                                            //
-          py::arg("ndim") = 0,                                      //
-          py::arg("dtype") = scalar_kind_t::f32_k,                  //
-          py::arg("metric") = metric_kind_t::ip_k,                  //
-          py::arg("connectivity") = default_connectivity(),         //
-          py::arg("expansion_add") = default_expansion_add(),       //
-          py::arg("expansion_search") = default_expansion_search(), //
-          py::arg("metric_pointer") = 0,                            //
-          py::arg("tune") = false                                   //
+    i.def(py::init(&make_index),                                           //
+          py::kw_only(),                                                   //
+          py::arg("ndim") = 0,                                             //
+          py::arg("dtype") = scalar_kind_t::f32_k,                         //
+          py::arg("metric") = metric_kind_t::ip_k,                         //
+          py::arg("connectivity") = default_connectivity(),                //
+          py::arg("expansion_add") = default_expansion_add(),              //
+          py::arg("expansion_search") = default_expansion_search(),        //
+          py::arg("metric_signature") = metric_signature_t::array_array_k, //
+          py::arg("metric_pointer") = 0,                                   //
+          py::arg("tune") = false                                          //
     );
 
     i.def(                         //
@@ -604,7 +678,8 @@ PYBIND11_MODULE(compiled, m) {
     i.def_property_readonly("connectivity", &punned_index_py_t::connectivity);
     i.def_property_readonly("capacity", &punned_index_py_t::capacity);
     i.def_property_readonly( //
-        "dtype", [](punned_index_py_t const& index) -> std::string { return scalar_kind_name(index.scalar_kind()); });
+        "dtype",
+        [](punned_index_py_t const& index) -> std::string { return scalar_kind_name(index.metric().scalar_kind_); });
     i.def_property_readonly( //
         "memory_usage", [](punned_index_py_t const& index) -> std::size_t { return index.memory_usage(); });
 
@@ -612,7 +687,8 @@ PYBIND11_MODULE(compiled, m) {
     i.def_property("expansion_search", &punned_index_py_t::expansion_search,
                    &punned_index_py_t::change_expansion_search);
 
-    i.def_property_readonly("labels", &get_labels<punned_index_py_t>);
+    i.def_property_readonly("labels", &get_all_labels<punned_index_py_t>);
+    i.def("get_labels", &get_labels<punned_index_py_t>, py::arg("limit") = std::numeric_limits<std::size_t>::max());
     i.def("__contains__", &punned_index_py_t::contains);
     i.def( //
         "__getitem__", &get_member<punned_index_py_t>, py::arg("label"), py::arg("dtype") = scalar_kind_t::f32_k);

package/python/scripts/bench.ipynb

@@ -6,7 +6,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install pandas nbformat jinja2 kaleido tqdm sklearn scikit-learn scipy matplotlib"
+    "!pip install faiss-cpu numba pandas plotly nbformat jinja2 kaleido tqdm sklearn scikit-learn scipy matplotlib"
    ]
   },
   {
@@ -75,7 +75,7 @@
    "outputs": [],
    "source": [
     "eval_default = Evaluation.for_dataset(data)\n",
-    "results_default = bench_speed(eval_default)\n",
+    "results_default = bench_speed(eval_default, jit=True)\n",
     "results_default"
    ]
   },
@@ -95,7 +95,7 @@
    "outputs": [],
    "source": [
     "eval_sliced = Evaluation.for_dataset(data, batch_size=256)\n",
-    "results_sliced = bench_speed(eval_sliced)\n",
+    "results_sliced = bench_speed(eval_sliced, jit=True)\n",
     "results_sliced"
    ]
   },
@@ -161,7 +161,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.3"
+   "version": "3.10.10"
   },
   "orig_nbformat": 4
  },

package/python/scripts/bench.py

@@ -21,15 +21,16 @@ def bench_speed(
     connectivity: int = DEFAULT_CONNECTIVITY,
     expansion_add: int = DEFAULT_EXPANSION_ADD,
     expansion_search: int = DEFAULT_EXPANSION_SEARCH,
+    jit: bool = False,
     train: bool = False,
 ) -> pd.DataFrame:
-
     # Build various indexes:
     indexes = []
     jit_options = [False]
-    dtype_options = ['f32', 'f16', 'f8']
+    if jit:
+        jit_options.append(jit)
+    dtype_options = ["f32", "f16", "f8"]
     for jit, dtype in itertools.product(jit_options, dtype_options):
-
         index = Index(
             ndim=eval.ndim,
             jit=jit,
@@ -37,48 +38,71 @@ def bench_speed(
             expansion_add=expansion_add,
             expansion_search=expansion_search,
             connectivity=connectivity,
-            path='USearch' + ['', '+JIT'][jit] + ':' + dtype,
+            path="USearch" + ["", "+JIT"][jit] + ":" + dtype,
         )
+
+        # Skip the cases, where JIT-ing is impossible
+        if jit and not index.jit:
+            continue
         indexes.append(index)
 
     # Add FAISS indexes to the mix:
     try:
         from index_faiss import IndexFAISS, IndexQuantizedFAISS
-        indexes.append(IndexFAISS(
-            ndim=eval.ndim,
-            expansion_add=expansion_add,
-            expansion_search=expansion_search,
-            connectivity=connectivity,
-            path='FAISS:f32',
-        ))
-        if train:
-            indexes.append(IndexQuantizedFAISS(
-                train=eval.tasks[0].vectors,
+
+        indexes.append(
+            IndexFAISS(
+                ndim=eval.ndim,
                 expansion_add=expansion_add,
                 expansion_search=expansion_search,
                 connectivity=connectivity,
-                path='FAISS+IVFPQ:f32',
-            ))
+                path="FAISS:f32",
+            )
+        )
+        if train:
+            indexes.append(
+                IndexQuantizedFAISS(
+                    train=eval.tasks[0].vectors,
+                    expansion_add=expansion_add,
+                    expansion_search=expansion_search,
+                    connectivity=connectivity,
+                    path="FAISS+IVFPQ:f32",
+                )
+            )
     except (ImportError, ModuleNotFoundError):
         pass
 
     # Time to evaluate:
     results = [eval(index) for index in indexes]
-    return pd.DataFrame({
-        'names': [i.path for i in indexes],
-        'add_per_second': [x['add_per_second'] for x in results],
-        'search_per_second': [x['search_per_second'] for x in results],
-        'recall_at_one': [x['recall_at_one'] for x in results],
-    })
+    return pd.DataFrame(
+        {
+            "names": [i.path for i in indexes],
+            "add_per_second": [x["add_per_second"] for x in results],
+            "search_per_second": [x["search_per_second"] for x in results],
+            "recall_at_one": [x["recall_at_one"] for x in results],
+        }
+    )
 
 
 def bench_params(
     count: int = 1_000_000,
     connectivities: int = range(10, 20),
     dimensions: List[int] = [
-        2, 3, 4, 8, 16, 32, 96, 100,
-        256, 384, 512, 768, 1024, 1536],
-
+        2,
+        3,
+        4,
+        8,
+        16,
+        32,
+        96,
+        100,
+        256,
+        384,
+        512,
+        768,
+        1024,
+        1536,
+    ],
     expansion_add: int = DEFAULT_EXPANSION_ADD,
     expansion_search: int = DEFAULT_EXPANSION_SEARCH,
 ) -> pd.DataFrame:
@@ -90,7 +114,6 @@ def bench_params(
 
     results = []
     for connectivity, ndim in itertools.product(connectivities, dimensions):
-
         task = AddTask(
             labels=np.arange(count, dtype=Label),
             vectors=np.random.rand(count, ndim).astype(np.float32),
@@ -102,8 +125,8 @@ def bench_params(
             expansion_search=expansion_search,
         )
         result = asdict(task(index))
-        result['ndim'] = dimensions
-        result['connectivity'] = connectivity
+        result["ndim"] = dimensions
+        result["connectivity"] = connectivity
         results.append(result)
 
     # return self._execute_tasks(

package/python/scripts/index_faiss.py

@@ -13,7 +13,6 @@ from usearch.index import (
 
 
 class IndexFAISS:
-
     def __init__(
         self,
         index: IndexHNSWFlat = None,
@@ -22,9 +21,9 @@ class IndexFAISS:
         expansion_add: int = DEFAULT_EXPANSION_ADD,
         expansion_search: int = DEFAULT_EXPANSION_SEARCH,
         path: Optional[os.PathLike] = None,
-        *args, **kwargs,
+        *args,
+        **kwargs,
     ):
-
         if index is None:
             index = IndexHNSWFlat(ndim, connectivity)
             index.hnsw.efConstruction = expansion_add
@@ -32,11 +31,11 @@ class IndexFAISS:
 
         self._faiss = index
         self._specs = {
-            'Class': 'usearch.IndexFAISS',
-            'Dimensions': ndim,
-            'Connectivity': connectivity,
-            'Expansion@Add': expansion_add,
-            'Expansion@Search': expansion_search,
+            "Class": "usearch.IndexFAISS",
+            "Dimensions": ndim,
+            "Connectivity": connectivity,
+            "Expansion@Add": expansion_add,
+            "Expansion@Search": expansion_search,
         }
 
         self.path = path
@@ -58,9 +57,11 @@ class IndexFAISS:
 
     @property
     def specs(self) -> dict:
-        self._specs.update({
-            'Size': len(self),
-        })
+        self._specs.update(
+            {
+                "Size": len(self),
+            }
+        )
         return self._specs
 
     def load(self, path: os.PathLike):
@@ -68,33 +69,32 @@ class IndexFAISS:
 
 
 class IndexQuantizedFAISS(IndexFAISS):
-
     def __init__(
         self,
         train: np.ndarray,
         connectivity: int = DEFAULT_CONNECTIVITY,
         expansion_add: int = DEFAULT_EXPANSION_ADD,
         expansion_search: int = DEFAULT_EXPANSION_SEARCH,
-        *args, **kwargs,
+        *args,
+        **kwargs,
     ):
-
         ndim = train.shape[1]
         super().__init__(
             ndim=ndim,
             connectivity=connectivity,
             expansion_add=expansion_add,
             expansion_search=expansion_search,
-            *args, **kwargs,
+            *args,
+            **kwargs,
         )
 
         nlist = 10000  # Number of inverted lists (number of partitions).
         nsegment = 16  # Number of segments for PQ (number of subquantizers).
-        nbit = 8       # Number of bits to encode each segment.
+        nbit = 8  # Number of bits to encode each segment.
 
         self._original_faiss = self._faiss
-        self._faiss = IndexIVFPQ(
-            self._faiss, ndim, nlist, nsegment, nbit)
+        self._faiss = IndexIVFPQ(self._faiss, ndim, nlist, nsegment, nbit)
 
         self._faiss.train(train)
         self._faiss.nprobe = 10
-        self._specs['Class'] = 'usearch.IndexQuantizedFAISS'
+        self._specs["Class"] = "usearch.IndexQuantizedFAISS"