sql: combine allocations of fetch and update col lists

Release note: None

Author: mgartner

pkg/sql/opt_exec_factory.go

@@ -1596,6 +1596,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,8 +1721,7 @@ 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.
 	ru, err := row.MakeUpdater(
@@ -2569,6 +2570,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.
 //