milvus/cpp/src/db/MemTableFile.cpp
zhiru 97d129fc38 Implemented add and serialize
Former-commit-id: 02d0ec1da6c441ff0c05d9933a487886dfbd0f96
2019-07-09 04:44:50 +08:00

108 lines
3.4 KiB
C++

#include "MemTableFile.h"
#include "Constants.h"
#include "Log.h"
#include "EngineFactory.h"
#include "metrics/Metrics.h"
#include <cmath>
namespace zilliz {
namespace milvus {
namespace engine {
MemTableFile::MemTableFile(const std::string& table_id,
const std::shared_ptr<meta::Meta>& meta,
const Options& options) :
table_id_(table_id),
meta_(meta),
options_(options) {
current_mem_ = 0;
auto status = CreateTableFile();
if (status.ok()) {
execution_engine_ = EngineFactory::Build(table_file_schema_.dimension_,
table_file_schema_.location_,
(EngineType)table_file_schema_.engine_type_);
}
}
Status MemTableFile::CreateTableFile() {
meta::TableFileSchema table_file_schema;
table_file_schema.table_id_ = table_id_;
auto status = meta_->CreateTableFile(table_file_schema);
if (status.ok()) {
table_file_schema_ = table_file_schema;
}
else {
std::string errMsg = "MemTableFile::CreateTableFile failed: " + status.ToString();
ENGINE_LOG_ERROR << errMsg;
}
return status;
}
Status MemTableFile::Add(const VectorSource::Ptr& source) {
if (table_file_schema_.dimension_ <= 0) {
std::string errMsg = "MemTableFile::Add: table_file_schema dimension = " +
std::to_string(table_file_schema_.dimension_) + ", table_id = " + table_file_schema_.table_id_;
ENGINE_LOG_ERROR << errMsg;
return Status::Error(errMsg);
}
size_t singleVectorMemSize = table_file_schema_.dimension_ * VECTOR_TYPE_SIZE;
size_t memLeft = GetMemLeft();
if (memLeft >= singleVectorMemSize) {
size_t numVectorsToAdd = std::ceil(memLeft / singleVectorMemSize);
size_t numVectorsAdded;
auto status = source->Add(execution_engine_, table_file_schema_, numVectorsToAdd, numVectorsAdded);
if (status.ok()) {
current_mem_ += (numVectorsAdded * singleVectorMemSize);
}
return status;
}
return Status::OK();
}
size_t MemTableFile::GetCurrentMem() {
return current_mem_;
}
size_t MemTableFile::GetMemLeft() {
return (MAX_TABLE_FILE_MEM - current_mem_);
}
bool MemTableFile::IsFull() {
size_t singleVectorMemSize = table_file_schema_.dimension_ * VECTOR_TYPE_SIZE;
return (GetMemLeft() < singleVectorMemSize);
}
Status MemTableFile::Serialize() {
auto start_time = METRICS_NOW_TIME;
auto size = GetCurrentMem();
execution_engine_->Serialize();
auto end_time = METRICS_NOW_TIME;
auto total_time = METRICS_MICROSECONDS(start_time, end_time);
table_file_schema_.size_ = size;
server::Metrics::GetInstance().DiskStoreIOSpeedGaugeSet((double)size/total_time);
table_file_schema_.file_type_ = (size >= options_.index_trigger_size) ?
meta::TableFileSchema::TO_INDEX : meta::TableFileSchema::RAW;
auto status = meta_->UpdateTableFile(table_file_schema_);
LOG(DEBUG) << "New " << ((table_file_schema_.file_type_ == meta::TableFileSchema::RAW) ? "raw" : "to_index")
<< " file " << table_file_schema_.file_id_ << " of size " << (double)size / (double)M << " M";
execution_engine_->Cache();
return status;
}
} // namespace engine
} // namespace milvus
} // namespace zilliz