Revert "collection.insert to return connection insert result directly (…

…#1943)" (#1946) This reverts commit 06799b3. Signed-off-by: Buqian Zheng <[email protected]>
milvus-io · Mar 4, 2024 · 6f4046c · 6f4046c
1 parent 06799b3
commit 6f4046c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 10 deletions.
diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
@@ -559,7 +559,7 @@ def batch_insert(
                 collection_name, entities, partition_name, timeout, **kwargs
             )
             rf = self._stub.Insert.future(request, timeout=timeout)
-            if kwargs.get("_async", False):
+            if kwargs.get("_async", False) is True:
                 cb = kwargs.get("_callback", None)
                 f = MutationFuture(rf, cb, timeout=timeout, **kwargs)
                 f.add_callback(ts_utils.update_ts_on_mutation(collection_name))

diff --git a/pymilvus/orm/collection.py b/pymilvus/orm/collection.py
@@ -493,15 +493,27 @@ def insert(
         conn = self._get_connection()
         if not row_based:
             check_insert_schema(self._schema, data)
-            data = Prepare.prepare_insert_data(data, self._schema)
-        return conn.batch_insert(
-            self._name,
-            data,
-            partition_name,
-            timeout=timeout,
-            schema=self._schema_dict,
-            **kwargs,
-        )
+            entities = Prepare.prepare_insert_data(data, self._schema)
+            res = conn.batch_insert(
+                self._name,
+                entities,
+                partition_name,
+                timeout=timeout,
+                schema=self._schema_dict,
+                **kwargs,
+            )
+            if kwargs.get("_async", False):
+                return MutationFuture(res)
+        else:
+            res = conn.insert_rows(
+                collection_name=self._name,
+                entities=data,
+                partition_name=partition_name,
+                timeout=timeout,
+                schema=self._schema_dict,
+                **kwargs,
+            )
+        return MutationResult(res)
 
     def delete(
         self,