expression: lazy init bufAllocator in baseBuiltinFunc

pkg/expression/bench_test.go

@@ -1338,6 +1338,10 @@ func genVecExprBenchCase(ctx BuildContext, funcName string, testCase vecExprBenc
 	}
 
 	output = chunk.New([]*types.FieldType{eType2FieldType(expr.GetType(ctx.GetEvalCtx()).EvalType())}, testCase.chunkSize, testCase.chunkSize)
+
+	if !expr.Vectorized() {
+		panic(fmt.Sprintf("func %s is not vectorized", funcName))
+	}
 	return expr, fts, input, output
 }
 
@@ -1352,6 +1356,7 @@ func testVectorizedEvalOneVec(t *testing.T, vecExprCases vecExprBenchCases) {
 				return fmt.Sprintf("func: %v, case %+v, row: %v, rowData: %v", funcName, testCase, row, input.GetRow(row).GetDatumRow(fts))
 			}
 			output2 := output.CopyConstruct()
+			require.True(t, expr.Vectorized(), "func %s is not vectorized", funcName)
 			require.NoErrorf(t, evalOneVec(ctx, expr, input, output, 0), "func: %v, case: %+v", funcName, testCase)
 			it := chunk.NewIterator4Chunk(input)
 			require.NoErrorf(t, evalOneColumn(ctx, expr, it, output2, 0), "func: %v, case: %+v", funcName, testCase)
@@ -1426,6 +1431,10 @@ func benchmarkVectorizedEvalOneVec(b *testing.B, vecExprCases vecExprBenchCases)
 				exprName = tmp[len(tmp)-1]
 			}
 
+			if !expr.Vectorized() {
+				panic(fmt.Sprintf("func %s is not vectorized", funcName))
+			}
+
 			b.Run(exprName+"-EvalOneVec", func(b *testing.B) {
 				b.ResetTimer()
 				for i := 0; i < b.N; i++ {
@@ -1518,6 +1527,9 @@ func genVecBuiltinFuncBenchCase(ctx BuildContext, funcName string, testCase vecE
 	if err != nil {
 		panic(err)
 	}
+	if !baseFunc.vectorized() || !baseFunc.isChildrenVectorized() {
+		panic(fmt.Sprintf("func %s is not vectorized", funcName))
+	}
 	result = chunk.NewColumn(eType2FieldType(testCase.retEvalType), testCase.chunkSize)
 	// Mess up the output to make sure vecEvalXXX to call ResizeXXX/ReserveXXX itself.
 	result.AppendNull()
@@ -1596,7 +1608,7 @@ func testVectorizedBuiltinFunc(t *testing.T, vecExprCases vecExprBenchCases) {
 				continue
 			}
 			// do not forget to implement the vectorized method.
-			require.Truef(t, baseFunc.vectorized(), "func: %v, case: %+v", baseFuncName, testCase)
+			require.Truef(t, baseFunc.vectorized() && baseFunc.isChildrenVectorized(), "func: %v, case: %+v", baseFuncName, testCase)
 			commentf := func(row int) string {
 				return fmt.Sprintf("func: %v, case %+v, row: %v, rowData: %v", baseFuncName, testCase, row, input.GetRow(row).GetDatumRow(fts))
 			}
@@ -1767,7 +1779,7 @@ func testVectorizedBuiltinFuncForRand(t *testing.T, vecExprCases vecExprBenchCas
 			tmp := strings.Split(baseFuncName, ".")
 			baseFuncName = tmp[len(tmp)-1]
 			// do not forget to implement the vectorized method.
-			require.Truef(t, baseFunc.vectorized(), "func: %v", baseFuncName)
+			require.Truef(t, baseFunc.vectorized() && baseFunc.isChildrenVectorized(), "func: %v", baseFuncName)
 			switch testCase.retEvalType {
 			case types.ETReal:
 				err := baseFunc.vecEvalReal(ctx, input, output)
@@ -1834,6 +1846,9 @@ func benchmarkVectorizedBuiltinFunc(b *testing.B, vecExprCases vecExprBenchCases
 			baseFuncName := fmt.Sprintf("%v", reflect.TypeOf(baseFunc))
 			tmp := strings.Split(baseFuncName, ".")
 			baseFuncName = tmp[len(tmp)-1]
+			if !baseFunc.vectorized() || !baseFunc.isChildrenVectorized() {
+				panic(fmt.Sprintf("func %s is not vectorized", funcName))
+			}
 
 			if !testAll && !testFunc[baseFuncName] && !testFunc[funcName] {
 				continue

pkg/expression/builtin.go

@@ -133,7 +133,6 @@ func newBaseBuiltinFunc(ctx BuildContext, funcName string, args []Expression, tp
 	}
 
 	bf := baseBuiltinFunc{
-		bufAllocator:           newLocalColumnPool(),
 		childrenVectorizedOnce: new(sync.Once),
 
 		args: args,
@@ -225,7 +224,6 @@ func newBaseBuiltinFuncWithTp(ctx BuildContext, funcName string, args []Expressi
 
 	fieldType := newReturnFieldTypeForBaseBuiltinFunc(funcName, retType, ec)
 	bf = baseBuiltinFunc{
-		bufAllocator:           newLocalColumnPool(),
 		childrenVectorizedOnce: new(sync.Once),
 
 		args: args,
@@ -286,7 +284,6 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args []
 
 	fieldType := newReturnFieldTypeForBaseBuiltinFunc(funcName, retType, ec)
 	bf = baseBuiltinFunc{
-		bufAllocator:           newLocalColumnPool(),
 		childrenVectorizedOnce: new(sync.Once),
 
 		args: args,
@@ -305,7 +302,6 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args []
 // do not check and compute collation.
 func newBaseBuiltinFuncWithFieldType(tp *types.FieldType, args []Expression) (baseBuiltinFunc, error) {
 	bf := baseBuiltinFunc{
-		bufAllocator:           newLocalColumnPool(),
 		childrenVectorizedOnce: new(sync.Once),
 
 		args: args,
@@ -397,6 +393,10 @@ func (b *baseBuiltinFunc) isChildrenVectorized() bool {
 				break
 			}
 		}
+		if b.childrenVectorized {
+			// only init this when all children are vectorized
+			b.bufAllocator = newLocalColumnPool()
+		}
 	})
 	return b.childrenVectorized
 }
@@ -437,7 +437,6 @@ func (b *baseBuiltinFunc) cloneFrom(from *baseBuiltinFunc) {
 	}
 	b.tp = from.tp
 	b.pbCode = from.pbCode
-	b.bufAllocator = newLocalColumnPool()
 	b.childrenVectorizedOnce = new(sync.Once)
 	if from.ctor != nil {
 		b.ctor = from.ctor.Clone()
@@ -481,7 +480,6 @@ func newBaseBuiltinCastFunc4String(ctx BuildContext, funcName string, args []Exp
 	var err error
 	if isExplicitCharset {
 		bf = baseBuiltinFunc{
-			bufAllocator:           newLocalColumnPool(),
 			childrenVectorizedOnce: new(sync.Once),
 
 			args: args,

pkg/expression/builtin_arithmetic_vec_test.go

@@ -229,6 +229,7 @@ func TestVectorizedDecimalErrOverflow(t *testing.T) {
 		input.AppendMyDecimal(1, dec2)
 		cols := []Expression{&Column{Index: 0, RetType: fts[0]}, &Column{Index: 1, RetType: fts[1]}}
 		baseFunc, err := funcs[tt.funcName].getFunction(ctx, cols)
+		require.True(t, baseFunc.vectorized() && baseFunc.isChildrenVectorized())
 		require.NoError(t, err)
 		result := chunk.NewColumn(eType2FieldType(types.ETDecimal), 1)
 		err = vecEvalType(ctx, baseFunc, types.ETDecimal, input, result)

pkg/expression/builtin_cast_vec_test.go

@@ -164,6 +164,7 @@ func TestVectorizedCastRealAsTime(t *testing.T) {
 		panic(err)
 	}
 	cast := &builtinCastRealAsTimeSig{baseFunc}
+	require.True(t, cast.vectorized() && cast.isChildrenVectorized())
 
 	inputChunk, expect := genCastRealAsTime()
 	inputs := []*chunk.Chunk{
@@ -256,6 +257,7 @@ func TestVectorizedCastStringAsDecimalWithUnsignedFlagInUnion(t *testing.T) {
 	// set the `UnsignedFlag` bit
 	baseCast.tp.AddFlag(mysql.UnsignedFlag)
 	cast := &builtinCastStringAsDecimalSig{baseCast}
+	require.True(t, cast.vectorized() && cast.isChildrenVectorized())
 
 	inputs := []*chunk.Chunk{
 		genCastStringAsDecimal(false),

pkg/expression/builtin_miscellaneous_vec_test.go

@@ -147,6 +147,7 @@ func TestSleepVectorized(t *testing.T) {
 	col0 := &Column{RetType: ft, Index: 0}
 	f, err := fc.getFunction(ctx, []Expression{col0})
 	require.NoError(t, err)
+	require.True(t, f.vectorized() && f.isChildrenVectorized())
 	input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 1024)
 	result := chunk.NewColumn(ft, 1024)
 	warnCnt := counter{}

pkg/expression/builtin_op_vec_test.go

@@ -169,6 +169,7 @@ func TestBuiltinUnaryMinusIntSig(t *testing.T) {
 	col0 := &Column{RetType: ft, Index: 0}
 	f, err := funcs[ast.UnaryMinus].getFunction(ctx, []Expression{col0})
 	require.NoError(t, err)
+	require.True(t, f.vectorized() && f.isChildrenVectorized())
 	input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 1024)
 	result := chunk.NewColumn(ft, 1024)
 

pkg/expression/builtin_other_test.go

@@ -339,6 +339,7 @@ func TestInFunc(t *testing.T) {
 	chk1 := chunk.NewChunkWithCapacity(nil, 1)
 	chk1.SetNumVirtualRows(1)
 	chk2 := chunk.NewChunkWithCapacity([]*types.FieldType{types.NewFieldType(mysql.TypeTiny)}, 1)
+	require.True(t, fn.vectorized() && fn.isChildrenVectorized())
 	err = vecEvalType(ctx, fn, types.ETInt, chk1, chk2.Column(0))
 	require.NoError(t, err)
 	require.Equal(t, int64(1), chk2.Column(0).GetInt64(0))

pkg/expression/builtin_other_vec_test.go

@@ -70,6 +70,7 @@ func TestInDecimal(t *testing.T) {
 	col1 := &Column{RetType: ft, Index: 1}
 	inFunc, err := funcs[ast.In].getFunction(ctx, []Expression{col0, col1})
 	require.NoError(t, err)
+	require.True(t, inFunc.vectorized() && inFunc.isChildrenVectorized())
 
 	input := chunk.NewChunkWithCapacity([]*types.FieldType{ft, ft}, 1024)
 	for i := range 1024 {
@@ -101,6 +102,7 @@ func TestGetParamVec(t *testing.T) {
 	col := &Column{RetType: ft, Index: 0}
 	fn, err := funcs[ast.GetParam].getFunction(ctx, []Expression{col})
 	require.NoError(t, err)
+	require.True(t, fn.vectorized() && fn.isChildrenVectorized())
 
 	input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 3)
 	for i := range params {

pkg/expression/builtin_regexp_vec_const_test.go

@@ -63,6 +63,7 @@ func genVecBuiltinRegexpBenchCaseForConstants(ctx BuildContext) (baseFunc builti
 func TestVectorizedBuiltinRegexpForConstants(t *testing.T) {
 	ctx := mock.NewContext()
 	bf, childrenFieldTypes, input, output := genVecBuiltinRegexpBenchCaseForConstants(ctx)
+	require.True(t, bf.vectorized() && bf.isChildrenVectorized())
 	err := vecEvalType(ctx, bf, types.ETInt, input, output)
 	require.NoError(t, err)
 	i64s := output.Int64s()
@@ -87,6 +88,9 @@ func TestVectorizedBuiltinRegexpForConstants(t *testing.T) {
 func BenchmarkVectorizedBuiltinRegexpForConstants(b *testing.B) {
 	ctx := mock.NewContext()
 	bf, _, input, output := genVecBuiltinRegexpBenchCaseForConstants(ctx)
+	if !bf.vectorized() || !bf.isChildrenVectorized() {
+		panic("builtinRegexpUTF8Sig is not vectorized")
+	}
 	b.Run("builtinRegexpUTF8Sig-Constants-VecBuiltinFunc", func(b *testing.B) {
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {

pkg/expression/builtin_time_vec_test.go

@@ -588,6 +588,7 @@ func TestVecMonth(t *testing.T) {
 	input.AppendTime(0, types.ZeroDate)
 
 	f, _, _, result := genVecBuiltinFuncBenchCase(ctx, ast.Month, vecExprBenchCase{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETDatetime}})
+	require.True(t, f.vectorized() && f.isChildrenVectorized())
 	require.True(t, ctx.GetSessionVars().SQLMode.HasStrictMode())
 	require.NoError(t, vecEvalType(ctx, f, types.ETInt, input, result))
 	require.Equal(t, 0, len(ctx.GetSessionVars().StmtCtx.GetWarnings()))

pkg/expression/builtin_vectorized_test.go

@@ -95,6 +95,9 @@ func genMockVecPlusIntBuiltinFunc(ctx BuildContext) (*mockVecPlusIntBuiltinFunc,
 		input.AppendInt64(0, int64(i))
 		input.AppendInt64(1, int64(i))
 	}
+	if !plus.isChildrenVectorized() {
+		panic("mockVecPlusIntBuiltinFunc's children should be vectorized")
+	}
 	return plus, input, buf
 }
 
@@ -469,6 +472,9 @@ func genMockRowDouble(ctx BuildContext, eType types.EvalType, enableVec bool) (b
 			input.AppendTime(0, types.NewTime(t, mysqlType, 0))
 		}
 	}
+	if !rowDouble.isChildrenVectorized() {
+		return nil, nil, nil, errors.New("mockBuiltinDouble's children should be vectorized")
+	}
 	return rowDouble, input, buf, nil
 }