From 1a0dca26eb9107a92db27733b56dee61b008761e Mon Sep 17 00:00:00 2001
From: Evan Yin <yyin@snapchat.com>
Date: Wed, 4 Aug 2021 17:59:26 -0700
Subject: [PATCH] [query] Add Graphite round function (#3648)

---
 .../graphite/native/builtin_functions.go      | 36 ++++++++++++-
 .../graphite/native/builtin_functions_test.go | 54 +++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/src/query/graphite/native/builtin_functions.go b/src/query/graphite/native/builtin_functions.go
index fc81efd3a0..441f804dad 100644
--- a/src/query/graphite/native/builtin_functions.go
+++ b/src/query/graphite/native/builtin_functions.go
@@ -613,6 +613,33 @@ func keepLastValue(ctx *common.Context, input singlePathSpec, limit int) (ts.Ser
 	return r, nil
 }
 
+func roundFunction(ctx *common.Context, input singlePathSpec, precision int) (ts.SeriesList, error) {
+	output := make([]*ts.Series, 0, len(input.Values))
+	for _, series := range input.Values {
+		numSteps := series.Len()
+		vals := ts.NewValues(ctx, series.MillisPerStep(), numSteps)
+		for i := 0; i < numSteps; i++ {
+			value := series.ValueAt(i)
+			if !math.IsNaN(value) {
+				value = roundTo(value, int32(precision))
+			}
+			vals.SetValueAt(i, value)
+		}
+		name := ""
+		if precision == 0 {
+			name = fmt.Sprintf("roundFunction(%s)", series.Name())
+		} else {
+			name = fmt.Sprintf("roundFunction(%s,%d)", series.Name(), precision)
+		}
+		newSeries := ts.NewSeries(ctx, name, series.StartTime(), vals)
+		output = append(output, newSeries)
+	}
+
+	r := ts.SeriesList(input)
+	r.Values = output
+	return r, nil
+}
+
 type comparator func(float64, float64) bool
 
 // equalFunc checks whether x is equal to y
@@ -1011,7 +1038,10 @@ func (impl *exponentialMovingAverageImpl) Evaluate(
 	}
 }
 
-func roundTo(n float64, decimals uint32) float64 {
+func roundTo(n float64, decimals int32) float64 {
+	if decimals < 0 {
+		return math.Round(n/math.Pow(10, math.Abs(float64(decimals)))) * math.Pow(10, math.Abs(float64(decimals)))
+	}
 	return math.Round(n*math.Pow(10, float64(decimals))) / math.Pow(10, float64(decimals))
 }
 
@@ -2942,6 +2972,9 @@ func init() {
 	MustRegisterFunction(removeEmptySeries).WithDefaultParams(map[uint8]interface{}{
 		2: 0.0, // xFilesFactor
 	})
+	MustRegisterFunction(roundFunction).WithDefaultParams(map[uint8]interface{}{
+		2: 0, // precision
+	})
 	MustRegisterFunction(scale)
 	MustRegisterFunction(scaleToSeconds)
 	MustRegisterFunction(sortBy).WithDefaultParams(map[uint8]interface{}{
@@ -3005,6 +3038,7 @@ func init() {
 	MustRegisterAliasedFunction("max", maxSeries)
 	MustRegisterAliasedFunction("min", minSeries)
 	MustRegisterAliasedFunction("randomWalk", randomWalkFunction)
+	MustRegisterAliasedFunction("round", roundFunction)
 	MustRegisterAliasedFunction("sum", sumSeries)
 	MustRegisterAliasedFunction("time", timeFunction)
 }
diff --git a/src/query/graphite/native/builtin_functions_test.go b/src/query/graphite/native/builtin_functions_test.go
index a0cacefe62..fb60764e4b 100644
--- a/src/query/graphite/native/builtin_functions_test.go
+++ b/src/query/graphite/native/builtin_functions_test.go
@@ -1396,6 +1396,60 @@ func TestKeepLastValue(t *testing.T) {
 	}
 }
 
+func TestRoundFunction(t *testing.T) {
+	ctx := common.NewTestContext()
+	defer func() { _ = ctx.Close() }()
+
+	tests := []struct {
+		name       string
+		inputs     []float64
+		outputs    []float64
+		outputName string
+		precision  int
+	}{
+		{
+			"foo",
+			[]float64{111.1, math.NaN(), 111.11, math.NaN(), 111.111},
+			[]float64{110, math.NaN(), 110, math.NaN(), 110},
+			"roundFunction(foo,-1)",
+			-1,
+		},
+		{
+			"foo",
+			[]float64{1.1, math.NaN(), 1.11, math.NaN(), 1.111},
+			[]float64{1, math.NaN(), 1, math.NaN(), 1},
+			"roundFunction(foo)",
+			0,
+		},
+		{
+			"foo",
+			[]float64{1.1, math.NaN(), 1.11, math.NaN(), 1.111},
+			[]float64{1.1, math.NaN(), 1.1, math.NaN(), 1.1},
+			"roundFunction(foo,1)",
+			1,
+		},
+		{
+			"foo",
+			[]float64{1.1, math.NaN(), 1.11, math.NaN(), 1.111},
+			[]float64{1.10, math.NaN(), 1.11, math.NaN(), 1.11},
+			"roundFunction(foo,2)",
+			2,
+		},
+	}
+
+	start := time.Now()
+	for _, test := range tests {
+		input := ts.NewSeries(ctx, test.name, start, common.NewTestSeriesValues(ctx, 100, test.inputs))
+		outputs, err := roundFunction(ctx, singlePathSpec{
+			Values: []*ts.Series{input},
+		}, test.precision)
+		expected := common.TestSeries{Name: test.outputName, Data: test.outputs}
+		require.NoError(t, err)
+		common.CompareOutputsAndExpected(t, 100, start,
+			[]common.TestSeries{expected}, outputs.Values)
+	}
+}
+
 func TestSustainedAbove(t *testing.T) {
 	ctx := common.NewTestContext()
 	defer func() { _ = ctx.Close() }()