Skip to content

Commit

Permalink
executor: support new aggregate function APPROX_PERCENTILE (#19799) (#…
Browse files Browse the repository at this point in the history
…20197)

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
  • Loading branch information
ti-srebot authored Oct 19, 2020
1 parent 233d278 commit cae843b
Show file tree
Hide file tree
Showing 18 changed files with 1,013 additions and 20 deletions.
3 changes: 3 additions & 0 deletions executor/aggfuncs/aggfunc_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,9 @@ func (s *testSuite) testAggFunc(c *C, p aggTest) {
if p.funcName == ast.AggFuncGroupConcat {
args = append(args, &expression.Constant{Value: types.NewStringDatum(" "), RetType: types.NewFieldType(mysql.TypeString)})
}
if p.funcName == ast.AggFuncApproxPercentile {
args = append(args, &expression.Constant{Value: types.NewIntDatum(50), RetType: types.NewFieldType(mysql.TypeLong)})
}
desc, err := aggregation.NewAggFuncDesc(s.ctx, p.funcName, args, false)
c.Assert(err, IsNil)
if p.orderBy {
Expand Down
5 changes: 5 additions & 0 deletions executor/aggfuncs/aggfuncs.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,11 @@ var (
_ AggFunc = (*approxCountDistinctPartial2)(nil)
_ AggFunc = (*approxCountDistinctFinal)(nil)

// All the AggFunc implementations for "APPROX_PERCENTILE" are listed here.
_ AggFunc = (*percentileOriginal4Int)(nil)
_ AggFunc = (*percentileOriginal4Real)(nil)
_ AggFunc = (*percentileOriginal4Decimal)(nil)

// All the AggFunc implementations for "FIRSTROW" are listed here.
_ AggFunc = (*firstRow4Decimal)(nil)
_ AggFunc = (*firstRow4Int)(nil)
Expand Down
45 changes: 43 additions & 2 deletions executor/aggfuncs/builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ import (
"github.com/pingcap/tidb/sessionctx/variable"
"github.com/pingcap/tidb/types"
"github.com/pingcap/tidb/util/chunk"
"github.com/pingcap/tidb/util/logutil"
"go.uber.org/zap"
)

// Build is used to build a specific AggFunc implementation according to the
Expand Down Expand Up @@ -53,12 +55,14 @@ func Build(ctx sessionctx.Context, aggFuncDesc *aggregation.AggFuncDesc, ordinal
return buildBitAnd(aggFuncDesc, ordinal)
case ast.AggFuncVarPop:
return buildVarPop(aggFuncDesc, ordinal)
case ast.AggFuncStddevPop:
return buildStdDevPop(aggFuncDesc, ordinal)
case ast.AggFuncJsonObjectAgg:
return buildJSONObjectAgg(aggFuncDesc, ordinal)
case ast.AggFuncApproxCountDistinct:
return buildApproxCountDistinct(aggFuncDesc, ordinal)
case ast.AggFuncStddevPop:
return buildStdDevPop(aggFuncDesc, ordinal)
case ast.AggFuncApproxPercentile:
return buildApproxPercentile(ctx, aggFuncDesc, ordinal)
case ast.AggFuncVarSamp:
return buildVarSamp(aggFuncDesc, ordinal)
case ast.AggFuncStddevSamp:
Expand Down Expand Up @@ -130,6 +134,43 @@ func buildApproxCountDistinct(aggFuncDesc *aggregation.AggFuncDesc, ordinal int)
return nil
}

func buildApproxPercentile(sctx sessionctx.Context, aggFuncDesc *aggregation.AggFuncDesc, ordinal int) AggFunc {
if aggFuncDesc.Mode == aggregation.DedupMode {
return nil
}

// Checked while building descriptor
percent, _, err := aggFuncDesc.Args[1].EvalInt(sctx, chunk.Row{})
if err != nil {
// Should not reach here
logutil.BgLogger().Error("Error happened when buildApproxPercentile", zap.Error(err))
return nil
}

base := basePercentile{percent: int(percent), baseAggFunc: baseAggFunc{args: aggFuncDesc.Args, ordinal: ordinal}}

switch aggFuncDesc.Mode {
case aggregation.CompleteMode, aggregation.Partial1Mode, aggregation.FinalMode:
switch aggFuncDesc.Args[0].GetType().EvalType() {
case types.ETInt:
return &percentileOriginal4Int{base}
case types.ETReal:
return &percentileOriginal4Real{base}
case types.ETDecimal:
return &percentileOriginal4Decimal{base}
case types.ETDatetime, types.ETTimestamp:
return &percentileOriginal4Time{base}
case types.ETDuration:
return &percentileOriginal4Duration{base}
default:
// Return NULL in any case
return &base
}
}

return nil
}

// buildCount builds the AggFunc implementation for function "COUNT".
func buildCount(aggFuncDesc *aggregation.AggFuncDesc, ordinal int) AggFunc {
// If mode is DedupMode, we return nil for not implemented.
Expand Down
Loading

0 comments on commit cae843b

Please sign in to comment.