avoid add vector to invalid group

Former-commit-id: f5d04e510b5d743488be82f492aad2b95929d9e1

cpp/src/server/RocksIdMapper.cpp

@@ -18,6 +18,8 @@ namespace zilliz {
 namespace vecwise {
 namespace server {
 
+static const std::string ROCKSDB_DEFAULT_GROUP = "default";
+
 RocksIdMapper::RocksIdMapper()
 : db_(nullptr) {
     OpenDb();
@@ -90,6 +92,40 @@ void RocksIdMapper::CloseDb() {
     }
 }
 
+//not thread-safe
+ServerError RocksIdMapper::AddGroup(const std::string& group) {
+    if(!IsGroupExist(group)) {
+        if(db_ == nullptr) {
+            return SERVER_NULL_POINTER;
+        }
+
+        try {//add group
+            rocksdb::ColumnFamilyHandle *cfh = nullptr;
+            rocksdb::Status s = db_->CreateColumnFamily(rocksdb::ColumnFamilyOptions(), group, &cfh);
+            if (!s.ok()) {
+                SERVER_LOG_ERROR << "ID mapper failed to create group:" << s.ToString();
+                return SERVER_UNEXPECTED_ERROR;
+            } else {
+                column_handles_.insert(std::make_pair(group, cfh));
+            }
+        } catch(std::exception& ex) {
+            SERVER_LOG_ERROR << "ID mapper failed to create group: " << ex.what();
+            return SERVER_UNEXPECTED_ERROR;
+        }
+    }
+
+    return SERVER_SUCCESS;
+}
+
+//not thread-safe
+bool RocksIdMapper::IsGroupExist(const std::string& group) const {
+    std::string group_name = group;
+    if(group_name.empty()){
+        group_name = ROCKSDB_DEFAULT_GROUP;
+    }
+    return (column_handles_.count(group_name) > 0 && column_handles_[group_name] != nullptr);
+}
+
 ServerError RocksIdMapper::Put(const std::string& nid, const std::string& sid, const std::string& group) {
     if(db_ == nullptr) {
         return SERVER_NULL_POINTER;
@@ -104,22 +140,12 @@ ServerError RocksIdMapper::Put(const std::string& nid, const std::string& sid, c
             return SERVER_UNEXPECTED_ERROR;
         }
     } else {
-        rocksdb::ColumnFamilyHandle *cfh = nullptr;
-        if(column_handles_.count(group) == 0) {
-            try {//add group
-                rocksdb::Status s = db_->CreateColumnFamily(rocksdb::ColumnFamilyOptions(), group, &cfh);
-                if (!s.ok()) {
-                    SERVER_LOG_ERROR << "ID mapper failed to create group:" << s.ToString();
-                } else {
-                    column_handles_.insert(std::make_pair(group, cfh));
-                }
-            } catch(std::exception& ex) {
-                std::cout << ex.what() << std::endl;
-            }
-        } else {
-            cfh = column_handles_[group];
+        //try create group
+        if(AddGroup(group) != SERVER_SUCCESS){
+            return SERVER_UNEXPECTED_ERROR;
         }
 
+        rocksdb::ColumnFamilyHandle *cfh = column_handles_[group];
         rocksdb::Status s = db_->Put(rocksdb::WriteOptions(), cfh, key, value);
         if (!s.ok()) {
             SERVER_LOG_ERROR << "ID mapper failed to put:" << s.ToString();

cpp/src/server/RocksIdMapper.h

@@ -23,6 +23,9 @@ public:
     RocksIdMapper();
     ~RocksIdMapper();
 
+    ServerError AddGroup(const std::string& group) override;
+    bool IsGroupExist(const std::string& group) const override;
+
     ServerError Put(const std::string& nid, const std::string& sid, const std::string& group = "") override;
     ServerError Put(const std::vector<std::string>& nid, const std::vector<std::string>& sid, const std::string& group = "") override;
 
@@ -38,7 +41,7 @@ private:
 
 private:
     rocksdb::DB* db_;
-    std::unordered_map<std::string, rocksdb::ColumnFamilyHandle*> column_handles_;
+    mutable std::unordered_map<std::string, rocksdb::ColumnFamilyHandle*> column_handles_;
 };
 
 }

cpp/src/server/VecIdMapper.cpp

@@ -39,6 +39,17 @@ SimpleIdMapper::~SimpleIdMapper() {
 
 }
 
+ServerError SimpleIdMapper::AddGroup(const std::string& group) {
+    if(id_groups_.count(group) == 0) {
+        id_groups_.insert(std::make_pair(group, ID_MAPPING()));
+    }
+}
+
+//not thread-safe
+bool SimpleIdMapper::IsGroupExist(const std::string& group) const {
+    return id_groups_.count(group) > 0;
+}
+
 //not thread-safe
 ServerError SimpleIdMapper::Put(const std::string& nid, const std::string& sid, const std::string& group) {
     ID_MAPPING& mapping = id_groups_[group];

cpp/src/server/VecIdMapper.h

@@ -25,6 +25,9 @@ public:
 
     virtual ~IVecIdMapper(){}
 
+    virtual ServerError AddGroup(const std::string& group) = 0;
+    virtual bool IsGroupExist(const std::string& group) const = 0;
+
     virtual ServerError Put(const std::string& nid, const std::string& sid, const std::string& group = "") = 0;
     virtual ServerError Put(const std::vector<std::string>& nid, const std::vector<std::string>& sid, const std::string& group = "") = 0;
 
@@ -41,6 +44,9 @@ public:
     SimpleIdMapper();
     ~SimpleIdMapper();
 
+    ServerError AddGroup(const std::string& group) override;
+    bool IsGroupExist(const std::string& group) const override;
+
     ServerError Put(const std::string& nid, const std::string& sid, const std::string& group = "") override;
     ServerError Put(const std::vector<std::string>& nid, const std::vector<std::string>& sid, const std::string& group = "") override;
 

cpp/src/server/VecServiceTask.cpp

@@ -87,6 +87,7 @@ BaseTaskPtr AddGroupTask::Create(int32_t dimension,
 
 ServerError AddGroupTask::OnExecute() {
     try {
+        IVecIdMapper::GetInstance()->AddGroup(group_id_);
         engine::meta::GroupSchema group_info;
         group_info.dimension = (size_t)dimension_;
         group_info.group_id = group_id_;
@@ -243,6 +244,13 @@ const AttribMap& AddVectorTask::GetVecAttrib() const {
 
 ServerError AddVectorTask::OnExecute() {
     try {
+        if(!IVecIdMapper::GetInstance()->IsGroupExist(group_id_)) {
+            error_code_ = SERVER_UNEXPECTED_ERROR;
+            error_msg_ = "group not exist";
+            SERVER_LOG_ERROR << error_msg_;
+            return error_code_;
+        }
+
         uint64_t vec_dim = GetVecDimension();
         std::vector<float> vec_f;
         vec_f.resize(vec_dim);