enhance: Cherry pick multiple prs from master branch (#3048)

XuanYang-cn · andrii-shchur · iamjoker021 · web-flow · commit 9b0f35d433de · 2025-10-27T15:14:02.000+08:00
- enhance: Replace ujson with orjson (#2993) - Add PK in the upsert response (#3035) - fix: Declare `pymilvus` as a namespace package. (#3040) See also: #1354, #3033, #3039 --------- Signed-off-by: yangxuan <xuan.yang@zilliz.com> Co-authored-by: Andrii Shchur <33606307+andrii-shchur@users.noreply.github.com> Co-authored-by: Andrii Shchur <andrii.shchur@datarobot.com> Co-authored-by: Joker021 <iamjoker021@gmail.com> Co-authored-by: John Sirois <john.sirois@gmail.com>
diff --git a/pymilvus/__init__.py b/pymilvus/__init__.py
@@ -10,6 +10,12 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 
+# Ensure `pymilvus` is a namespace package other distributions (like `pymilvus.model`) can
+# participate in.
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+
 from .client import __version__
 from .client.abstract import AnnSearchRequest, RRFRanker, WeightedRanker
 from .client.asynch import SearchFuture
diff --git a/pymilvus/client/abstract.py b/pymilvus/client/abstract.py
@@ -2,7 +2,7 @@
 import logging
 from typing import Any, Dict, List, Optional, Union
 
-import ujson
+import orjson
 
 from pymilvus.exceptions import DataTypeNotMatchException, ExceptionsMessage
 from pymilvus.settings import Config
@@ -60,7 +60,7 @@ def __pack(self, raw: Any):
         for type_param in raw.type_params:
             if type_param.key == "params":
                 try:
-                    self.params[type_param.key] = ujson.loads(type_param.value)
+                    self.params[type_param.key] = orjson.loads(type_param.value)
                 except Exception as e:
                     logger.error(
                         f"FieldSchema::__pack::65::Failed to load JSON type_param.value: {e}, original data: {type_param.value}"
@@ -92,7 +92,7 @@ def __pack(self, raw: Any):
         for index_param in raw.index_params:
             if index_param.key == "params":
                 try:
-                    index_dict[index_param.key] = ujson.loads(index_param.value)
+                    index_dict[index_param.key] = orjson.loads(index_param.value)
                 except Exception as e:
                     logger.error(
                         f"FieldSchema::__pack::92::Failed to load JSON index_param.value: {e}, original data: {index_param.value}"
diff --git a/pymilvus/client/entity_helper.py b/pymilvus/client/entity_helper.py
@@ -4,7 +4,7 @@
 from typing import Any, Dict, Iterable, List, Optional
 
 import numpy as np
-import ujson
+import orjson
 
 from pymilvus.exceptions import (
     DataNotMatchException,
@@ -201,7 +201,7 @@ def preprocess_numpy_types(obj: Any):
     if isinstance(obj, str):
         try:
             # Validate JSON string by parsing it
-            parsed_obj = ujson.loads(obj)
+            parsed_obj = orjson.loads(obj)
             # If it's a valid JSON string, validate dict keys if it's a dict
             if isinstance(parsed_obj, dict):
                 for k in parsed_obj:
@@ -225,7 +225,7 @@ def preprocess_numpy_types(obj: Any):
 
     processed_obj = preprocess_numpy_types(obj)
 
-    return ujson.dumps(processed_obj, ensure_ascii=False).encode(Config.EncodeProtocol)
+    return orjson.dumps(processed_obj)
 
 
 def convert_to_json_arr(objs: List[object], field_info: Any):
@@ -933,7 +933,7 @@ def check_append(field_data: Any, row_data: Dict):
                 row_data[field_data.field_name] = None
                 return
             try:
-                json_dict = ujson.loads(field_data.scalars.json_data.data[index])
+                json_dict = orjson.loads(field_data.scalars.json_data.data[index])
             except Exception as e:
                 logger.error(
                     f"extract_row_data_from_fields_data::Failed to load JSON data: {e}, original data: {field_data.scalars.json_data.data[index]}"
diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
@@ -4,7 +4,7 @@
 from typing import Any, Dict, Iterable, List, Mapping, Optional, Union
 
 import numpy as np
-import ujson
+import orjson
 
 from pymilvus.exceptions import DataNotMatchException, ExceptionsMessage, ParamError
 from pymilvus.grpc_gen import common_pb2
@@ -19,6 +19,7 @@
     isVectorDataType,
 )
 from pymilvus.orm.types import infer_dtype_by_scalar_data
+from pymilvus.settings import Config
 
 from . import __version__, blob, check, entity_helper, ts_utils, utils
 from .abstract import BaseRanker
@@ -161,7 +162,11 @@ def get_schema_from_collection_schema(
             for k, v in f.params.items():
                 kv_pair = common_types.KeyValuePair(
                     key=str(k) if k != "mmap_enabled" else "mmap.enabled",
-                    value=ujson.dumps(v) if not isinstance(v, str) else str(v),
+                    value=(
+                        orjson.dumps(v).decode(Config.EncodeProtocol)
+                        if not isinstance(v, str)
+                        else str(v)
+                    ),
                 )
                 field_schema.type_params.append(kv_pair)
 
@@ -2311,7 +2316,7 @@ def run_analyzer(
 
         if analyzer_params is not None:
             if isinstance(analyzer_params, dict):
-                req.analyzer_params = ujson.dumps(analyzer_params)
+                req.analyzer_params = orjson.dumps(analyzer_params).decode(Config.EncodeProtocol)
             else:
                 req.analyzer_params = analyzer_params
 
diff --git a/pymilvus/client/search_result.py b/pymilvus/client/search_result.py
@@ -2,7 +2,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
-import ujson
+import orjson
 
 from pymilvus.client.types import DataType
 from pymilvus.exceptions import MilvusException
@@ -152,7 +152,7 @@ def materialize(self):
                             json_data = field_data.scalars.json_data.data[idx]
                             try:
                                 json_dict_list = (
-                                    ujson.loads(json_data) if json_data is not None else None
+                                    orjson.loads(json_data) if json_data is not None else None
                                 )
                             except Exception as e:
                                 logger.error(
@@ -420,7 +420,7 @@ def _get_fields_by_range(
                 for item in res:
                     if item is not None:
                         try:
-                            json_dict_list.append(ujson.loads(item))
+                            json_dict_list.append(orjson.loads(item))
                         except Exception as e:
                             logger.error(
                                 f"SearchResult::_get_fields_by_range::Failed to load JSON item: {e}, original item: {item}"
@@ -809,7 +809,7 @@ def extract_struct_field_value(field_data: schema_pb2.FieldData, index: int) ->
             return field_data.scalars.string_data.data[index]
     elif field_data.type == DataType.JSON:
         if index < len(field_data.scalars.json_data.data):
-            return ujson.loads(field_data.scalars.json_data.data[index])
+            return orjson.loads(field_data.scalars.json_data.data[index])
     elif field_data.type == DataType.FLOAT_VECTOR:
         dim = field_data.vectors.dim
         start_idx = index * dim
diff --git a/pymilvus/client/types.py b/pymilvus/client/types.py
@@ -4,7 +4,7 @@
 from typing import Any, ClassVar, Dict, List, Optional, TypeVar, Union
 
 import numpy as np
-import ujson
+import orjson
 
 from pymilvus.exceptions import (
     AutoIDException,
@@ -1081,7 +1081,7 @@ def _extract_lazy_fields(self, index: int, field_data: Any, row_data: Dict) -> A
                 row_data[field_data.field_name] = None
                 return
             try:
-                json_dict = ujson.loads(field_data.scalars.json_data.data[index])
+                json_dict = orjson.loads(field_data.scalars.json_data.data[index])
             except Exception as e:
                 logger.error(
                     f"HybridExtraList::_extract_lazy_fields::Failed to load JSON data: {e}, original data: {field_data.scalars.json_data.data[index]}"
diff --git a/pymilvus/client/utils.py b/pymilvus/client/utils.py
@@ -5,10 +5,11 @@
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
 
-import ujson
+import orjson
 
 from pymilvus.exceptions import MilvusException, ParamError
 from pymilvus.grpc_gen.common_pb2 import Status
+from pymilvus.settings import Config
 
 from .constants import LOGICAL_BITS, LOGICAL_BITS_MASK
 from .types import DataType
@@ -308,7 +309,7 @@ def get_server_type(host: str):
 
 
 def dumps(v: Union[dict, str]) -> str:
-    return ujson.dumps(v) if isinstance(v, dict) else str(v)
+    return orjson.dumps(v).decode(Config.EncodeProtocol) if isinstance(v, dict) else str(v)
 
 
 class SciPyHelper:
diff --git a/pymilvus/milvus_client/async_milvus_client.py b/pymilvus/milvus_client/async_milvus_client.py
@@ -338,7 +338,7 @@ async def upsert(
             raise TypeError(msg)
 
         if len(data) == 0:
-            return {"upsert_count": 0}
+            return {"upsert_count": 0, "ids": []}
 
         conn = self._get_connection()
         # Upsert into the collection.
@@ -353,6 +353,7 @@ async def upsert(
             {
                 "upsert_count": res.upsert_count,
                 "cost": res.cost,
+                "ids": res.primary_keys,
             }
         )
 
diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py
@@ -272,7 +272,7 @@ def upsert(
             raise TypeError(msg)
 
         if len(data) == 0:
-            return {"upsert_count": 0}
+            return {"upsert_count": 0, "ids": []}
 
         conn = self._get_connection()
         # Upsert into the collection.
@@ -289,7 +289,7 @@ def upsert(
                 "cost": res.cost,
                 # milvus server supports upsert on autoid=ture from v2.4.15
                 # upsert on autoid=ture will return new ids for user
-                "primary_keys": res.primary_keys,
+                "ids": res.primary_keys,
             }
         )
 
diff --git a/pymilvus/orm/partition.py b/pymilvus/orm/partition.py
@@ -12,14 +12,15 @@
 
 from typing import Dict, List, Optional, TypeVar, Union
 
+import orjson
 import pandas as pd
-import ujson
 
 from pymilvus.client import utils
 from pymilvus.client.abstract import BaseRanker
 from pymilvus.client.search_result import SearchResult
 from pymilvus.client.types import Replica
 from pymilvus.exceptions import MilvusException
+from pymilvus.settings import Config
 
 from .mutation import MutationResult
 
@@ -57,13 +58,13 @@ def __init__(
             conn.create_partition(self._collection.name, self.name, **kwargs)
 
     def __repr__(self) -> str:
-        return ujson.dumps(
+        return orjson.dumps(
             {
                 "name": self.name,
                 "collection_name": self._collection.name,
                 "description": self.description,
             }
-        )
+        ).decode(Config.EncodeProtocol)
 
     def _get_connection(self):
         return self._collection._get_connection()
diff --git a/pymilvus/orm/schema.py b/pymilvus/orm/schema.py
@@ -13,8 +13,8 @@
 import copy
 from typing import Any, Dict, List, Optional, Union
 
+import orjson
 import pandas as pd
-import ujson
 from pandas.api.types import is_list_like, is_scalar
 
 from pymilvus.client.types import FunctionType
@@ -35,6 +35,7 @@
     SchemaNotReadyException,
 )
 from pymilvus.grpc_gen import schema_pb2 as schema_types
+from pymilvus.settings import Config
 
 from .constants import COMMON_TYPE_PARAMS
 from .types import (
@@ -477,7 +478,7 @@ def __init__(self, name: str, dtype: DataType, description: str = "", **kwargs)
 
         for key in ["analyzer_params", "multi_analyzer_params"]:
             if key in self._kwargs and isinstance(self._kwargs[key], dict):
-                self._kwargs[key] = ujson.dumps(self._kwargs[key])
+                self._kwargs[key] = orjson.dumps(self._kwargs[key]).decode(Config.EncodeProtocol)
 
         self._parse_type_params()
         self.is_function_output = False
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,9 +23,9 @@ dependencies=[
     "setuptools>69",
     "setuptools<70.1;python_version<='3.8'",
     "grpcio>=1.66.2,!=1.68.0,!=1.68.1,!=1.69.0,!=1.70.0,!=1.70.1,!=1.71.0,!=1.72.1,!=1.73.0",
+    "orjson>=3.10.15", # orjson 3.10.15 is the last version to support python3.8
     "protobuf>=5.27.2", # aligned with grpcio-tools generated codes
     "python-dotenv>=1.0.1, <2.0.0",
-    "ujson>=2.0.0",
     "pandas>=1.2.4",
     "numpy<1.25.0;python_version<='3.8'",
 ]
diff --git a/requirements.txt b/requirements.txt
@@ -12,7 +12,7 @@ pep517==0.10.0
 pyparsing==2.4.7
 six==1.16.0
 toml==0.10.2
-ujson>=2.0.0
+orjson
 urllib3==2.5.0
 m2r==0.3.1
 Sphinx==4.0.0
diff --git a/tests/test_search_result.py b/tests/test_search_result.py
@@ -4,7 +4,7 @@
 from typing import Dict
 
 import pytest
-import ujson
+import orjson
 from pymilvus.client.search_result import Hit, Hits, HybridHits, SearchResult
 from pymilvus.client.types import DataType
 from pymilvus.grpc_gen import schema_pb2
@@ -157,11 +157,11 @@ def test_search_result_with_fields_data(self, pk):
                                  )),
 
             schema_pb2.FieldData(type=DataType.JSON, field_name="normal_json_field", field_id=111,
-                scalars=schema_pb2.ScalarField(json_data=schema_pb2.JSONArray(data=[ujson.dumps({i: i for i in range(3)}).encode() for i in range(6)])),
+                scalars=schema_pb2.ScalarField(json_data=schema_pb2.JSONArray(data=[orjson.dumps({str(i): i for i in range(3)}) for i in range(6)])),
             ),
             schema_pb2.FieldData(type=DataType.JSON, field_name="$meta", field_id=112,
                 is_dynamic=True,
-                scalars=schema_pb2.ScalarField(json_data=schema_pb2.JSONArray(data=[ujson.dumps({str(i*100): i}).encode() for i in range(6)])),
+                scalars=schema_pb2.ScalarField(json_data=schema_pb2.JSONArray(data=[orjson.dumps({str(i*100): i}) for i in range(6)])),
             ),
 
             schema_pb2.FieldData(type=DataType.FLOAT_VECTOR, field_name="float_vector_field", field_id=113,

Original file line number	Diff line number	Diff line change
`@@ -338,7 +338,7 @@ async def upsert(`
`338`	`338`	`raise TypeError(msg)`
`339`	`339`
`340`	`340`	`if len(data) == 0:`
`341`		`- return {"upsert_count": 0}`
	`341`	`+ return {"upsert_count": 0, "ids": []}`
`342`	`342`
`343`	`343`	`conn = self._get_connection()`
`344`	`344`	`# Upsert into the collection.`
`@@ -353,6 +353,7 @@ async def upsert(`
`353`	`353`	`{`
`354`	`354`	`"upsert_count": res.upsert_count,`
`355`	`355`	`"cost": res.cost,`
	`356`	`+ "ids": res.primary_keys,`
`356`	`357`	`}`
`357`	`358`	`)`
`358`	`359`
Original file line number	Diff line number	Diff line change
`@@ -272,7 +272,7 @@ def upsert(`
`272`	`272`	`raise TypeError(msg)`
`273`	`273`
`274`	`274`	`if len(data) == 0:`
`275`		`- return {"upsert_count": 0}`
	`275`	`+ return {"upsert_count": 0, "ids": []}`
`276`	`276`
`277`	`277`	`conn = self._get_connection()`
`278`	`278`	`# Upsert into the collection.`
`@@ -289,7 +289,7 @@ def upsert(`
`289`	`289`	`"cost": res.cost,`
`290`	`290`	`# milvus server supports upsert on autoid=ture from v2.4.15`
`291`	`291`	`# upsert on autoid=ture will return new ids for user`
`292`		`- "primary_keys": res.primary_keys,`
	`292`	`+ "ids": res.primary_keys,`
`293`	`293`	`}`
`294`	`294`	`)`
`295`	`295`