Mmap4

python/test_pysdk/test_knn.py

@@ -1069,7 +1069,7 @@ def test_sparse_knn_with_index(self, check_data, suffix):
                                                index.IndexType.BMP,
                                                {"block_size": "8", "compress_type": "compress"}), ConflictType.Error)
 
-        # table_obj.optimize("idx1", {"topk": "3"})
+        table_obj.optimize("idx1", {"topk": "3"})
 
         res, extra_result = (table_obj
                              .output(["*", "_row_id", "_similarity"])
@@ -1592,7 +1592,7 @@ def test_sparse_knn_with_invalid_alpha_beta(self, check_data, alpha, beta, suffi
                                                index.IndexType.BMP,
                                                {"block_size": "8", "compress_type": "compress"}), ConflictType.Error)
 
-        # table_obj.optimize("idx1", {"topk": "3"})
+        table_obj.optimize("idx1", {"topk": "3"})
 
         with pytest.raises(InfinityException) as e:
             res, extra_result = (table_obj
@@ -1910,7 +1910,7 @@ def test_match_sparse_index_hint(self, check_data, suffix):
                                                index.IndexType.BMP,
                                                {"block_size": "8", "compress_type": "raww"}), ConflictType.Error)
 
-        # table_obj.optimize("idx1", {"topk": "3"})
+        table_obj.optimize("idx1", {"topk": "3"})
 
         res, extra_result = (table_obj
                              .output(["*", "_row_id", "_similarity"])

src/parser/type/serialize.h

@@ -57,6 +57,16 @@ inline std::string ReadBuf<std::string>(const char *buf) {
     return str;
 }
 
+template <>
+inline std::tuple<> ReadBuf<std::tuple<>>(const char *buf) {
+    return {};
+}
+
+template <>
+inline std::tuple<> ReadBufAdv<std::tuple<>>(const char *&buf) {
+    return {};
+}
+
 template <>
 inline std::string ReadBufAdv<std::string>(const char *&buf) {
     int32_t size = ReadBufAdv<int32_t>(buf);
@@ -101,6 +111,12 @@ inline void WriteBufAdv<std::string>(char *&buf, const std::string &value) {
     buf += len;
 }
 
+template <>
+inline void WriteBuf<std::tuple<>>(char *const buf, const std::tuple<> &) {}
+
+template <>
+inline void WriteBufAdv<std::tuple<>>(char *&buf, const std::tuple<> &) {}
+
 template <typename T>
 inline void WriteBufVecAdv(char *&buf, const T *data, size_t size) {
     static_assert(std::is_standard_layout_v<T>, "T must be POD");

src/storage/buffer/buffer_handle.cpp

@@ -64,14 +64,14 @@ BufferHandle::~BufferHandle() {
 const void *BufferHandle::GetData() const { return data_; }
 
 void *BufferHandle::GetDataMut() {
-    buffer_obj_->GetMutPointer();
+    data_ = buffer_obj_->GetMutPointer();
     return data_;
 }
 
 const FileWorker *BufferHandle::GetFileWorker() const { return buffer_obj_->file_worker(); }
 
 FileWorker *BufferHandle::GetFileWorkerMut() {
-    buffer_obj_->GetMutPointer();
+    data_ = buffer_obj_->GetMutPointer();
     return buffer_obj_->file_worker();
 }
 

src/storage/buffer/buffer_obj.cpp

@@ -72,6 +72,7 @@ BufferHandle BufferObj::Load() {
     std::unique_lock<std::mutex> locker(w_locker_);
     if (type_ == BufferType::kMmap) {
         switch (status_) {
+            case BufferStatus::kUnloaded:
             case BufferStatus::kLoaded: {
                 break;
             }
@@ -315,12 +316,20 @@ void BufferObj::LoadInner() {
     ++rc_;
 }
 
-void BufferObj::GetMutPointer() {
+void *BufferObj::GetMutPointer() {
     std::unique_lock<std::mutex> locker(w_locker_);
     if (type_ == BufferType::kTemp) {
         buffer_mgr_->RemoveTemp(this);
+    } else if (type_ == BufferType::kMmap) {
+        bool free_success = buffer_mgr_->RequestSpace(GetBufferSize());
+        if (!free_success) {
+            String error_message = "Out of memory.";
+            UnrecoverableError(error_message);
+        }
+        file_worker_->ReadFromFile(false);
     }
     type_ = BufferType::kEphemeral;
+    return file_worker_->GetData();
 }
 
 void BufferObj::UnloadInner() {
@@ -338,7 +347,7 @@ void BufferObj::UnloadInner() {
             type_ = BufferType::kMmap;
         } else if (type_ == BufferType::kMmap) {
             file_worker_->MmapNotNeed();
-            status_ = BufferStatus::kFreed;
+            status_ = BufferStatus::kUnloaded;
         } else {
             buffer_mgr_->PushGCQueue(this);
             status_ = BufferStatus::kUnloaded;

src/storage/buffer/buffer_obj.cppm

@@ -125,7 +125,7 @@ private:
     void LoadInner();
 
     // called when BufferHandle needs mutable pointer.
-    void GetMutPointer();
+    void *GetMutPointer();
 
     // called when BufferHandle destructs, to decrease rc_ by 1.
     void UnloadInner();

src/storage/buffer/file_worker/bmp_index_file_worker.cpp

@@ -121,7 +121,7 @@ bool BMPIndexFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, c
     return true;
 }
 
-void BMPIndexFileWorker::ReadFromFileImpl(SizeT file_size) {
+void BMPIndexFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     if (data_ != nullptr) {
         UnrecoverableError("Data is already allocated.");
     }
@@ -135,7 +135,11 @@ void BMPIndexFileWorker::ReadFromFileImpl(SizeT file_size) {
             } else {
                 using IndexT = std::decay_t<decltype(*index)>;
                 if constexpr (IndexT::kOwnMem) {
-                    index = new IndexT(IndexT::Load(*file_handle_));
+                    if (from_spill) {
+                        index = new IndexT(IndexT::Load(*file_handle_));
+                    } else {
+                        index = new IndexT(IndexT::LoadFromPtr(*file_handle_, file_size));
+                    }
                 } else {
                     UnrecoverableError("BMPIndexFileWorker::ReadFromFileImpl: index does not own memory");
                 }

src/storage/buffer/file_worker/bmp_index_file_worker.cppm

@@ -52,7 +52,7 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
 
     bool ReadFromMmapImpl(const void *ptr, SizeT size) override;
 

src/storage/buffer/file_worker/data_file_worker.cpp

@@ -98,7 +98,7 @@ bool DataFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const
     return true;
 }
 
-void DataFileWorker::ReadFromFileImpl(SizeT file_size) {
+void DataFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
 
     if (file_size < sizeof(u64) * 3) {
         Status status = Status::DataIOError(fmt::format("Incorrect file length {}.", file_size));
@@ -116,7 +116,7 @@ void DataFileWorker::ReadFromFileImpl(SizeT file_size) {
         RecoverableError(status);
     }
     if (magic_number != 0x00dd3344) {
-        Status status = Status::DataIOError(fmt::format("Read magic number which length isn't {}.", nbytes1));
+        Status status = Status::DataIOError(fmt::format("Read magic error, {} != 0x00dd3344.", magic_number));
         RecoverableError(status);
     }
 
@@ -153,7 +153,7 @@ bool DataFileWorker::ReadFromMmapImpl(const void *p, SizeT file_size) {
     const char *ptr = static_cast<const char *>(p);
     u64 magic_number = ReadBufAdv<u64>(ptr);
     if (magic_number != 0x00dd3344) {
-        Status status = Status::DataIOError(fmt::format("Read magic number which length isn't {}.", magic_number));
+        Status status = Status::DataIOError(fmt::format("Read magic error: {} != 0x00dd3344.", magic_number));
         RecoverableError(status);
     }
     u64 buffer_size = ReadBufAdv<u64>(ptr);

src/storage/buffer/file_worker/data_file_worker.cppm

@@ -46,7 +46,7 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
 
     bool ReadFromMmapImpl(const void *ptr, SizeT size) override;
 

src/storage/buffer/file_worker/emvb_index_file_worker.cpp

@@ -92,7 +92,7 @@ bool EMVBIndexFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success,
     return true;
 }
 
-void EMVBIndexFileWorker::ReadFromFileImpl(SizeT file_size) {
+void EMVBIndexFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     if (data_) {
         const auto error_message = "Data is already allocated.";
         UnrecoverableError(error_message);

src/storage/buffer/file_worker/emvb_index_file_worker.cppm

@@ -59,7 +59,7 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
 
 private:
     const EmbeddingInfo *GetEmbeddingInfo() const;

src/storage/buffer/file_worker/file_worker.cpp

@@ -123,7 +123,7 @@ void FileWorker::ReadFromFile(bool from_spill) {
     }
     file_handle_ = std::move(file_handle);
     DeferFn defer_fn2([&]() { file_handle_ = nullptr; });
-    ReadFromFileImpl(file_size);
+    ReadFromFileImpl(file_size, from_spill);
 }
 
 void FileWorker::MoveFile() {
@@ -210,8 +210,8 @@ void FileWorker::CleanupTempFile() const {
 }
 
 void FileWorker::Mmap() {
-    if (mmap_addr_ != nullptr) {
-        return;
+    if (mmap_addr_ != nullptr || mmap_data_ != nullptr) {
+        this->Munmap();
     }
     auto [defer_fn, read_path] = GetFilePathInner(false);
     bool use_object_cache = persistence_manager_ != nullptr;

src/storage/buffer/file_worker/file_worker.cppm

@@ -77,9 +77,8 @@ public:
 protected:
     virtual bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx = {}) = 0;
 
-    virtual void ReadFromFileImpl(SizeT file_size) = 0;
+    virtual void ReadFromFileImpl(SizeT file_size, bool from_spill) = 0;
 
-private:
     String ChooseFileDir(bool spill) const;
 
     Pair<Optional<DeferFn<std::function<void()>>>, String> GetFilePathInner(bool spill);

src/storage/buffer/file_worker/hnsw_file_worker.cpp

@@ -113,15 +113,20 @@ bool HnswFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const
             if constexpr (std::is_same_v<T, std::nullptr_t>) {
                 UnrecoverableError("Invalid index type.");
             } else {
-                index->Save(*file_handle_);
+                using IndexT = std::decay_t<decltype(*index)>;
+                if constexpr (IndexT::kOwnMem) {
+                    index->SaveToPtr(*file_handle_);
+                } else {
+                    UnrecoverableError("Invalid index type.");
+                }
             }
         },
         *hnsw_index);
     prepare_success = true;
     return true;
 }
 
-void HnswFileWorker::ReadFromFileImpl(SizeT file_size) {
+void HnswFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     if (data_ != nullptr) {
         UnrecoverableError("Data is already allocated.");
     }
@@ -134,10 +139,60 @@ void HnswFileWorker::ReadFromFileImpl(SizeT file_size) {
                 UnrecoverableError("Invalid index type.");
             } else {
                 using IndexT = std::decay_t<decltype(*index)>;
-                index = IndexT::Load(*file_handle_).release();
+                if constexpr (IndexT::kOwnMem) {
+                    if (from_spill) {
+                        index = IndexT::Load(*file_handle_).release();
+                    } else {
+                        index = IndexT::LoadFromPtr(*file_handle_, file_size).release();
+                    }
+                } else {
+                    UnrecoverableError("Invalid index type.");
+                }
+            }
+        },
+        *hnsw_index);
+}
+
+bool HnswFileWorker::ReadFromMmapImpl(const void *ptr, SizeT size) {
+    if (mmap_data_ != nullptr) {
+        UnrecoverableError("Mmap data is already allocated.");
+    }
+    mmap_data_ = reinterpret_cast<u8 *>(new AbstractHnsw(HnswIndexInMem::InitAbstractIndex(index_base_.get(), column_def_.get(), false)));
+    auto *hnsw_index = reinterpret_cast<AbstractHnsw *>(mmap_data_);
+    std::visit(
+        [&](auto &&index) {
+            using T = std::decay_t<decltype(index)>;
+            if constexpr (std::is_same_v<T, std::nullptr_t>) {
+                UnrecoverableError("Invalid index type.");
+            } else {
+                using IndexT = std::decay_t<decltype(*index)>;
+                if constexpr (!IndexT::kOwnMem) {
+                    const auto *p = static_cast<const char *>(ptr);
+                    index = IndexT::LoadFromPtr(p, size).release();
+                } else {
+                    UnrecoverableError("Invalid index type.");
+                }
+            }
+        },
+        *hnsw_index);
+    return true;
+}
+
+void HnswFileWorker::FreeFromMmapImpl() {
+    if (mmap_data_ == nullptr) {
+        UnrecoverableError("Mmap data is not allocated.");
+    }
+    auto *hnsw_index = reinterpret_cast<AbstractHnsw *>(mmap_data_);
+    std::visit(
+        [&](auto &&index) {
+            using T = std::decay_t<decltype(index)>;
+            if constexpr (!std::is_same_v<T, std::nullptr_t>) {
+                delete index;
             }
         },
         *hnsw_index);
+    delete hnsw_index;
+    mmap_data_ = nullptr;
 }
 
 } // namespace infinity

src/storage/buffer/file_worker/hnsw_file_worker.cppm

@@ -54,7 +54,11 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
+
+    bool ReadFromMmapImpl(const void *ptr, SizeT size) override;
+
+    void FreeFromMmapImpl() override;
 
 private:
     SizeT index_size_{};

src/storage/buffer/file_worker/ivf_index_file_worker.cpp

@@ -65,7 +65,7 @@ bool IVFIndexFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, c
     return true;
 }
 
-void IVFIndexFileWorker::ReadFromFileImpl(SizeT file_size) {
+void IVFIndexFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     if (!data_) [[likely]] {
         auto index = IVFIndexInChunk::GetNewIVFIndexInChunk(index_base_.get(), column_def_.get());
         index->ReadIndexInner(*file_handle_);

src/storage/buffer/file_worker/ivf_index_file_worker.cppm

@@ -54,7 +54,7 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
 };
 
 } // namespace infinity

src/storage/buffer/file_worker/raw_file_worker.cpp

@@ -74,7 +74,7 @@ bool RawFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const
     return true;
 }
 
-void RawFileWorker::ReadFromFileImpl(SizeT file_size) {
+void RawFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     buffer_size_ = file_handle_->FileSize();
     data_ = static_cast<void *>(new char[buffer_size_]);
     auto [nbytes, status1] = file_handle_->Read(data_, buffer_size_);

src/storage/buffer/file_worker/raw_file_worker.cppm

@@ -50,7 +50,7 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
 
 private:
     SizeT buffer_size_;

src/storage/buffer/file_worker/secondary_index_file_worker.cpp

@@ -71,7 +71,7 @@ bool SecondaryIndexFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_succ
     return true;
 }
 
-void SecondaryIndexFileWorker::ReadFromFileImpl(SizeT file_size) {
+void SecondaryIndexFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     if (!data_) [[likely]] {
         auto index = GetSecondaryIndexData(column_def_->type(), row_count_, false);
         index->ReadIndexInner(*file_handle_);

src/storage/buffer/file_worker/secondary_index_file_worker.cppm

@@ -60,7 +60,7 @@ public:
 protected:
     bool WriteToFileImpl(bool to_spill, bool &prepare_success, const FileWorkerSaveCtx &ctx) override;
 
-    void ReadFromFileImpl(SizeT file_size) override;
+    void ReadFromFileImpl(SizeT file_size, bool from_spill) override;
 
     const u32 row_count_{};
 };

src/storage/buffer/file_worker/var_file_worker.cpp

@@ -92,7 +92,7 @@ bool VarFileWorker::WriteToFileImpl(bool to_spill, bool &prepare_success, const
     return true;
 }
 
-void VarFileWorker::ReadFromFileImpl(SizeT file_size) {
+void VarFileWorker::ReadFromFileImpl(SizeT file_size, bool from_spill) {
     if (data_ != nullptr) {
         String error_message = "Data is not allocated.";
         UnrecoverableError(error_message);