提交 94057d9d 编写于 作者: X Xu Peng

refactor(db): refactor search with execution engine


Former-commit-id: 1625a8aff9582e61e6bb687b4da40f7721610f91
上级 473d603a
...@@ -118,17 +118,13 @@ Status DBImpl::search(const std::string& group_id, size_t k, size_t nq, ...@@ -118,17 +118,13 @@ Status DBImpl::search(const std::string& group_id, size_t k, size_t nq,
auto search_in_index = [&](meta::GroupFilesSchema& file_vec) -> void { auto search_in_index = [&](meta::GroupFilesSchema& file_vec) -> void {
for (auto &file : file_vec) { for (auto &file : file_vec) {
auto index = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(file.location); FaissExecutionEngine index(file.dimension, file.location);
if (!index) { index.Load();
LOG(DEBUG) << "Disk io from: " << file.location; auto file_size = index.PhysicalSize()/(1024*1024);
index = read_index(file.location.c_str());
zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->InsertItem(file.location, index);
}
auto file_size = index->dim * index->ntotal * 4 /(1024*1024);
search_set_size += file_size; search_set_size += file_size;
LOG(DEBUG) << "Search file_type " << file.file_type << " Of Size: " LOG(DEBUG) << "Search file_type " << file.file_type << " Of Size: "
<< file_size << " M"; << file_size << " M";
index->search(nq, vectors, k, output_distence, output_ids); index.Search(nq, vectors, k, output_distence, output_ids);
cluster(output_ids, output_distence); // cluster to each query cluster(output_ids, output_distence); // cluster to each query
memset(output_distence, 0, k * nq * sizeof(float)); memset(output_distence, 0, k * nq * sizeof(float));
memset(output_ids, 0, k * nq * sizeof(long)); memset(output_ids, 0, k * nq * sizeof(long));
......
...@@ -23,12 +23,20 @@ public: ...@@ -23,12 +23,20 @@ public:
virtual size_t Size() const = 0; virtual size_t Size() const = 0;
virtual size_t PhysicalSize() const = 0;
virtual Status Serialize() = 0; virtual Status Serialize() = 0;
virtual Status Load() = 0; virtual Status Load() = 0;
virtual Status Merge(const std::string& location) = 0; virtual Status Merge(const std::string& location) = 0;
virtual Status Search(long n,
const float *data,
long k,
float *distances,
long *labels) const = 0;
virtual std::shared_ptr<ExecutionEngine> BuildIndex(const std::string&) = 0; virtual std::shared_ptr<ExecutionEngine> BuildIndex(const std::string&) = 0;
virtual Status Cache() = 0; virtual Status Cache() = 0;
......
...@@ -39,6 +39,10 @@ size_t FaissExecutionEngine::Size() const { ...@@ -39,6 +39,10 @@ size_t FaissExecutionEngine::Size() const {
return (size_t)(Count() * pIndex_->d); return (size_t)(Count() * pIndex_->d);
} }
size_t FaissExecutionEngine::PhysicalSize() const {
return (size_t)(Size()*sizeof(float));
}
Status FaissExecutionEngine::Serialize() { Status FaissExecutionEngine::Serialize() {
write_index(pIndex_.get(), location_.c_str()); write_index(pIndex_.get(), location_.c_str());
return Status::OK(); return Status::OK();
...@@ -48,6 +52,8 @@ Status FaissExecutionEngine::Load() { ...@@ -48,6 +52,8 @@ Status FaissExecutionEngine::Load() {
auto index = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location_); auto index = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location_);
if (!index) { if (!index) {
index = read_index(location_); index = read_index(location_);
Cache();
LOG(DEBUG) << "Disk io from: " << location_;
} }
pIndex_ = index->data(); pIndex_ = index->data();
...@@ -85,6 +91,16 @@ std::shared_ptr<ExecutionEngine> FaissExecutionEngine::BuildIndex(const std::str ...@@ -85,6 +91,16 @@ std::shared_ptr<ExecutionEngine> FaissExecutionEngine::BuildIndex(const std::str
return new_ee; return new_ee;
} }
Status FaissExecutionEngine::Search(long n,
const float *data,
long k,
float *distances,
long *labels) const {
pIndex_->search(n, data, k, distances, labels);
return Status::OK();
}
Status FaissExecutionEngine::Cache() { Status FaissExecutionEngine::Cache() {
zilliz::vecwise::cache::CpuCacheMgr::GetInstance( zilliz::vecwise::cache::CpuCacheMgr::GetInstance(
)->InsertItem(location_, std::make_shared<Index>(pIndex_)); )->InsertItem(location_, std::make_shared<Index>(pIndex_));
......
...@@ -24,6 +24,8 @@ public: ...@@ -24,6 +24,8 @@ public:
virtual size_t Size() const override; virtual size_t Size() const override;
virtual size_t PhysicalSize() const override;
virtual Status Merge(const std::string& location) override; virtual Status Merge(const std::string& location) override;
virtual Status Serialize() override; virtual Status Serialize() override;
...@@ -31,6 +33,12 @@ public: ...@@ -31,6 +33,12 @@ public:
virtual Status Cache() override; virtual Status Cache() override;
virtual Status Search(long n,
const float *data,
long k,
float *distances,
long *labels) const override;
virtual std::shared_ptr<ExecutionEngine> BuildIndex(const std::string&) override; virtual std::shared_ptr<ExecutionEngine> BuildIndex(const std::string&) override;
protected: protected:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册