FaissExecutionEngine.cpp 1.1 KB
Newer Older
X
Xu Peng 已提交
1 2
#include <easylogging++.h>
#include <faiss/AutoTune.h>
X
Xu Peng 已提交
3
#include <wrapper/Index.h>
X
Xu Peng 已提交
4
#include <cache/CpuCacheMgr.h>
X
Xu Peng 已提交
5

X
Xu Peng 已提交
6
#include "FaissExecutionEngine.h"
X
Xu Peng 已提交
7 8 9 10 11 12 13

namespace zilliz {
namespace vecwise {
namespace engine {

const std::string IndexType = "IDMap,Flat";

X
Xu Peng 已提交
14
FaissExecutionEngine::FaissExecutionEngine(uint16_t dimension, const std::string& location)
X
Xu Peng 已提交
15 16
    : pIndex_(faiss::index_factory(dimension, IndexType.c_str())),
      location_(location) {
X
Xu Peng 已提交
17 18
}

X
Xu Peng 已提交
19
Status FaissExecutionEngine::AddWithIds(long n, const float *xdata, const long *xids) {
X
Xu Peng 已提交
20
    pIndex_->add_with_ids(n, xdata, xids);
X
Xu Peng 已提交
21 22 23
    return Status::OK();
}

X
Xu Peng 已提交
24
size_t FaissExecutionEngine::Count() const {
X
Xu Peng 已提交
25 26 27
    return (size_t)(pIndex_->ntotal);
}

X
Xu Peng 已提交
28
size_t FaissExecutionEngine::Size() const {
X
Xu Peng 已提交
29 30 31
    return (size_t)(Count() * pIndex_->d);
}

X
Xu Peng 已提交
32
Status FaissExecutionEngine::Serialize() {
X
Xu Peng 已提交
33 34
    write_index(pIndex_.get(), location_.c_str());
    return Status::OK();
X
Xu Peng 已提交
35 36
}

X
Xu Peng 已提交
37
Status FaissExecutionEngine::Cache() {
X
Xu Peng 已提交
38 39 40 41 42
    zilliz::vecwise::cache::CpuCacheMgr::GetInstance(
            )->InsertItem(location_, std::make_shared<Index>(pIndex_));

    return Status::OK();
}
X
Xu Peng 已提交
43 44 45 46

} // namespace engine
} // namespace vecwise
} // namespace zilliz