planner/core: fix bug when using generated column in aggregate statem…

…ent (pingcap#14111)

cmd/explaintest/r/generated_columns.result

@@ -172,3 +172,45 @@ Projection_4	10.00	root	test.person.name
 └─IndexLookUp_10	10.00	root	
   ├─IndexScan_8	10.00	cop[tikv]	table:person, index:city_no, range:[1,1], keep order:false, stats:pseudo
   └─TableScan_9	10.00	cop[tikv]	table:person, keep order:false, stats:pseudo
+DROP TABLE IF EXISTS t1;
+CREATE TABLE t1 (a INT,
+b INT GENERATED ALWAYS AS (-a) VIRTUAL,
+c INT GENERATED ALWAYS AS (-a) STORED,
+index (c));
+INSERT INTO t1 (a) VALUES (2), (1), (1), (3), (NULL);
+EXPLAIN SELECT sum(a) FROM t1 GROUP BY b;
+id	count	task	operator info
+HashAgg_5	8000.00	root	group by:Column#7, funcs:sum(Column#6)->Column#5
+└─Projection_12	10000.00	root	cast(test.t1.a)->Column#6, test.t1.b
+  └─TableReader_9	10000.00	root	data:TableScan_8
+    └─TableScan_8	10000.00	cop[tikv]	table:t1, range:[-inf,+inf], keep order:false, stats:pseudo
+EXPLAIN SELECT sum(a) FROM t1 GROUP BY c;
+id	count	task	operator info
+HashAgg_11	8000.00	root	group by:test.t1.c, funcs:sum(Column#6)->Column#5
+└─TableReader_12	8000.00	root	data:HashAgg_5
+  └─HashAgg_5	8000.00	cop[tikv]	group by:test.t1.c, funcs:sum(test.t1.a)->Column#6
+    └─TableScan_10	10000.00	cop[tikv]	table:t1, range:[-inf,+inf], keep order:false, stats:pseudo
+EXPLAIN SELECT sum(b) FROM t1 GROUP BY a;
+id	count	task	operator info
+HashAgg_5	8000.00	root	group by:Column#7, funcs:sum(Column#6)->Column#5
+└─Projection_12	10000.00	root	cast(test.t1.b)->Column#6, test.t1.a
+  └─TableReader_9	10000.00	root	data:TableScan_8
+    └─TableScan_8	10000.00	cop[tikv]	table:t1, range:[-inf,+inf], keep order:false, stats:pseudo
+EXPLAIN SELECT sum(b) FROM t1 GROUP BY c;
+id	count	task	operator info
+HashAgg_5	8000.00	root	group by:Column#9, funcs:sum(Column#8)->Column#5
+└─Projection_18	10000.00	root	cast(test.t1.b)->Column#8, test.t1.c
+  └─TableReader_11	10000.00	root	data:TableScan_10
+    └─TableScan_10	10000.00	cop[tikv]	table:t1, range:[-inf,+inf], keep order:false, stats:pseudo
+EXPLAIN SELECT sum(c) FROM t1 GROUP BY a;
+id	count	task	operator info
+HashAgg_9	8000.00	root	group by:test.t1.a, funcs:sum(Column#6)->Column#5
+└─TableReader_10	8000.00	root	data:HashAgg_5
+  └─HashAgg_5	8000.00	cop[tikv]	group by:test.t1.a, funcs:sum(test.t1.c)->Column#6
+    └─TableScan_8	10000.00	cop[tikv]	table:t1, range:[-inf,+inf], keep order:false, stats:pseudo
+EXPLAIN SELECT sum(c) FROM t1 GROUP BY b;
+id	count	task	operator info
+HashAgg_5	8000.00	root	group by:Column#7, funcs:sum(Column#6)->Column#5
+└─Projection_12	10000.00	root	cast(test.t1.c)->Column#6, test.t1.b
+  └─TableReader_9	10000.00	root	data:TableScan_8
+    └─TableScan_8	10000.00	cop[tikv]	table:t1, range:[-inf,+inf], keep order:false, stats:pseudo

cmd/explaintest/t/generated_columns.test

@@ -111,3 +111,20 @@ KEY(city_no));
 
 INSERT INTO person (name, address_info) VALUES ("John", CAST('{"city_no": 1}' AS JSON));
 EXPLAIN SELECT name FROM person where city_no=1;
+
+-- Virtual generated columns in aggregate statement
+-- ISSUE https://github.com/pingcap/tidb/issues/14072
+
+DROP TABLE IF EXISTS t1;
+CREATE TABLE t1 (a INT,
+                 b INT GENERATED ALWAYS AS (-a) VIRTUAL,
+                 c INT GENERATED ALWAYS AS (-a) STORED,
+                 index (c));
+INSERT INTO t1 (a) VALUES (2), (1), (1), (3), (NULL);
+
+EXPLAIN SELECT sum(a) FROM t1 GROUP BY b;
+EXPLAIN SELECT sum(a) FROM t1 GROUP BY c;
+EXPLAIN SELECT sum(b) FROM t1 GROUP BY a;
+EXPLAIN SELECT sum(b) FROM t1 GROUP BY c;
+EXPLAIN SELECT sum(c) FROM t1 GROUP BY a;
+EXPLAIN SELECT sum(c) FROM t1 GROUP BY b;

executor/aggregate_test.go

@@ -369,6 +369,17 @@ func (s *testSuiteAgg) TestAggregation(c *C) {
 	c.Assert(errors.Cause(err).Error(), Equals, "unsupported agg function: var_pop")
 	_, err = tk.Exec("select var_samp(a) from t")
 	c.Assert(errors.Cause(err).Error(), Equals, "unsupported agg function: var_samp")
+
+	// For issue #14072: wrong result when using generated column with aggregate statement
+	tk.MustExec("drop table if exists t1;")
+	tk.MustExec("create table t1 (a int, b int generated always as (-a) virtual, c int generated always as (-a) stored);")
+	tk.MustExec("insert into t1 (a) values (2), (1), (1), (3), (NULL);")
+	tk.MustQuery("select sum(a) from t1 group by b order by b;").Check(testkit.Rows("<nil>", "3", "2", "2"))
+	tk.MustQuery("select sum(a) from t1 group by c order by c;").Check(testkit.Rows("<nil>", "3", "2", "2"))
+	tk.MustQuery("select sum(b) from t1 group by a order by a;").Check(testkit.Rows("<nil>", "-2", "-2", "-3"))
+	tk.MustQuery("select sum(b) from t1 group by c order by c;").Check(testkit.Rows("<nil>", "-3", "-2", "-2"))
+	tk.MustQuery("select sum(c) from t1 group by a order by a;").Check(testkit.Rows("<nil>", "-2", "-2", "-3"))
+	tk.MustQuery("select sum(c) from t1 group by b order by b;").Check(testkit.Rows("<nil>", "-3", "-2", "-2"))
 }
 
 func (s *testSuiteAgg) TestAggPrune(c *C) {

planner/core/task.go

@@ -937,7 +937,7 @@ func (sel *PhysicalSelection) attach2Task(tasks ...task) task {
 	return t
 }
 
-// CheckAggCanPushCop checks whether the aggFuncs with groupByItems can
+// CheckAggCanPushCop checks whether the aggFuncs and groupByItems can
 // be pushed down to coprocessor.
 func CheckAggCanPushCop(sctx sessionctx.Context, aggFuncs []*aggregation.AggFuncDesc, groupByItems []expression.Expression, copToFlash bool) bool {
 	sc := sctx.GetSessionVars().StmtCtx
@@ -959,6 +959,9 @@ func CheckAggCanPushCop(sctx sessionctx.Context, aggFuncs []*aggregation.AggFunc
 			return false
 		}
 	}
+	if expression.ContainVirtualColumn(groupByItems) {
+		return false
+	}
 	_, _, remained := expression.ExpressionsToPB(sc, groupByItems, client)
 	if len(remained) > 0 {
 		return false