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