diff --git a/internal/proxy/task_upsert.go b/internal/proxy/task_upsert.go
index f8415f6d55..b89c09d954 100644
--- a/internal/proxy/task_upsert.go
+++ b/internal/proxy/task_upsert.go
@@ -198,8 +198,8 @@ func (it *upsertTask) insertPreExecute(ctx context.Context) error {
 	}
 
 	if it.partitionKeyMode {
-		pkFieldSchema, _ := it.schema.GetPkField()
-		it.partitionKeys, err = getPartitionKeyFieldData(pkFieldSchema, it.upsertMsg.InsertMsg)
+		fieldSchema, _ := typeutil.GetPartitionKeyFieldSchema(it.schema.CollectionSchema)
+		it.partitionKeys, err = getPartitionKeyFieldData(fieldSchema, it.upsertMsg.InsertMsg)
 		if err != nil {
 			log.Warn("get partition keys from insert request failed",
 				zap.String("collectionName", collectionName),