[release-20.0] Fix panic in user defined aggregation functions planning (#16398)

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

@@ -411,6 +411,9 @@ func (aggr Aggr) getPushColumnExprs() sqlparser.Exprs {
 		return sqlparser.Exprs{aggr.Original.Expr}
 	case opcode.AggregateCountStar:
 		return sqlparser.Exprs{sqlparser.NewIntLiteral("1")}
+	case opcode.AggregateUDF:
+		// AggregateUDFs can't be evaluated on the vtgate. So either we are able to push everything down, or we will have to fail the query.
+		return nil
 	default:
 		return aggr.Func.GetArgs()
 	}

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

@@ -69,7 +69,7 @@ type (
 	// Aggr encodes all information needed for aggregation functions
 	Aggr struct {
 		Original *sqlparser.AliasedExpr
-		Func     sqlparser.AggrFunc // if we are missing a Func, it means this is a AggregateAnyValue
+		Func     sqlparser.AggrFunc // if we are missing a Func, it means this is a AggregateAnyValue or AggregateUDF
 		OpCode   opcode.AggregateOpcode
 
 		// OriginalOpCode will contain opcode.AggregateUnassigned unless we are changing opcode while pushing them down

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

@@ -2158,6 +2158,54 @@
       ]
     }
   },
+  {
+    "comment": "User defined aggregation expression being used in order by of a query that is single sharded",
+    "query": "select col1, udf_aggr( col2 ) r from user where id = 1 group by col1 having r >= 0.3",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select col1, udf_aggr( col2 ) r from user where id = 1 group by col1 having r >= 0.3",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select col1, udf_aggr(col2) as r from `user` where 1 != 1 group by col1",
+        "Query": "select col1, udf_aggr(col2) as r from `user` where id = 1 group by col1 having udf_aggr(`user`.col2) >= 0.3",
+        "Table": "`user`",
+        "Values": [
+          "1"
+        ],
+        "Vindex": "user_index"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
+  {
+    "comment": "user defined aggregation such that it can pushed to mysql in a scatter route",
+    "query": "select id, udf_aggr( col2 ) r from user group by id",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select id, udf_aggr( col2 ) r from user group by id",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select id, udf_aggr(col2) as r from `user` where 1 != 1 group by id",
+        "Query": "select id, udf_aggr(col2) as r from `user` group by id",
+        "Table": "`user`"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
   {
     "comment": "distinct on text column with collation",
     "query": "select col, count(distinct textcol1) from user group by col",

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

@@ -19,6 +19,11 @@
     "query": "select id from user group by id, (select id from user_extra)",
     "plan": "VT12001: unsupported: subqueries in GROUP BY"
   },
+  {
+    "comment": "user defined functions used in having clause that needs evaluation on vtgate",
+    "query": "select col1, udf_aggr( col2 ) r from user group by col1 having r >= 0.3",
+    "plan": "VT12001: unsupported: Aggregate UDF 'udf_aggr(col2)' must be pushed down to MySQL"
+  },
   {
     "comment": "update changes primary vindex column",
     "query": "update user set id = 1 where id = 1",