Merge #24804 #24835

24804: opt: Enhance optsteps command to support exploration rules r=andy-kimball a=andy-kimball

This PR has 2 commits, for easier review:

1. Add AppliedRule event to memo: adds a new callback that the optimizer fires once it's done applying a rule; opsteps uses this to determine which part of the tree is affected by a rule

2. Enhance optsteps command to support exploration rules

24835: testutils: Update RPC error matcher for #22658 r=bdarnell a=bdarnell

Acceptance tests are (infrequently) failing with this message, when it
looks like this test has historically swallowed RPC errors here. I'm
not sure that's completely correct, but this change should restore the
test to its previous level of reliability.



Release note: None

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

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}

pkg/sql/opt/norm/factory.go

@@ -24,6 +24,21 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/types"
 )
 
+// MatchedRuleFunc defines the callback function for the NotifyOnMatchedRule
+// 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 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
 
 // Factory constructs a normalized expression tree within the memo. As each
@@ -57,10 +72,15 @@ type Factory struct {
 	// map, and will skip application of the rule.
 	ruleCycles map[memo.Fingerprint]bool
 
-	// onRuleMatch is the callback function that is invoked each time a normalize
+	// matchedRule is the callback function that is invoked each time a normalize
 	// rule has been matched by the factory. It can be set via a call to the
-	// SetOnRuleMatch method.
-	onRuleMatch func(ruleName opt.RuleName) bool
+	// NotifyOnMatchedRule method.
+	matchedRule MatchedRuleFunc
+
+	// appliedRule is the callback function which is invoked each time a normalize
+	// rule has been applied by the factory. It can be set via a call to the
+	// NotifyOnAppliedRule method.
+	appliedRule AppliedRuleFunc
 }
 
 // NewFactory returns a new Factory structure with a new, blank memo structure
@@ -77,17 +97,23 @@ func NewFactory(evalCtx *tree.EvalContext) *Factory {
 // expression tree becomes the output expression tree (because no transforms
 // are applied).
 func (f *Factory) DisableOptimizations() {
-	f.SetOnRuleMatch(func(opt.RuleName) bool { return false })
+	f.NotifyOnMatchedRule(func(opt.RuleName) bool { return false })
 }
 
-// SetOnRuleMatch sets a callback function which is invoked each time a
-// normalize rule has been matched by the factory. If the function returns
-// false, then the rule is not applied. By default, all rules are applied, but
-// callers can set the callback function to override the default behavior. In
+// 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, and all rules are applied by default. In
 // addition, callers can invoke the DisableOptimizations convenience method to
 // disable all rules.
-func (f *Factory) SetOnRuleMatch(onRuleMatch func(ruleName opt.RuleName) bool) {
-	f.onRuleMatch = onRuleMatch
+func (f *Factory) NotifyOnMatchedRule(matchedRule MatchedRuleFunc) {
+	f.matchedRule = matchedRule
+}
+
+// NotifyOnAppliedRule sets a callback function which is invoked each time a
+// normalize rule has been applied by the factory. If appliedRule is nil, then
+// no further notifications are sent.
+func (f *Factory) NotifyOnAppliedRule(appliedRule AppliedRuleFunc) {
+	f.appliedRule = appliedRule
 }
 
 // Memo returns the memo structure that the factory is operating upon.