github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/deselector.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  // deselectorOp consumes the input operator, and if resulting batches have a
    23  // selection vector, it coalesces them (meaning that tuples will be reordered
    24  // or omitted according to the selection vector). If the batches come with no
    25  // selection vector, it is a noop.
    26  type deselectorOp struct {
    27  	OneInputNode
    28  	NonExplainable
    29  	allocator  *colmem.Allocator
    30  	inputTypes []*types.T
    31  
    32  	output coldata.Batch
    33  }
    34  
    35  var _ colexecbase.Operator = &deselectorOp{}
    36  
    37  // NewDeselectorOp creates a new deselector operator on the given input
    38  // operator with the given column types.
    39  func NewDeselectorOp(
    40  	allocator *colmem.Allocator, input colexecbase.Operator, typs []*types.T,
    41  ) colexecbase.Operator {
    42  	return &deselectorOp{
    43  		OneInputNode: NewOneInputNode(input),
    44  		allocator:    allocator,
    45  		inputTypes:   typs,
    46  	}
    47  }
    48  
    49  func (p *deselectorOp) Init() {
    50  	p.input.Init()
    51  }
    52  
    53  func (p *deselectorOp) Next(ctx context.Context) coldata.Batch {
    54  	p.resetOutput()
    55  	batch := p.input.Next(ctx)
    56  	if batch.Selection() == nil {
    57  		return batch
    58  	}
    59  
    60  	sel := batch.Selection()
    61  	p.allocator.PerformOperation(p.output.ColVecs(), func() {
    62  		for i := range p.inputTypes {
    63  			toCol := p.output.ColVec(i)
    64  			fromCol := batch.ColVec(i)
    65  			toCol.Copy(
    66  				coldata.CopySliceArgs{
    67  					SliceArgs: coldata.SliceArgs{
    68  						Src:       fromCol,
    69  						Sel:       sel,
    70  						SrcEndIdx: batch.Length(),
    71  					},
    72  				},
    73  			)
    74  		}
    75  	})
    76  	p.output.SetLength(batch.Length())
    77  	return p.output
    78  }
    79  
    80  func (p *deselectorOp) resetOutput() {
    81  	if p.output == nil {
    82  		p.output = p.allocator.NewMemBatch(p.inputTypes)
    83  	} else {
    84  		p.output.ResetInternalBatch()
    85  	}
    86  }