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

     1  // Copyright 2019 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  // ordinalityOp is an operator that implements WITH ORDINALITY, which adds
    23  // an additional column to the result with an ordinal number.
    24  type ordinalityOp struct {
    25  	OneInputNode
    26  
    27  	allocator *colmem.Allocator
    28  	// outputIdx is the index of the column in which ordinalityOp will write the
    29  	// ordinal number.
    30  	outputIdx int
    31  	// counter is the number of tuples seen so far.
    32  	counter int64
    33  }
    34  
    35  var _ colexecbase.Operator = &ordinalityOp{}
    36  
    37  // NewOrdinalityOp returns a new WITH ORDINALITY operator.
    38  func NewOrdinalityOp(
    39  	allocator *colmem.Allocator, input colexecbase.Operator, outputIdx int,
    40  ) colexecbase.Operator {
    41  	input = newVectorTypeEnforcer(allocator, input, types.Int, outputIdx)
    42  	c := &ordinalityOp{
    43  		OneInputNode: NewOneInputNode(input),
    44  		allocator:    allocator,
    45  		outputIdx:    outputIdx,
    46  		counter:      1,
    47  	}
    48  	return c
    49  }
    50  
    51  func (c *ordinalityOp) Init() {
    52  	c.input.Init()
    53  }
    54  
    55  func (c *ordinalityOp) Next(ctx context.Context) coldata.Batch {
    56  	bat := c.input.Next(ctx)
    57  	if bat.Length() == 0 {
    58  		return coldata.ZeroBatch
    59  	}
    60  
    61  	outputVec := bat.ColVec(c.outputIdx)
    62  	if outputVec.MaybeHasNulls() {
    63  		// We need to make sure that there are no left over null values in the
    64  		// output vector.
    65  		outputVec.Nulls().UnsetNulls()
    66  	}
    67  	col := outputVec.Int64()
    68  	sel := bat.Selection()
    69  
    70  	if sel != nil {
    71  		// Bounds check elimination.
    72  		for _, i := range sel[:bat.Length()] {
    73  			col[i] = c.counter
    74  			c.counter++
    75  		}
    76  	} else {
    77  		// Bounds check elimination.
    78  		col = col[:bat.Length()]
    79  		for i := range col {
    80  			col[i] = c.counter
    81  			c.counter++
    82  		}
    83  	}
    84  
    85  	return bat
    86  }