From 73b2a3d4b7afa2f5eafe91bb19516a5cb300cb28 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 11 Dec 2019 13:55:18 +0200
Subject: [PATCH] app/vmselect/promql: return `lower` and `upper` bounds for
 the estimated percentile from `histogram_quantile` if third arg is passed

Updates https://github.com/prometheus/prometheus/issues/5706
---
 app/vmselect/promql/exec_test.go | 79 ++++++++++++++++++++++++++++++++
 app/vmselect/promql/transform.go | 59 ++++++++++++++++++------
 docs/ExtendedPromQL.md           |  1 +
 3 files changed, 126 insertions(+), 13 deletions(-)

diff --git a/app/vmselect/promql/exec_test.go b/app/vmselect/promql/exec_test.go
index 27e2d623f0..aadbbb33ca 100644
--- a/app/vmselect/promql/exec_test.go
+++ b/app/vmselect/promql/exec_test.go
@@ -2324,6 +2324,35 @@ func TestExecSuccess(t *testing.T) {
 		resultExpected := []netstorage.Result{r}
 		f(q, resultExpected)
 	})
+	t.Run(`histogram_quantile(single-value-valid-le, boundsLabel)`, func(t *testing.T) {
+		t.Parallel()
+		q := `sort(histogram_quantile(0.6, label_set(100, "le", "200"), "foobar"))`
+		r1 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{0, 0, 0, 0, 0, 0},
+			Timestamps: timestampsExpected,
+		}
+		r1.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("foobar"),
+			Value: []byte("lower"),
+		}}
+		r2 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{120, 120, 120, 120, 120, 120},
+			Timestamps: timestampsExpected,
+		}
+		r3 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{200, 200, 200, 200, 200, 200},
+			Timestamps: timestampsExpected,
+		}
+		r3.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("foobar"),
+			Value: []byte("upper"),
+		}}
+		resultExpected := []netstorage.Result{r1, r2, r3}
+		f(q, resultExpected)
+	})
 	t.Run(`histogram_quantile(single-value-valid-le-max-phi)`, func(t *testing.T) {
 		t.Parallel()
 		q := `histogram_quantile(1, (
@@ -2461,6 +2490,56 @@ func TestExecSuccess(t *testing.T) {
 		resultExpected := []netstorage.Result{r}
 		f(q, resultExpected)
 	})
+	t.Run(`histogram_quantile(normal-bucket-count, boundsLabel)`, func(t *testing.T) {
+		t.Parallel()
+		q := `sort(histogram_quantile(0.2,
+			label_set(0, "foo", "bar", "le", "10")
+			or label_set(100, "foo", "bar", "le", "30")
+			or label_set(300, "foo", "bar", "le", "+Inf"),
+			"xxx"
+		))`
+		r1 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{10, 10, 10, 10, 10, 10},
+			Timestamps: timestampsExpected,
+		}
+		r1.MetricName.Tags = []storage.Tag{
+			{
+				Key:   []byte("foo"),
+				Value: []byte("bar"),
+			},
+			{
+				Key:   []byte("xxx"),
+				Value: []byte("lower"),
+			},
+		}
+		r2 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{22, 22, 22, 22, 22, 22},
+			Timestamps: timestampsExpected,
+		}
+		r2.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("foo"),
+			Value: []byte("bar"),
+		}}
+		r3 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{30, 30, 30, 30, 30, 30},
+			Timestamps: timestampsExpected,
+		}
+		r3.MetricName.Tags = []storage.Tag{
+			{
+				Key:   []byte("foo"),
+				Value: []byte("bar"),
+			},
+			{
+				Key:   []byte("xxx"),
+				Value: []byte("upper"),
+			},
+		}
+		resultExpected := []netstorage.Result{r1, r2, r3}
+		f(q, resultExpected)
+	})
 	t.Run(`histogram_quantile(zero-bucket-count)`, func(t *testing.T) {
 		t.Parallel()
 		q := `histogram_quantile(0.6,
diff --git a/app/vmselect/promql/transform.go b/app/vmselect/promql/transform.go
index e89c0ddbb3..056e7a0a5b 100644
--- a/app/vmselect/promql/transform.go
+++ b/app/vmselect/promql/transform.go
@@ -400,17 +400,27 @@ func vmrangeBucketsToLE(tss []*timeseries) []*timeseries {
 
 func transformHistogramQuantile(tfa *transformFuncArg) ([]*timeseries, error) {
 	args := tfa.args
-	if err := expectTransformArgsNum(args, 2); err != nil {
-		return nil, err
+	if len(args) < 2 || len(args) > 3 {
+		return nil, fmt.Errorf("unexpected number of args; got %d; want 2...3", len(args))
 	}
 	phis, err := getScalar(args[0], 0)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("cannot parse phi: %s", err)
 	}
 
 	// Convert buckets with `vmrange` labels to buckets with `le` labels.
 	tss := vmrangeBucketsToLE(args[1])
 
+	// Parse boundsLabel. See https://github.com/prometheus/prometheus/issues/5706 for details.
+	var boundsLabel string
+	if len(args) > 2 {
+		s, err := getString(args[2], 2)
+		if err != nil {
+			return nil, fmt.Errorf("cannot parse boundsLabel (arg #3): %s", err)
+		}
+		boundsLabel = s
+	}
+
 	// Group metrics by all tags excluding "le"
 	type x struct {
 		le float64
@@ -453,10 +463,10 @@ func transformHistogramQuantile(tfa *transformFuncArg) ([]*timeseries, error) {
 		}
 		return nan
 	}
-	quantile := func(i int, phis []float64, xss []x) float64 {
+	quantile := func(i int, phis []float64, xss []x) (q, lower, upper float64) {
 		phi := phis[i]
 		if math.IsNaN(phi) {
-			return nan
+			return nan, nan, nan
 		}
 		// Fix broken buckets.
 		// They are already sorted by le, so their values must be in ascending order,
@@ -479,13 +489,13 @@ func transformHistogramQuantile(tfa *transformFuncArg) ([]*timeseries, error) {
 			xss = xss[:len(xss)-1]
 		}
 		if vLast == 0 || math.IsNaN(vLast) {
-			return nan
+			return nan, nan, nan
 		}
 		if phi < 0 {
-			return -inf
+			return -inf, -inf, xss[0].ts.Values[i]
 		}
 		if phi > 1 {
-			return inf
+			return inf, vLast, inf
 		}
 		vReq := vLast * phi
 		vPrev = 0
@@ -509,14 +519,17 @@ func transformHistogramQuantile(tfa *transformFuncArg) ([]*timeseries, error) {
 				continue
 			}
 			if math.IsInf(le, 0) {
-				return lastNonInf(i, xss)
+				vv := lastNonInf(i, xss)
+				return vv, vv, inf
 			}
 			if v == vPrev {
-				return lePrev
+				return lePrev, lePrev, v
 			}
-			return lePrev + (le-lePrev)*(vReq-vPrev)/(v-vPrev)
+			vv := lePrev + (le-lePrev)*(vReq-vPrev)/(v-vPrev)
+			return vv, lePrev, le
 		}
-		return lastNonInf(i, xss)
+		vv := lastNonInf(i, xss)
+		return vv, vv, inf
 	}
 	rvs := make([]*timeseries, 0, len(m))
 	for _, xss := range m {
@@ -524,10 +537,30 @@ func transformHistogramQuantile(tfa *transformFuncArg) ([]*timeseries, error) {
 			return xss[i].le < xss[j].le
 		})
 		dst := xss[0].ts
+		var tsLower, tsUpper *timeseries
+		if len(boundsLabel) > 0 {
+			tsLower = &timeseries{}
+			tsLower.CopyFromShallowTimestamps(dst)
+			tsLower.MetricName.RemoveTag(boundsLabel)
+			tsLower.MetricName.AddTag(boundsLabel, "lower")
+			tsUpper = &timeseries{}
+			tsUpper.CopyFromShallowTimestamps(dst)
+			tsUpper.MetricName.RemoveTag(boundsLabel)
+			tsUpper.MetricName.AddTag(boundsLabel, "upper")
+		}
 		for i := range dst.Values {
-			dst.Values[i] = quantile(i, phis, xss)
+			v, lower, upper := quantile(i, phis, xss)
+			dst.Values[i] = v
+			if len(boundsLabel) > 0 {
+				tsLower.Values[i] = lower
+				tsUpper.Values[i] = upper
+			}
 		}
 		rvs = append(rvs, dst)
+		if len(boundsLabel) > 0 {
+			rvs = append(rvs, tsLower)
+			rvs = append(rvs, tsUpper)
+		}
 	}
 	return rvs, nil
 }
diff --git a/docs/ExtendedPromQL.md b/docs/ExtendedPromQL.md
index bbd60eced2..e8fd955783 100644
--- a/docs/ExtendedPromQL.md
+++ b/docs/ExtendedPromQL.md
@@ -12,6 +12,7 @@ Try these extensions on [an editable Grafana dashboard](http://play-grafana.vict
 - `offset` may be put anywere in the query. For instance, `sum(foo) offset 24h`.
 - `offset` may be negative. For example, `q offset -1h`.
 - `default` binary operator. `q1 default q2` substitutes `NaN` values from `q1` with the corresponding values from `q2`.
+- `histogram_quantile` accepts optional third arg - `boundsLabel`. In this case it returns `lower` and `upper` bounds for the estimated percentile. See [this issue for details](https://github.com/prometheus/prometheus/issues/5706).
 - `if` binary operator. `q1 if q2` removes values from `q1` for `NaN` values from `q2`.
 - `ifnot` binary operator. `q1 ifnot q2` removes values from `q1` for non-`NaN` values from `q2`.
 - Trailing commas on all the lists are allowed - label filters, function args and with expressions. For instance, the following queries are valid: `m{foo="bar",}`, `f(a, b,)`, `WITH (x=y,) x`. This simplifies maintenance of multi-line queries.