expression, planner: support builtin function benchmark (#9252)

expression/builtin_info.go

@@ -384,7 +384,104 @@ type benchmarkFunctionClass struct {
 }
 
 func (c *benchmarkFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
-	return nil, errFunctionNotExists.GenWithStackByArgs("FUNCTION", "BENCHMARK")
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+
+	// Syntax: BENCHMARK(loop_count, expression)
+	// Define with same eval type of input arg to avoid unnecessary cast function.
+	sameEvalType := args[1].GetType().EvalType()
+	bf := newBaseBuiltinFuncWithTp(ctx, args, types.ETInt, types.ETInt, sameEvalType)
+	sig := &builtinBenchmarkSig{bf}
+	return sig, nil
+}
+
+type builtinBenchmarkSig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinBenchmarkSig) Clone() builtinFunc {
+	newSig := &builtinBenchmarkSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+// evalInt evals a builtinBenchmarkSig. It will execute expression repeatedly count times.
+// 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
+	}
+
+	// BENCHMARK() will return NULL if loop count < 0,
+	// behavior observed on MySQL 5.7.24.
+	if loopCount < 0 {
+		return 0, true, nil
+	}
+
+	// Eval loop count times based on arg type.
+	// BENCHMARK() will pass-through the eval error,
+	// behavior observed on MySQL 5.7.24.
+	var i int64
+	arg, ctx := b.args[1], b.ctx
+	switch evalType := arg.GetType().EvalType(); evalType {
+	case types.ETInt:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalInt(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	case types.ETReal:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalReal(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	case types.ETDecimal:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalDecimal(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	case types.ETString:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalString(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	case types.ETDatetime, types.ETTimestamp:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalTime(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	case types.ETDuration:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalDuration(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	case types.ETJson:
+		for ; i < loopCount; i++ {
+			_, isNull, err = arg.EvalJSON(ctx, row)
+			if err != nil {
+				return 0, isNull, err
+			}
+		}
+	default: // Should never go into here.
+		return 0, true, errors.Errorf("EvalType %v not implemented for builtin BENCHMARK()", evalType)
+	}
+
+	// Return value of BENCHMARK() is always 0.
+	return 0, false, nil
 }
 
 type charsetFunctionClass struct {

expression/builtin_info_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/parser/charset"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/mock"
 	"github.com/pingcap/tidb/util/printer"
@@ -120,10 +121,40 @@ func (s *testEvaluatorSuite) TestVersion(c *C) {
 
 func (s *testEvaluatorSuite) TestBenchMark(c *C) {
 	defer testleak.AfterTest(c)()
-	fc := funcs[ast.Benchmark]
-	f, err := fc.getFunction(s.ctx, s.datumsToConstants(types.MakeDatums(nil, nil)))
-	c.Assert(f, IsNil)
-	c.Assert(err, ErrorMatches, "*FUNCTION BENCHMARK does not exist")
+
+	cases := []struct {
+		LoopCount  int
+		Expression interface{}
+		Expected   int64
+		IsNil      bool
+	}{
+		{-3, 1, 0, true},
+		{0, 1, 0, false},
+		{3, 1, 0, false},
+		{3, 1.234, 0, false},
+		{3, types.NewDecFromFloatForTest(1.234), 0, false},
+		{3, "abc", 0, false},
+		{3, types.CurrentTime(mysql.TypeDatetime), 0, false},
+		{3, types.CurrentTime(mysql.TypeTimestamp), 0, false},
+		{3, types.CurrentTime(mysql.TypeDuration), 0, false},
+		{3, json.CreateBinary("[1]"), 0, false},
+	}
+
+	for _, t := range cases {
+		f, err := newFunctionForTest(s.ctx, ast.Benchmark, s.primitiveValsToConstants([]interface{}{
+			t.LoopCount,
+			t.Expression,
+		})...)
+		c.Assert(err, IsNil)
+
+		d, err := f.Eval(chunk.Row{})
+		c.Assert(err, IsNil)
+		if t.IsNil {
+			c.Assert(d.IsNull(), IsTrue)
+		} else {
+			c.Assert(d.GetInt64(), Equals, t.Expected)
+		}
+	}
 }
 
 func (s *testEvaluatorSuite) TestCharset(c *C) {

expression/function_traits.go

@@ -40,6 +40,14 @@ var unFoldableFunctions = map[string]struct{}{
 	ast.SetVar:    {},
 	ast.GetVar:    {},
 	ast.GetParam:  {},
+	ast.Benchmark: {},
+}
+
+// DisableFoldFunctions stores functions which prevent child scope functions from being constant folded.
+// Typically, these functions shall also exist in unFoldableFunctions, to stop from being folded when they themselves
+// are in child scope of an outer function, and the outer function is recursively folding its children.
+var DisableFoldFunctions = map[string]struct{}{
+	ast.Benchmark: {},
 }
 
 // DeferredFunctions stores non-deterministic functions, which can be deferred only when the plan cache is enabled.

expression/integration_test.go

@@ -2489,6 +2489,42 @@ func (s *testIntegrationSuite) TestInfoBuiltin(c *C) {
 	result.Check(testkit.Rows("1"))
 	result = tk.MustQuery("select row_count();")
 	result.Check(testkit.Rows("-1"))
+
+	// for benchmark
+	success := testkit.Rows("0")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (a int, b int)")
+	result = tk.MustQuery(`select benchmark(3, benchmark(2, length("abc")))`)
+	result.Check(success)
+	err := tk.ExecToErr(`select benchmark(3, length("a", "b"))`)
+	c.Assert(err, NotNil)
+	// Quoted from https://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_benchmark
+	// Although the expression can be a subquery, it must return a single column and at most a single row.
+	// For example, BENCHMARK(10, (SELECT * FROM t)) will fail if the table t has more than one column or
+	// more than one row.
+	oneColumnQuery := "select benchmark(10, (select a from t))"
+	twoColumnQuery := "select benchmark(10, (select * from t))"
+	// rows * columns:
+	// 0 * 1, success;
+	result = tk.MustQuery(oneColumnQuery)
+	result.Check(success)
+	// 0 * 2, error;
+	err = tk.ExecToErr(twoColumnQuery)
+	c.Assert(err, NotNil)
+	// 1 * 1, success;
+	tk.MustExec("insert t values (1, 2)")
+	result = tk.MustQuery(oneColumnQuery)
+	result.Check(success)
+	// 1 * 2, error;
+	err = tk.ExecToErr(twoColumnQuery)
+	c.Assert(err, NotNil)
+	// 2 * 1, error;
+	tk.MustExec("insert t values (3, 4)")
+	err = tk.ExecToErr(oneColumnQuery)
+	c.Assert(err, NotNil)
+	// 2 * 2, error.
+	err = tk.ExecToErr(twoColumnQuery)
+	c.Assert(err, NotNil)
 }
 
 func (s *testIntegrationSuite) TestControlBuiltin(c *C) {