Add: USearch engine

.github/workflows/benchmarks.yml

@@ -54,6 +54,7 @@ jobs:
           - sklearn
           - sptag
           - tinyknn
+          - usearch
           - vald
           - vearch
           - vespa

algos.yaml

@@ -142,6 +142,90 @@ float:
           arg-groups:
             - {"M": 96,  "efConstruction": 500}
           query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+    usearch-f32:
+      docker-tag: ann-benchmarks-usearch
+      module: ann_benchmarks.algorithms.usearch
+      constructor: USearch
+      base-args: ["@metric", "f32"]
+      run-groups:
+        M-4:
+          arg-groups:
+            - {"M": 4,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-8:
+          arg-groups:
+            - {"M": 8,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-12:
+          arg-groups:
+            - {"M": 12,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-16:
+          arg-groups:
+            - {"M": 16,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-24:
+          arg-groups:
+            - {"M": 24,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-36:
+          arg-groups:
+            - {"M": 36,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-48:
+          arg-groups:
+            - {"M": 48,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-64:
+          arg-groups:
+            - {"M": 64,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-96:
+          arg-groups:
+            - {"M": 96,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+    usearch-f8:
+      docker-tag: ann-benchmarks-usearch
+      module: ann_benchmarks.algorithms.usearch
+      constructor: USearch
+      base-args: ["@metric", "f8"]
+      run-groups:
+        M-4:
+          arg-groups:
+            - {"M": 4,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-8:
+          arg-groups:
+            - {"M": 8,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-12:
+          arg-groups:
+            - {"M": 12,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-16:
+          arg-groups:
+            - {"M": 16,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-24:
+          arg-groups:
+            - {"M": 24,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-36:
+          arg-groups:
+            - {"M": 36,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-48:
+          arg-groups:
+            - {"M": 48,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-64:
+          arg-groups:
+            - {"M": 64,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
+        M-96:
+          arg-groups:
+            - {"M": 96,  "efConstruction": 500}
+          query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]]
 
     hnsw(faiss):
      docker-tag: ann-benchmarks-faiss

ann_benchmarks/algorithms/usearch.py

@@ -0,0 +1,63 @@
+from usearch.index import Index, MetricKind, ScalarKind
+from usearch.numba import jit
+import numpy as np
+
+from .base import BaseANN
+
+class USearch(BaseANN):
+
+    def __init__(self, metric: str, accuracy: str, method_param: dict):
+        assert accuracy in ['f64', 'f32', 'f16', 'f8']
+        assert metric in ['angular', 'euclidean']
+        assert 'M' in method_param
+        assert 'efConstruction' in method_param
+
+        self._method_param = method_param
+        self._accuracy = {'f64': ScalarKind.F64, 'f32': ScalarKind.F32, 'f8': ScalarKind.F8}[accuracy]
+        self._metric = {'angular': MetricKind.Cos, 'euclidean': MetricKind.L2sq}[metric]
+
+    def __str__(self):
+        connectivity = self._method_param['M']
+        expansion_add = self._method_param['efConstruction']
+        return f'USearch(connecitivity={connectivity}, expansion_add={expansion_add})'
+
+    def fit(self, X):
+        connectivity = self._method_param['M']
+        expansion_add = self._method_param['efConstruction']
+        metric = jit(
+            X.shape[1],
+            self._metric,
+            self._accuracy
+        )
+
+        self._index = Index(
+            ndim=len(X[0]),
+            metric=metric,
+            dtype=self._accuracy,
+            connectivity=connectivity,
+            expansion_add=expansion_add
+        )
+
+        labels = np.arange(len(X), dtype=np.longlong)
+        self._index.add(labels, np.asarray(X))
+
+    def get_memory_usage(self) -> int:
+        if not hasattr(self, '_index'):
+            return 0
+
+        return self._index.memory_usage / 1024
+
+    def set_query_arguments(self, ef: int):
+        self._index.expansion_search = ef
+
+    def freeIndex(self):
+        del self._index
+
+    def query(self, v, n):
+        return self._index.search(np.expand_dims(v, axis=0), k=n)[0][0]
+
+    def batch_query(self, X, n):
+        self._batch_results = self._index.search(np.asarray(X), n)
+
+    def get_batch_results(self):
+        return self._batch_results

install/Dockerfile.usearch

@@ -0,0 +1,8 @@
+FROM ann-benchmarks
+
+RUN apt-get install -y python-setuptools python-pip
+RUN pip3 install pybind11 numpy setuptools numba
+RUN git clone https://github.com/unum-cloud/usearch.git && cd usearch && git submodule update --init --recursive
+RUN cd usearch && python3 setup.py install
+
+RUN python3 -c 'import usearch'