feat: add support for WITH ROLLUP (#15930)

go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go

@@ -101,14 +101,19 @@ func TestGroupBy(t *testing.T) {
 	mcmp, closer := start(t)
 	defer closer()
 	mcmp.Exec("insert into t3(id5, id6, id7) values(1,1,2), (2,2,4), (3,2,4), (4,1,2), (5,1,2), (6,3,6)")
-	// test ordering and group by int column
-	mcmp.AssertMatches("select id6, id7, count(*) k from t3 group by id6, id7 order by k", `[[INT64(3) INT64(6) INT64(1)] [INT64(2) INT64(4) INT64(2)] [INT64(1) INT64(2) INT64(3)]]`)
-	mcmp.AssertMatches("select id6+id7, count(*) k from t3 group by id6+id7 order by k", `[[INT64(9) INT64(1)] [INT64(6) INT64(2)] [INT64(3) INT64(3)]]`)
-
-	// Test the same queries in streaming mode
-	utils.Exec(t, mcmp.VtConn, "set workload = olap")
-	mcmp.AssertMatches("select id6, id7, count(*) k from t3 group by id6, id7 order by k", `[[INT64(3) INT64(6) INT64(1)] [INT64(2) INT64(4) INT64(2)] [INT64(1) INT64(2) INT64(3)]]`)
-	mcmp.AssertMatches("select id6+id7, count(*) k from t3 group by id6+id7 order by k", `[[INT64(9) INT64(1)] [INT64(6) INT64(2)] [INT64(3) INT64(3)]]`)
+
+	// run queries in both workloads
+	workloads := []string{"oltp", "olap"}
+	for _, workload := range workloads {
+		utils.Exec(t, mcmp.VtConn, fmt.Sprintf("set workload = %s", workload))
+		// test ordering and group by int column
+		mcmp.AssertMatches("select id6, id7, count(*) k from t3 group by id6, id7 order by k", `[[INT64(3) INT64(6) INT64(1)] [INT64(2) INT64(4) INT64(2)] [INT64(1) INT64(2) INT64(3)]]`)
+		mcmp.AssertMatches("select id6+id7, count(*) k from t3 group by id6+id7 order by k", `[[INT64(9) INT64(1)] [INT64(6) INT64(2)] [INT64(3) INT64(3)]]`)
+		if utils.BinaryIsAtLeastAtVersion(20, "vtgate") &&
+			utils.BinaryIsAtLeastAtVersion(20, "vttablet") {
+			mcmp.Exec("select id6, id7, count(*) k from t3 group by id6, id7 with rollup")
+		}
+	}
 }
 
 func TestEqualFilterOnScatter(t *testing.T) {

go/test/endtoend/vtgate/queries/random/query_gen.go

@@ -423,7 +423,7 @@ func (sg *selectGenerator) createGroupBy(tables []tableT) (grouping []column) {
 		if col.typ == "date" && !testFailingQueries {
 			continue
 		}
-		sg.sel.GroupBy = append(sg.sel.GroupBy, col.getASTExpr())
+		sg.sel.AddGroupBy(col.getASTExpr())
 
 		// add to select
 		if rand.IntN(2) < 1 {
@@ -478,8 +478,10 @@ func (sg *selectGenerator) createAggregations(tables []tableT) (aggregates []col
 // orders on all grouping expressions and on random SelectExprs
 func (sg *selectGenerator) createOrderBy() {
 	// always order on grouping expressions
-	for _, expr := range sg.sel.GroupBy {
-		sg.sel.OrderBy = append(sg.sel.OrderBy, sqlparser.NewOrder(expr, getRandomOrderDirection()))
+	if sg.sel.GroupBy != nil {
+		for _, expr := range sg.sel.GroupBy.Exprs {
+			sg.sel.OrderBy = append(sg.sel.OrderBy, sqlparser.NewOrder(expr, getRandomOrderDirection()))
+		}
 	}
 
 	// randomly order on SelectExprs

go/vt/sqlparser/ast.go

@@ -269,7 +269,7 @@ type (
 		Comments    *ParsedComments
 		SelectExprs SelectExprs
 		Where       *Where
-		GroupBy     GroupBy
+		GroupBy     *GroupBy
 		Having      *Where
 		Windows     NamedWindows
 		OrderBy     OrderBy
@@ -3451,7 +3451,10 @@ type ConvertType struct {
 }
 
 // GroupBy represents a GROUP BY clause.
-type GroupBy []Expr
+type GroupBy struct {
+	Exprs      []Expr
+	WithRollup bool
+}
 
 // OrderBy represents an ORDER By clause.
 type OrderBy []*Order

go/vt/sqlparser/ast_format.go

@@ -54,7 +54,8 @@ func (node *Select) Format(buf *TrackedBuffer) {
 
 	buf.astPrintf(node, "%v%v%v",
 		node.Where,
-		node.GroupBy, node.Having)
+		node.GroupBy,
+		node.Having)
 
 	if node.Windows != nil {
 		buf.astPrintf(node, " %v", node.Windows)
@@ -1922,12 +1923,18 @@ func (node *When) Format(buf *TrackedBuffer) {
 }
 
 // Format formats the node.
-func (node GroupBy) Format(buf *TrackedBuffer) {
+func (node *GroupBy) Format(buf *TrackedBuffer) {
+	if node == nil || len(node.Exprs) == 0 {
+		return
+	}
 	prefix := " group by "
-	for _, n := range node {
+	for _, n := range node.Exprs {
 		buf.astPrintf(node, "%s%v", prefix, n)
 		prefix = ", "
 	}
+	if node.WithRollup {
+		buf.literal(" with rollup")
+	}
 }
 
 // Format formats the node.