cmd/bosun: setvariant type to enable generics for expr funcs

Author: kylebrandt

cmd/bosun/expr/expr.go

@@ -412,7 +412,6 @@ func (e *State) union(a, b *Results, expression string) []*Union {
 	var group opentsdb.TagSet
 	for _, ra := range a.Results {
 		for _, rb := range b.Results {
-
 			if ra.Group.Equal(rb.Group) || len(ra.Group) == 0 || len(rb.Group) == 0 {
 				g := ra.Group
 				if len(ra.Group) == 0 {
@@ -759,7 +758,7 @@ func (e *State) walkFunc(node *parse.FuncNode, T miniprofiler.Timer) *Results {
 			} else {
 				argType = node.F.Args[i]
 			}
-			if f, ok := v.(float64); ok && argType == models.TypeNumberSet {
+			if f, ok := v.(float64); ok && (argType == models.TypeNumberSet || argType == models.TypeVariantSet) {
 				v = fromScalar(f)
 			}
 			in = append(in, reflect.ValueOf(v))

cmd/bosun/expr/expr_test.go

@@ -258,6 +258,79 @@ func TestQueryExpr(t *testing.T) {
 	}
 }
 
+func TestSetVariant(t *testing.T) {
+	series := `series("key1=a,key2=b", 0, 1, 1, 3)`
+	seriesAbs := `series("", 0, 1, 1, -3)`
+	tests := []exprInOut{
+		{
+			fmt.Sprintf(`addtags(addtags(%v, "key3=a"), "key4=b") + 1`, series),
+			Results{
+				Results: ResultSlice{
+					&Result{
+						Value: Series{
+							time.Unix(0, 0): 2,
+							time.Unix(1, 0): 4,
+						},
+						Group: opentsdb.TagSet{"key1": "a", "key2": "b", "key3": "a", "key4": "b"},
+					},
+				},
+			},
+			false,
+		},
+		{
+			fmt.Sprintf(`addtags(addtags(avg(%v + 1), "key3=a"), "key4=b") + 1`, series),
+			Results{
+				Results: ResultSlice{
+					&Result{
+						Value: Number(4),
+						Group: opentsdb.TagSet{"key1": "a", "key2": "b", "key3": "a", "key4": "b"},
+					},
+				},
+			},
+			false,
+		},
+		{
+			fmt.Sprintf(`avg(addtags(addtags(avg(%v + 1), "key3=a"), "key4=b")) + 1`, series),
+			Results{},
+			true,
+		},
+
+		{
+			fmt.Sprintf(`1 + abs(%v)`, seriesAbs),
+			Results{
+				Results: ResultSlice{
+					&Result{
+						Value: Series{
+							time.Unix(0, 0): 2,
+							time.Unix(1, 0): 4,
+						},
+						Group: opentsdb.TagSet{},
+					},
+				},
+			},
+			false,
+		},
+		{
+			fmt.Sprintf(`1 + abs(avg(%v))`, seriesAbs),
+			Results{
+				Results: ResultSlice{
+					&Result{
+						Value: Number(2),
+						Group: opentsdb.TagSet{},
+					},
+				},
+			},
+			false,
+		},
+	}
+	for _, test := range tests {
+		err := testExpression(test)
+		if err != nil {
+			t.Error(err)
+		}
+	}
+}
+
 func TestSeriesOperations(t *testing.T) {
 	seriesA := `series("key=a", 0, 1, 1, 2, 2, 1, 3, 4)`
 	seriesB := `series("key=a", 0, 1,       2, 0, 3, 4)`

cmd/bosun/expr/funcs.go

@@ -202,22 +202,22 @@ var builtins = map[string]parse.Func{
 
 	// Group functions
 	"addtags": {
-		Args:   []models.FuncType{models.TypeSeriesSet, models.TypeString},
-		Return: models.TypeSeriesSet,
-		Tags:   tagRename,
-		F:      AddTags,
+		Args:          []models.FuncType{models.TypeVariantSet, models.TypeString},
+		VariantReturn: true,
+		Tags:          tagRename,
+		F:             AddTags,
 	},
 	"rename": {
-		Args:   []models.FuncType{models.TypeSeriesSet, models.TypeString},
-		Return: models.TypeSeriesSet,
-		Tags:   tagRename,
-		F:      Rename,
+		Args:          []models.FuncType{models.TypeVariantSet, models.TypeString},
+		VariantReturn: true,
+		Tags:          tagRename,
+		F:             Rename,
 	},
 	"remove": {
-		Args:   []models.FuncType{models.TypeSeriesSet, models.TypeString},
-		Return: models.TypeSeriesSet,
-		Tags:   tagRemove,
-		F:      Remove,
+		Args:          []models.FuncType{models.TypeVariantSet, models.TypeString},
+		VariantReturn: true,
+		Tags:          tagRemove,
+		F:             Remove,
 	},
 	"t": {
 		Args:   []models.FuncType{models.TypeNumberSet, models.TypeString},
@@ -234,10 +234,10 @@ var builtins = map[string]parse.Func{
 	// Other functions
 
 	"abs": {
-		Args:   []models.FuncType{models.TypeNumberSet},
-		Return: models.TypeNumberSet,
-		Tags:   tagFirst,
-		F:      Abs,
+		Args:          []models.FuncType{models.TypeVariantSet},
+		VariantReturn: true,
+		Tags:          tagFirst,
+		F:             Abs,
 	},
 	"crop": {
 		Args:   []models.FuncType{models.TypeSeriesSet, models.TypeNumberSet, models.TypeNumberSet},
@@ -303,16 +303,16 @@ var builtins = map[string]parse.Func{
 		F:      Epoch,
 	},
 	"filter": {
-		Args:   []models.FuncType{models.TypeSeriesSet, models.TypeNumberSet},
-		Return: models.TypeSeriesSet,
-		Tags:   tagFirst,
-		F:      Filter,
+		Args:          []models.FuncType{models.TypeVariantSet, models.TypeNumberSet},
+		VariantReturn: true,
+		Tags:          tagFirst,
+		F:             Filter,
 	},
 	"limit": {
-		Args:   []models.FuncType{models.TypeNumberSet, models.TypeScalar},
-		Return: models.TypeNumberSet,
-		Tags:   tagFirst,
-		F:      Limit,
+		Args:          []models.FuncType{models.TypeVariantSet, models.TypeScalar},
+		VariantReturn: true,
+		Tags:          tagFirst,
+		F:             Limit,
 	},
 	"nv": {
 		Args:   []models.FuncType{models.TypeNumberSet, models.TypeScalar},
@@ -553,27 +553,27 @@ func Sort(e *State, T miniprofiler.Timer, series *Results, order string) (*Resul
 	return series, nil
 }
 
-func Limit(e *State, T miniprofiler.Timer, series *Results, v float64) (*Results, error) {
+func Limit(e *State, T miniprofiler.Timer, set *Results, v float64) (*Results, error) {
 	i := int(v)
-	if len(series.Results) > i {
-		series.Results = series.Results[:i]
+	if len(set.Results) > i {
+		set.Results = set.Results[:i]
 	}
-	return series, nil
+	return set, nil
 }
 
-func Filter(e *State, T miniprofiler.Timer, series *Results, number *Results) (*Results, error) {
+func Filter(e *State, T miniprofiler.Timer, set *Results, numberSet *Results) (*Results, error) {
 	var ns ResultSlice
-	for _, sr := range series.Results {
-		for _, nr := range number.Results {
+	for _, sr := range set.Results {
+		for _, nr := range numberSet.Results {
 			if sr.Group.Subset(nr.Group) || nr.Group.Subset(sr.Group) {
 				if nr.Value.Value().(Number) != 0 {
 					ns = append(ns, sr)
 				}
 			}
 		}
 	}
-	series.Results = ns
-	return series, nil
+	set.Results = ns
+	return set, nil
 }
 
 func Tail(e *State, T miniprofiler.Timer, series *Results, number *Results) (*Results, error) {
@@ -634,20 +634,20 @@ func Merge(e *State, T miniprofiler.Timer, series ...*Results) (*Results, error)
 	return res, nil
 }
 
-func Remove(e *State, T miniprofiler.Timer, seriesSet *Results, tagKey string) (*Results, error) {
+func Remove(e *State, T miniprofiler.Timer, set *Results, tagKey string) (*Results, error) {
 	seen := make(map[string]bool)
-	for _, r := range seriesSet.Results {
+	for _, r := range set.Results {
 		if _, ok := r.Group[tagKey]; ok {
 			delete(r.Group, tagKey)
 			if _, ok := seen[r.Group.String()]; ok {
-				return seriesSet, fmt.Errorf("duplicate group would result from removing tag key: %v", tagKey)
+				return set, fmt.Errorf("duplicate group would result from removing tag key: %v", tagKey)
 			}
 			seen[r.Group.String()] = true
 		} else {
-			return seriesSet, fmt.Errorf("tag key %v not found in result", tagKey)
+			return set, fmt.Errorf("tag key %v not found in result", tagKey)
 		}
 	}
-	return seriesSet, nil
+	return set, nil
 }
 
 func LeftJoin(e *State, T miniprofiler.Timer, keysCSV, columnsCSV string, rowData ...*Results) (*Results, error) {
@@ -832,11 +832,18 @@ func reduce(e *State, T miniprofiler.Timer, series *Results, F func(Series, ...f
 	return match(f, series, args...)
 }
 
-func Abs(e *State, T miniprofiler.Timer, series *Results) *Results {
-	for _, s := range series.Results {
-		s.Value = Number(math.Abs(float64(s.Value.Value().(Number))))
+func Abs(e *State, T miniprofiler.Timer, set *Results) *Results {
+	for _, s := range set.Results {
+		switch s.Type() {
+		case models.TypeNumberSet:
+			s.Value = Number(math.Abs(float64(s.Value.Value().(Number))))
+		case models.TypeSeriesSet:
+			for k, v := range s.Value.Value().(Series) {
+				s.Value.Value().(Series)[k] = math.Abs(v)
+			}
+		}
 	}
-	return series
+	return set
 }
 
 func Diff(e *State, T miniprofiler.Timer, series *Results) (r *Results, err error) {
@@ -1152,14 +1159,14 @@ func percentile(dps Series, args ...float64) (a float64) {
 	return x[int(i)]
 }
 
-func Rename(e *State, T miniprofiler.Timer, series *Results, s string) (*Results, error) {
+func Rename(e *State, T miniprofiler.Timer, set *Results, s string) (*Results, error) {
 	for _, section := range strings.Split(s, ",") {
 		kv := strings.Split(section, "=")
 		if len(kv) != 2 {
 			return nil, fmt.Errorf("error passing groups")
 		}
 		oldKey, newKey := kv[0], kv[1]
-		for _, res := range series.Results {
+		for _, res := range set.Results {
 			for tag, v := range res.Group {
 				if oldKey == tag {
 					if _, ok := res.Group[newKey]; ok {
@@ -1172,19 +1179,19 @@ func Rename(e *State, T miniprofiler.Timer, series *Results, s string) (*Results
 			}
 		}
 	}
-	return series, nil
+	return set, nil
 }
 
-func AddTags(e *State, T miniprofiler.Timer, series *Results, s string) (*Results, error) {
+func AddTags(e *State, T miniprofiler.Timer, set *Results, s string) (*Results, error) {
 	if s == "" {
-		return series, nil
+		return set, nil
 	}
 	tagSetToAdd, err := opentsdb.ParseTags(s)
 	if err != nil {
 		return nil, err
 	}
 	for tagKey, tagValue := range tagSetToAdd {
-		for _, res := range series.Results {
+		for _, res := range set.Results {
 			if res.Group == nil {
 				res.Group = make(opentsdb.TagSet)
 			}
@@ -1194,7 +1201,7 @@ func AddTags(e *State, T miniprofiler.Timer, series *Results, s string) (*Result
 			res.Group[tagKey] = tagValue
 		}
 	}
-	return series, nil
+	return set, nil
 }
 
 func Ungroup(e *State, T miniprofiler.Timer, d *Results) (*Results, error) {

cmd/bosun/expr/parse/node.go

@@ -70,13 +70,13 @@ type FuncNode struct {
 	NodeType
 	Pos
 	Name   string
-	F      Func
+	F      *Func
 	Args   []Node
 	Prefix string
 }
 
 func newFunc(pos Pos, name string, f Func) *FuncNode {
-	return &FuncNode{NodeType: NodeFunc, Pos: pos, Name: name, F: f}
+	return &FuncNode{NodeType: NodeFunc, Pos: pos, Name: name, F: &f}
 }
 
 func (f *FuncNode) append(arg Node) {
@@ -133,7 +133,12 @@ func (f *FuncNode) Check(t *Tree) error {
 		}
 		argType := arg.Return()
 		if funcType == models.TypeNumberSet && argType == models.TypeScalar {
-			// Scalars are promoted to NumberSets during execution.
+			argType = models.TypeNumberSet // Scalars are promoted to NumberSets during execution.
+		}
+		if funcType == models.TypeVariantSet {
+			if !(argType == models.TypeNumberSet || argType == models.TypeSeriesSet || argType == models.TypeScalar) {
+				return fmt.Errorf("parse: expected %v or %v for argument %v, got %v", models.TypeNumberSet, models.TypeSeriesSet, i, argType)
+			}
 		} else if funcType != argType {
 			return fmt.Errorf("parse: expected %v, got %v for argument %v (%v)", funcType, argType, i, arg.String())
 		}

cmd/bosun/expr/parse/parse.go

@@ -42,6 +42,7 @@ type Func struct {
 	MapFunc       bool // Func is only valid in map expressions
 	PrefixEnabled bool
 	PrefixKey     bool
+	VariantReturn bool
 	Check         func(*Tree, *FuncNode) error
 }
 
@@ -198,6 +199,12 @@ func (t *Tree) startParse(funcs []map[string]Func, lex *lexer) {
 	t.funcs = funcs
 	for _, funcMap := range funcs {
 		for name, f := range funcMap {
+			if f.VariantReturn {
+				if f.Tags == nil {
+					panic(fmt.Errorf("%v: expected Tags definition: got nil", name))
+				}
+				continue
+			}
 			switch f.Return {
 			case models.TypeSeriesSet, models.TypeNumberSet:
 				if f.Tags == nil {
@@ -374,7 +381,11 @@ func (t *Tree) Func() (f *FuncNode) {
 		switch token = t.next(); token.typ {
 		default:
 			t.backup()
-			f.append(t.O())
+			node := t.O()
+			f.append(node)
+			if len(f.Args) == 1 && f.F.VariantReturn {
+				f.F.Return = node.Return()
+			}
 		case itemTripleQuotedString:
 			f.append(newString(token.pos, token.val, token.val[3:len(token.val)-3]))
 		case itemString: