diff --git a/CHANGELOG.md b/CHANGELOG.md index 23759dd31471658fee6327f765a12ae6e9fe5fc8..8476231267c0005593c174d08e14733082a58cba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,8 @@ Please mark all changes in change log and use the issue from GitHub - \#3453 Fix server crashed during stability test - \#3482 Server crashed during adding entities - \#3490 Fix ut DBTest.DeleteEntitiesTest assert fail +- \#3511 SearchTask::nq() should not return 0 +- \#3514 Search failed with regex tag name ## Feature - \#2319 Redo metadata to support MVCC diff --git a/core/src/query/GeneralQuery.h b/core/src/query/GeneralQuery.h index eae08a2e1498198681d8f4cc1f9799649f0257a4..6c041abcef339bbb7624190301a66c83c984e17e 100644 --- a/core/src/query/GeneralQuery.h +++ b/core/src/query/GeneralQuery.h @@ -120,6 +120,7 @@ struct Query { std::vector field_names; std::set index_fields; std::unordered_map metric_types; + std::string index_type; }; using QueryPtr = std::shared_ptr; diff --git a/core/src/scheduler/selector/FaissFlatPass.cpp b/core/src/scheduler/selector/FaissFlatPass.cpp index c02a51eb431b4a68fe8e28d39fbacea6b50a84c1..3e11638da93c763b034bb39967a9b70fd23d87c3 100644 --- a/core/src/scheduler/selector/FaissFlatPass.cpp +++ b/core/src/scheduler/selector/FaissFlatPass.cpp @@ -44,6 +44,9 @@ FaissFlatPass::Run(const TaskPtr& task) { } auto search_task = std::static_pointer_cast(task); + if (search_task->IndexType() != knowhere::IndexEnum::INDEX_FAISS_IDMAP) { + return false; + } ResourcePtr res_ptr; if (!gpu_enable_) { diff --git a/core/src/scheduler/selector/FaissIVFFlatPass.cpp b/core/src/scheduler/selector/FaissIVFFlatPass.cpp index 2a510a859d2417ad41041c009df72ae470a7797a..f6f6b99790446d243d9c6c295e9496e3e7e169cf 100644 --- a/core/src/scheduler/selector/FaissIVFFlatPass.cpp +++ b/core/src/scheduler/selector/FaissIVFFlatPass.cpp @@ -12,6 +12,8 @@ #include "scheduler/selector/FaissIVFFlatPass.h" #include "cache/GpuCacheMgr.h" #include "config/ServerConfig.h" +#include "faiss/gpu/utils/DeviceUtils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "scheduler/SchedInst.h" #include "scheduler/Utils.h" #include "scheduler/task/SearchTask.h" @@ -45,13 +47,20 @@ FaissIVFFlatPass::Run(const TaskPtr& task) { } auto search_task = std::static_pointer_cast(task); + if (search_task->IndexType() != knowhere::IndexEnum::INDEX_FAISS_IVFFLAT) { + return false; + } ResourcePtr res_ptr; if (!gpu_enable_) { LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: gpu disable, specify cpu to search!"); res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); } else if (search_task->nq() < threshold_) { - LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nq < gpu_search_threshold, specify cpu to search!"); + LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nq < gpu_search_threshold, specify cpu to search! "); + res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + } else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get() > + faiss::gpu::getMaxKSelection()) { + LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!"); res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); } else { LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nq >= gpu_search_threshold, specify gpu %d to search!", diff --git a/core/src/scheduler/selector/FaissIVFPQPass.cpp b/core/src/scheduler/selector/FaissIVFPQPass.cpp index 98e0654da48c79caf6c10c4ee12a36f51bf66f5a..3bfe1a208157028d4ae67ffcf308ee40e882590d 100644 --- a/core/src/scheduler/selector/FaissIVFPQPass.cpp +++ b/core/src/scheduler/selector/FaissIVFPQPass.cpp @@ -12,6 +12,8 @@ #include "scheduler/selector/FaissIVFPQPass.h" #include "cache/GpuCacheMgr.h" #include "config/ServerConfig.h" +#include "faiss/gpu/utils/DeviceUtils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "scheduler/SchedInst.h" #include "scheduler/Utils.h" #include "scheduler/task/SearchTask.h" @@ -47,6 +49,9 @@ FaissIVFPQPass::Run(const TaskPtr& task) { } auto search_task = std::static_pointer_cast(task); + if (search_task->IndexType() != knowhere::IndexEnum::INDEX_FAISS_IVFPQ) { + return false; + } ResourcePtr res_ptr; if (!gpu_enable_) { @@ -55,6 +60,10 @@ FaissIVFPQPass::Run(const TaskPtr& task) { } else if (search_task->nq() < threshold_) { LOG_SERVER_DEBUG_ << LogOut("FaissIVFPQPass: nq < gpu_search_threshold, specify cpu to search!"); res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + } else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get() > + faiss::gpu::getMaxKSelection()) { + LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!"); + res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); } else { LOG_SERVER_DEBUG_ << LogOut("FaissIVFPQPass: nq >= gpu_search_threshold, specify gpu %d to search!", search_gpus_[idx_]); diff --git a/core/src/scheduler/selector/FaissIVFSQ8HPass.cpp b/core/src/scheduler/selector/FaissIVFSQ8HPass.cpp index d8e5288c07cd0bfbcba02ddf921ac76b9356b054..b788d04e2068857571510dba7d9e5c591faba20f 100644 --- a/core/src/scheduler/selector/FaissIVFSQ8HPass.cpp +++ b/core/src/scheduler/selector/FaissIVFSQ8HPass.cpp @@ -13,6 +13,8 @@ #include "scheduler/selector/FaissIVFSQ8HPass.h" #include "cache/GpuCacheMgr.h" #include "config/ServerConfig.h" +#include "faiss/gpu/utils/DeviceUtils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "scheduler/SchedInst.h" #include "scheduler/Utils.h" #include "scheduler/task/SearchTask.h" @@ -44,6 +46,9 @@ FaissIVFSQ8HPass::Run(const TaskPtr& task) { } auto search_task = std::static_pointer_cast(task); + if (search_task->IndexType() != knowhere::IndexEnum::INDEX_FAISS_IVFSQ8H) { + return false; + } ResourcePtr res_ptr; if (!gpu_enable_) { @@ -53,6 +58,10 @@ FaissIVFSQ8HPass::Run(const TaskPtr& task) { if (search_task->nq() < threshold_) { LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8HPass: nq < gpu_search_threshold, specify cpu to search!"); res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + } else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get() > + faiss::gpu::getMaxKSelection()) { + LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!"); + res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); } else { LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8HPass: nq >= gpu_search_threshold, specify gpu %d to search!", search_gpus_[idx_]); diff --git a/core/src/scheduler/selector/FaissIVFSQ8Pass.cpp b/core/src/scheduler/selector/FaissIVFSQ8Pass.cpp index 8d3fd2ee7eae70dc265c90b09681a571eb5e23b7..0496ccd4892bfb2e29ffac10e535e89a4020b7af 100644 --- a/core/src/scheduler/selector/FaissIVFSQ8Pass.cpp +++ b/core/src/scheduler/selector/FaissIVFSQ8Pass.cpp @@ -12,6 +12,8 @@ #include "scheduler/selector/FaissIVFSQ8Pass.h" #include "cache/GpuCacheMgr.h" #include "config/ServerConfig.h" +#include "faiss/gpu/utils/DeviceUtils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "scheduler/SchedInst.h" #include "scheduler/Utils.h" #include "scheduler/task/SearchTask.h" @@ -45,6 +47,9 @@ FaissIVFSQ8Pass::Run(const TaskPtr& task) { } auto search_task = std::static_pointer_cast(task); + if (search_task->IndexType() != knowhere::IndexEnum::INDEX_FAISS_IVFSQ8) { + return false; + } ResourcePtr res_ptr; if (!gpu_enable_) { @@ -53,6 +58,10 @@ FaissIVFSQ8Pass::Run(const TaskPtr& task) { } else if (search_task->nq() < threshold_) { LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8Pass: nq < gpu_search_threshold, specify cpu to search!"); res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + } else if (search_task->ExtraParam()[knowhere::IndexParams::nprobe].get() > + faiss::gpu::getMaxKSelection()) { + LOG_SERVER_DEBUG_ << LogOut("FaissIVFFlatPass: nprobe > gpu_max_nprobe_threshold, specify cpu to search!"); + res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); } else { LOG_SERVER_DEBUG_ << LogOut("FaissIVFSQ8Pass: nq >= gpu_search_threshold, specify gpu %d to search!", search_gpus_[idx_]); diff --git a/core/src/scheduler/task/SearchTask.cpp b/core/src/scheduler/task/SearchTask.cpp index ef636801913864736dd16157a904309d01e641c9..c6ff788be53707bb8ffa514b542881b140d69318 100644 --- a/core/src/scheduler/task/SearchTask.cpp +++ b/core/src/scheduler/task/SearchTask.cpp @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -220,8 +221,69 @@ SearchTask::MergeTopkToResultSet(const engine::ResultIds& src_ids, const engine: int64_t SearchTask::nq() { + if (query_ptr_) { + auto vector_query = query_ptr_->vectors.begin(); + if (vector_query != query_ptr_->vectors.end()) { + if (vector_query->second) { + auto vector_param = vector_query->second; + auto field_visitor = snapshot_->GetField(vector_query->second->field_name); + if (field_visitor) { + if (field_visitor->GetParams().contains(engine::PARAM_DIMENSION)) { + int64_t dim = field_visitor->GetParams()[engine::PARAM_DIMENSION]; + if (!vector_param->query_vector.float_data.empty()) { + return vector_param->query_vector.float_data.size() / dim; + } else if (!vector_param->query_vector.binary_data.empty()) { + return vector_param->query_vector.binary_data.size() * 8 / dim; + } + } + } + } + } + } return 0; } +milvus::json +SearchTask::ExtraParam() { + milvus::json param; + if (query_ptr_) { + auto vector_query = query_ptr_->vectors.begin(); + if (vector_query != query_ptr_->vectors.end()) { + if (vector_query->second) { + return vector_query->second->extra_params; + } + } + } + return param; +} + +std::string +SearchTask::IndexType() { + if (!index_type_.empty()) { + return index_type_; + } + auto seg_visitor = engine::SegmentVisitor::Build(snapshot_, segment_id_); + index_type_ = "FLAT"; + + if (seg_visitor) { + for (const auto& name : query_ptr_->index_fields) { + auto field_visitor = seg_visitor->GetFieldVisitor(name); + auto type = field_visitor->GetField()->GetFtype(); + if (!field_visitor) { + continue; + } + if (type == engine::DataType::VECTOR_FLOAT || type == engine::DataType::VECTOR_BINARY) { + auto fe_visitor = field_visitor->GetElementVisitor(engine::FieldElementType::FET_INDEX); + if (fe_visitor) { + auto element = fe_visitor->GetElement(); + index_type_ = element->GetTypeName(); + } + return index_type_; + } + } + } + return index_type_; +} + } // namespace scheduler } // namespace milvus diff --git a/core/src/scheduler/task/SearchTask.h b/core/src/scheduler/task/SearchTask.h index b73b697bcc0d517a7142350a78051e9368179ceb..bb846e231acd39d4a7f8b53e32677663479cfa94 100644 --- a/core/src/scheduler/task/SearchTask.h +++ b/core/src/scheduler/task/SearchTask.h @@ -52,6 +52,12 @@ class SearchTask : public Task { int64_t nq(); + milvus::json + ExtraParam(); + + std::string + IndexType(); + private: void CreateExecEngine(); @@ -63,6 +69,7 @@ class SearchTask : public Task { const engine::DBOptions& options_; query::QueryPtr query_ptr_; engine::snapshot::ID_TYPE segment_id_; + std::string index_type_; engine::ExecutionEnginePtr execution_engine_; diff --git a/core/src/server/web_impl/handler/WebRequestHandler.cpp b/core/src/server/web_impl/handler/WebRequestHandler.cpp index 3859161d1252269b103948916fbd61762746e680..6b9513a71aa518d8d4da2e49e5eb6a026d8bea9d 100644 --- a/core/src/server/web_impl/handler/WebRequestHandler.cpp +++ b/core/src/server/web_impl/handler/WebRequestHandler.cpp @@ -286,6 +286,12 @@ WebRequestHandler::GetCollectionMetaInfo(const std::string& collection_name, nlo field_json["extra_params"] = field.second.field_params_; json_out["fields"].push_back(field_json); } + if (schema.extra_params_.contains(engine::PARAM_SEGMENT_ROW_COUNT)) { + json_out[engine::PARAM_SEGMENT_ROW_COUNT] = schema.extra_params_[engine::PARAM_SEGMENT_ROW_COUNT]; + } + if (schema.extra_params_.contains(engine::PARAM_UID_AUTOGEN)) { + json_out[engine::PARAM_UID_AUTOGEN] = schema.extra_params_[engine::PARAM_UID_AUTOGEN]; + } return Status::OK(); }