From 12a7747a95acf8592fceae5eee933ea3921dafa4 Mon Sep 17 00:00:00 2001
From: yah01 <yah2er0ne@outlook.com>
Date: Wed, 9 Aug 2023 10:55:15 +0800
Subject: [PATCH] Fix miss index file while slice meta contains raw data only
 (#26190)

Signed-off-by: yah01 <yah2er0ne@outlook.com>
---
 internal/core/src/index/VectorMemIndex.cpp | 28 +++++++++++++++++-----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/internal/core/src/index/VectorMemIndex.cpp b/internal/core/src/index/VectorMemIndex.cpp
index 51db032634..7fc4cf8296 100644
--- a/internal/core/src/index/VectorMemIndex.cpp
+++ b/internal/core/src/index/VectorMemIndex.cpp
@@ -20,6 +20,7 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include "fmt/format.h"
 #include "index/Meta.h"
 #include "index/Utils.h"
@@ -105,6 +106,9 @@ VectorMemIndex::Load(const Config& config) {
     AssertInfo(index_files.has_value(),
                "index file paths is empty when load index");
 
+    std::unordered_set<std::string> pending_index_files(index_files->begin(),
+                                                        index_files->end());
+
     LOG_SEGCORE_INFO_ << "load index files: " << index_files.value().size();
 
     auto parallel_degree =
@@ -113,19 +117,20 @@ VectorMemIndex::Load(const Config& config) {
 
     // try to read slice meta first
     std::string slice_meta_filepath;
-    for (auto& file : index_files.value()) {
+    for (auto& file : pending_index_files) {
         auto file_name = file.substr(file.find_last_of('/') + 1);
         if (file_name == INDEX_FILE_SLICE_META) {
             slice_meta_filepath = file;
+            pending_index_files.erase(file);
             break;
         }
     }
 
-    if (slice_meta_filepath
-            .empty()) {  // no slice meta, we could simply load all these files
-        index_datas = file_manager_->LoadIndexToMemory(index_files.value());
-        AssembleIndexDatas(index_datas);
-    } else {  // load with the slice meta info, then we can load batch by batch
+    LOG_SEGCORE_INFO_ << "load with slice meta: "
+                      << !slice_meta_filepath.empty();
+
+    if (!slice_meta_filepath
+             .empty()) {  // load with the slice meta info, then we can load batch by batch
         std::string index_file_prefix = slice_meta_filepath.substr(
             0, slice_meta_filepath.find_last_of('/') + 1);
         std::vector<std::string> batch{};
@@ -153,6 +158,9 @@ VectorMemIndex::Load(const Config& config) {
                     auto data = batch_data[file_name];
                     new_field_data->FillFieldData(data->Data(), data->Size());
                 }
+                for (auto& file : batch) {
+                    pending_index_files.erase(file);
+                }
                 batch.clear();
             };
 
@@ -174,6 +182,14 @@ VectorMemIndex::Load(const Config& config) {
         }
     }
 
+    if (!pending_index_files.empty()) {
+        auto result = file_manager_->LoadIndexToMemory(std::vector<std::string>(
+            pending_index_files.begin(), pending_index_files.end()));
+        for (auto&& index_data : result) {
+            index_datas.insert(std::move(index_data));
+        }
+    }
+
     LOG_SEGCORE_INFO_ << "construct binary set...";
     BinarySet binary_set;
     for (auto& [key, data] : index_datas) {