From ec2f372ce41c3587d01fbe89c06757c84904e69b Mon Sep 17 00:00:00 2001
From: fixdb <1225233+fixdb@users.noreply.github.com>
Date: Thu, 1 Dec 2022 13:10:00 -0800
Subject: [PATCH 1/2] This is an automated cherry-pick of #39547

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>
---
 executor/tiflash_test.go | 29 +++++++++++++++++++++++++++++
 planner/core/task.go     |  9 +++++++++
 2 files changed, 38 insertions(+)

diff --git a/executor/tiflash_test.go b/executor/tiflash_test.go
index 97227e48f3246..54329bb93c893 100644
--- a/executor/tiflash_test.go
+++ b/executor/tiflash_test.go
@@ -1248,6 +1248,35 @@ func TestAggPushDownCountStar(t *testing.T) {
 	tk.MustQuery("select count(*) from c, o where c.c_id=o.c_id").Check(testkit.Rows("5"))
 }
 
+func TestGroupStreamAggOnTiFlash(t *testing.T) {
+	store := testkit.CreateMockStore(t, withMockTiFlash(2))
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists foo")
+	tk.MustExec("create table foo(a int, b int, c int, d int, primary key(a,b,c,d))")
+	tk.MustExec("alter table foo set tiflash replica 1")
+	tk.MustExec("insert into foo values(1,2,3,1),(1,2,3,6),(1,2,3,5)," +
+		"(1,2,3,2),(1,2,3,4),(1,2,3,7),(1,2,3,3),(1,2,3,0)")
+	tb := external.GetTableByName(t, tk, "test", "foo")
+	err := domain.GetDomain(tk.Session()).DDL().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)
+	require.NoError(t, err)
+	tk.MustExec("set @@tidb_allow_mpp=0")
+	sql := "select a,b,c,count(*) from foo group by a,b,c order by a,b,c"
+	tk.MustQuery(sql).Check(testkit.Rows("1 2 3 8"))
+	rows := tk.MustQuery("explain " + sql).Rows()
+
+	for _, row := range rows {
+		resBuff := bytes.NewBufferString("")
+		fmt.Fprintf(resBuff, "%s\n", row)
+		res := resBuff.String()
+		// StreamAgg with group keys on TiFlash is not supported
+		if strings.Contains(res, "tiflash") {
+			require.NotContains(t, res, "StreamAgg")
+		}
+	}
+}
+
 func TestTiflashEmptyDynamicPruneResult(t *testing.T) {
 	store, clean := testkit.CreateMockStore(t, withMockTiFlash(2))
 	defer clean()
diff --git a/planner/core/task.go b/planner/core/task.go
index bec15641decd1..cc0dee5a87c11 100644
--- a/planner/core/task.go
+++ b/planner/core/task.go
@@ -1662,11 +1662,20 @@ func (p *PhysicalStreamAgg) attach2Task(tasks ...task) task {
 			inputRows = t.count()
 			attachPlan2Task(p, t)
 		} else {
+<<<<<<< HEAD
 			copTaskType := cop.getStoreType()
 			partialAgg, finalAgg := p.newPartialAggregate(copTaskType, false)
 			if finalAgg != nil {
 				final = finalAgg.(*PhysicalStreamAgg)
 			}
+=======
+			storeType := cop.getStoreType()
+			// TiFlash doesn't support Stream Aggregation
+			if storeType == kv.TiFlash && len(p.GroupByItems) > 0 {
+				return invalidTask
+			}
+			partialAgg, finalAgg := p.newPartialAggregate(storeType, false)
+>>>>>>> 38b0ab7333 (planner: prohibit StreamAgg with group keys for TiFlash (#39547))
 			if partialAgg != nil {
 				if cop.tablePlan != nil {
 					cop.finishIndexPlan()

From 60e2d6cd8ae06aad06be343c8deb5b14ebb11f6d Mon Sep 17 00:00:00 2001
From: fixdb <1225233+fixdb@users.noreply.github.com>
Date: Tue, 10 Jan 2023 13:07:38 -0800
Subject: [PATCH 2/2] Resolve merge conflict

---
 executor/tiflash_test.go |  3 ++-
 planner/core/task.go     | 11 +++--------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/executor/tiflash_test.go b/executor/tiflash_test.go
index 54329bb93c893..999ea336acffa 100644
--- a/executor/tiflash_test.go
+++ b/executor/tiflash_test.go
@@ -1249,7 +1249,8 @@ func TestAggPushDownCountStar(t *testing.T) {
 }
 
 func TestGroupStreamAggOnTiFlash(t *testing.T) {
-	store := testkit.CreateMockStore(t, withMockTiFlash(2))
+	store, clean := testkit.CreateMockStore(t, withMockTiFlash(2))
+	defer clean()
 	tk := testkit.NewTestKit(t, store)
 
 	tk.MustExec("use test")
diff --git a/planner/core/task.go b/planner/core/task.go
index cc0dee5a87c11..075f1e21488bf 100644
--- a/planner/core/task.go
+++ b/planner/core/task.go
@@ -1662,20 +1662,15 @@ func (p *PhysicalStreamAgg) attach2Task(tasks ...task) task {
 			inputRows = t.count()
 			attachPlan2Task(p, t)
 		} else {
-<<<<<<< HEAD
-			copTaskType := cop.getStoreType()
-			partialAgg, finalAgg := p.newPartialAggregate(copTaskType, false)
-			if finalAgg != nil {
-				final = finalAgg.(*PhysicalStreamAgg)
-			}
-=======
 			storeType := cop.getStoreType()
 			// TiFlash doesn't support Stream Aggregation
 			if storeType == kv.TiFlash && len(p.GroupByItems) > 0 {
 				return invalidTask
 			}
 			partialAgg, finalAgg := p.newPartialAggregate(storeType, false)
->>>>>>> 38b0ab7333 (planner: prohibit StreamAgg with group keys for TiFlash (#39547))
+			if finalAgg != nil {
+				final = finalAgg.(*PhysicalStreamAgg)
+			}
 			if partialAgg != nil {
 				if cop.tablePlan != nil {
 					cop.finishIndexPlan()