From 94057d9d95fd556c2b5468258bad7f68ca470031 Mon Sep 17 00:00:00 2001
From: Xu Peng <xupeng3112@163.com>
Date: Tue, 30 Apr 2019 18:32:40 +0800
Subject: [PATCH] refactor(db): refactor search with execution engine

Former-commit-id: 1625a8aff9582e61e6bb687b4da40f7721610f91
---
 cpp/src/db/DBImpl.cpp               | 12 ++++--------
 cpp/src/db/ExecutionEngine.h        |  8 ++++++++
 cpp/src/db/FaissExecutionEngine.cpp | 16 ++++++++++++++++
 cpp/src/db/FaissExecutionEngine.h   |  8 ++++++++
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/cpp/src/db/DBImpl.cpp b/cpp/src/db/DBImpl.cpp
index 47833eac74..ec5f618381 100644
--- a/cpp/src/db/DBImpl.cpp
+++ b/cpp/src/db/DBImpl.cpp
@@ -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 {
             for (auto &file : file_vec) {
-                auto index = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(file.location);
-                if (!index) {
-                    LOG(DEBUG) << "Disk io from: " << file.location;
-                    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);
+                FaissExecutionEngine index(file.dimension, file.location);
+                index.Load();
+                auto file_size = index.PhysicalSize()/(1024*1024);
                 search_set_size += file_size;
                 LOG(DEBUG) << "Search file_type " << file.file_type << " Of Size: "
                     << 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
                 memset(output_distence, 0, k * nq * sizeof(float));
                 memset(output_ids, 0, k * nq * sizeof(long));
diff --git a/cpp/src/db/ExecutionEngine.h b/cpp/src/db/ExecutionEngine.h
index 1178e19e96..974402251a 100644
--- a/cpp/src/db/ExecutionEngine.h
+++ b/cpp/src/db/ExecutionEngine.h
@@ -23,12 +23,20 @@ public:
 
     virtual size_t Size() const = 0;
 
+    virtual size_t PhysicalSize() const = 0;
+
     virtual Status Serialize() = 0;
 
     virtual Status Load() = 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 Status Cache() = 0;
diff --git a/cpp/src/db/FaissExecutionEngine.cpp b/cpp/src/db/FaissExecutionEngine.cpp
index f6fe538a48..26bc161fe2 100644
--- a/cpp/src/db/FaissExecutionEngine.cpp
+++ b/cpp/src/db/FaissExecutionEngine.cpp
@@ -39,6 +39,10 @@ size_t FaissExecutionEngine::Size() const {
     return (size_t)(Count() * pIndex_->d);
 }
 
+size_t FaissExecutionEngine::PhysicalSize() const {
+    return (size_t)(Size()*sizeof(float));
+}
+
 Status FaissExecutionEngine::Serialize() {
     write_index(pIndex_.get(), location_.c_str());
     return Status::OK();
@@ -48,6 +52,8 @@ Status FaissExecutionEngine::Load() {
     auto index  = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location_);
     if (!index) {
         index = read_index(location_);
+        Cache();
+        LOG(DEBUG) << "Disk io from: " << location_;
     }
 
     pIndex_ = index->data();
@@ -85,6 +91,16 @@ std::shared_ptr<ExecutionEngine> FaissExecutionEngine::BuildIndex(const std::str
     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() {
     zilliz::vecwise::cache::CpuCacheMgr::GetInstance(
             )->InsertItem(location_, std::make_shared<Index>(pIndex_));
diff --git a/cpp/src/db/FaissExecutionEngine.h b/cpp/src/db/FaissExecutionEngine.h
index b8403c9602..a1208b00dd 100644
--- a/cpp/src/db/FaissExecutionEngine.h
+++ b/cpp/src/db/FaissExecutionEngine.h
@@ -24,6 +24,8 @@ public:
 
     virtual size_t Size() const override;
 
+    virtual size_t PhysicalSize() const override;
+
     virtual Status Merge(const std::string& location) override;
 
     virtual Status Serialize() override;
@@ -31,6 +33,12 @@ public:
 
     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;
 
 protected: