feat(db): add build_index api for ee

Former-commit-id: 75f268e4494d383f7a7b9e32c634a44f7beeba1c
pull/191/head
Xu Peng 2019-04-30 18:02:48 +08:00
parent 4cc5b8d523
commit 95adef2988
3 changed files with 48 additions and 2 deletions

View File

@ -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() {}

View File

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

View File

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