Merge pull request #109470 from cockroachdb/blathers/backport-release…

…-23.1-109394

release-23.1: typedesc: avoid slice copy when hydrating enums

pkg/sql/BUILD.bazel

@@ -667,6 +667,7 @@ go_test(
         "scatter_test.go",
         "schema_changer_helpers_test.go",
         "schema_changer_test.go",
+        "schema_resolver_test.go",
         "scrub_test.go",
         "sequence_test.go",
         "session_migration_test.go",

pkg/sql/catalog/typedesc/hydrate.go

@@ -137,16 +137,26 @@ func ensureTypeMetadataIsHydrated(
 		return
 	}
 	if e := maybeDesc.AsEnumTypeDescriptor(); e != nil {
-		n := e.NumEnumMembers()
-		tm.EnumData = &types.EnumMetadata{
-			LogicalRepresentations:  make([]string, n),
-			PhysicalRepresentations: make([][]byte, n),
-			IsMemberReadOnly:        make([]bool, n),
-		}
-		for i := 0; i < n; i++ {
-			tm.EnumData.LogicalRepresentations[i] = e.GetMemberLogicalRepresentation(i)
-			tm.EnumData.PhysicalRepresentations[i] = e.GetMemberPhysicalRepresentation(i)
-			tm.EnumData.IsMemberReadOnly[i] = e.IsMemberReadOnly(i)
+		if imm, ok := e.(*immutable); ok {
+			// Fast-path for immutable enum descriptors. We can use a pointer into the
+			// immutable descriptor's slices, since the TypMetadata is immutable too.
+			tm.EnumData = &types.EnumMetadata{
+				LogicalRepresentations:  imm.logicalReps,
+				PhysicalRepresentations: imm.physicalReps,
+				IsMemberReadOnly:        imm.readOnlyMembers,
+			}
+		} else {
+			n := e.NumEnumMembers()
+			tm.EnumData = &types.EnumMetadata{
+				LogicalRepresentations:  make([]string, n),
+				PhysicalRepresentations: make([][]byte, n),
+				IsMemberReadOnly:        make([]bool, n),
+			}
+			for i := 0; i < n; i++ {
+				tm.EnumData.LogicalRepresentations[i] = e.GetMemberLogicalRepresentation(i)
+				tm.EnumData.PhysicalRepresentations[i] = e.GetMemberPhysicalRepresentation(i)
+				tm.EnumData.IsMemberReadOnly[i] = e.IsMemberReadOnly(i)
+			}
 		}
 	}
 }

pkg/sql/schema_resolver_test.go

@@ -0,0 +1,69 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package sql
+
+import (
+	"context"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/security/username"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/lib/pq/oid"
+	"github.com/stretchr/testify/require"
+)
+
+// This benchmark tests the performance of resolving an enum with many (10,000)
+// values. This is a regression test for #109228.
+func BenchmarkResolveTypeByOID(b *testing.B) {
+	defer leaktest.AfterTest(b)()
+	defer log.Scope(b).Close(b)
+
+	s, sqlDB, kvDB := serverutils.StartServer(b, base.TestServerArgs{})
+	defer s.Stopper().Stop(context.Background())
+
+	query := strings.Builder{}
+	query.WriteString("CREATE TYPE typ AS ENUM ('v0'")
+	for i := 1; i < 10_000; i++ {
+		query.WriteString(", 'v")
+		query.WriteString(strconv.Itoa(i))
+		query.WriteString("'")
+	}
+	query.WriteString(")")
+	_, err := sqlDB.Exec(query.String())
+	require.NoError(b, err)
+
+	var typOID uint32
+	err = sqlDB.QueryRow("SELECT 'typ'::regtype::oid").Scan(&typOID)
+	require.NoError(b, err)
+
+	ctx := context.Background()
+	execCfg := s.ExecutorConfig().(ExecutorConfig)
+	sd := NewFakeSessionData(&execCfg.Settings.SV, "test")
+	sd.Database = "defaultdb"
+	planner, cleanup := newInternalPlanner("test", kv.NewTxn(ctx, kvDB, s.NodeID()),
+		username.RootUserName(), &MemoryMetrics{}, &execCfg, sd.SessionData,
+	)
+	defer cleanup()
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		typ, err := planner.schemaResolver.ResolveTypeByOID(ctx, oid.Oid(typOID))
+		require.NoError(b, err)
+		require.Equal(b, "typ", typ.Name())
+	}
+	b.StopTimer()
+}