perf(spanner): improve mutationProto allocations and performance (#12740)

rahul2393 · web-flow · commit 2a4add540520 · 2025-08-22T12:20:51.000+05:30
* perf(spanner): improve mutationProto allocations and performance

* fix test
diff --git a/spanner/mutation.go b/spanner/mutation.go
@@ -20,7 +20,6 @@ import (
 	"fmt"
 	"math/rand"
 	"reflect"
-	"time"
 
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
 	"google.golang.org/grpc/codes"
@@ -470,41 +469,39 @@ func (m Mutation) proto() (*sppb.Mutation, error) {
 // mutationsProto turns a spanner.Mutation array into a sppb.Mutation array,
 // it is convenient for sending batch mutations to Cloud Spanner.
 func mutationsProto(ms []*Mutation) ([]*sppb.Mutation, *sppb.Mutation, error) {
-	var selectedMutation *Mutation
-	var nonInsertMutations []*Mutation
-
-	l := make([]*sppb.Mutation, 0, len(ms))
-	for _, m := range ms {
-		if m.op != opInsert {
-			nonInsertMutations = append(nonInsertMutations, m)
-		}
-		if selectedMutation == nil {
-			selectedMutation = m
-		}
-		// Track the INSERT mutation with the highest number of values if only INSERT mutation were found
-		if selectedMutation.op == opInsert && m.op == opInsert && len(m.values) > len(selectedMutation.values) {
-			selectedMutation = m
-		}
-
-		// Convert the mutation to sppb.Mutation and add to the list
+	n := len(ms)
+	out := make([]*sppb.Mutation, 0, n)
+	if n == 0 {
+		return out, nil, nil
+	}
+	maxInsertIdx := -1
+	maxInsertVals := -1
+	nonInsertCount := 0
+	selectedNonInsertIdx := -1
+	for i, m := range ms {
 		pb, err := m.proto()
 		if err != nil {
 			return nil, nil, err
 		}
-		l = append(l, pb)
+		out = append(out, pb)
+		if m.op == opInsert {
+			if v := len(m.values); v >= maxInsertVals {
+				maxInsertVals, maxInsertIdx = v, i
+			}
+			continue
+		}
+		nonInsertCount++
+		if rand.Intn(nonInsertCount) == 0 {
+			selectedNonInsertIdx = i
+		}
 	}
-	if len(nonInsertMutations) > 0 {
-		selectedMutation = nonInsertMutations[rand.New(rand.NewSource(time.Now().UnixNano())).Intn(len(nonInsertMutations))]
+	if nonInsertCount > 0 {
+		return out, out[selectedNonInsertIdx], nil
 	}
-	if selectedMutation != nil {
-		m, err := selectedMutation.proto()
-		if err != nil {
-			return nil, nil, err
-		}
-		return l, m, nil
+	if maxInsertIdx >= 0 {
+		return out, out[maxInsertIdx], nil
 	}
-
-	return l, nil, nil
+	return out, nil, nil
 }
 
 // mutationGroupsProto turns a spanner.MutationGroup array into a
diff --git a/spanner/mutation_test.go b/spanner/mutation_test.go
@@ -885,3 +885,53 @@ func TestEncodeMutationGroupArray(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkMutationsProto(b *testing.B) {
+	type benchmarkCase struct {
+		name      string
+		mutations []*Mutation
+	}
+	benchmarkCases := []benchmarkCase{
+		{
+			name: "small number of mutations",
+			mutations: []*Mutation{
+				Insert("t_foo", []string{"col1", "col2"}, []interface{}{int64(1), int64(2)}),
+				Update("t_foo", []string{"col1", "col2"}, []interface{}{"one", []byte(nil)}),
+				InsertOrUpdate("t_foo", []string{"col1", "col2"}, []interface{}{1.0, 2.0}),
+				Replace("t_foo", []string{"col1", "col2"}, []interface{}{"one", 2.0}),
+				Delete("t_foo", Key{"foo"}),
+			},
+		},
+		{
+			name: "large number of mutations",
+			mutations: func() []*Mutation {
+				var mutations []*Mutation
+				for i := 0; i < 20; i++ {
+					mutations = append(mutations, Insert("t_foo", []string{"col1", "col2"}, []interface{}{int64(i), int64(i + 1)}))
+					mutations = append(mutations, Update("t_foo", []string{"col1", "col2"}, []interface{}{"one", []byte(nil)}))
+					mutations = append(mutations, InsertOrUpdate("t_foo", []string{"col1", "col2"}, []interface{}{1.0, 2.0}))
+					mutations = append(mutations, Replace("t_foo", []string{"col1", "col2"}, []interface{}{"one", 2.0}))
+					mutations = append(mutations, Delete("t_foo", Key{i}))
+				}
+				return mutations
+			}(),
+		},
+		{
+			name: "mixed type of mutations",
+			mutations: []*Mutation{
+				Insert("t_foo", []string{"col1", "col2"}, []interface{}{int64(1), int64(2)}),
+				Update("t_foo", []string{"col1", "col2"}, []interface{}{"one", []byte(nil)}),
+				Delete("t_foo", Key{"foo"}),
+				Insert("t_bar", []string{"col1"}, []interface{}{"bar"}),
+			},
+		},
+	}
+
+	for _, bc := range benchmarkCases {
+		b.Run(bc.name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				_, _, _ = mutationsProto(bc.mutations)
+			}
+		})
+	}
+}
