Horizon planning on operators

go/vt/vtgate/planbuilder/gen4_planner.go

@@ -497,7 +497,7 @@ func planHorizon(ctx *plancontext.PlanningContext, plan logicalPlan, in sqlparse
 			return nil, ctx.SemTable.NotSingleRouteErr
 		}
 		if isRoute && rb.isSingleShard() {
-			err = planSingleShardRoutePlan(node, rb)
+			err = planSingleRoutePlan(node, rb)
 		} else {
 			plan, err = planOrderByOnUnion(ctx, plan, node)
 		}

go/vt/vtgate/planbuilder/horizon_planning.go

@@ -44,7 +44,7 @@ func (hp *horizonPlanning) planHorizon(ctx *plancontext.PlanningContext, plan lo
 	}
 
 	if isRoute && rb.isSingleShard() {
-		err := planSingleShardRoutePlan(hp.sel, rb)
+		err := planSingleRoutePlan(hp.sel, rb)
 		if err != nil {
 			return nil, err
 		}
@@ -88,7 +88,7 @@ func (hp *horizonPlanning) planHorizon(ctx *plancontext.PlanningContext, plan lo
 		// if we already did sorting, we don't need to do it again
 		needsOrdering = needsOrdering && !hp.qp.CanPushDownSorting
 	case canShortcut:
-		err = planSingleShardRoutePlan(hp.sel, rb)
+		err = planSingleRoutePlan(hp.sel, rb)
 		if err != nil {
 			return nil, err
 		}
@@ -1097,7 +1097,7 @@ func exprHasVindex(semTable *semantics.SemTable, expr sqlparser.Expr, hasToBeUni
 	return false
 }
 
-func planSingleShardRoutePlan(sel sqlparser.SelectStatement, rb *routeGen4) error {
+func planSingleRoutePlan(sel sqlparser.SelectStatement, rb *routeGen4) error {
 	err := stripDownQuery(sel, rb.Select)
 	if err != nil {
 		return err

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -114,7 +114,7 @@ func transformHorizon(ctx *plancontext.PlanningContext, op *operators.Horizon, i
 		}
 		var plan logicalPlan
 		if isRoute && rb.isSingleShard() {
-			err = planSingleShardRoutePlan(node, rb)
+			err = planSingleRoutePlan(node, rb)
 			plan = rb
 		} else {
 			plan, err = planOrderByOnUnion(ctx, source, node)

go/vt/vtgate/planbuilder/operators/horizon_planning.go

@@ -17,23 +17,25 @@ limitations under the License.
 package operators
 
 import (
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/rewrite"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/ops"
 )
 
 var errNotHorizonPlanned = vterrors.VT12001("query cannot be fully operator planned")
 
-func planHorizons(in ops.Operator) (ops.Operator, error) {
+func planHorizons(ctx *plancontext.PlanningContext, in ops.Operator) (ops.Operator, error) {
 	return rewrite.TopDown(in, func(in ops.Operator) (ops.Operator, rewrite.TreeIdentity, rewrite.VisitRule, error) {
 		switch in := in.(type) {
 		case *Horizon:
-			op, err := planHorizon(in)
+			op, visit, err := planHorizon(ctx, in)
 			if err != nil {
 				return nil, rewrite.SameTree, rewrite.SkipChildren, err
 			}
-			return op, rewrite.NewTree, rewrite.VisitChildren, nil
+			return op, rewrite.NewTree, visit, nil
 		case *Route:
 			return in, rewrite.SameTree, rewrite.SkipChildren, nil
 		default:
@@ -42,18 +44,41 @@ func planHorizons(in ops.Operator) (ops.Operator, error) {
 	})
 }
 
-func planHorizon(in *Horizon) (ops.Operator, error) {
+func planHorizon(ctx *plancontext.PlanningContext, in *Horizon) (ops.Operator, rewrite.VisitRule, error) {
 	rb, isRoute := in.Source.(*Route)
 	if !isRoute {
-		return in, nil
+		return in, rewrite.VisitChildren, nil
 	}
 	if isRoute && rb.IsSingleShard() && in.Select.GetLimit() == nil {
-		return planSingleShardRoute(rb, in)
+		return planSingleRoute(rb, in)
 	}
 
-	return nil, errNotHorizonPlanned
+	sel, isSel := in.Select.(*sqlparser.Select)
+	if !isSel {
+		return nil, rewrite.VisitChildren, errNotHorizonPlanned
+	}
+
+	qp, err := CreateQPFromSelect(ctx, sel)
+	if err != nil {
+		return nil, rewrite.VisitChildren, err
+	}
+
+	needsOrdering := len(qp.OrderExprs) > 0
+	canShortcut := isRoute && sel.Having == nil && !needsOrdering
+
+	// If we still have a HAVING clause, it's because it could not be pushed to the WHERE,
+	// so it probably has aggregations
+	switch {
+	case qp.NeedsAggregation() || sel.Having != nil:
+		return nil, rewrite.VisitChildren, errNotHorizonPlanned
+	case canShortcut && !needsOrdering && !qp.NeedsDistinct() && in.Select.GetLimit() == nil:
+		return planSingleRoute(rb, in)
+	default:
+		return nil, rewrite.VisitChildren, errNotHorizonPlanned
+	}
 }
-func planSingleShardRoute(rb *Route, horizon *Horizon) (ops.Operator, error) {
+
+func planSingleRoute(rb *Route, horizon *Horizon) (ops.Operator, rewrite.VisitRule, error) {
 	rb.Source, horizon.Source = horizon, rb.Source
-	return rb, nil
+	return rb, rewrite.SkipChildren, nil
 }

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

@@ -68,7 +68,7 @@ func PlanQuery(ctx *plancontext.PlanningContext, selStmt sqlparser.Statement) (o
 
 	backup := Clone(op)
 
-	op, err = planHorizons(op)
+	op, err = planHorizons(ctx, op)
 	if err == errNotHorizonPlanned {
 		op = backup
 	} else if err != nil {