diff --git a/enginetest/queries/script_queries.go b/enginetest/queries/script_queries.go index d104891584..ac6a2f3de3 100644 --- a/enginetest/queries/script_queries.go +++ b/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{ diff --git a/sql/planbuilder/aggregates.go b/sql/planbuilder/aggregates.go index 55fd4cefd3..a0a61525e6 100644 --- a/sql/planbuilder/aggregates.go +++ b/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,19 +330,18 @@ 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) } } @@ -314,7 +349,16 @@ func (b *Builder) buildAggregateFunc(inScope *scope, name string, e *ast.FuncExp 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