[MMAP] Add in-memory (mmap) stream index mapping.

bindings/cpp/include/svs/runtime/flat_index.h

@@ -44,6 +44,16 @@ struct SVS_RUNTIME_API FlatIndex {
 
     virtual Status save(std::ostream& out) const noexcept = 0;
     static Status load(FlatIndex** index, std::istream& in, MetricType metric) noexcept;
+
+    // Load from a memory-mapped file.
+    // The file is expected to be in the format produced by save().
+    static Status
+    map_to_file(FlatIndex** index, const char* path, MetricType metric) noexcept;
+
+    // Load from a memory buffer.
+    // The buffer is expected to be in the format produced by save().
+    static Status
+    map_to_memory(FlatIndex** index, void* data, size_t size, MetricType metric) noexcept;
 };
 
 } // namespace v0

bindings/cpp/include/svs/runtime/vamana_index.h

@@ -95,6 +95,22 @@ struct SVS_RUNTIME_API VamanaIndex {
     static Status load(
         VamanaIndex** index, std::istream& in, MetricType metric, StorageKind storage_kind
     ) noexcept;
+
+    // Load from a memory-mapped file.
+    // The file is expected to be in the format produced by save().
+    static Status map_to_file(
+        VamanaIndex** index, const char* path, MetricType metric, StorageKind storage_kind
+    ) noexcept;
+
+    // Load from a memory buffer.
+    // The buffer is expected to be in the format produced by save().
+    static Status map_to_memory(
+        VamanaIndex** index,
+        void* data,
+        size_t size,
+        MetricType metric,
+        StorageKind storage_kind
+    ) noexcept;
 };
 
 struct SVS_RUNTIME_API VamanaIndexLeanVec : public VamanaIndex {

bindings/cpp/src/flat_index.cpp

@@ -115,5 +115,33 @@ Status FlatIndex::load(FlatIndex** index, std::istream& in, MetricType metric) n
         return Status_Ok;
     });
 }
+
+Status
+FlatIndex::map_to_file(FlatIndex** index, const char* path, MetricType metric) noexcept {
+    *index = nullptr;
+    return runtime_error_wrapper([&] {
+        std::filesystem::path fs_path(path);
+        auto is = std::make_unique<svs::io::mmstream>(fs_path);
+        std::unique_ptr<FlatIndexImpl> impl{
+            FlatIndexImpl::map_to_stream(std::move(is), metric)};
+        *index = new FlatIndexManager{std::move(impl)};
+        return Status_Ok;
+    });
+}
+
+Status FlatIndex::map_to_memory(
+    FlatIndex** index, void* data, size_t size, MetricType metric
+) noexcept {
+    *index = nullptr;
+    return runtime_error_wrapper([&] {
+        auto sp = std::span(reinterpret_cast<char*>(data), size);
+        auto is = std::make_unique<svs::io::ispanstream>(sp);
+        std::unique_ptr<FlatIndexImpl> impl{
+            FlatIndexImpl::map_to_stream(std::move(is), metric)};
+        *index = new FlatIndexManager{std::move(impl)};
+        return Status_Ok;
+    });
+}
+
 } // namespace runtime
 } // namespace svs

bindings/cpp/src/flat_index_impl.h

@@ -22,6 +22,7 @@
 
 #include <svs/core/data.h>
 #include <svs/core/distance.h>
+#include <svs/core/io/memstream.h>
 #include <svs/core/query_result.h>
 #include <svs/orchestrators/exhaustive.h>
 
@@ -111,12 +112,39 @@ class FlatIndexImpl {
         });
     }
 
+    static FlatIndexImpl*
+    map_to_stream(std::unique_ptr<std::istream>&& in, MetricType metric) {
+        if (!svs::io::is_memory_stream(*in)) {
+            throw StatusException{
+                ErrorCode::INVALID_ARGUMENT, "Provided stream is not a memory stream"};
+        }
+        auto threadpool = default_threadpool();
+        using storage_type = svs::runtime::storage::
+            StorageType_t<StorageKind::FP32, svs::io::MemoryStreamAllocator<float>>;
+
+        svs::DistanceDispatcher distance_dispatcher(to_svs_distance(metric));
+        return distance_dispatcher([&](auto&& distance) {
+            auto impl = new svs::Flat{svs::Flat::assemble<float, storage_type>(
+                *in, std::forward<decltype(distance)>(distance), std::move(threadpool)
+            )};
+
+            return new FlatIndexImpl(
+                std::unique_ptr<svs::Flat>{impl}, metric, std::move(in)
+            );
+        });
+    }
+
   protected:
     // Constructor used during loading
-    FlatIndexImpl(std::unique_ptr<svs::Flat>&& impl, MetricType metric)
+    FlatIndexImpl(
+        std::unique_ptr<svs::Flat>&& impl,
+        MetricType metric,
+        std::unique_ptr<std::istream> mapped_stream = nullptr
+    )
         : dim_{impl->dimensions()}
         , metric_type_{metric}
-        , impl_{std::move(impl)} {}
+        , impl_{std::move(impl)}
+        , mapped_stream_{std::move(mapped_stream)} {}
 
     void init_impl(data::ConstSimpleDataView<float> data) {
         auto threadpool = default_threadpool();
@@ -139,6 +167,8 @@ class FlatIndexImpl {
     size_t dim_;
     MetricType metric_type_;
     std::unique_ptr<svs::Flat> impl_;
+    // For memory-mapping, we need to keep the stream alive as long as the index is alive
+    std::unique_ptr<std::istream> mapped_stream_;
 };
 } // namespace runtime
 } // namespace svs

bindings/cpp/src/vamana_index.cpp

@@ -144,6 +144,38 @@ Status VamanaIndex::load(
     });
 }
 
+Status VamanaIndex::map_to_file(
+    VamanaIndex** index, const char* path, MetricType metric, StorageKind storage_kind
+) noexcept {
+    using Impl = VamanaIndexImpl;
+    *index = nullptr;
+    return runtime_error_wrapper([&] {
+        std::filesystem::path fs_path(path);
+        auto is = std::make_unique<svs::io::mmstream>(fs_path);
+        std::unique_ptr<Impl> impl{
+            Impl::map_to_stream(std::move(is), metric, storage_kind)};
+        *index = new VamanaIndexManagerBase<Impl>{std::move(impl)};
+    });
+}
+
+Status VamanaIndex::map_to_memory(
+    VamanaIndex** index,
+    void* data,
+    size_t size,
+    MetricType metric,
+    StorageKind storage_kind
+) noexcept {
+    using Impl = VamanaIndexImpl;
+    *index = nullptr;
+    return runtime_error_wrapper([&] {
+        auto sp = std::span(reinterpret_cast<char*>(data), size);
+        auto is = std::make_unique<svs::io::ispanstream>(sp);
+        std::unique_ptr<Impl> impl{
+            Impl::map_to_stream(std::move(is), metric, storage_kind)};
+        *index = new VamanaIndexManagerBase<Impl>{std::move(impl)};
+    });
+}
+
 #ifdef SVS_RUNTIME_HAVE_LVQ_LEANVEC
 // Specialization to build LeanVec-based Vamana index with specified leanvec dims
 Status VamanaIndexLeanVec::build(

bindings/cpp/src/vamana_index_impl.h

@@ -27,6 +27,7 @@
 #include <svs/core/data.h>
 #include <svs/core/distance.h>
 #include <svs/core/graph.h>
+#include <svs/core/io/memstream.h>
 #include <svs/core/query_result.h>
 #include <svs/extensions/vamana/scalar.h>
 #include <svs/lib/file.h>
@@ -387,14 +388,18 @@ class VamanaIndexImpl {
 
     // Constructor used during loading
     VamanaIndexImpl(
-        std::unique_ptr<svs::Vamana>&& impl, MetricType metric, StorageKind storage_kind
+        std::unique_ptr<svs::Vamana>&& impl,
+        MetricType metric,
+        StorageKind storage_kind,
+        std::unique_ptr<std::istream> mapped_stream = nullptr
     )
         : dim_{0}
         , metric_type_{metric}
         , storage_kind_{storage_kind}
         , build_params_{}
         , default_search_params_{}
-        , impl_{std::move(impl)} {
+        , impl_{std::move(impl)}
+        , mapped_stream_{std::move(mapped_stream)} {
         if (impl_) {
             dim_ = impl_->dimensions();
             const auto& buffer_config = impl_->get_search_parameters().buffer_config_;
@@ -410,11 +415,12 @@ class VamanaIndexImpl {
         }
     }
 
-    template <StorageKind Kind, typename Alloc>
+    template <StorageKind Kind, typename Alloc, typename... Args>
     static svs::Vamana* load_impl_t(
         storage::StorageType<Kind, Alloc>&& SVS_UNUSED(tag),
         std::istream& stream,
-        MetricType metric
+        MetricType metric,
+        Args&&... args
     ) {
         if constexpr (!storage::is_supported_storage_kind_v<Kind>) {
             throw StatusException(
@@ -425,7 +431,10 @@ class VamanaIndexImpl {
             auto threadpool = default_threadpool();
 
             return new svs::Vamana(svs::Vamana::assemble<float, storage_type>(
-                stream, to_svs_distance(metric), std::move(threadpool)
+                stream,
+                to_svs_distance(metric),
+                std::move(threadpool),
+                std::forward<Args>(args)...
             ));
         }
     }
@@ -447,6 +456,30 @@ class VamanaIndexImpl {
         );
     }
 
+    static VamanaIndexImpl* map_to_stream(
+        std::unique_ptr<std::istream>&& in, MetricType metric, StorageKind storage_kind
+    ) {
+        using map_allocator_type = svs::io::MemoryStreamAllocator<float>;
+        if (!svs::io::is_memory_stream(*in)) {
+            throw StatusException{
+                ErrorCode::INVALID_ARGUMENT, "Provided stream is not a memory stream"};
+        }
+        return storage::dispatch_storage_kind<map_allocator_type>(
+            storage_kind,
+            [&](auto&& tag, std::unique_ptr<std::istream>&& in, MetricType metric) {
+                using Tag = std::decay_t<decltype(tag)>;
+                auto impl = load_impl_t(
+                    std::forward<Tag>(tag), *in, metric, map_allocator_type{*in}
+                );
+                return new VamanaIndexImpl(
+                    std::unique_ptr<svs::Vamana>{impl}, metric, storage_kind, std::move(in)
+                );
+            },
+            std::move(in),
+            metric
+        );
+    }
+
     // Data members
   protected:
     size_t dim_;
@@ -455,6 +488,7 @@ class VamanaIndexImpl {
     VamanaIndex::BuildParams build_params_;
     VamanaIndex::SearchParams default_search_params_;
     std::unique_ptr<svs::Vamana> impl_;
+    std::unique_ptr<std::istream> mapped_stream_;
 };
 
 #ifdef SVS_RUNTIME_HAVE_LVQ_LEANVEC