Ejemplo n.º 1
0
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
}
Ejemplo n.º 2
0
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
}
Ejemplo n.º 3
0
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
}
Ejemplo n.º 4
0
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
}
Ejemplo n.º 5
0
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
}
Ejemplo n.º 6
0
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)
		}
	}
}
Ejemplo n.º 7
0
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
}