From da65a9d393c736dc057c84ef32d1f1a2796d5ec2 Mon Sep 17 00:00:00 2001
From: Filip Petkovski <filip.petkovsky@gmail.com>
Date: Mon, 20 Nov 2023 20:03:34 +0100
Subject: [PATCH] Fix group aggregate with NaN values (#328)

The group aggregate should produce an output value even when inputs are NaN.

Fixes https://github.com/thanos-io/promql-engine/issues/326.

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>
---
 engine/engine_test.go              |  9 +++++++++
 execution/aggregate/accumulator.go | 12 ++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/engine/engine_test.go b/engine/engine_test.go
index 540553c5..ad383bdd 100644
--- a/engine/engine_test.go
+++ b/engine/engine_test.go
@@ -1114,6 +1114,15 @@ func TestQueriesAgainstOldEngine(t *testing.T) {
 					http_requests_total{pod="nginx-2"} 2+2x18`,
 			query: `group by (pod) (http_requests_total)`,
 		},
+		{
+			// Issue https://github.com/thanos-io/promql-engine/issues/326.
+			name: "group by with NaN values",
+			load: `
+load 30s
+	http_requests_total{pod="nginx-1", route="/"} 1.00+1.00x4
+	http_requests_total{pod="nginx-2", route="/"}  1+2.00x4`,
+			query: `group by (pod, route) (atanh(-{__name__="http_requests_total"} offset -3m4s))`,
+		},
 		{
 			name: "resets",
 			load: `load 30s
diff --git a/execution/aggregate/accumulator.go b/execution/aggregate/accumulator.go
index b1ddf922..fa808359 100644
--- a/execution/aggregate/accumulator.go
+++ b/execution/aggregate/accumulator.go
@@ -79,9 +79,11 @@ func (s *sumAcc) Reset(_ float64) {
 }
 
 type genericAcc struct {
-	zeroVal         float64
-	value           float64
-	hasValue        bool
+	zeroVal  float64
+	value    float64
+	hasValue bool
+	skipNaNs bool
+
 	aggregate       func(float64, float64) float64
 	vectorAggregate func([]float64, []*histogram.FloatHistogram) float64
 }
@@ -113,6 +115,7 @@ func groupVecAggregate(_ []float64, _ []*histogram.FloatHistogram) float64 {
 
 func newMaxAcc() *genericAcc {
 	return &genericAcc{
+		skipNaNs:        true,
 		zeroVal:         math.MinInt64,
 		aggregate:       maxAggregate,
 		vectorAggregate: maxVecAggregate,
@@ -121,6 +124,7 @@ func newMaxAcc() *genericAcc {
 
 func newMinAcc() *genericAcc {
 	return &genericAcc{
+		skipNaNs:        true,
 		zeroVal:         math.MaxInt64,
 		aggregate:       minAggregate,
 		vectorAggregate: minVecAggregate,
@@ -136,7 +140,7 @@ func newGroupAcc() *genericAcc {
 }
 
 func (g *genericAcc) Add(v float64, _ *histogram.FloatHistogram) {
-	if math.IsNaN(v) {
+	if g.skipNaNs && math.IsNaN(v) {
 		return
 	}
 	if !g.hasValue {