vector: A new load balance to improve the vector search performance when tiflash replica >1 (#9636) (#9651)

close #9637

In this pr: [https://github.com/tikv/client-c/pull/193](https://github.com/tikv/client-c/pull/193)
add a bool paramater centralized_schedule for Client::buildBatchCopTasks
When set to true:  Regions are concentrated into as few BatchCopTasks as possible

This PR modifies the call to this function, setting the newly added parameter to true when the request is a ANN query.

Signed-off-by: Calvin Neo <calvinneo1995@gmail.com>

Co-authored-by: Yu Jin <yujin.yujin@gmail.com>
Co-authored-by: JaySon-Huang <tshent@qq.com>
Co-authored-by: Calvin Neo <CalvinNeo@users.noreply.github.com>
Co-authored-by: JaySon <tshent@qq.com>

Author: 4 people

contrib/client-c

@@ -188,6 +188,8 @@ struct ContextShared
 
     JointThreadInfoJeallocMapPtr joint_memory_allocation_map; /// Joint thread-wise alloc/dealloc map
 
+    std::unordered_set<uint64_t> store_id_blocklist; /// Those store id are blocked from batch cop request.
+
     class SessionKeyHash
     {
     public:
@@ -2218,6 +2220,44 @@ void Context::setMockMPPServerInfo(MockMPPServerInfo & info)
     mpp_server_info = info;
 }
 
+const std::unordered_set<uint64_t> * Context::getStoreIdBlockList() const
+{
+    return &shared->store_id_blocklist;
+}
+
+// NOLINTNEXTLINE(readability-convert-member-functions-to-static)
+bool Context::initializeStoreIdBlockList(const String & comma_sep_string)
+{
+#if SERVERLESS_PROXY == 1
+    std::istringstream iss(comma_sep_string);
+    std::string token;
+
+    while (std::getline(iss, token, ','))
+    {
+        try
+        {
+            uint64_t number = std::stoull(token);
+            shared->store_id_blocklist.insert(number);
+        }
+        catch (...)
+        {
+            // Keep empty
+            LOG_INFO(DB::Logger::get(), "StoreIdBlockList is not set, input_str={}", comma_sep_string);
+            shared->store_id_blocklist.clear();
+            return false;
+        }
+    }
+
+    if (!shared->store_id_blocklist.empty())
+        LOG_INFO(DB::Logger::get(), "StoreIdBlockList have been set, {}", shared->store_id_blocklist);
+
+    return true;
+#else
+    UNUSED(comma_sep_string);
+    return true;
+#endif
+}
+
 SessionCleaner::~SessionCleaner()
 {
     try
@@ -2251,4 +2291,5 @@ void SessionCleaner::run()
             break;
     }
 }
+
 } // namespace DB

dbms/src/Interpreters/Context.cpp

@@ -190,6 +190,7 @@ class Context
     }};
     using DatabasePtr = std::shared_ptr<IDatabase>;
     using Databases = std::map<String, std::shared_ptr<IDatabase>>;
+
     /// Use copy constructor or createGlobal() instead
     Context();
 
@@ -560,6 +561,9 @@ class Context
 
     void mockConfigLoaded() { is_config_loaded = true; }
 
+    bool initializeStoreIdBlockList(const String &);
+    const std::unordered_set<uint64_t> * getStoreIdBlockList() const;
+
 private:
     /** Check if the current client has access to the specified database.
       * If access is denied, throw an exception.

dbms/src/Interpreters/Context.h

@@ -849,6 +849,7 @@ struct SettingString
     void write(WriteBuffer & buf) const { writeBinary(value, buf); }
 
     String get() const { return value; }
+    const String & getRef() const { return value; }
 
 private:
     String value;

dbms/src/Interpreters/SettingsCommon.h

@@ -1512,6 +1512,9 @@ int Server::main(const std::vector<std::string> & /*args*/)
     global_context->setFormatSchemaPath(format_schema_path.path() + "/");
     format_schema_path.createDirectories();
 
+    // We do not support blocking store by id in OP mode currently.
+    global_context->initializeStoreIdBlockList("");
+
     LOG_INFO(log, "Loading metadata.");
     loadMetadataSystem(*global_context); // Load "system" database. Its engine keeps as Ordinary.
     /// After attaching system databases we can initialize system log.

dbms/src/Server/Server.cpp

@@ -105,7 +105,7 @@ std::vector<pingcap::coprocessor::BatchCopTask> StorageDisaggregated::buildBatch
         physical_table_ids.emplace_back(remote_table_range.first);
         ranges_for_each_physical_table.emplace_back(remote_table_range.second);
     }
-
+    bool has_ann_query = table_scan.getANNQueryInfo().query_type() != tipb::ANNQueryType::InvalidQueryType;
     pingcap::kv::Cluster * cluster = context.getTMTContext().getKVCluster();
     pingcap::kv::Backoffer bo(pingcap::kv::copBuildTaskMaxBackoff);
     pingcap::kv::StoreType store_type = pingcap::kv::StoreType::TiFlash;
@@ -116,9 +116,11 @@ std::vector<pingcap::coprocessor::BatchCopTask> StorageDisaggregated::buildBatch
         table_scan.isPartitionTableScan(),
         physical_table_ids,
         ranges_for_each_physical_table,
+        context.getStoreIdBlockList(),
         store_type,
         label_filter,
-        &Poco::Logger::get("pingcap/coprocessor"));
+        &Poco::Logger::get("pingcap/coprocessor"),
+        has_ann_query);
     LOG_DEBUG(log, "batch cop tasks(nums: {}) build finish for tiflash_storage node", batch_cop_tasks.size());
     return batch_cop_tasks;
 }