This is an automated cherry-pick of pingcap#46310

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>

Author: AilinKid

executor/test/tiflashtest/BUILD.bazel

@@ -0,0 +1,36 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+
+go_test(
+    name = "tiflashtest_test",
+    timeout = "moderate",
+    srcs = [
+        "main_test.go",
+        "tiflash_test.go",
+    ],
+    flaky = True,
+    race = "on",
+    shard_count = 39,
+    deps = [
+        "//config",
+        "//domain",
+        "//kv",
+        "//meta/autoid",
+        "//parser/terror",
+        "//planner/core",
+        "//store/mockstore",
+        "//store/mockstore/unistore",
+        "//testkit",
+        "//testkit/external",
+        "//util/dbterror/exeerrors",
+        "//util/memory",
+        "//util/tiflashcompute",
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
+        "@com_github_pingcap_kvproto//pkg/metapb",
+        "@com_github_stretchr_testify//require",
+        "@com_github_tikv_client_go_v2//testutils",
+        "@com_github_tikv_client_go_v2//tikv",
+        "@io_opencensus_go//stats/view",
+        "@org_uber_go_goleak//:goleak",
+    ],
+)

executor/tiflash_test.go

@@ -1248,6 +1248,37 @@ func TestAggPushDownCountStar(t *testing.T) {
 	tk.MustQuery("select count(*) from c, o where c.c_id=o.c_id").Check(testkit.Rows("5"))
 }
 
+func TestAggPushDownUnionAndMPP(t *testing.T) {
+	store := testkit.CreateMockStore(t, withMockTiFlash(2))
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("create table t (a int, b int)")
+	tk.MustExec("alter table t set tiflash replica 1")
+	tk.MustExec("insert into t values (1, 1);")
+	tk.MustExec("insert into t values (1, 1);")
+	tk.MustExec("insert into t values (1, 1);")
+	tk.MustExec("insert into t values (1, 1);")
+	tk.MustExec("insert into t values (1, 1);")
+	tk.MustExec("set @@tidb_allow_mpp=1;")
+	tk.MustExec("set @@tidb_enforce_mpp=1;")
+	tk.MustExec("set @@tidb_opt_agg_push_down=1")
+
+	tk.MustExec("create table c(c_id int)")
+	tk.MustExec("create table o(o_id int, c_id int)")
+	tk.MustExec("insert into c values(1),(1),(1),(1)")
+	tk.MustExec("insert into o values(1,1),(1,1),(1,2)")
+	tk.MustExec("alter table c set tiflash replica 1")
+	tk.MustExec("alter table o set tiflash replica 1")
+
+	tk.MustQuery("select a, count(*) from (select a, b from t " +
+		"union all " +
+		"select a, b from t" +
+		") t group by a order by a limit 10;").Check(testkit.Rows("1 10"))
+
+	tk.MustQuery("select o.o_id, count(*) from c, o where c.c_id=o.o_id group by o.o_id").Check(testkit.Rows("1 12"))
+}
+
 func TestGroupStreamAggOnTiFlash(t *testing.T) {
 	store, clean := testkit.CreateMockStore(t, withMockTiFlash(2))
 	defer clean()