github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/limit.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  )
    19  
    20  // limitOp is an operator that implements limit, returning only the first n
    21  // tuples from its input.
    22  type limitOp struct {
    23  	OneInputNode
    24  	closerHelper
    25  
    26  	limit int
    27  
    28  	// seen is the number of tuples seen so far.
    29  	seen int
    30  	// done is true if the limit has been reached.
    31  	done bool
    32  }
    33  
    34  var _ colexecbase.Operator = &limitOp{}
    35  var _ closableOperator = &limitOp{}
    36  
    37  // NewLimitOp returns a new limit operator with the given limit.
    38  func NewLimitOp(input colexecbase.Operator, limit int) colexecbase.Operator {
    39  	c := &limitOp{
    40  		OneInputNode: NewOneInputNode(input),
    41  		limit:        limit,
    42  	}
    43  	return c
    44  }
    45  
    46  func (c *limitOp) Init() {
    47  	c.input.Init()
    48  }
    49  
    50  func (c *limitOp) Next(ctx context.Context) coldata.Batch {
    51  	if c.done {
    52  		return coldata.ZeroBatch
    53  	}
    54  	bat := c.input.Next(ctx)
    55  	length := bat.Length()
    56  	if length == 0 {
    57  		return bat
    58  	}
    59  	newSeen := c.seen + length
    60  	if newSeen >= c.limit {
    61  		c.done = true
    62  		bat.SetLength(c.limit - c.seen)
    63  		return bat
    64  	}
    65  	c.seen = newSeen
    66  	return bat
    67  }
    68  
    69  // Close closes the limitOp's input.
    70  // TODO(asubiotto): Remove this method. It only exists so that we can call Close
    71  //  from some runTests subtests when not draining the input fully. The test
    72  //  should pass in the testing.T object used so that the caller can decide to
    73  //  explicitly close the input after checking the test.
    74  func (c *limitOp) IdempotentClose(ctx context.Context) error {
    75  	if !c.close() {
    76  		return nil
    77  	}
    78  	if closer, ok := c.input.(IdempotentCloser); ok {
    79  		return closer.IdempotentClose(ctx)
    80  	}
    81  	return nil
    82  }