From 95adef2988aa79e3a4ccb0d5771ae8121f51536a Mon Sep 17 00:00:00 2001
From: Xu Peng <xupeng3112@163.com>
Date: Tue, 30 Apr 2019 18:02:48 +0800
Subject: [PATCH] feat(db): add build_index api for ee

Former-commit-id: 75f268e4494d383f7a7b9e32c634a44f7beeba1c
---
 cpp/src/db/ExecutionEngine.h        |  7 ++++++
 cpp/src/db/FaissExecutionEngine.cpp | 38 +++++++++++++++++++++++++++--
 cpp/src/db/FaissExecutionEngine.h   |  5 ++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/cpp/src/db/ExecutionEngine.h b/cpp/src/db/ExecutionEngine.h
index 8750204509..1178e19e96 100644
--- a/cpp/src/db/ExecutionEngine.h
+++ b/cpp/src/db/ExecutionEngine.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <vector>
+#include <memory>
 
 #include "Status.h"
 
@@ -8,6 +9,8 @@ namespace zilliz {
 namespace vecwise {
 namespace engine {
 
+class ExecutionEngine;
+
 class ExecutionEngine {
 public:
 
@@ -22,8 +25,12 @@ public:
 
     virtual Status Serialize() = 0;
 
+    virtual Status Load() = 0;
+
     virtual Status Merge(const std::string& location) = 0;
 
+    virtual std::shared_ptr<ExecutionEngine> BuildIndex(const std::string&) = 0;
+
     virtual Status Cache() = 0;
 
     virtual ~ExecutionEngine() {}
diff --git a/cpp/src/db/FaissExecutionEngine.cpp b/cpp/src/db/FaissExecutionEngine.cpp
index 2bb0f99a6e..f6fe538a48 100644
--- a/cpp/src/db/FaissExecutionEngine.cpp
+++ b/cpp/src/db/FaissExecutionEngine.cpp
@@ -4,6 +4,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/index_io.h>
 #include <wrapper/Index.h>
+#include <wrapper/IndexBuilder.h>
 #include <cache/CpuCacheMgr.h>
 
 #include "FaissExecutionEngine.h"
@@ -12,10 +13,16 @@ namespace zilliz {
 namespace vecwise {
 namespace engine {
 
-const std::string IndexType = "IDMap,Flat";
+const std::string RawIndexType = "IDMap,Flat";
+const std::string BuildIndexType = "IDMap,Flat";
 
 FaissExecutionEngine::FaissExecutionEngine(uint16_t dimension, const std::string& location)
-    : pIndex_(faiss::index_factory(dimension, IndexType.c_str())),
+    : pIndex_(faiss::index_factory(dimension, RawIndexType.c_str())),
+      location_(location) {
+}
+
+FaissExecutionEngine::FaissExecutionEngine(std::shared_ptr<faiss::Index> index, const std::string& location)
+    : pIndex_(index),
       location_(location) {
 }
 
@@ -37,6 +44,16 @@ Status FaissExecutionEngine::Serialize() {
     return Status::OK();
 }
 
+Status FaissExecutionEngine::Load() {
+    auto index  = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location_);
+    if (!index) {
+        index = read_index(location_);
+    }
+
+    pIndex_ = index->data();
+    return Status::OK();
+}
+
 Status FaissExecutionEngine::Merge(const std::string& location) {
     if (location == location_) {
         return Status::Error("Cannot Merge Self");
@@ -51,6 +68,23 @@ Status FaissExecutionEngine::Merge(const std::string& location) {
     return Status::OK();
 }
 
+std::shared_ptr<ExecutionEngine> FaissExecutionEngine::BuildIndex(const std::string& location) {
+    auto opd = std::make_shared<Operand>();
+    opd->d = pIndex_->d;
+    opd->index_type = BuildIndexType;
+    IndexBuilderPtr pBuilder = GetIndexBuilder(opd);
+
+    auto from_index = dynamic_cast<faiss::IndexIDMap*>(pIndex_.get());
+
+    auto index = pBuilder->build_all(from_index->ntotal,
+            dynamic_cast<faiss::IndexFlat*>(from_index->index)->xb.data(),
+            from_index->id_map.data());
+
+    std::shared_ptr<ExecutionEngine> new_ee(new FaissExecutionEngine(index->data(), location));
+    new_ee->Serialize();
+    return new_ee;
+}
+
 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 a76e100e88..b8403c9602 100644
--- a/cpp/src/db/FaissExecutionEngine.h
+++ b/cpp/src/db/FaissExecutionEngine.h
@@ -16,6 +16,8 @@ namespace engine {
 class FaissExecutionEngine : public ExecutionEngine {
 public:
     FaissExecutionEngine(uint16_t dimension, const std::string& location);
+    FaissExecutionEngine(std::shared_ptr<faiss::Index> index, const std::string& location);
+
     virtual Status AddWithIds(long n, const float *xdata, const long *xids) override;
 
     virtual size_t Count() const override;
@@ -25,9 +27,12 @@ public:
     virtual Status Merge(const std::string& location) override;
 
     virtual Status Serialize() override;
+    virtual Status Load() override;
 
     virtual Status Cache() override;
 
+    virtual std::shared_ptr<ExecutionEngine> BuildIndex(const std::string&) override;
+
 protected:
     std::shared_ptr<faiss::Index> pIndex_;
     std::string location_;