[binder] hoist projections in certain cases where we can combine with top-level projection

Author: max-hoffman

enginetest/queries/imdb_plans.go

@@ -28,6 +28,7 @@ type Project struct {
 	UnaryNode
 	Projections []sql.Expression
 	CanDefer    bool
+	deps        sql.ColSet
 }
 
 var _ sql.Expressioner = (*Project)(nil)
@@ -94,6 +95,23 @@ func unwrapGetField(expr sql.Expression) *expression.GetField {
 	}
 }
 
+func ExprDeps(exprs []sql.Expression) sql.ColSet {
+	var deps sql.ColSet
+	for _, e := range exprs {
+		sql.Inspect(e, func(e sql.Expression) bool {
+			switch e := e.(type) {
+			case sql.IdExpression:
+				deps.Add(e.Id())
+			case *Subquery:
+				deps.Union(e.Correlated())
+			default:
+			}
+			return true
+		})
+	}
+	return deps
+}
+
 // Schema implements the Node interface.
 func (p *Project) Schema() sql.Schema {
 	var s = make(sql.Schema, len(p.Projections))

enginetest/queries/integration_plans.go

@@ -220,3 +220,28 @@ func (f *factory) buildDistinct(child sql.Node) sql.Node {
 	}
 	return plan.NewDistinct(child)
 }
+
+func (f *factory) buildSort(child sql.Node, exprs []sql.SortField, deps sql.ColSet, subquery bool) (sql.Node, error) {
+	{
+		// if Sort->Projection with no alias dependency, hoist projection
+		// for opportunity to squash with top-level projection
+		// sort -> proj -> child
+		// =>
+		// proj -> sort -> child
+		if p, ok := child.(*plan.Project); ok {
+			var aliases []sql.Expression
+			for _, p := range p.Projections {
+				if _, ok := p.(*expression.Alias); ok {
+					aliases = append(aliases, p)
+				}
+			}
+			aliasCols := plan.ExprDeps(aliases)
+			if !aliasCols.Intersects(deps) {
+				newP := plan.NewProject(p.Projections, plan.NewSort(exprs, p.Child))
+				return f.buildProject(newP, subquery)
+			}
+		}
+	}
+
+	return plan.NewSort(exprs, child), nil
+}

enginetest/queries/query_plans.go

@@ -22,7 +22,6 @@ import (
 
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/expression"
-	"github.com/dolthub/go-mysql-server/sql/plan"
 	"github.com/dolthub/go-mysql-server/sql/transform"
 	"github.com/dolthub/go-mysql-server/sql/types"
 )
@@ -199,6 +198,7 @@ func (b *Builder) buildOrderBy(inScope, orderByScope *scope) {
 		return
 	}
 	var sortFields sql.SortFields
+	var deps sql.ColSet
 	for _, c := range orderByScope.cols {
 		so := sql.Ascending
 		if c.descending {
@@ -213,8 +213,12 @@ func (b *Builder) buildOrderBy(inScope, orderByScope *scope) {
 			Order:  so,
 		}
 		sortFields = append(sortFields, sf)
+		deps.Add(sql.ColumnId(c.id))
+	}
+	sort, err := b.f.buildSort(inScope.node, sortFields, deps, inScope.refsSubquery)
+	if err != nil {
+		b.handleErr(err)
 	}
-	sort := plan.NewSort(sortFields, inScope.node)
 	inScope.node = sort
 	return
 }