usearch 1.2.2 → 2.0.0

package/include/usearch/index_dense.hpp

@@ -154,10 +154,10 @@ struct index_dense_metadata_result_t {
 };
 
 /**
- *  @brief  Extracts metadata from pre-constructed index on disk,
+ *  @brief  Extracts metadata from a pre-constructed index on disk,
  *          without loading it or mapping the whole binary file.
  */
-inline index_dense_metadata_result_t index_dense_metadata(char const* file_path) noexcept {
+inline index_dense_metadata_result_t index_dense_metadata_from_path(char const* file_path) noexcept {
     index_dense_metadata_result_t result;
     std::unique_ptr<std::FILE, int (*)(std::FILE*)> file(std::fopen(file_path, "rb"), &std::fclose);
     if (!file)
@@ -219,6 +219,56 @@ inline index_dense_metadata_result_t index_dense_metadata(char const* file_path)
     return result.failed("Not a dense USearch index!");
 }
 
+/**
+ *  @brief  Extracts metadata from a pre-constructed index serialized into an in-memory buffer.
+ */
+inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_mapped_file_t file,
+                                                                      std::size_t offset = 0) noexcept {
+    index_dense_metadata_result_t result;
+
+    // Read the header
+    if (offset + sizeof(index_dense_head_buffer_t) >= file.size())
+        return result.failed("End of file reached!");
+
+    byte_t* const file_data = file.data() + offset;
+    std::size_t const file_size = file.size() - offset;
+    std::memcpy(&result.head_buffer, file_data, sizeof(index_dense_head_buffer_t));
+
+    // Check if the file immediately starts with the index, instead of vectors
+    result.config.exclude_vectors = true;
+    if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
+        return result;
+
+    // Check if it starts with 32-bit
+    std::uint32_t dimensions_u32[2]{0};
+    std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32));
+    std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32);
+
+    std::uint64_t dimensions_u64[2]{0};
+    std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64));
+    std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64);
+
+    // Check if it starts with 32-bit
+    if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) {
+        std::memcpy(&result.head_buffer, file_data + offset_if_u32, sizeof(index_dense_head_buffer_t));
+        result.config.exclude_vectors = false;
+        result.config.use_64_bit_dimensions = false;
+        if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
+            return result;
+    }
+
+    // Check if it starts with 64-bit
+    if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) {
+        std::memcpy(&result.head_buffer, file_data + offset_if_u64, sizeof(index_dense_head_buffer_t));
+        result.config.exclude_vectors = false;
+        result.config.use_64_bit_dimensions = true;
+        if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0)
+            return result;
+    }
+
+    return result.failed("Not a dense USearch index!");
+}
+
 /**
  *  @brief  Oversimplified type-punned index for equidimensional vectors
  *          with automatic @b down-casting, hardware-specific @b SIMD metrics,
@@ -454,7 +504,7 @@ class index_dense_gt {
     }
 
     static index_dense_gt make(char const* path, bool view = false) {
-        index_dense_metadata_result_t meta = index_dense_metadata(path);
+        index_dense_metadata_result_t meta = index_dense_metadata_from_path(path);
         if (!meta)
             return {};
         metric_punned_t metric(meta.head.dimensions, meta.head.kind_metric, meta.head.kind_scalar);
@@ -501,6 +551,9 @@ class index_dense_gt {
 
     stats_t stats() const { return typed_->stats(); }
     stats_t stats(std::size_t level) const { return typed_->stats(level); }
+    stats_t stats(stats_t* stats_per_level, std::size_t max_level) const {
+        return typed_->stats(stats_per_level, max_level);
+    }
 
     dynamic_allocator_t const& allocator() const { return typed_->dynamic_allocator(); }
     key_t const& free_key() const { return free_key_; }
@@ -509,7 +562,7 @@ class index_dense_gt {
      *  @brief  A relatively accurate lower bound on the amount of memory consumed by the system.
      *          In practice it's error will be below 10%.
      *
-     *  @see    `stream_length` for the length of the binary serialized representation.
+     *  @see    `serialized_length` for the length of the binary serialized representation.
      */
     std::size_t memory_usage() const {
         return                                          //
@@ -531,31 +584,31 @@ class index_dense_gt {
 
     // clang-format off
     add_result_t add(key_t key, b1x8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_b1x8); }
-    add_result_t add(key_t key, i8_bits_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_i8); }
+    add_result_t add(key_t key, i8_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_i8); }
     add_result_t add(key_t key, f16_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f16); }
     add_result_t add(key_t key, f32_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f32); }
     add_result_t add(key_t key, f64_t const* vector, std::size_t thread = any_thread(), bool force_vector_copy = true) { return add_(key, vector, thread, force_vector_copy, casts_.from_f64); }
 
     search_result_t search(b1x8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_b1x8); }
-    search_result_t search(i8_bits_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_i8); }
+    search_result_t search(i8_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_i8); }
     search_result_t search(f16_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f16); }
     search_result_t search(f32_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f32); }
     search_result_t search(f64_t const* vector, std::size_t wanted, std::size_t thread = any_thread(), bool exact = false) const { return search_(vector, wanted, thread, exact, casts_.from_f64); }
 
-    bool get(key_t key, b1x8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_b1x8); }
-    bool get(key_t key, i8_bits_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_i8); }
-    bool get(key_t key, f16_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f16); }
-    bool get(key_t key, f32_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f32); }
-    bool get(key_t key, f64_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f64); }
+    std::size_t get(key_t key, b1x8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_b1x8); }
+    std::size_t get(key_t key, i8_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_i8); }
+    std::size_t get(key_t key, f16_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f16); }
+    std::size_t get(key_t key, f32_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f32); }
+    std::size_t get(key_t key, f64_t* vector, std::size_t vectors_count = 1) const { return get_(key, vector, vectors_count, casts_.to_f64); }
 
     cluster_result_t cluster(b1x8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_b1x8); }
-    cluster_result_t cluster(i8_bits_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_i8); }
+    cluster_result_t cluster(i8_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_i8); }
     cluster_result_t cluster(f16_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f16); }
     cluster_result_t cluster(f32_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f32); }
     cluster_result_t cluster(f64_t const* vector, std::size_t level, std::size_t thread = any_thread()) const { return cluster_(vector, level, thread, casts_.from_f64); }
 
     aggregated_distances_t distance_between(key_t key, b1x8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_b1x8); }
-    aggregated_distances_t distance_between(key_t key, i8_bits_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_i8); }
+    aggregated_distances_t distance_between(key_t key, i8_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_i8); }
     aggregated_distances_t distance_between(key_t key, f16_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f16); }
     aggregated_distances_t distance_between(key_t key, f32_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f32); }
     aggregated_distances_t distance_between(key_t key, f64_t const* vector, std::size_t thread = any_thread()) const { return distance_between_(key, vector, thread, casts_.to_f64); }
@@ -710,29 +763,11 @@ class index_dense_gt {
         std::iota(available_threads_.begin(), available_threads_.end(), 0ul);
     }
 
-    /**
-     *  @brief Saves the index to a file.
-     *  @param[in] path The path to the file.
-     *  @param[in] config Configuration parameters for exports.
-     *  @return Outcome descriptor explicitly convertible to boolean.
-     */
-    serialization_result_t save(output_file_t file, serialization_config_t config = {}) const {
-        serialization_result_t result = file.open_if_not();
-        if (result)
-            stream(
-                [&](void* buffer, std::size_t length) {
-                    result = file.write(buffer, length);
-                    return !!result;
-                },
-                config);
-        return result;
-    }
-
     /**
      *  @brief  Saves serialized binary index representation to a stream.
      */
     template <typename output_callback_at, typename progress_at = dummy_progress_t>
-    serialization_result_t stream(output_callback_at&& callback, serialization_config_t config = {}) const {
+    serialization_result_t save_to_stream(output_callback_at&& output, serialization_config_t config = {}) const {
 
         serialization_result_t result;
         std::uint64_t matrix_rows = 0;
@@ -745,7 +780,7 @@ class index_dense_gt {
                 std::uint32_t dimensions[2];
                 dimensions[0] = static_cast<std::uint32_t>(typed_->size());
                 dimensions[1] = static_cast<std::uint32_t>(metric_.bytes_per_vector());
-                if (!callback(&dimensions, sizeof(dimensions)))
+                if (!output(&dimensions, sizeof(dimensions)))
                     return result.failed("Failed to serialize into stream");
                 matrix_rows = dimensions[0];
                 matrix_cols = dimensions[1];
@@ -753,7 +788,7 @@ class index_dense_gt {
                 std::uint64_t dimensions[2];
                 dimensions[0] = static_cast<std::uint64_t>(typed_->size());
                 dimensions[1] = static_cast<std::uint64_t>(metric_.bytes_per_vector());
-                if (!callback(&dimensions, sizeof(dimensions)))
+                if (!output(&dimensions, sizeof(dimensions)))
                     return result.failed("Failed to serialize into stream");
                 matrix_rows = dimensions[0];
                 matrix_cols = dimensions[1];
@@ -762,7 +797,7 @@ class index_dense_gt {
             // Dump the vectors one after another
             for (std::uint64_t i = 0; i != matrix_rows; ++i) {
                 byte_t* vector = vectors_lookup_[i];
-                if (!callback(vector, matrix_cols))
+                if (!output(vector, matrix_cols))
                     return result.failed("Failed to serialize into stream");
             }
         }
@@ -791,25 +826,25 @@ class index_dense_gt {
             head.dimensions = dimensions();
             head.multi = multi();
 
-            if (!callback(&buffer, sizeof(buffer)))
+            if (!output(&buffer, sizeof(buffer)))
                 return result.failed("Failed to serialize into stream");
         }
 
         // Save the actual proximity graph
-        return typed_->stream(std::forward<output_callback_at>(callback));
+        return typed_->save_to_stream(std::forward<output_callback_at>(output));
     }
 
     /**
      *  @brief  Estimate the binary length (in bytes) of the serialized index.
      */
-    std::size_t stream_length(serialization_config_t config = {}) const noexcept {
+    std::size_t serialized_length(serialization_config_t config = {}) const noexcept {
         std::size_t dimensions_length = 0;
         std::size_t matrix_length = 0;
         if (!config.exclude_vectors) {
             dimensions_length = config.use_64_bit_dimensions ? sizeof(std::uint64_t) * 2 : sizeof(std::uint32_t) * 2;
             matrix_length = typed_->size() * metric_.bytes_per_vector();
         }
-        return dimensions_length + matrix_length + sizeof(index_dense_head_buffer_t) + typed_->stream_length();
+        return dimensions_length + matrix_length + sizeof(index_dense_head_buffer_t) + typed_->serialized_length();
     }
 
     /**
@@ -818,12 +853,14 @@ class index_dense_gt {
      *  @param[in] config Configuration parameters for imports.
      *  @return Outcome descriptor explicitly convertible to boolean.
      */
-    serialization_result_t load(input_file_t file, serialization_config_t config = {}) {
+    template <typename input_callback_at>
+    serialization_result_t load_from_stream(input_callback_at&& input, serialization_config_t config = {}) {
 
-        serialization_result_t result = file.open_if_not();
-        if (!result)
-            return result;
+        // Discard all previous memory allocations of `vectors_tape_allocator_`
+        reset();
 
+        // Infer the new index size
+        serialization_result_t result;
         std::uint64_t matrix_rows = 0;
         std::uint64_t matrix_cols = 0;
 
@@ -832,16 +869,14 @@ class index_dense_gt {
             // Save the matrix size
             if (!config.use_64_bit_dimensions) {
                 std::uint32_t dimensions[2];
-                result = file.read(&dimensions, sizeof(dimensions));
-                if (!result)
-                    return result;
+                if (!input(&dimensions, sizeof(dimensions)))
+                    return result.failed("Failed to read 32-bit dimensions of the matrix");
                 matrix_rows = dimensions[0];
                 matrix_cols = dimensions[1];
             } else {
                 std::uint64_t dimensions[2];
-                result = file.read(&dimensions, sizeof(dimensions));
-                if (!result)
-                    return result;
+                if (!input(&dimensions, sizeof(dimensions)))
+                    return result.failed("Failed to read 64-bit dimensions of the matrix");
                 matrix_rows = dimensions[0];
                 matrix_cols = dimensions[1];
             }
@@ -849,9 +884,8 @@ class index_dense_gt {
             vectors_lookup_.resize(matrix_rows);
             for (std::uint64_t slot = 0; slot != matrix_rows; ++slot) {
                 byte_t* vector = vectors_tape_allocator_.allocate(matrix_cols);
-                result = file.read(vector, matrix_cols);
-                if (!result)
-                    return result;
+                if (!input(vector, matrix_cols))
+                    return result.failed("Failed to read vectors");
                 vectors_lookup_[slot] = vector;
             }
         }
@@ -859,9 +893,8 @@ class index_dense_gt {
         // Load metadata and choose the right metric
         {
             index_dense_head_buffer_t buffer;
-            result = file.read(buffer, sizeof(buffer));
-            if (!result)
-                return result;
+            if (!input(buffer, sizeof(buffer)))
+                return result.failed("Failed to read the index ");
 
             index_dense_head_t head{buffer};
             if (std::memcmp(buffer, default_magic(), std::strlen(default_magic())) != 0)
@@ -882,7 +915,7 @@ class index_dense_gt {
         }
 
         // Pull the actual proximity graph
-        result = typed_->load(std::move(file));
+        result = typed_->load_from_stream(std::forward<input_callback_at>(input));
         if (!result)
             return result;
         if (typed_->size() != static_cast<std::size_t>(matrix_rows))
@@ -900,10 +933,14 @@ class index_dense_gt {
      */
     serialization_result_t view(memory_mapped_file_t file, std::size_t offset = 0, serialization_config_t config = {}) {
 
+        // Discard all previous memory allocations of `vectors_tape_allocator_`
+        reset();
+
         serialization_result_t result = file.open_if_not();
         if (!result)
             return result;
 
+        // Infer the new index size
         std::uint64_t matrix_rows = 0;
         std::uint64_t matrix_cols = 0;
         span_punned_t vectors_buffer;
@@ -976,6 +1013,109 @@ class index_dense_gt {
         return result;
     }
 
+    /**
+     *  @brief Saves the index to a file.
+     *  @param[in] path The path to the file.
+     *  @param[in] config Configuration parameters for exports.
+     *  @return Outcome descriptor explicitly convertible to boolean.
+     */
+    serialization_result_t save(output_file_t file, serialization_config_t config = {}) const {
+
+        serialization_result_t io_result = file.open_if_not();
+        if (!io_result)
+            return io_result;
+
+        serialization_result_t stream_result = save_to_stream(
+            [&](void const* buffer, std::size_t length) {
+                io_result = file.write(buffer, length);
+                return !!io_result;
+            },
+            config);
+
+        if (!stream_result)
+            return stream_result;
+        return io_result;
+    }
+
+    /**
+     *  @brief  Memory-maps the serialized binary index representation from disk,
+     *          @b without copying data into RAM, and fetching it on-demand.
+     */
+    serialization_result_t save(memory_mapped_file_t file, std::size_t offset = 0,
+                                serialization_config_t config = {}) const {
+
+        serialization_result_t io_result = file.open_if_not();
+        if (!io_result)
+            return io_result;
+
+        serialization_result_t stream_result = save_to_stream(
+            [&](void const* buffer, std::size_t length) {
+                if (offset + length > file.size())
+                    return false;
+                std::memcpy(file.data() + offset, buffer, length);
+                offset += length;
+                return true;
+            },
+            config);
+
+        return stream_result;
+    }
+
+    /**
+     *  @brief Parses the index from file to RAM.
+     *  @param[in] path The path to the file.
+     *  @param[in] config Configuration parameters for imports.
+     *  @return Outcome descriptor explicitly convertible to boolean.
+     */
+    serialization_result_t load(input_file_t file, serialization_config_t config = {}) {
+
+        serialization_result_t io_result = file.open_if_not();
+        if (!io_result)
+            return io_result;
+
+        serialization_result_t stream_result = load_from_stream(
+            [&](void* buffer, std::size_t length) {
+                io_result = file.read(buffer, length);
+                return !!io_result;
+            },
+            config);
+
+        if (!stream_result)
+            return stream_result;
+        return io_result;
+    }
+
+    /**
+     *  @brief  Memory-maps the serialized binary index representation from disk,
+     *          @b without copying data into RAM, and fetching it on-demand.
+     */
+    serialization_result_t load(memory_mapped_file_t file, std::size_t offset = 0, serialization_config_t config = {}) {
+
+        serialization_result_t io_result = file.open_if_not();
+        if (!io_result)
+            return io_result;
+
+        serialization_result_t stream_result = load_from_stream(
+            [&](void* buffer, std::size_t length) {
+                if (offset + length > file.size())
+                    return false;
+                std::memcpy(buffer, file.data() + offset, length);
+                offset += length;
+                return true;
+            },
+            config);
+
+        return stream_result;
+    }
+
+    serialization_result_t save(char const* file_path, serialization_config_t config = {}) const {
+        return save(output_file_t(file_path), config);
+    }
+
+    serialization_result_t load(char const* file_path, serialization_config_t config = {}) {
+        return load(input_file_t(file_path), config);
+    }
+
     /**
      *  @brief Checks if a vector with specified key is present.
      *  @return `true` if the key is present in the index, `false` otherwise.
@@ -1489,6 +1629,8 @@ class index_dense_gt {
 
         result.computed_distances = computed_distances;
         result.visited_members = visited_members;
+
+        (void)progress;
         return result;
     }
 
@@ -1663,11 +1805,6 @@ class index_dense_gt {
         return result;
     }
 
-    compressed_slot_t lookup_id_(key_t key) const {
-        shared_lock_t lock(slot_lookup_mutex_);
-        return slot_lookup_.at(key);
-    }
-
     void reindex_keys_() {
 
         // Estimate number of entries first
@@ -1735,13 +1872,13 @@ class index_dense_gt {
         casts_t result;
 
         result.from_b1x8 = cast_gt<b1x8_t, to_scalar_at>{};
-        result.from_i8 = cast_gt<i8_bits_t, to_scalar_at>{};
+        result.from_i8 = cast_gt<i8_t, to_scalar_at>{};
         result.from_f16 = cast_gt<f16_t, to_scalar_at>{};
         result.from_f32 = cast_gt<f32_t, to_scalar_at>{};
         result.from_f64 = cast_gt<f64_t, to_scalar_at>{};
 
         result.to_b1x8 = cast_gt<to_scalar_at, b1x8_t>{};
-        result.to_i8 = cast_gt<to_scalar_at, i8_bits_t>{};
+        result.to_i8 = cast_gt<to_scalar_at, i8_t>{};
         result.to_f16 = cast_gt<to_scalar_at, f16_t>{};
         result.to_f32 = cast_gt<to_scalar_at, f32_t>{};
         result.to_f64 = cast_gt<to_scalar_at, f64_t>{};
@@ -1754,7 +1891,7 @@ class index_dense_gt {
         case scalar_kind_t::f64_k: return make_casts_<f64_t>();
         case scalar_kind_t::f32_k: return make_casts_<f32_t>();
         case scalar_kind_t::f16_k: return make_casts_<f16_t>();
-        case scalar_kind_t::i8_k: return make_casts_<i8_bits_t>();
+        case scalar_kind_t::i8_k: return make_casts_<i8_t>();
         case scalar_kind_t::b1x8_k: return make_casts_<b1x8_t>();
         default: return {};
         }