usearch 0.18.8 → 0.19.1

package/python/usearch/index.py

@@ -1,22 +1,23 @@
+from __future__ import annotations
+
 # The purpose of this file is to provide Pythonic wrapper on top
 # the native precompiled CPython module. It improves compatibility
 # Python tooling, linters, and static analyzers. It also embeds JIT
 # into the primary `Index` class, connecting USearch with Numba.
 import os
-from math import sqrt
-from typing import Optional, Callable, Union, NamedTuple, List, Iterable
+from typing import Optional, Union, NamedTuple, List, Iterable
 
 import numpy as np
+from tqdm import tqdm
 
-from usearch.compiled import Index as _CompiledIndex
+from usearch.compiled import Index as _CompiledIndex, IndexMetadata
 from usearch.compiled import SparseIndex as _CompiledSetsIndex
 
-from usearch.compiled import MetricKind, ScalarKind
+from usearch.compiled import MetricKind, ScalarKind, MetricSignature
 from usearch.compiled import (
     DEFAULT_CONNECTIVITY,
     DEFAULT_EXPANSION_ADD,
     DEFAULT_EXPANSION_SEARCH,
-
     USES_OPENMP,
     USES_SIMSIMD,
     USES_NATIVE_F16,
@@ -33,32 +34,71 @@ SparseIndex = _CompiledSetsIndex
 Label = np.uint32
 
 
-def _normalize_dtype(dtype) -> ScalarKind:
-    if dtype is None:
-        return ScalarKind.F32
+def _normalize_dtype(dtype, metric: MetricKind = MetricKind.IP) -> ScalarKind:
+    if dtype is None or dtype == "":
+        return ScalarKind.B1 if metric in MetricKindBitwise else ScalarKind.F32
+
     if isinstance(dtype, ScalarKind):
         return dtype
 
     if isinstance(dtype, str):
+        dtype = dtype.lower()
+
+    _normalize = {
+        "f64": ScalarKind.F64,
+        "f32": ScalarKind.F32,
+        "f16": ScalarKind.F16,
+        "f8": ScalarKind.F8,
+        "b1": ScalarKind.B1,
+        "b1x8": ScalarKind.B1,
+        np.float64: ScalarKind.F64,
+        np.float32: ScalarKind.F32,
+        np.float16: ScalarKind.F16,
+        np.int8: ScalarKind.F8,
+    }
+    return _normalize[dtype]
+
+
+def _to_numpy_compatible_dtype(dtype: ScalarKind) -> ScalarKind:
+    _normalize = {
+        ScalarKind.F64: ScalarKind.F64,
+        ScalarKind.F32: ScalarKind.F32,
+        ScalarKind.F16: ScalarKind.F16,
+        ScalarKind.F8: ScalarKind.F16,
+        ScalarKind.B1: ScalarKind.B1,
+    }
+    return _normalize[dtype]
+
+
+def _to_numpy_dtype(dtype: ScalarKind):
+    _normalize = {
+        ScalarKind.F64: np.float64,
+        ScalarKind.F32: np.float32,
+        ScalarKind.F16: np.float16,
+        ScalarKind.F8: np.float16,
+        ScalarKind.B1: np.uint8,
+    }
+    return _normalize[dtype]
+
+
+def _normalize_metric(metric):
+    if metric is None:
+        return MetricKind.IP
+
+    if isinstance(metric, str):
         _normalize = {
-            'f64': ScalarKind.F64,
-            'f32': ScalarKind.F32,
-            'f16': ScalarKind.F16,
-            'f8': ScalarKind.F8,
-            'b1': ScalarKind.B1,
-        }
-        return _normalize[dtype.lower()]
-
-    if not isinstance(dtype, ScalarKind):
-        _normalize = {
-            np.float64: ScalarKind.F64,
-            np.float32: ScalarKind.F32,
-            np.float16: ScalarKind.F16,
-            np.int8: ScalarKind.F8,
+            "cos": MetricKind.Cos,
+            "ip": MetricKind.IP,
+            "l2_sq": MetricKind.L2sq,
+            "haversine": MetricKind.Haversine,
+            "perason": MetricKind.Pearson,
+            "hamming": MetricKind.Hamming,
+            "tanimoto": MetricKind.Tanimoto,
+            "sorensen": MetricKind.Sorensen,
         }
-        return _normalize[dtype]
+        return _normalize[metric.lower()]
 
-    return dtype
+    return metric
 
 
 class Matches(NamedTuple):
@@ -67,7 +107,7 @@ class Matches(NamedTuple):
     counts: Union[np.ndarray, int]
 
     @property
-    def is_batch(self) -> bool:
+    def is_multiple(self) -> bool:
         return isinstance(self.counts, np.ndarray)
 
     @property
@@ -79,9 +119,8 @@ class Matches(NamedTuple):
         return np.sum(self.counts)
 
     def to_list(self, row: Optional[int] = None) -> Union[List[dict], List[List[dict]]]:
-
-        if not self.is_batch:
-            assert row is None, 'Exporting a single sequence is only for batch requests'
+        if not self.is_multiple:
+            assert row is None, "Exporting a single sequence is only for batch requests"
             labels = self.labels
             distances = self.distances
 
@@ -94,17 +133,26 @@ class Matches(NamedTuple):
             distances = self.distances[row, :count]
 
         return [
-            {'label': int(label), 'distance': float(distance)}
+            {"label": int(label), "distance": float(distance)}
             for label, distance in zip(labels, distances)
         ]
 
     def recall_first(self, expected: np.ndarray) -> float:
-        best_matches = self.labels if not self.is_batch else self.labels[:, 0]
+        best_matches = self.labels if not self.is_multiple else self.labels[:, 0]
         return np.sum(best_matches == expected) / len(expected)
 
     def __repr__(self) -> str:
-        return f'usearch.Matches({self.total_matches})' if self.is_batch else \
-            f'usearch.Matches({self.total_matches} across {self.batch_size} queries)'
+        return (
+            f"usearch.Matches({self.total_matches})"
+            if self.is_multiple
+            else f"usearch.Matches({self.total_matches} across {self.batch_size} queries)"
+        )
+
+
+class CompiledMetric(NamedTuple):
+    pointer: int
+    kind: MetricKind
+    signature: MetricSignature
 
 
 class Index:
@@ -120,15 +168,13 @@ class Index:
     def __init__(
         self,
         ndim: int,
-        metric: Union[MetricKind, Callable, str] = MetricKind.IP,
+        metric: Union[str, MetricKind, CompiledMetric] = MetricKind.IP,
         dtype: Optional[str] = None,
         jit: bool = False,
-
         connectivity: int = DEFAULT_CONNECTIVITY,
         expansion_add: int = DEFAULT_EXPANSION_ADD,
         expansion_search: int = DEFAULT_EXPANSION_SEARCH,
         tune: bool = False,
-
         path: Optional[os.PathLike] = None,
         view: bool = False,
     ) -> None:
@@ -171,7 +217,7 @@ class Index:
 
         :param expansion_search: Traversal depth on queries, defaults to None
         :type expansion_search: Optional[int], optional
-            Hyper-parameter for the search depth when querying 
+            Hyper-parameter for the search depth when querying
             nearest neighbors. The original paper calls it "ef".
             Can be changed afterwards, as the `.expansion_search`.
 
@@ -184,64 +230,45 @@ class Index:
         :type view: bool, optional
         """
 
-        if metric in MetricKindBitwise:
-            assert dtype is None or dtype == ScalarKind.B1
-            dtype = ScalarKind.B1
-        else:
-            dtype = _normalize_dtype(dtype)
-
-
-        if metric is None:
-            metric = MetricKind.IP
-        elif isinstance(metric, str):
-            _normalize = {
-                'cos': MetricKind.Cos,
-                'ip': MetricKind.IP,
-                'l2_sq': MetricKind.L2sq,
-                'haversine': MetricKind.Haversine,
-                'perason': MetricKind.Pearson,
-                'hamming': MetricKind.Hamming,
-                'tanimoto': MetricKind.Tanimoto,
-                'sorensen': MetricKind.Sorensen,
-            }
-            metric = _normalize[metric.lower()]
-
-        if isinstance(metric, Callable):
-            self._metric_kind = MetricKind.Unknown
-            self._metric_jit = metric
-            self._metric_pointer = int(metric.address)
-
-        elif isinstance(metric, MetricKind):
-            if jit:
-
-                try:
-                    from usearch.numba import jit
-                except ImportError:
-                    raise ModuleNotFoundError(
-                        'To use JIT install Numba with `pip install numba`.'
-                        'Alternatively, reinstall usearch with `pip install usearch[jit]`')
-
-                self._metric_kind = metric
-                self._metric_jit = jit(
-                    ndim=ndim,
-                    metric=metric,
-                    dtype=dtype,
+        metric = _normalize_metric(metric)
+        if isinstance(metric, MetricKind) and jit:
+            try:
+                from usearch.numba import jit
+            except ImportError:
+                raise ModuleNotFoundError(
+                    "To use JIT install Numba with `pip install numba`."
+                    "Alternatively, reinstall with `pip install usearch[jit]`"
                 )
-                self._metric_pointer = self._metric_jit.address if \
-                    self._metric_jit else 0
-            else:
-                self._metric_kind = metric
-                self._metric_jit = None
-                self._metric_pointer = 0
+
+            metric = jit(
+                ndim=ndim,
+                metric=metric,
+                dtype=_normalize_dtype(dtype, metric),
+            )
+
+        if isinstance(metric, MetricKind):
+            self._metric_kind = metric
+            self._metric_jit = None
+            self._metric_pointer = 0
+            self._metric_signature = MetricSignature.ArrayArraySize
+        elif isinstance(metric, CompiledMetric):
+            self._metric_jit = metric
+            self._metric_kind = metric.kind
+            self._metric_pointer = metric.pointer
+            self._metric_signature = metric.signature
         else:
             raise ValueError(
-                'The `metric` must be Numba callback or a `MetricKind`')
+                "The `metric` must be a `CompiledMetric` or a `MetricKind`"
+            )
 
+        # Validate, that the right scalar type is defined
+        dtype = _normalize_dtype(dtype, self._metric_kind)
         self._compiled = _CompiledIndex(
             ndim=ndim,
+            dtype=dtype,
             metric=self._metric_kind,
             metric_pointer=self._metric_pointer,
-            dtype=dtype,
+            metric_signature=self._metric_signature,
             connectivity=connectivity,
             expansion_add=expansion_add,
             expansion_search=expansion_search,
@@ -255,17 +282,47 @@ class Index:
             else:
                 self._compiled.load(path)
 
+    @staticmethod
+    def metadata(path: os.PathLike) -> IndexMetadata:
+        return IndexMetadata(path)
+
+    @staticmethod
+    def restore(path: os.PathLike, view: bool = False) -> Index:
+        meta = Index.metadata(path)
+        bits_per_scalar = {
+            ScalarKind.F8: 8,
+            ScalarKind.F16: 16,
+            ScalarKind.F32: 32,
+            ScalarKind.F64: 64,
+            ScalarKind.B1: 1,
+        }[meta.scalar_kind]
+        ndim = meta.bytes_for_vectors * 8 // meta.size // bits_per_scalar
+        return Index(
+            ndim=ndim,
+            connectivity=meta.connectivity,
+            metric=meta.metric,
+            path=path,
+            view=view,
+        )
+
     def add(
-            self, labels, vectors, *,
-            copy: bool = True, threads: int = 0) -> Union[int, np.ndarray]:
+        self,
+        labels,
+        vectors,
+        *,
+        copy: bool = True,
+        threads: int = 0,
+        log: Union[str, bool] = False,
+        batch_size: int = 0,
+    ) -> Union[int, np.ndarray]:
         """Inserts one or move vectors into the index.
 
         For maximal performance the `labels` and `vectors`
         should conform to the Python's "buffer protocol" spec.
 
-        To index a single entry: 
+        To index a single entry:
             labels: int, vectors: np.ndarray.
-        To index many entries: 
+        To index many entries:
             labels: np.ndarray, vectors: np.ndarray.
 
         When working with extremely large indexes, you may want to
@@ -284,30 +341,70 @@ class Index:
         :return: Inserted label or labels
         :type: Union[int, np.ndarray]
         """
-        assert isinstance(vectors, np.ndarray), 'Expects a NumPy array'
-        assert vectors.ndim == 1 or vectors.ndim == 2, 'Expects a matrix or vector'
-        is_batch = vectors.ndim == 2
+        assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
+        assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
+        count_vectors = vectors.shape[0] if vectors.ndim == 2 else 1
+        is_multiple = count_vectors > 1
+        if not is_multiple:
+            vectors = vectors.flatten()
+
+        # Validate of generate teh labels
         generate_labels = labels is None
-
-        # If no `labels` were provided, generate some
         if generate_labels:
             start_id = len(self._compiled)
-            if is_batch:
-                labels = np.arange(start_id, start_id + vectors.shape[0])
+            if is_multiple:
+                labels = np.arange(start_id, start_id + count_vectors, dtype=Label)
             else:
                 labels = start_id
-        if isinstance(labels, np.ndarray):
-            labels = labels.astype(Label)
-        elif isinstance(labels, Iterable):
-            labels = np.array(labels, dtype=Label)
+        else:
+            if isinstance(labels, np.ndarray):
+                if not is_multiple:
+                    labels = int(labels[0])
+                else:
+                    labels = labels.astype(Label)
+        count_labels = len(labels) if isinstance(labels, Iterable) else 1
+        assert count_labels == count_vectors
+
+        # Split into batches and log progress, if needed
+        if batch_size:
+            labels = [
+                labels[start_row : start_row + batch_size]
+                for start_row in range(0, count_vectors, batch_size)
+            ]
+            vectors = [
+                vectors[start_row : start_row + batch_size, :]
+                for start_row in range(0, count_vectors, batch_size)
+            ]
+            tasks = zip(labels, vectors)
+            name = log if isinstance(log, str) else "Add"
+            pbar = tqdm(
+                tasks,
+                desc=name,
+                total=count_vectors,
+                unit="Vector",
+                disable=log is False,
+            )
+            for labels, vectors in tasks:
+                self._compiled.add(labels, vectors, copy=copy, threads=threads)
+                pbar.update(len(labels))
+
+            pbar.close()
 
-        self._compiled.add(labels, vectors, copy=copy, threads=threads)
+        else:
+            self._compiled.add(labels, vectors, copy=copy, threads=threads)
 
         return labels
 
     def search(
-            self, vectors, k: int = 10, *,
-            threads: int = 0, exact: bool = False) -> Matches:
+        self,
+        vectors,
+        k: int = 10,
+        *,
+        threads: int = 0,
+        exact: bool = False,
+        log: Union[str, bool] = False,
+        batch_size: int = 0,
+    ) -> Matches:
         """Performs approximate nearest neighbors search for one or more queries.
 
         :param vectors: Query vector or vectors.
@@ -323,26 +420,28 @@ class Index:
         :rtype: Matches
         """
         tuple_ = self._compiled.search(
-            vectors, k,
-            exact=exact, threads=threads,
+            vectors,
+            k,
+            exact=exact,
+            threads=threads,
         )
         return Matches(*tuple_)
 
     @property
     def specs(self) -> dict:
         return {
-            'Class': 'usearch.Index',
-            'Connectivity': self.connectivity,
-            'Size': self.size,
-            'Dimensions': self.ndim,
-            'Expansion@Add': self.expansion_add,
-            'Expansion@Search': self.expansion_search,
-            'OpenMP': USES_OPENMP,
-            'SimSIMD': USES_SIMSIMD,
-            'NativeF16': USES_NATIVE_F16,
-            'JIT': self.jit,
-            'DType': self.dtype,
-            'Path': self.path,
+            "Class": "usearch.Index",
+            "Connectivity": self.connectivity,
+            "Size": self.size,
+            "Dimensions": self.ndim,
+            "Expansion@Add": self.expansion_add,
+            "Expansion@Search": self.expansion_search,
+            "OpenMP": USES_OPENMP,
+            "SimSIMD": USES_SIMSIMD,
+            "NativeF16": USES_NATIVE_F16,
+            "JIT": self.jit,
+            "DType": self.dtype,
+            "Path": self.path,
         }
 
     def __len__(self) -> int:
@@ -396,33 +495,46 @@ class Index:
     def expansion_search(self, v: int):
         self._compiled.expansion_search = v
 
-    def save(self, path: os.PathLike):
+    def save(self, path: Optional[os.PathLike] = None):
+        path = path if path else self.path
+        if path is None:
+            raise Exception("Define `path` argument")
         self._compiled.save(path)
 
-    def load(self, path: os.PathLike):
+    def load(self, path: Optional[os.PathLike] = None):
+        path = path if path else self.path
+        if path is None:
+            raise Exception("Define `path` argument")
         self._compiled.load(path)
 
-    def view(self, path: os.PathLike):
+    def view(self, path: Optional[os.PathLike] = None):
+        path = path if path else self.path
+        if path is None:
+            raise Exception("Define `path` argument")
         self._compiled.view(path)
 
     def clear(self):
         self._compiled.clear()
 
+    def close(self):
+        self._compiled.close()
+
     @property
     def labels(self) -> np.ndarray:
         return self._compiled.labels
 
     def get_vectors(
-            self,
-            labels: np.ndarray,
-            dtype: ScalarKind = ScalarKind.F32) -> np.ndarray:
-
-        dtype = _normalize_dtype(dtype)
-        return np.vstack([self._compiled.__getitem__(l, dtype) for l in labels])
+        self, labels: np.ndarray, dtype: ScalarKind = ScalarKind.F32
+    ) -> np.ndarray:
+        dtype = _normalize_dtype(dtype, self._metric_kind)
+        get_dtype = _to_numpy_compatible_dtype(dtype)
+        vectors = np.vstack([self._compiled.__getitem__(l, get_dtype) for l in labels])
+        view_dtype = _to_numpy_dtype(dtype)
+        return vectors.view(view_dtype)
 
     @property
     def vectors(self) -> np.ndarray:
-        return self.get_vectors(self.labels)
+        return self.get_vectors(self.labels, self.dtype)
 
     def __delitem__(self, label: int):
         raise NotImplementedError()
@@ -431,4 +543,8 @@ class Index:
         return self._compiled.__contains__(label)
 
     def __getitem__(self, label: int) -> np.ndarray:
-        return self._compiled.__getitem__(label)
+        dtype = self.dtype
+        get_dtype = _to_numpy_compatible_dtype(dtype)
+        vector = self._compiled.__getitem__(label, get_dtype)
+        view_dtype = _to_numpy_dtype(dtype)
+        return None if vector is None else vector.view(view_dtype)

package/python/usearch/io.py

@@ -4,9 +4,8 @@ import numpy as np
 
 
 def load_matrix(
-        filename: str,
-        start_row: int = 0, count_rows: int = None,
-        view: bool = False) -> np.ndarray:
+    filename: str, start_row: int = 0, count_rows: int = None, view: bool = False
+) -> np.ndarray:
     """Read *.ibin, *.bbib, *.hbin, *.fbin, *.dbin files with matrices.
 
     :param filename: path to the matrix file
@@ -19,36 +18,37 @@ def load_matrix(
     """
     dtype = np.float32
     scalar_size = 4
-    if filename.endswith('.fbin'):
+    if filename.endswith(".fbin"):
         dtype = np.float32
         scalar_size = 4
-    elif filename.endswith('.dbin'):
+    elif filename.endswith(".dbin"):
         dtype = np.float64
         scalar_size = 8
-    elif filename.endswith('.hbin'):
+    elif filename.endswith(".hbin"):
         dtype = np.float16
         scalar_size = 2
-    elif filename.endswith('.ibin'):
+    elif filename.endswith(".ibin"):
         dtype = np.int32
         scalar_size = 4
-    elif filename.endswith('.bbin'):
-        dtype = np.int8
+    elif filename.endswith(".bbin"):
+        dtype = np.uint8
         scalar_size = 1
     else:
-        raise Exception('Unknown file type')
+        raise Exception("Unknown file type")
 
-    with open(filename, 'rb') as f:
+    with open(filename, "rb") as f:
         rows, cols = np.fromfile(f, count=2, dtype=np.int32)
         rows = (rows - start_row) if count_rows is None else count_rows
         row_offset = start_row * scalar_size * cols
 
         if view:
             return np.memmap(
-                f, dtype=dtype, mode='r', offset=8+row_offset, shape=(rows, cols))
+                f, dtype=dtype, mode="r", offset=8 + row_offset, shape=(rows, cols)
+            )
         else:
             return np.fromfile(
-                f, count=rows * cols,
-                dtype=dtype, offset=row_offset).reshape(rows, cols)
+                f, count=rows * cols, dtype=dtype, offset=row_offset
+            ).reshape(rows, cols)
 
 
 def save_matrix(vectors: np.ndarray, filename: str):
@@ -60,22 +60,22 @@ def save_matrix(vectors: np.ndarray, filename: str):
     :type filename: str
     """
     dtype = np.float32
-    if filename.endswith('.fbin'):
+    if filename.endswith(".fbin"):
         dtype = np.float32
-    elif filename.endswith('.dbin'):
+    elif filename.endswith(".dbin"):
         dtype = np.float64
-    elif filename.endswith('.hbin'):
+    elif filename.endswith(".hbin"):
         dtype = np.float16
-    elif filename.endswith('.ibin'):
+    elif filename.endswith(".ibin"):
         dtype = np.int32
-    elif filename.endswith('.bbin'):
-        dtype = np.int8
+    elif filename.endswith(".bbin"):
+        dtype = np.uint8
     else:
-        raise Exception('Unknown file type')
+        raise Exception("Unknown file type")
 
-    assert len(vectors.shape) == 2, 'Input array must have 2 dimensions'
-    with open(filename, 'wb') as f:
+    assert len(vectors.shape) == 2, "Input array must have 2 dimensions"
+    with open(filename, "wb") as f:
         count, dim = vectors.shape
-        f.write(struct.pack('<i', count))
-        f.write(struct.pack('<i', dim))
+        f.write(struct.pack("<i", count))
+        f.write(struct.pack("<i", dim))
         vectors.astype(dtype).flatten().tofile(f)

package/python/usearch/join.py

@@ -7,7 +7,6 @@ from index import Index, Matches, Label
 
 
 class bidict(dict):
-
     def __init__(self, *args, **kwargs):
         super(bidict, self).__init__(*args, **kwargs)
         self.inverse = {}
@@ -37,10 +36,10 @@ def index_in_array(array: np.ndarray, item: Label) -> int:
 
 
 def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
-    """Performs a Semantic Join, matching one entry from `a` with 
+    """Performs a Semantic Join, matching one entry from `a` with
     one entry from `b`, converging towards a Stable Marriage.
     Assuming the collections can be different in size, classical solution
-    doesn't provide 
+    doesn't provide
 
     :param resolution: Approximate matches per member to consider, defaults to 10
     :type resolution: int, optional
@@ -65,7 +64,6 @@ def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
     idle_cycles: int = 0
     hopeless_men_count = len(men) - len(women)
     while len(free_men) > hopeless_men_count:
-
         # In the worst case scenario we may need to match more candidates for every
         # remaining man. This, however, may drastically increase the runtime.
         if len(free_men) == idle_cycles: