github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/count.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  package colexec
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    20  )
    21  
    22  // countOp is an operator that counts the number of input rows it receives,
    23  // consuming its entire input and outputting a batch with a single integer
    24  // column containing a single integer, the count of rows received from the
    25  // upstream.
    26  type countOp struct {
    27  	OneInputNode
    28  
    29  	internalBatch coldata.Batch
    30  	done          bool
    31  	count         int64
    32  }
    33  
    34  var _ colexecbase.Operator = &countOp{}
    35  
    36  // NewCountOp returns a new count operator that counts the rows in its input.
    37  func NewCountOp(allocator *colmem.Allocator, input colexecbase.Operator) colexecbase.Operator {
    38  	c := &countOp{
    39  		OneInputNode: NewOneInputNode(input),
    40  	}
    41  	c.internalBatch = allocator.NewMemBatchWithSize([]*types.T{types.Int}, 1)
    42  	return c
    43  }
    44  
    45  func (c *countOp) Init() {
    46  	c.input.Init()
    47  	c.count = 0
    48  	c.done = false
    49  }
    50  
    51  func (c *countOp) Next(ctx context.Context) coldata.Batch {
    52  	c.internalBatch.ResetInternalBatch()
    53  	if c.done {
    54  		return coldata.ZeroBatch
    55  	}
    56  	for {
    57  		bat := c.input.Next(ctx)
    58  		length := bat.Length()
    59  		if length == 0 {
    60  			c.done = true
    61  			c.internalBatch.ColVec(0).Int64()[0] = c.count
    62  			c.internalBatch.SetLength(1)
    63  			return c.internalBatch
    64  		}
    65  		c.count += int64(length)
    66  	}
    67  }