expression: lazy init bufAllocator in baseBuiltinFunc (#65244)

ref pingcap/tidb#65003
This commit is contained in:
xufei
2025-12-27 09:09:26 +08:00
committed by GitHub
parent 3056f41c48
commit c4d80255cf
11 changed files with 40 additions and 8 deletions

View File

@ -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

View File

@ -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,

View File

@ -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)

View File

@ -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),

View File

@ -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{}

View File

@ -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)

View File

@ -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))

View File

@ -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 {

View File

@ -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++ {

View File

@ -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()))

View File

@ -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
}