Implement label_replace()

Implements part of https://github.com/prometheus/prometheus/issues/959.

promql/engine.go

@@ -613,6 +613,16 @@ func (ev *evaluator) evalMatrixBounds(e Expr) Matrix {
 	return ev.matrixSelectorBounds(ms)
 }
 
+// evalString attempts to evaluate e to a string value and errors otherwise.
+func (ev *evaluator) evalString(e Expr) *String {
+	val := ev.eval(e)
+	sv, ok := val.(*String)
+	if !ok {
+		ev.errorf("expected string but got %s", val.Type())
+	}
+	return sv
+}
+
 // evalOneOf evaluates e and errors unless the result is of one of the given types.
 func (ev *evaluator) evalOneOf(e Expr, t1, t2 ExprType) Value {
 	val := ev.eval(e)

promql/functions.go

@@ -16,6 +16,7 @@ package promql
 import (
 	"container/heap"
 	"math"
+	"regexp"
 	"sort"
 	"strconv"
 	"time"
@@ -603,6 +604,50 @@ func funcChanges(ev *evaluator, args Expressions) Value {
 	return out
 }
 
+// === label_replace(vector ExprVector, dst_label, replacement, src_labelname, regex ExprString) Vector ===
+func funcLabelReplace(ev *evaluator, args Expressions) Value {
+	var (
+		vector   = ev.evalVector(args[0])
+		dst      = clientmodel.LabelName(ev.evalString(args[1]).Value)
+		repl     = ev.evalString(args[2]).Value
+		src      = clientmodel.LabelName(ev.evalString(args[3]).Value)
+		regexStr = ev.evalString(args[4]).Value
+	)
+
+	regex, err := regexp.Compile(regexStr)
+	if err != nil {
+		ev.errorf("invalid regular expression in label_replace(): %s", regexStr)
+	}
+	if !clientmodel.LabelNameRE.MatchString(string(dst)) {
+		ev.errorf("invalid destination label name in label_replace(): %s", dst)
+	}
+
+	outSet := make(map[clientmodel.Fingerprint]struct{}, len(vector))
+	for _, el := range vector {
+		srcVal := string(el.Metric.Metric[src])
+		indexes := regex.FindStringSubmatchIndex(srcVal)
+		// If there is no match, no replacement should take place.
+		if indexes == nil {
+			continue
+		}
+		res := regex.ExpandString([]byte{}, repl, srcVal, indexes)
+		if len(res) == 0 {
+			el.Metric.Delete(dst)
+		} else {
+			el.Metric.Set(dst, clientmodel.LabelValue(res))
+		}
+
+		fp := el.Metric.Metric.Fingerprint()
+		if _, exists := outSet[fp]; exists {
+			ev.errorf("duplicated label set in output of label_replace(): %s", el.Metric.Metric)
+		} else {
+			outSet[fp] = struct{}{}
+		}
+	}
+
+	return vector
+}
+
 var functions = map[string]*Function{
 	"abs": {
 		Name:       "abs",
@@ -695,6 +740,12 @@ var functions = map[string]*Function{
 		ReturnType: ExprVector,
 		Call:       funcHistogramQuantile,
 	},
+	"label_replace": {
+		Name:       "label_replace",
+		ArgTypes:   []ExprType{ExprVector, ExprString, ExprString, ExprString, ExprString},
+		ReturnType: ExprVector,
+		Call:       funcLabelReplace,
+	},
 	"ln": {
 		Name:       "ln",
 		ArgTypes:   []ExprType{ExprVector},

promql/testdata/functions.test

@@ -71,24 +71,73 @@ load 5m
 	testcounter_reset_middle	0+10x4 0+10x5
 	http_requests{job="app-server", instance="1", group="canary"}		0+80x10
 
-# Deriv should return the same as rate in simple cases.
+# deriv should return the same as rate in simple cases.
 eval instant at 50m rate(http_requests{group="canary", instance="1", job="app-server"}[60m])
 	{group="canary", instance="1", job="app-server"} 0.26666666666666666
 
 eval instant at 50m deriv(http_requests{group="canary", instance="1", job="app-server"}[60m])
 	{group="canary", instance="1", job="app-server"} 0.26666666666666666
 
-# Deriv should return correct result.
+# deriv should return correct result.
 eval instant at 50m deriv(testcounter_reset_middle[100m])
 	{} 0.010606060606060607
 
-# Predict_linear should return correct result.
+# predict_linear should return correct result.
 eval instant at 50m predict_linear(testcounter_reset_middle[100m], 3600)
 	{} 88.181818181818185200
 
-# Predict_linear is syntactic sugar around deriv.
+# predict_linear is syntactic sugar around deriv.
 eval instant at 50m predict_linear(http_requests[50m], 3600) - (http_requests + deriv(http_requests[50m]) * 3600)
 	{group="canary", instance="1", job="app-server"} 0
 
 eval instant at 50m predict_linear(testcounter_reset_middle[100m], 3600) - (testcounter_reset_middle + deriv(testcounter_reset_middle[100m]) * 3600)
 	{} 0
+
+clear
+
+# Tests for label_replace.
+load 5m
+  testmetric{src="source-value-10",dst="original-destination-value"} 0
+  testmetric{src="source-value-20",dst="original-destination-value"} 1
+
+# label_replace does a substring match and replace.
+eval instant at 0m label_replace(testmetric, "dst", "destination-value-$1", "src", "value-(.*)")
+  testmetric{src="source-value-10",dst="destination-value-10"} 0
+  testmetric{src="source-value-20",dst="destination-value-20"} 1
+
+# label_replace works with multiple capture groups.
+eval instant at 0m label_replace(testmetric, "dst", "$1-value-$2", "src", "(.*)-value-(.*)")
+  testmetric{src="source-value-10",dst="source-value-10"} 0
+  testmetric{src="source-value-20",dst="source-value-20"} 1
+
+# label_replace does not overwrite the destination label if the source label
+# does not exist.
+eval instant at 0m label_replace(testmetric, "dst", "value-$1", "nonexistent-src", "source-value-(.*)")
+  testmetric{src="source-value-10",dst="original-destination-value"} 0
+  testmetric{src="source-value-20",dst="original-destination-value"} 1
+
+# label_replace overwrites the destination label if the source label is empty,
+# but matched.
+eval instant at 0m label_replace(testmetric, "dst", "value-$1", "nonexistent-src", "(.*)")
+  testmetric{src="source-value-10",dst="value-"} 0
+  testmetric{src="source-value-20",dst="value-"} 1
+
+# label_replace does not overwrite the destination label if the source label
+# is not matched.
+eval instant at 0m label_replace(testmetric, "dst", "value-$1", "src", "non-matching-regex")
+  testmetric{src="source-value-10",dst="original-destination-value"} 0
+  testmetric{src="source-value-20",dst="original-destination-value"} 1
+
+# label_replace drops labels that are set to empty values.
+eval instant at 0m label_replace(testmetric, "dst", "", "dst", "")
+  testmetric{src="source-value-10"} 0
+  testmetric{src="source-value-20"} 1
+
+# label_replace fails when the regex is invalid.
+eval_fail instant at 0m label_replace(testmetric, "dst", "value-$1", "src", "(.*")
+
+# label_replace fails when the destination label name is not a valid Prometheus label name.
+eval_fail instant at 0m label_replace(testmetric, "invalid-label-name", "", "src", "(.*)")
+
+# label_replace fails when there would be duplicated identical output label sets.
+eval_fail instant at 0m label_replace(testmetric, "src", "", "", "")