[feat][index] Optimize vector index flat add and remove performance.

yuhaijun999 · Apr 16, 2024 · 50750a8 · 50750a8
1 parent 641339b
commit 50750a8
Showing 1 changed file with 40 additions and 11 deletions.
diff --git a/src/vector/vector_index_flat.cc b/src/vector/vector_index_flat.cc
@@ -17,6 +17,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <future>
+#include <iterator>
 #include <memory>
 #include <string>
 #include <thread>
@@ -83,7 +84,7 @@ butil::Status VectorIndexFlat::AddOrUpsertWrapper(const std::vector<pb::common::
 }
 
 butil::Status VectorIndexFlat::AddOrUpsert(const std::vector<pb::common::VectorWithId>& vector_with_ids,
-                                           bool is_upsert) {
+                                           bool /*is_upsert*/) {
   if (vector_with_ids.empty()) {
     return butil::Status(pb::error::EILLEGAL_PARAMTETERS, "vector_with_ids is empty");
   }
@@ -98,9 +99,22 @@ butil::Status VectorIndexFlat::AddOrUpsert(const std::vector<pb::common::VectorW
   BvarLatencyGuard bvar_guard(&g_flat_upsert_latency);
   RWLockWriteGuard guard(&rw_lock_);
 
-  if (is_upsert) {
-    faiss::IDSelectorBatch sel(vector_with_ids.size(), ids.get());
-    index_id_map2_->remove_ids(sel);
+  // delete id exists.
+  if (!index_id_map2_->rev_map.empty()) {
+    std::vector<faiss::idx_t> internal_ids(&ids[0], &ids[vector_with_ids.size()]);
+    auto iter = internal_ids.begin();
+    while (iter != internal_ids.end()) {
+      if (0 == index_id_map2_->rev_map.count(*iter)) {
+        iter = internal_ids.erase(iter);
+      } else {
+        iter++;
+      }
+    }
+
+    if (!internal_ids.empty()) {
+      faiss::IDSelectorBatch sel(internal_ids.size(), internal_ids.data());
+      index_id_map2_->remove_ids(sel);
+    }
   }
   index_id_map2_->add_with_ids(vector_with_ids.size(), vector_values.get(), ids.get());
 
@@ -121,20 +135,35 @@ butil::Status VectorIndexFlat::Delete(const std::vector<int64_t>& delete_ids) {
   }
 
   const auto& ids = VectorIndexUtils::CastVectorId(delete_ids);
-  faiss::IDSelectorBatch sel(delete_ids.size(), ids.get());
 
   {
     BvarLatencyGuard bvar_guard(&g_flat_delete_latency);
     RWLockWriteGuard guard(&rw_lock_);
 
-    auto remove_count = index_id_map2_->remove_ids(sel);
-    if (0 == remove_count) {
-      DINGO_LOG(WARNING) << fmt::format("[vector_index.flat][id({})] remove not found vector id.", Id());
-      return butil::Status(pb::error::Errno::EVECTOR_INVALID, "remove not found vector id");
+    // delete id exists.
+    if (!index_id_map2_->rev_map.empty()) {
+      std::vector<faiss::idx_t> internal_ids(&ids[0], &ids[delete_ids.size()]);
+      auto iter = internal_ids.begin();
+      while (iter != internal_ids.end()) {
+        if (0 == index_id_map2_->rev_map.count(*iter)) {
+          iter = internal_ids.erase(iter);
+        } else {
+          iter++;
+        }
+      }
+
+      if (!internal_ids.empty()) {
+        faiss::IDSelectorBatch sel(internal_ids.size(), internal_ids.data());
+        auto remove_count = index_id_map2_->remove_ids(sel);
+        if (0 == remove_count) {
+          DINGO_LOG(WARNING) << fmt::format("[vector_index.flat][id({})] remove not found vector id.", Id());
+          return butil::Status(pb::error::Errno::EVECTOR_INVALID, "remove not found vector id");
+        }
+      }
     }
-  }
 
-  return butil::Status::OK();
+    return butil::Status::OK();
+  }
 }
 
 butil::Status VectorIndexFlat::Search(const std::vector<pb::common::VectorWithId>& vector_with_ids, uint32_t topk,