FaissExecutionEngine.cpp 4.3 KB
Newer Older
X
Xu Peng 已提交
1 2 3 4 5
/*******************************************************************************
 * Copyright 上海赜睿信息科技有限公司(Zilliz) - All Rights Reserved
 * Unauthorized copying of this file, via any medium is strictly prohibited.
 * Proprietary and confidential.
 ******************************************************************************/
X
Xu Peng 已提交
6 7 8
#ifndef FAISSEXECUTIONENGINE_CPP__
#define FAISSEXECUTIONENGINE_CPP__

X
Xu Peng 已提交
9 10
#include <easylogging++.h>
#include <faiss/AutoTune.h>
11 12 13
#include <faiss/MetaIndexes.h>
#include <faiss/IndexFlat.h>
#include <faiss/index_io.h>
X
Xu Peng 已提交
14
#include <wrapper/Index.h>
X
Xu Peng 已提交
15
#include <wrapper/IndexBuilder.h>
X
Xu Peng 已提交
16
#include <cache/CpuCacheMgr.h>
X
Xu Peng 已提交
17

X
Xu Peng 已提交
18
#include "FaissExecutionEngine.h"
X
Xu Peng 已提交
19 20 21 22 23

namespace zilliz {
namespace vecwise {
namespace engine {

24

X
Xu Peng 已提交
25 26
template<class IndexTrait>
FaissExecutionEngine<IndexTrait>::FaissExecutionEngine(uint16_t dimension, const std::string& location)
X
Xu Peng 已提交
27
    : pIndex_(faiss::index_factory(dimension, IndexTrait::RawIndexType)),
X
Xu Peng 已提交
28 29 30
      location_(location) {
}

X
Xu Peng 已提交
31 32
template<class IndexTrait>
FaissExecutionEngine<IndexTrait>::FaissExecutionEngine(std::shared_ptr<faiss::Index> index, const std::string& location)
X
Xu Peng 已提交
33
    : pIndex_(index),
X
Xu Peng 已提交
34
      location_(location) {
X
Xu Peng 已提交
35 36
}

X
Xu Peng 已提交
37 38
template<class IndexTrait>
Status FaissExecutionEngine<IndexTrait>::AddWithIds(long n, const float *xdata, const long *xids) {
X
Xu Peng 已提交
39
    pIndex_->add_with_ids(n, xdata, xids);
X
Xu Peng 已提交
40 41 42
    return Status::OK();
}

X
Xu Peng 已提交
43 44
template<class IndexTrait>
size_t FaissExecutionEngine<IndexTrait>::Count() const {
X
Xu Peng 已提交
45 46 47
    return (size_t)(pIndex_->ntotal);
}

X
Xu Peng 已提交
48 49
template<class IndexTrait>
size_t FaissExecutionEngine<IndexTrait>::Size() const {
X
Xu Peng 已提交
50 51 52
    return (size_t)(Count() * pIndex_->d);
}

X
Xu Peng 已提交
53 54
template<class IndexTrait>
size_t FaissExecutionEngine<IndexTrait>::PhysicalSize() const {
55 56 57
    return (size_t)(Size()*sizeof(float));
}

X
Xu Peng 已提交
58 59
template<class IndexTrait>
Status FaissExecutionEngine<IndexTrait>::Serialize() {
X
Xu Peng 已提交
60 61
    write_index(pIndex_.get(), location_.c_str());
    return Status::OK();
X
Xu Peng 已提交
62 63
}

X
Xu Peng 已提交
64 65
template<class IndexTrait>
Status FaissExecutionEngine<IndexTrait>::Load() {
X
Xu Peng 已提交
66
    auto index  = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location_);
67
    bool to_cache = false;
X
Xu Peng 已提交
68 69
    if (!index) {
        index = read_index(location_);
70
        to_cache = true;
71
        LOG(DEBUG) << "Disk io from: " << location_;
X
Xu Peng 已提交
72 73 74
    }

    pIndex_ = index->data();
75 76 77
    if (to_cache) {
        Cache();
    }
X
Xu Peng 已提交
78 79 80
    return Status::OK();
}

X
Xu Peng 已提交
81 82
template<class IndexTrait>
Status FaissExecutionEngine<IndexTrait>::Merge(const std::string& location) {
83 84 85 86 87 88 89 90 91 92 93 94 95
    if (location == location_) {
        return Status::Error("Cannot Merge Self");
    }
    auto to_merge = zilliz::vecwise::cache::CpuCacheMgr::GetInstance()->GetIndex(location);
    if (!to_merge) {
        to_merge = read_index(location);
    }
    auto file_index = dynamic_cast<faiss::IndexIDMap*>(to_merge->data().get());
    pIndex_->add_with_ids(file_index->ntotal, dynamic_cast<faiss::IndexFlat*>(file_index->index)->xb.data(),
            file_index->id_map.data());
    return Status::OK();
}

X
Xu Peng 已提交
96 97 98
template<class IndexTrait>
typename FaissExecutionEngine<IndexTrait>::Ptr
FaissExecutionEngine<IndexTrait>::BuildIndex(const std::string& location) {
X
Xu Peng 已提交
99 100
    auto opd = std::make_shared<Operand>();
    opd->d = pIndex_->d;
X
Xu Peng 已提交
101
    opd->index_type = IndexTrait::BuildIndexType;
X
Xu Peng 已提交
102 103 104 105 106 107 108 109
    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());

X
Xu Peng 已提交
110
    Ptr new_ee(new FaissExecutionEngine<IndexTrait>(index->data(), location));
X
Xu Peng 已提交
111 112 113 114
    new_ee->Serialize();
    return new_ee;
}

X
Xu Peng 已提交
115 116
template<class IndexTrait>
Status FaissExecutionEngine<IndexTrait>::Search(long n,
117 118 119 120 121 122 123 124 125
                                    const float *data,
                                    long k,
                                    float *distances,
                                    long *labels) const {

    pIndex_->search(n, data, k, distances, labels);
    return Status::OK();
}

X
Xu Peng 已提交
126 127
template<class IndexTrait>
Status FaissExecutionEngine<IndexTrait>::Cache() {
X
Xu Peng 已提交
128 129 130 131 132
    zilliz::vecwise::cache::CpuCacheMgr::GetInstance(
            )->InsertItem(location_, std::make_shared<Index>(pIndex_));

    return Status::OK();
}
X
Xu Peng 已提交
133

134

X
Xu Peng 已提交
135 136 137
} // namespace engine
} // namespace vecwise
} // namespace zilliz
X
Xu Peng 已提交
138 139

#endif