Add the load_state api

examples/collection.py

@@ -10,6 +10,7 @@
 # or implied. See the License for the specific language governing permissions and limitations under the License.
 
 from pymilvus import Collection, connections, FieldSchema, CollectionSchema, DataType, utility
+from pymilvus.client.types import LoadState
 
 import random
 import numpy as np
@@ -177,6 +178,7 @@ def test_create_collection():
     collection = Collection(name=name, schema=gen_default_fields())
     assert collection.is_empty is True
     assert collection.num_entities == 0
+    assert utility.load_state(name) == LoadState.NotLoad
     return name
 
 
@@ -197,6 +199,7 @@ def test_collection_only_name():
     collection.load()
     assert collection.is_empty is False
     assert collection.num_entities == default_nb
+    assert utility.load_state(name) == LoadState.Loaded
     collection.drop()
 
 

pymilvus/client/grpc_handler.py

@@ -23,6 +23,7 @@
 from .types import (
     Status,
     IndexState,
+    LoadState,
     DataType,
     CompactionState,
     State,
@@ -739,6 +740,14 @@ def get_loading_progress(self, collection_name, partition_names=None, timeout=No
             raise MilvusException(response.status.error_code, response.status.reason)
         return response.progress
 
+    @retry_on_rpc_failure()
+    def get_load_state(self, collection_name, partition_names=None, timeout=None):
+        request = Prepare.get_load_state(collection_name, partition_names)
+        response = self._stub.GetLoadState.future(request, timeout=timeout).result()
+        if response.status.error_code != 0:
+            raise MilvusException(response.status.error_code, response.status.reason)
+        return LoadState(response.state)
+
     @retry_on_rpc_failure()
     def load_partitions_progress(self, collection_name, partition_names, timeout=None):
         """ Return loading progress of partitions """

pymilvus/client/prepare.py

@@ -216,6 +216,14 @@ def get_loading_progress(cls, collection_name, partition_names=None):
             req.partition_names.extend(partition_names)
         return req
 
+    @classmethod
+    def get_load_state(cls, collection_name, partition_names=None):
+        check_pass_param(collection_name=collection_name, partition_name_array=partition_names)
+        req = milvus_types.GetLoadStateRequest(collection_name=collection_name)
+        if partition_names:
+            req.partition_names.extend(partition_names)
+        return req
+
     @classmethod
     def empty(cls):
         raise DeprecationWarning("no empty request later")

pymilvus/client/types.py

@@ -179,6 +179,25 @@ def __str__(self):
         return self._name_
 
 
+class LoadState(IntEnum):
+    """
+    NotExist:     collection or partition isn't existed
+    NotLoad:      collection or partition isn't loaded
+    Loading:      collection or partition is loading
+    Loaded:       collection or partition is loaded
+    """
+
+    NotExist = 0
+    NotLoad = 1
+    Loading = 2
+    Loaded = 3
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__}: {self._name_}>"
+
+    def __str__(self):
+        return self._name_
+
 class CompactionState:
     """
     in_executing:   number of plans in executing

pymilvus/grpc_gen/milvus_pb2_grpc.py

@@ -100,6 +100,11 @@ def __init__(self, channel):
                 request_serializer=milvus__pb2.GetLoadingProgressRequest.SerializeToString,
                 response_deserializer=milvus__pb2.GetLoadingProgressResponse.FromString,
                 )
+        self.GetLoadState = channel.unary_unary(
+                '/milvus.proto.milvus.MilvusService/GetLoadState',
+                request_serializer=milvus__pb2.GetLoadStateRequest.SerializeToString,
+                response_deserializer=milvus__pb2.GetLoadStateResponse.FromString,
+                )
         self.CreateAlias = channel.unary_unary(
                 '/milvus.proto.milvus.MilvusService/CreateAlias',
                 request_serializer=milvus__pb2.CreateAliasRequest.SerializeToString,
@@ -417,6 +422,12 @@ def GetLoadingProgress(self, request, context):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def GetLoadState(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def CreateAlias(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -764,6 +775,11 @@ def add_MilvusServiceServicer_to_server(servicer, server):
                     request_deserializer=milvus__pb2.GetLoadingProgressRequest.FromString,
                     response_serializer=milvus__pb2.GetLoadingProgressResponse.SerializeToString,
             ),
+            'GetLoadState': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetLoadState,
+                    request_deserializer=milvus__pb2.GetLoadStateRequest.FromString,
+                    response_serializer=milvus__pb2.GetLoadStateResponse.SerializeToString,
+            ),
             'CreateAlias': grpc.unary_unary_rpc_method_handler(
                     servicer.CreateAlias,
                     request_deserializer=milvus__pb2.CreateAliasRequest.FromString,
@@ -1273,6 +1289,23 @@ def GetLoadingProgress(request,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
+    @staticmethod
+    def GetLoadState(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/milvus.proto.milvus.MilvusService/GetLoadState',
+            milvus__pb2.GetLoadStateRequest.SerializeToString,
+            milvus__pb2.GetLoadStateResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
     @staticmethod
     def CreateAlias(request,
             target,

pymilvus/orm/utility.py

@@ -204,6 +204,43 @@ def loading_progress(collection_name, partition_names=None, using="default", tim
     }
 
 
+def load_state(collection_name, partition_names=None, using="default", timeout=None):
+    """ Show load state of collection or partitions.
+    :param collection_name: The name of collection is loading
+    :type  collection_name: str
+
+    :param partition_names: The names of partitions is loading
+    :type  partition_names: str list
+
+    :return LoadState:
+        The current state of collection or partitions.
+
+    :example:
+        >>> from pymilvus import Collection, connections, FieldSchema, CollectionSchema, DataType, utility
+        >>> from pymilvus.client.types import LoadState
+        >>> import pandas as pd
+        >>> import random
+        >>> connections.connect()
+        >>> assert utility.load_state("test_load_state") == LoadState.NotExist
+        >>> fields = [
+        ...     FieldSchema("film_id", DataType.INT64, is_primary=True),
+        ...     FieldSchema("films", DataType.FLOAT_VECTOR, dim=8),
+        ... ]
+        >>> schema = CollectionSchema(fields)
+        >>> collection = Collection("test_load_state", schema)
+        >>> assert utility.load_state("test_load_state") == LoadState.NotLoad
+        >>> data = pd.DataFrame({
+        ...     "film_id" : pd.Series(data=list(range(10, 20)), index=list(range(10))),
+        ...     "films": [[random.random() for _ in range(8)] for _ in range (10)],
+        ... })
+        >>> collection.insert(data)
+        >>> collection.create_index("films", {"index_type": "IVF_FLAT", "params": {"nlist": 8}, "metric_type": "L2"})
+        >>> collection.load(_async=True)
+        >>> assert utility.load_state("test_load_state") == LoadState.Loaded
+    """
+    return _get_connection(using).get_load_state(collection_name, partition_names, timeout=timeout)
+
+
 def wait_for_loading_complete(collection_name, partition_names=None, timeout=None, using="default"):
     """
     Block until loading is done or Raise Exception after timeout.