From ee00de891b30931290f234ede22930e40f9268f6 Mon Sep 17 00:00:00 2001
From: Manan Gupta <35839558+GuptaManan100@users.noreply.github.com>
Date: Tue, 16 Jul 2024 16:24:29 +0530
Subject: [PATCH] Fix panic in user defined aggregation functions planning
 (#16398)

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 .../planbuilder/operators/aggregator.go       |  3 ++
 .../planbuilder/operators/queryprojection.go  |  2 +-
 .../planbuilder/testdata/aggr_cases.json      | 48 +++++++++++++++++++
 .../testdata/unsupported_cases.json           |  5 ++
 4 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/go/vt/vtgate/planbuilder/operators/aggregator.go b/go/vt/vtgate/planbuilder/operators/aggregator.go
index 91eca26c0ea..13acaa64526 100644
--- a/go/vt/vtgate/planbuilder/operators/aggregator.go
+++ b/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()
 	}
diff --git a/go/vt/vtgate/planbuilder/operators/queryprojection.go b/go/vt/vtgate/planbuilder/operators/queryprojection.go
index 306a245e602..68227e7f4cb 100644
--- a/go/vt/vtgate/planbuilder/operators/queryprojection.go
+++ b/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
diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
index 5f4671ca51b..6800a636752 100644
--- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
+++ b/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",
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
index 251af436d27..d79b78e38ce 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.json
+++ b/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",