promql: faster range-query of label_replace and label_join

These functions act on the labels only, so don't need to go step by step
over the samples in a range query.

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>

promql/engine.go

@@ -1409,6 +1409,15 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 				break
 			}
 		}
+
+		// Special handling for functions that work on series not samples.
+		switch e.Func.Name {
+		case "label_replace":
+			return ev.evalLabelReplace(e.Args)
+		case "label_join":
+			return ev.evalLabelJoin(e.Args)
+		}
+
 		if !matrixArg {
 			// Does not have a matrix argument.
 			return ev.rangeEval(nil, func(v []parser.Value, _ [][]EvalSeriesHelper, enh *EvalNodeHelper) (Vector, annotations.Annotations) {

promql/functions.go

@@ -1321,59 +1321,47 @@ func funcChanges(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelp
 	return append(enh.Out, Sample{F: float64(changes)}), nil
 }
 
-// === label_replace(Vector parser.ValueTypeVector, dst_label, replacement, src_labelname, regex parser.ValueTypeString) (Vector, Annotations) ===
-func funcLabelReplace(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
+// label_replace function operates only on series; does not look at timestamps or values.
+func (ev *evaluator) evalLabelReplace(args parser.Expressions) (parser.Value, annotations.Annotations) {
 	var (
-		vector   = vals[0].(Vector)
 		dst      = stringFromArg(args[1])
 		repl     = stringFromArg(args[2])
 		src      = stringFromArg(args[3])
 		regexStr = stringFromArg(args[4])
 	)
 
-	if enh.regex == nil {
-		var err error
-		enh.regex, err = regexp.Compile("^(?:" + regexStr + ")$")
+	regex, err := regexp.Compile("^(?:" + regexStr + ")$")
 	if err != nil {
 		panic(fmt.Errorf("invalid regular expression in label_replace(): %s", regexStr))
 	}
 	if !model.LabelNameRE.MatchString(dst) {
 		panic(fmt.Errorf("invalid destination label name in label_replace(): %s", dst))
 	}
-		enh.Dmn = make(map[uint64]labels.Labels, len(enh.Out))
-	}
 
-	for _, el := range vector {
-		h := el.Metric.Hash()
-		var outMetric labels.Labels
-		if l, ok := enh.Dmn[h]; ok {
-			outMetric = l
-		} else {
-			srcVal := el.Metric.Get(src)
-			indexes := enh.regex.FindStringSubmatchIndex(srcVal)
-			if indexes == nil {
-				// If there is no match, no replacement should take place.
-				outMetric = el.Metric
-				enh.Dmn[h] = outMetric
-			} else {
-				res := enh.regex.ExpandString([]byte{}, repl, srcVal, indexes)
+	val, ws := ev.eval(args[0])
+	matrix := val.(Matrix)
+	lb := labels.NewBuilder(labels.EmptyLabels())
 
-				lb := labels.NewBuilder(el.Metric).Del(dst)
-				if len(res) > 0 {
+	for i, el := range matrix {
+		srcVal := el.Metric.Get(src)
+		indexes := regex.FindStringSubmatchIndex(srcVal)
+		if indexes != nil { // Only replace when regexp matches.
+			res := regex.ExpandString([]byte{}, repl, srcVal, indexes)
+			lb.Reset(el.Metric)
 			lb.Set(dst, string(res))
+			matrix[i].Metric = lb.Labels()
 		}
-				outMetric = lb.Labels()
-				enh.Dmn[h] = outMetric
 	}
+	if matrix.ContainsSameLabelset() {
+		ev.errorf("vector cannot contain metrics with the same labelset")
 	}
 
-		enh.Out = append(enh.Out, Sample{
-			Metric: outMetric,
-			F:      el.F,
-			H:      el.H,
-		})
+	return matrix, ws
 }
-	return enh.Out, nil
+
+// === label_replace(Vector parser.ValueTypeVector, dst_label, replacement, src_labelname, regex parser.ValueTypeString) (Vector, Annotations) ===
+func funcLabelReplace(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
+	panic("funcLabelReplace wrong implementation called")
 }
 
 // === Vector(s Scalar) (Vector, Annotations) ===
@@ -1385,19 +1373,13 @@ func funcVector(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelpe
 		}), nil
 }
 
-// === label_join(vector model.ValVector, dest_labelname, separator, src_labelname...) (Vector, Annotations) ===
-func funcLabelJoin(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
+// label_join function operates only on series; does not look at timestamps or values.
+func (ev *evaluator) evalLabelJoin(args parser.Expressions) (parser.Value, annotations.Annotations) {
 	var (
-		vector    = vals[0].(Vector)
 		dst       = stringFromArg(args[1])
 		sep       = stringFromArg(args[2])
 		srcLabels = make([]string, len(args)-3)
 	)
-
-	if enh.Dmn == nil {
-		enh.Dmn = make(map[uint64]labels.Labels, len(enh.Out))
-	}
-
 	for i := 3; i < len(args); i++ {
 		src := stringFromArg(args[i])
 		if !model.LabelName(src).IsValid() {
@@ -1406,42 +1388,27 @@ func funcLabelJoin(vals []parser.Value, args parser.Expressions, enh *EvalNodeHe
 		srcLabels[i-3] = src
 	}
 
-	if !model.LabelName(dst).IsValid() {
-		panic(fmt.Errorf("invalid destination label name in label_join(): %s", dst))
-	}
-
+	val, ws := ev.eval(args[0])
+	matrix := val.(Matrix)
 	srcVals := make([]string, len(srcLabels))
-	for _, el := range vector {
-		h := el.Metric.Hash()
-		var outMetric labels.Labels
-		if l, ok := enh.Dmn[h]; ok {
-			outMetric = l
-		} else {
+	lb := labels.NewBuilder(labels.EmptyLabels())
 
+	for i, el := range matrix {
 		for i, src := range srcLabels {
 			srcVals[i] = el.Metric.Get(src)
 		}
-
-			lb := labels.NewBuilder(el.Metric)
-
 		strval := strings.Join(srcVals, sep)
-			if strval == "" {
-				lb.Del(dst)
-			} else {
+		lb.Reset(el.Metric)
 		lb.Set(dst, strval)
+		matrix[i].Metric = lb.Labels()
 	}
 
-			outMetric = lb.Labels()
-			enh.Dmn[h] = outMetric
+	return matrix, ws
 }
 
-		enh.Out = append(enh.Out, Sample{
-			Metric: outMetric,
-			F:      el.F,
-			H:      el.H,
-		})
-	}
-	return enh.Out, nil
+// === label_join(vector model.ValVector, dest_labelname, separator, src_labelname...) (Vector, Annotations) ===
+func funcLabelJoin(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper) (Vector, annotations.Annotations) {
+	panic("funcLabelReplace wrong implementation called")
 }
 
 // Common code for date related functions.