usearch 0.1.0

package/src/javascript.cpp

@@ -0,0 +1,198 @@
+/**
+ *  @file javascript.cpp
+ *  @author Ashot Vardanian
+ *  @brief JavaScript bindings for Unum USearch.
+ *  @date 2023-04-26
+ *
+ *  @copyright Copyright (c) 2023
+ *
+ *  @see NodeJS docs: https://nodejs.org/api/addons.html#hello-world
+ *
+ */
+#include <napi.h>
+#include <node_api.h>
+
+#include "advanced.hpp"
+
+using namespace unum::usearch;
+using namespace unum;
+
+using label_t = std::uint32_t;
+using distance_t = punned_distance_t;
+using native_index_t = auto_index_gt<label_t>;
+
+class Index : public Napi::ObjectWrap<Index> {
+  public:
+    static Napi::Object Init(Napi::Env env, Napi::Object exports);
+    Index(Napi::CallbackInfo const& ctx);
+
+  private:
+    Napi::Value GetDimensions(Napi::CallbackInfo const& ctx);
+    Napi::Value GetSize(Napi::CallbackInfo const& ctx);
+    Napi::Value GetCapacity(Napi::CallbackInfo const& ctx);
+    Napi::Value GetConnectivity(Napi::CallbackInfo const& ctx);
+
+    void Add(Napi::CallbackInfo const& ctx);
+    Napi::Value Search(Napi::CallbackInfo const& ctx);
+
+    std::unique_ptr<native_index_t> native_;
+};
+
+Napi::Object Index::Init(Napi::Env env, Napi::Object exports) {
+    Napi::Function func = DefineClass( //
+        env, "Index",
+        {
+            InstanceMethod("dimensions", &Index::GetDimensions),
+            InstanceMethod("size", &Index::GetSize),
+            InstanceMethod("capacity", &Index::GetCapacity),
+            InstanceMethod("connectivity", &Index::GetConnectivity),
+            InstanceMethod("add", &Index::Add),
+            InstanceMethod("search", &Index::Search),
+        });
+
+    Napi::FunctionReference* constructor = new Napi::FunctionReference();
+    *constructor = Napi::Persistent(func);
+    env.SetInstanceData(constructor);
+
+    exports.Set("Index", func);
+    return exports;
+}
+
+Index::Index(Napi::CallbackInfo const& ctx) : Napi::ObjectWrap<Index>(ctx) {
+    Napi::Env env = ctx.Env();
+
+    int length = ctx.Length();
+    if (length == 0 || length >= 2 || !ctx[0].IsObject()) {
+        Napi::TypeError::New(
+            env, "Pass args as named objects: dimensions: uint, capacity: uint, metric: [ip, cos, l2_sq, haversine]")
+            .ThrowAsJavaScriptException();
+        return;
+    }
+
+    Napi::Object params = ctx[0].As<Napi::Object>();
+    std::size_t dimensions = params.Has("dimensions") ? params.Get("dimensions").As<Napi::Number>().Uint32Value() : 0;
+
+    config_t config;
+    if (params.Has("capacity"))
+        config.max_elements = params.Get("capacity").As<Napi::Number>().Uint32Value();
+    if (params.Has("connectivity"))
+        config.connectivity = params.Get("connectivity").As<Napi::Number>().Uint32Value();
+
+    accuracy_t accuracy = accuracy_t::f32_k;
+    if (params.Has("accuracy")) {
+        std::string accuracy_str = params.Get("connectivity").As<Napi::String>().Utf8Value();
+        if (accuracy_str == "f32")
+            accuracy = accuracy_t::f32_k;
+        else if (accuracy_str == "f64")
+            accuracy = accuracy_t::f64_k;
+        else if (accuracy_str == "f16")
+            accuracy = accuracy_t::f16_k;
+        else if (accuracy_str == "i8q100")
+            accuracy = accuracy_t::i8q100_k;
+        else {
+            Napi::TypeError::New(env, "Supported metrics are: [ip, cos, l2_sq, haversine]")
+                .ThrowAsJavaScriptException();
+            return;
+        }
+    }
+
+    // By default we use the Inner Product similarity
+    if (params.Has("metric")) {
+        std::string name = params.Get("metric").As<Napi::String>().Utf8Value();
+        if (name == "l2_sq" || name == "euclidean_sq") {
+            if (!dimensions) {
+                Napi::TypeError::New(env, "Please define the number of dimensions").ThrowAsJavaScriptException();
+                return;
+            }
+            native_.reset(new native_index_t(native_index_t::l2(dimensions, accuracy, config)));
+        } else if (name == "ip" || name == "inner" || name == "dot") {
+            if (!dimensions) {
+                Napi::TypeError::New(env, "Please define the number of dimensions").ThrowAsJavaScriptException();
+                return;
+            }
+            native_.reset(new native_index_t(native_index_t::ip(dimensions, accuracy, config)));
+        } else if (name == "cos" || name == "angular") {
+            if (!dimensions) {
+                Napi::TypeError::New(env, "Please define the number of dimensions").ThrowAsJavaScriptException();
+                return;
+            }
+            native_.reset(new native_index_t(native_index_t::cos(dimensions, accuracy, config)));
+        } else if (name == "haversine") {
+            native_.reset(new native_index_t(native_index_t::haversine(accuracy, config)));
+        } else {
+            Napi::TypeError::New(env, "Supported metrics are: [ip, cos, l2_sq, haversine]")
+                .ThrowAsJavaScriptException();
+            return;
+        }
+    } else {
+        if (!dimensions) {
+            Napi::TypeError::New(env, "Please define the number of dimensions").ThrowAsJavaScriptException();
+            return;
+        }
+        native_.reset(new native_index_t(native_index_t::ip(dimensions, accuracy, config)));
+    }
+}
+
+Napi::Value Index::GetDimensions(Napi::CallbackInfo const& ctx) {
+    return Napi::Number::New(ctx.Env(), native_->dimensions());
+}
+Napi::Value Index::GetSize(Napi::CallbackInfo const& ctx) { return Napi::Number::New(ctx.Env(), native_->size()); }
+Napi::Value Index::GetConnectivity(Napi::CallbackInfo const& ctx) {
+    return Napi::Number::New(ctx.Env(), native_->connectivity());
+}
+Napi::Value Index::GetCapacity(Napi::CallbackInfo const& ctx) {
+    return Napi::Number::New(ctx.Env(), native_->capacity());
+}
+
+void Index::Add(Napi::CallbackInfo const& ctx) {
+    Napi::Env env = ctx.Env();
+    if (ctx.Length() < 2 || !ctx[0].IsNumber() || !ctx[1].IsTypedArray())
+        return Napi::TypeError::New(env, "Expects an integral label and a float vector").ThrowAsJavaScriptException();
+
+    Napi::Number label_js = ctx[0].As<Napi::Number>();
+    Napi::Float32Array vector_js = ctx[1].As<Napi::Float32Array>();
+
+    label_t label = label_js.Uint32Value();
+    float const* vector = vector_js.Data();
+    std::size_t dimensions = static_cast<std::size_t>(vector_js.ElementLength());
+    if (dimensions != native_->dimensions())
+        return Napi::TypeError::New(env, "Wrong number of dimensions").ThrowAsJavaScriptException();
+
+    if (native_->size() + 1 >= native_->capacity())
+        native_->reserve(ceil2(native_->size() + 1));
+    native_->add(label, vector);
+}
+
+Napi::Value Index::Search(Napi::CallbackInfo const& ctx) {
+    Napi::Env env = ctx.Env();
+    if (ctx.Length() < 2 || !ctx[0].IsTypedArray() || !ctx[1].IsNumber()) {
+        Napi::TypeError::New(env, "Expects a float vector and the number of wanted results")
+            .ThrowAsJavaScriptException();
+        return {};
+    }
+
+    Napi::Float32Array vector_js = ctx[0].As<Napi::Float32Array>();
+    Napi::Number wanted_js = ctx[1].As<Napi::Number>();
+
+    float const* vector = vector_js.Data();
+    std::size_t dimensions = static_cast<std::size_t>(vector_js.ElementLength());
+    std::uint32_t wanted = wanted_js.Uint32Value();
+    if (dimensions != native_->dimensions()) {
+        Napi::TypeError::New(env, "Wrong number of dimensions").ThrowAsJavaScriptException();
+        return {};
+    }
+
+    Napi::Uint32Array matches_js = Napi::Uint32Array::New(env, wanted);
+    Napi::Float32Array distances_js = Napi::Float32Array::New(env, wanted);
+    std::size_t count = native_->search(vector, wanted, matches_js.Data(), distances_js.Data());
+
+    Napi::Object result_js = Napi::Object::New(env);
+    result_js.Set("labels", matches_js);
+    result_js.Set("distances", distances_js);
+    result_js.Set("count", Napi::Number::New(env, count));
+    return result_js;
+}
+
+Napi::Object InitAll(Napi::Env env, Napi::Object exports) { return Index::Init(env, exports); }
+
+NODE_API_MODULE(usearch, InitAll)

package/src/main.rs

@@ -0,0 +1,42 @@
+#[cxx::bridge]
+mod ffi {
+    
+    // Shared structs with fields visible to both languages.
+    struct Results {
+        size: usize,
+        labels: Vec<u32>,
+        distances: Vec<f32>,
+    }
+
+    // C++ types and signatures exposed to Rust.
+    unsafe extern "C++" {
+        include!("rust.h");
+
+        type Index;
+
+        fn new_index() -> UniquePtr<Index>;
+
+        fn dim(&self) -> u64;
+        fn connectivity(&self) -> u64;
+        fn size(&self) -> u64;
+        fn capacity(&self) -> u64;
+
+        fn add(&self, label: u64, vector: &[f32]);
+        fn search(&self, query: &[f32], count: u64) -> Results;
+    }
+}
+
+
+fn main() {
+    let client = ffi::new_index();
+
+    // Upload a blob.
+    let chunks = vec![b"fearless".to_vec(), b"concurrency".to_vec()];
+    let mut buf = MultiBuf { chunks, pos: 0 };
+    let blobid = client.add(&mut buf);
+    println!("blobid = {}", blobid);
+
+    // Read back the tags.
+    let search = client.search(blobid);
+    println!("tags = {:?}", search.tags);
+}

package/src/python.cpp

@@ -0,0 +1,226 @@
+/**
+ *  @file python.cpp
+ *  @author Ashot Vardanian
+ *  @brief Python bindings for Unum USearch.
+ *  @date 2023-04-26
+ *
+ *
+ *  https://pythoncapi.readthedocs.io/type_object.html
+ *  https://numpy.org/doc/stable/reference/c-api/types-and-structures.html
+ *  https://pythonextensionpatterns.readthedocs.io/en/latest/refcount.html
+ *  https://docs.python.org/3/extending/newtypes_tutorial.html#adding-data-and-methods-to-the-basic-example
+ *
+ *  @copyright Copyright (c) 2023
+ */
+#define PY_SSIZE_T_CLEAN
+#include <thread>
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+
+#include "advanced.hpp"
+
+using namespace unum::usearch;
+using namespace unum;
+
+namespace py = pybind11;
+
+using label_t = Py_ssize_t;
+using distance_t = punned_distance_t;
+using native_index_t = auto_index_gt<label_t>;
+
+static native_index_t make_index(   //
+    std::size_t dimensions,         //
+    std::size_t capacity,           //
+    std::string const& scalar_type, //
+    std::string const& metric,      //
+    std::size_t connectivity,       //
+    std::size_t expansion_add,      //
+    std::size_t expansion_search    //
+) {
+
+    config_t config;
+    config.expansion_add = expansion_add;
+    config.expansion_search = expansion_search;
+    config.connectivity = connectivity;
+    config.max_elements = capacity;
+    config.max_threads_add = std::thread::hardware_concurrency();
+    config.max_threads_search = std::thread::hardware_concurrency();
+
+    accuracy_t accuracy;
+    if (scalar_type == "f32")
+        accuracy = accuracy_t::f32_k;
+    else if (scalar_type == "f64")
+        accuracy = accuracy_t::f64_k;
+    else if (scalar_type == "f16")
+        accuracy = accuracy_t::f16_k;
+    else if (scalar_type == "i8q100")
+        accuracy = accuracy_t::i8q100_k;
+    else
+        throw std::runtime_error("Unknown type, choose: f32, f16, f64, i8q100");
+
+    if (metric == "l2_sq" || metric == "euclidean_sq")
+        return native_index_t::l2(dimensions, accuracy, config);
+    else if (metric == "ip" || metric == "inner" || metric == "dot")
+        return native_index_t::ip(dimensions, accuracy, config);
+    else if (metric == "cos" || metric == "angular")
+        return native_index_t::cos(dimensions, accuracy, config);
+    else if (metric == "haversine")
+        return native_index_t::haversine(accuracy, config);
+    else
+        throw std::runtime_error("Unknown distance, choose: l2_sq, ip, cos, hamming, jaccard");
+
+    return {};
+}
+
+static void add_to_index(native_index_t& index, py::buffer labels, py::buffer vectors, bool copy) {
+
+    py::buffer_info labels_info = labels.request();
+    py::buffer_info vectors_info = vectors.request();
+
+    if (labels_info.format != py::format_descriptor<label_t>::format())
+        throw std::runtime_error("Incompatible label type!");
+
+    if (labels_info.ndim != 1)
+        throw std::runtime_error("Labels must be placed in a single-dimensional array!");
+
+    if (vectors_info.ndim != 2)
+        throw std::runtime_error("Expects a matrix of vectors to add!");
+
+    ssize_t labels_count = labels_info.shape[0];
+    ssize_t vectors_count = vectors_info.shape[0];
+    ssize_t vectors_dimensions = vectors_info.shape[1];
+    if (vectors_dimensions != index.dimensions())
+        throw std::runtime_error("The number of vector dimensions doesn't match!");
+
+    if (labels_count != vectors_count)
+        throw std::runtime_error("Number of labels and vectors must match!");
+
+    if (index.size() + vectors_count >= index.capacity()) {
+        std::size_t next_capacity = ceil2(index.size() + vectors_count);
+        index.reserve(next_capacity);
+    }
+
+    char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
+    char const* labels_data = reinterpret_cast<char const*>(labels_info.ptr);
+
+    // https://docs.python.org/3/library/struct.html#format-characters
+    if (vectors_info.format == "e")
+        multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+            label_t label = *reinterpret_cast<label_t const*>(labels_data + task_idx * labels_info.strides[0]);
+            f16_converted_t const* vector =
+                reinterpret_cast<f16_converted_t const*>(vectors_data + task_idx * vectors_info.strides[0]);
+            index.add(label, vector, thread_idx, copy);
+        });
+    else if (vectors_info.format == "f")
+        multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+            label_t label = *reinterpret_cast<label_t const*>(labels_data + task_idx * labels_info.strides[0]);
+            float const* vector = reinterpret_cast<float const*>(vectors_data + task_idx * vectors_info.strides[0]);
+            index.add(label, vector, thread_idx, copy);
+        });
+    else if (vectors_info.format == "d")
+        multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+            label_t label = *reinterpret_cast<label_t const*>(labels_data + task_idx * labels_info.strides[0]);
+            double const* vector = reinterpret_cast<double const*>(vectors_data + task_idx * vectors_info.strides[0]);
+            index.add(label, vector, thread_idx, copy);
+        });
+    else
+        throw std::runtime_error("Incompatible scalars in the vectors matrix!");
+}
+
+/**
+ *  @param vectors Matrix of vectors to search for..
+ *  @param wanted Number of matches per request.
+ *
+ *  @return Tuple with:
+ *      1. matrix of neighbors,
+ *      2. matrix of distances,
+ *      3. array with match counts.
+ */
+static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssize_t wanted) {
+
+    py::buffer_info vectors_info = vectors.request();
+    if (vectors_info.ndim != 2)
+        throw std::runtime_error("Expects a matrix of vectors to add!");
+
+    ssize_t vectors_count = vectors_info.shape[0];
+    ssize_t vectors_dimensions = vectors_info.shape[1];
+    if (vectors_dimensions != index.dimensions())
+        throw std::runtime_error("The number of vector dimensions doesn't match!");
+
+    py::array_t<label_t> labels_py({vectors_count, wanted});
+    py::array_t<distance_t> distances_py({vectors_count, wanted});
+    py::array_t<Py_ssize_t> counts_py(vectors_count);
+    auto labels_py2d = labels_py.mutable_unchecked<2>();
+    auto distances_py2d = distances_py.mutable_unchecked<2>();
+    auto counts_py1d = counts_py.mutable_unchecked<1>();
+
+    char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
+
+    // https://docs.python.org/3/library/struct.html#format-characters
+    if (vectors_info.format == "e")
+        multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+            f16_converted_t const* vector = (f16_converted_t const*)(vectors_data + task_idx * vectors_info.strides[0]);
+            counts_py1d(task_idx) = static_cast<Py_ssize_t>(
+                index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
+        });
+    else if (vectors_info.format == "f")
+        multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+            float const* vector = (float const*)(vectors_data + task_idx * vectors_info.strides[0]);
+            counts_py1d(task_idx) = static_cast<Py_ssize_t>(
+                index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
+        });
+    else if (vectors_info.format == "d")
+        multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
+            double const* vector = (double const*)(vectors_data + task_idx * vectors_info.strides[0]);
+            counts_py1d(task_idx) = static_cast<Py_ssize_t>(
+                index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
+        });
+    else
+        throw std::runtime_error("Incompatible scalars in the query matrix!");
+
+    py::tuple results(3);
+    results[0] = labels_py;
+    results[1] = distances_py;
+    results[2] = counts_py;
+    return results;
+}
+
+static void save_index(native_index_t const& index, std::string const& path) { index.save(path.c_str()); }
+static void load_index(native_index_t& index, std::string const& path) { index.load(path.c_str()); }
+static void view_index(native_index_t& index, std::string const& path) { index.view(path.c_str()); }
+
+PYBIND11_MODULE(usearch, m) {
+    m.doc() = "Unum USearch Python bindings";
+
+    auto i = py::class_<native_index_t>(m, "Index");
+
+    i.def(py::init(&make_index),                                                 //
+          py::kw_only(),                                                         //
+          py::arg("dim"),                                                        //
+          py::arg("capacity") = 0,                                               //
+          py::arg("dtype") = std::string("f32"),                                 //
+          py::arg("metric") = std::string("ip"),                                 //
+          py::arg("connectivity") = config_t::connectivity_default_k,            //
+          py::arg("expansion_add") = config_t::expansion_construction_default_k, //
+          py::arg("expansion_search") = config_t::expansion_search_default_k     //
+    );
+
+    i.def(                     //
+        "add", &add_to_index,  //
+        py::arg("labels"),     //
+        py::arg("vectors"),    //
+        py::kw_only(),         //
+        py::arg("copy") = true //
+    );
+
+    i.def(                          //
+        "search", &search_in_index, //
+        py::arg("vectors"),         //
+        py::arg("count") = 10       //
+    );
+
+    i.def("save", &save_index, py::arg("path"));
+    i.def("load", &load_index, py::arg("path"));
+    i.def("view", &view_index, py::arg("path"));
+}

package/src/rust.cpp

@@ -0,0 +1,23 @@
+#include "rust.h"
+
+#include "advanced.hpp"
+
+using namespace unum::usearch;
+using namespace unum;
+
+class Index::impl {
+    friend Index;
+    downcasting_index_t native_;
+};
+
+Index::Index() : impl(new class Index::impl) {}
+
+void Index::add(uint64_t label, rust::Slice<float const> vector) const {}
+Results Index::search(rust::Slice<float const> vector, uint64_t count) const { return {}; }
+
+uint64_t Index::dim() const { return impl->native_.dim(); }
+uint64_t Index::connectivity() const { return impl->native_.connectivity(); }
+uint64_t Index::size() const { return impl->native_.size(); }
+uint64_t Index::capacity() const { return impl->native_.capacity(); }
+
+std::unique_ptr<Index> new_index() { return std::make_unique<Index>(); }

package/src/rust.h

@@ -0,0 +1,22 @@
+#pragma once
+#include "rust/cxx.h"
+
+struct Results;
+
+class Index {
+  public:
+    Index();
+    void add(uint64_t label, rust::Slice<float const> vector) const;
+    Results search(rust::Slice<float const> vector, uint64_t count) const;
+
+    uint64_t dim() const;
+    uint64_t connectivity() const;
+    uint64_t size() const;
+    uint64_t capacity() const;
+
+  private:
+    class impl;
+    std::shared_ptr<impl> impl;
+};
+
+std::unique_ptr<Index> new_index();

package/src/test.js

@@ -0,0 +1,17 @@
+var assert = require('assert');
+var usearch = require('bindings')('usearch');
+
+var index = new usearch.Index({ metric: 'l2_sq', connectivity: 16, dimensions: 2 })
+assert.equal(index.connectivity(), 16)
+assert.equal(index.dimensions(), 2)
+assert.equal(index.size(), 0)
+
+index.add(15, new Float32Array([10, 20]))
+index.add(16, new Float32Array([10, 25]))
+assert.equal(index.size(), 2)
+
+var results = index.search(new Float32Array([13, 14]), 2)
+assert.deepEqual(results.labels, new Uint32Array([15, 16]))
+assert.deepEqual(results.distances, new Float32Array([45, 130]))
+
+console.log('JavaScript tests passed!')

package/src/test.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+import usearch
+
+count_vectors = 100
+count_dimensions = 96
+
+index = usearch.Index(dim=96)
+print(index)
+
+vectors = np.random.uniform(
+    0, 0.3, (count_vectors, count_dimensions)).astype(np.float32)
+labels = np.array(range(count_vectors), dtype=np.longlong)
+print('will insert', labels.shape, vectors.shape)
+
+index.add(labels, vectors, copy=True)
+
+print('will search')
+results = index.search(vectors, 10)
+print('found', results[0].shape, results[1].shape, results[2].shape)