planner: fix wrong aggregate pruning for some cases (#25289) (#26033)

pingcap · Jul 14, 2021 · 4a6db53 · 4a6db53
1 parent 67b641d
commit 4a6db53
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/planner/core/rule_column_pruning.go b/planner/core/rule_column_pruning.go
@@ -89,13 +89,16 @@ func (la *LogicalAggregation) PruneColumns(parentUsedCols []*expression.Column)
 	used := expression.GetUsedList(parentUsedCols, la.Schema())
 
 	allFirstRow := true
+	allRemainFirstRow := true
 	for i := len(used) - 1; i >= 0; i-- {
 		if la.AggFuncs[i].Name != ast.AggFuncFirstRow {
 			allFirstRow = false
 		}
 		if !used[i] {
 			la.schema.Columns = append(la.schema.Columns[:i], la.schema.Columns[i+1:]...)
 			la.AggFuncs = append(la.AggFuncs[:i], la.AggFuncs[i+1:]...)
+		} else if la.AggFuncs[i].Name != ast.AggFuncFirstRow {
+			allRemainFirstRow = false
 		}
 	}
 	var selfUsedCols []*expression.Column
@@ -106,7 +109,7 @@ func (la *LogicalAggregation) PruneColumns(parentUsedCols []*expression.Column)
 		aggrFunc.OrderByItems, cols = pruneByItems(aggrFunc.OrderByItems)
 		selfUsedCols = append(selfUsedCols, cols...)
 	}
-	if len(la.AggFuncs) == 0 {
+	if len(la.AggFuncs) == 0 || (!allFirstRow && allRemainFirstRow) {
 		// If all the aggregate functions are pruned, we should add an aggregate function to maintain the info of row numbers.
 		// For all the aggregate functions except `first_row`, if we have an empty table defined as t(a,b),
 		// `select agg(a) from t` would always return one row, while `select agg(a) from t group by b` would return empty.
@@ -121,12 +124,12 @@ func (la *LogicalAggregation) PruneColumns(parentUsedCols []*expression.Column)
 		if err != nil {
 			return err
 		}
-		la.AggFuncs = []*aggregation.AggFuncDesc{newAgg}
+		la.AggFuncs = append(la.AggFuncs, newAgg)
 		col := &expression.Column{
 			UniqueID: la.ctx.GetSessionVars().AllocPlanColumnID(),
 			RetType:  newAgg.RetTp,
 		}
-		la.schema.Columns = []*expression.Column{col}
+		la.schema.Columns = append(la.schema.Columns, col)
 	}
 
 	if len(la.GroupByItems) > 0 {

diff --git a/planner/core/testdata/integration_suite_in.json b/planner/core/testdata/integration_suite_in.json
@@ -31,7 +31,8 @@
       "select count(1) from t join (select max(a) from t where false group by a) as tmp",
       "select count(1) from t join (select min(a) from t where false group by a) as tmp",
       "select count(1) from t join (select sum(a) from t where false group by a) as tmp",
-      "select count(1) from t join (select avg(a) from t where false group by a) as tmp"
+      "select count(1) from t join (select avg(a) from t where false group by a) as tmp",
+      "SELECT avg(2) FROM(SELECT min(c) FROM t JOIN(SELECT 1 c) d ORDER BY a) e"
     ]
   },
   {

diff --git a/planner/core/testdata/integration_suite_out.json b/planner/core/testdata/integration_suite_out.json
@@ -125,6 +125,12 @@
         "Res": [
           "0"
         ]
+      },
+      {
+        "SQL": "SELECT avg(2) FROM(SELECT min(c) FROM t JOIN(SELECT 1 c) d ORDER BY a) e",
+        "Res": [
+          "2.0000"
+        ]
       }
     ]
   },