Skip to content

Commit 182f5b9

Browse files
authored
fix: replace row affected 3.0 (#23187)
fix replace row affected
1 parent e0d63c4 commit 182f5b9

File tree

3 files changed

+210
-3
lines changed

3 files changed

+210
-3
lines changed

pkg/sql/colexec/multi_update/multi_update.go

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,8 @@ func (update *MultiUpdate) updateFlushS3Info(proc *process.Process, analyzer pro
311311
if err := batBufs[actionDelete].UnmarshalBinary(batData[i].GetByteSlice(batArea)); err != nil {
312312
return input, err
313313
}
314-
update.addDeleteAffectRows(tableType, rowCounts[i])
314+
// For REPLACE INTO, we don't count DELETE rows in affected rows
315+
// REPLACE INTO should only return the number of INSERT rows
315316
name := nameData[i].UnsafeGetString(nameArea)
316317

317318
crs := analyzer.GetOpCounterSet()
@@ -335,7 +336,11 @@ func (update *MultiUpdate) updateFlushS3Info(proc *process.Process, analyzer pro
335336
return input, err
336337
}
337338

338-
update.addInsertAffectRows(tableType, rowCounts[i])
339+
// For REPLACE INTO, we need to count INSERT rows based on the actual action
340+
// Always count INSERT rows for main table, regardless of update.ctr.action
341+
if tableType == UpdateMainTable {
342+
update.addAffectedRowsFunc(rowCounts[i])
343+
}
339344

340345
crs := analyzer.GetOpCounterSet()
341346
newCtx := perfcounter.AttachS3RequestKey(ctx, crs)

pkg/sql/colexec/multi_update/multi_update_partition_test.go

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,193 @@ func TestNewPartitionMultiUpdateFrom(t *testing.T) {
4040
require.Equal(t, ps.raw.Engine, op.(*PartitionMultiUpdate).raw.Engine)
4141
require.Equal(t, ps.tableID, op.(*PartitionMultiUpdate).tableID)
4242
}
43+
44+
func TestAddInsertAffectRows(t *testing.T) {
45+
tests := []struct {
46+
name string
47+
action actionType
48+
tableType UpdateTableType
49+
rowCount uint64
50+
expectRows uint64
51+
}{
52+
{
53+
name: "actionInsert with main table",
54+
action: actionInsert,
55+
tableType: UpdateMainTable,
56+
rowCount: 5,
57+
expectRows: 5,
58+
},
59+
{
60+
name: "actionUpdate with main table (REPLACE INTO)",
61+
action: actionUpdate,
62+
tableType: UpdateMainTable,
63+
rowCount: 3,
64+
expectRows: 3,
65+
},
66+
{
67+
name: "actionInsert with unique index table (should not count)",
68+
action: actionInsert,
69+
tableType: UpdateUniqueIndexTable,
70+
rowCount: 5,
71+
expectRows: 0,
72+
},
73+
{
74+
name: "actionUpdate with unique index table (should not count)",
75+
action: actionUpdate,
76+
tableType: UpdateUniqueIndexTable,
77+
rowCount: 3,
78+
expectRows: 0,
79+
},
80+
{
81+
name: "actionDelete with main table (should not count)",
82+
action: actionDelete,
83+
tableType: UpdateMainTable,
84+
rowCount: 2,
85+
expectRows: 0,
86+
},
87+
}
88+
89+
for _, tt := range tests {
90+
t.Run(tt.name, func(t *testing.T) {
91+
update := &MultiUpdate{
92+
ctr: container{
93+
action: tt.action,
94+
affectedRows: 0,
95+
},
96+
}
97+
update.addAffectedRowsFunc = update.doAddAffectedRows
98+
99+
update.addInsertAffectRows(tt.tableType, tt.rowCount)
100+
101+
require.Equal(t, tt.expectRows, update.ctr.affectedRows, "affected rows should match expected value")
102+
})
103+
}
104+
}
105+
106+
func TestAddDeleteAffectRows(t *testing.T) {
107+
tests := []struct {
108+
name string
109+
action actionType
110+
tableType UpdateTableType
111+
rowCount uint64
112+
expectRows uint64
113+
}{
114+
{
115+
name: "actionDelete with main table",
116+
action: actionDelete,
117+
tableType: UpdateMainTable,
118+
rowCount: 5,
119+
expectRows: 5,
120+
},
121+
{
122+
name: "actionUpdate with main table (should not count for REPLACE INTO)",
123+
action: actionUpdate,
124+
tableType: UpdateMainTable,
125+
rowCount: 3,
126+
expectRows: 0,
127+
},
128+
{
129+
name: "actionDelete with unique index table (should not count)",
130+
action: actionDelete,
131+
tableType: UpdateUniqueIndexTable,
132+
rowCount: 5,
133+
expectRows: 0,
134+
},
135+
{
136+
name: "actionUpdate with unique index table (should not count)",
137+
action: actionUpdate,
138+
tableType: UpdateUniqueIndexTable,
139+
rowCount: 3,
140+
expectRows: 0,
141+
},
142+
{
143+
name: "actionInsert with main table (should not count)",
144+
action: actionInsert,
145+
tableType: UpdateMainTable,
146+
rowCount: 2,
147+
expectRows: 0,
148+
},
149+
}
150+
151+
for _, tt := range tests {
152+
t.Run(tt.name, func(t *testing.T) {
153+
update := &MultiUpdate{
154+
ctr: container{
155+
action: tt.action,
156+
affectedRows: 0,
157+
},
158+
}
159+
update.addAffectedRowsFunc = update.doAddAffectedRows
160+
161+
update.addDeleteAffectRows(tt.tableType, tt.rowCount)
162+
163+
require.Equal(t, tt.expectRows, update.ctr.affectedRows, "affected rows should match expected value")
164+
})
165+
}
166+
}
167+
168+
func TestReplaceIntoAffectedRows(t *testing.T) {
169+
// Test REPLACE INTO scenario: should only count INSERT rows, not DELETE rows
170+
update := &MultiUpdate{
171+
ctr: container{
172+
action: actionUpdate, // REPLACE INTO uses actionUpdate
173+
affectedRows: 0,
174+
},
175+
}
176+
update.addAffectedRowsFunc = update.doAddAffectedRows
177+
178+
// Simulate REPLACE INTO: DELETE 2 rows, INSERT 2 rows
179+
// Should only count INSERT rows (2), not DELETE rows
180+
update.addDeleteAffectRows(UpdateMainTable, 2) // Should not count
181+
require.Equal(t, uint64(0), update.ctr.affectedRows, "DELETE rows should not be counted for REPLACE INTO")
182+
183+
update.addInsertAffectRows(UpdateMainTable, 2) // Should count
184+
require.Equal(t, uint64(2), update.ctr.affectedRows, "INSERT rows should be counted for REPLACE INTO")
185+
}
186+
187+
func TestUpdateAffectedRows(t *testing.T) {
188+
// Test UPDATE scenario: should only count INSERT rows (updated rows), not DELETE rows
189+
update := &MultiUpdate{
190+
ctr: container{
191+
action: actionUpdate, // UPDATE uses actionUpdate
192+
affectedRows: 0,
193+
},
194+
}
195+
update.addAffectedRowsFunc = update.doAddAffectedRows
196+
197+
// Simulate UPDATE: DELETE 3 rows, INSERT 3 rows
198+
// Should only count INSERT rows (3), not DELETE rows
199+
update.addDeleteAffectRows(UpdateMainTable, 3) // Should not count
200+
require.Equal(t, uint64(0), update.ctr.affectedRows, "DELETE rows should not be counted for UPDATE")
201+
202+
update.addInsertAffectRows(UpdateMainTable, 3) // Should count
203+
require.Equal(t, uint64(3), update.ctr.affectedRows, "INSERT rows should be counted for UPDATE")
204+
}
205+
206+
func TestInsertAffectedRows(t *testing.T) {
207+
// Test INSERT scenario: should count INSERT rows
208+
update := &MultiUpdate{
209+
ctr: container{
210+
action: actionInsert,
211+
affectedRows: 0,
212+
},
213+
}
214+
update.addAffectedRowsFunc = update.doAddAffectedRows
215+
216+
update.addInsertAffectRows(UpdateMainTable, 5)
217+
require.Equal(t, uint64(5), update.ctr.affectedRows, "INSERT rows should be counted")
218+
}
219+
220+
func TestDeleteAffectedRows(t *testing.T) {
221+
// Test DELETE scenario: should count DELETE rows
222+
update := &MultiUpdate{
223+
ctr: container{
224+
action: actionDelete,
225+
affectedRows: 0,
226+
},
227+
}
228+
update.addAffectedRowsFunc = update.doAddAffectedRows
229+
230+
update.addDeleteAffectRows(UpdateMainTable, 4)
231+
require.Equal(t, uint64(4), update.ctr.affectedRows, "DELETE rows should be counted")
232+
}

pkg/sql/colexec/multi_update/types.go

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,21 +190,33 @@ func (update *MultiUpdate) addInsertAffectRows(tableType UpdateTableType, rowCou
190190
if tableType != UpdateMainTable {
191191
return
192192
}
193+
// For REPLACE INTO, we always count INSERT rows, regardless of update.ctr.action
194+
// because REPLACE INTO should return at least the number of rows being inserted
193195
switch update.ctr.action {
194196
case actionInsert:
195197
update.addAffectedRowsFunc(rowCount)
198+
case actionUpdate:
199+
// For REPLACE INTO with both DELETE and INSERT, count INSERT rows
200+
update.addAffectedRowsFunc(rowCount)
196201
}
197202
}
198203

199204
func (update *MultiUpdate) addDeleteAffectRows(tableType UpdateTableType, rowCount uint64) {
200205
if tableType != UpdateMainTable {
201206
return
202207
}
208+
// For REPLACE INTO, we don't count DELETE rows in affected rows
209+
// REPLACE INTO should only return the number of INSERT rows
210+
// Only count DELETE rows for regular UPDATE operations
203211
switch update.ctr.action {
204212
case actionDelete:
213+
// Regular DELETE operation, count it
205214
update.addAffectedRowsFunc(rowCount)
206215
case actionUpdate:
207-
update.addAffectedRowsFunc(rowCount)
216+
// For UPDATE operations (not REPLACE INTO), count DELETE rows
217+
// But for REPLACE INTO, this should not be called or should be ignored
218+
// REPLACE INTO uses actionUpdate but should only count INSERT
219+
// So we don't count DELETE here for actionUpdate
208220
}
209221
}
210222

0 commit comments

Comments
 (0)