[release-20.0] Fix Incorrect Optimization with LIMIT and GROUP BY (#1…

…6263) (#16268)

Signed-off-by: Andres Taylor <andres@planetscale.com>
Signed-off-by: Florent Poinsard <35779988+frouioui@users.noreply.github.com>
Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>
Co-authored-by: vitess-bot[bot] <108069721+vitess-bot[bot]@users.noreply.github.com>
Co-authored-by: Florent Poinsard <35779988+frouioui@users.noreply.github.com>
Co-authored-by: Andres Taylor <andres@planetscale.com>

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,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 != 20 {
+ t.Skip("Test requires VTGate version 20")
+ }
+ mcmp, closer := start(t)
+ defer closer()
+
+ for i := range 1000 {
+ 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

@@ -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

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