Merge #43450

43450: opt: derive constant computed columns for index selection r=andy-kimball a=andy-kimball

The optimizer uses explicitly specified filter constraints to qualify
available indexes during the exploration phase. It also uses implicit
filter constraints derived from table check constraints.

This commit adds new implicit filter constraints based on constant
computed columns. Constant computed columns are based on other columns in
the table that are constrained to be constant by other filters. For
example:

  CREATE TABLE hashed (
    k STRING,
    hash INT AS (fnv32(k) % 4) STORED,
    INDEX hash_index (hash, k)
  )

  SELECT * FROM hashed WHERE k = 'andy'

Here, the value of the hash column can be computed at query build time,
and therefore "hash_index" selected as the lowest cost index. The resulting
plan would be:

  scan hashed@secondary
   ├── columns: k:1(string!null) hash:2(int)
   ├── constraint: /2/1/3: [/1/'andy' - /1/'andy']
   └── fd: ()-->(1)

This improved ability to select indexes is useful for implementing HASH
indexes, which scatter keys across N buckets (see Issue #39340).

Release note (sql change): The optimizer can now derive constant computed
columns during index selection. This enables more efficient HASH indexes.

Co-authored-by: Andrew Kimball <[email protected]>

pkg/sql/opt/constraint/span.go

@@ -13,6 +13,7 @@ package constraint
 import (
 	"bytes"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/errors"
 )
 
@@ -69,6 +70,24 @@ func (sp *Span) IsUnconstrained() bool {
 	return startUnconstrained && endUnconstrained
 }
 
+// HasSingleKey is true if the span contains exactly one key. This is true when
+// the start key is the same as the end key, and both boundaries are inclusive.
+func (sp *Span) HasSingleKey(evalCtx *tree.EvalContext) bool {
+	l := sp.start.Length()
+	if l == 0 || l != sp.end.Length() {
+		return false
+	}
+	if sp.startBoundary != IncludeBoundary || sp.endBoundary != IncludeBoundary {
+		return false
+	}
+	for i, n := 0, l; i < n; i++ {
+		if sp.start.Value(i).Compare(evalCtx, sp.end.Value(i)) != 0 {
+			return false
+		}
+	}
+	return true
+}
+
 // StartKey returns the start key.
 func (sp *Span) StartKey() Key {
 	return sp.start

pkg/sql/opt/constraint/span_test.go

@@ -17,6 +17,7 @@ import (
 	"math"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 )
 
@@ -114,6 +115,90 @@ func TestSpanUnconstrained(t *testing.T) {
 	}
 }
 
+func TestSpanSingleKey(t *testing.T) {
+	testCases := []struct {
+		start         Key
+		startBoundary SpanBoundary
+		end           Key
+		endBoundary   SpanBoundary
+		expected      bool
+	}{
+		{ // 0
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			true,
+		},
+		{ // 1
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			MakeKey(tree.NewDInt(2)), IncludeBoundary,
+			false,
+		},
+		{ // 2
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			MakeKey(tree.NewDInt(1)), ExcludeBoundary,
+			false,
+		},
+		{ // 3
+			MakeKey(tree.NewDInt(1)), ExcludeBoundary,
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			false,
+		},
+		{ // 4
+			EmptyKey, IncludeBoundary,
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			false,
+		},
+		{ // 5
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			EmptyKey, IncludeBoundary,
+			false,
+		},
+		{ // 6
+			MakeKey(tree.NewDInt(1)), IncludeBoundary,
+			MakeKey(tree.DNull), IncludeBoundary,
+			false,
+		},
+		{ // 7
+			MakeKey(tree.NewDString("a")), IncludeBoundary,
+			MakeKey(tree.NewDString("ab")), IncludeBoundary,
+			false,
+		},
+		{ // 8
+			MakeCompositeKey(tree.NewDString("cherry"), tree.NewDInt(1)), IncludeBoundary,
+			MakeCompositeKey(tree.NewDString("cherry"), tree.NewDInt(1)), IncludeBoundary,
+			true,
+		},
+		{ // 9
+			MakeCompositeKey(tree.NewDString("cherry"), tree.NewDInt(1)), IncludeBoundary,
+			MakeCompositeKey(tree.NewDString("mango"), tree.NewDInt(1)), IncludeBoundary,
+			false,
+		},
+		{ // 10
+			MakeCompositeKey(tree.NewDString("cherry")), IncludeBoundary,
+			MakeCompositeKey(tree.NewDString("cherry"), tree.NewDInt(1)), IncludeBoundary,
+			false,
+		},
+		{ // 11
+			MakeCompositeKey(tree.NewDString("cherry"), tree.NewDInt(1), tree.DNull), IncludeBoundary,
+			MakeCompositeKey(tree.NewDString("cherry"), tree.NewDInt(1), tree.DNull), IncludeBoundary,
+			true,
+		},
+	}
+
+	for i, tc := range testCases {
+		st := cluster.MakeTestingClusterSettings()
+		evalCtx := tree.MakeTestingEvalContext(st)
+
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			var sp Span
+			sp.Init(tc.start, tc.startBoundary, tc.end, tc.endBoundary)
+			if sp.HasSingleKey(&evalCtx) != tc.expected {
+				t.Errorf("expected: %v, actual: %v", tc.expected, !tc.expected)
+			}
+		})
+	}
+}
+
 func TestSpanCompare(t *testing.T) {
 	keyCtx := testKeyContext(1, 2)
 

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -653,7 +653,7 @@ func findRoundingFunction(typ *types.T, precision int) (*tree.FunctionProperties
 // constraint defined on the target table. The mutation operator will report
 // a constraint violation error if the value of the column is false.
 func (mb *mutationBuilder) addCheckConstraintCols() {
-	if mb.tab.CheckCount() > 0 {
+	if mb.tab.CheckCount() != 0 {
 		// Disambiguate names so that references in the constraint expression refer
 		// to the correct columns.
 		mb.disambiguateColumns()
@@ -810,17 +810,7 @@ func (mb *mutationBuilder) buildReturning(returning tree.ReturningExprs) {
 	//
 	inScope := mb.outScope.replace()
 	inScope.expr = mb.outScope.expr
-	inScope.cols = make([]scopeColumn, 0, mb.tab.ColumnCount())
-	for i, n := 0, mb.tab.ColumnCount(); i < n; i++ {
-		tabCol := mb.tab.Column(i)
-		inScope.cols = append(inScope.cols, scopeColumn{
-			name:   tabCol.ColName(),
-			table:  mb.alias,
-			typ:    tabCol.DatumType(),
-			id:     mb.tabID.ColumnID(i),
-			hidden: tabCol.IsHidden(),
-		})
-	}
+	inScope.appendColumnsFromTable(mb.md.TableMeta(mb.tabID), &mb.alias)
 
 	// extraAccessibleCols contains all the columns that the RETURNING
 	// clause can refer to in addition to the table columns. This is useful for

pkg/sql/opt/optbuilder/scope.go

@@ -166,6 +166,25 @@ func (s *scope) appendColumnsFromScope(src *scope) {
 	}
 }
 
+// appendColumnsFromTable adds all columns from the given table metadata to this
+// scope.
+func (s *scope) appendColumnsFromTable(tabMeta *opt.TableMeta, alias *tree.TableName) {
+	tab := tabMeta.Table
+	if s.cols == nil {
+		s.cols = make([]scopeColumn, 0, tab.ColumnCount())
+	}
+	for i, n := 0, tab.ColumnCount(); i < n; i++ {
+		tabCol := tab.Column(i)
+		s.cols = append(s.cols, scopeColumn{
+			name:   tabCol.ColName(),
+			table:  *alias,
+			typ:    tabCol.DatumType(),
+			id:     tabMeta.MetaID.ColumnID(i),
+			hidden: tabCol.IsHidden(),
+		})
+	}
+}
+
 // appendColumns adds newly bound variables to this scope.
 // The expressions in the new columns are reset to nil.
 func (s *scope) appendColumns(cols []scopeColumn) {

pkg/sql/opt/optbuilder/select.go

@@ -469,7 +469,8 @@ func (b *Builder) buildScan(
 		}
 		outScope.expr = b.factory.ConstructScan(&private)
 
-		b.addCheckConstraintsForTable(outScope, tabMeta, ordinals != nil /* allowMissingColumns */)
+		b.addCheckConstraintsForTable(tabMeta)
+		b.addComputedColsForTable(tabMeta)
 
 		if b.trackViewDeps {
 			dep := opt.ViewDep{DataSource: tab}
@@ -489,17 +490,12 @@ func (b *Builder) buildScan(
 // addCheckConstraintsForTable finds all the check constraints that apply to the
 // table and adds them to the table metadata. To do this, the scalar expression
 // of the check constraints are built here.
-//
-// If allowMissingColumns is true, we ignore check constraints that involve
-// columns not in the current scope (useful when we build a scan that doesn't
-// contain all table columns).
-func (b *Builder) addCheckConstraintsForTable(
-	scope *scope, tabMeta *opt.TableMeta, allowMissingColumns bool,
-) {
-	tab := tabMeta.Table
+func (b *Builder) addCheckConstraintsForTable(tabMeta *opt.TableMeta) {
 	// Find all the check constraints that apply to the table and add them
 	// to the table meta data. To do this, we must build them into scalar
 	// expressions.
+	tableScope := scope{builder: b}
+	tab := tabMeta.Table
 	for i, n := 0, tab.CheckCount(); i < n; i++ {
 		checkConstraint := tab.Check(i)
 
@@ -512,25 +508,40 @@ func (b *Builder) addCheckConstraintsForTable(
 			panic(err)
 		}
 
-		var texpr tree.TypedExpr
-		func() {
-			if allowMissingColumns {
-				// Swallow any undefined column errors.
-				defer func() {
-					if r := recover(); r != nil {
-						if err, ok := r.(error); ok {
-							if code := pgerror.GetPGCode(err); code == pgcode.UndefinedColumn {
-								return
-							}
-						}
-						panic(r)
-					}
-				}()
-			}
-			texpr = scope.resolveAndRequireType(expr, types.Bool)
-		}()
-		if texpr != nil {
-			tabMeta.AddConstraint(b.buildScalar(texpr, scope, nil, nil, nil))
+		if len(tableScope.cols) == 0 {
+			tableScope.appendColumnsFromTable(tabMeta, &tabMeta.Alias)
+		}
+
+		if texpr := tableScope.resolveAndRequireType(expr, types.Bool); texpr != nil {
+			scalar := b.buildScalar(texpr, &tableScope, nil, nil, nil)
+			tabMeta.AddConstraint(scalar)
+		}
+	}
+}
+
+// addComputedColsForTable finds all computed columns in the given table and
+// caches them in the table metadata as scalar expressions.
+func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
+	tableScope := scope{builder: b}
+	tab := tabMeta.Table
+	for i, n := 0, tab.ColumnCount(); i < n; i++ {
+		tabCol := tab.Column(i)
+		if !tabCol.IsComputed() {
+			continue
+		}
+		expr, err := parser.ParseExpr(tabCol.ComputedExprStr())
+		if err != nil {
+			continue
+		}
+
+		if len(tableScope.cols) == 0 {
+			tableScope.appendColumnsFromTable(tabMeta, &tabMeta.Alias)
+		}
+
+		if texpr := tableScope.resolveAndRequireType(expr, types.Any); texpr != nil {
+			colID := tabMeta.MetaID.ColumnID(i)
+			scalar := b.buildScalar(texpr, &tableScope, nil, nil, nil)
+			tabMeta.AddComputedCol(colID, scalar)
 		}
 	}
 }

pkg/sql/opt/optgen/cmd/optgen/factory_gen.go

@@ -158,19 +158,19 @@ func (g *factoryGen) genReplace() {
 	g.w.writeIndent("// ancestors via a calls to the corresponding factory Construct methods. Here\n")
 	g.w.writeIndent("// is example usage:\n")
 	g.w.writeIndent("//\n")
-	g.w.writeIndent("//   var replace func(e opt.Expr, replace ReplaceFunc) opt.Expr\n")
-	g.w.writeIndent("//   replace = func(e opt.Expr, replace ReplaceFunc) opt.Expr {\n")
+	g.w.writeIndent("//   var replace func(e opt.Expr) opt.Expr\n")
+	g.w.writeIndent("//   replace = func(e opt.Expr) opt.Expr {\n")
 	g.w.writeIndent("//     if e.Op() == opt.VariableOp {\n")
 	g.w.writeIndent("//       return getReplaceVar(e)\n")
 	g.w.writeIndent("//     }\n")
-	g.w.writeIndent("//     return e.Replace(e, replace)\n")
+	g.w.writeIndent("//     return factory.Replace(e, replace)\n")
 	g.w.writeIndent("//   }\n")
 	g.w.writeIndent("//   replace(root, replace)\n")
 	g.w.writeIndent("//\n")
 	g.w.writeIndent("// Here, all variables in the tree are being replaced by some other expression\n")
 	g.w.writeIndent("// in a pre-order traversal of the tree. Post-order traversal is trivially\n")
-	g.w.writeIndent("// achieved by moving the e.Replace call to the top of the replace function\n")
-	g.w.writeIndent("// rather than bottom.\n")
+	g.w.writeIndent("// achieved by moving the factory.Replace call to the top of the replace\n")
+	g.w.writeIndent("// function rather than bottom.\n")
 	g.w.nestIndent("func (f *Factory) Replace(e opt.Expr, replace ReplaceFunc) opt.Expr {\n")
 	g.w.writeIndent("switch t := e.(type) {\n")
 

pkg/sql/opt/optgen/cmd/optgen/testdata/factory

@@ -166,19 +166,19 @@ func (_f *Factory) ConstructKVOptionsItem(
 // ancestors via a calls to the corresponding factory Construct methods. Here
 // is example usage:
 //
-//   var replace func(e opt.Expr, replace ReplaceFunc) opt.Expr
-//   replace = func(e opt.Expr, replace ReplaceFunc) opt.Expr {
+//   var replace func(e opt.Expr) opt.Expr
+//   replace = func(e opt.Expr) opt.Expr {
 //     if e.Op() == opt.VariableOp {
 //       return getReplaceVar(e)
 //     }
-//     return e.Replace(e, replace)
+//     return factory.Replace(e, replace)
 //   }
 //   replace(root, replace)
 //
 // Here, all variables in the tree are being replaced by some other expression
 // in a pre-order traversal of the tree. Post-order traversal is trivially
-// achieved by moving the e.Replace call to the top of the replace function
-// rather than bottom.
+// achieved by moving the factory.Replace call to the top of the replace
+// function rather than bottom.
 func (f *Factory) Replace(e opt.Expr, replace ReplaceFunc) opt.Expr {
 	switch t := e.(type) {
 	case *memo.SelectExpr:

pkg/sql/opt/table_meta.go

@@ -138,6 +138,12 @@ type TableMeta struct {
 	// detail.
 	Constraints []ScalarExpr
 
+	// ComputedCols stores ScalarExprs for each computed column on the table,
+	// indexed by ColumnID. These will be used when building mutation statements
+	// and constraining indexes. See comment above GenerateConstrainedScans for
+	// more detail.
+	ComputedCols map[ColumnID]ScalarExpr
+
 	// anns annotates the table metadata with arbitrary data.
 	anns [maxTableAnnIDCount]interface{}
 }
@@ -182,6 +188,14 @@ func (tm *TableMeta) AddConstraint(constraint ScalarExpr) {
 	tm.Constraints = append(tm.Constraints, constraint)
 }
 
+// AddComputedCol adds a computed column expression to the table's metadata.
+func (tm *TableMeta) AddComputedCol(colID ColumnID, computedCol ScalarExpr) {
+	if tm.ComputedCols == nil {
+		tm.ComputedCols = make(map[ColumnID]ScalarExpr)
+	}
+	tm.ComputedCols[colID] = computedCol
+}
+
 // TableAnnotation returns the given annotation that is associated with the
 // given table. If the table has no such annotation, TableAnnotation returns
 // nil.