admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-30 15:35:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

avoid add vector to invalid group

Browse Source 
Former-commit-id: f5d04e510b5d743488be82f492aad2b95929d9e1
This commit is contained in:
groot 2019-05-15 19:57:31 +08:00 committed by jinhai
parent bef742c0a9
commit 65583daaed
5 changed files with 69 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
cpp/src/server/RocksIdMapper.cpp
View File

@ -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();

5
cpp/src/server/RocksIdMapper.h
View File

@ -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_;
};
}

11
cpp/src/server/VecIdMapper.cpp
View File

@ -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];

6
cpp/src/server/VecIdMapper.h
View File

@ -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;

8
cpp/src/server/VecServiceTask.cpp
View File

@ -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);