opt: add lookupjoin.Constraint struct

The `lookupjoin.Constraint` struct has been added to encapsulate
multiple data structures that represent a strategy for constraining a
lookup join.

Release note: None

pkg/sql/opt/lookupjoin/BUILD.bazel

@@ -12,9 +12,9 @@ go_library(
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/norm",
         "//pkg/sql/opt/props",
+        "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
-        "@com_github_cockroachdb_errors//:errors",
     ],
 )
 
@@ -29,7 +29,6 @@ go_test(
     deps = [
         "//pkg/settings/cluster",
         "//pkg/sql/opt",
-        "//pkg/sql/opt/constraint",
         "//pkg/sql/opt/exec/execbuilder",
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/norm",
@@ -38,6 +37,7 @@ go_test(
         "//pkg/sql/opt/testutils",
         "//pkg/sql/opt/testutils/testcat",
         "//pkg/sql/parser",
+        "//pkg/sql/sem/eval",
         "//pkg/sql/sem/tree",
         "//pkg/testutils",
         "//pkg/util/leaktest",

pkg/sql/opt/lookupjoin/constraint_builder.go

@@ -24,6 +24,49 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 )
 
+// Constraint is used to constrain a lookup join. There are two types of
+// constraints:
+//
+//   1. Constraints with KeyCols use columns from the input to directly
+//      constrain lookups into a target index.
+//   2. Constraints with a LookupExpr build multiple spans from an expression
+//      that is evaluated for each input row. These spans are used to perform
+//      lookups into a target index.
+//
+// A constraint is not constraining if both KeyCols and LookupExpr are empty.
+// See IsUnconstrained.
+type Constraint struct {
+	// KeyCols is an ordered list of columns from the left side of the join to
+	// be used as lookup join key columns. This list corresponds to the columns
+	// in RightSideCols. It will be nil if LookupExpr is non-nil.
+	KeyCols opt.ColList
+
+	// RightSideCols is an ordered list of prefix index columns that are
+	// constrained by this constraint. It corresponds 1:1 with the columns in
+	// KeyCols if KeyCols is non-nil. Otherwise, it includes the prefix of index
+	// columns constrained by LookupExpr.
+	RightSideCols opt.ColList
+
+	// LookupExpr is a lookup expression for multi-span lookup joins. It will be
+	// nil if KeyCols is non-nil.
+	LookupExpr memo.FiltersExpr
+
+	// InputProjections contains constant values and computed columns that must
+	// be projected on the lookup join's input.
+	InputProjections memo.ProjectionsExpr
+
+	// ConstFilters contains constant equalities and ranges in either KeyCols or
+	// LookupExpr that are used to aid selectivity estimation. See
+	// memo.LookupJoinPrivate.ConstFilters.
+	ConstFilters memo.FiltersExpr
+}
+
+// IsUnconstrained returns true if the constraint does not constrain a lookup
+// join.
+func (c *Constraint) IsUnconstrained() bool {
+	return len(c.KeyCols) == 0 && len(c.LookupExpr) == 0
+}
+
 // ConstraintBuilder determines how to constrain index key columns for a lookup
 // join. See Build for more details.
 type ConstraintBuilder struct {
@@ -73,33 +116,11 @@ func (b *ConstraintBuilder) Init(
 	}
 }
 
-// Build determines how to constrain index key columns for a lookup join. It
-// returns either a list of key columns or a lookup expression. It has several
-// return values:
-//
-//   keyCols          - An ordered list of columns from the left side of the
-//   					join that correspond to a prefix of columns in the given
-//   					index.
-//   lookupExpr       - A lookup expression for multi-span lookup joins.
-//   inputProjections - Projections of constant values and computed columns that
-//                      must be projected on the lookup join's input.
-//   constFilters     - Filters representing constant equalities and ranges in
-//                      either keyCols or lookupExpr that are used to aid
-//                      selectivity estimation.
-//   rightSideCols    - A list of constrained index columns.
-//
-// Build will return either non-nil keyCols or a non-nil lookupExpr, but not
-// both. If both keyCols and lookupExpr are nil, then the index cannot be used
-// for a lookup join.
+// Build returns a Constraint that constrains a lookup join on the given index.
+// The constraint returned may be unconstrained if no constraint could be built.
 func (b *ConstraintBuilder) Build(
 	index cat.Index, onFilters, optionalFilters memo.FiltersExpr,
-) (
-	keyCols opt.ColList,
-	lookupExpr memo.FiltersExpr,
-	inputProjections memo.ProjectionsExpr,
-	constFilters memo.FiltersExpr,
-	rightSideCols opt.ColList,
-) {
+) Constraint {
 	allFilters := append(onFilters, optionalFilters...)
 
 	// Check if the first column in the index either:
@@ -115,7 +136,7 @@ func (b *ConstraintBuilder) Build(
 	if _, ok := b.rightEq.Find(firstIdxCol); !ok {
 		if _, ok := b.findComputedColJoinEquality(b.table, firstIdxCol, b.rightEqSet); !ok {
 			if _, _, ok := b.findJoinFilterConstants(allFilters, firstIdxCol); !ok {
-				return
+				return Constraint{}
 			}
 		}
 	}
@@ -124,8 +145,11 @@ func (b *ConstraintBuilder) Build(
 	// an equality with another column or a constant.
 	numIndexKeyCols := index.LaxKeyColumnCount()
 
-	keyCols = make(opt.ColList, 0, numIndexKeyCols)
-	rightSideCols = make(opt.ColList, 0, numIndexKeyCols)
+	keyCols := make(opt.ColList, 0, numIndexKeyCols)
+	rightSideCols := make(opt.ColList, 0, numIndexKeyCols)
+	var inputProjections memo.ProjectionsExpr
+	var lookupExpr memo.FiltersExpr
+	var constFilters memo.FiltersExpr
 	shouldBuildMultiSpanLookupJoin := false
 
 	// All the lookup conditions must apply to the prefix of the index and so
@@ -237,12 +261,21 @@ func (b *ConstraintBuilder) Build(
 
 		// A multi-span lookup join with a lookup expression has no key columns
 		// and requires no projections on the input.
-		return nil /* keyCols */, lookupExpr, nil /* inputProjections */, constFilters, rightSideCols
+		return Constraint{
+			RightSideCols: rightSideCols,
+			LookupExpr:    lookupExpr,
+			ConstFilters:  constFilters,
+		}
 	}
 
 	// If we did not build a lookup expression, return the key columns we found,
 	// if any.
-	return keyCols, nil /* lookupExpr */, inputProjections, constFilters, rightSideCols
+	return Constraint{
+		KeyCols:          keyCols,
+		RightSideCols:    rightSideCols,
+		InputProjections: inputProjections,
+		ConstFilters:     constFilters,
+	}
 }
 
 // findComputedColJoinEquality returns the computed column expression of col and

pkg/sql/opt/lookupjoin/constraint_builder_test.go

@@ -147,36 +147,35 @@ func TestLookupConstraints(t *testing.T) {
 
 				var cb lookupjoin.ConstraintBuilder
 				cb.Init(&f, md, f.EvalContext(), rightTable, leftCols, rightCols, leftEq, rightEq)
-				keyCols, lookupExpr, inputProjections, _, _ :=
-					cb.Build(index, filters, optionalFilters)
+				lookupConstraint := cb.Build(index, filters, optionalFilters)
 				var b strings.Builder
-				if len(keyCols) == 0 && len(lookupExpr) == 0 {
+				if lookupConstraint.IsUnconstrained() {
 					b.WriteString("lookup join not possible")
 				}
-				if len(keyCols) > 0 {
+				if len(lookupConstraint.KeyCols) > 0 {
 					b.WriteString("key cols:\n")
-					for i := range keyCols {
+					for i := range lookupConstraint.KeyCols {
 						b.WriteString("  ")
 						b.WriteString(string(index.Column(i).ColName()))
 						b.WriteString(" = ")
-						b.WriteString(md.ColumnMeta(keyCols[i]).Alias)
+						b.WriteString(md.ColumnMeta(lookupConstraint.KeyCols[i]).Alias)
 						b.WriteString("\n")
 					}
 				}
-				if len(inputProjections) > 0 {
+				if len(lookupConstraint.InputProjections) > 0 {
 					b.WriteString("input projections:\n")
-					for i := range inputProjections {
-						col := inputProjections[i].Col
+					for i := range lookupConstraint.InputProjections {
+						col := lookupConstraint.InputProjections[i].Col
 						b.WriteString("  ")
 						b.WriteString(md.ColumnMeta(col).Alias)
 						b.WriteString(" = ")
-						b.WriteString(formatScalar(inputProjections[i].Element, &f, &evalCtx))
+						b.WriteString(formatScalar(lookupConstraint.InputProjections[i].Element, &f, &evalCtx))
 						b.WriteString("\n")
 					}
 				}
-				if len(lookupExpr) > 0 {
+				if len(lookupConstraint.LookupExpr) > 0 {
 					b.WriteString("lookup expression:\n  ")
-					b.WriteString(formatScalar(&lookupExpr, &f, &evalCtx))
+					b.WriteString(formatScalar(&lookupConstraint.LookupExpr, &f, &evalCtx))
 					b.WriteString("\n")
 				}
 				return b.String()

pkg/sql/opt/xform/join_funcs.go

@@ -370,9 +370,8 @@ func (c *CustomFuncs) generateLookupJoinsImpl(
 			return
 		}
 
-		keyCols, lookupExpr, inputProjections, constFilters, rightSideCols :=
-			cb.Build(index, onFilters, optionalFilters)
-		if len(keyCols) == 0 && len(lookupExpr) == 0 {
+		lookupConstraint := cb.Build(index, onFilters, optionalFilters)
+		if lookupConstraint.IsUnconstrained() {
 			// We couldn't find equality columns or a lookup expression to
 			// perform a lookup join on this index.
 			return
@@ -384,33 +383,33 @@ func (c *CustomFuncs) generateLookupJoinsImpl(
 		lookupJoin.Table = scanPrivate.Table
 		lookupJoin.Index = index.Ordinal()
 		lookupJoin.Locking = scanPrivate.Locking
-		lookupJoin.KeyCols = keyCols
-		lookupJoin.LookupExpr = lookupExpr
+		lookupJoin.KeyCols = lookupConstraint.KeyCols
+		lookupJoin.LookupExpr = lookupConstraint.LookupExpr
 
 		// Wrap the input in a Project if any projections are required. The
 		// lookup join will project away these synthesized columns.
-		if len(inputProjections) > 0 {
+		if len(lookupConstraint.InputProjections) > 0 {
 			lookupJoin.Input = c.e.f.ConstructProject(
 				lookupJoin.Input,
-				inputProjections,
+				lookupConstraint.InputProjections,
 				lookupJoin.Input.Relational().OutputCols,
 			)
 		}
 
 		tableFDs := memo.MakeTableFuncDep(md, scanPrivate.Table)
 		// A lookup join will drop any input row which contains NULLs, so a lax key
 		// is sufficient.
-		lookupJoin.LookupColsAreTableKey = tableFDs.ColsAreLaxKey(rightSideCols.ToSet())
+		lookupJoin.LookupColsAreTableKey = tableFDs.ColsAreLaxKey(lookupConstraint.RightSideCols.ToSet())
 
 		// Remove redundant filters from the ON condition if columns were
 		// constrained by equality filters or constant filters.
 		lookupJoin.On = onFilters
 		if len(lookupJoin.KeyCols) > 0 {
-			lookupJoin.On = memo.ExtractRemainingJoinFilters(lookupJoin.On, lookupJoin.KeyCols, rightSideCols)
+			lookupJoin.On = memo.ExtractRemainingJoinFilters(lookupJoin.On, lookupJoin.KeyCols, lookupConstraint.RightSideCols)
 		}
 		lookupJoin.On = lookupJoin.On.Difference(lookupJoin.LookupExpr)
-		lookupJoin.On = lookupJoin.On.Difference(constFilters)
-		lookupJoin.ConstFilters = constFilters
+		lookupJoin.On = lookupJoin.On.Difference(lookupConstraint.ConstFilters)
+		lookupJoin.ConstFilters = lookupConstraint.ConstFilters
 
 		// Add input columns and lookup expression columns, since these will be
 		// needed for all join types and cases.