scatter query with aggregation to check order by columns are in group by

Signed-off-by: Harshit Gangal <harshit@planetscale.com>

go/vt/vtgate/planbuilder/abstract/queryprojection.go

@@ -94,18 +94,6 @@ func CreateQPFromSelect(sel *sqlparser.Select) (*QueryProjection, error) {
 		qp.GroupByExprs = append(qp.GroupByExprs, GroupBy{Inner: expr, WeightStrExpr: weightStrExpr})
 	}
 
-	if qp.HasAggr || len(qp.GroupByExprs) > 0 {
-		expr := qp.getNonAggrExprNotMatchingGroupByExprs()
-		// if we have aggregation functions, non aggregating columns and GROUP BY,
-		// the non-aggregating expressions must all be listed in the GROUP BY list
-		if expr != nil {
-			if len(qp.GroupByExprs) > 0 {
-				return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.WrongFieldWithGroup, "Expression of SELECT list is not in GROUP BY clause and contains nonaggregated column '%s' which is not functionally dependent on columns in GROUP BY clause; this is incompatible with sql_mode=only_full_group_by", sqlparser.String(expr))
-			}
-			return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.MixOfGroupFuncAndFields, "In aggregated query without GROUP BY, expression of SELECT list contains nonaggregated column '%s'; this is incompatible with sql_mode=only_full_group_by", sqlparser.String(expr))
-		}
-	}
-
 	canPushDownSorting := true
 	for _, order := range sel.OrderBy {
 		expr, weightStrExpr, err := qp.getSimplifiedExpr(order.Expr, "order clause")
@@ -121,10 +109,24 @@ func CreateQPFromSelect(sel *sqlparser.Select) (*QueryProjection, error) {
 		})
 		canPushDownSorting = canPushDownSorting && !sqlparser.ContainsAggregation(weightStrExpr)
 	}
+
+	if qp.HasAggr || len(qp.GroupByExprs) > 0 {
+		expr := qp.getNonAggrExprNotMatchingGroupByExprs()
+		// if we have aggregation functions, non aggregating columns and GROUP BY,
+		// the non-aggregating expressions must all be listed in the GROUP BY list
+		if expr != nil {
+			if len(qp.GroupByExprs) > 0 {
+				return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.WrongFieldWithGroup, "Expression of SELECT list is not in GROUP BY clause and contains nonaggregated column '%s' which is not functionally dependent on columns in GROUP BY clause; this is incompatible with sql_mode=only_full_group_by", sqlparser.String(expr))
+			}
+			return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.MixOfGroupFuncAndFields, "In aggregated query without GROUP BY, expression of SELECT list contains nonaggregated column '%s'; this is incompatible with sql_mode=only_full_group_by", sqlparser.String(expr))
+		}
+	}
+
 	if qp.Distinct && !qp.HasAggr {
 		qp.GroupByExprs = nil
 	}
 	qp.CanPushDownSorting = canPushDownSorting
+
 	return qp, nil
 }
 
@@ -159,6 +161,21 @@ func (qp *QueryProjection) getNonAggrExprNotMatchingGroupByExprs() sqlparser.Exp
 			return expr.Col.Expr
 		}
 	}
+	for _, order := range qp.OrderExprs {
+		if sqlparser.IsAggregation(order.WeightStrExpr) {
+			continue
+		}
+		isGroupByOk := false
+		for _, groupByExpr := range qp.GroupByExprs {
+			if sqlparser.EqualsExpr(groupByExpr.WeightStrExpr, order.WeightStrExpr) {
+				isGroupByOk = true
+				break
+			}
+		}
+		if !isGroupByOk {
+			return order.Inner.Expr
+		}
+	}
 	return nil
 }
 

go/vt/vtgate/planbuilder/horizon_planning.go

@@ -473,11 +473,11 @@ func (hp *horizonPlanning) planDistinct() error {
 		// and then we might also add a distinct operator on top if it is needed
 		p.Select.MakeDistinct()
 		if !p.isSingleShard() && !selectHasUniqueVindex(hp.vschema, hp.semTable, hp.qp.SelectExprs) {
-			return hp.pushDistinct()
+			return hp.addDistinct()
 		}
 		return nil
 	case *joinGen4:
-		return hp.pushDistinct()
+		return hp.addDistinct()
 	case *orderedAggregate:
 		return hp.planDistinctOA(p)
 	default:
@@ -522,7 +522,7 @@ func (hp *horizonPlanning) planDistinctOA(currPlan *orderedAggregate) error {
 	return nil
 }
 
-func (hp *horizonPlanning) pushDistinct() error {
+func (hp *horizonPlanning) addDistinct() error {
 	eaggr := &engine.OrderedAggregate{}
 	oa := &orderedAggregate{
 		resultsBuilder: resultsBuilder{

go/vt/vtgate/planbuilder/testdata/postprocess_cases.txt

@@ -1002,7 +1002,7 @@ Gen4 plan same as above
 # Order by, invalid column number
 "select col from user order by 18446744073709551616"
 "error parsing column number: 18446744073709551616"
-Gen4 error: strconv.Atoi: parsing "18446744073709551616": value out of range
+Gen4 plan same as above
 
 # Order by, out of range column number
 "select col from user order by 2"
@@ -1713,31 +1713,7 @@ Gen4 error: In aggregated query without GROUP BY, expression of SELECT list cont
 # Scatter order by and aggregation: order by column must reference column from select list
 "select col, count(*) from user group by col order by c1"
 "unsupported: memory sort: order by must reference a column in the select list: c1 asc"
-{
-  "QueryType": "SELECT",
-  "Original": "select col, count(*) from user group by col order by c1",
-  "Instructions": {
-    "OperatorType": "Aggregate",
-    "Variant": "Ordered",
-    "Aggregates": "count(1)",
-    "GroupBy": "(0|2)",
-    "ResultColumns": 2,
-    "Inputs": [
-      {
-        "OperatorType": "Route",
-        "Variant": "SelectScatter",
-        "Keyspace": {
-          "Name": "user",
-          "Sharded": true
-        },
-        "FieldQuery": "select col, count(*), weight_string(col), c1, weight_string(c1) from `user` where 1 != 1 group by col",
-        "OrderBy": "(3|4) ASC, (0|2) ASC",
-        "Query": "select col, count(*), weight_string(col), c1, weight_string(c1) from `user` group by col order by c1 asc, col asc",
-        "Table": "`user`"
-      }
-    ]
-  }
-}
+Gen4 error: Expression of SELECT list is not in GROUP BY clause and contains nonaggregated column 'c1' which is not functionally dependent on columns in GROUP BY clause; this is incompatible with sql_mode=only_full_group_by
 
 # scatter aggregate with ambiguous aliases
 "select distinct a, b as a from user"

go/vt/vtgate/semantics/analyzer.go

@@ -198,7 +198,7 @@ func (a *analyzer) analyzeOrderByGroupByExprForLiteral(input sqlparser.Expr, cal
 	currScope := a.currentScope()
 	num, err := strconv.Atoi(l.Val)
 	if err != nil {
-		a.err = err
+		a.err = vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "error parsing column number: %s", l.Val)
 		return
 	}
 	if num < 1 || num > len(currScope.selectExprs) {