usearch 0.22.3 → 1.1.0

package/python/scripts/bench.ipynb

@@ -6,7 +6,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install faiss-cpu numba pandas plotly nbformat jinja2 kaleido tqdm sklearn scikit-learn scipy matplotlib"
+    "!pip install faiss-cpu numba pandas plotly nbformat jinja2 kaleido tqdm scikit-learn scipy matplotlib"
    ]
   },
   {
@@ -61,6 +61,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from usearch.eval import Dataset\n",
     "data = Dataset.build(\n",
     "    vectors='../../datasets/wiki_1M/base.1M.fbin',\n",
     "    queries='../../datasets/wiki_1M/query.public.100K.fbin',\n",
@@ -161,7 +162,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.10"
+   "version": "3.10.11"
   },
   "orig_nbformat": 4
  },

package/python/scripts/bench.py

@@ -5,7 +5,7 @@ from dataclasses import asdict
 import numpy as np
 import pandas as pd
 
-from usearch.index import Index, Label, MetricKind, ScalarKind
+from usearch.index import Index, Key, MetricKind, ScalarKind
 from usearch.numba import jit as njit
 from usearch.eval import Evaluation, AddTask
 from usearch.index import (
@@ -26,7 +26,7 @@ def bench_speed(
     # Build various indexes:
     indexes = []
     jit_options = [False, True] if jit else [False]
-    dtype_options = [ScalarKind.F32, ScalarKind.F16, ScalarKind.F8]
+    dtype_options = [ScalarKind.F32, ScalarKind.F16, ScalarKind.I8]
     for jit, dtype in itertools.product(jit_options, dtype_options):
         metric = MetricKind.IP
         if jit:
@@ -38,7 +38,7 @@ def bench_speed(
             expansion_add=expansion_add,
             expansion_search=expansion_search,
             connectivity=connectivity,
-            path="USearch" + ["", "+JIT"][jit] + ":" + dtype,
+            path="USearch" + ["", "+JIT"][jit] + ":" + str(dtype),
         )
 
         # Skip the cases, where JIT-ing is impossible
@@ -115,7 +115,7 @@ def bench_params(
     results = []
     for connectivity, ndim in itertools.product(connectivities, dimensions):
         task = AddTask(
-            labels=np.arange(count, dtype=Label),
+            keys=np.arange(count, dtype=Key),
             vectors=np.random.rand(count, ndim).astype(np.float32),
         )
         index = Index(

package/python/scripts/index_faiss.py

@@ -40,14 +40,14 @@ class IndexFAISS:
 
         self.path = path
 
-    def add(self, labels, vectors):
-        # Adding labels isn't supported for most index types
-        # self._faiss.add_with_ids(vectors, labels)
+    def add(self, keys, vectors):
+        # Adding keys isn't supported for most index types
+        # self._faiss.add_with_ids(vectors, keys)
         self._faiss.add(vectors)
 
     def search(self, queries, k: int) -> Matches:
-        distances, labels = self._faiss.search(queries, k)
-        return Matches(labels, distances, np.array([k] * queries.shape[0]))
+        distances, keys = self._faiss.search(queries, k)
+        return Matches(keys, distances, np.array([k] * queries.shape[0]))
 
     def __len__(self) -> int:
         return self._faiss.ntotal

package/python/scripts/join.py

@@ -110,28 +110,28 @@ if len(run_diagnostics) == 0 or run_diagnostics.lower() == "y":
         exact=exact,
     )
 
-    secs, a_self_recall = dt(lambda: a.search(a.vectors, **args).recall(a.labels))
+    secs, a_self_recall = dt(lambda: a.search(a.vectors, **args).recall(a.keys))
     print(
         "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
             k, a_name, a_self_recall * 100, secs
         )
     )
 
-    secs, b_self_recall = dt(lambda: b.search(b.vectors, **args).recall(b.labels))
+    secs, b_self_recall = dt(lambda: b.search(b.vectors, **args).recall(b.keys))
     print(
         "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
             k, b_name, b_self_recall * 100, secs
         )
     )
 
-    secs, ab_recall = dt(lambda: b.search(a.vectors, **args).recall(b.labels))
+    secs, ab_recall = dt(lambda: b.search(a.vectors, **args).recall(b.keys))
     print(
         "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
             k, a_name, b_name, ab_recall * 100, secs
         )
     )
 
-    secs, ba_recall = dt(lambda: a.search(b.vectors, **args).recall(a.labels))
+    secs, ba_recall = dt(lambda: a.search(b.vectors, **args).recall(a.keys))
     print(
         "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
             k, b_name, a_name, ba_recall * 100, secs

package/python/scripts/test.py

@@ -5,13 +5,16 @@ import numpy as np
 
 from usearch.io import load_matrix, save_matrix
 from usearch.eval import random_vectors
+from usearch.index import search
 
 from usearch.index import (
     Index,
-    SparseIndex,
+    Indexes,
     MetricKind,
     ScalarKind,
+    Match,
     Matches,
+    BatchMatches,
 )
 from usearch.index import (
     DEFAULT_CONNECTIVITY,
@@ -19,6 +22,11 @@ from usearch.index import (
     DEFAULT_EXPANSION_SEARCH,
 )
 
+# When running the tests concurrently in different VMs, one can get access violation errors.
+# To avoid that, we embed the OS name into the temporary files name.
+# https://github.com/unum-cloud/usearch/actions/runs/5705359698/job/15459945738
+temporary_filename = f"tmp-{os.name}"
+temporary_usearch_filename = temporary_filename + ".usearch"
 
 dimensions = [3, 97, 256]
 batch_sizes = [1, 77]
@@ -26,7 +34,7 @@ index_types = [
     ScalarKind.F32,
     ScalarKind.F64,
     ScalarKind.F16,
-    ScalarKind.F8,
+    ScalarKind.I8,
 ]
 numpy_types = [np.float32, np.float64, np.float16]
 
@@ -52,10 +60,10 @@ def test_serializing_fbin_matrix(rows: int, cols: int):
     :param int cols: The number of columns in the matrix.
     """
     original = np.random.rand(rows, cols).astype(np.float32)
-    save_matrix(original, "tmp.fbin")
-    reconstructed = load_matrix("tmp.fbin")
+    save_matrix(original, temporary_filename + ".fbin")
+    reconstructed = load_matrix(temporary_filename + ".fbin")
     assert np.allclose(original, reconstructed)
-    os.remove("tmp.fbin")
+    os.remove(temporary_filename + ".fbin")
 
 
 @pytest.mark.parametrize("rows", batch_sizes)
@@ -68,10 +76,31 @@ def test_serializing_ibin_matrix(rows: int, cols: int):
     :param int cols: The number of columns in the matrix.
     """
     original = np.random.randint(0, rows + 1, size=(rows, cols)).astype(np.int32)
-    save_matrix(original, "tmp.ibin")
-    reconstructed = load_matrix("tmp.ibin")
+    save_matrix(original, temporary_filename + ".ibin")
+    reconstructed = load_matrix(temporary_filename + ".ibin")
     assert np.allclose(original, reconstructed)
-    os.remove("tmp.ibin")
+    os.remove(temporary_filename + ".ibin")
+
+
+@pytest.mark.parametrize("rows", batch_sizes)
+@pytest.mark.parametrize("cols", dimensions)
+def test_exact_search(rows: int, cols: int):
+    """
+    Test exact search.
+
+    :param int rows: The number of rows in the matrix.
+    :param int cols: The number of columns in the matrix.
+    """
+    original = np.random.rand(rows, cols)
+    matches: BatchMatches = search(original, original, 10, exact=True)
+    top_matches = (
+        [int(m.keys[0]) for m in matches] if rows > 1 else int(matches.keys[0])
+    )
+    assert np.all(top_matches == np.arange(rows))
+
+    matches: Matches = search(original, original[0], 10, exact=True)
+    top_match = int(matches.keys[0])
+    assert top_match == 0
 
 
 @pytest.mark.parametrize("ndim", dimensions)
@@ -100,23 +129,26 @@ def test_index(
     vector = random_vectors(count=1, ndim=ndim, dtype=numpy_type).flatten()
     index.add(42, vector)
 
-    assert 42 in index
-    assert 42 in index.labels
-    assert 43 not in index
-    assert index[42] is not None
-    assert index[43] is None
+    assert len(index) == 1, "Size after addition"
+    assert 42 in index, "Presence in the index"
+    assert 42 in index.keys, "Presence among keys"
+    assert 43 not in index, "Presence in the index, false positive"
+    assert index[42] is not None, "Vector recovery"
+    assert index[43] is None, "Vector recovery, false positive"
     assert len(index[42]) == ndim
     if numpy_type != np.byte:
         assert np.allclose(index[42], vector, atol=0.1)
 
-    matches, distances, count = index.search(vector, 10)
-    assert len(index) == 1
-    assert len(matches) == count
-    assert len(distances) == count
-    assert count == 1
-    assert matches[0] == 42
-    assert distances[0] == pytest.approx(0, abs=1e-3)
+    matches: Matches = index.search(vector, 10)
+    assert len(matches.keys) == 1, "Number of matches"
+    assert len(matches.keys) == len(matches.distances), "Symmetric match sub-arrays"
+    assert len({match.key for match in matches}) == 1, "Iteration over matches"
+    assert matches[0].key == 42
+    assert matches[0].distance == pytest.approx(0, abs=1e-3)
+    assert matches.computed_distances != 0
+    assert matches.visited_members != 0
 
+    # Validating the index structure and metadata:
     assert index.max_level >= 0
     assert index.levels_stats.nodes >= 1
     assert index.level_stats(0).nodes == 1
@@ -129,30 +161,61 @@ def test_index(
     index.remove(43)
     assert len(index) == 1
 
-    index.save("tmp.usearch")
+    # Try inserting back
+    index.add(43, other_vector)
+    assert len(index) == 2
+    index.remove(43)
+    assert len(index) == 1
+
+    index.save(temporary_usearch_filename)
 
     # Re-populate cleared index
     index.clear()
     assert len(index) == 0
     index.add(42, vector)
     assert len(index) == 1
-    matches, distances, count = index.search(vector, 10)
-    assert count == 1
+    matches: Matches = index.search(vector, 10)
+    assert len(matches) == 1
+
+    index_copy = index.copy()
+    assert len(index_copy) == 1
+    assert len(index_copy[42]) == ndim
+    matches_copy: Matches = index_copy.search(vector, 10)
+    assert np.all(matches_copy.keys == matches.keys)
 
-    index.load("tmp.usearch")
+    index.load(temporary_usearch_filename)
     assert len(index) == 1
     assert len(index[42]) == ndim
 
-    index = Index.restore("tmp.usearch")
+    matches_loaded: Matches = index.search(vector, 10)
+    assert np.all(matches_loaded.keys == matches.keys)
+
+    index = Index.restore(temporary_usearch_filename, view=True)
     assert len(index) == 1
     assert len(index[42]) == ndim
 
-    index_copy = index.copy()
-    assert len(index_copy) == 1
-    assert len(index_copy[42]) == ndim
+    matches_viewed: Matches = index.search(vector, 10)
+    assert np.all(matches_viewed.keys == matches.keys)
 
     # Cleanup
-    os.remove("tmp.usearch")
+    index.reset()
+    os.remove(temporary_usearch_filename)
+
+    # Try opening a missing file
+    meta = Index.metadata(temporary_usearch_filename)
+    assert meta is None
+    index = Index.restore(temporary_usearch_filename)
+    assert index is None
+
+    # Try opening a corrupt file
+    with open(temporary_usearch_filename, "w") as file:
+        file.write("Some random string")
+    meta = Index.metadata(temporary_usearch_filename)
+    assert meta is None
+    index = Index.restore(temporary_usearch_filename)
+
+    assert index is None
+    os.remove(temporary_usearch_filename)
 
 
 @pytest.mark.parametrize("ndim", dimensions)
@@ -169,36 +232,50 @@ def test_index_batch(
 ):
     index = Index(ndim=ndim, metric=metric, dtype=index_type)
 
-    labels = np.arange(batch_size)
+    keys = np.arange(batch_size)
     vectors = random_vectors(count=batch_size, ndim=ndim, dtype=numpy_type)
 
-    index.add(labels, vectors, threads=2)
+    index.add(keys, vectors, threads=2)
     assert len(index) == batch_size
-    assert np.allclose(index.get_vectors(labels).astype(numpy_type), vectors, atol=0.1)
+    assert np.allclose(index.get_vectors(keys).astype(numpy_type), vectors, atol=0.1)
 
-    matches: Matches = index.search(vectors, 10, threads=2)
-    assert matches.labels.shape[0] == matches.distances.shape[0]
-    assert matches.counts.shape[0] == batch_size
-    assert np.all(np.sort(index.labels) == np.sort(labels))
+    matches: BatchMatches = index.search(vectors, 10, threads=2)
+    assert matches.keys.shape[0] == matches.distances.shape[0]
+    assert len(matches) == batch_size
+    assert np.all(np.sort(index.keys) == np.sort(keys))
 
-    assert index.max_level >= 0  # TODO: This should be 1
+    if batch_size > 1:
+        assert index.max_level >= 1
+    else:
+        assert index.max_level >= 0
     assert index.levels_stats.nodes >= batch_size
     assert index.level_stats(0).nodes == batch_size
 
-    index.save("tmp.usearch")
+    index.save(temporary_usearch_filename)
     index.clear()
     assert len(index) == 0
 
-    index.load("tmp.usearch")
+    index.load(temporary_usearch_filename)
     assert len(index) == batch_size
     assert len(index[0]) == ndim
 
-    index = Index.restore("tmp.usearch")
+    if batch_size > 1:
+        matches_loaded: BatchMatches = index.search(vectors, 10, threads=2)
+        for idx in range(len(matches_loaded)):
+            assert np.all(matches_loaded[idx].keys == matches[idx].keys)
+
+    index = Index.restore(temporary_usearch_filename, view=True)
     assert len(index) == batch_size
     assert len(index[0]) == ndim
 
+    if batch_size > 1:
+        matches_viewed: BatchMatches = index.search(vectors, 10, threads=2)
+        for idx in range(len(matches_viewed)):
+            assert np.all(matches_viewed[idx].keys == matches[idx].keys)
+
     # Cleanup
-    os.remove("tmp.usearch")
+    index.reset()
+    os.remove(temporary_usearch_filename)
 
 
 @pytest.mark.parametrize("metric", [MetricKind.L2sq])
@@ -213,33 +290,54 @@ def test_exact_recall(
 ):
     ndim: int = batch_size
     index = Index(ndim=ndim, metric=metric, dtype=index_type)
+    assert index.ndim == ndim
 
     vectors = np.zeros((batch_size, ndim), dtype=numpy_type)
     for i in range(batch_size):
         vectors[i, i] = 1
-    labels = np.arange(batch_size)
-    index.add(labels, vectors)
+    keys = np.arange(batch_size)
+    index.add(keys, vectors)
     assert len(index) == batch_size
 
     # Search one at a time
     for i in range(batch_size):
         matches: Matches = index.search(vectors[i], 10, exact=True)
-        found_labels = matches.labels
+        found_labels = matches.keys
         assert found_labels[0] == i
+        assert matches.computed_distances == len(index)
+        assert matches.visited_members == 0, "Exact search won't traverse the graph"
 
     # Search the whole batch
-    matches: Matches = index.search(vectors, 10, exact=True)
-    found_labels = matches.labels
-    for i in range(batch_size):
-        assert found_labels[i, 0] == i
+    if batch_size > 1:
+        matches: BatchMatches = index.search(vectors, 10, exact=True)
+        assert matches.computed_distances == len(index) * len(vectors)
+        assert matches.visited_members == 0, "Exact search won't traverse the graph"
 
-    # Match entries aginst themselves
+        found_labels = matches.keys
+        for i in range(batch_size):
+            assert found_labels[i, 0] == i
+
+    # Match entries against themselves
     index_copy: Index = index.copy()
     mapping: dict = index.join(index_copy, exact=True)
     for man, woman in mapping.items():
         assert man == woman, "Stable marriage failed"
 
 
+def test_indexes():
+    ndim = 10
+    index_a = Index(ndim=ndim)
+    index_b = Index(ndim=ndim)
+
+    vectors = random_vectors(count=3, ndim=ndim)
+    index_a.add(42, vectors[0])
+    index_b.add(43, vectors[1])
+
+    indexes = Indexes([index_a, index_b])
+    matches = indexes.search(vectors[2], 10)
+    assert len(matches) == 2
+
+
 @pytest.mark.parametrize("bits", dimensions)
 @pytest.mark.parametrize("metric", hash_metrics)
 @pytest.mark.parametrize("connectivity", connectivity_options)
@@ -252,22 +350,15 @@ def test_bitwise_index(
 ):
     index = Index(ndim=bits, metric=metric, connectivity=connectivity)
 
-    labels = np.arange(batch_size)
+    keys = np.arange(batch_size)
     byte_vectors = np.random.randint(2, size=(batch_size, bits))
     bit_vectors = np.packbits(byte_vectors, axis=1)
 
-    index.add(labels, bit_vectors)
-    assert np.allclose(index.get_vectors(labels), byte_vectors, atol=0.1)
-    assert np.all(index.get_vectors(labels, ScalarKind.B1) == bit_vectors)
+    index.add(keys, bit_vectors)
+    assert np.all(index.get_vectors(keys, ScalarKind.B1) == bit_vectors)
 
     index.search(bit_vectors, 10)
 
 
-@pytest.mark.parametrize("connectivity", connectivity_options)
-@pytest.mark.skipif(os.name == "nt", reason="Spurious behaviour on windows")
-def test_sets_index(connectivity: int):
-    index = SparseIndex(connectivity=connectivity)
-    index.add(10, np.array([10, 12, 15], dtype=np.uint32))
-    index.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
-    results = index.search(np.array([12, 15], dtype=np.uint32), 10)
-    assert list(results) == [10, 11]
+if __name__ == "__main__":
+    pytest.main(args=["python/scripts/test.py", "-s", "-x", "-v"])

package/python/scripts/test_jit.py

@@ -85,13 +85,12 @@ def test_index_numba(ndim: int, batch_size: int):
         )
         index = Index(ndim=ndim, metric=metric)
 
-        labels = np.arange(batch_size)
+        keys = np.arange(batch_size)
         vectors = random_vectors(count=batch_size, ndim=ndim)
 
-        index.add(labels, vectors)
-        matches, distances, count = index.search(vectors, 10)
-        assert matches.shape[0] == distances.shape[0]
-        assert count.shape[0] == batch_size
+        index.add(keys, vectors)
+        matches = index.search(vectors, 10)
+        assert len(matches) == batch_size
 
 
 @pytest.mark.parametrize("ndim", dimensions[-1:])
@@ -149,13 +148,12 @@ def test_index_cppyy(ndim: int, batch_size: int):
         )
         index = Index(ndim=ndim, metric=metric)
 
-        labels = np.arange(batch_size)
+        keys = np.arange(batch_size)
         vectors = random_vectors(count=batch_size, ndim=ndim)
 
-        index.add(labels, vectors)
-        matches, distances, count = index.search(vectors, 10)
-        assert matches.shape[0] == distances.shape[0]
-        assert count.shape[0] == batch_size
+        index.add(keys, vectors)
+        matches = index.search(vectors, 10)
+        assert len(matches) == batch_size
 
 
 @pytest.mark.parametrize("ndim", [8])
@@ -235,10 +233,10 @@ def test_index_peachpy(ndim: int, batch_size: int):
     )
     index = Index(ndim=ndim, metric=metric)
 
-    labels = np.arange(batch_size)
+    keys = np.arange(batch_size)
     vectors = random_vectors(count=batch_size, ndim=ndim)
 
-    index.add(labels, vectors)
+    index.add(keys, vectors)
     matches, distances, count = index.search(vectors, 10)
     assert matches.shape[0] == distances.shape[0]
     assert count.shape[0] == batch_size

package/python/scripts/test_sparse.py

@@ -0,0 +1,32 @@
+import os
+
+import pytest
+import numpy as np
+
+from usearch.index import (
+    SparseIndex,
+)
+from usearch.index import (
+    DEFAULT_CONNECTIVITY,
+    DEFAULT_EXPANSION_ADD,
+    DEFAULT_EXPANSION_SEARCH,
+)
+
+
+dimensions = [3, 97, 256]
+batch_sizes = [1, 77]
+connectivity_options = [3, 13, 50, DEFAULT_CONNECTIVITY]
+
+
+@pytest.mark.parametrize("connectivity", connectivity_options)
+@pytest.mark.skipif(os.name == "nt", reason="Spurious behaviour on windows")
+def test_sets_index(connectivity: int):
+    index = SparseIndex(connectivity=connectivity)
+    index.add(10, np.array([10, 12, 15], dtype=np.uint32))
+    index.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
+    results = index.search(np.array([12, 15], dtype=np.uint32), 10)
+    assert list(results) == [10, 11]
+
+
+if __name__ == "__main__":
+    pytest.main(args=["python/scripts/test.py", "-s", "-x", "-v"])

package/python/usearch/client.py

@@ -21,36 +21,33 @@ def _vector_to_ascii(vector: np.ndarray) -> Optional[str]:
 
 
 class IndexClient:
-
     def __init__(
-            self,
-            uri: str = '127.0.0.1',
-            port: int = 8545,
-            use_http: bool = True) -> None:
+        self, uri: str = "127.0.0.1", port: int = 8545, use_http: bool = True
+    ) -> None:
         self.client = Client(uri=uri, port=port, use_http=use_http)
 
-    def add_one(self, label: int, vector: np.ndarray):
-        assert isinstance(label, int)
+    def add_one(self, key: int, vector: np.ndarray):
+        assert isinstance(key, int)
         assert isinstance(vector, np.ndarray)
         vector = vector.flatten()
         ascii = _vector_to_ascii(vector)
         if ascii:
-            self.client.add_ascii(label=label, string=ascii)
+            self.client.add_ascii(key=key, string=ascii)
         else:
-            self.client.add_one(label=label, vectors=vector)
+            self.client.add_one(key=key, vectors=vector)
 
-    def add_many(self, labels: np.ndarray, vectors: np.ndarray):
-        assert isinstance(labels, int)
+    def add_many(self, keys: np.ndarray, vectors: np.ndarray):
+        assert isinstance(keys, int)
         assert isinstance(vectors, np.ndarray)
-        assert labels.ndim == 1 and vectors.ndim == 2
-        assert labels.shape[0] == vectors.shape[0]
-        self.client.add_many(labels=labels, vectors=vectors)
+        assert keys.ndim == 1 and vectors.ndim == 2
+        assert keys.shape[0] == vectors.shape[0]
+        self.client.add_many(keys=keys, vectors=vectors)
 
-    def add(self, labels: Union[np.ndarray, int], vectors: np.ndarray):
-        if isinstance(labels, int) or len(labels) == 1:
-            return self.add_one(labels, vectors)
+    def add(self, keys: Union[np.ndarray, int], vectors: np.ndarray):
+        if isinstance(keys, int) or len(keys) == 1:
+            return self.add_one(keys, vectors)
         else:
-            return self.add_many(labels, vectors)
+            return self.add_many(keys, vectors)
 
     def search_one(self, vector: np.ndarray, count: int) -> Matches:
         matches: List[dict] = []
@@ -64,31 +61,32 @@ class IndexClient:
         print(matches.data)
         matches = matches.json
 
-        labels = np.array((1, count), dtype=np.uint32)
+        keys = np.array((1, count), dtype=np.uint32)
         distances = np.array((1, count), dtype=np.float32)
         counts = np.array((1), dtype=np.uint32)
         for col, result in enumerate(matches):
-            labels[0, col] = result['label']
-            distances[0, col] = result['distance']
+            keys[0, col] = result["key"]
+            distances[0, col] = result["distance"]
         counts[0] = len(matches)
 
-        return labels, distances, counts
+        return keys, distances, counts
 
     def search_many(self, vectors: np.ndarray, count: int) -> Matches:
         batch_size: int = vectors.shape[0]
         list_of_matches: List[List[dict]] = self.client.search_many(
-            vectors=vectors, count=count)
+            vectors=vectors, count=count
+        )
 
-        labels = np.array((batch_size, count), dtype=np.uint32)
+        keys = np.array((batch_size, count), dtype=np.uint32)
         distances = np.array((batch_size, count), dtype=np.float32)
         counts = np.array((batch_size), dtype=np.uint32)
         for row, matches in enumerate(list_of_matches):
             for col, result in enumerate(matches):
-                labels[row, col] = result['label']
-                distances[row, col] = result['distance']
+                keys[row, col] = result["key"]
+                distances[row, col] = result["distance"]
             counts[row] = len(results)
 
-        return labels, distances, counts
+        return keys, distances, counts
 
     def search(self, vectors: np.ndarray, count: int) -> Matches:
         if vectors.ndim == 1 or (vectors.ndim == 2 and vectors.shape[0] == 1):
@@ -119,7 +117,7 @@ class IndexClient:
         raise NotImplementedError()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     index = IndexClient()
     index.add(42, np.array([0.4] * 256, dtype=np.float32))
     results = index.search(np.array([0.4] * 256, dtype=np.float32), 10)