refactor: change the planner to use the new AST

Signed-off-by: Andres Taylor <andres@planetscale.com>

go/vt/schemadiff/view.go

@@ -427,6 +427,6 @@ func (c *CreateViewEntity) identicalOtherThanName(other *CreateViewEntity) bool
 		c.IsReplace == other.IsReplace &&
 		sqlparser.Equals.RefOfDefiner(c.Definer, other.Definer) &&
 		sqlparser.Equals.Columns(c.Columns, other.Columns) &&
-		sqlparser.Equals.SelectStatement(c.Select, other.Select) &&
+		sqlparser.Equals.Statement(c.Select, other.Select) &&
 		sqlparser.Equals.RefOfParsedComments(c.Comments, other.Comments)
 }

go/vt/sqlparser/ast.go

@@ -82,6 +82,11 @@ type (
 		SetWith(with *With)
 	}
 
+	Distinctable interface {
+		MakeDistinct()
+		IsDistinct() bool
+	}
+
 	// SelectStatement any SELECT statement.
 	SelectStatement interface {
 		Statement
@@ -90,12 +95,11 @@ type (
 		Commented
 		ColumnResults
 		Withable
+		Distinctable
 		iSelectStatement()
 		GetLock() Lock
 		SetLock(lock Lock)
 		SetInto(into *SelectInto)
-		MakeDistinct()
-		IsDistinct() bool
 	}
 
 	// DDLStatement represents any DDL Statement

go/vt/sqlparser/ast_rewriting.go

@@ -532,7 +532,7 @@ func (er *astRewriter) existsRewrite(cursor *Cursor, node *ExistsExpr) {
 	sel.GroupBy = nil
 }
 
-// rewriteDistinctableAggr removed Distinct from Max and Min Aggregations as it does not impact the result. But, makes the plan simpler.
+// rewriteDistinctableAggr removed Distinctable from Max and Min Aggregations as it does not impact the result. But, makes the plan simpler.
 func (er *astRewriter) rewriteDistinctableAggr(cursor *Cursor, node DistinctableAggr) {
 	if !node.IsDistinct() {
 		return

go/vt/vtgate/planbuilder/ddl.go

@@ -196,7 +196,7 @@ func buildCreateViewCommon(
 	vschema plancontext.VSchema,
 	reservedVars *sqlparser.ReservedVars,
 	cfg dynamicconfig.DDL,
-	ddlSelect sqlparser.SelectStatement,
+	ddlSelect sqlparser.TableSubquery,
 	ddl sqlparser.DDLStatement,
 ) (key.Destination, *vindexes.Keyspace, error) {
 	// For Create View, we require that the keyspace exist and the select query can be satisfied within the keyspace itself

go/vt/vtgate/planbuilder/operators/SQL_builder.go

@@ -37,8 +37,9 @@ type (
 	}
 )
 
-func (qb *queryBuilder) asSelectStatement() sqlparser.SelectStatement {
-	return qb.stmt.(sqlparser.SelectStatement)
+func (qb *queryBuilder) asSelectStatement() sqlparser.TableSubquery {
+	return qb.stmt.(sqlparser.TableSubquery)
+
 }
 func (qb *queryBuilder) asOrderAndLimit() sqlparser.OrderAndLimit {
 	return qb.stmt.(sqlparser.OrderAndLimit)
@@ -191,7 +192,8 @@ func (qb *queryBuilder) pushUnionInsideDerived() {
 			As:   sqlparser.NewIdentifierCS("dt"),
 		}},
 	}
-	sel.SelectExprs = unionSelects(sqlparser.GetFirstSelect(selStmt).SelectExprs)
+	firstSelect := getFirstSelect(selStmt)
+	sel.SelectExprs = unionSelects(firstSelect.SelectExprs)
 	qb.stmt = sel
 }
 
@@ -208,9 +210,10 @@ func unionSelects(exprs sqlparser.SelectExprs) (selectExprs sqlparser.SelectExpr
 	return
 }
 
-func checkUnionColumnByName(column *sqlparser.ColName, sel sqlparser.SelectStatement) {
+func checkUnionColumnByName(column *sqlparser.ColName, sel sqlparser.TableSubquery) {
 	colName := column.Name.String()
-	exprs := sqlparser.GetFirstSelect(sel).SelectExprs
+	firstSelect := getFirstSelect(sel)
+	exprs := firstSelect.SelectExprs
 	offset := slices.IndexFunc(exprs, func(expr sqlparser.SelectExpr) bool {
 		switch ae := expr.(type) {
 		case *sqlparser.StarExpr:
@@ -244,8 +247,8 @@ func (qb *queryBuilder) unionWith(other *queryBuilder, distinct bool) {
 
 func (qb *queryBuilder) recursiveCteWith(other *queryBuilder, name, alias string, distinct bool, columns sqlparser.Columns) {
 	cteUnion := &sqlparser.Union{
-		Left:     qb.stmt.(sqlparser.SelectStatement),
-		Right:    other.stmt.(sqlparser.SelectStatement),
+		Left:     qb.stmt.(sqlparser.TableSubquery),
+		Right:    other.stmt.(sqlparser.TableSubquery),
 		Distinct: distinct,
 	}
 
@@ -393,7 +396,7 @@ func removeKeyspaceFromSelectExpr(expr sqlparser.SelectExpr) {
 	}
 }
 
-func stripDownQuery(from, to sqlparser.SelectStatement) {
+func stripDownQuery(from, to sqlparser.TableSubquery) {
 	switch node := from.(type) {
 	case *sqlparser.Select:
 		toNode, ok := to.(*sqlparser.Select)
@@ -450,7 +453,12 @@ func buildQuery(op Operator, qb *queryBuilder) {
 		buildUnion(op, qb)
 	case *Distinct:
 		buildQuery(op.Source, qb)
-		qb.asSelectStatement().MakeDistinct()
+		statement := qb.asSelectStatement()
+		d, ok := statement.(sqlparser.Distinctable)
+		if !ok {
+			panic(vterrors.VT13001("expected a select statement with distinct"))
+		}
+		d.MakeDistinct()
 	case *Update:
 		buildUpdate(op, qb)
 	case *Delete:

go/vt/vtgate/planbuilder/operators/aggregation_pushing.go

@@ -544,7 +544,7 @@ func splitAggrColumnsToLeftAndRight(
 
 	canPushDistinctAggr, distinctExprs := checkIfWeCanPush(ctx, aggregator)
 
-	// Distinct aggregation cannot be pushed down in the join.
+	// Distinctable aggregation cannot be pushed down in the join.
 	// We keep node of the distinct aggregation expression to be used later for ordering.
 	if !canPushDistinctAggr {
 		if len(distinctExprs) != 1 {

go/vt/vtgate/planbuilder/operators/ast_to_op.go

@@ -178,7 +178,7 @@ func createOperatorFromUnion(ctx *plancontext.PlanningContext, node *sqlparser.U
 	return newHorizon(union, node)
 }
 
-func translateQueryToOpForUnion(ctx *plancontext.PlanningContext, node sqlparser.SelectStatement) Operator {
+func translateQueryToOpForUnion(ctx *plancontext.PlanningContext, node sqlparser.TableSubquery) Operator {
 	op := translateQueryToOp(ctx, node)
 	if hz, ok := op.(*Horizon); ok {
 		hz.Truncate = true

go/vt/vtgate/planbuilder/operators/expressions.go

@@ -17,7 +17,10 @@ limitations under the License.
 package operators
 
 import (
+	"fmt"
+
 	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 )
@@ -121,3 +124,11 @@ func simplifyPredicates(ctx *plancontext.PlanningContext, in sqlparser.Expr) sql
 	}
 	return output
 }
+
+func getFirstSelect(selStmt sqlparser.TableSubquery) *sqlparser.Select {
+	firstSelect, err := sqlparser.GetFirstSelect(selStmt)
+	if err != nil {
+		panic(vterrors.VT12001(fmt.Sprintf("first UNION part not a SELECT: %v", err)))
+	}
+	return firstSelect
+}

go/vt/vtgate/planbuilder/operators/horizon.go

@@ -45,7 +45,7 @@ type Horizon struct {
 	// QP contains the QueryProjection for this op
 	QP *QueryProjection
 
-	Query sqlparser.SelectStatement
+	Query sqlparser.TableSubquery
 
 	// Columns needed to feed other plans
 	Columns       []*sqlparser.ColName
@@ -54,7 +54,7 @@ type Horizon struct {
 	Truncate bool
 }
 
-func newHorizon(src Operator, query sqlparser.SelectStatement) *Horizon {
+func newHorizon(src Operator, query sqlparser.TableSubquery) *Horizon {
 	return &Horizon{
 		unaryOperator: newUnaryOp(src),
 		Query:         query,
@@ -148,7 +148,7 @@ func (h *Horizon) FindCol(ctx *plancontext.PlanningContext, expr sqlparser.Expr,
 		return -1
 	}
 
-	for idx, se := range sqlparser.GetFirstSelect(h.Query).SelectExprs {
+	for idx, se := range getFirstSelect(h.Query).SelectExprs {
 		ae, ok := se.(*sqlparser.AliasedExpr)
 		if !ok {
 			panic(vterrors.VT09015())
@@ -174,7 +174,7 @@ func (h *Horizon) GetColumns(ctx *plancontext.PlanningContext) (exprs []*sqlpars
 }
 
 func (h *Horizon) GetSelectExprs(*plancontext.PlanningContext) sqlparser.SelectExprs {
-	return sqlparser.GetFirstSelect(h.Query).SelectExprs
+	return getFirstSelect(h.Query).SelectExprs
 }
 
 func (h *Horizon) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
@@ -185,7 +185,7 @@ func (h *Horizon) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
 }
 
 // TODO: REMOVE
-func (h *Horizon) selectStatement() sqlparser.SelectStatement {
+func (h *Horizon) selectStatement() sqlparser.TableSubquery {
 	return h.Query
 }
 

go/vt/vtgate/planbuilder/operators/horizon_expanding.go

@@ -95,7 +95,7 @@ func expandSelectHorizon(ctx *plancontext.PlanningContext, horizon *Horizon, sel
 
 	if qp.NeedsDistinct() {
 		op = newDistinct(op, qp, true)
-		extracted = append(extracted, "Distinct")
+		extracted = append(extracted, "Distinctable")
 	}
 
 	if sel.Having != nil {

go/vt/vtgate/planbuilder/operators/insert.go

@@ -394,7 +394,7 @@ func createInsertOperator(ctx *plancontext.PlanningContext, insStmt *sqlparser.I
 	case sqlparser.Values:
 		op = route
 		route.Source = insertRowsPlan(ctx, insOp, insStmt, rows)
-	case sqlparser.SelectStatement:
+	case sqlparser.TableSubquery:
 		op = insertSelectPlan(ctx, insOp, route, insStmt, rows)
 	}
 	if insStmt.Comments != nil {
@@ -408,7 +408,7 @@ func insertSelectPlan(
 	insOp *Insert,
 	routeOp *Route,
 	ins *sqlparser.Insert,
-	sel sqlparser.SelectStatement,
+	sel sqlparser.TableSubquery,
 ) *InsertSelection {
 	if columnMismatch(insOp.AutoIncrement, ins, sel) {
 		panic(vterrors.VT03006())
@@ -457,7 +457,7 @@ func insertSelectPlan(
 	return insertSelect
 }
 
-func columnMismatch(gen *Generate, ins *sqlparser.Insert, sel sqlparser.SelectStatement) bool {
+func columnMismatch(gen *Generate, ins *sqlparser.Insert, sel sqlparser.TableSubquery) bool {
 	origColCount := len(ins.Columns)
 	if gen != nil && gen.added {
 		// One column got added to the insert query ast for auto increment column.
@@ -468,7 +468,7 @@ func columnMismatch(gen *Generate, ins *sqlparser.Insert, sel sqlparser.SelectSt
 		return true
 	}
 	if origColCount > sel.GetColumnCount() {
-		sel := sqlparser.GetFirstSelect(sel)
+		sel := getFirstSelect(sel)
 		var hasStarExpr bool
 		for _, sExpr := range sel.SelectExprs {
 			if _, hasStarExpr = sExpr.(*sqlparser.StarExpr); hasStarExpr {

go/vt/vtgate/planbuilder/operators/phases.go

@@ -59,7 +59,7 @@ func (p Phase) String() string {
 	case addAggrOrdering:
 		return "optimize aggregations with ORDER BY"
 	case cleanOutPerfDistinct:
-		return "optimize Distinct operations"
+		return "optimize Distinctable operations"
 	case subquerySettling:
 		return "settle subqueries"
 	case dmlWithInput:

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -32,7 +32,7 @@ import (
 func planQuery(ctx *plancontext.PlanningContext, root Operator) Operator {
 	var selExpr sqlparser.SelectExprs
 	if horizon, isHorizon := root.(*Horizon); isHorizon {
-		sel := sqlparser.GetFirstSelect(horizon.Query)
+		sel := getFirstSelect(horizon.Query)
 		selExpr = sqlparser.Clone(sel.SelectExprs)
 	}
 
@@ -207,7 +207,7 @@ func pushOrExpandHorizon(ctx *plancontext.PlanningContext, in *Horizon) (Operato
 		!hasHaving &&
 		!needsOrdering &&
 		!qp.NeedsAggregation() &&
-		!in.selectStatement().IsDistinct() &&
+		!isDistinctAST(in.selectStatement()) &&
 		in.selectStatement().GetLimit() == nil
 
 	if canPush {
@@ -784,14 +784,21 @@ func isDistinct(op Operator) bool {
 	case *Union:
 		return op.distinct
 	case *Horizon:
-		return op.Query.IsDistinct()
+		return isDistinctAST(op.Query)
 	case *Limit:
 		return isDistinct(op.Source)
 	default:
 		return false
 	}
 }
 
+func isDistinctAST(s sqlparser.Statement) bool {
+	if d, ok := s.(sqlparser.Distinctable); ok {
+		return d.IsDistinct()
+	}
+	return false
+}
+
 func tryPushUnion(ctx *plancontext.PlanningContext, op *Union) (Operator, *ApplyResult) {
 	if res := compactUnion(op); res != NoRewrite {
 		return op, res

go/vt/vtgate/planbuilder/operators/queryprojection.go

@@ -203,7 +203,7 @@ func (qp *QueryProjection) addSelectExpressions(ctx *plancontext.PlanningContext
 func createQPFromUnion(ctx *plancontext.PlanningContext, union *sqlparser.Union) *QueryProjection {
 	qp := &QueryProjection{}
 
-	sel := sqlparser.GetFirstSelect(union)
+	sel := getFirstSelect(union)
 	qp.addSelectExpressions(ctx, sel)
 	qp.addOrderBy(ctx, union.OrderBy)
 
@@ -714,7 +714,7 @@ func CompareRefInt(a *int, b *int) bool {
 	return *a < *b
 }
 
-func CreateQPFromSelectStatement(ctx *plancontext.PlanningContext, stmt sqlparser.SelectStatement) *QueryProjection {
+func CreateQPFromSelectStatement(ctx *plancontext.PlanningContext, stmt sqlparser.TableSubquery) *QueryProjection {
 	switch sel := stmt.(type) {
 	case *sqlparser.Select:
 		return createQPFromSelect(ctx, sel)

go/vt/vtgate/planbuilder/operators/subquery_builder.go

@@ -115,7 +115,7 @@ func createSubqueryOp(
 // inspectStatement goes through all the predicates contained in the AST
 // and extracts subqueries into operators
 func (sqb *SubQueryBuilder) inspectStatement(ctx *plancontext.PlanningContext,
-	stmt sqlparser.SelectStatement,
+	stmt sqlparser.TableSubquery,
 ) (sqlparser.Exprs, []applyJoinColumn) {
 	switch stmt := stmt.(type) {
 	case *sqlparser.Select:

go/vt/vtgate/planbuilder/operators/subquery_planning.go

@@ -30,7 +30,7 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 )
 
-func isMergeable(ctx *plancontext.PlanningContext, query sqlparser.SelectStatement, op Operator) bool {
+func isMergeable(ctx *plancontext.PlanningContext, query sqlparser.TableSubquery, op Operator) bool {
 	validVindex := func(expr sqlparser.Expr) bool {
 		sc := findColumnVindex(ctx, op, expr)
 		return sc != nil && sc.IsUnique()
@@ -672,7 +672,7 @@ func (s *subqueryRouteMerger) rewriteASTExpression(ctx *plancontext.PlanningCont
 	if err != nil {
 		panic(err)
 	}
-	subqStmt, ok := stmt.(sqlparser.SelectStatement)
+	subqStmt, ok := stmt.(sqlparser.TableSubquery)
 	if !ok {
 		panic(vterrors.VT13001("subqueries should only be select statement"))
 	}
@@ -700,7 +700,7 @@ func (s *subqueryRouteMerger) rewriteASTExpression(ctx *plancontext.PlanningCont
 		if !deps.IsSolvedBy(subqID) {
 			cursor.Replace(exprFound)
 		}
-	}, nil).(sqlparser.SelectStatement)
+	}, nil).(sqlparser.TableSubquery)
 	if err != nil {
 		panic(err)
 	}

go/vt/vtgate/planbuilder/operators/union.go

@@ -150,7 +150,7 @@ func (u *Union) GetSelectFor(source int) *sqlparser.Select {
 	for {
 		switch op := src.(type) {
 		case *Horizon:
-			return sqlparser.GetFirstSelect(op.Query)
+			return getFirstSelect(op.Query)
 		case *Route:
 			src = op.Source
 		default: