From 3094e894c186486a3090a52ddc757fd1ae913013 Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Fri, 3 Jul 2026 17:36:41 +0800 Subject: [PATCH 1/2] fix(plan): bind grouped correlated columns --- pkg/sql/plan/base_binder.go | 47 ++++++++ pkg/sql/plan/build_test.go | 107 ++++++++++++++++++ pkg/sql/plan/having_binder.go | 17 ++- .../cases/subquery/subquery-with-any.result | 4 +- 4 files changed, 172 insertions(+), 3 deletions(-) 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 90b7341137b4e..cdefef409211d 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -825,6 +825,113 @@ 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"}, + {` + 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 + )`, "nation.n_regionkey"}, + } + + 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`, + } + + 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 462c8a370ee25..073c2d0bb96a0 100644 --- a/pkg/sql/plan/having_binder.go +++ b/pkg/sql/plan/having_binder.go @@ -102,7 +102,10 @@ func (b *HavingBinder) BindColRef(astExpr *tree.UnresolvedName, depth int32, isR return nil, err } - if _, ok := expr.Expr.(*plan.Expr_Corr); ok { + if corr, ok := expr.Expr.(*plan.Expr_Corr); ok { + if !b.builder.mysqlCompatible && b.corrColRefTargetsGroup(corr.Corr) { + return nil, b.newGroupByColumnError(astExpr) + } return nil, moerr.NewNYI(b.GetContext(), "correlated columns in aggregate function") } @@ -130,10 +133,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)) +} + func (b *HavingBinder) BindAggFunc(funcName string, astExpr *tree.FuncExpr, depth int32, isRoot bool) (*plan.Expr, error) { if b.insideAgg { return nil, moerr.NewSyntaxErrorf(b.GetContext(), "aggregate function %s calls cannot be nested", funcName) 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 ( From 11bf4a447b62090e8fd0207c9ac7be6ae3a78d3d Mon Sep 17 00:00:00 2001 From: GreatRiver Date: Mon, 6 Jul 2026 15:14:41 +0800 Subject: [PATCH 2/2] fix(plan): preserve correlated aggregate NYI --- pkg/sql/plan/build_test.go | 25 +++++++++++-------- pkg/sql/plan/having_binder.go | 5 +--- .../subquery/subquery-with-exists.result | 2 +- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/pkg/sql/plan/build_test.go b/pkg/sql/plan/build_test.go index 9ad6611e2cb8f..e85f062f8d5bb 100644 --- a/pkg/sql/plan/build_test.go +++ b/pkg/sql/plan/build_test.go @@ -892,16 +892,6 @@ func TestOnlyFullGroupByRejectsCorrelatedSubqueryOnUngroupedColumn(t *testing.T) FROM nation GROUP BY n_name HAVING (SELECT COUNT(*) FROM nation2 n2 WHERE n2.n_name = nation.n_comment) > 0`, "nation.n_comment"}, - {` - 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 - )`, "nation.n_regionkey"}, } for _, tt := range sqls { @@ -922,6 +912,21 @@ func TestOnlyFullGroupByPreservesCorrelatedAggregateNYI(t *testing.T) { 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 { diff --git a/pkg/sql/plan/having_binder.go b/pkg/sql/plan/having_binder.go index 073c2d0bb96a0..b14545254a8b2 100644 --- a/pkg/sql/plan/having_binder.go +++ b/pkg/sql/plan/having_binder.go @@ -102,10 +102,7 @@ func (b *HavingBinder) BindColRef(astExpr *tree.UnresolvedName, depth int32, isR return nil, err } - if corr, ok := expr.Expr.(*plan.Expr_Corr); ok { - if !b.builder.mysqlCompatible && b.corrColRefTargetsGroup(corr.Corr) { - return nil, b.newGroupByColumnError(astExpr) - } + if _, ok := expr.Expr.(*plan.Expr_Corr); ok { return nil, moerr.NewNYI(b.GetContext(), "correlated columns in aggregate function") } 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;