feat: use the UDF info when planning

Signed-off-by: Andres Taylor <[email protected]> Signed-off-by: Harshit Gangal <[email protected]>
vitessio · Apr 15, 2024 · 28e792a · 28e792a
1 parent 6b25965
commit 28e792a
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 2 deletions.
diff --git a/go/vt/vtgate/semantics/analyzer.go b/go/vt/vtgate/semantics/analyzer.go
@@ -78,6 +78,7 @@ func (a *analyzer) lateInit() {
 		aliasMapCache:   map[*sqlparser.Select]map[string]exprContainer{},
 		reAnalyze:       a.reAnalyze,
 		tables:          a.tables,
+		aggrUDFs:        a.si.GetAggregateUDFs(),
 	}
 	a.fk = &fkManager{
 		binder:   a.binder,

diff --git a/go/vt/vtgate/semantics/early_rewriter.go b/go/vt/vtgate/semantics/early_rewriter.go
@@ -41,6 +41,7 @@ type earlyRewriter struct {
 	// have happened, and we are introducing or changing the AST. We invoke it so all parts of the query have been
 	// typed, scoped and bound correctly
 	reAnalyze func(n sqlparser.SQLNode) error
+	aggrUDFs  []string
 }
 
 func (r *earlyRewriter) down(cursor *sqlparser.Cursor) error {
@@ -500,6 +501,15 @@ func (r *earlyRewriter) rewriteAliasesInGroupBy(node sqlparser.Expr, sel *sqlpar
 	return
 }
 
+func (at *aggrTracker) isAggregateUDF(name sqlparser.IdentifierCI) bool {
+	for _, aggrUDF := range at.aggrUDFs {
+		if name.EqualString(aggrUDF) {
+			return true
+		}
+	}
+	return false
+}
+
 func (r *earlyRewriter) rewriteAliasesInHaving(node sqlparser.Expr, sel *sqlparser.Select) (expr sqlparser.Expr, err error) {
 	currentScope := r.scoper.currentScope()
 	if currentScope.isUnion {
@@ -508,14 +518,22 @@ func (r *earlyRewriter) rewriteAliasesInHaving(node sqlparser.Expr, sel *sqlpars
 	}
 
 	aliases := r.getAliasMap(sel)
-	aggrTrack := &aggrTracker{}
+	aggrTrack := &aggrTracker{
+		insideAggr: false,
+		aggrUDFs:   r.aggrUDFs,
+	}
 	output := sqlparser.CopyOnRewrite(node, aggrTrack.down, func(cursor *sqlparser.CopyOnWriteCursor) {
 		var col *sqlparser.ColName
 
 		switch node := cursor.Node().(type) {
 		case sqlparser.AggrFunc:
 			aggrTrack.popAggr()
 			return
+		case *sqlparser.FuncExpr:
+			if aggrTrack.isAggregateUDF(node.Name) {
+				aggrTrack.popAggr()
+			}
+			return
 		case *sqlparser.ColName:
 			col = node
 		default:
@@ -565,14 +583,19 @@ func (r *earlyRewriter) rewriteAliasesInHaving(node sqlparser.Expr, sel *sqlpars
 
 type aggrTracker struct {
 	insideAggr bool
+	aggrUDFs   []string
 }
 
 func (at *aggrTracker) down(node, _ sqlparser.SQLNode) bool {
-	switch node.(type) {
+	switch node := node.(type) {
 	case *sqlparser.Subquery:
 		return false
 	case sqlparser.AggrFunc:
 		at.insideAggr = true
+	case *sqlparser.FuncExpr:
+		if at.isAggregateUDF(node.Name) {
+			at.insideAggr = true
+		}
 	}
 
 	return true

diff --git a/go/vt/vtgate/semantics/early_rewriter_test.go b/go/vt/vtgate/semantics/early_rewriter_test.go
@@ -536,6 +536,11 @@ func TestHavingColumnName(t *testing.T) {
 		expSQL:  "select id, sum(t1.foo) as foo from t1 having sum(foo) > 1",
 		expDeps: TS0,
 		warning: "Column 'foo' in having clause is ambiguous",
+	}, {
+		sql:     "select id, sum(t1.foo) as foo from t1 having custom_udf(foo) > 1",
+		expSQL:  "select id, sum(t1.foo) as foo from t1 having custom_udf(foo) > 1",
+		expDeps: TS0,
+		warning: "Column 'foo' in having clause is ambiguous",
 	}, {
 		sql:    "select id, sum(t1.foo) as XYZ from t1 having sum(XYZ) > 1",
 		expErr: "Invalid use of group function",