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 }