fix: select a single column consisting of a list of column names

Signed-off-by: lixinguo <xinguo.li@zilliz.com>
milvus-io · Aug 5, 2024 · 66c2362 · 66c2362
1 parent 9ff4b2c
commit 66c2362
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pymilvus/orm/schema.py b/pymilvus/orm/schema.py
@@ -510,7 +510,7 @@ def check_insert_schema(schema: CollectionSchema, data: Union[List[List], pd.Dat
             raise DataNotMatchException(message=msg)
         columns = list(data.columns)
         columns.remove(schema.primary_field)
-        data = data[[columns]]
+        data = data[columns]
 
     tmp_fields = copy.deepcopy(schema.fields)
     for i, field in enumerate(tmp_fields):
@@ -528,7 +528,7 @@ def check_upsert_schema(schema: CollectionSchema, data: Union[List[List], pd.Dat
         if schema.primary_field.name not in data or data[schema.primary_field.name].isnull().all():
             raise DataNotMatchException(message=ExceptionsMessage.UpsertPrimaryKeyEmpty)
         columns = list(data.columns)
-        data = data[[columns]]
+        data = data[columns]
 
     _check_data_schema_cnt(copy.deepcopy(schema.fields), data)
     _check_insert_data(data)