github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/interlock/aggfuncs/func_rank.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 aggfuncs
    15  
    16  import (
    17  	"unsafe"
    18  
    19  	"github.com/whtcorpsinc/milevadb/memex"
    20  	"github.com/whtcorpsinc/milevadb/stochastikctx"
    21  	"github.com/whtcorpsinc/milevadb/soliton/chunk"
    22  )
    23  
    24  const (
    25  	//DefPartialResult4RankSize is the size of partialResult4Rank
    26  	DefPartialResult4RankSize = int64(unsafe.Sizeof(partialResult4Rank{}))
    27  )
    28  
    29  type rank struct {
    30  	baseAggFunc
    31  	isDense bool
    32  	rowComparer
    33  }
    34  
    35  type partialResult4Rank struct {
    36  	curIdx   int64
    37  	lastRank int64
    38  	rows     []chunk.Event
    39  }
    40  
    41  func (r *rank) AllocPartialResult() (pr PartialResult, memDelta int64) {
    42  	return PartialResult(&partialResult4Rank{}), DefPartialResult4RankSize
    43  }
    44  
    45  func (r *rank) ResetPartialResult(pr PartialResult) {
    46  	p := (*partialResult4Rank)(pr)
    47  	p.curIdx = 0
    48  	p.lastRank = 0
    49  	p.rows = p.rows[:0]
    50  }
    51  
    52  func (r *rank) UFIDelatePartialResult(sctx stochastikctx.Context, rowsInGroup []chunk.Event, pr PartialResult) (memDelta int64, err error) {
    53  	p := (*partialResult4Rank)(pr)
    54  	p.rows = append(p.rows, rowsInGroup...)
    55  	memDelta += int64(len(rowsInGroup)) * DefEventSize
    56  	return memDelta, nil
    57  }
    58  
    59  func (r *rank) AppendFinalResult2Chunk(sctx stochastikctx.Context, pr PartialResult, chk *chunk.Chunk) error {
    60  	p := (*partialResult4Rank)(pr)
    61  	p.curIdx++
    62  	if p.curIdx == 1 {
    63  		p.lastRank = 1
    64  		chk.AppendInt64(r.ordinal, p.lastRank)
    65  		return nil
    66  	}
    67  	if r.compareEvents(p.rows[p.curIdx-2], p.rows[p.curIdx-1]) == 0 {
    68  		chk.AppendInt64(r.ordinal, p.lastRank)
    69  		return nil
    70  	}
    71  	if r.isDense {
    72  		p.lastRank++
    73  	} else {
    74  		p.lastRank = p.curIdx
    75  	}
    76  	chk.AppendInt64(r.ordinal, p.lastRank)
    77  	return nil
    78  }
    79  
    80  type rowComparer struct {
    81  	cmpFuncs []chunk.CompareFunc
    82  	defCausIdx   []int
    83  }
    84  
    85  func buildEventComparer(defcaus []*memex.DeferredCauset) rowComparer {
    86  	rc := rowComparer{}
    87  	rc.defCausIdx = make([]int, 0, len(defcaus))
    88  	rc.cmpFuncs = make([]chunk.CompareFunc, 0, len(defcaus))
    89  	for _, defCaus := range defcaus {
    90  		cmpFunc := chunk.GetCompareFunc(defCaus.RetType)
    91  		if cmpFunc == nil {
    92  			continue
    93  		}
    94  		rc.cmpFuncs = append(rc.cmpFuncs, chunk.GetCompareFunc(defCaus.RetType))
    95  		rc.defCausIdx = append(rc.defCausIdx, defCaus.Index)
    96  	}
    97  	return rc
    98  }
    99  
   100  func (rc *rowComparer) compareEvents(prev, curr chunk.Event) int {
   101  	for i, idx := range rc.defCausIdx {
   102  		res := rc.cmpFuncs[i](prev, idx, curr, idx)
   103  		if res != 0 {
   104  			return res
   105  		}
   106  	}
   107  	return 0
   108  }