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

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

@@ -332,6 +332,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

@@ -68,9 +68,15 @@ type (
 
 	// Aggr encodes all information needed for aggregation functions
 	Aggr struct {
+<<<<<<< HEAD
 		Original *sqlparser.AliasedExpr
 		Func     sqlparser.AggrFunc // if we are missing a Func, it means this is a AggregateAnyValue
 		OpCode   opcode.AggregateOpcode
+=======
+		Original *sqlparser.AliasedExpr // The original SQL expression for the aggregation
+		Func     sqlparser.AggrFunc     // The aggregation function (e.g., COUNT, SUM). If nil, it means AggregateAnyValue or AggregateUDF is used
+		OpCode   opcode.AggregateOpcode // The opcode representing the type of aggregation being performed
+>>>>>>> af8f42af13 (Fix panic in user defined aggregation functions planning (#16398))
 
 		// OriginalOpCode will contain opcode.AggregateUnassigned unless we are changing opcode while pushing them down
 		OriginalOpCode opcode.AggregateOpcode

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

@@ -2114,6 +2114,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

@@ -20,6 +20,7 @@
     "plan": "VT12001: unsupported: subqueries in GROUP BY"
   },
   {
+<<<<<<< HEAD
     "comment": "subqueries in delete",
     "query": "delete from user where col = (select id from unsharded)",
     "plan": "VT12001: unsupported: subqueries in DML"
@@ -43,6 +44,11 @@
     "comment": "scatter update with limit clause",
     "query": "update user_extra set val = 1 where (name = 'foo' or id = 1) limit 1",
     "plan": "VT12001: unsupported: multi shard UPDATE with LIMIT"
+=======
+    "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"
+>>>>>>> af8f42af13 (Fix panic in user defined aggregation functions planning (#16398))
   },
   {
     "comment": "update changes primary vindex column",