github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/interlock/executor_required_rows_test.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package interlock
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"math"
    20  	"math/rand"
    21  	"time"
    22  
    23  	"github.com/cznic/mathutil"
    24  	"github.com/whtcorpsinc/BerolinaSQL/allegrosql"
    25  	"github.com/whtcorpsinc/BerolinaSQL/ast"
    26  	. "github.com/whtcorpsinc/check"
    27  	causetembedded "github.com/whtcorpsinc/milevadb/causet/embedded"
    28  	"github.com/whtcorpsinc/milevadb/causet/soliton"
    29  	"github.com/whtcorpsinc/milevadb/memex"
    30  	"github.com/whtcorpsinc/milevadb/memex/aggregation"
    31  	"github.com/whtcorpsinc/milevadb/soliton/chunk"
    32  	"github.com/whtcorpsinc/milevadb/soliton/disk"
    33  	"github.com/whtcorpsinc/milevadb/soliton/memory"
    34  	"github.com/whtcorpsinc/milevadb/soliton/mock"
    35  	"github.com/whtcorpsinc/milevadb/stochastikctx"
    36  	"github.com/whtcorpsinc/milevadb/stochastikctx/variable"
    37  	"github.com/whtcorpsinc/milevadb/types"
    38  	"github.com/whtcorpsinc/milevadb/types/json"
    39  )
    40  
    41  type requiredEventsDataSource struct {
    42  	baseInterlockingDirectorate
    43  	totalEvents int
    44  	count       int
    45  	ctx         stochastikctx.Context
    46  
    47  	expectedEventsRet []int
    48  	numNextCalled     int
    49  
    50  	generator func(valType *types.FieldType) interface{}
    51  }
    52  
    53  func newRequiredEventsDataSourceWithGenerator(ctx stochastikctx.Context, totalEvents int, expectedEventsRet []int,
    54  	gen func(valType *types.FieldType) interface{}) *requiredEventsDataSource {
    55  	ds := newRequiredEventsDataSource(ctx, totalEvents, expectedEventsRet)
    56  	ds.generator = gen
    57  	return ds
    58  }
    59  
    60  func newRequiredEventsDataSource(ctx stochastikctx.Context, totalEvents int, expectedEventsRet []int) *requiredEventsDataSource {
    61  	// the schemaReplicant of output is fixed now, which is [Double, Long]
    62  	retTypes := []*types.FieldType{types.NewFieldType(allegrosql.TypeDouble), types.NewFieldType(allegrosql.TypeLonglong)}
    63  	defcaus := make([]*memex.DeferredCauset, len(retTypes))
    64  	for i := range retTypes {
    65  		defcaus[i] = &memex.DeferredCauset{Index: i, RetType: retTypes[i]}
    66  	}
    67  	schemaReplicant := memex.NewSchema(defcaus...)
    68  	baseInterDirc := newBaseInterlockingDirectorate(ctx, schemaReplicant, 0)
    69  	return &requiredEventsDataSource{baseInterDirc, totalEvents, 0, ctx, expectedEventsRet, 0, defaultGenerator}
    70  }
    71  
    72  func (r *requiredEventsDataSource) Next(ctx context.Context, req *chunk.Chunk) error {
    73  	defer func() {
    74  		if r.expectedEventsRet == nil {
    75  			r.numNextCalled++
    76  			return
    77  		}
    78  		rowsRet := req.NumEvents()
    79  		expected := r.expectedEventsRet[r.numNextCalled]
    80  		if rowsRet != expected {
    81  			panic(fmt.Sprintf("unexpected number of rows returned, obtain: %v, expected: %v", rowsRet, expected))
    82  		}
    83  		r.numNextCalled++
    84  	}()
    85  
    86  	req.Reset()
    87  	if r.count > r.totalEvents {
    88  		return nil
    89  	}
    90  	required := mathutil.Min(req.RequiredEvents(), r.totalEvents-r.count)
    91  	for i := 0; i < required; i++ {
    92  		req.AppendEvent(r.genOneEvent())
    93  	}
    94  	r.count += required
    95  	return nil
    96  }
    97  
    98  func (r *requiredEventsDataSource) genOneEvent() chunk.Event {
    99  	event := chunk.MutEventFromTypes(retTypes(r))
   100  	for i, tp := range retTypes(r) {
   101  		event.SetValue(i, r.generator(tp))
   102  	}
   103  	return event.ToEvent()
   104  }
   105  
   106  func defaultGenerator(valType *types.FieldType) interface{} {
   107  	switch valType.Tp {
   108  	case allegrosql.TypeLong, allegrosql.TypeLonglong:
   109  		return int64(rand.Int())
   110  	case allegrosql.TypeDouble:
   111  		return rand.Float64()
   112  	default:
   113  		panic("not implement")
   114  	}
   115  }
   116  
   117  func (r *requiredEventsDataSource) checkNumNextCalled() error {
   118  	if r.numNextCalled != len(r.expectedEventsRet) {
   119  		return fmt.Errorf("unexpected number of call on Next, obtain: %v, expected: %v",
   120  			r.numNextCalled, len(r.expectedEventsRet))
   121  	}
   122  	return nil
   123  }
   124  
   125  func (s *testInterDircSuite) TestLimitRequiredEvents(c *C) {
   126  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   127  	testCases := []struct {
   128  		totalEvents      int
   129  		limitOffset      int
   130  		limitCount       int
   131  		requiredEvents   []int
   132  		expectedEvents   []int
   133  		expectedEventsDS []int
   134  	}{
   135  		{
   136  			totalEvents:      20,
   137  			limitOffset:      0,
   138  			limitCount:       10,
   139  			requiredEvents:   []int{3, 5, 1, 500, 500},
   140  			expectedEvents:   []int{3, 5, 1, 1, 0},
   141  			expectedEventsDS: []int{3, 5, 1, 1},
   142  		},
   143  		{
   144  			totalEvents:      20,
   145  			limitOffset:      0,
   146  			limitCount:       25,
   147  			requiredEvents:   []int{9, 500},
   148  			expectedEvents:   []int{9, 11},
   149  			expectedEventsDS: []int{9, 11},
   150  		},
   151  		{
   152  			totalEvents:      100,
   153  			limitOffset:      50,
   154  			limitCount:       30,
   155  			requiredEvents:   []int{10, 5, 10, 20},
   156  			expectedEvents:   []int{10, 5, 10, 5},
   157  			expectedEventsDS: []int{60, 5, 10, 5},
   158  		},
   159  		{
   160  			totalEvents:      100,
   161  			limitOffset:      101,
   162  			limitCount:       10,
   163  			requiredEvents:   []int{10},
   164  			expectedEvents:   []int{0},
   165  			expectedEventsDS: []int{100, 0},
   166  		},
   167  		{
   168  			totalEvents:      maxChunkSize + 20,
   169  			limitOffset:      maxChunkSize + 1,
   170  			limitCount:       10,
   171  			requiredEvents:   []int{3, 3, 3, 100},
   172  			expectedEvents:   []int{3, 3, 3, 1},
   173  			expectedEventsDS: []int{maxChunkSize, 4, 3, 3, 1},
   174  		},
   175  	}
   176  
   177  	for _, testCase := range testCases {
   178  		sctx := defaultCtx()
   179  		ctx := context.Background()
   180  		ds := newRequiredEventsDataSource(sctx, testCase.totalEvents, testCase.expectedEventsDS)
   181  		exec := buildLimitInterDirc(sctx, ds, testCase.limitOffset, testCase.limitCount)
   182  		c.Assert(exec.Open(ctx), IsNil)
   183  		chk := newFirstChunk(exec)
   184  		for i := range testCase.requiredEvents {
   185  			chk.SetRequiredEvents(testCase.requiredEvents[i], sctx.GetStochastikVars().MaxChunkSize)
   186  			c.Assert(exec.Next(ctx, chk), IsNil)
   187  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   188  		}
   189  		c.Assert(exec.Close(), IsNil)
   190  		c.Assert(ds.checkNumNextCalled(), IsNil)
   191  	}
   192  }
   193  
   194  func buildLimitInterDirc(ctx stochastikctx.Context, src InterlockingDirectorate, offset, count int) InterlockingDirectorate {
   195  	n := mathutil.Min(count, ctx.GetStochastikVars().MaxChunkSize)
   196  	base := newBaseInterlockingDirectorate(ctx, src.Schema(), 0, src)
   197  	base.initCap = n
   198  	limitInterDirc := &LimitInterDirc{
   199  		baseInterlockingDirectorate: base,
   200  		begin:                       uint64(offset),
   201  		end:                         uint64(offset + count),
   202  	}
   203  	return limitInterDirc
   204  }
   205  
   206  func defaultCtx() stochastikctx.Context {
   207  	ctx := mock.NewContext()
   208  	ctx.GetStochastikVars().InitChunkSize = variable.DefInitChunkSize
   209  	ctx.GetStochastikVars().MaxChunkSize = variable.DefMaxChunkSize
   210  	ctx.GetStochastikVars().StmtCtx.MemTracker = memory.NewTracker(-1, ctx.GetStochastikVars().MemQuotaQuery)
   211  	ctx.GetStochastikVars().StmtCtx.DiskTracker = disk.NewTracker(-1, -1)
   212  	ctx.GetStochastikVars().SnapshotTS = uint64(1)
   213  	return ctx
   214  }
   215  
   216  func (s *testInterDircSuite) TestSortRequiredEvents(c *C) {
   217  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   218  	testCases := []struct {
   219  		totalEvents      int
   220  		groupBy          []int
   221  		requiredEvents   []int
   222  		expectedEvents   []int
   223  		expectedEventsDS []int
   224  	}{
   225  		{
   226  			totalEvents:      10,
   227  			groupBy:          []int{0},
   228  			requiredEvents:   []int{1, 5, 3, 10},
   229  			expectedEvents:   []int{1, 5, 3, 1},
   230  			expectedEventsDS: []int{10, 0},
   231  		},
   232  		{
   233  			totalEvents:      10,
   234  			groupBy:          []int{0, 1},
   235  			requiredEvents:   []int{1, 5, 3, 10},
   236  			expectedEvents:   []int{1, 5, 3, 1},
   237  			expectedEventsDS: []int{10, 0},
   238  		},
   239  		{
   240  			totalEvents:      maxChunkSize + 1,
   241  			groupBy:          []int{0},
   242  			requiredEvents:   []int{1, 5, 3, 10, maxChunkSize},
   243  			expectedEvents:   []int{1, 5, 3, 10, (maxChunkSize + 1) - 1 - 5 - 3 - 10},
   244  			expectedEventsDS: []int{maxChunkSize, 1, 0},
   245  		},
   246  		{
   247  			totalEvents:      3*maxChunkSize + 1,
   248  			groupBy:          []int{0},
   249  			requiredEvents:   []int{1, 5, 3, 10, maxChunkSize},
   250  			expectedEvents:   []int{1, 5, 3, 10, maxChunkSize},
   251  			expectedEventsDS: []int{maxChunkSize, maxChunkSize, maxChunkSize, 1, 0},
   252  		},
   253  	}
   254  
   255  	for _, testCase := range testCases {
   256  		sctx := defaultCtx()
   257  		ctx := context.Background()
   258  		ds := newRequiredEventsDataSource(sctx, testCase.totalEvents, testCase.expectedEventsDS)
   259  		byItems := make([]*soliton.ByItems, 0, len(testCase.groupBy))
   260  		for _, groupBy := range testCase.groupBy {
   261  			defCaus := ds.Schema().DeferredCausets[groupBy]
   262  			byItems = append(byItems, &soliton.ByItems{Expr: defCaus})
   263  		}
   264  		exec := buildSortInterDirc(sctx, byItems, ds)
   265  		c.Assert(exec.Open(ctx), IsNil)
   266  		chk := newFirstChunk(exec)
   267  		for i := range testCase.requiredEvents {
   268  			chk.SetRequiredEvents(testCase.requiredEvents[i], maxChunkSize)
   269  			c.Assert(exec.Next(ctx, chk), IsNil)
   270  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   271  		}
   272  		c.Assert(exec.Close(), IsNil)
   273  		c.Assert(ds.checkNumNextCalled(), IsNil)
   274  	}
   275  }
   276  
   277  func buildSortInterDirc(sctx stochastikctx.Context, byItems []*soliton.ByItems, src InterlockingDirectorate) InterlockingDirectorate {
   278  	sortInterDirc := SortInterDirc{
   279  		baseInterlockingDirectorate: newBaseInterlockingDirectorate(sctx, src.Schema(), 0, src),
   280  		ByItems:                     byItems,
   281  		schemaReplicant:             src.Schema(),
   282  	}
   283  	return &sortInterDirc
   284  }
   285  
   286  func (s *testInterDircSuite) TestTopNRequiredEvents(c *C) {
   287  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   288  	testCases := []struct {
   289  		totalEvents      int
   290  		topNOffset       int
   291  		topNCount        int
   292  		groupBy          []int
   293  		requiredEvents   []int
   294  		expectedEvents   []int
   295  		expectedEventsDS []int
   296  	}{
   297  		{
   298  			totalEvents:      10,
   299  			topNOffset:       0,
   300  			topNCount:        10,
   301  			groupBy:          []int{0},
   302  			requiredEvents:   []int{1, 1, 1, 1, 10},
   303  			expectedEvents:   []int{1, 1, 1, 1, 6},
   304  			expectedEventsDS: []int{10, 0},
   305  		},
   306  		{
   307  			totalEvents:      100,
   308  			topNOffset:       15,
   309  			topNCount:        11,
   310  			groupBy:          []int{0},
   311  			requiredEvents:   []int{1, 1, 1, 1, 10},
   312  			expectedEvents:   []int{1, 1, 1, 1, 7},
   313  			expectedEventsDS: []int{26, 100 - 26, 0},
   314  		},
   315  		{
   316  			totalEvents:      100,
   317  			topNOffset:       95,
   318  			topNCount:        10,
   319  			groupBy:          []int{0},
   320  			requiredEvents:   []int{1, 2, 3, 10},
   321  			expectedEvents:   []int{1, 2, 2, 0},
   322  			expectedEventsDS: []int{100, 0, 0},
   323  		},
   324  		{
   325  			totalEvents:      maxChunkSize + 20,
   326  			topNOffset:       1,
   327  			topNCount:        5,
   328  			groupBy:          []int{0, 1},
   329  			requiredEvents:   []int{1, 3, 7, 10},
   330  			expectedEvents:   []int{1, 3, 1, 0},
   331  			expectedEventsDS: []int{6, maxChunkSize, 14, 0},
   332  		},
   333  		{
   334  			totalEvents:      maxChunkSize + maxChunkSize + 20,
   335  			topNOffset:       maxChunkSize + 10,
   336  			topNCount:        8,
   337  			groupBy:          []int{0, 1},
   338  			requiredEvents:   []int{1, 2, 3, 5, 7},
   339  			expectedEvents:   []int{1, 2, 3, 2, 0},
   340  			expectedEventsDS: []int{maxChunkSize, 18, maxChunkSize, 2, 0},
   341  		},
   342  		{
   343  			totalEvents:      maxChunkSize*5 + 10,
   344  			topNOffset:       maxChunkSize*5 + 20,
   345  			topNCount:        10,
   346  			groupBy:          []int{0, 1},
   347  			requiredEvents:   []int{1, 2, 3},
   348  			expectedEvents:   []int{0, 0, 0},
   349  			expectedEventsDS: []int{maxChunkSize, maxChunkSize, maxChunkSize, maxChunkSize, maxChunkSize, 10, 0, 0},
   350  		},
   351  		{
   352  			totalEvents:      maxChunkSize + maxChunkSize + 10,
   353  			topNOffset:       10,
   354  			topNCount:        math.MaxInt64,
   355  			groupBy:          []int{0, 1},
   356  			requiredEvents:   []int{1, 2, 3, maxChunkSize, maxChunkSize},
   357  			expectedEvents:   []int{1, 2, 3, maxChunkSize, maxChunkSize - 1 - 2 - 3},
   358  			expectedEventsDS: []int{maxChunkSize, maxChunkSize, 10, 0, 0},
   359  		},
   360  	}
   361  
   362  	for _, testCase := range testCases {
   363  		sctx := defaultCtx()
   364  		ctx := context.Background()
   365  		ds := newRequiredEventsDataSource(sctx, testCase.totalEvents, testCase.expectedEventsDS)
   366  		byItems := make([]*soliton.ByItems, 0, len(testCase.groupBy))
   367  		for _, groupBy := range testCase.groupBy {
   368  			defCaus := ds.Schema().DeferredCausets[groupBy]
   369  			byItems = append(byItems, &soliton.ByItems{Expr: defCaus})
   370  		}
   371  		exec := buildTopNInterDirc(sctx, testCase.topNOffset, testCase.topNCount, byItems, ds)
   372  		c.Assert(exec.Open(ctx), IsNil)
   373  		chk := newFirstChunk(exec)
   374  		for i := range testCase.requiredEvents {
   375  			chk.SetRequiredEvents(testCase.requiredEvents[i], maxChunkSize)
   376  			c.Assert(exec.Next(ctx, chk), IsNil)
   377  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   378  		}
   379  		c.Assert(exec.Close(), IsNil)
   380  		c.Assert(ds.checkNumNextCalled(), IsNil)
   381  	}
   382  }
   383  
   384  func buildTopNInterDirc(ctx stochastikctx.Context, offset, count int, byItems []*soliton.ByItems, src InterlockingDirectorate) InterlockingDirectorate {
   385  	sortInterDirc := SortInterDirc{
   386  		baseInterlockingDirectorate: newBaseInterlockingDirectorate(ctx, src.Schema(), 0, src),
   387  		ByItems:                     byItems,
   388  		schemaReplicant:             src.Schema(),
   389  	}
   390  	return &TopNInterDirc{
   391  		SortInterDirc: sortInterDirc,
   392  		limit:         &causetembedded.PhysicalLimit{Count: uint64(count), Offset: uint64(offset)},
   393  	}
   394  }
   395  
   396  func (s *testInterDircSuite) TestSelectionRequiredEvents(c *C) {
   397  	gen01 := func() func(valType *types.FieldType) interface{} {
   398  		closureCount := 0
   399  		return func(valType *types.FieldType) interface{} {
   400  			switch valType.Tp {
   401  			case allegrosql.TypeLong, allegrosql.TypeLonglong:
   402  				ret := int64(closureCount % 2)
   403  				closureCount++
   404  				return ret
   405  			case allegrosql.TypeDouble:
   406  				return rand.Float64()
   407  			default:
   408  				panic("not implement")
   409  			}
   410  		}
   411  	}
   412  
   413  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   414  	testCases := []struct {
   415  		totalEvents       int
   416  		filtersOfDefCaus1 int
   417  		requiredEvents    []int
   418  		expectedEvents    []int
   419  		expectedEventsDS  []int
   420  		gen               func(valType *types.FieldType) interface{}
   421  	}{
   422  		{
   423  			totalEvents:      20,
   424  			requiredEvents:   []int{1, 2, 3, 4, 5, 20},
   425  			expectedEvents:   []int{1, 2, 3, 4, 5, 5},
   426  			expectedEventsDS: []int{20, 0},
   427  		},
   428  		{
   429  			totalEvents:       20,
   430  			filtersOfDefCaus1: 0,
   431  			requiredEvents:    []int{1, 3, 5, 7, 9},
   432  			expectedEvents:    []int{1, 3, 5, 1, 0},
   433  			expectedEventsDS:  []int{20, 0, 0},
   434  			gen:               gen01(),
   435  		},
   436  		{
   437  			totalEvents:       maxChunkSize + 20,
   438  			filtersOfDefCaus1: 1,
   439  			requiredEvents:    []int{1, 3, 5, maxChunkSize},
   440  			expectedEvents:    []int{1, 3, 5, maxChunkSize/2 - 1 - 3 - 5 + 10},
   441  			expectedEventsDS:  []int{maxChunkSize, 20, 0},
   442  			gen:               gen01(),
   443  		},
   444  	}
   445  
   446  	for _, testCase := range testCases {
   447  		sctx := defaultCtx()
   448  		ctx := context.Background()
   449  		var filters []memex.Expression
   450  		var ds *requiredEventsDataSource
   451  		if testCase.gen == nil {
   452  			// ignore filters
   453  			ds = newRequiredEventsDataSource(sctx, testCase.totalEvents, testCase.expectedEventsDS)
   454  		} else {
   455  			ds = newRequiredEventsDataSourceWithGenerator(sctx, testCase.totalEvents, testCase.expectedEventsDS, testCase.gen)
   456  			f, err := memex.NewFunction(
   457  				sctx, ast.EQ, types.NewFieldType(byte(types.ETInt)), ds.Schema().DeferredCausets[1], &memex.Constant{
   458  					Value:   types.NewCauset(testCase.filtersOfDefCaus1),
   459  					RetType: types.NewFieldType(allegrosql.TypeTiny),
   460  				})
   461  			c.Assert(err, IsNil)
   462  			filters = append(filters, f)
   463  		}
   464  		exec := buildSelectionInterDirc(sctx, filters, ds)
   465  		c.Assert(exec.Open(ctx), IsNil)
   466  		chk := newFirstChunk(exec)
   467  		for i := range testCase.requiredEvents {
   468  			chk.SetRequiredEvents(testCase.requiredEvents[i], maxChunkSize)
   469  			c.Assert(exec.Next(ctx, chk), IsNil)
   470  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   471  		}
   472  		c.Assert(exec.Close(), IsNil)
   473  		c.Assert(ds.checkNumNextCalled(), IsNil)
   474  	}
   475  }
   476  
   477  func buildSelectionInterDirc(ctx stochastikctx.Context, filters []memex.Expression, src InterlockingDirectorate) InterlockingDirectorate {
   478  	return &SelectionInterDirc{
   479  		baseInterlockingDirectorate: newBaseInterlockingDirectorate(ctx, src.Schema(), 0, src),
   480  		filters:                     filters,
   481  	}
   482  }
   483  
   484  func (s *testInterDircSuite) TestProjectionUnparallelRequiredEvents(c *C) {
   485  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   486  	testCases := []struct {
   487  		totalEvents      int
   488  		requiredEvents   []int
   489  		expectedEvents   []int
   490  		expectedEventsDS []int
   491  	}{
   492  		{
   493  			totalEvents:      20,
   494  			requiredEvents:   []int{1, 3, 5, 7, 9},
   495  			expectedEvents:   []int{1, 3, 5, 7, 4},
   496  			expectedEventsDS: []int{1, 3, 5, 7, 4},
   497  		},
   498  		{
   499  			totalEvents:      maxChunkSize + 10,
   500  			requiredEvents:   []int{1, 3, 5, 7, 9, maxChunkSize},
   501  			expectedEvents:   []int{1, 3, 5, 7, 9, maxChunkSize - 1 - 3 - 5 - 7 - 9 + 10},
   502  			expectedEventsDS: []int{1, 3, 5, 7, 9, maxChunkSize - 1 - 3 - 5 - 7 - 9 + 10},
   503  		},
   504  		{
   505  			totalEvents:      maxChunkSize*2 + 10,
   506  			requiredEvents:   []int{1, 7, 9, maxChunkSize, maxChunkSize + 10},
   507  			expectedEvents:   []int{1, 7, 9, maxChunkSize, maxChunkSize + 10 - 1 - 7 - 9},
   508  			expectedEventsDS: []int{1, 7, 9, maxChunkSize, maxChunkSize + 10 - 1 - 7 - 9},
   509  		},
   510  	}
   511  
   512  	for _, testCase := range testCases {
   513  		sctx := defaultCtx()
   514  		ctx := context.Background()
   515  		ds := newRequiredEventsDataSource(sctx, testCase.totalEvents, testCase.expectedEventsDS)
   516  		exprs := make([]memex.Expression, 0, len(ds.Schema().DeferredCausets))
   517  		if len(exprs) == 0 {
   518  			for _, defCaus := range ds.Schema().DeferredCausets {
   519  				exprs = append(exprs, defCaus)
   520  			}
   521  		}
   522  		exec := buildProjectionInterDirc(sctx, exprs, ds, 0)
   523  		c.Assert(exec.Open(ctx), IsNil)
   524  		chk := newFirstChunk(exec)
   525  		for i := range testCase.requiredEvents {
   526  			chk.SetRequiredEvents(testCase.requiredEvents[i], maxChunkSize)
   527  			c.Assert(exec.Next(ctx, chk), IsNil)
   528  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   529  		}
   530  		c.Assert(exec.Close(), IsNil)
   531  		c.Assert(ds.checkNumNextCalled(), IsNil)
   532  	}
   533  }
   534  
   535  func (s *testInterDircSuite) TestProjectionParallelRequiredEvents(c *C) {
   536  	c.Skip("not sblock because of goroutine schedule")
   537  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   538  	testCases := []struct {
   539  		totalEvents      int
   540  		numWorkers       int
   541  		requiredEvents   []int
   542  		expectedEvents   []int
   543  		expectedEventsDS []int
   544  	}{
   545  		{
   546  			totalEvents:      20,
   547  			numWorkers:       1,
   548  			requiredEvents:   []int{1, 2, 3, 4, 5, 6, 1, 1},
   549  			expectedEvents:   []int{1, 1, 2, 3, 4, 5, 4, 0},
   550  			expectedEventsDS: []int{1, 1, 2, 3, 4, 5, 4, 0},
   551  		},
   552  		{
   553  			totalEvents:      maxChunkSize * 2,
   554  			numWorkers:       1,
   555  			requiredEvents:   []int{7, maxChunkSize, maxChunkSize, maxChunkSize},
   556  			expectedEvents:   []int{7, 7, maxChunkSize, maxChunkSize - 14},
   557  			expectedEventsDS: []int{7, 7, maxChunkSize, maxChunkSize - 14, 0},
   558  		},
   559  		{
   560  			totalEvents:      20,
   561  			numWorkers:       2,
   562  			requiredEvents:   []int{1, 2, 3, 4, 5, 6, 1, 1, 1},
   563  			expectedEvents:   []int{1, 1, 1, 2, 3, 4, 5, 3, 0},
   564  			expectedEventsDS: []int{1, 1, 1, 2, 3, 4, 5, 3, 0},
   565  		},
   566  	}
   567  
   568  	for _, testCase := range testCases {
   569  		sctx := defaultCtx()
   570  		ctx := context.Background()
   571  		ds := newRequiredEventsDataSource(sctx, testCase.totalEvents, testCase.expectedEventsDS)
   572  		exprs := make([]memex.Expression, 0, len(ds.Schema().DeferredCausets))
   573  		if len(exprs) == 0 {
   574  			for _, defCaus := range ds.Schema().DeferredCausets {
   575  				exprs = append(exprs, defCaus)
   576  			}
   577  		}
   578  		exec := buildProjectionInterDirc(sctx, exprs, ds, testCase.numWorkers)
   579  		c.Assert(exec.Open(ctx), IsNil)
   580  		chk := newFirstChunk(exec)
   581  		for i := range testCase.requiredEvents {
   582  			chk.SetRequiredEvents(testCase.requiredEvents[i], maxChunkSize)
   583  			c.Assert(exec.Next(ctx, chk), IsNil)
   584  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   585  
   586  			// wait projectionInputFetcher blocked on fetching data
   587  			// from child in the background.
   588  			time.Sleep(time.Millisecond * 25)
   589  		}
   590  		c.Assert(exec.Close(), IsNil)
   591  		c.Assert(ds.checkNumNextCalled(), IsNil)
   592  	}
   593  }
   594  
   595  func buildProjectionInterDirc(ctx stochastikctx.Context, exprs []memex.Expression, src InterlockingDirectorate, numWorkers int) InterlockingDirectorate {
   596  	return &ProjectionInterDirc{
   597  		baseInterlockingDirectorate: newBaseInterlockingDirectorate(ctx, src.Schema(), 0, src),
   598  		numWorkers:                  int64(numWorkers),
   599  		evaluatorSuit:               memex.NewEvaluatorSuite(exprs, false),
   600  	}
   601  }
   602  
   603  func divGenerator(factor int) func(valType *types.FieldType) interface{} {
   604  	closureCountInt := 0
   605  	closureCountDouble := 0
   606  	return func(valType *types.FieldType) interface{} {
   607  		switch valType.Tp {
   608  		case allegrosql.TypeLong, allegrosql.TypeLonglong:
   609  			ret := int64(closureCountInt / factor)
   610  			closureCountInt++
   611  			return ret
   612  		case allegrosql.TypeDouble:
   613  			ret := float64(closureCountInt / factor)
   614  			closureCountDouble++
   615  			return ret
   616  		default:
   617  			panic("not implement")
   618  		}
   619  	}
   620  }
   621  
   622  func (s *testInterDircSuite) TestStreamAggRequiredEvents(c *C) {
   623  	maxChunkSize := defaultCtx().GetStochastikVars().MaxChunkSize
   624  	testCases := []struct {
   625  		totalEvents      int
   626  		aggFunc          string
   627  		requiredEvents   []int
   628  		expectedEvents   []int
   629  		expectedEventsDS []int
   630  		gen              func(valType *types.FieldType) interface{}
   631  	}{
   632  		{
   633  			totalEvents:      1000000,
   634  			aggFunc:          ast.AggFuncSum,
   635  			requiredEvents:   []int{1, 2, 3, 4, 5, 6, 7},
   636  			expectedEvents:   []int{1, 2, 3, 4, 5, 6, 7},
   637  			expectedEventsDS: []int{maxChunkSize},
   638  			gen:              divGenerator(1),
   639  		},
   640  		{
   641  			totalEvents:      maxChunkSize * 3,
   642  			aggFunc:          ast.AggFuncAvg,
   643  			requiredEvents:   []int{1, 3},
   644  			expectedEvents:   []int{1, 2},
   645  			expectedEventsDS: []int{maxChunkSize, maxChunkSize, maxChunkSize, 0},
   646  			gen:              divGenerator(maxChunkSize),
   647  		},
   648  		{
   649  			totalEvents:      maxChunkSize*2 - 1,
   650  			aggFunc:          ast.AggFuncMax,
   651  			requiredEvents:   []int{maxChunkSize/2 + 1},
   652  			expectedEvents:   []int{maxChunkSize/2 + 1},
   653  			expectedEventsDS: []int{maxChunkSize, maxChunkSize - 1},
   654  			gen:              divGenerator(2),
   655  		},
   656  	}
   657  
   658  	for _, testCase := range testCases {
   659  		sctx := defaultCtx()
   660  		ctx := context.Background()
   661  		ds := newRequiredEventsDataSourceWithGenerator(sctx, testCase.totalEvents, testCase.expectedEventsDS, testCase.gen)
   662  		childDefCauss := ds.Schema().DeferredCausets
   663  		schemaReplicant := memex.NewSchema(childDefCauss...)
   664  		groupBy := []memex.Expression{childDefCauss[1]}
   665  		aggFunc, err := aggregation.NewAggFuncDesc(sctx, testCase.aggFunc, []memex.Expression{childDefCauss[0]}, true)
   666  		c.Assert(err, IsNil)
   667  		aggFuncs := []*aggregation.AggFuncDesc{aggFunc}
   668  		exec := buildStreamAggInterlockingDirectorate(sctx, ds, schemaReplicant, aggFuncs, groupBy)
   669  		c.Assert(exec.Open(ctx), IsNil)
   670  		chk := newFirstChunk(exec)
   671  		for i := range testCase.requiredEvents {
   672  			chk.SetRequiredEvents(testCase.requiredEvents[i], maxChunkSize)
   673  			c.Assert(exec.Next(ctx, chk), IsNil)
   674  			c.Assert(chk.NumEvents(), Equals, testCase.expectedEvents[i])
   675  		}
   676  		c.Assert(exec.Close(), IsNil)
   677  		c.Assert(ds.checkNumNextCalled(), IsNil)
   678  	}
   679  }
   680  
   681  func (s *testInterDircSuite) TestMergeJoinRequiredEvents(c *C) {
   682  	justReturn1 := func(valType *types.FieldType) interface{} {
   683  		switch valType.Tp {
   684  		case allegrosql.TypeLong, allegrosql.TypeLonglong:
   685  			return int64(1)
   686  		case allegrosql.TypeDouble:
   687  			return float64(1)
   688  		default:
   689  			panic("not support")
   690  		}
   691  	}
   692  	joinTypes := []causetembedded.JoinType{causetembedded.RightOuterJoin, causetembedded.LeftOuterJoin,
   693  		causetembedded.LeftOuterSemiJoin, causetembedded.AntiLeftOuterSemiJoin}
   694  	for _, joinType := range joinTypes {
   695  		ctx := defaultCtx()
   696  		required := make([]int, 100)
   697  		for i := range required {
   698  			required[i] = rand.Int()%ctx.GetStochastikVars().MaxChunkSize + 1
   699  		}
   700  		innerSrc := newRequiredEventsDataSourceWithGenerator(ctx, 1, nil, justReturn1)             // just return one event: (1, 1)
   701  		outerSrc := newRequiredEventsDataSourceWithGenerator(ctx, 10000000, required, justReturn1) // always return (1, 1)
   702  		exec := buildMergeJoinInterDirc(ctx, joinType, innerSrc, outerSrc)
   703  		c.Assert(exec.Open(context.Background()), IsNil)
   704  
   705  		chk := newFirstChunk(exec)
   706  		for i := range required {
   707  			chk.SetRequiredEvents(required[i], ctx.GetStochastikVars().MaxChunkSize)
   708  			c.Assert(exec.Next(context.Background(), chk), IsNil)
   709  		}
   710  		c.Assert(exec.Close(), IsNil)
   711  		c.Assert(outerSrc.checkNumNextCalled(), IsNil)
   712  	}
   713  }
   714  
   715  func genTestChunk4VecGroupChecker(chkEvents []int, sameNum int) (expr []memex.Expression, inputs []*chunk.Chunk) {
   716  	chkNum := len(chkEvents)
   717  	numEvents := 0
   718  	inputs = make([]*chunk.Chunk, chkNum)
   719  	fts := make([]*types.FieldType, 1)
   720  	fts[0] = types.NewFieldType(allegrosql.TypeLonglong)
   721  	for i := 0; i < chkNum; i++ {
   722  		inputs[i] = chunk.New(fts, chkEvents[i], chkEvents[i])
   723  		numEvents += chkEvents[i]
   724  	}
   725  	var numGroups int
   726  	if numEvents%sameNum == 0 {
   727  		numGroups = numEvents / sameNum
   728  	} else {
   729  		numGroups = numEvents/sameNum + 1
   730  	}
   731  
   732  	rand.Seed(time.Now().Unix())
   733  	nullPos := rand.Intn(numGroups)
   734  	cnt := 0
   735  	val := rand.Int63()
   736  	for i := 0; i < chkNum; i++ {
   737  		defCaus := inputs[i].DeferredCauset(0)
   738  		defCaus.ResizeInt64(chkEvents[i], false)
   739  		i64s := defCaus.Int64s()
   740  		for j := 0; j < chkEvents[i]; j++ {
   741  			if cnt == sameNum {
   742  				val = rand.Int63()
   743  				cnt = 0
   744  				nullPos--
   745  			}
   746  			if nullPos == 0 {
   747  				defCaus.SetNull(j, true)
   748  			} else {
   749  				i64s[j] = val
   750  			}
   751  			cnt++
   752  		}
   753  	}
   754  
   755  	expr = make([]memex.Expression, 1)
   756  	expr[0] = &memex.DeferredCauset{
   757  		RetType: &types.FieldType{Tp: allegrosql.TypeLonglong, Flen: allegrosql.MaxIntWidth},
   758  		Index:   0,
   759  	}
   760  	return
   761  }
   762  
   763  func (s *testInterDircSuite) TestVecGroupChecker(c *C) {
   764  	testCases := []struct {
   765  		chunkEvents    []int
   766  		expectedGroups int
   767  		expectedFlag   []bool
   768  		sameNum        int
   769  	}{
   770  		{
   771  			chunkEvents:    []int{1024, 1},
   772  			expectedGroups: 1025,
   773  			expectedFlag:   []bool{false, false},
   774  			sameNum:        1,
   775  		},
   776  		{
   777  			chunkEvents:    []int{1024, 1},
   778  			expectedGroups: 1,
   779  			expectedFlag:   []bool{false, true},
   780  			sameNum:        1025,
   781  		},
   782  		{
   783  			chunkEvents:    []int{1, 1},
   784  			expectedGroups: 1,
   785  			expectedFlag:   []bool{false, true},
   786  			sameNum:        2,
   787  		},
   788  		{
   789  			chunkEvents:    []int{1, 1},
   790  			expectedGroups: 2,
   791  			expectedFlag:   []bool{false, false},
   792  			sameNum:        1,
   793  		},
   794  		{
   795  			chunkEvents:    []int{2, 2},
   796  			expectedGroups: 2,
   797  			expectedFlag:   []bool{false, false},
   798  			sameNum:        2,
   799  		},
   800  		{
   801  			chunkEvents:    []int{2, 2},
   802  			expectedGroups: 1,
   803  			expectedFlag:   []bool{false, true},
   804  			sameNum:        4,
   805  		},
   806  	}
   807  
   808  	ctx := mock.NewContext()
   809  	for _, testCase := range testCases {
   810  		expr, inputChks := genTestChunk4VecGroupChecker(testCase.chunkEvents, testCase.sameNum)
   811  		groupChecker := newVecGroupChecker(ctx, expr)
   812  		groupNum := 0
   813  		for i, inputChk := range inputChks {
   814  			flag, err := groupChecker.splitIntoGroups(inputChk)
   815  			c.Assert(err, IsNil)
   816  			c.Assert(flag, Equals, testCase.expectedFlag[i])
   817  			if flag {
   818  				groupNum += groupChecker.groupCount - 1
   819  			} else {
   820  				groupNum += groupChecker.groupCount
   821  			}
   822  		}
   823  		c.Assert(groupNum, Equals, testCase.expectedGroups)
   824  	}
   825  }
   826  
   827  func buildMergeJoinInterDirc(ctx stochastikctx.Context, joinType causetembedded.JoinType, innerSrc, outerSrc InterlockingDirectorate) InterlockingDirectorate {
   828  	if joinType == causetembedded.RightOuterJoin {
   829  		innerSrc, outerSrc = outerSrc, innerSrc
   830  	}
   831  
   832  	innerDefCauss := innerSrc.Schema().DeferredCausets
   833  	outerDefCauss := outerSrc.Schema().DeferredCausets
   834  	j := causetembedded.BuildMergeJoinCauset(ctx, joinType, outerDefCauss, innerDefCauss)
   835  
   836  	j.SetChildren(&mockCauset{exec: outerSrc}, &mockCauset{exec: innerSrc})
   837  	defcaus := append(append([]*memex.DeferredCauset{}, outerDefCauss...), innerDefCauss...)
   838  	schemaReplicant := memex.NewSchema(defcaus...)
   839  	j.SetSchema(schemaReplicant)
   840  
   841  	j.CompareFuncs = make([]memex.CompareFunc, 0, len(j.LeftJoinKeys))
   842  	for i := range j.LeftJoinKeys {
   843  		j.CompareFuncs = append(j.CompareFuncs, memex.GetCmpFunction(nil, j.LeftJoinKeys[i], j.RightJoinKeys[i]))
   844  	}
   845  
   846  	b := newInterlockingDirectorateBuilder(ctx, nil)
   847  	return b.build(j)
   848  }
   849  
   850  type mockCauset struct {
   851  	MockPhysicalCauset
   852  	exec InterlockingDirectorate
   853  }
   854  
   855  func (mp *mockCauset) GetInterlockingDirectorate() InterlockingDirectorate {
   856  	return mp.exec
   857  }
   858  
   859  func (mp *mockCauset) Schema() *memex.Schema {
   860  	return mp.exec.Schema()
   861  }
   862  
   863  func (s *testInterDircSuite) TestVecGroupCheckerDATARACE(c *C) {
   864  	ctx := mock.NewContext()
   865  
   866  	mTypes := []byte{allegrosql.TypeVarString, allegrosql.TypeNewDecimal, allegrosql.TypeJSON}
   867  	for _, mType := range mTypes {
   868  		exprs := make([]memex.Expression, 1)
   869  		exprs[0] = &memex.DeferredCauset{
   870  			RetType: &types.FieldType{Tp: mType},
   871  			Index:   0,
   872  		}
   873  		vgc := newVecGroupChecker(ctx, exprs)
   874  
   875  		fts := []*types.FieldType{types.NewFieldType(mType)}
   876  		chk := chunk.New(fts, 1, 1)
   877  		vgc.allocateBuffer = func(evalType types.EvalType, capacity int) (*chunk.DeferredCauset, error) {
   878  			return chk.DeferredCauset(0), nil
   879  		}
   880  		vgc.releaseBuffer = func(defCausumn *chunk.DeferredCauset) {}
   881  
   882  		switch mType {
   883  		case allegrosql.TypeVarString:
   884  			chk.DeferredCauset(0).ReserveString(1)
   885  			chk.DeferredCauset(0).AppendString("abc")
   886  		case allegrosql.TypeNewDecimal:
   887  			chk.DeferredCauset(0).ResizeDecimal(1, false)
   888  			chk.DeferredCauset(0).Decimals()[0] = *types.NewDecFromInt(123)
   889  		case allegrosql.TypeJSON:
   890  			chk.DeferredCauset(0).ReserveJSON(1)
   891  			j := new(json.BinaryJSON)
   892  			c.Assert(j.UnmarshalJSON([]byte(fmt.Sprintf(`{"%v":%v}`, 123, 123))), IsNil)
   893  			chk.DeferredCauset(0).AppendJSON(*j)
   894  		}
   895  
   896  		_, err := vgc.splitIntoGroups(chk)
   897  		c.Assert(err, IsNil)
   898  
   899  		switch mType {
   900  		case allegrosql.TypeVarString:
   901  			c.Assert(vgc.firstEventCausets[0].GetString(), Equals, "abc")
   902  			c.Assert(vgc.lastEventCausets[0].GetString(), Equals, "abc")
   903  			chk.DeferredCauset(0).ReserveString(1)
   904  			chk.DeferredCauset(0).AppendString("edf")
   905  			c.Assert(vgc.firstEventCausets[0].GetString(), Equals, "abc")
   906  			c.Assert(vgc.lastEventCausets[0].GetString(), Equals, "abc")
   907  		case allegrosql.TypeNewDecimal:
   908  			c.Assert(vgc.firstEventCausets[0].GetMysqlDecimal().String(), Equals, "123")
   909  			c.Assert(vgc.lastEventCausets[0].GetMysqlDecimal().String(), Equals, "123")
   910  			chk.DeferredCauset(0).ResizeDecimal(1, false)
   911  			chk.DeferredCauset(0).Decimals()[0] = *types.NewDecFromInt(456)
   912  			c.Assert(vgc.firstEventCausets[0].GetMysqlDecimal().String(), Equals, "123")
   913  			c.Assert(vgc.lastEventCausets[0].GetMysqlDecimal().String(), Equals, "123")
   914  		case allegrosql.TypeJSON:
   915  			c.Assert(vgc.firstEventCausets[0].GetMysqlJSON().String(), Equals, `{"123": 123}`)
   916  			c.Assert(vgc.lastEventCausets[0].GetMysqlJSON().String(), Equals, `{"123": 123}`)
   917  			chk.DeferredCauset(0).ReserveJSON(1)
   918  			j := new(json.BinaryJSON)
   919  			c.Assert(j.UnmarshalJSON([]byte(fmt.Sprintf(`{"%v":%v}`, 456, 456))), IsNil)
   920  			chk.DeferredCauset(0).AppendJSON(*j)
   921  			c.Assert(vgc.firstEventCausets[0].GetMysqlJSON().String(), Equals, `{"123": 123}`)
   922  			c.Assert(vgc.lastEventCausets[0].GetMysqlJSON().String(), Equals, `{"123": 123}`)
   923  		}
   924  	}
   925  }