github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/count_agg_tmpl.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // {{/*
    12  // +build execgen_template
    13  //
    14  // This file is the execgen template for count_agg.eg.go. It's formatted in a
    15  // special way, so it's both valid Go and a valid text/template input. This
    16  // permits editing this file with editor support.
    17  //
    18  // */}}
    19  
    20  package colexec
    21  
    22  import (
    23  	"unsafe"
    24  
    25  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    26  	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
    27  )
    28  
    29  // {{range .}}
    30  
    31  func newCount_KINDAggAlloc(allocator *colmem.Allocator, allocSize int64) aggregateFuncAlloc {
    32  	return &count_KINDAggAlloc{allocator: allocator, allocSize: allocSize}
    33  }
    34  
    35  // count_KINDAgg supports either COUNT(*) or COUNT(col) aggregate.
    36  type count_KINDAgg struct {
    37  	groups []bool
    38  	vec    []int64
    39  	nulls  *coldata.Nulls
    40  	curIdx int
    41  	curAgg int64
    42  }
    43  
    44  var _ aggregateFunc = &count_KINDAgg{}
    45  
    46  const sizeOfCount_KINDAgg = int64(unsafe.Sizeof(count_KINDAgg{}))
    47  
    48  func (a *count_KINDAgg) Init(groups []bool, vec coldata.Vec) {
    49  	a.groups = groups
    50  	a.vec = vec.Int64()
    51  	a.nulls = vec.Nulls()
    52  	a.Reset()
    53  }
    54  
    55  func (a *count_KINDAgg) Reset() {
    56  	a.curIdx = -1
    57  	a.curAgg = 0
    58  	a.nulls.UnsetNulls()
    59  }
    60  
    61  func (a *count_KINDAgg) CurrentOutputIndex() int {
    62  	return a.curIdx
    63  }
    64  
    65  func (a *count_KINDAgg) SetOutputIndex(idx int) {
    66  	if a.curIdx != -1 {
    67  		a.curIdx = idx
    68  		a.nulls.UnsetNullsAfter(idx + 1)
    69  	}
    70  }
    71  
    72  func (a *count_KINDAgg) Compute(b coldata.Batch, inputIdxs []uint32) {
    73  	inputLen := b.Length()
    74  	sel := b.Selection()
    75  
    76  	// {{if not (eq .Kind "Rows")}}
    77  	// If this is a COUNT(col) aggregator and there are nulls in this batch,
    78  	// we must check each value for nullity. Note that it is only legal to do a
    79  	// COUNT aggregate on a single column.
    80  	nulls := b.ColVec(int(inputIdxs[0])).Nulls()
    81  	if nulls.MaybeHasNulls() {
    82  		if sel != nil {
    83  			for _, i := range sel[:inputLen] {
    84  				_ACCUMULATE_COUNT(a, nulls, i, true)
    85  			}
    86  		} else {
    87  			for i := range a.groups[:inputLen] {
    88  				_ACCUMULATE_COUNT(a, nulls, i, true)
    89  			}
    90  		}
    91  	} else
    92  	// {{end}}
    93  	{
    94  		if sel != nil {
    95  			for _, i := range sel[:inputLen] {
    96  				_ACCUMULATE_COUNT(a, nulls, i, false)
    97  			}
    98  		} else {
    99  			for i := range a.groups[:inputLen] {
   100  				_ACCUMULATE_COUNT(a, nulls, i, false)
   101  			}
   102  		}
   103  	}
   104  }
   105  
   106  func (a *count_KINDAgg) Flush() {
   107  	a.vec[a.curIdx] = a.curAgg
   108  	a.curIdx++
   109  }
   110  
   111  func (a *count_KINDAgg) HandleEmptyInputScalar() {
   112  	a.vec[0] = 0
   113  }
   114  
   115  type count_KINDAggAlloc struct {
   116  	allocator *colmem.Allocator
   117  	allocSize int64
   118  	aggFuncs  []count_KINDAgg
   119  }
   120  
   121  var _ aggregateFuncAlloc = &count_KINDAggAlloc{}
   122  
   123  func (a *count_KINDAggAlloc) newAggFunc() aggregateFunc {
   124  	if len(a.aggFuncs) == 0 {
   125  		a.allocator.AdjustMemoryUsage(sizeOfCount_KINDAgg * a.allocSize)
   126  		a.aggFuncs = make([]count_KINDAgg, a.allocSize)
   127  	}
   128  	f := &a.aggFuncs[0]
   129  	a.aggFuncs = a.aggFuncs[1:]
   130  	return f
   131  }
   132  
   133  // {{end}}
   134  
   135  // {{/*
   136  // _ACCUMULATE_COUNT aggregates the value at index i into the count aggregate.
   137  // _COL_WITH_NULLS indicates whether we should be paying attention to NULLs.
   138  func _ACCUMULATE_COUNT(a *countAgg, nulls *coldata.Nulls, i int, _COL_WITH_NULLS bool) { // */}}
   139  	// {{define "accumulateCount" -}}
   140  
   141  	if a.groups[i] {
   142  		if a.curIdx != -1 {
   143  			a.vec[a.curIdx] = a.curAgg
   144  		}
   145  		a.curIdx++
   146  		a.curAgg = int64(0)
   147  	}
   148  	var y int64
   149  	// {{if .ColWithNulls}}
   150  	y = int64(0)
   151  	if !nulls.NullAt(i) {
   152  		y = 1
   153  	}
   154  	// {{else}}
   155  	y = int64(1)
   156  	// {{end}}
   157  	a.curAgg += y
   158  	// {{end}}
   159  
   160  	// {{/*
   161  } // */}}