#4075 Improve performance for create large amount of partitions (#4083)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

CHANGELOG.md

@@ -8,6 +8,7 @@ Please mark all change in change log and use the issue from GitHub
 -   \#3903 The performance of IVF_SQ8H in 0.10.3 is degraded.
 -   \#3906 Change DeleteTask state when it is loaded to avoid server crash.
 -   \#4012 Milvus hangs when continually creating and dropping partitions.
+-   \#4075 Improve performance for create large amount of partitions
 
 ## Feature
 -   \#3773 Support IVF_PQ to run on FPGA.

core/src/db/DB.h

@@ -96,6 +96,9 @@ class DB {
     virtual Status
     ShowPartitions(const std::string& collection_id, std::vector<meta::CollectionSchema>& partition_schema_array) = 0;
 
+    virtual Status
+    CountPartitions(const std::string& collection_id, int64_t& partition_count) = 0;
+
     virtual Status
     InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) = 0;
 

core/src/db/DBImpl.cpp

@@ -687,6 +687,15 @@ DBImpl::ShowPartitions(const std::string& collection_id, std::vector<meta::Colle
     return meta_ptr_->ShowPartitions(collection_id, partition_schema_array);
 }
 
+Status
+DBImpl::CountPartitions(const std::string& collection_id, int64_t& partition_count) {
+    if (!initialized_.load(std::memory_order_acquire)) {
+        return SHUTDOWN_ERROR;
+    }
+
+    return meta_ptr_->CountPartitions(collection_id, partition_count);
+}
+
 Status
 DBImpl::InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) {
     //    LOG_ENGINE_DEBUG_ << "Insert " << n << " vectors to cache";

core/src/db/DBImpl.h

@@ -107,6 +107,9 @@ class DBImpl : public DB, public server::CacheConfigHandler, public server::Engi
     ShowPartitions(const std::string& collection_id,
                    std::vector<meta::CollectionSchema>& partition_schema_array) override;
 
+    Status
+    CountPartitions(const std::string& collection_id, int64_t& partition_count) override;
+
     Status
     InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) override;
 

core/src/db/meta/Meta.h

@@ -107,7 +107,7 @@ class Meta {
     DropCollectionIndex(const std::string& collection_id) = 0;
 
     virtual Status
-    CreatePartition(const std::string& collection_name, const std::string& partition_name, const std::string& tag,
+    CreatePartition(const std::string& collection_id, const std::string& partition_name, const std::string& tag,
                     uint64_t lsn) = 0;
 
     virtual Status
@@ -117,10 +117,13 @@ class Meta {
     DropPartition(const std::string& partition_name) = 0;
 
     virtual Status
-    ShowPartitions(const std::string& collection_name, std::vector<meta::CollectionSchema>& partition_schema_array) = 0;
+    ShowPartitions(const std::string& collection_id, std::vector<meta::CollectionSchema>& partition_schema_array) = 0;
 
     virtual Status
-    GetPartitionName(const std::string& collection_name, const std::string& tag, std::string& partition_name) = 0;
+    CountPartitions(const std::string& collection_id, int64_t& partition_count) = 0;
+
+    virtual Status
+    GetPartitionName(const std::string& collection_id, const std::string& tag, std::string& partition_name) = 0;
 
     virtual Status
     FilesToSearch(const std::string& collection_id, FilesHolder& files_holder) = 0;

core/src/db/meta/MySQLMetaImpl.cpp

@@ -1617,6 +1617,38 @@ MySQLMetaImpl::ShowPartitions(const std::string& collection_id,
     return Status::OK();
 }
 
+Status
+MySQLMetaImpl::CountPartitions(const std::string& collection_id, int64_t& partition_count) {
+    try {
+        server::MetricCollector metric;
+        mysqlpp::StoreQueryResult res;
+        {
+            mysqlpp::ScopedConnection connectionPtr(*mysql_connection_pool_, safe_grab_);
+
+            bool is_null_connection = (connectionPtr == nullptr);
+            if (is_null_connection) {
+                return Status(DB_ERROR, "Failed to connect to meta server(mysql)");
+            }
+
+            mysqlpp::Query statement = connectionPtr->query();
+            statement << "SELECT count(*) FROM " << META_TABLES << " WHERE owner_table = " << mysqlpp::quote
+                      << collection_id << " AND state <> " << std::to_string(CollectionSchema::TO_DELETE) << ";";
+
+            LOG_ENGINE_DEBUG_ << "CountPartitions: " << statement.str();
+
+            res = statement.store();
+        }  // Scoped Connection
+
+        for (auto& resRow : res) {
+            partition_count = resRow["count"];
+        }
+    } catch (std::exception& e) {
+        return HandleException("Failed to count partitions", e.what());
+    }
+
+    return Status::OK();
+}
+
 Status
 MySQLMetaImpl::GetPartitionName(const std::string& collection_id, const std::string& tag, std::string& partition_name) {
     try {

core/src/db/meta/MySQLMetaImpl.h

@@ -104,6 +104,9 @@ class MySQLMetaImpl : public Meta {
     ShowPartitions(const std::string& collection_id,
                    std::vector<meta::CollectionSchema>& partition_schema_array) override;
 
+    Status
+    CountPartitions(const std::string& collection_id, int64_t& partition_count) override;
+
     Status
     GetPartitionName(const std::string& collection_id, const std::string& tag, std::string& partition_name) override;
 

core/src/db/meta/SqliteMetaImpl.cpp

@@ -1023,6 +1023,25 @@ SqliteMetaImpl::ShowPartitions(const std::string& collection_id,
     return Status::OK();
 }
 
+Status
+SqliteMetaImpl::CountPartitions(const std::string& collection_id, int64_t& partition_count) {
+    try {
+        partition_count = 0;
+        auto ret = ConnectorPtr->select(count(),
+                                        where(c(&CollectionSchema::owner_collection_) == collection_id and
+                                              c(&CollectionSchema::state_) != (int)CollectionSchema::TO_DELETE));
+        if (ret.size() == 1) {
+            partition_count = ret[0];
+        } else {
+            return Status(DB_NOT_FOUND, "Collection " + collection_id + " not found");
+        }
+    } catch (std::exception& e) {
+        return HandleException("Encounter exception when count partitions", e.what());
+    }
+
+    return Status::OK();
+}
+
 Status
 SqliteMetaImpl::GetPartitionName(const std::string& collection_id, const std::string& tag,
                                  std::string& partition_name) {

core/src/db/meta/SqliteMetaImpl.h

@@ -106,6 +106,9 @@ class SqliteMetaImpl : public Meta {
     ShowPartitions(const std::string& collection_id,
                    std::vector<meta::CollectionSchema>& partition_schema_array) override;
 
+    Status
+    CountPartitions(const std::string& collection_id, int64_t& partition_count) override;
+
     Status
     GetPartitionName(const std::string& collection_id, const std::string& tag, std::string& partition_name) override;
 

core/src/server/delivery/request/CreatePartitionRequest.cpp

@@ -80,9 +80,9 @@ CreatePartitionRequest::OnExecute() {
         }
 
         // check partition total count
-        std::vector<engine::meta::CollectionSchema> schema_array;
-        status = DBWrapper::DB()->ShowPartitions(collection_name_, schema_array);
-        if (schema_array.size() >= MAX_PARTITION_LIMIT) {
+        int64_t partition_count = 0;
+        status = DBWrapper::DB()->CountPartitions(collection_name_, partition_count);
+        if (partition_count >= MAX_PARTITION_LIMIT) {
             return Status(SERVER_UNSUPPORTED_ERROR, "The number of partitions exceeds the upper limit(4096)");
         }