From 808ce815e49fc07f014a63d97a61d97e7eade9c5 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Sat, 10 Dec 2022 22:47:30 -0800
Subject: [PATCH] app/vmselect/promql: allow passing `inf` arg into functions,
 which accept numeric limit on the number of output time series

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3461
---
 app/vmselect/promql/aggr.go      | 18 ++++++++++++------
 app/vmselect/promql/exec_test.go | 24 ++++++++++++++++++++++++
 app/vmselect/promql/rollup.go    |  2 +-
 app/vmselect/promql/transform.go |  9 +++------
 docs/CHANGELOG.md                |  1 +
 5 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/app/vmselect/promql/aggr.go b/app/vmselect/promql/aggr.go
index 07573fe3ce..9baa1efdc4 100644
--- a/app/vmselect/promql/aggr.go
+++ b/app/vmselect/promql/aggr.go
@@ -748,7 +748,7 @@ func getIntK(k float64, kMax int) int {
 	if math.IsNaN(k) {
 		return 0
 	}
-	kn := int(k)
+	kn := floatToIntBounded(k)
 	if kn < 0 {
 		return 0
 	}
@@ -999,14 +999,10 @@ func aggrFuncLimitK(afa *aggrFuncArg) ([]*timeseries, error) {
 	if err := expectTransformArgsNum(args, 2); err != nil {
 		return nil, err
 	}
-	limits, err := getScalar(args[0], 0)
+	limit, err := getIntNumber(args[0], 0)
 	if err != nil {
 		return nil, fmt.Errorf("cannot obtain limit arg: %w", err)
 	}
-	limit := 0
-	if len(limits) > 0 {
-		limit = int(limits[0])
-	}
 	if limit < 0 {
 		limit = 0
 	}
@@ -1155,3 +1151,13 @@ func lessWithNaNs(a, b float64) bool {
 	}
 	return a < b
 }
+
+func floatToIntBounded(f float64) int {
+	if f > math.MaxInt {
+		return math.MaxInt
+	}
+	if f < math.MinInt {
+		return math.MinInt
+	}
+	return int(f)
+}
diff --git a/app/vmselect/promql/exec_test.go b/app/vmselect/promql/exec_test.go
index 8f286eb45e..b669ef8ecf 100644
--- a/app/vmselect/promql/exec_test.go
+++ b/app/vmselect/promql/exec_test.go
@@ -5559,6 +5559,30 @@ func TestExecSuccess(t *testing.T) {
 		resultExpected := []netstorage.Result{r1, r2}
 		f(q, resultExpected)
 	})
+	t.Run(`limitk(inf)`, func(t *testing.T) {
+		t.Parallel()
+		q := `sort(limitk(inf, label_set(10, "foo", "bar") or label_set(time()/150, "baz", "sss")))`
+		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"),
+		}}
+		r2 := netstorage.Result{
+			MetricName: metricNameExpected,
+			Values:     []float64{6.666666666666667, 8, 9.333333333333334, 10.666666666666666, 12, 13.333333333333334},
+			Timestamps: timestampsExpected,
+		}
+		r2.MetricName.Tags = []storage.Tag{{
+			Key:   []byte("baz"),
+			Value: []byte("sss"),
+		}}
+		resultExpected := []netstorage.Result{r1, r2}
+		f(q, resultExpected)
+	})
 	t.Run(`any()`, func(t *testing.T) {
 		t.Parallel()
 		q := `any(label_set(10, "__name__", "x", "foo", "bar") or label_set(time()/150, "__name__", "y", "baz", "sss"))`
diff --git a/app/vmselect/promql/rollup.go b/app/vmselect/promql/rollup.go
index 587c54137b..06603aac27 100644
--- a/app/vmselect/promql/rollup.go
+++ b/app/vmselect/promql/rollup.go
@@ -2127,7 +2127,7 @@ func getIntNumber(arg interface{}, argNum int) (int, error) {
 	}
 	n := 0
 	if len(v) > 0 {
-		n = int(v[0])
+		n = floatToIntBounded(v[0])
 	}
 	return n, nil
 }
diff --git a/app/vmselect/promql/transform.go b/app/vmselect/promql/transform.go
index 752775b4ed..80bc96bef1 100644
--- a/app/vmselect/promql/transform.go
+++ b/app/vmselect/promql/transform.go
@@ -371,14 +371,10 @@ func transformBucketsLimit(tfa *transformFuncArg) ([]*timeseries, error) {
 	if err := expectTransformArgsNum(args, 2); err != nil {
 		return nil, err
 	}
-	limits, err := getScalar(args[0], 1)
+	limit, err := getIntNumber(args[0], 0)
 	if err != nil {
 		return nil, err
 	}
-	limit := 0
-	if len(limits) > 0 {
-		limit = int(limits[0])
-	}
 	if limit <= 0 {
 		return nil, nil
 	}
@@ -390,6 +386,7 @@ func transformBucketsLimit(tfa *transformFuncArg) ([]*timeseries, error) {
 	if len(tss) == 0 {
 		return nil, nil
 	}
+	pointsCount := len(tss[0].Values)
 
 	// Group timeseries by all MetricGroup+tags excluding `le` tag.
 	type x struct {
@@ -437,7 +434,7 @@ func transformBucketsLimit(tfa *transformFuncArg) ([]*timeseries, error) {
 		sort.Slice(leGroup, func(i, j int) bool {
 			return leGroup[i].le < leGroup[j].le
 		})
-		for n := range limits {
+		for n := 0; n < pointsCount; n++ {
 			prevValue := float64(0)
 			for i := range leGroup {
 				xx := &leGroup[i]
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index b90d79a437..8078517338 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -64,6 +64,7 @@ The following tip changes can be tested by building VictoriaMetrics components f
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): allow changing timezones for the requested data. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3075).
 * FEATURE: [vmui](https://docs.victoriametrics.com/#vmui): provide fast path for hiding results for all the queries except the given one by clicking `eye` icon with `ctrl` key pressed. See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3446).
 * FEATURE: [MetricsQL](https://docs.victoriametrics.com/MetricsQL.html): add `range_trim_spikes(phi, q)` function for trimming `phi` percent of the largest spikes per each time series returned by `q`. See [these docs](https://docs.victoriametrics.com/MetricsQL.html#range_trim_spikes).
+* FEATURE: [MetricsQL](https://docs.victoriametrics.com/MetricsQL.html): allow passing `inf` arg into [limitk](https://docs.victoriametrics.com/MetricsQL.html#limitk), [topk](https://docs.victoriametrics.com/MetricsQL.html#topk), [bottomk](https://docs.victoriametrics.com/MetricsQL.html) and other functions, which accept numeric arg, which limits the number of output time series. See [this feature request](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/3461).
 * FEATURE: [vmgateway](https://docs.victoriametrics.com/vmgateway.html): add support for JWT token signature verification. See [these docs](https://docs.victoriametrics.com/vmgateway.html#jwt-signature-verification) for details.
 * FEATURE: put the version of VictoriaMetrics in the first message of [query trace](https://docs.victoriametrics.com/#query-tracing). This should simplify debugging.