refactor(db): using index factory in MemManager

Former-commit-id: 808f882176b0f933c609475bbd3930b3402525a8
pull/191/head
Xu Peng 2019-04-17 10:32:57 +08:00
parent 45a74e1560
commit ada3903352
2 changed files with 13 additions and 19 deletions

View File

@ -1,5 +1,6 @@
#include <faiss/IndexFlat.h> /* #include <faiss/IndexFlat.h> */
#include <faiss/MetaIndexes.h> /* #include <faiss/MetaIndexes.h> */
#include <faiss/AutoTune.h>
#include <faiss/index_io.h> #include <faiss/index_io.h>
#include <iostream> #include <iostream>
#include <sstream> #include <sstream>
@ -19,20 +20,19 @@ MemVectors::MemVectors(const std::string& group_id,
_file_location(file_location), _file_location(file_location),
_pIdGenerator(new SimpleIDGenerator()), _pIdGenerator(new SimpleIDGenerator()),
_dimension(dimension), _dimension(dimension),
_pInnerIndex(new faiss::IndexFlat(_dimension)), pIndex_(faiss::index_factory(_dimension, "IDMap,Flat")) {
_pIdMapIndex(new faiss::IndexIDMap(_pInnerIndex)) {
} }
void MemVectors::add(size_t n_, const float* vectors_, IDNumbers& vector_ids_) { void MemVectors::add(size_t n_, const float* vectors_, IDNumbers& vector_ids_) {
_pIdGenerator->getNextIDNumbers(n_, vector_ids_); _pIdGenerator->getNextIDNumbers(n_, vector_ids_);
_pIdMapIndex->add_with_ids(n_, vectors_, &vector_ids_[0]); pIndex_->add_with_ids(n_, vectors_, &vector_ids_[0]);
for(auto i=0 ; i<n_; i++) { for(auto i=0 ; i<n_; i++) {
vector_ids_.push_back(i); vector_ids_.push_back(i);
} }
} }
size_t MemVectors::total() const { size_t MemVectors::total() const {
return _pIdMapIndex->ntotal; return pIndex_->ntotal;
} }
size_t MemVectors::approximate_size() const { size_t MemVectors::approximate_size() const {
@ -42,10 +42,10 @@ size_t MemVectors::approximate_size() const {
Status MemVectors::serialize(std::string& group_id) { Status MemVectors::serialize(std::string& group_id) {
/* std::stringstream ss; */ /* std::stringstream ss; */
/* ss << "/tmp/test/" << _pIdGenerator->getNextIDNumber(); */ /* ss << "/tmp/test/" << _pIdGenerator->getNextIDNumber(); */
/* faiss::write_index(_pIdMapIndex, ss.str().c_str()); */ /* faiss::write_index(pIndex_, ss.str().c_str()); */
/* std::cout << _pIdMapIndex->ntotal << std::endl; */ /* std::cout << pIndex_->ntotal << std::endl; */
/* std::cout << _file_location << std::endl; */ /* std::cout << _file_location << std::endl; */
faiss::write_index(_pIdMapIndex, _file_location.c_str()); faiss::write_index(pIndex_, _file_location.c_str());
group_id = group_id_; group_id = group_id_;
return Status::OK(); return Status::OK();
} }
@ -55,13 +55,9 @@ MemVectors::~MemVectors() {
delete _pIdGenerator; delete _pIdGenerator;
_pIdGenerator = nullptr; _pIdGenerator = nullptr;
} }
if (_pIdMapIndex != nullptr) { if (pIndex_ != nullptr) {
delete _pIdMapIndex; delete pIndex_;
_pIdMapIndex = nullptr; pIndex_ = nullptr;
}
if (_pInnerIndex != nullptr) {
delete _pInnerIndex;
_pInnerIndex = nullptr;
} }
} }

View File

@ -10,7 +10,6 @@
#include "Status.h" #include "Status.h"
namespace faiss { namespace faiss {
class IndexIDMap;
class Index; class Index;
} }
@ -50,8 +49,7 @@ private:
const std::string _file_location; const std::string _file_location;
IDGenerator* _pIdGenerator; IDGenerator* _pIdGenerator;
size_t _dimension; size_t _dimension;
faiss::Index* _pInnerIndex; faiss::Index* pIndex_;
faiss::IndexIDMap* _pIdMapIndex;
}; // MemVectors }; // MemVectors