usearch 2.2.0 → 2.3.0

package/python/usearch/index.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from inspect import signature
 
 # The purpose of this file is to provide Pythonic wrapper on top
 # the native precompiled CPython module. It improves compatibility
@@ -8,6 +9,7 @@ import os
 import math
 from dataclasses import dataclass
 from typing import (
+    Any,
     Optional,
     Union,
     NamedTuple,
@@ -72,6 +74,17 @@ NoneType = type(None)
 
 PathOrBuffer = Union[str, os.PathLike, bytes]
 
+ProgressCallback = Callable[[int, int], bool]
+
+
+def _match_signature(
+    func: Callable[[Any], Any], arg_types: List[type], ret_type: type
+) -> bool:
+    assert callable(func), "Not callable"
+    sig = signature(func)
+    param_types = [param.annotation for param in sig.parameters.values()]
+    return param_types == arg_types and sig.return_annotation == ret_type
+
 
 def _normalize_dtype(
     dtype,
@@ -153,12 +166,16 @@ def _search_in_compiled(
     vectors: np.ndarray,
     *,
     log: Union[str, bool],
-    batch_size: int,
+    progress: Optional[ProgressCallback],
     **kwargs,
 ) -> Union[Matches, BatchMatches]:
     #
     assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
     assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
+    assert not progress or _match_signature(
+        progress, [int, int], bool
+    ), "Invalid callback"
+
     if vectors.ndim == 1:
         vectors = vectors.reshape(1, len(vectors))
     count_vectors = vectors.shape[0]
@@ -168,42 +185,25 @@ def _search_in_compiled(
     ) -> Union[BatchMatches, Matches]:
         return batch_matches[0] if count_vectors == 1 else batch_matches
 
-    if log and batch_size == 0:
-        batch_size = int(math.ceil(count_vectors / 100))
-
-    if batch_size:
-        tasks = [
-            vectors[start_row : start_row + batch_size, :]
-            for start_row in range(0, count_vectors, batch_size)
-        ]
-        tasks_matches = []
+    # Create progress bar if needed
+    if log:
         name = log if isinstance(log, str) else "Search"
         pbar = tqdm(
-            tasks,
             desc=name,
             total=count_vectors,
             unit="vector",
-            disable=log is False,
         )
-        for vectors in tasks:
-            tuple_ = compiled_callable(vectors, **kwargs)
-            tasks_matches.append(BatchMatches(*tuple_))
-            pbar.update(vectors.shape[0])
 
-        pbar.close()
-        return distil_batch(
-            BatchMatches(
-                keys=np.vstack([m.keys for m in tasks_matches]),
-                distances=np.vstack([m.distances for m in tasks_matches]),
-                counts=np.concatenate([m.counts for m in tasks_matches], axis=None),
-                visited_members=sum([m.visited_members for m in tasks_matches]),
-                computed_distances=sum([m.computed_distances for m in tasks_matches]),
-            )
-        )
+        def update_progress_bar(processed: int, total: int) -> bool:
+            pbar.update(processed - pbar.n)
+            return progress if progress else True
 
+        tuple_ = compiled_callable(vectors, progress=update_progress_bar, **kwargs)
+        pbar.close()
     else:
         tuple_ = compiled_callable(vectors, **kwargs)
-        return distil_batch(BatchMatches(*tuple_))
+
+    return distil_batch(BatchMatches(*tuple_))
 
 
 def _add_to_compiled(
@@ -214,9 +214,13 @@ def _add_to_compiled(
     copy: bool,
     threads: int,
     log: Union[str, bool],
-    batch_size: int,
+    progress: Optional[ProgressCallback],
 ) -> Union[int, np.ndarray]:
+    #
     assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
+    assert not progress or _match_signature(
+        progress, [int, int], bool
+    ), "Invalid callback"
     assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
     if vectors.ndim == 1:
         vectors = vectors.reshape(1, len(vectors))
@@ -235,37 +239,29 @@ def _add_to_compiled(
 
     assert len(keys) == count_vectors
 
-    # If logging is requested, and batch size is undefined, set it to grow 1% at a time:
-    if log and batch_size == 0:
-        batch_size = int(math.ceil(count_vectors / 100))
-
-    # Split into batches and log progress, if needed
-    if batch_size:
-        keys = [
-            keys[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(keys, vectors)
+    # Create progress bar if needed
+    if log:
         name = log if isinstance(log, str) else "Add"
         pbar = tqdm(
-            tasks,
             desc=name,
             total=count_vectors,
             unit="vector",
-            disable=log is False,
         )
-        for keys, vectors in tasks:
-            compiled.add_many(keys, vectors, copy=copy, threads=threads)
-            pbar.update(len(keys))
 
-        pbar.close()
+        def update_progress_bar(processed: int, total: int) -> bool:
+            pbar.update(processed - pbar.n)
+            return progress(processed, total) if progress else True
 
+        compiled.add_many(
+            keys,
+            vectors,
+            copy=copy,
+            threads=threads,
+            progress=update_progress_bar,
+        )
+        pbar.close()
     else:
-        compiled.add_many(keys, vectors, copy=copy, threads=threads)
+        compiled.add_many(keys, vectors, copy=copy, threads=threads, progress=progress)
 
     return keys
 
@@ -625,7 +621,7 @@ class Index:
         copy: bool = True,
         threads: int = 0,
         log: Union[str, bool] = False,
-        batch_size: int = 0,
+        progress: Optional[ProgressCallback] = None,
     ) -> Union[int, np.ndarray]:
         """Inserts one or move vectors into the index.
 
@@ -651,8 +647,8 @@ class Index:
         :type threads: int, defaults to 0
         :param log: Whether to print the progress bar
         :type log: Union[str, bool], defaults to False
-        :param batch_size: Number of vectors to process at once
-        :type batch_size: int, defaults to 0
+        :param progress: Callback to report stats of the progress and control it
+        :type progress: Optional[ProgressCallback], defaults to None
         :return: Inserted key or keys
         :type: Union[int, np.ndarray]
         """
@@ -663,7 +659,7 @@ class Index:
             copy=copy,
             threads=threads,
             log=log,
-            batch_size=batch_size,
+            progress=progress,
         )
 
     def search(
@@ -675,7 +671,7 @@ class Index:
         threads: int = 0,
         exact: bool = False,
         log: Union[str, bool] = False,
-        batch_size: int = 0,
+        progress: Optional[ProgressCallback] = None,
     ) -> Union[Matches, BatchMatches]:
         """
         Performs approximate nearest neighbors search for one or more queries.
@@ -690,8 +686,8 @@ class Index:
         :type exact: bool, defaults to False
         :param log: Whether to print the progress bar, default to False
         :type log: Union[str, bool], optional
-        :param batch_size: Number of vectors to process at once
-        :type batch_size: int, defaults to 0
+        :param progress: Callback to report stats of the progress and control it
+        :type progress: Optional[ProgressCallback], defaults to None
         :return: Matches for one or more queries
         :rtype: Union[Matches, BatchMatches]
         """
@@ -701,11 +697,11 @@ class Index:
             vectors,
             # Batch scheduling:
             log=log,
-            batch_size=batch_size,
             # Search constraints:
             count=count,
             exact=exact,
             threads=threads,
+            progress=progress,
         )
 
     def contains(self, keys: KeyOrKeysLike) -> Union[bool, np.ndarray]:
@@ -932,37 +928,57 @@ class Index:
         self._compiled.expansion_search = v
 
     def save(
-        self, path_or_buffer: Union[str, os.PathLike, NoneType] = None
+        self,
+        path_or_buffer: Union[str, os.PathLike, NoneType] = None,
+        progress: Optional[ProgressCallback] = None,
     ) -> Optional[bytes]:
+        assert not progress or _match_signature(
+            progress, [int, int], bool
+        ), "Invalid callback signature"
+
         path_or_buffer = path_or_buffer if path_or_buffer else self.path
         if path_or_buffer is None:
-            return self._compiled.save_index_to_buffer()
+            return self._compiled.save_index_to_buffer(progress)
         else:
-            self._compiled.save_index_to_path(os.fspath(path_or_buffer))
+            self._compiled.save_index_to_path(os.fspath(path_or_buffer), progress)
+
+    def load(
+        self,
+        path_or_buffer: Union[str, os.PathLike, bytes, NoneType] = None,
+        progress: Optional[ProgressCallback] = None,
+    ):
+        assert not progress or _match_signature(
+            progress, [int, int], bool
+        ), "Invalid callback signature"
 
-    def load(self, path_or_buffer: Union[str, os.PathLike, bytes, NoneType] = None):
         path_or_buffer = path_or_buffer if path_or_buffer else self.path
         if path_or_buffer is None:
             raise Exception("Define the source")
         if isinstance(path_or_buffer, bytearray):
             path_or_buffer = bytes(path_or_buffer)
         if isinstance(path_or_buffer, bytes):
-            self._compiled.load_index_from_buffer(path_or_buffer)
+            self._compiled.load_index_from_buffer(path_or_buffer, progress)
         else:
-            self._compiled.load_index_from_path(os.fspath(path_or_buffer))
+            self._compiled.load_index_from_path(os.fspath(path_or_buffer), progress)
 
     def view(
-        self, path_or_buffer: Union[str, os.PathLike, bytes, bytearray, NoneType] = None
+        self,
+        path_or_buffer: Union[str, os.PathLike, bytes, bytearray, NoneType] = None,
+        progress: Optional[ProgressCallback] = None,
     ):
+        assert not progress or _match_signature(
+            progress, [int, int], bool
+        ), "Invalid callback signature"
+
         path_or_buffer = path_or_buffer if path_or_buffer else self.path
         if path_or_buffer is None:
             raise Exception("Define the source")
         if isinstance(path_or_buffer, bytearray):
             path_or_buffer = bytes(path_or_buffer)
         if isinstance(path_or_buffer, bytes):
-            self._compiled.view_index_from_buffer(path_or_buffer)
+            self._compiled.view_index_from_buffer(path_or_buffer, progress)
         else:
-            self._compiled.view_index_from_path(os.fspath(path_or_buffer))
+            self._compiled.view_index_from_path(os.fspath(path_or_buffer), progress)
 
     def clear(self):
         """Erases all the vectors from the index, preserving the space for future insertions."""
@@ -993,6 +1009,7 @@ class Index:
         other: Index,
         max_proposals: int = 0,
         exact: bool = False,
+        progress: Optional[ProgressCallback] = None,
     ) -> Dict[Key, Key]:
         """Performs "Semantic Join" or pairwise matching between `self` & `other` index.
         Is different from `search`, as no collisions are allowed in resulting pairs.
@@ -1005,13 +1022,20 @@ class Index:
         :type max_proposals: int, optional
         :param exact: Controls if underlying `search` should be exact, defaults to False
         :type exact: bool, optional
+        :param progress: Callback to report stats of the progress and control it
+        :type progress: Optional[ProgressCallback], defaults to None
         :return: Mapping from keys of `self` to keys of `other`
         :rtype: Dict[Key, Key]
         """
+        assert not progress or _match_signature(
+            progress, [int, int], bool
+        ), "Invalid callback signature"
+
         return self._compiled.join(
             other=other._compiled,
             max_proposals=max_proposals,
             exact=exact,
+            progress=progress,
         )
 
     def cluster(
@@ -1023,7 +1047,7 @@ class Index:
         max_count: Optional[int] = None,
         threads: int = 0,
         log: Union[str, bool] = False,
-        batch_size: int = 0,
+        progress: Optional[ProgressCallback] = None,
     ) -> Clustering:
         """
         Clusters already indexed or provided `vectors`, mapping them to various centroids.
@@ -1037,11 +1061,15 @@ class Index:
         :type threads: int, defaults to 0
         :param log: Whether to print the progress bar
         :type log: Union[str, bool], defaults to False
-        :param batch_size: Number of vectors to process at once, defaults to 0
-        :type batch_size: int, defaults to 0
+        :param progress: Callback to report stats of the progress and control it
+        :type progress: Optional[ProgressCallback], defaults to None
         :return: Matches for one or more queries
         :rtype: Union[Matches, BatchMatches]
         """
+        assert not progress or _match_signature(
+            progress, [int, int], bool
+        ), "Invalid callback signature"
+
         if min_count is None:
             min_count = 0
         if max_count is None:
@@ -1054,6 +1082,7 @@ class Index:
                 min_count=min_count,
                 max_count=max_count,
                 threads=threads,
+                progress=progress,
             )
         else:
             if keys is None:
@@ -1066,6 +1095,7 @@ class Index:
                 min_count=min_count,
                 max_count=max_count,
                 threads=threads,
+                progress=progress,
             )
 
         batch_matches = BatchMatches(*results)
@@ -1231,17 +1261,18 @@ class Indexes:
         *,
         threads: int = 0,
         exact: bool = False,
+        progress: Optional[ProgressCallback] = None,
     ):
         return _search_in_compiled(
             self._compiled.search_many,
             vectors,
             # Batch scheduling:
             log=False,
-            batch_size=None,
             # Search constraints:
             count=count,
             exact=exact,
             threads=threads,
+            progress=progress,
         )
 
 
@@ -1254,7 +1285,7 @@ def search(
     exact: bool = False,
     threads: int = 0,
     log: Union[str, bool] = False,
-    batch_size: int = 0,
+    progress: Optional[ProgressCallback] = None,
 ) -> Union[Matches, BatchMatches]:
     """Shortcut for search, that can avoid index construction. Particularly useful for
     tiny datasets, where brute-force exact search works fast enough.
@@ -1279,12 +1310,14 @@ def search(
     :type exact: bool, optional
     :param log: Whether to print the progress bar, default to False
     :type log: Union[str, bool], optional
-    :param batch_size: Number of vectors to process at once, defaults to 0
-    :type batch_size: int, optional
-
+    :param progress: Callback to report stats of the progress and control it
+    :type progress: Optional[ProgressCallback], defaults to None
     :return: Matches for one or more queries
     :rtype: Union[Matches, BatchMatches]
     """
+    assert not progress or _match_signature(
+        progress, [int, int], bool
+    ), "Invalid callback signature"
     assert dataset.ndim == 2, "Dataset must be a matrix, with a vector in each row"
 
     if not exact:
@@ -1298,14 +1331,14 @@ def search(
             dataset,
             threads=threads,
             log=log,
-            batch_size=batch_size,
+            progress=progress,
         )
         return index.search(
             query,
             count,
             threads=threads,
             log=log,
-            batch_size=batch_size,
+            progress=progress,
         )
 
     metric = _normalize_metric(metric)
@@ -1329,8 +1362,8 @@ def search(
             dataset,
             query,
             metric_kind=metric_kind,
-            metric_pointer=metric_pointer,
             metric_signature=metric_signature,
+            metric_pointer=metric_pointer,
             **kwargs,
         )
 
@@ -1339,8 +1372,8 @@ def search(
         query,
         # Batch scheduling:
         log=log,
-        batch_size=batch_size,
         # Search constraints:
         count=count,
         threads=threads,
+        progress=progress,
     )