[gen4]: small refactoring around Compact

go/vt/vtgate/planbuilder/gen4_planner.go

@@ -206,6 +206,10 @@ func newBuildSelectPlan(
 	if err != nil {
 		return nil, nil, err
 	}
+	logical, err = operators.Compact(ctx, logical)
+	if err != nil {
+		return nil, nil, err
+	}
 	err = operators.CheckValid(logical)
 	if err != nil {
 		return nil, nil, err

go/vt/vtgate/planbuilder/operators/compact.go

@@ -23,34 +23,39 @@ import (
 
 // Compact will optimise the operator tree into a smaller but equivalent version
 func Compact(ctx *plancontext.PlanningContext, op Operator) (Operator, error) {
-	switch op := op.(type) {
-	case *Union:
-		compactConcatenate(op)
-	case *Filter:
-		if len(op.Predicates) == 0 {
-			return op.Source, nil
+	newOp, _, err := rewriteBottomUp(ctx, op, func(ctx *plancontext.PlanningContext, op Operator) (Operator, bool, error) {
+		newOp, ok := op.(compactable)
+		if !ok {
+			return op, false, nil
 		}
-	case *Join:
-		return compactJoin(ctx, op)
-	}
+		return newOp.compact(ctx)
+	})
+	return newOp, err
+}
 
-	return op, nil
+func (f *Filter) compact(*plancontext.PlanningContext) (Operator, bool, error) {
+	if len(f.Predicates) == 0 {
+		return f.Source, true, nil
+	}
+	return f, false, nil
 }
 
-func compactConcatenate(op *Union) {
+func (u *Union) compact(*plancontext.PlanningContext) (Operator, bool, error) {
 	var newSources []Operator
 	var newSels []*sqlparser.Select
-	for i, source := range op.Sources {
-		other, isConcat := source.(*Union)
-		if !isConcat {
+	anythingChanged := false
+	for i, source := range u.Sources {
+		other, isUnion := source.(*Union)
+		if !isUnion {
 			newSources = append(newSources, source)
-			newSels = append(newSels, op.SelectStmts[i])
+			newSels = append(newSels, u.SelectStmts[i])
 			continue
 		}
+		anythingChanged = true
 		switch {
 		case len(other.Ordering) == 0 && !other.Distinct:
 			fallthrough
-		case op.Distinct:
+		case u.Distinct:
 			// if the current UNION is a DISTINCT, we can safely ignore everything from children UNIONs, except LIMIT
 			newSources = append(newSources, other.Sources...)
 			newSels = append(newSels, other.SelectStmts...)
@@ -60,30 +65,33 @@ func compactConcatenate(op *Union) {
 			newSels = append(newSels, nil)
 		}
 	}
-	op.Sources = newSources
-	op.SelectStmts = newSels
+	if anythingChanged {
+		u.Sources = newSources
+		u.SelectStmts = newSels
+	}
+	return u, anythingChanged, nil
 }
 
-func compactJoin(ctx *plancontext.PlanningContext, op *Join) (Operator, error) {
-	if op.LeftJoin {
+func (j *Join) compactJoin(ctx *plancontext.PlanningContext) (Operator, bool, error) {
+	if j.LeftJoin {
 		// we can't merge outer joins into a single QG
-		return op, nil
+		return j, false, nil
 	}
 
-	lqg, lok := op.LHS.(*QueryGraph)
-	rqg, rok := op.RHS.(*QueryGraph)
+	lqg, lok := j.LHS.(*QueryGraph)
+	rqg, rok := j.RHS.(*QueryGraph)
 	if !lok || !rok {
-		return op, nil
+		return j, false, nil
 	}
 
 	newOp := &QueryGraph{
 		Tables:     append(lqg.Tables, rqg.Tables...),
 		innerJoins: append(lqg.innerJoins, rqg.innerJoins...),
 		NoDeps:     sqlparser.AndExpressions(lqg.NoDeps, rqg.NoDeps),
 	}
-	err := newOp.collectPredicate(ctx, op.Predicate)
+	err := newOp.collectPredicate(ctx, j.Predicate)
 	if err != nil {
-		return nil, err
+		return nil, false, err
 	}
-	return newOp, nil
+	return newOp, true, nil
 }

go/vt/vtgate/planbuilder/operators/node.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 
 	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 )
 
@@ -90,3 +91,31 @@ func checkSize(inputs []Operator, shouldBe int) {
 		panic(fmt.Sprintf("BUG: got the wrong number of inputs: got %d, expected %d", len(inputs), shouldBe))
 	}
 }
+
+type rewriterFunc func(*plancontext.PlanningContext, Operator) (newOp Operator, changed bool, err error)
+
+func rewriteBottomUp(ctx *plancontext.PlanningContext, root Operator, rewriter rewriterFunc) (Operator, bool, error) {
+	oldInputs := root.Inputs()
+	anythingChanged := false
+	newInputs := make([]Operator, len(oldInputs))
+	for i, operator := range oldInputs {
+		in, changed, err := rewriteBottomUp(ctx, operator, rewriter)
+		if err != nil {
+			return nil, false, err
+		}
+		if changed {
+			anythingChanged = true
+		}
+		newInputs[i] = in
+	}
+
+	if anythingChanged {
+		root = root.Clone(newInputs)
+	}
+
+	newOp, b, err := rewriter(ctx, root)
+	if err != nil {
+		return nil, false, err
+	}
+	return newOp, anythingChanged || b, nil
+}

go/vt/vtgate/planbuilder/operators/operator.go

@@ -68,6 +68,11 @@ type (
 		CheckValid() error
 	}
 
+	compactable interface {
+		// implement this interface for operators that have easy to see optimisations
+		compact(ctx *plancontext.PlanningContext) (Operator, bool, error)
+	}
+
 	// helper type that implements Inputs() returning nil
 	noInputs struct{}
 )
@@ -209,7 +214,7 @@ func CreateLogicalOperatorFromAST(ctx *plancontext.PlanningContext, selStmt sqlp
 	if err != nil {
 		return nil, err
 	}
-	return Compact(ctx, op)
+	return op, nil
 }
 
 func createOperatorFromUnion(ctx *plancontext.PlanningContext, node *sqlparser.Union) (Operator, error) {

go/vt/vtgate/planbuilder/operators/operator_test.go

@@ -105,6 +105,8 @@ func TestOperator(t *testing.T) {
 			ctx := plancontext.NewPlanningContext(nil, semTable, nil, 0)
 			optree, err := CreateLogicalOperatorFromAST(ctx, stmt)
 			require.NoError(t, err)
+			optree, err = Compact(ctx, optree)
+			require.NoError(t, err)
 			output := testString(optree)
 			assert.Equal(t, tc.expected, output)
 			if t.Failed() {