From a466599a918f82e6949f8e84e608e6bde679676e Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Thu, 11 Jul 2019 10:31:40 -0400
Subject: [PATCH 1/2] opt: add EXPLAIN (OPT, CATALOG)

Adding a `CATALOG` flag to `EXPLAIN (OPT)` which also formats all the
cat.Table objects. This is useful for testing the opt_catalog
implementation.

Release note: None
---
 pkg/ccl/logictestccl/testdata/logic_test/zone | 141 ++++++++++++++++++
 pkg/sql/apply_join.go                         |   2 +-
 pkg/sql/opt/bench/bench_test.go               |   3 +-
 pkg/sql/opt/exec/execbuilder/builder.go       |  17 ++-
 pkg/sql/opt/exec/execbuilder/format.go        |   2 +-
 pkg/sql/opt/exec/execbuilder/statement.go     |  20 ++-
 pkg/sql/opt/exec/execbuilder/testdata/catalog |  80 ++++++++++
 pkg/sql/opt/exec/execbuilder/testdata/explain |  57 +++++++
 .../idxconstraint/index_constraints_test.go   |   2 +-
 pkg/sql/opt_index_selection.go                |   2 +-
 pkg/sql/plan_opt.go                           |   2 +-
 pkg/sql/sem/tree/eval_test.go                 |   2 +-
 pkg/sql/sem/tree/explain.go                   |   2 +
 13 files changed, 321 insertions(+), 11 deletions(-)
 create mode 100644 pkg/sql/opt/exec/execbuilder/testdata/catalog

diff --git a/pkg/ccl/logictestccl/testdata/logic_test/zone b/pkg/ccl/logictestccl/testdata/logic_test/zone
index 5ba78f5c4708..72e0f43f9762 100644
--- a/pkg/ccl/logictestccl/testdata/logic_test/zone
+++ b/pkg/ccl/logictestccl/testdata/logic_test/zone
@@ -30,6 +30,28 @@ scan  ·      ·
 ·     table  t@secondary
 ·     spans  /10-/11
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * FROM t
+----
+TABLE t
+ ├── k int not null
+ ├── v string
+ ├── INDEX primary
+ │    ├── k int not null
+ │    └── ZONE
+ │         └── constraints: [+region=test,+dc=dc2]
+ ├── INDEX secondary
+ │    ├── k int not null
+ │    ├── v string (storing)
+ │    └── ZONE
+ │         └── constraints: [+region=test,+dc=dc1]
+ └── INDEX tertiary
+      ├── k int not null
+      ├── v string (storing)
+      └── ZONE
+           └── constraints: [+region=test,+dc=dc2]
+scan t@secondary
+
 # ------------------------------------------------------------------------------
 # Move secondary to dc3 and put tertiary in dc1 and ensure that gateway matches
 # tertiary instead of secondary. Regression for #35546.
@@ -48,6 +70,28 @@ scan  ·      ·
 ·     table  t@tertiary
 ·     spans  /10-/11
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * FROM t
+----
+TABLE t
+ ├── k int not null
+ ├── v string
+ ├── INDEX primary
+ │    ├── k int not null
+ │    └── ZONE
+ │         └── constraints: [+region=test,+dc=dc2]
+ ├── INDEX secondary
+ │    ├── k int not null
+ │    ├── v string (storing)
+ │    └── ZONE
+ │         └── constraints: [+region=test,+dc=dc3]
+ └── INDEX tertiary
+      ├── k int not null
+      ├── v string (storing)
+      └── ZONE
+           └── constraints: [+region=test,+dc=dc1]
+scan t@tertiary
+
 # ------------------------------------------------------------------------------
 # Swap secondary and tertiary localities and ensure invalidation occurs.
 # Regression for #35546.
@@ -84,6 +128,28 @@ scan  ·      ·
 ·     table  t@primary
 ·     spans  /10-/10/#
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * FROM t
+----
+TABLE t
+ ├── k int not null
+ ├── v string
+ ├── INDEX primary
+ │    ├── k int not null
+ │    └── ZONE
+ │         └── constraints: [+region=test,+dc=dc1]
+ ├── INDEX secondary
+ │    ├── k int not null
+ │    ├── v string (storing)
+ │    └── ZONE
+ │         └── constraints: [+region=test,+dc=dc2]
+ └── INDEX tertiary
+      ├── k int not null
+      ├── v string (storing)
+      └── ZONE
+           └── constraints: [+region=test,+dc=dc3]
+scan t
+
 # ------------------------------------------------------------------------------
 # Use PREPARE to make sure that the prepared plan is invalidated when the
 # secondary index's constraints change.
@@ -155,6 +221,31 @@ scan  ·      ·
 ·     table  t@tertiary
 ·     spans  /10-/11
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * FROM t
+----
+TABLE t
+ ├── k int not null
+ ├── v string
+ ├── INDEX primary
+ │    ├── k int not null
+ │    └── ZONE
+ │         ├── constraints: [+region=test]
+ │         └── lease preference: [+region=test,+dc=dc2]
+ ├── INDEX secondary
+ │    ├── k int not null
+ │    ├── v string (storing)
+ │    └── ZONE
+ │         ├── constraints: [+region=test]
+ │         └── lease preference: [+region=test,+dc=dc3]
+ └── INDEX tertiary
+      ├── k int not null
+      ├── v string (storing)
+      └── ZONE
+           ├── constraints: [+region=test]
+           └── lease preference: [+region=test,+dc=dc1]
+scan t@tertiary
+
 # ------------------------------------------------------------------------------
 # Ensure that an index constrained to a region is preferred over an index that
 # merely has a lease preference in that region (since lease preferences can
@@ -180,6 +271,31 @@ scan  ·      ·
 ·     table  t@secondary
 ·     spans  /10-/11
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * FROM t
+----
+TABLE t
+ ├── k int not null
+ ├── v string
+ ├── INDEX primary
+ │    ├── k int not null
+ │    └── ZONE
+ │         ├── constraints: [+region=test]
+ │         └── lease preference: [+region=test,+dc=dc1]
+ ├── INDEX secondary
+ │    ├── k int not null
+ │    ├── v string (storing)
+ │    └── ZONE
+ │         ├── constraints: [+region=test,+dc=dc1]
+ │         └── lease preference: [+region=test,+dc=dc3]
+ └── INDEX tertiary
+      ├── k int not null
+      ├── v string (storing)
+      └── ZONE
+           ├── constraints: [+region=test]
+           └── lease preference: [+region=test,+dc=dc1]
+scan t@secondary
+
 # ------------------------------------------------------------------------------
 # Use PREPARE to make sure that the prepared plan is invalidated when the
 # secondary index's lease preferences change.
@@ -246,6 +362,31 @@ scan  ·      ·
 ·     table  t36642@tertiary
 ·     spans  /10-/11
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * FROM t
+----
+TABLE t
+ ├── k int not null
+ ├── v string
+ ├── INDEX primary
+ │    ├── k int not null
+ │    └── ZONE
+ │         ├── constraints: [+region=test]
+ │         └── lease preference: [+region=test,+dc=dc1]
+ ├── INDEX secondary
+ │    ├── k int not null
+ │    ├── v string (storing)
+ │    └── ZONE
+ │         ├── constraints: [+region=test]
+ │         └── lease preference: [+region=test,+dc=dc2]
+ └── INDEX tertiary
+      ├── k int not null
+      ├── v string (storing)
+      └── ZONE
+           ├── constraints: [+region=test]
+           └── lease preference: [+region=test,+dc=dc1]
+scan t
+
 
 # ------------------------------------------------------------------------------
 # Regression for issue #36644. Allow matching constraints for leading locality
diff --git a/pkg/sql/apply_join.go b/pkg/sql/apply_join.go
index 8895c1851727..f259c984785b 100644
--- a/pkg/sql/apply_join.go
+++ b/pkg/sql/apply_join.go
@@ -262,7 +262,7 @@ func (a *applyJoinNode) Next(params runParams) (bool, error) {
 		}
 
 		execFactory := makeExecFactory(params.p)
-		eb := execbuilder.New(&execFactory, factory.Memo(), newRightSide, params.EvalContext())
+		eb := execbuilder.New(&execFactory, factory.Memo(), nil /* catalog */, newRightSide, params.EvalContext())
 		eb.DisableTelemetry()
 		p, err := eb.Build()
 		if err != nil {
diff --git a/pkg/sql/opt/bench/bench_test.go b/pkg/sql/opt/bench/bench_test.go
index a967a688cf70..1689a07478c2 100644
--- a/pkg/sql/opt/bench/bench_test.go
+++ b/pkg/sql/opt/bench/bench_test.go
@@ -584,7 +584,8 @@ func (h *harness) runUsingAPI(tb testing.TB, bmType BenchmarkType, usePrepared b
 
 	root := execMemo.RootExpr()
 	execFactory := stubFactory{}
-	if _, err = execbuilder.New(&execFactory, execMemo, root, &h.evalCtx).Build(); err != nil {
+	eb := execbuilder.New(&execFactory, execMemo, nil /* catalog */, root, &h.evalCtx)
+	if _, err = eb.Build(); err != nil {
 		tb.Fatalf("%v", err)
 	}
 }
diff --git a/pkg/sql/opt/exec/execbuilder/builder.go b/pkg/sql/opt/exec/execbuilder/builder.go
index 901f78cf9055..586d81f67b38 100644
--- a/pkg/sql/opt/exec/execbuilder/builder.go
+++ b/pkg/sql/opt/exec/execbuilder/builder.go
@@ -12,6 +12,7 @@ package execbuilder
 
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
+	"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/sem/tree"
@@ -25,6 +26,7 @@ import (
 type Builder struct {
 	factory            exec.Factory
 	mem                *memo.Memo
+	catalog            cat.Catalog
 	e                  opt.Expr
 	disableTelemetry   bool
 	evalCtx            *tree.EvalContext
@@ -53,12 +55,23 @@ type Builder struct {
 // New constructs an instance of the execution node builder using the
 // given factory to construct nodes. The Build method will build the execution
 // node tree from the given optimized expression tree.
-func New(factory exec.Factory, mem *memo.Memo, e opt.Expr, evalCtx *tree.EvalContext) *Builder {
+//
+// catalog is only needed if the statement contains an EXPLAIN (OPT, CATALOG).
+func New(
+	factory exec.Factory, mem *memo.Memo, catalog cat.Catalog, e opt.Expr, evalCtx *tree.EvalContext,
+) *Builder {
 	var nameGen *memo.ExprNameGenerator
 	if evalCtx != nil && evalCtx.SessionData.SaveTablesPrefix != "" {
 		nameGen = memo.NewExprNameGenerator(evalCtx.SessionData.SaveTablesPrefix)
 	}
-	return &Builder{factory: factory, mem: mem, e: e, evalCtx: evalCtx, nameGen: nameGen}
+	return &Builder{
+		factory: factory,
+		mem:     mem,
+		catalog: catalog,
+		e:       e,
+		evalCtx: evalCtx,
+		nameGen: nameGen,
+	}
 }
 
 // DisableTelemetry prevents the execbuilder from updating telemetry counters.
diff --git a/pkg/sql/opt/exec/execbuilder/format.go b/pkg/sql/opt/exec/execbuilder/format.go
index 4dc52ffbb610..9adde2b60023 100644
--- a/pkg/sql/opt/exec/execbuilder/format.go
+++ b/pkg/sql/opt/exec/execbuilder/format.go
@@ -49,7 +49,7 @@ func fmtInterceptor(f *memo.ExprFmtCtx, tp treeprinter.Node, nd opt.Expr) bool {
 	}
 
 	// Build the scalar expression and format it as a single tree node.
-	bld := New(nil /* factory */, f.Memo, nd, nil /* evalCtx */)
+	bld := New(nil /* factory */, f.Memo, nil /* catalog */, nd, nil /* evalCtx */)
 	md := f.Memo.Metadata()
 	ivh := tree.MakeIndexedVarHelper(nil /* container */, md.NumColumns())
 	expr, err := bld.BuildScalar(&ivh)
diff --git a/pkg/sql/opt/exec/execbuilder/statement.go b/pkg/sql/opt/exec/execbuilder/statement.go
index 2052be9ed241..292b70503908 100644
--- a/pkg/sql/opt/exec/execbuilder/statement.go
+++ b/pkg/sql/opt/exec/execbuilder/statement.go
@@ -11,10 +11,14 @@
 package execbuilder
 
 import (
+	"bytes"
+
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
+	"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/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/util/treeprinter"
 )
 
 func (b *Builder) buildCreateTable(ct *memo.CreateTableExpr) (execPlan, error) {
@@ -60,9 +64,21 @@ func (b *Builder) buildExplain(explain *memo.ExplainExpr) (execPlan, error) {
 		}
 
 		// Format the plan here and pass it through to the exec factory.
+
+		// If catalog option was passed, show catalog object details for all tables.
+		var planText bytes.Buffer
+		if explain.Options.Flags.Contains(tree.ExplainFlagCatalog) {
+			for _, t := range b.mem.Metadata().AllTables() {
+				tp := treeprinter.New()
+				cat.FormatTable(b.catalog, t.Table, tp)
+				planText.WriteString(tp.String())
+			}
+			// TODO(radu): add views, sequences
+		}
+
 		f := memo.MakeExprFmtCtx(fmtFlags, b.mem)
 		f.FormatExpr(explain.Input)
-		planText := f.Buffer.String()
+		planText.WriteString(f.Buffer.String())
 
 		// If we're going to display the environment, there's a bunch of queries we
 		// need to run to get that information, and we can't run them from here, so
@@ -73,7 +89,7 @@ func (b *Builder) buildExplain(explain *memo.ExplainExpr) (execPlan, error) {
 		}
 
 		var err error
-		node, err = b.factory.ConstructExplainOpt(planText, envOpts)
+		node, err = b.factory.ConstructExplainOpt(planText.String(), envOpts)
 		if err != nil {
 			return execPlan{}, err
 		}
diff --git a/pkg/sql/opt/exec/execbuilder/testdata/catalog b/pkg/sql/opt/exec/execbuilder/testdata/catalog
new file mode 100644
index 000000000000..35091bd2c425
--- /dev/null
+++ b/pkg/sql/opt/exec/execbuilder/testdata/catalog
@@ -0,0 +1,80 @@
+# LogicTest: local-opt
+
+statement ok
+CREATE TABLE xyz (
+  x INT PRIMARY KEY,
+  y INT,
+  z INT,
+  INDEX foo (z, y)
+)
+
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from xyz
+----
+TABLE xyz
+ ├── x int not null
+ ├── y int
+ ├── z int
+ ├── INDEX primary
+ │    └── x int not null
+ └── INDEX foo
+      ├── z int
+      ├── y int
+      └── x int not null
+scan xyz
+
+statement ok
+CREATE TABLE abcdef (
+    a INT NOT NULL,
+    b INT,
+    c INT DEFAULT (10),
+    d INT AS (abcdef.b + c + 1) STORED,
+    e INT AS (a) STORED,
+    f INT CHECK (f > 2)
+)
+
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from abcdef
+----
+TABLE abcdef
+ ├── a int not null
+ ├── b int
+ ├── c int default (10:::INT8)
+ ├── d int as ((b + c) + 1) stored
+ ├── e int as (a) stored
+ ├── f int
+ ├── rowid int not null default (unique_rowid()) [hidden]
+ ├── CHECK (f > 2)
+ └── INDEX primary
+      └── rowid int not null default (unique_rowid()) [hidden]
+scan abcdef
+
+statement ok
+CREATE TABLE uvwxy (
+    u INT,
+    v INT,
+    w INT,
+    x INT,
+    y INT,
+    PRIMARY KEY (u,v),
+    FAMILY (u,v,w),
+    FAMILY (x),
+    FAMILY (y)
+)
+
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from uvwxy
+----
+TABLE uvwxy
+ ├── u int not null
+ ├── v int not null
+ ├── w int
+ ├── x int
+ ├── y int
+ ├── FAMILY fam_0_u_v_w (u, v, w)
+ ├── FAMILY fam_1_x (x)
+ ├── FAMILY fam_2_y (y)
+ └── INDEX primary
+      ├── u int not null
+      └── v int not null
+scan uvwxy
diff --git a/pkg/sql/opt/exec/execbuilder/testdata/explain b/pkg/sql/opt/exec/execbuilder/testdata/explain
index 51b1f5026dc9..3a502a5b0178 100644
--- a/pkg/sql/opt/exec/execbuilder/testdata/explain
+++ b/pkg/sql/opt/exec/execbuilder/testdata/explain
@@ -704,6 +704,63 @@ sort
            ├── key: (3)
            └── fd: ()-->(1)
 
+query T
+EXPLAIN (OPT,CATALOG) SELECT * FROM tc WHERE a = 10 ORDER BY b
+----
+TABLE tc
+ ├── a int
+ ├── b int
+ ├── rowid int not null default (unique_rowid()) [hidden]
+ ├── INDEX primary
+ │    └── rowid int not null default (unique_rowid()) [hidden]
+ └── INDEX c
+      ├── a int
+      └── rowid int not null default (unique_rowid()) [hidden]
+sort
+ └── index-join tc
+      └── scan tc@c
+           └── constraint: /1/3: [/10 - /10]
+
+query T
+EXPLAIN (OPT,VERBOSE,CATALOG) SELECT * FROM tc JOIN t ON k=a
+----
+TABLE tc
+ ├── a int
+ ├── b int
+ ├── rowid int not null default (unique_rowid()) [hidden]
+ ├── 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
+ └── INDEX primary
+      └── k int not null
+inner-join (hash)
+ ├── columns: a:1 b:2 k:4 v:5
+ ├── stats: [rows=1000, distinct(1)=100, null(1)=0, distinct(4)=100, null(4)=0]
+ ├── cost: 2220.05
+ ├── fd: (4)-->(5), (1)==(4), (4)==(1)
+ ├── prune: (2,5)
+ ├── scan tc
+ │    ├── columns: a:1 b:2
+ │    ├── stats: [rows=1000, distinct(1)=100, null(1)=10]
+ │    ├── cost: 1100.02
+ │    ├── prune: (1,2)
+ │    └── interesting orderings: (+1)
+ ├── scan t
+ │    ├── columns: k:4 v:5
+ │    ├── stats: [rows=1000, distinct(4)=1000, null(4)=0]
+ │    ├── cost: 1080.02
+ │    ├── key: (4)
+ │    ├── fd: (4)-->(5)
+ │    ├── prune: (4,5)
+ │    └── interesting orderings: (+4)
+ └── filters
+      └── k = a [outer=(1,4), constraints=(/1: (/NULL - ]; /4: (/NULL - ]), fd=(1)==(4), (4)==(1)]
+
 query T
 EXPLAIN (OPT) SELECT * FROM tc WHERE a + 2 * b > 1 ORDER BY a*b
 ----
diff --git a/pkg/sql/opt/idxconstraint/index_constraints_test.go b/pkg/sql/opt/idxconstraint/index_constraints_test.go
index 46aa70ebc544..4947be8346ba 100644
--- a/pkg/sql/opt/idxconstraint/index_constraints_test.go
+++ b/pkg/sql/opt/idxconstraint/index_constraints_test.go
@@ -161,7 +161,7 @@ func TestIndexConstraints(t *testing.T) {
 				}
 				remainingFilter := ic.RemainingFilters()
 				if !remainingFilter.IsTrue() {
-					execBld := execbuilder.New(nil /* execFactory */, f.Memo(), &remainingFilter, &evalCtx)
+					execBld := execbuilder.New(nil /* execFactory */, f.Memo(), nil /* catalog */, &remainingFilter, &evalCtx)
 					expr, err := execBld.BuildScalar(&iVarHelper)
 					if err != nil {
 						return fmt.Sprintf("error: %v\n", err)
diff --git a/pkg/sql/opt_index_selection.go b/pkg/sql/opt_index_selection.go
index 6aa7c204903d..1b8b24769b0a 100644
--- a/pkg/sql/opt_index_selection.go
+++ b/pkg/sql/opt_index_selection.go
@@ -249,7 +249,7 @@ func (p *planner) selectIndex(
 		if rem.IsTrue() {
 			s.filter = nil
 		} else {
-			execBld := execbuilder.New(nil /* execFactory */, optimizer.Memo(), &rem, nil /* evalCtx */)
+			execBld := execbuilder.New(nil /* execFactory */, optimizer.Memo(), nil /* catalog */, &rem, nil /* evalCtx */)
 			s.filter, err = execBld.BuildScalar(&s.filterVars)
 			if err != nil {
 				return nil, err
diff --git a/pkg/sql/plan_opt.go b/pkg/sql/plan_opt.go
index bc2b1e7eb28a..76a96b9223cb 100644
--- a/pkg/sql/plan_opt.go
+++ b/pkg/sql/plan_opt.go
@@ -149,7 +149,7 @@ func (p *planner) makeOptimizerPlan(ctx context.Context) (_ *planTop, isCorrelat
 	// Build the plan tree.
 	root := execMemo.RootExpr()
 	execFactory := makeExecFactory(p)
-	plan, err := execbuilder.New(&execFactory, execMemo, root, p.EvalContext()).Build()
+	plan, err := execbuilder.New(&execFactory, execMemo, &opc.catalog, root, p.EvalContext()).Build()
 	if err != nil {
 		return nil, isCorrelated, err
 	}
diff --git a/pkg/sql/sem/tree/eval_test.go b/pkg/sql/sem/tree/eval_test.go
index 88afa2ba2111..e2fd94aa6ca5 100644
--- a/pkg/sql/sem/tree/eval_test.go
+++ b/pkg/sql/sem/tree/eval_test.go
@@ -85,7 +85,7 @@ func optBuildScalar(evalCtx *tree.EvalContext, e tree.TypedExpr) (tree.TypedExpr
 		return nil, err
 	}
 
-	bld := execbuilder.New(nil /* factory */, o.Memo(), o.Memo().RootExpr(), evalCtx)
+	bld := execbuilder.New(nil /* factory */, o.Memo(), nil /* catalog */, o.Memo().RootExpr(), evalCtx)
 	ivh := tree.MakeIndexedVarHelper(nil /* container */, 0)
 
 	expr, err := bld.BuildScalar(&ivh)
diff --git a/pkg/sql/sem/tree/explain.go b/pkg/sql/sem/tree/explain.go
index 2dccebca78f4..467b1d2962fb 100644
--- a/pkg/sql/sem/tree/explain.go
+++ b/pkg/sql/sem/tree/explain.go
@@ -112,6 +112,7 @@ const (
 	ExplainFlagNoOptimize
 	ExplainFlagAnalyze
 	ExplainFlagEnv
+	ExplainFlagCatalog
 )
 
 var explainFlagStrings = map[string]int{
@@ -123,6 +124,7 @@ var explainFlagStrings = map[string]int{
 	"nooptimize":  ExplainFlagNoOptimize,
 	"analyze":     ExplainFlagAnalyze,
 	"env":         ExplainFlagEnv,
+	"catalog":     ExplainFlagCatalog,
 }
 
 // ParseOptions parses the options for an EXPLAIN statement.

From 39c2ec8378397c7ce6687da1a9a164bd539f4f2a Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Wed, 10 Jul 2019 15:28:31 -0400
Subject: [PATCH 2/2] opt: add Index.PartitionByListPrefixes to catalog

Add a catalog function that returns index regions of interest based on
PARTITION BY LIST values. This information will be used for an "index
skip scan".

Informs #38031.

Release note: None
---
 .../testdata/logic_test/partitioning          | 169 +++++++++++++++++-
 pkg/sql/opt/cat/index.go                      |  34 ++++
 pkg/sql/opt/cat/utils.go                      |  11 ++
 pkg/sql/opt/testutils/testcat/create_table.go |  17 +-
 pkg/sql/opt/testutils/testcat/test_catalog.go |  64 +++++++
 pkg/sql/opt/testutils/testcat/testdata/table  |  76 ++++++++
 pkg/sql/opt_catalog.go                        |  29 +++
 pkg/sql/sem/tree/create.go                    |  18 +-
 8 files changed, 402 insertions(+), 16 deletions(-)

diff --git a/pkg/ccl/logictestccl/testdata/logic_test/partitioning b/pkg/ccl/logictestccl/testdata/logic_test/partitioning
index 1ca7eaba4573..cf9dbf7e9329 100644
--- a/pkg/ccl/logictestccl/testdata/logic_test/partitioning
+++ b/pkg/ccl/logictestccl/testdata/logic_test/partitioning
@@ -1,4 +1,4 @@
-# LogicTest: local
+# LogicTest: local-opt
 
 statement error syntax
 CREATE TABLE t (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST ()
@@ -399,6 +399,21 @@ ok1  CREATE TABLE ok1 (
    PARTITION p2 VALUES IN ((2))
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok1
+----
+TABLE ok1
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      └── partition by list prefixes
+           ├── (1)
+           └── (2)
+scan ok1
+
 statement ok
 CREATE TABLE ok2 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST (a) (
     PARTITION p1 VALUES IN ((1)),
@@ -419,6 +434,21 @@ ok2  CREATE TABLE ok2 (
    PARTITION p2 VALUES IN ((2))
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok2
+----
+TABLE ok2
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      └── partition by list prefixes
+           ├── (1)
+           └── (2)
+scan ok2
+
 statement ok
 CREATE TABLE ok3 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST (a) (
     PARTITION p1 VALUES IN (1),
@@ -439,6 +469,20 @@ ok3  CREATE TABLE ok3 (
    PARTITION p2 VALUES IN ((DEFAULT))
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok3
+----
+TABLE ok3
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      └── partition by list prefixes
+           └── (1)
+scan ok3
+
 statement ok
 CREATE TABLE ok4 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST (a, b) (
     PARTITION p1 VALUES IN ((1, 1)),
@@ -463,6 +507,22 @@ ok4  CREATE TABLE ok4 (
    PARTITION p4 VALUES IN ((DEFAULT, DEFAULT))
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok4
+----
+TABLE ok4
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      └── partition by list prefixes
+           ├── (1, 1)
+           ├── (1)
+           └── (2, 3)
+scan ok4
+
 statement ok
 CREATE TABLE ok5 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST (a) (
     PARTITION p1 VALUES IN (1) PARTITION BY LIST (b) (
@@ -475,6 +535,21 @@ CREATE TABLE ok5 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST (a)
     PARTITION p3 VALUES IN (DEFAULT)
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok5
+----
+TABLE ok5
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      └── partition by list prefixes
+           ├── (1)
+           └── (2)
+scan ok5
+
 query TT
 SHOW CREATE TABLE ok5
 ----
@@ -515,6 +590,18 @@ ok6  CREATE TABLE ok6 (
    PARTITION p2 VALUES FROM (1) TO (2)
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok6
+----
+TABLE ok6
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      └── b int not null
+scan ok6
+
 statement ok
 CREATE TABLE ok7 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY RANGE (a) (
     PARTITION p1 VALUES FROM ((0)) TO (((1)))
@@ -533,6 +620,18 @@ ok7  CREATE TABLE ok7 (
    PARTITION p1 VALUES FROM (0) TO (1)
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok7
+----
+TABLE ok7
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      └── b int not null
+scan ok7
+
 statement ok
 CREATE TABLE ok8 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY RANGE (a) (
     PARTITION p1 VALUES FROM (MINVALUE) TO (1),
@@ -555,6 +654,18 @@ ok8  CREATE TABLE ok8 (
    PARTITION p3 VALUES FROM (2) TO (MAXVALUE)
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok8
+----
+TABLE ok8
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      └── b int not null
+scan ok8
+
 statement ok
 CREATE TABLE ok9 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY RANGE (a, b) (
     PARTITION p1 VALUES FROM (MINVALUE, MINVALUE) TO (1, MAXVALUE),
@@ -579,6 +690,18 @@ ok9  CREATE TABLE ok9 (
    PARTITION p4 VALUES FROM (3, MAXVALUE) TO (MAXVALUE, MAXVALUE)
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok9
+----
+TABLE ok9
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      └── b int not null
+scan ok9
+
 statement ok
 CREATE TABLE ok10 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY RANGE (a, b) (
     PARTITION p1 VALUES FROM (MINVALUE, MINVALUE) TO (1, 1),
@@ -605,6 +728,18 @@ ok10  CREATE TABLE ok10 (
    PARTITION p5 VALUES FROM (3, 4) TO (MAXVALUE, MAXVALUE)
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok10
+----
+TABLE ok10
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      └── b int not null
+scan ok10
+
 statement ok
 CREATE TABLE ok11 (a INT, b INT, c INT, PRIMARY KEY (a, b, c)) PARTITION BY LIST (a) (
     PARTITION p1 VALUES IN (1) PARTITION BY LIST (b) (
@@ -639,6 +774,22 @@ ok11  CREATE TABLE ok11 (
    )
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok11
+----
+TABLE ok11
+ ├── a int not null
+ ├── b int not null
+ ├── c int not null
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      ├── c int not null
+      └── partition by list prefixes
+           ├── (1)
+           └── (6)
+scan ok11
+
 statement ok
 CREATE TABLE IF NOT EXISTS ok12 (a INT, b INT, c INT, PRIMARY KEY (a, b)) PARTITION BY LIST (a) (
     PARTITION pu VALUES IN (NULL),
@@ -662,6 +813,22 @@ ok12  CREATE TABLE ok12 (
    PARTITION p2 VALUES IN ((2))
 )
 
+query T
+EXPLAIN (OPT, CATALOG) SELECT * from ok12
+----
+TABLE ok12
+ ├── a int not null
+ ├── b int not null
+ ├── c int
+ └── INDEX primary
+      ├── a int not null
+      ├── b int not null
+      └── partition by list prefixes
+           ├── (NULL)
+           ├── (1)
+           └── (2)
+scan ok12
+
 # Verify that creating a partition that includes NULL does not change the
 # implicit NOT NULL contrainst of a primary key.
 statement error null value in column "a" violates not-null constraint
diff --git a/pkg/sql/opt/cat/index.go b/pkg/sql/opt/cat/index.go
index 790fc24902d0..9509bf106a27 100644
--- a/pkg/sql/opt/cat/index.go
+++ b/pkg/sql/opt/cat/index.go
@@ -128,6 +128,40 @@ type Index interface {
 
 	// Span returns the KV span associated with the index.
 	Span() roachpb.Span
+
+	// PartitionByListPrefixes returns values that correspond to PARTITION BY LIST
+	// values. Specifically, it returns a list of tuples where each tuple contains
+	// values for a prefix of index columns (indicating a region of the index).
+	// Each tuple corresponds to a configured partition or subpartition.
+	//
+	// Note: this function decodes and allocates datums; use sparingly.
+	//
+	// Example:
+	//
+	// CREATE INDEX idx ON t(region,subregion,val) PARTITION BY LIST (region,subregion) (
+	//     PARTITION westcoast VALUES IN (('us', 'seattle'), ('us', 'cali')),
+	//     PARTITION us VALUES IN (('us', DEFAULT)),
+	//     PARTITION eu VALUES IN (('eu', DEFAULT)),
+	//     PARTITION default VALUES IN (DEFAULT)
+	// );
+	//
+	// PartitionByListPrefixes() returns
+	//  ('us', 'seattle'),
+	//  ('us', 'cali'),
+	//  ('us'),
+	//  ('eu').
+	//
+	// The intended use of this function is for index skip scans. Each tuple
+	// corresponds to a region of the index that we can constrain further. In the
+	// example above: if we have a val=1 filter, instead of a full index scan we
+	// can skip most of the data under /us/cali and /us/seattle by scanning spans:
+	//   [                 - /us/cali      )
+	//   [ /us/cali/1      - /us/cali/1    ]
+	//   [ /us/cali\x00    - /us/seattle   )
+	//   [ /us/seattle/1   - /us/seattle/1 ]
+	//   [ /us/seattle\x00 -               ]
+	//
+	PartitionByListPrefixes() []tree.Datums
 }
 
 // IndexColumn describes a single column that is part of an index definition.
diff --git a/pkg/sql/opt/cat/utils.go b/pkg/sql/opt/cat/utils.go
index 6ea38dffefb7..71dbba72c984 100644
--- a/pkg/sql/opt/cat/utils.go
+++ b/pkg/sql/opt/cat/utils.go
@@ -133,6 +133,9 @@ func FindTableColumnByName(tab Table, name tree.Name) int {
 // and testing.
 func FormatTable(cat Catalog, tab Table, tp treeprinter.Node) {
 	child := tp.Childf("TABLE %s", tab.Name().TableName)
+	if tab.IsVirtualTable() {
+		child.Child("virtual table")
+	}
 
 	var buf bytes.Buffer
 	for i := 0; i < tab.DeletableColumnCount(); i++ {
@@ -207,6 +210,14 @@ func formatCatalogIndex(tab Table, ord int, tp treeprinter.Node) {
 	}
 
 	FormatZone(idx.Zone(), child)
+
+	partPrefixes := idx.PartitionByListPrefixes()
+	if len(partPrefixes) != 0 {
+		c := child.Child("partition by list prefixes")
+		for i := range partPrefixes {
+			c.Child(partPrefixes[i].String())
+		}
+	}
 }
 
 // formatColPrefix returns a string representation of a list of columns. The
diff --git a/pkg/sql/opt/testutils/testcat/create_table.go b/pkg/sql/opt/testutils/testcat/create_table.go
index 20cb0a1c0167..66fcd79048b5 100644
--- a/pkg/sql/opt/testutils/testcat/create_table.go
+++ b/pkg/sql/opt/testutils/testcat/create_table.go
@@ -133,6 +133,12 @@ func (tc *Catalog) CreateTable(stmt *tree.CreateTable) *Table {
 	} else if !tab.IsVirtual {
 		tab.addPrimaryColumnIndex("rowid")
 	}
+	if stmt.PartitionBy != nil {
+		if len(tab.Indexes) == 0 {
+			panic("cannot partition virtual table")
+		}
+		tab.Indexes[0].partitionBy = stmt.PartitionBy
+	}
 
 	// Add check constraints.
 	for _, def := range stmt.Defs {
@@ -368,11 +374,12 @@ func (tt *Table) addColumn(def *tree.ColumnTableDef) {
 
 func (tt *Table) addIndex(def *tree.IndexTableDef, typ indexType) *Index {
 	idx := &Index{
-		IdxName:  tt.makeIndexName(def.Name, typ),
-		Unique:   typ != nonUniqueIndex,
-		Inverted: def.Inverted,
-		IdxZone:  &config.ZoneConfig{},
-		table:    tt,
+		IdxName:     tt.makeIndexName(def.Name, typ),
+		Unique:      typ != nonUniqueIndex,
+		Inverted:    def.Inverted,
+		IdxZone:     &config.ZoneConfig{},
+		table:       tt,
+		partitionBy: def.PartitionBy,
 	}
 
 	// Look for name suffixes indicating this is a mutation index.
diff --git a/pkg/sql/opt/testutils/testcat/test_catalog.go b/pkg/sql/opt/testutils/testcat/test_catalog.go
index 0061597cb178..7cbaa257cfc5 100644
--- a/pkg/sql/opt/testutils/testcat/test_catalog.go
+++ b/pkg/sql/opt/testutils/testcat/test_catalog.go
@@ -690,6 +690,10 @@ type Index struct {
 
 	// table is a back reference to the table this index is on.
 	table *Table
+
+	// partitionBy is the partitioning clause that corresponds to this index. Used
+	// to implement PartitionByListPrefixes.
+	partitionBy *tree.PartitionBy
 }
 
 // ID is part of the cat.Index interface.
@@ -752,6 +756,66 @@ func (ti *Index) Span() roachpb.Span {
 	panic("not implemented")
 }
 
+// PartitionByListPrefixes is part of the cat.Index interface.
+func (ti *Index) PartitionByListPrefixes() []tree.Datums {
+	p := ti.partitionBy
+	if p == nil {
+		return nil
+	}
+	if len(p.List) == 0 {
+		return nil
+	}
+	var res []tree.Datums
+	semaCtx := tree.MakeSemaContext()
+	evalCtx := tree.MakeTestingEvalContext(cluster.MakeTestingClusterSettings())
+	for i := range p.Fields {
+		if i >= len(ti.Columns) || p.Fields[i] != ti.Columns[i].ColName() {
+			panic("partition by columns must be a prefix of the index columns")
+		}
+	}
+	for i := range p.List {
+		// Exprs contains a list of values.
+		for _, e := range p.List[i].Exprs {
+			var vals []tree.Expr
+			switch t := e.(type) {
+			case *tree.Tuple:
+				vals = t.Exprs
+			default:
+				vals = []tree.Expr{e}
+			}
+
+			// Cut off at DEFAULT, if present.
+			for i := range vals {
+				if _, ok := vals[i].(tree.DefaultVal); ok {
+					vals = vals[:i]
+				}
+			}
+			if len(vals) == 0 {
+				continue
+			}
+			d := make(tree.Datums, len(vals))
+			for i := range vals {
+				c := tree.CastExpr{Expr: vals[i], Type: ti.Columns[i].DatumType()}
+				cTyped, err := c.TypeCheck(&semaCtx, nil)
+				if err != nil {
+					panic(err)
+				}
+				d[i], err = cTyped.Eval(&evalCtx)
+				if err != nil {
+					panic(err)
+				}
+			}
+
+			// TODO(radu): split into multiple prefixes if Subpartition is also by list.
+			// Note that this functionality should be kept in sync with the real catalog
+			// implementation (opt_catalog.go).
+
+			res = append(res, d)
+		}
+	}
+	return res
+}
+
 // Column implements the cat.Column interface for testing purposes.
 type Column struct {
 	Ordinal      int
diff --git a/pkg/sql/opt/testutils/testcat/testdata/table b/pkg/sql/opt/testutils/testcat/testdata/table
index 8f0b6f1897b1..8f0be78a2974 100644
--- a/pkg/sql/opt/testutils/testcat/testdata/table
+++ b/pkg/sql/opt/testutils/testcat/testdata/table
@@ -86,3 +86,79 @@ TABLE a
  └── INDEX a_a_key
       ├── a int
       └── rowid int not null default (unique_rowid()) [hidden] (storing)
+
+exec-ddl
+CREATE TABLE system.vtable (a INT, b INT)
+----
+
+exec-ddl
+SHOW CREATE system.vtable
+----
+TABLE vtable
+ ├── virtual table
+ ├── a int
+ └── b int
+
+exec-ddl
+CREATE TABLE part1 (a INT PRIMARY KEY, b INT) PARTITION BY LIST (a) (
+  PARTITION p1 VALUES IN (1),
+  PARTITION p2 VALUES IN (3, 4, 5),
+  PARTITION p3 VALUES IN (DEFAULT)
+)
+----
+
+exec-ddl
+SHOW CREATE part1
+----
+TABLE part1
+ ├── a int not null
+ ├── b int
+ └── INDEX primary
+      ├── a int not null
+      └── partition by list prefixes
+           ├── (1)
+           ├── (3)
+           ├── (4)
+           └── (5)
+
+exec-ddl
+CREATE TABLE part2 (
+  a STRING,
+  b STRING,
+  c INT,
+  PRIMARY KEY(a,b,c),
+  INDEX (c) PARTITION BY LIST (c) (
+    PARTITION pi1 VALUES IN (1),
+    PARTITION pi2 VALUES IN (3, 4)
+  )
+) PARTITION BY LIST (a, b) (
+  PARTITION p1 VALUES IN (('foo', 'bar'), ('foo', 'baz'), ('qux', 'qux')),
+  PARTITION p2 VALUES IN (('waldo', DEFAULT)),
+  PARTITION p3 VALUES IN (DEFAULT)
+)
+----
+
+exec-ddl
+SHOW CREATE part2
+----
+TABLE part2
+ ├── a string not null
+ ├── b string not null
+ ├── c int not null
+ ├── INDEX primary
+ │    ├── a string not null
+ │    ├── b string not null
+ │    ├── c int not null
+ │    └── partition by list prefixes
+ │         ├── ('foo', 'bar')
+ │         ├── ('foo', 'baz')
+ │         ├── ('qux', 'qux')
+ │         └── ('waldo')
+ └── INDEX secondary
+      ├── c int not null
+      ├── a string not null
+      ├── b string not null
+      └── partition by list prefixes
+           ├── (1)
+           ├── (3)
+           └── (4)
diff --git a/pkg/sql/opt_catalog.go b/pkg/sql/opt_catalog.go
index a7bbf6ffdd6c..12077df57d0c 100644
--- a/pkg/sql/opt_catalog.go
+++ b/pkg/sql/opt_catalog.go
@@ -1015,6 +1015,35 @@ func (oi *optIndex) Ordinal() int {
 	return oi.indexOrdinal
 }
 
+// PartitionByListPrefixes is part of the cat.Index interface.
+func (oi *optIndex) PartitionByListPrefixes() []tree.Datums {
+	list := oi.desc.Partitioning.List
+	if len(list) == 0 {
+		return nil
+	}
+	res := make([]tree.Datums, 0, len(list))
+	var a sqlbase.DatumAlloc
+	for i := range list {
+		for _, valueEncBuf := range list[i].Values {
+			t, _, err := sqlbase.DecodePartitionTuple(
+				&a, &oi.tab.desc.TableDescriptor, oi.desc, &oi.desc.Partitioning,
+				valueEncBuf, nil, /* prefixDatums */
+			)
+			if err != nil {
+				panic(errors.NewAssertionErrorWithWrappedErrf(err, "while decoding partition tuple"))
+			}
+			// Ignore the DEFAULT case, where there is nothing to return.
+			if len(t.Datums) > 0 {
+				res = append(res, t.Datums)
+			}
+			// TODO(radu): split into multiple prefixes if Subpartition is also by list.
+			// Note that this functionality should be kept in sync with the test catalog
+			// implementation (test_catalog.go).
+		}
+	}
+	return res
+}
+
 type optTableStat struct {
 	createdAt      time.Time
 	columnOrdinals []int
diff --git a/pkg/sql/sem/tree/create.go b/pkg/sql/sem/tree/create.go
index 35006c17577c..f37c40c8ac96 100644
--- a/pkg/sql/sem/tree/create.go
+++ b/pkg/sql/sem/tree/create.go
@@ -156,9 +156,11 @@ type TableDef interface {
 	SetName(name Name)
 }
 
-func (*ColumnTableDef) tableDef() {}
-func (*IndexTableDef) tableDef()  {}
-func (*FamilyTableDef) tableDef() {}
+func (*ColumnTableDef) tableDef()               {}
+func (*IndexTableDef) tableDef()                {}
+func (*FamilyTableDef) tableDef()               {}
+func (*ForeignKeyConstraintTableDef) tableDef() {}
+func (*CheckConstraintTableDef) tableDef()      {}
 
 // TableDefs represents a list of table definitions.
 type TableDefs []TableDef
@@ -569,7 +571,9 @@ type ConstraintTableDef interface {
 	constraintTableDef()
 }
 
-func (*UniqueConstraintTableDef) constraintTableDef() {}
+func (*UniqueConstraintTableDef) constraintTableDef()     {}
+func (*ForeignKeyConstraintTableDef) constraintTableDef() {}
+func (*CheckConstraintTableDef) constraintTableDef()      {}
 
 // UniqueConstraintTableDef represents a unique constraint within a CREATE
 // TABLE statement.
@@ -715,12 +719,6 @@ func (node *ForeignKeyConstraintTableDef) SetName(name Name) {
 	node.Name = name
 }
 
-func (*ForeignKeyConstraintTableDef) tableDef()           {}
-func (*ForeignKeyConstraintTableDef) constraintTableDef() {}
-
-func (*CheckConstraintTableDef) tableDef()           {}
-func (*CheckConstraintTableDef) constraintTableDef() {}
-
 // CheckConstraintTableDef represents a check constraint within a CREATE
 // TABLE statement.
 type CheckConstraintTableDef struct {