expression: implement vectorized evaluation for builtinBenchmarkSig (…

…pingcap#13643)

expression/builtin_info.go

@@ -455,17 +455,27 @@ func (c *benchmarkFunctionClass) getFunction(ctx sessionctx.Context, args []Expr
 	// Syntax: BENCHMARK(loop_count, expression)
 	// Define with same eval type of input arg to avoid unnecessary cast function.
 	sameEvalType := args[1].GetType().EvalType()
+	// constLoopCount is used by VecEvalInt
+	// since non-constant loop count would be different between rows, and cannot be vectorized.
+	var constLoopCount int64
+	con, ok := args[0].(*Constant)
+	if ok && con.Value.Kind() == types.KindInt64 {
+		if lc, isNull, err := con.EvalInt(ctx, chunk.Row{}); err == nil && !isNull {
+			constLoopCount = lc
+		}
+	}
 	bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETInt, types.ETInt, sameEvalType)
-	sig := &builtinBenchmarkSig{bf}
+	sig := &builtinBenchmarkSig{bf, constLoopCount}
 	return sig, nil
 }
 
 type builtinBenchmarkSig struct {
 	baseBuiltinFunc
+	constLoopCount int64
 }
 
 func (b *builtinBenchmarkSig) Clone() builtinFunc {
-	newSig := &builtinBenchmarkSig{}
+	newSig := &builtinBenchmarkSig{constLoopCount: b.constLoopCount}
 	newSig.cloneFrom(&b.baseBuiltinFunc)
 	return newSig
 }
@@ -474,9 +484,16 @@ func (b *builtinBenchmarkSig) Clone() builtinFunc {
 // See https://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_benchmark
 func (b *builtinBenchmarkSig) evalInt(row chunk.Row) (int64, bool, error) {
 	// Get loop count.
-	loopCount, isNull, err := b.args[0].EvalInt(b.ctx, row)
-	if isNull || err != nil {
-		return 0, isNull, err
+	var loopCount int64
+	var isNull bool
+	var err error
+	if b.constLoopCount > 0 {
+		loopCount = b.constLoopCount
+	} else {
+		loopCount, isNull, err = b.args[0].EvalInt(b.ctx, row)
+		if isNull || err != nil {
+			return 0, isNull, err
+		}
 	}
 
 	// BENCHMARK() will return NULL if loop count < 0,

expression/builtin_info_test.go

@@ -167,6 +167,10 @@ func (s *testEvaluatorSuite) TestBenchMark(c *C) {
 		} else {
 			c.Assert(d.GetInt64(), Equals, t.Expected)
 		}
+
+		// test clone
+		b1 := f.Clone().(*ScalarFunction).Function.(*builtinBenchmarkSig)
+		c.Assert(b1.constLoopCount, Equals, int64(t.LoopCount))
 	}
 }
 

expression/builtin_info_vec.go

@@ -210,11 +210,73 @@ func (b *builtinFoundRowsSig) vecEvalInt(input *chunk.Chunk, result *chunk.Colum
 }
 
 func (b *builtinBenchmarkSig) vectorized() bool {
-	return false
+	return b.constLoopCount > 0
 }
 
 func (b *builtinBenchmarkSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
-	return errors.Errorf("not implemented")
+	n := input.NumRows()
+	loopCount := b.constLoopCount
+	arg, ctx := b.args[1], b.ctx
+	evalType := arg.GetType().EvalType()
+	buf, err := b.bufAllocator.get(evalType, n)
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(buf)
+
+	var k int64
+	switch evalType {
+	case types.ETInt:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalInt(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	case types.ETReal:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalReal(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	case types.ETDecimal:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalDecimal(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	case types.ETString:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalString(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	case types.ETDatetime, types.ETTimestamp:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalTime(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	case types.ETDuration:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalDuration(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	case types.ETJson:
+		for ; k < loopCount; k++ {
+			if err = arg.VecEvalJSON(ctx, input, buf); err != nil {
+				return err
+			}
+		}
+	default: // Should never go into here.
+		return errors.Errorf("EvalType %v not implemented for builtin BENCHMARK()", evalType)
+	}
+
+	// Return value of BENCHMARK() is always 0.
+	// even if args[1].IsNull(i)
+	result.ResizeInt64(n, false)
+
+	return nil
 }
 
 func (b *builtinLastInsertIDSig) vectorized() bool {

expression/builtin_info_vec_test.go

@@ -20,6 +20,7 @@ import (
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/parser/ast"
+	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
 )
@@ -90,6 +91,24 @@ var vecBuiltinInfoCases = map[string][]vecExprBenchCase{
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{}},
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt}},
 	},
+	ast.Benchmark: {
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETInt},
+			constants: []*Constant{{Value: types.NewIntDatum(10), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETReal},
+			constants: []*Constant{{Value: types.NewIntDatum(11), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETDecimal},
+			constants: []*Constant{{Value: types.NewIntDatum(12), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETString},
+			constants: []*Constant{{Value: types.NewIntDatum(13), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETDatetime},
+			constants: []*Constant{{Value: types.NewIntDatum(14), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETTimestamp},
+			constants: []*Constant{{Value: types.NewIntDatum(15), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETDuration},
+			constants: []*Constant{{Value: types.NewIntDatum(16), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETInt, types.ETJson},
+			constants: []*Constant{{Value: types.NewIntDatum(17), RetType: types.NewFieldType(mysql.TypeLonglong)}, nil}},
+	},
 }
 
 func (s *testEvaluatorSuite) TestVectorizedBuiltinInfoFunc(c *C) {