enhance: fix unittest

Signed-off-by: chyezh <[email protected]>

internal/core/src/common/protobuf_utils.cpp

@@ -1,10 +1,23 @@
+// 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");
@@ -19,3 +32,14 @@ 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

internal/core/src/common/protobuf_utils.h

@@ -42,33 +42,42 @@ RepeatedKeyValToMap(
 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:
-    static ProtoLayoutPtr
-    CreateEmpty() {
-        return std::make_unique<ProtoLayout>();
-    }
+    ProtoLayout();
 
-    ProtoLayout() : blob(nullptr), size(0) {
-    }
+    ProtoLayout(const ProtoLayout&) = delete;
 
-    ~ProtoLayout() {
-        if (blob != nullptr) {
-            delete[] static_cast<uint8_t*>(blob);
-            blob = nullptr;
-        }
-    }
+    ProtoLayout(ProtoLayout&&) = delete;
 
-    void
-    Initialize(size_t initial_size) {
-        assert(size == 0);
-        auto new_blob = new uint8_t[initial_size];
-        this->blob = new_blob;
-        this->size = initial_size;
+    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_);
     }
 
-    void* blob;
-    size_t size;
+ private:
+    void* blob_;
+    size_t size_;
 };
 
 }  //namespace milvus

internal/core/src/common/protobuf_utils_c.h

@@ -1,15 +1,33 @@
+// 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;
 
-typedef ProtoLayout* ProtoLayoutInterface;
+// 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();
 

internal/core/src/index/BitmapIndex.cpp

@@ -298,8 +298,8 @@ BitmapIndex<T>::Upload(const Config& config) {
     for (auto& file : remote_path_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        file_manager_->GetAddedTotalMemSize(), std::move(index_files));
+    return CreateIndexResult::New(file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 template <typename T>

internal/core/src/index/CreateIndexResult.cpp

@@ -1,6 +1,26 @@
+// 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)
@@ -22,9 +42,8 @@ CreateIndexResult::SerializeAt(milvus::ProtoLayout* layout) {
         serialized_info->set_file_name(info.file_name);
         serialized_info->set_file_size(info.file_size);
     }
-    layout->Initialize(result.ByteSizeLong());
-    AssertInfo(result.SerializeToArray(layout->blob, layout->size),
-               "unmarshal CreateIndexResult failed");
+    AssertInfo(layout->SerializeAndHoldProto(result),
+               "marshal CreateIndexResult failed");
 }
 
 std::vector<std::string>

internal/core/src/index/CreateIndexResult.h

@@ -1,3 +1,14 @@
+// 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>
@@ -17,15 +28,32 @@ class SerializedIndexFileInfo {
     int64_t file_size;
 };
 
+class CreateIndexResult;
+
+using CreateIndexResultPtr = std::unique_ptr<CreateIndexResult>;
+
 class CreateIndexResult {
  public:
-    CreateIndexResult(
-        int64_t mem_size,
+    // 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);
 
@@ -39,10 +67,11 @@ class CreateIndexResult {
     GetSerializedSize() const;
 
  private:
+    CreateIndexResult(
+        int64_t mem_size,
+        std::vector<SerializedIndexFileInfo>&& serialized_index_infos);
+
     int64_t mem_size_;
     std::vector<SerializedIndexFileInfo> serialized_index_infos_;
 };
-
-using CreateIndexResultPtr = std::unique_ptr<CreateIndexResult>;
-
 }  // namespace milvus::index

internal/core/src/index/InvertedIndexTantivy.cpp

@@ -167,8 +167,8 @@ InvertedIndexTantivy<T>::Upload(const Config& config) {
     for (auto& file : remote_mem_path_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        mem_file_manager_->GetAddedTotalMemSize(), std::move(index_files));
+    return CreateIndexResult::New(mem_file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 template <typename T>

internal/core/src/index/ScalarIndexSort.cpp

@@ -162,8 +162,8 @@ ScalarIndexSort<T>::Upload(const Config& config) {
     for (auto& file : remote_paths_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        file_manager_->GetAddedTotalMemSize(), std::move(index_files));
+    return CreateIndexResult::New(file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 template <typename T>

internal/core/src/index/StringIndexMarisa.cpp

@@ -185,8 +185,8 @@ StringIndexMarisa::Upload(const Config& config) {
     for (auto& file : remote_paths_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        file_manager_->GetAddedTotalMemSize(), std::move(index_files));
+    return CreateIndexResult::New(file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 void

internal/core/src/index/TextMatchIndex.cpp

@@ -113,8 +113,8 @@ TextMatchIndex::Upload(const Config& config) {
     for (auto& file : remote_mem_path_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        mem_file_manager_->GetAddedTotalMemSize(), std::move(index_files));
+    return CreateIndexResult::New(mem_file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 void

internal/core/src/index/VectorDiskIndex.cpp

@@ -129,8 +129,8 @@ VectorDiskAnnIndex<T>::Upload(const Config& config) {
     for (auto& file : remote_paths_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        file_manager_->GetAddedTotalFileSize(), std::move(index_files));
+    return CreateIndexResult::New(file_manager_->GetAddedTotalFileSize(),
+                                  std::move(index_files));
 }
 
 template <typename T>

internal/core/src/index/VectorMemIndex.cpp

@@ -103,8 +103,8 @@ VectorMemIndex<T>::Upload(const Config& config) {
     for (auto& file : remote_paths_to_size) {
         index_files.emplace_back(file.first, file.second);
     }
-    return std::make_unique<CreateIndexResult>(
-        file_manager_->GetAddedTotalMemSize(), std::move(index_files));
+    return CreateIndexResult::New(file_manager_->GetAddedTotalMemSize(),
+                                  std::move(index_files));
 }
 
 template <typename T>

internal/core/unittest/test_indexing.cpp

@@ -548,7 +548,7 @@ TEST_P(IndexTest, Mmap) {
     milvus::index::IndexBasePtr new_index;
     milvus::index::VectorIndex* vec_index = nullptr;
 
-    auto create_index_result= index->Upload();
+    auto create_index_result = index->Upload();
     index.reset();
 
     new_index = milvus::index::IndexFactory::GetInstance().CreateIndex(
@@ -704,7 +704,7 @@ TEST_P(IndexTest, GetVector_EmptySparseVector) {
 
     auto create_index_result = index->Upload();
     index.reset();
-    auto index_files = create_index_result->GetIndexFiles();    
+    auto index_files = create_index_result->GetIndexFiles();
     new_index = milvus::index::IndexFactory::GetInstance().CreateIndex(
         create_index_info, file_manager_context);
     load_conf = generate_load_conf(index_type, metric_type, 0);
@@ -962,4 +962,23 @@ TEST(Indexing, SearchDiskAnnWithBFloat16) {
     SearchResult result;
     EXPECT_NO_THROW(vec_index->Query(xq_dataset, search_info, nullptr, result));
 }
-#endif
+#endif
+
+TEST(Indexing, CreateIndexResult) {
+    using milvus::index::CreateIndexResult;
+    using milvus::index::SerializedIndexFileInfo;
+    auto result =
+        CreateIndexResult::New(16,
+                               std::vector<SerializedIndexFileInfo>{
+                                   SerializedIndexFileInfo{"file1", 100},
+                                   SerializedIndexFileInfo{"file2", 200}});
+    result->AppendSerializedIndexFileInfo(
+        SerializedIndexFileInfo{"file3", 300});
+    auto files = result->GetIndexFiles();
+    ASSERT_EQ(files.size(), 3);
+    ASSERT_EQ(files[0], "file1");
+    ASSERT_EQ(files[1], "file2");
+    ASSERT_EQ(files[2], "file3");
+    ASSERT_EQ(result->GetMemSize(), 16);
+    ASSERT_EQ(result->GetSerializedSize(), 600);
+}

internal/core/unittest/test_utils.cpp

@@ -199,3 +199,9 @@ TEST(Util, dis_closer) {
     EXPECT_FALSE(milvus::query::dis_closer(0.1, 0.2, "IP"));
     EXPECT_FALSE(milvus::query::dis_closer(0.1, 0.1, "IP"));
 }
+
+TEST(Util, ProtoLayout) {
+    milvus::ProtoLayout layout;
+    milvus::proto::cgo::CreateIndexResult result;
+    EXPECT_TRUE(layout.SerializeAndHoldProto(result));
+}

internal/datacoord/index_meta.go

@@ -45,6 +45,7 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
 	"github.com/milvus-io/milvus/pkg/util/indexparams"
 	"github.com/milvus-io/milvus/pkg/util/metricsinfo"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/timerecord"
 	"github.com/milvus-io/milvus/pkg/util/tsoutil"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
@@ -154,6 +155,9 @@ func (m *indexMeta) reloadFromKV() error {
 		return err
 	}
 	for _, segIdx := range segmentIndexes {
+		if segIdx.IndexMemSize == 0 {
+			segIdx.IndexMemSize = segIdx.IndexSerializedSize * paramtable.Get().DataCoordCfg.IndexMemSizeEstimateMultiplier.GetAsUint64()
+		}
 		m.updateSegmentIndex(segIdx)
 		metrics.FlushedSegmentFileNum.WithLabelValues(metrics.IndexFileLabel).Observe(float64(len(segIdx.IndexFileKeys)))
 	}