[release-18.0] Fix Incorrect Optimization with LIMIT and GROUP BY (vitessio#16263) (vitessio#16266)

Signed-off-by: Andres Taylor <[email protected]>
Signed-off-by: Florent Poinsard <[email protected]>
Co-authored-by: Andrés Taylor <[email protected]>

Author: vitess-bot[bot]

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

@@ -18,6 +18,7 @@ package aggregation
 
 import (
 	"fmt"
+	"math/rand"
 	"slices"
 	"sort"
 	"strings"
@@ -66,6 +67,22 @@ func start(t *testing.T) (utils.MySQLCompare, func()) {
 	}
 }
 
+func TestAggrWithLimit(t *testing.T) {
+	version, err := cluster.GetMajorVersion("vtgate")
+	require.NoError(t, err)
+	if version != 18 {
+		t.Skip("Test requires VTGate version 18")
+	}
+	mcmp, closer := start(t)
+	defer closer()
+
+	for i := 0; i < 1000; i++ {
+		r := rand.Intn(50)
+		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

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

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

@@ -147,7 +147,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`"
                   }
                 ]