enhance: record memory size (uncompressed) item for index

- Current milvus use a serialized index size(compressed) for estimate resource for loading. - Add a new field MemSize (before compressing) for index to estimate resource. Signed-off-by: chyezh <[email protected]>
milvus-io · Dec 30, 2024 · 5a89096 · 5a89096
1 parent dd81a13
commit 5a89096
Show file tree

Hide file tree

Showing 67 changed files with 1,683 additions and 1,324 deletions.
diff --git a/cmd/tools/migration/meta/210_to_220.go b/cmd/tools/migration/meta/210_to_220.go
@@ -251,21 +251,21 @@ func combineToSegmentIndexesMeta220(segmentIndexes SegmentIndexesMeta210, indexB
 			}
 
 			segmentIndexModel := &model.SegmentIndex{
-				SegmentID:      segID,
-				CollectionID:   record.GetCollectionID(),
-				PartitionID:    record.GetPartitionID(),
-				NumRows:        buildMeta.GetReq().GetNumRows(),
-				IndexID:        indexID,
-				BuildID:        record.GetBuildID(),
-				NodeID:         buildMeta.GetNodeID(),
-				IndexVersion:   buildMeta.GetIndexVersion(),
-				IndexState:     buildMeta.GetState(),
-				FailReason:     buildMeta.GetFailReason(),
-				IsDeleted:      buildMeta.GetMarkDeleted(),
-				CreatedUTCTime: record.GetCreateTime(),
-				IndexFileKeys:  fileKeys,
-				IndexSize:      buildMeta.GetSerializeSize(),
-				WriteHandoff:   buildMeta.GetState() == commonpb.IndexState_Finished,
+				SegmentID:           segID,
+				CollectionID:        record.GetCollectionID(),
+				PartitionID:         record.GetPartitionID(),
+				NumRows:             buildMeta.GetReq().GetNumRows(),
+				IndexID:             indexID,
+				BuildID:             record.GetBuildID(),
+				NodeID:              buildMeta.GetNodeID(),
+				IndexVersion:        buildMeta.GetIndexVersion(),
+				IndexState:          buildMeta.GetState(),
+				FailReason:          buildMeta.GetFailReason(),
+				IsDeleted:           buildMeta.GetMarkDeleted(),
+				CreatedUTCTime:      record.GetCreateTime(),
+				IndexFileKeys:       fileKeys,
+				IndexSerializedSize: buildMeta.GetSerializeSize(),
+				WriteHandoff:        buildMeta.GetState() == commonpb.IndexState_Finished,
 			}
 			segmentIndexModels.AddRecord(segID, indexID, segmentIndexModel)
 		}

diff --git a/internal/core/src/common/protobuf_utils.cpp b/internal/core/src/common/protobuf_utils.cpp
@@ -0,0 +1,45 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License
+
+#include "common/protobuf_utils.h"
+#include "common/protobuf_utils_c.h"
+
+// Make a static_assert to ensure that the size and alignment of the C++ and C
+static_assert(
+    sizeof(milvus::ProtoLayout) == sizeof(ProtoLayout),
+    "Size of milvus::ProtoLayout is not equal to size of ProtoLayoutInterface");
+
+// Make a static_assert to ensure that the size and alignment of the C++ and C
+static_assert(alignof(milvus::ProtoLayout) == alignof(ProtoLayout),
+              "Alignment of milvus::ProtoLayout is not equal to alignment of "
+              "ProtoLayoutInterface");
+
+ProtoLayoutInterface
+CreateProtoLayout() {
+    auto ptr = new milvus::ProtoLayout();
+    return reinterpret_cast<ProtoLayoutInterface>(ptr);
+}
+
+void
+ReleaseProtoLayout(ProtoLayoutInterface proto) {
+    delete reinterpret_cast<milvus::ProtoLayout*>(proto);
+}
+
+namespace milvus {
+ProtoLayout::ProtoLayout() : blob_(nullptr), size_(0) {
+}
+
+ProtoLayout::~ProtoLayout() {
+    if (blob_ != nullptr) {
+        delete[] static_cast<uint8_t*>(blob_);
+    }
+}
+}  // namespace milvus
diff --git a/internal/core/src/common/protobuf_utils.h b/internal/core/src/common/protobuf_utils.h
@@ -38,4 +38,46 @@ RepeatedKeyValToMap(
     }
     return mapping;
 }
+
+class ProtoLayout;
+using ProtoLayoutPtr = std::unique_ptr<ProtoLayout>;
+
+// ProtoLayout is a c++ type for esaier resource management at C-side.
+// It's always keep same memory layout with ProtoLayout at C side for cgo call.
+class ProtoLayout {
+ public:
+    ProtoLayout();
+
+    ProtoLayout(const ProtoLayout&) = delete;
+
+    ProtoLayout(ProtoLayout&&) = delete;
+
+    ProtoLayout&
+    operator=(const ProtoLayout&) = delete;
+
+    ProtoLayout&
+    operator=(ProtoLayout&&) = delete;
+
+    ~ProtoLayout();
+
+    // Serialize the proto into bytes and hold it in the layout.
+    // Return false if failure.
+    template <typename T>
+    bool
+    SerializeAndHoldProto(T& proto) {
+        if (blob_ != nullptr || size_ != 0) {
+            throw std::runtime_error(
+                "ProtoLayout should always be empty "
+                "before calling SerializeAndHoldProto");
+        }
+        size_ = proto.ByteSizeLong();
+        blob_ = new uint8_t[size_];
+        return proto.SerializeToArray(blob_, size_);
+    }
+
+ private:
+    void* blob_;
+    size_t size_;
+};
+
 }  //namespace milvus
diff --git a/internal/core/src/common/protobuf_utils_c.h b/internal/core/src/common/protobuf_utils_c.h
@@ -0,0 +1,40 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License
+
+#pragma once
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// ProtoLayout is a common ffi type for cgo call with serialized protobuf message.
+// It's always keep same memory layout with milvus::ProtoLayout at C++ side.
+typedef struct ProtoLayout {
+    void* blob;
+    size_t size;
+} ProtoLayout;
+
+// ProtoLayoutInterface is the pointer alias for ProtoLayout.
+// It should always created by CreateProtoLayout and released by ReleaseProtoLayout.
+typedef struct ProtoLayout* ProtoLayoutInterface;
+
+// CreateProtoLayout is used to create an empty ProtoLayout.
+// When you want to create a ProtoLayout at go-side, and return some data from C-side.
+// You should use this API.
+ProtoLayoutInterface
+CreateProtoLayout();
+
+void
+ReleaseProtoLayout(ProtoLayoutInterface proto);
+
+#ifdef __cplusplus
+}
+
+#endif
diff --git a/internal/core/src/index/BitmapIndex.cpp b/internal/core/src/index/BitmapIndex.cpp
@@ -286,18 +286,20 @@ BitmapIndex<T>::Serialize(const Config& config) {
 }
 
 template <typename T>
-BinarySet
+CreateIndexResultPtr
 BitmapIndex<T>::Upload(const Config& config) {
     auto binary_set = Serialize(config);
 
     file_manager_->AddFile(binary_set);
 
     auto remote_path_to_size = file_manager_->GetRemotePathsToFileSize();
-    BinarySet ret;
+    std::vector<SerializedIndexFileInfo> index_files;
+    index_files.reserve(remote_path_to_size.size());
     for (auto& file : remote_path_to_size) {
-        ret.Append(file.first, nullptr, file.second);
+        index_files.emplace_back(file.first, file.second);
     }
-    return ret;
+    return CreateIndexResult::New(file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 template <typename T>

diff --git a/internal/core/src/index/BitmapIndex.h b/internal/core/src/index/BitmapIndex.h
@@ -114,7 +114,7 @@ class BitmapIndex : public ScalarIndex<T> {
         return Count();
     }
 
-    BinarySet
+    CreateIndexResultPtr
     Upload(const Config& config = {}) override;
 
     const bool

diff --git a/internal/core/src/index/CreateIndexResult.cpp b/internal/core/src/index/CreateIndexResult.cpp
@@ -0,0 +1,72 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License
+
+#include "index/CreateIndexResult.h"
+
+namespace milvus::index {
+
+CreateIndexResultPtr
+CreateIndexResult::New(
+    int64_t mem_size,
+    std::vector<SerializedIndexFileInfo>&& serialized_index_infos) {
+    return std::unique_ptr<CreateIndexResult>(
+        new CreateIndexResult(mem_size, std::move(serialized_index_infos)));
+}
+
+CreateIndexResult::CreateIndexResult(
+    int64_t mem_size,
+    std::vector<SerializedIndexFileInfo>&& serialized_index_infos)
+    : mem_size_(mem_size), serialized_index_infos_(serialized_index_infos) {
+}
+
+void
+CreateIndexResult::AppendSerializedIndexFileInfo(
+    SerializedIndexFileInfo&& info) {
+    serialized_index_infos_.push_back(std::move(info));
+}
+
+void
+CreateIndexResult::SerializeAt(milvus::ProtoLayout* layout) {
+    milvus::proto::cgo::CreateIndexResult result;
+    result.set_mem_size(mem_size_);
+    for (auto& info : serialized_index_infos_) {
+        auto serialized_info = result.add_serialized_index_infos();
+        serialized_info->set_file_name(info.file_name);
+        serialized_info->set_file_size(info.file_size);
+    }
+    AssertInfo(layout->SerializeAndHoldProto(result),
+               "marshal CreateIndexResult failed");
+}
+
+std::vector<std::string>
+CreateIndexResult::GetIndexFiles() const {
+    std::vector<std::string> files;
+    for (auto& info : serialized_index_infos_) {
+        files.push_back(info.file_name);
+    }
+    return files;
+}
+
+int64_t
+CreateIndexResult::GetMemSize() const {
+    return mem_size_;
+}
+
+int64_t
+CreateIndexResult::GetSerializedSize() const {
+    int64_t size = 0;
+    for (auto& info : serialized_index_infos_) {
+        size += info.file_size;
+    }
+    return size;
+}
+
+}  // namespace milvus::index
diff --git a/internal/core/src/index/CreateIndexResult.h b/internal/core/src/index/CreateIndexResult.h
@@ -0,0 +1,77 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License
+
+#pragma once
+
+#include <string>
+#include <vector>
+#include "common/protobuf_utils.h"
+#include "pb/cgo_msg.pb.h"
+
+namespace milvus::index {
+
+class SerializedIndexFileInfo {
+ public:
+    SerializedIndexFileInfo(const std::string& file_name, int64_t file_size)
+        : file_name(file_name), file_size(file_size) {
+    }
+
+    std::string file_name;
+    int64_t file_size;
+};
+
+class CreateIndexResult;
+
+using CreateIndexResultPtr = std::unique_ptr<CreateIndexResult>;
+
+class CreateIndexResult {
+ public:
+    // Create a new CreateIndexResult instance.
+    static CreateIndexResultPtr
+    New(int64_t mem_size,
+        std::vector<SerializedIndexFileInfo>&& serialized_index_infos);
+
+    CreateIndexResult(const CreateIndexResult&) = delete;
+
+    CreateIndexResult(CreateIndexResult&&) = delete;
+
+    CreateIndexResult&
+    operator=(const CreateIndexResult&) = delete;
+
+    CreateIndexResult&
+    operator=(CreateIndexResult&&) = delete;
+
+    // Append a new serialized index file info into the result.
+    void
+    AppendSerializedIndexFileInfo(SerializedIndexFileInfo&& info);
+
+    // Serialize the result into the target proto layout.
+    void
+    SerializeAt(milvus::ProtoLayout* layout);
+
+    std::vector<std::string>
+    GetIndexFiles() const;
+
+    int64_t
+    GetMemSize() const;
+
+    int64_t
+    GetSerializedSize() const;
+
+ private:
+    CreateIndexResult(
+        int64_t mem_size,
+        std::vector<SerializedIndexFileInfo>&& serialized_index_infos);
+
+    int64_t mem_size_;
+    std::vector<SerializedIndexFileInfo> serialized_index_infos_;
+};
+}  // namespace milvus::index
diff --git a/internal/core/src/index/HybridScalarIndex.cpp b/internal/core/src/index/HybridScalarIndex.cpp
@@ -298,17 +298,17 @@ HybridScalarIndex<T>::SerializeIndexType() {
 }
 
 template <typename T>
-BinarySet
+CreateIndexResultPtr
 HybridScalarIndex<T>::Upload(const Config& config) {
     auto internal_index = GetInternalIndex();
     auto index_ret = internal_index->Upload(config);
 
     auto index_type_ret = SerializeIndexType();
 
     for (auto& [key, value] : index_type_ret.binary_map_) {
-        index_ret.Append(key, value);
+        index_ret->AppendSerializedIndexFileInfo(
+            SerializedIndexFileInfo(key, value->size));
     }
-
     return index_ret;
 }
 

diff --git a/internal/core/src/index/HybridScalarIndex.h b/internal/core/src/index/HybridScalarIndex.h
@@ -153,7 +153,7 @@ class HybridScalarIndex : public ScalarIndex<T> {
         return internal_index_->HasRawData();
     }
 
-    BinarySet
+    CreateIndexResultPtr
     Upload(const Config& config = {}) override;
 
  private: