usearch 1.2.2 → 1.3.0

package/python/lib.cpp

@@ -759,17 +759,86 @@ static void compact_index(dense_index_py_t& index, std::size_t threads) {
     index.compact(executor_default_t{threads});
 }
 
+static py::dict index_metadata(index_dense_metadata_result_t const& meta) {
+    py::dict result;
+    result["matrix_included"] = !meta.config.exclude_vectors;
+    result["matrix_uses_64_bit_dimensions"] = meta.config.use_64_bit_dimensions;
+
+    index_dense_head_t const& head = meta.head;
+    result["version"] = std::to_string(head.version_major) + "." + //
+                        std::to_string(head.version_minor) + "." + //
+                        std::to_string(head.version_patch);
+
+    result["kind_metric"] = metric_kind_t(head.kind_metric);
+    result["kind_scalar"] = scalar_kind_t(head.kind_scalar);
+    result["kind_key"] = scalar_kind_t(head.kind_key);
+    result["kind_compressed_slot"] = scalar_kind_t(head.kind_compressed_slot);
+
+    result["count_present"] = std::uint64_t(head.count_present);
+    result["count_deleted"] = std::uint64_t(head.count_deleted);
+    result["dimensions"] = std::uint64_t(head.dimensions);
+
+    return result;
+}
+
 // clang-format off
-template <typename index_at> void save_index(index_at const& index, std::string const& path) { index.save(path.c_str()).error.raise(); }
-template <typename index_at> void load_index(index_at& index, std::string const& path) { index.load(path.c_str()).error.raise(); }
-template <typename index_at> void view_index(index_at& index, std::string const& path) { index.view(path.c_str()).error.raise(); }
+template <typename index_at> void save_index_to_path(index_at const& index, std::string const& path) { index.save(path.c_str()).error.raise(); }
+template <typename index_at> void load_index_from_path(index_at& index, std::string const& path) { index.load(path.c_str()).error.raise(); }
+template <typename index_at> void view_index_from_path(index_at& index, std::string const& path) { index.view(path.c_str()).error.raise(); }
 template <typename index_at> void reset_index(index_at& index) { index.reset(); }
 template <typename index_at> void clear_index(index_at& index) { index.clear(); }
 template <typename index_at> std::size_t max_level(index_at const &index) { return index.max_level(); }
+template <typename index_at> std::size_t serialized_length(index_at const &index) { return index.serialized_length(); }
 template <typename index_at> typename index_at::stats_t compute_stats(index_at const &index) { return index.stats(); }
 template <typename index_at> typename index_at::stats_t compute_level_stats(index_at const &index, std::size_t level) { return index.stats(level); }
 // clang-format on
 
+template <typename py_bytes_at> memory_mapped_file_t memory_map_from_bytes(py_bytes_at&& bytes) {
+    py::buffer_info info(py::buffer(bytes).request());
+    return {(byte_t*)(info.ptr), static_cast<std::size_t>(info.size)};
+}
+
+template <typename index_at> py::object save_index_to_buffer(index_at const& index) {
+    std::size_t serialized_length = index.serialized_length();
+
+    // Create an empty bytearray object using CPython API
+    PyObject* byte_array = PyByteArray_FromStringAndSize(nullptr, 0);
+    if (!byte_array)
+        throw std::runtime_error("Could not allocate bytearray object");
+
+    // Resize the bytearray object to the desired length
+    if (PyByteArray_Resize(byte_array, static_cast<Py_ssize_t>(serialized_length)) != 0) {
+        Py_XDECREF(byte_array);
+        throw std::runtime_error("Could not resize bytearray object");
+    }
+
+    char* buffer = PyByteArray_AS_STRING(byte_array);
+    memory_mapped_file_t memory_map((byte_t*)buffer, serialized_length);
+    serialization_result_t result = index.save(std::move(memory_map));
+
+    if (!result) {
+        Py_XDECREF(byte_array);
+        result.error.raise();
+    }
+
+    return py::reinterpret_steal<py::object>(byte_array);
+}
+
+template <typename index_at> void load_index_from_buffer(index_at& index, py::bytes const& buffer) {
+    index.load(memory_map_from_bytes(buffer)).error.raise();
+}
+template <typename index_at> void view_index_from_buffer(index_at& index, py::bytes const& buffer) {
+    index.view(memory_map_from_bytes(buffer)).error.raise();
+}
+
+template <typename index_at> std::vector<typename index_at::stats_t> compute_levels_stats(index_at const& index) {
+    using stats_t = typename index_at::stats_t;
+    std::size_t max_level = index.max_level();
+    std::vector<stats_t> result(max_level + 1);
+    index.stats(result.data(), max_level);
+    return result;
+}
+
 template <typename internal_at, typename external_at = internal_at, typename index_at = void>
 static py::tuple get_typed_vectors_for_keys(index_at const& index, py::buffer keys) {
 
@@ -865,30 +934,16 @@ PYBIND11_MODULE(compiled, m) {
         .value("I16", scalar_kind_t::i16_k)
         .value("I8", scalar_kind_t::i8_k);
 
-    m.def("index_dense_metadata", [](std::string const& path) -> py::dict {
-        index_dense_metadata_result_t meta = index_dense_metadata(path.c_str());
+    m.def("index_dense_metadata_from_path", [](std::string const& path) -> py::dict {
+        index_dense_metadata_result_t meta = index_dense_metadata_from_path(path.c_str());
         forward_error(meta);
+        return index_metadata(meta);
+    });
 
-        index_dense_head_t const& head = meta.head;
-
-        py::dict result;
-        result["matrix_included"] = !meta.config.exclude_vectors;
-        result["matrix_uses_64_bit_dimensions"] = meta.config.use_64_bit_dimensions;
-
-        result["version"] = std::to_string(head.version_major) + "." + //
-                            std::to_string(head.version_minor) + "." + //
-                            std::to_string(head.version_patch);
-
-        result["kind_metric"] = metric_kind_t(head.kind_metric);
-        result["kind_scalar"] = scalar_kind_t(head.kind_scalar);
-        result["kind_key"] = scalar_kind_t(head.kind_key);
-        result["kind_compressed_slot"] = scalar_kind_t(head.kind_compressed_slot);
-
-        result["count_present"] = std::uint64_t(head.count_present);
-        result["count_deleted"] = std::uint64_t(head.count_deleted);
-        result["dimensions"] = std::uint64_t(head.dimensions);
-
-        return result;
+    m.def("index_dense_metadata_from_buffer", [](py::bytes const& buffer) -> py::dict {
+        index_dense_metadata_result_t meta = index_dense_metadata_from_buffer(memory_map_from_bytes(buffer));
+        forward_error(meta);
+        return index_metadata(meta);
     });
 
     m.def("exact_search", &search_many_brute_force,                        //
@@ -1034,8 +1089,9 @@ PYBIND11_MODULE(compiled, m) {
                             [](dense_index_py_t const& index) -> std::size_t { return index.metric().dimensions(); });
     i.def_property_readonly( //
         "dtype", [](dense_index_py_t const& index) -> scalar_kind_t { return index.scalar_kind(); });
-    i.def_property_readonly( //
-        "memory_usage", [](dense_index_py_t const& index) -> std::size_t { return index.memory_usage(); });
+
+    i.def_property_readonly("serialized_length", &dense_index_py_t::serialized_length);
+    i.def_property_readonly("memory_usage", &dense_index_py_t::memory_usage);
 
     i.def_property("expansion_add", &dense_index_py_t::expansion_add, &dense_index_py_t::change_expansion_add);
     i.def_property("expansion_search", &dense_index_py_t::expansion_search, &dense_index_py_t::change_expansion_search);
@@ -1138,9 +1194,14 @@ PYBIND11_MODULE(compiled, m) {
         },
         py::arg("offset"));
 
-    i.def("save", &save_index<dense_index_py_t>, py::arg("path"));
-    i.def("load", &load_index<dense_index_py_t>, py::arg("path"));
-    i.def("view", &view_index<dense_index_py_t>, py::arg("path"));
+    i.def("save_index_to_path", &save_index_to_path<dense_index_py_t>);
+    i.def("load_index_from_path", &load_index_from_path<dense_index_py_t>);
+    i.def("view_index_from_path", &view_index_from_path<dense_index_py_t>);
+
+    i.def("save_index_to_buffer", &save_index_to_buffer<dense_index_py_t>);
+    i.def("load_index_from_buffer", &load_index_from_buffer<dense_index_py_t>);
+    i.def("view_index_from_buffer", &view_index_from_buffer<dense_index_py_t>);
+
     i.def("reset", &reset_index<dense_index_py_t>);
     i.def("clear", &clear_index<dense_index_py_t>);
     i.def("copy", &copy_index, py::kw_only(), py::arg("copy") = true);
@@ -1155,7 +1216,8 @@ PYBIND11_MODULE(compiled, m) {
     i_stats.def_readonly("allocated_bytes", &punned_index_stats_t::allocated_bytes);
 
     i.def_property_readonly("max_level", &max_level<dense_index_py_t>);
-    i.def_property_readonly("levels_stats", &compute_stats<dense_index_py_t>);
+    i.def_property_readonly("stats", &compute_stats<dense_index_py_t>);
+    i.def_property_readonly("levels_stats", &compute_levels_stats<dense_index_py_t>);
     i.def("level_stats", &compute_level_stats<dense_index_py_t>, py::arg("level"));
 
     auto is = py::class_<dense_indexes_py_t>(m, "Indexes");

package/python/scripts/test_index.py

@@ -126,35 +126,63 @@ def test_index_stats(batch_size):
     index.add(keys, vectors, threads=threads)
 
     assert index.max_level >= 0
-    assert index.levels_stats.nodes >= batch_size
+    assert index.stats.nodes >= batch_size
+    assert index.levels_stats[0].nodes == batch_size
     assert index.level_stats(0).nodes == batch_size
 
+    assert index.levels_stats[index.max_level].nodes > 0
+
 
 @pytest.mark.parametrize("ndim", [1, 3, 8, 32, 256, 4096])
-@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+@pytest.mark.parametrize("batch_size", [0, 1, 7, 1024])
 @pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
 def test_index_save_load_restore_copy(ndim, quantization, batch_size):
     index = Index(ndim=ndim, dtype=quantization, multi=False)
-    keys = np.arange(batch_size)
-    vectors = random_vectors(count=batch_size, ndim=ndim)
-    index.add(keys, vectors, threads=threads)
+
+    if batch_size > 0:
+        keys = np.arange(batch_size)
+        vectors = random_vectors(count=batch_size, ndim=ndim)
+        index.add(keys, vectors, threads=threads)
 
     index.save("tmp.usearch")
     index.clear()
     assert len(index) == 0
+    assert os.path.exists("tmp.usearch")
 
     index.load("tmp.usearch")
     assert len(index) == batch_size
-    assert len(index[0].flatten()) == ndim
+    if batch_size > 0:
+        assert len(index[0].flatten()) == ndim
 
-    index = Index.restore("tmp.usearch", view=True)
+    index_meta = Index.metadata("tmp.usearch")
+    assert index_meta is not None
+
+    index = Index.restore("tmp.usearch", view=False)
     assert len(index) == batch_size
-    assert len(index[0].flatten()) == ndim
+    if batch_size > 0:
+        assert len(index[0].flatten()) == ndim
 
     copied_index = index.copy()
     assert len(copied_index) == len(index)
-    assert np.allclose(np.vstack(copied_index.get(keys)), np.vstack(index.get(keys)))
-
+    if batch_size > 0:
+        assert np.allclose(
+            np.vstack(copied_index.get(keys)), np.vstack(index.get(keys))
+        )
+
+    # Perform the same operations in RAM, without touching the filesystem
+    serialized_index = index.save()
+    deserialized_metadata = Index.metadata(serialized_index)
+    assert deserialized_metadata is not None
+
+    deserialized_index = Index.restore(serialized_index)
+    assert len(deserialized_index) == len(index)
+    assert set(np.array(deserialized_index.keys)) == set(np.array(index.keys))
+    if batch_size > 0:
+        assert np.allclose(
+            np.vstack(deserialized_index.get(keys)), np.vstack(index.get(keys))
+        )
+
+    deserialized_index.reset()
     index.reset()
     os.remove("tmp.usearch")
 

package/python/usearch/index.py

@@ -21,14 +21,21 @@ from typing import (
 import numpy as np
 from tqdm import tqdm
 
-from usearch.compiled import Index as _CompiledIndex
-from usearch.compiled import Indexes as _CompiledIndexes
-from usearch.compiled import IndexStats as _CompiledIndexStats
+# Precompiled symbols that won't be exposed directly:
+from usearch.compiled import (
+    Index as _CompiledIndex,
+    Indexes as _CompiledIndexes,
+    IndexStats as _CompiledIndexStats,
+    index_dense_metadata_from_path as _index_dense_metadata_from_path,
+    index_dense_metadata_from_buffer as _index_dense_metadata_from_buffer,
+    exact_search as _exact_search,
+)
 
-from usearch.compiled import index_dense_metadata as _index_dense_metadata
-from usearch.compiled import exact_search as _exact_search
-from usearch.compiled import MetricKind, ScalarKind, MetricSignature
+# Precompiled symbols that will be exposed
 from usearch.compiled import (
+    MetricKind,
+    ScalarKind,
+    MetricSignature,
     DEFAULT_CONNECTIVITY,
     DEFAULT_EXPANSION_ADD,
     DEFAULT_EXPANSION_SEARCH,
@@ -60,6 +67,10 @@ DTypeLike = Union[str, ScalarKind]
 
 MetricLike = Union[str, MetricKind, CompiledMetric]
 
+NoneType = type(None)
+
+PathOrBuffer = Union[str, os.PathLike, bytes]
+
 
 def _normalize_dtype(dtype, metric: MetricKind = MetricKind.Cos) -> ScalarKind:
     if dtype is None or dtype == "":
@@ -559,29 +570,38 @@ class Index:
                 self._compiled.load(path)
 
     @staticmethod
-    def metadata(path: os.PathLike) -> Optional[dict]:
-        path = os.fspath(path)
-        if not os.path.exists(path):
-            return None
+    def metadata(path_or_buffer: PathOrBuffer) -> Optional[dict]:
         try:
-            return _index_dense_metadata(path)
-        except Exception:
-            return None
+            if isinstance(path_or_buffer, bytearray):
+                path_or_buffer = bytes(path_or_buffer)
+            if isinstance(path_or_buffer, bytes):
+                return _index_dense_metadata_from_buffer(path_or_buffer)
+            else:
+                path_or_buffer = os.fspath(path_or_buffer)
+                if not os.path.exists(path_or_buffer):
+                    return None
+                return _index_dense_metadata_from_path(path_or_buffer)
+        except Exception as e:
+            raise e
 
     @staticmethod
-    def restore(path: os.PathLike, view: bool = False) -> Optional[Index]:
-        path = os.fspath(path)
-        meta = Index.metadata(path)
+    def restore(path_or_buffer: PathOrBuffer, view: bool = False) -> Optional[Index]:
+        meta = Index.metadata(path_or_buffer)
         if not meta:
             return None
-        return Index(
+
+        index = Index(
             ndim=meta["dimensions"],
             dtype=meta["kind_scalar"],
             metric=meta["kind_metric"],
-            path=path,
-            view=view,
         )
 
+        if view:
+            index.view(path_or_buffer)
+        else:
+            index.load(path_or_buffer)
+        return index
+
     def __len__(self) -> int:
         return self._compiled.__len__()
 
@@ -837,6 +857,10 @@ class Index:
     def ndim(self) -> int:
         return self._compiled.ndim
 
+    @property
+    def serialized_length(self) -> int:
+        return self._compiled.serialized_length
+
     @property
     def metric(self) -> Union[MetricKind, CompiledMetric]:
         return self._metric_jit if self._metric_jit else self._metric_kind
@@ -895,23 +919,34 @@ class Index:
     def expansion_search(self, v: int):
         self._compiled.expansion_search = v
 
-    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(os.fspath(path))
-
-    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(os.fspath(path))
-
-    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(os.fspath(path))
+    def save(self, path_or_buffer: Union[str, os.PathLike, NoneType] = None) -> Optional[bytes]:
+        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()
+        else:
+            self._compiled.save_index_to_path(os.fspath(path_or_buffer))
+
+    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)
+        else:
+            self._compiled.load_index_from_path(os.fspath(path_or_buffer))
+
+    def view(self, path_or_buffer: Union[str, os.PathLike, bytes, bytearray, 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.view_index_from_buffer(path_or_buffer)
+        else:
+            self._compiled.view_index_from_path(os.fspath(path_or_buffer))
 
     def clear(self):
         """Erases all the vectors from the index, preserving the space for future insertions."""
@@ -1049,12 +1084,27 @@ class Index:
         return self._compiled.max_level + 1
 
     @property
-    def levels_stats(self) -> _CompiledIndexStats:
+    def stats(self) -> _CompiledIndexStats:
         """Get the accumulated statistics for the entire multi-level graph.
 
         :return: Statistics for the entire multi-level graph.
         :rtype: _CompiledIndexStats
 
+        Statistics:
+            - ``nodes`` (int): The number of nodes in that level.
+            - ``edges`` (int): The number of edges in that level.
+            - ``max_edges`` (int): The maximum possible number of edges in that level.
+            - ``allocated_bytes`` (int): The amount of allocated memory for that level.
+        """
+        return self._compiled.stats
+
+    @property
+    def levels_stats(self) -> List[_CompiledIndexStats]:
+        """Get the accumulated statistics for every level graph.
+
+        :return: Statistics for every level graph.
+        :rtype: List[_CompiledIndexStats]
+
         Statistics:
             - ``nodes`` (int): The number of nodes in that level.
             - ``edges`` (int): The number of edges in that level.

package/rust/README.md

@@ -11,7 +11,7 @@ cargo add usearch
 ```rust
 
 let options = IndexOptions {
-            dimensions: 5,
+            dimensions: 3,
             metric: MetricKind::IP,
             quantization: ScalarKind::F16,
             connectivity: 0,
@@ -36,7 +36,7 @@ assert_eq!(index.size(), 2);
 
 // Read back the tags
 let results = index.search(&first, 10).unwrap();
-assert_eq!(results.count, 2);
+assert_eq!(results.keys.len(), 2);
 ```
 
 ## Multi-Threading

package/rust/lib.cpp

@@ -47,12 +47,23 @@ size_t Index::dimensions() const { return index_->dimensions(); }
 size_t Index::connectivity() const { return index_->connectivity(); }
 size_t Index::size() const { return index_->size(); }
 size_t Index::capacity() const { return index_->capacity(); }
+size_t Index::serialized_length() const { return index_->serialized_length(); }
 
-void Index::save(rust::Str path) const { index_->save(std::string(path).c_str()); }
-void Index::load(rust::Str path) const { index_->load(std::string(path).c_str()); }
-void Index::view(rust::Str path) const { index_->view(std::string(path).c_str()); }
+void Index::save(rust::Str path) const { index_->save(output_file_t(std::string(path).c_str())).error.raise(); }
+void Index::load(rust::Str path) const { index_->load(input_file_t(std::string(path).c_str())).error.raise(); }
+void Index::view(rust::Str path) const { index_->view(memory_mapped_file_t(std::string(path).c_str())).error.raise(); }
 
-scalar_kind_t quantization(rust::Str quant) { return scalar_kind_from_name(quant.data(), quant.size()); }
+void Index::save_to_buffer(rust::Slice<uint8_t> buffer) const {
+    index_->save(memory_mapped_file_t((byte_t*)buffer.data(), buffer.size())).error.raise();
+}
+
+void Index::load_from_buffer(rust::Slice<uint8_t const> buffer) const {
+    index_->load(memory_mapped_file_t((byte_t*)buffer.data(), buffer.size())).error.raise();
+}
+
+void Index::view_from_buffer(rust::Slice<uint8_t const> buffer) const {
+    index_->view(memory_mapped_file_t((byte_t*)buffer.data(), buffer.size())).error.raise();
+}
 
 std::unique_ptr<Index> wrap(index_t&& index) {
     std::unique_ptr<index_t> punned_ptr;

package/rust/lib.hpp

@@ -32,11 +32,16 @@ class Index {
     size_t connectivity() const;
     size_t size() const;
     size_t capacity() const;
+    size_t serialized_length() const;
 
     void save(rust::Str path) const;
     void load(rust::Str path) const;
     void view(rust::Str path) const;
 
+    void save_to_buffer(rust::Slice<uint8_t> buffer) const;
+    void load_from_buffer(rust::Slice<uint8_t const> buffer) const;
+    void view_from_buffer(rust::Slice<uint8_t const> buffer) const;
+
   private:
     std::unique_ptr<index_dense_t> index_;
 };