diff --git a/cpp/src/server/RocksIdMapper.cpp b/cpp/src/server/RocksIdMapper.cpp index 2dba544243eb4bce817350b13c228ec7e7054bff..1f24f6c9da254d009cf521d6a7b9c628fd779416 100644 --- a/cpp/src/server/RocksIdMapper.cpp +++ b/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(); diff --git a/cpp/src/server/RocksIdMapper.h b/cpp/src/server/RocksIdMapper.h index 8c7315590384b2964403ac5b78381e1ee43191d3..70b77eee04b1986726af2277982fc64f5c97a5df 100644 --- a/cpp/src/server/RocksIdMapper.h +++ b/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& nid, const std::vector& sid, const std::string& group = "") override; @@ -38,7 +41,7 @@ private: private: rocksdb::DB* db_; - std::unordered_map column_handles_; + mutable std::unordered_map column_handles_; }; } diff --git a/cpp/src/server/VecIdMapper.cpp b/cpp/src/server/VecIdMapper.cpp index ecf5058a4b46f75924508b06c6f69938afb40ddb..b8bea7b348708cee53b97a47dc7eaa3a780eacc6 100644 --- a/cpp/src/server/VecIdMapper.cpp +++ b/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]; diff --git a/cpp/src/server/VecIdMapper.h b/cpp/src/server/VecIdMapper.h index 9bb6d500da10b4562826b1004adb2774d5cb3b08..f3c2bdde27d7c3f741a82250cafb9c60defc4949 100644 --- a/cpp/src/server/VecIdMapper.h +++ b/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& nid, const std::vector& 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& nid, const std::vector& sid, const std::string& group = "") override; diff --git a/cpp/src/server/VecServiceTask.cpp b/cpp/src/server/VecServiceTask.cpp index 9ebe2c1bdfd3de7e475bb163f136eac6692d2521..9c996cb3a14ea5ea390c57c1accfb6bb2f586982 100644 --- a/cpp/src/server/VecServiceTask.cpp +++ b/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 vec_f; vec_f.resize(vec_dim);