Fix Incorrect Optimization with LIMIT and GROUP BY (#16263)

Signed-off-by: Andres Taylor <[email protected]>

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

@@ -18,6 +18,7 @@ package aggregation
 
 import (
 	"fmt"
+	"math/rand/v2"
 	"slices"
 	"sort"
 	"strings"
@@ -69,6 +70,18 @@ func start(t *testing.T) (utils.MySQLCompare, func()) {
 	}
 }
 
+func TestAggrWithLimit(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 21, "vtgate")
+	mcmp, closer := start(t)
+	defer closer()
+
+	for i := range 1000 {
+		r := rand.IntN(10)
+		mcmp.Exec(fmt.Sprintf("insert into aggr_test(id, val1, val2) values(%d, 'a', %d)", i, r))
+	}
+	mcmp.Exec("select val2, count(*) from aggr_test group by val2 order by count(*), val2 limit 10")
+}
+
 func TestAggregateTypes(t *testing.T) {
 	mcmp, closer := start(t)
 	defer closer()

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

@@ -490,6 +490,11 @@ func setUpperLimit(in *Limit) (Operator, *ApplyResult) {
 		case *Join, *ApplyJoin, *SubQueryContainer, *SubQuery:
 			// we can't push limits down on either side
 			return SkipChildren
+		case *Aggregator:
+			if len(op.Grouping) > 0 {
+				// we can't push limits down if we have a group by
+				return SkipChildren
+			}
 		case *Route:
 			newSrc := &Limit{
 				Source: op.Source,
@@ -498,9 +503,8 @@ func setUpperLimit(in *Limit) (Operator, *ApplyResult) {
 			op.Source = newSrc
 			result = result.Merge(Rewrote("push upper limit under route"))
 			return SkipChildren
-		default:
-			return VisitChildren
 		}
+		return VisitChildren
 	}
 
 	TopDown(in.Source, TableID, visitor, shouldVisit)

go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json

@@ -149,7 +149,7 @@
                     },
                     "FieldQuery": "select a, b, count(*) as k, weight_string(a) from `user` where 1 != 1 group by a, weight_string(a)",
                     "OrderBy": "(0|3) ASC",
-                    "Query": "select a, b, count(*) as k, weight_string(a) from `user` group by a, weight_string(a) order by a asc limit :__upper_limit",
+                    "Query": "select a, b, count(*) as k, weight_string(a) from `user` group by a, weight_string(a) order by a asc",
                     "Table": "`user`"
                   }
                 ]