plan: change the way that we find column in group-by items in cbo phase. (#1960)

plan/physical_plan_builder.go

@@ -627,36 +627,24 @@ func (p *Aggregation) convert2PhysicalPlanStream(prop *requiredProperty) (*physi
 	agg.HasGby = len(p.GroupByItems) > 0
 	agg.SetSchema(p.schema)
 	// TODO: Consider distinct key.
-	gbyCols := make([]*expression.Column, 0, len(p.GroupByItems)+1)
 	info := &physicalPlanInfo{cost: math.MaxFloat64}
-	// TODO: extract columns in monotonic functions.
-	for _, item := range p.GroupByItems {
-		if col, ok := item.(*expression.Column); ok {
-			if !col.Correlated {
-				gbyCols = append(gbyCols, col)
-			}
-		} else {
-			// group by a + b is not interested in any order.
-			return info, nil
-		}
+	gbyCols := p.groupByCols
+	if len(gbyCols) != len(p.GroupByItems) {
+		// group by a + b is not interested in any order.
+		return info, nil
 	}
 	isSortKey := make([]bool, len(gbyCols))
 	newProp := &requiredProperty{
 		props: make([]*columnProp, 0, len(gbyCols)),
 	}
 	for _, pro := range prop.props {
-		var matchedCol *expression.Column
-		for i, col := range gbyCols {
-			if !col.Correlated && col.ColName.L == pro.col.ColName.L {
-				isSortKey[i] = true
-				matchedCol = col
-			}
-		}
-		if matchedCol == nil {
+		idx := p.getGbyColIndex(pro.col)
+		if idx == -1 {
 			return info, nil
 		}
+		isSortKey[idx] = true
 		// We should add columns in aggregation in order to keep index right.
-		newProp.props = append(newProp.props, &columnProp{col: matchedCol, desc: pro.desc})
+		newProp.props = append(newProp.props, &columnProp{col: gbyCols[idx], desc: pro.desc})
 	}
 	newProp.sortKeyLen = len(newProp.props)
 	for i, col := range gbyCols {

plan/plan_test.go

@@ -871,6 +871,10 @@ func (s *testPlanSuite) TestCBO(c *C) {
 			sql:  "select count(*) from t group by a",
 			best: "Table(t)->StreamAgg",
 		},
+		{
+			sql:  "select count(*) from t group by a order by a",
+			best: "Table(t)->StreamAgg->Trim",
+		},
 		{
 			sql:  "select count(distinct e) from t where c = 1 group by d",
 			best: "Index(t.c_d_e)[[1,1]]->StreamAgg",

plan/predicate_push_down.go

@@ -538,11 +538,14 @@ func (p *Union) PredicatePushDown(predicates []expression.Expression) (ret []exp
 	return
 }
 
-// CheckColumnIsGroupByItem check whether a column is a group-by item.
-func (p *Aggregation) CheckColumnIsGroupByItem(col *expression.Column) bool {
+// getGbyColIndex gets the column's index in the group-by columns.
+func (p *Aggregation) getGbyColIndex(col *expression.Column) int {
 	id := p.GetSchema().GetIndex(col)
 	colOriginal, isColumn := p.AggFuncs[id].GetArgs()[0].(*expression.Column)
-	return isColumn && expression.Schema(p.groupByCols).GetIndex(colOriginal) != -1
+	if !isColumn {
+		return -1
+	}
+	return expression.Schema(p.groupByCols).GetIndex(colOriginal)
 }
 
 // PredicatePushDown implements LogicalPlan PredicatePushDown interface.
@@ -565,7 +568,7 @@ func (p *Aggregation) PredicatePushDown(predicates []expression.Expression) (ret
 			extractedCols, _ := extractColumn(cond, nil, nil)
 			ok := true
 			for _, col := range extractedCols {
-				if !p.CheckColumnIsGroupByItem(col) {
+				if p.getGbyColIndex(col) == -1 {
 					ok = false
 					break
 				}