planner: fix correlated aggregates which should be evaluated in outer query

executor/testdata/agg_suite_out.json

@@ -49,21 +49,21 @@
     "Name": "TestIssue12759HashAggCalledByApply",
     "Cases": [
       [
-        "Projection_28 1.00 root  Column#3, Column#6, Column#9, Column#12",
+        "Projection_28 1.00 root  Column#9, Column#10, Column#11, Column#12",
         "└─Apply_30 1.00 root  CARTESIAN left outer join",
         "  ├─Apply_32(Build) 1.00 root  CARTESIAN left outer join",
         "  │ ├─Apply_34(Build) 1.00 root  CARTESIAN left outer join",
-        "  │ │ ├─HashAgg_39(Build) 1.00 root  funcs:sum(Column#22)->Column#3, funcs:firstrow(Column#23)->test.test.a",
+        "  │ │ ├─HashAgg_39(Build) 1.00 root  funcs:sum(Column#22)->Column#9, funcs:firstrow(Column#23)->test.test.a",
         "  │ │ │ └─TableReader_40 1.00 root  data:HashAgg_35",
         "  │ │ │   └─HashAgg_35 1.00 cop[tikv]  funcs:sum(test.test.a)->Column#22, funcs:firstrow(test.test.a)->Column#23",
         "  │ │ │     └─TableFullScan_38 10000.00 cop[tikv] table:tt keep order:false, stats:pseudo",
-        "  │ │ └─Projection_43(Probe) 1.00 root  <nil>->Column#6",
+        "  │ │ └─Projection_43(Probe) 1.00 root  <nil>->Column#10",
         "  │ │   └─Limit_44 1.00 root  offset:0, count:1",
         "  │ │     └─TableReader_50 1.00 root  data:Limit_49",
         "  │ │       └─Limit_49 1.00 cop[tikv]  offset:0, count:1",
         "  │ │         └─Selection_48 1.00 cop[tikv]  eq(test.test.a, test.test.a)",
         "  │ │           └─TableFullScan_47 1000.00 cop[tikv] table:test keep order:false, stats:pseudo",
-        "  │ └─Projection_54(Probe) 1.00 root  <nil>->Column#9",
+        "  │ └─Projection_54(Probe) 1.00 root  <nil>->Column#11",
         "  │   └─Limit_55 1.00 root  offset:0, count:1",
         "  │     └─TableReader_61 1.00 root  data:Limit_60",
         "  │       └─Limit_60 1.00 cop[tikv]  offset:0, count:1",

planner/cascades/stringer_test.go

@@ -86,6 +86,6 @@ func (s *testStringerSuite) TestGroupStringer(c *C) {
 			output[i].SQL = sql
 			output[i].Result = ToString(group)
 		})
-		c.Assert(ToString(group), DeepEquals, output[i].Result)
+		c.Assert(ToString(group), DeepEquals, output[i].Result, Commentf("case:%v, sql:%s", i, sql))
 	}
 }

planner/cascades/testdata/integration_suite_out.json

@@ -1035,10 +1035,10 @@
       {
         "SQL": "select sum(a), (select t1.a from t1 where t1.a = t2.a limit 1), (select t1.b from t1 where t1.b = t2.b limit 1) from t2",
         "Plan": [
-          "Projection_30 1.00 root  Column#3, test.t1.a, test.t1.b",
+          "Projection_30 1.00 root  Column#7, test.t1.a, test.t1.b",
           "└─Apply_32 1.00 root  CARTESIAN left outer join",
           "  ├─Apply_34(Build) 1.00 root  CARTESIAN left outer join",
-          "  │ ├─HashAgg_39(Build) 1.00 root  funcs:sum(Column#8)->Column#3, funcs:firstrow(Column#9)->test.t2.a, funcs:firstrow(Column#10)->test.t2.b",
+          "  │ ├─HashAgg_39(Build) 1.00 root  funcs:sum(Column#8)->Column#7, funcs:firstrow(Column#9)->test.t2.a, funcs:firstrow(Column#10)->test.t2.b",
           "  │ │ └─TableReader_40 1.00 root  data:HashAgg_41",
           "  │ │   └─HashAgg_41 1.00 cop[tikv]  funcs:sum(test.t2.a)->Column#8, funcs:firstrow(test.t2.a)->Column#9, funcs:firstrow(test.t2.b)->Column#10",
           "  │ │     └─TableFullScan_38 10000.00 cop[tikv] table:t2 keep order:false, stats:pseudo",

planner/cascades/testdata/stringer_suite_out.json

@@ -290,7 +290,7 @@
         "SQL": "select a = (select a from t t2 where t1.b = t2.b order by a limit 1) from t t1",
         "Result": [
           "Group#0 Schema:[Column#25]",
-          "    Projection_2 input:[Group#1], eq(test.t.a, test.t.a)->Column#25",
+          "    Projection_3 input:[Group#1], eq(test.t.a, test.t.a)->Column#25",
           "Group#1 Schema:[test.t.a,test.t.b,test.t.a]",
           "    Apply_9 input:[Group#2,Group#3], left outer join",
           "Group#2 Schema:[test.t.a,test.t.b], UniqueKey:[test.t.a]",

planner/core/expression_rewriter.go

@@ -339,11 +339,24 @@ func (er *expressionRewriter) Enter(inNode ast.Node) (ast.Node, bool) {
 		if er.aggrMap != nil {
 			index, ok = er.aggrMap[v]
 		}
-		if !ok {
-			er.err = ErrInvalidGroupFuncUse
+		if ok {
+			// index < 0 indicates this is a correlated aggregate belonging to outer query,
+			// for which a correlated column will be created later, so we append a null constant
+			// as a temporary result expression.
+			if index < 0 {
+				er.ctxStackAppend(expression.NewNull(), types.EmptyName)
+			} else {
+				// index >= 0 indicates this is a regular aggregate column
+				er.ctxStackAppend(er.schema.Columns[index], er.names[index])
+			}
 			return inNode, true
 		}
-		er.ctxStackAppend(er.schema.Columns[index], er.names[index])
+		// replace correlated aggregate in sub-query with its corresponding correlated column
+		if col, ok := er.b.correlatedAggMapper[v]; ok {
+			er.ctxStackAppend(col, types.EmptyName)
+			return inNode, true
+		}
+		er.err = ErrInvalidGroupFuncUse
 		return inNode, true
 	case *ast.ColumnNameExpr:
 		if index, ok := er.b.colMapper[v]; ok {

planner/core/integration_test.go

@@ -2012,6 +2012,71 @@ func (s *testIntegrationSuite) TestOrderByNotInSelectDistinct(c *C) {
 	tk.MustQuery("select distinct v1 as z from ttest order by v1+z").Check(testkit.Rows("1", "4"))
 }
 
+func (s *testIntegrationSuite) TestCorrelatedAggregate(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+
+	// #18350
+	tk.MustExec("DROP TABLE IF EXISTS tab, tab2")
+	tk.MustExec("CREATE TABLE tab(i INT)")
+	tk.MustExec("CREATE TABLE tab2(j INT)")
+	tk.MustExec("insert into tab values(1),(2),(3)")
+	tk.MustExec("insert into tab2 values(1),(2),(3),(15)")
+	tk.MustQuery(`SELECT m.i,
+       (SELECT COUNT(n.j)
+           FROM tab2 WHERE j=15) AS o
+    FROM tab m, tab2 n GROUP BY 1 order by m.i`).Check(testkit.Rows("1 4", "2 4", "3 4"))
+	tk.MustQuery(`SELECT
+         (SELECT COUNT(n.j)
+             FROM tab2 WHERE j=15) AS o
+    FROM tab m, tab2 n order by m.i`).Check(testkit.Rows("12"))
+
+	// #17748
+	tk.MustExec("drop table if exists t1, t2")
+	tk.MustExec("create table t1 (a int, b int)")
+	tk.MustExec("create table t2 (m int, n int)")
+	tk.MustExec("insert into t1 values (2,2), (2,2), (3,3), (3,3), (3,3), (4,4)")
+	tk.MustExec("insert into t2 values (1,11), (2,22), (3,32), (4,44), (4,44)")
+	tk.MustExec("set @@sql_mode='TRADITIONAL'")
+
+	tk.MustQuery(`select count(*) c, a,
+		( select group_concat(count(a)) from t2 where m = a )
+		from t1 group by a order by a`).
+		Check(testkit.Rows("2 2 2", "3 3 3", "1 4 1,1"))
+
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (a int, b int)")
+	tk.MustExec("insert into t values (1,1),(2,1),(2,2),(3,1),(3,2),(3,3)")
+	// from SELECT fields
+	tk.MustQuery("select (select count(a)) from t").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select (select (select count(a)))) from t").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select (select count(n.a)) from t m order by count(m.b)) from t n").Check(testkit.Rows("6"))
+	// from WHERE
+	tk.MustQuery("select (select count(n.a) from t where count(n.a)=3) from t n").Check(testkit.Rows("<nil>"))
+	tk.MustQuery("select (select count(a) from t where count(distinct n.a)=3) from t n").Check(testkit.Rows("6"))
+	// from HAVING
+	tk.MustQuery("select (select count(n.a) from t having count(n.a)=6 limit 1) from t n").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select count(n.a) from t having count(distinct n.b)=3 limit 1) from t n").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select sum(distinct n.a) from t having count(distinct n.b)=3 limit 1) from t n").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select sum(distinct n.a) from t having count(distinct n.b)=6 limit 1) from t n").Check(testkit.Rows("<nil>"))
+	// from ORDER BY
+	tk.MustQuery("select (select count(n.a) from t order by count(n.b) limit 1) from t n").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select count(distinct n.b) from t order by count(n.b) limit 1) from t n").Check(testkit.Rows("3"))
+	// from TableRefsClause
+	tk.MustQuery("select (select cnt from (select count(a) cnt) s) from t").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select count(cnt) from (select count(a) cnt) s) from t").Check(testkit.Rows("1"))
+	// from sub-query inside aggregate
+	tk.MustQuery("select (select sum((select count(a)))) from t").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select sum((select count(a))+sum(a))) from t").Check(testkit.Rows("20"))
+	// from GROUP BY
+	tk.MustQuery("select (select count(a) from t group by count(n.a)) from t n").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select count(distinct a) from t group by count(n.a)) from t n").Check(testkit.Rows("3"))
+
+	// Nested aggregate (correlated aggregate inside aggregate)
+	tk.MustQuery("select (select sum(count(a))) from t").Check(testkit.Rows("6"))
+	tk.MustQuery("select (select sum(sum(a))) from t").Check(testkit.Rows("14"))
+}
+
 func (s *testIntegrationSuite) TestCorrelatedColumnAggFuncPushDown(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("use test;")