opt: Enhance optsteps command to support exploration rules

The current version of the optsteps command doesn't handle exploration
rules very well. This is because exploration rules are not guaranteed to
produce a lower cost tree. Unless extra measures are taken, the returned
ExprView would not included the changed portion of the Memo, since
ExprView only shows the lowest cost path through the Memo.

The solution is to figure out which portion(s) of the tree are affected
by a transformation, and then use an alternate Coster that fools the
optimizer into thinking that any new expression(s) are the lowest cost.
It does this by assigning an infinite cost to the "real" best
expressions. The alternate expressions are now shown in the optsteps
output, but are deemphasized using different header text in order to
show these are not truly the lowest cost expressions.

Release note: None

pkg/sql/opt/memo/best_expr.go

@@ -82,6 +82,15 @@ func (be *BestExpr) Operator() opt.Operator {
 	return be.op
 }
 
+// Expr returns the memo expression referenced by this best expression. Note
+// that if the best expression is an enforcer (like a Sort), then the memo
+// expression is wrapped by the enforcer (maybe even by multiple enforcers).
+// This means that the same ExprID can be returned by different best expressions
+// in the same group, each of which would have a different Operator type.
+func (be *BestExpr) Expr() ExprID {
+	return be.eid
+}
+
 // Group returns the memo group which contains this best expression.
 func (be *BestExpr) Group() GroupID {
 	return be.eid.Group

pkg/sql/opt/memo/expr.go

@@ -39,6 +39,10 @@ type ExprID struct {
 	Expr  ExprOrdinal
 }
 
+// InvalidExprID is the uninitialized ExprID that never points to a valid
+// expression.
+var InvalidExprID = ExprID{}
+
 // MakeNormExprID returns the id of the normalized expression for the given
 // group.
 func MakeNormExprID(group GroupID) ExprID {

pkg/sql/opt/memo/expr_view.go

@@ -120,7 +120,7 @@ func (ev ExprView) Physical() *PhysicalProps {
 	if ev.best == normBestOrdinal {
 		panic("physical properties are not available when traversing the normalized tree")
 	}
-	return ev.mem.LookupPhysicalProps(ev.lookupBestExpr().required)
+	return ev.mem.LookupPhysicalProps(ev.bestExpr().required)
 }
 
 // Group returns the memo group containing this expression.
@@ -138,7 +138,7 @@ func (ev ExprView) Child(nth int) ExprView {
 		group := ev.ChildGroup(nth)
 		return MakeNormExprView(ev.mem, group)
 	}
-	return MakeExprView(ev.mem, ev.lookupBestExpr().Child(nth))
+	return MakeExprView(ev.mem, ev.bestExpr().Child(nth))
 }
 
 // ChildCount returns the number of expressions that are inputs to this
@@ -147,7 +147,7 @@ func (ev ExprView) ChildCount() int {
 	if ev.best == normBestOrdinal {
 		return ev.mem.NormExpr(ev.group).ChildCount()
 	}
-	return ev.lookupBestExpr().ChildCount()
+	return ev.bestExpr().ChildCount()
 }
 
 // ChildGroup returns the memo group containing the nth child of this parent
@@ -156,7 +156,7 @@ func (ev ExprView) ChildGroup(nth int) GroupID {
 	if ev.best == normBestOrdinal {
 		return ev.mem.NormExpr(ev.group).ChildGroup(ev.mem, nth)
 	}
-	return ev.lookupBestExpr().Child(nth).group
+	return ev.bestExpr().Child(nth).group
 }
 
 // Private returns any private data associated with this expression, or nil if
@@ -165,7 +165,7 @@ func (ev ExprView) Private() interface{} {
 	if ev.best == normBestOrdinal {
 		return ev.mem.NormExpr(ev.group).Private(ev.mem)
 	}
-	return ev.mem.Expr(ev.lookupBestExpr().eid).Private(ev.mem)
+	return ev.mem.Expr(ev.bestExpr().eid).Private(ev.mem)
 }
 
 // Metadata returns the metadata that's specific to this expression tree. Some
@@ -175,11 +175,21 @@ func (ev ExprView) Metadata() *opt.Metadata {
 	return ev.mem.metadata
 }
 
-func (ev ExprView) lookupChildGroup(nth int) *group {
+// Cost returns the cost of executing this expression tree, as estimated by the
+// optimizer. It is not available when the ExprView is traversing the normalized
+// expression tree.
+func (ev ExprView) Cost() Cost {
+	if ev.best == normBestOrdinal {
+		panic("Cost is not available when traversing the normalized tree")
+	}
+	return ev.mem.bestExpr(BestExprID{group: ev.group, ordinal: ev.best}).cost
+}
+
+func (ev ExprView) childGroup(nth int) *group {
 	return ev.mem.group(ev.ChildGroup(nth))
 }
 
-func (ev ExprView) lookupBestExpr() *BestExpr {
+func (ev ExprView) bestExpr() *BestExpr {
 	return ev.mem.group(ev.group).bestExpr(ev.best)
 }
 
@@ -328,7 +338,7 @@ func (ev ExprView) formatRelational(tp treeprinter.Node, flags ExprFmtFlags) {
 	}
 
 	if !flags.HasFlags(ExprFmtHideCost) && ev.best != normBestOrdinal {
-		tp.Childf("cost: %.2f", ev.lookupBestExpr().cost)
+		tp.Childf("cost: %.2f", ev.bestExpr().cost)
 	}
 
 	// Format weak keys.

pkg/sql/opt/memo/logical_props_factory.go

@@ -116,7 +116,7 @@ func (f logicalPropsFactory) constructScanProps(ev ExprView) LogicalProps {
 func (f logicalPropsFactory) constructSelectProps(ev ExprView) LogicalProps {
 	props := LogicalProps{Relational: &RelationalProps{}}
 
-	inputProps := ev.lookupChildGroup(0).logical.Relational
+	inputProps := ev.childGroup(0).logical.Relational
 
 	// Inherit input properties as starting point.
 	*props.Relational = *inputProps
@@ -129,7 +129,7 @@ func (f logicalPropsFactory) constructSelectProps(ev ExprView) LogicalProps {
 func (f logicalPropsFactory) constructProjectProps(ev ExprView) LogicalProps {
 	props := LogicalProps{Relational: &RelationalProps{}}
 
-	inputProps := ev.lookupChildGroup(0).logical.Relational
+	inputProps := ev.childGroup(0).logical.Relational
 
 	// Use output columns from projection list.
 	props.Relational.OutputCols = opt.ColListToSet(ev.Child(1).Private().(opt.ColList))
@@ -154,8 +154,8 @@ func (f logicalPropsFactory) constructProjectProps(ev ExprView) LogicalProps {
 func (f logicalPropsFactory) constructJoinProps(ev ExprView) LogicalProps {
 	props := LogicalProps{Relational: &RelationalProps{}}
 
-	leftProps := ev.lookupChildGroup(0).logical.Relational
-	rightProps := ev.lookupChildGroup(1).logical.Relational
+	leftProps := ev.childGroup(0).logical.Relational
+	rightProps := ev.childGroup(1).logical.Relational
 
 	// Output columns are union of columns from left and right inputs, except
 	// in case of semi and anti joins, which only project the left columns.
@@ -199,7 +199,7 @@ func (f logicalPropsFactory) constructJoinProps(ev ExprView) LogicalProps {
 func (f logicalPropsFactory) constructGroupByProps(ev ExprView) LogicalProps {
 	props := LogicalProps{Relational: &RelationalProps{}}
 
-	inputProps := ev.lookupChildGroup(0).logical.Relational
+	inputProps := ev.childGroup(0).logical.Relational
 
 	// Output columns are the union of grouping columns with columns from the
 	// aggregate projection list.
@@ -241,8 +241,8 @@ func (f logicalPropsFactory) constructGroupByProps(ev ExprView) LogicalProps {
 func (f logicalPropsFactory) constructSetProps(ev ExprView) LogicalProps {
 	props := LogicalProps{Relational: &RelationalProps{}}
 
-	leftProps := ev.lookupChildGroup(0).logical.Relational
-	rightProps := ev.lookupChildGroup(1).logical.Relational
+	leftProps := ev.childGroup(0).logical.Relational
+	rightProps := ev.childGroup(1).logical.Relational
 	colMap := *ev.Private().(*SetOpColMap)
 	if len(colMap.Out) != len(colMap.Left) || len(colMap.Out) != len(colMap.Right) {
 		panic(fmt.Errorf("lists in SetOpColMap are not all the same length. new:%d, left:%d, right:%d",
@@ -311,7 +311,7 @@ func (f logicalPropsFactory) constructMax1RowProps(ev ExprView) LogicalProps {
 func (f logicalPropsFactory) passThroughRelationalProps(ev ExprView, childIdx int) LogicalProps {
 	// Properties are immutable after construction, so just inherit relational
 	// props pointer from child.
-	return LogicalProps{Relational: ev.lookupChildGroup(childIdx).logical.Relational}
+	return LogicalProps{Relational: ev.childGroup(childIdx).logical.Relational}
 }
 
 func (f logicalPropsFactory) constructScalarProps(ev ExprView) LogicalProps {
@@ -326,7 +326,7 @@ func (f logicalPropsFactory) constructScalarProps(ev ExprView) LogicalProps {
 
 	// By default, union outer cols from all children, both relational and scalar.
 	for i := 0; i < ev.ChildCount(); i++ {
-		logical := &ev.lookupChildGroup(i).logical
+		logical := &ev.childGroup(i).logical
 		if logical.Scalar != nil {
 			props.Scalar.OuterCols.UnionWith(logical.Scalar.OuterCols)
 		} else {

pkg/sql/opt/memo/memo.go

@@ -188,8 +188,8 @@ func (m *Memo) GroupProperties(group GroupID) *LogicalProps {
 	return &m.groups[group].logical
 }
 
-// GroupByFingerprint returns the group of the expression that has the
-// given fingerprint.
+// GroupByFingerprint returns the group of the expression that has the given
+// fingerprint.
 func (m *Memo) GroupByFingerprint(f Fingerprint) GroupID {
 	return m.exprMap[f]
 }
@@ -252,7 +252,6 @@ func (m *Memo) MemoizeNormExpr(evalCtx *tree.EvalContext, norm Expr) GroupID {
 	if m.exprMap[norm.Fingerprint()] != 0 {
 		panic("normalized expression has been entered into the memo more than once")
 	}
-
 	mgrp := m.newGroup(norm)
 	ev := MakeNormExprView(m, mgrp.id)
 	logPropsFactory := logicalPropsFactory{evalCtx: evalCtx}
@@ -274,8 +273,7 @@ func (m *Memo) MemoizeDenormExpr(group GroupID, denorm Expr) {
 		}
 	} else {
 		// Add the denormalized expression to the memo.
-		mgrp := m.group(group)
-		mgrp.addExpr(denorm)
+		m.group(group).addExpr(denorm)
 		m.exprMap[denorm.Fingerprint()] = group
 	}
 }

pkg/sql/opt/norm/factory.go

@@ -25,17 +25,18 @@ import (
 )
 
 // MatchedRuleFunc defines the callback function for the NotifyOnMatchedRule
-// event. It is invoked each time a normalization rule has been matched by the
-// optimizer. The name of the matched rule is passed as a parameter. If the
-// function returns false, then the rule is not applied (i.e. it's skipped).
+// event supported by the optimizer and factory. It is invoked each time an
+// optimization rule (Normalize or Explore) has been matched. The name of the
+// matched rule is passed as a parameter. If the function returns false, then
+// the rule is not applied (i.e. skipped).
 type MatchedRuleFunc func(ruleName opt.RuleName) bool
 
-// AppliedRuleFunc defines the callback function for the AppliedRuleFunc event
-// supported by the Optimizer and Factory. It is invoked each time an
-// optimization rule (Normalize or Explore) has been applied by the optimizer.
-// The function is called with the name of the rule and the memo group it
-// affected. If the rule was an exploration rule, then the added parameter
-// gives the number of expressions added to the group by the rule.
+// AppliedRuleFunc defines the callback function for the NotifyOnAppliedRule
+// event supported by the optimizer and factory. It is invoked each time an
+// optimization rule (Normalize or Explore) has been applied. The function is
+// called with the name of the rule and the memo group it affected. If the rule
+// was an exploration rule, then the added parameter gives the number of
+// expressions added to the group by the rule.
 type AppliedRuleFunc func(ruleName opt.RuleName, group memo.GroupID, added int)
 
 //go:generate optgen -out factory.og.go factory ../ops/*.opt rules/*.opt
@@ -101,9 +102,9 @@ func (f *Factory) DisableOptimizations() {
 
 // NotifyOnMatchedRule sets a callback function which is invoked each time a
 // normalize rule has been matched by the factory. If matchedRule is nil, then
-// no further notifications are sent. If no callback function is set, then all
-// rules are applied by default. In addition, callers can invoke the
-// DisableOptimizations convenience method to disable all rules.
+// no further notifications are sent, and all rules are applied by default. In
+// addition, callers can invoke the DisableOptimizations convenience method to
+// disable all rules.
 func (f *Factory) NotifyOnMatchedRule(matchedRule MatchedRuleFunc) {
 	f.matchedRule = matchedRule
 }