added support for group by with weight string needed. Add truncation …

…if required after horizon planning at top most plan. all select columns to be available irrepective of if they are same expression (for correctness)

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

go/test/endtoend/vtgate/gen4/gen4_test.go

@@ -48,6 +48,27 @@ func TestOrderBy(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestGroupBy(t *testing.T) {
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	// insert some data.
+	checkedExec(t, conn, `insert into t1(id, col) values (1, 123),(2, 12),(3, 13),(4, 1234)`)
+	checkedExec(t, conn, `insert into t2(id, tcol1, tcol2) values (1, 'A', 'A'),(2, 'B', 'C'),(3, 'A', 'C'),(4, 'C', 'A'),(5, 'A', 'A'),(6, 'B', 'C'),(7, 'B', 'A'),(8, 'C', 'B')`)
+
+	// Gen4 only supported query.
+	assertMatches(t, conn, `select tcol2, tcol1, count(id) from t2 group by tcol2, tcol1`,
+		`[[VARCHAR("A") VARCHAR("A") INT64(2)] [VARCHAR("A") VARCHAR("B") INT64(1)] [VARCHAR("A") VARCHAR("C") INT64(1)] [VARCHAR("B") VARCHAR("C") INT64(1)] [VARCHAR("C") VARCHAR("A") INT64(1)] [VARCHAR("C") VARCHAR("B") INT64(2)]]`)
+
+	assertMatches(t, conn, `select tcol1, tcol1 from t2 order by tcol1`,
+		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")] [VARCHAR("C") VARCHAR("C")]]`)
+
+	assertMatches(t, conn, `select tcol1, tcol1 from t1 join t2 on t1.id = t2.id order by tcol1`,
+		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")]]`)
+}
+
 func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
 	t.Helper()
 	qr := checkedExec(t, conn, query)

go/test/endtoend/vtgate/gen4/main_test.go

@@ -35,6 +35,13 @@ var (
 	col bigint,
 	primary key(id)
 ) Engine=InnoDB;
+
+create table t2(
+	id bigint,
+	tcol1 varchar(50),
+	tcol2 varchar(50),
+	primary key(id)
+) Engine=InnoDB;
 `
 
 	VSchema = `
@@ -53,6 +60,20 @@ var (
           "name": "xxhash"
         }
       ]
+    },
+    "t2": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ],
+      "columns": [
+        {
+          "name": "tcol1",
+          "type": "VARCHAR"
+        }
+      ]
     }
   }
 }`
@@ -84,7 +105,7 @@ func TestMain(m *testing.M) {
 		}
 
 		// Start vtgate
-		clusterInstance.VtGateExtraArgs = []string{"-planner_version", "Gen4Fallback"} // enable Gen4 planner.
+		clusterInstance.VtGateExtraArgs = []string{"-planner_version", "Gen4"} // enable Gen4 planner.
 		err = clusterInstance.StartVtgate()
 		if err != nil {
 			return 1

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

@@ -37,7 +37,7 @@ type (
 	QueryProjection struct {
 		SelectExprs  []SelectExpr
 		HasAggr      bool
-		GroupByExprs sqlparser.Exprs
+		GroupByExprs []GroupBy
 		OrderExprs   []OrderBy
 	}
 
@@ -46,6 +46,12 @@ type (
 		Inner         *sqlparser.Order
 		WeightStrExpr sqlparser.Expr
 	}
+
+	// GroupBy contains the expression to used in group by and also if grouping is needed at VTGate level then what the weight_string function expression to be sent down for evaluation.
+	GroupBy struct {
+		Inner         sqlparser.Expr
+		WeightStrExpr sqlparser.Expr
+	}
 )
 
 // CreateQPFromSelect created the QueryProjection for the input *sqlparser.Select
@@ -85,10 +91,10 @@ func CreateQPFromSelect(sel *sqlparser.Select) (*QueryProjection, error) {
 		return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.MixOfGroupFuncAndFields, "Mixing of aggregation and non-aggregation columns is not allowed if there is no GROUP BY clause")
 	}
 
-	for _, expr := range sel.GroupBy {
+	for _, group := range sel.GroupBy {
 		// todo dont ignore weightstringexpr
-		e, _ := qp.getSimplifiedExpr(expr)
-		qp.GroupByExprs = append(qp.GroupByExprs, e)
+		expr, weightStrExpr := qp.getSimplifiedExpr(group)
+		qp.GroupByExprs = append(qp.GroupByExprs, GroupBy{Inner: expr, WeightStrExpr: weightStrExpr})
 	}
 
 	for _, order := range sel.OrderBy {
@@ -176,7 +182,7 @@ func (qp *QueryProjection) toString() string {
 	}
 
 	for _, expr := range qp.GroupByExprs {
-		out.Grouping = append(out.Grouping, sqlparser.String(expr))
+		out.Grouping = append(out.Grouping, sqlparser.String(expr.Inner))
 	}
 	for _, expr := range qp.OrderExprs {
 		out.OrderBy = append(out.OrderBy, sqlparser.String(expr.Inner))

go/vt/vtgate/planbuilder/route_planning.go

@@ -174,31 +174,49 @@ func planHorizon(sel *sqlparser.Select, plan logicalPlan, semTable *semantics.Se
 		return nil, err
 	}
 
+	var needsTruncation bool
 	if qp.HasAggr {
-		plan, err = planAggregations(qp, plan, semTable)
+		plan, needsTruncation, err = planAggregations(qp, plan, semTable)
 		if err != nil {
 			return nil, err
 		}
 	} else {
 		for _, e := range qp.SelectExprs {
-			if _, _, err := pushProjection(e.Col, plan, semTable, true); err != nil {
+			if _, _, err := pushProjection(e.Col, plan, semTable, true, false); err != nil {
 				return nil, err
 			}
 		}
 	}
 
 	if len(qp.OrderExprs) > 0 {
-		plan, err = planOrderBy(qp, qp.OrderExprs, plan, semTable)
+		var colAdded bool
+		plan, colAdded, err = planOrderBy(qp, qp.OrderExprs, plan, semTable)
 		if err != nil {
 			return nil, err
 		}
+		needsTruncation = needsTruncation || colAdded
 	}
 
 	if qp.HasAggr {
-		plan, err = planOrderByUsingGroupBy(qp, plan, semTable)
+		var colAdded bool
+		plan, colAdded, err = planOrderByUsingGroupBy(qp, plan, semTable)
 		if err != nil {
 			return nil, err
 		}
+		needsTruncation = needsTruncation || colAdded
+	}
+
+	if needsTruncation {
+		switch p := plan.(type) {
+		case *route:
+			p.eroute.SetTruncateColumnCount(sel.GetColumnCount())
+		case *orderedAggregate:
+			p.eaggr.SetTruncateColumnCount(sel.GetColumnCount())
+		case *memorySort:
+			p.truncater.SetTruncateColumnCount(sel.GetColumnCount())
+		default:
+			return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "plan type not known for column truncation: %T", plan)
+		}
 	}
 
 	return plan, nil