github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/builtin_funcs.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/colexecbase/colexecerror"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    23  )
    24  
    25  type defaultBuiltinFuncOperator struct {
    26  	OneInputNode
    27  	allocator    *colmem.Allocator
    28  	evalCtx      *tree.EvalContext
    29  	funcExpr     *tree.FuncExpr
    30  	columnTypes  []*types.T
    31  	argumentCols []int
    32  	outputIdx    int
    33  	outputType   *types.T
    34  	converter    func(tree.Datum) (interface{}, error)
    35  
    36  	row tree.Datums
    37  	da  sqlbase.DatumAlloc
    38  }
    39  
    40  var _ colexecbase.Operator = &defaultBuiltinFuncOperator{}
    41  
    42  func (b *defaultBuiltinFuncOperator) Init() {
    43  	b.input.Init()
    44  }
    45  
    46  func (b *defaultBuiltinFuncOperator) Next(ctx context.Context) coldata.Batch {
    47  	batch := b.input.Next(ctx)
    48  	n := batch.Length()
    49  	if n == 0 {
    50  		return coldata.ZeroBatch
    51  	}
    52  
    53  	sel := batch.Selection()
    54  	output := batch.ColVec(b.outputIdx)
    55  	if output.MaybeHasNulls() {
    56  		// We need to make sure that there are no left over null values in the
    57  		// output vector.
    58  		output.Nulls().UnsetNulls()
    59  	}
    60  	b.allocator.PerformOperation(
    61  		[]coldata.Vec{output},
    62  		func() {
    63  			for i := 0; i < n; i++ {
    64  				rowIdx := i
    65  				if sel != nil {
    66  					rowIdx = sel[i]
    67  				}
    68  
    69  				hasNulls := false
    70  
    71  				for j := range b.argumentCols {
    72  					col := batch.ColVec(b.argumentCols[j])
    73  					b.row[j] = PhysicalTypeColElemToDatum(col, rowIdx, &b.da, b.columnTypes[b.argumentCols[j]])
    74  					hasNulls = hasNulls || b.row[j] == tree.DNull
    75  				}
    76  
    77  				var (
    78  					res tree.Datum
    79  					err error
    80  				)
    81  				// Some functions cannot handle null arguments.
    82  				if hasNulls && !b.funcExpr.CanHandleNulls() {
    83  					res = tree.DNull
    84  				} else {
    85  					res, err = b.funcExpr.ResolvedOverload().Fn(b.evalCtx, b.row)
    86  					if err != nil {
    87  						colexecerror.ExpectedError(err)
    88  					}
    89  				}
    90  
    91  				// Convert the datum into a physical type and write it out.
    92  				if res == tree.DNull {
    93  					batch.ColVec(b.outputIdx).Nulls().SetNull(rowIdx)
    94  				} else {
    95  					converted, err := b.converter(res)
    96  					if err != nil {
    97  						colexecerror.InternalError(err)
    98  					}
    99  					coldata.SetValueAt(output, converted, rowIdx)
   100  				}
   101  			}
   102  		},
   103  	)
   104  	// Although we didn't change the length of the batch, it is necessary to set
   105  	// the length anyway (this helps maintaining the invariant of flat bytes).
   106  	batch.SetLength(n)
   107  	return batch
   108  }
   109  
   110  // NewBuiltinFunctionOperator returns an operator that applies builtin functions.
   111  func NewBuiltinFunctionOperator(
   112  	allocator *colmem.Allocator,
   113  	evalCtx *tree.EvalContext,
   114  	funcExpr *tree.FuncExpr,
   115  	columnTypes []*types.T,
   116  	argumentCols []int,
   117  	outputIdx int,
   118  	input colexecbase.Operator,
   119  ) (colexecbase.Operator, error) {
   120  	switch funcExpr.ResolvedOverload().SpecializedVecBuiltin {
   121  	case tree.SubstringStringIntInt:
   122  		input = newVectorTypeEnforcer(allocator, input, types.String, outputIdx)
   123  		return newSubstringOperator(
   124  			allocator, columnTypes, argumentCols, outputIdx, input,
   125  		), nil
   126  	default:
   127  		outputType := funcExpr.ResolvedType()
   128  		input = newVectorTypeEnforcer(allocator, input, outputType, outputIdx)
   129  		return &defaultBuiltinFuncOperator{
   130  			OneInputNode: NewOneInputNode(input),
   131  			allocator:    allocator,
   132  			evalCtx:      evalCtx,
   133  			funcExpr:     funcExpr,
   134  			outputIdx:    outputIdx,
   135  			columnTypes:  columnTypes,
   136  			outputType:   outputType,
   137  			converter:    getDatumToPhysicalFn(outputType),
   138  			row:          make(tree.Datums, len(argumentCols)),
   139  			argumentCols: argumentCols,
   140  		}, nil
   141  	}
   142  }