Merge pull request #67775 from rytaft/backport21.1-67702

release-21.1: opt,sql: add EXPLAIN (OPT, MEMO) to show optimizer memo

pkg/sql/opt/bench/bench_test.go

@@ -475,7 +475,13 @@ func (h *harness) runSimple(tb testing.TB, query benchQuery, phase Phase) {
 
 	root := execMemo.RootExpr()
 	eb := execbuilder.New(
-		exec.StubFactory{}, execMemo, nil /* catalog */, root, &h.evalCtx, true, /* allowAutoCommit */
+		exec.StubFactory{},
+		&h.optimizer,
+		execMemo,
+		nil, /* catalog */
+		root,
+		&h.evalCtx,
+		true, /* allowAutoCommit */
 	)
 	if _, err = eb.Build(); err != nil {
 		tb.Fatalf("%v", err)
@@ -521,7 +527,13 @@ func (h *harness) runPrepared(tb testing.TB, phase Phase) {
 
 	root := execMemo.RootExpr()
 	eb := execbuilder.New(
-		exec.StubFactory{}, execMemo, nil /* catalog */, root, &h.evalCtx, true, /* allowAutoCommit */
+		exec.StubFactory{},
+		&h.optimizer,
+		execMemo,
+		nil, /* catalog */
+		root,
+		&h.evalCtx,
+		true, /* allowAutoCommit */
 	)
 	if _, err := eb.Build(); err != nil {
 		tb.Fatalf("%v", err)

pkg/sql/opt/exec/execbuilder/builder.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
@@ -32,6 +33,7 @@ const ParallelScanResultThreshold = 10000
 // expression tree (opt.Expr).
 type Builder struct {
 	factory          exec.Factory
+	optimizer        *xform.Optimizer
 	mem              *memo.Memo
 	catalog          cat.Catalog
 	e                opt.Expr
@@ -105,6 +107,7 @@ type Builder struct {
 // transaction.
 func New(
 	factory exec.Factory,
+	optimizer *xform.Optimizer,
 	mem *memo.Memo,
 	catalog cat.Catalog,
 	e opt.Expr,
@@ -113,6 +116,7 @@ func New(
 ) *Builder {
 	b := &Builder{
 		factory:                factory,
+		optimizer:              optimizer,
 		mem:                    mem,
 		catalog:                catalog,
 		e:                      e,

pkg/sql/opt/exec/execbuilder/cascades.go

@@ -293,7 +293,7 @@ func (cb *cascadeBuilder) planCascade(
 	}
 
 	// 5. Execbuild the optimized expression.
-	eb := New(execFactory, factory.Memo(), cb.b.catalog, optimizedExpr, evalCtx, allowAutoCommit)
+	eb := New(execFactory, &o, factory.Memo(), cb.b.catalog, optimizedExpr, evalCtx, allowAutoCommit)
 	if bufferRef != nil {
 		// Set up the With binding.
 		eb.addBuiltWithExpr(cascadeInputWithID, bufferColMap, bufferRef)

pkg/sql/opt/exec/execbuilder/format.go

@@ -35,7 +35,15 @@ func fmtInterceptor(f *memo.ExprFmtCtx, scalar opt.ScalarExpr) string {
 	}
 
 	// Build the scalar expression and format it as a single string.
-	bld := New(nil /* factory */, f.Memo, nil /* catalog */, scalar, nil /* evalCtx */, false /* allowAutoCommit */)
+	bld := New(
+		nil, /* factory */
+		nil, /* optimizer */
+		f.Memo,
+		nil, /* catalog */
+		scalar,
+		nil,   /* evalCtx */
+		false, /* allowAutoCommit */
+	)
 	expr, err := bld.BuildScalar()
 	if err != nil {
 		// Not all scalar operators are supported (e.g. Projections).

pkg/sql/opt/exec/execbuilder/relational.go

@@ -818,7 +818,7 @@ func (b *Builder) buildApplyJoin(join memo.RelExpr) (execPlan, error) {
 			return nil, err
 		}
 
-		eb := New(ef, f.Memo(), b.catalog, newRightSide, b.evalCtx, false /* allowAutoCommit */)
+		eb := New(ef, &o, f.Memo(), b.catalog, newRightSide, b.evalCtx, false /* allowAutoCommit */)
 		eb.disableTelemetry = true
 		plan, err := eb.Build()
 		if err != nil {

pkg/sql/opt/exec/execbuilder/statement.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/util/treeprinter"
@@ -93,6 +94,11 @@ func (b *Builder) buildExplainOpt(explain *memo.ExplainExpr) (execPlan, error) {
 		// TODO(radu): add views, sequences
 	}
 
+	// If MEMO option was passed, show the memo.
+	if explain.Options.Flags[tree.ExplainFlagMemo] {
+		planText.WriteString(b.optimizer.FormatMemo(xform.FmtPretty))
+	}
+
 	f := memo.MakeExprFmtCtx(fmtFlags, b.mem, b.catalog)
 	f.FormatExpr(explain.Input)
 	planText.WriteString(f.Buffer.String())
@@ -122,7 +128,9 @@ func (b *Builder) buildExplain(explain *memo.ExplainExpr) (execPlan, error) {
 		explain.StmtType,
 		func(ef exec.ExplainFactory) (exec.Plan, error) {
 			// Create a separate builder for the explain query.
-			explainBld := New(ef, b.mem, b.catalog, explain.Input, b.evalCtx, b.initialAllowAutoCommit)
+			explainBld := New(
+				ef, b.optimizer, b.mem, b.catalog, explain.Input, b.evalCtx, b.initialAllowAutoCommit,
+			)
 			explainBld.disableTelemetry = true
 			return explainBld.Build()
 		},

pkg/sql/opt/exec/execbuilder/testdata/explain

@@ -1567,3 +1567,124 @@ vectorized: true
 • insert fast path
   into: t2(x)
   auto commit
+
+# Tests for EXPLAIN (OPT, MEMO).
+query T
+EXPLAIN (OPT, MEMO) SELECT * FROM tc JOIN t ON k=a
+----
+memo (optimized, ~16KB, required=[presentation: info:10])
+ ├── G1: (explain G2 [presentation: a:1,b:2,k:6,v:7])
+ │    └── [presentation: info:10]
+ │         ├── best: (explain G2="[presentation: a:1,b:2,k:6,v:7]" [presentation: a:1,b:2,k:6,v:7])
+ │         └── cost: 2271.55
+ ├── G2: (inner-join G3 G4 G5) (inner-join G3 G4 G5) (inner-join G4 G3 G5) (merge-join G3 G4 G6 inner-join,+1,+6) (lookup-join G3 G6 t,keyCols=[1],outCols=(1,2,6,7)) (merge-join G3 G4 G6 inner-join,+1,+6) (lookup-join G3 G6 t,keyCols=[1],outCols=(1,2,6,7)) (merge-join G4 G3 G6 inner-join,+6,+1) (lookup-join G7 G6 tc,keyCols=[3],outCols=(1,2,6,7))
+ │    └── [presentation: a:1,b:2,k:6,v:7]
+ │         ├── best: (inner-join G3 G4 G5)
+ │         └── cost: 2271.54
+ ├── G3: (scan tc,cols=(1,2))
+ │    ├── [ordering: +1]
+ │    │    ├── best: (sort G3)
+ │    │    └── cost: 1365.34
+ │    └── []
+ │         ├── best: (scan tc,cols=(1,2))
+ │         └── cost: 1126.01
+ ├── G4: (scan t,cols=(6,7))
+ │    ├── [ordering: +6]
+ │    │    ├── best: (scan t,cols=(6,7))
+ │    │    └── cost: 1105.61
+ │    └── []
+ │         ├── best: (scan t,cols=(6,7))
+ │         └── cost: 1105.61
+ ├── G5: (filters G8)
+ ├── G6: (filters)
+ ├── G7: (lookup-join G4 G6 tc@c,keyCols=[6],outCols=(1,3,6,7))
+ │    └── []
+ │         ├── best: (lookup-join G4 G6 tc@c,keyCols=[6],outCols=(1,3,6,7))
+ │         └── cost: 23174.72
+ ├── G8: (eq G9 G10)
+ ├── G9: (variable k)
+ └── G10: (variable a)
+inner-join (hash)
+ ├── scan tc
+ ├── scan t
+ └── filters
+      └── k = a
+
+query T
+EXPLAIN (OPT, MEMO, VERBOSE, CATALOG) SELECT * FROM tc JOIN t ON k=a
+----
+TABLE tc
+ ├── a int
+ ├── b int
+ ├── rowid int not null default (unique_rowid()) [hidden]
+ ├── crdb_internal_mvcc_timestamp decimal [hidden] [system]
+ ├── tableoid oid [hidden] [system]
+ ├── PRIMARY INDEX primary
+ │    └── rowid int not null default (unique_rowid()) [hidden]
+ └── INDEX c
+      ├── a int
+      └── rowid int not null default (unique_rowid()) [hidden]
+TABLE t
+ ├── k int not null
+ ├── v int
+ ├── crdb_internal_mvcc_timestamp decimal [hidden] [system]
+ ├── tableoid oid [hidden] [system]
+ └── PRIMARY INDEX primary
+      └── k int not null
+memo (optimized, ~16KB, required=[presentation: info:10])
+ ├── G1: (explain G2 [presentation: a:1,b:2,k:6,v:7])
+ │    └── [presentation: info:10]
+ │         ├── best: (explain G2="[presentation: a:1,b:2,k:6,v:7]" [presentation: a:1,b:2,k:6,v:7])
+ │         └── cost: 2271.55
+ ├── G2: (inner-join G3 G4 G5) (inner-join G3 G4 G5) (inner-join G4 G3 G5) (merge-join G3 G4 G6 inner-join,+1,+6) (lookup-join G3 G6 t,keyCols=[1],outCols=(1,2,6,7)) (merge-join G3 G4 G6 inner-join,+1,+6) (lookup-join G3 G6 t,keyCols=[1],outCols=(1,2,6,7)) (merge-join G4 G3 G6 inner-join,+6,+1) (lookup-join G7 G6 tc,keyCols=[3],outCols=(1,2,6,7))
+ │    └── [presentation: a:1,b:2,k:6,v:7]
+ │         ├── best: (inner-join G3 G4 G5)
+ │         └── cost: 2271.54
+ ├── G3: (scan tc,cols=(1,2))
+ │    ├── [ordering: +1]
+ │    │    ├── best: (sort G3)
+ │    │    └── cost: 1365.34
+ │    └── []
+ │         ├── best: (scan tc,cols=(1,2))
+ │         └── cost: 1126.01
+ ├── G4: (scan t,cols=(6,7))
+ │    ├── [ordering: +6]
+ │    │    ├── best: (scan t,cols=(6,7))
+ │    │    └── cost: 1105.61
+ │    └── []
+ │         ├── best: (scan t,cols=(6,7))
+ │         └── cost: 1105.61
+ ├── G5: (filters G8)
+ ├── G6: (filters)
+ ├── G7: (lookup-join G4 G6 tc@c,keyCols=[6],outCols=(1,3,6,7))
+ │    └── []
+ │         ├── best: (lookup-join G4 G6 tc@c,keyCols=[6],outCols=(1,3,6,7))
+ │         └── cost: 23174.72
+ ├── G8: (eq G9 G10)
+ ├── G9: (variable k)
+ └── G10: (variable a)
+inner-join (hash)
+ ├── columns: a:1 b:2 k:6 v:7
+ ├── multiplicity: left-rows(zero-or-one), right-rows(zero-or-more)
+ ├── stats: [rows=990, distinct(1)=99, null(1)=0, distinct(6)=99, null(6)=0]
+ ├── cost: 2271.54
+ ├── fd: (6)-->(7), (1)==(6), (6)==(1)
+ ├── prune: (2,7)
+ ├── scan tc
+ │    ├── columns: a:1 b:2
+ │    ├── stats: [rows=1000, distinct(1)=100, null(1)=10]
+ │    ├── cost: 1126.01
+ │    ├── prune: (1,2)
+ │    ├── interesting orderings: (+1)
+ │    └── unfiltered-cols: (1-5)
+ ├── scan t
+ │    ├── columns: k:6 v:7
+ │    ├── stats: [rows=1000, distinct(6)=1000, null(6)=0]
+ │    ├── cost: 1105.61
+ │    ├── key: (6)
+ │    ├── fd: (6)-->(7)
+ │    ├── prune: (6,7)
+ │    ├── interesting orderings: (+6)
+ │    └── unfiltered-cols: (6-9)
+ └── filters
+      └── k:6 = a:1 [outer=(1,6), constraints=(/1: (/NULL - ]; /6: (/NULL - ]), fd=(1)==(6), (6)==(1)]

pkg/sql/opt/idxconstraint/index_constraints_test.go

@@ -134,8 +134,8 @@ func TestIndexConstraints(t *testing.T) {
 				remainingFilter := ic.RemainingFilters()
 				if !remainingFilter.IsTrue() {
 					execBld := execbuilder.New(
-						nil /* execFactory */, f.Memo(), nil /* catalog */, &remainingFilter,
-						&evalCtx, false, /* allowAutoCommit */
+						nil /* execFactory */, nil /* optimizer */, f.Memo(), nil, /* catalog */
+						&remainingFilter, &evalCtx, false, /* allowAutoCommit */
 					)
 					expr, err := execBld.BuildScalar()
 					if err != nil {

pkg/sql/opt/partialidx/implicator_test.go

@@ -112,8 +112,8 @@ func TestImplicator(t *testing.T) {
 				buf.WriteString("none")
 			} else {
 				execBld := execbuilder.New(
-					nil /* factory */, f.Memo(), nil /* catalog */, &remainingFilters,
-					&evalCtx, false, /* allowAutoCommit */
+					nil /* factory */, nil /* optimizer */, f.Memo(), nil, /* catalog */
+					&remainingFilters, &evalCtx, false, /* allowAutoCommit */
 				)
 				expr, err := execBld.BuildScalar()
 				if err != nil {

pkg/sql/plan_opt.go

@@ -550,7 +550,7 @@ func (opc *optPlanningCtx) runExecBuilder(
 	var containsFullIndexScan bool
 	if !planTop.instrumentation.ShouldBuildExplainPlan() {
 		// No instrumentation.
-		bld := execbuilder.New(f, mem, &opc.catalog, mem.RootExpr(), evalCtx, allowAutoCommit)
+		bld := execbuilder.New(f, &opc.optimizer, mem, &opc.catalog, mem.RootExpr(), evalCtx, allowAutoCommit)
 		plan, err := bld.Build()
 		if err != nil {
 			return err
@@ -562,7 +562,9 @@ func (opc *optPlanningCtx) runExecBuilder(
 	} else {
 		// Create an explain factory and record the explain.Plan.
 		explainFactory := explain.NewFactory(f)
-		bld := execbuilder.New(explainFactory, mem, &opc.catalog, mem.RootExpr(), evalCtx, allowAutoCommit)
+		bld := execbuilder.New(
+			explainFactory, &opc.optimizer, mem, &opc.catalog, mem.RootExpr(), evalCtx, allowAutoCommit,
+		)
 		plan, err := bld.Build()
 		if err != nil {
 			return err

pkg/sql/sem/tree/eval_test.go

@@ -103,7 +103,7 @@ func optBuildScalar(evalCtx *tree.EvalContext, e tree.Expr) (tree.TypedExpr, err
 	}
 
 	bld := execbuilder.New(
-		nil /* factory */, o.Memo(), nil /* catalog */, o.Memo().RootExpr(),
+		nil /* factory */, &o, o.Memo(), nil /* catalog */, o.Memo().RootExpr(),
 		evalCtx, false, /* allowAutoCommit */
 	)
 	expr, err := bld.BuildScalar()

pkg/sql/sem/tree/explain.go

@@ -111,6 +111,7 @@ const (
 	ExplainFlagJSON
 	ExplainFlagStages
 	ExplainFlagDeps
+	ExplainFlagMemo
 	numExplainFlags = iota
 )
 
@@ -122,6 +123,7 @@ var explainFlagStrings = [...]string{
 	ExplainFlagJSON:    "JSON",
 	ExplainFlagStages:  "STAGES",
 	ExplainFlagDeps:    "DEPS",
+	ExplainFlagMemo:    "MEMO",
 }
 
 var explainFlagStringMap = func() map[string]ExplainFlag {