github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/interlock/aggregate.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  	"bytes"
    18  	"context"
    19  	"fmt"
    20  	"sync"
    21  
    22  	"github.com/cznic/mathutil"
    23  	"github.com/whtcorpsinc/errors"
    24  	"github.com/whtcorpsinc/failpoint"
    25  	"github.com/whtcorpsinc/BerolinaSQL/allegrosql"
    26  	"github.com/whtcorpsinc/milevadb/interlock/aggfuncs"
    27  	"github.com/whtcorpsinc/milevadb/memex"
    28  	"github.com/whtcorpsinc/milevadb/stochastikctx"
    29  	"github.com/whtcorpsinc/milevadb/stochastikctx/stmtctx"
    30  	"github.com/whtcorpsinc/milevadb/types"
    31  	"github.com/whtcorpsinc/milevadb/types/json"
    32  	"github.com/whtcorpsinc/milevadb/soliton/chunk"
    33  	"github.com/whtcorpsinc/milevadb/soliton/codec"
    34  	"github.com/whtcorpsinc/milevadb/soliton/execdetails"
    35  	"github.com/whtcorpsinc/milevadb/soliton/logutil"
    36  	"github.com/whtcorpsinc/milevadb/soliton/memory"
    37  	"github.com/whtcorpsinc/milevadb/soliton/set"
    38  	"github.com/twmb/murmur3"
    39  	"go.uber.org/zap"
    40  )
    41  
    42  type aggPartialResultMapper map[string][]aggfuncs.PartialResult
    43  
    44  // baseHashAggWorker stores the common attributes of HashAggFinalWorker and HashAggPartialWorker.
    45  type baseHashAggWorker struct {
    46  	ctx          stochastikctx.Context
    47  	finishCh     <-chan struct{}
    48  	aggFuncs     []aggfuncs.AggFunc
    49  	maxChunkSize int
    50  }
    51  
    52  func newBaseHashAggWorker(ctx stochastikctx.Context, finishCh <-chan struct{}, aggFuncs []aggfuncs.AggFunc, maxChunkSize int) baseHashAggWorker {
    53  	return baseHashAggWorker{
    54  		ctx:          ctx,
    55  		finishCh:     finishCh,
    56  		aggFuncs:     aggFuncs,
    57  		maxChunkSize: maxChunkSize,
    58  	}
    59  }
    60  
    61  // HashAggPartialWorker indicates the partial workers of parallel hash agg execution,
    62  // the number of the worker can be set by `milevadb_hashagg_partial_concurrency`.
    63  type HashAggPartialWorker struct {
    64  	baseHashAggWorker
    65  
    66  	inputCh           chan *chunk.Chunk
    67  	outputChs         []chan *HashAggIntermData
    68  	globalOutputCh    chan *AfFinalResult
    69  	giveBackCh        chan<- *HashAggInput
    70  	partialResultsMap aggPartialResultMapper
    71  	groupByItems      []memex.Expression
    72  	groupKey          [][]byte
    73  	// chk stores the input data from child,
    74  	// and is reused by childInterDirc and partial worker.
    75  	chk        *chunk.Chunk
    76  	memTracker *memory.Tracker
    77  }
    78  
    79  // HashAggFinalWorker indicates the final workers of parallel hash agg execution,
    80  // the number of the worker can be set by `milevadb_hashagg_final_concurrency`.
    81  type HashAggFinalWorker struct {
    82  	baseHashAggWorker
    83  
    84  	rowBuffer           []types.Causet
    85  	mublockEvent          chunk.MutEvent
    86  	partialResultMap    aggPartialResultMapper
    87  	groupSet            set.StringSet
    88  	inputCh             chan *HashAggIntermData
    89  	outputCh            chan *AfFinalResult
    90  	finalResultHolderCh chan *chunk.Chunk
    91  	groupKeys           [][]byte
    92  }
    93  
    94  // AfFinalResult indicates aggregation functions final result.
    95  type AfFinalResult struct {
    96  	chk        *chunk.Chunk
    97  	err        error
    98  	giveBackCh chan *chunk.Chunk
    99  }
   100  
   101  // HashAggInterDirc deals with all the aggregate functions.
   102  // It is built from the Aggregate Causet. When Next() is called, it reads all the data from Src
   103  // and uFIDelates all the items in PartialAggFuncs.
   104  // The parallel execution flow is as the following graph shows:
   105  //
   106  //                            +-------------+
   107  //                            | Main Thread |
   108  //                            +------+------+
   109  //                                   ^
   110  //                                   |
   111  //                                   +
   112  //                              +-+-            +-+
   113  //                              | |    ......   | |  finalOutputCh
   114  //                              +++-            +-+
   115  //                               ^
   116  //                               |
   117  //                               +---------------+
   118  //                               |               |
   119  //                 +--------------+             +--------------+
   120  //                 | final worker |     ......  | final worker |
   121  //                 +------------+-+             +-+------------+
   122  //                              ^                 ^
   123  //                              |                 |
   124  //                             +-+  +-+  ......  +-+
   125  //                             | |  | |          | |
   126  //                             ...  ...          ...    partialOutputChs
   127  //                             | |  | |          | |
   128  //                             +++  +++          +++
   129  //                              ^    ^            ^
   130  //          +-+                 |    |            |
   131  //          | |        +--------o----+            |
   132  // inputCh  +-+        |        +-----------------+---+
   133  //          | |        |                              |
   134  //          ...    +---+------------+            +----+-----------+
   135  //          | |    | partial worker |   ......   | partial worker |
   136  //          +++    +--------------+-+            +-+--------------+
   137  //           |                     ^                ^
   138  //           |                     |                |
   139  //      +----v---------+          +++ +-+          +++
   140  //      | data fetcher | +------> | | | |  ......  | |   partialInputChs
   141  //      +--------------+          +-+ +-+          +-+
   142  type HashAggInterDirc struct {
   143  	baseInterlockingDirectorate
   144  
   145  	sc               *stmtctx.StatementContext
   146  	PartialAggFuncs  []aggfuncs.AggFunc
   147  	FinalAggFuncs    []aggfuncs.AggFunc
   148  	partialResultMap aggPartialResultMapper
   149  	groupSet         set.StringSet
   150  	groupKeys        []string
   151  	cursor4GroupKey  int
   152  	GroupByItems     []memex.Expression
   153  	groupKeyBuffer   [][]byte
   154  
   155  	finishCh         chan struct{}
   156  	finalOutputCh    chan *AfFinalResult
   157  	partialOutputChs []chan *HashAggIntermData
   158  	inputCh          chan *HashAggInput
   159  	partialInputChs  []chan *chunk.Chunk
   160  	partialWorkers   []HashAggPartialWorker
   161  	finalWorkers     []HashAggFinalWorker
   162  	defaultVal       *chunk.Chunk
   163  	childResult      *chunk.Chunk
   164  
   165  	// isChildReturnEmpty indicates whether the child interlock only returns an empty input.
   166  	isChildReturnEmpty bool
   167  	// After we support parallel execution for aggregation functions with distinct,
   168  	// we can remove this attribute.
   169  	isUnparallelInterDirc bool
   170  	prepared         bool
   171  	executed         bool
   172  
   173  	memTracker *memory.Tracker // track memory usage.
   174  }
   175  
   176  // HashAggInput indicates the input of hash agg exec.
   177  type HashAggInput struct {
   178  	chk *chunk.Chunk
   179  	// giveBackCh is bound with specific partial worker,
   180  	// it's used to reuse the `chk`,
   181  	// and tell the data-fetcher which partial worker it should send data to.
   182  	giveBackCh chan<- *chunk.Chunk
   183  }
   184  
   185  // HashAggIntermData indicates the intermediate data of aggregation execution.
   186  type HashAggIntermData struct {
   187  	groupKeys        []string
   188  	cursor           int
   189  	partialResultMap aggPartialResultMapper
   190  }
   191  
   192  // getPartialResultBatch fetches a batch of partial results from HashAggIntermData.
   193  func (d *HashAggIntermData) getPartialResultBatch(sc *stmtctx.StatementContext, prs [][]aggfuncs.PartialResult, aggFuncs []aggfuncs.AggFunc, maxChunkSize int) (_ [][]aggfuncs.PartialResult, groupKeys []string, reachEnd bool) {
   194  	keyStart := d.cursor
   195  	for ; d.cursor < len(d.groupKeys) && len(prs) < maxChunkSize; d.cursor++ {
   196  		prs = append(prs, d.partialResultMap[d.groupKeys[d.cursor]])
   197  	}
   198  	if d.cursor == len(d.groupKeys) {
   199  		reachEnd = true
   200  	}
   201  	return prs, d.groupKeys[keyStart:d.cursor], reachEnd
   202  }
   203  
   204  // Close implements the InterlockingDirectorate Close interface.
   205  func (e *HashAggInterDirc) Close() error {
   206  	if e.isUnparallelInterDirc {
   207  		e.memTracker.Consume(-e.childResult.MemoryUsage())
   208  		e.childResult = nil
   209  		e.groupSet = nil
   210  		e.partialResultMap = nil
   211  		return e.baseInterlockingDirectorate.Close()
   212  	}
   213  	// `Close` may be called after `Open` without calling `Next` in test.
   214  	if !e.prepared {
   215  		close(e.inputCh)
   216  		for _, ch := range e.partialOutputChs {
   217  			close(ch)
   218  		}
   219  		for _, ch := range e.partialInputChs {
   220  			close(ch)
   221  		}
   222  		close(e.finalOutputCh)
   223  	}
   224  	close(e.finishCh)
   225  	for _, ch := range e.partialOutputChs {
   226  		for range ch {
   227  		}
   228  	}
   229  	for _, ch := range e.partialInputChs {
   230  		for chk := range ch {
   231  			e.memTracker.Consume(-chk.MemoryUsage())
   232  		}
   233  	}
   234  	for range e.finalOutputCh {
   235  	}
   236  	e.executed = false
   237  
   238  	if e.runtimeStats != nil {
   239  		var partialConcurrency, finalConcurrency int
   240  		if e.isUnparallelInterDirc {
   241  			partialConcurrency = 0
   242  			finalConcurrency = 0
   243  		} else {
   244  			partialConcurrency = cap(e.partialWorkers)
   245  			finalConcurrency = cap(e.finalWorkers)
   246  		}
   247  		partialConcurrencyInfo := execdetails.NewConcurrencyInfo("PartialConcurrency", partialConcurrency)
   248  		finalConcurrencyInfo := execdetails.NewConcurrencyInfo("FinalConcurrency", finalConcurrency)
   249  		runtimeStats := &execdetails.RuntimeStatsWithConcurrencyInfo{}
   250  		runtimeStats.SetConcurrencyInfo(partialConcurrencyInfo, finalConcurrencyInfo)
   251  		e.ctx.GetStochastikVars().StmtCtx.RuntimeStatsDefCausl.RegisterStats(e.id, runtimeStats)
   252  	}
   253  	return e.baseInterlockingDirectorate.Close()
   254  }
   255  
   256  // Open implements the InterlockingDirectorate Open interface.
   257  func (e *HashAggInterDirc) Open(ctx context.Context) error {
   258  	if err := e.baseInterlockingDirectorate.Open(ctx); err != nil {
   259  		return err
   260  	}
   261  	e.prepared = false
   262  
   263  	e.memTracker = memory.NewTracker(e.id, -1)
   264  	e.memTracker.AttachTo(e.ctx.GetStochastikVars().StmtCtx.MemTracker)
   265  
   266  	if e.isUnparallelInterDirc {
   267  		e.initForUnparallelInterDirc()
   268  		return nil
   269  	}
   270  	e.initForParallelInterDirc(e.ctx)
   271  	return nil
   272  }
   273  
   274  func (e *HashAggInterDirc) initForUnparallelInterDirc() {
   275  	e.groupSet = set.NewStringSet()
   276  	e.partialResultMap = make(aggPartialResultMapper)
   277  	e.groupKeyBuffer = make([][]byte, 0, 8)
   278  	e.childResult = newFirstChunk(e.children[0])
   279  	e.memTracker.Consume(e.childResult.MemoryUsage())
   280  }
   281  
   282  func (e *HashAggInterDirc) initForParallelInterDirc(ctx stochastikctx.Context) {
   283  	stochastikVars := e.ctx.GetStochastikVars()
   284  	finalConcurrency := stochastikVars.HashAggFinalConcurrency()
   285  	partialConcurrency := stochastikVars.HashAggPartialConcurrency()
   286  	e.isChildReturnEmpty = true
   287  	e.finalOutputCh = make(chan *AfFinalResult, finalConcurrency)
   288  	e.inputCh = make(chan *HashAggInput, partialConcurrency)
   289  	e.finishCh = make(chan struct{}, 1)
   290  
   291  	e.partialInputChs = make([]chan *chunk.Chunk, partialConcurrency)
   292  	for i := range e.partialInputChs {
   293  		e.partialInputChs[i] = make(chan *chunk.Chunk, 1)
   294  	}
   295  	e.partialOutputChs = make([]chan *HashAggIntermData, finalConcurrency)
   296  	for i := range e.partialOutputChs {
   297  		e.partialOutputChs[i] = make(chan *HashAggIntermData, partialConcurrency)
   298  	}
   299  
   300  	e.partialWorkers = make([]HashAggPartialWorker, partialConcurrency)
   301  	e.finalWorkers = make([]HashAggFinalWorker, finalConcurrency)
   302  
   303  	// Init partial workers.
   304  	for i := 0; i < partialConcurrency; i++ {
   305  		w := HashAggPartialWorker{
   306  			baseHashAggWorker: newBaseHashAggWorker(e.ctx, e.finishCh, e.PartialAggFuncs, e.maxChunkSize),
   307  			inputCh:           e.partialInputChs[i],
   308  			outputChs:         e.partialOutputChs,
   309  			giveBackCh:        e.inputCh,
   310  			globalOutputCh:    e.finalOutputCh,
   311  			partialResultsMap: make(aggPartialResultMapper),
   312  			groupByItems:      e.GroupByItems,
   313  			chk:               newFirstChunk(e.children[0]),
   314  			groupKey:          make([][]byte, 0, 8),
   315  			memTracker:        e.memTracker,
   316  		}
   317  		e.memTracker.Consume(w.chk.MemoryUsage())
   318  		e.partialWorkers[i] = w
   319  
   320  		input := &HashAggInput{
   321  			chk:        newFirstChunk(e.children[0]),
   322  			giveBackCh: w.inputCh,
   323  		}
   324  		e.memTracker.Consume(input.chk.MemoryUsage())
   325  		e.inputCh <- input
   326  	}
   327  
   328  	// Init final workers.
   329  	for i := 0; i < finalConcurrency; i++ {
   330  		e.finalWorkers[i] = HashAggFinalWorker{
   331  			baseHashAggWorker:   newBaseHashAggWorker(e.ctx, e.finishCh, e.FinalAggFuncs, e.maxChunkSize),
   332  			partialResultMap:    make(aggPartialResultMapper),
   333  			groupSet:            set.NewStringSet(),
   334  			inputCh:             e.partialOutputChs[i],
   335  			outputCh:            e.finalOutputCh,
   336  			finalResultHolderCh: make(chan *chunk.Chunk, 1),
   337  			rowBuffer:           make([]types.Causet, 0, e.Schema().Len()),
   338  			mublockEvent:          chunk.MutEventFromTypes(retTypes(e)),
   339  			groupKeys:           make([][]byte, 0, 8),
   340  		}
   341  		e.finalWorkers[i].finalResultHolderCh <- newFirstChunk(e)
   342  	}
   343  }
   344  
   345  func (w *HashAggPartialWorker) getChildInput() bool {
   346  	select {
   347  	case <-w.finishCh:
   348  		return false
   349  	case chk, ok := <-w.inputCh:
   350  		if !ok {
   351  			return false
   352  		}
   353  		w.chk.SwapDeferredCausets(chk)
   354  		w.giveBackCh <- &HashAggInput{
   355  			chk:        chk,
   356  			giveBackCh: w.inputCh,
   357  		}
   358  	}
   359  	return true
   360  }
   361  
   362  func recoveryHashAgg(output chan *AfFinalResult, r interface{}) {
   363  	err := errors.Errorf("%v", r)
   364  	output <- &AfFinalResult{err: errors.Errorf("%v", r)}
   365  	logutil.BgLogger().Error("parallel hash aggregation panicked", zap.Error(err), zap.Stack("stack"))
   366  }
   367  
   368  func (w *HashAggPartialWorker) run(ctx stochastikctx.Context, waitGroup *sync.WaitGroup, finalConcurrency int) {
   369  	needShuffle, sc := false, ctx.GetStochastikVars().StmtCtx
   370  	defer func() {
   371  		if r := recover(); r != nil {
   372  			recoveryHashAgg(w.globalOutputCh, r)
   373  		}
   374  		if needShuffle {
   375  			w.shuffleIntermData(sc, finalConcurrency)
   376  		}
   377  		w.memTracker.Consume(-w.chk.MemoryUsage())
   378  		waitGroup.Done()
   379  	}()
   380  	for {
   381  		if !w.getChildInput() {
   382  			return
   383  		}
   384  		if err := w.uFIDelatePartialResult(ctx, sc, w.chk, len(w.partialResultsMap)); err != nil {
   385  			w.globalOutputCh <- &AfFinalResult{err: err}
   386  			return
   387  		}
   388  		// The intermData can be promised to be not empty if reaching here,
   389  		// so we set needShuffle to be true.
   390  		needShuffle = true
   391  	}
   392  }
   393  
   394  func (w *HashAggPartialWorker) uFIDelatePartialResult(ctx stochastikctx.Context, sc *stmtctx.StatementContext, chk *chunk.Chunk, finalConcurrency int) (err error) {
   395  	w.groupKey, err = getGroupKey(w.ctx, chk, w.groupKey, w.groupByItems)
   396  	if err != nil {
   397  		return err
   398  	}
   399  
   400  	partialResults := w.getPartialResult(sc, w.groupKey, w.partialResultsMap)
   401  	numEvents := chk.NumEvents()
   402  	rows := make([]chunk.Event, 1)
   403  	for i := 0; i < numEvents; i++ {
   404  		for j, af := range w.aggFuncs {
   405  			rows[0] = chk.GetEvent(i)
   406  			if _, err := af.UFIDelatePartialResult(ctx, rows, partialResults[i][j]); err != nil {
   407  				return err
   408  			}
   409  		}
   410  	}
   411  	return nil
   412  }
   413  
   414  // shuffleIntermData shuffles the intermediate data of partial workers to corresponded final workers.
   415  // We only support parallel execution for single-machine, so process of encode and decode can be skipped.
   416  func (w *HashAggPartialWorker) shuffleIntermData(sc *stmtctx.StatementContext, finalConcurrency int) {
   417  	groupKeysSlice := make([][]string, finalConcurrency)
   418  	for groupKey := range w.partialResultsMap {
   419  		finalWorkerIdx := int(murmur3.Sum32([]byte(groupKey))) % finalConcurrency
   420  		if groupKeysSlice[finalWorkerIdx] == nil {
   421  			groupKeysSlice[finalWorkerIdx] = make([]string, 0, len(w.partialResultsMap)/finalConcurrency)
   422  		}
   423  		groupKeysSlice[finalWorkerIdx] = append(groupKeysSlice[finalWorkerIdx], groupKey)
   424  	}
   425  
   426  	for i := range groupKeysSlice {
   427  		if groupKeysSlice[i] == nil {
   428  			continue
   429  		}
   430  		w.outputChs[i] <- &HashAggIntermData{
   431  			groupKeys:        groupKeysSlice[i],
   432  			partialResultMap: w.partialResultsMap,
   433  		}
   434  	}
   435  }
   436  
   437  // getGroupKey evaluates the group items and args of aggregate functions.
   438  func getGroupKey(ctx stochastikctx.Context, input *chunk.Chunk, groupKey [][]byte, groupByItems []memex.Expression) ([][]byte, error) {
   439  	numEvents := input.NumEvents()
   440  	avlGroupKeyLen := mathutil.Min(len(groupKey), numEvents)
   441  	for i := 0; i < avlGroupKeyLen; i++ {
   442  		groupKey[i] = groupKey[i][:0]
   443  	}
   444  	for i := avlGroupKeyLen; i < numEvents; i++ {
   445  		groupKey = append(groupKey, make([]byte, 0, 10*len(groupByItems)))
   446  	}
   447  
   448  	for _, item := range groupByItems {
   449  		tp := item.GetType()
   450  		buf, err := memex.GetDeferredCauset(tp.EvalType(), numEvents)
   451  		if err != nil {
   452  			return nil, err
   453  		}
   454  
   455  		if err := memex.EvalExpr(ctx, item, input, buf); err != nil {
   456  			memex.PutDeferredCauset(buf)
   457  			return nil, err
   458  		}
   459  		// This check is used to avoid error during the execution of `EncodeDecimal`.
   460  		if item.GetType().Tp == allegrosql.TypeNewDecimal {
   461  			newTp := *tp
   462  			newTp.Flen = 0
   463  			tp = &newTp
   464  		}
   465  		groupKey, err = codec.HashGroupKey(ctx.GetStochastikVars().StmtCtx, input.NumEvents(), buf, groupKey, tp)
   466  		if err != nil {
   467  			memex.PutDeferredCauset(buf)
   468  			return nil, err
   469  		}
   470  		memex.PutDeferredCauset(buf)
   471  	}
   472  	return groupKey, nil
   473  }
   474  
   475  func (w baseHashAggWorker) getPartialResult(sc *stmtctx.StatementContext, groupKey [][]byte, mapper aggPartialResultMapper) [][]aggfuncs.PartialResult {
   476  	n := len(groupKey)
   477  	partialResults := make([][]aggfuncs.PartialResult, n)
   478  	for i := 0; i < n; i++ {
   479  		var ok bool
   480  		if partialResults[i], ok = mapper[string(groupKey[i])]; ok {
   481  			continue
   482  		}
   483  		for _, af := range w.aggFuncs {
   484  			partialResult, _ := af.AllocPartialResult()
   485  			partialResults[i] = append(partialResults[i], partialResult)
   486  		}
   487  		mapper[string(groupKey[i])] = partialResults[i]
   488  	}
   489  	return partialResults
   490  }
   491  
   492  func (w *HashAggFinalWorker) getPartialInput() (input *HashAggIntermData, ok bool) {
   493  	select {
   494  	case <-w.finishCh:
   495  		return nil, false
   496  	case input, ok = <-w.inputCh:
   497  		if !ok {
   498  			return nil, false
   499  		}
   500  	}
   501  	return
   502  }
   503  
   504  func (w *HashAggFinalWorker) consumeIntermData(sctx stochastikctx.Context) (err error) {
   505  	var (
   506  		input            *HashAggIntermData
   507  		ok               bool
   508  		intermDataBuffer [][]aggfuncs.PartialResult
   509  		groupKeys        []string
   510  		sc               = sctx.GetStochastikVars().StmtCtx
   511  	)
   512  	for {
   513  		if input, ok = w.getPartialInput(); !ok {
   514  			return nil
   515  		}
   516  		if intermDataBuffer == nil {
   517  			intermDataBuffer = make([][]aggfuncs.PartialResult, 0, w.maxChunkSize)
   518  		}
   519  		// Consume input in batches, size of every batch is less than w.maxChunkSize.
   520  		for reachEnd := false; !reachEnd; {
   521  			intermDataBuffer, groupKeys, reachEnd = input.getPartialResultBatch(sc, intermDataBuffer[:0], w.aggFuncs, w.maxChunkSize)
   522  			groupKeysLen := len(groupKeys)
   523  			w.groupKeys = w.groupKeys[:0]
   524  			for i := 0; i < groupKeysLen; i++ {
   525  				w.groupKeys = append(w.groupKeys, []byte(groupKeys[i]))
   526  			}
   527  			finalPartialResults := w.getPartialResult(sc, w.groupKeys, w.partialResultMap)
   528  			for i, groupKey := range groupKeys {
   529  				if !w.groupSet.Exist(groupKey) {
   530  					w.groupSet.Insert(groupKey)
   531  				}
   532  				prs := intermDataBuffer[i]
   533  				for j, af := range w.aggFuncs {
   534  					if _, err = af.MergePartialResult(sctx, prs[j], finalPartialResults[i][j]); err != nil {
   535  						return err
   536  					}
   537  				}
   538  			}
   539  		}
   540  	}
   541  }
   542  
   543  func (w *HashAggFinalWorker) getFinalResult(sctx stochastikctx.Context) {
   544  	result, finished := w.receiveFinalResultHolder()
   545  	if finished {
   546  		return
   547  	}
   548  	w.groupKeys = w.groupKeys[:0]
   549  	for groupKey := range w.groupSet {
   550  		w.groupKeys = append(w.groupKeys, []byte(groupKey))
   551  	}
   552  	partialResults := w.getPartialResult(sctx.GetStochastikVars().StmtCtx, w.groupKeys, w.partialResultMap)
   553  	for i := 0; i < len(w.groupSet); i++ {
   554  		for j, af := range w.aggFuncs {
   555  			if err := af.AppendFinalResult2Chunk(sctx, partialResults[i][j], result); err != nil {
   556  				logutil.BgLogger().Error("HashAggFinalWorker failed to append final result to Chunk", zap.Error(err))
   557  			}
   558  		}
   559  		if len(w.aggFuncs) == 0 {
   560  			result.SetNumVirtualEvents(result.NumEvents() + 1)
   561  		}
   562  		if result.IsFull() {
   563  			w.outputCh <- &AfFinalResult{chk: result, giveBackCh: w.finalResultHolderCh}
   564  			result, finished = w.receiveFinalResultHolder()
   565  			if finished {
   566  				return
   567  			}
   568  		}
   569  	}
   570  	w.outputCh <- &AfFinalResult{chk: result, giveBackCh: w.finalResultHolderCh}
   571  }
   572  
   573  func (w *HashAggFinalWorker) receiveFinalResultHolder() (*chunk.Chunk, bool) {
   574  	select {
   575  	case <-w.finishCh:
   576  		return nil, true
   577  	case result, ok := <-w.finalResultHolderCh:
   578  		return result, !ok
   579  	}
   580  }
   581  
   582  func (w *HashAggFinalWorker) run(ctx stochastikctx.Context, waitGroup *sync.WaitGroup) {
   583  	defer func() {
   584  		if r := recover(); r != nil {
   585  			recoveryHashAgg(w.outputCh, r)
   586  		}
   587  		waitGroup.Done()
   588  	}()
   589  	if err := w.consumeIntermData(ctx); err != nil {
   590  		w.outputCh <- &AfFinalResult{err: err}
   591  	}
   592  	w.getFinalResult(ctx)
   593  }
   594  
   595  // Next implements the InterlockingDirectorate Next interface.
   596  func (e *HashAggInterDirc) Next(ctx context.Context, req *chunk.Chunk) error {
   597  	req.Reset()
   598  	if e.isUnparallelInterDirc {
   599  		return e.unparallelInterDirc(ctx, req)
   600  	}
   601  	return e.parallelInterDirc(ctx, req)
   602  }
   603  
   604  func (e *HashAggInterDirc) fetchChildData(ctx context.Context) {
   605  	var (
   606  		input *HashAggInput
   607  		chk   *chunk.Chunk
   608  		ok    bool
   609  		err   error
   610  	)
   611  	defer func() {
   612  		if r := recover(); r != nil {
   613  			recoveryHashAgg(e.finalOutputCh, r)
   614  		}
   615  		for i := range e.partialInputChs {
   616  			close(e.partialInputChs[i])
   617  		}
   618  	}()
   619  	for {
   620  		select {
   621  		case <-e.finishCh:
   622  			return
   623  		case input, ok = <-e.inputCh:
   624  			if !ok {
   625  				return
   626  			}
   627  			chk = input.chk
   628  		}
   629  		mSize := chk.MemoryUsage()
   630  		err = Next(ctx, e.children[0], chk)
   631  		if err != nil {
   632  			e.finalOutputCh <- &AfFinalResult{err: err}
   633  			e.memTracker.Consume(-mSize)
   634  			return
   635  		}
   636  		if chk.NumEvents() == 0 {
   637  			e.memTracker.Consume(-mSize)
   638  			return
   639  		}
   640  		e.memTracker.Consume(chk.MemoryUsage() - mSize)
   641  		input.giveBackCh <- chk
   642  	}
   643  }
   644  
   645  func (e *HashAggInterDirc) waitPartialWorkerAndCloseOutputChs(waitGroup *sync.WaitGroup) {
   646  	waitGroup.Wait()
   647  	close(e.inputCh)
   648  	for input := range e.inputCh {
   649  		e.memTracker.Consume(-input.chk.MemoryUsage())
   650  	}
   651  	for _, ch := range e.partialOutputChs {
   652  		close(ch)
   653  	}
   654  }
   655  
   656  func (e *HashAggInterDirc) waitFinalWorkerAndCloseFinalOutput(waitGroup *sync.WaitGroup) {
   657  	waitGroup.Wait()
   658  	close(e.finalOutputCh)
   659  }
   660  
   661  func (e *HashAggInterDirc) prepare4ParallelInterDirc(ctx context.Context) {
   662  	go e.fetchChildData(ctx)
   663  
   664  	partialWorkerWaitGroup := &sync.WaitGroup{}
   665  	partialWorkerWaitGroup.Add(len(e.partialWorkers))
   666  	for i := range e.partialWorkers {
   667  		go e.partialWorkers[i].run(e.ctx, partialWorkerWaitGroup, len(e.finalWorkers))
   668  	}
   669  	go e.waitPartialWorkerAndCloseOutputChs(partialWorkerWaitGroup)
   670  
   671  	finalWorkerWaitGroup := &sync.WaitGroup{}
   672  	finalWorkerWaitGroup.Add(len(e.finalWorkers))
   673  	for i := range e.finalWorkers {
   674  		go e.finalWorkers[i].run(e.ctx, finalWorkerWaitGroup)
   675  	}
   676  	go e.waitFinalWorkerAndCloseFinalOutput(finalWorkerWaitGroup)
   677  }
   678  
   679  // HashAggInterDirc employs one input reader, M partial workers and N final workers to execute parallelly.
   680  // The parallel execution flow is:
   681  // 1. input reader reads data from child interlock and send them to partial workers.
   682  // 2. partial worker receives the input data, uFIDelates the partial results, and shuffle the partial results to the final workers.
   683  // 3. final worker receives partial results from all the partial workers, evaluates the final results and sends the final results to the main thread.
   684  func (e *HashAggInterDirc) parallelInterDirc(ctx context.Context, chk *chunk.Chunk) error {
   685  	if !e.prepared {
   686  		e.prepare4ParallelInterDirc(ctx)
   687  		e.prepared = true
   688  	}
   689  
   690  	failpoint.Inject("parallelHashAggError", func(val failpoint.Value) {
   691  		if val.(bool) {
   692  			failpoint.Return(errors.New("HashAggInterDirc.parallelInterDirc error"))
   693  		}
   694  	})
   695  
   696  	if e.executed {
   697  		return nil
   698  	}
   699  	for {
   700  		result, ok := <-e.finalOutputCh
   701  		if !ok {
   702  			e.executed = true
   703  			if e.isChildReturnEmpty && e.defaultVal != nil {
   704  				chk.Append(e.defaultVal, 0, 1)
   705  			}
   706  			return nil
   707  		}
   708  		if result.err != nil {
   709  			return result.err
   710  		}
   711  		chk.SwapDeferredCausets(result.chk)
   712  		result.chk.Reset()
   713  		result.giveBackCh <- result.chk
   714  		if chk.NumEvents() > 0 {
   715  			e.isChildReturnEmpty = false
   716  			return nil
   717  		}
   718  	}
   719  }
   720  
   721  // unparallelInterDirc executes hash aggregation algorithm in single thread.
   722  func (e *HashAggInterDirc) unparallelInterDirc(ctx context.Context, chk *chunk.Chunk) error {
   723  	// In this stage we consider all data from src as a single group.
   724  	if !e.prepared {
   725  		err := e.execute(ctx)
   726  		if err != nil {
   727  			return err
   728  		}
   729  		if (len(e.groupSet) == 0) && len(e.GroupByItems) == 0 {
   730  			// If no groupby and no data, we should add an empty group.
   731  			// For example:
   732  			// "select count(c) from t;" should return one event [0]
   733  			// "select count(c) from t group by c1;" should return empty result set.
   734  			e.groupSet.Insert("")
   735  			e.groupKeys = append(e.groupKeys, "")
   736  		}
   737  		e.prepared = true
   738  	}
   739  	chk.Reset()
   740  
   741  	// Since we return e.maxChunkSize rows every time, so we should not traverse
   742  	// `groupSet` because of its randomness.
   743  	for ; e.cursor4GroupKey < len(e.groupKeys); e.cursor4GroupKey++ {
   744  		partialResults := e.getPartialResults(e.groupKeys[e.cursor4GroupKey])
   745  		if len(e.PartialAggFuncs) == 0 {
   746  			chk.SetNumVirtualEvents(chk.NumEvents() + 1)
   747  		}
   748  		for i, af := range e.PartialAggFuncs {
   749  			if err := af.AppendFinalResult2Chunk(e.ctx, partialResults[i], chk); err != nil {
   750  				return err
   751  			}
   752  		}
   753  		if chk.IsFull() {
   754  			e.cursor4GroupKey++
   755  			return nil
   756  		}
   757  	}
   758  	return nil
   759  }
   760  
   761  // execute fetches Chunks from src and uFIDelate each aggregate function for each event in Chunk.
   762  func (e *HashAggInterDirc) execute(ctx context.Context) (err error) {
   763  	for {
   764  		mSize := e.childResult.MemoryUsage()
   765  		err := Next(ctx, e.children[0], e.childResult)
   766  		e.memTracker.Consume(e.childResult.MemoryUsage() - mSize)
   767  		if err != nil {
   768  			return err
   769  		}
   770  
   771  		failpoint.Inject("unparallelHashAggError", func(val failpoint.Value) {
   772  			if val.(bool) {
   773  				failpoint.Return(errors.New("HashAggInterDirc.unparallelInterDirc error"))
   774  			}
   775  		})
   776  
   777  		// no more data.
   778  		if e.childResult.NumEvents() == 0 {
   779  			return nil
   780  		}
   781  
   782  		e.groupKeyBuffer, err = getGroupKey(e.ctx, e.childResult, e.groupKeyBuffer, e.GroupByItems)
   783  		if err != nil {
   784  			return err
   785  		}
   786  
   787  		for j := 0; j < e.childResult.NumEvents(); j++ {
   788  			groupKey := string(e.groupKeyBuffer[j]) // do memory copy here, because e.groupKeyBuffer may be reused.
   789  			if !e.groupSet.Exist(groupKey) {
   790  				e.groupSet.Insert(groupKey)
   791  				e.groupKeys = append(e.groupKeys, groupKey)
   792  			}
   793  			partialResults := e.getPartialResults(groupKey)
   794  			for i, af := range e.PartialAggFuncs {
   795  				memDelta, err := af.UFIDelatePartialResult(e.ctx, []chunk.Event{e.childResult.GetEvent(j)}, partialResults[i])
   796  				if err != nil {
   797  					return err
   798  				}
   799  				e.memTracker.Consume(memDelta)
   800  			}
   801  		}
   802  	}
   803  }
   804  
   805  func (e *HashAggInterDirc) getPartialResults(groupKey string) []aggfuncs.PartialResult {
   806  	partialResults, ok := e.partialResultMap[groupKey]
   807  	if !ok {
   808  		partialResults = make([]aggfuncs.PartialResult, 0, len(e.PartialAggFuncs))
   809  		for _, af := range e.PartialAggFuncs {
   810  			partialResult, memDelta := af.AllocPartialResult()
   811  			partialResults = append(partialResults, partialResult)
   812  			e.memTracker.Consume(memDelta)
   813  		}
   814  		e.partialResultMap[groupKey] = partialResults
   815  	}
   816  	return partialResults
   817  }
   818  
   819  // StreamAggInterDirc deals with all the aggregate functions.
   820  // It assumes all the input data is sorted by group by key.
   821  // When Next() is called, it will return a result for the same group.
   822  type StreamAggInterDirc struct {
   823  	baseInterlockingDirectorate
   824  
   825  	executed bool
   826  	// isChildReturnEmpty indicates whether the child interlock only returns an empty input.
   827  	isChildReturnEmpty bool
   828  	defaultVal         *chunk.Chunk
   829  	groupChecker       *vecGroupChecker
   830  	inputIter          *chunk.Iterator4Chunk
   831  	inputEvent           chunk.Event
   832  	aggFuncs           []aggfuncs.AggFunc
   833  	partialResults     []aggfuncs.PartialResult
   834  	groupEvents          []chunk.Event
   835  	childResult        *chunk.Chunk
   836  
   837  	memTracker *memory.Tracker // track memory usage.
   838  }
   839  
   840  // Open implements the InterlockingDirectorate Open interface.
   841  func (e *StreamAggInterDirc) Open(ctx context.Context) error {
   842  	if err := e.baseInterlockingDirectorate.Open(ctx); err != nil {
   843  		return err
   844  	}
   845  	e.childResult = newFirstChunk(e.children[0])
   846  	e.executed = false
   847  	e.isChildReturnEmpty = true
   848  	e.inputIter = chunk.NewIterator4Chunk(e.childResult)
   849  	e.inputEvent = e.inputIter.End()
   850  
   851  	e.partialResults = make([]aggfuncs.PartialResult, 0, len(e.aggFuncs))
   852  	for _, aggFunc := range e.aggFuncs {
   853  		partialResult, memDelta := aggFunc.AllocPartialResult()
   854  		e.partialResults = append(e.partialResults, partialResult)
   855  		e.memTracker.Consume(memDelta)
   856  	}
   857  
   858  	// bytesLimit <= 0 means no limit, for now we just track the memory footprint
   859  	e.memTracker = memory.NewTracker(e.id, -1)
   860  	e.memTracker.AttachTo(e.ctx.GetStochastikVars().StmtCtx.MemTracker)
   861  	e.memTracker.Consume(e.childResult.MemoryUsage())
   862  	return nil
   863  }
   864  
   865  // Close implements the InterlockingDirectorate Close interface.
   866  func (e *StreamAggInterDirc) Close() error {
   867  	e.memTracker.Consume(-e.childResult.MemoryUsage())
   868  	e.childResult = nil
   869  	e.groupChecker.reset()
   870  	return e.baseInterlockingDirectorate.Close()
   871  }
   872  
   873  // Next implements the InterlockingDirectorate Next interface.
   874  func (e *StreamAggInterDirc) Next(ctx context.Context, req *chunk.Chunk) (err error) {
   875  	req.Reset()
   876  	for !e.executed && !req.IsFull() {
   877  		err = e.consumeOneGroup(ctx, req)
   878  		if err != nil {
   879  			e.executed = true
   880  			return err
   881  		}
   882  	}
   883  	return nil
   884  }
   885  
   886  func (e *StreamAggInterDirc) consumeOneGroup(ctx context.Context, chk *chunk.Chunk) (err error) {
   887  	if e.groupChecker.isExhausted() {
   888  		if err = e.consumeCurGroupEventsAndFetchChild(ctx, chk); err != nil {
   889  			return err
   890  		}
   891  		if !e.executed {
   892  			_, err := e.groupChecker.splitIntoGroups(e.childResult)
   893  			if err != nil {
   894  				return err
   895  			}
   896  		} else {
   897  			return nil
   898  		}
   899  	}
   900  	begin, end := e.groupChecker.getNextGroup()
   901  	for i := begin; i < end; i++ {
   902  		e.groupEvents = append(e.groupEvents, e.childResult.GetEvent(i))
   903  	}
   904  
   905  	for meetLastGroup := end == e.childResult.NumEvents(); meetLastGroup; {
   906  		meetLastGroup = false
   907  		if err = e.consumeCurGroupEventsAndFetchChild(ctx, chk); err != nil || e.executed {
   908  			return err
   909  		}
   910  
   911  		isFirstGroupSameAsPrev, err := e.groupChecker.splitIntoGroups(e.childResult)
   912  		if err != nil {
   913  			return err
   914  		}
   915  
   916  		if isFirstGroupSameAsPrev {
   917  			begin, end = e.groupChecker.getNextGroup()
   918  			for i := begin; i < end; i++ {
   919  				e.groupEvents = append(e.groupEvents, e.childResult.GetEvent(i))
   920  			}
   921  			meetLastGroup = end == e.childResult.NumEvents()
   922  		}
   923  	}
   924  
   925  	err = e.consumeGroupEvents()
   926  	if err != nil {
   927  		return err
   928  	}
   929  
   930  	return e.appendResult2Chunk(chk)
   931  }
   932  
   933  func (e *StreamAggInterDirc) consumeGroupEvents() error {
   934  	if len(e.groupEvents) == 0 {
   935  		return nil
   936  	}
   937  
   938  	for i, aggFunc := range e.aggFuncs {
   939  		memDelta, err := aggFunc.UFIDelatePartialResult(e.ctx, e.groupEvents, e.partialResults[i])
   940  		if err != nil {
   941  			return err
   942  		}
   943  		e.memTracker.Consume(memDelta)
   944  	}
   945  	e.groupEvents = e.groupEvents[:0]
   946  	return nil
   947  }
   948  
   949  func (e *StreamAggInterDirc) consumeCurGroupEventsAndFetchChild(ctx context.Context, chk *chunk.Chunk) (err error) {
   950  	// Before fetching a new batch of input, we should consume the last group.
   951  	err = e.consumeGroupEvents()
   952  	if err != nil {
   953  		return err
   954  	}
   955  
   956  	mSize := e.childResult.MemoryUsage()
   957  	err = Next(ctx, e.children[0], e.childResult)
   958  	e.memTracker.Consume(e.childResult.MemoryUsage() - mSize)
   959  	if err != nil {
   960  		return err
   961  	}
   962  
   963  	// No more data.
   964  	if e.childResult.NumEvents() == 0 {
   965  		if !e.isChildReturnEmpty {
   966  			err = e.appendResult2Chunk(chk)
   967  		} else if e.defaultVal != nil {
   968  			chk.Append(e.defaultVal, 0, 1)
   969  		}
   970  		e.executed = true
   971  		return err
   972  	}
   973  	// Reach here, "e.childrenResults[0].NumEvents() > 0" is guaranteed.
   974  	e.isChildReturnEmpty = false
   975  	e.inputEvent = e.inputIter.Begin()
   976  	return nil
   977  }
   978  
   979  // appendResult2Chunk appends result of all the aggregation functions to the
   980  // result chunk, and reset the evaluation context for each aggregation.
   981  func (e *StreamAggInterDirc) appendResult2Chunk(chk *chunk.Chunk) error {
   982  	for i, aggFunc := range e.aggFuncs {
   983  		err := aggFunc.AppendFinalResult2Chunk(e.ctx, e.partialResults[i], chk)
   984  		if err != nil {
   985  			return err
   986  		}
   987  		aggFunc.ResetPartialResult(e.partialResults[i])
   988  	}
   989  	if len(e.aggFuncs) == 0 {
   990  		chk.SetNumVirtualEvents(chk.NumEvents() + 1)
   991  	}
   992  	return nil
   993  }
   994  
   995  // vecGroupChecker is used to split a given chunk according to the `group by` memex in a vectorized manner
   996  // It is usually used for streamAgg
   997  type vecGroupChecker struct {
   998  	ctx          stochastikctx.Context
   999  	GroupByItems []memex.Expression
  1000  
  1001  	// groupOffset holds the offset of the last event in each group of the current chunk
  1002  	groupOffset []int
  1003  	// groupCount is the count of groups in the current chunk
  1004  	groupCount int
  1005  	// nextGroupID records the group id of the next group to be consumed
  1006  	nextGroupID int
  1007  
  1008  	// lastGroupKeyOfPrevChk is the groupKey of the last group of the previous chunk
  1009  	lastGroupKeyOfPrevChk []byte
  1010  	// firstGroupKey and lastGroupKey are used to causetstore the groupKey of the first and last group of the current chunk
  1011  	firstGroupKey []byte
  1012  	lastGroupKey  []byte
  1013  
  1014  	// firstEventCausets and lastEventCausets causetstore the results of the memex evaluation for the first and last rows of the current chunk in causet
  1015  	// They are used to encode to get firstGroupKey and lastGroupKey
  1016  	firstEventCausets []types.Causet
  1017  	lastEventCausets  []types.Causet
  1018  
  1019  	// sameGroup is used to check whether the current event belongs to the same group as the previous event
  1020  	sameGroup []bool
  1021  
  1022  	// set these functions for testing
  1023  	allocateBuffer func(evalType types.EvalType, capacity int) (*chunk.DeferredCauset, error)
  1024  	releaseBuffer  func(buf *chunk.DeferredCauset)
  1025  }
  1026  
  1027  func newVecGroupChecker(ctx stochastikctx.Context, items []memex.Expression) *vecGroupChecker {
  1028  	return &vecGroupChecker{
  1029  		ctx:          ctx,
  1030  		GroupByItems: items,
  1031  		groupCount:   0,
  1032  		nextGroupID:  0,
  1033  		sameGroup:    make([]bool, 1024),
  1034  	}
  1035  }
  1036  
  1037  // splitIntoGroups splits a chunk into multiple groups which the event in the same group have the same groupKey
  1038  // `isFirstGroupSameAsPrev` indicates whether the groupKey of the first group of the newly passed chunk is equal to the groupKey of the last group left before
  1039  // TODO: Since all the group by items are only a defCausumn reference, guaranteed by building projection below aggregation, we can directly compare data in a chunk.
  1040  func (e *vecGroupChecker) splitIntoGroups(chk *chunk.Chunk) (isFirstGroupSameAsPrev bool, err error) {
  1041  	// The numEvents can not be zero. `fetchChild` is called before `splitIntoGroups` is called.
  1042  	// if numEvents == 0, it will be returned in `fetchChild`. See `fetchChild` for more details.
  1043  	numEvents := chk.NumEvents()
  1044  
  1045  	e.reset()
  1046  	e.nextGroupID = 0
  1047  	if len(e.GroupByItems) == 0 {
  1048  		e.groupOffset = append(e.groupOffset, numEvents)
  1049  		e.groupCount = 1
  1050  		return true, nil
  1051  	}
  1052  
  1053  	for _, item := range e.GroupByItems {
  1054  		err = e.getFirstAndLastEventCauset(item, chk, numEvents)
  1055  		if err != nil {
  1056  			return false, err
  1057  		}
  1058  	}
  1059  	e.firstGroupKey, err = codec.EncodeValue(e.ctx.GetStochastikVars().StmtCtx, e.firstGroupKey, e.firstEventCausets...)
  1060  	if err != nil {
  1061  		return false, err
  1062  	}
  1063  
  1064  	e.lastGroupKey, err = codec.EncodeValue(e.ctx.GetStochastikVars().StmtCtx, e.lastGroupKey, e.lastEventCausets...)
  1065  	if err != nil {
  1066  		return false, err
  1067  	}
  1068  
  1069  	if len(e.lastGroupKeyOfPrevChk) == 0 {
  1070  		isFirstGroupSameAsPrev = false
  1071  	} else {
  1072  		if bytes.Equal(e.lastGroupKeyOfPrevChk, e.firstGroupKey) {
  1073  			isFirstGroupSameAsPrev = true
  1074  		} else {
  1075  			isFirstGroupSameAsPrev = false
  1076  		}
  1077  	}
  1078  
  1079  	if length := len(e.lastGroupKey); len(e.lastGroupKeyOfPrevChk) >= length {
  1080  		e.lastGroupKeyOfPrevChk = e.lastGroupKeyOfPrevChk[:length]
  1081  	} else {
  1082  		e.lastGroupKeyOfPrevChk = make([]byte, length)
  1083  	}
  1084  	copy(e.lastGroupKeyOfPrevChk, e.lastGroupKey)
  1085  
  1086  	if bytes.Equal(e.firstGroupKey, e.lastGroupKey) {
  1087  		e.groupOffset = append(e.groupOffset, numEvents)
  1088  		e.groupCount = 1
  1089  		return isFirstGroupSameAsPrev, nil
  1090  	}
  1091  
  1092  	if cap(e.sameGroup) < numEvents {
  1093  		e.sameGroup = make([]bool, 0, numEvents)
  1094  	}
  1095  	e.sameGroup = append(e.sameGroup, false)
  1096  	for i := 1; i < numEvents; i++ {
  1097  		e.sameGroup = append(e.sameGroup, true)
  1098  	}
  1099  
  1100  	for _, item := range e.GroupByItems {
  1101  		err = e.evalGroupItemsAndResolveGroups(item, chk, numEvents)
  1102  		if err != nil {
  1103  			return false, err
  1104  		}
  1105  	}
  1106  
  1107  	for i := 1; i < numEvents; i++ {
  1108  		if !e.sameGroup[i] {
  1109  			e.groupOffset = append(e.groupOffset, i)
  1110  		}
  1111  	}
  1112  	e.groupOffset = append(e.groupOffset, numEvents)
  1113  	e.groupCount = len(e.groupOffset)
  1114  	return isFirstGroupSameAsPrev, nil
  1115  }
  1116  
  1117  func (e *vecGroupChecker) getFirstAndLastEventCauset(item memex.Expression, chk *chunk.Chunk, numEvents int) (err error) {
  1118  	var firstEventCauset, lastEventCauset types.Causet
  1119  	tp := item.GetType()
  1120  	eType := tp.EvalType()
  1121  	switch eType {
  1122  	case types.ETInt:
  1123  		firstEventVal, firstEventIsNull, err := item.EvalInt(e.ctx, chk.GetEvent(0))
  1124  		if err != nil {
  1125  			return err
  1126  		}
  1127  		lastEventVal, lastEventIsNull, err := item.EvalInt(e.ctx, chk.GetEvent(numEvents-1))
  1128  		if err != nil {
  1129  			return err
  1130  		}
  1131  		if !firstEventIsNull {
  1132  			firstEventCauset.SetInt64(firstEventVal)
  1133  		} else {
  1134  			firstEventCauset.SetNull()
  1135  		}
  1136  		if !lastEventIsNull {
  1137  			lastEventCauset.SetInt64(lastEventVal)
  1138  		} else {
  1139  			lastEventCauset.SetNull()
  1140  		}
  1141  	case types.ETReal:
  1142  		firstEventVal, firstEventIsNull, err := item.EvalReal(e.ctx, chk.GetEvent(0))
  1143  		if err != nil {
  1144  			return err
  1145  		}
  1146  		lastEventVal, lastEventIsNull, err := item.EvalReal(e.ctx, chk.GetEvent(numEvents-1))
  1147  		if err != nil {
  1148  			return err
  1149  		}
  1150  		if !firstEventIsNull {
  1151  			firstEventCauset.SetFloat64(firstEventVal)
  1152  		} else {
  1153  			firstEventCauset.SetNull()
  1154  		}
  1155  		if !lastEventIsNull {
  1156  			lastEventCauset.SetFloat64(lastEventVal)
  1157  		} else {
  1158  			lastEventCauset.SetNull()
  1159  		}
  1160  	case types.ETDecimal:
  1161  		firstEventVal, firstEventIsNull, err := item.EvalDecimal(e.ctx, chk.GetEvent(0))
  1162  		if err != nil {
  1163  			return err
  1164  		}
  1165  		lastEventVal, lastEventIsNull, err := item.EvalDecimal(e.ctx, chk.GetEvent(numEvents-1))
  1166  		if err != nil {
  1167  			return err
  1168  		}
  1169  		if !firstEventIsNull {
  1170  			// make a copy to avoid DATA RACE
  1171  			firstCauset := types.MyDecimal{}
  1172  			err := firstCauset.FromString(firstEventVal.ToString())
  1173  			if err != nil {
  1174  				return err
  1175  			}
  1176  			firstEventCauset.SetMysqlDecimal(&firstCauset)
  1177  		} else {
  1178  			firstEventCauset.SetNull()
  1179  		}
  1180  		if !lastEventIsNull {
  1181  			// make a copy to avoid DATA RACE
  1182  			lastCauset := types.MyDecimal{}
  1183  			err := lastCauset.FromString(lastEventVal.ToString())
  1184  			if err != nil {
  1185  				return err
  1186  			}
  1187  			lastEventCauset.SetMysqlDecimal(&lastCauset)
  1188  		} else {
  1189  			lastEventCauset.SetNull()
  1190  		}
  1191  	case types.ETDatetime, types.ETTimestamp:
  1192  		firstEventVal, firstEventIsNull, err := item.EvalTime(e.ctx, chk.GetEvent(0))
  1193  		if err != nil {
  1194  			return err
  1195  		}
  1196  		lastEventVal, lastEventIsNull, err := item.EvalTime(e.ctx, chk.GetEvent(numEvents-1))
  1197  		if err != nil {
  1198  			return err
  1199  		}
  1200  		if !firstEventIsNull {
  1201  			firstEventCauset.SetMysqlTime(firstEventVal)
  1202  		} else {
  1203  			firstEventCauset.SetNull()
  1204  		}
  1205  		if !lastEventIsNull {
  1206  			lastEventCauset.SetMysqlTime(lastEventVal)
  1207  		} else {
  1208  			lastEventCauset.SetNull()
  1209  		}
  1210  	case types.ETDuration:
  1211  		firstEventVal, firstEventIsNull, err := item.EvalDuration(e.ctx, chk.GetEvent(0))
  1212  		if err != nil {
  1213  			return err
  1214  		}
  1215  		lastEventVal, lastEventIsNull, err := item.EvalDuration(e.ctx, chk.GetEvent(numEvents-1))
  1216  		if err != nil {
  1217  			return err
  1218  		}
  1219  		if !firstEventIsNull {
  1220  			firstEventCauset.SetMysqlDuration(firstEventVal)
  1221  		} else {
  1222  			firstEventCauset.SetNull()
  1223  		}
  1224  		if !lastEventIsNull {
  1225  			lastEventCauset.SetMysqlDuration(lastEventVal)
  1226  		} else {
  1227  			lastEventCauset.SetNull()
  1228  		}
  1229  	case types.ETJson:
  1230  		firstEventVal, firstEventIsNull, err := item.EvalJSON(e.ctx, chk.GetEvent(0))
  1231  		if err != nil {
  1232  			return err
  1233  		}
  1234  		lastEventVal, lastEventIsNull, err := item.EvalJSON(e.ctx, chk.GetEvent(numEvents-1))
  1235  		if err != nil {
  1236  			return err
  1237  		}
  1238  		if !firstEventIsNull {
  1239  			// make a copy to avoid DATA RACE
  1240  			firstEventCauset.SetMysqlJSON(firstEventVal.Copy())
  1241  		} else {
  1242  			firstEventCauset.SetNull()
  1243  		}
  1244  		if !lastEventIsNull {
  1245  			// make a copy to avoid DATA RACE
  1246  			lastEventCauset.SetMysqlJSON(lastEventVal.Copy())
  1247  		} else {
  1248  			lastEventCauset.SetNull()
  1249  		}
  1250  	case types.ETString:
  1251  		firstEventVal, firstEventIsNull, err := item.EvalString(e.ctx, chk.GetEvent(0))
  1252  		if err != nil {
  1253  			return err
  1254  		}
  1255  		lastEventVal, lastEventIsNull, err := item.EvalString(e.ctx, chk.GetEvent(numEvents-1))
  1256  		if err != nil {
  1257  			return err
  1258  		}
  1259  		if !firstEventIsNull {
  1260  			// make a copy to avoid DATA RACE
  1261  			firstCauset := string([]byte(firstEventVal))
  1262  			firstEventCauset.SetString(firstCauset, tp.DefCauslate)
  1263  		} else {
  1264  			firstEventCauset.SetNull()
  1265  		}
  1266  		if !lastEventIsNull {
  1267  			// make a copy to avoid DATA RACE
  1268  			lastCauset := string([]byte(lastEventVal))
  1269  			lastEventCauset.SetString(lastCauset, tp.DefCauslate)
  1270  		} else {
  1271  			lastEventCauset.SetNull()
  1272  		}
  1273  	default:
  1274  		err = errors.New(fmt.Sprintf("invalid eval type %v", eType))
  1275  		return err
  1276  	}
  1277  
  1278  	e.firstEventCausets = append(e.firstEventCausets, firstEventCauset)
  1279  	e.lastEventCausets = append(e.lastEventCausets, lastEventCauset)
  1280  	return err
  1281  }
  1282  
  1283  // evalGroupItemsAndResolveGroups evaluates the chunk according to the memex item.
  1284  // And resolve the rows into groups according to the evaluation results
  1285  func (e *vecGroupChecker) evalGroupItemsAndResolveGroups(item memex.Expression, chk *chunk.Chunk, numEvents int) (err error) {
  1286  	tp := item.GetType()
  1287  	eType := tp.EvalType()
  1288  	if e.allocateBuffer == nil {
  1289  		e.allocateBuffer = memex.GetDeferredCauset
  1290  	}
  1291  	if e.releaseBuffer == nil {
  1292  		e.releaseBuffer = memex.PutDeferredCauset
  1293  	}
  1294  	defCaus, err := e.allocateBuffer(eType, numEvents)
  1295  	if err != nil {
  1296  		return err
  1297  	}
  1298  	defer e.releaseBuffer(defCaus)
  1299  	err = memex.EvalExpr(e.ctx, item, chk, defCaus)
  1300  	if err != nil {
  1301  		return err
  1302  	}
  1303  
  1304  	previousIsNull := defCaus.IsNull(0)
  1305  	switch eType {
  1306  	case types.ETInt:
  1307  		vals := defCaus.Int64s()
  1308  		for i := 1; i < numEvents; i++ {
  1309  			isNull := defCaus.IsNull(i)
  1310  			if e.sameGroup[i] {
  1311  				switch {
  1312  				case !previousIsNull && !isNull:
  1313  					if vals[i] != vals[i-1] {
  1314  						e.sameGroup[i] = false
  1315  					}
  1316  				case isNull != previousIsNull:
  1317  					e.sameGroup[i] = false
  1318  				}
  1319  			}
  1320  			previousIsNull = isNull
  1321  		}
  1322  	case types.ETReal:
  1323  		vals := defCaus.Float64s()
  1324  		for i := 1; i < numEvents; i++ {
  1325  			isNull := defCaus.IsNull(i)
  1326  			if e.sameGroup[i] {
  1327  				switch {
  1328  				case !previousIsNull && !isNull:
  1329  					if vals[i] != vals[i-1] {
  1330  						e.sameGroup[i] = false
  1331  					}
  1332  				case isNull != previousIsNull:
  1333  					e.sameGroup[i] = false
  1334  				}
  1335  			}
  1336  			previousIsNull = isNull
  1337  		}
  1338  	case types.ETDecimal:
  1339  		vals := defCaus.Decimals()
  1340  		for i := 1; i < numEvents; i++ {
  1341  			isNull := defCaus.IsNull(i)
  1342  			if e.sameGroup[i] {
  1343  				switch {
  1344  				case !previousIsNull && !isNull:
  1345  					if vals[i].Compare(&vals[i-1]) != 0 {
  1346  						e.sameGroup[i] = false
  1347  					}
  1348  				case isNull != previousIsNull:
  1349  					e.sameGroup[i] = false
  1350  				}
  1351  			}
  1352  			previousIsNull = isNull
  1353  		}
  1354  	case types.ETDatetime, types.ETTimestamp:
  1355  		vals := defCaus.Times()
  1356  		for i := 1; i < numEvents; i++ {
  1357  			isNull := defCaus.IsNull(i)
  1358  			if e.sameGroup[i] {
  1359  				switch {
  1360  				case !previousIsNull && !isNull:
  1361  					if vals[i].Compare(vals[i-1]) != 0 {
  1362  						e.sameGroup[i] = false
  1363  					}
  1364  				case isNull != previousIsNull:
  1365  					e.sameGroup[i] = false
  1366  				}
  1367  			}
  1368  			previousIsNull = isNull
  1369  		}
  1370  	case types.ETDuration:
  1371  		vals := defCaus.GoDurations()
  1372  		for i := 1; i < numEvents; i++ {
  1373  			isNull := defCaus.IsNull(i)
  1374  			if e.sameGroup[i] {
  1375  				switch {
  1376  				case !previousIsNull && !isNull:
  1377  					if vals[i] != vals[i-1] {
  1378  						e.sameGroup[i] = false
  1379  					}
  1380  				case isNull != previousIsNull:
  1381  					e.sameGroup[i] = false
  1382  				}
  1383  			}
  1384  			previousIsNull = isNull
  1385  		}
  1386  	case types.ETJson:
  1387  		var previousKey, key json.BinaryJSON
  1388  		if !previousIsNull {
  1389  			previousKey = defCaus.GetJSON(0)
  1390  		}
  1391  		for i := 1; i < numEvents; i++ {
  1392  			isNull := defCaus.IsNull(i)
  1393  			if !isNull {
  1394  				key = defCaus.GetJSON(i)
  1395  			}
  1396  			if e.sameGroup[i] {
  1397  				if isNull == previousIsNull {
  1398  					if !isNull && json.CompareBinary(previousKey, key) != 0 {
  1399  						e.sameGroup[i] = false
  1400  					}
  1401  				} else {
  1402  					e.sameGroup[i] = false
  1403  				}
  1404  			}
  1405  			if !isNull {
  1406  				previousKey = key
  1407  			}
  1408  			previousIsNull = isNull
  1409  		}
  1410  	case types.ETString:
  1411  		previousKey := codec.ConvertByDefCauslationStr(defCaus.GetString(0), tp)
  1412  		for i := 1; i < numEvents; i++ {
  1413  			key := codec.ConvertByDefCauslationStr(defCaus.GetString(i), tp)
  1414  			isNull := defCaus.IsNull(i)
  1415  			if e.sameGroup[i] {
  1416  				if isNull != previousIsNull || previousKey != key {
  1417  					e.sameGroup[i] = false
  1418  				}
  1419  			}
  1420  			previousKey = key
  1421  			previousIsNull = isNull
  1422  		}
  1423  	default:
  1424  		err = errors.New(fmt.Sprintf("invalid eval type %v", eType))
  1425  	}
  1426  	if err != nil {
  1427  		return err
  1428  	}
  1429  
  1430  	return err
  1431  }
  1432  
  1433  func (e *vecGroupChecker) getNextGroup() (begin, end int) {
  1434  	if e.nextGroupID == 0 {
  1435  		begin = 0
  1436  	} else {
  1437  		begin = e.groupOffset[e.nextGroupID-1]
  1438  	}
  1439  	end = e.groupOffset[e.nextGroupID]
  1440  	e.nextGroupID++
  1441  	return begin, end
  1442  }
  1443  
  1444  func (e *vecGroupChecker) isExhausted() bool {
  1445  	return e.nextGroupID >= e.groupCount
  1446  }
  1447  
  1448  func (e *vecGroupChecker) reset() {
  1449  	if e.groupOffset != nil {
  1450  		e.groupOffset = e.groupOffset[:0]
  1451  	}
  1452  	if e.sameGroup != nil {
  1453  		e.sameGroup = e.sameGroup[:0]
  1454  	}
  1455  	if e.firstGroupKey != nil {
  1456  		e.firstGroupKey = e.firstGroupKey[:0]
  1457  	}
  1458  	if e.lastGroupKey != nil {
  1459  		e.lastGroupKey = e.lastGroupKey[:0]
  1460  	}
  1461  	if e.firstEventCausets != nil {
  1462  		e.firstEventCausets = e.firstEventCausets[:0]
  1463  	}
  1464  	if e.lastEventCausets != nil {
  1465  		e.lastEventCausets = e.lastEventCausets[:0]
  1466  	}
  1467  }