github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/is_null_ops.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  // isNullProjOp is an Operator that projects into outputIdx Vec whether the
    23  // corresponding value in colIdx Vec is NULL (i.e. it performs IS NULL check).
    24  // If negate is true, it does the opposite - it performs IS NOT NULL check.
    25  type isNullProjOp struct {
    26  	OneInputNode
    27  	allocator *colmem.Allocator
    28  	colIdx    int
    29  	outputIdx int
    30  	negate    bool
    31  }
    32  
    33  func newIsNullProjOp(
    34  	allocator *colmem.Allocator, input colexecbase.Operator, colIdx, outputIdx int, negate bool,
    35  ) colexecbase.Operator {
    36  	input = newVectorTypeEnforcer(allocator, input, types.Bool, outputIdx)
    37  	return &isNullProjOp{
    38  		OneInputNode: NewOneInputNode(input),
    39  		allocator:    allocator,
    40  		colIdx:       colIdx,
    41  		outputIdx:    outputIdx,
    42  		negate:       negate,
    43  	}
    44  }
    45  
    46  var _ colexecbase.Operator = &isNullProjOp{}
    47  
    48  func (o *isNullProjOp) Init() {
    49  	o.input.Init()
    50  }
    51  
    52  func (o *isNullProjOp) Next(ctx context.Context) coldata.Batch {
    53  	batch := o.input.Next(ctx)
    54  	n := batch.Length()
    55  	if n == 0 {
    56  		return coldata.ZeroBatch
    57  	}
    58  	vec := batch.ColVec(o.colIdx)
    59  	nulls := vec.Nulls()
    60  	projVec := batch.ColVec(o.outputIdx)
    61  	projCol := projVec.Bool()
    62  	if projVec.MaybeHasNulls() {
    63  		// We need to make sure that there are no left over null values in the
    64  		// output vector.
    65  		projVec.Nulls().UnsetNulls()
    66  	}
    67  	if nulls.MaybeHasNulls() {
    68  		if sel := batch.Selection(); sel != nil {
    69  			sel = sel[:n]
    70  			for _, i := range sel {
    71  				projCol[i] = nulls.NullAt(i) != o.negate
    72  			}
    73  		} else {
    74  			projCol = projCol[:n]
    75  			for i := range projCol {
    76  				projCol[i] = nulls.NullAt(i) != o.negate
    77  			}
    78  		}
    79  	} else {
    80  		// There are no NULLs, so we don't need to check each index for nullity.
    81  		result := o.negate
    82  		if sel := batch.Selection(); sel != nil {
    83  			sel = sel[:n]
    84  			for _, i := range sel {
    85  				projCol[i] = result
    86  			}
    87  		} else {
    88  			projCol = projCol[:n]
    89  			for i := range projCol {
    90  				projCol[i] = result
    91  			}
    92  		}
    93  	}
    94  	return batch
    95  }
    96  
    97  // isNullSelOp is an Operator that selects all the tuples that have a NULL
    98  // value in colIdx Vec. If negate is true, then it does the opposite -
    99  // selecting all the tuples that have a non-NULL value in colIdx Vec.
   100  type isNullSelOp struct {
   101  	OneInputNode
   102  	colIdx int
   103  	negate bool
   104  }
   105  
   106  func newIsNullSelOp(input colexecbase.Operator, colIdx int, negate bool) colexecbase.Operator {
   107  	return &isNullSelOp{
   108  		OneInputNode: NewOneInputNode(input),
   109  		colIdx:       colIdx,
   110  		negate:       negate,
   111  	}
   112  }
   113  
   114  var _ colexecbase.Operator = &isNullSelOp{}
   115  
   116  func (o *isNullSelOp) Init() {
   117  	o.input.Init()
   118  }
   119  
   120  func (o *isNullSelOp) Next(ctx context.Context) coldata.Batch {
   121  	for {
   122  		batch := o.input.Next(ctx)
   123  		n := batch.Length()
   124  		if n == 0 {
   125  			return batch
   126  		}
   127  		var idx int
   128  		vec := batch.ColVec(o.colIdx)
   129  		nulls := vec.Nulls()
   130  		if nulls.MaybeHasNulls() {
   131  			// There might be NULLs in the Vec, so we'll need to iterate over all
   132  			// tuples.
   133  			if sel := batch.Selection(); sel != nil {
   134  				sel = sel[:n]
   135  				for _, i := range sel {
   136  					if nulls.NullAt(i) != o.negate {
   137  						sel[idx] = i
   138  						idx++
   139  					}
   140  				}
   141  			} else {
   142  				batch.SetSelection(true)
   143  				sel := batch.Selection()[:n]
   144  				for i := range sel {
   145  					if nulls.NullAt(i) != o.negate {
   146  						sel[idx] = i
   147  						idx++
   148  					}
   149  				}
   150  			}
   151  			if idx > 0 {
   152  				batch.SetLength(idx)
   153  				return batch
   154  			}
   155  		} else {
   156  			// There are no NULLs, so we don't need to check each index for nullity.
   157  			if o.negate {
   158  				// o.negate is true, so we select all tuples, i.e. we don't need to
   159  				// modify the batch and can just return it.
   160  				return batch
   161  			}
   162  			// o.negate is false, so we omit all tuples from this batch and move onto
   163  			// the next one.
   164  		}
   165  	}
   166  }