milvus/cpp/src/storage/memvectors.cpp
Xu Peng d7f5e351a2 feat(storage): add id generators and memvectors
Former-commit-id: 78f7406456556882247c1ac4ccc9bb928699f896
2019-04-15 20:49:30 +08:00

85 lines
2.0 KiB
C++

#include <IndexFlat.h>
#include <MetaIndexes.h>
#include <index_io.h>
#include "memvectors.h"
namespace vecengine {
MemVectors::MemVectors(size_t dimension_, const std::string& file_location_) :
_file_location(file_location_),
_pIdGenerator(new SimpleIDGenerator()),
_dimension(dimension_),
_pInnerIndex(new faiss::IndexFlat(_dimension)),
_pIdMapIndex = new faiss::IndexIDMap(_pInnerIndex) {
}
IDNumbers&& MemVectors::add(size_t n, const float* vectors) {
IDNumbers&& ids = _pIdGenerator->getNextIDNumbers(n);
_pIdMapIndex->add_with_ids(n, vectors, pIds, &ids[0]);
return ids;
}
size_t MemVectors::total() const {
return _pIdMapIndex->ntotal;
}
size_t MemVectors::approximate_size() const {
return total() * _dimension;
}
void MemVectors::serialize() {
faiss::write_index(_pIdMapIndex, _file_location);
}
MemVectors::~MemVectors() {
if (_pIdGenerator != nullptr) {
delete _pIdGenerator;
_pIdGenerator = nullptr;
}
if (_pIdMapIndex != nullptr) {
delete _pIdMapIndex;
_pIdMapIndex = nullptr;
}
if (_pInnerIndex != nullptr) {
delete _pInnerIndex;
_pInnerIndex = nullptr;
}
}
/*
* MemManager
*/
MemVectors* MemManager::get_mem_by_group(const std::string& group_id_) {
auto memIt = _memMap.find(group_id_);
if memIt != _memMap.end() {
return &(memIt->second);
}
// PXU TODO:
// 1. Read Group meta info
// 2. Initalize MemVectors base meta info
return nullptr;
/* GroupMetaInfo info; */
/* bool succ = env->getGroupMeta(group_id, &info); */
/* if (!succ) { */
/* return nullptr; */
/* } */
/* _memMap[group_id] = MemVectors(info.dimension, info.next_file_location); */
/* return &(_memMap[group_id]); */
}
IDNumbers&& MemManager::add_vectors_no_lock(const std::string& group_id_,
size_t n,
const float* vectors) {
auto mem = get_group_mem(group_id_);
if (mem == nullptr) {
return IDNumbers();
}
return mem->add(n, vectors);
}
} // namespace vecengine