Skip to content

Commit c4d8025

Browse files
authored
expression: lazy init bufAllocator in baseBuiltinFunc (#65244)
ref #65003
1 parent 3056f41 commit c4d8025

11 files changed

+40
-8
lines changed

pkg/expression/bench_test.go

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1338,6 +1338,10 @@ func genVecExprBenchCase(ctx BuildContext, funcName string, testCase vecExprBenc
13381338
}
13391339

13401340
output = chunk.New([]*types.FieldType{eType2FieldType(expr.GetType(ctx.GetEvalCtx()).EvalType())}, testCase.chunkSize, testCase.chunkSize)
1341+
1342+
if !expr.Vectorized() {
1343+
panic(fmt.Sprintf("func %s is not vectorized", funcName))
1344+
}
13411345
return expr, fts, input, output
13421346
}
13431347

@@ -1352,6 +1356,7 @@ func testVectorizedEvalOneVec(t *testing.T, vecExprCases vecExprBenchCases) {
13521356
return fmt.Sprintf("func: %v, case %+v, row: %v, rowData: %v", funcName, testCase, row, input.GetRow(row).GetDatumRow(fts))
13531357
}
13541358
output2 := output.CopyConstruct()
1359+
require.True(t, expr.Vectorized(), "func %s is not vectorized", funcName)
13551360
require.NoErrorf(t, evalOneVec(ctx, expr, input, output, 0), "func: %v, case: %+v", funcName, testCase)
13561361
it := chunk.NewIterator4Chunk(input)
13571362
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)
14261431
exprName = tmp[len(tmp)-1]
14271432
}
14281433

1434+
if !expr.Vectorized() {
1435+
panic(fmt.Sprintf("func %s is not vectorized", funcName))
1436+
}
1437+
14291438
b.Run(exprName+"-EvalOneVec", func(b *testing.B) {
14301439
b.ResetTimer()
14311440
for i := 0; i < b.N; i++ {
@@ -1518,6 +1527,9 @@ func genVecBuiltinFuncBenchCase(ctx BuildContext, funcName string, testCase vecE
15181527
if err != nil {
15191528
panic(err)
15201529
}
1530+
if !baseFunc.vectorized() || !baseFunc.isChildrenVectorized() {
1531+
panic(fmt.Sprintf("func %s is not vectorized", funcName))
1532+
}
15211533
result = chunk.NewColumn(eType2FieldType(testCase.retEvalType), testCase.chunkSize)
15221534
// Mess up the output to make sure vecEvalXXX to call ResizeXXX/ReserveXXX itself.
15231535
result.AppendNull()
@@ -1596,7 +1608,7 @@ func testVectorizedBuiltinFunc(t *testing.T, vecExprCases vecExprBenchCases) {
15961608
continue
15971609
}
15981610
// do not forget to implement the vectorized method.
1599-
require.Truef(t, baseFunc.vectorized(), "func: %v, case: %+v", baseFuncName, testCase)
1611+
require.Truef(t, baseFunc.vectorized() && baseFunc.isChildrenVectorized(), "func: %v, case: %+v", baseFuncName, testCase)
16001612
commentf := func(row int) string {
16011613
return fmt.Sprintf("func: %v, case %+v, row: %v, rowData: %v", baseFuncName, testCase, row, input.GetRow(row).GetDatumRow(fts))
16021614
}
@@ -1767,7 +1779,7 @@ func testVectorizedBuiltinFuncForRand(t *testing.T, vecExprCases vecExprBenchCas
17671779
tmp := strings.Split(baseFuncName, ".")
17681780
baseFuncName = tmp[len(tmp)-1]
17691781
// do not forget to implement the vectorized method.
1770-
require.Truef(t, baseFunc.vectorized(), "func: %v", baseFuncName)
1782+
require.Truef(t, baseFunc.vectorized() && baseFunc.isChildrenVectorized(), "func: %v", baseFuncName)
17711783
switch testCase.retEvalType {
17721784
case types.ETReal:
17731785
err := baseFunc.vecEvalReal(ctx, input, output)
@@ -1834,6 +1846,9 @@ func benchmarkVectorizedBuiltinFunc(b *testing.B, vecExprCases vecExprBenchCases
18341846
baseFuncName := fmt.Sprintf("%v", reflect.TypeOf(baseFunc))
18351847
tmp := strings.Split(baseFuncName, ".")
18361848
baseFuncName = tmp[len(tmp)-1]
1849+
if !baseFunc.vectorized() || !baseFunc.isChildrenVectorized() {
1850+
panic(fmt.Sprintf("func %s is not vectorized", funcName))
1851+
}
18371852

18381853
if !testAll && !testFunc[baseFuncName] && !testFunc[funcName] {
18391854
continue

pkg/expression/builtin.go

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,6 @@ func newBaseBuiltinFunc(ctx BuildContext, funcName string, args []Expression, tp
133133
}
134134

135135
bf := baseBuiltinFunc{
136-
bufAllocator: newLocalColumnPool(),
137136
childrenVectorizedOnce: new(sync.Once),
138137

139138
args: args,
@@ -225,7 +224,6 @@ func newBaseBuiltinFuncWithTp(ctx BuildContext, funcName string, args []Expressi
225224

226225
fieldType := newReturnFieldTypeForBaseBuiltinFunc(funcName, retType, ec)
227226
bf = baseBuiltinFunc{
228-
bufAllocator: newLocalColumnPool(),
229227
childrenVectorizedOnce: new(sync.Once),
230228

231229
args: args,
@@ -286,7 +284,6 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args []
286284

287285
fieldType := newReturnFieldTypeForBaseBuiltinFunc(funcName, retType, ec)
288286
bf = baseBuiltinFunc{
289-
bufAllocator: newLocalColumnPool(),
290287
childrenVectorizedOnce: new(sync.Once),
291288

292289
args: args,
@@ -305,7 +302,6 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args []
305302
// do not check and compute collation.
306303
func newBaseBuiltinFuncWithFieldType(tp *types.FieldType, args []Expression) (baseBuiltinFunc, error) {
307304
bf := baseBuiltinFunc{
308-
bufAllocator: newLocalColumnPool(),
309305
childrenVectorizedOnce: new(sync.Once),
310306

311307
args: args,
@@ -397,6 +393,10 @@ func (b *baseBuiltinFunc) isChildrenVectorized() bool {
397393
break
398394
}
399395
}
396+
if b.childrenVectorized {
397+
// only init this when all children are vectorized
398+
b.bufAllocator = newLocalColumnPool()
399+
}
400400
})
401401
return b.childrenVectorized
402402
}
@@ -437,7 +437,6 @@ func (b *baseBuiltinFunc) cloneFrom(from *baseBuiltinFunc) {
437437
}
438438
b.tp = from.tp
439439
b.pbCode = from.pbCode
440-
b.bufAllocator = newLocalColumnPool()
441440
b.childrenVectorizedOnce = new(sync.Once)
442441
if from.ctor != nil {
443442
b.ctor = from.ctor.Clone()
@@ -481,7 +480,6 @@ func newBaseBuiltinCastFunc4String(ctx BuildContext, funcName string, args []Exp
481480
var err error
482481
if isExplicitCharset {
483482
bf = baseBuiltinFunc{
484-
bufAllocator: newLocalColumnPool(),
485483
childrenVectorizedOnce: new(sync.Once),
486484

487485
args: args,

pkg/expression/builtin_arithmetic_vec_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ func TestVectorizedDecimalErrOverflow(t *testing.T) {
229229
input.AppendMyDecimal(1, dec2)
230230
cols := []Expression{&Column{Index: 0, RetType: fts[0]}, &Column{Index: 1, RetType: fts[1]}}
231231
baseFunc, err := funcs[tt.funcName].getFunction(ctx, cols)
232+
require.True(t, baseFunc.vectorized() && baseFunc.isChildrenVectorized())
232233
require.NoError(t, err)
233234
result := chunk.NewColumn(eType2FieldType(types.ETDecimal), 1)
234235
err = vecEvalType(ctx, baseFunc, types.ETDecimal, input, result)

pkg/expression/builtin_cast_vec_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ func TestVectorizedCastRealAsTime(t *testing.T) {
164164
panic(err)
165165
}
166166
cast := &builtinCastRealAsTimeSig{baseFunc}
167+
require.True(t, cast.vectorized() && cast.isChildrenVectorized())
167168

168169
inputChunk, expect := genCastRealAsTime()
169170
inputs := []*chunk.Chunk{
@@ -256,6 +257,7 @@ func TestVectorizedCastStringAsDecimalWithUnsignedFlagInUnion(t *testing.T) {
256257
// set the `UnsignedFlag` bit
257258
baseCast.tp.AddFlag(mysql.UnsignedFlag)
258259
cast := &builtinCastStringAsDecimalSig{baseCast}
260+
require.True(t, cast.vectorized() && cast.isChildrenVectorized())
259261

260262
inputs := []*chunk.Chunk{
261263
genCastStringAsDecimal(false),

pkg/expression/builtin_miscellaneous_vec_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ func TestSleepVectorized(t *testing.T) {
147147
col0 := &Column{RetType: ft, Index: 0}
148148
f, err := fc.getFunction(ctx, []Expression{col0})
149149
require.NoError(t, err)
150+
require.True(t, f.vectorized() && f.isChildrenVectorized())
150151
input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 1024)
151152
result := chunk.NewColumn(ft, 1024)
152153
warnCnt := counter{}

pkg/expression/builtin_op_vec_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,7 @@ func TestBuiltinUnaryMinusIntSig(t *testing.T) {
169169
col0 := &Column{RetType: ft, Index: 0}
170170
f, err := funcs[ast.UnaryMinus].getFunction(ctx, []Expression{col0})
171171
require.NoError(t, err)
172+
require.True(t, f.vectorized() && f.isChildrenVectorized())
172173
input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 1024)
173174
result := chunk.NewColumn(ft, 1024)
174175

pkg/expression/builtin_other_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,7 @@ func TestInFunc(t *testing.T) {
339339
chk1 := chunk.NewChunkWithCapacity(nil, 1)
340340
chk1.SetNumVirtualRows(1)
341341
chk2 := chunk.NewChunkWithCapacity([]*types.FieldType{types.NewFieldType(mysql.TypeTiny)}, 1)
342+
require.True(t, fn.vectorized() && fn.isChildrenVectorized())
342343
err = vecEvalType(ctx, fn, types.ETInt, chk1, chk2.Column(0))
343344
require.NoError(t, err)
344345
require.Equal(t, int64(1), chk2.Column(0).GetInt64(0))

pkg/expression/builtin_other_vec_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ func TestInDecimal(t *testing.T) {
7070
col1 := &Column{RetType: ft, Index: 1}
7171
inFunc, err := funcs[ast.In].getFunction(ctx, []Expression{col0, col1})
7272
require.NoError(t, err)
73+
require.True(t, inFunc.vectorized() && inFunc.isChildrenVectorized())
7374

7475
input := chunk.NewChunkWithCapacity([]*types.FieldType{ft, ft}, 1024)
7576
for i := range 1024 {
@@ -101,6 +102,7 @@ func TestGetParamVec(t *testing.T) {
101102
col := &Column{RetType: ft, Index: 0}
102103
fn, err := funcs[ast.GetParam].getFunction(ctx, []Expression{col})
103104
require.NoError(t, err)
105+
require.True(t, fn.vectorized() && fn.isChildrenVectorized())
104106

105107
input := chunk.NewChunkWithCapacity([]*types.FieldType{ft}, 3)
106108
for i := range params {

pkg/expression/builtin_regexp_vec_const_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ func genVecBuiltinRegexpBenchCaseForConstants(ctx BuildContext) (baseFunc builti
6363
func TestVectorizedBuiltinRegexpForConstants(t *testing.T) {
6464
ctx := mock.NewContext()
6565
bf, childrenFieldTypes, input, output := genVecBuiltinRegexpBenchCaseForConstants(ctx)
66+
require.True(t, bf.vectorized() && bf.isChildrenVectorized())
6667
err := vecEvalType(ctx, bf, types.ETInt, input, output)
6768
require.NoError(t, err)
6869
i64s := output.Int64s()
@@ -87,6 +88,9 @@ func TestVectorizedBuiltinRegexpForConstants(t *testing.T) {
8788
func BenchmarkVectorizedBuiltinRegexpForConstants(b *testing.B) {
8889
ctx := mock.NewContext()
8990
bf, _, input, output := genVecBuiltinRegexpBenchCaseForConstants(ctx)
91+
if !bf.vectorized() || !bf.isChildrenVectorized() {
92+
panic("builtinRegexpUTF8Sig is not vectorized")
93+
}
9094
b.Run("builtinRegexpUTF8Sig-Constants-VecBuiltinFunc", func(b *testing.B) {
9195
b.ResetTimer()
9296
for i := 0; i < b.N; i++ {

pkg/expression/builtin_time_vec_test.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -588,6 +588,7 @@ func TestVecMonth(t *testing.T) {
588588
input.AppendTime(0, types.ZeroDate)
589589

590590
f, _, _, result := genVecBuiltinFuncBenchCase(ctx, ast.Month, vecExprBenchCase{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETDatetime}})
591+
require.True(t, f.vectorized() && f.isChildrenVectorized())
591592
require.True(t, ctx.GetSessionVars().SQLMode.HasStrictMode())
592593
require.NoError(t, vecEvalType(ctx, f, types.ETInt, input, result))
593594
require.Equal(t, 0, len(ctx.GetSessionVars().StmtCtx.GetWarnings()))

0 commit comments

Comments
 (0)