vam: sqrt() (#5599)

Author: mattnibs

runtime/sam/expr/function/math.go

@@ -170,9 +170,18 @@ type Sqrt struct {
 }
 
 func (s *Sqrt) Call(_ super.Allocator, args []super.Value) super.Value {
-	x, ok := coerce.ToFloat(args[0], super.TypeFloat64)
+	val := args[0].Under()
+	if id := val.Type().ID(); id == super.IDNull {
+		return val
+	} else if !super.IsNumber(id) {
+		return s.zctx.WrapError("sqrt: number argument required", val)
+	}
+	if val.IsNull() {
+		return super.NullFloat64
+	}
+	x, ok := coerce.ToFloat(val, super.TypeFloat64)
 	if !ok {
-		return s.zctx.WrapError("sqrt: not a number", args[0])
+		return s.zctx.WrapError("sqrt: not a number", val)
 	}
 	return super.NewFloat64(math.Sqrt(x))
 }

runtime/vam/expr/function/function.go

@@ -70,6 +70,8 @@ func New(zctx *super.Context, name string, narg int) (expr.Function, field.Path,
 	case "split":
 		argmin, argmax = 2, 2
 		f = &Split{zctx}
+	case "sqrt":
+		f = &Sqrt{zctx}
 	case "strftime":
 		argmin, argmax = 2, 2
 		f = &Strftime{zctx: zctx}

runtime/vam/expr/function/math.go

@@ -178,3 +178,27 @@ func (l *Log) Call(args ...vector.Any) vector.Any {
 	}
 	return out
 }
+
+// https://github.com/brimdata/super/blob/main/docs/language/functions.md#sqrt
+type Sqrt struct {
+	zctx *super.Context
+}
+
+func (s *Sqrt) Call(args ...vector.Any) vector.Any {
+	vec := vector.Under(args[0])
+	if id := vec.Type().ID(); id == super.IDNull {
+		return vec
+	} else if !super.IsNumber(id) {
+		return vector.NewWrappedError(s.zctx, "sqrt: number argument required", vec)
+	}
+	vec = cast.To(s.zctx, vec, super.TypeFloat64)
+	vals := make([]float64, vec.Len())
+	for i := range vec.Len() {
+		v, isnull := vector.FloatValue(vec, i)
+		if isnull {
+			continue
+		}
+		vals[i] = math.Sqrt(v)
+	}
+	return vector.NewFloat(super.TypeFloat64, vals, vector.NullsOf(vec))
+}

runtime/ztests/expr/function/sqrt.yaml

@@ -0,0 +1,21 @@
+zed: sqrt(this)
+
+vector: true
+
+input: |
+  4
+  2.
+  1e10
+  -1
+  null(float64)
+  null
+  "9"
+
+output: |
+  2.
+  1.4142135623730951
+  100000.
+  NaN
+  null(float64)
+  null
+  error({message:"sqrt: number argument required",on:"9"})