Merge #146655

146655: opt,sql: assorted optimizations r=mgartner a=mgartner

#### opt: add benchmarks for sysbench UPDATE statements

Release note: None

#### execbuilder: add global empty buildScalarCtx

Release note: None

#### sql: refactor MakeUpdater

Some conditionals in the `needsUpdate` have been moved out of the
closure. Also, `needsUpdate` no longer allocates a slice of column IDs.
Instead, it iterates directly over a set of column IDs.

Release note: None

#### opt: convert needsUpdate closure to function

Release note: None

#### sql: lazily allocate includeIndexes

Release note: None

#### sql: combine index entry allocations

Release note: None

#### sql: combine allocations of fetch and update col lists

Release note: None

#### sql: combine index slice allocations for updates and upserts

Release note: None

#### sql/row: lazily compute index column direction slices

The encoding directions of the primary and secondary indexes are only
used in mutations when pretty-printing keys when tracing is enabled and
when `CheckRowSize` emits a log event. These encoding directions are now
lazily computed to eliminate their overhead in the common case.

Release note: None


Co-authored-by: Marcus Gartner <[email protected]>

Author: craig[bot]

pkg/sql/opt/bench/bench_test.go

@@ -122,6 +122,15 @@ type benchQuery struct {
 var schemas = []string{
 	`CREATE TABLE kv (k BIGINT NOT NULL PRIMARY KEY, v BYTES NOT NULL)`,
 	`
+	CREATE TABLE sbtest (
+		id INT8 PRIMARY KEY,
+		k INT8 NOT NULL DEFAULT 0,
+		c CHAR(120) NOT NULL DEFAULT '',
+		pad CHAR(60) NOT NULL DEFAULT '',
+		INDEX (k)
+	)
+	`,
+	`
 	CREATE TABLE customer
 	(
 		c_id           integer        not null,
@@ -389,6 +398,18 @@ var queries = [...]benchQuery{
 		args:  []interface{}{},
 	},
 
+	{
+		name:  "sysbench-update-index",
+		query: `UPDATE sbtest SET k=k+1 WHERE id=$1`,
+		args:  []interface{}{10},
+	},
+
+	{
+		name:  "sysbench-update-non-index",
+		query: `UPDATE sbtest SET c=$2 WHERE id=$1`,
+		args:  []interface{}{10, "foo"},
+	},
+
 	// 1. Table with many columns.
 	// 2. Multi-column primary key.
 	// 3. Mutiple indexes to consider.

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

@@ -541,15 +541,14 @@ func (b *Builder) buildValuesRows(values *memo.ValuesExpr) ([][]tree.TypedExpr,
 	numCols := len(values.Cols)
 
 	rows := makeTypedExprMatrix(len(values.Rows), numCols)
-	scalarCtx := buildScalarCtx{}
 	for i := range rows {
 		tup := values.Rows[i].(*memo.TupleExpr)
 		if len(tup.Elems) != numCols {
 			return nil, fmt.Errorf("inconsistent row length %d vs %d", len(tup.Elems), numCols)
 		}
 		var err error
 		for j := 0; j < numCols; j++ {
-			rows[i][j], err = b.buildScalar(&scalarCtx, tup.Elems[j])
+			rows[i][j], err = b.buildScalar(&emptyBuildScalarCtx, tup.Elems[j])
 			if err != nil {
 				return nil, err
 			}
@@ -3565,11 +3564,10 @@ func (b *Builder) buildCall(c *memo.CallExpr) (_ execPlan, outputCols colOrdMap,
 
 	// Build the argument expressions.
 	var args tree.TypedExprs
-	ctx := buildScalarCtx{}
 	if len(udf.Args) > 0 {
 		args = make(tree.TypedExprs, len(udf.Args))
 		for i := range udf.Args {
-			args[i], err = b.buildScalar(&ctx, udf.Args[i])
+			args[i], err = b.buildScalar(&emptyBuildScalarCtx, udf.Args[i])
 			if err != nil {
 				return execPlan{}, colOrdMap{}, err
 			}
@@ -3993,8 +3991,7 @@ func (b *Builder) buildVectorSearch(
 		}
 	}
 	outColOrds, outColMap := b.getColumns(search.Cols, search.Table)
-	ctx := buildScalarCtx{}
-	queryVector, err := b.buildScalar(&ctx, search.QueryVector)
+	queryVector, err := b.buildScalar(&emptyBuildScalarCtx, search.QueryVector)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}

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

@@ -40,6 +40,8 @@ type buildScalarCtx struct {
 	ivarMap colOrdMap
 }
 
+var emptyBuildScalarCtx buildScalarCtx
+
 type buildFunc func(b *Builder, ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.TypedExpr, error)
 
 var scalarBuildFuncMap [opt.NumOperators]buildFunc

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

@@ -219,8 +219,7 @@ func (b *Builder) buildAlterTableSplit(
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
-	scalarCtx := buildScalarCtx{}
-	expiration, err := b.buildScalar(&scalarCtx, split.Expiration)
+	expiration, err := b.buildScalar(&emptyBuildScalarCtx, split.Expiration)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
@@ -295,12 +294,11 @@ func (b *Builder) buildAlterRangeRelocate(
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
-	scalarCtx := buildScalarCtx{}
-	toStoreID, err := b.buildScalar(&scalarCtx, relocate.ToStoreID)
+	toStoreID, err := b.buildScalar(&emptyBuildScalarCtx, relocate.ToStoreID)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
-	fromStoreID, err := b.buildScalar(&scalarCtx, relocate.FromStoreID)
+	fromStoreID, err := b.buildScalar(&emptyBuildScalarCtx, relocate.FromStoreID)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
@@ -325,8 +323,7 @@ func (b *Builder) buildControlJobs(
 		return execPlan{}, colOrdMap{}, err
 	}
 
-	scalarCtx := buildScalarCtx{}
-	reason, err := b.buildScalar(&scalarCtx, ctl.Reason)
+	reason, err := b.buildScalar(&emptyBuildScalarCtx, ctl.Reason)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
@@ -432,8 +429,7 @@ func (b *Builder) buildExport(
 		return execPlan{}, colOrdMap{}, err
 	}
 
-	scalarCtx := buildScalarCtx{}
-	fileName, err := b.buildScalar(&scalarCtx, export.FileName)
+	fileName, err := b.buildScalar(&emptyBuildScalarCtx, export.FileName)
 	if err != nil {
 		return execPlan{}, colOrdMap{}, err
 	}
@@ -442,7 +438,7 @@ func (b *Builder) buildExport(
 	for i, o := range export.Options {
 		opts[i].Key = o.Key
 		var err error
-		opts[i].Value, err = b.buildScalar(&scalarCtx, o.Value)
+		opts[i].Value, err = b.buildScalar(&emptyBuildScalarCtx, o.Value)
 		if err != nil {
 			return execPlan{}, colOrdMap{}, err
 		}

pkg/sql/opt_exec_factory.go

@@ -1401,6 +1401,25 @@ func ordinalsToIndexes(table cat.Table, ords cat.IndexOrdinals) []catalog.Index
 	return retval
 }
 
+func ordinalsToIndexes2(
+	table cat.Table, a, b cat.IndexOrdinals,
+) ([]catalog.Index, []catalog.Index) {
+	lenA, lenB := len(a), len(b)
+	if lenA+lenB == 0 {
+		return nil, nil
+	}
+
+	indexes := make([]catalog.Index, lenA+lenB)
+	indexesA, indexesB := indexes[:lenA:lenA], indexes[lenA:]
+	for i, idx := range a {
+		indexesA[i] = table.Index(idx).(*optIndex).idx
+	}
+	for i, idx := range b {
+		indexesB[i] = table.Index(idx).(*optIndex).idx
+	}
+	return indexesA, indexesB
+}
+
 func (ef *execFactory) ConstructInsert(
 	input exec.Node,
 	table cat.Table,
@@ -1596,6 +1615,8 @@ func (ef *execFactory) ConstructUpdate(
 	}
 
 	// If rows are not needed, no columns are returned.
+	// TODO(mgartner): Combine returnCols allocations with allocations for
+	// fetchCols and updateCols in constructUpdateRun.
 	var returnCols []catalog.Column
 	if rowsNeeded {
 		returnCols = makeColList(table, returnColOrdSet)
@@ -1719,15 +1740,15 @@ func (ef *execFactory) constructUpdateRun(
 	lockedIndexes cat.IndexOrdinals,
 ) error {
 	tabDesc := table.(*optTable).desc
-	fetchCols := makeColList(table, fetchColOrdSet)
-	updateCols := makeColList(table, updateColOrdSet)
+	fetchCols, updateCols := makeColList2(table, fetchColOrdSet, updateColOrdSet)
 
 	// Create the table updater.
+	tombstoneIdxs, lockIdxs := ordinalsToIndexes2(table, uniqueWithTombstoneIndexes, lockedIndexes)
 	ru, err := row.MakeUpdater(
 		ef.planner.ExecCfg().Codec,
 		tabDesc,
-		ordinalsToIndexes(table, uniqueWithTombstoneIndexes),
-		ordinalsToIndexes(table, lockedIndexes),
+		tombstoneIdxs,
+		lockIdxs,
 		updateCols,
 		fetchCols,
 		row.UpdaterDefault,
@@ -1798,11 +1819,12 @@ func (ef *execFactory) ConstructUpsert(
 	}
 
 	// Create the table updater, which does the bulk of the update-related work.
+	tombstoneIdxs, lockIdxs := ordinalsToIndexes2(table, uniqueWithTombstoneIndexes, lockedIndexes)
 	ru, err := row.MakeUpdater(
 		ef.planner.ExecCfg().Codec,
 		tabDesc,
-		ordinalsToIndexes(table, uniqueWithTombstoneIndexes),
-		ordinalsToIndexes(table, lockedIndexes),
+		tombstoneIdxs,
+		lockIdxs,
 		updateCols,
 		fetchCols,
 		row.UpdaterDefault,
@@ -2569,6 +2591,27 @@ func makeColList(table cat.Table, cols exec.TableColumnOrdinalSet) []catalog.Col
 	return ret
 }
 
+// makeColList2 is similar to makeColList, but it takes two sets of ordinals and
+// allocates a single slice which is split into two.
+func makeColList2(
+	table cat.Table, a, b exec.TableColumnOrdinalSet,
+) ([]catalog.Column, []catalog.Column) {
+	tab := table.(optCatalogTableInterface)
+	lenA, lenB := a.Len(), b.Len()
+	cols := make([]catalog.Column, 0, lenA+lenB)
+	listA, listB := cols[:0:lenA], cols[lenA:lenA]
+	for i, n := 0, table.ColumnCount(); i < n; i++ {
+		col := tab.getCol(i)
+		if a.Contains(i) {
+			listA = append(listA, col)
+		}
+		if b.Contains(i) {
+			listB = append(listB, col)
+		}
+	}
+	return listA, listB
+}
+
 // makePublicToReturnColumnIndexMapping returns a map from the ordinals
 // of the table's public columns to ordinals in the returnColDescs slice.
 //

pkg/sql/row/deleter.go

@@ -179,10 +179,10 @@ func (rd *Deleter) DeleteRow(
 			if oth.IsSet() {
 				oth.DelWithCPut(ctx, b, &rd.key, expValue, traceKV)
 			} else {
-				delWithCPutFn(ctx, b, &rd.key, expValue, traceKV, rd.Helper.primIndexValDirs)
+				delWithCPutFn(ctx, b, &rd.key, expValue, traceKV, &rd.Helper, primaryIndexDirs)
 			}
 		} else {
-			delFn(ctx, b, &rd.key, !rd.primaryLocked /* needsLock */, traceKV, rd.Helper.primIndexValDirs)
+			delFn(ctx, b, &rd.key, !rd.primaryLocked /* needsLock */, traceKV, &rd.Helper, primaryIndexDirs)
 		}
 
 		rd.key = nil
@@ -218,7 +218,7 @@ func (rd *Deleter) DeleteRow(
 		for _, e := range entries {
 			if err = rd.Helper.deleteIndexEntry(
 				ctx, b, index, &e.Key, alreadyLocked, rd.Helper.sd.BufferedWritesUseLockingOnNonUniqueIndexes,
-				traceKV, rd.Helper.secIndexValDirs[i],
+				traceKV, secondaryIndexDirs(i),
 			); err != nil {
 				return err
 			}

pkg/sql/row/helper.go

@@ -91,9 +91,11 @@ type RowHelper struct {
 	UniqueWithTombstoneIndexes intsets.Fast
 	indexEntries               map[catalog.Index][]rowenc.IndexEntry
 
-	// Computed during initialization for pretty-printing.
-	primIndexValDirs []encoding.Direction
-	secIndexValDirs  [][]encoding.Direction
+	// Lazily computed for pretty-printing and CheckRowSize.
+	dirs struct {
+		primary   []encoding.Direction
+		secondary [][]encoding.Direction
+	}
 
 	// Computed and cached.
 	PrimaryIndexKeyPrefix []byte
@@ -126,28 +128,50 @@ func NewRowHelper(
 	for _, index := range uniqueWithTombstoneIndexes {
 		uniqueWithTombstoneIndexesSet.Add(index.Ordinal())
 	}
-	rh := RowHelper{
+	return RowHelper{
 		Codec:                      codec,
 		TableDesc:                  desc,
 		Indexes:                    indexes,
 		UniqueWithTombstoneIndexes: uniqueWithTombstoneIndexesSet,
 		sd:                         sd,
 		metrics:                    metrics,
+		maxRowSizeLog:              uint32(maxRowSizeLog.Get(sv)),
+		maxRowSizeErr:              uint32(maxRowSizeErr.Get(sv)),
 	}
+}
 
-	// Pre-compute the encoding directions of the index key values for
-	// pretty-printing in traces.
-	rh.primIndexValDirs = catalogkeys.IndexKeyValDirs(rh.TableDesc.GetPrimaryIndex())
+// lazyIndexDirs represents encoding directions of an index. Those directions
+// may not have been, and may never be computed. The value of -2 represents
+// empty encoding directions. The value of -1 represents the encoding directions
+// of the primary index, otherwise a value i represents the encoding directions
+// of the i-th secondary index.
+type lazyIndexDirs int
 
-	rh.secIndexValDirs = make([][]encoding.Direction, len(rh.Indexes))
-	for i := range rh.Indexes {
-		rh.secIndexValDirs[i] = catalogkeys.IndexKeyValDirs(rh.Indexes[i])
-	}
+const (
+	emptyIndexDirs   lazyIndexDirs = -2
+	primaryIndexDirs lazyIndexDirs = -1
+)
 
-	rh.maxRowSizeLog = uint32(maxRowSizeLog.Get(sv))
-	rh.maxRowSizeErr = uint32(maxRowSizeErr.Get(sv))
+func secondaryIndexDirs(i int) lazyIndexDirs { return lazyIndexDirs(i) }
 
-	return rh
+func (d lazyIndexDirs) compute(rh *RowHelper) []encoding.Direction {
+	switch d {
+	case emptyIndexDirs:
+		return nil
+	case primaryIndexDirs:
+		if rh.dirs.primary == nil {
+			rh.dirs.primary = catalogkeys.IndexKeyValDirs(rh.TableDesc.GetPrimaryIndex())
+		}
+		return rh.dirs.primary
+	default:
+		if rh.dirs.secondary == nil {
+			rh.dirs.secondary = make([][]encoding.Direction, len(rh.Indexes))
+			for i := range rh.Indexes {
+				rh.dirs.secondary[i] = catalogkeys.IndexKeyValDirs(rh.Indexes[i])
+			}
+		}
+		return rh.dirs.secondary[d]
+	}
 }
 
 // encodeIndexes encodes the primary and secondary index keys. The
@@ -432,7 +456,7 @@ func (rh *RowHelper) CheckRowSize(
 		RowSize:    size,
 		TableID:    uint32(rh.TableDesc.GetID()),
 		FamilyID:   uint32(family),
-		PrimaryKey: keys.PrettyPrint(rh.primIndexValDirs, *key),
+		PrimaryKey: keys.PrettyPrint(primaryIndexDirs.compute(rh), *key),
 	}
 	if rh.sd.Internal && shouldErr {
 		// Internal work should never err and always log if violating either limit.
@@ -471,11 +495,12 @@ func delFn(
 	key *roachpb.Key,
 	needsLock bool,
 	traceKV bool,
-	keyEncodingDirs []encoding.Direction,
+	rh *RowHelper,
+	dirs lazyIndexDirs,
 ) {
 	if needsLock {
 		if traceKV {
-			if keyEncodingDirs != nil {
+			if keyEncodingDirs := dirs.compute(rh); keyEncodingDirs != nil {
 				log.VEventf(ctx, 2, "Del (locking) %s", keys.PrettyPrint(keyEncodingDirs, *key))
 			} else {
 				log.VEventf(ctx, 2, "Del (locking) %s", *key)
@@ -484,7 +509,7 @@ func delFn(
 		b.DelMustAcquireExclusiveLock(key)
 	} else {
 		if traceKV {
-			if keyEncodingDirs != nil {
+			if keyEncodingDirs := dirs.compute(rh); keyEncodingDirs != nil {
 				log.VEventf(ctx, 2, "Del %s", keys.PrettyPrint(keyEncodingDirs, *key))
 			} else {
 				log.VEventf(ctx, 2, "Del %s", *key)
@@ -500,10 +525,11 @@ func delWithCPutFn(
 	key *roachpb.Key,
 	expVal []byte,
 	traceKV bool,
-	keyEncodingDirs []encoding.Direction,
+	rh *RowHelper,
+	dirs lazyIndexDirs,
 ) {
 	if traceKV {
-		if keyEncodingDirs != nil {
+		if keyEncodingDirs := dirs.compute(rh); keyEncodingDirs != nil {
 			log.VEventf(ctx, 2, "CPut %s -> nil (delete)", keys.PrettyPrint(keyEncodingDirs, *key))
 		} else {
 			log.VEventf(ctx, 2, "CPut %s -> nil (delete)", *key)
@@ -520,7 +546,7 @@ func (rh *RowHelper) deleteIndexEntry(
 	alreadyLocked bool,
 	lockNonUnique bool,
 	traceKV bool,
-	valDirs []encoding.Direction,
+	dirs lazyIndexDirs,
 ) error {
 	needsLock := !alreadyLocked && (index.IsUnique() || lockNonUnique)
 	if index.UseDeletePreservingEncoding() {
@@ -537,7 +563,7 @@ func (rh *RowHelper) deleteIndexEntry(
 			b.Put(key, deleteEncoding)
 		}
 	} else {
-		delFn(ctx, b, key, needsLock, traceKV, valDirs)
+		delFn(ctx, b, key, needsLock, traceKV, rh, dirs)
 	}
 	return nil
 }