Merge #60821

60821: opt: simplify optbuilder arbiter logic r=mgartner a=mgartner

#### opt: add arbiterSet to track arbiter indexes and unique constraints

This commit introduces a `arbiterSet` which can contain both indexes and
unique constraints. The `ForEach` function of this set abstracts the
type of arbiter so that consumers of the set need not be concerned with
the type. This cleans up code that builds expressions for each arbiter.
Note that this will allow further cleanup within `buildInputForUpsert`,
once that function supports partial unique constraints.

Release note: None

#### opt: move left-join building logic to mutation_builder_arbiter.go

This commit move the logic for building `UPSERT` left-joins to
`mutation_builder_arbiter.go`, making it consistent with the structure
of `INSERT ON CONFLICT DO NOTHING` logic. Some duplicate code for
building DistinctOn expressions was also removed - both `UPSERT` and
`INSERT ON CONFLICT DO NOTHING` build essentially the same DistinctOn
expressions.

Some of the code for dealing with `UPSERT` arbiter indexes and arbiter
constraints is still repetitive. This will be cleaned up in a follow-up
commit that adds `UPSERT` support for partial unique constraint
arbiters.

Release note: None


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

pkg/sql/opt/optbuilder/BUILD.bazel

@@ -4,6 +4,7 @@ go_library(
     name = "optbuilder",
     srcs = [
         "alter_table.go",
+        "arbiter_set.go",
         "builder.go",
         "create_table.go",
         "create_view.go",

pkg/sql/opt/optbuilder/arbiter_set.go

@@ -0,0 +1,128 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package optbuilder
+
+import (
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/util"
+)
+
+// arbiterSet represents a set of arbiters. Unique indexes or constraints can be
+// arbiters. This set provides an abstraction on top of both types of arbiters
+// so that consumers of the set do not need to be concerned with the underlying
+// arbiter type.
+type arbiterSet struct {
+	mb *mutationBuilder
+
+	// indexes contains the index arbiters in the set, as ordinals into the
+	// table's indexes.
+	indexes util.FastIntSet
+
+	// uniqueConstraints contains the unique constraint arbiters in the set, as
+	// ordinals into the table's unique constraints.
+	uniqueConstraints util.FastIntSet
+}
+
+// makeArbiterSet returns an initialized arbiterSet.
+func makeArbiterSet(mb *mutationBuilder) arbiterSet {
+	return arbiterSet{
+		mb: mb,
+	}
+}
+
+// makeSingleIndexArbiterSet returns an initialized arbiterSet with the given
+// index as the sole arbiter in the set.
+func makeSingleIndexArbiterSet(mb *mutationBuilder, idx cat.IndexOrdinal) arbiterSet {
+	a := makeArbiterSet(mb)
+	a.AddIndex(idx)
+	return a
+}
+
+// makeSingleUniqueConstraintArbiterSet returns an initialized arbiterSet with
+// the given unique constraint as the sole arbiter in the set.
+func makeSingleUniqueConstraintArbiterSet(mb *mutationBuilder, uniq cat.UniqueOrdinal) arbiterSet {
+	a := makeArbiterSet(mb)
+	a.AddUniqueConstraint(uniq)
+	return a
+}
+
+// AddIndex adds an index arbiter to the set.
+func (a *arbiterSet) AddIndex(idx cat.IndexOrdinal) {
+	a.indexes.Add(idx)
+}
+
+// AddUniqueConstraint adds a unique constraint arbiter to the set.
+func (a *arbiterSet) AddUniqueConstraint(uniq cat.UniqueOrdinal) {
+	a.uniqueConstraints.Add(uniq)
+}
+
+// Empty returns true if the set is empty.
+func (a *arbiterSet) Empty() bool {
+	return a.indexes.Empty() && a.uniqueConstraints.Empty()
+}
+
+// Len returns the number of the arbiters in the set.
+func (a *arbiterSet) Len() int {
+	return a.indexes.Len() + a.uniqueConstraints.Len()
+}
+
+// IndexOrdinals returns a slice of all the index ordinals in the set.
+func (a *arbiterSet) IndexOrdinals() []int {
+	return a.indexes.Ordered()
+}
+
+// UniqueConstraintOrdinals returns a slice of all the unique constraint
+// ordinals in the set.
+func (a *arbiterSet) UniqueConstraintOrdinals() []int {
+	return a.uniqueConstraints.Ordered()
+}
+
+// ForEach calls a function for every arbiter in the set. The function is called
+// with the following arguments:
+//
+//   - name is the name of the index or unique constraint.
+//   - conflictOrds is the set of table column ordinals of the arbiter. For
+//     index arbiters, this is the lax key columns. For constraint arbiters,
+//     this is all the columns in the constraint.
+//   - pred is the partial predicate expression of the arbiter, if the arbiter
+//     is a partial index or partial constraint. If the arbiter is not partial,
+//     pred is nil.
+//
+func (a *arbiterSet) ForEach(f func(name string, conflictOrds util.FastIntSet, pred tree.Expr)) {
+	// Call the callback for each index arbiter.
+	a.indexes.ForEach(func(i int) {
+		index := a.mb.tab.Index(i)
+
+		conflictOrds := getIndexLaxKeyOrdinals(index)
+
+		var pred tree.Expr
+		if _, isPartial := index.Predicate(); isPartial {
+			pred = a.mb.parsePartialIndexPredicateExpr(i)
+		}
+
+		f(string(index.Name()), conflictOrds, pred)
+	})
+
+	// Call the callback for each unique constraint arbiter.
+	a.uniqueConstraints.ForEach(func(i int) {
+		uniqueConstraint := a.mb.tab.Unique(i)
+
+		conflictOrds := getUniqueConstraintOrdinals(a.mb.tab, uniqueConstraint)
+
+		var pred tree.Expr
+		if _, isPartial := uniqueConstraint.Predicate(); isPartial {
+			pred = a.mb.parseUniqueConstraintPredicateExpr(i)
+		}
+
+		f(uniqueConstraint.Name(), conflictOrds, pred)
+	})
+}

pkg/sql/opt/optbuilder/distinct.go

@@ -55,8 +55,8 @@ func (b *Builder) constructDistinct(inScope *scope) memo.RelExpr {
 // operator rather than the DistinctOn operator (see the UpsertDistinctOn
 // operator comment for details on the differences). The errorOnDup parameter
 // controls whether multiple rows in the same distinct group trigger an error.
-// This can only take on a value in the EnsureDistinctOn and
-// EnsureUpsertDistinctOn cases.
+// If empty, no error is triggered. This can only take on a value in the
+// EnsureDistinctOn and EnsureUpsertDistinctOn cases.
 func (b *Builder) buildDistinctOn(
 	distinctOnCols opt.ColSet, inScope *scope, nullsAreDistinct bool, errorOnDup string,
 ) (outScope *scope) {