usearch 0.19.0 → 0.19.2

package/python/usearch/index.py

@@ -8,6 +8,7 @@ import os
 from typing import Optional, Union, NamedTuple, List, Iterable
 
 import numpy as np
+from tqdm import tqdm
 
 from usearch.compiled import Index as _CompiledIndex, IndexMetadata
 from usearch.compiled import SparseIndex as _CompiledSetsIndex
@@ -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,
+            "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.lower()]
+        return _normalize[metric.lower()]
 
-    if not isinstance(dtype, ScalarKind):
-        _normalize = {
-            np.float64: ScalarKind.F64,
-            np.float32: ScalarKind.F32,
-            np.float16: ScalarKind.F16,
-            np.int8: ScalarKind.F8,
-        }
-        return _normalize[dtype]
-
-    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,7 +119,7 @@ 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:
+        if not self.is_multiple:
             assert row is None, "Exporting a single sequence is only for batch requests"
             labels = self.labels
             distances = self.distances
@@ -98,13 +138,13 @@ class Matches(NamedTuple):
         ]
 
     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
+            if self.is_multiple
             else f"usearch.Matches({self.total_matches} across {self.batch_size} queries)"
         )
 
@@ -127,10 +167,9 @@ class Index:
 
     def __init__(
         self,
-        ndim: int,
+        ndim: int = 0,
         metric: Union[str, MetricKind, CompiledMetric] = MetricKind.IP,
-        dtype: Optional[str] = None,
-        jit: bool = False,
+        dtype: Optional[Union[str, ScalarKind]] = None,
         connectivity: int = DEFAULT_CONNECTIVITY,
         expansion_add: int = DEFAULT_EXPANSION_ADD,
         expansion_search: int = DEFAULT_EXPANSION_SEARCH,
@@ -142,6 +181,10 @@ class Index:
 
         :param ndim: Number of vector dimensions
         :type ndim: int
+            Required for some metrics, optional for others.
+            Haversine, for example, only applies to 2-dimensional latitude/longitude
+            coordinates. Angular (Cos) and Euclidean (L2sq), obviously, apply to
+            vectors with arbitrary number of dimensions.
 
         :param metric: Distance function, defaults to MetricKind.IP
         :type metric: Union[MetricKind, Callable, str], optional
@@ -151,18 +194,12 @@ class Index:
             Not every kind is JIT-able. For Jaccard distance, use `SparseIndex`.
 
         :param dtype: Scalar type for internal vector storage, defaults to None
-        :type dtype: str, optional
+        :type dtype: Optional[Union[str, ScalarKind]], optional
             For continuous metrics can be: f16, f32, f64, or f8.
             For bitwise metrics it's implementation-defined, and can't change.
             Example: you can use the `f16` index with `f32` vectors in Euclidean space,
             which will be automatically downcasted.
 
-        :param jit: Enable Numba to JIT compile the metric, defaults to False
-        :type jit: bool, optional
-            This can result in up-to 3x performance difference on very large vectors
-            and very recent hardware, as the Python module is compiled with high
-            compatibility in mind and avoids very fancy assembly instructions.
-
         :param connectivity: Connections per node in HNSW, defaults to None
         :type connectivity: Optional[int], optional
             Hyper-parameter for the number of Graph connections
@@ -190,43 +227,7 @@ 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
-
-        if 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, 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]`"
-                )
-
-            metric = jit(
-                ndim=ndim,
-                metric=metric,
-                dtype=dtype,
-            )
-
+        metric = _normalize_metric(metric)
         if isinstance(metric, MetricKind):
             self._metric_kind = metric
             self._metric_jit = None
@@ -242,6 +243,8 @@ class Index:
                 "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,
@@ -285,7 +288,14 @@ class Index:
         )
 
     def add(
-        self, labels, vectors, *, copy: bool = True, threads: int = 0
+        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.
 
@@ -315,27 +325,67 @@ class Index:
         """
         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
-        generate_labels = labels is None
+        count_vectors = vectors.shape[0] if vectors.ndim == 2 else 1
+        is_multiple = count_vectors > 1
+        if not is_multiple:
+            vectors = vectors.flatten()
 
-        # If no `labels` were provided, generate some
+        # Validate of generate teh labels
+        generate_labels = labels is None
         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))
 
-        self._compiled.add(labels, vectors, copy=copy, threads=threads)
+            pbar.close()
+
+        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
+        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.
 
@@ -427,18 +477,30 @@ 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
@@ -446,12 +508,15 @@ class Index:
     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])
+        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()
@@ -460,4 +525,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

@@ -31,7 +31,7 @@ def load_matrix(
         dtype = np.int32
         scalar_size = 4
     elif filename.endswith(".bbin"):
-        dtype = np.int8
+        dtype = np.uint8
         scalar_size = 1
     else:
         raise Exception("Unknown file type")
@@ -69,7 +69,7 @@ def save_matrix(vectors: np.ndarray, filename: str):
     elif filename.endswith(".ibin"):
         dtype = np.int32
     elif filename.endswith(".bbin"):
-        dtype = np.int8
+        dtype = np.uint8
     else:
         raise Exception("Unknown file type")
 

package/python/usearch/numba.py

@@ -4,31 +4,19 @@
 # into the primary `Index` class, connecting USearch with Numba.
 from math import sqrt
 
-from numba import cfunc, types, carray
-
 from usearch.index import MetricKind, ScalarKind, MetricSignature, CompiledMetric
 
 
-signature_i8args = types.float32(types.CPointer(types.int8), types.CPointer(types.int8))
-
-signature_f16args = types.float32(
-    types.CPointer(types.float16), types.CPointer(types.float16)
-)
-
-signature_f32args = types.float32(
-    types.CPointer(types.float32), types.CPointer(types.float32)
-)
-
-signature_f64args = types.float32(
-    types.CPointer(types.float64), types.CPointer(types.float64)
-)
-
-
 def jit(
-    ndim: int, metric: MetricKind = MetricKind.Cos, dtype: ScalarKind = ScalarKind.F32
+    ndim: int,
+    metric: MetricKind = MetricKind.Cos,
+    dtype: ScalarKind = ScalarKind.F32,
 ) -> CompiledMetric:
-    """JIT-compiles a distance metric specifically tuned for the target hardware
-    and number of dimensions.
+    """JIT-compiles the metric for target hardware and number of dimensions.
+
+    This can result in up-to 3x performance difference on very large vectors
+    and very recent hardware, as the Python module is compiled with high
+    compatibility in mind and avoids very fancy assembly instructions.
 
     Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
     of `ctypes` to support half-precision.
@@ -37,6 +25,24 @@ def jit(
     assert isinstance(metric, MetricKind)
     assert isinstance(dtype, ScalarKind)
 
+    from numba import cfunc, types, carray
+
+    signature_i8args = types.float32(
+        types.CPointer(types.int8), types.CPointer(types.int8)
+    )
+
+    signature_f16args = types.float32(
+        types.CPointer(types.float16), types.CPointer(types.float16)
+    )
+
+    signature_f32args = types.float32(
+        types.CPointer(types.float32), types.CPointer(types.float32)
+    )
+
+    signature_f64args = types.float32(
+        types.CPointer(types.float64), types.CPointer(types.float64)
+    )
+
     numba_supported_types = (
         ScalarKind.F8,
         # Half-precision is still unsupported
@@ -46,7 +52,7 @@ def jit(
         ScalarKind.F64,
     )
     if dtype not in numba_supported_types:
-        return None
+        return metric
 
     scalar_kind_to_accumulator_type = {
         ScalarKind.F8: types.int32,
@@ -109,7 +115,7 @@ def jit(
 
     pointer = cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
     return CompiledMetric(
-        pointer=pointer,
+        pointer=pointer.address,
         kind=metric,
         signature=MetricSignature.ArrayArray,
     )

package/setup.py

@@ -5,97 +5,94 @@ from setuptools import setup
 from pybind11.setup_helpers import Pybind11Extension
 
 
-compile_args = ['-std=c++11', '-O3']
+compile_args = []
 link_args = []
 macros_args = [
-    ('USEARCH_USE_NATIVE_F16', '0'),
-    ('USEARCH_USE_SIMSIMD', '0'),
+    ("USEARCH_USE_NATIVE_F16", "0"),
+    ("USEARCH_USE_SIMSIMD", "0"),
 ]
 
-if sys.platform == 'linux':
-    compile_args.append('-fopenmp')
-    link_args.append('-lgomp')
-    compile_args.append('-Wno-unknown-pragmas')
-    macros_args.append(('USEARCH_USE_OPENMP', '1'))
+if sys.platform == "linux":
+    compile_args.append("-std=c++11")
+    compile_args.append("-O3")
+    compile_args.append("-Wno-unknown-pragmas")
 
-if sys.platform == 'darwin':
+    macros_args.append(("USEARCH_USE_OPENMP", "1"))
+    compile_args.append("-fopenmp")
+    link_args.append("-lgomp")
+
+if sys.platform == "darwin":
     # MacOS 10.15 or higher is needed for `aligned_alloc` support.
     # https://github.com/unum-cloud/usearch/actions/runs/4975434891/jobs/8902603392
-    compile_args.append('-mmacosx-version-min=10.15')
-    compile_args.append('-Wno-unknown-pragmas')
+    compile_args.append("-mmacosx-version-min=10.15")
+    compile_args.append("-std=c++11")
+    compile_args.append("-O3")
+    compile_args.append("-Wno-unknown-pragmas")
 
-if sys.platform == 'win32':
-    compile_args = ['/std:c++14', '/O2']
+    # Linking OpenMP requires additional preparion in CIBuildWheel
+    # macros_args.append(("USEARCH_USE_OPENMP", "1"))
+    # compile_args.append("-Xpreprocessor -fopenmp")
+    # link_args.append("-Xpreprocessor -lomp")
 
+if sys.platform == "win32":
+    compile_args.append("/std:c++14")
+    compile_args.append("/O2")
 
 ext_modules = [
     Pybind11Extension(
-        'usearch.compiled',
-        ['python/lib.cpp'],
+        "usearch.compiled",
+        ["python/lib.cpp"],
         extra_compile_args=compile_args,
         extra_link_args=link_args,
         define_macros=macros_args,
     ),
 ]
 
-__version__ = open('VERSION', 'r').read().strip()
-__lib_name__ = 'usearch'
+__version__ = open("VERSION", "r").read().strip()
+__lib_name__ = "usearch"
 
 
 this_directory = os.path.abspath(os.path.dirname(__file__))
-with open(os.path.join(this_directory, 'README.md')) as f:
+with open(os.path.join(this_directory, "README.md")) as f:
     long_description = f.read()
 
 
 setup(
-
     name=__lib_name__,
     version=__version__,
-    packages=['usearch'],
-    package_dir={'usearch': 'python/usearch'},
-
-    description='Smaller & Faster Single-File Vector Search Engine from Unum',
+    packages=["usearch"],
+    package_dir={"usearch": "python/usearch"},
+    description="Smaller & Faster Single-File Vector Search Engine from Unum",
     long_description=long_description,
-    long_description_content_type='text/markdown',
-    license='Apache-2.0',
-
+    long_description_content_type="text/markdown",
+    license="Apache-2.0",
     classifiers=[
-        'Development Status :: 5 - Production/Stable',
-
-        'Natural Language :: English',
-        'Intended Audience :: Developers',
-        'Intended Audience :: Information Technology',
-        'License :: OSI Approved :: Apache Software License',
-
-        'Programming Language :: C++',
-        'Programming Language :: Python :: 3 :: Only',
-        'Programming Language :: Python :: Implementation :: CPython',
-        'Programming Language :: Java',
-        'Programming Language :: JavaScript',
-        'Programming Language :: Objective C',
-        'Programming Language :: Rust',
-        'Programming Language :: Other',
-
-        'Operating System :: MacOS',
-        'Operating System :: Unix',
-        'Operating System :: Microsoft :: Windows',
-
-        'Topic :: System :: Clustering',
-        'Topic :: Database :: Database Engines/Servers',
-        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        "Development Status :: 5 - Production/Stable",
+        "Natural Language :: English",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Information Technology",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: C++",
+        "Programming Language :: Python :: 3 :: Only",
+        "Programming Language :: Python :: Implementation :: CPython",
+        "Programming Language :: Java",
+        "Programming Language :: JavaScript",
+        "Programming Language :: Objective C",
+        "Programming Language :: Rust",
+        "Programming Language :: Other",
+        "Operating System :: MacOS",
+        "Operating System :: Unix",
+        "Operating System :: Microsoft :: Windows",
+        "Topic :: System :: Clustering",
+        "Topic :: Database :: Database Engines/Servers",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
-
-    include_dirs=[
-        'include',
-        'fp16/include', 'robin-map/include', 'simsimd/include'],
-
+    include_dirs=["include", "fp16/include", "robin-map/include", "simsimd/include"],
     ext_modules=ext_modules,
     install_requires=[
-        'numpy',
-        'pandas',
+        "numpy",
+        "pandas",
+        "tqdm",
         'ucall; python_version >= "3.9"',
     ],
-    extras_require={
-        'jit': ['numba', 'cppyy'],
-    },
 )

package/swift/Test.swift

@@ -6,16 +6,20 @@
 //
 
 import Foundation
+import XCTest
 import USearch
 
 @available(iOS 13, macOS 10.15, tvOS 13.0, watchOS 6.0, *)
-func testUnit() throws {
-    let index = USearchIndex.make(metric:USearchMetric.l2sq, dimensions: 4, connectivity: 8, quantization:USearchScalar.F32)
-    let vectorA: [Float32] = [0.3, 0.5, 1.2, 1.4]
-    let vectorB: [Float32] = [0.4, 0.2, 1.2, 1.1]
-    index.add(label: 42, vector: vectorA[...])
-    index.add(label: 43, vector: vectorB[...])
-
-    let results = index.search(vector: vectorA[...], count: 10)
-    assert(results.0[0] == 42)
+class Test: XCTestCase {
+    func testUnit() throws {
+        let index = USearchIndex.make(metric:USearchMetric.l2sq, dimensions: 4, connectivity: 8, quantization:USearchScalar.F32)
+        let vectorA: [Float32] = [0.3, 0.5, 1.2, 1.4]
+        let vectorB: [Float32] = [0.4, 0.2, 1.2, 1.1]
+        index.reserve(2)
+        index.add(label: 42, vector: vectorA[...])
+        index.add(label: 43, vector: vectorB[...])
+        
+        let results = index.search(vector: vectorA[...], count: 10)
+        assert(results.0[0] == 42)
+    }
 }