mirror of https://github.com/milvus-io/milvus.git
feat(db): add build_index api for ee
Former-commit-id: 75f268e4494d383f7a7b9e32c634a44f7beeba1cpull/191/head
parent
4cc5b8d523
commit
95adef2988
|
@ -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() {}
|
||||
|
|
|
@ -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_;
|
||||
|
|
Loading…
Reference in New Issue