enhance: record memory size (uncompressed) item for index (#38770)

issue: #38715 - 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 · Jan 14, 2025 · 3e788f0 · 3e788f0
1 parent 5e38f01
commit 3e788f0
Show file tree

Hide file tree

Showing 70 changed files with 2,548 additions and 1,973 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,15 @@ BitmapIndex<T>::Serialize(const Config& config) {
 }
 
 template <typename T>
-BinarySet
+IndexStatsPtr
 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;
-    for (auto& file : remote_path_to_size) {
-        ret.Append(file.first, nullptr, file.second);
-    }
-    return ret;
+    return IndexStats::NewFromSizeMap(file_manager_->GetAddedTotalMemSize(),
+                                      remote_path_to_size);
 }
 
 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
+    IndexStatsPtr
     Upload(const Config& config = {}) override;
 
     const bool

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
+IndexStatsPtr
 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
+    IndexStatsPtr
     Upload(const Config& config = {}) override;
 
  private:

diff --git a/internal/core/src/index/Index.h b/internal/core/src/index/Index.h
@@ -26,6 +26,7 @@
 #include "common/Tracer.h"
 #include "common/Types.h"
 #include "index/Meta.h"
+#include "index/IndexStats.h"
 
 namespace milvus::index {
 
@@ -57,7 +58,7 @@ class IndexBase {
     virtual int64_t
     Count() = 0;
 
-    virtual BinarySet
+    virtual IndexStatsPtr
     Upload(const Config& config = {}) = 0;
 
     virtual const bool

diff --git a/internal/core/src/index/IndexStats.cpp b/internal/core/src/index/IndexStats.cpp
@@ -0,0 +1,81 @@
+// 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/IndexStats.h"
+
+namespace milvus::index {
+
+IndexStatsPtr
+IndexStats::NewFromSizeMap(int64_t mem_size,
+                           std::map<std::string, int64_t>& index_size_map) {
+    std::vector<SerializedIndexFileInfo> serialized_index_infos;
+    serialized_index_infos.reserve(index_size_map.size());
+    for (auto& file : index_size_map) {
+        serialized_index_infos.emplace_back(file.first, file.second);
+    }
+    return IndexStats::New(mem_size, std::move(serialized_index_infos));
+}
+
+IndexStatsPtr
+IndexStats::New(int64_t mem_size,
+                std::vector<SerializedIndexFileInfo>&& serialized_index_infos) {
+    return std::unique_ptr<IndexStats>(
+        new IndexStats(mem_size, std::move(serialized_index_infos)));
+}
+
+IndexStats::IndexStats(
+    int64_t mem_size,
+    std::vector<SerializedIndexFileInfo>&& serialized_index_infos)
+    : mem_size_(mem_size), serialized_index_infos_(serialized_index_infos) {
+}
+
+void
+IndexStats::AppendSerializedIndexFileInfo(SerializedIndexFileInfo&& info) {
+    serialized_index_infos_.push_back(std::move(info));
+}
+
+void
+IndexStats::SerializeAt(milvus::ProtoLayout* layout) {
+    milvus::proto::cgo::IndexStats 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 IndexStats failed");
+}
+
+std::vector<std::string>
+IndexStats::GetIndexFiles() const {
+    std::vector<std::string> files;
+    for (auto& info : serialized_index_infos_) {
+        files.push_back(info.file_name);
+    }
+    return files;
+}
+
+int64_t
+IndexStats::GetMemSize() const {
+    return mem_size_;
+}
+
+int64_t
+IndexStats::GetSerializedSize() const {
+    int64_t size = 0;
+    for (auto& info : serialized_index_infos_) {
+        size += info.file_size;
+    }
+    return size;
+}
+
+}  // namespace milvus::index