Resolve correlated aggregates at their outer scope

enginetest/queries/script_queries.go

@@ -6014,6 +6014,80 @@ CREATE TABLE tab3 (
 			},
 		},
 	},
+	{
+		Name: "correlated subquery references outer aggregate",
+		SetUpScript: []string{
+			"CREATE TABLE correlated_aggregate_scope (id INT PRIMARY KEY, grp INT, val INT);",
+			"INSERT INTO correlated_aggregate_scope VALUES (1, 1, 1), (2, 1, 2), (3, 2, 1);",
+			"CREATE TABLE correlated_aggregate_probe (probe INT);",
+			"INSERT INTO correlated_aggregate_probe VALUES (1);",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "SELECT grp, SUM(DISTINCT val) FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_scope b WHERE SUM(DISTINCT a.val) = b.val) ORDER BY grp;",
+				Expected: []sql.Row{{2, float64(1)}},
+			},
+			{
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_scope b WHERE SUM(DISTINCT a.val) = b.val) ORDER BY grp;",
+				Expected: []sql.Row{{2}},
+			},
+			{
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 WHERE COALESCE(SUM(a.val), 0) = 1) ORDER BY grp;",
+				Expected: []sql.Row{{2}},
+			},
+			{
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 WHERE SUM(val) = 1) ORDER BY grp;",
+				Expected: []sql.Row{{2}},
+			},
+			{
+				// TODO: https://github.com/dolthub/go-mysql-server/issues/3814
+				// The aggregate crosses two subqueries but still belongs to the scope that provides a.val.
+				Skip:     true,
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_probe b WHERE EXISTS (SELECT 1 WHERE SUM(a.val) = b.probe)) ORDER BY grp;",
+				Expected: []sql.Row{{2}},
+			},
+			{
+				// TODO: https://github.com/dolthub/go-mysql-server/issues/3814
+				// The middle a alias shadows outer a, so the deepest aggregate must belong to the middle query.
+				Skip:     true,
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_scope a HAVING EXISTS (SELECT 1 WHERE SUM(a.val) = 4)) ORDER BY grp;",
+				Expected: []sql.Row{{1}, {2}},
+			},
+			{
+				// TODO: https://github.com/dolthub/go-mysql-server/issues/3814
+				// The inner table has no val column, so normal name resolution should fall back to outer a.val.
+				Skip:     true,
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_probe b WHERE SUM(val) = b.probe) ORDER BY grp;",
+				Expected: []sql.Row{{2}},
+			},
+			{
+				// An unqualified local val shadows the outer column, so this aggregate stays in the subquery.
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_scope b HAVING SUM(val) = 4) ORDER BY grp;",
+				Expected: []sql.Row{{1}, {2}},
+			},
+			{
+				// Joins do not change ownership when every aggregate argument comes from the outer scope.
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT 1 FROM correlated_aggregate_probe b JOIN correlated_aggregate_probe c ON b.probe = c.probe WHERE SUM(a.val) = b.probe) ORDER BY grp;",
+				Expected: []sql.Row{{2}},
+			},
+			{
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp HAVING EXISTS (SELECT SUM(a.val) FROM correlated_aggregate_scope a HAVING SUM(a.val) = 4) ORDER BY grp;",
+				Expected: []sql.Row{{1}, {2}},
+			},
+			{
+				// Mixing outer a.val with inner b.val makes the aggregate belong to the inner query.
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp, a.val HAVING EXISTS (SELECT SUM(a.val + b.val) FROM correlated_aggregate_scope b HAVING SUM(a.val + b.val) > 0) ORDER BY grp;",
+				Expected: []sql.Row{{1}, {1}, {2}},
+			},
+			{
+				// TODO: https://github.com/dolthub/go-mysql-server/issues/3814
+				// The deepest aggregate mixes top-level a.val and middle-level b.probe, so it belongs to the middle query.
+				Skip:     true,
+				Query:    "SELECT grp FROM correlated_aggregate_scope a GROUP BY grp, a.val HAVING EXISTS (SELECT 1 FROM correlated_aggregate_probe b HAVING EXISTS (SELECT 1 WHERE SUM(a.val + b.probe) > 0)) ORDER BY grp;",
+				Expected: []sql.Row{{1}, {1}, {2}},
+			},
+		},
+	},
 	{
 		Name: "having clause without groupby clause, all rows implicitly form a single aggregate group",
 		SetUpScript: []string{

sql/planbuilder/aggregates.go

@@ -79,6 +79,42 @@ func (g *groupBy) getAggRef(name string) sql.Expression {
 	return ret.scalarGf()
 }
 
+// outerAggregateScope returns the outer scope that owns every column argument, if one exists.
+func (s *scope) outerAggregateScope(ctx *sql.Context, args []sql.Expression) *scope {
+	var columnIds sql.ColSet
+	var correlated sql.ColSet
+	if subquery := s.nearestSubquery(); subquery != nil {
+		correlated = subquery.correlated
+	}
+	for _, arg := range args {
+		sql.Inspect(ctx, arg, func(ctx *sql.Context, expr sql.Expression) bool {
+			gf, ok := expr.(*expression.GetField)
+			if !ok {
+				return true
+			}
+			columnIds.Add(gf.Id())
+			return false
+		})
+	}
+	if columnIds.Empty() || !columnIds.SubsetOf(correlated) {
+		return nil
+	}
+	for parent := s.parent; parent != nil; parent = parent.parent {
+		if columnIds.SubsetOf(parent.colset) {
+			return parent
+		}
+	}
+	return nil
+}
+
+// isCorrelatedColumn reports whether the column is supplied by an outer query to the active subquery.
+func (s *scope) isCorrelatedColumn(id sql.ColumnId) bool {
+	if subquery := s.nearestSubquery(); subquery != nil {
+		return subquery.correlated.Contains(id)
+	}
+	return false
+}
+
 type aggregateInfo struct {
 	ast.Expr
 }
@@ -294,27 +330,35 @@ func (b *Builder) buildAggregateFunc(inScope *scope, name string, e *ast.FuncExp
 		b.handleErr(err)
 	}
 
-	inScope.initGroupBy()
-	gb := inScope.groupBy
-
 	if strings.EqualFold(name, "count") {
 		if _, ok := e.Exprs[0].(*ast.StarExpr); ok {
-			return b.buildCountStarAggregate(e, gb)
+			inScope.initGroupBy()
+			return b.buildCountStarAggregate(e, inScope.groupBy)
 		}
 	}
 
 	if strings.EqualFold(name, "jsonarray") {
 		// TODO we don't have any tests for this
 		if _, ok := e.Exprs[0].(*ast.StarExpr); ok {
-			return b.buildJsonArrayStarAggregate(gb)
+			inScope.initGroupBy()
+			return b.buildJsonArrayStarAggregate(inScope.groupBy)
 		}
 	}
 
 	if strings.EqualFold(name, "any_value") {
 		b.qFlags.Set(sql.QFlagAnyAgg)
 	}
 
-	args := b.buildAggFunctionArgs(inScope, e, gb)
+	args := b.buildAggFunctionArgs(inScope, e)
+	var gb *groupBy
+	if outerScope := inScope.outerAggregateScope(b.ctx, args); outerScope != nil {
+		outerScope.initGroupBy()
+		gb = outerScope.groupBy
+	} else {
+		inScope.initGroupBy()
+		gb = inScope.groupBy
+	}
+	b.addAggFunctionArgs(gb, args)
 	agg := b.newAggregation(e, name, args)
 
 	if name == "count" {
@@ -328,7 +372,6 @@ func (b *Builder) buildAggregateFunc(inScope *scope, name string, e *ast.FuncExp
 		// if we've already computed use reference here
 		return gf
 	}
-
 	col := scopeColumn{col: aggName, scalar: agg, typ: aggType, nullable: agg.IsNullable(b.ctx)}
 	id := gb.outScope.newColumn(col)
 
@@ -380,38 +423,48 @@ func (b *Builder) newAggregation(e *ast.FuncExpr, name string, args []sql.Expres
 }
 
 // buildAggFunctionArgs builds the arguments for an aggregate function
-func (b *Builder) buildAggFunctionArgs(inScope *scope, e *ast.FuncExpr, gb *groupBy) []sql.Expression {
+func (b *Builder) buildAggFunctionArgs(inScope *scope, e *ast.FuncExpr) []sql.Expression {
 	var args []sql.Expression
 	for _, arg := range e.Exprs {
 		e := b.selectExprToExpression(inScope, arg)
 		// if GetField is an alias, alias must be masking a column
-		if gf, ok := e.(*expression.GetField); ok && gf.TableId() == 0 {
+		if gf, ok := e.(*expression.GetField); ok && gf.TableId() == 0 && !inScope.isCorrelatedColumn(gf.Id()) {
 			e = b.selectExprToExpression(inScope.parent, arg)
 		}
 		switch e := e.(type) {
 		case *expression.GetField:
-			if e.TableId() == 0 {
+			if e.TableId() == 0 && !inScope.isCorrelatedColumn(e.Id()) {
 				b.handleErr(fmt.Errorf("failed to resolve aggregate column argument: %s", e))
 			}
 			args = append(args, e)
-			col := scopeColumn{tableId: e.TableID(), db: e.Database(), table: e.Table(), col: e.Name(), scalar: e, typ: e.Type(b.ctx), nullable: e.IsNullable(b.ctx)}
-			gb.addInCol(col)
 		case *expression.Star:
 			err := sql.ErrStarUnsupported.New()
 			b.handleErr(err)
 		case *plan.Subquery:
 			args = append(args, e)
-			col := scopeColumn{col: e.QueryString, scalar: e, typ: e.Type(b.ctx)}
-			gb.addInCol(col)
 		default:
 			args = append(args, e)
-			col := scopeColumn{col: e.String(), scalar: e, typ: e.Type(b.ctx)}
-			gb.addInCol(col)
 		}
 	}
 	return args
 }
 
+// addAggFunctionArgs records aggregate inputs in the scope that owns the aggregate.
+func (b *Builder) addAggFunctionArgs(gb *groupBy, args []sql.Expression) {
+	for _, arg := range args {
+		var col scopeColumn
+		switch arg := arg.(type) {
+		case *expression.GetField:
+			col = scopeColumn{tableId: arg.TableID(), db: arg.Database(), table: arg.Table(), col: arg.Name(), scalar: arg, typ: arg.Type(b.ctx), nullable: arg.IsNullable(b.ctx)}
+		case *plan.Subquery:
+			col = scopeColumn{col: arg.QueryString, scalar: arg, typ: arg.Type(b.ctx)}
+		default:
+			col = scopeColumn{col: arg.String(), scalar: arg, typ: arg.Type(b.ctx)}
+		}
+		gb.addInCol(col)
+	}
+}
+
 // buildJsonArrayStarAggregate builds a JSON_ARRAY(*) aggregate function
 func (b *Builder) buildJsonArrayStarAggregate(gb *groupBy) sql.Expression {
 	var agg sql.Aggregation
@@ -843,6 +896,7 @@ func (b *Builder) analyzeHaving(fromScope, projScope *scope, having *ast.Where)
 	ast.Walk(func(node ast.SQLNode) (bool, error) {
 		switch n := node.(type) {
 		case *ast.Subquery:
+			b.analyzeOuterAggregates(fromScope, n.Select)
 			return false, nil
 		case *ast.FuncExpr:
 			name := n.Name.Lowered()
@@ -880,6 +934,90 @@ func (b *Builder) analyzeHaving(fromScope, projScope *scope, having *ast.Where)
 	}, having.Expr)
 }
 
+// analyzeOuterAggregates registers aggregates whose arguments all belong to the enclosing query.
+func (b *Builder) analyzeOuterAggregates(fromScope *scope, node ast.SelectStatement) {
+	selectStmt, ok := node.(*ast.Select)
+	if !ok {
+		return
+	}
+	localTables := localTableNames(selectStmt.From)
+	ast.Walk(func(node ast.SQLNode) (bool, error) {
+		switch n := node.(type) {
+		case *ast.Subquery:
+			return false, nil
+		case *ast.FuncExpr:
+			name := n.Name.Lowered()
+			isAggregate, err := IsAggregateFunc(b.ctx, name)
+			if err != nil {
+				b.handleErr(err)
+			}
+			if isAggregate && aggregateArgsBelongToScope(fromScope, n, localTables, len(selectStmt.From) == 0) {
+				_ = b.buildAggregateFunc(fromScope, name, n)
+				return false, nil
+			}
+			return true, nil
+		}
+		return true, nil
+	}, node)
+}
+
+// localTableNames returns the table names and aliases introduced by a subquery's FROM clause.
+func localTableNames(from ast.TableExprs) map[string]struct{} {
+	names := make(map[string]struct{})
+	for _, tableExpr := range from {
+		ast.Walk(func(node ast.SQLNode) (bool, error) {
+			aliased, ok := node.(*ast.AliasedTableExpr)
+			if !ok {
+				return true, nil
+			}
+			name := strings.ToLower(aliased.As.String())
+			if name == "" {
+				if tableName, ok := aliased.Expr.(ast.TableName); ok {
+					name = strings.ToLower(tableName.Name.String())
+				}
+			}
+			if name != "" {
+				names[name] = struct{}{}
+			}
+			return false, nil
+		}, tableExpr)
+	}
+	return names
+}
+
+// aggregateArgsBelongToScope reports whether every column argument resolves only in the given enclosing scope.
+func aggregateArgsBelongToScope(fromScope *scope, fn *ast.FuncExpr, localTables map[string]struct{}, allowUnqualified bool) bool {
+	hasColumn := false
+	belongs := true
+	for _, arg := range fn.Exprs {
+		ast.Walk(func(node ast.SQLNode) (bool, error) {
+			col, ok := node.(*ast.ColName)
+			if !ok {
+				return true, nil
+			}
+			hasColumn = true
+			dbName := strings.ToLower(col.Qualifier.DbQualifier.String())
+			tblName := strings.ToLower(col.Qualifier.Name.String())
+			colName := strings.ToLower(col.Name.String())
+			if tblName == "" {
+				if !allowUnqualified {
+					belongs = false
+					return false, nil
+				}
+			} else if _, local := localTables[tblName]; local {
+				belongs = false
+				return false, nil
+			}
+			_, ok = fromScope.resolveColumn(dbName, tblName, colName, false, false)
+			if !ok {
+				belongs = false
+			}
+			return false, nil
+		}, arg)
+	}
+	return hasColumn && belongs
+}
+
 func (b *Builder) buildInnerProj(fromScope, projScope *scope) *scope {
 	outScope := fromScope
 	var proj []sql.Expression