diff --git a/CHANGELOG.md b/CHANGELOG.md index b2719a2781..8c0820834d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ Please mark all change in change log and use the issue from GitHub - \#1660 IVF PQ CPU support deleted vectors searching ## Improvement +- \#342 Knowhere and Wrapper refactor - \#1537 Optimize raw vector and uids read/write - \#1546 Move Config.cpp to config directory - \#1547 Rename storage/file to storage/disk and rename classes diff --git a/core/src/CMakeLists.txt b/core/src/CMakeLists.txt index 173e1aa293..479c01ef2c 100644 --- a/core/src/CMakeLists.txt +++ b/core/src/CMakeLists.txt @@ -116,7 +116,7 @@ set(storage_files aux_source_directory(${MILVUS_ENGINE_SRC}/utils utils_files) -aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper wrapper_files) +aux_source_directory(${MILVUS_ENGINE_SRC}/index/archive wrapper_files) aux_source_directory(${MILVUS_ENGINE_SRC}/tracing tracing_files) diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index 921c920617..bd8e6dd215 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -22,21 +22,31 @@ #include "cache/GpuCacheMgr.h" #include "config/Config.h" #include "db/Utils.h" +#include "index/archive/VecIndex.h" #include "knowhere/common/Config.h" +#include "knowhere/index/vector_index/ConfAdapter.h" +#include "knowhere/index/vector_index/ConfAdapterMgr.h" +#include "knowhere/index/vector_index/IndexBinaryIDMAP.h" +#include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/VecIndex.h" +#include "knowhere/index/vector_index/VecIndexFactory.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/gpu/GPUIndex.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/gpu/Quantizer.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" +#endif +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "metrics/Metrics.h" #include "scheduler/Utils.h" #include "utils/CommonUtil.h" #include "utils/Exception.h" #include "utils/Log.h" +#include "utils/Status.h" #include "utils/TimeRecorder.h" #include "utils/ValidationUtil.h" -#include "wrapper/BinVecImpl.h" -#include "wrapper/ConfAdapter.h" -#include "wrapper/ConfAdapterMgr.h" -#include "wrapper/VecImpl.h" -#include "wrapper/VecIndex.h" -//#define ON_SEARCH namespace milvus { namespace engine { @@ -74,12 +84,13 @@ MappingMetricType(MetricType metric_type, milvus::json& conf) { } bool -IsBinaryIndexType(IndexType type) { - return type == IndexType::FAISS_BIN_IDMAP || type == IndexType::FAISS_BIN_IVFLAT_CPU; +IsBinaryIndexType(knowhere::IndexType type) { + return type == knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP || type == knowhere::IndexEnum::INDEX_FAISS_BIN_IVFFLAT; } } // namespace +#ifdef MILVUS_GPU_VERSION class CachedQuantizer : public cache::DataObj { public: explicit CachedQuantizer(knowhere::QuantizerPtr data) : data_(std::move(data)) { @@ -98,6 +109,7 @@ class CachedQuantizer : public cache::DataObj { private: knowhere::QuantizerPtr data_; }; +#endif ExecutionEngineImpl::ExecutionEngineImpl(uint16_t dimension, const std::string& location, EngineType index_type, MetricType metric_type, const milvus::json& index_params) @@ -118,24 +130,22 @@ ExecutionEngineImpl::ExecutionEngineImpl(uint16_t dimension, const std::string& conf[knowhere::meta::DIM] = dimension; MappingMetricType(metric_type, conf); ENGINE_LOG_DEBUG << "Index params: " << conf.dump(); - auto adapter = AdapterMgr::GetInstance().GetAdapter(index_->GetType()); - if (!adapter->CheckTrain(conf)) { + auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_->index_type()); + if (!adapter->CheckTrain(conf, index_->index_mode())) { throw Exception(DB_ERROR, "Illegal index params"); } - ErrorCode ec = KNOWHERE_UNEXPECTED_ERROR; - if (auto bf_index = std::dynamic_pointer_cast(index_)) { - ec = bf_index->Build(conf); - } else if (auto bf_bin_index = std::dynamic_pointer_cast(index_)) { - ec = bf_bin_index->Build(conf); - } - if (ec != KNOWHERE_SUCCESS) { - throw Exception(DB_ERROR, "Build index error"); + fiu_do_on("ExecutionEngineImpl.throw_exception", throw Exception(DB_ERROR, "")); + if (auto bf_index = std::dynamic_pointer_cast(index_)) { + bf_index->Train(knowhere::DatasetPtr(), conf); + } else if (auto bf_bin_index = std::dynamic_pointer_cast(index_)) { + bf_bin_index->Train(knowhere::DatasetPtr(), conf); } } -ExecutionEngineImpl::ExecutionEngineImpl(VecIndexPtr index, const std::string& location, EngineType index_type, - MetricType metric_type, const milvus::json& index_params) +ExecutionEngineImpl::ExecutionEngineImpl(knowhere::VecIndexPtr index, const std::string& location, + EngineType index_type, MetricType metric_type, + const milvus::json& index_params) : index_(std::move(index)), location_(location), index_type_(index_type), @@ -143,105 +153,93 @@ ExecutionEngineImpl::ExecutionEngineImpl(VecIndexPtr index, const std::string& l index_params_(index_params) { } -VecIndexPtr +knowhere::VecIndexPtr ExecutionEngineImpl::CreatetVecIndex(EngineType type) { + knowhere::VecIndexFactory& vec_index_factory = knowhere::VecIndexFactory::GetInstance(); + knowhere::IndexMode mode = knowhere::IndexMode::MODE_CPU; #ifdef MILVUS_GPU_VERSION server::Config& config = server::Config::GetInstance(); bool gpu_resource_enable = true; config.GetGpuResourceConfigEnable(gpu_resource_enable); fiu_do_on("ExecutionEngineImpl.CreatetVecIndex.gpu_res_disabled", gpu_resource_enable = false); + if (gpu_resource_enable) { + mode = knowhere::IndexMode::MODE_GPU; + } #endif - fiu_do_on("ExecutionEngineImpl.CreatetVecIndex.invalid_type", type = EngineType::INVALID); - std::shared_ptr index; + fiu_do_on("ExecutionEngineImpl.CreateVecIndex.invalid_type", type = EngineType::INVALID); + knowhere::VecIndexPtr index = nullptr; switch (type) { case EngineType::FAISS_IDMAP: { - index = GetVecIndexFactory(IndexType::FAISS_IDMAP); + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_IDMAP, mode); break; } case EngineType::FAISS_IVFFLAT: { -#ifdef MILVUS_GPU_VERSION - if (gpu_resource_enable) - index = GetVecIndexFactory(IndexType::FAISS_IVFFLAT_MIX); - else -#endif - index = GetVecIndexFactory(IndexType::FAISS_IVFFLAT_CPU); + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT, mode); + break; + } + case EngineType::FAISS_PQ: { + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_IVFPQ, mode); break; } case EngineType::FAISS_IVFSQ8: { -#ifdef MILVUS_GPU_VERSION - if (gpu_resource_enable) - index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_MIX); - else -#endif - index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_CPU); - break; - } - case EngineType::NSG_MIX: { - index = GetVecIndexFactory(IndexType::NSG_MIX); + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_IVFSQ8, mode); break; } #ifdef CUSTOMIZATION #ifdef MILVUS_GPU_VERSION case EngineType::FAISS_IVFSQ8H: { - if (gpu_resource_enable) { - index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_HYBRID); - } else { - throw Exception(DB_ERROR, "No GPU resources for IVFSQ8H"); - } + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H, mode); break; } #endif #endif - case EngineType::FAISS_PQ: { -#ifdef MILVUS_GPU_VERSION - if (gpu_resource_enable) - index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_MIX); - else -#endif - index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_CPU); - break; - } - case EngineType::SPTAG_KDT: { - index = GetVecIndexFactory(IndexType::SPTAG_KDT_RNT_CPU); - break; - } - case EngineType::SPTAG_BKT: { - index = GetVecIndexFactory(IndexType::SPTAG_BKT_RNT_CPU); - break; - } - case EngineType::HNSW: { - index = GetVecIndexFactory(IndexType::HNSW); - break; - } case EngineType::FAISS_BIN_IDMAP: { - index = GetVecIndexFactory(IndexType::FAISS_BIN_IDMAP); + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP, mode); break; } case EngineType::FAISS_BIN_IVFFLAT: { - index = GetVecIndexFactory(IndexType::FAISS_BIN_IVFLAT_CPU); + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_BIN_IVFFLAT, mode); + break; + } + case EngineType::NSG_MIX: { + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_NSG, mode); + break; + } + case EngineType::SPTAG_KDT: { + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_SPTAG_KDT_RNT, mode); + break; + } + case EngineType::SPTAG_BKT: { + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_SPTAG_BKT_RNT, mode); + break; + } + case EngineType::HNSW: { + index = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_HNSW, mode); break; } default: { - ENGINE_LOG_ERROR << "Unsupported index type"; + ENGINE_LOG_ERROR << "Unsupported index type " << (int)type; return nullptr; } } + if (index == nullptr) { + std::string err_msg = "Invalid index type " + std::to_string((int)type) + " mod " + std::to_string((int)mode); + ENGINE_LOG_ERROR << err_msg; + throw Exception(DB_ERROR, err_msg); + } return index; } void ExecutionEngineImpl::HybridLoad() const { - if (index_type_ != EngineType::FAISS_IVFSQ8H) { - return; - } - - if (index_->GetType() == IndexType::FAISS_IDMAP) { - ENGINE_LOG_WARNING << "HybridLoad with type FAISS_IDMAP, ignore"; - return; - } - #ifdef MILVUS_GPU_VERSION + auto hybrid_index = std::dynamic_pointer_cast(index_); + if (hybrid_index == nullptr) { + ENGINE_LOG_WARNING << "HybridLoad only support with IVFSQHybrid"; + return; + } + const std::string key = location_ + ".quantizer"; server::Config& config = server::Config::GetInstance(); @@ -267,7 +265,7 @@ ExecutionEngineImpl::HybridLoad() const { } if (device_id != NOT_FOUND) { - index_->SetQuantizer(quantizer); + hybrid_index->SetQuantizer(quantizer); return; } } @@ -286,12 +284,12 @@ ExecutionEngineImpl::HybridLoad() const { auto best_device_id = gpus[best_index]; milvus::json quantizer_conf{{knowhere::meta::DEVICEID, best_device_id}, {"mode", 1}}; - auto quantizer = index_->LoadQuantizer(quantizer_conf); + auto quantizer = hybrid_index->LoadQuantizer(quantizer_conf); ENGINE_LOG_DEBUG << "Quantizer params: " << quantizer_conf.dump(); if (quantizer == nullptr) { ENGINE_LOG_ERROR << "quantizer is nullptr"; } - index_->SetQuantizer(quantizer); + hybrid_index->SetQuantizer(quantizer); auto cache_quantizer = std::make_shared(quantizer); cache::GpuCacheMgr::GetInstance(best_device_id)->InsertItem(key, cache_quantizer); } @@ -300,25 +298,27 @@ ExecutionEngineImpl::HybridLoad() const { void ExecutionEngineImpl::HybridUnset() const { - if (index_type_ != EngineType::FAISS_IVFSQ8H) { +#ifdef MILVUS_GPU_VERSION + auto hybrid_index = std::dynamic_pointer_cast(index_); + if (hybrid_index == nullptr) { return; } - if (index_->GetType() == IndexType::FAISS_IDMAP) { - return; - } - index_->UnsetQuantizer(); + hybrid_index->UnsetQuantizer(); +#endif } Status ExecutionEngineImpl::AddWithIds(int64_t n, const float* xdata, const int64_t* xids) { - auto status = index_->Add(n, xdata, xids); - return status; + auto dataset = knowhere::GenDatasetWithIds(n, index_->Dim(), xdata, xids); + index_->Add(dataset, knowhere::Config()); + return Status::OK(); } Status ExecutionEngineImpl::AddWithIds(int64_t n, const uint8_t* xdata, const int64_t* xids) { - auto status = index_->Add(n, xdata, xids); - return status; + auto dataset = knowhere::GenDatasetWithIds(n, index_->Dim(), xdata, xids); + index_->Add(dataset, knowhere::Config()); + return Status::OK(); } size_t @@ -336,7 +336,7 @@ ExecutionEngineImpl::Dimension() const { ENGINE_LOG_ERROR << "ExecutionEngineImpl: index is null, return dimension " << dim_; return dim_; } - return index_->Dimension(); + return index_->Dim(); } size_t @@ -369,21 +369,25 @@ Status ExecutionEngineImpl::Load(bool to_cache) { // TODO(zhiru): refactor - index_ = std::static_pointer_cast(cache::CpuCacheMgr::GetInstance()->GetIndex(location_)); + index_ = std::static_pointer_cast(cache::CpuCacheMgr::GetInstance()->GetIndex(location_)); bool already_in_cache = (index_ != nullptr); if (!already_in_cache) { std::string segment_dir; utils::GetParentPath(location_, segment_dir); auto segment_reader_ptr = std::make_shared(segment_dir); + knowhere::VecIndexFactory& vec_index_factory = knowhere::VecIndexFactory::GetInstance(); if (utils::IsRawIndexType((int32_t)index_type_)) { - index_ = index_type_ == EngineType::FAISS_IDMAP ? GetVecIndexFactory(IndexType::FAISS_IDMAP) - : GetVecIndexFactory(IndexType::FAISS_BIN_IDMAP); + if (index_type_ == EngineType::FAISS_IDMAP) { + index_ = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_IDMAP); + } else { + index_ = vec_index_factory.CreateVecIndex(knowhere::IndexEnum::INDEX_FAISS_BIN_IDMAP); + } milvus::json conf{{knowhere::meta::DEVICEID, gpu_num_}, {knowhere::meta::DIM, dim_}}; MappingMetricType(metric_type_, conf); - auto adapter = AdapterMgr::GetInstance().GetAdapter(index_->GetType()); + auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_->index_type()); ENGINE_LOG_DEBUG << "Index params: " << conf.dump(); - if (!adapter->CheckTrain(conf)) { + if (!adapter->CheckTrain(conf, index_->index_mode())) { throw Exception(DB_ERROR, "Illegal index params"); } @@ -413,26 +417,21 @@ ExecutionEngineImpl::Load(bool to_cache) { } } - ErrorCode ec = KNOWHERE_UNEXPECTED_ERROR; if (index_type_ == EngineType::FAISS_IDMAP) { + auto bf_index = std::static_pointer_cast(index_); std::vector float_vectors; float_vectors.resize(vectors_data.size() / sizeof(float)); memcpy(float_vectors.data(), vectors_data.data(), vectors_data.size()); - ec = std::static_pointer_cast(index_)->Build(conf); - if (ec != KNOWHERE_SUCCESS) { - return status; - } - status = std::static_pointer_cast(index_)->AddWithoutIds(vectors->GetCount(), - float_vectors.data(), Config()); - status = std::static_pointer_cast(index_)->SetBlacklist(concurrent_bitset_ptr); + bf_index->Train(knowhere::DatasetPtr(), conf); + auto dataset = knowhere::GenDataset(vectors->GetCount(), this->dim_, float_vectors.data()); + bf_index->AddWithoutIds(dataset, conf); + bf_index->SetBlacklist(concurrent_bitset_ptr); } else if (index_type_ == EngineType::FAISS_BIN_IDMAP) { - ec = std::static_pointer_cast(index_)->Build(conf); - if (ec != KNOWHERE_SUCCESS) { - return status; - } - status = std::static_pointer_cast(index_)->AddWithoutIds(vectors->GetCount(), - vectors_data.data(), Config()); - status = std::static_pointer_cast(index_)->SetBlacklist(concurrent_bitset_ptr); + auto bin_bf_index = std::static_pointer_cast(index_); + bin_bf_index->Train(knowhere::DatasetPtr(), conf); + auto dataset = knowhere::GenDataset(vectors->GetCount(), this->dim_, vectors_data.data()); + bin_bf_index->AddWithoutIds(dataset, conf); + bin_bf_index->SetBlacklist(concurrent_bitset_ptr); } int64_t index_size = vectors->Size(); // vector data size + vector ids size @@ -444,7 +443,6 @@ ExecutionEngineImpl::Load(bool to_cache) { } ENGINE_LOG_DEBUG << "Finished loading raw data from segment " << segment_dir; - } else { try { // size_t physical_size = PhysicalSize(); @@ -554,7 +552,8 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) { #endif #ifdef MILVUS_GPU_VERSION - auto index = std::static_pointer_cast(cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_)); + auto data_obj_ptr = cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_); + auto index = std::static_pointer_cast(data_obj_ptr); bool already_in_cache = (index != nullptr); if (already_in_cache) { index_ = index; @@ -565,7 +564,7 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) { } try { - index_ = index_->CopyToGpu(device_id); + index_ = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config()); ENGINE_LOG_DEBUG << "CPU to GPU" << device_id; } catch (std::exception& e) { ENGINE_LOG_ERROR << e.what(); @@ -587,7 +586,7 @@ ExecutionEngineImpl::CopyToIndexFileToGpu(uint64_t device_id) { // the ToIndexData is only a placeholder, cpu-copy-to-gpu action is performed in if (index_) { gpu_num_ = device_id; - auto to_index_data = std::make_shared(index_->Size()); + auto to_index_data = std::make_shared(index_->Size()); cache::DataObjPtr obj = std::static_pointer_cast(to_index_data); milvus::cache::GpuCacheMgr::GetInstance(device_id)->InsertItem(location_ + "_placeholder", obj); } @@ -597,7 +596,8 @@ ExecutionEngineImpl::CopyToIndexFileToGpu(uint64_t device_id) { Status ExecutionEngineImpl::CopyToCpu() { - auto index = std::static_pointer_cast(cache::CpuCacheMgr::GetInstance()->GetIndex(location_)); +#ifdef MILVUS_GPU_VERSION + auto index = std::static_pointer_cast(cache::CpuCacheMgr::GetInstance()->GetIndex(location_)); bool already_in_cache = (index != nullptr); if (already_in_cache) { index_ = index; @@ -608,7 +608,7 @@ ExecutionEngineImpl::CopyToCpu() { } try { - index_ = index_->CopyToCpu(); + index_ = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); ENGINE_LOG_DEBUG << "GPU to CPU"; } catch (std::exception& e) { ENGINE_LOG_ERROR << e.what(); @@ -620,14 +620,18 @@ ExecutionEngineImpl::CopyToCpu() { Cache(); } return Status::OK(); +#else + ENGINE_LOG_ERROR << "Calling ExecutionEngineImpl::CopyToCpu when using CPU version"; + return Status(DB_ERROR, "Calling ExecutionEngineImpl::CopyToCpu when using CPU version"); +#endif } ExecutionEnginePtr ExecutionEngineImpl::BuildIndex(const std::string& location, EngineType engine_type) { ENGINE_LOG_DEBUG << "Build index file: " << location << " from: " << location_; - auto from_index = std::dynamic_pointer_cast(index_); - auto bin_from_index = std::dynamic_pointer_cast(index_); + auto from_index = std::dynamic_pointer_cast(index_); + auto bin_from_index = std::dynamic_pointer_cast(index_); if (from_index == nullptr && bin_from_index == nullptr) { ENGINE_LOG_ERROR << "ExecutionEngineImpl: from_index is null, failed to build index"; return nullptr; @@ -644,24 +648,36 @@ ExecutionEngineImpl::BuildIndex(const std::string& location, EngineType engine_t conf[knowhere::meta::DEVICEID] = gpu_num_; MappingMetricType(metric_type_, conf); ENGINE_LOG_DEBUG << "Index params: " << conf.dump(); - auto adapter = AdapterMgr::GetInstance().GetAdapter(to_index->GetType()); - if (!adapter->CheckTrain(conf)) { + auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(to_index->index_type()); + if (!adapter->CheckTrain(conf, to_index->index_mode())) { throw Exception(DB_ERROR, "Illegal index params"); } ENGINE_LOG_DEBUG << "Index config: " << conf.dump(); - auto status = Status::OK(); std::vector uids; faiss::ConcurrentBitsetPtr blacklist; if (from_index) { - status = to_index->BuildAll(Count(), from_index->GetRawVectors(), from_index->GetRawIds(), conf); + auto dataset = + knowhere::GenDatasetWithIds(Count(), Dimension(), from_index->GetRawVectors(), from_index->GetRawIds()); + to_index->BuildAll(dataset, conf); uids = from_index->GetUids(); from_index->GetBlacklist(blacklist); } else if (bin_from_index) { - status = to_index->BuildAll(Count(), bin_from_index->GetRawVectors(), bin_from_index->GetRawIds(), conf); + auto dataset = knowhere::GenDatasetWithIds(Count(), Dimension(), bin_from_index->GetRawVectors(), + bin_from_index->GetRawIds()); + to_index->BuildAll(dataset, conf); uids = bin_from_index->GetUids(); bin_from_index->GetBlacklist(blacklist); } + +#ifdef MILVUS_GPU_VERSION + /* for GPU index, need copy back to CPU */ + if (to_index->index_mode() == knowhere::IndexMode::MODE_GPU) { + auto device_index = std::dynamic_pointer_cast(to_index); + to_index = device_index->CopyGpuToCpu(conf); + } +#endif + to_index->SetUids(uids); ENGINE_LOG_DEBUG << "Set " << to_index->GetUids().size() << "uids for " << location; if (blacklist != nullptr) { @@ -669,23 +685,31 @@ ExecutionEngineImpl::BuildIndex(const std::string& location, EngineType engine_t ENGINE_LOG_DEBUG << "Set blacklist for index " << location; } - if (!status.ok()) { - throw Exception(DB_ERROR, status.message()); - } - ENGINE_LOG_DEBUG << "Finish build index: " << location; return std::make_shared(to_index, location, engine_type, metric_type_, index_params_); } -// map offsets to ids void -MapUids(const std::vector& uids, int64_t* labels, size_t num) { +MapAndCopyResult(const knowhere::DatasetPtr& dataset, const std::vector& uids, int64_t nq, + int64_t k, float* distances, int64_t* labels) { + int64_t* res_ids = dataset->Get(knowhere::meta::IDS); + float* res_dist = dataset->Get(knowhere::meta::DISTANCE); + + memcpy(distances, res_dist, sizeof(float) * nq * k); + + /* map offsets to ids */ + int64_t num = nq * k; for (int64_t i = 0; i < num; ++i) { - int64_t& offset = labels[i]; + int64_t offset = res_ids[i]; if (offset != -1) { - offset = uids[offset]; + labels[i] = uids[offset]; + } else { + labels[i] = -1; } } + + free(res_ids); + free(res_dist); } Status @@ -752,9 +776,9 @@ ExecutionEngineImpl::Search(int64_t n, const float* data, int64_t k, const milvu milvus::json conf = extra_params; conf[knowhere::meta::TOPK] = k; - auto adapter = AdapterMgr::GetInstance().GetAdapter(index_->GetType()); + auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_->index_type()); ENGINE_LOG_DEBUG << "Search params: " << conf.dump(); - if (!adapter->CheckSearch(conf, index_->GetType())) { + if (!adapter->CheckSearch(conf, index_->index_type(), index_->index_mode())) { throw Exception(DB_ERROR, "Illegal search params"); } @@ -762,24 +786,20 @@ ExecutionEngineImpl::Search(int64_t n, const float* data, int64_t k, const milvu HybridLoad(); } - rc.RecordSection("search prepare"); - auto status = index_->Search(n, data, distances, labels, conf); - rc.RecordSection("search done"); + rc.RecordSection("query prepare"); + auto dataset = knowhere::GenDataset(n, index_->Dim(), data); + auto result = index_->Query(dataset, conf); + rc.RecordSection("query done"); - // map offsets to ids ENGINE_LOG_DEBUG << "get uids " << index_->GetUids().size() << " from index " << location_; - MapUids(index_->GetUids(), labels, n * k); - + MapAndCopyResult(result, index_->GetUids(), n, k, distances, labels); rc.RecordSection("map uids " + std::to_string(n * k)); if (hybrid) { HybridUnset(); } - if (!status.ok()) { - ENGINE_LOG_ERROR << "Search error:" << status.message(); - } - return status; + return Status::OK(); } Status @@ -794,9 +814,9 @@ ExecutionEngineImpl::Search(int64_t n, const uint8_t* data, int64_t k, const mil milvus::json conf = extra_params; conf[knowhere::meta::TOPK] = k; - auto adapter = AdapterMgr::GetInstance().GetAdapter(index_->GetType()); + auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_->index_type()); ENGINE_LOG_DEBUG << "Search params: " << conf.dump(); - if (!adapter->CheckSearch(conf, index_->GetType())) { + if (!adapter->CheckSearch(conf, index_->index_type(), index_->index_mode())) { throw Exception(DB_ERROR, "Illegal search params"); } @@ -804,24 +824,20 @@ ExecutionEngineImpl::Search(int64_t n, const uint8_t* data, int64_t k, const mil HybridLoad(); } - rc.RecordSection("search prepare"); - auto status = index_->Search(n, data, distances, labels, conf); - rc.RecordSection("search done"); + rc.RecordSection("query prepare"); + auto dataset = knowhere::GenDataset(n, index_->Dim(), data); + auto result = index_->Query(dataset, conf); + rc.RecordSection("query done"); - // map offsets to ids ENGINE_LOG_DEBUG << "get uids " << index_->GetUids().size() << " from index " << location_; - MapUids(index_->GetUids(), labels, n * k); - + MapAndCopyResult(result, index_->GetUids(), n, k, distances, labels); rc.RecordSection("map uids " + std::to_string(n * k)); if (hybrid) { HybridUnset(); } - if (!status.ok()) { - ENGINE_LOG_ERROR << "Search error:" << status.message(); - } - return status; + return Status::OK(); } Status @@ -836,9 +852,9 @@ ExecutionEngineImpl::Search(int64_t n, const std::vector& ids, int64_t milvus::json conf = extra_params; conf[knowhere::meta::TOPK] = k; - auto adapter = AdapterMgr::GetInstance().GetAdapter(index_->GetType()); + auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_->index_type()); ENGINE_LOG_DEBUG << "Search params: " << conf.dump(); - if (!adapter->CheckSearch(conf, index_->GetType())) { + if (!adapter->CheckSearch(conf, index_->index_type(), index_->index_mode())) { throw Exception(DB_ERROR, "Illegal search params"); } @@ -887,15 +903,13 @@ ExecutionEngineImpl::Search(int64_t n, const std::vector& ids, int64_t rc.RecordSection("get offset"); - auto status = Status::OK(); if (!offsets.empty()) { - status = index_->SearchById(offsets.size(), offsets.data(), distances, labels, conf); - rc.RecordSection("search done"); + auto dataset = knowhere::GenDatasetWithIds(offsets.size(), index_->Dim(), nullptr, offsets.data()); + auto result = index_->QueryById(dataset, conf); + rc.RecordSection("query by id done"); - // map offsets to ids ENGINE_LOG_DEBUG << "get uids " << index_->GetUids().size() << " from index " << location_; - MapUids(uids, labels, offsets.size() * k); - + MapAndCopyResult(result, uids, offsets.size(), k, distances, labels); rc.RecordSection("map uids " + std::to_string(offsets.size() * k)); } @@ -903,10 +917,7 @@ ExecutionEngineImpl::Search(int64_t n, const std::vector& ids, int64_t HybridUnset(); } - if (!status.ok()) { - ENGINE_LOG_ERROR << "Search error:" << status.message(); - } - return status; + return Status::OK(); } Status @@ -922,16 +933,16 @@ ExecutionEngineImpl::GetVectorByID(const int64_t& id, float* vector, bool hybrid // Only one id for now std::vector ids{id}; - auto status = index_->GetVectorById(1, ids.data(), vector, milvus::json()); + auto dataset = knowhere::GenDatasetWithIds(1, index_->Dim(), nullptr, ids.data()); + auto result = index_->GetVectorById(dataset, knowhere::Config()); + float* res_vec = (float*)(result->Get(knowhere::meta::TENSOR)); + memcpy(vector, res_vec, sizeof(float) * 1 * index_->Dim()); if (hybrid) { HybridUnset(); } - if (!status.ok()) { - ENGINE_LOG_ERROR << "Search error:" << status.message(); - } - return status; + return Status::OK(); } Status @@ -949,16 +960,16 @@ ExecutionEngineImpl::GetVectorByID(const int64_t& id, uint8_t* vector, bool hybr // Only one id for now std::vector ids{id}; - auto status = index_->GetVectorById(1, ids.data(), vector, milvus::json()); + auto dataset = knowhere::GenDatasetWithIds(1, index_->Dim(), nullptr, ids.data()); + auto result = index_->GetVectorById(dataset, knowhere::Config()); + uint8_t* res_vec = (uint8_t*)(result->Get(knowhere::meta::TENSOR)); + memcpy(vector, res_vec, sizeof(uint8_t) * 1 * index_->Dim()); if (hybrid) { HybridUnset(); } - if (!status.ok()) { - ENGINE_LOG_ERROR << "Search error:" << status.message(); - } - return status; + return Status::OK(); } Status diff --git a/core/src/db/engine/ExecutionEngineImpl.h b/core/src/db/engine/ExecutionEngineImpl.h index 10a3968c51..d8ef30a6aa 100644 --- a/core/src/db/engine/ExecutionEngineImpl.h +++ b/core/src/db/engine/ExecutionEngineImpl.h @@ -19,7 +19,7 @@ #include #include "ExecutionEngine.h" -#include "wrapper/VecIndex.h" +#include "knowhere/index/vector_index/VecIndex.h" namespace milvus { namespace engine { @@ -29,8 +29,8 @@ class ExecutionEngineImpl : public ExecutionEngine { ExecutionEngineImpl(uint16_t dimension, const std::string& location, EngineType index_type, MetricType metric_type, const milvus::json& index_params); - ExecutionEngineImpl(VecIndexPtr index, const std::string& location, EngineType index_type, MetricType metric_type, - const milvus::json& index_params); + ExecutionEngineImpl(knowhere::VecIndexPtr index, const std::string& location, EngineType index_type, + MetricType metric_type, const milvus::json& index_params); Status AddWithIds(int64_t n, const float* xdata, const int64_t* xids) override; @@ -108,10 +108,10 @@ class ExecutionEngineImpl : public ExecutionEngine { } private: - VecIndexPtr + knowhere::VecIndexPtr CreatetVecIndex(EngineType type); - VecIndexPtr + knowhere::VecIndexPtr Load(const std::string& location); void @@ -121,7 +121,7 @@ class ExecutionEngineImpl : public ExecutionEngine { HybridUnset() const; protected: - VecIndexPtr index_ = nullptr; + knowhere::VecIndexPtr index_ = nullptr; EngineType index_type_; MetricType metric_type_; diff --git a/core/src/db/insert/MemTable.cpp b/core/src/db/insert/MemTable.cpp index 2ef361d927..6080ee321d 100644 --- a/core/src/db/insert/MemTable.cpp +++ b/core/src/db/insert/MemTable.cpp @@ -9,8 +9,6 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. -#include "db/insert/MemTable.h" - #include #include #include @@ -20,9 +18,10 @@ #include "cache/CpuCacheMgr.h" #include "db/OngoingFileChecker.h" #include "db/Utils.h" +#include "db/insert/MemTable.h" +#include "knowhere/index/vector_index/VecIndex.h" #include "segment/SegmentReader.h" #include "utils/Log.h" -#include "wrapper/VecIndex.h" namespace milvus { namespace engine { @@ -243,11 +242,11 @@ MemTable::ApplyDeletes() { } // Get all index that contains blacklist in cache - std::vector indexes; + std::vector indexes; std::vector blacklists; for (auto& file : segment_files) { - auto index = - std::static_pointer_cast(cache::CpuCacheMgr::GetInstance()->GetIndex(file.location_)); + auto data_obj_ptr = cache::CpuCacheMgr::GetInstance()->GetIndex(file.location_); + auto index = std::static_pointer_cast(data_obj_ptr); faiss::ConcurrentBitsetPtr blacklist = nullptr; if (index != nullptr) { index->GetBlacklist(blacklist); diff --git a/core/src/wrapper/KnowhereResource.cpp b/core/src/index/archive/KnowhereResource.cpp similarity index 98% rename from core/src/wrapper/KnowhereResource.cpp rename to core/src/index/archive/KnowhereResource.cpp index 3e4888b7d0..905d40069f 100644 --- a/core/src/wrapper/KnowhereResource.cpp +++ b/core/src/index/archive/KnowhereResource.cpp @@ -9,7 +9,7 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. -#include "wrapper/KnowhereResource.h" +#include "index/archive/KnowhereResource.h" #ifdef MILVUS_GPU_VERSION #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif diff --git a/core/src/wrapper/KnowhereResource.h b/core/src/index/archive/KnowhereResource.h similarity index 100% rename from core/src/wrapper/KnowhereResource.h rename to core/src/index/archive/KnowhereResource.h diff --git a/core/src/index/archive/VecIndex.cpp b/core/src/index/archive/VecIndex.cpp new file mode 100644 index 0000000000..0d0ecf2347 --- /dev/null +++ b/core/src/index/archive/VecIndex.cpp @@ -0,0 +1,172 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include "config/Config.h" +#include "index/archive/VecIndex.h" +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/VecIndex.h" +#include "knowhere/index/vector_index/VecIndexFactory.h" +#include "storage/disk/DiskIOReader.h" +#include "storage/disk/DiskIOWriter.h" +#include "storage/s3/S3IOReader.h" +#include "storage/s3/S3IOWriter.h" +#include "utils/Log.h" +#include "utils/TimeRecorder.h" + +namespace milvus { +namespace engine { + +knowhere::VecIndexPtr +LoadVecIndex(const knowhere::IndexType& type, const knowhere::BinarySet& index_binary, int64_t size) { + knowhere::VecIndexFactory& vec_index_factory = knowhere::VecIndexFactory::GetInstance(); + auto index = vec_index_factory.CreateVecIndex(type, knowhere::IndexMode::MODE_CPU); + if (index == nullptr) + return nullptr; + // else + index->Load(index_binary); + index->set_size(size); + return index; +} + +knowhere::VecIndexPtr +read_index(const std::string& location) { + milvus::TimeRecorder recorder("read_index"); + knowhere::BinarySet load_data_list; + + bool s3_enable = false; + milvus::server::Config& config = milvus::server::Config::GetInstance(); + config.GetStorageConfigS3Enable(s3_enable); + + std::shared_ptr reader_ptr; + if (s3_enable) { + reader_ptr = std::make_shared(); + } else { + reader_ptr = std::make_shared(); + } + + recorder.RecordSection("Start"); + reader_ptr->open(location); + + size_t length = reader_ptr->length(); + if (length <= 0) { + return nullptr; + } + + size_t rp = 0; + reader_ptr->seekg(0); + + int32_t current_type = 0; + reader_ptr->read(¤t_type, sizeof(current_type)); + rp += sizeof(current_type); + reader_ptr->seekg(rp); + + while (rp < length) { + size_t meta_length; + reader_ptr->read(&meta_length, sizeof(meta_length)); + rp += sizeof(meta_length); + reader_ptr->seekg(rp); + + auto meta = new char[meta_length]; + reader_ptr->read(meta, meta_length); + rp += meta_length; + reader_ptr->seekg(rp); + + size_t bin_length; + reader_ptr->read(&bin_length, sizeof(bin_length)); + rp += sizeof(bin_length); + reader_ptr->seekg(rp); + + auto bin = new uint8_t[bin_length]; + reader_ptr->read(bin, bin_length); + rp += bin_length; + reader_ptr->seekg(rp); + + auto binptr = std::make_shared(); + binptr.reset(bin); + load_data_list.Append(std::string(meta, meta_length), binptr, bin_length); + delete[] meta; + } + reader_ptr->close(); + + double span = recorder.RecordSection("End"); + double rate = length * 1000000.0 / span / 1024 / 1024; + STORAGE_LOG_DEBUG << "read_index(" << location << ") rate " << rate << "MB/s"; + + return LoadVecIndex(knowhere::OldIndexTypeToStr(current_type), load_data_list, length); +} + +milvus::Status +write_index(knowhere::VecIndexPtr index, const std::string& location) { + try { + milvus::TimeRecorder recorder("write_index"); + + auto binaryset = index->Serialize(knowhere::Config()); + int32_t index_type = knowhere::StrToOldIndexType(index->index_type()); + + bool s3_enable = false; + milvus::server::Config& config = milvus::server::Config::GetInstance(); + config.GetStorageConfigS3Enable(s3_enable); + + std::shared_ptr writer_ptr; + if (s3_enable) { + writer_ptr = std::make_shared(); + } else { + writer_ptr = std::make_shared(); + } + + recorder.RecordSection("Start"); + writer_ptr->open(location); + + writer_ptr->write(&index_type, sizeof(index_type)); + + for (auto& iter : binaryset.binary_map_) { + auto meta = iter.first.c_str(); + size_t meta_length = iter.first.length(); + writer_ptr->write(&meta_length, sizeof(meta_length)); + writer_ptr->write((void*)meta, meta_length); + + auto binary = iter.second; + int64_t binary_length = binary->size; + writer_ptr->write(&binary_length, sizeof(binary_length)); + writer_ptr->write((void*)binary->data.get(), binary_length); + } + writer_ptr->close(); + + double span = recorder.RecordSection("End"); + double rate = writer_ptr->length() * 1000000.0 / span / 1024 / 1024; + STORAGE_LOG_DEBUG << "write_index(" << location << ") rate " << rate << "MB/s"; + } catch (knowhere::KnowhereException& e) { + WRAPPER_LOG_ERROR << e.what(); + return milvus::Status(milvus::KNOWHERE_UNEXPECTED_ERROR, e.what()); + } catch (std::exception& e) { + WRAPPER_LOG_ERROR << e.what(); + std::string estring(e.what()); + if (estring.find("No space left on device") != estring.npos) { + WRAPPER_LOG_ERROR << "No space left on the device"; + return milvus::Status(milvus::KNOWHERE_NO_SPACE, "No space left on the device"); + } else { + return milvus::Status(milvus::KNOWHERE_ERROR, e.what()); + } + } + return milvus::Status::OK(); +} + +} // namespace engine +} // namespace milvus diff --git a/core/src/index/archive/VecIndex.h b/core/src/index/archive/VecIndex.h new file mode 100644 index 0000000000..92037019f1 --- /dev/null +++ b/core/src/index/archive/VecIndex.h @@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include + +#include "cache/DataObj.h" +#include "knowhere/common/BinarySet.h" +#include "knowhere/index/vector_index/VecIndex.h" +#include "utils/Status.h" + +namespace milvus { +namespace engine { + +extern milvus::Status +write_index(knowhere::VecIndexPtr index, const std::string& location); + +extern knowhere::VecIndexPtr +read_index(const std::string& location); + +} // namespace engine +} // namespace milvus diff --git a/core/src/index/build.sh b/core/src/index/build.sh old mode 100644 new mode 100755 diff --git a/core/src/index/knowhere/CMakeLists.txt b/core/src/index/knowhere/CMakeLists.txt index 1aae53d0d4..02add8ef97 100644 --- a/core/src/index/knowhere/CMakeLists.txt +++ b/core/src/index/knowhere/CMakeLists.txt @@ -21,31 +21,35 @@ if (NOT TARGET SPTAGLibStatic) endif () set(external_srcs - knowhere/adapter/SptagAdapter.cpp knowhere/common/Exception.cpp knowhere/common/Timer.cpp ) set(index_srcs - knowhere/index/preprocessor/Normalize.cpp - knowhere/index/vector_index/IndexSPTAG.cpp - knowhere/index/vector_index/IndexIDMAP.cpp - knowhere/index/vector_index/IndexIVF.cpp - knowhere/index/vector_index/IndexBinaryIVF.cpp - knowhere/index/vector_index/FaissBaseBinaryIndex.cpp - knowhere/index/vector_index/IndexBinaryIDMAP.cpp - knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp - knowhere/index/vector_index/IndexNSG.cpp - knowhere/index/vector_index/IndexHNSW.cpp - knowhere/index/vector_index/nsg/NSG.cpp - knowhere/index/vector_index/nsg/NSGIO.cpp - knowhere/index/vector_index/nsg/NSGHelper.cpp - knowhere/index/vector_index/nsg/Distance.cpp - knowhere/index/vector_index/IndexIVFSQ.cpp - knowhere/index/vector_index/IndexIVFPQ.cpp - knowhere/index/vector_index/FaissBaseIndex.cpp + knowhere/index/vector_index/adapter/SptagAdapter.cpp + knowhere/index/vector_index/adapter/VectorAdapter.cpp knowhere/index/vector_index/helpers/FaissIO.cpp knowhere/index/vector_index/helpers/IndexParameter.cpp + knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp + knowhere/index/vector_index/impl/nsg/Distance.cpp + knowhere/index/vector_index/impl/nsg/NSG.cpp + knowhere/index/vector_index/impl/nsg/NSGHelper.cpp + knowhere/index/vector_index/impl/nsg/NSGIO.cpp + knowhere/index/vector_index/ConfAdapter.cpp + knowhere/index/vector_index/ConfAdapterMgr.cpp + knowhere/index/vector_index/FaissBaseBinaryIndex.cpp + knowhere/index/vector_index/FaissBaseIndex.cpp + knowhere/index/vector_index/IndexBinaryIDMAP.cpp + knowhere/index/vector_index/IndexBinaryIVF.cpp + knowhere/index/vector_index/IndexHNSW.cpp + knowhere/index/vector_index/IndexIDMAP.cpp + knowhere/index/vector_index/IndexIVF.cpp + knowhere/index/vector_index/IndexIVFPQ.cpp + knowhere/index/vector_index/IndexIVFSQ.cpp + knowhere/index/vector_index/IndexNSG.cpp + knowhere/index/vector_index/IndexSPTAG.cpp + knowhere/index/vector_index/IndexType.cpp + knowhere/index/vector_index/VecIndexFactory.cpp ) set(depend_libs @@ -82,13 +86,13 @@ if (KNOWHERE_GPU_VERSION) ) set(index_srcs ${index_srcs} - knowhere/index/vector_index/IndexGPUIVF.cpp + knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp + knowhere/index/vector_index/gpu/IndexGPUIVF.cpp + knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp + knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp + knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp knowhere/index/vector_index/helpers/Cloner.cpp knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp - knowhere/index/vector_index/IndexGPUIVFSQ.cpp - knowhere/index/vector_index/IndexIVFSQHybrid.cpp - knowhere/index/vector_index/IndexGPUIVFPQ.cpp - knowhere/index/vector_index/IndexGPUIDMAP.cpp ) endif () @@ -110,7 +114,7 @@ set(INDEX_INCLUDE_DIRS ${INDEX_SOURCE_DIR}/knowhere ${INDEX_SOURCE_DIR}/thirdparty ${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService - ${ARROW_INCLUDE_DIR} +# ${ARROW_INCLUDE_DIR} ${FAISS_INCLUDE_DIR} ${OPENBLAS_INCLUDE_DIR} ${LAPACK_INCLUDE_DIR} diff --git a/core/src/index/knowhere/knowhere/common/BinarySet.h b/core/src/index/knowhere/knowhere/common/BinarySet.h index 199096a794..e4acb99cef 100644 --- a/core/src/index/knowhere/knowhere/common/BinarySet.h +++ b/core/src/index/knowhere/knowhere/common/BinarySet.h @@ -7,24 +7,21 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include - #include #include #include #include #include -#include "Id.h" - +namespace milvus { namespace knowhere { struct Binary { - ID id; std::shared_ptr data; int64_t size = 0; }; @@ -76,3 +73,4 @@ class BinarySet { }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Config.h b/core/src/index/knowhere/knowhere/common/Config.h index 4ab9ac0e67..c0dda5f3ce 100644 --- a/core/src/index/knowhere/knowhere/common/Config.h +++ b/core/src/index/knowhere/knowhere/common/Config.h @@ -13,8 +13,10 @@ #include "src/utils/Json.h" +namespace milvus { namespace knowhere { using Config = milvus::json; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Dataset.h b/core/src/index/knowhere/knowhere/common/Dataset.h index 9694ecc444..2d742799f9 100644 --- a/core/src/index/knowhere/knowhere/common/Dataset.h +++ b/core/src/index/knowhere/knowhere/common/Dataset.h @@ -7,74 +7,21 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once -#include +#include #include #include #include -#include #include -#include #include +namespace milvus { namespace knowhere { -struct BaseValue; -using BasePtr = std::unique_ptr; -struct BaseValue { - virtual ~BaseValue() = default; - - // virtual BasePtr - // Clone() const = 0; -}; - -template -struct AnyValue : public BaseValue { - T data_; - - template - explicit AnyValue(U&& value) : data_(std::forward(value)) { - } - - // BasePtr - // Clone() const { - // return BasePtr(data_); - // } -}; - -struct Value { - std::type_index type_; - BasePtr data_; - - template ::type, Value>::value, U>::type> - explicit Value(U&& value) - : data_(new AnyValue::type>(std::forward(value))), - type_(std::type_index(typeid(typename std::decay::type))) { - } - - template - bool - Is() const { - return type_ == std::type_index(typeid(U)); - } - - template - U& - AnyCast() { - if (!Is()) { - std::stringstream ss; - ss << "Can't cast t " << type_.name() << " to " << typeid(U).name(); - throw std::logic_error(ss.str()); - } - - auto derived = dynamic_cast*>(data_.get()); - return derived->data_; - } -}; +using Value = std::any; using ValuePtr = std::shared_ptr; class Dataset { @@ -85,18 +32,16 @@ class Dataset { void Set(const std::string& k, T&& v) { std::lock_guard lk(mutex_); - auto value = std::make_shared(std::forward(v)); - data_[k] = value; + data_[k] = std::make_shared(std::forward(v)); } template T Get(const std::string& k) { std::lock_guard lk(mutex_); - auto finder = data_.find(k); - if (finder != data_.end()) { - return finder->second->AnyCast(); - } else { + try { + return std::any_cast(*(data_.at(k))); + } catch (...) { throw std::logic_error("Can't find this key"); } } @@ -113,3 +58,4 @@ class Dataset { using DatasetPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Exception.cpp b/core/src/index/knowhere/knowhere/common/Exception.cpp index 776e7cd7cf..7e849f7d90 100644 --- a/core/src/index/knowhere/knowhere/common/Exception.cpp +++ b/core/src/index/knowhere/knowhere/common/Exception.cpp @@ -7,13 +7,14 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include "Log.h" #include "knowhere/common/Exception.h" +namespace milvus { namespace knowhere { KnowhereException::KnowhereException(const std::string& msg) : msg(msg) { @@ -41,3 +42,4 @@ KnowhereException::what() const noexcept { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Exception.h b/core/src/index/knowhere/knowhere/common/Exception.h index 520493d1bc..76a08d1a90 100644 --- a/core/src/index/knowhere/knowhere/common/Exception.h +++ b/core/src/index/knowhere/knowhere/common/Exception.h @@ -7,13 +7,14 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include +namespace milvus { namespace knowhere { class KnowhereException : public std::exception { @@ -45,3 +46,4 @@ class KnowhereException : public std::exception { } while (false) } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Id.h b/core/src/index/knowhere/knowhere/common/Id.h deleted file mode 100644 index 83e609df9e..0000000000 --- a/core/src/index/knowhere/knowhere/common/Id.h +++ /dev/null @@ -1,50 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include -#include - -namespace knowhere { - -class ID { - public: - constexpr static int64_t kIDSize = 20; - - public: - const int32_t* - data() const { - return content_; - } - - int32_t* - mutable_data() { - return content_; - } - - bool - IsValid() const; - - std::string - ToString() const; - - bool - operator==(const ID& that) const; - - bool - operator<(const ID& that) const; - - protected: - int32_t content_[5] = {}; -}; - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/common/Log.h b/core/src/index/knowhere/knowhere/common/Log.h index ddd2d7bc45..510799f28c 100644 --- a/core/src/index/knowhere/knowhere/common/Log.h +++ b/core/src/index/knowhere/knowhere/common/Log.h @@ -7,12 +7,13 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include "easyloggingpp/easylogging++.h" +namespace milvus { namespace knowhere { #define KNOWHERE_DOMAIN_NAME "[KNOWHERE] " @@ -26,3 +27,4 @@ namespace knowhere { #define KNOWHERE_LOG_FATAL LOG(FATAL) << KNOWHERE_DOMAIN_NAME } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Timer.cpp b/core/src/index/knowhere/knowhere/common/Timer.cpp index 97aaf1d119..3480163481 100644 --- a/core/src/index/knowhere/knowhere/common/Timer.cpp +++ b/core/src/index/knowhere/knowhere/common/Timer.cpp @@ -7,13 +7,14 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include // TODO(linxj): using Log instead #include "knowhere/common/Log.h" #include "knowhere/common/Timer.h" +namespace milvus { namespace knowhere { TimeRecorder::TimeRecorder(const std::string& header, int64_t log_level) : header_(header), log_level_(log_level) { @@ -81,3 +82,4 @@ TimeRecorder::ElapseFromBegin(const std::string& msg) { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/common/Timer.h b/core/src/index/knowhere/knowhere/common/Timer.h index 6846edb54d..1e98c4ae99 100644 --- a/core/src/index/knowhere/knowhere/common/Timer.h +++ b/core/src/index/knowhere/knowhere/common/Timer.h @@ -7,13 +7,14 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include +namespace milvus { namespace knowhere { class TimeRecorder { @@ -45,3 +46,4 @@ class TimeRecorder { }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/wrapper/WrapperException.h b/core/src/index/knowhere/knowhere/common/Typedef.h similarity index 67% rename from core/src/wrapper/WrapperException.h rename to core/src/index/knowhere/knowhere/common/Typedef.h index a5bc0326ed..8b43a8159d 100644 --- a/core/src/wrapper/WrapperException.h +++ b/core/src/index/knowhere/knowhere/common/Typedef.h @@ -1,3 +1,4 @@ + // Copyright (C) 2019-2020 Zilliz. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance @@ -7,25 +8,22 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once -#include #include +#include namespace milvus { -namespace engine { +namespace knowhere { -class WrapperException : public std::exception { - public: - explicit WrapperException(const std::string& msg); +using MetricType = std::string; +// using IndexType = std::string; +using IDType = int64_t; +using FloatType = float; +using BinaryType = uint8_t; +using GraphType = std::vector>; - const char* - what() const noexcept override; - - const std::string msg; -}; - -} // namespace engine +} // namespace knowhere } // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/Index.h b/core/src/index/knowhere/knowhere/index/Index.h index 7fbcdeb81a..1ef8918c7d 100644 --- a/core/src/index/knowhere/knowhere/index/Index.h +++ b/core/src/index/knowhere/knowhere/index/Index.h @@ -7,56 +7,57 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include -#include "IndexModel.h" -#include "IndexType.h" +#include "cache/DataObj.h" #include "knowhere/common/BinarySet.h" #include "knowhere/common/Config.h" -#include "knowhere/common/Dataset.h" -#include "knowhere/index/preprocessor/Preprocessor.h" +namespace milvus { namespace knowhere { -class Index { +class Index : public milvus::cache::DataObj { public: virtual BinarySet - Serialize() = 0; + Serialize(const Config& config = Config()) = 0; virtual void - Load(const BinarySet& index_binary) = 0; + Load(const BinarySet&) = 0; - // @throw - virtual DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) = 0; - - public: - IndexType - idx_type() const { - return idx_type_; + int64_t + Size() override { + return size_; } void - set_idx_type(IndexType idx_type) { - idx_type_ = idx_type; + set_size(const int64_t& size) { + size_ = size; } - virtual void - set_preprocessor(PreprocessorPtr preprocessor) { - } - - virtual void - set_index_model(IndexModelPtr model) { - } - - private: - IndexType idx_type_; + protected: + int64_t size_ = -1; }; using IndexPtr = std::shared_ptr; +// todo: remove from knowhere +class ToIndexData : public milvus::cache::DataObj { + public: + explicit ToIndexData(int64_t size) : size_(size) { + } + + int64_t + Size() override { + return size_; + } + + private: + int64_t size_ = 0; +}; + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/IndexType.h b/core/src/index/knowhere/knowhere/index/IndexType.h deleted file mode 100644 index 6c0be8151a..0000000000 --- a/core/src/index/knowhere/knowhere/index/IndexType.h +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -namespace knowhere { - -enum class IndexType { - kUnknown = 0, - kVecIdxBegin = 100, - kVecIVFFlat = kVecIdxBegin, - kVecIdxEnd, -}; - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/preprocessor/Normalize.cpp b/core/src/index/knowhere/knowhere/index/preprocessor/Normalize.cpp deleted file mode 100644 index a0f4c45636..0000000000 --- a/core/src/index/knowhere/knowhere/index/preprocessor/Normalize.cpp +++ /dev/null @@ -1,51 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. -// -//#include "knowhere/index/vector_index/definitions.h" -//#include "knowhere/common/config.h" -#include "knowhere/index/preprocessor/Normalize.h" -// -// -// -// namespace knowhere { -// -// DatasetPtr -// NormalizePreprocessor::Preprocess(const DatasetPtr &dataset) { -// // TODO: wrap dataset->tensor -// auto tensor = dataset->tensor()[0]; -// auto p_data = (float *)tensor->raw_mutable_data(); -// auto dimension = tensor->shape()[1]; -// auto rows = tensor->shape()[0]; -// -//#pragma omp parallel for -// for (auto i = 0; i < rows; ++i) { -// Normalize(&(p_data[i * dimension]), dimension); -// } -//} -// -// void -// NormalizePreprocessor::Normalize(float *arr, int64_t dimension) { -// double vector_length = 0; -// for (auto j = 0; j < dimension; j++) { -// double val = arr[j]; -// vector_length += val * val; -// } -// vector_length = std::sqrt(vector_length); -// if (vector_length < 1e-6) { -// auto val = (float) (1.0 / std::sqrt((double) dimension)); -// for (int j = 0; j < dimension; j++) arr[j] = val; -// } else { -// for (int j = 0; j < dimension; j++) arr[j] = (float) (arr[j] / vector_length); -// } -//} -// -//} // namespace knowhere -// diff --git a/core/src/index/knowhere/knowhere/index/preprocessor/Normalize.h b/core/src/index/knowhere/knowhere/index/preprocessor/Normalize.h deleted file mode 100644 index e94a8fcd74..0000000000 --- a/core/src/index/knowhere/knowhere/index/preprocessor/Normalize.h +++ /dev/null @@ -1,37 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. -// -//#pragma once -// -//#include -//#include "preprocessor.h" -// -// -// -// namespace knowhere { -// -// class NormalizePreprocessor : public Preprocessor { -// public: -// DatasetPtr -// Preprocess(const DatasetPtr &input) override; -// -// private: -// -// void -// Normalize(float *arr, int64_t dimension); -//}; -// -// -// using NormalizePreprocessorPtr = std::shared_ptr; -// -// -//} // namespace knowhere -// diff --git a/core/src/index/knowhere/knowhere/index/preprocessor/Preprocessor.h b/core/src/index/knowhere/knowhere/index/preprocessor/Preprocessor.h index a6d78cebf6..6cda87122f 100644 --- a/core/src/index/knowhere/knowhere/index/preprocessor/Preprocessor.h +++ b/core/src/index/knowhere/knowhere/index/preprocessor/Preprocessor.h @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once @@ -15,6 +15,7 @@ #include "knowhere/common/Dataset.h" +namespace milvus { namespace knowhere { class Preprocessor { @@ -26,3 +27,4 @@ class Preprocessor { using PreprocessorPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapter.cpp similarity index 86% rename from core/src/wrapper/ConfAdapter.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/ConfAdapter.cpp index 7026c530a9..3a13da5afd 100644 --- a/core/src/wrapper/ConfAdapter.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapter.cpp @@ -9,21 +9,17 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. -#include "wrapper/ConfAdapter.h" +#include "knowhere/index/vector_index/ConfAdapter.h" -#include #include #include #include #include -#include "WrapperException.h" -#include "config/Config.h" #include "knowhere/index/vector_index/helpers/IndexParameter.h" -#include "utils/Log.h" namespace milvus { -namespace engine { +namespace knowhere { #if CUDA_VERSION > 9000 #define GPU_MAX_NRPOBE 2048 @@ -70,19 +66,16 @@ namespace engine { } bool -ConfAdapter::CheckTrain(milvus::json& oricfg) { +ConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static std::vector METRICS{knowhere::Metric::L2, knowhere::Metric::IP}; - CheckIntByRange(knowhere::meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM); CheckStrByValues(knowhere::Metric::TYPE, METRICS); - return true; } bool -ConfAdapter::CheckSearch(milvus::json& oricfg, const IndexType& type) { +ConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) { CheckIntByRange(knowhere::meta::TOPK, DEFAULT_MIN_K - 1, DEFAULT_MAX_K); - return true; } @@ -100,7 +93,7 @@ MatchNlist(const int64_t& size, const int64_t& nlist, const int64_t& per_nlist) } bool -IVFConfAdapter::CheckTrain(milvus::json& oricfg) { +IVFConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static int64_t MAX_NLIST = 999999; static int64_t MIN_NLIST = 1; @@ -111,42 +104,43 @@ IVFConfAdapter::CheckTrain(milvus::json& oricfg) { // CheckIntByRange(knowhere::meta::ROWS, nlist, DEFAULT_MAX_ROWS); // auto tune params - oricfg[knowhere::IndexParams::nlist] = MatchNlist(oricfg[knowhere::meta::ROWS].get(), - oricfg[knowhere::IndexParams::nlist].get(), 16384); + int64_t nq = oricfg[knowhere::meta::ROWS].get(); + int64_t nlist = oricfg[knowhere::IndexParams::nlist].get(); + oricfg[knowhere::IndexParams::nlist] = MatchNlist(nq, nlist, 16384); // Best Practice // static int64_t MIN_POINTS_PER_CENTROID = 40; // static int64_t MAX_POINTS_PER_CENTROID = 256; // CheckIntByRange(knowhere::meta::ROWS, MIN_POINTS_PER_CENTROID * nlist, MAX_POINTS_PER_CENTROID * nlist); - return ConfAdapter::CheckTrain(oricfg); + return ConfAdapter::CheckTrain(oricfg, mode); } bool -IVFConfAdapter::CheckSearch(milvus::json& oricfg, const IndexType& type) { +IVFConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) { static int64_t MIN_NPROBE = 1; static int64_t MAX_NPROBE = 999999; // todo(linxj): [1, nlist] - if (type == IndexType::FAISS_IVFPQ_GPU || type == IndexType::FAISS_IVFSQ8_GPU || - type == IndexType::FAISS_IVFSQ8_HYBRID || type == IndexType::FAISS_IVFFLAT_GPU) { + if (mode == IndexMode::MODE_GPU) { CheckIntByRange(knowhere::IndexParams::nprobe, MIN_NPROBE, GPU_MAX_NRPOBE); } else { CheckIntByRange(knowhere::IndexParams::nprobe, MIN_NPROBE, MAX_NPROBE); } + CheckIntByRange(knowhere::IndexParams::nprobe, MIN_NPROBE, MAX_NPROBE); - return ConfAdapter::CheckSearch(oricfg, type); + return ConfAdapter::CheckSearch(oricfg, type, mode); } bool -IVFSQConfAdapter::CheckTrain(milvus::json& oricfg) { +IVFSQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static int64_t DEFAULT_NBITS = 8; oricfg[knowhere::IndexParams::nbits] = DEFAULT_NBITS; - return IVFConfAdapter::CheckTrain(oricfg); + return IVFConfAdapter::CheckTrain(oricfg, mode); } bool -IVFPQConfAdapter::CheckTrain(milvus::json& oricfg) { +IVFPQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static int64_t DEFAULT_NBITS = 8; static int64_t MAX_NLIST = 999999; static int64_t MIN_NLIST = 1; @@ -155,17 +149,11 @@ IVFPQConfAdapter::CheckTrain(milvus::json& oricfg) { oricfg[knowhere::IndexParams::nbits] = DEFAULT_NBITS; -#ifdef MILVUS_GPU_VERSION - Status s; - bool enable_gpu = false; - server::Config& config = server::Config::GetInstance(); - s = config.GetGpuResourceConfigEnable(enable_gpu); - if (s.ok()) { + if (mode == IndexMode::MODE_GPU) { CheckStrByValues(knowhere::Metric::TYPE, GPU_METRICS); } else { CheckStrByValues(knowhere::Metric::TYPE, CPU_METRICS); } -#endif CheckIntByRange(knowhere::meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM); CheckIntByRange(knowhere::meta::ROWS, DEFAULT_MIN_ROWS, DEFAULT_MAX_ROWS); CheckIntByRange(knowhere::IndexParams::nlist, MIN_NLIST, MAX_NLIST); @@ -214,7 +202,7 @@ IVFPQConfAdapter::GetValidMList(int64_t dimension, std::vector& resset) } bool -NSGConfAdapter::CheckTrain(milvus::json& oricfg) { +NSGConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static int64_t MIN_KNNG = 5; static int64_t MAX_KNNG = 300; static int64_t MIN_SEARCH_LENGTH = 10; @@ -242,17 +230,17 @@ NSGConfAdapter::CheckTrain(milvus::json& oricfg) { } bool -NSGConfAdapter::CheckSearch(milvus::json& oricfg, const IndexType& type) { +NSGConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) { static int64_t MIN_SEARCH_LENGTH = 1; static int64_t MAX_SEARCH_LENGTH = 300; CheckIntByRange(knowhere::IndexParams::search_length, MIN_SEARCH_LENGTH, MAX_SEARCH_LENGTH); - return ConfAdapter::CheckSearch(oricfg, type); + return ConfAdapter::CheckSearch(oricfg, type, mode); } bool -HNSWConfAdapter::CheckTrain(milvus::json& oricfg) { +HNSWConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static int64_t MIN_EFCONSTRUCTION = 100; static int64_t MAX_EFCONSTRUCTION = 500; static int64_t MIN_M = 5; @@ -262,20 +250,20 @@ HNSWConfAdapter::CheckTrain(milvus::json& oricfg) { CheckIntByRange(knowhere::IndexParams::efConstruction, MIN_EFCONSTRUCTION, MAX_EFCONSTRUCTION); CheckIntByRange(knowhere::IndexParams::M, MIN_M, MAX_M); - return ConfAdapter::CheckTrain(oricfg); + return ConfAdapter::CheckTrain(oricfg, mode); } bool -HNSWConfAdapter::CheckSearch(milvus::json& oricfg, const IndexType& type) { +HNSWConfAdapter::CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) { static int64_t MAX_EF = 4096; CheckIntByRange(knowhere::IndexParams::ef, oricfg[knowhere::meta::TOPK], MAX_EF); - return ConfAdapter::CheckSearch(oricfg, type); + return ConfAdapter::CheckSearch(oricfg, type, mode); } bool -BinIDMAPConfAdapter::CheckTrain(milvus::json& oricfg) { +BinIDMAPConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static std::vector METRICS{knowhere::Metric::HAMMING, knowhere::Metric::JACCARD, knowhere::Metric::TANIMOTO, knowhere::Metric::SUBSTRUCTURE, knowhere::Metric::SUPERSTRUCTURE}; @@ -287,7 +275,7 @@ BinIDMAPConfAdapter::CheckTrain(milvus::json& oricfg) { } bool -BinIVFConfAdapter::CheckTrain(milvus::json& oricfg) { +BinIVFConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) { static std::vector METRICS{knowhere::Metric::HAMMING, knowhere::Metric::JACCARD, knowhere::Metric::TANIMOTO}; static int64_t MAX_NLIST = 999999; @@ -308,5 +296,6 @@ BinIVFConfAdapter::CheckTrain(milvus::json& oricfg) { return true; } -} // namespace engine + +} // namespace knowhere } // namespace milvus diff --git a/core/src/wrapper/ConfAdapter.h b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapter.h similarity index 59% rename from core/src/wrapper/ConfAdapter.h rename to core/src/index/knowhere/knowhere/index/vector_index/ConfAdapter.h index 0b2e8b30fe..a46da52cbf 100644 --- a/core/src/wrapper/ConfAdapter.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapter.h @@ -14,49 +14,41 @@ #include #include -#include "VecIndex.h" -#include "utils/Json.h" +#include "knowhere/common/Config.h" +#include "knowhere/index/vector_index/IndexType.h" namespace milvus { -namespace engine { +namespace knowhere { class ConfAdapter { public: virtual bool - CheckTrain(milvus::json& oricfg); + CheckTrain(Config& oricfg, const IndexMode mode); virtual bool - CheckSearch(milvus::json& oricfg, const IndexType& type); - - // todo(linxj): refactor in next release. - // - // virtual bool - // CheckTrain(milvus::json&, IndexMode&) = 0; - // - // virtual bool - // CheckSearch(milvus::json&, const IndexType&, IndexMode&) = 0; + CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode); }; using ConfAdapterPtr = std::shared_ptr; class IVFConfAdapter : public ConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; bool - CheckSearch(milvus::json& oricfg, const IndexType& type) override; + CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override; }; class IVFSQConfAdapter : public IVFConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; }; class IVFPQConfAdapter : public IVFConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; static void GetValidMList(int64_t dimension, std::vector& resset); @@ -65,32 +57,32 @@ class IVFPQConfAdapter : public IVFConfAdapter { class NSGConfAdapter : public IVFConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; bool - CheckSearch(milvus::json& oricfg, const IndexType& type) override; + CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override; }; class BinIDMAPConfAdapter : public ConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; }; class BinIVFConfAdapter : public IVFConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; }; class HNSWConfAdapter : public ConfAdapter { public: bool - CheckTrain(milvus::json& oricfg) override; + CheckTrain(Config& oricfg, const IndexMode mode) override; bool - CheckSearch(milvus::json& oricfg, const IndexType& type) override; + CheckSearch(Config& oricfg, const IndexType type, const IndexMode mode) override; }; -} // namespace engine +} // namespace knowhere } // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp new file mode 100644 index 0000000000..91cd5d05be --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.cpp @@ -0,0 +1,52 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License. + +#include "knowhere/index/vector_index/ConfAdapterMgr.h" + +#include "knowhere/common/Exception.h" +#include "knowhere/common/Log.h" + +namespace milvus { +namespace knowhere { + +ConfAdapterPtr +AdapterMgr::GetAdapter(const IndexType type) { + if (!init_) + RegisterAdapter(); + + try { + return table_.at(type)(); + } catch (...) { + KNOWHERE_THROW_MSG("Can not find this type of confadapter"); + } +} + +#define REGISTER_CONF_ADAPTER(T, TYPE, NAME) static AdapterMgr::register_t reg_##NAME##_(TYPE) + +void +AdapterMgr::RegisterAdapter() { + init_ = true; + + REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_FAISS_IDMAP, idmap_adapter); + REGISTER_CONF_ADAPTER(IVFConfAdapter, IndexEnum::INDEX_FAISS_IVFFLAT, ivf_adapter); + REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexEnum::INDEX_FAISS_IVFPQ, ivfpq_adapter); + REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexEnum::INDEX_FAISS_IVFSQ8, ivfsq8_adapter); + REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexEnum::INDEX_FAISS_IVFSQ8H, ivfsq8h_adapter); + REGISTER_CONF_ADAPTER(BinIDMAPConfAdapter, IndexEnum::INDEX_FAISS_BIN_IDMAP, idmap_bin_adapter); + REGISTER_CONF_ADAPTER(BinIDMAPConfAdapter, IndexEnum::INDEX_FAISS_BIN_IVFFLAT, ivf_bin_adapter); + REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexEnum::INDEX_NSG, nsg_adapter); + REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_SPTAG_KDT_RNT, sptag_kdt_adapter); + REGISTER_CONF_ADAPTER(ConfAdapter, IndexEnum::INDEX_SPTAG_BKT_RNT, sptag_bkt_adapter); + REGISTER_CONF_ADAPTER(HNSWConfAdapter, IndexEnum::INDEX_HNSW, hnsw_adapter); +} + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/wrapper/ConfAdapterMgr.h b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h similarity index 69% rename from core/src/wrapper/ConfAdapterMgr.h rename to core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h index 316ee5b29e..15ca5c49cc 100644 --- a/core/src/wrapper/ConfAdapterMgr.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/ConfAdapterMgr.h @@ -12,21 +12,21 @@ #pragma once #include -#include #include +#include -#include "ConfAdapter.h" -#include "VecIndex.h" +#include "knowhere/index/vector_index/ConfAdapter.h" +#include "knowhere/index/vector_index/IndexType.h" namespace milvus { -namespace engine { +namespace knowhere { class AdapterMgr { public: template struct register_t { - explicit register_t(const IndexType& key) { - AdapterMgr::GetInstance().table_.emplace(key, [] { return std::make_shared(); }); + explicit register_t(const IndexType type) { + AdapterMgr::GetInstance().table_[type] = ([] { return std::make_shared(); }); } }; @@ -37,15 +37,15 @@ class AdapterMgr { } ConfAdapterPtr - GetAdapter(const IndexType& indexType); + GetAdapter(const IndexType indexType); void RegisterAdapter(); protected: bool init_ = false; - std::map > table_; + std::unordered_map> table_; }; -} // namespace engine +} // namespace knowhere } // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.cpp b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.cpp index cadc119068..5915b0cd8b 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.cpp @@ -11,31 +11,24 @@ #include -#include - #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/FaissBaseBinaryIndex.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +namespace milvus { namespace knowhere { -FaissBaseBinaryIndex::FaissBaseBinaryIndex(std::shared_ptr index) : index_(std::move(index)) { -} - BinarySet -FaissBaseBinaryIndex::SerializeImpl() { +FaissBaseBinaryIndex::SerializeImpl(const IndexType& type) { try { faiss::IndexBinary* index = index_.get(); - // SealImpl(); - MemoryIOWriter writer; faiss::write_index_binary(index, &writer); auto data = std::make_shared(); data.reset(writer.data_); BinarySet res_set; - // TODO(linxj): use virtual func Name() instead of raw string. res_set.Append("BinaryIVF", data, writer.rp); return res_set; } catch (std::exception& e) { @@ -44,7 +37,7 @@ FaissBaseBinaryIndex::SerializeImpl() { } void -FaissBaseBinaryIndex::LoadImpl(const BinarySet& index_binary) { +FaissBaseBinaryIndex::LoadImpl(const BinarySet& index_binary, const IndexType& type) { auto binary = index_binary.GetByName("BinaryIVF"); MemoryIOReader reader; @@ -52,8 +45,8 @@ FaissBaseBinaryIndex::LoadImpl(const BinarySet& index_binary) { reader.data_ = binary->data.get(); faiss::IndexBinary* index = faiss::read_index_binary(&reader); - index_.reset(index); } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.h index a067a17a85..f4d8784749 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.h @@ -12,26 +12,31 @@ #pragma once #include +#include #include #include "knowhere/common/BinarySet.h" #include "knowhere/common/Dataset.h" +#include "knowhere/index/vector_index/IndexType.h" +namespace milvus { namespace knowhere { class FaissBaseBinaryIndex { protected: - explicit FaissBaseBinaryIndex(std::shared_ptr index); + explicit FaissBaseBinaryIndex(std::shared_ptr index) : index_(std::move(index)) { + } virtual BinarySet - SerializeImpl(); + SerializeImpl(const IndexType& type); virtual void - LoadImpl(const BinarySet& index_binary); + LoadImpl(const BinarySet& index_binary, const IndexType& type); public: std::shared_ptr index_ = nullptr; }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp index fbfcdf8ac9..18bff724f0 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp @@ -7,31 +7,25 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include -#include #include "knowhere/common/Exception.h" -#include "knowhere/common/Log.h" #include "knowhere/index/vector_index/FaissBaseIndex.h" -#include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/IndexType.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +namespace milvus { namespace knowhere { -FaissBaseIndex::FaissBaseIndex(std::shared_ptr index) : index_(std::move(index)) { -} - BinarySet -FaissBaseIndex::SerializeImpl() { +FaissBaseIndex::SerializeImpl(const IndexType& type) { try { fiu_do_on("FaissBaseIndex.SerializeImpl.throw_exception", throw std::exception()); faiss::Index* index = index_.get(); - // SealImpl(); - MemoryIOWriter writer; faiss::write_index(index, &writer); auto data = std::make_shared(); @@ -47,32 +41,18 @@ FaissBaseIndex::SerializeImpl() { } void -FaissBaseIndex::LoadImpl(const BinarySet& index_binary) { - auto binary = index_binary.GetByName("IVF"); +FaissBaseIndex::LoadImpl(const BinarySet& binary_set, const IndexType& type) { + auto binary = binary_set.GetByName("IVF"); MemoryIOReader reader; reader.total = binary->size; reader.data_ = binary->data.get(); faiss::Index* index = faiss::read_index(&reader); - index_.reset(index); SealImpl(); } -void -FaissBaseIndex::SealImpl() { -#ifdef CUSTOMIZATION - faiss::Index* index = index_.get(); - auto idx = dynamic_cast(index); - if (idx != nullptr) { - // To be deleted - KNOWHERE_LOG_DEBUG << "Test before to_readonly:" - << " IVF READONLY " << std::boolalpha << idx->is_readonly(); - idx->to_readonly(); - } -#endif -} - } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.h index 79cc0f65b7..0da0123877 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/FaissBaseIndex.h @@ -7,33 +7,39 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include +#include #include #include "knowhere/common/BinarySet.h" +#include "knowhere/index/vector_index/IndexType.h" +namespace milvus { namespace knowhere { class FaissBaseIndex { protected: - explicit FaissBaseIndex(std::shared_ptr index); + explicit FaissBaseIndex(std::shared_ptr index) : index_(std::move(index)) { + } virtual BinarySet - SerializeImpl(); + SerializeImpl(const IndexType& type); virtual void - LoadImpl(const BinarySet& index_binary); + LoadImpl(const BinarySet&, const IndexType& type); virtual void - SealImpl(); + SealImpl() { /* do nothing */ + } public: std::shared_ptr index_ = nullptr; }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp index bf60f2deee..7f70a5ce85 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp @@ -17,41 +17,43 @@ #include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +namespace milvus { namespace knowhere { BinarySet -BinaryIDMAP::Serialize() { +BinaryIDMAP::Serialize(const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } std::lock_guard lk(mutex_); - return SerializeImpl(); + return SerializeImpl(index_type_); } void BinaryIDMAP::Load(const BinarySet& index_binary) { std::lock_guard lk(mutex_); - LoadImpl(index_binary); + LoadImpl(index_binary, index_type_); } DatasetPtr -BinaryIDMAP::Search(const DatasetPtr& dataset, const Config& config) { +BinaryIDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } - GETBINARYTENSOR(dataset) + GETTENSOR(dataset_ptr) - auto elems = rows * config[meta::TOPK].get(); + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; size_t p_id_size = sizeof(int64_t) * elems; size_t p_dist_size = sizeof(float) * elems; auto p_id = (int64_t*)malloc(p_id_size); auto p_dist = (float*)malloc(p_dist_size); - search_impl(rows, (uint8_t*)p_data, config[meta::TOPK].get(), p_dist, p_id, Config()); + QueryImpl(rows, (uint8_t*)p_data, k, p_dist, p_id, Config()); auto ret_ds = std::make_shared(); if (index_->metric_type == faiss::METRIC_Hamming) { @@ -70,44 +72,65 @@ BinaryIDMAP::Search(const DatasetPtr& dataset, const Config& config) { return ret_ds; } -void -BinaryIDMAP::search_impl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, - const Config& cfg) { - int32_t* pdistances = (int32_t*)distances; - index_->search(n, (uint8_t*)data, k, pdistances, labels, bitset_); +DatasetPtr +BinaryIDMAP::QueryById(const DatasetPtr& dataset_ptr, const Config& config) { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); + } + + auto dim = dataset_ptr->Get(meta::DIM); + auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); + + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; + size_t p_id_size = sizeof(int64_t) * elems; + size_t p_dist_size = sizeof(float) * elems; + auto p_id = (int64_t*)malloc(p_id_size); + auto p_dist = (float*)malloc(p_dist_size); + + auto* pdistances = (int32_t*)p_dist; + index_->search_by_id(rows, p_data, k, pdistances, p_id, bitset_); + + auto ret_ds = std::make_shared(); + if (index_->metric_type == faiss::METRIC_Hamming) { + auto pf_dist = (float*)malloc(p_dist_size); + int32_t* pi_dist = (int32_t*)p_dist; + for (int i = 0; i < elems; i++) { + *(pf_dist + i) = (float)(*(pi_dist + i)); + } + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, pf_dist); + free(p_dist); + } else { + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, p_dist); + } + + return ret_ds; } void -BinaryIDMAP::Add(const DatasetPtr& dataset, const Config& config) { +BinaryIDMAP::Add(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } std::lock_guard lk(mutex_); - GETBINARYTENSOR(dataset) + GETTENSORWITHIDS(dataset_ptr) - auto p_ids = dataset->Get(meta::IDS); index_->add_with_ids(rows, (uint8_t*)p_data, p_ids); } void -BinaryIDMAP::Train(const Config& config) { - const char* type = "BFlat"; - auto index = faiss::index_binary_factory(config[meta::DIM].get(), type, - GetMetricType(config[Metric::TYPE].get())); +BinaryIDMAP::Train(const DatasetPtr& dataset_ptr, const Config& config) { + const char* desc = "BFlat"; + int64_t dim = config[meta::DIM].get(); + faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get()); + auto index = faiss::index_binary_factory(dim, desc, metric_type); index_.reset(index); } -int64_t -BinaryIDMAP::Count() { - return index_->ntotal; -} - -int64_t -BinaryIDMAP::Dimension() { - return index_->d; -} - const uint8_t* BinaryIDMAP::GetRawVectors() { try { @@ -130,18 +153,13 @@ BinaryIDMAP::GetRawIds() { } void -BinaryIDMAP::Seal() { - // do nothing -} - -void -BinaryIDMAP::AddWithoutId(const DatasetPtr& dataset, const Config& config) { +BinaryIDMAP::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } std::lock_guard lk(mutex_); - GETBINARYTENSOR(dataset) + GETTENSOR(dataset_ptr) std::vector new_ids(rows); for (int i = 0; i < rows; ++i) { @@ -152,15 +170,15 @@ BinaryIDMAP::AddWithoutId(const DatasetPtr& dataset, const Config& config) { } DatasetPtr -BinaryIDMAP::GetVectorById(const DatasetPtr& dataset, const Config& config) { +BinaryIDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } - // GETBINARYTENSOR(dataset) - // auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - auto elems = dataset->Get(meta::DIM); + // GETBINARYTENSOR(dataset_ptr) + // auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); + auto elems = dataset_ptr->Get(meta::DIM); size_t p_x_size = sizeof(uint8_t) * elems; auto p_x = (uint8_t*)malloc(p_x_size); @@ -172,51 +190,12 @@ BinaryIDMAP::GetVectorById(const DatasetPtr& dataset, const Config& config) { return ret_ds; } -DatasetPtr -BinaryIDMAP::SearchById(const DatasetPtr& dataset, const Config& config) { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - - auto dim = dataset->Get(meta::DIM); - auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - - auto elems = rows * config[meta::TOPK].get(); - size_t p_id_size = sizeof(int64_t) * elems; - size_t p_dist_size = sizeof(float) * elems; - auto p_id = (int64_t*)malloc(p_id_size); - auto p_dist = (float*)malloc(p_dist_size); - - auto* pdistances = (int32_t*)p_dist; - index_->search_by_id(rows, p_data, config[meta::TOPK].get(), pdistances, p_id, bitset_); - - auto ret_ds = std::make_shared(); - if (index_->metric_type == faiss::METRIC_Hamming) { - auto pf_dist = (float*)malloc(p_dist_size); - int32_t* pi_dist = (int32_t*)p_dist; - for (int i = 0; i < elems; i++) { - *(pf_dist + i) = (float)(*(pi_dist + i)); - } - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, pf_dist); - free(p_dist); - } else { - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, p_dist); - } - - return ret_ds; -} - void -BinaryIDMAP::SetBlacklist(faiss::ConcurrentBitsetPtr list) { - bitset_ = std::move(list); -} - -void -BinaryIDMAP::GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - list = bitset_; +BinaryIDMAP::QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, + const Config& config) { + int32_t* pdistances = (int32_t*)distances; + index_->search(n, (uint8_t*)data, k, pdistances, labels, bitset_); } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h index bdee7fed87..ab3d08b6ef 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.h @@ -11,81 +11,84 @@ #pragma once -#include #include #include #include #include -#include "FaissBaseBinaryIndex.h" -#include "VectorIndex.h" +#include "knowhere/index/vector_index/FaissBaseBinaryIndex.h" +#include "knowhere/index/vector_index/VecIndex.h" +namespace milvus { namespace knowhere { -class BinaryIDMAP : public VectorIndex, public FaissBaseBinaryIndex { +class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex { public: BinaryIDMAP() : FaissBaseBinaryIndex(nullptr) { + index_type_ = IndexEnum::INDEX_FAISS_BIN_IDMAP; } explicit BinaryIDMAP(std::shared_ptr index) : FaissBaseBinaryIndex(std::move(index)) { + index_type_ = IndexEnum::INDEX_FAISS_BIN_IDMAP; } BinarySet - Serialize() override; + Serialize(const Config& config = Config()) override; void Load(const BinarySet& index_binary) override; + void + Train(const DatasetPtr&, const Config&) override; + + void + Add(const DatasetPtr&, const Config&) override; + + void + AddWithoutIds(const DatasetPtr&, const Config&) override; + DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; + Query(const DatasetPtr&, const Config&) override; - void - Add(const DatasetPtr& dataset, const Config& config) override; - - void - AddWithoutId(const DatasetPtr& dataset, const Config& config); - - void - Train(const Config& config); + DatasetPtr + QueryById(const DatasetPtr& dataset_ptr, const Config& config) override; int64_t - Count() override; + Count() override { + return index_->ntotal; + } int64_t - Dimension() override; + Dim() override { + return index_->d; + } - void - Seal() override; + int64_t + Size() override { + if (size_ != -1) { + return size_; + } + return Count() * Dim() * sizeof(uint8_t); + } - const uint8_t* + DatasetPtr + GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) override; + + virtual const uint8_t* GetRawVectors(); - const int64_t* + virtual const int64_t* GetRawIds(); - DatasetPtr - GetVectorById(const DatasetPtr& dataset, const Config& config) override; - - DatasetPtr - SearchById(const DatasetPtr& dataset, const Config& config) override; - - void - SetBlacklist(faiss::ConcurrentBitsetPtr list); - - void - GetBlacklist(faiss::ConcurrentBitsetPtr& list); - protected: virtual void - search_impl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& cfg); + QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& config); protected: std::mutex mutex_; - - private: - faiss::ConcurrentBitsetPtr bitset_ = nullptr; }; using BinaryIDMAPPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp index f52b919e6b..f2a0ceac25 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp @@ -17,47 +17,49 @@ #include #include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Log.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +namespace milvus { namespace knowhere { using stdclock = std::chrono::high_resolution_clock; BinarySet -BinaryIVF::Serialize() { +BinaryIVF::Serialize(const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } std::lock_guard lk(mutex_); - return SerializeImpl(); + return SerializeImpl(index_type_); } void BinaryIVF::Load(const BinarySet& index_binary) { std::lock_guard lk(mutex_); - LoadImpl(index_binary); + LoadImpl(index_binary, index_type_); } DatasetPtr -BinaryIVF::Search(const DatasetPtr& dataset, const Config& config) { +BinaryIVF::Query(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - GETBINARYTENSOR(dataset) + GETTENSOR(dataset_ptr) try { - auto elems = rows * config[meta::TOPK].get(); + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; size_t p_id_size = sizeof(int64_t) * elems; size_t p_dist_size = sizeof(float) * elems; auto p_id = (int64_t*)malloc(p_id_size); auto p_dist = (float*)malloc(p_dist_size); - search_impl(rows, (uint8_t*)p_data, config[meta::TOPK].get(), p_dist, p_id, config); + QueryImpl(rows, (uint8_t*)p_data, k, p_dist, p_id, config); auto ret_ds = std::make_shared(); if (index_->metric_type == faiss::METRIC_Hamming) { @@ -81,80 +83,73 @@ BinaryIVF::Search(const DatasetPtr& dataset, const Config& config) { } } -void -BinaryIVF::search_impl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, - const Config& cfg) { - auto params = GenParams(cfg); - auto ivf_index = dynamic_cast(index_.get()); - ivf_index->nprobe = params->nprobe; - int32_t* pdistances = (int32_t*)distances; - stdclock::time_point before = stdclock::now(); - - // todo: remove static cast (zhiru) - static_cast(index_.get())->search(n, (uint8_t*)data, k, pdistances, labels, bitset_); - - stdclock::time_point after = stdclock::now(); - double search_cost = (std::chrono::duration(after - before)).count(); - KNOWHERE_LOG_DEBUG << "IVF search cost: " << search_cost - << ", quantization cost: " << faiss::indexIVF_stats.quantization_time - << ", data search cost: " << faiss::indexIVF_stats.search_time; - faiss::indexIVF_stats.quantization_time = 0; - faiss::indexIVF_stats.search_time = 0; -} - -std::shared_ptr -BinaryIVF::GenParams(const Config& config) { - auto params = std::make_shared(); - params->nprobe = config[IndexParams::nprobe]; - // params->max_codes = config["max_code"]; - return params; -} - -IndexModelPtr -BinaryIVF::Train(const DatasetPtr& dataset, const Config& config) { - GETBINARYTENSOR(dataset) - auto p_ids = dataset->Get(meta::IDS); - - faiss::IndexBinary* coarse_quantizer = - new faiss::IndexBinaryFlat(dim, GetMetricType(config[Metric::TYPE].get())); - auto index = std::make_shared(coarse_quantizer, dim, config[IndexParams::nlist], - GetMetricType(config[Metric::TYPE].get())); - index->train(rows, (uint8_t*)p_data); - index->add_with_ids(rows, (uint8_t*)p_data, p_ids); - index_ = index; - return nullptr; -} - -int64_t -BinaryIVF::Count() { - return index_->ntotal; -} - -int64_t -BinaryIVF::Dimension() { - return index_->d; -} - -void -BinaryIVF::Add(const DatasetPtr& dataset, const Config& config) { - KNOWHERE_THROW_MSG("not support yet"); -} - -void -BinaryIVF::Seal() { - // do nothing -} - DatasetPtr -BinaryIVF::GetVectorById(const DatasetPtr& dataset, const Config& config) { +BinaryIVF::QueryById(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - // GETBINARYTENSOR(dataset) - // auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - auto elems = dataset->Get(meta::DIM); + auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); + + try { + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; + + size_t p_id_size = sizeof(int64_t) * elems; + size_t p_dist_size = sizeof(float) * elems; + auto p_id = (int64_t*)malloc(p_id_size); + auto p_dist = (float*)malloc(p_dist_size); + + int32_t* pdistances = (int32_t*)p_dist; + index_->search_by_id(rows, p_data, k, pdistances, p_id, bitset_); + + auto ret_ds = std::make_shared(); + if (index_->metric_type == faiss::METRIC_Hamming) { + auto pf_dist = (float*)malloc(p_dist_size); + int32_t* pi_dist = (int32_t*)p_dist; + for (int i = 0; i < elems; i++) { + *(pf_dist + i) = (float)(*(pi_dist + i)); + } + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, pf_dist); + free(p_dist); + } else { + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, p_dist); + } + + return ret_ds; + } catch (faiss::FaissException& e) { + KNOWHERE_THROW_MSG(e.what()); + } catch (std::exception& e) { + KNOWHERE_THROW_MSG(e.what()); + } +} + +void +BinaryIVF::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSORWITHIDS(dataset_ptr) + + int64_t nlist = config[IndexParams::nlist]; + faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get()); + faiss::IndexBinary* coarse_quantizer = new faiss::IndexBinaryFlat(dim, metric_type); + auto index = std::make_shared(coarse_quantizer, dim, nlist, metric_type); + index->train(rows, (uint8_t*)p_data); + index->add_with_ids(rows, (uint8_t*)p_data, p_ids); + index_ = index; +} + +DatasetPtr +BinaryIVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) { + if (!index_ || !index_->is_trained) { + KNOWHERE_THROW_MSG("index not initialize or trained"); + } + + // GETBINARYTENSOR(dataset_ptr) + // auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); + auto elems = dataset_ptr->Get(meta::DIM); try { size_t p_x_size = sizeof(uint8_t) * elems; @@ -172,57 +167,34 @@ BinaryIVF::GetVectorById(const DatasetPtr& dataset, const Config& config) { } } -DatasetPtr -BinaryIVF::SearchById(const DatasetPtr& dataset, const Config& config) { - if (!index_ || !index_->is_trained) { - KNOWHERE_THROW_MSG("index not initialize or trained"); - } - - auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - - try { - auto elems = rows * config[meta::TOPK].get(); - - size_t p_id_size = sizeof(int64_t) * elems; - size_t p_dist_size = sizeof(float) * elems; - auto p_id = (int64_t*)malloc(p_id_size); - auto p_dist = (float*)malloc(p_dist_size); - - int32_t* pdistances = (int32_t*)p_dist; - index_->search_by_id(rows, p_data, config[meta::TOPK].get(), pdistances, p_id, bitset_); - - auto ret_ds = std::make_shared(); - if (index_->metric_type == faiss::METRIC_Hamming) { - auto pf_dist = (float*)malloc(p_dist_size); - int32_t* pi_dist = (int32_t*)p_dist; - for (int i = 0; i < elems; i++) { - *(pf_dist + i) = (float)(*(pi_dist + i)); - } - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, pf_dist); - free(p_dist); - } else { - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, p_dist); - } - - return ret_ds; - } catch (faiss::FaissException& e) { - KNOWHERE_THROW_MSG(e.what()); - } catch (std::exception& e) { - KNOWHERE_THROW_MSG(e.what()); - } +std::shared_ptr +BinaryIVF::GenParams(const Config& config) { + auto params = std::make_shared(); + params->nprobe = config[IndexParams::nprobe]; + // params->max_codes = config["max_code"]; + return params; } void -BinaryIVF::SetBlacklist(faiss::ConcurrentBitsetPtr list) { - bitset_ = std::move(list); -} +BinaryIVF::QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, + const Config& config) { + auto params = GenParams(config); + auto ivf_index = dynamic_cast(index_.get()); + ivf_index->nprobe = params->nprobe; + int32_t* pdistances = (int32_t*)distances; + stdclock::time_point before = stdclock::now(); -void -BinaryIVF::GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - list = bitset_; + // todo: remove static cast (zhiru) + static_cast(index_.get())->search(n, (uint8_t*)data, k, pdistances, labels, bitset_); + + stdclock::time_point after = stdclock::now(); + double search_cost = (std::chrono::duration(after - before)).count(); + KNOWHERE_LOG_DEBUG << "IVF search cost: " << search_cost + << ", quantization cost: " << faiss::indexIVF_stats.quantization_time + << ", data search cost: " << faiss::indexIVF_stats.search_time; + faiss::indexIVF_stats.quantization_time = 0; + faiss::indexIVF_stats.search_time = 0; } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h index b400cbd701..67848cbf0a 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h @@ -16,71 +16,88 @@ #include #include -#include -#include "FaissBaseBinaryIndex.h" -#include "VectorIndex.h" -#include "faiss/IndexIVF.h" +#include +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/FaissBaseBinaryIndex.h" +#include "knowhere/index/vector_index/VecIndex.h" + +namespace milvus { namespace knowhere { -class BinaryIVF : public VectorIndex, public FaissBaseBinaryIndex { +class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex { public: BinaryIVF() : FaissBaseBinaryIndex(nullptr) { + index_type_ = IndexEnum::INDEX_FAISS_BIN_IVFFLAT; } explicit BinaryIVF(std::shared_ptr index) : FaissBaseBinaryIndex(std::move(index)) { + index_type_ = IndexEnum::INDEX_FAISS_BIN_IVFFLAT; } BinarySet - Serialize() override; + Serialize(const Config& config = Config()) override; + + void + BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override { + Train(dataset_ptr, config); + } void Load(const BinarySet& index_binary) override; + void + Train(const DatasetPtr& dataset_ptr, const Config& config) override; + + void + Add(const DatasetPtr& dataset_ptr, const Config& config) override { + KNOWHERE_THROW_MSG("not support yet"); + } + + void + AddWithoutIds(const DatasetPtr&, const Config&) override { + KNOWHERE_THROW_MSG("AddWithoutIds is not supported"); + } + DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; + Query(const DatasetPtr& dataset_ptr, const Config& config) override; - void - Add(const DatasetPtr& dataset, const Config& config) override; - - void - Seal() override; - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; + DatasetPtr + QueryById(const DatasetPtr& dataset_ptr, const Config& config) override; int64_t - Count() override; + Count() override { + return index_->ntotal; + } int64_t - Dimension() override; + Dim() override { + return index_->d; + } + + int64_t + Size() override { + if (size_ != -1) { + return size_; + } + return Count() * Dim() * sizeof(uint8_t); + } DatasetPtr - GetVectorById(const DatasetPtr& dataset, const Config& config); - - DatasetPtr - SearchById(const DatasetPtr& dataset, const Config& config); - - void - SetBlacklist(faiss::ConcurrentBitsetPtr list); - - void - GetBlacklist(faiss::ConcurrentBitsetPtr& list); + GetVectorById(const DatasetPtr& dataset_ptr, const Config& config); protected: virtual std::shared_ptr GenParams(const Config& config); virtual void - search_impl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& cfg); + QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& config); protected: std::mutex mutex_; - - private: - faiss::ConcurrentBitsetPtr bitset_ = nullptr; }; using BinaryIVFIndexPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.h deleted file mode 100644 index 1d7789ca80..0000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.h +++ /dev/null @@ -1,86 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include -#include - -#include "IndexIVF.h" -#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" - -namespace knowhere { - -class GPUIndex { - public: - explicit GPUIndex(const int& device_id) : gpu_id_(device_id) { - } - - GPUIndex(const int& device_id, const ResPtr& resource) : gpu_id_(device_id), res_(resource) { - } - - virtual VectorIndexPtr - CopyGpuToCpu(const Config& config) = 0; - - virtual VectorIndexPtr - CopyGpuToGpu(const int64_t& device_id, const Config& config) = 0; - - void - SetGpuDevice(const int& gpu_id); - - const int64_t& - GetGpuDevice(); - - protected: - int64_t gpu_id_; - ResWPtr res_; -}; - -class GPUIVF : public IVF, public GPUIndex { - public: - explicit GPUIVF(const int& device_id) : IVF(), GPUIndex(device_id) { - } - - explicit GPUIVF(std::shared_ptr index, const int64_t& device_id, ResPtr& resource) - : IVF(std::move(index)), GPUIndex(device_id, resource) { - } - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; - - void - Add(const DatasetPtr& dataset, const Config& config) override; - - void - set_index_model(IndexModelPtr model) override; - - // DatasetPtr Search(const DatasetPtr &dataset, const Config &config) override; - VectorIndexPtr - CopyGpuToCpu(const Config& config) override; - - VectorIndexPtr - CopyGpuToGpu(const int64_t& device_id, const Config& config) override; - - // VectorIndexPtr - // Clone() final; - - protected: - void - search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override; - - BinarySet - SerializeImpl() override; - - void - LoadImpl(const BinarySet& index_binary) override; -}; - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp index 596de9e8c5..cbae93da68 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.cpp @@ -20,11 +20,12 @@ #include "hnswlib/hnswalg.h" #include "hnswlib/space_ip.h" #include "hnswlib/space_l2.h" -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Log.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +namespace milvus { namespace knowhere { void @@ -36,7 +37,7 @@ normalize_vector(float* data, float* norm_array, size_t dim) { } BinarySet -IndexHNSW::Serialize() { +IndexHNSW::Serialize(const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize or trained"); } @@ -74,12 +75,62 @@ IndexHNSW::Load(const BinarySet& index_binary) { } } +void +IndexHNSW::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) + + hnswlib::SpaceInterface* space; + if (config[Metric::TYPE] == Metric::L2) { + space = new hnswlib::L2Space(dim); + } else if (config[Metric::TYPE] == Metric::IP) { + space = new hnswlib::InnerProductSpace(dim); + normalize = true; + } + index_ = std::make_shared>(space, rows, config[IndexParams::M].get(), + config[IndexParams::efConstruction].get()); +} + +void +IndexHNSW::Add(const DatasetPtr& dataset_ptr, const Config& config) { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); + } + + std::lock_guard lk(mutex_); + + GETTENSORWITHIDS(dataset_ptr) + + // if (normalize) { + // std::vector ep_norm_vector(Dim()); + // normalize_vector((float*)(p_data), ep_norm_vector.data(), Dim()); + // index_->addPoint((void*)(ep_norm_vector.data()), p_ids[0]); + // #pragma omp parallel for + // for (int i = 1; i < rows; ++i) { + // std::vector norm_vector(Dim()); + // normalize_vector((float*)(p_data + Dim() * i), norm_vector.data(), Dim()); + // index_->addPoint((void*)(norm_vector.data()), p_ids[i]); + // } + // } else { + // index_->addPoint((void*)(p_data), p_ids[0]); + // #pragma omp parallel for + // for (int i = 1; i < rows; ++i) { + // index_->addPoint((void*)(p_data + Dim() * i), p_ids[i]); + // } + // } + + index_->addPoint(p_data, p_ids[0]); +#pragma omp parallel for + for (int i = 1; i < rows; ++i) { + index_->addPoint(((float*)p_data + Dim() * i), p_ids[i]); + } +} + DatasetPtr -IndexHNSW::Search(const DatasetPtr& dataset, const Config& config) { +IndexHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - GETTENSOR(dataset) + GETTENSOR(dataset_ptr) size_t id_size = sizeof(int64_t) * config[meta::TOPK].get(); size_t dist_size = sizeof(float) * config[meta::TOPK].get(); @@ -93,11 +144,11 @@ IndexHNSW::Search(const DatasetPtr& dataset, const Config& config) { #pragma omp parallel for for (unsigned int i = 0; i < rows; ++i) { std::vector

ret; - const float* single_query = p_data + i * Dimension(); + const float* single_query = (float*)p_data + i * Dim(); // if (normalize) { - // std::vector norm_vector(Dimension()); - // normalize_vector((float*)(single_query), norm_vector.data(), Dimension()); + // std::vector norm_vector(Dim()); + // normalize_vector((float*)(single_query), norm_vector.data(), Dim()); // ret = index_->searchKnn((float*)(norm_vector.data()), config[meta::TOPK].get(), compare); // } else { // ret = index_->searchKnn((float*)single_query, config[meta::TOPK].get(), compare); @@ -130,64 +181,6 @@ IndexHNSW::Search(const DatasetPtr& dataset, const Config& config) { return ret_ds; } -IndexModelPtr -IndexHNSW::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) - - hnswlib::SpaceInterface* space; - if (config[Metric::TYPE] == Metric::L2) { - space = new hnswlib::L2Space(dim); - } else if (config[Metric::TYPE] == Metric::IP) { - space = new hnswlib::InnerProductSpace(dim); - normalize = true; - } - index_ = std::make_shared>(space, rows, config[IndexParams::M].get(), - config[IndexParams::efConstruction].get()); - - return nullptr; -} - -void -IndexHNSW::Add(const DatasetPtr& dataset, const Config& config) { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - - std::lock_guard lk(mutex_); - - GETTENSOR(dataset) - auto p_ids = dataset->Get(meta::IDS); - - // if (normalize) { - // std::vector ep_norm_vector(Dimension()); - // normalize_vector((float*)(p_data), ep_norm_vector.data(), Dimension()); - // index_->addPoint((void*)(ep_norm_vector.data()), p_ids[0]); - // #pragma omp parallel for - // for (int i = 1; i < rows; ++i) { - // std::vector norm_vector(Dimension()); - // normalize_vector((float*)(p_data + Dimension() * i), norm_vector.data(), Dimension()); - // index_->addPoint((void*)(norm_vector.data()), p_ids[i]); - // } - // } else { - // index_->addPoint((void*)(p_data), p_ids[0]); - // #pragma omp parallel for - // for (int i = 1; i < rows; ++i) { - // index_->addPoint((void*)(p_data + Dimension() * i), p_ids[i]); - // } - // } - - index_->addPoint((void*)(p_data), p_ids[0]); -#pragma omp parallel for - for (int i = 1; i < rows; ++i) { - index_->addPoint((void*)(p_data + Dimension() * i), p_ids[i]); - } -} - -void -IndexHNSW::Seal() { - // do nothing -} - int64_t IndexHNSW::Count() { if (!index_) { @@ -197,7 +190,7 @@ IndexHNSW::Count() { } int64_t -IndexHNSW::Dimension() { +IndexHNSW::Dim() { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } @@ -205,3 +198,4 @@ IndexHNSW::Dimension() { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.h index 20ea1685d0..9a21797f0f 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexHNSW.h @@ -16,44 +16,43 @@ #include "hnswlib/hnswlib.h" -#include "knowhere/index/vector_index/VectorIndex.h" +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/VecIndex.h" +namespace milvus { namespace knowhere { -class IndexHNSW : public VectorIndex { +class IndexHNSW : public VecIndex { public: + IndexHNSW() { + index_type_ = IndexEnum::INDEX_HNSW; + } + BinarySet - Serialize() override; + Serialize(const Config& config = Config()) override; void Load(const BinarySet& index_binary) override; + void + Train(const DatasetPtr& dataset_ptr, const Config& config) override; + + void + Add(const DatasetPtr& dataset_ptr, const Config& config) override; + + void + AddWithoutIds(const DatasetPtr&, const Config&) override { + KNOWHERE_THROW_MSG("Incremental index is not supported"); + } + DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; - - // void - // set_preprocessor(PreprocessorPtr preprocessor) override; - // - // void - // set_index_model(IndexModelPtr model) override; - // - // PreprocessorPtr - // BuildPreprocessor(const DatasetPtr& dataset, const Config& config) override; - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; - - void - Add(const DatasetPtr& dataset, const Config& config) override; - - void - Seal() override; + Query(const DatasetPtr& dataset_ptr, const Config& config) override; int64_t Count() override; int64_t - Dimension() override; + Dim() override; private: bool normalize = false; @@ -62,3 +61,4 @@ class IndexHNSW : public VectorIndex { }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp index 6cac1539d2..786976b54f 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp @@ -7,7 +7,9 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include "knowhere/index/vector_index/IndexIDMAP.h" #include #include @@ -15,95 +17,72 @@ #include #include #include - #ifdef MILVUS_GPU_VERSION - #include - #endif #include #include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" - +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #ifdef MILVUS_GPU_VERSION - -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" - #endif +namespace milvus { namespace knowhere { BinarySet -IDMAP::Serialize() { +IDMAP::Serialize(const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } std::lock_guard lk(mutex_); - return SerializeImpl(); + return SerializeImpl(index_type_); } void -IDMAP::Load(const BinarySet& index_binary) { +IDMAP::Load(const BinarySet& binary_set) { std::lock_guard lk(mutex_); - LoadImpl(index_binary); -} - -DatasetPtr -IDMAP::Search(const DatasetPtr& dataset, const Config& config) { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - GETTENSOR(dataset) - - auto elems = rows * config[meta::TOPK].get(); - size_t p_id_size = sizeof(int64_t) * elems; - size_t p_dist_size = sizeof(float) * elems; - auto p_id = (int64_t*)malloc(p_id_size); - auto p_dist = (float*)malloc(p_dist_size); - - search_impl(rows, (float*)p_data, config[meta::TOPK].get(), p_dist, p_id, Config()); - - auto ret_ds = std::make_shared(); - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, p_dist); - return ret_ds; + LoadImpl(binary_set, index_type_); } void -IDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) { - index_->search(n, (float*)data, k, distances, labels, bitset_); +IDMAP::Train(const DatasetPtr& dataset_ptr, const Config& config) { + const char* desc = "IDMap,Flat"; + int64_t dim = config[meta::DIM].get(); + faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get()); + auto index = faiss::index_factory(dim, desc, metric_type); + index_.reset(index); } void -IDMAP::Add(const DatasetPtr& dataset, const Config& config) { +IDMAP::Add(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } std::lock_guard lk(mutex_); - GETTENSOR(dataset) - - auto p_ids = dataset->Get(meta::IDS); + GETTENSORWITHIDS(dataset_ptr) index_->add_with_ids(rows, (float*)p_data, p_ids); } void -IDMAP::AddWithoutId(const DatasetPtr& dataset, const Config& config) { +IDMAP::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } std::lock_guard lk(mutex_); - auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::TENSOR); + auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::TENSOR); + // TODO: caiyd need check std::vector new_ids(rows); for (int i = 0; i < rows; ++i) { new_ids[i] = i; @@ -112,14 +91,71 @@ IDMAP::AddWithoutId(const DatasetPtr& dataset, const Config& config) { index_->add_with_ids(rows, (float*)p_data, new_ids.data()); } -int64_t -IDMAP::Count() { - return index_->ntotal; +DatasetPtr +IDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config) { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); + } + GETTENSOR(dataset_ptr) + + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; + size_t p_id_size = sizeof(int64_t) * elems; + size_t p_dist_size = sizeof(float) * elems; + auto p_id = (int64_t*)malloc(p_id_size); + auto p_dist = (float*)malloc(p_dist_size); + + QueryImpl(rows, (float*)p_data, k, p_dist, p_id, Config()); + + auto ret_ds = std::make_shared(); + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, p_dist); + return ret_ds; } -int64_t -IDMAP::Dimension() { - return index_->d; +DatasetPtr +IDMAP::QueryById(const DatasetPtr& dataset_ptr, const Config& config) { + if (!index_) { + KNOWHERE_THROW_MSG("index not initialize"); + } + // GETTENSOR(dataset) + auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); + + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; + size_t p_id_size = sizeof(int64_t) * elems; + size_t p_dist_size = sizeof(float) * elems; + auto p_id = (int64_t*)malloc(p_id_size); + auto p_dist = (float*)malloc(p_dist_size); + + // todo: enable search by id (zhiru) + // auto blacklist = dataset_ptr->Get("bitset"); + // index_->searchById(rows, (float*)p_data, config[meta::TOPK].get(), p_dist, p_id, blacklist); + index_->search_by_id(rows, p_data, k, p_dist, p_id, bitset_); + + auto ret_ds = std::make_shared(); + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, p_dist); + return ret_ds; +} + +VecIndexPtr +IDMAP::CopyCpuToGpu(const int64_t device_id, const Config& config) { +#ifdef MILVUS_GPU_VERSION + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { + ResScope rs(res, device_id, false); + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); + + std::shared_ptr device_index; + device_index.reset(gpu_index); + return std::make_shared(device_index, device_id, res); + } else { + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); + } +#else + KNOWHERE_THROW_MSG("Calling IDMAP::CopyCpuToGpu when we are using CPU version"); +#endif } const float* @@ -143,57 +179,15 @@ IDMAP::GetRawIds() { } } -void -IDMAP::Train(const Config& config) { - const char* type = "IDMap,Flat"; - auto index = faiss::index_factory(config[meta::DIM].get(), type, - GetMetricType(config[Metric::TYPE].get())); - index_.reset(index); -} - -// VectorIndexPtr -// IDMAP::Clone() { -// std::lock_guard lk(mutex_); -// -// auto clone_index = faiss::clone_index(index_.get()); -// std::shared_ptr new_index; -// new_index.reset(clone_index); -// return std::make_shared(new_index); -//} - -VectorIndexPtr -IDMAP::CopyCpuToGpu(const int64_t& device_id, const Config& config) { -#ifdef MILVUS_GPU_VERSION - - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { - ResScope rs(res, device_id, false); - auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); - - std::shared_ptr device_index; - device_index.reset(gpu_index); - return std::make_shared(device_index, device_id, res); - } else { - KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); - } -#else - KNOWHERE_THROW_MSG("Calling IDMAP::CopyCpuToGpu when we are using CPU version"); -#endif -} - -void -IDMAP::Seal() { - // do nothing -} - DatasetPtr -IDMAP::GetVectorById(const DatasetPtr& dataset, const Config& config) { +IDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_) { KNOWHERE_THROW_MSG("index not initialize"); } // GETTENSOR(dataset) - // auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - auto elems = dataset->Get(meta::DIM); + // auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); + auto elems = dataset_ptr->Get(meta::DIM); size_t p_x_size = sizeof(float) * elems; auto p_x = (float*)malloc(p_x_size); @@ -205,40 +199,10 @@ IDMAP::GetVectorById(const DatasetPtr& dataset, const Config& config) { return ret_ds; } -DatasetPtr -IDMAP::SearchById(const DatasetPtr& dataset, const Config& config) { - if (!index_) { - KNOWHERE_THROW_MSG("index not initialize"); - } - // GETTENSOR(dataset) - auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - - auto elems = rows * config[meta::TOPK].get(); - size_t p_id_size = sizeof(int64_t) * elems; - size_t p_dist_size = sizeof(float) * elems; - auto p_id = (int64_t*)malloc(p_id_size); - auto p_dist = (float*)malloc(p_dist_size); - - // todo: enable search by id (zhiru) - // auto blacklist = dataset->Get("bitset"); - // index_->searchById(rows, (float*)p_data, config[meta::TOPK].get(), p_dist, p_id, blacklist); - index_->search_by_id(rows, p_data, config[meta::TOPK].get(), p_dist, p_id, bitset_); - - auto ret_ds = std::make_shared(); - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, p_dist); - return ret_ds; -} - void -IDMAP::SetBlacklist(faiss::ConcurrentBitsetPtr list) { - bitset_ = std::move(list); -} - -void -IDMAP::GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - list = bitset_; +IDMAP::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config) { + index_->search(n, (float*)data, k, distances, labels, bitset_); } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h index 64f740b7f1..6d3aa39531 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h @@ -7,58 +7,65 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once -#include "IndexIVF.h" - -#include #include #include +#include "knowhere/index/vector_index/FaissBaseIndex.h" +#include "knowhere/index/vector_index/VecIndex.h" + +namespace milvus { namespace knowhere { -class IDMAP : public VectorIndex, public FaissBaseIndex { +class IDMAP : public VecIndex, public FaissBaseIndex { public: IDMAP() : FaissBaseIndex(nullptr) { + index_type_ = IndexEnum::INDEX_FAISS_IDMAP; } explicit IDMAP(std::shared_ptr index) : FaissBaseIndex(std::move(index)) { + index_type_ = IndexEnum::INDEX_FAISS_IDMAP; } BinarySet - Serialize() override; + Serialize(const Config& config = Config()) override; void - Load(const BinarySet& index_binary) override; + Load(const BinarySet&) override; void - Train(const Config& config); + Train(const DatasetPtr&, const Config&) override; + + void + Add(const DatasetPtr&, const Config&) override; + + void + AddWithoutIds(const DatasetPtr&, const Config&) override; DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; + Query(const DatasetPtr&, const Config&) override; + + DatasetPtr + QueryById(const DatasetPtr& dataset, const Config& config) override; int64_t - Count() override; - - // VectorIndexPtr - // Clone() override; + Count() override { + return index_->ntotal; + } int64_t - Dimension() override; + Dim() override { + return index_->d; + } - void - Add(const DatasetPtr& dataset, const Config& config) override; + DatasetPtr + GetVectorById(const DatasetPtr& dataset, const Config& config) override; - void - AddWithoutId(const DatasetPtr& dataset, const Config& config); - - VectorIndexPtr - CopyCpuToGpu(const int64_t& device_id, const Config& config); - - void - Seal() override; + VecIndexPtr + CopyCpuToGpu(const int64_t, const Config&); virtual const float* GetRawVectors(); @@ -66,29 +73,15 @@ class IDMAP : public VectorIndex, public FaissBaseIndex { virtual const int64_t* GetRawIds(); - DatasetPtr - GetVectorById(const DatasetPtr& dataset, const Config& config); - - DatasetPtr - SearchById(const DatasetPtr& dataset, const Config& config); - - void - SetBlacklist(faiss::ConcurrentBitsetPtr list); - - void - GetBlacklist(faiss::ConcurrentBitsetPtr& list); - protected: virtual void - search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg); + QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&); protected: std::mutex mutex_; - - private: - faiss::ConcurrentBitsetPtr bitset_ = nullptr; }; using IDMAPPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp index 674638bcca..46db3bc2fb 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include @@ -30,91 +30,92 @@ #include #include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Log.h" -#ifdef MILVUS_GPU_VERSION -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#endif #include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" +#endif +namespace milvus { namespace knowhere { using stdclock = std::chrono::high_resolution_clock; -IndexModelPtr -IVF::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) - - faiss::Index* coarse_quantizer = new faiss::IndexFlatL2(dim); - auto index = std::make_shared(coarse_quantizer, dim, config[IndexParams::nlist].get(), - GetMetricType(config[Metric::TYPE].get())); - index->train(rows, (float*)p_data); - - // TODO(linxj): override here. train return model or not. - return std::make_shared(index); -} - -void -IVF::Add(const DatasetPtr& dataset, const Config& config) { +BinarySet +IVF::Serialize(const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } std::lock_guard lk(mutex_); - GETTENSOR(dataset) + return SerializeImpl(index_type_); +} - auto p_ids = dataset->Get(meta::IDS); +void +IVF::Load(const BinarySet& binary_set) { + std::lock_guard lk(mutex_); + LoadImpl(binary_set, index_type_); +} + +void +IVF::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) + + faiss::Index* coarse_quantizer = new faiss::IndexFlatL2(dim); + int64_t nlist = config[IndexParams::nlist].get(); + faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get()); + auto index = std::make_shared(coarse_quantizer, dim, nlist, metric_type); + index->train(rows, (float*)p_data); + + index_.reset(faiss::clone_index(index.get())); +} + +void +IVF::Add(const DatasetPtr& dataset_ptr, const Config& config) { + if (!index_ || !index_->is_trained) { + KNOWHERE_THROW_MSG("index not initialize or trained"); + } + + std::lock_guard lk(mutex_); + GETTENSORWITHIDS(dataset_ptr) index_->add_with_ids(rows, (float*)p_data, p_ids); } void -IVF::AddWithoutIds(const DatasetPtr& dataset, const Config& config) { +IVF::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } std::lock_guard lk(mutex_); - GETTENSOR(dataset) - + GETTENSOR(dataset_ptr) index_->add(rows, (float*)p_data); } -BinarySet -IVF::Serialize() { - if (!index_ || !index_->is_trained) { - KNOWHERE_THROW_MSG("index not initialize or trained"); - } - - std::lock_guard lk(mutex_); - return SerializeImpl(); -} - -void -IVF::Load(const BinarySet& index_binary) { - std::lock_guard lk(mutex_); - LoadImpl(index_binary); -} - DatasetPtr -IVF::Search(const DatasetPtr& dataset, const Config& config) { +IVF::Query(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - GETTENSOR(dataset) + GETTENSOR(dataset_ptr) try { fiu_do_on("IVF.Search.throw_std_exception", throw std::exception()); fiu_do_on("IVF.Search.throw_faiss_exception", throw faiss::FaissException("")); - auto elems = rows * config[meta::TOPK].get(); + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; size_t p_id_size = sizeof(int64_t) * elems; size_t p_dist_size = sizeof(float) * elems; auto p_id = (int64_t*)malloc(p_id_size); auto p_dist = (float*)malloc(p_dist_size); - search_impl(rows, (float*)p_data, config[meta::TOPK].get(), p_dist, p_id, config); + QueryImpl(rows, (float*)p_data, k, p_dist, p_id, config); // std::stringstream ss_res_id, ss_res_dist; // for (int i = 0; i < 10; ++i) { @@ -140,166 +141,18 @@ IVF::Search(const DatasetPtr& dataset, const Config& config) { } } -void -IVF::set_index_model(IndexModelPtr model) { - std::lock_guard lk(mutex_); - - auto rel_model = std::static_pointer_cast(model); - - // Deep copy here. - index_.reset(faiss::clone_index(rel_model->index_.get())); -} - -std::shared_ptr -IVF::GenParams(const Config& config) { - auto params = std::make_shared(); - - params->nprobe = config[IndexParams::nprobe]; - // params->max_codes = config["max_codes"]; - - return params; -} - -int64_t -IVF::Count() { - return index_->ntotal; -} - -int64_t -IVF::Dimension() { - return index_->d; -} - -void -IVF::GenGraph(const float* data, const int64_t& k, Graph& graph, const Config& config) { - int64_t K = k + 1; - auto ntotal = Count(); - - size_t dim = config[meta::DIM]; - auto batch_size = 1000; - auto tail_batch_size = ntotal % batch_size; - auto batch_search_count = ntotal / batch_size; - auto total_search_count = tail_batch_size == 0 ? batch_search_count : batch_search_count + 1; - - std::vector res_dis(K * batch_size); - graph.resize(ntotal); - Graph res_vec(total_search_count); - for (int i = 0; i < total_search_count; ++i) { - auto b_size = (i == (total_search_count - 1)) && tail_batch_size != 0 ? tail_batch_size : batch_size; - - auto& res = res_vec[i]; - res.resize(K * b_size); - - auto xq = data + batch_size * dim * i; - search_impl(b_size, (float*)xq, K, res_dis.data(), res.data(), config); - - for (int j = 0; j < b_size; ++j) { - auto& node = graph[batch_size * i + j]; - node.resize(k); - auto start_pos = j * K + 1; - for (int m = 0, cursor = start_pos; m < k && cursor < start_pos + k; ++m, ++cursor) { - node[m] = res[cursor]; - } - } - } -} - -void -IVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) { - auto params = GenParams(cfg); - auto ivf_index = dynamic_cast(index_.get()); - ivf_index->nprobe = params->nprobe; - stdclock::time_point before = stdclock::now(); - ivf_index->search(n, (float*)data, k, distances, labels, bitset_); - stdclock::time_point after = stdclock::now(); - double search_cost = (std::chrono::duration(after - before)).count(); - KNOWHERE_LOG_DEBUG << "IVF search cost: " << search_cost - << ", quantization cost: " << faiss::indexIVF_stats.quantization_time - << ", data search cost: " << faiss::indexIVF_stats.search_time; - faiss::indexIVF_stats.quantization_time = 0; - faiss::indexIVF_stats.search_time = 0; -} - -VectorIndexPtr -IVF::CopyCpuToGpu(const int64_t& device_id, const Config& config) { -#ifdef MILVUS_GPU_VERSION - - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { - ResScope rs(res, device_id, false); - auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); - - std::shared_ptr device_index; - device_index.reset(gpu_index); - return std::make_shared(device_index, device_id, res); - } else { - KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); - } - -#else - KNOWHERE_THROW_MSG("Calling IVF::CopyCpuToGpu when we are using CPU version"); -#endif -} - -// VectorIndexPtr -// IVF::Clone() { -// std::lock_guard lk(mutex_); -// -// auto clone_index = faiss::clone_index(index_.get()); -// std::shared_ptr new_index; -// new_index.reset(clone_index); -// return Clone_impl(new_index); -//} -// -// VectorIndexPtr -// IVF::Clone_impl(const std::shared_ptr& index) { -// return std::make_shared(index); -//} - -void -IVF::Seal() { - if (!index_ || !index_->is_trained) { - KNOWHERE_THROW_MSG("index not initialize or trained"); - } - SealImpl(); -} - DatasetPtr -IVF::GetVectorById(const DatasetPtr& dataset, const Config& config) { +IVF::QueryById(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - auto p_data = dataset->Get(meta::IDS); - auto elems = dataset->Get(meta::DIM); + auto rows = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); try { - size_t p_x_size = sizeof(float) * elems; - auto p_x = (float*)malloc(p_x_size); - - auto index_ivf = std::static_pointer_cast(index_); - index_ivf->get_vector_by_id(1, p_data, p_x, bitset_); - - auto ret_ds = std::make_shared(); - ret_ds->Set(meta::TENSOR, p_x); - return ret_ds; - } catch (faiss::FaissException& e) { - KNOWHERE_THROW_MSG(e.what()); - } catch (std::exception& e) { - KNOWHERE_THROW_MSG(e.what()); - } -} - -DatasetPtr -IVF::SearchById(const DatasetPtr& dataset, const Config& config) { - if (!index_ || !index_->is_trained) { - KNOWHERE_THROW_MSG("index not initialize or trained"); - } - - auto rows = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); - - try { - auto elems = rows * config[meta::TOPK].get(); + int64_t k = config[meta::TOPK].get(); + auto elems = rows * k; size_t p_id_size = sizeof(int64_t) * elems; size_t p_dist_size = sizeof(float) * elems; @@ -307,9 +160,9 @@ IVF::SearchById(const DatasetPtr& dataset, const Config& config) { auto p_dist = (float*)malloc(p_dist_size); // todo: enable search by id (zhiru) - // auto blacklist = dataset->Get("bitset"); + // auto blacklist = dataset_ptr->Get("bitset"); auto index_ivf = std::static_pointer_cast(index_); - index_ivf->search_by_id(rows, p_data, config[meta::TOPK].get(), p_dist, p_id, bitset_); + index_ivf->search_by_id(rows, p_data, k, p_dist, p_id, bitset_); // std::stringstream ss_res_id, ss_res_dist; // for (int i = 0; i < 10; ++i) { @@ -335,37 +188,129 @@ IVF::SearchById(const DatasetPtr& dataset, const Config& config) { } } -void -IVF::SetBlacklist(faiss::ConcurrentBitsetPtr list) { - bitset_ = std::move(list); -} - -void -IVF::GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - list = bitset_; -} - -IVFIndexModel::IVFIndexModel(std::shared_ptr index) : FaissBaseIndex(std::move(index)) { -} - -BinarySet -IVFIndexModel::Serialize() { +DatasetPtr +IVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { - KNOWHERE_THROW_MSG("indexmodel not initialize or trained"); + KNOWHERE_THROW_MSG("index not initialize or trained"); + } + + auto p_data = dataset_ptr->Get(meta::IDS); + auto elems = dataset_ptr->Get(meta::DIM); + + try { + size_t p_x_size = sizeof(float) * elems; + auto p_x = (float*)malloc(p_x_size); + + auto index_ivf = std::static_pointer_cast(index_); + index_ivf->get_vector_by_id(1, p_data, p_x, bitset_); + + auto ret_ds = std::make_shared(); + ret_ds->Set(meta::TENSOR, p_x); + return ret_ds; + } catch (faiss::FaissException& e) { + KNOWHERE_THROW_MSG(e.what()); + } catch (std::exception& e) { + KNOWHERE_THROW_MSG(e.what()); } - std::lock_guard lk(mutex_); - return SerializeImpl(); } void -IVFIndexModel::Load(const BinarySet& binary_set) { - std::lock_guard lk(mutex_); - LoadImpl(binary_set); +IVF::Seal() { + if (!index_ || !index_->is_trained) { + KNOWHERE_THROW_MSG("index not initialize or trained"); + } + SealImpl(); +} + +VecIndexPtr +IVF::CopyCpuToGpu(const int64_t device_id, const Config& config) { +#ifdef MILVUS_GPU_VERSION + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { + ResScope rs(res, device_id, false); + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); + + std::shared_ptr device_index; + device_index.reset(gpu_index); + return std::make_shared(device_index, device_id, res); + } else { + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); + } + +#else + KNOWHERE_THROW_MSG("Calling IVF::CopyCpuToGpu when we are using CPU version"); +#endif } void -IVFIndexModel::SealImpl() { - // do nothing +IVF::GenGraph(const float* data, const int64_t k, GraphType& graph, const Config& config) { + int64_t K = k + 1; + auto ntotal = Count(); + + size_t dim = config[meta::DIM]; + auto batch_size = 1000; + auto tail_batch_size = ntotal % batch_size; + auto batch_search_count = ntotal / batch_size; + auto total_search_count = tail_batch_size == 0 ? batch_search_count : batch_search_count + 1; + + std::vector res_dis(K * batch_size); + graph.resize(ntotal); + GraphType res_vec(total_search_count); + for (int i = 0; i < total_search_count; ++i) { + auto b_size = (i == (total_search_count - 1)) && tail_batch_size != 0 ? tail_batch_size : batch_size; + + auto& res = res_vec[i]; + res.resize(K * b_size); + + auto xq = data + batch_size * dim * i; + QueryImpl(b_size, (float*)xq, K, res_dis.data(), res.data(), config); + + for (int j = 0; j < b_size; ++j) { + auto& node = graph[batch_size * i + j]; + node.resize(k); + auto start_pos = j * K + 1; + for (int m = 0, cursor = start_pos; m < k && cursor < start_pos + k; ++m, ++cursor) { + node[m] = res[cursor]; + } + } + } +} + +std::shared_ptr +IVF::GenParams(const Config& config) { + auto params = std::make_shared(); + params->nprobe = config[IndexParams::nprobe]; + // params->max_codes = config["max_codes"]; + return params; +} + +void +IVF::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config) { + auto params = GenParams(config); + auto ivf_index = dynamic_cast(index_.get()); + ivf_index->nprobe = params->nprobe; + stdclock::time_point before = stdclock::now(); + ivf_index->search(n, (float*)data, k, distances, labels, bitset_); + stdclock::time_point after = stdclock::now(); + double search_cost = (std::chrono::duration(after - before)).count(); + KNOWHERE_LOG_DEBUG << "IVF search cost: " << search_cost + << ", quantization cost: " << faiss::indexIVF_stats.quantization_time + << ", data search cost: " << faiss::indexIVF_stats.search_time; + faiss::indexIVF_stats.quantization_time = 0; + faiss::indexIVF_stats.search_time = 0; +} + +void +IVF::SealImpl() { +#ifdef CUSTOMIZATION + faiss::Index* index = index_.get(); + auto idx = dynamic_cast(index); + if (idx != nullptr) { + // To be deleted + KNOWHERE_LOG_DEBUG << "Test before to_readonly: IVF READONLY " << std::boolalpha << idx->is_readonly(); + idx->to_readonly(); + } +#endif } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.h index bc557fddab..cb688eddbb 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.h @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once @@ -16,111 +16,75 @@ #include #include -#include "FaissBaseIndex.h" -#include "VectorIndex.h" -#include "faiss/IndexIVF.h" -#include "faiss/utils/ConcurrentBitset.h" +#include +#include "knowhere/common/Typedef.h" +#include "knowhere/index/vector_index/FaissBaseIndex.h" +#include "knowhere/index/vector_index/VecIndex.h" + +namespace milvus { namespace knowhere { -using Graph = std::vector>; - -class IVF : public VectorIndex, public FaissBaseIndex { +class IVF : public VecIndex, public FaissBaseIndex { public: IVF() : FaissBaseIndex(nullptr) { + index_type_ = IndexEnum::INDEX_FAISS_IVFFLAT; } explicit IVF(std::shared_ptr index) : FaissBaseIndex(std::move(index)) { + index_type_ = IndexEnum::INDEX_FAISS_IVFFLAT; } - // VectorIndexPtr - // Clone() override; - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; + BinarySet + Serialize(const Config& config = Config()) override; void - set_index_model(IndexModelPtr model) override; + Load(const BinarySet&) override; void - Add(const DatasetPtr& dataset, const Config& config) override; + Train(const DatasetPtr&, const Config&) override; void - AddWithoutIds(const DatasetPtr& dataset, const Config& config); + Add(const DatasetPtr&, const Config&) override; + + void + AddWithoutIds(const DatasetPtr&, const Config&) override; DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; + Query(const DatasetPtr&, const Config&) override; - void - GenGraph(const float* data, const int64_t& k, Graph& graph, const Config& config); - - BinarySet - Serialize() override; - - void - Load(const BinarySet& index_binary) override; + DatasetPtr + QueryById(const DatasetPtr& dataset, const Config& config) override; int64_t - Count() override; + Count() override { + return index_->ntotal; + } int64_t - Dimension() override; - - void - Seal() override; - - virtual VectorIndexPtr - CopyCpuToGpu(const int64_t& device_id, const Config& config); + Dim() override { + return index_->d; + } DatasetPtr GetVectorById(const DatasetPtr& dataset, const Config& config) override; - DatasetPtr - SearchById(const DatasetPtr& dataset, const Config& config) override; + virtual void + Seal(); - void - SetBlacklist(faiss::ConcurrentBitsetPtr list); + virtual VecIndexPtr + CopyCpuToGpu(const int64_t, const Config&); - void - GetBlacklist(faiss::ConcurrentBitsetPtr& list); + virtual void + GenGraph(const float* data, const int64_t k, GraphType& graph, const Config& config); protected: virtual std::shared_ptr - GenParams(const Config& config); - - // virtual VectorIndexPtr - // Clone_impl(const std::shared_ptr& index); + GenParams(const Config&); virtual void - search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg); + QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&); - protected: - std::mutex mutex_; - - private: - faiss::ConcurrentBitsetPtr bitset_ = nullptr; -}; - -using IVFIndexPtr = std::shared_ptr; - -class GPUIVF; -class IVFIndexModel : public IndexModel, public FaissBaseIndex { - friend IVF; - friend GPUIVF; - - public: - explicit IVFIndexModel(std::shared_ptr index); - - IVFIndexModel() : FaissBaseIndex(nullptr) { - } - - BinarySet - Serialize() override; - - void - Load(const BinarySet& binary) override; - - protected: void SealImpl() override; @@ -128,6 +92,7 @@ class IVFIndexModel : public IndexModel, public FaissBaseIndex { std::mutex mutex_; }; -using IVFIndexModelPtr = std::shared_ptr; +using IVFPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp index b7278107b5..c6a7e4de5a 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp @@ -7,31 +7,33 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include #include #include +#include #ifdef MILVUS_GPU_VERSION #include #endif -#include -#include -#include - -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" -#ifdef MILVUS_GPU_VERSION -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#endif +#include "knowhere/common/Log.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h" +#endif +namespace milvus { namespace knowhere { -IndexModelPtr -IVFPQ::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) +void +IVFPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, GetMetricType(config[Metric::TYPE].get())); auto index = std::make_shared(coarse_quantizer, dim, config[IndexParams::nlist].get(), @@ -39,27 +41,11 @@ IVFPQ::Train(const DatasetPtr& dataset, const Config& config) { config[IndexParams::nbits].get()); index->train(rows, (float*)p_data); - return std::make_shared(index); + index_.reset(faiss::clone_index(index.get())); } -std::shared_ptr -IVFPQ::GenParams(const Config& config) { - auto params = std::make_shared(); - params->nprobe = config[IndexParams::nprobe]; - // params->scan_table_threshold = config["scan_table_threhold"] - // params->polysemous_ht = config["polysemous_ht"] - // params->max_codes = config["max_codes"] - - return params; -} - -// VectorIndexPtr -// IVFPQ::Clone_impl(const std::shared_ptr& index) { -// return std::make_shared(index); -//} - -VectorIndexPtr -IVFPQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { +VecIndexPtr +IVFPQ::CopyCpuToGpu(const int64_t device_id, const Config& config) { #ifdef MILVUS_GPU_VERSION if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { ResScope rs(res, device_id, false); @@ -76,4 +62,16 @@ IVFPQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { #endif } +std::shared_ptr +IVFPQ::GenParams(const Config& config) { + auto params = std::make_shared(); + params->nprobe = config[IndexParams::nprobe]; + // params->scan_table_threshold = config["scan_table_threhold"] + // params->polysemous_ht = config["polysemous_ht"] + // params->max_codes = config["max_codes"] + + return params; +} + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h index 89cdb5d685..8582a6ea93 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h @@ -7,36 +7,40 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include -#include "IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVF.h" +namespace milvus { namespace knowhere { class IVFPQ : public IVF { public: - explicit IVFPQ(std::shared_ptr index) : IVF(std::move(index)) { + IVFPQ() : IVF() { + index_type_ = IndexEnum::INDEX_FAISS_IVFPQ; } - IVFPQ() = default; + explicit IVFPQ(std::shared_ptr index) : IVF(std::move(index)) { + index_type_ = IndexEnum::INDEX_FAISS_IVFPQ; + } - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; + void + Train(const DatasetPtr&, const Config&) override; - VectorIndexPtr - CopyCpuToGpu(const int64_t& device_id, const Config& config) override; + VecIndexPtr + CopyCpuToGpu(const int64_t, const Config&) override; protected: std::shared_ptr GenParams(const Config& config) override; - - // VectorIndexPtr - // Clone_impl(const std::shared_ptr& index) override; }; +using IVFPQPtr = std::shared_ptr; + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp index e37abaf321..39d26ad8da 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp @@ -7,31 +7,33 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include +#include #ifdef MILVUS_GPU_VERSION #include #include #endif +#include #include -#include -#include - -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" - +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #ifdef MILVUS_GPU_VERSION -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif +namespace milvus { namespace knowhere { -IndexModelPtr -IVFSQ::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) +void +IVFSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) std::stringstream index_type; index_type << "IVF" << config[IndexParams::nlist] << "," @@ -40,20 +42,12 @@ IVFSQ::Train(const DatasetPtr& dataset, const Config& config) { faiss::index_factory(dim, index_type.str().c_str(), GetMetricType(config[Metric::TYPE].get())); build_index->train(rows, (float*)p_data); - std::shared_ptr ret_index; - ret_index.reset(build_index); - return std::make_shared(ret_index); + index_.reset(faiss::clone_index(build_index)); } -// VectorIndexPtr -// IVFSQ::Clone_impl(const std::shared_ptr& index) { -// return std::make_shared(index); -//} - -VectorIndexPtr -IVFSQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { +VecIndexPtr +IVFSQ::CopyCpuToGpu(const int64_t device_id, const Config& config) { #ifdef MILVUS_GPU_VERSION - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { ResScope rs(res, device_id, false); @@ -65,10 +59,10 @@ IVFSQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { } else { KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); } - #else KNOWHERE_THROW_MSG("Calling IVFSQ::CopyCpuToGpu when we are using CPU version"); #endif } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.h index 5ceb8cf478..927ceb90f1 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.h @@ -7,33 +7,36 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include -#include "IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVF.h" +namespace milvus { namespace knowhere { class IVFSQ : public IVF { public: - explicit IVFSQ(std::shared_ptr index) : IVF(std::move(index)) { + IVFSQ() : IVF() { + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8; } - IVFSQ() = default; + explicit IVFSQ(std::shared_ptr index) : IVF(std::move(index)) { + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8; + } - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; + void + Train(const DatasetPtr&, const Config&) override; - VectorIndexPtr - CopyCpuToGpu(const int64_t& device_id, const Config& config) override; - - protected: - // VectorIndexPtr - // Clone_impl(const std::shared_ptr& index) override; + VecIndexPtr + CopyCpuToGpu(const int64_t, const Config&) override; }; +using IVFSQPtr = std::shared_ptr; + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp index 3f3d725bdd..9049c349d2 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp @@ -7,43 +7,42 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "knowhere/index/vector_index/IndexNSG.h" - -#include "knowhere/adapter/VectorAdapter.h" -#include "knowhere/common/Exception.h" -#include "knowhere/common/Timer.h" - -#ifdef MILVUS_GPU_VERSION - -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/helpers/Cloner.h" - -#endif +// or implied. See the License for the specific language governing permissions and limitations under the License #include +#include "knowhere/common/Exception.h" +#include "knowhere/common/Timer.h" #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexIVF.h" -#include "knowhere/index/vector_index/nsg/NSG.h" -#include "knowhere/index/vector_index/nsg/NSGIO.h" +#include "knowhere/index/vector_index/IndexNSG.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/impl/nsg/NSG.h" +#include "knowhere/index/vector_index/impl/nsg/NSGIO.h" +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" +#endif + +namespace milvus { namespace knowhere { BinarySet -NSG::Serialize() { +NSG::Serialize(const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } try { fiu_do_on("NSG.Serialize.throw_exception", throw std::exception()); - algo::NsgIndex* index = index_.get(); + std::lock_guard lk(mutex_); + impl::NsgIndex* index = index_.get(); MemoryIOWriter writer; - algo::write_index(index, writer); + impl::write_index(index, writer); auto data = std::make_shared(); data.reset(writer.data_); @@ -59,13 +58,14 @@ void NSG::Load(const BinarySet& index_binary) { try { fiu_do_on("NSG.Load.throw_exception", throw std::exception()); + std::lock_guard lk(mutex_); auto binary = index_binary.GetByName("NSG"); MemoryIOReader reader; reader.total = binary->size; reader.data_ = binary->data.get(); - auto index = algo::read_index(reader); + auto index = impl::read_index(reader); index_.reset(index); } catch (std::exception& e) { KNOWHERE_THROW_MSG(e.what()); @@ -73,73 +73,75 @@ NSG::Load(const BinarySet& index_binary) { } DatasetPtr -NSG::Search(const DatasetPtr& dataset, const Config& config) { +NSG::Query(const DatasetPtr& dataset_ptr, const Config& config) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - GETTENSOR(dataset) + GETTENSOR(dataset_ptr) - auto elems = rows * config[meta::TOPK].get(); - size_t p_id_size = sizeof(int64_t) * elems; - size_t p_dist_size = sizeof(float) * elems; - auto p_id = (int64_t*)malloc(p_id_size); - auto p_dist = (float*)malloc(p_dist_size); + try { + auto elems = rows * config[meta::TOPK].get(); + size_t p_id_size = sizeof(int64_t) * elems; + size_t p_dist_size = sizeof(float) * elems; + auto p_id = (int64_t*)malloc(p_id_size); + auto p_dist = (float*)malloc(p_dist_size); - algo::SearchParams s_params; - s_params.search_length = config[IndexParams::search_length]; - index_->Search((float*)p_data, rows, dim, config[meta::TOPK].get(), p_dist, p_id, s_params); + impl::SearchParams s_params; + s_params.search_length = config[IndexParams::search_length]; + s_params.k = config[meta::TOPK]; + { + std::lock_guard lk(mutex_); + index_->Search((float*)p_data, rows, dim, config[meta::TOPK].get(), p_dist, p_id, s_params); + } - auto ret_ds = std::make_shared(); - ret_ds->Set(meta::IDS, p_id); - ret_ds->Set(meta::DISTANCE, p_dist); - return ret_ds; + auto ret_ds = std::make_shared(); + ret_ds->Set(meta::IDS, p_id); + ret_ds->Set(meta::DISTANCE, p_dist); + return ret_ds; + } catch (std::exception& e) { + KNOWHERE_THROW_MSG(e.what()); + } } -IndexModelPtr -NSG::Train(const DatasetPtr& dataset, const Config& config) { +void +NSG::Train(const DatasetPtr& dataset_ptr, const Config& config) { auto idmap = std::make_shared(); - idmap->Train(config); - idmap->AddWithoutId(dataset, config); - Graph knng; + idmap->Train(dataset_ptr, config); + idmap->AddWithoutIds(dataset_ptr, config); + impl::Graph knng; const float* raw_data = idmap->GetRawVectors(); + const int64_t device_id = config[knowhere::meta::DEVICEID].get(); + const int64_t k = config[IndexParams::knng].get(); #ifdef MILVUS_GPU_VERSION - if (config[knowhere::meta::DEVICEID].get() == -1) { + if (device_id == -1) { auto preprocess_index = std::make_shared(); - auto model = preprocess_index->Train(dataset, config); - preprocess_index->set_index_model(model); - preprocess_index->AddWithoutIds(dataset, config); - preprocess_index->GenGraph(raw_data, config[IndexParams::knng].get(), knng, config); + preprocess_index->Train(dataset_ptr, config); + preprocess_index->AddWithoutIds(dataset_ptr, config); + preprocess_index->GenGraph(raw_data, k, knng, config); } else { - auto gpu_idx = cloner::CopyCpuToGpu(idmap, config[knowhere::meta::DEVICEID].get(), config); + auto gpu_idx = cloner::CopyCpuToGpu(idmap, device_id, config); auto gpu_idmap = std::dynamic_pointer_cast(gpu_idx); - gpu_idmap->GenGraph(raw_data, config[IndexParams::knng].get(), knng, config); + gpu_idmap->GenGraph(raw_data, k, knng, config); } #else auto preprocess_index = std::make_shared(); - auto model = preprocess_index->Train(dataset, config); - preprocess_index->set_index_model(model); - preprocess_index->AddWithoutIds(dataset, config); - preprocess_index->GenGraph(raw_data, config[IndexParams::knng].get(), knng, config); + preprocess_index->Train(dataset_ptr, config); + preprocess_index->AddWithoutIds(dataset_ptr, config); + preprocess_index->GenGraph(raw_data, k, knng, config); #endif - algo::BuildParams b_params; + impl::BuildParams b_params; b_params.candidate_pool_size = config[IndexParams::candidate]; b_params.out_degree = config[IndexParams::out_degree]; b_params.search_length = config[IndexParams::search_length]; - auto p_ids = dataset->Get(meta::IDS); + auto p_ids = dataset_ptr->Get(meta::IDS); - GETTENSOR(dataset) - index_ = std::make_shared(dim, rows); + GETTENSOR(dataset_ptr) + index_ = std::make_shared(dim, rows); index_->SetKnnGraph(knng); index_->Build_with_ids(rows, (float*)p_data, (int64_t*)p_ids, b_params); - return nullptr; // TODO(linxj): support serialize -} - -void -NSG::Add(const DatasetPtr& dataset, const Config& config) { - // do nothing } int64_t @@ -148,18 +150,9 @@ NSG::Count() { } int64_t -NSG::Dimension() { +NSG::Dim() { return index_->dimension; } -// VectorIndexPtr -// NSG::Clone() { -// KNOWHERE_THROW_MSG("not support"); -//} - -void -NSG::Seal() { - // do nothing -} - } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.h index baffc7b1bf..3c43447160 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.h @@ -7,52 +7,73 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include -#include "VectorIndex.h" +#include "knowhere/common/Exception.h" +#include "knowhere/common/Log.h" +#include "knowhere/index/vector_index/VecIndex.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { class NsgIndex; } -class NSG : public VectorIndex { +class NSG : public VecIndex { public: - explicit NSG(const int64_t& gpu_num) : gpu_(gpu_num) { + explicit NSG(const int64_t& gpu_num = -1) : gpu_(gpu_num) { + if (gpu_ >= 0) { + index_mode_ = IndexMode::MODE_GPU; + } + index_type_ = IndexEnum::INDEX_NSG; } - NSG() = default; - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; - DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; - void - Add(const DatasetPtr& dataset, const Config& config) override; BinarySet - Serialize() override; + Serialize(const Config& config = Config()) override; + void - Load(const BinarySet& index_binary) override; + Load(const BinarySet&) override; + + void + BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override { + Train(dataset_ptr, config); + } + + void + Train(const DatasetPtr&, const Config&) override; + + void + Add(const DatasetPtr&, const Config&) override { + KNOWHERE_THROW_MSG("Incremental index is not supported"); + } + + void + AddWithoutIds(const DatasetPtr&, const Config&) override { + KNOWHERE_THROW_MSG("Addwithoutids is not supported"); + } + + DatasetPtr + Query(const DatasetPtr&, const Config&) override; + int64_t Count() override; + int64_t - Dimension() override; - // VectorIndexPtr - // Clone() override; - void - Seal() override; + Dim() override; private: - std::shared_ptr index_; + std::mutex mutex_; int64_t gpu_; + std::shared_ptr index_; }; using NSGIndexPtr = std::shared_ptr(); } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp index 5066790967..51d3a1baf8 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp @@ -19,29 +19,29 @@ #undef mkdir -#include "knowhere/adapter/SptagAdapter.h" -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexSPTAG.h" -#include "knowhere/index/vector_index/helpers/Definitions.h" +#include "knowhere/index/vector_index/adapter/SptagAdapter.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h" +namespace milvus { namespace knowhere { CPUSPTAGRNG::CPUSPTAGRNG(const std::string& IndexType) { if (IndexType == "KDT") { index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float); index_ptr_->SetParameter("DistCalcMethod", "L2"); - index_type_ = SPTAG::IndexAlgoType::KDT; + index_type_ = IndexEnum::INDEX_SPTAG_KDT_RNT; } else { index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::BKT, SPTAG::VectorValueType::Float); index_ptr_->SetParameter("DistCalcMethod", "L2"); - index_type_ = SPTAG::IndexAlgoType::BKT; + index_type_ = IndexEnum::INDEX_SPTAG_BKT_RNT; } } BinarySet -CPUSPTAGRNG::Serialize() { +CPUSPTAGRNG::Serialize(const Config& config) { std::string index_config; std::vector index_blobs; @@ -72,15 +72,15 @@ CPUSPTAGRNG::Serialize() { metadata1.reset(static_cast(index_blobs[4].Data())); auto metadata2 = std::make_shared(); metadata2.reset(static_cast(index_blobs[5].Data())); - auto config = std::make_shared(); - config.reset(static_cast((void*)cstr)); + auto x_cfg = std::make_shared(); + x_cfg.reset(static_cast((void*)cstr)); binary_set.Append("samples", sample, index_blobs[0].Length()); binary_set.Append("tree", tree, index_blobs[1].Length()); binary_set.Append("deleteid", deleteid, index_blobs[3].Length()); binary_set.Append("metadata1", metadata1, index_blobs[4].Length()); binary_set.Append("metadata2", metadata2, index_blobs[5].Length()); - binary_set.Append("config", config, length); + binary_set.Append("config", x_cfg, length); binary_set.Append("graph", graph, index_blobs[2].Length()); return binary_set; @@ -115,16 +115,11 @@ CPUSPTAGRNG::Load(const BinarySet& binary_set) { index_ptr_->LoadIndex(index_config, index_blobs); } -// PreprocessorPtr -// CPUKDTRNG::BuildPreprocessor(const DatasetPtr &dataset, const Config &config) { -// return std::make_shared(); -//} - -IndexModelPtr +void CPUSPTAGRNG::Train(const DatasetPtr& origin, const Config& train_config) { SetParameters(train_config); - DatasetPtr dataset = origin; // TODO(linxj): copy or reference? + DatasetPtr dataset = origin; // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine // && preprocessor_) { @@ -134,24 +129,6 @@ CPUSPTAGRNG::Train(const DatasetPtr& origin, const Config& train_config) { auto vectorset = ConvertToVectorSet(dataset); auto metaset = ConvertToMetadataSet(dataset); index_ptr_->BuildIndex(vectorset, metaset); - - // TODO: return IndexModelPtr - return nullptr; -} - -void -CPUSPTAGRNG::Add(const DatasetPtr& origin, const Config& add_config) { - // SetParameters(add_config); - // DatasetPtr dataset = origin->Clone(); - // - // // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine - // // && preprocessor_) { - // // preprocessor_->Preprocess(dataset); - // //} - // - // auto vectorset = ConvertToVectorSet(dataset); - // auto metaset = ConvertToMetadataSet(dataset); - // index_ptr_->AddIndex(vectorset, metaset); } void @@ -159,7 +136,7 @@ CPUSPTAGRNG::SetParameters(const Config& config) { #define Assign(param_name, str_name) \ index_ptr_->SetParameter(str_name, std::to_string(build_cfg[param_name].get())) - if (index_type_ == SPTAG::IndexAlgoType::KDT) { + if (index_type_ == IndexEnum::INDEX_SPTAG_KDT_RNT) { auto build_cfg = SPTAGParameterMgr::GetInstance().GetKDTParameters(); Assign("kdtnumber", "KDTNumber"); @@ -204,17 +181,17 @@ CPUSPTAGRNG::SetParameters(const Config& config) { } DatasetPtr -CPUSPTAGRNG::Search(const DatasetPtr& dataset, const Config& config) { +CPUSPTAGRNG::Query(const DatasetPtr& dataset_ptr, const Config& config) { SetParameters(config); - auto p_data = dataset->Get(meta::TENSOR); + float* p_data = (float*)dataset_ptr->Get(meta::TENSOR); for (auto i = 0; i < 10; ++i) { for (auto j = 0; j < 10; ++j) { std::cout << p_data[i * 10 + j] << " "; } std::cout << std::endl; } - std::vector query_results = ConvertToQueryResult(dataset, config); + std::vector query_results = ConvertToQueryResult(dataset_ptr, config); #pragma omp parallel for for (auto i = 0; i < query_results.size(); ++i) { @@ -232,28 +209,24 @@ CPUSPTAGRNG::Count() { } int64_t -CPUSPTAGRNG::Dimension() { +CPUSPTAGRNG::Dim() { return index_ptr_->GetFeatureDim(); } -// VectorIndexPtr -// CPUSPTAGRNG::Clone() { -// KNOWHERE_THROW_MSG("not support"); -//} +// void +// CPUSPTAGRNG::Add(const DatasetPtr& origin, const Config& add_config) { +// SetParameters(add_config); +// DatasetPtr dataset = origin->Clone(); -void -CPUSPTAGRNG::Seal() { - return; // do nothing -} +// // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine +// // && preprocessor_) { +// // preprocessor_->Preprocess(dataset); +// //} -BinarySet -CPUSPTAGRNGIndexModel::Serialize() { - // KNOWHERE_THROW_MSG("not support"); // not support -} - -void -CPUSPTAGRNGIndexModel::Load(const BinarySet& binary) { - // KNOWHERE_THROW_MSG("not support"); // not support -} +// auto vectorset = ConvertToVectorSet(dataset); +// auto metaset = ConvertToMetadataSet(dataset); +// index_ptr_->AddIndex(vectorset, metaset); +// } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h index 0fee16ee6a..25361ae857 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h @@ -17,71 +17,58 @@ #include #include -#include "VectorIndex.h" -#include "knowhere/index/IndexModel.h" +#include "knowhere/index/vector_index/VecIndex.h" +namespace milvus { namespace knowhere { -class CPUSPTAGRNG : public VectorIndex { +class CPUSPTAGRNG : public VecIndex { public: explicit CPUSPTAGRNG(const std::string& IndexType); public: BinarySet - Serialize() override; - - // VectorIndexPtr - // Clone() override; + Serialize(const Config& config = Config()) override; void Load(const BinarySet& index_array) override; - public: - // PreprocessorPtr - // BuildPreprocessor(const DatasetPtr &dataset, const Config &config) override; + void + BuildAll(const DatasetPtr& dataset_ptr, const Config& config) override { + Train(dataset_ptr, config); + } + + void + Train(const DatasetPtr& dataset_ptr, const Config& config) override; + + void + Add(const DatasetPtr&, const Config&) override { + KNOWHERE_THROW_MSG("Incremental index is not supported"); + } + + void + AddWithoutIds(const DatasetPtr&, const Config&) override { + KNOWHERE_THROW_MSG("Incremental index is not supported"); + } + + DatasetPtr + Query(const DatasetPtr& dataset_ptr, const Config& config) override; int64_t Count() override; int64_t - Dimension() override; - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; - - void - Add(const DatasetPtr& dataset, const Config& config) override; - - DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; - - void - Seal() override; + Dim() override; private: void SetParameters(const Config& config); private: - PreprocessorPtr preprocessor_; std::shared_ptr index_ptr_; - SPTAG::IndexAlgoType index_type_; }; using CPUSPTAGRNGPtr = std::shared_ptr; -class CPUSPTAGRNGIndexModel : public IndexModel { - public: - BinarySet - Serialize() override; - - void - Load(const BinarySet& binary) override; - - private: - std::shared_ptr index_; -}; - -using CPUSPTAGRNGIndexModelPtr = std::shared_ptr; - } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexType.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexType.cpp new file mode 100644 index 0000000000..27ec3444d1 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexType.cpp @@ -0,0 +1,75 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include "knowhere/index/vector_index/IndexType.h" +#include +#include "knowhere/common/Exception.h" + +namespace milvus { +namespace knowhere { + +/* for compatible with 0.7.0 */ +static std::unordered_map old_index_type_str_map = { + {(int32_t)OldIndexType::INVALID, "INVALID"}, + {(int32_t)OldIndexType::FAISS_IDMAP, IndexEnum::INDEX_FAISS_IDMAP}, + {(int32_t)OldIndexType::FAISS_IVFFLAT_CPU, IndexEnum::INDEX_FAISS_IVFFLAT}, + {(int32_t)OldIndexType::FAISS_IVFFLAT_GPU, IndexEnum::INDEX_FAISS_IVFFLAT}, + {(int32_t)OldIndexType::FAISS_IVFFLAT_MIX, IndexEnum::INDEX_FAISS_IVFFLAT}, + {(int32_t)OldIndexType::FAISS_IVFPQ_CPU, IndexEnum::INDEX_FAISS_IVFPQ}, + {(int32_t)OldIndexType::FAISS_IVFPQ_GPU, IndexEnum::INDEX_FAISS_IVFPQ}, + {(int32_t)OldIndexType::FAISS_IVFPQ_MIX, IndexEnum::INDEX_FAISS_IVFPQ}, + {(int32_t)OldIndexType::FAISS_IVFSQ8_MIX, IndexEnum::INDEX_FAISS_IVFSQ8}, + {(int32_t)OldIndexType::FAISS_IVFSQ8_CPU, IndexEnum::INDEX_FAISS_IVFSQ8}, + {(int32_t)OldIndexType::FAISS_IVFSQ8_GPU, IndexEnum::INDEX_FAISS_IVFSQ8}, + {(int32_t)OldIndexType::FAISS_IVFSQ8_HYBRID, IndexEnum::INDEX_FAISS_IVFSQ8H}, + {(int32_t)OldIndexType::NSG_MIX, IndexEnum::INDEX_NSG}, + {(int32_t)OldIndexType::SPTAG_KDT_RNT_CPU, IndexEnum::INDEX_SPTAG_KDT_RNT}, + {(int32_t)OldIndexType::SPTAG_BKT_RNT_CPU, IndexEnum::INDEX_SPTAG_BKT_RNT}, + {(int32_t)OldIndexType::HNSW, IndexEnum::INDEX_HNSW}, + {(int32_t)OldIndexType::FAISS_BIN_IDMAP, IndexEnum::INDEX_FAISS_BIN_IDMAP}, + {(int32_t)OldIndexType::FAISS_BIN_IVFLAT_CPU, IndexEnum::INDEX_FAISS_BIN_IVFFLAT}, +}; + +static std::unordered_map str_old_index_type_map = { + {"", (int32_t)OldIndexType::INVALID}, + {IndexEnum::INDEX_FAISS_IDMAP, (int32_t)OldIndexType::FAISS_IDMAP}, + {IndexEnum::INDEX_FAISS_IVFFLAT, (int32_t)OldIndexType::FAISS_IVFFLAT_CPU}, + {IndexEnum::INDEX_FAISS_IVFPQ, (int32_t)OldIndexType::FAISS_IVFPQ_CPU}, + {IndexEnum::INDEX_FAISS_IVFSQ8, (int32_t)OldIndexType::FAISS_IVFSQ8_CPU}, + {IndexEnum::INDEX_FAISS_IVFSQ8H, (int32_t)OldIndexType::FAISS_IVFSQ8_HYBRID}, + {IndexEnum::INDEX_NSG, (int32_t)OldIndexType::NSG_MIX}, + {IndexEnum::INDEX_SPTAG_KDT_RNT, (int32_t)OldIndexType::SPTAG_KDT_RNT_CPU}, + {IndexEnum::INDEX_SPTAG_BKT_RNT, (int32_t)OldIndexType::SPTAG_BKT_RNT_CPU}, + {IndexEnum::INDEX_HNSW, (int32_t)OldIndexType::HNSW}, + {IndexEnum::INDEX_FAISS_BIN_IDMAP, (int32_t)OldIndexType::FAISS_BIN_IDMAP}, + {IndexEnum::INDEX_FAISS_BIN_IVFFLAT, (int32_t)OldIndexType::FAISS_BIN_IVFLAT_CPU}, +}; + +std::string +OldIndexTypeToStr(const int32_t type) { + try { + return old_index_type_str_map.at(type); + } catch (...) { + KNOWHERE_THROW_MSG("Invalid index type " + std::to_string(type)); + } +} + +int32_t +StrToOldIndexType(const std::string& str) { + try { + return str_old_index_type_map.at(str); + } catch (...) { + KNOWHERE_THROW_MSG("Invalid index str " + str); + } +} + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexType.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexType.h new file mode 100644 index 0000000000..1f2b26866d --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexType.h @@ -0,0 +1,67 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include + +namespace milvus { +namespace knowhere { + +/* used in 0.7.0 */ +enum class OldIndexType { + INVALID = 0, + FAISS_IDMAP = 1, + FAISS_IVFFLAT_CPU, + FAISS_IVFFLAT_GPU, + FAISS_IVFFLAT_MIX, // build on gpu and search on cpu + FAISS_IVFPQ_CPU, + FAISS_IVFPQ_GPU, + SPTAG_KDT_RNT_CPU, + FAISS_IVFSQ8_MIX, + FAISS_IVFSQ8_CPU, + FAISS_IVFSQ8_GPU, + FAISS_IVFSQ8_HYBRID, // only support build on gpu. + NSG_MIX, + FAISS_IVFPQ_MIX, + SPTAG_BKT_RNT_CPU, + HNSW, + FAISS_BIN_IDMAP = 100, + FAISS_BIN_IVFLAT_CPU = 101, +}; + +using IndexType = std::string; + +/* used in 0.8.0 */ +namespace IndexEnum { +constexpr const char* INVALID = ""; +constexpr const char* INDEX_FAISS_IDMAP = "IDMAP"; +constexpr const char* INDEX_FAISS_IVFFLAT = "IVF_FLAT"; +constexpr const char* INDEX_FAISS_IVFPQ = "IVF_PQ"; +constexpr const char* INDEX_FAISS_IVFSQ8 = "IVF_SQ8"; +constexpr const char* INDEX_FAISS_IVFSQ8H = "IVF_SQ8_HYBRID"; +constexpr const char* INDEX_FAISS_BIN_IDMAP = "BIN_IDMAP"; +constexpr const char* INDEX_FAISS_BIN_IVFFLAT = "BIN_IVF_FLAT"; +constexpr const char* INDEX_NSG = "NSG"; +constexpr const char* INDEX_SPTAG_KDT_RNT = "SPTAG_KDT_RNT"; +constexpr const char* INDEX_SPTAG_BKT_RNT = "SPTAG_BKT_RNT"; +constexpr const char* INDEX_HNSW = "HNSW"; +} // namespace IndexEnum + +enum class IndexMode { MODE_CPU = 0, MODE_GPU = 1 }; + +extern std::string +OldIndexTypeToStr(const int32_t type); + +extern int32_t +StrToOldIndexType(const std::string& str); +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h new file mode 100644 index 0000000000..3f2d4b7500 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/VecIndex.h @@ -0,0 +1,121 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include +#include +#include + +#include "knowhere/common/Dataset.h" +#include "knowhere/common/Typedef.h" +#include "knowhere/index/Index.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "segment/Types.h" + +namespace milvus { +namespace knowhere { + +class VecIndex : public Index { + public: + virtual void + BuildAll(const DatasetPtr& dataset_ptr, const Config& config) { + Train(dataset_ptr, config); + Add(dataset_ptr, config); + } + + virtual void + Train(const DatasetPtr& dataset, const Config& config) = 0; + + virtual void + Add(const DatasetPtr& dataset, const Config& config) = 0; + + virtual void + AddWithoutIds(const DatasetPtr& dataset, const Config& config) = 0; + + virtual DatasetPtr + Query(const DatasetPtr& dataset, const Config& config) = 0; + + virtual DatasetPtr + QueryById(const DatasetPtr& dataset, const Config& config) { + return nullptr; + } + + // virtual DatasetPtr + // QueryByRange(const DatasetPtr&, const Config&) = 0; + // + // virtual MetricType + // metric_type() = 0; + + virtual int64_t + Dim() = 0; + + virtual int64_t + Count() = 0; + + virtual IndexType + index_type() const { + return index_type_; + } + + virtual IndexMode + index_mode() const { + return index_mode_; + } + + virtual DatasetPtr + GetVectorById(const DatasetPtr& dataset, const Config& config) { + return nullptr; + } + + virtual void + GetBlacklist(faiss::ConcurrentBitsetPtr& bitset_ptr) { + bitset_ptr = bitset_; + } + + virtual void + SetBlacklist(faiss::ConcurrentBitsetPtr bitset_ptr) { + bitset_ = std::move(bitset_ptr); + } + + virtual const std::vector& + GetUids() const { + return uids_; + } + + virtual void + SetUids(std::vector& uids) { + uids_.clear(); + uids_.swap(uids); + } + + int64_t + Size() override { + if (size_ != -1) { + return size_; + } + return Count() * Dim() * sizeof(FloatType); + } + + protected: + IndexType index_type_ = ""; + IndexMode index_mode_ = IndexMode::MODE_CPU; + faiss::ConcurrentBitsetPtr bitset_ = nullptr; + + private: + std::vector uids_; +}; + +using VecIndexPtr = std::shared_ptr; + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/VecIndexFactory.cpp b/core/src/index/knowhere/knowhere/index/vector_index/VecIndexFactory.cpp new file mode 100644 index 0000000000..5a512870ed --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/VecIndexFactory.cpp @@ -0,0 +1,87 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include "knowhere/index/vector_index/VecIndexFactory.h" + +#include "knowhere/common/Exception.h" +#include "knowhere/common/Log.h" +#include "knowhere/index/vector_index/IndexBinaryIDMAP.h" +#include "knowhere/index/vector_index/IndexBinaryIVF.h" +#include "knowhere/index/vector_index/IndexHNSW.h" +#include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexNSG.h" +#include "knowhere/index/vector_index/IndexSPTAG.h" + +#ifdef MILVUS_GPU_VERSION +#include +#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/helpers/Cloner.h" +#endif + +namespace milvus { +namespace knowhere { + +VecIndexPtr +VecIndexFactory::CreateVecIndex(const IndexType& type, const IndexMode mode) { + auto gpu_device = -1; // TODO: remove hardcode here, get from invoker + if (type == IndexEnum::INDEX_FAISS_IDMAP) { + return std::make_shared(); + } else if (type == IndexEnum::INDEX_FAISS_IVFFLAT) { +#ifdef MILVUS_GPU_VERSION + if (mode == IndexMode::MODE_GPU) { + return std::make_shared(gpu_device); + } +#endif + return std::make_shared(); + } else if (type == IndexEnum::INDEX_FAISS_IVFPQ) { +#ifdef MILVUS_GPU_VERSION + if (mode == IndexMode::MODE_GPU) { + return std::make_shared(gpu_device); + } +#endif + return std::make_shared(); + } else if (type == IndexEnum::INDEX_FAISS_IVFSQ8) { +#ifdef MILVUS_GPU_VERSION + if (mode == IndexMode::MODE_GPU) { + return std::make_shared(gpu_device); + } +#endif + return std::make_shared(); +#ifdef MILVUS_GPU_VERSION + } else if (type == IndexEnum::INDEX_FAISS_IVFSQ8H) { + return std::make_shared(gpu_device); +#endif + } else if (type == IndexEnum::INDEX_FAISS_BIN_IDMAP) { + return std::make_shared(); + } else if (type == IndexEnum::INDEX_FAISS_BIN_IVFFLAT) { + return std::make_shared(); + } else if (type == IndexEnum::INDEX_NSG) { + return std::make_shared(-1); + } else if (type == IndexEnum::INDEX_SPTAG_KDT_RNT) { + return std::make_shared("KDT"); + } else if (type == IndexEnum::INDEX_SPTAG_BKT_RNT) { + return std::make_shared("BKT"); + } else if (type == IndexEnum::INDEX_HNSW) { + return std::make_shared(); + } else { + return nullptr; + } +} + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/IndexModel.h b/core/src/index/knowhere/knowhere/index/vector_index/VecIndexFactory.h similarity index 51% rename from core/src/index/knowhere/knowhere/index/IndexModel.h rename to core/src/index/knowhere/knowhere/index/vector_index/VecIndexFactory.h index d28de3d015..7a09b54bff 100644 --- a/core/src/index/knowhere/knowhere/index/IndexModel.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/VecIndexFactory.h @@ -7,24 +7,35 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include -#include "knowhere/common/BinarySet.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/VecIndex.h" + +namespace milvus { namespace knowhere { -class IndexModel { - public: - virtual BinarySet - Serialize() = 0; +class VecIndexFactory { + private: + VecIndexFactory() = default; + VecIndexFactory(const VecIndexFactory&) = delete; + VecIndexFactory + operator=(const VecIndexFactory&) = delete; - virtual void - Load(const BinarySet& binary) = 0; + public: + static VecIndexFactory& + GetInstance() { + static VecIndexFactory inst; + return inst; + } + + knowhere::VecIndexPtr + CreateVecIndex(const IndexType& type, const IndexMode mode = IndexMode::MODE_CPU); }; -using IndexModelPtr = std::shared_ptr; - } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/VectorIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/VectorIndex.h deleted file mode 100644 index 5e2f4c514c..0000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/VectorIndex.h +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include -#include - -#include "knowhere/common/Config.h" -#include "knowhere/common/Dataset.h" -#include "knowhere/index/Index.h" -#include "knowhere/index/preprocessor/Preprocessor.h" -#include "knowhere/index/vector_index/helpers/IndexParameter.h" -#include "segment/Types.h" - -namespace knowhere { - -class VectorIndex; -using VectorIndexPtr = std::shared_ptr; - -class VectorIndex : public Index { - public: - virtual PreprocessorPtr - BuildPreprocessor(const DatasetPtr& dataset, const Config& config) { - return nullptr; - } - - virtual IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) { - return nullptr; - } - - virtual DatasetPtr - GetVectorById(const DatasetPtr& dataset, const Config& config) { - return nullptr; - } - - virtual DatasetPtr - SearchById(const DatasetPtr& dataset, const Config& config) { - return nullptr; - } - - virtual void - Add(const DatasetPtr& dataset, const Config& config) = 0; - - virtual void - Seal() = 0; - - // TODO(linxj): Deprecated - // virtual VectorIndexPtr - // Clone() = 0; - - virtual int64_t - Count() = 0; - - virtual int64_t - Dimension() = 0; - - virtual const std::vector& - GetUids() const { - return uids_; - } - - virtual void - SetUids(std::vector& uids) { - uids_.clear(); - uids_.swap(uids); - } - - private: - std::vector uids_; -}; - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/adapter/SptagAdapter.cpp b/core/src/index/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.cpp similarity index 77% rename from core/src/index/knowhere/knowhere/adapter/SptagAdapter.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.cpp index 9a12093dc4..030aa8665f 100644 --- a/core/src/index/knowhere/knowhere/adapter/SptagAdapter.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.cpp @@ -9,16 +9,16 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. -#include "knowhere/adapter/SptagAdapter.h" - -#include "VectorAdapter.h" +#include "knowhere/index/vector_index/adapter/SptagAdapter.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +namespace milvus { namespace knowhere { std::shared_ptr -ConvertToMetadataSet(const DatasetPtr& dataset) { - auto elems = dataset->Get(meta::ROWS); - auto p_data = dataset->Get(meta::IDS); +ConvertToMetadataSet(const DatasetPtr& dataset_ptr) { + auto elems = dataset_ptr->Get(meta::ROWS); + auto p_data = dataset_ptr->Get(meta::IDS); auto p_offset = (int64_t*)malloc(sizeof(int64_t) * (elems + 1)); for (auto i = 0; i <= elems; ++i) p_offset[i] = i * 8; @@ -31,8 +31,8 @@ ConvertToMetadataSet(const DatasetPtr& dataset) { } std::shared_ptr -ConvertToVectorSet(const DatasetPtr& dataset) { - GETTENSOR(dataset); +ConvertToVectorSet(const DatasetPtr& dataset_ptr) { + GETTENSOR(dataset_ptr); size_t num_bytes = rows * dim * sizeof(float); SPTAG::ByteArray byte_array((uint8_t*)p_data, num_bytes, false); @@ -41,13 +41,13 @@ ConvertToVectorSet(const DatasetPtr& dataset) { } std::vector -ConvertToQueryResult(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset); +ConvertToQueryResult(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr); - std::vector query_results(rows, - SPTAG::QueryResult(nullptr, config[meta::TOPK].get(), true)); + int64_t k = config[meta::TOPK].get(); + std::vector query_results(rows, SPTAG::QueryResult(nullptr, k, true)); for (auto i = 0; i < rows; ++i) { - query_results[i].SetTarget(&p_data[i * dim]); + query_results[i].SetTarget((float*)p_data + i * dim); } return query_results; @@ -81,3 +81,4 @@ ConvertToDataset(std::vector query_results) { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/adapter/SptagAdapter.h b/core/src/index/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.h similarity index 82% rename from core/src/index/knowhere/knowhere/adapter/SptagAdapter.h rename to core/src/index/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.h index 23663296ac..a8ff4eaf58 100644 --- a/core/src/index/knowhere/knowhere/adapter/SptagAdapter.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.h @@ -18,18 +18,20 @@ #include "knowhere/common/Config.h" #include "knowhere/common/Dataset.h" +namespace milvus { namespace knowhere { std::shared_ptr -ConvertToVectorSet(const DatasetPtr& dataset); +ConvertToVectorSet(const DatasetPtr& dataset_ptr); std::shared_ptr -ConvertToMetadataSet(const DatasetPtr& dataset); +ConvertToMetadataSet(const DatasetPtr& dataset_ptr); std::vector -ConvertToQueryResult(const DatasetPtr& dataset, const Config& config); +ConvertToQueryResult(const DatasetPtr& dataset_ptr, const Config& config); DatasetPtr ConvertToDataset(std::vector query_results); } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/adapter/VectorAdapter.h b/core/src/index/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.cpp similarity index 52% rename from core/src/index/knowhere/knowhere/adapter/VectorAdapter.h rename to core/src/index/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.cpp index ae89038855..a9ba05298a 100644 --- a/core/src/index/knowhere/knowhere/adapter/VectorAdapter.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.cpp @@ -9,22 +9,33 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. -#pragma once +#include -#include #include "knowhere/common/Dataset.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "knowhere/index/vector_index/helpers/IndexParameter.h" +namespace milvus { namespace knowhere { -#define GETTENSOR(dataset) \ - auto dim = dataset->Get(meta::DIM); \ - auto rows = dataset->Get(meta::ROWS); \ - auto p_data = dataset->Get(meta::TENSOR); +DatasetPtr +GenDatasetWithIds(const int64_t nb, const int64_t dim, const void* xb, const int64_t* ids) { + auto ret_ds = std::make_shared(); + ret_ds->Set(meta::ROWS, nb); + ret_ds->Set(meta::DIM, dim); + ret_ds->Set(meta::TENSOR, xb); + ret_ds->Set(meta::IDS, ids); + return ret_ds; +} -#define GETBINARYTENSOR(dataset) \ - auto dim = dataset->Get(meta::DIM); \ - auto rows = dataset->Get(meta::ROWS); \ - auto p_data = dataset->Get(meta::TENSOR); +DatasetPtr +GenDataset(const int64_t nb, const int64_t dim, const void* xb) { + auto ret_ds = std::make_shared(); + ret_ds->Set(meta::ROWS, nb); + ret_ds->Set(meta::DIM, dim); + ret_ds->Set(meta::TENSOR, xb); + return ret_ds; +} } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.h b/core/src/index/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.h new file mode 100644 index 0000000000..44086796b5 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/adapter/VectorAdapter.h @@ -0,0 +1,39 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License. + +#pragma once + +#include +#include "knowhere/common/Dataset.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" + +namespace milvus { +namespace knowhere { + +#define GETTENSOR(dataset_ptr) \ + int64_t dim = dataset_ptr->Get(meta::DIM); \ + int64_t rows = dataset_ptr->Get(meta::ROWS); \ + const void* p_data = dataset_ptr->Get(meta::TENSOR); + +#define GETTENSORWITHIDS(dataset_ptr) \ + int64_t dim = dataset_ptr->Get(meta::DIM); \ + int64_t rows = dataset_ptr->Get(meta::ROWS); \ + const void* p_data = dataset_ptr->Get(meta::TENSOR); \ + const int64_t* p_ids = dataset_ptr->Get(meta::IDS); + +extern DatasetPtr +GenDatasetWithIds(const int64_t nb, const int64_t dim, const void* xb, const int64_t* ids); + +extern DatasetPtr +GenDataset(const int64_t nb, const int64_t dim, const void* xb); + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/gpu/GPUIndex.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/GPUIndex.h new file mode 100644 index 0000000000..1e079efd7c --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/GPUIndex.h @@ -0,0 +1,50 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include "knowhere/index/vector_index/VecIndex.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" + +namespace milvus { +namespace knowhere { + +class GPUIndex { + public: + explicit GPUIndex(const int& device_id) : gpu_id_(device_id) { + } + + GPUIndex(const int& device_id, const ResPtr& resource) : gpu_id_(device_id), res_(resource) { + } + + virtual VecIndexPtr + CopyGpuToCpu(const Config&) = 0; + + virtual VecIndexPtr + CopyGpuToGpu(const int64_t, const Config&) = 0; + + void + SetGpuDevice(const int& gpu_id) { + gpu_id_ = gpu_id; + } + + const int64_t + GetGpuDevice() { + return gpu_id_; + } + + protected: + int64_t gpu_id_; + ResWPtr res_; +}; + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp similarity index 80% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp index d659f8fef4..c976520527 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp @@ -7,31 +7,28 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include #include #include -#include -#include - #ifdef MILVUS_GPU_VERSION - #include - #endif +#include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +namespace milvus { namespace knowhere { -VectorIndexPtr +VecIndexPtr GPUIDMAP::CopyGpuToCpu(const Config& config) { std::lock_guard lk(mutex_); @@ -43,19 +40,8 @@ GPUIDMAP::CopyGpuToCpu(const Config& config) { return std::make_shared(new_index); } -// VectorIndexPtr -// GPUIDMAP::Clone() { -// auto cpu_idx = CopyGpuToCpu(Config()); -// -// if (auto idmap = std::dynamic_pointer_cast(cpu_idx)) { -// return idmap->CopyCpuToGpu(gpu_id_, Config()); -// } else { -// KNOWHERE_THROW_MSG("IndexType not Support GpuClone"); -// } -//} - BinarySet -GPUIDMAP::SerializeImpl() { +GPUIDMAP::SerializeImpl(const IndexType& type) { try { fiu_do_on("GPUIDMP.SerializeImpl.throw_exception", throw std::exception()); MemoryIOWriter writer; @@ -79,7 +65,7 @@ GPUIDMAP::SerializeImpl() { } void -GPUIDMAP::LoadImpl(const BinarySet& index_binary) { +GPUIDMAP::LoadImpl(const BinarySet& index_binary, const IndexType& type) { auto binary = index_binary.GetByName("IVF"); MemoryIOReader reader; { @@ -101,30 +87,30 @@ GPUIDMAP::LoadImpl(const BinarySet& index_binary) { } } -VectorIndexPtr -GPUIDMAP::CopyGpuToGpu(const int64_t& device_id, const Config& config) { +VecIndexPtr +GPUIDMAP::CopyGpuToGpu(const int64_t device_id, const Config& config) { auto cpu_index = CopyGpuToCpu(config); return std::static_pointer_cast(cpu_index)->CopyCpuToGpu(device_id, config); } -float* +const float* GPUIDMAP::GetRawVectors() { KNOWHERE_THROW_MSG("Not support"); } -int64_t* +const int64_t* GPUIDMAP::GetRawIds() { KNOWHERE_THROW_MSG("Not support"); } void -GPUIDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) { +GPUIDMAP::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config) { ResScope rs(res_, gpu_id_); index_->search(n, (float*)data, k, distances, labels); } void -GPUIDMAP::GenGraph(const float* data, const int64_t& k, Graph& graph, const Config& config) { +GPUIDMAP::GenGraph(const float* data, const int64_t k, GraphType& graph, const Config& config) { int64_t K = k + 1; auto ntotal = Count(); @@ -144,7 +130,7 @@ GPUIDMAP::GenGraph(const float* data, const int64_t& k, Graph& graph, const Conf res.resize(K * b_size); auto xq = data + batch_size * dim * i; - search_impl(b_size, (float*)xq, K, res_dis.data(), res.data(), config); + QueryImpl(b_size, (float*)xq, K, res_dis.data(), res.data(), config); for (int j = 0; j < b_size; ++j) { auto& node = graph[batch_size * i + j]; @@ -158,3 +144,4 @@ GPUIDMAP::GenGraph(const float* data, const int64_t& k, Graph& graph, const Conf } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.h similarity index 59% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.h index c0121f67a4..af3462fbb0 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.h @@ -7,55 +7,56 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once -#include "IndexGPUIVF.h" -#include "IndexIDMAP.h" -#include "IndexIVF.h" - #include #include #include +#include "knowhere/index/vector_index/IndexIDMAP.h" +#include "knowhere/index/vector_index/gpu/GPUIndex.h" + +namespace milvus { namespace knowhere { +using Graph = std::vector>; + class GPUIDMAP : public IDMAP, public GPUIndex { public: explicit GPUIDMAP(std::shared_ptr index, const int64_t& device_id, ResPtr& res) : IDMAP(std::move(index)), GPUIndex(device_id, res) { + index_mode_ = IndexMode::MODE_GPU; } - VectorIndexPtr - CopyGpuToCpu(const Config& config) override; + VecIndexPtr + CopyGpuToCpu(const Config&) override; - float* + VecIndexPtr + CopyGpuToGpu(const int64_t, const Config&) override; + + const float* GetRawVectors() override; - int64_t* + const int64_t* GetRawIds() override; - // VectorIndexPtr - // Clone() override; - - VectorIndexPtr - CopyGpuToGpu(const int64_t& device_id, const Config& config) override; - void - GenGraph(const float* data, const int64_t& k, Graph& graph, const Config& config); + GenGraph(const float*, const int64_t, GraphType&, const Config&); protected: - void - search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override; - BinarySet - SerializeImpl() override; + SerializeImpl(const IndexType&) override; void - LoadImpl(const BinarySet& index_binary) override; + LoadImpl(const BinarySet&, const IndexType&) override; + + void + QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&) override; }; using GPUIDMAPPtr = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp similarity index 65% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp index 1feb08e4f1..19f8229f6c 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp @@ -7,66 +7,87 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include -#include #include #include #include #include #include +#include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" #include "knowhere/index/vector_index/helpers/Cloner.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +namespace milvus { namespace knowhere { -IndexModelPtr -GPUIVF::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) +void +GPUIVF::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) gpu_id_ = config[knowhere::meta::DEVICEID]; - auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); - if (temp_resource != nullptr) { - ResScope rs(temp_resource, gpu_id_, true); + auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (gpu_res != nullptr) { + ResScope rs(gpu_res, gpu_id_, true); faiss::gpu::GpuIndexIVFFlatConfig idx_config; idx_config.device = gpu_id_; - faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim, config[IndexParams::nlist], - GetMetricType(config[Metric::TYPE].get()), idx_config); + int32_t nlist = config[IndexParams::nlist]; + faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get()); + faiss::gpu::GpuIndexIVFFlat device_index(gpu_res->faiss_res.get(), dim, nlist, metric_type, idx_config); device_index.train(rows, (float*)p_data); std::shared_ptr host_index = nullptr; host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index)); - return std::make_shared(host_index); + auto device_index1 = faiss::gpu::index_cpu_to_gpu(gpu_res->faiss_res.get(), gpu_id_, host_index.get()); + index_.reset(device_index1); + res_ = gpu_res; } else { KNOWHERE_THROW_MSG("Build IVF can't get gpu resource"); } } void -GPUIVF::set_index_model(IndexModelPtr model) { - std::lock_guard lk(mutex_); - - auto host_index = std::static_pointer_cast(model); - if (auto gpures = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) { - ResScope rs(gpures, gpu_id_, false); - auto device_index = faiss::gpu::index_cpu_to_gpu(gpures->faiss_res.get(), gpu_id_, host_index->index_.get()); - index_.reset(device_index); - res_ = gpures; +GPUIVF::Add(const DatasetPtr& dataset_ptr, const Config& config) { + if (auto spt = res_.lock()) { + ResScope rs(res_, gpu_id_); + IVF::Add(dataset_ptr, config); } else { - KNOWHERE_THROW_MSG("load index model error, can't get gpu_resource"); + KNOWHERE_THROW_MSG("Add IVF can't get gpu resource"); } } +VecIndexPtr +GPUIVF::CopyGpuToCpu(const Config& config) { + std::lock_guard lk(mutex_); + + if (auto device_idx = std::dynamic_pointer_cast(index_)) { + faiss::Index* device_index = index_.get(); + faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); + } else { + return std::make_shared(index_); + } +} + +VecIndexPtr +GPUIVF::CopyGpuToGpu(const int64_t device_id, const Config& config) { + auto host_index = CopyGpuToCpu(config); + return std::static_pointer_cast(host_index)->CopyCpuToGpu(device_id, config); +} + BinarySet -GPUIVF::SerializeImpl() { +GPUIVF::SerializeImpl(const IndexType& type) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } @@ -94,8 +115,8 @@ GPUIVF::SerializeImpl() { } void -GPUIVF::LoadImpl(const BinarySet& index_binary) { - auto binary = index_binary.GetByName("IVF"); +GPUIVF::LoadImpl(const BinarySet& binary_set, const IndexType& type) { + auto binary = binary_set.GetByName("IVF"); MemoryIOReader reader; { reader.total = binary->size; @@ -117,7 +138,7 @@ GPUIVF::LoadImpl(const BinarySet& index_binary) { } void -GPUIVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config) { +GPUIVF::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config) { std::lock_guard lk(mutex_); auto device_index = std::dynamic_pointer_cast(index_); @@ -131,52 +152,5 @@ GPUIVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, i } } -VectorIndexPtr -GPUIVF::CopyGpuToCpu(const Config& config) { - std::lock_guard lk(mutex_); - - if (auto device_idx = std::dynamic_pointer_cast(index_)) { - faiss::Index* device_index = index_.get(); - faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); - - std::shared_ptr new_index; - new_index.reset(host_index); - return std::make_shared(new_index); - } else { - return std::make_shared(index_); - } -} - -// VectorIndexPtr -// GPUIVF::Clone() { -// auto cpu_idx = CopyGpuToCpu(Config()); -// return knowhere::cloner::CopyCpuToGpu(cpu_idx, gpu_id_, Config()); -//} - -VectorIndexPtr -GPUIVF::CopyGpuToGpu(const int64_t& device_id, const Config& config) { - auto host_index = CopyGpuToCpu(config); - return std::static_pointer_cast(host_index)->CopyCpuToGpu(device_id, config); -} - -void -GPUIVF::Add(const DatasetPtr& dataset, const Config& config) { - if (auto spt = res_.lock()) { - ResScope rs(res_, gpu_id_); - IVF::Add(dataset, config); - } else { - KNOWHERE_THROW_MSG("Add IVF can't get gpu resource"); - } -} - -void -GPUIndex::SetGpuDevice(const int& gpu_id) { - gpu_id_ = gpu_id; -} - -const int64_t& -GPUIndex::GetGpuDevice() { - return gpu_id_; -} - } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.h new file mode 100644 index 0000000000..49d1b3eef0 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.h @@ -0,0 +1,60 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License + +#pragma once + +#include +#include + +#include "knowhere/index/vector_index/IndexIVF.h" +#include "knowhere/index/vector_index/gpu/GPUIndex.h" + +namespace milvus { +namespace knowhere { + +class GPUIVF : public IVF, public GPUIndex { + public: + explicit GPUIVF(const int& device_id) : IVF(), GPUIndex(device_id) { + index_mode_ = IndexMode::MODE_GPU; + } + + explicit GPUIVF(std::shared_ptr index, const int64_t device_id, ResPtr& res) + : IVF(std::move(index)), GPUIndex(device_id, res) { + index_mode_ = IndexMode::MODE_GPU; + } + + void + Train(const DatasetPtr&, const Config&) override; + + void + Add(const DatasetPtr&, const Config&) override; + + VecIndexPtr + CopyGpuToCpu(const Config&) override; + + VecIndexPtr + CopyGpuToGpu(const int64_t, const Config&) override; + + protected: + BinarySet + SerializeImpl(const IndexType&) override; + + void + LoadImpl(const BinarySet&, const IndexType&) override; + + void + QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&) override; +}; + +using GPUIVFPtr = std::shared_ptr; + +} // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp similarity index 64% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp index dfe3602931..6706c173a1 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp @@ -7,44 +7,60 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License + +#include #include #include #include #include -#include -#include - -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +namespace milvus { namespace knowhere { -IndexModelPtr -GPUIVFPQ::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) +void +GPUIVFPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) gpu_id_ = config[knowhere::meta::DEVICEID]; - auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); - if (temp_resource != nullptr) { - ResScope rs(temp_resource, gpu_id_, true); - auto device_index = new faiss::gpu::GpuIndexIVFPQ( - temp_resource->faiss_res.get(), dim, config[IndexParams::nlist].get(), config[IndexParams::m], - config[IndexParams::nbits], - GetMetricType(config[Metric::TYPE].get())); // IP not support + auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (gpu_res != nullptr) { + ResScope rs(gpu_res, gpu_id_, true); + auto device_index = + new faiss::gpu::GpuIndexIVFPQ(gpu_res->faiss_res.get(), dim, config[IndexParams::nlist].get(), + config[IndexParams::m], config[IndexParams::nbits], + GetMetricType(config[Metric::TYPE].get())); // IP not support device_index->train(rows, (float*)p_data); std::shared_ptr host_index = nullptr; host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index)); - return std::make_shared(host_index); + + auto device_index1 = faiss::gpu::index_cpu_to_gpu(gpu_res->faiss_res.get(), gpu_id_, host_index.get()); + index_.reset(device_index1); + res_ = gpu_res; } else { KNOWHERE_THROW_MSG("Build IVFPQ can't get gpu resource"); } } +VecIndexPtr +GPUIVFPQ::CopyGpuToCpu(const Config& config) { + std::lock_guard lk(mutex_); + + faiss::Index* device_index = index_.get(); + faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); +} + std::shared_ptr GPUIVFPQ::GenParams(const Config& config) { auto params = std::make_shared(); @@ -56,16 +72,5 @@ GPUIVFPQ::GenParams(const Config& config) { return params; } -VectorIndexPtr -GPUIVFPQ::CopyGpuToCpu(const Config& config) { - std::lock_guard lk(mutex_); - - faiss::Index* device_index = index_.get(); - faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); - - std::shared_ptr new_index; - new_index.reset(host_index); - return std::make_shared(new_index); -} - } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h similarity index 61% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h index e07c30bc32..df83efbe79 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h @@ -7,37 +7,41 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include -#include "IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +namespace milvus { namespace knowhere { class GPUIVFPQ : public GPUIVF { public: explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) { + index_type_ = IndexEnum::INDEX_FAISS_IVFPQ; } - GPUIVFPQ(std::shared_ptr index, const int64_t& device_id, ResPtr& resource) - : GPUIVF(std::move(index), device_id, resource) { + GPUIVFPQ(std::shared_ptr index, const int64_t device_id, ResPtr& res) + : GPUIVF(std::move(index), device_id, res) { + index_type_ = IndexEnum::INDEX_FAISS_IVFPQ; } - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; + void + Train(const DatasetPtr&, const Config&) override; - public: - VectorIndexPtr - CopyGpuToCpu(const Config& config) override; + VecIndexPtr + CopyGpuToCpu(const Config&) override; protected: - // TODO(linxj): remove GenParams. std::shared_ptr GenParams(const Config& config) override; }; +using GPUIVFPQPtr = std::shared_ptr; + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp similarity index 61% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp index f1390c0fcb..36e5dc13bd 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include @@ -15,43 +15,40 @@ #include #include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +namespace milvus { namespace knowhere { -IndexModelPtr -GPUIVFSQ::Train(const DatasetPtr& dataset, const Config& config) { - GETTENSOR(dataset) +void +GPUIVFSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) gpu_id_ = config[knowhere::meta::DEVICEID]; std::stringstream index_type; index_type << "IVF" << config[IndexParams::nlist] << "," << "SQ" << config[IndexParams::nbits]; - auto build_index = - faiss::index_factory(dim, index_type.str().c_str(), GetMetricType(config[Metric::TYPE].get())); + faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get()); + auto build_index = faiss::index_factory(dim, index_type.str().c_str(), metric_type); - auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); - if (temp_resource != nullptr) { - ResScope rs(temp_resource, gpu_id_, true); - auto device_index = faiss::gpu::index_cpu_to_gpu(temp_resource->faiss_res.get(), gpu_id_, build_index); + auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (gpu_res != nullptr) { + ResScope rs(gpu_res, gpu_id_, true); + auto device_index = faiss::gpu::index_cpu_to_gpu(gpu_res->faiss_res.get(), gpu_id_, build_index); device_index->train(rows, (float*)p_data); - std::shared_ptr host_index = nullptr; - host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index)); - - delete device_index; - delete build_index; - - return std::make_shared(host_index); + index_.reset(device_index); + res_ = gpu_res; } else { KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource"); } } -VectorIndexPtr +VecIndexPtr GPUIVFSQ::CopyGpuToCpu(const Config& config) { std::lock_guard lk(mutex_); @@ -64,3 +61,4 @@ GPUIVFSQ::CopyGpuToCpu(const Config& config) { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h similarity index 62% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h index 462eff2447..df886046a6 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h @@ -7,31 +7,37 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include -#include "IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +namespace milvus { namespace knowhere { class GPUIVFSQ : public GPUIVF { public: explicit GPUIVFSQ(const int& device_id) : GPUIVF(device_id) { + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8; } - explicit GPUIVFSQ(std::shared_ptr index, const int64_t& device_id, ResPtr& resource) - : GPUIVF(std::move(index), device_id, resource) { + explicit GPUIVFSQ(std::shared_ptr index, const int64_t device_id, ResPtr& res) + : GPUIVF(std::move(index), device_id, res) { + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8; } - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; + void + Train(const DatasetPtr&, const Config&) override; - VectorIndexPtr - CopyGpuToCpu(const Config& config) override; + VecIndexPtr + CopyGpuToCpu(const Config&) override; }; +using GPUIVFSQPtr = std::shared_ptr; + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp similarity index 74% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp index 1b8002bc26..9cb24f0ffd 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp @@ -8,12 +8,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" -#include "knowhere/adapter/VectorAdapter.h" -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/helpers/FaissIO.h" +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include @@ -22,16 +17,20 @@ #include #include +#include "knowhere/common/Exception.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/helpers/FaissIO.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" + +namespace milvus { namespace knowhere { #ifdef CUSTOMIZATION -// std::mutex g_mutex; - -IndexModelPtr -IVFSQHybrid::Train(const DatasetPtr& dataset, const Config& config) { - // std::lock_guard lk(g_mutex); - GETTENSOR(dataset) +void +IVFSQHybrid::Train(const DatasetPtr& dataset_ptr, const Config& config) { + GETTENSOR(dataset_ptr) gpu_id_ = config[knowhere::meta::DEVICEID]; std::stringstream index_type; @@ -40,10 +39,10 @@ IVFSQHybrid::Train(const DatasetPtr& dataset, const Config& config) { auto build_index = faiss::index_factory(dim, index_type.str().c_str(), GetMetricType(config[Metric::TYPE].get())); - auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); - if (temp_resource != nullptr) { - ResScope rs(temp_resource, gpu_id_, true); - auto device_index = faiss::gpu::index_cpu_to_gpu(temp_resource->faiss_res.get(), gpu_id_, build_index); + auto gpu_res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (gpu_res != nullptr) { + ResScope rs(gpu_res, gpu_id_, true); + auto device_index = faiss::gpu::index_cpu_to_gpu(gpu_res->faiss_res.get(), gpu_id_, build_index); device_index->train(rows, (float*)p_data); std::shared_ptr host_index = nullptr; @@ -52,15 +51,18 @@ IVFSQHybrid::Train(const DatasetPtr& dataset, const Config& config) { delete device_index; delete build_index; - return std::make_shared(host_index); + device_index = faiss::gpu::index_cpu_to_gpu(gpu_res->faiss_res.get(), gpu_id_, host_index.get()); + index_.reset(device_index); + res_ = gpu_res; + gpu_mode_ = 2; } else { KNOWHERE_THROW_MSG("Build IVFSQHybrid can't get gpu resource"); } } -VectorIndexPtr +VecIndexPtr IVFSQHybrid::CopyGpuToCpu(const Config& config) { - if (gpu_mode == 0) { + if (gpu_mode_ == 0) { return std::make_shared(index_); } std::lock_guard lk(mutex_); @@ -80,8 +82,8 @@ IVFSQHybrid::CopyGpuToCpu(const Config& config) { return std::make_shared(new_index); } -VectorIndexPtr -IVFSQHybrid::CopyCpuToGpu(const int64_t& device_id, const Config& config) { +VecIndexPtr +IVFSQHybrid::CopyCpuToGpu(const int64_t device_id, const Config& config) { if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { ResScope rs(res, device_id, false); faiss::gpu::GpuClonerOptions option; @@ -98,135 +100,8 @@ IVFSQHybrid::CopyCpuToGpu(const int64_t& device_id, const Config& config) { } } -void -IVFSQHybrid::LoadImpl(const BinarySet& index_binary) { - FaissBaseIndex::LoadImpl(index_binary); // load on cpu - auto* ivf_index = dynamic_cast(index_.get()); - ivf_index->backup_quantizer(); - gpu_mode = 0; -} - -void -IVFSQHybrid::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, - const Config& cfg) { - // std::lock_guard lk(g_mutex); - // static int64_t search_count; - // ++search_count; - - if (gpu_mode == 2) { - GPUIVF::search_impl(n, data, k, distances, labels, cfg); - // index_->search(n, (float*)data, k, distances, labels); - } else if (gpu_mode == 1) { // hybrid - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(quantizer_gpu_id_)) { - ResScope rs(res, quantizer_gpu_id_, true); - IVF::search_impl(n, data, k, distances, labels, cfg); - } else { - KNOWHERE_THROW_MSG("Hybrid Search Error, can't get gpu: " + std::to_string(quantizer_gpu_id_) + "resource"); - } - } else if (gpu_mode == 0) { - IVF::search_impl(n, data, k, distances, labels, cfg); - } -} - -QuantizerPtr -IVFSQHybrid::LoadQuantizer(const Config& config) { - // std::lock_guard lk(g_mutex); - - auto gpu_id = config[knowhere::meta::DEVICEID].get(); - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) { - ResScope rs(res, gpu_id, false); - faiss::gpu::GpuClonerOptions option; - option.allInGpu = true; - - auto index_composition = new faiss::IndexComposition; - index_composition->index = index_.get(); - index_composition->quantizer = nullptr; - index_composition->mode = 1; // only 1 - - auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id, index_composition, &option); - delete gpu_index; - - auto q = std::make_shared(); - - auto& q_ptr = index_composition->quantizer; - q->size = q_ptr->d * q_ptr->getNumVecs() * sizeof(float); - q->quantizer = q_ptr; - q->gpu_id = gpu_id; - res_ = res; - gpu_mode = 1; - return q; - } else { - KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id) + "resource"); - } -} - -void -IVFSQHybrid::SetQuantizer(const QuantizerPtr& q) { - // std::lock_guard lk(g_mutex); - - auto ivf_quantizer = std::dynamic_pointer_cast(q); - if (ivf_quantizer == nullptr) { - KNOWHERE_THROW_MSG("Quantizer type error"); - } - - faiss::IndexIVF* ivf_index = dynamic_cast(index_.get()); - - faiss::gpu::GpuIndexFlat* is_gpu_flat_index = dynamic_cast(ivf_index->quantizer); - if (is_gpu_flat_index == nullptr) { - // delete ivf_index->quantizer; - ivf_index->quantizer = ivf_quantizer->quantizer; - } - quantizer_gpu_id_ = ivf_quantizer->gpu_id; - gpu_mode = 1; -} - -void -IVFSQHybrid::UnsetQuantizer() { - // std::lock_guard lk(g_mutex); - - auto* ivf_index = dynamic_cast(index_.get()); - if (ivf_index == nullptr) { - KNOWHERE_THROW_MSG("Index type error"); - } - - ivf_index->quantizer = nullptr; - quantizer_gpu_id_ = -1; -} - -VectorIndexPtr -IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& config) { - // std::lock_guard lk(g_mutex); - - int64_t gpu_id = config[knowhere::meta::DEVICEID]; - - if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) { - ResScope rs(res, gpu_id, false); - faiss::gpu::GpuClonerOptions option; - option.allInGpu = true; - - auto ivf_quantizer = std::dynamic_pointer_cast(q); - if (ivf_quantizer == nullptr) - KNOWHERE_THROW_MSG("quantizer type not faissivfquantizer"); - - auto index_composition = new faiss::IndexComposition; - index_composition->index = index_.get(); - index_composition->quantizer = ivf_quantizer->quantizer; - index_composition->mode = 2; // only 2 - - auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id, index_composition, &option); - std::shared_ptr new_idx; - new_idx.reset(gpu_index); - auto sq_idx = std::make_shared(new_idx, gpu_id, res); - return sq_idx; - } else { - KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id) + "resource"); - } -} - -std::pair -IVFSQHybrid::CopyCpuToGpuWithQuantizer(const int64_t& device_id, const Config& config) { - // std::lock_guard lk(g_mutex); - +std::pair +IVFSQHybrid::CopyCpuToGpuWithQuantizer(const int64_t device_id, const Config& config) { if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { ResScope rs(res, device_id, false); faiss::gpu::GpuClonerOptions option; @@ -253,29 +128,102 @@ IVFSQHybrid::CopyCpuToGpuWithQuantizer(const int64_t& device_id, const Config& c } } -void -IVFSQHybrid::set_index_model(IndexModelPtr model) { - std::lock_guard lk(mutex_); +VecIndexPtr +IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& quantizer_ptr, const Config& config) { + int64_t gpu_id = config[knowhere::meta::DEVICEID]; - auto host_index = std::static_pointer_cast(model); - if (auto gpures = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) { - ResScope rs(gpures, gpu_id_, false); - auto device_index = faiss::gpu::index_cpu_to_gpu(gpures->faiss_res.get(), gpu_id_, host_index->index_.get()); - index_.reset(device_index); - res_ = gpures; - gpu_mode = 2; + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) { + ResScope rs(res, gpu_id, false); + faiss::gpu::GpuClonerOptions option; + option.allInGpu = true; + + auto ivf_quantizer = std::dynamic_pointer_cast(quantizer_ptr); + if (ivf_quantizer == nullptr) + KNOWHERE_THROW_MSG("quantizer type not faissivfquantizer"); + + auto index_composition = new faiss::IndexComposition; + index_composition->index = index_.get(); + index_composition->quantizer = ivf_quantizer->quantizer; + index_composition->mode = 2; // only 2 + + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id, index_composition, &option); + std::shared_ptr new_idx; + new_idx.reset(gpu_index); + auto sq_idx = std::make_shared(new_idx, gpu_id, res); + return sq_idx; } else { - KNOWHERE_THROW_MSG("load index model error, can't get gpu_resource"); + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id) + "resource"); } } +QuantizerPtr +IVFSQHybrid::LoadQuantizer(const Config& config) { + auto gpu_id = config[knowhere::meta::DEVICEID].get(); + + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id)) { + ResScope rs(res, gpu_id, false); + faiss::gpu::GpuClonerOptions option; + option.allInGpu = true; + + auto index_composition = new faiss::IndexComposition; + index_composition->index = index_.get(); + index_composition->quantizer = nullptr; + index_composition->mode = 1; // only 1 + + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id, index_composition, &option); + delete gpu_index; + + auto q = std::make_shared(); + + auto& q_ptr = index_composition->quantizer; + q->size = q_ptr->d * q_ptr->getNumVecs() * sizeof(float); + q->quantizer = q_ptr; + q->gpu_id = gpu_id; + res_ = res; + gpu_mode_ = 1; + return q; + } else { + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu: " + std::to_string(gpu_id) + "resource"); + } +} + +void +IVFSQHybrid::SetQuantizer(const QuantizerPtr& quantizer_ptr) { + auto ivf_quantizer = std::dynamic_pointer_cast(quantizer_ptr); + if (ivf_quantizer == nullptr) { + KNOWHERE_THROW_MSG("Quantizer type error"); + } + + faiss::IndexIVF* ivf_index = dynamic_cast(index_.get()); + + faiss::gpu::GpuIndexFlat* is_gpu_flat_index = dynamic_cast(ivf_index->quantizer); + if (is_gpu_flat_index == nullptr) { + // delete ivf_index->quantizer; + ivf_index->quantizer = ivf_quantizer->quantizer; + } + quantizer_gpu_id_ = ivf_quantizer->gpu_id; + gpu_mode_ = 1; +} + +void +IVFSQHybrid::UnsetQuantizer() { + auto* ivf_index = dynamic_cast(index_.get()); + if (ivf_index == nullptr) { + KNOWHERE_THROW_MSG("Index type error"); + } + + ivf_index->quantizer = nullptr; + quantizer_gpu_id_ = -1; +} + BinarySet -IVFSQHybrid::SerializeImpl() { +IVFSQHybrid::SerializeImpl(const IndexType& type) { if (!index_ || !index_->is_trained) { KNOWHERE_THROW_MSG("index not initialize or trained"); } - fiu_do_on("IVFSQHybrid.SerializeImpl.zero_gpu_mode", gpu_mode = 0); - if (gpu_mode == 0) { + + fiu_do_on("IVFSQHybrid.SerializeImpl.zero_gpu_mode", gpu_mode_ = 0); + if (gpu_mode_ == 0) { MemoryIOWriter writer; faiss::write_index(index_.get(), &writer); @@ -286,13 +234,39 @@ IVFSQHybrid::SerializeImpl() { res_set.Append("IVF", data, writer.rp); return res_set; - } else if (gpu_mode == 2) { - return GPUIVF::SerializeImpl(); + } else if (gpu_mode_ == 2) { + return GPUIVF::SerializeImpl(type); } else { KNOWHERE_THROW_MSG("Can't serialize IVFSQ8Hybrid"); } } +void +IVFSQHybrid::LoadImpl(const BinarySet& binary_set, const IndexType& type) { + FaissBaseIndex::LoadImpl(binary_set, index_type_); // load on cpu + auto* ivf_index = dynamic_cast(index_.get()); + ivf_index->backup_quantizer(); + gpu_mode_ = 0; +} + +void +IVFSQHybrid::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, + const Config& config) { + if (gpu_mode_ == 2) { + GPUIVF::QueryImpl(n, data, k, distances, labels, config); + // index_->search(n, (float*)data, k, distances, labels); + } else if (gpu_mode_ == 1) { // hybrid + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(quantizer_gpu_id_)) { + ResScope rs(res, quantizer_gpu_id_, true); + IVF::QueryImpl(n, data, k, distances, labels, config); + } else { + KNOWHERE_THROW_MSG("Hybrid Search Error, can't get gpu: " + std::to_string(quantizer_gpu_id_) + "resource"); + } + } else if (gpu_mode_ == 0) { + IVF::QueryImpl(n, data, k, distances, labels, config); + } +} + FaissIVFQuantizer::~FaissIVFQuantizer() { if (quantizer != nullptr) { delete quantizer; @@ -302,4 +276,6 @@ FaissIVFQuantizer::~FaissIVFQuantizer() { } #endif + } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h similarity index 60% rename from core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h index 8db8f42ac3..ea4a317c79 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once @@ -17,9 +17,10 @@ #include #include -#include "IndexGPUIVFSQ.h" -#include "Quantizer.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/gpu/Quantizer.h" +namespace milvus { namespace knowhere { #ifdef CUSTOMIZATION @@ -34,22 +35,37 @@ using FaissIVFQuantizerPtr = std::shared_ptr; class IVFSQHybrid : public GPUIVFSQ { public: explicit IVFSQHybrid(const int& device_id) : GPUIVFSQ(device_id) { - gpu_mode = 0; + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H; + gpu_mode_ = 0; } explicit IVFSQHybrid(std::shared_ptr index) : GPUIVFSQ(-1) { + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H; index_ = index; - gpu_mode = 0; + gpu_mode_ = 0; } explicit IVFSQHybrid(std::shared_ptr index, const int64_t& device_id, ResPtr& resource) : GPUIVFSQ(index, device_id, resource) { - gpu_mode = 2; + index_type_ = IndexEnum::INDEX_FAISS_IVFSQ8H; + gpu_mode_ = 2; } public: void - set_index_model(IndexModelPtr model) override; + Train(const DatasetPtr&, const Config&) override; + + VecIndexPtr + CopyGpuToCpu(const Config&) override; + + VecIndexPtr + CopyCpuToGpu(const int64_t, const Config&) override; + + std::pair + CopyCpuToGpuWithQuantizer(const int64_t, const Config&); + + VecIndexPtr + LoadData(const knowhere::QuantizerPtr&, const Config&); QuantizerPtr LoadQuantizer(const Config& conf); @@ -60,35 +76,24 @@ class IVFSQHybrid : public GPUIVFSQ { void UnsetQuantizer(); - VectorIndexPtr - LoadData(const knowhere::QuantizerPtr& q, const Config& conf); - - std::pair - CopyCpuToGpuWithQuantizer(const int64_t& device_id, const Config& config); - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; - - VectorIndexPtr - CopyGpuToCpu(const Config& config) override; - - VectorIndexPtr - CopyCpuToGpu(const int64_t& device_id, const Config& config) override; - + protected: BinarySet - SerializeImpl(); - - protected: - void - search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override; + SerializeImpl(const IndexType&) override; void - LoadImpl(const BinarySet& index_binary) override; + LoadImpl(const BinarySet&, const IndexType&) override; + + void + QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&) override; protected: - int64_t gpu_mode = 0; // 0,1,2 + int64_t gpu_mode_ = 0; // 0,1,2 int64_t quantizer_gpu_id_ = -1; }; + +using IVFSQHybridPtr = std::shared_ptr; + #endif } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/Quantizer.h b/core/src/index/knowhere/knowhere/index/vector_index/gpu/Quantizer.h similarity index 91% rename from core/src/index/knowhere/knowhere/index/vector_index/Quantizer.h rename to core/src/index/knowhere/knowhere/index/vector_index/gpu/Quantizer.h index fd5752a1f3..89f1e03d79 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/Quantizer.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/gpu/Quantizer.h @@ -7,13 +7,14 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include #include "knowhere/common/Config.h" +namespace milvus { namespace knowhere { struct Quantizer { @@ -29,3 +30,4 @@ using QuantizerPtr = std::shared_ptr; // using QuantizerConfig = std::shared_ptr; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp index 330ee79978..b2aca67128 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp @@ -7,24 +7,26 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include "knowhere/index/vector_index/helpers/Cloner.h" #include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" #include "knowhere/index/vector_index/IndexIDMAP.h" #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/gpu/GPUIndex.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" +namespace milvus { namespace knowhere { namespace cloner { -VectorIndexPtr -CopyGpuToCpu(const VectorIndexPtr& index, const Config& config) { +VecIndexPtr +CopyGpuToCpu(const VecIndexPtr& index, const Config& config) { if (auto device_index = std::dynamic_pointer_cast(index)) { - VectorIndexPtr result = device_index->CopyGpuToCpu(config); + VecIndexPtr result = device_index->CopyGpuToCpu(config); auto uids = index->GetUids(); result->SetUids(uids); return result; @@ -33,9 +35,9 @@ CopyGpuToCpu(const VectorIndexPtr& index, const Config& config) { } } -VectorIndexPtr -CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config& config) { - VectorIndexPtr result; +VecIndexPtr +CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& config) { + VecIndexPtr result; auto uids = index->GetUids(); #ifdef CUSTOMIZATION if (auto device_index = std::dynamic_pointer_cast(index)) { @@ -60,7 +62,7 @@ CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { result = cpu_index->CopyCpuToGpu(device_id, config); } else { - KNOWHERE_THROW_MSG("this index type not support tranfer to gpu"); + KNOWHERE_THROW_MSG("this index type not support transfer to gpu"); } result->SetUids(uids); @@ -69,3 +71,4 @@ CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config } // namespace cloner } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.h index 2418d0e161..8252d27d4a 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.h @@ -7,21 +7,22 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once -#include "knowhere/index/vector_index/VectorIndex.h" +#include "knowhere/index/vector_index/VecIndex.h" +namespace milvus { namespace knowhere { namespace cloner { -// TODO(linxj): rename CopyToGpu -extern VectorIndexPtr -CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config& config); +extern VecIndexPtr +CopyCpuToGpu(const VecIndexPtr& index, const int64_t device_id, const Config& config); -extern VectorIndexPtr -CopyGpuToCpu(const VectorIndexPtr& index, const Config& config); +extern VecIndexPtr +CopyGpuToCpu(const VecIndexPtr& index, const Config& config); } // namespace cloner } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Definitions.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Definitions.h deleted file mode 100644 index feac9890da..0000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Definitions.h +++ /dev/null @@ -1,22 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -namespace knowhere { -namespace definition { - -#define META_ROWS ("rows") -#define META_DIM ("dimension") -#define META_K ("k") - -} // namespace definition -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp index d379598185..191ab8d10c 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp @@ -7,13 +7,14 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #include #include +namespace milvus { namespace knowhere { FaissGpuResourceMgr& @@ -120,3 +121,4 @@ FaissGpuResourceMgr::Dump() { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h index 7911f2f315..926140d2d7 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once @@ -18,8 +18,9 @@ #include -#include "utils/BlockingQueue.h" +#include "src/utils/BlockingQueue.h" +namespace milvus { namespace knowhere { struct Resource { @@ -125,3 +126,4 @@ class ResScope { }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp index 1aa8b0fbdc..58142508a7 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp @@ -7,12 +7,13 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include "knowhere/index/vector_index/helpers/FaissIO.h" +namespace milvus { namespace knowhere { // TODO(linxj): Get From Config File @@ -27,6 +28,7 @@ MemoryIOWriter::operator()(const void* ptr, size_t size, size_t nitems) { rp = size * nitems; data_ = new uint8_t[total]; memcpy((void*)(data_), ptr, rp); + return nitems; } if (total_need > total) { @@ -59,3 +61,4 @@ MemoryIOReader::operator()(void* ptr, size_t size, size_t nitems) { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.h index f246649d07..306a082191 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissIO.h @@ -7,12 +7,13 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include +namespace milvus { namespace knowhere { struct MemoryIOWriter : public faiss::IOWriter { @@ -46,3 +47,4 @@ struct MemoryIOReader : public faiss::IOReader { }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp index 1d810c023d..2d63a3f4a0 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp @@ -14,6 +14,7 @@ #include +namespace milvus { namespace knowhere { faiss::MetricType @@ -44,3 +45,4 @@ GetMetricType(const std::string& type) { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h index a21a27aca0..b9f6b361fe 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h @@ -14,6 +14,7 @@ #include #include +namespace milvus { namespace knowhere { namespace meta { @@ -60,3 +61,4 @@ extern faiss::MetricType GetMetricType(const std::string& type); } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp index b32fab0f8b..968002f389 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp @@ -13,6 +13,7 @@ #include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h" +namespace milvus { namespace knowhere { const Config& @@ -65,3 +66,4 @@ SPTAGParameterMgr::SPTAGParameterMgr() { } } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h index a92dfca90a..300ddc7db6 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h @@ -20,7 +20,9 @@ #include "IndexParameter.h" #include "knowhere/common/Config.h" +namespace milvus { namespace knowhere { + class SPTAGParameterMgr { public: const Config& @@ -50,3 +52,4 @@ class SPTAGParameterMgr { }; } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/Distance.cpp b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Distance.cpp similarity index 97% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/Distance.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Distance.cpp index 40ed6aeb3c..c6f4a8fa81 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/Distance.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Distance.cpp @@ -7,14 +7,15 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include -#include "knowhere/index/vector_index/nsg/Distance.h" +#include "knowhere/index/vector_index/impl/nsg/Distance.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { float DistanceL2::Compare(const float* a, const float* b, unsigned size) const { @@ -235,5 +236,6 @@ DistanceIP::Compare(const float* a, const float* b, unsigned size) const { // return faiss::fvec_inner_product(a,b,size); //} -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/Distance.h b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Distance.h similarity index 88% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/Distance.h rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Distance.h index d38c479dc0..9561d95936 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/Distance.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Distance.h @@ -7,12 +7,13 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { struct Distance { virtual float @@ -29,5 +30,6 @@ struct DistanceIP : public Distance { Compare(const float* a, const float* b, unsigned size) const override; }; -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSG.cpp b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSG.cpp similarity index 82% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/NSG.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSG.cpp index e220b42c64..95537fc628 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSG.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSG.cpp @@ -7,31 +7,38 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include #include #include #include #include +#include #include #include "knowhere/common/Exception.h" #include "knowhere/common/Log.h" #include "knowhere/common/Timer.h" -#include "knowhere/index/vector_index/nsg/NSG.h" -#include "knowhere/index/vector_index/nsg/NSGHelper.h" - -//#include +#include "knowhere/index/vector_index/impl/nsg/NSG.h" +#include "knowhere/index/vector_index/impl/nsg/NSGHelper.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { unsigned int seed = 100; -NsgIndex::NsgIndex(const size_t& dimension, const size_t& n, const std::string& metric) +NsgIndex::NsgIndex(const size_t& dimension, const size_t& n, std::string metric) : dimension(dimension), ntotal(n), metric_type(metric) { - distance_ = new DistanceL2; // hardcode here + // switch (metric) { + // case METRICTYPE::L2: + // break; + // case METRICTYPE::IP: + // distance_ = new DistanceIP; + // break; + // } + distance_ = new DistanceL2; } NsgIndex::~NsgIndex() { @@ -53,7 +60,6 @@ NsgIndex::Build_with_ids(size_t nb, const float* data, const int64_t* ids, const candidate_pool_size = parameters.candidate_pool_size; TimeRecorder rc("NSG", 1); - InitNavigationPoint(); rc.RecordSection("init"); @@ -62,17 +68,17 @@ NsgIndex::Build_with_ids(size_t nb, const float* data, const int64_t* ids, const CheckConnectivity(); rc.RecordSection("Connect"); + rc.ElapseFromBegin("finish"); + + is_trained = true; int total_degree = 0; for (size_t i = 0; i < ntotal; ++i) { total_degree += nsg[i].size(); } - KNOWHERE_LOG_DEBUG << "Graph physical size: " << total_degree * sizeof(node_t) / 1024 / 1024 << "m"; KNOWHERE_LOG_DEBUG << "Average degree: " << total_degree / ntotal; - is_trained = true; - // Debug code // for (size_t i = 0; i < ntotal; i++) { // auto& x = nsg[i]; @@ -125,7 +131,7 @@ NsgIndex::GetNeighbors(const float* query, std::vector& resset, std::v size_t buffer_size = search_length; if (buffer_size > ntotal) { - // TODO: throw exception here. + KNOWHERE_THROW_MSG("Build Error, search_length > ntotal"); } resset.resize(search_length); @@ -226,7 +232,7 @@ NsgIndex::GetNeighbors(const float* query, std::vector& resset, std::v size_t buffer_size = search_length; if (buffer_size > ntotal) { - // TODO: throw exception here. + KNOWHERE_THROW_MSG("Build Error, search_length > ntotal"); } // std::vector init_ids; @@ -322,7 +328,7 @@ NsgIndex::GetNeighbors(const float* query, std::vector& resset, Graph& size_t buffer_size = params ? params->search_length : search_length; if (buffer_size > ntotal) { - // TODO: throw exception here. + KNOWHERE_THROW_MSG("Build Error, search_length > ntotal"); } // std::vector init_ids; @@ -445,6 +451,7 @@ NsgIndex::Link() { // std::cout << "id: " << fullset[k].id << ", dis: " << fullset[k].distance << std::endl; // } } + knng.clear(); // Debug code // for (size_t i = 0; i < ntotal; i++) @@ -457,14 +464,11 @@ NsgIndex::Link() { // std::cout << std::endl; // } - knng.clear(); - std::vector mutex_vec(ntotal); #pragma omp for schedule(dynamic, 100) for (unsigned n = 0; n < ntotal; ++n) { InterInsert(n, mutex_vec, cut_graph_dist); } - delete[] cut_graph_dist; } @@ -527,7 +531,7 @@ NsgIndex::InterInsert(unsigned n, std::vector& mutex_vec, float* cut if (nsn_dist_pool[j] == -1) break; - // 保证至少有一条边能连回来 + // At least one edge can be connected back if (n == nsn_id_pool[j]) { duplicate = true; break; @@ -693,41 +697,139 @@ NsgIndex::FindUnconnectedNode(boost::dynamic_bitset<>& has_linked, int64_t& root nsg[root].push_back(id); } +void +NsgIndex::GetNeighbors(const float* query, node_t* I, float* D, SearchParams* params) { + size_t buffer_size = params ? params->search_length : search_length; + + if (buffer_size > ntotal) { + KNOWHERE_THROW_MSG("Search Error, search_length > ntotal"); + } + + std::vector resset(buffer_size); + std::vector init_ids(buffer_size); + boost::dynamic_bitset<> has_calculated_dist{ntotal, 0}; + + { + /* + * copy navigation-point neighbor, pick random node if less than buffer size + */ + size_t count = 0; + + // Get all neighbors + for (size_t i = 0; i < init_ids.size() && i < nsg[navigation_point].size(); ++i) { + init_ids[i] = nsg[navigation_point][i]; + has_calculated_dist[init_ids[i]] = true; + ++count; + } + while (count < buffer_size) { + node_t id = rand_r(&seed) % ntotal; + if (has_calculated_dist[id]) + continue; // duplicate id + init_ids[count] = id; + ++count; + has_calculated_dist[id] = true; + } + } + + { + // init resset and sort by distance + for (size_t i = 0; i < init_ids.size(); ++i) { + node_t id = init_ids[i]; + + if (id >= static_cast(ntotal)) { + KNOWHERE_THROW_MSG("Search Error, id > ntotal"); + } + + float dist = distance_->Compare(ori_data_ + id * dimension, query, dimension); + resset[i] = Neighbor(id, dist, false); + } + std::sort(resset.begin(), resset.end()); // sort by distance + + // search nearest neighbor + size_t cursor = 0; + while (cursor < buffer_size) { + size_t nearest_updated_pos = buffer_size; + + if (!resset[cursor].has_explored) { + resset[cursor].has_explored = true; + + node_t start_pos = resset[cursor].id; + auto& wait_for_search_node_vec = nsg[start_pos]; + for (size_t i = 0; i < wait_for_search_node_vec.size(); ++i) { + node_t id = wait_for_search_node_vec[i]; + if (has_calculated_dist[id]) + continue; + has_calculated_dist[id] = true; + + float dist = distance_->Compare(query, ori_data_ + dimension * id, dimension); + + if (dist >= resset[buffer_size - 1].distance) + continue; + + //// difference from other GetNeighbors + Neighbor nn(id, dist, false); + /////////////////////////////////////// + + size_t pos = InsertIntoPool(resset.data(), buffer_size, nn); // replace with a closer node + if (pos < nearest_updated_pos) + nearest_updated_pos = pos; + + //>> Debug code + ///// + // std::cout << "pos: " << pos << ", nn: " << nn.id << ":" << nn.distance << ", nup: " << + // nearest_updated_pos << std::endl; + ///// + + // trick: avoid search query search_length < init_ids.size() ... + if (buffer_size + 1 < resset.size()) + ++buffer_size; + } + } + if (cursor >= nearest_updated_pos) { + cursor = nearest_updated_pos; // re-search from new pos + } else { + ++cursor; + } + } + } + + if ((resset.size() - params->k) >= 0) { + for (size_t i = 0; i < params->k; ++i) { + I[i] = resset[i].id; + D[i] = resset[i].distance; + } + } else { + size_t i = 0; + for (; i < resset.size(); ++i) { + I[i] = resset[i].id; + D[i] = resset[i].distance; + } + for (; i < params->k; ++i) { + I[i] = -1; + D[i] = -1; + } + } +} + void NsgIndex::Search(const float* query, const unsigned& nq, const unsigned& dim, const unsigned& k, float* dist, int64_t* ids, SearchParams& params) { - std::vector> resset(nq); + // if (k >= 45) { + // params.search_length = k; + // } + + TimeRecorder rc("nsgsearch", 1); - TimeRecorder rc("NsgIndex::search", 1); if (nq == 1) { - GetNeighbors(query, resset[0], nsg, ¶ms); + GetNeighbors(query, ids, dist, ¶ms); } else { #pragma omp parallel for for (unsigned int i = 0; i < nq; ++i) { const float* single_query = query + i * dim; - GetNeighbors(single_query, resset[i], nsg, ¶ms); + GetNeighbors(single_query, ids + i * k, dist + i * k, ¶ms); } } - rc.RecordSection("search"); - for (unsigned int i = 0; i < nq; ++i) { - int64_t var = resset[i].size() - k; - if (var >= 0) { - for (unsigned int j = 0; j < k; ++j) { - ids[i * k + j] = ids_[resset[i][j].id]; - dist[i * k + j] = resset[i][j].distance; - } - } else { - for (unsigned int j = 0; j < resset[i].size(); ++j) { - ids[i * k + j] = ids_[resset[i][j].id]; - dist[i * k + j] = resset[i][j].distance; - } - for (unsigned int j = resset[i].size(); j < k; ++j) { - ids[i * k + j] = -1; - dist[i * k + j] = -1; - } - } - } - rc.RecordSection("merge"); + rc.ElapseFromBegin("seach finish"); } void @@ -735,5 +837,6 @@ NsgIndex::SetKnnGraph(Graph& g) { knng = std::move(g); } -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSG.h b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSG.h similarity index 91% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/NSG.h rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSG.h index fba15220a9..d3670d505d 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSG.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSG.h @@ -7,7 +7,7 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once @@ -22,8 +22,9 @@ #include "Neighbor.h" #include "knowhere/common/Config.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { using node_t = int64_t; @@ -35,6 +36,7 @@ struct BuildParams { struct SearchParams { size_t search_length; + size_t k; }; using Graph = std::vector>; @@ -43,7 +45,7 @@ class NsgIndex { public: size_t dimension; size_t ntotal; // totabl nb of indexed vectors - std::string metric_type; // todo(linxj) IP + std::string metric_type; // L2 | IP Distance* distance_; float* ori_data_; @@ -63,7 +65,7 @@ class NsgIndex { size_t out_degree; public: - explicit NsgIndex(const size_t& dimension, const size_t& n, const std::string& metric = "L2"); + explicit NsgIndex(const size_t& dimension, const size_t& n, std::string metric = "L2"); NsgIndex() = default; @@ -105,10 +107,14 @@ class NsgIndex { void GetNeighbors(const float* query, std::vector& resset, std::vector& fullset); - // search and navigation-point + // navigation-point void GetNeighbors(const float* query, std::vector& resset, Graph& graph, SearchParams* param = nullptr); + // used by search + void + GetNeighbors(const float* query, node_t* I, float* D, SearchParams* params); + void Link(); @@ -131,5 +137,6 @@ class NsgIndex { FindUnconnectedNode(boost::dynamic_bitset<>& flags, int64_t& root); }; -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGHelper.cpp b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGHelper.cpp similarity index 91% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGHelper.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGHelper.cpp index 6325f44072..334cc5d02c 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGHelper.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGHelper.cpp @@ -7,14 +7,15 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include -#include "knowhere/index/vector_index/nsg/NSGHelper.h" +#include "knowhere/index/vector_index/impl/nsg/NSGHelper.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { // TODO: impl search && insert && return insert pos. why not just find and swap? int @@ -61,5 +62,6 @@ InsertIntoPool(Neighbor* addr, unsigned K, Neighbor nn) { return right; } -}; // namespace algo +}; // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGHelper.h b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGHelper.h similarity index 84% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGHelper.h rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGHelper.h index 604d682806..f4be91db37 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGHelper.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGHelper.h @@ -7,17 +7,19 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include "Neighbor.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { extern int InsertIntoPool(Neighbor* addr, unsigned K, Neighbor nn); -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGIO.cpp b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGIO.cpp similarity index 92% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGIO.cpp rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGIO.cpp index 7c9e7bd027..880cca71ea 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGIO.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGIO.cpp @@ -7,14 +7,15 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #include -#include "knowhere/index/vector_index/nsg/NSGIO.h" +#include "knowhere/index/vector_index/impl/nsg/NSGIO.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { void write_index(NsgIndex* index, MemoryIOWriter& writer) { @@ -59,5 +60,6 @@ read_index(MemoryIOReader& reader) { return index; } -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGIO.h b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGIO.h similarity index 86% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGIO.h rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGIO.h index be0734d595..e40092d7c8 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/NSGIO.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/NSGIO.h @@ -7,15 +7,16 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include "NSG.h" #include "knowhere/index/vector_index/helpers/FaissIO.h" +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { extern void write_index(NsgIndex* index, MemoryIOWriter& writer); @@ -23,5 +24,6 @@ write_index(NsgIndex* index, MemoryIOWriter& writer); extern NsgIndex* read_index(MemoryIOReader& reader); -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/knowhere/knowhere/index/vector_index/nsg/Neighbor.h b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Neighbor.h similarity index 73% rename from core/src/index/knowhere/knowhere/index/vector_index/nsg/Neighbor.h rename to core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Neighbor.h index 31a9eea04b..6644f8f262 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/nsg/Neighbor.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/impl/nsg/Neighbor.h @@ -7,14 +7,15 @@ // // Unless required by applicable law or agreed to in writing, software distributed under the License // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. +// or implied. See the License for the specific language governing permissions and limitations under the License #pragma once #include +namespace milvus { namespace knowhere { -namespace algo { +namespace impl { using node_t = int64_t; @@ -38,19 +39,8 @@ struct Neighbor { } }; -// struct SimpleNeighbor { -// node_t id; // offset of node in origin data -// float distance; -// -// SimpleNeighbor() = default; -// explicit SimpleNeighbor(node_t id, float distance) : id{id}, distance{distance}{} -// -// inline bool operator<(const Neighbor &other) const { -// return distance < other.distance; -// } -//}; - typedef std::lock_guard LockGuard; -} // namespace algo +} // namespace impl } // namespace knowhere +} // namespace milvus diff --git a/core/src/index/thirdparty/hnswlib/hnswalg.h b/core/src/index/thirdparty/hnswlib/hnswalg.h index 0d87d37cc0..3ac3dd4dc1 100644 --- a/core/src/index/thirdparty/hnswlib/hnswalg.h +++ b/core/src/index/thirdparty/hnswlib/hnswalg.h @@ -10,234 +10,318 @@ #include "knowhere/index/vector_index/helpers/FaissIO.h" namespace hnswlib { - typedef unsigned int tableint; - typedef unsigned int linklistsizeint; - template - class HierarchicalNSW : public AlgorithmInterface { - public: +typedef unsigned int tableint; +typedef unsigned int linklistsizeint; - HierarchicalNSW(SpaceInterface *s) { +template +class HierarchicalNSW : public AlgorithmInterface { + +public: + HierarchicalNSW(SpaceInterface *s) { + } + + HierarchicalNSW(SpaceInterface *s, const std::string &location, bool nmslib = false, size_t max_elements=0) { + loadIndex(location, s, max_elements); + } + + HierarchicalNSW(SpaceInterface *s, size_t max_elements, size_t M = 16, size_t ef_construction = 200, size_t random_seed = 100) : + link_list_locks_(max_elements), element_levels_(max_elements) { + // linxj + space = s; + if (auto x = dynamic_cast(s)) { + metric_type_ = 0; + } else if (auto x = dynamic_cast(s)) { + metric_type_ = 1; + } else { + metric_type_ = 100; } - HierarchicalNSW(SpaceInterface *s, const std::string &location, bool nmslib = false, size_t max_elements=0) { - loadIndex(location, s, max_elements); - } + max_elements_ = max_elements; - HierarchicalNSW(SpaceInterface *s, size_t max_elements, size_t M = 16, size_t ef_construction = 200, size_t random_seed = 100) : - link_list_locks_(max_elements), element_levels_(max_elements) { - // linxj - space = s; - if (auto x = dynamic_cast(s)) { - metric_type_ = 0; - } else if (auto x = dynamic_cast(s)) { - metric_type_ = 1; + has_deletions_=false; + data_size_ = s->get_data_size(); + fstdistfunc_ = s->get_dist_func(); + dist_func_param_ = s->get_dist_func_param(); + M_ = M; + maxM_ = M_; + maxM0_ = M_ * 2; + ef_construction_ = std::max(ef_construction,M_); + ef_ = 10; + + level_generator_.seed(random_seed); + + size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint); + size_data_per_element_ = size_links_level0_ + data_size_ + sizeof(labeltype); + offsetData_ = size_links_level0_; + label_offset_ = size_links_level0_ + data_size_; + offsetLevel0_ = 0; + + data_level0_memory_ = (char *) malloc(max_elements_ * size_data_per_element_); + if (data_level0_memory_ == nullptr) + throw std::runtime_error("Not enough memory"); + + cur_element_count = 0; + + visited_list_pool_ = new VisitedListPool(1, max_elements); + + + + //initializations for special treatment of the first node + enterpoint_node_ = -1; + maxlevel_ = -1; + + linkLists_ = (char **) malloc(sizeof(void *) * max_elements_); + if (linkLists_ == nullptr) + throw std::runtime_error("Not enough memory: HierarchicalNSW failed to allocate linklists"); + size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint); + mult_ = 1 / log(1.0 * M_); + revSize_ = 1.0 / mult_; + } + + struct CompareByFirst { + constexpr bool operator()(std::pair const &a, + std::pair const &b) const noexcept { + return a.first < b.first; + } + }; + + ~HierarchicalNSW() { + + free(data_level0_memory_); + for (tableint i = 0; i < cur_element_count; i++) { + if (element_levels_[i] > 0) + free(linkLists_[i]); + } + free(linkLists_); + delete visited_list_pool_; + + // linxj: delete + delete space; + } + + // linxj: use for free resource + SpaceInterface *space; + size_t metric_type_; // 0:l2, 1:ip + + size_t max_elements_; + size_t cur_element_count; + size_t size_data_per_element_; + size_t size_links_per_element_; + + size_t M_; + size_t maxM_; + size_t maxM0_; + size_t ef_construction_; + + double mult_, revSize_; + int maxlevel_; + + + VisitedListPool *visited_list_pool_; + std::mutex cur_element_count_guard_; + + std::vector link_list_locks_; + tableint enterpoint_node_; + + + size_t size_links_level0_; + size_t offsetData_, offsetLevel0_; + + + char *data_level0_memory_; + char **linkLists_; + std::vector element_levels_; + + size_t data_size_; + + bool has_deletions_; + + + size_t label_offset_; + DISTFUNC fstdistfunc_; + void *dist_func_param_; + std::unordered_map label_lookup_; + + std::default_random_engine level_generator_; + + inline labeltype getExternalLabel(tableint internal_id) const { + labeltype return_label; + memcpy(&return_label,(data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_), sizeof(labeltype)); + return return_label; + } + + inline void setExternalLabel(tableint internal_id, labeltype label) const { + memcpy((data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_), &label, sizeof(labeltype)); + } + + inline labeltype *getExternalLabeLp(tableint internal_id) const { + return (labeltype *) (data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_); + } + + inline char *getDataByInternalId(tableint internal_id) const { + return (data_level0_memory_ + internal_id * size_data_per_element_ + offsetData_); + } + + int getRandomLevel(double reverse_size) { + std::uniform_real_distribution distribution(0.0, 1.0); + double r = -log(distribution(level_generator_)) * reverse_size; + return (int) r; + } + + std::priority_queue, std::vector>, CompareByFirst> + searchBaseLayer(tableint ep_id, const void *data_point, int layer) { + VisitedList *vl = visited_list_pool_->getFreeVisitedList(); + vl_type *visited_array = vl->mass; + vl_type visited_array_tag = vl->curV; + + std::priority_queue, std::vector>, CompareByFirst> top_candidates; + std::priority_queue, std::vector>, CompareByFirst> candidateSet; + + dist_t lowerBound; + if (!isMarkedDeleted(ep_id)) { + dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_); + top_candidates.emplace(dist, ep_id); + lowerBound = dist; + candidateSet.emplace(-dist, ep_id); + } else { + lowerBound = std::numeric_limits::max(); + candidateSet.emplace(-lowerBound, ep_id); + } + visited_array[ep_id] = visited_array_tag; + + while (!candidateSet.empty()) { + std::pair curr_el_pair = candidateSet.top(); + if ((-curr_el_pair.first) > lowerBound) { + break; + } + candidateSet.pop(); + + tableint curNodeNum = curr_el_pair.second; + + std::unique_lock lock(link_list_locks_[curNodeNum]); + + int *data;// = (int *)(linkList0_ + curNodeNum * size_links_per_element0_); + if (layer == 0) { + data = (int*)get_linklist0(curNodeNum); } else { - metric_type_ = 100; + data = (int*)get_linklist(curNodeNum, layer); + // data = (int *) (linkLists_[curNodeNum] + (layer - 1) * size_links_per_element_); } - - max_elements_ = max_elements; - - has_deletions_=false; - data_size_ = s->get_data_size(); - fstdistfunc_ = s->get_dist_func(); - dist_func_param_ = s->get_dist_func_param(); - M_ = M; - maxM_ = M_; - maxM0_ = M_ * 2; - ef_construction_ = std::max(ef_construction,M_); - ef_ = 10; - - level_generator_.seed(random_seed); - - size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint); - size_data_per_element_ = size_links_level0_ + data_size_ + sizeof(labeltype); - offsetData_ = size_links_level0_; - label_offset_ = size_links_level0_ + data_size_; - offsetLevel0_ = 0; - - data_level0_memory_ = (char *) malloc(max_elements_ * size_data_per_element_); - if (data_level0_memory_ == nullptr) - throw std::runtime_error("Not enough memory"); - - cur_element_count = 0; - - visited_list_pool_ = new VisitedListPool(1, max_elements); - - - - //initializations for special treatment of the first node - enterpoint_node_ = -1; - maxlevel_ = -1; - - linkLists_ = (char **) malloc(sizeof(void *) * max_elements_); - if (linkLists_ == nullptr) - throw std::runtime_error("Not enough memory: HierarchicalNSW failed to allocate linklists"); - size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint); - mult_ = 1 / log(1.0 * M_); - revSize_ = 1.0 / mult_; - } - - struct CompareByFirst { - constexpr bool operator()(std::pair const &a, - std::pair const &b) const noexcept { - return a.first < b.first; - } - }; - - ~HierarchicalNSW() { - - free(data_level0_memory_); - for (tableint i = 0; i < cur_element_count; i++) { - if (element_levels_[i] > 0) - free(linkLists_[i]); - } - free(linkLists_); - delete visited_list_pool_; - - // linxj: delete - delete space; - } - - // linxj: use for free resource - SpaceInterface *space; - size_t metric_type_; // 0:l2, 1:ip - - size_t max_elements_; - size_t cur_element_count; - size_t size_data_per_element_; - size_t size_links_per_element_; - - size_t M_; - size_t maxM_; - size_t maxM0_; - size_t ef_construction_; - - double mult_, revSize_; - int maxlevel_; - - - VisitedListPool *visited_list_pool_; - std::mutex cur_element_count_guard_; - - std::vector link_list_locks_; - tableint enterpoint_node_; - - - size_t size_links_level0_; - size_t offsetData_, offsetLevel0_; - - - char *data_level0_memory_; - char **linkLists_; - std::vector element_levels_; - - size_t data_size_; - - bool has_deletions_; - - - size_t label_offset_; - DISTFUNC fstdistfunc_; - void *dist_func_param_; - std::unordered_map label_lookup_; - - std::default_random_engine level_generator_; - - inline labeltype getExternalLabel(tableint internal_id) const { - labeltype return_label; - memcpy(&return_label,(data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_), sizeof(labeltype)); - return return_label; - } - - inline void setExternalLabel(tableint internal_id, labeltype label) const { - memcpy((data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_), &label, sizeof(labeltype)); - } - - inline labeltype *getExternalLabeLp(tableint internal_id) const { - return (labeltype *) (data_level0_memory_ + internal_id * size_data_per_element_ + label_offset_); - } - - inline char *getDataByInternalId(tableint internal_id) const { - return (data_level0_memory_ + internal_id * size_data_per_element_ + offsetData_); - } - - int getRandomLevel(double reverse_size) { - std::uniform_real_distribution distribution(0.0, 1.0); - double r = -log(distribution(level_generator_)) * reverse_size; - return (int) r; - } - - std::priority_queue, std::vector>, CompareByFirst> - searchBaseLayer(tableint ep_id, const void *data_point, int layer) { - VisitedList *vl = visited_list_pool_->getFreeVisitedList(); - vl_type *visited_array = vl->mass; - vl_type visited_array_tag = vl->curV; - - std::priority_queue, std::vector>, CompareByFirst> top_candidates; - std::priority_queue, std::vector>, CompareByFirst> candidateSet; - - dist_t lowerBound; - if (!isMarkedDeleted(ep_id)) { - dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_); - top_candidates.emplace(dist, ep_id); - lowerBound = dist; - candidateSet.emplace(-dist, ep_id); - } else { - lowerBound = std::numeric_limits::max(); - candidateSet.emplace(-lowerBound, ep_id); - } - visited_array[ep_id] = visited_array_tag; - - while (!candidateSet.empty()) { - std::pair curr_el_pair = candidateSet.top(); - if ((-curr_el_pair.first) > lowerBound) { - break; - } - candidateSet.pop(); - - tableint curNodeNum = curr_el_pair.second; - - std::unique_lock lock(link_list_locks_[curNodeNum]); - - int *data;// = (int *)(linkList0_ + curNodeNum * size_links_per_element0_); - if (layer == 0) { - data = (int*)get_linklist0(curNodeNum); - } else { - data = (int*)get_linklist(curNodeNum, layer); -// data = (int *) (linkLists_[curNodeNum] + (layer - 1) * size_links_per_element_); - } - size_t size = getListCount((linklistsizeint*)data); - tableint *datal = (tableint *) (data + 1); + size_t size = getListCount((linklistsizeint*)data); + tableint *datal = (tableint *) (data + 1); #ifdef USE_SSE - _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0); - _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0); - _mm_prefetch(getDataByInternalId(*datal), _MM_HINT_T0); - _mm_prefetch(getDataByInternalId(*(datal + 1)), _MM_HINT_T0); + _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0); + _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0); + _mm_prefetch(getDataByInternalId(*datal), _MM_HINT_T0); + _mm_prefetch(getDataByInternalId(*(datal + 1)), _MM_HINT_T0); #endif - for (size_t j = 0; j < size; j++) { - tableint candidate_id = *(datal + j); -// if (candidate_id == 0) continue; + for (size_t j = 0; j < size; j++) { + tableint candidate_id = *(datal + j); + // if (candidate_id == 0) continue; #ifdef USE_SSE - _mm_prefetch((char *) (visited_array + *(datal + j + 1)), _MM_HINT_T0); - _mm_prefetch(getDataByInternalId(*(datal + j + 1)), _MM_HINT_T0); + _mm_prefetch((char *) (visited_array + *(datal + j + 1)), _MM_HINT_T0); + _mm_prefetch(getDataByInternalId(*(datal + j + 1)), _MM_HINT_T0); #endif - if (visited_array[candidate_id] == visited_array_tag) continue; + if (visited_array[candidate_id] == visited_array_tag) continue; + visited_array[candidate_id] = visited_array_tag; + char *currObj1 = (getDataByInternalId(candidate_id)); + + dist_t dist1 = fstdistfunc_(data_point, currObj1, dist_func_param_); + if (top_candidates.size() < ef_construction_ || lowerBound > dist1) { + candidateSet.emplace(-dist1, candidate_id); +#ifdef USE_SSE + _mm_prefetch(getDataByInternalId(candidateSet.top().second), _MM_HINT_T0); +#endif + + if (!isMarkedDeleted(candidate_id)) + top_candidates.emplace(dist1, candidate_id); + + if (top_candidates.size() > ef_construction_) + top_candidates.pop(); + + if (!top_candidates.empty()) + lowerBound = top_candidates.top().first; + } + } + } + visited_list_pool_->releaseVisitedList(vl); + + return top_candidates; + } + + template + std::priority_queue, std::vector>, CompareByFirst> + searchBaseLayerST(tableint ep_id, const void *data_point, size_t ef) const { + VisitedList *vl = visited_list_pool_->getFreeVisitedList(); + vl_type *visited_array = vl->mass; + vl_type visited_array_tag = vl->curV; + + std::priority_queue, std::vector>, CompareByFirst> top_candidates; + std::priority_queue, std::vector>, CompareByFirst> candidate_set; + + dist_t lowerBound; + if (!has_deletions || !isMarkedDeleted(ep_id)) { + dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_); + lowerBound = dist; + top_candidates.emplace(dist, ep_id); + candidate_set.emplace(-dist, ep_id); + } else { + lowerBound = std::numeric_limits::max(); + candidate_set.emplace(-lowerBound, ep_id); + } + + visited_array[ep_id] = visited_array_tag; + + while (!candidate_set.empty()) { + + std::pair current_node_pair = candidate_set.top(); + + if ((-current_node_pair.first) > lowerBound) { + break; + } + candidate_set.pop(); + + tableint current_node_id = current_node_pair.second; + int *data = (int *) get_linklist0(current_node_id); + size_t size = getListCount((linklistsizeint*)data); + // bool cur_node_deleted = isMarkedDeleted(current_node_id); + +#ifdef USE_SSE + _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0); + _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0); + _mm_prefetch(data_level0_memory_ + (*(data + 1)) * size_data_per_element_ + offsetData_, _MM_HINT_T0); + _mm_prefetch((char *) (data + 2), _MM_HINT_T0); +#endif + + for (size_t j = 1; j <= size; j++) { + int candidate_id = *(data + j); + // if (candidate_id == 0) continue; +#ifdef USE_SSE + _mm_prefetch((char *) (visited_array + *(data + j + 1)), _MM_HINT_T0); + _mm_prefetch(data_level0_memory_ + (*(data + j + 1)) * size_data_per_element_ + offsetData_, + _MM_HINT_T0);//////////// +#endif + if (!(visited_array[candidate_id] == visited_array_tag)) { + visited_array[candidate_id] = visited_array_tag; - char *currObj1 = (getDataByInternalId(candidate_id)); - dist_t dist1 = fstdistfunc_(data_point, currObj1, dist_func_param_); - if (top_candidates.size() < ef_construction_ || lowerBound > dist1) { - candidateSet.emplace(-dist1, candidate_id); + char *currObj1 = (getDataByInternalId(candidate_id)); + dist_t dist = fstdistfunc_(data_point, currObj1, dist_func_param_); + + if (top_candidates.size() < ef || lowerBound > dist) { + candidate_set.emplace(-dist, candidate_id); #ifdef USE_SSE - _mm_prefetch(getDataByInternalId(candidateSet.top().second), _MM_HINT_T0); + _mm_prefetch(data_level0_memory_ + candidate_set.top().second * size_data_per_element_ + + offsetLevel0_,/////////// + _MM_HINT_T0);//////////////////////// #endif - if (!isMarkedDeleted(candidate_id)) - top_candidates.emplace(dist1, candidate_id); + if (!has_deletions || !isMarkedDeleted(candidate_id)) + top_candidates.emplace(dist, candidate_id); - if (top_candidates.size() > ef_construction_) + if (top_candidates.size() > ef) top_candidates.pop(); if (!top_candidates.empty()) @@ -245,915 +329,808 @@ namespace hnswlib { } } } - visited_list_pool_->releaseVisitedList(vl); - - return top_candidates; } - template - std::priority_queue, std::vector>, CompareByFirst> - searchBaseLayerST(tableint ep_id, const void *data_point, size_t ef) const { - VisitedList *vl = visited_list_pool_->getFreeVisitedList(); - vl_type *visited_array = vl->mass; - vl_type visited_array_tag = vl->curV; + visited_list_pool_->releaseVisitedList(vl); + return top_candidates; + } - std::priority_queue, std::vector>, CompareByFirst> top_candidates; - std::priority_queue, std::vector>, CompareByFirst> candidate_set; + void getNeighborsByHeuristic2( + std::priority_queue, std::vector>, CompareByFirst> &top_candidates, + const size_t M) { + if (top_candidates.size() < M) { + return; + } + std::priority_queue> queue_closest; + std::vector> return_list; + while (top_candidates.size() > 0) { + queue_closest.emplace(-top_candidates.top().first, top_candidates.top().second); + top_candidates.pop(); + } - dist_t lowerBound; - if (!has_deletions || !isMarkedDeleted(ep_id)) { - dist_t dist = fstdistfunc_(data_point, getDataByInternalId(ep_id), dist_func_param_); - lowerBound = dist; - top_candidates.emplace(dist, ep_id); - candidate_set.emplace(-dist, ep_id); - } else { - lowerBound = std::numeric_limits::max(); - candidate_set.emplace(-lowerBound, ep_id); - } - - visited_array[ep_id] = visited_array_tag; - - while (!candidate_set.empty()) { - - std::pair current_node_pair = candidate_set.top(); - - if ((-current_node_pair.first) > lowerBound) { + while (queue_closest.size()) { + if (return_list.size() >= M) + break; + std::pair curent_pair = queue_closest.top(); + dist_t dist_to_query = -curent_pair.first; + queue_closest.pop(); + bool good = true; + for (std::pair second_pair : return_list) { + dist_t curdist = + fstdistfunc_(getDataByInternalId(second_pair.second), + getDataByInternalId(curent_pair.second), + dist_func_param_);; + if (curdist < dist_to_query) { + good = false; break; } - candidate_set.pop(); - - tableint current_node_id = current_node_pair.second; - int *data = (int *) get_linklist0(current_node_id); - size_t size = getListCount((linklistsizeint*)data); -// bool cur_node_deleted = isMarkedDeleted(current_node_id); - -#ifdef USE_SSE - _mm_prefetch((char *) (visited_array + *(data + 1)), _MM_HINT_T0); - _mm_prefetch((char *) (visited_array + *(data + 1) + 64), _MM_HINT_T0); - _mm_prefetch(data_level0_memory_ + (*(data + 1)) * size_data_per_element_ + offsetData_, _MM_HINT_T0); - _mm_prefetch((char *) (data + 2), _MM_HINT_T0); -#endif - - for (size_t j = 1; j <= size; j++) { - int candidate_id = *(data + j); -// if (candidate_id == 0) continue; -#ifdef USE_SSE - _mm_prefetch((char *) (visited_array + *(data + j + 1)), _MM_HINT_T0); - _mm_prefetch(data_level0_memory_ + (*(data + j + 1)) * size_data_per_element_ + offsetData_, - _MM_HINT_T0);//////////// -#endif - if (!(visited_array[candidate_id] == visited_array_tag)) { - - visited_array[candidate_id] = visited_array_tag; - - char *currObj1 = (getDataByInternalId(candidate_id)); - dist_t dist = fstdistfunc_(data_point, currObj1, dist_func_param_); - - if (top_candidates.size() < ef || lowerBound > dist) { - candidate_set.emplace(-dist, candidate_id); -#ifdef USE_SSE - _mm_prefetch(data_level0_memory_ + candidate_set.top().second * size_data_per_element_ + - offsetLevel0_,/////////// - _MM_HINT_T0);//////////////////////// -#endif - - if (!has_deletions || !isMarkedDeleted(candidate_id)) - top_candidates.emplace(dist, candidate_id); - - if (top_candidates.size() > ef) - top_candidates.pop(); - - if (!top_candidates.empty()) - lowerBound = top_candidates.top().first; - } - } - } + } + if (good) { + return_list.push_back(curent_pair); } - visited_list_pool_->releaseVisitedList(vl); - return top_candidates; + } - void getNeighborsByHeuristic2( - std::priority_queue, std::vector>, CompareByFirst> &top_candidates, - const size_t M) { - if (top_candidates.size() < M) { - return; - } - std::priority_queue> queue_closest; - std::vector> return_list; - while (top_candidates.size() > 0) { - queue_closest.emplace(-top_candidates.top().first, top_candidates.top().second); - top_candidates.pop(); - } + for (std::pair curent_pair : return_list) { - while (queue_closest.size()) { - if (return_list.size() >= M) - break; - std::pair curent_pair = queue_closest.top(); - dist_t dist_to_query = -curent_pair.first; - queue_closest.pop(); - bool good = true; - for (std::pair second_pair : return_list) { - dist_t curdist = - fstdistfunc_(getDataByInternalId(second_pair.second), - getDataByInternalId(curent_pair.second), - dist_func_param_);; - if (curdist < dist_to_query) { - good = false; - break; - } - } - if (good) { - return_list.push_back(curent_pair); - } + top_candidates.emplace(-curent_pair.first, curent_pair.second); + } + } - } + linklistsizeint *get_linklist0(tableint internal_id) const { + return (linklistsizeint *) (data_level0_memory_ + internal_id * size_data_per_element_ + offsetLevel0_); + }; - for (std::pair curent_pair : return_list) { + linklistsizeint *get_linklist0(tableint internal_id, char *data_level0_memory_) const { + return (linklistsizeint *) (data_level0_memory_ + internal_id * size_data_per_element_ + offsetLevel0_); + }; - top_candidates.emplace(-curent_pair.first, curent_pair.second); - } + linklistsizeint *get_linklist(tableint internal_id, int level) const { + return (linklistsizeint *) (linkLists_[internal_id] + (level - 1) * size_links_per_element_); + }; + + void mutuallyConnectNewElement(const void *data_point, tableint cur_c, + std::priority_queue, std::vector>, CompareByFirst> top_candidates, + int level) { + + size_t Mcurmax = level ? maxM_ : maxM0_; + getNeighborsByHeuristic2(top_candidates, M_); + if (top_candidates.size() > M_) + throw std::runtime_error("Should be not be more than M_ candidates returned by the heuristic"); + + std::vector selectedNeighbors; + selectedNeighbors.reserve(M_); + while (top_candidates.size() > 0) { + selectedNeighbors.push_back(top_candidates.top().second); + top_candidates.pop(); } + { + linklistsizeint *ll_cur; + if (level == 0) + ll_cur = get_linklist0(cur_c); + else + ll_cur = get_linklist(cur_c, level); - linklistsizeint *get_linklist0(tableint internal_id) const { - return (linklistsizeint *) (data_level0_memory_ + internal_id * size_data_per_element_ + offsetLevel0_); - }; - - linklistsizeint *get_linklist0(tableint internal_id, char *data_level0_memory_) const { - return (linklistsizeint *) (data_level0_memory_ + internal_id * size_data_per_element_ + offsetLevel0_); - }; - - linklistsizeint *get_linklist(tableint internal_id, int level) const { - return (linklistsizeint *) (linkLists_[internal_id] + (level - 1) * size_links_per_element_); - }; - - void mutuallyConnectNewElement(const void *data_point, tableint cur_c, - std::priority_queue, std::vector>, CompareByFirst> top_candidates, - int level) { - - size_t Mcurmax = level ? maxM_ : maxM0_; - getNeighborsByHeuristic2(top_candidates, M_); - if (top_candidates.size() > M_) - throw std::runtime_error("Should be not be more than M_ candidates returned by the heuristic"); - - std::vector selectedNeighbors; - selectedNeighbors.reserve(M_); - while (top_candidates.size() > 0) { - selectedNeighbors.push_back(top_candidates.top().second); - top_candidates.pop(); + if (*ll_cur) { + throw std::runtime_error("The newly inserted element should have blank link list"); } - - { - linklistsizeint *ll_cur; - if (level == 0) - ll_cur = get_linklist0(cur_c); - else - ll_cur = get_linklist(cur_c, level); - - if (*ll_cur) { - throw std::runtime_error("The newly inserted element should have blank link list"); - } - setListCount(ll_cur,selectedNeighbors.size()); - tableint *data = (tableint *) (ll_cur + 1); + setListCount(ll_cur,selectedNeighbors.size()); + tableint *data = (tableint *) (ll_cur + 1); - for (size_t idx = 0; idx < selectedNeighbors.size(); idx++) { - if (data[idx]) - throw std::runtime_error("Possible memory corruption"); - if (level > element_levels_[selectedNeighbors[idx]]) - throw std::runtime_error("Trying to make a link on a non-existent level"); - - data[idx] = selectedNeighbors[idx]; - - } - } for (size_t idx = 0; idx < selectedNeighbors.size(); idx++) { - - std::unique_lock lock(link_list_locks_[selectedNeighbors[idx]]); - - - linklistsizeint *ll_other; - if (level == 0) - ll_other = get_linklist0(selectedNeighbors[idx]); - else - ll_other = get_linklist(selectedNeighbors[idx], level); - - size_t sz_link_list_other = getListCount(ll_other); - - if (sz_link_list_other > Mcurmax) - throw std::runtime_error("Bad value of sz_link_list_other"); - if (selectedNeighbors[idx] == cur_c) - throw std::runtime_error("Trying to connect an element to itself"); + if (data[idx]) + throw std::runtime_error("Possible memory corruption"); if (level > element_levels_[selectedNeighbors[idx]]) throw std::runtime_error("Trying to make a link on a non-existent level"); - tableint *data = (tableint *) (ll_other + 1); - if (sz_link_list_other < Mcurmax) { - data[sz_link_list_other] = cur_c; - setListCount(ll_other, sz_link_list_other + 1); - } else { - // finding the "weakest" element to replace it with the new one - dist_t d_max = fstdistfunc_(getDataByInternalId(cur_c), getDataByInternalId(selectedNeighbors[idx]), - dist_func_param_); - // Heuristic: - std::priority_queue, std::vector>, CompareByFirst> candidates; - candidates.emplace(d_max, cur_c); - - for (size_t j = 0; j < sz_link_list_other; j++) { - candidates.emplace( - fstdistfunc_(getDataByInternalId(data[j]), getDataByInternalId(selectedNeighbors[idx]), - dist_func_param_), data[j]); - } - - getNeighborsByHeuristic2(candidates, Mcurmax); - - int indx = 0; - while (candidates.size() > 0) { - data[indx] = candidates.top().second; - candidates.pop(); - indx++; - } - setListCount(ll_other, indx); - // Nearest K: - /*int indx = -1; - for (int j = 0; j < sz_link_list_other; j++) { - dist_t d = fstdistfunc_(getDataByInternalId(data[j]), getDataByInternalId(rez[idx]), dist_func_param_); - if (d > d_max) { - indx = j; - d_max = d; - } - } - if (indx >= 0) { - data[indx] = cur_c; - } */ - } + data[idx] = selectedNeighbors[idx]; } } + for (size_t idx = 0; idx < selectedNeighbors.size(); idx++) { - std::mutex global; - size_t ef_; - - void setEf(size_t ef) { - ef_ = ef; - } + std::unique_lock lock(link_list_locks_[selectedNeighbors[idx]]); - std::priority_queue> searchKnnInternal(void *query_data, int k) { - std::priority_queue> top_candidates; - if (cur_element_count == 0) return top_candidates; - tableint currObj = enterpoint_node_; - dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_); + linklistsizeint *ll_other; + if (level == 0) + ll_other = get_linklist0(selectedNeighbors[idx]); + else + ll_other = get_linklist(selectedNeighbors[idx], level); - for (size_t level = maxlevel_; level > 0; level--) { - bool changed = true; - while (changed) { - changed = false; - int *data; - data = (int *) get_linklist(currObj,level); - int size = getListCount(data); - tableint *datal = (tableint *) (data + 1); - for (int i = 0; i < size; i++) { - tableint cand = datal[i]; - if (cand < 0 || cand > max_elements_) - throw std::runtime_error("cand error"); - dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_); + size_t sz_link_list_other = getListCount(ll_other); - if (d < curdist) { - curdist = d; - currObj = cand; - changed = true; - } - } - } - } + if (sz_link_list_other > Mcurmax) + throw std::runtime_error("Bad value of sz_link_list_other"); + if (selectedNeighbors[idx] == cur_c) + throw std::runtime_error("Trying to connect an element to itself"); + if (level > element_levels_[selectedNeighbors[idx]]) + throw std::runtime_error("Trying to make a link on a non-existent level"); - if (has_deletions_) { - std::priority_queue> top_candidates1=searchBaseLayerST(currObj, query_data, - ef_); - top_candidates.swap(top_candidates1); - } - else{ - std::priority_queue> top_candidates1=searchBaseLayerST(currObj, query_data, - ef_); - top_candidates.swap(top_candidates1); - } - - while (top_candidates.size() > k) { - top_candidates.pop(); - } - return top_candidates; - }; - - void resizeIndex(size_t new_max_elements){ - if (new_max_elements(new_max_elements).swap(link_list_locks_); - - - // Reallocate base layer - char * data_level0_memory_new = (char *) malloc(new_max_elements * size_data_per_element_); - if (data_level0_memory_new == nullptr) - throw std::runtime_error("Not enough memory: resizeIndex failed to allocate base layer"); - memcpy(data_level0_memory_new, data_level0_memory_,cur_element_count * size_data_per_element_); - free(data_level0_memory_); - data_level0_memory_=data_level0_memory_new; - - // Reallocate all other layers - char ** linkLists_new = (char **) malloc(sizeof(void *) * new_max_elements); - if (linkLists_new == nullptr) - throw std::runtime_error("Not enough memory: resizeIndex failed to allocate other layers"); - memcpy(linkLists_new, linkLists_,cur_element_count * sizeof(void *)); - free(linkLists_); - linkLists_=linkLists_new; - - max_elements_=new_max_elements; - - } - - void saveIndex(knowhere::MemoryIOWriter& output) { - // write l2/ip calculator - writeBinaryPOD(output, metric_type_); - writeBinaryPOD(output, data_size_); - writeBinaryPOD(output, *((size_t *) dist_func_param_)); - - writeBinaryPOD(output, offsetLevel0_); - writeBinaryPOD(output, max_elements_); - writeBinaryPOD(output, cur_element_count); - writeBinaryPOD(output, size_data_per_element_); - writeBinaryPOD(output, label_offset_); - writeBinaryPOD(output, offsetData_); - writeBinaryPOD(output, maxlevel_); - writeBinaryPOD(output, enterpoint_node_); - writeBinaryPOD(output, maxM_); - - writeBinaryPOD(output, maxM0_); - writeBinaryPOD(output, M_); - writeBinaryPOD(output, mult_); - writeBinaryPOD(output, ef_construction_); - - output.write(data_level0_memory_, cur_element_count * size_data_per_element_); - - for (size_t i = 0; i < cur_element_count; i++) { - unsigned int linkListSize = element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; - writeBinaryPOD(output, linkListSize); - if (linkListSize) - output.write(linkLists_[i], linkListSize); - } -// output.close(); - } - - void loadIndex(knowhere::MemoryIOReader& input, size_t max_elements_i = 0) { - auto totoal_filesize = input.total; - - // linxj: init with metrictype - size_t dim = 100; - readBinaryPOD(input, metric_type_); - readBinaryPOD(input, data_size_); - readBinaryPOD(input, dim); - if (metric_type_ == 0) { - space = new hnswlib::L2Space(dim); - } else if (metric_type_ == 1) { - space = new hnswlib::InnerProductSpace(dim); + tableint *data = (tableint *) (ll_other + 1); + if (sz_link_list_other < Mcurmax) { + data[sz_link_list_other] = cur_c; + setListCount(ll_other, sz_link_list_other + 1); } else { - // throw exception + // finding the "weakest" element to replace it with the new one + dist_t d_max = fstdistfunc_(getDataByInternalId(cur_c), getDataByInternalId(selectedNeighbors[idx]), + dist_func_param_); + // Heuristic: + std::priority_queue, std::vector>, CompareByFirst> candidates; + candidates.emplace(d_max, cur_c); + + for (size_t j = 0; j < sz_link_list_other; j++) { + candidates.emplace( + fstdistfunc_(getDataByInternalId(data[j]), getDataByInternalId(selectedNeighbors[idx]), + dist_func_param_), data[j]); + } + + getNeighborsByHeuristic2(candidates, Mcurmax); + + int indx = 0; + while (candidates.size() > 0) { + data[indx] = candidates.top().second; + candidates.pop(); + indx++; + } + setListCount(ll_other, indx); + // Nearest K: + /*int indx = -1; + for (int j = 0; j < sz_link_list_other; j++) { + dist_t d = fstdistfunc_(getDataByInternalId(data[j]), getDataByInternalId(rez[idx]), dist_func_param_); + if (d > d_max) { + indx = j; + d_max = d; + } + } + if (indx >= 0) { + data[indx] = cur_c; + } */ } - fstdistfunc_ = space->get_dist_func(); - dist_func_param_ = space->get_dist_func_param(); - readBinaryPOD(input, offsetLevel0_); - readBinaryPOD(input, max_elements_); - readBinaryPOD(input, cur_element_count); + } + } - size_t max_elements=max_elements_i; - if(max_elements < cur_element_count) - max_elements = max_elements_; - max_elements_ = max_elements; - readBinaryPOD(input, size_data_per_element_); - readBinaryPOD(input, label_offset_); - readBinaryPOD(input, offsetData_); - readBinaryPOD(input, maxlevel_); - readBinaryPOD(input, enterpoint_node_); + std::mutex global; + size_t ef_; - readBinaryPOD(input, maxM_); - readBinaryPOD(input, maxM0_); - readBinaryPOD(input, M_); - readBinaryPOD(input, mult_); - readBinaryPOD(input, ef_construction_); + void setEf(size_t ef) { + ef_ = ef; + } -// data_size_ = s->get_data_size(); -// fstdistfunc_ = s->get_dist_func(); -// dist_func_param_ = s->get_dist_func_param(); + std::priority_queue> searchKnnInternal(void *query_data, int k) { + std::priority_queue> top_candidates; + if (cur_element_count == 0) return top_candidates; + tableint currObj = enterpoint_node_; + dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_); -// auto pos= input.rp; + for (size_t level = maxlevel_; level > 0; level--) { + bool changed = true; + while (changed) { + changed = false; + int *data; + data = (int *) get_linklist(currObj,level); + int size = getListCount(data); + tableint *datal = (tableint *) (data + 1); + for (int i = 0; i < size; i++) { + tableint cand = datal[i]; + if (cand < 0 || cand > max_elements_) + throw std::runtime_error("cand error"); + dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_); - -// /// Optional - check if index is ok: -// -// input.seekg(cur_element_count * size_data_per_element_,input.cur); -// for (size_t i = 0; i < cur_element_count; i++) { -// if(input.tellg() < 0 || input.tellg()>=total_filesize){ -// throw std::runtime_error("Index seems to be corrupted or unsupported"); -// } -// -// unsigned int linkListSize; -// readBinaryPOD(input, linkListSize); -// if (linkListSize != 0) { -// input.seekg(linkListSize,input.cur); -// } -// } -// -// // throw exception if it either corrupted or old index -// if(input.tellg()!=total_filesize) -// throw std::runtime_error("Index seems to be corrupted or unsupported"); -// -// input.clear(); -// -// /// Optional check end -// -// input.seekg(pos,input.beg); - - - data_level0_memory_ = (char *) malloc(max_elements * size_data_per_element_); - if (data_level0_memory_ == nullptr) - throw std::runtime_error("Not enough memory: loadIndex failed to allocate level0"); - input.read(data_level0_memory_, cur_element_count * size_data_per_element_); - - - - - size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint); - - - size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint); - std::vector(max_elements).swap(link_list_locks_); - - - visited_list_pool_ = new VisitedListPool(1, max_elements); - - - linkLists_ = (char **) malloc(sizeof(void *) * max_elements); - if (linkLists_ == nullptr) - throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklists"); - element_levels_ = std::vector(max_elements); - revSize_ = 1.0 / mult_; - ef_ = 10; - for (size_t i = 0; i < cur_element_count; i++) { - label_lookup_[getExternalLabel(i)]=i; - unsigned int linkListSize; - readBinaryPOD(input, linkListSize); - if (linkListSize == 0) { - element_levels_[i] = 0; - - linkLists_[i] = nullptr; - } else { - element_levels_[i] = linkListSize / size_links_per_element_; - linkLists_[i] = (char *) malloc(linkListSize); - if (linkLists_[i] == nullptr) - throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklist"); - input.read(linkLists_[i], linkListSize); + if (d < curdist) { + curdist = d; + currObj = cand; + changed = true; + } } } - - has_deletions_=false; - - for (size_t i = 0; i < cur_element_count; i++) { - if(isMarkedDeleted(i)) - has_deletions_=true; - } - - return; } - void saveIndex(const std::string &location) { - std::ofstream output(location, std::ios::binary); - std::streampos position; - - writeBinaryPOD(output, offsetLevel0_); - writeBinaryPOD(output, max_elements_); - writeBinaryPOD(output, cur_element_count); - writeBinaryPOD(output, size_data_per_element_); - writeBinaryPOD(output, label_offset_); - writeBinaryPOD(output, offsetData_); - writeBinaryPOD(output, maxlevel_); - writeBinaryPOD(output, enterpoint_node_); - writeBinaryPOD(output, maxM_); - - writeBinaryPOD(output, maxM0_); - writeBinaryPOD(output, M_); - writeBinaryPOD(output, mult_); - writeBinaryPOD(output, ef_construction_); - - output.write(data_level0_memory_, cur_element_count * size_data_per_element_); - - for (size_t i = 0; i < cur_element_count; i++) { - unsigned int linkListSize = element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; - writeBinaryPOD(output, linkListSize); - if (linkListSize) - output.write(linkLists_[i], linkListSize); - } - output.close(); + if (has_deletions_) { + std::priority_queue> top_candidates1=searchBaseLayerST(currObj, query_data, + ef_); + top_candidates.swap(top_candidates1); + } + else{ + std::priority_queue> top_candidates1=searchBaseLayerST(currObj, query_data, + ef_); + top_candidates.swap(top_candidates1); } - void loadIndex(const std::string &location, SpaceInterface *s, size_t max_elements_i=0) { + while (top_candidates.size() > k) { + top_candidates.pop(); + } + return top_candidates; + }; + + void resizeIndex(size_t new_max_elements){ + if (new_max_elements(new_max_elements).swap(link_list_locks_); - data_size_ = s->get_data_size(); - fstdistfunc_ = s->get_dist_func(); - dist_func_param_ = s->get_dist_func_param(); + // Reallocate base layer + char * data_level0_memory_new = (char *) malloc(new_max_elements * size_data_per_element_); + if (data_level0_memory_new == nullptr) + throw std::runtime_error("Not enough memory: resizeIndex failed to allocate base layer"); + memcpy(data_level0_memory_new, data_level0_memory_,cur_element_count * size_data_per_element_); + free(data_level0_memory_); + data_level0_memory_=data_level0_memory_new; - auto pos=input.tellg(); - - - /// Optional - check if index is ok: + // Reallocate all other layers + char ** linkLists_new = (char **) malloc(sizeof(void *) * new_max_elements); + if (linkLists_new == nullptr) + throw std::runtime_error("Not enough memory: resizeIndex failed to allocate other layers"); + memcpy(linkLists_new, linkLists_,cur_element_count * sizeof(void *)); + free(linkLists_); + linkLists_=linkLists_new; - input.seekg(cur_element_count * size_data_per_element_,input.cur); - for (size_t i = 0; i < cur_element_count; i++) { - if(input.tellg() < 0 || input.tellg()>=total_filesize){ - throw std::runtime_error("Index seems to be corrupted or unsupported"); - } + max_elements_=new_max_elements; - unsigned int linkListSize; - readBinaryPOD(input, linkListSize); - if (linkListSize != 0) { - input.seekg(linkListSize,input.cur); - } + } + + void saveIndex(milvus::knowhere::MemoryIOWriter& output) { + // write l2/ip calculator + writeBinaryPOD(output, metric_type_); + writeBinaryPOD(output, data_size_); + writeBinaryPOD(output, *((size_t *) dist_func_param_)); + + writeBinaryPOD(output, offsetLevel0_); + writeBinaryPOD(output, max_elements_); + writeBinaryPOD(output, cur_element_count); + writeBinaryPOD(output, size_data_per_element_); + writeBinaryPOD(output, label_offset_); + writeBinaryPOD(output, offsetData_); + writeBinaryPOD(output, maxlevel_); + writeBinaryPOD(output, enterpoint_node_); + writeBinaryPOD(output, maxM_); + + writeBinaryPOD(output, maxM0_); + writeBinaryPOD(output, M_); + writeBinaryPOD(output, mult_); + writeBinaryPOD(output, ef_construction_); + + output.write(data_level0_memory_, cur_element_count * size_data_per_element_); + + for (size_t i = 0; i < cur_element_count; i++) { + unsigned int linkListSize = element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; + writeBinaryPOD(output, linkListSize); + if (linkListSize) + output.write(linkLists_[i], linkListSize); + } + // output.close(); + } + + void loadIndex(milvus::knowhere::MemoryIOReader& input, size_t max_elements_i = 0) { + auto totoal_filesize = input.total; + + // linxj: init with metrictype + size_t dim = 100; + readBinaryPOD(input, metric_type_); + readBinaryPOD(input, data_size_); + readBinaryPOD(input, dim); + if (metric_type_ == 0) { + space = new hnswlib::L2Space(dim); + } else if (metric_type_ == 1) { + space = new hnswlib::InnerProductSpace(dim); + } else { + // throw exception + } + fstdistfunc_ = space->get_dist_func(); + dist_func_param_ = space->get_dist_func_param(); + + readBinaryPOD(input, offsetLevel0_); + readBinaryPOD(input, max_elements_); + readBinaryPOD(input, cur_element_count); + + size_t max_elements=max_elements_i; + if(max_elements < cur_element_count) + max_elements = max_elements_; + max_elements_ = max_elements; + readBinaryPOD(input, size_data_per_element_); + readBinaryPOD(input, label_offset_); + readBinaryPOD(input, offsetData_); + readBinaryPOD(input, maxlevel_); + readBinaryPOD(input, enterpoint_node_); + + readBinaryPOD(input, maxM_); + readBinaryPOD(input, maxM0_); + readBinaryPOD(input, M_); + readBinaryPOD(input, mult_); + readBinaryPOD(input, ef_construction_); + + + // data_size_ = s->get_data_size(); + // fstdistfunc_ = s->get_dist_func(); + // dist_func_param_ = s->get_dist_func_param(); + + // auto pos= input.rp; + + + // /// Optional - check if index is ok: + // + // input.seekg(cur_element_count * size_data_per_element_,input.cur); + // for (size_t i = 0; i < cur_element_count; i++) { + // if(input.tellg() < 0 || input.tellg()>=total_filesize){ + // throw std::runtime_error("Index seems to be corrupted or unsupported"); + // } + // + // unsigned int linkListSize; + // readBinaryPOD(input, linkListSize); + // if (linkListSize != 0) { + // input.seekg(linkListSize,input.cur); + // } + // } + // + // // throw exception if it either corrupted or old index + // if(input.tellg()!=total_filesize) + // throw std::runtime_error("Index seems to be corrupted or unsupported"); + // + // input.clear(); + // + // /// Optional check end + // + // input.seekg(pos,input.beg); + + + data_level0_memory_ = (char *) malloc(max_elements * size_data_per_element_); + if (data_level0_memory_ == nullptr) + throw std::runtime_error("Not enough memory: loadIndex failed to allocate level0"); + input.read(data_level0_memory_, cur_element_count * size_data_per_element_); + + + + + size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint); + + + size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint); + std::vector(max_elements).swap(link_list_locks_); + + + visited_list_pool_ = new VisitedListPool(1, max_elements); + + + linkLists_ = (char **) malloc(sizeof(void *) * max_elements); + if (linkLists_ == nullptr) + throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklists"); + element_levels_ = std::vector(max_elements); + revSize_ = 1.0 / mult_; + ef_ = 10; + for (size_t i = 0; i < cur_element_count; i++) { + label_lookup_[getExternalLabel(i)]=i; + unsigned int linkListSize; + readBinaryPOD(input, linkListSize); + if (linkListSize == 0) { + element_levels_[i] = 0; + + linkLists_[i] = nullptr; + } else { + element_levels_[i] = linkListSize / size_links_per_element_; + linkLists_[i] = (char *) malloc(linkListSize); + if (linkLists_[i] == nullptr) + throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklist"); + input.read(linkLists_[i], linkListSize); } + } - // throw exception if it either corrupted or old index - if(input.tellg()!=total_filesize) + has_deletions_=false; + + for (size_t i = 0; i < cur_element_count; i++) { + if(isMarkedDeleted(i)) + has_deletions_=true; + } + + return; + } + + void saveIndex(const std::string &location) { + std::ofstream output(location, std::ios::binary); + std::streampos position; + + writeBinaryPOD(output, offsetLevel0_); + writeBinaryPOD(output, max_elements_); + writeBinaryPOD(output, cur_element_count); + writeBinaryPOD(output, size_data_per_element_); + writeBinaryPOD(output, label_offset_); + writeBinaryPOD(output, offsetData_); + writeBinaryPOD(output, maxlevel_); + writeBinaryPOD(output, enterpoint_node_); + writeBinaryPOD(output, maxM_); + + writeBinaryPOD(output, maxM0_); + writeBinaryPOD(output, M_); + writeBinaryPOD(output, mult_); + writeBinaryPOD(output, ef_construction_); + + output.write(data_level0_memory_, cur_element_count * size_data_per_element_); + + for (size_t i = 0; i < cur_element_count; i++) { + unsigned int linkListSize = element_levels_[i] > 0 ? size_links_per_element_ * element_levels_[i] : 0; + writeBinaryPOD(output, linkListSize); + if (linkListSize) + output.write(linkLists_[i], linkListSize); + } + output.close(); + } + + void loadIndex(const std::string &location, SpaceInterface *s, size_t max_elements_i=0) { + std::ifstream input(location, std::ios::binary); + + if (!input.is_open()) + throw std::runtime_error("Cannot open file"); + + // get file size: + input.seekg(0,input.end); + std::streampos total_filesize=input.tellg(); + input.seekg(0,input.beg); + + readBinaryPOD(input, offsetLevel0_); + readBinaryPOD(input, max_elements_); + readBinaryPOD(input, cur_element_count); + + size_t max_elements=max_elements_i; + if(max_elements < cur_element_count) + max_elements = max_elements_; + max_elements_ = max_elements; + readBinaryPOD(input, size_data_per_element_); + readBinaryPOD(input, label_offset_); + readBinaryPOD(input, offsetData_); + readBinaryPOD(input, maxlevel_); + readBinaryPOD(input, enterpoint_node_); + + readBinaryPOD(input, maxM_); + readBinaryPOD(input, maxM0_); + readBinaryPOD(input, M_); + readBinaryPOD(input, mult_); + readBinaryPOD(input, ef_construction_); + + data_size_ = s->get_data_size(); + fstdistfunc_ = s->get_dist_func(); + dist_func_param_ = s->get_dist_func_param(); + + auto pos=input.tellg(); + + /// Optional - check if index is ok: + + input.seekg(cur_element_count * size_data_per_element_,input.cur); + for (size_t i = 0; i < cur_element_count; i++) { + if(input.tellg() < 0 || input.tellg()>=total_filesize){ throw std::runtime_error("Index seems to be corrupted or unsupported"); - - input.clear(); - - /// Optional check end - - input.seekg(pos,input.beg); - - - data_level0_memory_ = (char *) malloc(max_elements * size_data_per_element_); - if (data_level0_memory_ == nullptr) - throw std::runtime_error("Not enough memory: loadIndex failed to allocate level0"); - input.read(data_level0_memory_, cur_element_count * size_data_per_element_); - - - - - size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint); - - - size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint); - std::vector(max_elements).swap(link_list_locks_); - - - visited_list_pool_ = new VisitedListPool(1, max_elements); - - - linkLists_ = (char **) malloc(sizeof(void *) * max_elements); - if (linkLists_ == nullptr) - throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklists"); - element_levels_ = std::vector(max_elements); - revSize_ = 1.0 / mult_; - ef_ = 10; - for (size_t i = 0; i < cur_element_count; i++) { - label_lookup_[getExternalLabel(i)]=i; - unsigned int linkListSize; - readBinaryPOD(input, linkListSize); - if (linkListSize == 0) { - element_levels_[i] = 0; - - linkLists_[i] = nullptr; - } else { - element_levels_[i] = linkListSize / size_links_per_element_; - linkLists_[i] = (char *) malloc(linkListSize); - if (linkLists_[i] == nullptr) - throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklist"); - input.read(linkLists_[i], linkListSize); - } } - has_deletions_=false; - - for (size_t i = 0; i < cur_element_count; i++) { - if(isMarkedDeleted(i)) - has_deletions_=true; + unsigned int linkListSize; + readBinaryPOD(input, linkListSize); + if (linkListSize != 0) { + input.seekg(linkListSize,input.cur); } - - input.close(); - - return; } - template - std::vector getDataByLabel(labeltype label) + // throw exception if it either corrupted or old index + if(input.tellg()!=total_filesize) + throw std::runtime_error("Index seems to be corrupted or unsupported"); + + input.clear(); + + /// Optional check end + + input.seekg(pos,input.beg); + + data_level0_memory_ = (char *) malloc(max_elements * size_data_per_element_); + if (data_level0_memory_ == nullptr) + throw std::runtime_error("Not enough memory: loadIndex failed to allocate level0"); + input.read(data_level0_memory_, cur_element_count * size_data_per_element_); + + size_links_per_element_ = maxM_ * sizeof(tableint) + sizeof(linklistsizeint); + + size_links_level0_ = maxM0_ * sizeof(tableint) + sizeof(linklistsizeint); + std::vector(max_elements).swap(link_list_locks_); + + visited_list_pool_ = new VisitedListPool(1, max_elements); + + linkLists_ = (char **) malloc(sizeof(void *) * max_elements); + if (linkLists_ == nullptr) + throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklists"); + element_levels_ = std::vector(max_elements); + revSize_ = 1.0 / mult_; + ef_ = 10; + for (size_t i = 0; i < cur_element_count; i++) { + label_lookup_[getExternalLabel(i)]=i; + unsigned int linkListSize; + readBinaryPOD(input, linkListSize); + if (linkListSize == 0) { + element_levels_[i] = 0; + linkLists_[i] = nullptr; + } else { + element_levels_[i] = linkListSize / size_links_per_element_; + linkLists_[i] = (char *) malloc(linkListSize); + if (linkLists_[i] == nullptr) + throw std::runtime_error("Not enough memory: loadIndex failed to allocate linklist"); + input.read(linkLists_[i], linkListSize); + } + } + + has_deletions_=false; + + for (size_t i = 0; i < cur_element_count; i++) { + if(isMarkedDeleted(i)) + has_deletions_=true; + } + + input.close(); + return; + } + + template + std::vector getDataByLabel(labeltype label) { + tableint label_c; + auto search = label_lookup_.find(label); + if (search == label_lookup_.end() || isMarkedDeleted(search->second)) { + throw std::runtime_error("Label not found"); + } + label_c = search->second; + + char* data_ptrv = getDataByInternalId(label_c); + size_t dim = *((size_t *) dist_func_param_); + std::vector data; + data_t* data_ptr = (data_t*) data_ptrv; + for (int i = 0; i < dim; i++) { + data.push_back(*data_ptr); + data_ptr += 1; + } + return data; + } + + static const unsigned char DELETE_MARK = 0x01; + // static const unsigned char REUSE_MARK = 0x10; + /** + * Marks an element with the given label deleted, does NOT really change the current graph. + * @param label + */ + void markDelete(labeltype label) + { + has_deletions_=true; + auto search = label_lookup_.find(label); + if (search == label_lookup_.end()) { + throw std::runtime_error("Label not found"); + } + markDeletedInternal(search->second); + } + + /** + * Uses the first 8 bits of the memory for the linked list to store the mark, + * whereas maxM0_ has to be limited to the lower 24 bits, however, still large enough in almost all cases. + * @param internalId + */ + void markDeletedInternal(tableint internalId) { + unsigned char *ll_cur = ((unsigned char *)get_linklist0(internalId))+2; + *ll_cur |= DELETE_MARK; + } + + /** + * Remove the deleted mark of the node. + * @param internalId + */ + void unmarkDeletedInternal(tableint internalId) { + unsigned char *ll_cur = ((unsigned char *)get_linklist0(internalId))+2; + *ll_cur &= ~DELETE_MARK; + } + + /** + * Checks the first 8 bits of the memory to see if the element is marked deleted. + * @param internalId + * @return + */ + bool isMarkedDeleted(tableint internalId) const { + unsigned char *ll_cur = ((unsigned char*)get_linklist0(internalId))+2; + return *ll_cur & DELETE_MARK; + } + + unsigned short int getListCount(linklistsizeint * ptr) const { + return *((unsigned short int *)ptr); + } + + void setListCount(linklistsizeint * ptr, unsigned short int size) const { + *((unsigned short int*)(ptr))=*((unsigned short int *)&size); + } + + void addPoint(const void *data_point, labeltype label) { + addPoint(data_point, label,-1); + } + + tableint addPoint(const void *data_point, labeltype label, int level) { + tableint cur_c = 0; { - tableint label_c; + std::unique_lock lock(cur_element_count_guard_); + if (cur_element_count >= max_elements_) { + throw std::runtime_error("The number of elements exceeds the specified limit"); + }; + + cur_c = cur_element_count; + cur_element_count++; + auto search = label_lookup_.find(label); - if (search == label_lookup_.end() || isMarkedDeleted(search->second)) { - throw std::runtime_error("Label not found"); + if (search != label_lookup_.end()) { + std::unique_lock lock_el(link_list_locks_[search->second]); + has_deletions_ = true; + markDeletedInternal(search->second); } - label_c = search->second; - - char* data_ptrv = getDataByInternalId(label_c); - size_t dim = *((size_t *) dist_func_param_); - std::vector data; - data_t* data_ptr = (data_t*) data_ptrv; - for (int i = 0; i < dim; i++) { - data.push_back(*data_ptr); - data_ptr += 1; - } - return data; + label_lookup_[label] = cur_c; } - static const unsigned char DELETE_MARK = 0x01; -// static const unsigned char REUSE_MARK = 0x10; - /** - * Marks an element with the given label deleted, does NOT really change the current graph. - * @param label - */ - void markDelete(labeltype label) - { - has_deletions_=true; - auto search = label_lookup_.find(label); - if (search == label_lookup_.end()) { - throw std::runtime_error("Label not found"); - } - markDeletedInternal(search->second); + std::unique_lock lock_el(link_list_locks_[cur_c]); + int curlevel = getRandomLevel(mult_); + if (level > 0) + curlevel = level; + + element_levels_[cur_c] = curlevel; + + std::unique_lock templock(global); + int maxlevelcopy = maxlevel_; + if (curlevel <= maxlevelcopy) + templock.unlock(); + tableint currObj = enterpoint_node_; + tableint enterpoint_copy = enterpoint_node_; + + memset(data_level0_memory_ + cur_c * size_data_per_element_ + offsetLevel0_, 0, size_data_per_element_); + + // Initialisation of the data and label + memcpy(getExternalLabeLp(cur_c), &label, sizeof(labeltype)); + memcpy(getDataByInternalId(cur_c), data_point, data_size_); + + if (curlevel) { + linkLists_[cur_c] = (char *) malloc(size_links_per_element_ * curlevel + 1); + if (linkLists_[cur_c] == nullptr) + throw std::runtime_error("Not enough memory: addPoint failed to allocate linklist"); + memset(linkLists_[cur_c], 0, size_links_per_element_ * curlevel + 1); } - /** - * Uses the first 8 bits of the memory for the linked list to store the mark, - * whereas maxM0_ has to be limited to the lower 24 bits, however, still large enough in almost all cases. - * @param internalId - */ - void markDeletedInternal(tableint internalId) { - unsigned char *ll_cur = ((unsigned char *)get_linklist0(internalId))+2; - *ll_cur |= DELETE_MARK; - } + if ((signed)currObj != -1) { - /** - * Remove the deleted mark of the node. - * @param internalId - */ - void unmarkDeletedInternal(tableint internalId) { - unsigned char *ll_cur = ((unsigned char *)get_linklist0(internalId))+2; - *ll_cur &= ~DELETE_MARK; - } + if (curlevel < maxlevelcopy) { - /** - * Checks the first 8 bits of the memory to see if the element is marked deleted. - * @param internalId - * @return - */ - bool isMarkedDeleted(tableint internalId) const { - unsigned char *ll_cur = ((unsigned char*)get_linklist0(internalId))+2; - return *ll_cur & DELETE_MARK; - } + dist_t curdist = fstdistfunc_(data_point, getDataByInternalId(currObj), dist_func_param_); + for (int level = maxlevelcopy; level > curlevel; level--) { + bool changed = true; + while (changed) { + changed = false; + unsigned int *data; + std::unique_lock lock(link_list_locks_[currObj]); + data = get_linklist(currObj,level); + int size = getListCount(data); - unsigned short int getListCount(linklistsizeint * ptr) const { - return *((unsigned short int *)ptr); - } - - void setListCount(linklistsizeint * ptr, unsigned short int size) const { - *((unsigned short int*)(ptr))=*((unsigned short int *)&size); - } - - void addPoint(const void *data_point, labeltype label) { - addPoint(data_point, label,-1); - } - - tableint addPoint(const void *data_point, labeltype label, int level) { - tableint cur_c = 0; - { - std::unique_lock lock(cur_element_count_guard_); - if (cur_element_count >= max_elements_) { - throw std::runtime_error("The number of elements exceeds the specified limit"); - }; - - cur_c = cur_element_count; - cur_element_count++; - - auto search = label_lookup_.find(label); - if (search != label_lookup_.end()) { - std::unique_lock lock_el(link_list_locks_[search->second]); - has_deletions_ = true; - markDeletedInternal(search->second); - } - label_lookup_[label] = cur_c; - } - - std::unique_lock lock_el(link_list_locks_[cur_c]); - int curlevel = getRandomLevel(mult_); - if (level > 0) - curlevel = level; - - element_levels_[cur_c] = curlevel; - - - std::unique_lock templock(global); - int maxlevelcopy = maxlevel_; - if (curlevel <= maxlevelcopy) - templock.unlock(); - tableint currObj = enterpoint_node_; - tableint enterpoint_copy = enterpoint_node_; - - - memset(data_level0_memory_ + cur_c * size_data_per_element_ + offsetLevel0_, 0, size_data_per_element_); - - // Initialisation of the data and label - memcpy(getExternalLabeLp(cur_c), &label, sizeof(labeltype)); - memcpy(getDataByInternalId(cur_c), data_point, data_size_); - - - if (curlevel) { - linkLists_[cur_c] = (char *) malloc(size_links_per_element_ * curlevel + 1); - if (linkLists_[cur_c] == nullptr) - throw std::runtime_error("Not enough memory: addPoint failed to allocate linklist"); - memset(linkLists_[cur_c], 0, size_links_per_element_ * curlevel + 1); - } - - if ((signed)currObj != -1) { - - if (curlevel < maxlevelcopy) { - - dist_t curdist = fstdistfunc_(data_point, getDataByInternalId(currObj), dist_func_param_); - for (int level = maxlevelcopy; level > curlevel; level--) { - - - bool changed = true; - while (changed) { - changed = false; - unsigned int *data; - std::unique_lock lock(link_list_locks_[currObj]); - data = get_linklist(currObj,level); - int size = getListCount(data); - - tableint *datal = (tableint *) (data + 1); - for (int i = 0; i < size; i++) { - tableint cand = datal[i]; - if (cand < 0 || cand > max_elements_) - throw std::runtime_error("cand error"); - dist_t d = fstdistfunc_(data_point, getDataByInternalId(cand), dist_func_param_); - if (d < curdist) { - curdist = d; - currObj = cand; - changed = true; - } + tableint *datal = (tableint *) (data + 1); + for (int i = 0; i < size; i++) { + tableint cand = datal[i]; + if (cand < 0 || cand > max_elements_) + throw std::runtime_error("cand error"); + dist_t d = fstdistfunc_(data_point, getDataByInternalId(cand), dist_func_param_); + if (d < curdist) { + curdist = d; + currObj = cand; + changed = true; } } } } + } - bool epDeleted = isMarkedDeleted(enterpoint_copy); - for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { - if (level > maxlevelcopy || level < 0) // possible? - throw std::runtime_error("Level error"); + bool epDeleted = isMarkedDeleted(enterpoint_copy); + for (int level = std::min(curlevel, maxlevelcopy); level >= 0; level--) { + if (level > maxlevelcopy || level < 0) // possible? + throw std::runtime_error("Level error"); - std::priority_queue, std::vector>, CompareByFirst> top_candidates = searchBaseLayer( - currObj, data_point, level); - if (epDeleted) { - top_candidates.emplace(fstdistfunc_(data_point, getDataByInternalId(enterpoint_copy), dist_func_param_), enterpoint_copy); - if (top_candidates.size() > ef_construction_) - top_candidates.pop(); - } - mutuallyConnectNewElement(data_point, cur_c, top_candidates, level); - - currObj = top_candidates.top().second; + std::priority_queue, std::vector>, CompareByFirst> top_candidates = searchBaseLayer( + currObj, data_point, level); + if (epDeleted) { + top_candidates.emplace(fstdistfunc_(data_point, getDataByInternalId(enterpoint_copy), dist_func_param_), enterpoint_copy); + if (top_candidates.size() > ef_construction_) + top_candidates.pop(); } + mutuallyConnectNewElement(data_point, cur_c, top_candidates, level); - - } else { - // Do nothing for the first element - enterpoint_node_ = 0; - maxlevel_ = curlevel; - + currObj = top_candidates.top().second; } - - //Releasing lock for the maximum level - if (curlevel > maxlevelcopy) { - enterpoint_node_ = cur_c; - maxlevel_ = curlevel; - } - return cur_c; - }; - - std::priority_queue> - searchKnn(const void *query_data, size_t k) const { - std::priority_queue> result; - if (cur_element_count == 0) return result; - - tableint currObj = enterpoint_node_; - dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_); - - for (int level = maxlevel_; level > 0; level--) { - bool changed = true; - while (changed) { - changed = false; - unsigned int *data; - - data = (unsigned int *) get_linklist(currObj, level); - int size = getListCount(data); - tableint *datal = (tableint *) (data + 1); - for (int i = 0; i < size; i++) { - tableint cand = datal[i]; - if (cand < 0 || cand > max_elements_) - throw std::runtime_error("cand error"); - dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_); - - if (d < curdist) { - curdist = d; - currObj = cand; - changed = true; - } - } - } - } - - std::priority_queue, std::vector>, CompareByFirst> top_candidates; - if (has_deletions_) { - std::priority_queue, std::vector>, CompareByFirst> top_candidates1=searchBaseLayerST( - currObj, query_data, std::max(ef_, k)); - top_candidates.swap(top_candidates1); - } - else{ - std::priority_queue, std::vector>, CompareByFirst> top_candidates1=searchBaseLayerST( - currObj, query_data, std::max(ef_, k)); - top_candidates.swap(top_candidates1); - } - while (top_candidates.size() > k) { - top_candidates.pop(); - } - while (top_candidates.size() > 0) { - std::pair rez = top_candidates.top(); - result.push(std::pair(rez.first, getExternalLabel(rez.second))); - top_candidates.pop(); - } - return result; - }; - - template - std::vector> - searchKnn(const void* query_data, size_t k, Comp comp) { - std::vector> result; - if (cur_element_count == 0) return result; - - auto ret = searchKnn(query_data, k); - - while (!ret.empty()) { - result.push_back(ret.top()); - ret.pop(); - } - - std::sort(result.begin(), result.end(), comp); - - return result; + } else { + // Do nothing for the first element + enterpoint_node_ = 0; + maxlevel_ = curlevel; } + //Releasing lock for the maximum level + if (curlevel > maxlevelcopy) { + enterpoint_node_ = cur_c; + maxlevel_ = curlevel; + } + return cur_c; }; + std::priority_queue> + searchKnn(const void *query_data, size_t k) const { + std::priority_queue> result; + if (cur_element_count == 0) return result; + + tableint currObj = enterpoint_node_; + dist_t curdist = fstdistfunc_(query_data, getDataByInternalId(enterpoint_node_), dist_func_param_); + + for (int level = maxlevel_; level > 0; level--) { + bool changed = true; + while (changed) { + changed = false; + unsigned int *data; + + data = (unsigned int *) get_linklist(currObj, level); + int size = getListCount(data); + tableint *datal = (tableint *) (data + 1); + for (int i = 0; i < size; i++) { + tableint cand = datal[i]; + if (cand < 0 || cand > max_elements_) + throw std::runtime_error("cand error"); + dist_t d = fstdistfunc_(query_data, getDataByInternalId(cand), dist_func_param_); + + if (d < curdist) { + curdist = d; + currObj = cand; + changed = true; + } + } + } + } + + std::priority_queue, std::vector>, CompareByFirst> top_candidates; + if (has_deletions_) { + std::priority_queue, std::vector>, CompareByFirst> top_candidates1=searchBaseLayerST( + currObj, query_data, std::max(ef_, k)); + top_candidates.swap(top_candidates1); + } + else{ + std::priority_queue, std::vector>, CompareByFirst> top_candidates1=searchBaseLayerST( + currObj, query_data, std::max(ef_, k)); + top_candidates.swap(top_candidates1); + } + while (top_candidates.size() > k) { + top_candidates.pop(); + } + while (top_candidates.size() > 0) { + std::pair rez = top_candidates.top(); + result.push(std::pair(rez.first, getExternalLabel(rez.second))); + top_candidates.pop(); + } + return result; + }; + + template + std::vector> + searchKnn(const void* query_data, size_t k, Comp comp) { + std::vector> result; + if (cur_element_count == 0) return result; + + auto ret = searchKnn(query_data, k); + + while (!ret.empty()) { + result.push_back(ret.top()); + ret.pop(); + } + + std::sort(result.begin(), result.end(), comp); + + return result; + } +}; + } diff --git a/core/src/index/unittest/CMakeLists.txt b/core/src/index/unittest/CMakeLists.txt index b898f2d702..329a991355 100644 --- a/core/src/index/unittest/CMakeLists.txt +++ b/core/src/index/unittest/CMakeLists.txt @@ -30,6 +30,7 @@ set(util_srcs ${MILVUS_THIRDPARTY_SRC}/easyloggingpp/easylogging++.cc ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexType.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/common/Exception.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/common/Timer.cpp ${INDEX_SOURCE_DIR}/unittest/utils.cpp @@ -52,23 +53,23 @@ endif () # set(ivf_srcs - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVF.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseBinaryIndex.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVF.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp ) if (KNOWHERE_GPU_VERSION) set(ivf_srcs ${ivf_srcs} - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFSQ.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVFPQ.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/gpu/IndexIVFSQHybrid.cpp ) endif () if (NOT TARGET test_ivf) @@ -96,8 +97,7 @@ target_link_libraries(test_binaryidmap ${depend_libs} ${unittest_libs} ${basic_l # set(sptag_srcs - ${INDEX_SOURCE_DIR}/knowhere/knowhere/adapter/SptagAdapter.cpp - ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/preprocessor/Normalize.cpp + ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/adapter/SptagAdapter.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp ) diff --git a/core/src/index/unittest/Helper.h b/core/src/index/unittest/Helper.h index c56e8e272e..aec319442f 100644 --- a/core/src/index/unittest/Helper.h +++ b/core/src/index/unittest/Helper.h @@ -15,12 +15,14 @@ #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #ifdef MILVUS_GPU_VERSION -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" #endif int DEVICEID = 0; @@ -32,35 +34,37 @@ constexpr int64_t PINMEM = 1024 * 1024 * 200; constexpr int64_t TEMPMEM = 1024 * 1024 * 300; constexpr int64_t RESNUM = 2; -knowhere::IVFIndexPtr -IndexFactory(const std::string& type) { - if (type == "IVF") { - return std::make_shared(); - } else if (type == "IVFPQ") { - return std::make_shared(); - } else if (type == "IVFSQ") { - return std::make_shared(); +milvus::knowhere::IVFPtr +IndexFactory(const milvus::knowhere::IndexType& type, const milvus::knowhere::IndexMode mode) { + if (mode == milvus::knowhere::IndexMode::MODE_CPU) { + if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT) { + return std::make_shared(); + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ) { + return std::make_shared(); + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8) { + return std::make_shared(); + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H) { + std::cout << "IVFSQ8H does not support MODE_CPU" << std::endl; + } else { + std::cout << "Invalid IndexType " << type << std::endl; + } #ifdef MILVUS_GPU_VERSION - } else if (type == "GPUIVF") { - return std::make_shared(DEVICEID); - } else if (type == "GPUIVFPQ") { - return std::make_shared(DEVICEID); - } else if (type == "GPUIVFSQ") { - return std::make_shared(DEVICEID); -#ifdef CUSTOMIZATION - } else if (type == "IVFSQHybrid") { - return std::make_shared(DEVICEID); -#endif + } else { + if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT) { + return std::make_shared(DEVICEID); + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ) { + return std::make_shared(DEVICEID); + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8) { + return std::make_shared(DEVICEID); + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H) { + return std::make_shared(DEVICEID); + } else { + std::cout << "Invalid IndexType " << type << std::endl; + } #endif } } -enum class ParameterType { - ivf, - ivfpq, - ivfsq, -}; - class ParamGenerator { public: static ParamGenerator& @@ -69,35 +73,41 @@ class ParamGenerator { return instance; } - knowhere::Config - Gen(const ParameterType& type) { - if (type == ParameterType::ivf) { - return knowhere::Config{ - {knowhere::meta::DIM, DIM}, - {knowhere::meta::TOPK, K}, - {knowhere::IndexParams::nlist, 100}, - {knowhere::IndexParams::nprobe, 4}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}, - {knowhere::meta::DEVICEID, DEVICEID}, + milvus::knowhere::Config + Gen(const milvus::knowhere::IndexType& type) { + if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT) { + return milvus::knowhere::Config{ + {milvus::knowhere::meta::DIM, DIM}, + {milvus::knowhere::meta::TOPK, K}, + {milvus::knowhere::IndexParams::nlist, 100}, + {milvus::knowhere::IndexParams::nprobe, 4}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}, + {milvus::knowhere::meta::DEVICEID, DEVICEID}, }; - } else if (type == ParameterType::ivfpq) { - return knowhere::Config{ - {knowhere::meta::DIM, DIM}, - {knowhere::meta::TOPK, K}, - {knowhere::IndexParams::nlist, 100}, - {knowhere::IndexParams::nprobe, 4}, - {knowhere::IndexParams::m, 4}, - {knowhere::IndexParams::nbits, 8}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}, - {knowhere::meta::DEVICEID, DEVICEID}, + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ) { + return milvus::knowhere::Config{ + {milvus::knowhere::meta::DIM, DIM}, + {milvus::knowhere::meta::TOPK, K}, + {milvus::knowhere::IndexParams::nlist, 100}, + {milvus::knowhere::IndexParams::nprobe, 4}, + {milvus::knowhere::IndexParams::m, 4}, + {milvus::knowhere::IndexParams::nbits, 8}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}, + {milvus::knowhere::meta::DEVICEID, DEVICEID}, }; - } else if (type == ParameterType::ivfsq) { - return knowhere::Config{ - {knowhere::meta::DIM, DIM}, {knowhere::meta::TOPK, K}, - {knowhere::IndexParams::nlist, 100}, {knowhere::IndexParams::nprobe, 4}, - {knowhere::IndexParams::nbits, 8}, {knowhere::Metric::TYPE, knowhere::Metric::L2}, - {knowhere::meta::DEVICEID, DEVICEID}, + } else if (type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8 || + type == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H) { + return milvus::knowhere::Config{ + {milvus::knowhere::meta::DIM, DIM}, + {milvus::knowhere::meta::TOPK, K}, + {milvus::knowhere::IndexParams::nlist, 100}, + {milvus::knowhere::IndexParams::nprobe, 4}, + {milvus::knowhere::IndexParams::nbits, 8}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}, + {milvus::knowhere::meta::DEVICEID, DEVICEID}, }; + } else { + std::cout << "Invalid index type " << type << std::endl; } } }; @@ -109,14 +119,14 @@ class TestGpuIndexBase : public ::testing::Test { void SetUp() override { #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); #endif } void TearDown() override { #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().Free(); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif } }; diff --git a/core/src/index/unittest/test_binaryidmap.cpp b/core/src/index/unittest/test_binaryidmap.cpp index a2bb7a06c9..2188068f35 100644 --- a/core/src/index/unittest/test_binaryidmap.cpp +++ b/core/src/index/unittest/test_binaryidmap.cpp @@ -26,14 +26,14 @@ class BinaryIDMAPTest : public BinaryDataGen, public TestWithParam void SetUp() override { Init_with_binary_default(); - index_ = std::make_shared(); + index_ = std::make_shared(); } void TearDown() override{}; protected: - knowhere::BinaryIDMAPPtr index_ = nullptr; + milvus::knowhere::BinaryIDMAPPtr index_ = nullptr; }; INSTANTIATE_TEST_CASE_P(METRICParameters, BinaryIDMAPTest, @@ -43,26 +43,26 @@ TEST_P(BinaryIDMAPTest, binaryidmap_basic) { ASSERT_TRUE(!xb.empty()); std::string MetricType = GetParam(); - knowhere::Config conf{ - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::Metric::TYPE, MetricType}, + milvus::knowhere::Config conf{ + {milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::Metric::TYPE, MetricType}, }; - index_->Train(conf); + index_->Train(base_dataset, conf); index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); ASSERT_TRUE(index_->GetRawVectors() != nullptr); ASSERT_TRUE(index_->GetRawIds() != nullptr); - auto result = index_->Search(query_dataset, conf); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); // PrintResult(result, nq, k); auto binaryset = index_->Serialize(); - auto new_index = std::make_shared(); + auto new_index = std::make_shared(); new_index->Load(binaryset); - auto result2 = index_->Search(query_dataset, conf); + auto result2 = index_->Query(query_dataset, conf); AssertAnns(result2, nq, k); // PrintResult(re_result, nq, k); @@ -72,7 +72,7 @@ TEST_P(BinaryIDMAPTest, binaryidmap_basic) { } index_->SetBlacklist(concurrent_bitset_ptr); - auto result3 = index_->Search(query_dataset, conf); + auto result3 = index_->Query(query_dataset, conf); AssertAnns(result3, nq, k, CheckMode::CHECK_NOT_EQUAL); // auto result4 = index_->SearchById(id_dataset, conf); @@ -80,7 +80,7 @@ TEST_P(BinaryIDMAPTest, binaryidmap_basic) { } TEST_P(BinaryIDMAPTest, binaryidmap_serialize) { - auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { + auto serialize = [](const std::string& filename, milvus::knowhere::BinaryPtr& bin, uint8_t* ret) { FileIOWriter writer(filename); writer(static_cast(bin->data.get()), bin->size); @@ -89,21 +89,21 @@ TEST_P(BinaryIDMAPTest, binaryidmap_serialize) { }; std::string MetricType = GetParam(); - knowhere::Config conf{ - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::Metric::TYPE, MetricType}, + milvus::knowhere::Config conf{ + {milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::Metric::TYPE, MetricType}, }; { // serialize index - index_->Train(conf); - index_->Add(base_dataset, knowhere::Config()); - auto re_result = index_->Search(query_dataset, conf); + index_->Train(base_dataset, conf); + index_->Add(base_dataset, milvus::knowhere::Config()); + auto re_result = index_->Query(query_dataset, conf); AssertAnns(re_result, nq, k); // PrintResult(re_result, nq, k); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); auto binaryset = index_->Serialize(); auto bin = binaryset.GetByName("BinaryIVF"); @@ -118,8 +118,8 @@ TEST_P(BinaryIDMAPTest, binaryidmap_serialize) { index_->Load(binaryset); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); // PrintResult(result, nq, k); } diff --git a/core/src/index/unittest/test_binaryivf.cpp b/core/src/index/unittest/test_binaryivf.cpp index 3934a6810c..9dc5acd115 100644 --- a/core/src/index/unittest/test_binaryivf.cpp +++ b/core/src/index/unittest/test_binaryivf.cpp @@ -10,16 +10,13 @@ // or implied. See the License for the specific language governing permissions and limitations under the License. #include - #include #include -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" - #include "knowhere/index/vector_index/IndexBinaryIVF.h" - +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "unittest/Helper.h" #include "unittest/utils.h" @@ -37,12 +34,12 @@ class BinaryIVFTest : public BinaryDataGen, public TestWithParam { // nq = 1000; // k = 1000; // Generate(DIM, NB, NQ); - index_ = std::make_shared(); + index_ = std::make_shared(); - knowhere::Config temp_conf{ - {knowhere::meta::DIM, dim}, {knowhere::meta::TOPK, k}, - {knowhere::IndexParams::nlist, 100}, {knowhere::IndexParams::nprobe, 10}, - {knowhere::Metric::TYPE, MetricType}, + milvus::knowhere::Config temp_conf{ + {milvus::knowhere::meta::DIM, dim}, {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::IndexParams::nlist, 100}, {milvus::knowhere::IndexParams::nprobe, 10}, + {milvus::knowhere::Metric::TYPE, MetricType}, }; conf = temp_conf; } @@ -53,8 +50,8 @@ class BinaryIVFTest : public BinaryDataGen, public TestWithParam { protected: std::string index_type; - knowhere::Config conf; - knowhere::BinaryIVFIndexPtr index_ = nullptr; + milvus::knowhere::Config conf; + milvus::knowhere::BinaryIVFIndexPtr index_ = nullptr; }; INSTANTIATE_TEST_CASE_P(METRICParameters, BinaryIVFTest, @@ -70,10 +67,10 @@ TEST_P(BinaryIVFTest, binaryivf_basic) { // index_->set_index_model(model); // index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + auto result = index_->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); faiss::ConcurrentBitsetPtr concurrent_bitset_ptr = std::make_shared(nb); @@ -82,10 +79,10 @@ TEST_P(BinaryIVFTest, binaryivf_basic) { } index_->SetBlacklist(concurrent_bitset_ptr); - auto result2 = index_->Search(query_dataset, conf); + auto result2 = index_->Query(query_dataset, conf); AssertAnns(result2, nq, k, CheckMode::CHECK_NOT_EQUAL); - auto result3 = index_->SearchById(id_dataset, conf); + auto result3 = index_->QueryById(id_dataset, conf); AssertAnns(result3, nq, k, CheckMode::CHECK_NOT_EQUAL); // auto result4 = index_->GetVectorById(xid_dataset, conf); @@ -93,7 +90,7 @@ TEST_P(BinaryIVFTest, binaryivf_basic) { } TEST_P(BinaryIVFTest, binaryivf_serialize) { - auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { + auto serialize = [](const std::string& filename, milvus::knowhere::BinaryPtr& bin, uint8_t* ret) { FileIOWriter writer(filename); writer(static_cast(bin->data.get()), bin->size); @@ -120,8 +117,8 @@ TEST_P(BinaryIVFTest, binaryivf_serialize) { // // index_->set_index_model(model); // index_->Add(base_dataset, conf); - // auto result = index_->Search(query_dataset, conf); - // AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + // auto result = index_->Query(query_dataset, conf); + // AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // } { @@ -143,9 +140,9 @@ TEST_P(BinaryIVFTest, binaryivf_serialize) { index_->Load(binaryset); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); } } diff --git a/core/src/index/unittest/test_common.cpp b/core/src/index/unittest/test_common.cpp index 7857fb5229..bdebac1eb2 100644 --- a/core/src/index/unittest/test_common.cpp +++ b/core/src/index/unittest/test_common.cpp @@ -18,7 +18,7 @@ /*Some unittest for knowhere/common, mainly for improve code coverage.*/ TEST(COMMON_TEST, dataset_test) { - knowhere::Dataset set; + milvus::knowhere::Dataset set; int64_t v1 = 111; set.Set("key1", v1); @@ -31,14 +31,14 @@ TEST(COMMON_TEST, dataset_test) { TEST(COMMON_TEST, knowhere_exception) { const std::string msg = "test"; - knowhere::KnowhereException ex(msg); + milvus::knowhere::KnowhereException ex(msg); ASSERT_EQ(ex.what(), msg); } TEST(COMMON_TEST, time_recoder) { InitLog(); - knowhere::TimeRecorder recoder("COMMTEST", 0); + milvus::knowhere::TimeRecorder recoder("COMMTEST", 0); sleep(1); double span = recoder.ElapseFromBegin("get time"); ASSERT_GE(span, 1.0); diff --git a/core/src/index/unittest/test_customized_index.cpp b/core/src/index/unittest/test_customized_index.cpp index f55da9acdf..b0ce1dae17 100644 --- a/core/src/index/unittest/test_customized_index.cpp +++ b/core/src/index/unittest/test_customized_index.cpp @@ -13,12 +13,11 @@ #include #include #include -#include "knowhere/index/vector_index/helpers/Cloner.h" - -#include "unittest/Helper.h" -#include "unittest/utils.h" #include "knowhere/common/Timer.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "unittest/Helper.h" +#include "unittest/utils.h" class SingleIndexTest : public DataGen, public TestGpuIndexBase { protected: @@ -38,38 +37,38 @@ class SingleIndexTest : public DataGen, public TestGpuIndexBase { } protected: - std::string index_type; - knowhere::IVFIndexPtr index_ = nullptr; + milvus::knowhere::IndexType index_type_; + milvus::knowhere::IndexMode index_mode_; + milvus::knowhere::IVFPtr index_ = nullptr; }; #ifdef CUSTOMIZATION TEST_F(SingleIndexTest, IVFSQHybrid) { assert(!xb.empty()); - index_type = "IVFSQHybrid"; - index_ = IndexFactory(index_type); - auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); + index_type_ = milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H; + index_mode_ = milvus::knowhere::IndexMode::MODE_GPU; + index_ = IndexFactory(index_type_, index_mode_); + + auto conf = ParamGenerator::GetInstance().Gen(index_type_); fiu_init(0); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); + index_->Train(base_dataset, conf); index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); auto binaryset = index_->Serialize(); { // copy cpu to gpu - auto cpu_idx = std::make_shared(DEVICEID); + auto cpu_idx = std::make_shared(DEVICEID); cpu_idx->Load(binaryset); { for (int i = 0; i < 3; ++i) { auto gpu_idx = cpu_idx->CopyCpuToGpu(DEVICEID, conf); - auto result = gpu_idx->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + auto result = gpu_idx->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); } } @@ -77,7 +76,7 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { { // quantization already in gpu, only copy data - auto cpu_idx = std::make_shared(DEVICEID); + auto cpu_idx = std::make_shared(DEVICEID); cpu_idx->Load(binaryset); ASSERT_ANY_THROW(cpu_idx->CopyCpuToGpuWithQuantizer(-1, conf)); @@ -85,36 +84,36 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { auto gpu_idx = pair.first; auto quantization = pair.second; - auto result = gpu_idx->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + auto result = gpu_idx->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); - milvus::json quantizer_conf{{knowhere::meta::DEVICEID, DEVICEID}, {"mode", 2}}; + milvus::json quantizer_conf{{milvus::knowhere::meta::DEVICEID, DEVICEID}, {"mode", 2}}; for (int i = 0; i < 2; ++i) { - auto hybrid_idx = std::make_shared(DEVICEID); + auto hybrid_idx = std::make_shared(DEVICEID); hybrid_idx->Load(binaryset); auto new_idx = hybrid_idx->LoadData(quantization, quantizer_conf); - auto result = new_idx->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + auto result = new_idx->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); } } { // quantization already in gpu, only set quantization - auto cpu_idx = std::make_shared(DEVICEID); + auto cpu_idx = std::make_shared(DEVICEID); cpu_idx->Load(binaryset); auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf); auto quantization = pair.second; for (int i = 0; i < 2; ++i) { - auto hybrid_idx = std::make_shared(DEVICEID); + auto hybrid_idx = std::make_shared(DEVICEID); hybrid_idx->Load(binaryset); hybrid_idx->SetQuantizer(quantization); - auto result = hybrid_idx->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + auto result = hybrid_idx->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); hybrid_idx->UnsetQuantizer(); } @@ -127,7 +126,7 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { // index_type = "IVFSQHybrid"; // index_ = IndexFactory(index_type); // auto base = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq); -// auto conf = std::dynamic_pointer_cast(base); +// auto conf = std::dynamic_pointer_cast(base); // conf->nlist = 16384; // conf->k = k; // conf->nprobe = 10; @@ -145,12 +144,12 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { // // // -// auto cpu_idx = std::make_shared(DEVICEID); +// auto cpu_idx = std::make_shared(DEVICEID); // cpu_idx->Load(binaryset); // auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf); // auto quantizer = pair.second; // -// auto quantizer_conf = std::make_shared(); +// auto quantizer_conf = std::make_shared(); // quantizer_conf->mode = 2; // only copy data // quantizer_conf->gpu_id = DEVICEID; // @@ -164,7 +163,7 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { // } // }; // -// auto hybrid_qt_idx = std::make_shared(DEVICEID); +// auto hybrid_qt_idx = std::make_shared(DEVICEID); // hybrid_qt_idx->Load(binaryset); // auto SetQuantizerDoSearch = [&](int64_t search_count) { // for (int i = 0; i < search_count; ++i) { @@ -176,7 +175,7 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { // } // }; // -// auto hybrid_data_idx = std::make_shared(DEVICEID); +// auto hybrid_data_idx = std::make_shared(DEVICEID); // hybrid_data_idx->Load(binaryset); // auto LoadDataDoSearch = [&](int64_t search_count, bool do_search = false) { // for (int i = 0; i < search_count; ++i) { @@ -188,7 +187,7 @@ TEST_F(SingleIndexTest, IVFSQHybrid) { // } // }; // -// knowhere::TimeRecorder tc(""); +// milvus::knowhere::TimeRecorder tc(""); // CopyAllToGpu(2000/2, false); // tc.RecordSection("CopyAllToGpu witout search"); // CopyAllToGpu(400/2, true); diff --git a/core/src/index/unittest/test_gpuresource.cpp b/core/src/index/unittest/test_gpuresource.cpp index 71cec5078a..15fb35e1e5 100644 --- a/core/src/index/unittest/test_gpuresource.cpp +++ b/core/src/index/unittest/test_gpuresource.cpp @@ -20,13 +20,14 @@ #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/helpers/Cloner.h" #include "unittest/Helper.h" @@ -53,8 +54,9 @@ class GPURESTEST : public DataGen, public TestGpuIndexBase { } protected: - std::string index_type; - knowhere::IVFIndexPtr index_ = nullptr; + milvus::knowhere::IndexType index_type_; + milvus::knowhere::IndexMode index_mode_; + milvus::knowhere::IVFPtr index_ = nullptr; int64_t* ids = nullptr; float* dis = nullptr; @@ -65,28 +67,27 @@ TEST_F(GPURESTEST, copyandsearch) { // search and copy at the same time printf("==================\n"); - index_type = "GPUIVF"; - index_ = IndexFactory(index_type); + index_type_ = milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT; + index_mode_ = milvus::knowhere::IndexMode::MODE_GPU; + index_ = IndexFactory(index_type_, index_mode_); - auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); + auto conf = ParamGenerator::GetInstance().Gen(index_type_); + index_->Train(base_dataset, conf); index_->Add(base_dataset, conf); - auto result = index_->Search(query_dataset, conf); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); - cpu_idx->Seal(); - auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + auto cpu_idx = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); + milvus::knowhere::IVFPtr ivf_idx = std::dynamic_pointer_cast(cpu_idx); + ivf_idx->Seal(); + auto search_idx = milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); constexpr int64_t search_count = 50; constexpr int64_t load_count = 15; auto search_func = [&] { // TimeRecorder tc("search&load"); for (int i = 0; i < search_count; ++i) { - search_idx->Search(query_dataset, conf); + search_idx->Query(query_dataset, conf); // if (i > search_count - 6 || i == 0) // tc.RecordSection("search once"); } @@ -95,17 +96,17 @@ TEST_F(GPURESTEST, copyandsearch) { auto load_func = [&] { // TimeRecorder tc("search&load"); for (int i = 0; i < load_count; ++i) { - knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); // if (i > load_count -5 || i < 5) // tc.RecordSection("Copy to gpu"); } // tc.ElapseFromBegin("load finish"); }; - knowhere::TimeRecorder tc("Basic"); - knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + milvus::knowhere::TimeRecorder tc("Basic"); + milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); tc.RecordSection("Copy to gpu once"); - search_idx->Search(query_dataset, conf); + search_idx->Query(query_dataset, conf); tc.RecordSection("Search once"); search_func(); tc.RecordSection("Search total cost"); @@ -120,33 +121,29 @@ TEST_F(GPURESTEST, copyandsearch) { } TEST_F(GPURESTEST, trainandsearch) { - index_type = "GPUIVF"; - index_ = IndexFactory(index_type); + index_type_ = milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT; + index_mode_ = milvus::knowhere::IndexMode::MODE_GPU; + index_ = IndexFactory(index_type_, index_mode_); - auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - auto model = index_->Train(base_dataset, conf); - auto new_index = IndexFactory(index_type); - new_index->set_index_model(model); - new_index->Add(base_dataset, conf); - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config()); - cpu_idx->Seal(); - auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + auto conf = ParamGenerator::GetInstance().Gen(index_type_); + index_->Train(base_dataset, conf); + index_->Add(base_dataset, conf); + auto cpu_idx = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); + milvus::knowhere::IVFPtr ivf_idx = std::dynamic_pointer_cast(cpu_idx); + ivf_idx->Seal(); + auto search_idx = milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); constexpr int train_count = 5; constexpr int search_count = 200; auto train_stage = [&] { for (int i = 0; i < train_count; ++i) { - auto model = index_->Train(base_dataset, conf); - auto test_idx = IndexFactory(index_type); - test_idx->set_index_model(model); - test_idx->Add(base_dataset, conf); + index_->Train(base_dataset, conf); + index_->Add(base_dataset, conf); } }; - auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) { + auto search_stage = [&](milvus::knowhere::VecIndexPtr& search_idx) { for (int i = 0; i < search_count; ++i) { - auto result = search_idx->Search(query_dataset, conf); + auto result = search_idx->Query(query_dataset, conf); AssertAnns(result, nq, k); } }; @@ -173,7 +170,7 @@ TEST_F(GPURESTEST, trainandsearch) { } { // search parallel - auto search_idx_2 = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + auto search_idx_2 = milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); std::thread search_1(search_stage, std::ref(search_idx)); std::thread search_2(search_stage, std::ref(search_idx_2)); search_1.join(); @@ -186,10 +183,10 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) { assert(!xb.empty()); { - index_ = std::make_shared(-1); - ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), -1); - std::dynamic_pointer_cast(index_)->SetGpuDevice(DEVICEID); - ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), DEVICEID); + index_ = std::make_shared(-1); + ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), -1); + std::dynamic_pointer_cast(index_)->SetGpuDevice(DEVICEID); + ASSERT_EQ(std::dynamic_pointer_cast(index_)->GetGpuDevice(), DEVICEID); auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq); auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); @@ -200,7 +197,7 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) { EXPECT_EQ(index_->Count(), nb); EXPECT_EQ(index_->Dimension(), dim); - // knowhere::TimeRecorder tc("knowere GPUIVF"); + // milvus::knowhere::TimeRecorder tc("knowere GPUIVF"); for (int i = 0; i < search_count; ++i) { index_->Search(query_dataset, conf); if (i > search_count - 6 || i < 5) @@ -208,7 +205,7 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) { } // tc.ElapseFromBegin("search all"); } - knowhere::FaissGpuResourceMgr::GetInstance().Dump(); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Dump(); // { // // ori faiss IVF-Search @@ -219,7 +216,7 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) { // device_index.train(nb, xb.data()); // device_index.add(nb, xb.data()); // - // knowhere::TimeRecorder tc("ori IVF"); + // milvus::knowhere::TimeRecorder tc("ori IVF"); // for (int i = 0; i < search_count; ++i) { // device_index.search(nq, xq.data(), k, dis, ids); // if (i > search_count - 6 || i < 5) @@ -235,11 +232,11 @@ TEST_F(GPURESTEST, gpuivfsq) { index_type = "GPUIVFSQ"; index_ = IndexFactory(index_type); - auto conf = std::make_shared(); + auto conf = std::make_shared(); conf->nlist = 1638; conf->d = dim; conf->gpu_id = DEVICEID; - conf->metric_type = knowhere::METRICTYPE::L2; + conf->metric_type = milvus::knowhere::METRICTYPE::L2; conf->k = k; conf->nbits = 8; conf->nprobe = 1; @@ -252,11 +249,11 @@ TEST_F(GPURESTEST, gpuivfsq) { // auto result = index_->Search(query_dataset, conf); // AssertAnns(result, nq, k); - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + auto cpu_idx = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); cpu_idx->Seal(); - knowhere::TimeRecorder tc("knowhere GPUSQ8"); - auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + milvus::knowhere::TimeRecorder tc("knowhere GPUSQ8"); + auto search_idx = milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); tc.RecordSection("Copy to gpu"); for (int i = 0; i < search_count; ++i) { search_idx->Search(query_dataset, conf); @@ -287,7 +284,7 @@ TEST_F(GPURESTEST, gpuivfsq) { faiss::gpu::GpuClonerOptions option; option.allInGpu = true; - knowhere::TimeRecorder tc("ori GPUSQ8"); + milvus::knowhere::TimeRecorder tc("ori GPUSQ8"); faiss::Index* search_idx = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, cpu_index, &option); tc.RecordSection("Copy to gpu"); for (int i = 0; i < search_count; ++i) { diff --git a/core/src/index/unittest/test_idmap.cpp b/core/src/index/unittest/test_idmap.cpp index 69d2c0ced9..0a13a2528c 100644 --- a/core/src/index/unittest/test_idmap.cpp +++ b/core/src/index/unittest/test_idmap.cpp @@ -17,7 +17,7 @@ #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexIDMAP.h" #ifdef MILVUS_GPU_VERSION -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" #include "knowhere/index/vector_index/helpers/Cloner.h" #endif #include "Helper.h" @@ -30,7 +30,7 @@ class IDMAPTest : public DataGen, public TestGpuIndexBase { TestGpuIndexBase::SetUp(); Init_with_default(); - index_ = std::make_shared(); + index_ = std::make_shared(); } void @@ -39,42 +39,42 @@ class IDMAPTest : public DataGen, public TestGpuIndexBase { } protected: - knowhere::IDMAPPtr index_ = nullptr; + milvus::knowhere::IDMAPPtr index_ = nullptr; }; TEST_F(IDMAPTest, idmap_basic) { ASSERT_TRUE(!xb.empty()); - knowhere::Config conf{ - {knowhere::meta::DIM, dim}, {knowhere::meta::TOPK, k}, {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + milvus::knowhere::Config conf{{milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}}; // null faiss index { ASSERT_ANY_THROW(index_->Serialize()); - ASSERT_ANY_THROW(index_->Search(query_dataset, conf)); + ASSERT_ANY_THROW(index_->Query(query_dataset, conf)); ASSERT_ANY_THROW(index_->Add(nullptr, conf)); - ASSERT_ANY_THROW(index_->AddWithoutId(nullptr, conf)); + ASSERT_ANY_THROW(index_->AddWithoutIds(nullptr, conf)); } - index_->Train(conf); + index_->Train(base_dataset, conf); index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); ASSERT_TRUE(index_->GetRawVectors() != nullptr); ASSERT_TRUE(index_->GetRawIds() != nullptr); - auto result = index_->Search(query_dataset, conf); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); // PrintResult(result, nq, k); - index_->Seal(); auto binaryset = index_->Serialize(); - auto new_index = std::make_shared(); + auto new_index = std::make_shared(); new_index->Load(binaryset); - auto result2 = index_->Search(query_dataset, conf); + auto result2 = index_->Query(query_dataset, conf); AssertAnns(result2, nq, k); // PrintResult(re_result, nq, k); - auto result3 = index_->SearchById(id_dataset, conf); + auto result3 = index_->QueryById(id_dataset, conf); AssertAnns(result3, nq, k); auto result4 = index_->GetVectorById(xid_dataset, conf); @@ -86,10 +86,10 @@ TEST_F(IDMAPTest, idmap_basic) { } index_->SetBlacklist(concurrent_bitset_ptr); - auto result_bs_1 = index_->Search(query_dataset, conf); + auto result_bs_1 = index_->Query(query_dataset, conf); AssertAnns(result_bs_1, nq, k, CheckMode::CHECK_NOT_EQUAL); - auto result_bs_2 = index_->SearchById(id_dataset, conf); + auto result_bs_2 = index_->QueryById(id_dataset, conf); AssertAnns(result_bs_2, nq, k, CheckMode::CHECK_NOT_EQUAL); auto result_bs_3 = index_->GetVectorById(xid_dataset, conf); @@ -97,7 +97,7 @@ TEST_F(IDMAPTest, idmap_basic) { } TEST_F(IDMAPTest, idmap_serialize) { - auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { + auto serialize = [](const std::string& filename, milvus::knowhere::BinaryPtr& bin, uint8_t* ret) { FileIOWriter writer(filename); writer(static_cast(bin->data.get()), bin->size); @@ -105,18 +105,19 @@ TEST_F(IDMAPTest, idmap_serialize) { reader(ret, bin->size); }; - knowhere::Config conf{ - {knowhere::meta::DIM, dim}, {knowhere::meta::TOPK, k}, {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + milvus::knowhere::Config conf{{milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}}; { // serialize index - index_->Train(conf); - index_->Add(base_dataset, knowhere::Config()); - auto re_result = index_->Search(query_dataset, conf); + index_->Train(base_dataset, conf); + index_->Add(base_dataset, milvus::knowhere::Config()); + auto re_result = index_->Query(query_dataset, conf); AssertAnns(re_result, nq, k); // PrintResult(re_result, nq, k); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); auto binaryset = index_->Serialize(); auto bin = binaryset.GetByName("IVF"); @@ -131,8 +132,8 @@ TEST_F(IDMAPTest, idmap_serialize) { index_->Load(binaryset); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); // PrintResult(result, nq, k); } @@ -142,16 +143,17 @@ TEST_F(IDMAPTest, idmap_serialize) { TEST_F(IDMAPTest, copy_test) { ASSERT_TRUE(!xb.empty()); - knowhere::Config conf{ - {knowhere::meta::DIM, dim}, {knowhere::meta::TOPK, k}, {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + milvus::knowhere::Config conf{{milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}}; - index_->Train(conf); + index_->Train(base_dataset, conf); index_->Add(base_dataset, conf); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); ASSERT_TRUE(index_->GetRawVectors() != nullptr); ASSERT_TRUE(index_->GetRawIds() != nullptr); - auto result = index_->Search(query_dataset, conf); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); // PrintResult(result, nq, k); @@ -164,14 +166,14 @@ TEST_F(IDMAPTest, copy_test) { { // cpu to gpu - ASSERT_ANY_THROW(knowhere::cloner::CopyCpuToGpu(index_, -1, conf)); - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); - auto clone_result = clone_index->Search(query_dataset, conf); + ASSERT_ANY_THROW(milvus::knowhere::cloner::CopyCpuToGpu(index_, -1, conf)); + auto clone_index = milvus::knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); + auto clone_result = clone_index->Query(query_dataset, conf); AssertAnns(clone_result, nq, k); - ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawVectors(); }, - knowhere::KnowhereException); - ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawIds(); }, - knowhere::KnowhereException); + ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawVectors(); }, + milvus::knowhere::KnowhereException); + ASSERT_THROW({ std::static_pointer_cast(clone_index)->GetRawIds(); }, + milvus::knowhere::KnowhereException); fiu_init(0); fiu_enable("GPUIDMP.SerializeImpl.throw_exception", 1, nullptr, 0); @@ -180,7 +182,7 @@ TEST_F(IDMAPTest, copy_test) { auto binary = clone_index->Serialize(); clone_index->Load(binary); - auto new_result = clone_index->Search(query_dataset, conf); + auto new_result = clone_index->Query(query_dataset, conf); AssertAnns(new_result, nq, k); // auto clone_gpu_idx = clone_index->Clone(); @@ -188,17 +190,17 @@ TEST_F(IDMAPTest, copy_test) { // AssertAnns(clone_gpu_res, nq, k); // gpu to cpu - auto host_index = knowhere::cloner::CopyGpuToCpu(clone_index, conf); - auto host_result = host_index->Search(query_dataset, conf); + auto host_index = milvus::knowhere::cloner::CopyGpuToCpu(clone_index, conf); + auto host_result = host_index->Query(query_dataset, conf); AssertAnns(host_result, nq, k); - ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawVectors() != nullptr); - ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawIds() != nullptr); + ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawVectors() != nullptr); + ASSERT_TRUE(std::static_pointer_cast(host_index)->GetRawIds() != nullptr); // gpu to gpu - auto device_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); + auto device_index = milvus::knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf); auto new_device_index = - std::static_pointer_cast(device_index)->CopyGpuToGpu(DEVICEID, conf); - auto device_result = new_device_index->Search(query_dataset, conf); + std::static_pointer_cast(device_index)->CopyGpuToGpu(DEVICEID, conf); + auto device_result = new_device_index->Query(query_dataset, conf); AssertAnns(device_result, nq, k); } } diff --git a/core/src/index/unittest/test_ivf.cpp b/core/src/index/unittest/test_ivf.cpp index b30cffae51..c2de43fab5 100644 --- a/core/src/index/unittest/test_ivf.cpp +++ b/core/src/index/unittest/test_ivf.cpp @@ -17,27 +17,24 @@ #include #ifdef MILVUS_GPU_VERSION - #include - #endif -#include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" #include "knowhere/common/Timer.h" - #include "knowhere/index/vector_index/IndexIVF.h" #include "knowhere/index/vector_index/IndexIVFPQ.h" #include "knowhere/index/vector_index/IndexIVFSQ.h" +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #ifdef MILVUS_GPU_VERSION - -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVF.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFPQ.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIVFSQ.h" +#include "knowhere/index/vector_index/gpu/IndexIVFSQHybrid.h" #include "knowhere/index/vector_index/helpers/Cloner.h" - +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif #include "unittest/Helper.h" @@ -47,81 +44,80 @@ using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; -class IVFTest : public DataGen, public TestWithParam<::std::tuple> { +class IVFTest : public DataGen, + public TestWithParam<::std::tuple> { protected: void SetUp() override { #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); #endif - std::tie(index_type, parameter_type_) = GetParam(); + std::tie(index_type_, index_mode_) = GetParam(); // Init_with_default(); // nb = 1000000; // nq = 1000; // k = 1000; Generate(DIM, NB, NQ); - index_ = IndexFactory(index_type); - conf = ParamGenerator::GetInstance().Gen(parameter_type_); - // KNOWHERE_LOG_DEBUG << "conf: " << conf->dump(); + index_ = IndexFactory(index_type_, index_mode_); + conf_ = ParamGenerator::GetInstance().Gen(index_type_); + // conf_->Dump(); } void TearDown() override { #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().Free(); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif } protected: - std::string index_type; - knowhere::Config conf; - knowhere::IVFIndexPtr index_ = nullptr; - ParameterType parameter_type_; + milvus::knowhere::IndexType index_type_; + milvus::knowhere::IndexMode index_mode_; + milvus::knowhere::Config conf_; + milvus::knowhere::IVFPtr index_ = nullptr; }; -INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest, - Values( +INSTANTIATE_TEST_CASE_P( + IVFParameters, IVFTest, + Values( #ifdef MILVUS_GPU_VERSION - std::make_tuple("GPUIVF", ParameterType::ivf), - std::make_tuple("GPUIVFPQ", ParameterType::ivfpq), - std::make_tuple("GPUIVFSQ", ParameterType::ivfsq), + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT, milvus::knowhere::IndexMode::MODE_GPU), + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ, milvus::knowhere::IndexMode::MODE_GPU), + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8, milvus::knowhere::IndexMode::MODE_GPU), #ifdef CUSTOMIZATION - std::make_tuple("IVFSQHybrid", ParameterType::ivfsq), + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H, milvus::knowhere::IndexMode::MODE_GPU), #endif #endif - std::make_tuple("IVF", ParameterType::ivf), std::make_tuple("IVFPQ", ParameterType::ivfpq), - std::make_tuple("IVFSQ", ParameterType::ivfsq))); + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8, milvus::knowhere::IndexMode::MODE_CPU))); TEST_P(IVFTest, ivf_basic) { assert(!xb.empty()); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - // null faiss index - ASSERT_ANY_THROW(index_->Add(base_dataset, conf)); - ASSERT_ANY_THROW(index_->AddWithoutIds(base_dataset, conf)); + ASSERT_ANY_THROW(index_->Add(base_dataset, conf_)); + ASSERT_ANY_THROW(index_->AddWithoutIds(base_dataset, conf_)); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); + index_->Train(base_dataset, conf_); + index_->Add(base_dataset, conf_); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); + EXPECT_EQ(index_->Dim(), dim); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + auto result = index_->Query(query_dataset, conf_); + AssertAnns(result, nq, conf_[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); - if (index_type.find("GPU") == std::string::npos && index_type.find("Hybrid") == std::string::npos && - index_type.find("PQ") == std::string::npos) { - auto result2 = index_->SearchById(id_dataset, conf); + if (index_mode_ == milvus::knowhere::IndexMode::MODE_CPU && + index_type_ != milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ) { + auto result2 = index_->QueryById(id_dataset, conf_); AssertAnns(result2, nq, k); - if (index_type.find("IVFSQ") == std::string::npos) { - auto result3 = index_->GetVectorById(xid_dataset, conf); + if (index_type_ != milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8) { + auto result3 = index_->GetVectorById(xid_dataset, conf_); AssertVec(result3, base_dataset, xid_dataset, 1, dim); } else { - auto result3 = index_->GetVectorById(xid_dataset, conf); + auto result3 = index_->GetVectorById(xid_dataset, conf_); /* for SQ8, sometimes the mean diff can bigger than 20% */ // AssertVec(result3, base_dataset, xid_dataset, 1, dim, CheckMode::CHECK_APPROXIMATE_EQUAL); } @@ -132,26 +128,26 @@ TEST_P(IVFTest, ivf_basic) { } index_->SetBlacklist(concurrent_bitset_ptr); - auto result_bs_1 = index_->Search(query_dataset, conf); + auto result_bs_1 = index_->Query(query_dataset, conf_); AssertAnns(result_bs_1, nq, k, CheckMode::CHECK_NOT_EQUAL); // PrintResult(result, nq, k); - auto result_bs_2 = index_->SearchById(id_dataset, conf); + auto result_bs_2 = index_->QueryById(id_dataset, conf_); AssertAnns(result_bs_2, nq, k, CheckMode::CHECK_NOT_EQUAL); // PrintResult(result, nq, k); - auto result_bs_3 = index_->GetVectorById(xid_dataset, conf); + auto result_bs_3 = index_->GetVectorById(xid_dataset, conf_); AssertVec(result_bs_3, base_dataset, xid_dataset, 1, dim, CheckMode::CHECK_NOT_EQUAL); } #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().Dump(); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Dump(); #endif } TEST_P(IVFTest, ivf_serialize) { fiu_init(0); - auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) { + auto serialize = [](const std::string& filename, milvus::knowhere::BinaryPtr& bin, uint8_t* ret) { FileIOWriter writer(filename); writer(static_cast(bin->data.get()), bin->size); @@ -159,45 +155,10 @@ TEST_P(IVFTest, ivf_serialize) { reader(ret, bin->size); }; - { - // null faisss index serialize - ASSERT_ANY_THROW(index_->Serialize()); - knowhere::IVFIndexModel model(nullptr); - ASSERT_ANY_THROW(model.Serialize()); - } - - { - // serialize index-model - auto model = index_->Train(base_dataset, conf); - auto binaryset = model->Serialize(); - auto bin = binaryset.GetByName("IVF"); - - fiu_enable("FaissBaseIndex.SerializeImpl.throw_exception", 1, nullptr, 0); - ASSERT_ANY_THROW(model->Serialize()); - fiu_disable("FaissBaseIndex.SerializeImpl.throw_exception"); - - std::string filename = "/tmp/ivf_test_model_serialize.bin"; - auto load_data = new uint8_t[bin->size]; - serialize(filename, bin, load_data); - - binaryset.clear(); - auto data = std::make_shared(); - data.reset(load_data); - binaryset.Append("IVF", data, bin->size); - - model->Load(binaryset); - - index_->set_index_model(model); - index_->Add(base_dataset, conf); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); - } - { // serialize index - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); + index_->Train(base_dataset, conf_); + index_->Add(base_dataset, conf_); auto binaryset = index_->Serialize(); auto bin = binaryset.GetByName("IVF"); @@ -212,9 +173,9 @@ TEST_P(IVFTest, ivf_serialize) { index_->Load(binaryset); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf_); + AssertAnns(result, nq, conf_[milvus::knowhere::meta::TOPK]); } } @@ -223,21 +184,17 @@ TEST_P(IVFTest, ivf_serialize) { TEST_P(IVFTest, clone_test) { assert(!xb.empty()); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); + index_->Train(base_dataset, conf_); + index_->Add(base_dataset, conf_); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf_); + AssertAnns(result, nq, conf_[milvus::knowhere::meta::TOPK]); // PrintResult(result, nq, k); - auto AssertEqual = [&](knowhere::DatasetPtr p1, knowhere::DatasetPtr p2) { - auto ids_p1 = p1->Get(knowhere::meta::IDS); - auto ids_p2 = p2->Get(knowhere::meta::IDS); + auto AssertEqual = [&](milvus::knowhere::DatasetPtr p1, milvus::knowhere::DatasetPtr p2) { + auto ids_p1 = p1->Get(milvus::knowhere::meta::IDS); + auto ids_p2 = p2->Get(milvus::knowhere::meta::IDS); for (int i = 0; i < nq * k; ++i) { EXPECT_EQ(*((int64_t*)(ids_p2) + i), *((int64_t*)(ids_p1) + i)); @@ -254,7 +211,7 @@ TEST_P(IVFTest, clone_test) { // EXPECT_NO_THROW({ // auto clone_index = index_->Clone(); // auto clone_result = clone_index->Search(query_dataset, conf); - // //AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + // //AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); // AssertEqual(result, clone_result); // std::cout << "inplace clone [" << index_type << "] success" << std::endl; // }); @@ -268,48 +225,33 @@ TEST_P(IVFTest, clone_test) { { // copy from gpu to cpu - std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "GPUIVFPQ", "IVFSQHybrid"}; - auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); - if (finder != support_idx_vec.cend()) { + if (index_mode_ == milvus::knowhere::IndexMode::MODE_GPU) { EXPECT_NO_THROW({ - auto clone_index = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); - auto clone_result = clone_index->Search(query_dataset, conf); + auto clone_index = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); + auto clone_result = clone_index->Query(query_dataset, conf_); AssertEqual(result, clone_result); - std::cout << "clone G <=> C [" << index_type << "] success" << std::endl; + std::cout << "clone G <=> C [" << index_type_ << "] success" << std::endl; }); } else { EXPECT_THROW( { - std::cout << "clone G <=> C [" << index_type << "] failed" << std::endl; - auto clone_index = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + std::cout << "clone G <=> C [" << index_type_ << "] failed" << std::endl; + auto clone_index = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); }, - knowhere::KnowhereException); + milvus::knowhere::KnowhereException); } } - if (index_type == "IVFSQHybrid") { - return; - } - { // copy to gpu - std::vector support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ", "IVFPQ", "GPUIVFPQ"}; - auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); - if (finder != support_idx_vec.cend()) { + if (index_type_ != milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H) { EXPECT_NO_THROW({ - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); - auto clone_result = clone_index->Search(query_dataset, conf); + auto clone_index = milvus::knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, milvus::knowhere::Config()); + auto clone_result = clone_index->Query(query_dataset, conf_); AssertEqual(result, clone_result); - std::cout << "clone C <=> G [" << index_type << "] success" << std::endl; + std::cout << "clone C <=> G [" << index_type_ << "] success" << std::endl; }); - EXPECT_ANY_THROW(knowhere::cloner::CopyCpuToGpu(index_, -1, knowhere::Config())); - } else { - EXPECT_THROW( - { - std::cout << "clone C <=> G [" << index_type << "] failed" << std::endl; - auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config()); - }, - knowhere::KnowhereException); + EXPECT_ANY_THROW(milvus::knowhere::cloner::CopyCpuToGpu(index_, -1, milvus::knowhere::Config())); } } } @@ -319,71 +261,61 @@ TEST_P(IVFTest, clone_test) { #ifdef MILVUS_GPU_VERSION #ifdef CUSTOMIZATION TEST_P(IVFTest, gpu_seal_test) { - std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ"}; - auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); - if (finder == support_idx_vec.cend()) { + if (index_mode_ != milvus::knowhere::IndexMode::MODE_GPU) { return; } - assert(!xb.empty()); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - - ASSERT_ANY_THROW(index_->Search(query_dataset, conf)); + ASSERT_ANY_THROW(index_->Query(query_dataset, conf_)); ASSERT_ANY_THROW(index_->Seal()); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); + index_->Train(base_dataset, conf_); + index_->Add(base_dataset, conf_); EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - auto result = index_->Search(query_dataset, conf); - AssertAnns(result, nq, conf[knowhere::meta::TOPK]); + EXPECT_EQ(index_->Dim(), dim); + auto result = index_->Query(query_dataset, conf_); + AssertAnns(result, nq, conf_[milvus::knowhere::meta::TOPK]); fiu_init(0); fiu_enable("IVF.Search.throw_std_exception", 1, nullptr, 0); - ASSERT_ANY_THROW(index_->Search(query_dataset, conf)); + ASSERT_ANY_THROW(index_->Query(query_dataset, conf_)); fiu_disable("IVF.Search.throw_std_exception"); fiu_enable("IVF.Search.throw_faiss_exception", 1, nullptr, 0); - ASSERT_ANY_THROW(index_->Search(query_dataset, conf)); + ASSERT_ANY_THROW(index_->Query(query_dataset, conf_)); fiu_disable("IVF.Search.throw_faiss_exception"); - auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + auto cpu_idx = milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); + milvus::knowhere::IVFPtr ivf_idx = std::dynamic_pointer_cast(cpu_idx); - knowhere::TimeRecorder tc("CopyToGpu"); - knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + milvus::knowhere::TimeRecorder tc("CopyToGpu"); + milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); auto without_seal = tc.RecordSection("Without seal"); - cpu_idx->Seal(); + ivf_idx->Seal(); tc.RecordSection("seal cost"); - knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config()); + milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, milvus::knowhere::Config()); auto with_seal = tc.RecordSection("With seal"); ASSERT_GE(without_seal, with_seal); // copy to GPU with invalid device id - ASSERT_ANY_THROW(knowhere::cloner::CopyCpuToGpu(cpu_idx, -1, knowhere::Config())); + ASSERT_ANY_THROW(milvus::knowhere::cloner::CopyCpuToGpu(cpu_idx, -1, milvus::knowhere::Config())); } #endif TEST_P(IVFTest, invalid_gpu_source) { - std::vector support_idx_vec{"GPUIVF", "GPUIVFPQ", "GPUIVFSQ"}; - auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); - if (finder == support_idx_vec.cend()) { + if (index_mode_ != milvus::knowhere::IndexMode::MODE_GPU) { return; } - auto invalid_conf = ParamGenerator::GetInstance().Gen(parameter_type_); - invalid_conf[knowhere::meta::DEVICEID] = -1; + auto invalid_conf = ParamGenerator::GetInstance().Gen(index_type_); + invalid_conf[milvus::knowhere::meta::DEVICEID] = -1; - if (index_type == "GPUIVF") { + if (index_type_ == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFFLAT) { // null faiss index - knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config()); + milvus::knowhere::cloner::CopyGpuToCpu(index_, milvus::knowhere::Config()); } - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - auto binaryset = model->Serialize(); + index_->Train(base_dataset, conf_); fiu_init(0); fiu_enable("GPUIVF.SerializeImpl.throw_exception", 1, nullptr, 0); @@ -391,48 +323,42 @@ TEST_P(IVFTest, invalid_gpu_source) { fiu_disable("GPUIVF.SerializeImpl.throw_exception"); fiu_enable("GPUIVF.search_impl.invald_index", 1, nullptr, 0); - ASSERT_ANY_THROW(index_->Search(base_dataset, invalid_conf)); + ASSERT_ANY_THROW(index_->Query(base_dataset, invalid_conf)); fiu_disable("GPUIVF.search_impl.invald_index"); - auto ivf_index = std::dynamic_pointer_cast(index_); + auto ivf_index = std::dynamic_pointer_cast(index_); if (ivf_index) { - auto gpu_index = std::dynamic_pointer_cast(ivf_index); + auto gpu_index = std::dynamic_pointer_cast(ivf_index); gpu_index->SetGpuDevice(-1); ASSERT_EQ(gpu_index->GetGpuDevice(), -1); - ASSERT_ANY_THROW(index_->set_index_model(model)); } - ASSERT_ANY_THROW(index_->Load(binaryset)); + // ASSERT_ANY_THROW(index_->Load(binaryset)); ASSERT_ANY_THROW(index_->Train(base_dataset, invalid_conf)); } #ifdef CUSTOMIZATION TEST_P(IVFTest, IVFSQHybrid_test) { - std::vector support_idx_vec{"IVFSQHybrid"}; - auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); - if (finder == support_idx_vec.cend()) { + if (index_type_ != milvus::knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H) { return; } fiu_init(0); - knowhere::cloner::CopyGpuToCpu(index_, conf); - ASSERT_ANY_THROW(knowhere::cloner::CopyCpuToGpu(index_, -1, conf)); + milvus::knowhere::cloner::CopyGpuToCpu(index_, conf_); + ASSERT_ANY_THROW(milvus::knowhere::cloner::CopyCpuToGpu(index_, -1, conf_)); fiu_enable("FaissGpuResourceMgr.GetRes.ret_null", 1, nullptr, 0); - ASSERT_ANY_THROW(index_->Train(base_dataset, conf)); - ASSERT_ANY_THROW(index_->CopyCpuToGpu(DEVICEID, conf)); + ASSERT_ANY_THROW(index_->Train(base_dataset, conf_)); + ASSERT_ANY_THROW(index_->CopyCpuToGpu(DEVICEID, conf_)); fiu_disable("FaissGpuResourceMgr.GetRes.ret_null"); - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - - auto index = std::dynamic_pointer_cast(index_); + index_->Train(base_dataset, conf_); + auto index = std::dynamic_pointer_cast(index_); ASSERT_TRUE(index != nullptr); ASSERT_ANY_THROW(index->UnsetQuantizer()); ASSERT_ANY_THROW(index->SetQuantizer(nullptr)); } - #endif #endif diff --git a/core/unittest/wrapper/test_knowhere.cpp b/core/src/index/unittest/test_knowhere.cpp similarity index 100% rename from core/unittest/wrapper/test_knowhere.cpp rename to core/src/index/unittest/test_knowhere.cpp index d27910f226..2d6a783247 100644 --- a/core/unittest/wrapper/test_knowhere.cpp +++ b/core/src/index/unittest/test_knowhere.cpp @@ -13,8 +13,8 @@ #include "wrapper/KnowhereResource.h" #include "wrapper/utils.h" -#include #include +#include #include TEST_F(KnowhereTest, KNOWHERE_RESOURCE_TEST) { diff --git a/core/src/index/unittest/test_nsg/CMakeLists.txt b/core/src/index/unittest/test_nsg/CMakeLists.txt index 82dbb5110c..d2ccc7d9e5 100644 --- a/core/src/index/unittest/test_nsg/CMakeLists.txt +++ b/core/src/index/unittest/test_nsg/CMakeLists.txt @@ -12,8 +12,8 @@ else () message(FATAL_ERROR "no OpenMP supprot") endif () -include_directories(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg) -aux_source_directory(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg nsg_src) +include_directories(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/impl/nsg) +aux_source_directory(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/impl/nsg nsg_src) set(interface_src ${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexNSG.cpp diff --git a/core/src/index/unittest/test_nsg/test_nsg.cpp b/core/src/index/unittest/test_nsg/test_nsg.cpp index 074e3f2736..867c84ce9b 100644 --- a/core/src/index/unittest/test_nsg/test_nsg.cpp +++ b/core/src/index/unittest/test_nsg/test_nsg.cpp @@ -9,23 +9,24 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. +#include +#include #include #include #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/FaissBaseIndex.h" #include "knowhere/index/vector_index/IndexNSG.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #ifdef MILVUS_GPU_VERSION -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" +#include "knowhere/index/vector_index/gpu/IndexGPUIDMAP.h" #include "knowhere/index/vector_index/helpers/Cloner.h" #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" #endif #include "knowhere/common/Timer.h" -#include "knowhere/index/vector_index/nsg/NSGIO.h" +#include "knowhere/index/vector_index/impl/nsg/NSGIO.h" -#include -#include #include "unittest/utils.h" using ::testing::Combine; @@ -41,37 +42,37 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test { // Init_with_default(); #ifdef MILVUS_GPU_VERSION int64_t MB = 1024 * 1024; - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, MB * 200, MB * 600, 1); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, MB * 200, MB * 600, 1); #endif Generate(256, 1000000 / 100, 1); - index_ = std::make_shared(); + index_ = std::make_shared(); - train_conf = knowhere::Config{{knowhere::meta::DIM, 256}, - {knowhere::IndexParams::nlist, 163}, - {knowhere::IndexParams::nprobe, 8}, - {knowhere::IndexParams::knng, 20}, - {knowhere::IndexParams::search_length, 40}, - {knowhere::IndexParams::out_degree, 30}, - {knowhere::IndexParams::candidate, 100}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}}; + train_conf = milvus::knowhere::Config{{milvus::knowhere::meta::DIM, 256}, + {milvus::knowhere::IndexParams::nlist, 163}, + {milvus::knowhere::IndexParams::nprobe, 8}, + {milvus::knowhere::IndexParams::knng, 20}, + {milvus::knowhere::IndexParams::search_length, 40}, + {milvus::knowhere::IndexParams::out_degree, 30}, + {milvus::knowhere::IndexParams::candidate, 100}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}}; - search_conf = knowhere::Config{ - {knowhere::meta::TOPK, k}, - {knowhere::IndexParams::search_length, 30}, + search_conf = milvus::knowhere::Config{ + {milvus::knowhere::meta::TOPK, k}, + {milvus::knowhere::IndexParams::search_length, 30}, }; } void TearDown() override { #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().Free(); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif } protected: - std::shared_ptr index_; - knowhere::Config train_conf; - knowhere::Config search_conf; + std::shared_ptr index_; + milvus::knowhere::Config train_conf; + milvus::knowhere::Config search_conf; }; TEST_F(NSGInterfaceTest, basic_test) { @@ -79,14 +80,14 @@ TEST_F(NSGInterfaceTest, basic_test) { fiu_init(0); // untrained index { - ASSERT_ANY_THROW(index_->Search(query_dataset, search_conf)); + ASSERT_ANY_THROW(index_->Query(query_dataset, search_conf)); ASSERT_ANY_THROW(index_->Serialize()); } - // train_conf->gpu_id = knowhere::INVALID_VALUE; + // train_conf->gpu_id = milvus::knowhere::INVALID_VALUE; // auto model_invalid_gpu = index_->Train(base_dataset, train_conf); - train_conf[knowhere::meta::DEVICEID] = DEVICEID; - auto model = index_->Train(base_dataset, train_conf); - auto result = index_->Search(query_dataset, search_conf); + train_conf[milvus::knowhere::meta::DEVICEID] = DEVICEID; + index_->Train(base_dataset, train_conf); + auto result = index_->Query(query_dataset, search_conf); AssertAnns(result, nq, k); auto binaryset = index_->Serialize(); @@ -96,7 +97,7 @@ TEST_F(NSGInterfaceTest, basic_test) { fiu_disable("NSG.Serialize.throw_exception"); } - auto new_index = std::make_shared(); + auto new_index = std::make_shared(); new_index->Load(binaryset); { fiu_enable("NSG.Load.throw_exception", 1, nullptr, 0); @@ -104,23 +105,23 @@ TEST_F(NSGInterfaceTest, basic_test) { fiu_disable("NSG.Load.throw_exception"); } - auto new_result = new_index->Search(query_dataset, search_conf); + auto new_result = new_index->Query(query_dataset, search_conf); AssertAnns(result, nq, k); ASSERT_EQ(index_->Count(), nb); - ASSERT_EQ(index_->Dimension(), dim); - // ASSERT_THROW({ index_->Clone(); }, knowhere::KnowhereException); - ASSERT_NO_THROW({ - index_->Add(base_dataset, knowhere::Config()); - index_->Seal(); - }); + ASSERT_EQ(index_->Dim(), dim); + // ASSERT_THROW({ index_->Clone(); }, milvus::knowhere::KnowhereException); + // ASSERT_NO_THROW({ + // index_->Add(base_dataset, milvus::knowhere::Config()); + // index_->Seal(); + // }); } TEST_F(NSGInterfaceTest, comparetest) { - knowhere::algo::DistanceL2 distanceL2; - knowhere::algo::DistanceIP distanceIP; + milvus::knowhere::impl::DistanceL2 distanceL2; + milvus::knowhere::impl::DistanceIP distanceIP; - knowhere::TimeRecorder tc("Compare"); + milvus::knowhere::TimeRecorder tc("Compare"); for (int i = 0; i < 1000; ++i) { distanceL2.Compare(xb.data(), xq.data(), 256); } @@ -233,7 +234,7 @@ TEST_F(NSGInterfaceTest, comparetest) { // } // } // } -// printf("R@1 = %.4f\n", n_1 / float(nq));; +// printf("R@1 = %.4f\n", n_1 / float(nq)); // printf("R@10 = %.4f\n", n_10 / float(nq)); // printf("R@100 = %.4f\n", n_100 / float(nq)); //} diff --git a/core/src/index/unittest/test_sptag.cpp b/core/src/index/unittest/test_sptag.cpp index bcaebb7da8..a7a51ef1b0 100644 --- a/core/src/index/unittest/test_sptag.cpp +++ b/core/src/index/unittest/test_sptag.cpp @@ -10,14 +10,13 @@ // or implied. See the License for the specific language governing permissions and limitations under the License. #include - #include #include -#include "knowhere/adapter/SptagAdapter.h" -#include "knowhere/adapter/VectorAdapter.h" + #include "knowhere/common/Exception.h" #include "knowhere/index/vector_index/IndexSPTAG.h" -#include "knowhere/index/vector_index/helpers/Definitions.h" +#include "knowhere/index/vector_index/adapter/SptagAdapter.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include "unittest/utils.h" @@ -31,18 +30,18 @@ class SPTAGTest : public DataGen, public TestWithParam { SetUp() override { IndexType = GetParam(); Generate(128, 100, 5); - index_ = std::make_shared(IndexType); + index_ = std::make_shared(IndexType); if (IndexType == "KDT") { - conf = knowhere::Config{ - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, 10}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}, + conf = milvus::knowhere::Config{ + {milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, 10}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}, }; } else { - conf = knowhere::Config{ - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, 10}, - {knowhere::Metric::TYPE, knowhere::Metric::L2}, + conf = milvus::knowhere::Config{ + {milvus::knowhere::meta::DIM, dim}, + {milvus::knowhere::meta::TOPK, 10}, + {milvus::knowhere::Metric::TYPE, milvus::knowhere::Metric::L2}, }; } @@ -50,8 +49,8 @@ class SPTAGTest : public DataGen, public TestWithParam { } protected: - knowhere::Config conf; - std::shared_ptr index_ = nullptr; + milvus::knowhere::Config conf; + std::shared_ptr index_ = nullptr; std::string IndexType; }; @@ -61,18 +60,14 @@ INSTANTIATE_TEST_CASE_P(SPTAGParameters, SPTAGTest, Values("KDT", "BKT")); TEST_P(SPTAGTest, sptag_basic) { assert(!xb.empty()); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - - auto model = index_->Train(base_dataset, conf); - index_->set_index_model(model); - index_->Add(base_dataset, conf); - auto result = index_->Search(query_dataset, conf); + index_->Train(base_dataset, conf); + // index_->Add(base_dataset, conf); + auto result = index_->Query(query_dataset, conf); AssertAnns(result, nq, k); { - auto ids = result->Get(knowhere::meta::IDS); - auto dist = result->Get(knowhere::meta::DISTANCE); + auto ids = result->Get(milvus::knowhere::meta::IDS); + auto dist = result->Get(milvus::knowhere::meta::DISTANCE); std::stringstream ss_id; std::stringstream ss_dist; @@ -89,36 +84,22 @@ TEST_P(SPTAGTest, sptag_basic) { std::cout << "id\n" << ss_id.str() << std::endl; std::cout << "dist\n" << ss_dist.str() << std::endl; } - - // Though these functions do nothing, use them to improve code coverage - { - index_->Seal(); - knowhere::CPUSPTAGRNGIndexModel index_model; - // Function Serialize's implementation do'nt have return value, - // which will cause undefined behavior. - // index_model.Serialize(); - index_model.Load(knowhere::BinarySet()); - } } TEST_P(SPTAGTest, sptag_serialize) { assert(!xb.empty()); - auto preprocessor = index_->BuildPreprocessor(base_dataset, conf); - index_->set_preprocessor(preprocessor); - - auto model = index_->Train(base_dataset, conf); - - index_->Add(base_dataset, conf); + index_->Train(base_dataset, conf); + // index_->Add(base_dataset, conf); auto binaryset = index_->Serialize(); - auto new_index = std::make_shared(IndexType); + auto new_index = std::make_shared(IndexType); new_index->Load(binaryset); - auto result = new_index->Search(query_dataset, conf); + auto result = new_index->Query(query_dataset, conf); AssertAnns(result, nq, k); PrintResult(result, nq, k); ASSERT_EQ(new_index->Count(), nb); - ASSERT_EQ(new_index->Dimension(), dim); - // ASSERT_THROW({ new_index->Clone(); }, knowhere::KnowhereException); + ASSERT_EQ(new_index->Dim(), dim); + // ASSERT_THROW({ new_index->Clone(); }, milvus::knowhere::KnowhereException); // ASSERT_NO_THROW({ new_index->Seal(); }); { @@ -136,7 +117,7 @@ TEST_P(SPTAGTest, sptag_serialize) { ++fileno; } - knowhere::BinarySet load_data_list; + milvus::knowhere::BinarySet load_data_list; for (int i = 0; i < filename_list.size() && i < meta_list.size(); ++i) { auto bin_size = meta_list[i].second; FileIOReader reader(filename_list[i]); @@ -148,9 +129,9 @@ TEST_P(SPTAGTest, sptag_serialize) { load_data_list.Append(meta_list[i].first, data, bin_size); } - auto new_index = std::make_shared(IndexType); + auto new_index = std::make_shared(IndexType); new_index->Load(load_data_list); - auto result = new_index->Search(query_dataset, conf); + auto result = new_index->Query(query_dataset, conf); AssertAnns(result, nq, k); PrintResult(result, nq, k); } diff --git a/core/src/index/unittest/test_vecindex.cpp b/core/src/index/unittest/test_vecindex.cpp new file mode 100644 index 0000000000..287fdb1e62 --- /dev/null +++ b/core/src/index/unittest/test_vecindex.cpp @@ -0,0 +1,121 @@ +// Copyright (C) 2019-2020 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied. See the License for the specific language governing permissions and limitations under the License. + +#include + +#include "knowhere/index/vector_index/IndexType.h" +#include "knowhere/index/vector_index/VecIndex.h" +#include "knowhere/index/vector_index/VecIndexFactory.h" +#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" + +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/helpers/Cloner.h" +#endif + +#include "unittest/Helper.h" +#include "unittest/utils.h" + +using ::testing::Combine; +using ::testing::TestWithParam; +using ::testing::Values; + +class VecIndexTest : public DataGen, public Tuple> { + protected: + void + SetUp() override { +#ifdef MILVUS_GPU_VERSION + milvus::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); +#endif + std::tie(index_type_, index_mode_, parameter_type_) = GetParam(); + Generate(DIM, NB, NQ); + index_ = milvus::knowhere::VecIndexFactory::GetInstance().CreateVecIndex(index_type_, index_mode_); + conf = ParamGenerator::GetInstance().Gen(parameter_type_); + } + + void + TearDown() override { +#ifdef MILVUS_GPU_VERSION + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Free(); +#endif + } + + protected: + milvus::knowhere::IndexType index_type_; + milvus::knowhere::IndexMode index_mode_; + ParameterType parameter_type_; + milvus::knowhere::Config conf; + milvus::knowhere::VecIndexPtr index_ = nullptr; +}; + +INSTANTIATE_TEST_CASE_P( + IVFParameters, IVFTest, + Values( +#ifdef MILVUS_GPU_VERSION + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFFLAT, milvus::knowhere::IndexMode::MODE_GPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFPQ, milvus::knowhere::IndexMode::MODE_GPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFSQ8, milvus::knowhere::IndexMode::MODE_GPU), +#ifdef CUSTOMIZATION + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFSQ8H, milvus::knowhere::IndexMode::MODE_GPU), +#endif +#endif + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFFLAT, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFPQ, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_FAISS_IVFSQ8, milvus::knowhere::IndexMode::MODE_CPU) + std::make_tuple(milvus::knowhere::IndexType::INDEX_NSG, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_HNSW, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_SPTAG_KDT_RNT, milvus::knowhere::IndexMode::MODE_CPU), + std::make_tuple(milvus::knowhere::IndexType::INDEX_SPTAG_BKT_RNT, milvus::knowhere::IndexMode::MODE_CPU))); + +TEST_P(VecIndexTest, basic) { + assert(!xb.empty()); + KNOWHERE_LOG_DEBUG << "conf: " << conf->dump(); + + index_->BuildAll(base_dataset, conf); + EXPECT_EQ(index_->Dim(), dim); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->index_type(), index_type_); + EXPECT_EQ(index_->index_mode(), index_mode_); + + auto result = index_->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); + PrintResult(result, nq, k); +} + +TEST_P(VecIndexTest, serialize) { + index_->BuildAll(base_dataset, conf); + EXPECT_EQ(index_->Dim(), dim); + EXPECT_EQ(index_->Count(), nb); + EXPECT_EQ(index_->index_type(), index_type_); + EXPECT_EQ(index_->index_mode(), index_mode_); + auto result = index_->Query(query_dataset, conf); + AssertAnns(result, nq, conf[milvus::knowhere::meta::TOPK]); + + auto binaryset = index_->Serialize(); + auto new_index = milvus::knowhere::VecIndexFactory::GetInstance().CreateVecIndex(index_type_, index_mode_); + new_index->Load(binaryset); + EXPECT_EQ(index_->Dim(), new_index->Dim()); + EXPECT_EQ(index_->Count(), new_index->Count()); + EXPECT_EQ(index_->index_type(), new_index->index_type()); + EXPECT_EQ(index_->index_mode(), new_index->index_mode()); + auto new_result = new_index_->Query(query_dataset, conf); + AssertAnns(new_result, nq, conf[milvus::knowhere::meta::TOPK]); +} + +// todo +#ifdef MILVUS_GPU_VERSION +TEST_P(VecIndexTest, copytogpu) { + // todo +} + +TEST_P(VecIndexTest, copytocpu) { + // todo +} +#endif diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/src/index/unittest/test_wrapper.cpp similarity index 98% rename from core/unittest/wrapper/test_wrapper.cpp rename to core/src/index/unittest/test_wrapper.cpp index 65e672e8c6..503768d31d 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/src/index/unittest/test_wrapper.cpp @@ -45,13 +45,11 @@ class KnowhereWrapperTest std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam(); GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); - knowhere::Config tempconf{ - {knowhere::Metric::TYPE, knowhere::Metric::L2}, - {knowhere::meta::ROWS, nb}, - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::meta::DEVICEID, DEVICEID} - }; + knowhere::Config tempconf{{knowhere::Metric::TYPE, knowhere::Metric::L2}, + {knowhere::meta::ROWS, nb}, + {knowhere::meta::DIM, dim}, + {knowhere::meta::TOPK, k}, + {knowhere::meta::DEVICEID, DEVICEID}}; index_ = GetVecIndexFactory(index_type); conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf); diff --git a/core/src/index/unittest/utils.cpp b/core/src/index/unittest/utils.cpp index a1d3d746d4..8caabe685b 100644 --- a/core/src/index/unittest/utils.cpp +++ b/core/src/index/unittest/utils.cpp @@ -10,7 +10,7 @@ // or implied. See the License for the specific language governing permissions and limitations under the License. #include "unittest/utils.h" -#include "knowhere/adapter/VectorAdapter.h" +#include "knowhere/index/vector_index/adapter/VectorAdapter.h" #include #include @@ -52,7 +52,7 @@ DataGen::Generate(const int& dim, const int& nb, const int& nq) { query_dataset = generate_query_dataset(nq, dim, xq.data()); id_dataset = generate_id_dataset(nq, ids.data()); xid_dataset = generate_id_dataset(nq, xids.data()); - xid_dataset->Set(knowhere::meta::DIM, (int64_t)dim); + xid_dataset->Set(milvus::knowhere::meta::DIM, (int64_t)dim); } void @@ -66,8 +66,8 @@ BinaryDataGen::Generate(const int& dim, const int& nb, const int& nq) { assert(xb.size() == (size_t)dim_x * nb); assert(xq.size() == (size_t)dim_x * nq); - base_dataset = generate_binary_dataset(nb, dim, xb.data(), ids.data()); - query_dataset = generate_binary_query_dataset(nq, dim, xq.data()); + base_dataset = generate_dataset(nb, dim, xb.data(), ids.data()); + query_dataset = generate_query_dataset(nq, dim, xq.data()); id_dataset = generate_id_dataset(nq, ids.data()); xid_dataset = generate_id_dataset(nq, xids.data()); } @@ -176,55 +176,36 @@ FileIOWriter::operator()(void* ptr, size_t size) { return size; } -knowhere::DatasetPtr -generate_dataset(int64_t nb, int64_t dim, const float* xb, const int64_t* ids) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - ret_ds->Set(knowhere::meta::IDS, ids); +milvus::knowhere::DatasetPtr +generate_dataset(int64_t nb, int64_t dim, const void* xb, const int64_t* ids) { + auto ret_ds = std::make_shared(); + ret_ds->Set(milvus::knowhere::meta::ROWS, nb); + ret_ds->Set(milvus::knowhere::meta::DIM, dim); + ret_ds->Set(milvus::knowhere::meta::TENSOR, xb); + ret_ds->Set(milvus::knowhere::meta::IDS, ids); return ret_ds; } -knowhere::DatasetPtr -generate_binary_dataset(int64_t nb, int64_t dim, const uint8_t* xb, const int64_t* ids) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - ret_ds->Set(knowhere::meta::IDS, ids); +milvus::knowhere::DatasetPtr +generate_query_dataset(int64_t nb, int64_t dim, const void* xb) { + auto ret_ds = std::make_shared(); + ret_ds->Set(milvus::knowhere::meta::ROWS, nb); + ret_ds->Set(milvus::knowhere::meta::DIM, dim); + ret_ds->Set(milvus::knowhere::meta::TENSOR, xb); return ret_ds; } -knowhere::DatasetPtr -generate_query_dataset(int64_t nb, int64_t dim, const float* xb) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - return ret_ds; -} - -knowhere::DatasetPtr +milvus::knowhere::DatasetPtr generate_id_dataset(int64_t nb, const int64_t* ids) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::IDS, ids); - return ret_ds; -} - -knowhere::DatasetPtr -generate_binary_query_dataset(int64_t nb, int64_t dim, const uint8_t* xb) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); + auto ret_ds = std::make_shared(); + ret_ds->Set(milvus::knowhere::meta::ROWS, nb); + ret_ds->Set(milvus::knowhere::meta::IDS, ids); return ret_ds; } void -AssertAnns(const knowhere::DatasetPtr& result, const int nq, const int k, const CheckMode check_mode) { - auto ids = result->Get(knowhere::meta::IDS); +AssertAnns(const milvus::knowhere::DatasetPtr& result, const int nq, const int k, const CheckMode check_mode) { + auto ids = result->Get(milvus::knowhere::meta::IDS); for (auto i = 0; i < nq; i++) { switch (check_mode) { case CheckMode::CHECK_EQUAL: @@ -241,11 +222,11 @@ AssertAnns(const knowhere::DatasetPtr& result, const int nq, const int k, const } void -AssertVec(const knowhere::DatasetPtr& result, const knowhere::DatasetPtr& base_dataset, - const knowhere::DatasetPtr& id_dataset, const int n, const int dim, const CheckMode check_mode) { - auto base = base_dataset->Get(knowhere::meta::TENSOR); - auto ids = id_dataset->Get(knowhere::meta::IDS); - auto x = result->Get(knowhere::meta::TENSOR); +AssertVec(const milvus::knowhere::DatasetPtr& result, const milvus::knowhere::DatasetPtr& base_dataset, + const milvus::knowhere::DatasetPtr& id_dataset, const int n, const int dim, const CheckMode check_mode) { + float* base = (float*)base_dataset->Get(milvus::knowhere::meta::TENSOR); + auto ids = id_dataset->Get(milvus::knowhere::meta::IDS); + auto x = result->Get(milvus::knowhere::meta::TENSOR); for (auto i = 0; i < n; i++) { auto id = ids[i]; for (auto j = 0; j < dim; j++) { @@ -273,11 +254,11 @@ AssertVec(const knowhere::DatasetPtr& result, const knowhere::DatasetPtr& base_d } void -AssertBinVeceq(const knowhere::DatasetPtr& result, const knowhere::DatasetPtr& base_dataset, - const knowhere::DatasetPtr& id_dataset, const int n, const int dim) { - auto base = base_dataset->Get(knowhere::meta::TENSOR); - auto ids = id_dataset->Get(knowhere::meta::IDS); - auto x = result->Get(knowhere::meta::TENSOR); +AssertBinVeceq(const milvus::knowhere::DatasetPtr& result, const milvus::knowhere::DatasetPtr& base_dataset, + const milvus::knowhere::DatasetPtr& id_dataset, const int n, const int dim) { + auto base = base_dataset->Get(milvus::knowhere::meta::TENSOR); + auto ids = id_dataset->Get(milvus::knowhere::meta::IDS); + auto x = result->Get(milvus::knowhere::meta::TENSOR); for (auto i = 0; i < 1; i++) { auto id = ids[i]; for (auto j = 0; j < dim; j++) { @@ -287,9 +268,9 @@ AssertBinVeceq(const knowhere::DatasetPtr& result, const knowhere::DatasetPtr& b } void -PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) { - auto ids = result->Get(knowhere::meta::IDS); - auto dist = result->Get(knowhere::meta::DISTANCE); +PrintResult(const milvus::knowhere::DatasetPtr& result, const int& nq, const int& k) { + auto ids = result->Get(milvus::knowhere::meta::IDS); + auto dist = result->Get(milvus::knowhere::meta::DISTANCE); std::stringstream ss_id; std::stringstream ss_dist; diff --git a/core/src/index/unittest/utils.h b/core/src/index/unittest/utils.h index c94e199e59..57218c7827 100644 --- a/core/src/index/unittest/utils.h +++ b/core/src/index/unittest/utils.h @@ -28,7 +28,7 @@ class DataGen { void Generate(const int& dim, const int& nb, const int& nq); - knowhere::DatasetPtr + milvus::knowhere::DatasetPtr GenQuery(const int& nq); protected: @@ -40,10 +40,10 @@ class DataGen { std::vector xq; std::vector ids; std::vector xids; - knowhere::DatasetPtr base_dataset = nullptr; - knowhere::DatasetPtr query_dataset = nullptr; - knowhere::DatasetPtr id_dataset = nullptr; - knowhere::DatasetPtr xid_dataset = nullptr; + milvus::knowhere::DatasetPtr base_dataset = nullptr; + milvus::knowhere::DatasetPtr query_dataset = nullptr; + milvus::knowhere::DatasetPtr id_dataset = nullptr; + milvus::knowhere::DatasetPtr xid_dataset = nullptr; }; class BinaryDataGen { @@ -54,7 +54,7 @@ class BinaryDataGen { void Generate(const int& dim, const int& nb, const int& nq); - knowhere::DatasetPtr + milvus::knowhere::DatasetPtr GenQuery(const int& nq); protected: @@ -66,10 +66,10 @@ class BinaryDataGen { std::vector xq; std::vector ids; std::vector xids; - knowhere::DatasetPtr base_dataset = nullptr; - knowhere::DatasetPtr query_dataset = nullptr; - knowhere::DatasetPtr id_dataset = nullptr; - knowhere::DatasetPtr xid_dataset = nullptr; + milvus::knowhere::DatasetPtr base_dataset = nullptr; + milvus::knowhere::DatasetPtr query_dataset = nullptr; + milvus::knowhere::DatasetPtr id_dataset = nullptr; + milvus::knowhere::DatasetPtr xid_dataset = nullptr; }; extern void @@ -96,21 +96,15 @@ GenBinaryBase(const int64_t& dim, const int64_t& nb, uint8_t* xb, int64_t* ids); extern void InitLog(); -knowhere::DatasetPtr -generate_dataset(int64_t nb, int64_t dim, const float* xb, const int64_t* ids); +milvus::knowhere::DatasetPtr +generate_dataset(int64_t nb, int64_t dim, const void* xb, const int64_t* ids); -knowhere::DatasetPtr -generate_binary_dataset(int64_t nb, int64_t dim, const uint8_t* xb, const int64_t* ids); +milvus::knowhere::DatasetPtr +generate_query_dataset(int64_t nb, int64_t dim, const void* xb); -knowhere::DatasetPtr -generate_query_dataset(int64_t nb, int64_t dim, const float* xb); - -knowhere::DatasetPtr +milvus::knowhere::DatasetPtr generate_id_dataset(int64_t nb, const int64_t* ids); -knowhere::DatasetPtr -generate_binary_query_dataset(int64_t nb, int64_t dim, const uint8_t* xb); - enum class CheckMode { CHECK_EQUAL = 0, CHECK_NOT_EQUAL = 1, @@ -118,21 +112,21 @@ enum class CheckMode { }; void -AssertAnns(const knowhere::DatasetPtr& result, const int nq, const int k, +AssertAnns(const milvus::knowhere::DatasetPtr& result, const int nq, const int k, const CheckMode check_mode = CheckMode::CHECK_EQUAL); void -AssertVec(const knowhere::DatasetPtr& result, const knowhere::DatasetPtr& base_dataset, - const knowhere::DatasetPtr& id_dataset, const int n, const int dim, +AssertVec(const milvus::knowhere::DatasetPtr& result, const milvus::knowhere::DatasetPtr& base_dataset, + const milvus::knowhere::DatasetPtr& id_dataset, const int n, const int dim, const CheckMode check_mode = CheckMode::CHECK_EQUAL); void -AssertBinVeceq(const knowhere::DatasetPtr& result, const knowhere::DatasetPtr& base_dataset, - const knowhere::DatasetPtr& id_dataset, const int n, const int dim, +AssertBinVeceq(const milvus::knowhere::DatasetPtr& result, const milvus::knowhere::DatasetPtr& base_dataset, + const milvus::knowhere::DatasetPtr& id_dataset, const int n, const int dim, const CheckMode check_mode = CheckMode::CHECK_EQUAL); void -PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k); +PrintResult(const milvus::knowhere::DatasetPtr& result, const int& nq, const int& k); struct FileIOWriter { std::fstream fs; diff --git a/core/src/server/Server.cpp b/core/src/server/Server.cpp index d90eadafa4..339d470e24 100644 --- a/core/src/server/Server.cpp +++ b/core/src/server/Server.cpp @@ -16,6 +16,7 @@ #include #include "config/Config.h" +#include "index/archive/KnowhereResource.h" #include "metrics/Metrics.h" #include "scheduler/SchedInst.h" #include "server/DBWrapper.h" @@ -28,7 +29,6 @@ #include "utils/LogUtil.h" #include "utils/SignalUtil.h" #include "utils/TimeRecorder.h" -#include "wrapper/KnowhereResource.h" namespace milvus { namespace server { diff --git a/core/src/utils/ValidationUtil.cpp b/core/src/utils/ValidationUtil.cpp index 7db74b8c91..97a5ba88ae 100644 --- a/core/src/utils/ValidationUtil.cpp +++ b/core/src/utils/ValidationUtil.cpp @@ -13,9 +13,9 @@ #include "Log.h" #include "db/Utils.h" #include "db/engine/ExecutionEngine.h" -#include "index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h" +#include "knowhere/index/vector_index/ConfAdapter.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "utils/StringHelpFunctions.h" -#include "wrapper/ConfAdapter.h" #include @@ -212,7 +212,7 @@ ValidationUtil::ValidateIndexParams(const milvus::json& index_params, const engi // special check for 'm' parameter std::vector resset; - milvus::engine::IVFPQConfAdapter::GetValidMList(table_schema.dimension_, resset); + milvus::knowhere::IVFPQConfAdapter::GetValidMList(table_schema.dimension_, resset); int64_t m_value = index_params[index_params, knowhere::IndexParams::m]; if (resset.empty()) { std::string msg = "Invalid table dimension, unable to get reasonable values for 'm'"; diff --git a/core/src/wrapper/BinVecImpl.cpp b/core/src/wrapper/BinVecImpl.cpp deleted file mode 100644 index ee219bfb4f..0000000000 --- a/core/src/wrapper/BinVecImpl.cpp +++ /dev/null @@ -1,291 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/BinVecImpl.h" - -#include "WrapperException.h" -#include "index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.h" -#include "knowhere/adapter/VectorAdapter.h" -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexBinaryIDMAP.h" -#include "utils/Log.h" - -namespace milvus { -namespace engine { - -Status -BinVecImpl::BuildAll(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const uint8_t* xt) { - try { - dim = cfg[knowhere::meta::DIM]; - - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - ret_ds->Set(knowhere::meta::IDS, ids); - - index_->Train(ret_ds, cfg); - } 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 -BinVecImpl::Search(const int64_t& nq, const uint8_t* xq, float* dist, int64_t* ids, const Config& cfg) { - try { - int64_t k = cfg[knowhere::meta::TOPK]; - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nq); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xq); - - auto res = index_->Search(ret_ds, cfg); - //{ - // auto& ids = ids_array; - // auto& dists = dis_array; - // std::stringstream ss_id; - // std::stringstream ss_dist; - // for (auto i = 0; i < 10; i++) { - // for (auto j = 0; j < k; ++j) { - // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - // } - // ss_id << std::endl; - // ss_dist << std::endl; - // } - // std::cout << "id\n" << ss_id.str() << std::endl; - // std::cout << "dist\n" << ss_dist.str() << std::endl; - //} - - // auto p_ids = ids_array->data()->GetValues(1, 0); - // auto p_dist = dis_array->data()->GetValues(1, 0); - - // TODO(linxj): avoid copy here. - auto res_ids = res->Get(knowhere::meta::IDS); - auto res_dist = res->Get(knowhere::meta::DISTANCE); - memcpy(ids, res_ids, sizeof(int64_t) * nq * k); - memcpy(dist, res_dist, sizeof(float) * nq * k); - free(res_ids); - free(res_dist); - } 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 -BinVecImpl::Add(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg) { - try { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - ret_ds->Set(knowhere::meta::IDS, ids); - - index_->Add(ret_ds, cfg); - } 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 -BinVecImpl::CopyToGpu(const int64_t& device_id, const Config& cfg) { - const char* errmsg = "Binary Index not support CopyToGpu"; - WRAPPER_LOG_ERROR << errmsg; - throw WrapperException(errmsg); -} - -VecIndexPtr -BinVecImpl::CopyToCpu(const Config& cfg) { - const char* errmsg = "Binary Index not support CopyToCpu"; - WRAPPER_LOG_ERROR << errmsg; - throw WrapperException(errmsg); -} - -Status -BinVecImpl::SetBlacklist(faiss::ConcurrentBitsetPtr list) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->SetBlacklist(list); - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->SetBlacklist(list); - } - return Status::OK(); -} - -Status -BinVecImpl::GetVectorById(const int64_t n, const int64_t* ids, uint8_t* x, const Config& cfg) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else { - throw WrapperException("not support"); - } - try { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, n); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::IDS, ids); - - auto res = index_->GetVectorById(ret_ds, cfg); - - // TODO(linxj): avoid copy here. - auto res_x = res->Get(knowhere::meta::TENSOR); - memcpy(x, res_x, sizeof(uint8_t) * n * dim); - free(res_x); - } 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 -BinVecImpl::SearchById(const int64_t& nq, const int64_t* xq, float* dist, int64_t* ids, const Config& cfg) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else { - throw WrapperException("not support"); - } - try { - int64_t k = cfg[knowhere::meta::TOPK]; - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nq); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::IDS, xq); - - auto res = index_->SearchById(ret_ds, cfg); - //{ - // auto& ids = ids_array; - // auto& dists = dis_array; - // std::stringstream ss_id; - // std::stringstream ss_dist; - // for (auto i = 0; i < 10; i++) { - // for (auto j = 0; j < k; ++j) { - // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - // } - // ss_id << std::endl; - // ss_dist << std::endl; - // } - // std::cout << "id\n" << ss_id.str() << std::endl; - // std::cout << "dist\n" << ss_dist.str() << std::endl; - //} - - // auto p_ids = ids_array->data()->GetValues(1, 0); - // auto p_dist = dis_array->data()->GetValues(1, 0); - - // TODO(linxj): avoid copy here. - auto res_ids = res->Get(knowhere::meta::IDS); - auto res_dist = res->Get(knowhere::meta::DISTANCE); - memcpy(ids, res_ids, sizeof(int64_t) * nq * k); - memcpy(dist, res_dist, sizeof(float) * nq * k); - free(res_ids); - free(res_dist); - } 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 -BinVecImpl::GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->GetBlacklist(list); - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->GetBlacklist(list); - } - return Status::OK(); -} - -ErrorCode -BinBFIndex::Build(const Config& cfg) { - try { - dim = cfg[knowhere::meta::DIM]; - std::static_pointer_cast(index_)->Train(cfg); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return KNOWHERE_UNEXPECTED_ERROR; - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return KNOWHERE_ERROR; - } - return KNOWHERE_SUCCESS; -} - -Status -BinBFIndex::BuildAll(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const uint8_t* xt) { - try { - dim = cfg[knowhere::meta::DIM]; - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - ret_ds->Set(knowhere::meta::IDS, ids); - - std::static_pointer_cast(index_)->Train(cfg); - index_->Add(ret_ds, cfg); - } 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(); -} - -const uint8_t* -BinBFIndex::GetRawVectors() { - auto raw_index = std::dynamic_pointer_cast(index_); - if (raw_index) { - return raw_index->GetRawVectors(); - } - return nullptr; -} - -const int64_t* -BinBFIndex::GetRawIds() { - return std::static_pointer_cast(index_)->GetRawIds(); -} - -Status -BinBFIndex::AddWithoutIds(const int64_t& nb, const uint8_t* xb, const Config& cfg) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - std::static_pointer_cast(index_)->AddWithoutId(ret_ds, cfg); - return Status::OK(); -} -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/BinVecImpl.h b/core/src/wrapper/BinVecImpl.h deleted file mode 100644 index 2d3d8d0836..0000000000 --- a/core/src/wrapper/BinVecImpl.h +++ /dev/null @@ -1,84 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include -#include - -#include "VecImpl.h" - -namespace milvus { -namespace engine { - -class BinVecImpl : public VecIndexImpl { - public: - explicit BinVecImpl(std::shared_ptr index, const IndexType& type) - : VecIndexImpl(std::move(index), type) { - } - - Status - BuildAll(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const uint8_t* xt) override; - Status - Search(const int64_t& nq, const uint8_t* xq, float* dist, int64_t* ids, const Config& cfg) override; - - Status - Add(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg) override; - - VecIndexPtr - CopyToGpu(const int64_t& device_id, const Config& cfg) override; - - Status - SetBlacklist(faiss::ConcurrentBitsetPtr list) override; - - Status - GetBlacklist(faiss::ConcurrentBitsetPtr& list) override; - - // Status - // SearchById(const int64_t& nq, const uint8_t* xq, faiss::ConcurrentBitsetPtr bitset, float* dist, int64_t* ids, - // const Config& cfg) override; - - Status - GetVectorById(const int64_t n, const int64_t* xid, uint8_t* x, const Config& cfg) override; - - Status - SearchById(const int64_t& nq, const int64_t* xq, float* dist, int64_t* ids, const Config& cfg) override; - - VecIndexPtr - CopyToCpu(const Config& cfg) override; -}; - -class BinBFIndex : public BinVecImpl { - public: - explicit BinBFIndex(std::shared_ptr index) - : BinVecImpl(std::move(index), IndexType::FAISS_BIN_IDMAP) { - } - - ErrorCode - Build(const Config& cfg); - - Status - BuildAll(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const uint8_t* xt) override; - - const uint8_t* - GetRawVectors(); - - const int64_t* - GetRawIds(); - - Status - AddWithoutIds(const int64_t& nb, const uint8_t* xb, const Config& cfg); -}; - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/ConfAdapterMgr.cpp b/core/src/wrapper/ConfAdapterMgr.cpp deleted file mode 100644 index 298a3a588e..0000000000 --- a/core/src/wrapper/ConfAdapterMgr.cpp +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/ConfAdapterMgr.h" - -#include "utils/Exception.h" - -namespace milvus { -namespace engine { - -ConfAdapterPtr -AdapterMgr::GetAdapter(const IndexType& indexType) { - if (!init_) - RegisterAdapter(); - - auto it = table_.find(indexType); - if (it != table_.end()) { - return it->second(); - } else { - throw Exception(KNOWHERE_INVALID_ARGUMENT, "Can not find this type of confadapter"); - } -} - -#define REGISTER_CONF_ADAPTER(T, KEY, NAME) static AdapterMgr::register_t reg_##NAME##_(KEY) - -void -AdapterMgr::RegisterAdapter() { - init_ = true; - - REGISTER_CONF_ADAPTER(ConfAdapter, IndexType::FAISS_IDMAP, idmap); - REGISTER_CONF_ADAPTER(BinIDMAPConfAdapter, IndexType::FAISS_BIN_IDMAP, idmap_bin); - - REGISTER_CONF_ADAPTER(IVFConfAdapter, IndexType::FAISS_IVFFLAT_CPU, ivf_cpu); - REGISTER_CONF_ADAPTER(IVFConfAdapter, IndexType::FAISS_IVFFLAT_GPU, ivf_gpu); - REGISTER_CONF_ADAPTER(IVFConfAdapter, IndexType::FAISS_IVFFLAT_MIX, ivf_mix); - REGISTER_CONF_ADAPTER(BinIVFConfAdapter, IndexType::FAISS_BIN_IVFLAT_CPU, ivf_bin_cpu); - - REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexType::FAISS_IVFSQ8_CPU, ivfsq8_cpu); - REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexType::FAISS_IVFSQ8_GPU, ivfsq8_gpu); - REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexType::FAISS_IVFSQ8_MIX, ivfsq8_mix); - REGISTER_CONF_ADAPTER(IVFSQConfAdapter, IndexType::FAISS_IVFSQ8_HYBRID, ivfsq8_h); - - REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_CPU, ivfpq_cpu); - REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_GPU, ivfpq_gpu); - REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix); - - REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix); - - REGISTER_CONF_ADAPTER(ConfAdapter, IndexType::SPTAG_KDT_RNT_CPU, sptag_kdt); - REGISTER_CONF_ADAPTER(ConfAdapter, IndexType::SPTAG_BKT_RNT_CPU, sptag_bkt); - - REGISTER_CONF_ADAPTER(HNSWConfAdapter, IndexType::HNSW, hnsw); -} - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/DataTransfer.cpp b/core/src/wrapper/DataTransfer.cpp deleted file mode 100644 index 6dfa8b1d43..0000000000 --- a/core/src/wrapper/DataTransfer.cpp +++ /dev/null @@ -1,41 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/DataTransfer.h" - -#include - -#include "knowhere/adapter/VectorAdapter.h" - -namespace milvus { -namespace engine { - -knowhere::DatasetPtr -GenDatasetWithIds(const int64_t& nb, const int64_t& dim, const float* xb, const int64_t* ids) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - ret_ds->Set(knowhere::meta::IDS, ids); - return ret_ds; -} - -knowhere::DatasetPtr -GenDataset(const int64_t& nb, const int64_t& dim, const float* xb) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::DIM, dim); - ret_ds->Set(knowhere::meta::TENSOR, xb); - return ret_ds; -} - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/DataTransfer.h b/core/src/wrapper/DataTransfer.h deleted file mode 100644 index 74e3bb0849..0000000000 --- a/core/src/wrapper/DataTransfer.h +++ /dev/null @@ -1,26 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include "knowhere/common/Dataset.h" - -namespace milvus { -namespace engine { - -extern knowhere::DatasetPtr -GenDatasetWithIds(const int64_t& nb, const int64_t& dim, const float* xb, const int64_t* ids); - -extern knowhere::DatasetPtr -GenDataset(const int64_t& nb, const int64_t& dim, const float* xb); - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/VecImpl.cpp b/core/src/wrapper/VecImpl.cpp deleted file mode 100644 index 71fc59c083..0000000000 --- a/core/src/wrapper/VecImpl.cpp +++ /dev/null @@ -1,379 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/VecImpl.h" - -#include "DataTransfer.h" -#include "knowhere/adapter/VectorAdapter.h" -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexIDMAP.h" -#include "utils/Log.h" -#include "wrapper/WrapperException.h" -#include "wrapper/gpu/GPUVecImpl.h" - -#ifdef MILVUS_GPU_VERSION - -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" -#include "knowhere/index/vector_index/helpers/Cloner.h" - -#endif - -#include -/* - * no parameter check in this layer. - * only responsible for index combination - */ - -namespace milvus { -namespace engine { - -Status -VecIndexImpl::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[knowhere::meta::DIM]; - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - fiu_do_on("VecIndexImpl.BuildAll.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("VecIndexImpl.BuildAll.throw_std_exception", throw std::exception()); - - 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); - } 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 -VecIndexImpl::Add(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg) { - try { - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - fiu_do_on("VecIndexImpl.Add.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("VecIndexImpl.Add.throw_std_exception", throw std::exception()); - index_->Add(dataset, cfg); - } 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 -VecIndexImpl::Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids, const Config& cfg) { - try { - int64_t k = cfg[knowhere::meta::TOPK]; - auto dataset = GenDataset(nq, dim, xq); - - fiu_do_on("VecIndexImpl.Search.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("VecIndexImpl.Search.throw_std_exception", throw std::exception()); - - auto res = index_->Search(dataset, cfg); - //{ - // auto& ids = ids_array; - // auto& dists = dis_array; - // std::stringstream ss_id; - // std::stringstream ss_dist; - // for (auto i = 0; i < 10; i++) { - // for (auto j = 0; j < k; ++j) { - // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - // } - // ss_id << std::endl; - // ss_dist << std::endl; - // } - // std::cout << "id\n" << ss_id.str() << std::endl; - // std::cout << "dist\n" << ss_dist.str() << std::endl; - //} - - // auto p_ids = ids_array->data()->GetValues(1, 0); - // auto p_dist = dis_array->data()->GetValues(1, 0); - - // TODO(linxj): avoid copy here. - auto res_ids = res->Get(knowhere::meta::IDS); - auto res_dist = res->Get(knowhere::meta::DISTANCE); - memcpy(ids, res_ids, sizeof(int64_t) * nq * k); - memcpy(dist, res_dist, sizeof(float) * nq * k); - free(res_ids); - free(res_dist); - } 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(); -} - -knowhere::BinarySet -VecIndexImpl::Serialize() { - type = ConvertToCpuIndexType(type); - return index_->Serialize(); -} - -Status -VecIndexImpl::Load(const knowhere::BinarySet& index_binary) { - index_->Load(index_binary); - dim = Dimension(); - return Status::OK(); -} - -int64_t -VecIndexImpl::Dimension() { - return index_->Dimension(); -} - -int64_t -VecIndexImpl::Count() { - return index_->Count(); -} - -IndexType -VecIndexImpl::GetType() const { - return type; -} - -VecIndexPtr -VecIndexImpl::CopyToGpu(const int64_t& device_id, const Config& cfg) { -// TODO(linxj): exception handle -#ifdef MILVUS_GPU_VERSION - auto gpu_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, cfg); - auto new_index = std::make_shared(gpu_index, ConvertToGpuIndexType(type)); - new_index->dim = dim; - return new_index; -#else - WRAPPER_LOG_ERROR << "Calling VecIndexImpl::CopyToGpu when we are using CPU version"; - throw WrapperException("Calling VecIndexImpl::CopyToGpu when we are using CPU version"); -#endif -} - -VecIndexPtr -VecIndexImpl::CopyToCpu(const Config& cfg) { -// TODO(linxj): exception handle -#ifdef MILVUS_GPU_VERSION - auto cpu_index = knowhere::cloner::CopyGpuToCpu(index_, cfg); - auto new_index = std::make_shared(cpu_index, ConvertToCpuIndexType(type)); - new_index->dim = dim; - return new_index; -#else - WRAPPER_LOG_ERROR << "Calling VecIndexImpl::CopyToCpu when we are using CPU version"; - throw WrapperException("Calling VecIndexImpl::CopyToCpu when we are using CPU version"); -#endif -} - -// VecIndexPtr -// VecIndexImpl::Clone() { -// // TODO(linxj): exception handle -// auto clone_index = std::make_shared(index_->Clone(), type); -// clone_index->dim = dim; -// return clone_index; -//} - -int64_t -VecIndexImpl::GetDeviceId() { -#ifdef MILVUS_GPU_VERSION - if (auto device_idx = std::dynamic_pointer_cast(index_)) { - return device_idx->GetGpuDevice(); - } -#else - // else - return -1; // -1 == cpu -#endif -} - -Status -VecIndexImpl::GetVectorById(const int64_t n, const int64_t* xid, float* x, const Config& cfg) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else { - throw WrapperException("not support"); - } - - try { - auto dataset = std::make_shared(); - dataset->Set(knowhere::meta::ROWS, n); - dataset->Set(knowhere::meta::DIM, dim); - dataset->Set(knowhere::meta::IDS, xid); - - auto res = index_->GetVectorById(dataset, cfg); - - // TODO(linxj): avoid copy here. - auto res_x = res->Get(knowhere::meta::TENSOR); - memcpy(x, res_x, sizeof(float) * n * dim); - free(res_x); - } 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 -VecIndexImpl::SearchById(const int64_t& nq, const int64_t* xq, float* dist, int64_t* ids, const Config& cfg) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - } else { - throw WrapperException("not support"); - } - - try { - int64_t k = cfg[knowhere::meta::TOPK]; - auto dataset = std::make_shared(); - dataset->Set(knowhere::meta::ROWS, nq); - dataset->Set(knowhere::meta::DIM, dim); - dataset->Set(knowhere::meta::IDS, xq); - - auto res = index_->SearchById(dataset, cfg); - //{ - // auto& ids = ids_array; - // auto& dists = dis_array; - // std::stringstream ss_id; - // std::stringstream ss_dist; - // for (auto i = 0; i < 10; i++) { - // for (auto j = 0; j < k; ++j) { - // ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; - // ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; - // } - // ss_id << std::endl; - // ss_dist << std::endl; - // } - // std::cout << "id\n" << ss_id.str() << std::endl; - // std::cout << "dist\n" << ss_dist.str() << std::endl; - //} - - // auto p_ids = ids_array->data()->GetValues(1, 0); - // auto p_dist = dis_array->data()->GetValues(1, 0); - - // TODO(linxj): avoid copy here. - auto res_ids = res->Get(knowhere::meta::IDS); - auto res_dist = res->Get(knowhere::meta::DISTANCE); - memcpy(ids, res_ids, sizeof(int64_t) * nq * k); - memcpy(dist, res_dist, sizeof(float) * nq * k); - free(res_ids); - free(res_dist); - } 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 -VecIndexImpl::SetBlacklist(faiss::ConcurrentBitsetPtr list) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->SetBlacklist(list); - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->SetBlacklist(list); - } - return Status::OK(); -} - -Status -VecIndexImpl::GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->GetBlacklist(list); - } else if (auto raw_index = std::dynamic_pointer_cast(index_)) { - raw_index->GetBlacklist(list); - } - return Status::OK(); -} - -Status -VecIndexImpl::SetUids(std::vector& uids) { - index_->SetUids(uids); - return Status::OK(); -} - -const std::vector& -VecIndexImpl::GetUids() const { - return index_->GetUids(); -} - -const float* -BFIndex::GetRawVectors() { - auto raw_index = std::dynamic_pointer_cast(index_); - if (raw_index) { - return raw_index->GetRawVectors(); - } - return nullptr; -} - -const int64_t* -BFIndex::GetRawIds() { - return std::static_pointer_cast(index_)->GetRawIds(); -} - -ErrorCode -BFIndex::Build(const Config& cfg) { - try { - fiu_do_on("BFIndex.Build.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("BFIndex.Build.throw_std_exception", throw std::exception()); - dim = cfg[knowhere::meta::DIM]; - std::static_pointer_cast(index_)->Train(cfg); - } catch (knowhere::KnowhereException& e) { - WRAPPER_LOG_ERROR << e.what(); - return KNOWHERE_UNEXPECTED_ERROR; - } catch (std::exception& e) { - WRAPPER_LOG_ERROR << e.what(); - return KNOWHERE_ERROR; - } - return KNOWHERE_SUCCESS; -} - -Status -BFIndex::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[knowhere::meta::DIM]; - auto dataset = GenDatasetWithIds(nb, dim, xb, ids); - fiu_do_on("BFIndex.BuildAll.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("BFIndex.BuildAll.throw_std_exception", throw std::exception()); - - std::static_pointer_cast(index_)->Train(cfg); - index_->Add(dataset, cfg); - } 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 -BFIndex::AddWithoutIds(const int64_t& nb, const float* xb, const Config& cfg) { - auto ret_ds = std::make_shared(); - ret_ds->Set(knowhere::meta::ROWS, nb); - ret_ds->Set(knowhere::meta::TENSOR, xb); - std::static_pointer_cast(index_)->AddWithoutId(ret_ds, cfg); - return Status::OK(); -} - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/VecImpl.h b/core/src/wrapper/VecImpl.h deleted file mode 100644 index b0cd27a82f..0000000000 --- a/core/src/wrapper/VecImpl.h +++ /dev/null @@ -1,131 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include -#include -#include - -#include "VecIndex.h" -#include "knowhere/index/vector_index/VectorIndex.h" - -namespace milvus { -namespace engine { - -class VecIndexImpl : public VecIndex { - public: - explicit VecIndexImpl(std::shared_ptr index, const IndexType& type) - : index_(std::move(index)), type(type) { - } - - Status - BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) override; - - VecIndexPtr - CopyToGpu(const int64_t& device_id, const Config& cfg) override; - - VecIndexPtr - CopyToCpu(const Config& cfg) override; - - IndexType - GetType() const override; - - int64_t - Dimension() override; - - int64_t - Count() override; - - Status - Add(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg) override; - - knowhere::BinarySet - Serialize() override; - - Status - Load(const knowhere::BinarySet& index_binary) override; - - // VecIndexPtr - // Clone() override; - - int64_t - GetDeviceId() override; - - Status - Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids, const Config& cfg) override; - - Status - GetVectorById(const int64_t n, const int64_t* xid, float* x, const Config& cfg) override; - - Status - SearchById(const int64_t& nq, const int64_t* xq, float* dist, int64_t* ids, const Config& cfg) override; - - Status - SetBlacklist(faiss::ConcurrentBitsetPtr list) override; - - Status - GetBlacklist(faiss::ConcurrentBitsetPtr& list) override; - - Status - SetUids(std::vector& uids) override; - - const std::vector& - GetUids() const override; - - protected: - int64_t dim = 0; - - IndexType type = IndexType::INVALID; - - std::shared_ptr index_ = nullptr; -}; - -class BFIndex : public VecIndexImpl { - public: - explicit BFIndex(std::shared_ptr index) - : VecIndexImpl(std::move(index), IndexType::FAISS_IDMAP) { - } - - ErrorCode - Build(const Config& cfg); - - const float* - GetRawVectors(); - - Status - BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) override; - - const int64_t* - GetRawIds(); - - Status - AddWithoutIds(const int64_t& nb, const float* xb, const Config& cfg); -}; - -class ToIndexData : public cache::DataObj { - public: - explicit ToIndexData(int64_t size) : size_(size) { - } - - int64_t - Size() override { - return size_; - } - - private: - int64_t size_ = 0; -}; - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/VecIndex.cpp b/core/src/wrapper/VecIndex.cpp deleted file mode 100644 index 868567ca36..0000000000 --- a/core/src/wrapper/VecIndex.cpp +++ /dev/null @@ -1,337 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/VecIndex.h" - -#include "config/Config.h" -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexBinaryIDMAP.h" -#include "knowhere/index/vector_index/IndexBinaryIVF.h" -#include "knowhere/index/vector_index/IndexHNSW.h" -#include "knowhere/index/vector_index/IndexIDMAP.h" -#include "knowhere/index/vector_index/IndexIVF.h" -#include "knowhere/index/vector_index/IndexIVFPQ.h" -#include "knowhere/index/vector_index/IndexIVFSQ.h" -#include "knowhere/index/vector_index/IndexNSG.h" -#include "knowhere/index/vector_index/IndexSPTAG.h" -#include "storage/disk/DiskIOReader.h" -#include "storage/disk/DiskIOWriter.h" -#include "storage/s3/S3IOReader.h" -#include "storage/s3/S3IOWriter.h" -#include "utils/Exception.h" -#include "utils/Log.h" -#include "utils/TimeRecorder.h" -#include "wrapper/BinVecImpl.h" -#include "wrapper/VecImpl.h" - -#ifdef MILVUS_GPU_VERSION -#include - -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" -#include "knowhere/index/vector_index/IndexGPUIVFSQ.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" -#include "wrapper/gpu/GPUVecImpl.h" -#endif - -#include - -namespace milvus { -namespace engine { - -int64_t -VecIndex::Size() { - return size_; -} - -void -VecIndex::set_size(int64_t size) { - size_ = size; -} - -VecIndexPtr -GetVecIndexFactory(const IndexType& type, const Config& cfg) { - std::shared_ptr index; - auto gpu_device = -1; // TODO(linxj): remove hardcode here - switch (type) { - case IndexType::FAISS_IDMAP: { - index = std::make_shared(); - return std::make_shared(index); - } - case IndexType::FAISS_BIN_IDMAP: { - index = std::make_shared(); - return std::make_shared(index); - } - case IndexType::FAISS_IVFFLAT_CPU: { - index = std::make_shared(); - break; - } - case IndexType::FAISS_BIN_IVFLAT_CPU: { - index = std::make_shared(); - return std::make_shared(index, type); - } - case IndexType::FAISS_IVFPQ_CPU: { - index = std::make_shared(); - break; - } - case IndexType::SPTAG_KDT_RNT_CPU: { - index = std::make_shared("KDT"); - break; - } - case IndexType::SPTAG_BKT_RNT_CPU: { - index = std::make_shared("BKT"); - break; - } - case IndexType::FAISS_IVFSQ8_CPU: { - index = std::make_shared(); - break; - } - case IndexType::HNSW: { - index = std::make_shared(); - break; - } - -#ifdef MILVUS_GPU_VERSION - case IndexType::FAISS_IVFFLAT_GPU: { - index = std::make_shared(gpu_device); - break; - } - case IndexType::FAISS_IVFFLAT_MIX: { - index = std::make_shared(gpu_device); - return std::make_shared(index, IndexType::FAISS_IVFFLAT_MIX); - } - case IndexType::FAISS_IVFPQ_GPU: { - index = std::make_shared(gpu_device); - break; - } - case IndexType::FAISS_IVFPQ_MIX: { - index = std::make_shared(gpu_device); - return std::make_shared(index, IndexType::FAISS_IVFPQ_MIX); - } - case IndexType::FAISS_IVFSQ8_MIX: { - index = std::make_shared(gpu_device); - return std::make_shared(index, IndexType::FAISS_IVFSQ8_MIX); - } - case IndexType::FAISS_IVFSQ8_GPU: { - index = std::make_shared(gpu_device); - break; - } - -#endif -#ifdef CUSTOMIZATION -#ifdef MILVUS_GPU_VERSION - case IndexType::FAISS_IVFSQ8_HYBRID: { - server::Config& config = server::Config::GetInstance(); - bool gpu_resource_enable = true; - config.GetGpuResourceConfigEnable(gpu_resource_enable); - if (gpu_resource_enable) { - index = std::make_shared(gpu_device); - fiu_do_on("GetVecIndexFactory.IVFSQHybrid.mock", index = std::make_shared()); - return std::make_shared(index, IndexType::FAISS_IVFSQ8_HYBRID); - } else { - throw Exception(DB_ERROR, "No GPU resources for IndexType::FAISS_IVFSQ8_HYBRID"); - } - } -#endif -#endif - case IndexType::NSG_MIX: { // TODO(linxj): bug. - index = std::make_shared(gpu_device); - break; - } - default: { return nullptr; } - } - return std::make_shared(index, type); -} - -VecIndexPtr -LoadVecIndex(const IndexType& index_type, const knowhere::BinarySet& index_binary, int64_t size) { - auto index = GetVecIndexFactory(index_type); - if (index == nullptr) - return nullptr; - // else - index->Load(index_binary); - index->set_size(size); - return index; -} - -VecIndexPtr -read_index(const std::string& location) { - fiu_return_on("read_null_index", nullptr); - fiu_do_on("vecIndex.throw_read_exception", throw std::exception()); - TimeRecorder recorder("read_index"); - knowhere::BinarySet load_data_list; - - bool s3_enable = false; - server::Config& config = server::Config::GetInstance(); - config.GetStorageConfigS3Enable(s3_enable); - - std::shared_ptr reader_ptr; - if (s3_enable) { - reader_ptr = std::make_shared(); - } else { - reader_ptr = std::make_shared(); - } - - recorder.RecordSection("Start"); - reader_ptr->open(location); - - size_t length = reader_ptr->length(); - if (length <= 0) { - return nullptr; - } - - size_t rp = 0; - reader_ptr->seekg(0); - - auto current_type = IndexType::INVALID; - reader_ptr->read(¤t_type, sizeof(current_type)); - rp += sizeof(current_type); - reader_ptr->seekg(rp); - - while (rp < length) { - size_t meta_length; - reader_ptr->read(&meta_length, sizeof(meta_length)); - rp += sizeof(meta_length); - reader_ptr->seekg(rp); - - auto meta = new char[meta_length]; - reader_ptr->read(meta, meta_length); - rp += meta_length; - reader_ptr->seekg(rp); - - size_t bin_length; - reader_ptr->read(&bin_length, sizeof(bin_length)); - rp += sizeof(bin_length); - reader_ptr->seekg(rp); - - auto bin = new uint8_t[bin_length]; - reader_ptr->read(bin, bin_length); - rp += bin_length; - reader_ptr->seekg(rp); - - auto binptr = std::make_shared(); - binptr.reset(bin); - load_data_list.Append(std::string(meta, meta_length), binptr, bin_length); - delete[] meta; - } - - reader_ptr->close(); - - double span = recorder.RecordSection("End"); - double rate = length * 1000000.0 / span / 1024 / 1024; - STORAGE_LOG_DEBUG << "read_index(" << location << ") rate " << rate << "MB/s"; - - return LoadVecIndex(current_type, load_data_list, length); -} - -Status -write_index(VecIndexPtr index, const std::string& location) { - try { - TimeRecorder recorder("write_index"); - - auto binaryset = index->Serialize(); - auto index_type = index->GetType(); - - fiu_do_on("VecIndex.write_index.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("VecIndex.write_index.throw_std_exception", throw std::exception()); - fiu_do_on("VecIndex.write_index.throw_no_space_exception", - throw Exception(SERVER_INVALID_ARGUMENT, "No space left on device")); - - bool s3_enable = false; - server::Config& config = server::Config::GetInstance(); - config.GetStorageConfigS3Enable(s3_enable); - - std::shared_ptr writer_ptr; - if (s3_enable) { - writer_ptr = std::make_shared(); - } else { - writer_ptr = std::make_shared(); - } - - recorder.RecordSection("Start"); - writer_ptr->open(location); - - writer_ptr->write(&index_type, sizeof(IndexType)); - - for (auto& iter : binaryset.binary_map_) { - auto meta = iter.first.c_str(); - size_t meta_length = iter.first.length(); - writer_ptr->write(&meta_length, sizeof(meta_length)); - writer_ptr->write((void*)meta, meta_length); - - auto binary = iter.second; - int64_t binary_length = binary->size; - writer_ptr->write(&binary_length, sizeof(binary_length)); - writer_ptr->write((void*)binary->data.get(), binary_length); - } - - writer_ptr->close(); - - double span = recorder.RecordSection("End"); - double rate = writer_ptr->length() * 1000000.0 / span / 1024 / 1024; - STORAGE_LOG_DEBUG << "write_index(" << location << ") rate " << rate << "MB/s"; - } 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(); - std::string estring(e.what()); - if (estring.find("No space left on device") != estring.npos) { - WRAPPER_LOG_ERROR << "No space left on the device"; - return Status(KNOWHERE_NO_SPACE, "No space left on the device"); - } else { - return Status(KNOWHERE_ERROR, e.what()); - } - } - return Status::OK(); -} - -IndexType -ConvertToCpuIndexType(const IndexType& type) { - // TODO(linxj): add IDMAP - switch (type) { - case IndexType::FAISS_IVFFLAT_GPU: - case IndexType::FAISS_IVFFLAT_MIX: { - return IndexType::FAISS_IVFFLAT_CPU; - } - case IndexType::FAISS_IVFSQ8_GPU: - case IndexType::FAISS_IVFSQ8_MIX: { - return IndexType::FAISS_IVFSQ8_CPU; - } - case IndexType::FAISS_IVFPQ_GPU: - case IndexType::FAISS_IVFPQ_MIX: { - return IndexType::FAISS_IVFPQ_CPU; - } - default: { return type; } - } -} - -IndexType -ConvertToGpuIndexType(const IndexType& type) { - switch (type) { - case IndexType::FAISS_IVFFLAT_MIX: - case IndexType::FAISS_IVFFLAT_CPU: { - return IndexType::FAISS_IVFFLAT_GPU; - } - case IndexType::FAISS_IVFSQ8_MIX: - case IndexType::FAISS_IVFSQ8_CPU: { - return IndexType::FAISS_IVFSQ8_GPU; - } - case IndexType::FAISS_IVFPQ_MIX: - case IndexType::FAISS_IVFPQ_CPU: { - return IndexType::FAISS_IVFPQ_GPU; - } - default: { return type; } - } -} -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/VecIndex.h b/core/src/wrapper/VecIndex.h deleted file mode 100644 index 7a98e971fe..0000000000 --- a/core/src/wrapper/VecIndex.h +++ /dev/null @@ -1,231 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include - -#include -#include -#include -#include -#include - -#include "cache/DataObj.h" -#include "knowhere/common/BinarySet.h" -#include "knowhere/common/Config.h" -#include "knowhere/index/vector_index/Quantizer.h" -#include "segment/Types.h" -#include "utils/Log.h" -#include "utils/Status.h" - -namespace milvus { -namespace engine { - -using json = nlohmann::json; -using Config = json; - -// TODO(linxj): replace with string, Do refactor serialization -enum class IndexType { - INVALID = 0, - FAISS_IDMAP = 1, - FAISS_IVFFLAT_CPU, - FAISS_IVFFLAT_GPU, - FAISS_IVFFLAT_MIX, // build on gpu and search on cpu - FAISS_IVFPQ_CPU, - FAISS_IVFPQ_GPU, - SPTAG_KDT_RNT_CPU, - FAISS_IVFSQ8_MIX, - FAISS_IVFSQ8_CPU, - FAISS_IVFSQ8_GPU, - FAISS_IVFSQ8_HYBRID, // only support build on gpu. - NSG_MIX, - FAISS_IVFPQ_MIX, - SPTAG_BKT_RNT_CPU, - HNSW, - FAISS_BIN_IDMAP = 100, - FAISS_BIN_IVFLAT_CPU = 101, -}; - -class VecIndex; - -using VecIndexPtr = std::shared_ptr; - -class VecIndex : public cache::DataObj { - public: - virtual Status - BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt = 0, - const float* xt = nullptr) = 0; - - virtual Status - BuildAll(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg, const int64_t& nt = 0, - const uint8_t* xt = nullptr) { - ENGINE_LOG_ERROR << "BuildAll with uint8_t not support"; - return Status::OK(); - } - - virtual Status - Add(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg = Config()) = 0; - - virtual Status - Add(const int64_t& nb, const uint8_t* xb, const int64_t* ids, const Config& cfg = Config()) { - ENGINE_LOG_ERROR << "Add with uint8_t not support"; - return Status::OK(); - } - - virtual Status - Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids, const Config& cfg = Config()) = 0; - - virtual Status - Search(const int64_t& nq, const uint8_t* xq, float* dist, int64_t* ids, const Config& cfg = Config()) { - ENGINE_LOG_ERROR << "Search with uint8_t not support"; - return Status::OK(); - } - - virtual VecIndexPtr - CopyToGpu(const int64_t& device_id, const Config& cfg = Config()) = 0; - - virtual VecIndexPtr - CopyToCpu(const Config& cfg = Config()) = 0; - - // TODO(linxj): Deprecated - // virtual VecIndexPtr - // Clone() = 0; - - virtual int64_t - GetDeviceId() = 0; - - virtual IndexType - GetType() const = 0; - - virtual int64_t - Dimension() = 0; - - virtual int64_t - Count() = 0; - - int64_t - Size() override; - - void - set_size(int64_t size); - - virtual knowhere::BinarySet - Serialize() = 0; - - virtual Status - Load(const knowhere::BinarySet& index_binary) = 0; - - // TODO(linxj): refactor later - //////////////// - virtual knowhere::QuantizerPtr - LoadQuantizer(const Config& conf) { - ENGINE_LOG_ERROR << "LoadQuantizer virtual funciton called."; - return nullptr; - } - - virtual VecIndexPtr - LoadData(const knowhere::QuantizerPtr& q, const Config& conf) { - return nullptr; - } - - virtual Status - SetQuantizer(const knowhere::QuantizerPtr& q) { - return Status::OK(); - } - - virtual Status - UnsetQuantizer() { - return Status::OK(); - } - - virtual std::pair - CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg = Config()) { - return std::make_pair(nullptr, nullptr); - } - //////////////// - - virtual Status - GetVectorById(const int64_t n, const int64_t* xid, float* x, const Config& cfg = Config()) { - ENGINE_LOG_ERROR << "GetVectorById not support"; - return Status::OK(); - } - - virtual Status - GetVectorById(const int64_t n, const int64_t* xid, uint8_t* x, const Config& cfg = Config()) { - ENGINE_LOG_ERROR << "GetVectorById not support"; - return Status::OK(); - } - - virtual Status - SearchById(const int64_t& nq, const int64_t* xq, float* dist, int64_t* ids, const Config& cfg = Config()) { - ENGINE_LOG_ERROR << "SearchById not support"; - return Status::OK(); - } - // - // virtual Status - // SearchById(const int64_t& nq, const uint8_t * xq, faiss::ConcurrentBitsetPtr bitset, float* dist, int64_t* - // ids, const Config& cfg = Config()){ - // ENGINE_LOG_ERROR << "SearchById with uint8_t not support"; - // return Status::OK(); - // } - - virtual Status - SetBlacklist(faiss::ConcurrentBitsetPtr list) { - // ENGINE_LOG_ERROR << "SetBlacklist not support"; - return Status::OK(); - } - - virtual Status - GetBlacklist(faiss::ConcurrentBitsetPtr& list) { - // ENGINE_LOG_ERROR << "GetBlacklist not support"; - ENGINE_LOG_WARNING << "Deletion on unsupported index type"; - return Status::OK(); - } - - virtual Status - SetUids(std::vector& uids) { - ENGINE_LOG_ERROR << "SetUIDArray not support"; - return Status::OK(); - } - - virtual const std::vector& - GetUids() const { - ENGINE_LOG_ERROR << "GetUIDArray not support"; - } - - private: - int64_t size_ = 0; -}; - -extern Status -write_index(VecIndexPtr index, const std::string& location); - -extern VecIndexPtr -read_index(const std::string& location); - -VecIndexPtr -read_index(const std::string& location, knowhere::BinarySet& index_binary); - -extern VecIndexPtr -GetVecIndexFactory(const IndexType& type, const Config& cfg = Config()); - -extern VecIndexPtr -LoadVecIndex(const IndexType& index_type, const knowhere::BinarySet& index_binary, int64_t size); - -extern IndexType -ConvertToCpuIndexType(const IndexType& type); - -extern IndexType -ConvertToGpuIndexType(const IndexType& type); - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/WrapperException.cpp b/core/src/wrapper/WrapperException.cpp deleted file mode 100644 index 9a8ebde02d..0000000000 --- a/core/src/wrapper/WrapperException.cpp +++ /dev/null @@ -1,26 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/WrapperException.h" - -namespace milvus { -namespace engine { - -WrapperException::WrapperException(const std::string& msg) : msg(msg) { -} - -const char* -WrapperException::what() const noexcept { - return msg.c_str(); -} - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/gpu/GPUVecImpl.cpp b/core/src/wrapper/gpu/GPUVecImpl.cpp deleted file mode 100644 index edf8f05c85..0000000000 --- a/core/src/wrapper/gpu/GPUVecImpl.cpp +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "wrapper/gpu/GPUVecImpl.h" - -#include "knowhere/common/Exception.h" -#include "knowhere/index/vector_index/IndexGPUIDMAP.h" -#include "knowhere/index/vector_index/IndexGPUIVF.h" -#include "knowhere/index/vector_index/IndexIDMAP.h" -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" -#include "knowhere/index/vector_index/helpers/Cloner.h" -#include "src/wrapper/DataTransfer.h" -#include "utils/Log.h" -#include "wrapper/VecImpl.h" - -#include -/* - * no parameter check in this layer. - * only responible for index combination - */ - -namespace milvus { -namespace engine { - -// 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 { - fiu_do_on("IVFMixIndex.BuildAll.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("IVFMixIndex.BuildAll.throw_std_exception", throw std::exception()); - - dim = cfg[knowhere::meta::DIM]; - 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(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(); -} - -#ifdef CUSTOMIZATION -knowhere::QuantizerPtr -IVFHybridIndex::LoadQuantizer(const Config& conf) { - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - return new_idx->LoadQuantizer(conf); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); - return nullptr; - } -} - -Status -IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) { - try { - fiu_do_on("IVFHybridIndex.SetQuantizer.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("IVFHybridIndex.SetQuantizer.throw_std_exception", throw std::exception()); - - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - new_idx->SetQuantizer(q); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); - return Status(KNOWHERE_ERROR, "not support"); - } - } 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 { - fiu_do_on("IVFHybridIndex.UnsetQuantizer.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("IVFHybridIndex.UnsetQuantizer.throw_std_exception", throw std::exception()); - - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - new_idx->UnsetQuantizer(); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type); - return Status(KNOWHERE_ERROR, "not support"); - } - } 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 { - fiu_do_on("IVFHybridIndex.LoadData.throw_knowhere_exception", throw knowhere::KnowhereException("")); - fiu_do_on("IVFHybridIndex.LoadData.throw_std_exception", throw std::exception()); - - // TODO(linxj): Hardcode here - if (auto new_idx = std::dynamic_pointer_cast(index_)) { - return std::make_shared(new_idx->LoadData(q, conf), type); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support 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 -IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) { - try { - fiu_do_on("IVFHybridIndex.CopyToGpuWithQuantizer.throw_knowhere_exception", - throw knowhere::KnowhereException("")); - fiu_do_on("IVFHybridIndex.CopyToGpuWithQuantizer.throw_std_exception", throw std::exception()); - - // TODO(linxj): Hardcode here - if (auto hybrid_idx = std::dynamic_pointer_cast(index_)) { - auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg); - auto new_idx = std::make_shared(pair.first, type); - return std::make_pair(new_idx, pair.second); - } else { - WRAPPER_LOG_ERROR << "Hybrid mode not support 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); -} -#endif - -} // namespace engine -} // namespace milvus diff --git a/core/src/wrapper/gpu/GPUVecImpl.h b/core/src/wrapper/gpu/GPUVecImpl.h deleted file mode 100644 index b55c4505cf..0000000000 --- a/core/src/wrapper/gpu/GPUVecImpl.h +++ /dev/null @@ -1,61 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include -#include - -#include "knowhere/index/vector_index/VectorIndex.h" -#include "wrapper/VecImpl.h" -#include "wrapper/VecIndex.h" - -namespace milvus { -namespace engine { - -class IVFMixIndex : public VecIndexImpl { - public: - explicit IVFMixIndex(std::shared_ptr index, const IndexType& type) - : VecIndexImpl(std::move(index), type) { - } - - Status - BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt, - const float* xt) override; - - Status - Load(const knowhere::BinarySet& index_binary) override; -}; - -class IVFHybridIndex : public IVFMixIndex { - public: - explicit IVFHybridIndex(std::shared_ptr index, const IndexType& type) - : IVFMixIndex(std::move(index), type) { - } - - knowhere::QuantizerPtr - LoadQuantizer(const Config& conf) override; - - Status - SetQuantizer(const knowhere::QuantizerPtr& q) override; - - Status - UnsetQuantizer() override; - - std::pair - CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) override; - - VecIndexPtr - LoadData(const knowhere::QuantizerPtr& q, const Config& conf) override; -}; - -} // namespace engine -} // namespace milvus diff --git a/core/unittest/CMakeLists.txt b/core/unittest/CMakeLists.txt index 6674fdf032..5a2bfc3689 100644 --- a/core/unittest/CMakeLists.txt +++ b/core/unittest/CMakeLists.txt @@ -94,7 +94,7 @@ set(server_delivery_files ${server_delivery_request_files}) aux_source_directory(${MILVUS_ENGINE_SRC}/utils utils_files) -aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper wrapper_files) +aux_source_directory(${MILVUS_ENGINE_SRC}/index/archive wrapper_files) aux_source_directory(${MILVUS_ENGINE_SRC}/tracing tracing_files) aux_source_directory(${MILVUS_ENGINE_SRC}/storage storage_main_files) @@ -190,14 +190,9 @@ if (MILVUS_GPU_VERSION) cudart cublas ) - aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper/gpu wrapper_gpu_files) - set(common_files ${common_files} - ${wrapper_gpu_files} - ) endif () add_subdirectory(db) -add_subdirectory(wrapper) add_subdirectory(metrics) add_subdirectory(scheduler) add_subdirectory(server) diff --git a/core/unittest/db/test_engine.cpp b/core/unittest/db/test_engine.cpp index 4a99b8fb9a..3ff8b476bb 100644 --- a/core/unittest/db/test_engine.cpp +++ b/core/unittest/db/test_engine.cpp @@ -134,21 +134,21 @@ TEST_F(EngineTest, FACTORY_TEST) { { fiu_init(0); - // test ExecutionEngineImpl constructor when create vecindex failed - FIU_ENABLE_FIU("ExecutionEngineImpl.CreatetVecIndex.invalid_type"); + // test ExecutionEngineImpl constructor when create VecIndex failed + FIU_ENABLE_FIU("ExecutionEngineImpl.CreateVecIndex.invalid_type"); ASSERT_ANY_THROW(milvus::engine::EngineFactory::Build( 512, "/tmp/milvus_index_1", milvus::engine::EngineType::SPTAG_KDT, milvus::engine::MetricType::L2, index_params)); - fiu_disable("ExecutionEngineImpl.CreatetVecIndex.invalid_type"); + fiu_disable("ExecutionEngineImpl.CreateVecIndex.invalid_type"); } { - // test ExecutionEngineImpl constructor when build BFindex failed - FIU_ENABLE_FIU("BFIndex.Build.throw_knowhere_exception"); + // test ExecutionEngineImpl constructor when build failed + FIU_ENABLE_FIU("ExecutionEngineImpl.throw_exception"); ASSERT_ANY_THROW(milvus::engine::EngineFactory::Build( 512, "/tmp/milvus_index_1", milvus::engine::EngineType::SPTAG_KDT, milvus::engine::MetricType::L2, index_params)); - fiu_disable("BFIndex.Build.throw_knowhere_exception"); + fiu_disable("ExecutionEngineImpl.throw_exception"); } } @@ -165,9 +165,6 @@ TEST_F(EngineTest, ENGINE_IMPL_TEST) { ASSERT_EQ(engine_ptr->IndexMetricType(), milvus::engine::MetricType::IP); ASSERT_ANY_THROW(engine_ptr->BuildIndex(INIT_PATH, milvus::engine::EngineType::INVALID)); - FIU_ENABLE_FIU("VecIndexImpl.BuildAll.throw_knowhere_exception"); - ASSERT_ANY_THROW(engine_ptr->BuildIndex(INIT_PATH, milvus::engine::EngineType::SPTAG_KDT)); - fiu_disable("VecIndexImpl.BuildAll.throw_knowhere_exception"); auto engine_build = engine_ptr->BuildIndex("/tmp/milvus_index_2", milvus::engine::EngineType::FAISS_IVFSQ8); ASSERT_NE(engine_build, nullptr); diff --git a/core/unittest/db/utils.cpp b/core/unittest/db/utils.cpp index 3d9863f159..68ec958c21 100644 --- a/core/unittest/db/utils.cpp +++ b/core/unittest/db/utils.cpp @@ -133,7 +133,7 @@ BaseTest::SetUp() { auto trace_context = std::make_shared(mock_span); dummy_context_->SetTraceContext(trace_context); #ifdef MILVUS_GPU_VERSION - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(0, 1024 * 1024 * 200, 1024 * 1024 * 300, 2); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(0, 1024 * 1024 * 200, 1024 * 1024 * 300, 2); #endif } @@ -142,7 +142,7 @@ BaseTest::TearDown() { milvus::cache::CpuCacheMgr::GetInstance()->ClearCache(); #ifdef MILVUS_GPU_VERSION milvus::cache::GpuCacheMgr::GetInstance(0)->ClearCache(); - knowhere::FaissGpuResourceMgr::GetInstance().Free(); + milvus::knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif } diff --git a/core/unittest/scheduler/test_scheduler.cpp b/core/unittest/scheduler/test_scheduler.cpp index 6d33150c25..5b5d8bb1bf 100644 --- a/core/unittest/scheduler/test_scheduler.cpp +++ b/core/unittest/scheduler/test_scheduler.cpp @@ -22,74 +22,73 @@ #include "scheduler/task/TestTask.h" #include "scheduler/tasklabel/SpecResLabel.h" #include "utils/Error.h" -#include "wrapper/VecIndex.h" namespace milvus { namespace scheduler { -class MockVecIndex : public engine::VecIndex { - public: - virtual Status - BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const engine::Config& cfg, const int64_t& nt = 0, - const float* xt = nullptr) { - } - - // engine::VecIndexPtr - // Clone() override { - // return milvus::engine::VecIndexPtr(); - // } - - int64_t - GetDeviceId() override { - return 0; - } - - engine::IndexType - GetType() const override { - return engine::IndexType::INVALID; - } - - virtual Status - Add(const int64_t& nb, const float* xb, const int64_t* ids, const engine::Config& cfg = engine::Config()) { - } - - virtual Status - Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids, - const engine::Config& cfg = engine::Config()) { - } - - engine::VecIndexPtr - CopyToGpu(const int64_t& device_id, const engine::Config& cfg) override { - } - - engine::VecIndexPtr - CopyToCpu(const engine::Config& cfg) override { - } - - virtual int64_t - Dimension() { - return dimension_; - } - - virtual int64_t - Count() { - return ntotal_; - } - - virtual knowhere::BinarySet - Serialize() { - knowhere::BinarySet binset; - return binset; - } - - virtual Status - Load(const knowhere::BinarySet& index_binary) { - } - - public: - int64_t dimension_ = 512; - int64_t ntotal_ = 0; -}; +//class MockVecIndex : public engine::VecIndex { +// public: +// virtual Status +// BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const engine::Config& cfg, const int64_t& nt = 0, +// const float* xt = nullptr) { +// } +// +// // engine::VecIndexPtr +// // Clone() override { +// // return milvus::engine::VecIndexPtr(); +// // } +// +// int64_t +// GetDeviceId() override { +// return 0; +// } +// +// engine::IndexType +// GetType() const override { +// return engine::IndexType::INVALID; +// } +// +// virtual Status +// Add(const int64_t& nb, const float* xb, const int64_t* ids, const engine::Config& cfg = engine::Config()) { +// } +// +// virtual Status +// Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids, +// const engine::Config& cfg = engine::Config()) { +// } +// +// engine::VecIndexPtr +// CopyToGpu(const int64_t& device_id, const engine::Config& cfg) override { +// } +// +// engine::VecIndexPtr +// CopyToCpu(const engine::Config& cfg) override { +// } +// +// virtual int64_t +// Dimension() { +// return dimension_; +// } +// +// virtual int64_t +// Count() { +// return ntotal_; +// } +// +// virtual knowhere::BinarySet +// Serialize() { +// knowhere::BinarySet binset; +// return binset; +// } +// +// virtual Status +// Load(const knowhere::BinarySet& index_binary) { +// } +// +// public: +// int64_t dimension_ = 512; +// int64_t ntotal_ = 0; +//}; class SchedulerTest : public testing::Test { protected: @@ -209,47 +208,47 @@ class SchedulerTest2 : public testing::Test { // ASSERT_EQ(res_mgr_->GetResource(ResourceType::GPU, 1)->task_table().Size(), NUM); //} -TEST(SchedulerTestResource, SPECIFIED_RESOURCE_TEST) { - auto mock_index_ptr = std::make_shared(); - milvus::engine::Config config; - auto quantizer_ptr = mock_index_ptr->LoadQuantizer(config); - ASSERT_EQ(quantizer_ptr, nullptr); - - auto vec_index_ptr = mock_index_ptr->LoadData(quantizer_ptr, config); - ASSERT_EQ(vec_index_ptr, nullptr); - - auto s = mock_index_ptr->SetQuantizer(quantizer_ptr); - ASSERT_TRUE(s.ok()); - - s = mock_index_ptr->UnsetQuantizer(); - ASSERT_TRUE(s.ok()); - - auto res = mock_index_ptr->CopyToGpuWithQuantizer(0, config); - ASSERT_EQ(res.first, nullptr); - ASSERT_EQ(res.second, nullptr); - - using IndexType = milvus::engine::IndexType; - auto index = GetVecIndexFactory(IndexType::SPTAG_KDT_RNT_CPU, config); - ASSERT_EQ(index->GetType(), IndexType::SPTAG_KDT_RNT_CPU); - - index = GetVecIndexFactory(IndexType::SPTAG_BKT_RNT_CPU, config); - ASSERT_EQ(index->GetType(), IndexType::SPTAG_BKT_RNT_CPU); - -#ifdef MILVUS_GPU_VERSION - index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_GPU, config); - ASSERT_EQ(index->GetType(), IndexType::FAISS_IVFPQ_GPU); -#endif - - index = GetVecIndexFactory(IndexType::NSG_MIX, config); - ASSERT_EQ(index->GetType(), IndexType::NSG_MIX); - - index = GetVecIndexFactory(IndexType::INVALID, config); - ASSERT_EQ(index, nullptr); - - knowhere::BinarySet empty_set; - auto res_ptr = LoadVecIndex(IndexType::INVALID, empty_set, 0); - ASSERT_EQ(res_ptr, nullptr); -} +//TEST(SchedulerTestResource, SPECIFIED_RESOURCE_TEST) { +// auto mock_index_ptr = std::make_shared(); +// milvus::engine::Config config; +// auto quantizer_ptr = mock_index_ptr->LoadQuantizer(config); +// ASSERT_EQ(quantizer_ptr, nullptr); +// +// auto vec_index_ptr = mock_index_ptr->LoadData(quantizer_ptr, config); +// ASSERT_EQ(vec_index_ptr, nullptr); +// +// auto s = mock_index_ptr->SetQuantizer(quantizer_ptr); +// ASSERT_TRUE(s.ok()); +// +// s = mock_index_ptr->UnsetQuantizer(); +// ASSERT_TRUE(s.ok()); +// +// auto res = mock_index_ptr->CopyToGpuWithQuantizer(0, config); +// ASSERT_EQ(res.first, nullptr); +// ASSERT_EQ(res.second, nullptr); +// +// using IndexType = milvus::engine::IndexType; +// auto index = GetVecIndexFactory(IndexType::SPTAG_KDT_RNT_CPU, config); +// ASSERT_EQ(index->GetType(), IndexType::SPTAG_KDT_RNT_CPU); +// +// index = GetVecIndexFactory(IndexType::SPTAG_BKT_RNT_CPU, config); +// ASSERT_EQ(index->GetType(), IndexType::SPTAG_BKT_RNT_CPU); +// +//#ifdef MILVUS_GPU_VERSION +// index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_GPU, config); +// ASSERT_EQ(index->GetType(), IndexType::FAISS_IVFPQ_GPU); +//#endif +// +// index = GetVecIndexFactory(IndexType::NSG_MIX, config); +// ASSERT_EQ(index->GetType(), IndexType::NSG_MIX); +// +// index = GetVecIndexFactory(IndexType::INVALID, config); +// ASSERT_EQ(index, nullptr); +// +// knowhere::BinarySet empty_set; +// auto res_ptr = LoadVecIndex(IndexType::INVALID, empty_set, 0); +// ASSERT_EQ(res_ptr, nullptr); +//} TEST_F(SchedulerTest, schedule) { scheduler_->Dump(); diff --git a/core/unittest/server/test_cache.cpp b/core/unittest/server/test_cache.cpp index f4e85a64a0..8e022cb38e 100644 --- a/core/unittest/server/test_cache.cpp +++ b/core/unittest/server/test_cache.cpp @@ -12,11 +12,11 @@ #include #include #include -#include "utils/Error.h" -#include "wrapper/VecIndex.h" #include "cache/CpuCacheMgr.h" #include "cache/GpuCacheMgr.h" +#include "knowhere/index/vector_index/VecIndex.h" +#include "utils/Error.h" namespace { @@ -33,59 +33,39 @@ class LessItemCacheMgr : public milvus::cache::CacheMgr(256, 1000000); + milvus::knowhere::VecIndexPtr mock_index = std::make_shared(256, 1000000); milvus::cache::DataObjPtr data_obj = std::static_pointer_cast(mock_index); cpu_mgr->InsertItem("index_" + std::to_string(i), data_obj); } @@ -169,7 +150,7 @@ TEST(CacheTest, CPU_CACHE_TEST) { cpu_mgr->SetCapacity(g_num * gbyte); // each vector is 1k byte, total size less than 6G - milvus::engine::VecIndexPtr mock_index = std::make_shared(256, 6000000); + milvus::knowhere::VecIndexPtr mock_index = std::make_shared(256, 6000000); milvus::cache::DataObjPtr data_obj = std::static_pointer_cast(mock_index); cpu_mgr->InsertItem("index_6g", data_obj); ASSERT_TRUE(cpu_mgr->ItemExists("index_6g")); @@ -195,7 +176,7 @@ TEST(CacheTest, GPU_CACHE_TEST) { for (int i = 0; i < 20; i++) { // each vector is 1k byte - milvus::engine::VecIndexPtr mock_index = std::make_shared(256, 1000); + milvus::knowhere::VecIndexPtr mock_index = std::make_shared(256, 1000); milvus::cache::DataObjPtr data_obj = std::static_pointer_cast(mock_index); gpu_mgr->InsertItem("index_" + std::to_string(i), data_obj); } @@ -208,7 +189,7 @@ TEST(CacheTest, GPU_CACHE_TEST) { for (auto i = 0; i < 3; i++) { // TODO(myh): use gpu index to mock // each vector is 1k byte, total size less than 2G - milvus::engine::VecIndexPtr mock_index = std::make_shared(256, 2000000); + milvus::knowhere::VecIndexPtr mock_index = std::make_shared(256, 2000000); milvus::cache::DataObjPtr data_obj = std::static_pointer_cast(mock_index); std::cout << data_obj->Size() << std::endl; gpu_mgr->InsertItem("index_" + std::to_string(i), data_obj); @@ -245,7 +226,7 @@ TEST(CacheTest, INVALID_TEST) { LessItemCacheMgr mgr; for (int i = 0; i < 20; i++) { // each vector is 1k byte - milvus::engine::VecIndexPtr mock_index = std::make_shared(256, 2); + milvus::knowhere::VecIndexPtr mock_index = std::make_shared(256, 2); milvus::cache::DataObjPtr data_obj = std::static_pointer_cast(mock_index); mgr.InsertItem("index_" + std::to_string(i), data_obj); } diff --git a/core/unittest/server/test_web.cpp b/core/unittest/server/test_web.cpp index 9aa53d780d..3a87d641c3 100644 --- a/core/unittest/server/test_web.cpp +++ b/core/unittest/server/test_web.cpp @@ -41,7 +41,7 @@ #include "unittest/server/utils.h" #include "utils/CommonUtil.h" #include "version.h" -#include "wrapper/VecIndex.h" +//#include "knowhere/VecIndex.h" static const char* TABLE_NAME = "test_web"; diff --git a/core/unittest/wrapper/CMakeLists.txt b/core/unittest/wrapper/CMakeLists.txt deleted file mode 100644 index 487477f3da..0000000000 --- a/core/unittest/wrapper/CMakeLists.txt +++ /dev/null @@ -1,52 +0,0 @@ -#------------------------------------------------------------------------------- -# Copyright (C) 2019-2020 Zilliz. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software distributed under the License -# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing permissions and limitations under the License. -#------------------------------------------------------------------------------- - -set(test_files - test_knowhere.cpp - test_binindex.cpp - test_wrapper.cpp) -if (MILVUS_GPU_VERSION) - set(test_files ${test_files} - test_hybrid_index.cpp) -endif () - -set(wrapper_files - ${MILVUS_ENGINE_SRC}/wrapper/DataTransfer.cpp - ${MILVUS_ENGINE_SRC}/wrapper/VecImpl.cpp - ${MILVUS_ENGINE_SRC}/wrapper/VecIndex.cpp - ) - -set(storage_files - ${MILVUS_ENGINE_SRC}/storage/disk/DiskIOReader.cpp - ${MILVUS_ENGINE_SRC}/storage/disk/DiskIOWriter.cpp - ) - -set(util_files - utils.cpp - ${MILVUS_THIRDPARTY_SRC}/easyloggingpp/easylogging++.cc - ${MILVUS_ENGINE_SRC}/utils/Status.cpp - ) - -add_executable(test_wrapper - ${test_files} - ${wrapper_files} - ${storage_files} - ${util_files} - ${common_files}) - -target_link_libraries(test_wrapper - knowhere - metrics - ${unittest_libs}) - -install(TARGETS test_wrapper DESTINATION unittest) diff --git a/core/unittest/wrapper/test_binindex.cpp b/core/unittest/wrapper/test_binindex.cpp deleted file mode 100644 index 9c6f626475..0000000000 --- a/core/unittest/wrapper/test_binindex.cpp +++ /dev/null @@ -1,111 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include -#include "knowhere/index/vector_index/helpers/IndexParameter.h" -#include "wrapper/VecIndex.h" -#include "wrapper/utils.h" - -#include - -using ::testing::Combine; -using ::testing::TestWithParam; -using ::testing::Values; - -class BinIndexTest : public BinDataGen, - public TestWithParam<::std::tuple> { - protected: - void - SetUp() override { - std::tie(index_type, dim, nb, nq, k) = GetParam(); - Generate(dim, nb, nq, k); - - knowhere::Config tempconf{ - {knowhere::Metric::TYPE, knowhere::Metric::TANIMOTO}, - {knowhere::meta::ROWS, nb}, - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k} - }; - - index_ = GetVecIndexFactory(index_type); - conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf); - searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf); - } - - void - TearDown() override { - } - - protected: - milvus::engine::IndexType index_type; - milvus::engine::VecIndexPtr index_ = nullptr; - knowhere::Config conf; - knowhere::Config searchconf; -}; - -INSTANTIATE_TEST_CASE_P(WrapperParam, BinIndexTest, - Values( - //["Index type", "dim", "nb", "nq", "k", "build config", "search config"] - std::make_tuple(milvus::engine::IndexType::FAISS_BIN_IDMAP, 64, 1000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_BIN_IVFLAT_CPU, DIM, NB, 10, 10))); - -TEST_P(BinIndexTest, BASE_TEST) { - EXPECT_EQ(index_->GetType(), index_type); - // conf->Dump(); - // searchconf->Dump(); - - auto elems = nq * k; - std::vector res_ids(elems); - std::vector res_dis(elems); - - index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); - AssertResult(res_ids, res_dis); -} - -TEST_P(BinIndexTest, SERIALIZE_TEST) { - EXPECT_EQ(index_->GetType(), index_type); - - auto elems = nq * k; - std::vector res_ids(elems); - std::vector res_dis(elems); - index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); - AssertResult(res_ids, res_dis); - - { - auto binary = index_->Serialize(); - auto type = index_->GetType(); - auto new_index = GetVecIndexFactory(type); - new_index->Load(binary); - EXPECT_EQ(new_index->Dimension(), index_->Dimension()); - EXPECT_EQ(new_index->Count(), index_->Count()); - - std::vector res_ids(elems); - std::vector res_dis(elems); - new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); - AssertResult(res_ids, res_dis); - } - - { - std::string file_location = "/tmp/knowhere"; - write_index(index_, file_location); - auto new_index = milvus::engine::read_index(file_location); - EXPECT_EQ(new_index->GetType(), ConvertToCpuIndexType(index_type)); - EXPECT_EQ(new_index->Dimension(), index_->Dimension()); - EXPECT_EQ(new_index->Count(), index_->Count()); - - std::vector res_ids(elems); - std::vector res_dis(elems); - new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); - AssertResult(res_ids, res_dis); - } -} diff --git a/core/unittest/wrapper/test_hybrid_index.cpp b/core/unittest/wrapper/test_hybrid_index.cpp deleted file mode 100644 index 0488163107..0000000000 --- a/core/unittest/wrapper/test_hybrid_index.cpp +++ /dev/null @@ -1,208 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h" -#include "knowhere/index/vector_index/helpers/IndexParameter.h" -#include "wrapper/VecIndex.h" -#include "wrapper/utils.h" - -#include -#include -#include -#include "knowhere/index/vector_index/IndexIVFSQHybrid.h" - -using ::testing::Combine; -using ::testing::TestWithParam; -using ::testing::Values; - -class KnowhereHybrid : public DataGenBase, public ::testing::Test { - protected: - void - SetUp() override { - knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); - - dim = 128; - nb = 10000; - nq = 100; - k = 100; - GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); - } - - void - TearDown() override { - knowhere::FaissGpuResourceMgr::GetInstance().Free(); - } - - protected: - milvus::engine::IndexType index_type; - milvus::engine::VecIndexPtr index_ = nullptr; - knowhere::Config conf; -}; - -#ifdef CUSTOMIZATION -TEST_F(KnowhereHybrid, test_interface) { - assert(!xb.empty()); - - knowhere::Config tempconf{ - {knowhere::Metric::TYPE, knowhere::Metric::L2}, - {knowhere::meta::ROWS, nb}, - {knowhere::meta::DIM, dim}, - {knowhere::meta::TOPK, k}, - {knowhere::meta::DEVICEID, DEVICEID} - }; - - index_type = milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID; - index_ = GetVecIndexFactory(index_type); - conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf); - auto search_cfg = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf); - - auto elems = nq * k; - std::vector res_ids(elems); - std::vector res_dis(elems); - - index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); - AssertResult(res_ids, res_dis); - EXPECT_EQ(index_->Count(), nb); - EXPECT_EQ(index_->Dimension(), dim); - - auto binaryset = index_->Serialize(); - { - // cpu -> gpu - auto cpu_idx = GetVecIndexFactory(index_type); - cpu_idx->Load(binaryset); - { - for (int i = 0; i < 2; ++i) { - auto gpu_idx = cpu_idx->CopyToGpu(DEVICEID, conf); - gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); - AssertResult(res_ids, res_dis); - } - } - } - - { - // quantization already in gpu, only copy data - auto cpu_idx = GetVecIndexFactory(index_type); - cpu_idx->Load(binaryset); - - auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf); - auto gpu_idx = pair.first; - auto quantization = pair.second; - - gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); - AssertResult(res_ids, res_dis); - - milvus::json quantizer_conf { - {"mode", 2}, - {knowhere::meta::DEVICEID, DEVICEID} - }; - for (int i = 0; i < 2; ++i) { - auto hybrid_idx = GetVecIndexFactory(index_type); - hybrid_idx->Load(binaryset); - - hybrid_idx->LoadData(quantization, quantizer_conf); - hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); - AssertResult(res_ids, res_dis); - } - } - - { - // quantization already in gpu, only set quantization - auto cpu_idx = GetVecIndexFactory(index_type); - cpu_idx->Load(binaryset); - - auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf); - auto quantization = pair.second; - - for (int i = 0; i < 2; ++i) { - auto hybrid_idx = GetVecIndexFactory(index_type); - hybrid_idx->Load(binaryset); - - hybrid_idx->SetQuantizer(quantization); - hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg); - AssertResult(res_ids, res_dis); - hybrid_idx->UnsetQuantizer(); - } - } -} - -// TEST_F(KnowhereHybrid, test_invalid_index) { -// assert(!xb.empty()); -// fiu_init(0); - -// index_type = milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID; -// fiu_enable("GetVecIndexFactory.IVFSQHybrid.mock", 1, nullptr, 0); -// index_ = GetVecIndexFactory(index_type); -// fiu_disable("GetVecIndexFactory.IVFSQHybrid.mock"); -// conf = ParamGenerator::GetInstance().Gen(index_type); - -// auto elems = nq * k; -// std::vector res_ids(elems); -// std::vector res_dis(elems); - -// conf->gpu_id = DEVICEID; -// conf->d = dim; -// conf->k = k; - -// auto status = index_->BuildAll(nb, xb.data(), ids.data(), conf); -// ASSERT_FALSE(status.ok()); - -// auto quantizer = std::make_shared(); -// auto quantizer_cfg = std::make_shared(); -// index_->LoadQuantizer(quantizer_cfg); -// status = index_->SetQuantizer(quantizer); -// ASSERT_FALSE(status.ok()); - -// fiu_enable("IVFHybridIndex.SetQuantizer.throw_knowhere_exception", 1, nullptr, 0); -// status = index_->SetQuantizer(quantizer); -// ASSERT_FALSE(status.ok()); -// fiu_disable("IVFHybridIndex.SetQuantizer.throw_knowhere_exception"); - -// fiu_enable("IVFHybridIndex.SetQuantizer.throw_std_exception", 1, nullptr, 0); -// status = index_->SetQuantizer(quantizer); -// ASSERT_FALSE(status.ok()); -// fiu_disable("IVFHybridIndex.SetQuantizer.throw_std_exception"); - -// status = index_->UnsetQuantizer(); -// ASSERT_FALSE(status.ok()); -// fiu_enable("IVFHybridIndex.UnsetQuantizer.throw_knowhere_exception", 1, nullptr, 0); -// status = index_->UnsetQuantizer(); -// ASSERT_FALSE(status.ok()); -// fiu_disable("IVFHybridIndex.UnsetQuantizer.throw_knowhere_exception"); - -// fiu_enable("IVFHybridIndex.UnsetQuantizer.throw_std_exception", 1, nullptr, 0); -// status = index_->UnsetQuantizer(); -// ASSERT_FALSE(status.ok()); -// fiu_disable("IVFHybridIndex.UnsetQuantizer.throw_std_exception"); - -// auto vecindex = index_->LoadData(quantizer, quantizer_cfg); -// ASSERT_EQ(vecindex, nullptr); - -// fiu_enable("IVFHybridIndex.LoadData.throw_std_exception", 1, nullptr, 0); -// vecindex = index_->LoadData(quantizer, quantizer_cfg); -// ASSERT_EQ(vecindex, nullptr); -// fiu_disable("IVFHybridIndex.LoadData.throw_std_exception"); - -// fiu_enable("IVFHybridIndex.LoadData.throw_knowhere_exception", 1, nullptr, 0); -// vecindex = index_->LoadData(quantizer, quantizer_cfg); -// ASSERT_EQ(vecindex, nullptr); -// fiu_disable("IVFHybridIndex.LoadData.throw_knowhere_exception"); - -// index_->CopyToGpuWithQuantizer(DEVICEID, conf); -// fiu_enable("IVFHybridIndex.LoadData.throw_knowhere_exception", 1, nullptr, 0); -// index_->CopyToGpuWithQuantizer(DEVICEID, conf); -// fiu_disable("IVFHybridIndex.LoadData.throw_knowhere_exception"); -// fiu_enable("IVFHybridIndex.LoadData.throw_std_exception", 1, nullptr, 0); -// index_->CopyToGpuWithQuantizer(DEVICEID, conf); -// fiu_disable("IVFHybridIndex.LoadData.throw_std_exception"); -// } - -#endif diff --git a/core/unittest/wrapper/utils.cpp b/core/unittest/wrapper/utils.cpp deleted file mode 100644 index 5cb89a50d1..0000000000 --- a/core/unittest/wrapper/utils.cpp +++ /dev/null @@ -1,217 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#include -#include -#include - -#include "utils/CommonUtil.h" -#include "wrapper/utils.h" - -namespace { -static const char* CONFIG_STR = - "# All the following configurations are default values.\n" - "\n" - "server_config:\n" - " address: 0.0.0.0 # milvus server ip address (IPv4)\n" - " port: 19530 # port range: 1025 ~ 65534\n" - " deploy_mode: single \n" - " time_zone: UTC+8\n" - "\n" - "db_config:\n" - " backend_url: sqlite://:@:/ # URI format: dialect://username:password@host:port/database\n" - " # Replace 'dialect' with 'mysql' or 'sqlite'\n" - "\n" - "storage_config:\n" - " primary_path: /tmp/milvus # path used to store data and meta\n" - " secondary_path: # path used to store data only, split by semicolon\n" - "\n" - "metric_config:\n" - " enable_monitor: false # enable monitoring or not\n" - " address: 127.0.0.1\n" - " port: 8080 # port prometheus used to fetch metrics\n" - "\n" - "cache_config:\n" - " cpu_cache_capacity: 4 # GB, CPU memory used for cache\n" - " cpu_cache_threshold: 0.85 # percentage of data kept when cache cleanup triggered\n" - " insert_buffer_size: 4 # GB, maximum insert buffer size allowed\n" - " cache_insert_data: false # whether load inserted data into cache\n" - "\n" - "engine_config:\n" - " blas_threshold: 20\n" - "\n" -#ifdef MILVUS_GPU_VERSION - "gpu_resource_config:\n" - " enable: true # whether to enable GPU resources\n" - " cache_capacity: 4 # GB, size of GPU memory per card used for cache, must be a positive integer\n" - " search_resources: # define the GPU devices used for search computation, must be in format gpux\n" - " - gpu0\n" - " build_index_resources: # define the GPU devices used for index building, must be in format gpux\n" - " - gpu0\n" -#endif - "\n"; - -void -WriteToFile(const std::string& file_path, const char* content) { - std::fstream fs(file_path.c_str(), std::ios_base::out); - - // write data to file - fs << content; - fs.close(); -} - -} // namespace - -void -KnowhereTest::SetUp() { - std::string config_path(CONFIG_PATH); - milvus::server::CommonUtil::CreateDirectory(config_path); - WriteToFile(config_path + CONFIG_FILE, CONFIG_STR); -} - -void -KnowhereTest::TearDown() { - std::string config_path(CONFIG_PATH); - milvus::server::CommonUtil::DeleteDirectory(config_path); -} - -void -DataGenBase::GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids, const int& k, - int64_t* gt_ids, float* gt_dis) { - for (auto i = 0; i < nb; ++i) { - for (auto j = 0; j < dim; ++j) { - // p_data[i * d + j] = float(base + i); - xb[i * dim + j] = drand48(); - } - xb[dim * i] += i / 1000.; - ids[i] = i; - } - for (size_t i = 0; i < nq * dim; ++i) { - xq[i] = xb[i]; - } - - faiss::IndexFlatL2 index(dim); - // index.add_with_ids(nb, xb, ids); - index.add(nb, xb); - index.search(nq, xq, k, gt_dis, gt_ids); -} - -void -DataGenBase::GenData(const int& dim, const int& nb, const int& nq, std::vector& xb, std::vector& xq, - std::vector& ids, const int& k, std::vector& gt_ids, - std::vector& gt_dis) { - xb.clear(); - xq.clear(); - ids.clear(); - gt_ids.clear(); - gt_dis.clear(); - xb.resize(nb * dim); - xq.resize(nq * dim); - ids.resize(nb); - gt_ids.resize(nq * k); - gt_dis.resize(nq * k); - GenData(dim, nb, nq, xb.data(), xq.data(), ids.data(), k, gt_ids.data(), gt_dis.data()); -} - -void -DataGenBase::AssertResult(const std::vector& ids, const std::vector& dis) { - EXPECT_EQ(ids.size(), nq * k); - EXPECT_EQ(dis.size(), nq * k); - - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(ids[i * k], gt_ids[i * k]); - // EXPECT_EQ(dis[i * k], gt_dis[i * k]); - } - - int match = 0; - for (int i = 0; i < nq; ++i) { - for (int j = 0; j < k; ++j) { - for (int l = 0; l < k; ++l) { - if (ids[i * nq + j] == gt_ids[i * nq + l]) - match++; - } - } - } - - auto precision = float(match) / (nq * k); - EXPECT_GT(precision, 0.5); - std::cout << std::endl << "Precision: " << precision << ", match: " << match << ", total: " << nq * k << std::endl; -} - -void -BinDataGen::GenData(const int& dim, - const int& nb, - const int& nq, - uint8_t* xb, - uint8_t* xq, - int64_t* ids, - const int& k, - int64_t* gt_ids, - float* gt_dis) { - for (auto i = 0; i < nb; ++i) { - for (auto j = 0; j < dim; ++j) { - // p_data[i * d + j] = float(base + i); - xb[i * dim + j] = (uint8_t)lrand48(); - } - ids[i] = i; - } - for (int64_t i = 0; i < nq * dim; ++i) { - xq[i] = xb[i]; - } -} - -void -BinDataGen::GenData(const int& dim, - const int& nb, - const int& nq, - std::vector& xb, - std::vector& xq, - std::vector& ids, - const int& k, - std::vector& gt_ids, - std::vector& gt_dis) { - xb.clear(); - xq.clear(); - ids.clear(); - gt_ids.clear(); - gt_dis.clear(); - xb.resize(nb * dim); - xq.resize(nq * dim); - ids.resize(nb); - gt_ids.resize(nq * k); - gt_dis.resize(nq * k); - GenData(dim, nb, nq, xb.data(), xq.data(), ids.data(), k, gt_ids.data(), gt_dis.data()); - assert(xb.size() == (size_t)dim * nb); - assert(xq.size() == (size_t)dim * nq); -} - -void -BinDataGen::AssertResult(const std::vector& ids, const std::vector& dis) { - EXPECT_EQ(ids.size(), nq * k); - EXPECT_EQ(dis.size(), nq * k); - - for (auto i = 0; i < nq; i++) { - EXPECT_EQ(ids[i * k], i); - // EXPECT_EQ(dis[i * k], gt_dis[i * k]); - } -} - -void -BinDataGen::Generate(const int& dim, const int& nb, const int& nq, const int& k) { - this->nb = nb; - this->nq = nq; - this->dim = dim; - this->k = k; - - int64_t dim_x = dim / 8; - - GenData(dim_x, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); -} diff --git a/core/unittest/wrapper/utils.h b/core/unittest/wrapper/utils.h deleted file mode 100644 index 527d2c306f..0000000000 --- a/core/unittest/wrapper/utils.h +++ /dev/null @@ -1,199 +0,0 @@ -// Copyright (C) 2019-2020 Zilliz. All rights reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software distributed under the License -// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -// or implied. See the License for the specific language governing permissions and limitations under the License. - -#pragma once - -#include - -#include -#include -#include -#include -#include - -#include "knowhere/index/vector_index/helpers/IndexParameter.h" -#include "wrapper/ConfAdapter.h" -#include "wrapper/ConfAdapterMgr.h" -#include "wrapper/VecIndex.h" -#include "wrapper/utils.h" - -class DataGenBase; - -using DataGenPtr = std::shared_ptr; - -constexpr int64_t DIM = 128; -constexpr int64_t NB = 100000; -constexpr int64_t NQ = 10; -constexpr int64_t DEVICEID = 0; -constexpr int64_t PINMEM = 1024 * 1024 * 200; -constexpr int64_t TEMPMEM = 1024 * 1024 * 300; -constexpr int64_t RESNUM = 2; - -static const char* CONFIG_PATH = "/tmp/milvus_test"; -static const char* CONFIG_FILE = "/server_config.yaml"; - -class KnowhereTest : public ::testing::Test { - protected: - void - SetUp() override; - void - TearDown() override; -}; - -class DataGenBase { - public: - virtual void - GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids, const int& k, - int64_t* gt_ids, float* gt_dis); - - virtual void - GenData(const int& dim, const int& nb, const int& nq, std::vector& xb, std::vector& xq, - std::vector& ids, const int& k, std::vector& gt_ids, std::vector& gt_dis); - - void - AssertResult(const std::vector& ids, const std::vector& dis); - - int dim = DIM; - int nb = NB; - int nq = NQ; - int k = 10; - std::vector xb; - std::vector xq; - std::vector ids; - - // Ground Truth - std::vector gt_ids; - std::vector gt_dis; -}; - -class BinDataGen { - public: - virtual void - GenData(const int& dim, const int& nb, const int& nq, uint8_t* xb, uint8_t* xq, int64_t* ids, const int& k, - int64_t* gt_ids, float* gt_dis); - - virtual void - GenData(const int& dim, const int& nb, const int& nq, std::vector& xb, std::vector& xq, - std::vector& ids, const int& k, std::vector& gt_ids, std::vector& gt_dis); - - void - AssertResult(const std::vector& ids, const std::vector& dis); - - void - Generate(const int& dim, const int& nb, const int& nq, const int& k); - - int dim = DIM; - int nb = NB; - int nq = NQ; - int k = 10; - std::vector xb; - std::vector xq; - std::vector ids; - - // Ground Truth - std::vector gt_ids; - std::vector gt_dis; -}; - -class ParamGenerator { - public: - static ParamGenerator& - GetInstance() { - static ParamGenerator instance; - return instance; - } - - knowhere::Config - GenSearchConf(const milvus::engine::IndexType& type, const milvus::json& conf) { - auto search_cfg = conf; - switch (type) { - case milvus::engine::IndexType::FAISS_IDMAP: { - break; - } - case milvus::engine::IndexType::FAISS_BIN_IVFLAT_CPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_CPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_GPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: - case milvus::engine::IndexType::FAISS_IVFPQ_CPU: - case milvus::engine::IndexType::FAISS_IVFPQ_GPU: - case milvus::engine::IndexType::FAISS_IVFPQ_MIX: - case milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID: - case milvus::engine::IndexType::FAISS_IVFSQ8_CPU: - case milvus::engine::IndexType::FAISS_IVFSQ8_GPU: - case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: { - search_cfg[knowhere::IndexParams::nprobe] = 16; - break; - } - case milvus::engine::IndexType::NSG_MIX: { - search_cfg[knowhere::IndexParams::search_length] = 20; - break; - } - case milvus::engine::IndexType::HNSW: { - search_cfg[knowhere::IndexParams::ef] = conf[knowhere::meta::TOPK].get() + 10; - break; - } - } - auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type); - adapter->CheckSearch(search_cfg, type); - return search_cfg; - } - - knowhere::Config - GenBuild(const milvus::engine::IndexType& type, const milvus::json& conf) { - auto build_cfg = conf; - switch (type) { - case milvus::engine::IndexType::FAISS_IDMAP: { - case milvus::engine::IndexType::FAISS_BIN_IDMAP: { - break; - } - case milvus::engine::IndexType::FAISS_BIN_IVFLAT_CPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_CPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_GPU: - case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: - case milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID: - case milvus::engine::IndexType::FAISS_IVFSQ8_CPU: - case milvus::engine::IndexType::FAISS_IVFSQ8_GPU: - case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: { - build_cfg[knowhere::IndexParams::nlist] = 16; - break; - } - case milvus::engine::IndexType::FAISS_IVFPQ_CPU: - case milvus::engine::IndexType::FAISS_IVFPQ_GPU: - case milvus::engine::IndexType::FAISS_IVFPQ_MIX: { - build_cfg[knowhere::IndexParams::nlist] = 16; - build_cfg[knowhere::IndexParams::m] = 8; - break; - } - case milvus::engine::IndexType::NSG_MIX: { - build_cfg[knowhere::IndexParams::knng] = 10; - build_cfg[knowhere::IndexParams::search_length] = 20; - build_cfg[knowhere::IndexParams::out_degree] = 10; - build_cfg[knowhere::IndexParams::candidate] = 50; - break; - } - case milvus::engine::IndexType::HNSW: { - build_cfg[knowhere::IndexParams::efConstruction] = 100; - build_cfg[knowhere::IndexParams::M] = 12; - break; - } - } - } - auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type); - adapter->CheckTrain(build_cfg); - return build_cfg; - } -}; - -// class SanityCheck : public DataGenBase { -// public: -// void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, long *ids, -// const int &k, long *gt_ids, float *gt_dis) override; -//};