usearch 0.18.8 → 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);