support the report value in the dml and dql request

Signed-off-by: SimFG <[email protected]>

pymilvus/client/abstract.py

@@ -4,7 +4,7 @@
 import ujson
 
 from pymilvus.exceptions import DataTypeNotMatchException, ExceptionsMessage, MilvusException
-from pymilvus.grpc_gen import schema_pb2
+from pymilvus.grpc_gen import schema_pb2, common_pb2
 from pymilvus.settings import Config
 
 from . import entity_helper
@@ -195,6 +195,7 @@ def __init__(self, raw: Any):
         self._timestamp = 0
         self._succ_index = []
         self._err_index = []
+        self._cost = 0
 
         self._pack(raw)
 
@@ -234,10 +235,15 @@ def succ_index(self):
     def err_index(self):
         return self._err_index
 
+    @property
+    def cost(self):
+        return self._cost
+
     def __str__(self):
         return (
             f"(insert count: {self._insert_cnt}, delete count: {self._delete_cnt}, upsert count: {self._upsert_cnt}, "
-            f"timestamp: {self._timestamp}, success count: {self.succ_count}, err count: {self.err_count})"
+            f"timestamp: {self._timestamp}, success count: {self.succ_count}, err count: {self.err_count}, "
+            f"cost: {self._cost})"
         )
 
     __repr__ = __str__
@@ -262,6 +268,8 @@ def _pack(self, raw: Any):
         self._timestamp = raw.timestamp
         self._succ_index = raw.succ_index
         self._err_index = raw.err_index
+        self._cost = raw.status.extra_info["report_value"] \
+            if raw.status and raw.status.extra_info else 0
 
 
 class SequenceIterator:
@@ -374,10 +382,14 @@ def __str__(self):
 class SearchResult(list):
     """nq results: List[Hits]"""
 
-    def __init__(self, res: schema_pb2.SearchResultData, round_decimal: Optional[int] = None):
+    def __init__(self, res: schema_pb2.SearchResultData,
+                 round_decimal: Optional[int] = None,
+                 status: Optional[common_pb2.Status] = None):
         self._nq = res.num_queries
         all_topks = res.topks
 
+        self.cost = status.extra_info["report_value"] if status and status.extra_info else 0
+
         output_fields = res.output_fields
         fields_data = res.fields_data
 
@@ -497,7 +509,7 @@ def __iter__(self) -> SequenceIterator:
 
     def __str__(self) -> str:
         """Only print at most 10 query results"""
-        return str(list(map(str, self[:10])))
+        return f"cost: {self.cost}, part data: {str(list(map(str, self[:10])))}"
 
     __repr__ = __str__
 

pymilvus/client/asynch.py

@@ -164,7 +164,7 @@ def exception(self):
 class SearchFuture(Future):
     def on_response(self, response: milvus_pb2.SearchResults):
         check_status(response.status)
-        return SearchResult(response.results)
+        return SearchResult(response.results, status=response.status)
 
 
 class MutationFuture(Future):

pymilvus/client/entity_helper.py

@@ -313,7 +313,6 @@ def pack_field_value_to_field_data(
                     message="invalid input for float32 vector, expect np.ndarray with dtype=float32"
                 )
             f_value = field_value.view(np.float32).tolist()
-
         field_data.vectors.dim = len(f_value)
         field_data.vectors.float_vector.data.extend(f_value)
 

pymilvus/client/grpc_handler.py

@@ -43,6 +43,7 @@
     CompactionPlans,
     CompactionState,
     DataType,
+    ExtraList,
     GrantInfo,
     Group,
     IndexState,
@@ -56,6 +57,7 @@
     State,
     Status,
     UserInfo,
+    get_cost_extra,
 )
 from .utils import (
     check_invalid_binary_vector,
@@ -731,7 +733,7 @@ def _execute_search(
             response = self._stub.Search(request, timeout=timeout)
             check_status(response.status)
             round_decimal = kwargs.get("round_decimal", -1)
-            return SearchResult(response.results, round_decimal)
+            return SearchResult(response.results, round_decimal, status=response.status)
 
         except Exception as e:
             if kwargs.get("_async", False):
@@ -750,7 +752,7 @@ def _execute_hybrid_search(
             response = self._stub.HybridSearch(request, timeout=timeout)
             check_status(response.status)
             round_decimal = kwargs.get("round_decimal", -1)
-            return SearchResult(response.results, round_decimal)
+            return SearchResult(response.results, round_decimal, status=response.status)
 
         except Exception as e:
             if kwargs.get("_async", False):
@@ -1503,7 +1505,7 @@ def query(
                 response.fields_data, index, dynamic_fields
             )
             results.append(entity_row_data)
-        return results
+        return ExtraList(results, extra=get_cost_extra(response.status))
 
     @retry_on_rpc_failure()
     def load_balance(

pymilvus/client/types.py

@@ -1,6 +1,6 @@
 import time
 from enum import IntEnum
-from typing import Any, ClassVar, Dict, List, TypeVar, Union
+from typing import Any, ClassVar, Dict, List, Optional, TypeVar, Union
 
 from pymilvus.exceptions import (
     AutoIDException,
@@ -763,7 +763,6 @@ def groups(self):
 
 class ResourceGroupInfo:
     def __init__(self, resource_group: Any) -> None:
-
         self._name = resource_group.name
         self._capacity = resource_group.capacity
         self._num_available_node = resource_group.num_available_node
@@ -888,3 +887,33 @@ def hostname(self) -> str:
 Attributes:
     resource_group (str): The name of the resource group that can be transferred to or from.
 """
+
+
+class ExtraList(list):
+    """
+    A list that can hold extra information.
+    Attributes:
+        extra (dict): The extra information of the list.
+    Example:
+        ExtraList([1, 2, 3], extra={"total": 3})
+    """
+
+    def __init__(self, *args, extra: Optional[Dict] = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.extra = extra or {}
+
+    def __str__(self) -> str:
+        """Only print at most 10 query results"""
+        return f"part data: {str(list(map(str, self[:10])))}, extra_info: {self.extra}"
+
+    __repr__ = __str__
+
+
+def get_cost_from_status(status: Optional[common_pb2.Status] = None):
+    return status.extra_info["report_value"] if status and status.extra_info else 0
+
+def get_cost_extra(status: Optional[common_pb2.Status] = None):
+    return {"cost": get_cost_from_status(status)}
+
+def construct_cost_extra(cost: int):
+    return {"cost": cost}

pymilvus/milvus_client/milvus_client.py

@@ -5,7 +5,13 @@
 from uuid import uuid4
 
 from pymilvus.client.constants import DEFAULT_CONSISTENCY_LEVEL
-from pymilvus.client.types import ExceptionsMessage, LoadState
+from pymilvus.client.types import (
+    ExceptionsMessage,
+    LoadState,
+    ExtraList,
+    get_cost_extra,
+    construct_cost_extra,
+)
 from pymilvus.exceptions import (
     DataTypeNotMatchException,
     MilvusException,
@@ -216,7 +222,11 @@ def insert(
             )
         except Exception as ex:
             raise ex from ex
-        return {"insert_count": res.insert_count, "ids": res.primary_keys}
+        return {
+            "insert_count": res.insert_count,
+            "ids": res.primary_keys,
+            "cost": res.cost,
+        }
 
     def upsert(
         self,
@@ -263,7 +273,10 @@ def upsert(
         except Exception as ex:
             raise ex from ex
 
-        return {"upsert_count": res.upsert_count}
+        return {
+            "upsert_count": res.upsert_count,
+            "cost": res.cost,
+        }
 
     def search(
         self,
@@ -325,7 +338,7 @@ def search(
                 query_result.append(hit.to_dict())
             ret.append(query_result)
 
-        return ret
+        return ExtraList(ret, extra=construct_cost_extra(res.cost))
 
     def query(
         self,
@@ -543,7 +556,7 @@ def delete(
         if ret_pks:
             return ret_pks
 
-        return {"delete_count": res.delete_count}
+        return {"delete_count": res.delete_count, "cost": res.cost}
 
     def get_collection_stats(self, collection_name: str, timeout: Optional[float] = None) -> Dict:
         conn = self._get_connection()

pymilvus/orm/mutation.py

@@ -53,6 +53,10 @@ def succ_index(self):
     def err_index(self):
         return self._mr.err_index if self._mr else []
 
+    @property
+    def cost(self):
+        return self._mr.cost if self._mr else 0
+
     def __str__(self) -> str:
         """
         Return the information of mutation result