You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

241 lines
7.3 KiB

// Copyright 2019 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,
// See the License for the specific language governing permissions and
// limitations under the License.
package executor
import (
"context"
"fmt"
"math/rand"
"github.com/cznic/mathutil"
. "github.com/pingcap/check"
"github.com/pingcap/parser/model"
"github.com/pingcap/parser/mysql"
"github.com/pingcap/tidb/distsql"
"github.com/pingcap/tidb/expression"
"github.com/pingcap/tidb/kv"
"github.com/pingcap/tidb/planner/core"
"github.com/pingcap/tidb/sessionctx"
"github.com/pingcap/tidb/statistics"
"github.com/pingcap/tidb/table/tables"
"github.com/pingcap/tidb/types"
"github.com/pingcap/tidb/util/chunk"
"github.com/pingcap/tipb/go-tipb"
)
type requiredRowsSelectResult struct {
retTypes []*types.FieldType
totalRows int
count int
expectedRowsRet []int
numNextCalled int
}
func (r *requiredRowsSelectResult) Fetch(context.Context) {}
func (r *requiredRowsSelectResult) NextRaw(context.Context) ([]byte, error) { return nil, nil }
func (r *requiredRowsSelectResult) Close() error { return nil }
func (r *requiredRowsSelectResult) Next(ctx context.Context, chk *chunk.Chunk) error {
defer func() {
if r.numNextCalled >= len(r.expectedRowsRet) {
return
}
rowsRet := chk.NumRows()
expected := r.expectedRowsRet[r.numNextCalled]
if rowsRet != expected {
panic(fmt.Sprintf("unexpected number of rows returned, obtain: %v, expected: %v", rowsRet, expected))
}
r.numNextCalled++
}()
chk.Reset()
if r.count > r.totalRows {
return nil
}
required := mathutil.Min(chk.RequiredRows(), r.totalRows-r.count)
for i := 0; i < required; i++ {
chk.AppendRow(r.genOneRow())
}
r.count += required
return nil
}
func (r *requiredRowsSelectResult) genOneRow() chunk.Row {
row := chunk.MutRowFromTypes(r.retTypes)
for i := range r.retTypes {
row.SetValue(i, r.genValue(r.retTypes[i]))
}
return row.ToRow()
}
func (r *requiredRowsSelectResult) genValue(valType *types.FieldType) interface{} {
switch valType.Tp {
case mysql.TypeLong, mysql.TypeLonglong:
return int64(rand.Int())
case mysql.TypeDouble:
return rand.Float64()
default:
panic("not implement")
}
}
func mockDistsqlSelectCtxSet(totalRows int, expectedRowsRet []int) context.Context {
ctx := context.Background()
ctx = context.WithValue(ctx, "totalRows", totalRows)
ctx = context.WithValue(ctx, "expectedRowsRet", expectedRowsRet)
return ctx
}
func mockDistsqlSelectCtxGet(ctx context.Context) (totalRows int, expectedRowsRet []int) {
totalRows = ctx.Value("totalRows").(int)
expectedRowsRet = ctx.Value("expectedRowsRet").([]int)
return
}
func mockSelectResult(ctx context.Context, sctx sessionctx.Context, kvReq *kv.Request,
fieldTypes []*types.FieldType, fb *statistics.QueryFeedback, copPlanIDs []int) (distsql.SelectResult, error) {
totalRows, expectedRowsRet := mockDistsqlSelectCtxGet(ctx)
return &requiredRowsSelectResult{
retTypes: fieldTypes,
totalRows: totalRows,
expectedRowsRet: expectedRowsRet,
}, nil
}
func buildTableReader(sctx sessionctx.Context) Executor {
e := &TableReaderExecutor{
baseExecutor: buildMockBaseExec(sctx),
table: &tables.TableCommon{},
dagPB: buildMockDAGRequest(sctx),
selectResultHook: selectResultHook{mockSelectResult},
}
return e
}
func buildMockDAGRequest(sctx sessionctx.Context) *tipb.DAGRequest {
builder := newExecutorBuilder(sctx, nil)
req, _, err := builder.constructDAGReq([]core.PhysicalPlan{&core.PhysicalTableScan{
Columns: []*model.ColumnInfo{},
Table: &model.TableInfo{ID: 12345, PKIsHandle: false},
Desc: false,
}}, kv.TiKV)
if err != nil {
panic(err)
}
return req
}
func buildMockBaseExec(sctx sessionctx.Context) baseExecutor {
retTypes := []*types.FieldType{types.NewFieldType(mysql.TypeDouble), types.NewFieldType(mysql.TypeLonglong)}
cols := make([]*expression.Column, len(retTypes))
for i := range retTypes {
cols[i] = &expression.Column{Index: i, RetType: retTypes[i]}
}
schema := expression.NewSchema(cols...)
baseExec := newBaseExecutor(sctx, schema, 0)
return baseExec
}
func (s *testExecSuite) TestTableReaderRequiredRows(c *C) {
maxChunkSize := defaultCtx().GetSessionVars().MaxChunkSize
testCases := []struct {
totalRows int
requiredRows []int
expectedRows []int
expectedRowsDS []int
}{
{
totalRows: 10,
requiredRows: []int{1, 5, 3, 10},
expectedRows: []int{1, 5, 3, 1},
expectedRowsDS: []int{1, 5, 3, 1},
},
{
totalRows: maxChunkSize + 1,
requiredRows: []int{1, 5, 3, 10, maxChunkSize},
expectedRows: []int{1, 5, 3, 10, (maxChunkSize + 1) - 1 - 5 - 3 - 10},
expectedRowsDS: []int{1, 5, 3, 10, (maxChunkSize + 1) - 1 - 5 - 3 - 10},
},
{
totalRows: 3*maxChunkSize + 1,
requiredRows: []int{3, 10, maxChunkSize},
expectedRows: []int{3, 10, maxChunkSize},
expectedRowsDS: []int{3, 10, maxChunkSize},
},
}
for _, testCase := range testCases {
sctx := defaultCtx()
ctx := mockDistsqlSelectCtxSet(testCase.totalRows, testCase.expectedRowsDS)
exec := buildTableReader(sctx)
c.Assert(exec.Open(ctx), IsNil)
chk := newFirstChunk(exec)
for i := range testCase.requiredRows {
chk.SetRequiredRows(testCase.requiredRows[i], maxChunkSize)
c.Assert(exec.Next(ctx, chk), IsNil)
c.Assert(chk.NumRows(), Equals, testCase.expectedRows[i])
}
c.Assert(exec.Close(), IsNil)
}
}
func buildIndexReader(sctx sessionctx.Context) Executor {
e := &IndexReaderExecutor{
baseExecutor: buildMockBaseExec(sctx),
dagPB: buildMockDAGRequest(sctx),
index: &model.IndexInfo{},
selectResultHook: selectResultHook{mockSelectResult},
}
return e
}
func (s *testExecSuite) TestIndexReaderRequiredRows(c *C) {
maxChunkSize := defaultCtx().GetSessionVars().MaxChunkSize
testCases := []struct {
totalRows int
requiredRows []int
expectedRows []int
expectedRowsDS []int
}{
{
totalRows: 10,
requiredRows: []int{1, 5, 3, 10},
expectedRows: []int{1, 5, 3, 1},
expectedRowsDS: []int{1, 5, 3, 1},
},
{
totalRows: maxChunkSize + 1,
requiredRows: []int{1, 5, 3, 10, maxChunkSize},
expectedRows: []int{1, 5, 3, 10, (maxChunkSize + 1) - 1 - 5 - 3 - 10},
expectedRowsDS: []int{1, 5, 3, 10, (maxChunkSize + 1) - 1 - 5 - 3 - 10},
},
{
totalRows: 3*maxChunkSize + 1,
requiredRows: []int{3, 10, maxChunkSize},
expectedRows: []int{3, 10, maxChunkSize},
expectedRowsDS: []int{3, 10, maxChunkSize},
},
}
for _, testCase := range testCases {
sctx := defaultCtx()
ctx := mockDistsqlSelectCtxSet(testCase.totalRows, testCase.expectedRowsDS)
exec := buildIndexReader(sctx)
c.Assert(exec.Open(ctx), IsNil)
chk := newFirstChunk(exec)
for i := range testCase.requiredRows {
chk.SetRequiredRows(testCase.requiredRows[i], maxChunkSize)
c.Assert(exec.Next(ctx, chk), IsNil)
c.Assert(chk.NumRows(), Equals, testCase.expectedRows[i])
}
c.Assert(exec.Close(), IsNil)
}
}