use aggregation engine over distinct engine when overlapping order by…

… used

Signed-off-by: Harshit Gangal <harshit@planetscale.com>

go/vt/vtgate/planbuilder/operators/queryprojection.go

@@ -387,8 +387,13 @@ func (qp *QueryProjection) addOrderBy(ctx *plancontext.PlanningContext, orderBy
 
 func (qp *QueryProjection) calculateDistinct(ctx *plancontext.PlanningContext) {
  if qp.Distinct && !qp.HasAggr {
- // grouping and distinct both lead to unique results, so we don't need
- qp.groupByExprs = nil
+ if qp.useGroupingOverDistinct(ctx) {
+ // if order by exists with overlap with select expressions, we can use the aggregation with ordering over distinct.
+ qp.Distinct = false
+ } else {
+ // grouping and distinct both lead to unique results, so we don't need
+ qp.groupByExprs = nil
+ }
  }
 
  if qp.HasAggr && len(qp.groupByExprs) == 0 {
@@ -850,6 +855,45 @@ func (qp *QueryProjection) GetColumnCount() int {
  return len(qp.SelectExprs) - qp.AddedColumn
 }
 
+func (qp *QueryProjection) orderByOverlapWithSelectExpr(ctx *plancontext.PlanningContext) bool {
+ for _, expr := range qp.OrderExprs {
+ idx, _ := qp.FindSelectExprIndexForExpr(ctx, expr.SimplifiedExpr)
+ if idx != nil {
+ return true
+ }
+ }
+ return false
+}
+
+func (qp *QueryProjection) useGroupingOverDistinct(ctx *plancontext.PlanningContext) bool {
+ if !qp.orderByOverlapWithSelectExpr(ctx) {
+ return false
+ }
+ var gbs []GroupBy
+ for idx, selExpr := range qp.SelectExprs {
+ ae, err := selExpr.GetAliasedExpr()
+ if err != nil {
+ // not an alias Expr, cannot continue forward.
+ return false
+ }
+ sExpr := qp.GetSimplifiedExpr(ae.Expr)
+ // check if the grouping already exists on that column.
+ found := slices.IndexFunc(qp.groupByExprs, func(gb GroupBy) bool {
+ return ctx.SemTable.EqualsExprWithDeps(gb.SimplifiedExpr, sExpr)
+ })
+ if found != -1 {
+ continue
+ }
+ groupBy := NewGroupBy(ae.Expr, sExpr, ae)
+ selectExprIdx := idx
+ groupBy.InnerIndex = &selectExprIdx
+
+ gbs = append(gbs, groupBy)
+ }
+ qp.groupByExprs = append(qp.groupByExprs, gbs...)
+ return true
+}
+
 func checkForInvalidGroupingExpressions(expr sqlparser.Expr) error {
  return sqlparser.Walk(func(node sqlparser.SQLNode) (bool, error) {
  if _, isAggregate := node.(sqlparser.AggrFunc); isAggregate {

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

@@ -106,28 +106,21 @@
  "QueryType": "SELECT",
  "Original": "SELECT DISTINCT c FROM sbtest30 WHERE id BETWEEN 1 AND 10 ORDER BY c",
  "Instructions": {
- "OperatorType": "Sort",
- "Variant": "Memory",
- "OrderBy": "0 ASC COLLATE latin1_swedish_ci",
+ "OperatorType": "Aggregate",
+ "Variant": "Ordered",
+ "GroupBy": "0 COLLATE latin1_swedish_ci",
  "Inputs": [
  {
- "OperatorType": "Distinct",
- "Collations": [
- "0: latin1_swedish_ci"
- ],
- "Inputs": [
- {
- "OperatorType": "Route",
- "Variant": "Scatter",
- "Keyspace": {
- "Name": "main",
- "Sharded": true
- },
- "FieldQuery": "select c from sbtest30 where 1 != 1",
- "Query": "select distinct c from sbtest30 where id between 1 and 10",
- "Table": "sbtest30"
- }
- ]
+ "OperatorType": "Route",
+ "Variant": "Scatter",
+ "Keyspace": {
+ "Name": "main",
+ "Sharded": true
+ },
+ "FieldQuery": "select c from sbtest30 where 1 != 1 group by c",
+ "OrderBy": "0 ASC COLLATE latin1_swedish_ci",
+ "Query": "select c from sbtest30 where id between 1 and 10 group by c order by c asc",
+ "Table": "sbtest30"
  }
  ]
  },

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

@@ -2081,5 +2081,105 @@
  "user.user"
  ]
  }
+ },
+ {
+ "comment": "distinct with order by using aggregation engine",
+ "query": "select distinct col from user where id between :vtg1 and :vtg2 order by col asc",
+ "plan": {
+ "QueryType": "SELECT",
+ "Original": "select distinct col from user where id between :vtg1 and :vtg2 order by col asc",
+ "Instructions": {
+ "OperatorType": "Aggregate",
+ "Variant": "Ordered",
+ "GroupBy": "0",
+ "Inputs": [
+ {
+ "OperatorType": "Route",
+ "Variant": "Scatter",
+ "Keyspace": {
+ "Name": "user",
+ "Sharded": true
+ },
+ "FieldQuery": "select col from `user` where 1 != 1 group by col",
+ "OrderBy": "0 ASC",
+ "Query": "select col from `user` where id between :vtg1 and :vtg2 group by col order by col asc",
+ "Table": "`user`"
+ }
+ ]
+ },
+ "TablesUsed": [
+ "user.user"
+ ]
+ }
+ },
+ {
+ "comment": "distinct with order by having additional non-order by columns in the selection using aggregation engine",
+ "query": "select distinct foo, col from user where id between :vtg1 and :vtg2 order by col asc",
+ "plan": {
+ "QueryType": "SELECT",
+ "Original": "select distinct foo, col from user where id between :vtg1 and :vtg2 order by col asc",
+ "Instructions": {
+ "OperatorType": "Aggregate",
+ "Variant": "Ordered",
+ "GroupBy": "1, (0|2)",
+ "ResultColumns": 2,
+ "Inputs": [
+ {
+ "OperatorType": "Route",
+ "Variant": "Scatter",
+ "Keyspace": {
+ "Name": "user",
+ "Sharded": true
+ },
+ "FieldQuery": "select foo, col, weight_string(foo) from `user` where 1 != 1 group by col, foo, weight_string(foo)",
+ "OrderBy": "1 ASC, (0|2) ASC",
+ "Query": "select foo, col, weight_string(foo) from `user` where id between :vtg1 and :vtg2 group by col, foo, weight_string(foo) order by col asc, foo asc",
+ "Table": "`user`"
+ }
+ ]
+ },
+ "TablesUsed": [
+ "user.user"
+ ]
+ }
+ },
+ {
+ "comment": "distinct with order by having no overalap with the selection columns - using distinct engine",
+ "query": "select distinct foo from user where id between :vtg1 and :vtg2 order by col asc",
+ "plan": {
+ "QueryType": "SELECT",
+ "Original": "select distinct foo from user where id between :vtg1 and :vtg2 order by col asc",
+ "Instructions": {
+ "OperatorType": "Sort",
+ "Variant": "Memory",
+ "OrderBy": "1 ASC",
+ "ResultColumns": 1,
+ "Inputs": [
+ {
+ "OperatorType": "Distinct",
+ "Collations": [
+ "(0:2)",
+ "1"
+ ],
+ "Inputs": [
+ {
+ "OperatorType": "Route",
+ "Variant": "Scatter",
+ "Keyspace": {
+ "Name": "user",
+ "Sharded": true
+ },
+ "FieldQuery": "select foo, col, weight_string(foo) from `user` where 1 != 1",
+ "Query": "select distinct foo, col, weight_string(foo) from `user` where id between :vtg1 and :vtg2",
+ "Table": "`user`"
+ }
+ ]
+ }
+ ]
+ },
+ "TablesUsed": [
+ "user.user"
+ ]
+ }
  }
 ]