Merge cockroachdb#58110

58110: opt: add OptionalColList type r=RaduBerinde a=RaduBerinde

I am open to better suggestions for the name.

#### opt: change ColSetToList() func to a ToList() method

Release note: None

#### 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

Co-authored-by: Radu Berinde <[email protected]>

pkg/sql/opt/colset.go

@@ -99,6 +99,15 @@ func (s ColSet) SingleColumn() ColumnID {
 	return col
 }
 
+// ToList converts the set to a ColList, in column ID order.
+func (s ColSet) ToList() ColList {
+	res := make(ColList, 0, s.Len())
+	s.ForEach(func(x ColumnID) {
+		res = append(res, x)
+	})
+	return res
+}
+
 // TranslateColSet is used to translate a ColSet from one set of column IDs
 // to an equivalent set. This is relevant for set operations such as UNION,
 // INTERSECT and EXCEPT, and can be used to map a ColSet defined on the left

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,13 +66,68 @@ func (cl ColList) Equals(other ColList) bool {
 	return true
 }
 
-// ColSetToList converts a column id set to a list, in column id order.
-func ColSetToList(set ColSet) ColList {
-	res := make(ColList, 0, set.Len())
-	set.ForEach(func(x ColumnID) {
-		res = append(res, x)
-	})
-	return res
+// 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

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

@@ -193,7 +193,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))
@@ -203,7 +203,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))
 		}
@@ -302,7 +302,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)
@@ -724,7 +724,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)
@@ -737,11 +737,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/exec/execbuilder/relational.go

@@ -1197,7 +1197,7 @@ func (b *Builder) buildDistinct(distinct memo.RelExpr) (execPlan, error) {
 	if input.outputCols.Len() == outCols.Len() {
 		return ep, nil
 	}
-	return b.ensureColumns(ep, opt.ColSetToList(outCols), distinct.ProvidedPhysical().Ordering)
+	return b.ensureColumns(ep, outCols.ToList(), distinct.ProvidedPhysical().Ordering)
 }
 
 func (b *Builder) buildGroupByInput(groupBy memo.RelExpr) (execPlan, error) {

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

@@ -282,7 +282,7 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 
 	default:
 		// Fall back to writing output columns in column id order.
-		colList = opt.ColSetToList(e.Relational().OutputCols)
+		colList = e.Relational().OutputCols.ToList()
 	}
 
 	f.formatColumns(e, tp, colList, required.Presentation)
@@ -294,7 +294,7 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 		*UpsertDistinctOnExpr, *EnsureUpsertDistinctOnExpr:
 		private := e.Private().(*GroupingPrivate)
 		if !f.HasFlags(ExprFmtHideColumns) && !private.GroupingCols.Empty() {
-			f.formatColList(e, tp, "grouping columns:", opt.ColSetToList(private.GroupingCols))
+			f.formatColList(e, tp, "grouping columns:", private.GroupingCols.ToList())
 		}
 		if !f.HasFlags(ExprFmtHidePhysProps) && !private.Ordering.Any() {
 			tp.Childf("internal-ordering: %s", private.Ordering)
@@ -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)
 }