From 6dab55faea0afa9f59e4a5f6214732ef61da87ff Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Fri, 23 Sep 2022 12:33:47 -0700
Subject: [PATCH] bench: add a benchmark for update with assigment casts

Release note: None
---
 pkg/bench/bench_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/pkg/bench/bench_test.go b/pkg/bench/bench_test.go
index 6b3b92fcfbce..05c4d573d415 100644
--- a/pkg/bench/bench_test.go
+++ b/pkg/bench/bench_test.go
@@ -388,6 +388,7 @@ func BenchmarkSQL(b *testing.B) {
 			runBenchmarkInsertSecondaryIndex,
 			runBenchmarkTrackChoices,
 			runBenchmarkUpdate,
+			runBenchmarkUpdateWithAssignmentCast,
 			runBenchmarkUpsert,
 		} {
 			fnName := runtime.FuncForPC(reflect.ValueOf(runFn).Pointer()).Name()
@@ -572,6 +573,44 @@ func runBenchmarkUpdate(b *testing.B, db *sqlutils.SQLRunner, count int) {
 	b.StopTimer()
 }
 
+// runBenchmarkUpdateWithAssignmentCast benchmarks updating count random rows in
+// a table where we need to perform an assigment cast to get the updated values.
+func runBenchmarkUpdateWithAssignmentCast(b *testing.B, db *sqlutils.SQLRunner, count int) {
+	defer func() {
+		db.Exec(b, `DROP TABLE IF EXISTS bench.update`)
+	}()
+
+	const rows = 10000
+	db.Exec(b, `CREATE TABLE bench.update (k INT PRIMARY KEY, v INT)`)
+
+	var buf bytes.Buffer
+	buf.WriteString(`INSERT INTO bench.update VALUES `)
+	for i := 0; i < rows; i++ {
+		if i > 0 {
+			buf.WriteString(", ")
+		}
+		fmt.Fprintf(&buf, "(%d, %d)", i, i)
+	}
+	db.Exec(b, buf.String())
+
+	s := rand.New(rand.NewSource(5432))
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		buf.Reset()
+		buf.WriteString(`UPDATE bench.update SET v = (v + 1.0)::FLOAT WHERE k IN (`)
+		for j := 0; j < count; j++ {
+			if j > 0 {
+				buf.WriteString(", ")
+			}
+			fmt.Fprintf(&buf, `%d`, s.Intn(rows))
+		}
+		buf.WriteString(`)`)
+		db.Exec(b, buf.String())
+	}
+	b.StopTimer()
+}
+
 // runBenchmarkUpsert benchmarks upserting count rows in a table.
 func runBenchmarkUpsert(b *testing.B, db *sqlutils.SQLRunner, count int) {
 	defer func() {