make sure to add weight_string for the right expression (#16325)

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

go/vt/vtgate/planbuilder/operators/aggregation_pushing.go

@@ -177,6 +177,7 @@ func pushAggregations(ctx *plancontext.PlanningContext, aggregator *Aggregator,
 		if !aggr.Distinct || canPushDistinctAggr {
 			aggrBelowRoute.Aggregations = append(aggrBelowRoute.Aggregations, aggr)
 			aggregateTheAggregate(aggregator, i)
+			aggregator.Aggregations[i].PushedDown = true
 			continue
 		}
 

go/vt/vtgate/planbuilder/operators/aggregator.go

@@ -391,8 +391,16 @@ func (a *Aggregator) planOffsets(ctx *plancontext.PlanningContext) Operator {
 		if !aggr.NeedsWeightString(ctx) {
 			continue
 		}
-		arg := aggr.getPushColumn()
-		offset := a.internalAddColumn(ctx, aeWrap(weightStringFor(arg)), true)
+		var offset int
+		if aggr.PushedDown {
+			// if we have already pushed down aggregation, we need to use
+			// the weight string of the aggregation and not the argument
+			offset = a.internalAddColumn(ctx, aeWrap(weightStringFor(aggr.Func)), false)
+		} else {
+			// If we have not pushed down the aggregation, we need the weight_string of the argument
+			arg := aggr.getPushColumn()
+			offset = a.internalAddColumn(ctx, aeWrap(weightStringFor(arg)), true)
+		}
 		a.Aggregations[idx].WSOffset = offset
 	}
 	return nil

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

@@ -84,6 +84,8 @@ type (
 		WSOffset  int // Offset for the weight string of the column
 
 		SubQueryExpression []*SubQuery // Subqueries associated with this aggregation
+
+		PushedDown bool // Whether the aggregation has been pushed down to the next layer
 	}
 )
 

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

@@ -667,9 +667,9 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select col1, min(col2) as `min(distinct col2)`, weight_string(col1), weight_string(col2) from `user` where 1 != 1 group by col1, weight_string(col1), weight_string(col2)",
+            "FieldQuery": "select col1, min(col2) as `min(distinct col2)`, weight_string(col1), weight_string(min(col2)) from `user` where 1 != 1 group by col1, weight_string(col1)",
             "OrderBy": "(0|2) ASC",
-            "Query": "select col1, min(col2) as `min(distinct col2)`, weight_string(col1), weight_string(col2) from `user` group by col1, weight_string(col1), weight_string(col2) order by col1 asc",
+            "Query": "select col1, min(col2) as `min(distinct col2)`, weight_string(col1), weight_string(min(col2)) from `user` group by col1, weight_string(col1) order by col1 asc",
             "Table": "`user`"
           }
         ]
@@ -2605,9 +2605,9 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select col1, min(id) as `min(distinct id)`, col3, weight_string(col1), weight_string(id), weight_string(col3) from `user` where 1 != 1 group by col1, col3, weight_string(col1), weight_string(id), weight_string(col3)",
+            "FieldQuery": "select col1, min(id) as `min(distinct id)`, col3, weight_string(col1), weight_string(min(id)), weight_string(col3) from `user` where 1 != 1 group by col1, col3, weight_string(col1), weight_string(col3)",
             "OrderBy": "(0|3) ASC, (2|5) ASC",
-            "Query": "select col1, min(id) as `min(distinct id)`, col3, weight_string(col1), weight_string(id), weight_string(col3) from `user` group by col1, col3, weight_string(col1), weight_string(id), weight_string(col3) order by col1 asc, col3 asc",
+            "Query": "select col1, min(id) as `min(distinct id)`, col3, weight_string(col1), weight_string(min(id)), weight_string(col3) from `user` group by col1, col3, weight_string(col1), weight_string(col3) order by col1 asc, col3 asc",
             "Table": "`user`"
           }
         ]
@@ -4023,9 +4023,9 @@
                   "Name": "user",
                   "Sharded": true
                 },
-                "FieldQuery": "select foo, max(baz) as bazo, weight_string(foo), weight_string(baz) from (select foo, baz from `user` where 1 != 1) as f where 1 != 1 group by foo, weight_string(foo), weight_string(baz)",
+                "FieldQuery": "select foo, max(baz) as bazo, weight_string(foo), weight_string(max(baz)) from (select foo, baz from `user` where 1 != 1) as f where 1 != 1 group by foo, weight_string(foo)",
                 "OrderBy": "(0|2) ASC",
-                "Query": "select foo, max(baz) as bazo, weight_string(foo), weight_string(baz) from (select foo, baz from `user`) as f group by foo, weight_string(foo), weight_string(baz) order by foo asc",
+                "Query": "select foo, max(baz) as bazo, weight_string(foo), weight_string(max(baz)) from (select foo, baz from `user`) as f group by foo, weight_string(foo) order by foo asc",
                 "Table": "`user`"
               }
             ]
@@ -4949,9 +4949,9 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select min(textcol1), max(textcol2), textcol1, textcol1, weight_string(textcol2) from `user` where 1 != 1 group by textcol1, weight_string(textcol2)",
+            "FieldQuery": "select min(textcol1), max(textcol2), textcol1, textcol1, weight_string(max(textcol2)) from `user` where 1 != 1 group by textcol1",
             "OrderBy": "2 ASC COLLATE latin1_swedish_ci",
-            "Query": "select min(textcol1), max(textcol2), textcol1, textcol1, weight_string(textcol2) from `user` group by textcol1, weight_string(textcol2) order by textcol1 asc",
+            "Query": "select min(textcol1), max(textcol2), textcol1, textcol1, weight_string(max(textcol2)) from `user` group by textcol1 order by textcol1 asc",
             "Table": "`user`"
           }
         ]
@@ -4981,9 +4981,9 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select col, min(textcol1), max(textcol2), textcol1, textcol1, weight_string(textcol2) from `user` where 1 != 1 group by col, textcol1, weight_string(textcol2)",
+            "FieldQuery": "select col, min(textcol1), max(textcol2), textcol1, textcol1, weight_string(max(textcol2)) from `user` where 1 != 1 group by col, textcol1",
             "OrderBy": "0 ASC, 3 ASC COLLATE latin1_swedish_ci",
-            "Query": "select col, min(textcol1), max(textcol2), textcol1, textcol1, weight_string(textcol2) from `user` group by col, textcol1, weight_string(textcol2) order by col asc, textcol1 asc",
+            "Query": "select col, min(textcol1), max(textcol2), textcol1, textcol1, weight_string(max(textcol2)) from `user` group by col, textcol1 order by col asc, textcol1 asc",
             "Table": "`user`"
           }
         ]
@@ -5868,9 +5868,9 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select foo, min(bar) as `min(distinct bar)`, baz, baz, max(toto) as `max(distinct toto)`, weight_string(foo), weight_string(bar), weight_string(baz), weight_string(toto) from `user` where 1 != 1 group by foo, baz, weight_string(foo), weight_string(bar), weight_string(baz), weight_string(toto)",
+            "FieldQuery": "select foo, min(bar) as `min(distinct bar)`, baz, baz, max(toto) as `max(distinct toto)`, weight_string(foo), weight_string(min(bar)), weight_string(baz), weight_string(max(toto)) from `user` where 1 != 1 group by foo, baz, weight_string(foo), weight_string(baz)",
             "OrderBy": "(0|5) ASC, (2|7) ASC",
-            "Query": "select foo, min(bar) as `min(distinct bar)`, baz, baz, max(toto) as `max(distinct toto)`, weight_string(foo), weight_string(bar), weight_string(baz), weight_string(toto) from `user` group by foo, baz, weight_string(foo), weight_string(bar), weight_string(baz), weight_string(toto) order by foo asc, baz asc",
+            "Query": "select foo, min(bar) as `min(distinct bar)`, baz, baz, max(toto) as `max(distinct toto)`, weight_string(foo), weight_string(min(bar)), weight_string(baz), weight_string(max(toto)) from `user` group by foo, baz, weight_string(foo), weight_string(baz) order by foo asc, baz asc",
             "Table": "`user`"
           }
         ]
@@ -6424,7 +6424,7 @@
                       {
                         "OperatorType": "Aggregate",
                         "Variant": "Scalar",
-                        "Aggregates": "min(0|2) AS min_id, max(1|2) AS max_id",
+                        "Aggregates": "min(0|2) AS min_id, max(1|3) AS max_id",
                         "ResultColumns": 2,
                         "Inputs": [
                           {
@@ -6434,9 +6434,9 @@
                               "Name": "user",
                               "Sharded": true
                             },
-                            "FieldQuery": "select min(bb.id) as min_id, max(bb.id) as max_id, weight_string(bb.id) from `user` as bb where 1 != 1 group by weight_string(bb.id)",
+                            "FieldQuery": "select min(bb.id) as min_id, max(bb.id) as max_id, weight_string(min(bb.id)), weight_string(max(bb.id)) from `user` as bb where 1 != 1",
                             "OrderBy": "0 ASC COLLATE utf8mb4_0900_ai_ci",
-                            "Query": "select min(bb.id) as min_id, max(bb.id) as max_id, weight_string(bb.id) from `user` as bb group by weight_string(bb.id) order by min(bb.id) asc",
+                            "Query": "select min(bb.id) as min_id, max(bb.id) as max_id, weight_string(min(bb.id)), weight_string(max(bb.id)) from `user` as bb order by min(bb.id) asc",
                             "Table": "`user`"
                           }
                         ]
@@ -6909,8 +6909,8 @@
               "Name": "user",
               "Sharded": true
             },
-            "FieldQuery": "select max((select max(col2) from `user` as u1 where 1 != 1)), weight_string((select max(col2) from `user` as u1 where 1 != 1)) from `user` as u2 where 1 != 1 group by weight_string((select max(col2) from `user` as u1 where 1 != 1))",
-            "Query": "select max((select max(col2) from `user` as u1 where u1.id = u2.id)), weight_string((select max(col2) from `user` as u1 where u1.id = u2.id)) from `user` as u2 group by weight_string((select max(col2) from `user` as u1 where u1.id = u2.id))",
+            "FieldQuery": "select max((select max(col2) from `user` as u1 where 1 != 1)), weight_string(max((select max(col2) from `user` as u1 where 1 != 1))) from `user` as u2 where 1 != 1",
+            "Query": "select max((select max(col2) from `user` as u1 where u1.id = u2.id)), weight_string(max((select max(col2) from `user` as u1 where u1.id = u2.id))) from `user` as u2",
             "Table": "`user`"
           }
         ]

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

@@ -488,9 +488,9 @@
                   "Name": "user",
                   "Sharded": true
                 },
-                "FieldQuery": "select foo, max(baz) as bazo, weight_string(foo), weight_string(baz) from (select foo, baz from `user` where 1 != 1) as f where 1 != 1 group by foo, weight_string(foo), weight_string(baz)",
+                "FieldQuery": "select foo, max(baz) as bazo, weight_string(foo), weight_string(max(baz)) from (select foo, baz from `user` where 1 != 1) as f where 1 != 1 group by foo, weight_string(foo)",
                 "OrderBy": "(0|2) ASC",
-                "Query": "select foo, max(baz) as bazo, weight_string(foo), weight_string(baz) from (select foo, baz from `user`) as f group by foo, weight_string(foo), weight_string(baz) order by foo asc",
+                "Query": "select foo, max(baz) as bazo, weight_string(foo), weight_string(max(baz)) from (select foo, baz from `user`) as f group by foo, weight_string(foo) order by foo asc",
                 "Table": "`user`"
               }
             ]

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

@@ -4350,8 +4350,8 @@
                   "Name": "user",
                   "Sharded": true
                 },
-                "FieldQuery": "select max(music.id), weight_string(music.id) from music where 1 != 1 group by weight_string(music.id)",
-                "Query": "select max(music.id), weight_string(music.id) from music where music.user_id in ::__vals group by weight_string(music.id)",
+                "FieldQuery": "select max(music.id), weight_string(max(music.id)) from music where 1 != 1",
+                "Query": "select max(music.id), weight_string(max(music.id)) from music where music.user_id in ::__vals",
                 "Table": "music",
                 "Values": [
                   "(5, 6)"

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

@@ -1908,8 +1908,8 @@
                   "Name": "main",
                   "Sharded": true
                 },
-                "FieldQuery": "select max(total_revenue), weight_string(total_revenue) from revenue0 where 1 != 1 group by weight_string(total_revenue)",
-                "Query": "select max(total_revenue), weight_string(total_revenue) from revenue0 group by weight_string(total_revenue)",
+                "FieldQuery": "select max(total_revenue), weight_string(max(total_revenue)) from revenue0 where 1 != 1",
+                "Query": "select max(total_revenue), weight_string(max(total_revenue)) from revenue0",
                 "Table": "revenue0"
               }
             ]