[bugfix] avg_over_time calculation fix (grafana#6252)

Author: ruslan-mikhailov

CHANGELOG.md

@@ -1,4 +1,5 @@
 ## main / unreleased
+* [BUGFIX] Correct avg_over_time calculation [#6252](https://github.com/grafana/tempo/pull/6252) (@ruslan-mikhailov)
 
 # v2.10.0-rc.0
 

pkg/traceql/engine_metrics_average.go

@@ -268,6 +268,7 @@ func (k *averageSeries) addWeigthedMean(interval int, mean float64, weight float
 	meanDelta -= currentMean.compensation
 
 	currentMean.add(meanDelta)
+	currentMean.weight = sumWeights
 	k.values[interval] = currentMean
 }
 

pkg/traceql/engine_metrics_test.go

@@ -1294,6 +1294,43 @@ func TestAvgOverTimeForDurationWithoutAggregation(t *testing.T) {
 	assert.Equal(t, 200., avg.Values[2]*float64(time.Second))
 }
 
+func TestAvgOverTimeCombine(t *testing.T) {
+	req := &tempopb.QueryRangeRequest{
+		Start: 1,
+		End:   uint64(3 * time.Second),
+		Step:  uint64(1 * time.Second),
+		Query: "{ } | avg_over_time(duration)",
+	}
+
+	// Three set of spans with different data for a single time interval
+	// Each set will be passed to level 2 independently
+	in1 := []Span{
+		newMockSpan(nil).WithStartTime(uint64(1 * time.Second)).WithDuration(uint64(100 * time.Second)),
+	}
+
+	in2 := []Span{
+		newMockSpan(nil).WithStartTime(uint64(1 * time.Second)).WithDuration(uint64(200 * time.Second)),
+		newMockSpan(nil).WithStartTime(uint64(1 * time.Second)).WithDuration(uint64(400 * time.Second)),
+	}
+
+	in3 := []Span{
+		newMockSpan(nil).WithStartTime(uint64(1 * time.Second)).WithDuration(uint64(500 * time.Second)),
+		newMockSpan(nil).WithStartTime(uint64(1 * time.Second)).WithDuration(uint64(600 * time.Second)),
+		newMockSpan(nil).WithStartTime(uint64(1 * time.Second)).WithDuration(uint64(700 * time.Second)),
+	}
+
+	result, seriesCount, err := runTraceQLMetric(req, in1, in2, in3)
+	require.NoError(t, err)
+	require.Equal(t, len(result), seriesCount)
+
+	avg := result[LabelsFromArgs("__name__", "avg_over_time").MapKey()]
+
+	expectedResult := (100 + 200 + 400 + 500 + 600 + 700) / 6.
+	assert.Equal(t, expectedResult, avg.Values[0])
+	assert.True(t, math.IsNaN(avg.Values[1]), "expected NaN, got %f", avg.Values[1])
+	assert.True(t, math.IsNaN(avg.Values[2]), "expected NaN, got %f", avg.Values[2])
+}
+
 func TestAvgOverTimeForSpanAttribute(t *testing.T) {
 	req := &tempopb.QueryRangeRequest{
 		Start: 1,

tempodb/tempodb_metrics_test.go

@@ -244,8 +244,7 @@ var queryRangeTestCases = []struct {
 				},
 			},
 		},
-		expectedL2: nil,
-		expectedL3: []*tempopb.TimeSeries{
+		expectedL2: []*tempopb.TimeSeries{
 			{
 				Labels: []common_v1.KeyValue{tempopb.MakeKeyValueString("__name__", "avg_over_time")},
 				Samples: []tempopb.Sample{
@@ -255,6 +254,29 @@ var queryRangeTestCases = []struct {
 					{TimestampMs: 60_000, Value: 240 / 5.0},  // sum from 46 to 50 is 240
 				},
 			},
+			{
+				Labels: []common_v1.KeyValue{
+					tempopb.MakeKeyValueString("__name__", "avg_over_time"),
+					tempopb.MakeKeyValueString("__meta_type", "__count"),
+				},
+				Samples: []tempopb.Sample{
+					{TimestampMs: 15_000, Value: 2 * 15},
+					{TimestampMs: 30_000, Value: 2 * 15},
+					{TimestampMs: 45_000, Value: 2 * 15},
+					{TimestampMs: 60_000, Value: 2 * 5},
+				},
+			},
+		},
+		expectedL3: []*tempopb.TimeSeries{
+			{
+				Labels: []common_v1.KeyValue{tempopb.MakeKeyValueString("__name__", "avg_over_time")},
+				Samples: []tempopb.Sample{
+					{TimestampMs: 15_000, Value: 120 / 15.0},
+					{TimestampMs: 30_000, Value: 345 / 15.0},
+					{TimestampMs: 45_000, Value: 570 / 15.0},
+					{TimestampMs: 60_000, Value: 240 / 5.0},
+				},
+			},
 		},
 	},
 	{