solve conflicts

Signed-off-by: NWang Xiangyu <xy.wang@zilliz.com>

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

core/src/query/GeneralQuery.h

@@ -120,6 +120,7 @@ struct Query {
     std::vector<std::string> field_names;
     std::set<std::string> index_fields;
     std::unordered_map<std::string, std::string> metric_types;
+    std::string index_type;
 };
 using QueryPtr = std::shared_ptr<Query>;
 

core/src/scheduler/selector/FaissFlatPass.cpp

@@ -44,6 +44,9 @@ FaissFlatPass::Run(const TaskPtr& task) {
     }
 
     auto search_task = std::static_pointer_cast<SearchTask>(task);
+    if (search_task->IndexType() != knowhere::IndexEnum::INDEX_FAISS_IDMAP) {
+        return false;
+    }
 
     ResourcePtr res_ptr;
     if (!gpu_enable_) {

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<SearchTask>(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<int64_t>() >
+               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!",

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<SearchTask>(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<int64_t>() >
+               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_]);

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<SearchTask>(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<int64_t>() >
+               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_]);

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<SearchTask>(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<int64_t>() >
+               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_]);

core/src/scheduler/task/SearchTask.cpp

@@ -13,6 +13,7 @@
 
 #include <fiu/fiu-local.h>
 
+#include <src/index/thirdparty/faiss/IndexFlat.h>
 #include <algorithm>
 #include <memory>
 #include <string>
@@ -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

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_;
 

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();
 }