From fdc2a9d1d72fb24796eec33675bf431a79afb218 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Thu, 12 Mar 2020 19:13:28 +0200
Subject: [PATCH] app/vmselect: add `label_map(q, label, srcValue1, dstValue1,
 ... srcValueN, dstValueN)` function to MetricsQL

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/369
---
 app/vmselect/promql/exec_test.go | 43 ++++++++++++++++++++++++++++++++
 app/vmselect/promql/transform.go | 31 +++++++++++++++++++++++
 docs/MetricsQL.md                |  1 +
 lib/metricsql/transform.go       |  1 +
 4 files changed, 76 insertions(+)

diff --git a/app/vmselect/promql/exec_test.go b/app/vmselect/promql/exec_test.go
index 2b3f4e035..0560dee79 100644
--- a/app/vmselect/promql/exec_test.go
+++ b/app/vmselect/promql/exec_test.go
@@ -974,6 +974,49 @@ func TestExecSuccess(t *testing.T) {
 		resultExpected := []netstorage.Result{r}
 		f(q, resultExpected)
 	})
+	t.Run(`label_map(match)`, func(t *testing.T) {
+		t.Parallel()
+		q := `sort(label_map((
+			label_set(time(), "label", "v1"),
+			label_set(time()+100, "label", "v2"),
+			label_set(time()+200, "label", "v3"),
+			label_set(time()+300, "x", "y"),
+		), "label", "v1", "foo", "v2", "bar"))`
+		r1 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1000, 1200, 1400, 1600, 1800, 2000},
+			Timestamps: timestampsExpected,
+		}
+		r1.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("label"),
+			Value: []byte("foo"),
+		}}
+		r2 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1100, 1300, 1500, 1700, 1900, 2100},
+			Timestamps: timestampsExpected,
+		}
+		r2.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("label"),
+			Value: []byte("bar"),
+		}}
+		r3 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1200, 1400, 1600, 1800, 2000, 2200},
+			Timestamps: timestampsExpected,
+		}
+		r4 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{1300, 1500, 1700, 1900, 2100, 2300},
+			Timestamps: timestampsExpected,
+		}
+		r4.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("x"),
+			Value: []byte("y"),
+		}}
+		resultExpected := []netstorage.Result{r1, r2, r3, r4}
+		f(q, resultExpected)
+	})
 	t.Run(`label_copy(new_tag)`, func(t *testing.T) {
 		t.Parallel()
 		q := `label_copy(
diff --git a/app/vmselect/promql/transform.go b/app/vmselect/promql/transform.go
index 54c1dbf7e..0f177d64f 100644
--- a/app/vmselect/promql/transform.go
+++ b/app/vmselect/promql/transform.go
@@ -59,6 +59,7 @@ var transformFuncs = map[string]transformFunc{
 
 	// New funcs
 	"label_set":          transformLabelSet,
+	"label_map":          transformLabelMap,
 	"label_del":          transformLabelDel,
 	"label_keep":         transformLabelKeep,
 	"label_copy":         transformLabelCopy,
@@ -1026,6 +1027,36 @@ func transformLabelSet(tfa *transformFuncArg) ([]*timeseries, error) {
 	return rvs, nil
 }
 
+func transformLabelMap(tfa *transformFuncArg) ([]*timeseries, error) {
+	args := tfa.args
+	if len(args) < 2 {
+		return nil, fmt.Errorf(`not enough args; got %d; want at least %d`, len(args), 2)
+	}
+	label, err := getString(args[1], 1)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read label name: %s", err)
+	}
+	srcValues, dstValues, err := getStringPairs(args[2:])
+	if err != nil {
+		return nil, err
+	}
+	m := make(map[string]string, len(srcValues))
+	for i, srcValue := range srcValues {
+		m[srcValue] = dstValues[i]
+	}
+	rvs := args[0]
+	for _, ts := range rvs {
+		mn := &ts.MetricName
+		dstValue := getDstValue(mn, label)
+		value := m[string(*dstValue)]
+		*dstValue = append((*dstValue)[:0], value...)
+		if len(value) == 0 {
+			mn.RemoveTag(label)
+		}
+	}
+	return rvs, nil
+}
+
 func transformLabelCopy(tfa *transformFuncArg) ([]*timeseries, error) {
 	return transformLabelCopyExt(tfa, false)
 }
diff --git a/docs/MetricsQL.md b/docs/MetricsQL.md
index de4cfc506..08b163e8c 100644
--- a/docs/MetricsQL.md
+++ b/docs/MetricsQL.md
@@ -46,6 +46,7 @@ This functionality can be tried at [an editable Grafana dashboard](http://play-g
 - Functions for label manipulation:
   - `alias(q, name)` for setting metric name across all the time series `q`.
   - `label_set(q, label1, value1, ... labelN, valueN)` for setting the given values for the given labels on `q`.
+  - `label_map(q, label, srcValue1, dstValue1, ... srcValueN, dstValueN)` for mapping `label` values from `src*` to `dst*`.
   - `label_del(q, label1, ... labelN)` for deleting the given labels from `q`.
   - `label_keep(q, label1, ... labelN)` for deleting all the labels except the given labels from `q`.
   - `label_copy(q, src_label1, dst_label1, ... src_labelN, dst_labelN)` for copying label values from `src_*` to `dst_*`.
diff --git a/lib/metricsql/transform.go b/lib/metricsql/transform.go
index 082eaeeb5..5b5d6e03a 100644
--- a/lib/metricsql/transform.go
+++ b/lib/metricsql/transform.go
@@ -38,6 +38,7 @@ var transformFuncs = map[string]bool{
 
 	// New funcs from MetricsQL
 	"label_set":          true,
+	"label_map":          true,
 	"label_del":          true,
 	"label_keep":         true,
 	"label_copy":         true,