// Copyright 2016 PingCAP, Inc. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // See the License for the specific language governing permissions and // limitations under the License. package aggregation import ( "bytes" "strings" "github.com/pingcap/errors" "github.com/pingcap/parser/ast" "github.com/pingcap/tidb/expression" "github.com/pingcap/tidb/kv" "github.com/pingcap/tidb/sessionctx/stmtctx" "github.com/pingcap/tidb/types" "github.com/pingcap/tidb/util/chunk" "github.com/pingcap/tipb/go-tipb" ) // Aggregation stands for aggregate functions. type Aggregation interface { // Update during executing. Update(evalCtx *AggEvaluateContext, sc *stmtctx.StatementContext, row chunk.Row) error // GetPartialResult will called by coprocessor to get partial results. For avg function, partial results will return // sum and count values at the same time. GetPartialResult(evalCtx *AggEvaluateContext) []types.Datum // GetResult will be called when all data have been processed. GetResult(evalCtx *AggEvaluateContext) types.Datum // CreateContext creates a new AggEvaluateContext for the aggregation function. CreateContext(sc *stmtctx.StatementContext) *AggEvaluateContext // ResetContext resets the content of the evaluate context. ResetContext(sc *stmtctx.StatementContext, evalCtx *AggEvaluateContext) } // NewDistAggFunc creates new Aggregate function for mock tikv. func NewDistAggFunc(expr *tipb.Expr, fieldTps []*types.FieldType, sc *stmtctx.StatementContext) (Aggregation, error) { args := make([]expression.Expression, 0, len(expr.Children)) for _, child := range expr.Children { arg, err := expression.PBToExpr(child, fieldTps, sc) if err != nil { return nil, err } args = append(args, arg) } switch expr.Tp { case tipb.ExprType_Sum: return &sumFunction{aggFunction: newAggFunc(ast.AggFuncSum, args, false)}, nil case tipb.ExprType_Count: return &countFunction{aggFunction: newAggFunc(ast.AggFuncCount, args, false)}, nil case tipb.ExprType_Avg: return &avgFunction{aggFunction: newAggFunc(ast.AggFuncAvg, args, false)}, nil case tipb.ExprType_GroupConcat: return &concatFunction{aggFunction: newAggFunc(ast.AggFuncGroupConcat, args, false)}, nil case tipb.ExprType_Max: return &maxMinFunction{aggFunction: newAggFunc(ast.AggFuncMax, args, false), isMax: true}, nil case tipb.ExprType_Min: return &maxMinFunction{aggFunction: newAggFunc(ast.AggFuncMin, args, false)}, nil case tipb.ExprType_First: return &firstRowFunction{aggFunction: newAggFunc(ast.AggFuncFirstRow, args, false)}, nil case tipb.ExprType_Agg_BitOr: return &bitOrFunction{aggFunction: newAggFunc(ast.AggFuncBitOr, args, false)}, nil case tipb.ExprType_Agg_BitXor: return &bitXorFunction{aggFunction: newAggFunc(ast.AggFuncBitXor, args, false)}, nil case tipb.ExprType_Agg_BitAnd: return &bitAndFunction{aggFunction: newAggFunc(ast.AggFuncBitAnd, args, false)}, nil } return nil, errors.Errorf("Unknown aggregate function type %v", expr.Tp) } // AggEvaluateContext is used to store intermediate result when calculating aggregate functions. type AggEvaluateContext struct { DistinctChecker *distinctChecker Count int64 Value types.Datum Buffer *bytes.Buffer // Buffer is used for group_concat. GotFirstRow bool // It will check if the agg has met the first row key. } // AggFunctionMode stands for the aggregation function's mode. type AggFunctionMode int // |-----------------|--------------|--------------| // | AggFunctionMode | input | output | // |-----------------|--------------|--------------| // | CompleteMode | origin data | final result | // | FinalMode | partial data | final result | // | Partial1Mode | origin data | partial data | // | Partial2Mode | partial data | partial data | // | DedupMode | origin data | origin data | // |-----------------|--------------|--------------| const ( CompleteMode AggFunctionMode = iota FinalMode Partial1Mode Partial2Mode DedupMode ) type aggFunction struct { *AggFuncDesc } func newAggFunc(funcName string, args []expression.Expression, hasDistinct bool) aggFunction { agg := &AggFuncDesc{HasDistinct: hasDistinct} agg.Name = funcName agg.Args = args return aggFunction{AggFuncDesc: agg} } // CreateContext implements Aggregation interface. func (af *aggFunction) CreateContext(sc *stmtctx.StatementContext) *AggEvaluateContext { evalCtx := &AggEvaluateContext{} if af.HasDistinct { evalCtx.DistinctChecker = createDistinctChecker(sc) } return evalCtx } func (af *aggFunction) ResetContext(sc *stmtctx.StatementContext, evalCtx *AggEvaluateContext) { if af.HasDistinct { evalCtx.DistinctChecker = createDistinctChecker(sc) } evalCtx.Value.SetNull() } func (af *aggFunction) updateSum(sc *stmtctx.StatementContext, evalCtx *AggEvaluateContext, row chunk.Row) error { a := af.Args[0] value, err := a.Eval(row) if err != nil { return err } if value.IsNull() { return nil } if af.HasDistinct { d, err1 := evalCtx.DistinctChecker.Check([]types.Datum{value}) if err1 != nil { return err1 } if !d { return nil } } evalCtx.Value, err = calculateSum(sc, evalCtx.Value, value) if err != nil { return err } evalCtx.Count++ return nil } // NeedCount indicates whether the aggregate function should record count. func NeedCount(name string) bool { return name == ast.AggFuncCount || name == ast.AggFuncAvg } // NeedValue indicates whether the aggregate function should record value. func NeedValue(name string) bool { switch name { case ast.AggFuncSum, ast.AggFuncAvg, ast.AggFuncFirstRow, ast.AggFuncMax, ast.AggFuncMin, ast.AggFuncGroupConcat, ast.AggFuncBitOr, ast.AggFuncBitAnd, ast.AggFuncBitXor, ast.AggFuncApproxPercentile: return true default: return false } } // IsAllFirstRow checks whether functions in `aggFuncs` are all FirstRow. func IsAllFirstRow(aggFuncs []*AggFuncDesc) bool { for _, fun := range aggFuncs { if fun.Name != ast.AggFuncFirstRow { return false } } return true } // CheckAggPushDown checks whether an agg function can be pushed to storage. func CheckAggPushDown(aggFunc *AggFuncDesc, storeType kv.StoreType) bool { if len(aggFunc.OrderByItems) > 0 { return false } if aggFunc.Name == ast.AggFuncApproxPercentile { return false } ret := true switch storeType { case kv.TiFlash: ret = CheckAggPushFlash(aggFunc) } if ret { ret = expression.IsPushDownEnabled(strings.ToLower(aggFunc.Name), storeType) } return ret } // CheckAggPushFlash checks whether an agg function can be pushed to flash storage. func CheckAggPushFlash(aggFunc *AggFuncDesc) bool { switch aggFunc.Name { case ast.AggFuncSum, ast.AggFuncCount, ast.AggFuncMin, ast.AggFuncMax, ast.AggFuncAvg, ast.AggFuncFirstRow, ast.AggFuncApproxCountDistinct: return true } return false }