mirror of https://github.com/milvus-io/milvus.git
fix
xiaojun.lin
parent
e01eb29fff
commit
38da8dbb1d
|
|
@ -248,130 +248,5 @@ BFIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const
|
||||||
return Status::OK();
|
return Status::OK();
|
||||||
}
|
}
|
||||||
|
|
||||||
// TODO(linxj): add lock here.
|
|
||||||
Status
|
|
||||||
IVFMixIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt,
|
|
||||||
const float* xt) {
|
|
||||||
try {
|
|
||||||
dim = cfg->d;
|
|
||||||
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
|
|
||||||
|
|
||||||
auto preprocessor = index_->BuildPreprocessor(dataset, cfg);
|
|
||||||
index_->set_preprocessor(preprocessor);
|
|
||||||
auto model = index_->Train(dataset, cfg);
|
|
||||||
index_->set_index_model(model);
|
|
||||||
index_->Add(dataset, cfg);
|
|
||||||
|
|
||||||
if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
|
|
||||||
auto host_index = device_index->CopyGpuToCpu(Config());
|
|
||||||
index_ = host_index;
|
|
||||||
type = ConvertToCpuIndexType(type);
|
|
||||||
} else {
|
|
||||||
WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed";
|
|
||||||
return Status(KNOWHERE_ERROR, "Build IVFMIXIndex Failed");
|
|
||||||
}
|
|
||||||
} catch (knowhere::KnowhereException& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
|
|
||||||
} catch (std::exception& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
return Status(KNOWHERE_ERROR, e.what());
|
|
||||||
}
|
|
||||||
return Status::OK();
|
|
||||||
}
|
|
||||||
|
|
||||||
Status
|
|
||||||
IVFMixIndex::Load(const knowhere::BinarySet& index_binary) {
|
|
||||||
index_->Load(index_binary);
|
|
||||||
dim = Dimension();
|
|
||||||
return Status::OK();
|
|
||||||
}
|
|
||||||
|
|
||||||
knowhere::QuantizerPtr
|
|
||||||
IVFHybridIndex::LoadQuantizer(const Config& conf) {
|
|
||||||
// TODO(linxj): Hardcode here
|
|
||||||
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
|
|
||||||
return new_idx->LoadQuantizer(conf);
|
|
||||||
} else {
|
|
||||||
WRAPPER_LOG_ERROR << "Hybrid mode not supported for index type: " << int(type);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Status
|
|
||||||
IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) {
|
|
||||||
try {
|
|
||||||
// TODO(linxj): Hardcode here
|
|
||||||
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
|
|
||||||
new_idx->SetQuantizer(q);
|
|
||||||
} else {
|
|
||||||
WRAPPER_LOG_ERROR << "Hybrid mode not supported for index type: " << int(type);
|
|
||||||
return Status(KNOWHERE_ERROR, "not supported");
|
|
||||||
}
|
|
||||||
} catch (knowhere::KnowhereException& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
|
|
||||||
} catch (std::exception& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
return Status(KNOWHERE_ERROR, e.what());
|
|
||||||
}
|
|
||||||
return Status::OK();
|
|
||||||
}
|
|
||||||
|
|
||||||
Status
|
|
||||||
IVFHybridIndex::UnsetQuantizer() {
|
|
||||||
try {
|
|
||||||
// TODO(linxj): Hardcode here
|
|
||||||
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
|
|
||||||
new_idx->UnsetQuantizer();
|
|
||||||
} else {
|
|
||||||
WRAPPER_LOG_ERROR << "Hybrid mode not supported for index type: " << int(type);
|
|
||||||
return Status(KNOWHERE_ERROR, "not supported");
|
|
||||||
}
|
|
||||||
} catch (knowhere::KnowhereException& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
|
|
||||||
} catch (std::exception& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
return Status(KNOWHERE_ERROR, e.what());
|
|
||||||
}
|
|
||||||
return Status::OK();
|
|
||||||
}
|
|
||||||
|
|
||||||
VecIndexPtr
|
|
||||||
IVFHybridIndex::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) {
|
|
||||||
try {
|
|
||||||
// TODO(linxj): Hardcode here
|
|
||||||
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
|
|
||||||
return std::make_shared<IVFHybridIndex>(new_idx->LoadData(q, conf), type);
|
|
||||||
} else {
|
|
||||||
WRAPPER_LOG_ERROR << "Hybrid mode not supported for index type: " << int(type);
|
|
||||||
}
|
|
||||||
} catch (knowhere::KnowhereException& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
} catch (std::exception& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
}
|
|
||||||
return nullptr;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::pair<VecIndexPtr, knowhere::QuantizerPtr>
|
|
||||||
IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) {
|
|
||||||
try {
|
|
||||||
// TODO(linxj): Hardcode here
|
|
||||||
if (auto hybrid_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
|
|
||||||
auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg);
|
|
||||||
auto new_idx = std::make_shared<IVFHybridIndex>(pair.first, type);
|
|
||||||
return std::make_pair(new_idx, pair.second);
|
|
||||||
} else {
|
|
||||||
WRAPPER_LOG_ERROR << "Hybrid mode not supported for index type: " << int(type);
|
|
||||||
}
|
|
||||||
} catch (knowhere::KnowhereException& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
} catch (std::exception& e) {
|
|
||||||
WRAPPER_LOG_ERROR << e.what();
|
|
||||||
}
|
|
||||||
return std::make_pair(nullptr, nullptr);
|
|
||||||
}
|
|
||||||
|
|
||||||
} // namespace engine
|
} // namespace engine
|
||||||
} // namespace milvus
|
} // namespace milvus
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue