Add get_load_state

Signed-off-by: junjie.jiangjjj <junjie.jiang@zilliz.com>

examples/hello_milvus.py

@@ -15,6 +15,7 @@
     FieldSchema, CollectionSchema, DataType,
     Collection,
 )
+from pymilvus.client.types import LoadState
 
 from milvus_lite.server_manager import server_manager_instance
 uri = server_manager_instance.start_and_get_uri("./110.db")
@@ -114,6 +115,7 @@
 # Before conducting a search or a query, you need to load the data in `hello_milvus` into memory.
 print(fmt.format("Start loading"))
 hello_milvus.load()
+assert utility.load_state("hello_milvus") == LoadState.Loaded
 
 # -----------------------------------------------------------------------------
 # search based on vector similarity

src/milvus_local.cpp

@@ -34,7 +34,6 @@ namespace milvus::local {
         CHECK_STATUS(CheckCollectionName(string_util::Trim(collection_name)), \
                      "");                                                     \
         if (!storage_.CollectionExist(collection_name)) {                     \
-            LOG_ERROR("Collecton {} not existed", collection_name);           \
             return Status::CollectionNotFound();                              \
         }                                                                     \
     } while (0)
@@ -44,7 +43,6 @@ namespace milvus::local {
         CHECK_STATUS(CheckCollectionName(string_util::Trim(collection_name)), \
                      "");                                                     \
         if (storage_.CollectionExist(collection_name)) {                      \
-            LOG_ERROR("Collecton {} already existed", collection_name);       \
             return Status::CollectionAlreadExist();                           \
         }                                                                     \
     } while (0)
@@ -153,6 +151,16 @@ MilvusLocal::ReleaseCollection(const std::string& collection_name) {
     return Status::SegcoreErr();
 }
 
+Status
+MilvusLocal::GetLoadState(const std::string &collection_name) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    CHECK_COLLECTION_EXIST(collection_name);
+    if (!index_.HasLoaded(collection_name)) {
+        return Status::CollectionNotLoaded();
+    }
+    return Status::CollectionLoaded();
+}
+
 Status
 MilvusLocal::CreateCollection(const std::string& collection_name,
                               const std::string& pk_name,

src/milvus_local.h

@@ -49,6 +49,9 @@ class MilvusLocal final : NonCopyableNonMovable {
                      const std::string& pk_name,
                      const std::string& schema_proto);
 
+    Status
+    GetLoadState(const std::string& collection_name);
+
     Status
     GetCollection(const std::string& collection_name,
                   std::string* schema_proto);

src/milvus_proxy.cpp

@@ -55,6 +55,28 @@ MilvusProxy::LoadCollection(const std::string& collection_name) {
     return milvus_local_.LoadCollection(collection_name);
 }
 
+Status
+MilvusProxy::GetLoadState(
+    const std::string& collection_name,
+    ::milvus::proto::milvus::GetLoadStateResponse* response) {
+    auto s = milvus_local_.GetLoadState(collection_name);
+    if (s.Code() == ErrCollectionNotFound) {
+        response->set_state(::milvus::proto::common::LoadState::LoadStateNotExist);
+        return Status::Ok();
+    }
+
+    if (s.Code() == ErrCollectionNotLoaded) {
+        response->set_state(::milvus::proto::common::LoadState::LoadStateNotLoad);
+        return Status::Ok();        
+    }
+
+    if (s.Code() == ErrCollectionLoaded) {
+        response->set_state(::milvus::proto::common::LoadState::LoadStateLoaded);
+        return Status::Ok();
+    }
+    return s;
+}
+
 Status
 MilvusProxy::ReleaseCollection(const std::string& collection_name) {
     // Alignment error code with milvus

src/milvus_proxy.h

@@ -45,6 +45,9 @@ class MilvusProxy : NonCopyableNonMovable {
     Status
     HasCollection(const std::string& collection_name);
 
+    Status
+    GetLoadState(const std::string& collection_name, ::milvus::proto::milvus::GetLoadStateResponse* response);
+
     bool
     DropCollection(const std::string& collection_name);
 

src/milvus_service_impl.cpp

@@ -263,4 +263,13 @@ MilvusServiceImpl::GetCollectionStatistics(
     return ::grpc::Status::OK;
 }
 
+::grpc::Status
+MilvusServiceImpl::GetLoadState(::grpc::ServerContext *context,
+             const ::milvus::proto::milvus::GetLoadStateRequest *request,
+             ::milvus::proto::milvus::GetLoadStateResponse *response) {
+    auto s = proxy_.GetLoadState(request->collection_name(), response);
+    Status2Response(s, response->mutable_status());
+    return ::grpc::Status::OK;
+}
+
 }  // namespace milvus::local

src/milvus_service_impl.h

@@ -146,6 +146,8 @@ class MilvusServiceImpl final
         ::milvus::proto::milvus::GetCollectionStatisticsResponse* response)
         override;
 
+    ::grpc::Status GetLoadState(::grpc::ServerContext* context, const ::milvus::proto::milvus::GetLoadStateRequest* request, ::milvus::proto::milvus::GetLoadStateResponse* response) override;
+
  private:
     ::milvus::local::MilvusProxy proxy_;
 };