github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/offset.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  )
    19  
    20  // offsetOp is an operator that implements offset, returning everything
    21  // after the first n tuples in its input.
    22  type offsetOp struct {
    23  	OneInputNode
    24  
    25  	offset int
    26  
    27  	// seen is the number of tuples seen so far.
    28  	seen int
    29  }
    30  
    31  var _ colexecbase.Operator = &offsetOp{}
    32  
    33  // NewOffsetOp returns a new offset operator with the given offset.
    34  func NewOffsetOp(input colexecbase.Operator, offset int) colexecbase.Operator {
    35  	c := &offsetOp{
    36  		OneInputNode: NewOneInputNode(input),
    37  		offset:       offset,
    38  	}
    39  	return c
    40  }
    41  
    42  func (c *offsetOp) Init() {
    43  	c.input.Init()
    44  }
    45  
    46  func (c *offsetOp) Next(ctx context.Context) coldata.Batch {
    47  	for {
    48  		bat := c.input.Next(ctx)
    49  		length := bat.Length()
    50  		if length == 0 {
    51  			return bat
    52  		}
    53  
    54  		c.seen += length
    55  
    56  		delta := c.seen - c.offset
    57  		// If the current batch encompasses the offset "boundary",
    58  		// add the elements after the boundary to the selection vector.
    59  		if delta > 0 && delta < length {
    60  			sel := bat.Selection()
    61  			outputStartIdx := length - delta
    62  			if sel != nil {
    63  				copy(sel, sel[outputStartIdx:length])
    64  			} else {
    65  				bat.SetSelection(true)
    66  				sel = bat.Selection()[:delta] // slice for bounds check elimination
    67  				for i := range sel {
    68  					sel[i] = outputStartIdx + i
    69  				}
    70  			}
    71  			bat.SetLength(delta)
    72  		}
    73  
    74  		if c.seen > c.offset {
    75  			return bat
    76  		}
    77  	}
    78  }
    79  
    80  // Reset resets the offsetOp for another run. Primarily used for
    81  // benchmarks.
    82  func (c *offsetOp) Reset() {
    83  	c.seen = 0
    84  }