diff --git a/internal/proxy/task_delete.go b/internal/proxy/task_delete.go
index 281b3f700b590..c91ee9b181b16 100644
--- a/internal/proxy/task_delete.go
+++ b/internal/proxy/task_delete.go
@@ -389,15 +389,10 @@ func (dr *deleteRunner) Init(ctx context.Context) error {
 }
 
 func (dr *deleteRunner) Run(ctx context.Context) error {
-	// Make sure partitionIDs aren't empty
-	if len(dr.partitionIDs) == 0 {
-		dr.partitionIDs = []UniqueID{common.AllPartitionsID}
-	}
-
 	isSimple, pk, numRow := getPrimaryKeysFromPlan(dr.schema.CollectionSchema, dr.plan)
 	if isSimple {
 		// if could get delete.primaryKeys from delete expr
-		err := dr.simpleDelete(ctx, pk, numRow, dr.partitionIDs[0])
+		err := dr.simpleDelete(ctx, pk, numRow)
 		if err != nil {
 			return err
 		}
@@ -589,7 +584,11 @@ func (dr *deleteRunner) complexDelete(ctx context.Context, plan *planpb.PlanNode
 	return nil
 }
 
-func (dr *deleteRunner) simpleDelete(ctx context.Context, pk *schemapb.IDs, numRow int64, partitionID UniqueID) error {
+func (dr *deleteRunner) simpleDelete(ctx context.Context, pk *schemapb.IDs, numRow int64) error {
+	partitionID := common.AllPartitionsID
+	if len(dr.partitionIDs) == 1 {
+		partitionID = dr.partitionIDs[0]
+	}
 	log.Ctx.Debug("get primary keys from expr",
 		zap.Int64("len of primary keys", numRow),
 		zap.Int64("collectionID", dr.collectionID),