implement nullif() function (#6448)

Author: nwt

runtime/sam/expr/function/function.go

@@ -115,6 +115,9 @@ func New(sctx *super.Context, name string, narg int) (expr.Function, error) {
 		argmax = 0
 		argmin = 0
 		f = &Now{}
+	case "nullif":
+		argmin, argmax = 2, 2
+		f = newNullIf()
 	case "parse_sup":
 		f = newParseSUP(sctx)
 	case "parse_uri":

runtime/sam/expr/function/nullif.go

@@ -0,0 +1,29 @@
+package function
+
+import (
+	"github.com/brimdata/super"
+	"github.com/brimdata/super/order"
+	"github.com/brimdata/super/runtime/sam/expr"
+)
+
+type NullIf struct {
+	compareFn expr.CompareFn
+}
+
+func newNullIf() *NullIf {
+	return &NullIf{expr.NewValueCompareFn(order.Asc, order.NullsLast)}
+}
+
+func (n *NullIf) Call(args []super.Value) super.Value {
+	val0, val1 := args[0].Under(), args[1].Under()
+	if val0.IsError() {
+		return val0
+	}
+	if val1.IsError() {
+		return val1
+	}
+	if n.compareFn(val0, val1) == 0 {
+		return super.NewValue(val0.Type(), nil)
+	}
+	return val0
+}

runtime/vam/expr/compare.go

@@ -23,6 +23,10 @@ func NewCompare(sctx *super.Context, op string, lhs, rhs Evaluator) *Compare {
 	return &Compare{sctx, vector.CompareOpFromString(op), lhs, rhs}
 }
 
+func (c *Compare) Compare(vec0, vec1 vector.Any) vector.Any {
+	return c.eval(vec0, vec1)
+}
+
 func (c *Compare) Eval(val vector.Any) vector.Any {
 	return vector.Apply(true, c.eval, c.lhs.Eval(val), c.rhs.Eval(val))
 }

runtime/vam/expr/function/function.go

@@ -87,6 +87,9 @@ func New(sctx *super.Context, name string, narg int) (expr.Function, error) {
 	case "network_of":
 		argmax = 2
 		f = &NetworkOf{sctx}
+	case "nullif":
+		argmin, argmax = 2, 2
+		f = newNullIf(sctx)
 	case "parse_sup":
 		f = newParseSUP(sctx)
 	case "parse_uri":

runtime/vam/expr/function/nullif.go

@@ -0,0 +1,39 @@
+package function
+
+import (
+	"github.com/brimdata/super"
+	"github.com/brimdata/super/runtime/vam/expr"
+	"github.com/brimdata/super/vector"
+	"github.com/brimdata/super/vector/bitvec"
+)
+
+type NullIf struct {
+	compare *expr.Compare
+}
+
+func newNullIf(sctx *super.Context) *NullIf {
+	return &NullIf{expr.NewCompare(sctx, "==", nil, nil)}
+}
+
+func (n *NullIf) Call(vecs ...vector.Any) vector.Any {
+	if vecs[0].Type().Kind() == super.ErrorKind {
+		return vecs[0]
+	}
+	if vecs[1].Type().Kind() == super.ErrorKind {
+		return vecs[1]
+	}
+	result := n.compare.Compare(vecs[0], vecs[1])
+	if result.Type().Kind() == super.ErrorKind {
+		return vecs[0]
+	}
+	nulls := vector.NullsOf(vecs[0]).Clone()
+	for i := range result.Len() {
+		if b, _ := vector.BoolValue(result, i); b {
+			if nulls.IsZero() {
+				nulls = bitvec.NewFalse(vecs[0].Len())
+			}
+			nulls.Set(i)
+		}
+	}
+	return vector.CopyAndSetNulls(vecs[0], nulls)
+}

runtime/ztests/expr/function/nullif.yaml

@@ -0,0 +1,33 @@
+spq: nullif(this[0], this[1])
+
+vector: true
+
+input: |
+  [1, 1]
+  [1, 1::int32]
+  [1::int32, 1]
+  [1, 2]
+  [1, "s"]
+  [1, null]
+  [1, error(0)]
+  [null, 1]
+  [null, null]
+  [null, error(0)]
+  [error(0), 1]
+  [error(0), null]
+  [error(0), error(1)]
+
+output: |
+  null::int64
+  null::int64
+  null::int32
+  1
+  1
+  1
+  error(0)
+  null::int64
+  null
+  null::error(int64)
+  error(0)
+  error(0)
+  error(0)

vector/bitvec/bits.go

@@ -2,6 +2,7 @@ package bitvec
 
 import (
 	"math/bits"
+	"slices"
 	"strings"
 )
 
@@ -28,6 +29,10 @@ func NewTrue(n uint32) Bits {
 	return b
 }
 
+func (b Bits) Clone() Bits {
+	return Bits{slices.Clone(b.bits), b.length}
+}
+
 func (b Bits) IsZero() bool {
 	return b.length == 0
 }

vector/bool.go

@@ -1,6 +1,8 @@
 package vector
 
 import (
+	"fmt"
+
 	"github.com/brimdata/super"
 	"github.com/brimdata/super/scode"
 	"github.com/brimdata/super/vector/bitvec"
@@ -81,7 +83,7 @@ func BoolValue(vec Any, slot uint32) (bool, bool) {
 	case *View:
 		return BoolValue(vec.Any, vec.Index[slot])
 	}
-	panic(vec)
+	panic(fmt.Sprintf("%#v", vec))
 }
 
 func NullsOf(v Any) bitvec.Bits {
@@ -211,6 +213,14 @@ func CopyAndSetNulls(v Any, nulls bitvec.Bits) Any {
 		return &copy
 	case *Union:
 		return NewUnion(v.Typ, v.Tags, v.Values, nulls)
+	case *View:
+		newNulls := bitvec.NewFalse(uint32(len(v.Index)))
+		for i, idx := range v.Index {
+			if nulls.IsSet(uint32(i)) {
+				newNulls.Set(idx)
+			}
+		}
+		return NewView(CopyAndSetNulls(v.Any, newNulls), v.Index)
 	default:
 		panic(v)
 	}