Fixed regression for grouping by integer functions

The regression in the linked issue was found to be occurring from adding weight_string function due to order by as introduced in vitessio#7678.
MySQL does not support the generated query -

```sql
select ascii(val1) as a, count(*), weight_string(ascii(val1)) from aggr_test group by a order by a asc
```

In order to fix this, we also add the weight_string function to the group by clause as follows -

```
select ascii(val1) as a, count(*), weight_string(ascii(val1)) from aggr_test group by a, weight_string(ascii(val1)) order by a asc
```

Co-authored-by: Manan Gupta <manan@planetscale.com>

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

go/vt/vtgate/endtoend/aggr_test.go

@@ -54,4 +54,34 @@ func TestAggregateTypes(t *testing.T) {
 	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)]]`; got != want {
 		t.Errorf("select:\n%v want\n%v", got, want)
 	}
+
+	qr = exec(t, conn, "select ascii(val1) as a, count(*) from aggr_test group by a")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[INT32(65) INT64(1)] [INT32(69) INT64(1)] [INT32(97) INT64(1)] [INT32(98) INT64(1)] [INT32(99) INT64(2)] [INT32(100) INT64(1)] [INT32(101) INT64(1)]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	qr = exec(t, conn, "select ascii(val1) as a, count(*) from aggr_test group by a order by a")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[INT32(65) INT64(1)] [INT32(69) INT64(1)] [INT32(97) INT64(1)] [INT32(98) INT64(1)] [INT32(99) INT64(2)] [INT32(100) INT64(1)] [INT32(101) INT64(1)]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	qr = exec(t, conn, "select ascii(val1) as a, count(*) from aggr_test group by a order by 2, a")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[INT32(65) INT64(1)] [INT32(69) INT64(1)] [INT32(97) INT64(1)] [INT32(98) INT64(1)] [INT32(100) INT64(1)] [INT32(101) INT64(1)] [INT32(99) INT64(2)]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	qr = exec(t, conn, "select val1 as a, count(*) from aggr_test group by a")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("a") INT64(2)] [VARCHAR("b") INT64(1)] [VARCHAR("c") INT64(2)] [VARCHAR("d") INT64(1)] [VARCHAR("e") INT64(2)]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	qr = exec(t, conn, "select val1 as a, count(*) from aggr_test group by a order by a")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("a") INT64(2)] [VARCHAR("b") INT64(1)] [VARCHAR("c") INT64(2)] [VARCHAR("d") INT64(1)] [VARCHAR("e") INT64(2)]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
+
+	qr = exec(t, conn, "select val1 as a, count(*) from aggr_test group by a order by 2, a")
+	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("b") INT64(1)] [VARCHAR("d") INT64(1)] [VARCHAR("a") INT64(2)] [VARCHAR("c") INT64(2)] [VARCHAR("e") INT64(2)]]`; got != want {
+		t.Errorf("select:\n%v want\n%v", got, want)
+	}
 }

go/vt/vtgate/engine/ordered_aggregate.go

@@ -48,6 +48,9 @@ type OrderedAggregate struct {
 	// the aggregation key.
 	Keys []int
 
+	// Keeps track if the keys above were added because of GroupBy or not
+	FromGroupBy []bool
+
 	// TruncateColumnCount specifies the number of columns to return
 	// in the final result. Rest of the columns are truncated
 	// from the result received. If 0, no truncation happens.

go/vt/vtgate/engine/route.go

@@ -130,6 +130,9 @@ type OrderbyParams struct {
 	// It is set to -1 if such a column is not added to the query
 	WeightStringCol int
 	Desc            bool
+
+	// v3 specific boolean. Used to also add weight strings originating from GroupBys to the Group by clause
+	FromGroupBy bool
 }
 
 func (obp OrderbyParams) String() string {

go/vt/vtgate/executor_select_test.go

@@ -1430,7 +1430,7 @@ func TestSelectScatterAggregate(t *testing.T) {
 	require.NoError(t, err)
 
 	wantQueries := []*querypb.BoundQuery{{
-		Sql:           "select col, sum(foo), weight_string(col) from `user` group by col order by col asc",
+		Sql:           "select col, sum(foo), weight_string(col) from `user` group by col, weight_string(col) order by col asc",
 		BindVariables: map[string]*querypb.BindVariable{},
 	}}
 	for _, conn := range conns {
@@ -1489,7 +1489,7 @@ func TestStreamSelectScatterAggregate(t *testing.T) {
 	require.NoError(t, err)
 
 	wantQueries := []*querypb.BoundQuery{{
-		Sql:           "select col, sum(foo), weight_string(col) from `user` group by col order by col asc",
+		Sql:           "select col, sum(foo), weight_string(col) from `user` group by col, weight_string(col) order by col asc",
 		BindVariables: map[string]*querypb.BindVariable{},
 	}}
 	for _, conn := range conns {

go/vt/vtgate/planbuilder/concatenate.go

@@ -70,7 +70,7 @@ func (c *concatenate) SupplyCol(col *sqlparser.ColName) (rc *resultColumn, colNu
 	panic("implement me")
 }
 
-func (c *concatenate) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
+func (c *concatenate) SupplyWeightString(int, bool) (weightcolNumber int, err error) {
 	panic("implement me")
 }
 

go/vt/vtgate/planbuilder/grouping.go

@@ -77,6 +77,7 @@ func planGroupBy(pb *primitiveBuilder, input logicalPlan, groupBy sqlparser.Grou
 				return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: in scatter query: only simple references allowed")
 			}
 			node.eaggr.Keys = append(node.eaggr.Keys, colNumber)
+			node.eaggr.FromGroupBy = append(node.eaggr.FromGroupBy, true)
 		}
 		// Append the distinct aggregate if any.
 		if node.extraDistinct != nil {
@@ -110,6 +111,7 @@ func planDistinct(input logicalPlan) (logicalPlan, error) {
 				return newDistinct(node), nil
 			}
 			node.eaggr.Keys = append(node.eaggr.Keys, i)
+			node.eaggr.FromGroupBy = append(node.eaggr.FromGroupBy, false)
 		}
 		newInput, err := planDistinct(node.input)
 		if err != nil {

go/vt/vtgate/planbuilder/join.go

@@ -212,20 +212,20 @@ func (jb *join) SupplyCol(col *sqlparser.ColName) (rc *resultColumn, colNumber i
 }
 
 // SupplyWeightString implements the logicalPlan interface
-func (jb *join) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
+func (jb *join) SupplyWeightString(colNumber int, alsoAddToGroupBy bool) (weightcolNumber int, err error) {
 	rc := jb.resultColumns[colNumber]
 	if weightcolNumber, ok := jb.weightStrings[rc]; ok {
 		return weightcolNumber, nil
 	}
 	routeNumber := rc.column.Origin().Order()
 	if jb.isOnLeft(routeNumber) {
-		sourceCol, err := jb.Left.SupplyWeightString(-jb.ejoin.Cols[colNumber] - 1)
+		sourceCol, err := jb.Left.SupplyWeightString(-jb.ejoin.Cols[colNumber]-1, alsoAddToGroupBy)
 		if err != nil {
 			return 0, err
 		}
 		jb.ejoin.Cols = append(jb.ejoin.Cols, -sourceCol-1)
 	} else {
-		sourceCol, err := jb.Right.SupplyWeightString(jb.ejoin.Cols[colNumber] - 1)
+		sourceCol, err := jb.Right.SupplyWeightString(jb.ejoin.Cols[colNumber]-1, alsoAddToGroupBy)
 		if err != nil {
 			return 0, err
 		}

go/vt/vtgate/planbuilder/join2.go

@@ -73,7 +73,7 @@ func (j *joinV4) SupplyCol(col *sqlparser.ColName) (rc *resultColumn, colNumber
 }
 
 // SupplyWeightString implements the logicalPlan interface
-func (j *joinV4) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
+func (j *joinV4) SupplyWeightString(int, bool) (weightcolNumber int, err error) {
 	panic("implement me")
 }
 

go/vt/vtgate/planbuilder/logical_plan.go

@@ -70,7 +70,7 @@ type logicalPlan interface {
 
 	// SupplyWeightString must supply a weight_string expression of the
 	// specified column. It returns an error if we cannot supply a weight column for it.
-	SupplyWeightString(colNumber int) (weightcolNumber int, err error)
+	SupplyWeightString(colNumber int, alsoAddToGroupBy bool) (weightcolNumber int, err error)
 
 	// Primitive returns the underlying primitive.
 	// This function should only be called after Wireup is finished.
@@ -170,8 +170,8 @@ func (bc *logicalPlanCommon) SupplyCol(col *sqlparser.ColName) (rc *resultColumn
 	return bc.input.SupplyCol(col)
 }
 
-func (bc *logicalPlanCommon) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
-	return bc.input.SupplyWeightString(colNumber)
+func (bc *logicalPlanCommon) SupplyWeightString(colNumber int, alsoAddToGroupBy bool) (weightcolNumber int, err error) {
+	return bc.input.SupplyWeightString(colNumber, alsoAddToGroupBy)
 }
 
 // Rewrite implements the logicalPlan interface
@@ -239,12 +239,14 @@ func (rsb *resultsBuilder) SupplyCol(col *sqlparser.ColName) (rc *resultColumn,
 	return rc, colNumber
 }
 
-func (rsb *resultsBuilder) SupplyWeightString(colNumber int) (weightcolNumber int, err error) {
+func (rsb *resultsBuilder) SupplyWeightString(colNumber int, alsoAddToGroupBy bool) (weightcolNumber int, err error) {
 	rc := rsb.resultColumns[colNumber]
-	if weightcolNumber, ok := rsb.weightStrings[rc]; ok {
+	var ok bool
+	weightcolNumber, ok = rsb.weightStrings[rc]
+	if !alsoAddToGroupBy && ok {
 		return weightcolNumber, nil
 	}
-	weightcolNumber, err = rsb.input.SupplyWeightString(colNumber)
+	weightcolNumber, err = rsb.input.SupplyWeightString(colNumber, alsoAddToGroupBy)
 	if err != nil {
 		return 0, nil
 	}

go/vt/vtgate/planbuilder/memory_sort.go

@@ -40,7 +40,7 @@ type memorySort struct {
 }
 
 // newMemorySort builds a new memorySort.
-func newMemorySort(plan logicalPlan, orderBy sqlparser.OrderBy) (*memorySort, error) {
+func newMemorySort(plan logicalPlan, orderBy v3OrderBy) (*memorySort, error) {
 	eMemorySort := &engine.MemorySort{}
 	ms := &memorySort{
 		resultsBuilder: newResultsBuilder(plan, eMemorySort),
@@ -113,15 +113,9 @@ func (ms *memorySort) Wireup(plan logicalPlan, jt *jointab) error {
 		rc := ms.resultColumns[orderby.Col]
 		// Add a weight_string column if we know that the column is a textual column or if its type is unknown
 		if sqltypes.IsText(rc.column.typ) || rc.column.typ == sqltypes.Null {
-			// If a weight string was previously requested, reuse it.
-			if weightcolNumber, ok := ms.weightStrings[rc]; ok {
-				ms.eMemorySort.OrderBy[i].WeightStringCol = weightcolNumber
-				continue
-			}
-			weightcolNumber, err := ms.input.SupplyWeightString(orderby.Col)
+			weightcolNumber, err := ms.input.SupplyWeightString(orderby.Col, orderby.FromGroupBy)
 			if err != nil {
-				_, isUnsupportedErr := err.(UnsupportedSupplyWeightString)
-				if isUnsupportedErr {
+				if _, isUnsupportedErr := err.(UnsupportedSupplyWeightString); isUnsupportedErr {
 					continue
 				}
 				return err

go/vt/vtgate/planbuilder/merge_sort.go

@@ -68,13 +68,8 @@ func (ms *mergeSort) Wireup(plan logicalPlan, jt *jointab) error {
 		rc := ms.resultColumns[orderby.Col]
 		// Add a weight_string column if we know that the column is a textual column or if its type is unknown
 		if sqltypes.IsText(rc.column.typ) || rc.column.typ == sqltypes.Null {
-			// If a weight string was previously requested, reuse it.
-			if colNumber, ok := ms.weightStrings[rc]; ok {
-				rb.eroute.OrderBy[i].WeightStringCol = colNumber
-				continue
-			}
 			var err error
-			rb.eroute.OrderBy[i].WeightStringCol, err = rb.SupplyWeightString(orderby.Col)
+			rb.eroute.OrderBy[i].WeightStringCol, err = rb.SupplyWeightString(orderby.Col, orderby.FromGroupBy)
 			if err != nil {
 				_, isUnsupportedErr := err.(UnsupportedSupplyWeightString)
 				if isUnsupportedErr {

go/vt/vtgate/planbuilder/ordered_aggregate.go

@@ -334,11 +334,7 @@ func (oa *orderedAggregate) Wireup(plan logicalPlan, jt *jointab) error {
 	for i, colNumber := range oa.eaggr.Keys {
 		rc := oa.resultColumns[colNumber]
 		if sqltypes.IsText(rc.column.typ) {
-			if weightcolNumber, ok := oa.weightStrings[rc]; ok {
-				oa.eaggr.Keys[i] = weightcolNumber
-				continue
-			}
-			weightcolNumber, err := oa.input.SupplyWeightString(colNumber)
+			weightcolNumber, err := oa.input.SupplyWeightString(colNumber, oa.eaggr.FromGroupBy[i])
 			if err != nil {
 				_, isUnsupportedErr := err.(UnsupportedSupplyWeightString)
 				if isUnsupportedErr {

go/vt/vtgate/planbuilder/ordering.go

@@ -22,7 +22,14 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/engine"
 )
 
-func planOrdering(pb *primitiveBuilder, input logicalPlan, orderBy sqlparser.OrderBy) (logicalPlan, error) {
+type v3Order struct {
+	*sqlparser.Order
+	fromGroupBy bool
+}
+
+type v3OrderBy []*v3Order
+
+func planOrdering(pb *primitiveBuilder, input logicalPlan, orderBy v3OrderBy) (logicalPlan, error) {
 	switch node := input.(type) {
 	case *subquery, *vindexFunc:
 		if len(orderBy) == 0 {
@@ -56,7 +63,7 @@ func planOrdering(pb *primitiveBuilder, input logicalPlan, orderBy sqlparser.Ord
 	return nil, vterrors.Errorf(vtrpc.Code_INTERNAL, "[BUG] unreachable %T.ordering", input)
 }
 
-func planOAOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, oa *orderedAggregate) (logicalPlan, error) {
+func planOAOrdering(pb *primitiveBuilder, orderBy v3OrderBy, oa *orderedAggregate) (logicalPlan, error) {
 	// The requested order must be such that the ordering can be done
 	// before the group by, which will allow us to push it down to the
 	// route. This is actually true in most use cases, except for situations
@@ -78,7 +85,7 @@ func planOAOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, oa *ordered
 	// referenced tracks the keys referenced by the order by clause.
 	referenced := make([]bool, len(oa.eaggr.Keys))
 	postSort := false
-	selOrderBy := make(sqlparser.OrderBy, 0, len(orderBy))
+	selOrderBy := make(v3OrderBy, 0, len(orderBy))
 	for _, order := range orderBy {
 		// Identify the order by column.
 		var orderByCol *column
@@ -110,6 +117,7 @@ func planOAOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, oa *ordered
 
 			found = true
 			referenced[j] = true
+			order.fromGroupBy = oa.eaggr.FromGroupBy[j]
 			selOrderBy = append(selOrderBy, order)
 			break
 		}
@@ -128,12 +136,19 @@ func planOAOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, oa *ordered
 		if err != nil {
 			return nil, vterrors.Wrapf(err, "generating order by clause")
 		}
-		selOrderBy = append(selOrderBy, &sqlparser.Order{Expr: col, Direction: sqlparser.AscOrder})
+		selOrderBy = append(selOrderBy, &v3Order{
+			Order:       &sqlparser.Order{Expr: col, Direction: sqlparser.AscOrder},
+			fromGroupBy: oa.eaggr.FromGroupBy[i],
+		})
 	}
 
 	// Append the distinct aggregate if any.
 	if oa.extraDistinct != nil {
-		selOrderBy = append(selOrderBy, &sqlparser.Order{Expr: oa.extraDistinct, Direction: sqlparser.AscOrder})
+		el := &sqlparser.Order{Expr: oa.extraDistinct, Direction: sqlparser.AscOrder}
+		selOrderBy = append(selOrderBy, &v3Order{
+			Order:       el,
+			fromGroupBy: true,
+		})
 	}
 
 	// Push down the order by.
@@ -151,7 +166,7 @@ func planOAOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, oa *ordered
 	return oa, nil
 }
 
-func planJoinOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, node *join) (logicalPlan, error) {
+func planJoinOrdering(pb *primitiveBuilder, orderBy v3OrderBy, node *join) (logicalPlan, error) {
 	isSpecial := false
 	switch len(orderBy) {
 	case 0:
@@ -226,7 +241,7 @@ func planJoinOrdering(pb *primitiveBuilder, orderBy sqlparser.OrderBy, node *joi
 	return node, nil
 }
 
-func planRouteOrdering(orderBy sqlparser.OrderBy, node *route) (logicalPlan, error) {
+func planRouteOrdering(orderBy v3OrderBy, node *route) (logicalPlan, error) {
 	switch len(orderBy) {
 	case 0:
 		return node, nil
@@ -240,14 +255,14 @@ func planRouteOrdering(orderBy sqlparser.OrderBy, node *route) (logicalPlan, err
 			}
 		}
 		if isSpecial {
-			node.Select.AddOrder(orderBy[0])
+			node.Select.AddOrder(orderBy[0].Order)
 			return node, nil
 		}
 	}
 
 	if node.isSingleShard() {
 		for _, order := range orderBy {
-			node.Select.AddOrder(order)
+			node.Select.AddOrder(order.Order)
 		}
 		return node, nil
 	}
@@ -293,10 +308,11 @@ func planRouteOrdering(orderBy sqlparser.OrderBy, node *route) (logicalPlan, err
 			Col:             colNumber,
 			WeightStringCol: -1,
 			Desc:            order.Direction == sqlparser.DescOrder,
+			FromGroupBy:     order.fromGroupBy,
 		}
 		node.eroute.OrderBy = append(node.eroute.OrderBy, ob)
 
-		node.Select.AddOrder(order)
+		node.Select.AddOrder(order.Order)
 	}
 	return newMergeSort(node), nil
 }

go/vt/vtgate/planbuilder/postprocess.go

@@ -54,7 +54,16 @@ func (pb *primitiveBuilder) pushOrderBy(orderBy sqlparser.OrderBy) error {
 	if err := pb.st.ResolveSymbols(orderBy); err != nil {
 		return err
 	}
-	plan, err := planOrdering(pb, pb.plan, orderBy)
+
+	var v3OrderBylist v3OrderBy
+
+	if orderBy != nil {
+		v3OrderBylist = make(v3OrderBy, 0, len(orderBy))
+		for _, order := range orderBy {
+			v3OrderBylist = append(v3OrderBylist, &v3Order{Order: order})
+		}
+	}
+	plan, err := planOrdering(pb, pb.plan, v3OrderBylist)
 	if err != nil {
 		return err
 	}