[Cherry-pick] Backport 6 commits to 2.6.4

- enhance: Migrate six APIs to MilvusClient (#3045)
- fix: Add support for numpy ndarrays in Array fields (#3069) (#3070)
- Fix: MilvusClient.insert() does not pass **kwargs to underlying insert_rows() call (#3076)
- feat: add detailed traceback to error_handler decorator
- fix(async): include event loop ID in connection alias to prevent reusing closed connections (#3086)
- fix: fixed the key names for retrieving segment info in the MilvusClient (#3098)

Co-authored-by: XuanYang-cn <[email protected]>
Co-authored-by: zhuwenxing <[email protected]>
Co-authored-by: tianhang <[email protected]>
Co-authored-by: jac <[email protected]>
Co-authored-by: wt <[email protected]>

Author: 6 people

.gitignore

@@ -42,3 +42,4 @@ uv.lock
 
 # AI rules
 WARP.md
+CLAUDE.md

pymilvus/bulk_writer/constants.py

@@ -62,7 +62,7 @@
     DataType.BFLOAT16_VECTOR.name: lambda x, dim: float16_vector_validator(x, dim, True),
     DataType.SPARSE_FLOAT_VECTOR.name: lambda x: sparse_vector_validator(x),
     DataType.INT8_VECTOR.name: lambda x, dim: int8_vector_validator(x, dim),
-    DataType.ARRAY.name: lambda x, cap: isinstance(x, list) and len(x) <= cap,
+    DataType.ARRAY.name: lambda x, cap: (isinstance(x, (list, np.ndarray)) and len(x) <= cap),
 }
 
 NUMPY_TYPE_CREATOR = {

pymilvus/client/entity_helper.py

@@ -246,6 +246,10 @@ def convert_to_array_arr(objs: List[Any], field_info: Any):
 
 
 def convert_to_array(obj: List[Any], field_info: Any):
+    # Convert numpy ndarray to list if needed
+    if isinstance(obj, np.ndarray):
+        obj = obj.tolist()
+
     field_data = schema_types.ScalarField()
     element_type = field_info.get("element_type", None)
     if element_type == DataType.BOOL:

pymilvus/client/grpc_handler.py

@@ -1098,14 +1098,15 @@ def hybrid_search(
         )
 
     @retry_on_rpc_failure()
-    def get_query_segment_info(self, collection_name: str, timeout: float = 30, **kwargs):
+    def get_query_segment_info(
+        self, collection_name: str, timeout: float = 30, **kwargs
+    ) -> List[milvus_types.QuerySegmentInfo]:
         req = Prepare.get_query_segment_info_request(collection_name)
         response = self._stub.GetQuerySegmentInfo(
             req, timeout=timeout, metadata=_api_level_md(**kwargs)
         )
-        status = response.status
-        check_status(status)
-        return response.infos  # todo: A wrapper class of QuerySegmentInfo
+        check_status(response.status)
+        return response.infos
 
     @retry_on_rpc_failure()
     def create_alias(
@@ -1665,18 +1666,18 @@ def get_flush_state(
         response = self._stub.GetFlushState(req, timeout=timeout, metadata=_api_level_md(**kwargs))
         status = response.status
         check_status(status)
-        return response.flushed  # todo: A wrapper class of PersistentSegmentInfo
+        return response.flushed
 
     @retry_on_rpc_failure()
     def get_persistent_segment_infos(
         self, collection_name: str, timeout: Optional[float] = None, **kwargs
-    ):
+    ) -> List[milvus_types.PersistentSegmentInfo]:
         req = Prepare.get_persistent_segment_info_request(collection_name)
         response = self._stub.GetPersistentSegmentInfo(
             req, timeout=timeout, metadata=_api_level_md(**kwargs)
         )
         check_status(response.status)
-        return response.infos  # todo: A wrapper class of PersistentSegmentInfo
+        return response.infos
 
     def _wait_for_flushed(
         self,

pymilvus/client/prepare.py

@@ -543,6 +543,10 @@ def _process_struct_field(
             struct_sub_field_info: Two-level dict [struct_name][field_name] -> field info
             struct_sub_fields_data: Two-level dict [struct_name][field_name] -> FieldData
         """
+        # Convert numpy ndarray to list if needed
+        if isinstance(values, np.ndarray):
+            values = values.tolist()
+
         if not isinstance(values, list):
             msg = f"Field '{field_name}': Expected list, got {type(values).__name__}"
             raise TypeError(msg)

pymilvus/client/types.py

@@ -1,5 +1,6 @@
 import logging
 import time
+from dataclasses import dataclass
 from enum import IntEnum
 from typing import Any, ClassVar, Dict, List, Optional, TypeVar, Union
 
@@ -1372,3 +1373,20 @@ def __str__(self) -> str:
         return str(self.tokens)
 
     __repr__ = __str__
+
+
+@dataclass
+class SegmentInfo:
+    segment_id: int
+    collection_id: int
+    collection_name: str
+    num_rows: int
+    is_sorted: bool
+    state: common_pb2.SegmentState
+    level: common_pb2.SegmentLevel
+    storage_version: int
+
+
+@dataclass
+class LoadedSegmentInfo(SegmentInfo):
+    mem_size: int

pymilvus/decorators.py

@@ -3,6 +3,7 @@
 import functools
 import logging
 import time
+import traceback
 from typing import Any, Callable, Optional
 
 import grpc
@@ -216,25 +217,37 @@ async def async_handler(*args, **kwargs):
                     return await func(*args, **kwargs)
                 except MilvusException as e:
                     record_dict["RPC error"] = str(datetime.datetime.now())
-                    LOGGER.error(f"RPC error: [{inner_name}], {e}, <Time:{record_dict}>")
+                    tb_str = traceback.format_exc()
+                    LOGGER.error(
+                        f"RPC error: [{inner_name}], {e}, <Time:{record_dict}>\n"
+                        f"Traceback:\n{tb_str}"
+                    )
                     raise e from e
                 except grpc.FutureTimeoutError as e:
                     record_dict["gRPC timeout"] = str(datetime.datetime.now())
+                    tb_str = traceback.format_exc()
                     LOGGER.error(
                         f"grpc Timeout: [{inner_name}], <{e.__class__.__name__}: "
-                        f"{e.code()}, {e.details()}>, <Time:{record_dict}>"
+                        f"{e.code()}, {e.details()}>, <Time:{record_dict}>\n"
+                        f"Traceback:\n{tb_str}"
                     )
                     raise e from e
                 except grpc.RpcError as e:
                     record_dict["gRPC error"] = str(datetime.datetime.now())
+                    tb_str = traceback.format_exc()
                     LOGGER.error(
                         f"grpc RpcError: [{inner_name}], <{e.__class__.__name__}: "
-                        f"{e.code()}, {e.details()}>, <Time:{record_dict}>"
+                        f"{e.code()}, {e.details()}>, <Time:{record_dict}>\n"
+                        f"Traceback:\n{tb_str}"
                     )
                     raise e from e
                 except Exception as e:
                     record_dict["Exception"] = str(datetime.datetime.now())
-                    LOGGER.error(f"Unexpected error: [{inner_name}], {e}, <Time: {record_dict}>")
+                    tb_str = traceback.format_exc()
+                    LOGGER.error(
+                        f"Unexpected error: [{inner_name}], {e}, <Time: {record_dict}>\n"
+                        f"Traceback:\n{tb_str}"
+                    )
                     raise MilvusException(message=f"Unexpected error, message=<{e!s}>") from e
 
             return async_handler
@@ -250,25 +263,37 @@ def handler(*args, **kwargs):
                 return func(*args, **kwargs)
             except MilvusException as e:
                 record_dict["RPC error"] = str(datetime.datetime.now())
-                LOGGER.error(f"RPC error: [{inner_name}], {e}, <Time:{record_dict}>")
+                tb_str = traceback.format_exc()
+                LOGGER.error(
+                    f"RPC error: [{inner_name}], {e}, <Time:{record_dict}>\n"
+                    f"Traceback:\n{tb_str}"
+                )
                 raise e from e
             except grpc.FutureTimeoutError as e:
                 record_dict["gRPC timeout"] = str(datetime.datetime.now())
+                tb_str = traceback.format_exc()
                 LOGGER.error(
                     f"grpc Timeout: [{inner_name}], <{e.__class__.__name__}: "
-                    f"{e.code()}, {e.details()}>, <Time:{record_dict}>"
+                    f"{e.code()}, {e.details()}>, <Time:{record_dict}>\n"
+                    f"Traceback:\n{tb_str}"
                 )
                 raise e from e
             except grpc.RpcError as e:
                 record_dict["gRPC error"] = str(datetime.datetime.now())
+                tb_str = traceback.format_exc()
                 LOGGER.error(
                     f"grpc RpcError: [{inner_name}], <{e.__class__.__name__}: "
-                    f"{e.code()}, {e.details()}>, <Time:{record_dict}>"
+                    f"{e.code()}, {e.details()}>, <Time:{record_dict}>\n"
+                    f"Traceback:\n{tb_str}"
                 )
                 raise e from e
             except Exception as e:
                 record_dict["Exception"] = str(datetime.datetime.now())
-                LOGGER.error(f"Unexpected error: [{inner_name}], {e}, <Time: {record_dict}>")
+                tb_str = traceback.format_exc()
+                LOGGER.error(
+                    f"Unexpected error: [{inner_name}], {e}, <Time: {record_dict}>\n"
+                    f"Traceback:\n{tb_str}"
+                )
                 raise MilvusException(message=f"Unexpected error, message=<{e!s}>") from e
 
         return handler

pymilvus/milvus_client/_utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import hashlib
 import logging
 
@@ -33,8 +34,15 @@ def create_connection(
             md5.update(token.encode())
             auth_fmt = f"{md5.hexdigest()}"
 
-        # different uri, auth, db_name cannot share the same connection
-        not_empty = [v for v in [use_async_fmt, uri, db_name, auth_fmt] if v]
+        # For async connections, include event loop ID in alias to prevent
+        # reusing connections from closed event loops
+        loop_id_fmt = ""
+        if use_async:
+            loop = asyncio.get_running_loop()
+            loop_id_fmt = f"loop{id(loop)}"
+
+        # different uri, auth, db_name, and event loop (for async) cannot share the same connection
+        not_empty = [v for v in [use_async_fmt, uri, db_name, auth_fmt, loop_id_fmt] if v]
         using = "-".join(not_empty)
 
     if connections.has_connection(using):

pymilvus/milvus_client/async_milvus_client.py

@@ -287,7 +287,7 @@ async def insert(
         conn = self._get_connection()
         # Insert into the collection.
         res = await conn.insert_rows(
-            collection_name, data, partition_name=partition_name, timeout=timeout
+            collection_name, data, partition_name=partition_name, timeout=timeout, **kwargs
         )
         return OmitZeroDict(
             {

pymilvus/milvus_client/milvus_client.py

@@ -6,11 +6,14 @@
 from pymilvus.client.embedding_list import EmbeddingList
 from pymilvus.client.search_iterator import SearchIteratorV2
 from pymilvus.client.types import (
+    CompactionPlans,
     ExceptionsMessage,
+    LoadedSegmentInfo,
     LoadState,
     OmitZeroDict,
     ReplicaInfo,
     ResourceGroupConfig,
+    SegmentInfo,
 )
 from pymilvus.client.utils import convert_struct_fields_to_user_format, get_params, is_vector_type
 from pymilvus.exceptions import (
@@ -21,7 +24,6 @@
     PrimaryKeyException,
     ServerVersionIncompatibleException,
 )
-from pymilvus.orm import utility
 from pymilvus.orm.collection import CollectionSchema, FieldSchema, Function, FunctionScore
 from pymilvus.orm.connections import connections
 from pymilvus.orm.constants import FIELDS, METRIC_TYPE, TYPE, UNLIMITED
@@ -67,7 +69,7 @@ def __init__(
         self._using = create_connection(
             uri, token, db_name, user=user, password=password, timeout=timeout, **kwargs
         )
-        self.is_self_hosted = bool(utility.get_server_type(using=self._using) == "milvus")
+        self.is_self_hosted = bool(self.get_server_type() == "milvus")
 
     def create_collection(
         self,
@@ -219,7 +221,7 @@ def insert(
         # Insert into the collection.
         try:
             res = conn.insert_rows(
-                collection_name, data, partition_name=partition_name, timeout=timeout
+                collection_name, data, partition_name=partition_name, timeout=timeout, **kwargs
             )
         except Exception as ex:
             raise ex from ex
@@ -1808,3 +1810,117 @@ def update_replicate_configuration(
             timeout=timeout,
             **kwargs,
         )
+
+    def flush_all(self, timeout: Optional[float] = None, **kwargs) -> None:
+        """Flush all collections.
+
+        Args:
+            timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC.
+            **kwargs: Additional arguments.
+        """
+        self._get_connection().flush_all(timeout=timeout, **kwargs)
+
+    def get_flush_all_state(self, timeout: Optional[float] = None, **kwargs) -> bool:
+        """Get the flush all state.
+
+        Args:
+            timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC.
+            **kwargs: Additional arguments.
+
+        Returns:
+            bool: True if flush all operation is completed, False otherwise.
+        """
+        return self._get_connection().get_flush_all_state(timeout=timeout, **kwargs)
+
+    def list_loaded_segments(
+        self,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ) -> List[LoadedSegmentInfo]:
+        """List loaded segments for a collection.
+
+        Args:
+            collection_name (str): The name of the collection.
+            timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC.
+            **kwargs: Additional arguments.
+
+        Returns:
+            List[LoadedSegmentInfo]: A list of loaded segment information.
+        """
+        infos = self._get_connection().get_query_segment_info(
+            collection_name, timeout=timeout, **kwargs
+        )
+        return [
+            LoadedSegmentInfo(
+                info.segmentID,
+                info.collectionID,
+                collection_name,
+                info.num_rows,
+                info.is_sorted,
+                info.state,
+                info.level,
+                info.storage_version,
+                info.mem_size,
+            )
+            for info in infos
+        ]
+
+    def list_persistent_segments(
+        self,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ) -> List[SegmentInfo]:
+        """List persistent segments for a collection.
+
+        Args:
+            collection_name (str): The name of the collection.
+            timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC.
+            **kwargs: Additional arguments.
+
+        Returns:
+            List[SegmentInfo]: A list of persistent segment information.
+        """
+        infos = self._get_connection().get_persistent_segment_infos(
+            collection_name, timeout=timeout, **kwargs
+        )
+        return [
+            SegmentInfo(
+                info.segmentID,
+                info.collectionID,
+                collection_name,
+                info.num_rows,
+                info.is_sorted,
+                info.state,
+                info.level,
+                info.storage_version,
+            )
+            for info in infos
+        ]
+
+    def get_server_type(self):
+        """Get the server type.
+
+        Returns:
+            str: The server type (e.g., "milvus", "zilliz").
+        """
+        return self._get_connection().get_server_type()
+
+    def get_compaction_plans(
+        self,
+        job_id: int,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ) -> CompactionPlans:
+        """Get compaction plans for a specific job.
+
+        Args:
+            job_id (int): The ID of the compaction job.
+            timeout (Optional[float]): An optional duration of time in seconds to allow for the RPC.
+            **kwargs: Additional arguments.
+
+        Returns:
+            CompactionPlans: The compaction plans for the specified job.
+        """
+        return self._get_connection().get_compaction_plans(job_id, timeout=timeout, **kwargs)