opt: speed up execbuilder phase

pkg/sql/opt/bench/bench_test.go

@@ -1584,3 +1584,70 @@ func BenchmarkSlowQueries(b *testing.B) {
 		})
 	}
 }
+
+// BenchmarkExecBuild measures the time that the execbuilder phase takes. It
+// does not include any other phases.
+func BenchmarkExecBuild(b *testing.B) {
+	type testCase struct {
+		query  benchQuery
+		schema []string
+	}
+	var testCases []testCase
+
+	// Add the basic queries.
+	for _, query := range queriesToTest(b) {
+		testCases = append(testCases, testCase{query, schemas})
+	}
+
+	// Add the slow queries.
+	p := datapathutils.TestDataPath(b, "slow-schemas.sql")
+	slowSchemas, err := os.ReadFile(p)
+	if err != nil {
+		b.Fatalf("%v", err)
+	}
+	for _, query := range slowQueries {
+		testCases = append(testCases, testCase{query, []string{string(slowSchemas)}})
+	}
+
+	for _, tc := range testCases {
+		h := newHarness(b, tc.query, tc.schema)
+
+		stmt, err := parser.ParseOne(tc.query.query)
+		if err != nil {
+			b.Fatalf("%v", err)
+		}
+
+		h.optimizer.Init(context.Background(), &h.evalCtx, h.testCat)
+		bld := optbuilder.New(h.ctx, &h.semaCtx, &h.evalCtx, h.testCat, h.optimizer.Factory(), stmt.AST)
+		if err = bld.Build(); err != nil {
+			b.Fatalf("%v", err)
+		}
+
+		if _, err := h.optimizer.Optimize(); err != nil {
+			panic(err)
+		}
+
+		execMemo := h.optimizer.Memo()
+		root := execMemo.RootExpr()
+
+		b.Run(tc.query.name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				eb := execbuilder.New(
+					context.Background(),
+					explain.NewPlanGistFactory(exec.StubFactory{}),
+					&h.optimizer,
+					execMemo,
+					nil, /* catalog */
+					root,
+					&h.semaCtx,
+					&h.evalCtx,
+					true,  /* allowAutoCommit */
+					false, /* isANSIDML */
+				)
+				if _, err := eb.Build(); err != nil {
+					b.Fatalf("%v", err)
+				}
+			}
+		})
+	}
+}

pkg/sql/opt/exec/execbuilder/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "builder.go",
         "cascades.go",
+        "col_ord_map.go",
         "format.go",
         "mutation.go",
         "relational.go",
@@ -65,6 +66,7 @@ go_test(
     name = "execbuilder_test",
     size = "small",
     srcs = [
+        "col_ord_map_test.go",
         "main_test.go",
         "mutation_test.go",
     ],

pkg/sql/opt/exec/execbuilder/builder.go

@@ -61,6 +61,7 @@ type Builder struct {
 	disableTelemetry bool
 	semaCtx          *tree.SemaContext
 	evalCtx          *eval.Context
+	colOrdsAlloc     colOrdMapAllocator
 
 	// subqueries accumulates information about subqueries that are part of scalar
 	// expressions we built. Each entry is associated with a tree.Subquery
@@ -205,6 +206,7 @@ func New(
 		initialAllowAutoCommit: allowAutoCommit,
 		IsANSIDML:              isANSIDML,
 	}
+	b.colOrdsAlloc.Init(mem.Metadata().MaxColumn())
 	if evalCtx != nil {
 		sd := evalCtx.SessionData()
 		if sd.SaveTablesPrefix != "" {
@@ -230,7 +232,7 @@ func New(
 // Build constructs the execution node tree and returns its root node if no
 // error occurred.
 func (b *Builder) Build() (_ exec.Plan, err error) {
-	plan, err := b.build(b.e)
+	plan, _, err := b.build(b.e)
 	if err != nil {
 		return nil, err
 	}
@@ -257,7 +259,7 @@ func (b *Builder) wrapFunction(fnName string) (tree.ResolvableFunctionReference,
 	return tree.WrapFunction(fnName), nil
 }
 
-func (b *Builder) build(e opt.Expr) (_ execPlan, err error) {
+func (b *Builder) build(e opt.Expr) (_ execPlan, outputCols colOrdMap, err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			// This code allows us to propagate errors without adding lots of checks
@@ -274,7 +276,7 @@ func (b *Builder) build(e opt.Expr) (_ execPlan, err error) {
 
 	rel, ok := e.(memo.RelExpr)
 	if !ok {
-		return execPlan{}, errors.AssertionFailedf(
+		return execPlan{}, colOrdMap{}, errors.AssertionFailedf(
 			"building execution for non-relational operator %s", redact.Safe(e.Op()),
 		)
 	}
@@ -297,12 +299,12 @@ func (b *Builder) BuildScalar() (tree.TypedExpr, error) {
 	if !ok {
 		return nil, errors.AssertionFailedf("BuildScalar cannot be called for non-scalar operator %s", redact.Safe(b.e.Op()))
 	}
-	var ctx buildScalarCtx
 	md := b.mem.Metadata()
-	ctx.ivh = tree.MakeIndexedVarHelper(&mdVarContainer{md: md}, md.NumColumns())
+	cols := b.colOrdsAlloc.Alloc()
 	for i := 0; i < md.NumColumns(); i++ {
-		ctx.ivarMap.Set(i+1, i)
+		cols.Set(opt.ColumnID(i+1), i)
 	}
+	ctx := makeBuildScalarCtx(cols)
 	return b.buildScalar(&ctx, scalar)
 }
 
@@ -320,11 +322,11 @@ type builtWithExpr struct {
 	id opt.WithID
 	// outputCols maps the output ColumnIDs of the With expression to the ordinal
 	// positions they are output to. See execPlan.outputCols for more details.
-	outputCols opt.ColMap
+	outputCols colOrdMap
 	bufferNode exec.Node
 }
 
-func (b *Builder) addBuiltWithExpr(id opt.WithID, outputCols opt.ColMap, bufferNode exec.Node) {
+func (b *Builder) addBuiltWithExpr(id opt.WithID, outputCols colOrdMap, bufferNode exec.Node) {
 	b.withExprs = append(b.withExprs, builtWithExpr{
 		id:         id,
 		outputCols: outputCols,

pkg/sql/opt/exec/execbuilder/cascades.go

@@ -113,7 +113,7 @@ type cascadeBuilder struct {
 	mutationBuffer exec.Node
 	// mutationBufferCols maps With column IDs from the original memo to buffer
 	// node column ordinals; see builtWithExpr.outputCols.
-	mutationBufferCols opt.ColMap
+	mutationBufferCols colOrdMap
 
 	// colMeta remembers the metadata of the With columns from the original memo.
 	colMeta []opt.ColumnMeta
@@ -142,10 +142,9 @@ func makeCascadeBuilder(b *Builder, mutationWithID opt.WithID) (*cascadeBuilder,
 	// Remember the column metadata, as we will need to recreate it in the new
 	// memo.
 	md := b.mem.Metadata()
-	cb.colMeta = make([]opt.ColumnMeta, 0, cb.mutationBufferCols.Len())
-	cb.mutationBufferCols.ForEach(func(key, val int) {
-		id := opt.ColumnID(key)
-		cb.colMeta = append(cb.colMeta, *md.ColumnMeta(id))
+	cb.colMeta = make([]opt.ColumnMeta, 0, cb.mutationBufferCols.MaxOrd())
+	cb.mutationBufferCols.ForEach(func(col opt.ColumnID, ord int) {
+		cb.colMeta = append(cb.colMeta, *md.ColumnMeta(col))
 	})
 
 	return cb, nil
@@ -198,7 +197,7 @@ func (cb *cascadeBuilder) planCascade(
 	var relExpr memo.RelExpr
 	// bufferColMap is the mapping between the column IDs in the new memo and
 	// the column ordinal in the buffer node.
-	var bufferColMap opt.ColMap
+	var bufferColMap colOrdMap
 	if bufferRef == nil {
 		// No input buffering.
 		var err error
@@ -219,15 +218,19 @@ func (cb *cascadeBuilder) planCascade(
 	} else {
 		// Set up metadata for the buffer columns.
 
+		// Allocate a map with enough capacity to store the new columns being
+		// added below.
+		bufferColMap = newColOrdMap(md.MaxColumn() + opt.ColumnID(len(cb.colMeta)))
+
 		// withColRemap is the mapping between the With column IDs in the original
 		// memo and the corresponding column IDs in the new memo.
 		var withColRemap opt.ColMap
 		var withCols opt.ColSet
 		for i := range cb.colMeta {
 			id := md.AddColumn(cb.colMeta[i].Alias, cb.colMeta[i].Type)
 			withCols.Add(id)
-			ordinal, _ := cb.mutationBufferCols.Get(int(cb.colMeta[i].MetaID))
-			bufferColMap.Set(int(id), ordinal)
+			ordinal, _ := cb.mutationBufferCols.Get(cb.colMeta[i].MetaID)
+			bufferColMap.Set(id, ordinal)
 			withColRemap.Set(int(cb.colMeta[i].MetaID), int(id))
 		}