From b0122ebd46ea23979e217fdd2ed7a3ea5103999c Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Tue, 26 Jan 2021 00:17:37 +0800
Subject: [PATCH] planner: do not push down the aggregation function with
 correlated column (#21453) (#21546)

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 expression/util.go               | 15 +++++++++++++++
 planner/core/integration_test.go | 11 +++++++++++
 planner/core/task.go             |  3 ++-
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/expression/util.go b/expression/util.go
index 664d7c88dd73c..61b7eacdb8221 100644
--- a/expression/util.go
+++ b/expression/util.go
@@ -880,6 +880,21 @@ func ContainVirtualColumn(exprs []Expression) bool {
 	return false
 }
 
+// ContainCorrelatedColumn checks if the expressions contain a correlated column
+func ContainCorrelatedColumn(exprs []Expression) bool {
+	for _, expr := range exprs {
+		switch v := expr.(type) {
+		case *CorrelatedColumn:
+			return true
+		case *ScalarFunction:
+			if ContainCorrelatedColumn(v.GetArgs()) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // ContainMutableConst checks if the expressions contain a lazy constant.
 func ContainMutableConst(ctx sessionctx.Context, exprs []Expression) bool {
 	// Treat all constants immutable if plan cache is not enabled for this query.
diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index a7c59062a35b8..ead195345d250 100644
--- a/planner/core/integration_test.go
+++ b/planner/core/integration_test.go
@@ -1625,6 +1625,17 @@ func (s *testIntegrationSuite) TestUpdateMultiUpdatePK(c *C) {
 	tk.MustQuery("SELECT * FROM t").Check(testkit.Rows("2 12"))
 }
 
+func (s *testIntegrationSuite) TestCorrelatedColumnAggFuncPushDown(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test;")
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (a int, b int);")
+	tk.MustExec("insert into t values (1,1);")
+	tk.MustQuery("select (select count(n.a + a) from t) from t n;").Check(testkit.Rows(
+		"1",
+	))
+}
+
 func (s *testIntegrationSuite) TestIssue22105(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 
diff --git a/planner/core/task.go b/planner/core/task.go
index 10aa66b6d8f13..f4ba96cc9e0e2 100644
--- a/planner/core/task.go
+++ b/planner/core/task.go
@@ -1034,7 +1034,8 @@ func CheckAggCanPushCop(sctx sessionctx.Context, aggFuncs []*aggregation.AggFunc
 	sc := sctx.GetSessionVars().StmtCtx
 	client := sctx.GetClient()
 	for _, aggFunc := range aggFuncs {
-		if expression.ContainVirtualColumn(aggFunc.Args) {
+		// if the aggFunc contain VirtualColumn or CorrelatedColumn, it can not be pushed down.
+		if expression.ContainVirtualColumn(aggFunc.Args) || expression.ContainCorrelatedColumn(aggFunc.Args) {
 			return false
 		}
 		pb := aggregation.AggFuncToPBExpr(sc, client, aggFunc)