github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/batch/stream.go (about)

     1  package batch
     2  
     3  import (
     4  	promchunk "github.com/cortexproject/cortex/pkg/chunk/encoding"
     5  )
     6  
     7  // batchStream deals with iteratoring through multiple, non-overlapping batches,
     8  // and building new slices of non-overlapping batches.  Designed to be used
     9  // without allocations.
    10  type batchStream []promchunk.Batch
    11  
    12  // reset, hasNext, next, atTime etc are all inlined in go1.11.
    13  
    14  func (bs *batchStream) reset() {
    15  	for i := range *bs {
    16  		(*bs)[i].Index = 0
    17  	}
    18  }
    19  
    20  func (bs *batchStream) hasNext() bool {
    21  	return len(*bs) > 0
    22  }
    23  
    24  func (bs *batchStream) next() {
    25  	(*bs)[0].Index++
    26  	if (*bs)[0].Index >= (*bs)[0].Length {
    27  		*bs = (*bs)[1:]
    28  	}
    29  }
    30  
    31  func (bs *batchStream) atTime() int64 {
    32  	return (*bs)[0].Timestamps[(*bs)[0].Index]
    33  }
    34  
    35  func (bs *batchStream) at() (int64, float64) {
    36  	b := &(*bs)[0]
    37  	return b.Timestamps[b.Index], b.Values[b.Index]
    38  }
    39  
    40  func mergeStreams(left, right batchStream, result batchStream, size int) batchStream {
    41  	// Reset the Index and Length of existing batches.
    42  	for i := range result {
    43  		result[i].Index = 0
    44  		result[i].Length = 0
    45  	}
    46  	resultLen := 1 // Number of batches in the final result.
    47  	b := &result[0]
    48  
    49  	// This function adds a new batch to the result
    50  	// if the current batch being appended is full.
    51  	checkForFullBatch := func() {
    52  		if b.Index == size {
    53  			// The batch reached it intended size.
    54  			// Add another batch the the result
    55  			// and use it for further appending.
    56  
    57  			// The Index is the place at which new sample
    58  			// has to be appended, hence it tells the length.
    59  			b.Length = b.Index
    60  			resultLen++
    61  			if resultLen > len(result) {
    62  				// It is possible that result can grow longer
    63  				// then the one provided.
    64  				result = append(result, promchunk.Batch{})
    65  			}
    66  			b = &result[resultLen-1]
    67  		}
    68  	}
    69  
    70  	for left.hasNext() && right.hasNext() {
    71  		checkForFullBatch()
    72  		t1, t2 := left.atTime(), right.atTime()
    73  		if t1 < t2 {
    74  			b.Timestamps[b.Index], b.Values[b.Index] = left.at()
    75  			left.next()
    76  		} else if t1 > t2 {
    77  			b.Timestamps[b.Index], b.Values[b.Index] = right.at()
    78  			right.next()
    79  		} else {
    80  			b.Timestamps[b.Index], b.Values[b.Index] = left.at()
    81  			left.next()
    82  			right.next()
    83  		}
    84  		b.Index++
    85  	}
    86  
    87  	// This function adds all the samples from the provided
    88  	// batchStream into the result in the same order.
    89  	addToResult := func(bs batchStream) {
    90  		for ; bs.hasNext(); bs.next() {
    91  			checkForFullBatch()
    92  			b.Timestamps[b.Index], b.Values[b.Index] = bs.at()
    93  			b.Index++
    94  			b.Length++
    95  		}
    96  	}
    97  
    98  	addToResult(left)
    99  	addToResult(right)
   100  
   101  	// The Index is the place at which new sample
   102  	// has to be appended, hence it tells the length.
   103  	b.Length = b.Index
   104  
   105  	// The provided 'result' slice might be bigger
   106  	// than the actual result, hence return the subslice.
   107  	result = result[:resultLen]
   108  	result.reset()
   109  	return result
   110  }