301 lines
7.7 KiB
Go
301 lines
7.7 KiB
Go
// Copyright 2023 PingCAP, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package testutil
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/base64"
|
|
"fmt"
|
|
"math/big"
|
|
"sort"
|
|
"sync/atomic"
|
|
|
|
"github.com/pingcap/tidb/pkg/executor/internal/exec"
|
|
"github.com/pingcap/tidb/pkg/expression"
|
|
"github.com/pingcap/tidb/pkg/parser/mysql"
|
|
plannercore "github.com/pingcap/tidb/pkg/planner/core"
|
|
"github.com/pingcap/tidb/pkg/planner/property"
|
|
"github.com/pingcap/tidb/pkg/sessionctx"
|
|
"github.com/pingcap/tidb/pkg/types"
|
|
"github.com/pingcap/tidb/pkg/util/chunk"
|
|
"github.com/pingcap/tidb/pkg/util/memory"
|
|
"github.com/pingcap/tidb/pkg/util/stringutil"
|
|
)
|
|
|
|
// MockDataSourceParameters mpcks data source parameters
|
|
type MockDataSourceParameters struct {
|
|
Ctx sessionctx.Context
|
|
DataSchema *expression.Schema
|
|
GenDataFunc func(row int, typ *types.FieldType) any
|
|
Ndvs []int
|
|
Orders []bool
|
|
Rows int
|
|
}
|
|
|
|
// MockDataSource mocks data source
|
|
type MockDataSource struct {
|
|
GenData []*chunk.Chunk
|
|
Chunks []*chunk.Chunk
|
|
P MockDataSourceParameters
|
|
exec.BaseExecutor
|
|
ChunkPtr int
|
|
}
|
|
|
|
// GenColDatums get column datums
|
|
func (mds *MockDataSource) GenColDatums(col int) (results []any) {
|
|
typ := mds.RetFieldTypes()[col]
|
|
order := false
|
|
if col < len(mds.P.Orders) {
|
|
order = mds.P.Orders[col]
|
|
}
|
|
rows := mds.P.Rows
|
|
ndv := 0
|
|
if col < len(mds.P.Ndvs) {
|
|
ndv = mds.P.Ndvs[col]
|
|
}
|
|
results = make([]any, 0, rows)
|
|
if ndv == 0 {
|
|
if mds.P.GenDataFunc == nil {
|
|
for i := 0; i < rows; i++ {
|
|
results = append(results, mds.RandDatum(typ))
|
|
}
|
|
} else {
|
|
for i := 0; i < rows; i++ {
|
|
results = append(results, mds.P.GenDataFunc(i, typ))
|
|
}
|
|
}
|
|
} else {
|
|
datumSet := make(map[string]bool, ndv)
|
|
datums := make([]any, 0, ndv)
|
|
for len(datums) < ndv {
|
|
d := mds.RandDatum(typ)
|
|
str := fmt.Sprintf("%v", d)
|
|
if datumSet[str] {
|
|
continue
|
|
}
|
|
datumSet[str] = true
|
|
datums = append(datums, d)
|
|
}
|
|
|
|
for i := 0; i < rows; i++ {
|
|
val, err := rand.Int(rand.Reader, big.NewInt(int64(ndv)))
|
|
if err != nil {
|
|
panic("Fail to generate int number")
|
|
}
|
|
results = append(results, datums[val.Int64()])
|
|
}
|
|
}
|
|
|
|
if order {
|
|
sort.Slice(results, func(i, j int) bool {
|
|
switch typ.GetType() {
|
|
case mysql.TypeLong, mysql.TypeLonglong:
|
|
return results[i].(int64) < results[j].(int64)
|
|
case mysql.TypeDouble:
|
|
return results[i].(float64) < results[j].(float64)
|
|
case mysql.TypeVarString:
|
|
return results[i].(string) < results[j].(string)
|
|
default:
|
|
panic("not implement")
|
|
}
|
|
})
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// RandDatum rand datum
|
|
func (*MockDataSource) RandDatum(typ *types.FieldType) any {
|
|
val, _ := rand.Int(rand.Reader, big.NewInt(1000000))
|
|
switch typ.GetType() {
|
|
case mysql.TypeLong, mysql.TypeLonglong:
|
|
return val.Int64()
|
|
case mysql.TypeFloat:
|
|
floatVal, _ := val.Float64()
|
|
return float32(floatVal / 1000)
|
|
case mysql.TypeDouble:
|
|
floatVal, _ := val.Float64()
|
|
return floatVal / 1000
|
|
case mysql.TypeNewDecimal:
|
|
var d types.MyDecimal
|
|
return d.FromInt(val.Int64())
|
|
case mysql.TypeVarString:
|
|
buff := make([]byte, 10)
|
|
_, err := rand.Read(buff)
|
|
if err != nil {
|
|
panic("rand.Read returns error")
|
|
}
|
|
return base64.RawURLEncoding.EncodeToString(buff)
|
|
default:
|
|
panic("not implement")
|
|
}
|
|
}
|
|
|
|
// PrepareChunks prepares chunks
|
|
func (mds *MockDataSource) PrepareChunks() {
|
|
mds.Chunks = make([]*chunk.Chunk, len(mds.GenData))
|
|
for i := range mds.Chunks {
|
|
mds.Chunks[i] = mds.GenData[i].CopyConstruct()
|
|
}
|
|
mds.ChunkPtr = 0
|
|
}
|
|
|
|
// Next get next chunk
|
|
func (mds *MockDataSource) Next(_ context.Context, req *chunk.Chunk) error {
|
|
if mds.ChunkPtr >= len(mds.Chunks) {
|
|
req.Reset()
|
|
return nil
|
|
}
|
|
dataChk := mds.Chunks[mds.ChunkPtr]
|
|
dataChk.SwapColumns(req)
|
|
|
|
mds.ChunkPtr++
|
|
return nil
|
|
}
|
|
|
|
// MockPhysicalPlan is used to return a specified executor in when build.
|
|
// It is mainly used for testing.
|
|
type MockPhysicalPlan interface {
|
|
plannercore.PhysicalPlan
|
|
GetExecutor() exec.Executor
|
|
}
|
|
|
|
// MockDataPhysicalPlan mocks physical plan
|
|
type MockDataPhysicalPlan struct {
|
|
MockPhysicalPlan
|
|
DataSchema *expression.Schema
|
|
Exec exec.Executor
|
|
}
|
|
|
|
// GetExecutor gets executor
|
|
func (mp *MockDataPhysicalPlan) GetExecutor() exec.Executor {
|
|
return mp.Exec
|
|
}
|
|
|
|
// Schema returns schema
|
|
func (mp *MockDataPhysicalPlan) Schema() *expression.Schema {
|
|
return mp.DataSchema
|
|
}
|
|
|
|
// ExplainID returns explain id
|
|
func (*MockDataPhysicalPlan) ExplainID() fmt.Stringer {
|
|
return stringutil.MemoizeStr(func() string {
|
|
return "mockData_0"
|
|
})
|
|
}
|
|
|
|
// ID returns 0
|
|
func (*MockDataPhysicalPlan) ID() int {
|
|
return 0
|
|
}
|
|
|
|
// Stats returns nil
|
|
func (*MockDataPhysicalPlan) Stats() *property.StatsInfo {
|
|
return nil
|
|
}
|
|
|
|
// QueryBlockOffset returns 0
|
|
func (*MockDataPhysicalPlan) QueryBlockOffset() int {
|
|
return 0
|
|
}
|
|
|
|
// MemoryUsage of mockDataPhysicalPlan is only for testing
|
|
func (*MockDataPhysicalPlan) MemoryUsage() (sum int64) {
|
|
return
|
|
}
|
|
|
|
// BuildMockDataPhysicalPlan builds MockDataPhysicalPlan
|
|
func BuildMockDataPhysicalPlan(_ sessionctx.Context, srcExec exec.Executor) *MockDataPhysicalPlan {
|
|
return &MockDataPhysicalPlan{
|
|
DataSchema: srcExec.Schema(),
|
|
Exec: srcExec,
|
|
}
|
|
}
|
|
|
|
// BuildMockDataSource builds MockDataSource
|
|
func BuildMockDataSource(opt MockDataSourceParameters) *MockDataSource {
|
|
baseExec := exec.NewBaseExecutor(opt.Ctx, opt.DataSchema, 0)
|
|
m := &MockDataSource{
|
|
BaseExecutor: baseExec,
|
|
ChunkPtr: 0,
|
|
P: opt,
|
|
GenData: nil,
|
|
Chunks: nil,
|
|
}
|
|
rTypes := exec.RetTypes(m)
|
|
colData := make([][]any, len(rTypes))
|
|
for i := 0; i < len(rTypes); i++ {
|
|
colData[i] = m.GenColDatums(i)
|
|
}
|
|
|
|
m.GenData = make([]*chunk.Chunk, (m.P.Rows+m.MaxChunkSize()-1)/m.MaxChunkSize())
|
|
for i := range m.GenData {
|
|
m.GenData[i] = chunk.NewChunkWithCapacity(exec.RetTypes(m), m.MaxChunkSize())
|
|
}
|
|
|
|
for i := 0; i < m.P.Rows; i++ {
|
|
idx := i / m.MaxChunkSize()
|
|
retTypes := exec.RetTypes(m)
|
|
for colIdx := 0; colIdx < len(rTypes); colIdx++ {
|
|
switch retTypes[colIdx].GetType() {
|
|
case mysql.TypeLong, mysql.TypeLonglong:
|
|
m.GenData[idx].AppendInt64(colIdx, colData[colIdx][i].(int64))
|
|
case mysql.TypeFloat:
|
|
m.GenData[idx].AppendFloat32(colIdx, colData[colIdx][i].(float32))
|
|
case mysql.TypeDouble:
|
|
m.GenData[idx].AppendFloat64(colIdx, colData[colIdx][i].(float64))
|
|
case mysql.TypeNewDecimal:
|
|
m.GenData[idx].AppendMyDecimal(colIdx, colData[colIdx][i].(*types.MyDecimal))
|
|
case mysql.TypeVarString:
|
|
m.GenData[idx].AppendString(colIdx, colData[colIdx][i].(string))
|
|
default:
|
|
panic("not implement")
|
|
}
|
|
}
|
|
}
|
|
return m
|
|
}
|
|
|
|
// BuildMockDataSourceWithIndex builds MockDataSourceWithIndex
|
|
func BuildMockDataSourceWithIndex(opt MockDataSourceParameters, index []int) *MockDataSource {
|
|
opt.Orders = make([]bool, len(opt.DataSchema.Columns))
|
|
for _, idx := range index {
|
|
opt.Orders[idx] = true
|
|
}
|
|
return BuildMockDataSource(opt)
|
|
}
|
|
|
|
// MockActionOnExceed is for test.
|
|
type MockActionOnExceed struct {
|
|
memory.BaseOOMAction
|
|
triggeredNum atomic.Int32
|
|
}
|
|
|
|
// Action add the triggered number.
|
|
func (m *MockActionOnExceed) Action(*memory.Tracker) {
|
|
m.triggeredNum.Add(1)
|
|
}
|
|
|
|
// GetPriority get the priority of the Action.
|
|
func (*MockActionOnExceed) GetPriority() int64 {
|
|
return memory.DefLogPriority
|
|
}
|
|
|
|
// GetTriggeredNum get the triggered number of the Action
|
|
func (m *MockActionOnExceed) GetTriggeredNum() int {
|
|
return int(m.triggeredNum.Load())
|
|
}
|