Repository: pingcap/tidb (Go)

Issue:
DO with group functions is not working in TiDB.
## Bug Report

Please answer these questions before submitting your issue. Thanks!

### 1. Minimal reproduce step (Required)
```sql
 DO GROUP_CONCAT(NULLIF(ELT(1, @e), 2.0) ORDER BY 1);
```

### 2. What did you expect to see? (Required)
```sql
mysql> DO GROUP_CONCAT(NULLIF(ELT(1, @e), 2.0) ORDER BY 1);
Query OK, 0 rows affected (0.00 sec)
```
### 3. What did you see instead (Required)
```sql
mysql> DO GROUP_CONCAT(NULLIF(ELT(1, @e), 2.0) ORDER BY 1);
ERROR 1111 (HY000): Invalid use of group function
```
### 4. What is your TiDB version? (Required)
```sql
| Release Version: v5.4.0-alpha-264-g6efa36df6
Edition: Community
Git Commit Hash: 6efa36df6cff325106f67ecfe3d79816ba37ca6a
Git Branch: master
UTC Build Time: 2021-12-28 02:03:55
GoVersion: go1.17.2
Race Enabled: false
TiKV Min Version: v3.0.0-60965b006877ca7234adaced7890d7b029ed1306
Check Table Before Drop: false |
```




Discussion:
It's a rare usage of the statement, we move down the severity.
Definition: DO executes the expressions but does not return any results. In most respects, DO is shorthand for SELECT expr, ..., but has the advantage that it is slightly faster when you do not care about the result.

Reason: In normal build-phase, the aggFunc is built ahead of other select fields, informing the latter one to skip rewrite progress of the agg field via aggMapper. While we miss this when handling do-statement.



Patch:
diff --git a/planner/core/logical_plan_builder.go b/planner/core/logical_plan_builder.go
index 499af00131e79..6e05f37c22334 100644
--- a/planner/core/logical_plan_builder.go
+++ b/planner/core/logical_plan_builder.go
@@ -2206,6 +2206,20 @@ func (b *PlanBuilder) resolveHavingAndOrderBy(sel *ast.SelectStmt, p LogicalPlan
 	return havingAggMapper, extractor.aggMapper, nil
 }
 
+func (b *PlanBuilder) extractAggFuncsInExprs(exprs []ast.ExprNode) ([]*ast.AggregateFuncExpr, map[*ast.AggregateFuncExpr]int) {
+	extractor := &AggregateFuncExtractor{skipAggMap: b.correlatedAggMapper}
+	for _, expr := range exprs {
+		expr.Accept(extractor)
+	}
+	aggList := extractor.AggFuncs
+	totalAggMapper := make(map[*ast.AggregateFuncExpr]int, len(aggList))
+
+	for i, agg := range aggList {
+		totalAggMapper[agg] = i
+	}
+	return aggList, totalAggMapper
+}
+
 func (b *PlanBuilder) extractAggFuncsInSelectFields(fields []*ast.SelectField) ([]*ast.AggregateFuncExpr, map[*ast.AggregateFuncExpr]int) {
 	extractor := &AggregateFuncExtractor{skipAggMap: b.correlatedAggMapper}
 	for _, f := range fields {
diff --git a/planner/core/planbuilder.go b/planner/core/planbuilder.go
index 75e9a6677f868..8133976912756 100644
--- a/planner/core/planbuilder.go
+++ b/planner/core/planbuilder.go
@@ -787,8 +787,39 @@ func (b *PlanBuilder) buildDo(ctx context.Context, v *ast.DoStmt) (Plan, error)
 	proj := LogicalProjection{Exprs: make([]expression.Expression, 0, len(v.Exprs))}.Init(b.ctx, b.getSelectOffset())
 	proj.names = make([]*types.FieldName, len(v.Exprs))
 	schema := expression.NewSchema(make([]*expression.Column, 0, len(v.Exprs))...)
+
+	// Since do statement only contain expression list, and it may contain aggFunc, detecting to build the aggMapper firstly.
+	var (
+		err      error
+		aggFuncs []*ast.AggregateFuncExpr
+		totalMap map[*ast.AggregateFuncExpr]int
+	)
+	hasAgg := b.detectAggInExprNode(v.Exprs)
+	needBuildAgg := hasAgg
+	if hasAgg {
+		if b.buildingRecursivePartForCTE {
+			return nil, ErrCTERecursiveForbidsAggregation.GenWithStackByArgs(b.genCTETableNameForError())
+		}
+
+		aggFuncs, totalMap = b.extractAggFuncsInExprs(v.Exprs)
+
+		if len(aggFuncs) == 0 {
+			needBuildAgg = false
+		}
+	}
+	if needBuildAgg {
+		var aggIndexMap map[int]int
+		p, aggIndexMap, err = b.buildAggregation(ctx, p, aggFuncs, nil, nil)
+		if err != nil {
+			return nil, err
+		}
+		for agg, idx := range totalMap {
+			totalMap[agg] = aggIndexMap[idx]
+		}
+	}
+
 	for _, astExpr := range v.Exprs {
-		expr, np, err := b.rewrite(ctx, astExpr, p, nil, true)
+		expr, np, err := b.rewrite(ctx, astExpr, p, totalMap, true)
 		if err != nil {
 			return nil, err
 		}
@@ -907,6 +938,16 @@ func (b *PlanBuilder) buildCreateBindPlan(v *ast.CreateBindingStmt) (Plan, error
 	return p, nil
 }
 
+// detectAggInExprNode detects an aggregate function in its exprs.
+func (b *PlanBuilder) detectAggInExprNode(exprs []ast.ExprNode) bool {
+	for _, expr := range exprs {
+		if ast.HasAggFlag(expr) {
+			return true
+		}
+	}
+	return false
+}
+
 // detectSelectAgg detects an aggregate function or GROUP BY clause.
 func (b *PlanBuilder) detectSelectAgg(sel *ast.SelectStmt) bool {
 	if sel.GroupBy != nil {


Test patch:
diff --git a/executor/simple_test.go b/executor/simple_test.go
index a637ad72efb16..d81e0d662be6c 100644
--- a/executor/simple_test.go
+++ b/executor/simple_test.go
@@ -54,6 +54,17 @@ func TestDo(t *testing.T) {
 	tk.MustQuery("select * from t").Check(testkit.Rows("1", "2"))
 }
 
+func TestDoWithAggFunc(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("DO sum(1)")
+	tk.MustExec("DO avg(@e+@f)")
+	tk.MustExec("DO GROUP_CONCAT(NULLIF(ELT(1, @e), 2.0) ORDER BY 1)")
+}
+
 func TestSetRoleAllCorner(t *testing.T) {
 	store, clean := testkit.CreateMockStore(t)
 	defer clean()