func (a *aggPushDownSolver) makeNewAgg(aggFuncs []expression.AggregationFunction, gbyCols []*expression.Column) *Aggregation {
agg := &Aggregation{
GroupByItems: expression.Schema2Exprs(gbyCols),
baseLogicalPlan: newBaseLogicalPlan(Agg, a.alloc),
groupByCols: gbyCols,
}
agg.initIDAndContext(a.ctx)
var newAggFuncs []expression.AggregationFunction
schema := make(expression.Schema, 0, len(aggFuncs))
for _, aggFunc := range aggFuncs {
var newFuncs []expression.AggregationFunction
newFuncs, schema = a.decompose(aggFunc, schema, agg.GetID())
newAggFuncs = append(newAggFuncs, newFuncs...)
for _, arg := range aggFunc.GetArgs() {
agg.correlated = agg.correlated || arg.IsCorrelated()
}
}
for _, gbyCol := range gbyCols {
firstRow := expression.NewAggFunction(ast.AggFuncFirstRow, []expression.Expression{gbyCol.Clone()}, false)
newAggFuncs = append(newAggFuncs, firstRow)
schema = append(schema, gbyCol.Clone().(*expression.Column))
}
agg.AggFuncs = newAggFuncs
agg.SetSchema(schema)
return agg
}
func (b *planBuilder) buildAggregation(p LogicalPlan, aggFuncList []*ast.AggregateFuncExpr, gby []expression.Expression, correlated bool) LogicalPlan {
agg := &Aggregation{
AggFuncs: make([]expression.AggregationFunction, 0, len(aggFuncList)),
baseLogicalPlan: newBaseLogicalPlan(Agg, b.allocator)}
agg.initID()
agg.correlated = p.IsCorrelated() || correlated
addChild(agg, p)
schema := make([]*expression.Column, 0, len(aggFuncList))
for i, aggFunc := range aggFuncList {
var newArgList []expression.Expression
for _, arg := range aggFunc.Args {
newArg, np, correlated, err := b.rewrite(arg, p, nil, true)
if err != nil {
b.err = errors.Trace(err)
return nil
}
p = np
agg.correlated = correlated || agg.correlated
newArgList = append(newArgList, newArg)
}
agg.AggFuncs = append(agg.AggFuncs, expression.NewAggFunction(aggFunc.F, newArgList, aggFunc.Distinct))
schema = append(schema, &expression.Column{FromID: agg.id,
ColName: model.NewCIStr(fmt.Sprintf("%s_col_%d", agg.id, i)),
Position: i,
IsAggOrSubq: true,
RetType: aggFunc.GetType()})
}
agg.GroupByItems = gby
agg.SetSchema(schema)
return agg
}
func (p *physicalTableSource) addAggregation(agg *PhysicalAggregation) expression.Schema {
if p.client == nil {
return nil
}
for _, f := range agg.AggFuncs {
pb := aggFuncToPBExpr(p.client, f)
if pb == nil {
p.clear()
return nil
}
p.AggFuncs = append(p.AggFuncs, pb)
}
for _, item := range agg.GroupByItems {
pb := groupByItemToPB(p.client, item)
if pb == nil {
p.clear()
return nil
}
p.GbyItems = append(p.GbyItems, pb)
}
p.Aggregated = true
gk := types.NewFieldType(mysql.TypeBlob)
gk.Charset = charset.CharsetBin
gk.Collate = charset.CollationBin
p.AggFields = append(p.AggFields, gk)
var schema expression.Schema
cursor := 0
schema = append(schema, &expression.Column{Index: cursor, ColName: model.NewCIStr(fmt.Sprint(agg.GroupByItems))})
agg.GroupByItems = []expression.Expression{schema[cursor]}
newAggFuncs := make([]expression.AggregationFunction, len(agg.AggFuncs))
for i, aggFun := range agg.AggFuncs {
fun := expression.NewAggFunction(aggFun.GetName(), nil, false)
var args []expression.Expression
colName := model.NewCIStr(fmt.Sprint(aggFun.GetArgs()))
if needCount(fun) {
cursor++
schema = append(schema, &expression.Column{Index: cursor, ColName: colName})
args = append(args, schema[cursor])
ft := types.NewFieldType(mysql.TypeLonglong)
ft.Flen = 21
ft.Charset = charset.CharsetBin
ft.Collate = charset.CollationBin
p.AggFields = append(p.AggFields, ft)
}
if needValue(fun) {
cursor++
schema = append(schema, &expression.Column{Index: cursor, ColName: colName})
args = append(args, schema[cursor])
p.AggFields = append(p.AggFields, agg.schema[i].GetType())
}
fun.SetArgs(args)
fun.SetMode(expression.FinalMode)
newAggFuncs[i] = fun
}
agg.AggFuncs = newAggFuncs
return schema
}
func (b *planBuilder) buildAggregation(p LogicalPlan, aggFuncList []*ast.AggregateFuncExpr, gbyItems []expression.Expression) (LogicalPlan, map[int]int) {
agg := &Aggregation{
AggFuncs: make([]expression.AggregationFunction, 0, len(aggFuncList)),
baseLogicalPlan: newBaseLogicalPlan(Agg, b.allocator)}
agg.self = agg
agg.initIDAndContext(b.ctx)
agg.correlated = p.IsCorrelated()
for _, item := range gbyItems {
agg.correlated = agg.correlated || item.IsCorrelated()
}
addChild(agg, p)
schema := make(expression.Schema, 0, len(aggFuncList))
// aggIdxMap maps the old index to new index after applying common aggregation functions elimination.
aggIndexMap := make(map[int]int)
for i, aggFunc := range aggFuncList {
var newArgList []expression.Expression
for _, arg := range aggFunc.Args {
newArg, np, err := b.rewrite(arg, p, nil, true)
if err != nil {
b.err = errors.Trace(err)
return nil, nil
}
p = np
agg.correlated = agg.correlated || newArg.IsCorrelated()
newArgList = append(newArgList, newArg)
}
newFunc := expression.NewAggFunction(aggFunc.F, newArgList, aggFunc.Distinct)
combined := false
for j, oldFunc := range agg.AggFuncs {
if oldFunc.Equal(newFunc, b.ctx) {
aggIndexMap[i] = j
combined = true
break
}
}
if !combined {
position := len(agg.AggFuncs)
aggIndexMap[i] = position
agg.AggFuncs = append(agg.AggFuncs, newFunc)
schema = append(schema, &expression.Column{
FromID: agg.id,
ColName: model.NewCIStr(fmt.Sprintf("%s_col_%d", agg.id, position)),
Position: position,
IsAggOrSubq: true,
RetType: aggFunc.GetType()})
}
}
agg.GroupByItems = gbyItems
agg.SetSchema(schema)
agg.collectGroupByColumns()
return agg, aggIndexMap
}
func (b *planBuilder) buildAggregation(p Plan, aggFuncList []*ast.AggregateFuncExpr, gby *ast.GroupByClause) Plan {
newAggFuncList := make([]expression.AggregationFunction, 0, len(aggFuncList))
agg := &Aggregation{}
agg.id = b.allocID(agg)
agg.correlated = p.IsCorrelated()
addChild(agg, p)
schema := make([]*expression.Column, 0, len(aggFuncList))
for i, aggFunc := range aggFuncList {
var newArgList []expression.Expression
for _, arg := range aggFunc.Args {
newArg, np, correlated, err := b.rewrite(arg, p, nil)
if err != nil {
b.err = errors.Trace(err)
return nil
}
p = np
agg.correlated = correlated || agg.correlated
newArgList = append(newArgList, newArg)
}
newAggFuncList = append(newAggFuncList, expression.NewAggFunction(aggFunc.F, newArgList, aggFunc.Distinct))
schema = append(schema, &expression.Column{FromID: agg.id,
ColName: model.NewCIStr(fmt.Sprintf("%s_col_%d", agg.id, i))})
}
var gbyExprList []expression.Expression
if gby != nil {
gbyExprList = make([]expression.Expression, 0, len(gby.Items))
for _, gbyItem := range gby.Items {
gbyExpr, np, correlated, err := b.rewrite(gbyItem.Expr, p, nil)
if err != nil {
b.err = errors.Trace(err)
return nil
}
p = np
agg.correlated = correlated || agg.correlated
gbyExprList = append(gbyExprList, gbyExpr)
}
}
agg.AggFuncs = newAggFuncList
agg.GroupByItems = gbyExprList
agg.SetSchema(schema)
return agg
}
func (s *testSuite) TestStreamAgg(c *C) {
col := &expression.Column{
Index: 1,
}
gbyCol := &expression.Column{
Index: 0,
}
sumAgg := expression.NewAggFunction(ast.AggFuncSum, []expression.Expression{col}, false)
cntAgg := expression.NewAggFunction(ast.AggFuncCount, []expression.Expression{col}, false)
avgAgg := expression.NewAggFunction(ast.AggFuncAvg, []expression.Expression{col}, false)
maxAgg := expression.NewAggFunction(ast.AggFuncMax, []expression.Expression{col}, false)
cases := []struct {
aggFunc expression.AggregationFunction
result string
input [][]interface{}
result1 []string
}{
{
sumAgg,
"<nil>",
[][]interface{}{
{0, 1}, {0, nil}, {1, 2}, {1, 3},
},
[]string{
"1", "5",
},
},
{
cntAgg,
"0",
[][]interface{}{
{0, 1}, {0, nil}, {1, 2}, {1, 3},
},
[]string{
"1", "2",
},
},
{
avgAgg,
"<nil>",
[][]interface{}{
{0, 1}, {0, nil}, {1, 2}, {1, 3},
},
[]string{
"1.0000", "2.5000",
},
},
{
maxAgg,
"<nil>",
[][]interface{}{
{0, 1}, {0, nil}, {1, 2}, {1, 3},
},
[]string{
"1", "3",
},
},
}
for _, ca := range cases {
mock := &MockExec{}
e := &executor.StreamAggExec{
AggFuncs: []expression.AggregationFunction{ca.aggFunc},
Src: mock,
}
row, err := e.Next()
c.Check(err, IsNil)
c.Check(row, NotNil)
c.Assert(fmt.Sprintf("%v", row.Data[0].GetValue()), Equals, ca.result)
e.GroupByItems = append(e.GroupByItems, gbyCol)
e.Close()
row, err = e.Next()
c.Check(err, IsNil)
c.Check(row, IsNil)
e.Close()
for _, input := range ca.input {
data := types.MakeDatums(input...)
mock.Rows = append(mock.Rows, &executor.Row{Data: data})
}
for _, res := range ca.result1 {
row, err = e.Next()
c.Check(err, IsNil)
c.Check(row, NotNil)
c.Assert(fmt.Sprintf("%v", row.Data[0].GetValue()), Equals, res)
}
}
}
func (p *physicalTableSource) addAggregation(ctx context.Context, agg *PhysicalAggregation) expression.Schema {
if p.client == nil {
return nil
}
sc := ctx.GetSessionVars().StmtCtx
for _, f := range agg.AggFuncs {
pb := aggFuncToPBExpr(sc, p.client, f)
if pb == nil {
// When we fail to convert any agg function to PB struct, we should clear the environments.
p.clearForAggPushDown()
return nil
}
p.AggFuncsPB = append(p.AggFuncsPB, pb)
p.aggFuncs = append(p.aggFuncs, f.Clone())
}
for _, item := range agg.GroupByItems {
pb := groupByItemToPB(sc, p.client, item)
if pb == nil {
// When we fail to convert any group-by item to PB struct, we should clear the environments.
p.clearForAggPushDown()
return nil
}
p.GbyItemsPB = append(p.GbyItemsPB, pb)
p.gbyItems = append(p.gbyItems, item.Clone())
}
p.Aggregated = true
gk := types.NewFieldType(mysql.TypeBlob)
gk.Charset = charset.CharsetBin
gk.Collate = charset.CollationBin
p.AggFields = append(p.AggFields, gk)
var schema expression.Schema
cursor := 0
schema = append(schema, &expression.Column{Index: cursor, ColName: model.NewCIStr(fmt.Sprint(agg.GroupByItems))})
agg.GroupByItems = []expression.Expression{schema[cursor]}
newAggFuncs := make([]expression.AggregationFunction, len(agg.AggFuncs))
for i, aggFun := range agg.AggFuncs {
fun := expression.NewAggFunction(aggFun.GetName(), nil, false)
var args []expression.Expression
colName := model.NewCIStr(fmt.Sprint(aggFun.GetArgs()))
if needCount(fun) {
cursor++
schema = append(schema, &expression.Column{Index: cursor, ColName: colName})
args = append(args, schema[cursor])
ft := types.NewFieldType(mysql.TypeLonglong)
ft.Flen = 21
ft.Charset = charset.CharsetBin
ft.Collate = charset.CollationBin
p.AggFields = append(p.AggFields, ft)
}
if needValue(fun) {
cursor++
schema = append(schema, &expression.Column{Index: cursor, ColName: colName})
args = append(args, schema[cursor])
p.AggFields = append(p.AggFields, agg.schema[i].GetType())
}
fun.SetArgs(args)
fun.SetMode(expression.FinalMode)
newAggFuncs[i] = fun
}
agg.AggFuncs = newAggFuncs
return schema
}
