opt: add OptionalColList type

For mutations, we use ColLists that map 1-to-1 to table columns, and
where some of the entries in the list can be 0. This is really meant
to represent a mapping of columns and it is an abuse of the ColList
type (which is supposed to be a simple list of columns).

We add a separate OptionalColList type which has the desired
semantics. This helps clarify things a bit (in particular it is now
obvious which of the lists are "real" lists and which are "mappings").
It will also allow a different ToSet() method which makes sense for
this structure (i.e. doesn't put 0 in the set).

Release note: None

pkg/sql/opt/column_meta.go

@@ -26,33 +26,12 @@ func (c ColumnID) index() int {
 	return int(c - 1)
 }
 
-// ColList is a list of column ids.
+// ColList is a list of column IDs.
 //
 // TODO(radu): perhaps implement a FastIntList with the same "small"
 // representation as FastIntMap but with a slice for large cases.
 type ColList []ColumnID
 
-// ColumnMeta stores information about one of the columns stored in the
-// metadata.
-type ColumnMeta struct {
-	// MetaID is the identifier for this column that is unique within the query
-	// metadata.
-	MetaID ColumnID
-
-	// Alias is the best-effort name of this column. Since the same column in a
-	// query can have multiple names (using aliasing), one of those is chosen to
-	// be used for pretty-printing and debugging. This might be different than
-	// what is stored in the physical properties and is presented to end users.
-	Alias string
-
-	// Type is the scalar SQL type of this column.
-	Type *types.T
-
-	// Table is the base table to which this column belongs.
-	// If the column was synthesized (i.e. no base table), then it is 0.
-	Table TableID
-}
-
 // ToSet converts a column id list to a column id set.
 func (cl ColList) ToSet() ColSet {
 	var r ColSet
@@ -87,6 +66,70 @@ func (cl ColList) Equals(other ColList) bool {
 	return true
 }
 
+// OptionalColList is a list of column IDs where some of the IDs can be unset.
+// It is used when the columns map 1-to-1 to a known list of objects (e.g. table
+// columns).
+type OptionalColList []ColumnID
+
+// IsEmpty returns true if all columns in the list are unset.
+func (ocl OptionalColList) IsEmpty() bool {
+	for i := range ocl {
+		if ocl[i] != 0 {
+			return false
+		}
+	}
+	return true
+}
+
+// Equals returns true if this column list is identical to another list.
+func (ocl OptionalColList) Equals(other OptionalColList) bool {
+	if len(ocl) != len(other) {
+		return false
+	}
+	for i := range ocl {
+		if ocl[i] != other[i] {
+			return false
+		}
+	}
+	return true
+}
+
+// Find searches for a column in the list and returns its index in the list (if
+// successful).
+func (ocl OptionalColList) Find(col ColumnID) (idx int, ok bool) {
+	if col == 0 {
+		// An OptionalColList cannot contain zero column IDs.
+		return -1, false
+	}
+	for i := range ocl {
+		if ocl[i] == col {
+			return i, true
+		}
+	}
+	return -1, false
+}
+
+// ColumnMeta stores information about one of the columns stored in the
+// metadata.
+type ColumnMeta struct {
+	// MetaID is the identifier for this column that is unique within the query
+	// metadata.
+	MetaID ColumnID
+
+	// Alias is the best-effort name of this column. Since the same column in a
+	// query can have multiple names (using aliasing), one of those is chosen to
+	// be used for pretty-printing and debugging. This might be different than
+	// what is stored in the physical properties and is presented to end users.
+	Alias string
+
+	// Type is the scalar SQL type of this column.
+	Type *types.T
+
+	// Table is the base table to which this column belongs.
+	// If the column was synthesized (i.e. no base table), then it is 0.
+	Table TableID
+}
+
 // ColMap provides a 1:1 mapping from one column id to another. It is used by
 // operators that need to match columns from its inputs.
 type ColMap = util.FastIntMap

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

@@ -185,7 +185,7 @@ func (b *Builder) tryBuildFastPathInsert(ins *memo.InsertExpr) (_ execPlan, ok b
 
 		out := &fkChecks[i]
 		out.InsertCols = make([]exec.TableColumnOrdinal, len(lookupJoin.KeyCols))
-		findCol := func(cols opt.ColList, col opt.ColumnID) int {
+		findCol := func(cols opt.OptionalColList, col opt.ColumnID) int {
 			res, ok := cols.Find(col)
 			if !ok {
 				panic(errors.AssertionFailedf("cannot find column %d", col))
@@ -195,7 +195,7 @@ func (b *Builder) tryBuildFastPathInsert(ins *memo.InsertExpr) (_ execPlan, ok b
 		for i, keyCol := range lookupJoin.KeyCols {
 			// The keyCol comes from the WithScan operator. We must find the matching
 			// column in the mutation input.
-			withColOrd := findCol(withScan.OutCols, keyCol)
+			withColOrd := findCol(opt.OptionalColList(withScan.OutCols), keyCol)
 			inputCol := withScan.InCols[withColOrd]
 			out.InsertCols[i] = exec.TableColumnOrdinal(findCol(ins.InsertCols, inputCol))
 		}
@@ -294,7 +294,7 @@ func (b *Builder) buildUpdate(upd *memo.UpdateExpr) (execPlan, error) {
 	// to passthrough those columns so the projection above can use
 	// them.
 	if upd.NeedResults() {
-		colList = appendColsWhenPresent(colList, upd.PassthroughCols)
+		colList = append(colList, upd.PassthroughCols...)
 	}
 	colList = appendColsWhenPresent(colList, upd.CheckCols)
 	colList = appendColsWhenPresent(colList, upd.PartialIndexPutCols)
@@ -712,7 +712,7 @@ func (b *Builder) buildDeleteRange(
 
 // appendColsWhenPresent appends non-zero column IDs from the src list into the
 // dst list, and returns the possibly grown list.
-func appendColsWhenPresent(dst, src opt.ColList) opt.ColList {
+func appendColsWhenPresent(dst opt.ColList, src opt.OptionalColList) opt.ColList {
 	for _, col := range src {
 		if col != 0 {
 			dst = append(dst, col)
@@ -725,11 +725,8 @@ func appendColsWhenPresent(dst, src opt.ColList) opt.ColList {
 // column ID in the given list. This is used with mutation operators, which
 // maintain lists that correspond to the target table, with zero column IDs
 // indicating columns that are not involved in the mutation.
-func ordinalSetFromColList(colList opt.ColList) util.FastIntSet {
+func ordinalSetFromColList(colList opt.OptionalColList) util.FastIntSet {
 	var res util.FastIntSet
-	if colList == nil {
-		return res
-	}
 	for i, col := range colList {
 		if col != 0 {
 			res.Add(i)

pkg/sql/opt/memo/check_expr.go

@@ -285,7 +285,7 @@ func (m *Memo) CheckExpr(e opt.Expr) {
 	}
 }
 
-func (m *Memo) checkColListLen(colList opt.ColList, expectedLen int, listName string) {
+func (m *Memo) checkColListLen(colList opt.OptionalColList, expectedLen int, listName string) {
 	if len(colList) != expectedLen {
 		panic(errors.AssertionFailedf("column list %s expected length = %d, actual length = %d",
 			listName, log.Safe(expectedLen), len(colList)))

pkg/sql/opt/memo/expr_format.go

@@ -510,8 +510,8 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 			}
 			f.formatArbiters(tp, t.Arbiters, t.Table)
 			f.formatMutationCols(e, tp, "insert-mapping:", t.InsertCols, t.Table)
-			f.formatColList(e, tp, "check columns:", t.CheckCols)
-			f.formatColList(e, tp, "partial index put columns:", t.PartialIndexPutCols)
+			f.formatOptionalColList(e, tp, "check columns:", t.CheckCols)
+			f.formatOptionalColList(e, tp, "partial index put columns:", t.PartialIndexPutCols)
 			f.formatMutationCommon(tp, &t.MutationPrivate)
 		}
 
@@ -520,11 +520,11 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 			if len(colList) == 0 {
 				tp.Child("columns: <none>")
 			}
-			f.formatColList(e, tp, "fetch columns:", t.FetchCols)
+			f.formatOptionalColList(e, tp, "fetch columns:", t.FetchCols)
 			f.formatMutationCols(e, tp, "update-mapping:", t.UpdateCols, t.Table)
-			f.formatColList(e, tp, "check columns:", t.CheckCols)
-			f.formatColList(e, tp, "partial index put columns:", t.PartialIndexPutCols)
-			f.formatColList(e, tp, "partial index del columns:", t.PartialIndexDelCols)
+			f.formatOptionalColList(e, tp, "check columns:", t.CheckCols)
+			f.formatOptionalColList(e, tp, "partial index put columns:", t.PartialIndexPutCols)
+			f.formatOptionalColList(e, tp, "partial index del columns:", t.PartialIndexDelCols)
 			f.formatMutationCommon(tp, &t.MutationPrivate)
 		}
 
@@ -536,16 +536,16 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 			if t.CanaryCol != 0 {
 				f.formatArbiters(tp, t.Arbiters, t.Table)
 				f.formatColList(e, tp, "canary column:", opt.ColList{t.CanaryCol})
-				f.formatColList(e, tp, "fetch columns:", t.FetchCols)
+				f.formatOptionalColList(e, tp, "fetch columns:", t.FetchCols)
 				f.formatMutationCols(e, tp, "insert-mapping:", t.InsertCols, t.Table)
 				f.formatMutationCols(e, tp, "update-mapping:", t.UpdateCols, t.Table)
 				f.formatMutationCols(e, tp, "return-mapping:", t.ReturnCols, t.Table)
 			} else {
 				f.formatMutationCols(e, tp, "upsert-mapping:", t.InsertCols, t.Table)
 			}
-			f.formatColList(e, tp, "check columns:", t.CheckCols)
-			f.formatColList(e, tp, "partial index put columns:", t.PartialIndexPutCols)
-			f.formatColList(e, tp, "partial index del columns:", t.PartialIndexDelCols)
+			f.formatOptionalColList(e, tp, "check columns:", t.CheckCols)
+			f.formatOptionalColList(e, tp, "partial index put columns:", t.PartialIndexPutCols)
+			f.formatOptionalColList(e, tp, "partial index del columns:", t.PartialIndexDelCols)
 			f.formatMutationCommon(tp, &t.MutationPrivate)
 		}
 
@@ -554,8 +554,8 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 			if len(colList) == 0 {
 				tp.Child("columns: <none>")
 			}
-			f.formatColList(e, tp, "fetch columns:", t.FetchCols)
-			f.formatColList(e, tp, "partial index del columns:", t.PartialIndexDelCols)
+			f.formatOptionalColList(e, tp, "fetch columns:", t.FetchCols)
+			f.formatOptionalColList(e, tp, "partial index del columns:", t.PartialIndexDelCols)
 			f.formatMutationCommon(tp, &t.MutationPrivate)
 		}
 
@@ -1160,6 +1160,23 @@ func (f *ExprFmtCtx) formatColList(
 	nd RelExpr, tp treeprinter.Node, heading string, colList opt.ColList,
 ) {
 	if len(colList) > 0 {
+		notNullCols := nd.Relational().NotNullCols
+		f.Buffer.Reset()
+		f.Buffer.WriteString(heading)
+		for _, col := range colList {
+			f.space()
+			f.formatCol("" /* label */, col, notNullCols)
+		}
+		tp.Child(f.Buffer.String())
+	}
+}
+
+// formatOptionalColList constructs a new treeprinter child containing the
+// specified list of optional columns formatted using the formatCol method.
+func (f *ExprFmtCtx) formatOptionalColList(
+	nd RelExpr, tp treeprinter.Node, heading string, colList opt.OptionalColList,
+) {
+	if !colList.IsEmpty() {
 		notNullCols := nd.Relational().NotNullCols
 		f.Buffer.Reset()
 		f.Buffer.WriteString(heading)
@@ -1180,9 +1197,9 @@ func (f *ExprFmtCtx) formatColList(
 //   a:1 => x:4
 //
 func (f *ExprFmtCtx) formatMutationCols(
-	nd RelExpr, tp treeprinter.Node, heading string, colList opt.ColList, tabID opt.TableID,
+	nd RelExpr, tp treeprinter.Node, heading string, colList opt.OptionalColList, tabID opt.TableID,
 ) {
-	if len(colList) == 0 {
+	if colList.IsEmpty() {
 		return
 	}
 

pkg/sql/opt/memo/interner.go

@@ -437,6 +437,15 @@ func (h *hasher) HashColList(val opt.ColList) {
 	h.hash = hash
 }
 
+func (h *hasher) HashOptionalColList(val opt.OptionalColList) {
+	hash := h.hash
+	for _, id := range val {
+		hash ^= internHash(id)
+		hash *= prime64
+	}
+	h.hash = hash
+}
+
 func (h *hasher) HashOrdering(val opt.Ordering) {
 	hash := h.hash
 	for _, id := range val {
@@ -836,6 +845,10 @@ func (h *hasher) IsColListEqual(l, r opt.ColList) bool {
 	return l.Equals(r)
 }
 
+func (h *hasher) IsOptionalColListEqual(l, r opt.OptionalColList) bool {
+	return l.Equals(r)
+}
+
 func (h *hasher) IsOrderingEqual(l, r opt.Ordering) bool {
 	return l.Equals(r)
 }

pkg/sql/opt/norm/prune_cols_funcs.go

@@ -48,7 +48,7 @@ func (c *CustomFuncs) NeededMutationCols(
 	var cols opt.ColSet
 
 	// Add all input columns referenced by the mutation private.
-	addCols := func(list opt.ColList) {
+	addCols := func(list opt.OptionalColList) {
 		for _, id := range list {
 			if id != 0 {
 				cols.Add(id)
@@ -63,7 +63,7 @@ func (c *CustomFuncs) NeededMutationCols(
 	addCols(private.PartialIndexPutCols)
 	addCols(private.PartialIndexDelCols)
 	addCols(private.ReturnCols)
-	addCols(private.PassthroughCols)
+	addCols(opt.OptionalColList(private.PassthroughCols))
 	if private.CanaryCol != 0 {
 		cols.Add(private.CanaryCol)
 	}
@@ -318,13 +318,11 @@ func (c *CustomFuncs) PruneMutationFetchCols(
 // that are not in the neededCols set to zero. This indicates that those input
 // columns are not needed by this mutation list.
 func (c *CustomFuncs) filterMutationList(
-	tabID opt.TableID, inList opt.ColList, neededCols opt.ColSet,
-) opt.ColList {
-	newList := make(opt.ColList, len(inList))
+	tabID opt.TableID, inList opt.OptionalColList, neededCols opt.ColSet,
+) opt.OptionalColList {
+	newList := make(opt.OptionalColList, len(inList))
 	for i, c := range inList {
-		if !neededCols.Contains(tabID.ColumnID(i)) {
-			newList[i] = 0
-		} else {
+		if c != 0 && neededCols.Contains(tabID.ColumnID(i)) {
 			newList[i] = c
 		}
 	}
@@ -585,7 +583,7 @@ func DerivePruneCols(e memo.RelExpr) opt.ColSet {
 		// Find the columns that would need to be fetched, if no returning
 		// clause were present.
 		withoutReturningPrivate := *e.Private().(*memo.MutationPrivate)
-		withoutReturningPrivate.ReturnCols = opt.ColList{}
+		withoutReturningPrivate.ReturnCols = opt.OptionalColList{}
 		neededCols := neededMutationFetchCols(e.Memo(), e.Op(), &withoutReturningPrivate)
 
 		// Only the "free" RETURNING columns can be pruned away (i.e. the columns
@@ -639,25 +637,19 @@ func (c *CustomFuncs) PruneMutationReturnCols(
 	private *memo.MutationPrivate, needed opt.ColSet,
 ) *memo.MutationPrivate {
 	newPrivate := *private
-	newReturnCols := make(opt.ColList, len(private.ReturnCols))
-	newPassthroughCols := make(opt.ColList, 0, len(private.PassthroughCols))
 	tabID := c.mem.Metadata().TableMeta(private.Table).MetaID
 
 	// Prune away the ReturnCols that are unused.
-	for i := range private.ReturnCols {
-		if needed.Contains(tabID.ColumnID(i)) {
-			newReturnCols[i] = private.ReturnCols[i]
-		}
-	}
+	newPrivate.ReturnCols = c.filterMutationList(tabID, private.ReturnCols, needed)
 
+	newPassthroughCols := make(opt.ColList, 0, len(private.PassthroughCols))
 	// Prune away the PassthroughCols that are unused.
 	for _, passthroughCol := range private.PassthroughCols {
 		if passthroughCol != 0 && needed.Contains(passthroughCol) {
 			newPassthroughCols = append(newPassthroughCols, passthroughCol)
 		}
 	}
 
-	newPrivate.ReturnCols = newReturnCols
 	newPrivate.PassthroughCols = newPassthroughCols
 	return &newPrivate
 }