SSE: Add is_nan() and other like functions and allow '_' in func names (#43555)

* SSE: Add is_nan() and allow '_' in func names

is_null() infn(), is_inf(), is_number()

pkg/expr/mathexp/funcs.go

@@ -21,14 +21,38 @@ var builtins = map[string]parse.Func{
 		Return: parse.TypeScalar,
 		F:      nan,
 	},
+	"is_nan": {
+		Args:          []parse.ReturnType{parse.TypeVariantSet},
+		VariantReturn: true,
+		F:             isNaN,
+	},
 	"inf": {
 		Return: parse.TypeScalar,
 		F:      inf,
 	},
+	"infn": {
+		Return: parse.TypeScalar,
+		F:      infn,
+	},
+	"is_inf": {
+		Args:          []parse.ReturnType{parse.TypeVariantSet},
+		VariantReturn: true,
+		F:             isInf,
+	},
 	"null": {
 		Return: parse.TypeScalar,
 		F:      null,
 	},
+	"is_null": {
+		Args:          []parse.ReturnType{parse.TypeVariantSet},
+		VariantReturn: true,
+		F:             isNull,
+	},
+	"is_number": {
+		Args:          []parse.ReturnType{parse.TypeVariantSet},
+		VariantReturn: true,
+		F:             isNumber,
+	},
 }
 
 // abs returns the absolute value for each result in NumberSet, SeriesSet, or Scalar
@@ -57,6 +81,43 @@ func log(e *State, varSet Results) (Results, error) {
 	return newRes, nil
 }
 
+// isNaN returns 1 if the value for each result in NumberSet, SeriesSet, or Scalar is NaN, else 0.
+func isNaN(e *State, varSet Results) (Results, error) {
+	newRes := Results{}
+	for _, res := range varSet.Values {
+		newVal, err := perFloat(e, res, func(f float64) float64 {
+			if math.IsNaN(f) {
+				return 1
+			}
+			return 0
+		})
+		if err != nil {
+			return newRes, err
+		}
+		newRes.Values = append(newRes.Values, newVal)
+	}
+	return newRes, nil
+}
+
+// isInf returns 1 if the value for each result in NumberSet, SeriesSet, or Scalar is a
+// positive or negative Inf, else 0.
+func isInf(e *State, varSet Results) (Results, error) {
+	newRes := Results{}
+	for _, res := range varSet.Values {
+		newVal, err := perFloat(e, res, func(f float64) float64 {
+			if math.IsInf(f, 0) {
+				return 1
+			}
+			return 0
+		})
+		if err != nil {
+			return newRes, err
+		}
+		newRes.Values = append(newRes.Values, newVal)
+	}
+	return newRes, nil
+}
+
 // nan returns a scalar nan value
 func nan(e *State) Results {
 	aNaN := math.NaN()
@@ -69,11 +130,59 @@ func inf(e *State) Results {
 	return NewScalarResults(e.RefID, &aInf)
 }
 
+// infn returns a scalar negative infinity value
+func infn(e *State) Results {
+	aInf := math.Inf(-1)
+	return NewScalarResults(e.RefID, &aInf)
+}
+
 // null returns a null scalar value
 func null(e *State) Results {
 	return NewScalarResults(e.RefID, nil)
 }
 
+// isNull returns 1 if the value for each result in NumberSet, SeriesSet, or Scalar is null, else 0.
+func isNull(e *State, varSet Results) (Results, error) {
+	newRes := Results{}
+	for _, res := range varSet.Values {
+		newVal, err := perNullableFloat(e, res, func(f *float64) *float64 {
+			nF := float64(0)
+			if f == nil {
+				nF = 1
+			}
+			return &nF
+		})
+		if err != nil {
+			return newRes, err
+		}
+		newRes.Values = append(newRes.Values, newVal)
+	}
+	return newRes, nil
+}
+
+// isNumber returns 1 if the value for each result in NumberSet, SeriesSet, or Scalar is a real number, else 0.
+// Therefore 0 is returned if the value Inf+, Inf-, NaN, or Null.
+func isNumber(e *State, varSet Results) (Results, error) {
+	newRes := Results{}
+	for _, res := range varSet.Values {
+		newVal, err := perNullableFloat(e, res, func(f *float64) *float64 {
+			nF := float64(1)
+			if f == nil || math.IsInf(*f, 0) || math.IsNaN(*f) {
+				nF = 0
+			}
+			return &nF
+		})
+		if err != nil {
+			return newRes, err
+		}
+		newRes.Values = append(newRes.Values, newVal)
+	}
+	return newRes, nil
+}
+
+// perFloat passes the non-null value of a Scalar/Number or each value point of a Series to floatF.
+// The return Value type will be the same type provided to function, (e.g. a Series input returns a series).
+// If input values are null the function is not called and NaN is returned for each value.
 func perFloat(e *State, val Value, floatF func(x float64) float64) (Value, error) {
 	var newVal Value
 	switch val.Type() {
@@ -113,3 +222,34 @@ func perFloat(e *State, val Value, floatF func(x float64) float64) (Value, error
 
 	return newVal, nil
 }
+
+// perNullableFloat is like perFloat, but takes and returns float pointers instead of floats.
+// This is for instead for functions that need specific null handling.
+// The input float pointer should not be modified in the floatF func.
+func perNullableFloat(e *State, val Value, floatF func(x *float64) *float64) (Value, error) {
+	var newVal Value
+	switch val.Type() {
+	case parse.TypeNumberSet:
+		n := NewNumber(e.RefID, val.GetLabels())
+		f := val.(Number).GetFloat64Value()
+		n.SetValue(floatF(f))
+		newVal = n
+	case parse.TypeScalar:
+		f := val.(Scalar).GetFloat64Value()
+		newVal = NewScalar(e.RefID, floatF(f))
+	case parse.TypeSeriesSet:
+		resSeries := val.(Series)
+		newSeries := NewSeries(e.RefID, resSeries.GetLabels(), resSeries.Len())
+		for i := 0; i < resSeries.Len(); i++ {
+			t, f := resSeries.GetPoint(i)
+			if err := newSeries.SetPoint(i, t, floatF(f)); err != nil {
+				return newSeries, err
+			}
+		}
+		newVal = newSeries
+	default:
+		// TODO: Should we deal with TypeString, TypeVariantSet?
+	}
+
+	return newVal, nil
+}

pkg/expr/mathexp/funcs_test.go

@@ -1,20 +1,21 @@
 package mathexp
 
 import (
+	"math"
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-func TestFunc(t *testing.T) {
+func TestAbsFunc(t *testing.T) {
 	var tests = []struct {
 		name      string
 		expr      string
 		vars      Vars
-		newErrIs  assert.ErrorAssertionFunc
-		execErrIs assert.ErrorAssertionFunc
-		resultIs  assert.ComparisonAssertionFunc
+		newErrIs  require.ErrorAssertionFunc
+		execErrIs require.ErrorAssertionFunc
+		resultIs  require.ComparisonAssertionFunc
 		results   Results
 	}{
 		{
@@ -27,18 +28,18 @@ func TestFunc(t *testing.T) {
 					},
 				},
 			},
-			newErrIs:  assert.NoError,
-			execErrIs: assert.NoError,
-			resultIs:  assert.Equal,
+			newErrIs:  require.NoError,
+			execErrIs: require.NoError,
+			resultIs:  require.Equal,
 			results:   Results{[]Value{makeNumber("", nil, float64Pointer(7))}},
 		},
 		{
 			name:      "abs on scalar",
 			expr:      "abs(-1)",
 			vars:      Vars{},
-			newErrIs:  assert.NoError,
-			execErrIs: assert.NoError,
-			resultIs:  assert.Equal,
+			newErrIs:  require.NoError,
+			execErrIs: require.NoError,
+			resultIs:  require.Equal,
 			results:   Results{[]Value{NewScalar("", float64Pointer(1.0))}},
 		},
 		{
@@ -55,9 +56,9 @@ func TestFunc(t *testing.T) {
 					},
 				},
 			},
-			newErrIs:  assert.NoError,
-			execErrIs: assert.NoError,
-			resultIs:  assert.Equal,
+			newErrIs:  require.NoError,
+			execErrIs: require.NoError,
+			resultIs:  require.Equal,
 			results: Results{
 				[]Value{
 					makeSeries("", nil, tp{
@@ -72,7 +73,7 @@ func TestFunc(t *testing.T) {
 			name:     "abs on string - should error",
 			expr:     `abs("hi")`,
 			vars:     Vars{},
-			newErrIs: assert.Error,
+			newErrIs: require.Error,
 		},
 	}
 	for _, tt := range tests {
@@ -87,3 +88,74 @@ func TestFunc(t *testing.T) {
 		})
 	}
 }
+
+func TestIsNumberFunc(t *testing.T) {
+	var tests = []struct {
+		name    string
+		expr    string
+		vars    Vars
+		results Results
+	}{
+		{
+			name: "is_number on number type with real number value",
+			expr: "is_number($A)",
+			vars: Vars{
+				"A": Results{
+					[]Value{
+						makeNumber("", nil, float64Pointer(6)),
+					},
+				},
+			},
+			results: Results{[]Value{makeNumber("", nil, float64Pointer(1))}},
+		},
+		{
+			name: "is_number on number type with null value",
+			expr: "is_number($A)",
+			vars: Vars{
+				"A": Results{
+					[]Value{
+						makeNumber("", nil, nil),
+					},
+				},
+			},
+			results: Results{[]Value{makeNumber("", nil, float64Pointer(0))}},
+		},
+		{
+			name: "is_number on on series",
+			expr: "is_number($A)",
+			vars: Vars{
+				"A": Results{
+					[]Value{
+						makeSeries("", nil,
+							tp{time.Unix(5, 0), float64Pointer(5)},
+							tp{time.Unix(10, 0), nil},
+							tp{time.Unix(15, 0), float64Pointer(math.NaN())},
+							tp{time.Unix(20, 0), float64Pointer(math.Inf(-1))},
+							tp{time.Unix(25, 0), float64Pointer(math.Inf(0))}),
+					},
+				},
+			},
+			results: Results{
+				[]Value{
+					makeSeries("", nil,
+						tp{time.Unix(5, 0), float64Pointer(1)},
+						tp{time.Unix(10, 0), float64Pointer(0)},
+						tp{time.Unix(15, 0), float64Pointer(0)},
+						tp{time.Unix(20, 0), float64Pointer(0)},
+						tp{time.Unix(25, 0), float64Pointer(0)}),
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			e, err := New(tt.expr)
+			require.NoError(t, err)
+			if e != nil {
+				res, err := e.Execute("", tt.vars)
+				require.NoError(t, err)
+				require.Equal(t, tt.results, res)
+			}
+		})
+	}
+}

pkg/expr/mathexp/parse/lex.go

@@ -276,7 +276,7 @@ func lexSymbol(l *lexer) stateFn {
 func lexFunc(l *lexer) stateFn {
 	for {
 		switch r := l.next(); {
-		case unicode.IsLetter(r):
+		case unicode.IsLetter(r) || r == '_':
 			// absorb
 		default:
 			l.backup()