提交 95adef29 编写于 作者: X Xu Peng

feat(db): add build_index api for ee


Former-commit-id: 75f268e4494d383f7a7b9e32c634a44f7beeba1c
上级 4cc5b8d5
#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() {}
......
......@@ -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_));
......
......@@ -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_;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册