diff --git a/pkg/sql/plan/base_binder.go b/pkg/sql/plan/base_binder.go index 82ea9315da7e7..f5f84be158ae6 100644 --- a/pkg/sql/plan/base_binder.go +++ b/pkg/sql/plan/base_binder.go @@ -423,6 +423,23 @@ func (b *baseBinder) baseBindColRef(astExpr *tree.UnresolvedName, depth int32, i } } + if groupPos, ok := b.correlatedGroupByColPos(depth, name, table, col); ok { + expr = &plan.Expr{ + Typ: b.ctx.groups[groupPos].Typ, + Expr: &plan.Expr_Corr{ + Corr: &plan.CorrColRef{ + RelPos: b.ctx.groupTag, + ColPos: groupPos, + Depth: depth, + }, + }, + } + if err != nil { + errutil.ReportError(b.GetContext(), err) + } + return + } + if isEnumOrSetPlanType(typ) { if err != nil { errutil.ReportError(b.GetContext(), err) @@ -513,6 +530,36 @@ func (b *baseBinder) baseBindColRef(astExpr *tree.UnresolvedName, depth int32, i return } +func (b *baseBinder) correlatedGroupByColPos(depth int32, astName, table, col string) (int32, bool) { + if depth == 0 || b.ctx == nil || len(b.ctx.groupByAst) == 0 { + return 0, false + } + if pos, ok := b.ctx.groupByAst[astName]; ok && int(pos) < len(b.ctx.groups) { + return pos, true + } + if table != "" { + if pos, ok := b.ctx.groupByAst[table+"."+col]; ok && int(pos) < len(b.ctx.groups) { + return pos, true + } + } + return 0, false +} + +func (b *baseBinder) corrColRefTargetsGroup(corr *plan.CorrColRef) bool { + if corr == nil { + return false + } + ctx := b.ctx + for depth := int32(0); depth < corr.Depth && ctx != nil; depth++ { + ctx = ctx.parent + } + return ctx != nil && ctx.groupTag > 0 && corr.RelPos == ctx.groupTag +} + +func (b *baseBinder) corrColRefTargetsCurrentGroup(corr *plan.CorrColRef) bool { + return corr != nil && b.ctx != nil && b.ctx.groupTag > 0 && corr.RelPos == b.ctx.groupTag +} + func (b *baseBinder) baseBindSubquery(astExpr *tree.Subquery, isRoot bool) (*Expr, error) { if b.ctx == nil { return nil, moerr.NewInvalidInput(b.GetContext(), "field reference doesn't support SUBQUERY") diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 0fc4916166bbc..f7c3b285e39d1 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -825,6 +825,118 @@ func TestSingleTableSQLBuilder(t *testing.T) { runTestShouldError(mock, t, sqls) } +func TestOnlyFullGroupByAllowsCorrelatedSubqueryOnGroupedColumn(t *testing.T) { + sqls := []string{ + ` + SELECT n_name, + (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_name) AS c + FROM nation + GROUP BY n_name + ORDER BY n_name`, + ` + SELECT n_name + FROM nation + GROUP BY n_name + HAVING (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_name) > 0`, + ` + SELECT n_name + FROM nation + GROUP BY n_name + ORDER BY (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_name)`, + ` + SELECT n.n_name, + (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = n.n_name) AS c + FROM nation n + GROUP BY n.n_name`, + ` + SELECT n_name, + (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_name) AS c + FROM nation + GROUP BY 1`, + ` + SELECT n_name AS name, + (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_name) AS c + FROM nation + GROUP BY name`, + ` + SELECT n_regionkey + FROM nation + GROUP BY n_regionkey + HAVING EXISTS ( + SELECT n_name + FROM nation2 + GROUP BY n_name + HAVING COUNT(*) > nation.n_regionkey + )`, + } + + for _, sql := range sqls { + mock := NewMockOptimizer(false) + _, err := runOneStmt(mock, t, sql) + require.NoError(t, err, sql) + } +} + +func TestOnlyFullGroupByRejectsCorrelatedSubqueryOnUngroupedColumn(t *testing.T) { + sqls := []struct { + sql string + errContains string + }{ + {` + SELECT n_name, + (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_comment) AS c + FROM nation + GROUP BY n_name`, "nation.n_comment"}, + {` + SELECT n_name + FROM nation + GROUP BY n_name + HAVING (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_comment) > 0`, "nation.n_comment"}, + } + + for _, tt := range sqls { + mock := NewMockOptimizer(false) + _, err := runOneStmt(mock, t, tt.sql) + require.Error(t, err, tt.sql) + require.Contains(t, err.Error(), tt.errContains) + } +} + +func TestOnlyFullGroupByPreservesCorrelatedAggregateNYI(t *testing.T) { + sqls := []string{ + ` + SELECT (SELECT COUNT(DISTINCT nation.n_comment)) + FROM nation + GROUP BY n_name`, + ` + SELECT n_name + FROM nation + WHERE (SELECT AVG(nation.n_regionkey) FROM nation2) = 1`, + ` + SELECT n_name, + (SELECT COUNT(nation.n_name) FROM nation2) AS c + FROM nation + GROUP BY n_name`, + ` + SELECT n_regionkey + FROM nation + GROUP BY n_regionkey + HAVING EXISTS ( + SELECT n_name + FROM nation2 + GROUP BY n_name + HAVING SUM(nation.n_regionkey) > 0 + )`, + } + + for _, sql := range sqls { + mock := NewMockOptimizer(false) + _, err := runOneStmt(mock, t, sql) + require.Error(t, err, sql) + require.Contains(t, err.Error(), "correlated columns in aggregate function") + } +} + // test join table plan building func TestJoinTableSqlBuilder(t *testing.T) { mock := NewMockOptimizer(false) diff --git a/pkg/sql/plan/having_binder.go b/pkg/sql/plan/having_binder.go index 530cb212039f8..3b0df8ef095e7 100644 --- a/pkg/sql/plan/having_binder.go +++ b/pkg/sql/plan/having_binder.go @@ -132,10 +132,20 @@ func (b *HavingBinder) BindColRef(astExpr *tree.UnresolvedName, depth int32, isR }, }, nil } else { - return nil, moerr.NewSyntaxErrorf(b.GetContext(), "column %q must appear in the GROUP BY clause or be used in an aggregate function", tree.String(astExpr, dialect.MYSQL)) + if expr, err := b.baseBindColRef(astExpr, depth, isRoot); err == nil { + if corr, ok := expr.Expr.(*plan.Expr_Corr); ok && + (b.corrColRefTargetsCurrentGroup(corr.Corr) || b.corrColRefTargetsGroup(corr.Corr)) { + return expr, nil + } + } + return nil, b.newGroupByColumnError(astExpr) } } +func (b *HavingBinder) newGroupByColumnError(astExpr *tree.UnresolvedName) error { + return moerr.NewSyntaxErrorf(b.GetContext(), "column %q must appear in the GROUP BY clause or be used in an aggregate function", tree.String(astExpr, dialect.MYSQL)) +} + // validateCountArgs validates COUNT function arguments against MySQL semantics. // - COUNT(*), COUNT(expr): always valid // - COUNT(expr1, expr2, ...): only valid with DISTINCT diff --git a/test/distributed/cases/subquery/subquery-with-any.result b/test/distributed/cases/subquery/subquery-with-any.result index f88ade28ba76a..cde2c466b54d1 100755 --- a/test/distributed/cases/subquery/subquery-with-any.result +++ b/test/distributed/cases/subquery/subquery-with-any.result @@ -356,7 +356,9 @@ FROM t1 GROUP BY field1, field2 HAVING COUNT(*) < ANY (SELECT fieldB FROM t2 WHERE fieldA = field1); -SQL syntax error: column "field1" must appear in the GROUP BY clause or be used in an aggregate function +field1 field2 +1 1 +1 3 DROP TABLE IF EXISTS t1; DROP TABLE IF EXISTS t2; CREATE TABLE t1 ( diff --git a/test/distributed/cases/subquery/subquery-with-exists.result b/test/distributed/cases/subquery/subquery-with-exists.result index 0ac19469971a6..3c6b9731b4306 100755 --- a/test/distributed/cases/subquery/subquery-with-exists.result +++ b/test/distributed/cases/subquery/subquery-with-exists.result @@ -298,7 +298,7 @@ a 4 SELECT a FROM t1 GROUP BY a HAVING EXISTS(SELECT c FROM t2 GROUP BY c HAVING SUM(a) = c); -SQL syntax error: column "a" must appear in the GROUP BY clause or be used in an aggregate function +correlated columns in aggregate function is not yet implemented SELECT a FROM t1 WHERE a < 3 AND EXISTS(SELECT c FROM t2 GROUP BY c HAVING SUM(a) != c) GROUP BY a;