cache sql.Schema in plan.Project and plan.TableAlias (#3158)

Author: jycor

memory/index.go

@@ -298,8 +298,9 @@ func (idx *Index) Reversible() bool {
 	return true
 }
 
-func (idx Index) copy() *Index {
-	return &idx
+func (idx *Index) copy() *Index {
+	newIdx := *idx
+	return &newIdx
 }
 
 // columnIndexes returns the indexes in the given schema for the fields in this index

sql/analyzer/common_test.go

@@ -157,6 +157,8 @@ func runTestCases(t *testing.T, ctx *sql.Context, testCases []analyzerFnTestCase
 			if expected == nil {
 				expected = tt.node
 			}
+			// Schema of certain nodes aren't filled until needed
+			expected.Schema()
 
 			assertNodesEqualWithDiff(t, expected, result)
 		})

sql/analyzer/index_analyzer_test.go

@@ -131,17 +131,18 @@ type dummyIdx struct {
 
 var _ sql.Index = (*dummyIdx)(nil)
 
-func (i dummyIdx) CanSupport(context *sql.Context, r ...sql.Range) bool {
+func (i *dummyIdx) CanSupport(context *sql.Context, r ...sql.Range) bool {
 	return true
 }
 
-func (i dummyIdx) Expressions() []string {
+func (i *dummyIdx) Expressions() []string {
 	var exprs []string
 	for _, e := range i.expr {
 		exprs = append(exprs, e.String())
 	}
 	return exprs
 }
+
 func (i *dummyIdx) ID() string                            { return i.id }
 func (i *dummyIdx) Database() string                      { return i.database }
 func (i *dummyIdx) Table() string                         { return i.table }

sql/index_registry_test.go

@@ -443,6 +443,7 @@ func (i dummyIdx) Expressions() []string {
 	}
 	return exprs
 }
+
 func (i dummyIdx) ID() string              { return i.id }
 func (i dummyIdx) Database() string        { return i.database }
 func (i dummyIdx) Table() string           { return i.table }

sql/memo/rel_props.go

@@ -131,13 +131,13 @@ func newRelProps(rel RelExpr) *relProps {
 }
 
 // idxExprsColumns returns the column names used in an index's expressions.
-// TODO: this is unstable as long as periods in Index.Expressions()
-// identifiers are ambiguous.
+// Identifiers are ambiguous.
 func idxExprsColumns(idx sql.Index) []string {
-	columns := make([]string, len(idx.Expressions()))
-	for i, e := range idx.Expressions() {
-		parts := strings.Split(e, ".")
-		columns[i] = strings.ToLower(parts[1])
+	exprs := idx.Expressions()
+	columns := make([]string, len(exprs))
+	for i, e := range exprs {
+		colName := e[strings.IndexRune(e, '.')+1:]
+		columns[i] = strings.ToLower(colName)
 	}
 	return columns
 }
@@ -793,15 +793,12 @@ func sortedColsForRel(rel RelExpr) sql.Schema {
 		var ret sql.Schema
 		for _, e := range r.InnerScan.Table.Index().Expressions() {
 			// TODO columns can have "." characters, this will miss cases
-			parts := strings.Split(e, ".")
-			var name string
-			if len(parts) == 2 {
-				name = parts[1]
-			} else {
+			idx := strings.IndexRune(e, '.')
+			if idx == -1 {
 				return nil
 			}
 			ret = append(ret, &sql.Column{
-				Name:     strings.ToLower(name),
+				Name:     strings.ToLower(e[idx+1:]),
 				Source:   strings.ToLower(r.InnerScan.Table.Name()),
 				Nullable: true},
 			)

sql/plan/project.go

@@ -35,6 +35,8 @@ type Project struct {
 	// a RowIter.
 	IncludesNestedIters bool
 	deps                sql.ColSet
+
+	sch sql.Schema
 }
 
 var _ sql.Expressioner = (*Project)(nil)
@@ -122,14 +124,16 @@ func ExprDeps(exprs ...sql.Expression) sql.ColSet {
 
 // Schema implements the Node interface.
 func (p *Project) Schema() sql.Schema {
-	var s = make(sql.Schema, len(p.Projections))
-	for i, expr := range p.Projections {
-		s[i] = transform.ExpressionToColumn(expr, AliasSubqueryString(expr))
-		if gf := unwrapGetField(expr); gf != nil {
-			s[i].Default = findDefault(p.Child, gf)
+	if p.sch == nil {
+		p.sch = make(sql.Schema, len(p.Projections))
+		for i, expr := range p.Projections {
+			p.sch[i] = transform.ExpressionToColumn(expr, AliasSubqueryString(expr))
+			if gf := unwrapGetField(expr); gf != nil {
+				p.sch[i].Default = findDefault(p.Child, gf)
+			}
 		}
 	}
-	return s
+	return p.sch
 }
 
 // Resolved implements the Resolvable interface.
@@ -190,6 +194,7 @@ func (p *Project) WithChildren(children ...sql.Node) (sql.Node, error) {
 	}
 	np := *p
 	np.Child = children[0]
+	np.sch = nil
 	return &np, nil
 }
 
@@ -205,6 +210,7 @@ func (p *Project) WithExpressions(exprs ...sql.Expression) (sql.Node, error) {
 	}
 	np := *p
 	np.Projections = exprs
+	np.sch = nil
 	return &np, nil
 }
 

sql/plan/tablealias.go

@@ -25,6 +25,7 @@ type TableAlias struct {
 	comment string
 	id      sql.TableId
 	cols    sql.ColSet
+	sch     sql.Schema
 }
 
 var _ sql.RenameableNode = (*TableAlias)(nil)
@@ -33,7 +34,10 @@ var _ sql.CollationCoercible = (*TableAlias)(nil)
 
 // NewTableAlias returns a new Table alias node.
 func NewTableAlias(name string, node sql.Node) *TableAlias {
-	ret := &TableAlias{UnaryNode: &UnaryNode{Child: node}, name: name}
+	ret := &TableAlias{
+		UnaryNode: &UnaryNode{Child: node},
+		name:      name,
+	}
 	if tin, ok := node.(TableIdNode); ok {
 		ret.id = tin.Id()
 		ret.cols = tin.Columns()
@@ -87,14 +91,16 @@ func (t *TableAlias) Comment() string {
 // Schema implements the Node interface. TableAlias alters the schema of its child element to rename the source of
 // columns to the alias.
 func (t *TableAlias) Schema() sql.Schema {
-	childSchema := t.Child.Schema()
-	copy := make(sql.Schema, len(childSchema))
-	for i, col := range childSchema {
-		colCopy := *col
-		colCopy.Source = t.name
-		copy[i] = &colCopy
+	if t.sch == nil {
+		childSchema := t.Child.Schema()
+		t.sch = make(sql.Schema, len(childSchema))
+		for i, col := range childSchema {
+			newCol := *col
+			newCol.Source = t.name
+			t.sch[i] = &newCol
+		}
 	}
-	return copy
+	return t.sch
 }
 
 // WithChildren implements the Node interface.
@@ -118,21 +124,22 @@ func (t *TableAlias) CollationCoercibility(ctx *sql.Context) (collation sql.Coll
 	return sql.Collation_binary, 7
 }
 
-func (t TableAlias) String() string {
+func (t *TableAlias) String() string {
 	pr := sql.NewTreePrinter()
 	_ = pr.WriteNode("TableAlias(%s)", t.name)
 	_ = pr.WriteChildren(t.Child.String())
 	return pr.String()
 }
 
-func (t TableAlias) DebugString() string {
+func (t *TableAlias) DebugString() string {
 	pr := sql.NewTreePrinter()
 	_ = pr.WriteNode("TableAlias(%s)", t.name)
 	_ = pr.WriteChildren(sql.DebugString(t.Child))
 	return pr.String()
 }
 
-func (t TableAlias) WithName(name string) sql.Node {
-	t.name = name
-	return &t
+func (t *TableAlias) WithName(name string) sql.Node {
+	nt := *t
+	nt.name = name
+	return &nt
 }

sql/planbuilder/ddl.go

@@ -395,11 +395,11 @@ func (b *Builder) getIndexDefs(table sql.Table) sql.IndexDefs {
 	if !isIdxTbl {
 		return nil
 	}
-	var idxDefs sql.IndexDefs
 	idxs, err := idxTbl.GetIndexes(b.ctx)
 	if err != nil {
 		b.handleErr(err)
 	}
+	idxDefs := make(sql.IndexDefs, 0, len(idxs))
 	for _, idx := range idxs {
 		if idx.IsGenerated() {
 			continue
@@ -412,10 +412,10 @@ func (b *Builder) getIndexDefs(table sql.Table) sql.IndexDefs {
 				constraint = sql.IndexConstraint_Unique
 			}
 		}
-		columns := make([]sql.IndexColumn, len(idx.Expressions()))
-		for i, col := range idx.Expressions() {
-			// TODO: find a better way to get only the column name if the table is present
-			col = strings.TrimPrefix(col, idxTbl.Name()+".")
+		exprs := idx.Expressions()
+		columns := make([]sql.IndexColumn, len(exprs))
+		for i, col := range exprs {
+			col = col[strings.IndexByte(col, '.')+1:]
 			columns[i] = sql.IndexColumn{Name: col}
 		}
 		idxDefs = append(idxDefs, &sql.IndexDef{