From c11c58a5554bde96cae2083ed556ad3830d9977f Mon Sep 17 00:00:00 2001
From: wxyu <xy.wang@zilliz.com>
Date: Wed, 14 Aug 2019 21:24:20 +0800
Subject: [PATCH] MS-355 Add copy interface in ExcutionEngine

Former-commit-id: 9320dde8c2452982c552f604289d2661e5c1907e
---
 cpp/CHANGELOG.md                          |  1 +
 cpp/src/db/engine/ExecutionEngine.h       |  4 ++++
 cpp/src/db/engine/ExecutionEngineImpl.cpp | 26 +++++++++++++++++++++++
 cpp/src/db/engine/ExecutionEngineImpl.h   | 10 ++++++---
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/cpp/CHANGELOG.md b/cpp/CHANGELOG.md
index 27c2a4c1..6f1c49e9 100644
--- a/cpp/CHANGELOG.md
+++ b/cpp/CHANGELOG.md
@@ -15,6 +15,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-348 - Add ResourceFactory Test
 - MS-350 - Remove knowhere submodule
 - MS-354 - Add task class and interface in scheduler
+- MS-355 - Add copy interface in ExcutionEngine
 
 ## New Feature
 - MS-343 - Implement ResourceMgr
diff --git a/cpp/src/db/engine/ExecutionEngine.h b/cpp/src/db/engine/ExecutionEngine.h
index 5871be7f..a122231a 100644
--- a/cpp/src/db/engine/ExecutionEngine.h
+++ b/cpp/src/db/engine/ExecutionEngine.h
@@ -42,6 +42,10 @@ public:
 
     virtual Status Load(bool to_cache = true) = 0;
 
+    virtual Status CopyToGpu(uint64_t device_id) = 0;
+
+    virtual Status CopyToCpu() = 0;
+
     virtual Status Merge(const std::string& location) = 0;
 
     virtual Status Search(long n,
diff --git a/cpp/src/db/engine/ExecutionEngineImpl.cpp b/cpp/src/db/engine/ExecutionEngineImpl.cpp
index 21878377..037c52cf 100644
--- a/cpp/src/db/engine/ExecutionEngineImpl.cpp
+++ b/cpp/src/db/engine/ExecutionEngineImpl.cpp
@@ -143,6 +143,32 @@ Status ExecutionEngineImpl::Load(bool to_cache) {
     return Status::OK();
 }
 
+Status ExecutionEngineImpl::CopyToGpu(uint64_t device_id) {
+    try {
+        index_ = index_->CopyToGpu(device_id);
+        ENGINE_LOG_DEBUG << "CPU to GPU" << device_id;
+    } catch (knowhere::KnowhereException &e) {
+        ENGINE_LOG_ERROR << e.what();
+        return Status::Error(e.what());
+    } catch (std::exception &e) {
+        return Status::Error(e.what());
+    }
+    return Status::OK();
+}
+
+Status ExecutionEngineImpl::CopyToCpu() {
+    try {
+        index_ = index_->CopyToCpu();
+        ENGINE_LOG_DEBUG << "GPU to CPU";
+    } catch (knowhere::KnowhereException &e) {
+        ENGINE_LOG_ERROR << e.what();
+        return Status::Error(e.what());
+    } catch (std::exception &e) {
+        return Status::Error(e.what());
+    }
+    return Status::OK();
+}
+
 Status ExecutionEngineImpl::Merge(const std::string &location) {
     if (location == location_) {
         return Status::Error("Cannot Merge Self");
diff --git a/cpp/src/db/engine/ExecutionEngineImpl.h b/cpp/src/db/engine/ExecutionEngineImpl.h
index 12579d7c..cb50af23 100644
--- a/cpp/src/db/engine/ExecutionEngineImpl.h
+++ b/cpp/src/db/engine/ExecutionEngineImpl.h
@@ -18,7 +18,7 @@ namespace engine {
 
 
 class ExecutionEngineImpl : public ExecutionEngine {
- public:
+public:
 
     ExecutionEngineImpl(uint16_t dimension,
                         const std::string &location,
@@ -42,6 +42,10 @@ class ExecutionEngineImpl : public ExecutionEngine {
 
     Status Load(bool to_cache) override;
 
+    Status CopyToGpu(uint64_t device_id) override;
+
+    Status CopyToCpu() override;
+
     Status Merge(const std::string &location) override;
 
     Status Search(long n,
@@ -56,12 +60,12 @@ class ExecutionEngineImpl : public ExecutionEngine {
 
     Status Init() override;
 
- private:
+private:
     VecIndexPtr CreatetVecIndex(EngineType type);
 
     VecIndexPtr Load(const std::string &location);
 
- protected:
+protected:
     VecIndexPtr index_ = nullptr;
     EngineType build_type;
     EngineType current_type;
-- 
GitLab