From 6e328dea25bebcc7458a0016452149c27b7e3375 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20Taylor?= <andres@planetscale.com>
Date: Wed, 26 Jun 2024 15:49:53 +0200
Subject: [PATCH] Fix Incorrect Optimization with LIMIT and GROUP BY (#16263)

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 .../vtgate/queries/aggregation/aggregation_test.go  | 13 +++++++++++++
 .../vtgate/planbuilder/operators/query_planning.go  |  8 ++++++--
 .../planbuilder/testdata/memory_sort_cases.json     |  2 +-
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go b/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go
index 8150f974d47..6020046ab82 100644
--- a/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go
+++ b/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()
diff --git a/go/vt/vtgate/planbuilder/operators/query_planning.go b/go/vt/vtgate/planbuilder/operators/query_planning.go
index 2f3259394e2..533d740f300 100644
--- a/go/vt/vtgate/planbuilder/operators/query_planning.go
+++ b/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)
diff --git a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json b/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json
index 3281feacc76..f26b160ef69 100644
--- a/go/vt/vtgate/planbuilder/testdata/memory_sort_cases.json
+++ b/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`"
                   }
                 ]