[cherry-pick-1.1-dev] fix load data stmt affect rows count calculatio…

…n incorrectly (#15435) fix load data stmt affect rows count calculation incorrectly Approved by: @nnsgmsone, @m-schen, @sukki37
matrixorigin · Apr 10, 2024 · 6688544 · 6688544
1 parent d301800
commit 6688544
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/pkg/sql/colexec/mergeblock/mergeblock_test.go b/pkg/sql/colexec/mergeblock/mergeblock_test.go
@@ -119,6 +119,7 @@ func TestMergeBlock(t *testing.T) {
 			IsFirst: false,
 			IsLast:  false,
 		},
+		AddAffectedRows: true,
 	}
 	resetChildren(&argument1, batch1)
 
@@ -210,6 +211,7 @@ func TestArgument_GetMetaLocBat(t *testing.T) {
 			IsFirst: false,
 			IsLast:  false,
 		},
+		AddAffectedRows: true,
 	}
 
 	proc := testutil.NewProc()

diff --git a/pkg/sql/colexec/mergeblock/types.go b/pkg/sql/colexec/mergeblock/types.go
@@ -44,6 +44,7 @@ type Argument struct {
 	Tbl engine.Relation
 	// 2. partition sub tables
 	PartitionSources []engine.Relation
+	AddAffectedRows  bool
 	affectedRows     uint64
 	container        *Container
 
@@ -164,8 +165,10 @@ func (arg *Argument) Split(proc *process.Process, bat *batch.Batch) error {
 	hasObject := false
 	for i := range tblIdx { // append s3 writer returned blk info
 		if tblIdx[i] >= 0 {
-			blkInfo := catalog.DecodeBlockInfo(blkInfosVec.GetBytesAt(i))
-			arg.affectedRows += uint64(blkInfo.MetaLocation().Rows())
+			if arg.AddAffectedRows {
+				blkInfo := catalog.DecodeBlockInfo(blkInfosVec.GetBytesAt(i))
+				arg.affectedRows += uint64(blkInfo.MetaLocation().Rows())
+			}
 			vector.AppendBytes(arg.container.mp[int(tblIdx[i])].Vecs[0],
 				blkInfosVec.GetBytesAt(i), false, proc.GetMPool())
 			hasObject = true
@@ -176,7 +179,9 @@ func (arg *Argument) Split(proc *process.Process, bat *batch.Batch) error {
 				return err
 			}
 			newBat.Cnt = 1
-			arg.affectedRows += uint64(newBat.RowCount())
+			if arg.AddAffectedRows {
+				arg.affectedRows += uint64(newBat.RowCount())
+			}
 			arg.container.mp2[idx] = append(arg.container.mp2[idx], newBat)
 		}
 	}

diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go
@@ -1410,6 +1410,7 @@ func (c *Compile) compilePlanScope(ctx context.Context, step int32, curNodeIdx i
 					Arg: &mergeblock.Argument{
 						Tbl:              insertArg.InsertCtx.Rel,
 						PartitionSources: insertArg.InsertCtx.PartitionSources,
+						AddAffectedRows:  insertArg.InsertCtx.AddAffectedRows,
 					},
 				})
 				ss = []*Scope{rs}
@@ -1477,6 +1478,7 @@ func (c *Compile) compilePlanScope(ctx context.Context, step int32, curNodeIdx i
 					Arg: &mergeblock.Argument{
 						Tbl:              insertArg.InsertCtx.Rel,
 						PartitionSources: insertArg.InsertCtx.PartitionSources,
+						AddAffectedRows:  insertArg.InsertCtx.AddAffectedRows,
 					},
 				})
 				ss = []*Scope{rs}