From 2ff9a36f52817b98aaa459ac3d1900e54e7da08a Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Mon, 23 Nov 2020 11:53:57 +0100
Subject: [PATCH 1/2] Fixes vector grouping injection.

During an improvement I introduced a bug where we would pre-aggregate some range vector
aggregations, although the result is not the same.

We can only inject vector grouping in range vector if the operation is associative through grouping.

This is not the case for max/min/stddev/stddvar/quantile/avg_over_time.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
---
 pkg/logql/ast.go      | 15 ++++++++++++++-
 pkg/logql/ast_test.go | 38 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/pkg/logql/ast.go b/pkg/logql/ast.go
index fc80c4133834..0f378fe37913 100644
--- a/pkg/logql/ast.go
+++ b/pkg/logql/ast.go
@@ -711,12 +711,25 @@ func (e *vectorAggregationExpr) Selector() LogSelectorExpr {
 func (e *vectorAggregationExpr) Extractor() (log.SampleExtractor, error) {
 	// inject in the range vector extractor the outer groups to improve performance.
 	// This is only possible if the operation is a sum. Anything else needs all labels.
-	if r, ok := e.left.(*rangeAggregationExpr); ok && e.operation == OpTypeSum {
+	if r, ok := e.left.(*rangeAggregationExpr); ok && canInjectVectorGrouping(e.operation, r.operation) {
 		return r.extractor(e.grouping, len(e.grouping.groups) == 0)
 	}
 	return e.left.Extractor()
 }
 
+// canInjectVectorGrouping tells if a vector operation can inject grouping into the nested range vector.
+func canInjectVectorGrouping(vecOp, rangeOp string) bool {
+	if vecOp != OpTypeSum {
+		return false
+	}
+	switch rangeOp {
+	case OpRangeTypeBytes, OpRangeTypeBytesRate, OpRangeTypeSum, OpRangeTypeRate, OpRangeTypeCount:
+		return true
+	default:
+		return false
+	}
+}
+
 func (e *vectorAggregationExpr) String() string {
 	var params []string
 	if e.params != 0 {
diff --git a/pkg/logql/ast_test.go b/pkg/logql/ast_test.go
index 5bed2bf473a0..6ff709916dbe 100644
--- a/pkg/logql/ast_test.go
+++ b/pkg/logql/ast_test.go
@@ -4,7 +4,6 @@ import (
 	"testing"
 
 	"github.com/grafana/loki/pkg/logql/log"
-
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -340,3 +339,40 @@ func mustNewRegexParser(re string) log.Stage {
 	}
 	return r
 }
+
+func Test_canInjectVectorGrouping(t *testing.T) {
+
+	tests := []struct {
+		vecOp   string
+		rangeOp string
+		want    bool
+	}{
+		{OpTypeSum, OpRangeTypeBytes, true},
+		{OpTypeSum, OpRangeTypeBytesRate, true},
+		{OpTypeSum, OpRangeTypeSum, true},
+		{OpTypeSum, OpRangeTypeRate, true},
+		{OpTypeSum, OpRangeTypeCount, true},
+
+		{OpTypeSum, OpRangeTypeAvg, false},
+		{OpTypeSum, OpRangeTypeMax, false},
+		{OpTypeSum, OpRangeTypeQuantile, false},
+		{OpTypeSum, OpRangeTypeStddev, false},
+		{OpTypeSum, OpRangeTypeStdvar, false},
+		{OpTypeSum, OpRangeTypeMin, false},
+		{OpTypeSum, OpRangeTypeMax, false},
+
+		{OpTypeAvg, OpRangeTypeBytes, false},
+		{OpTypeCount, OpRangeTypeBytesRate, false},
+		{OpTypeBottomK, OpRangeTypeSum, false},
+		{OpTypeMax, OpRangeTypeRate, false},
+		{OpTypeMin, OpRangeTypeCount, false},
+		{OpTypeTopK, OpRangeTypeCount, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.vecOp+"_"+tt.rangeOp, func(t *testing.T) {
+			if got := canInjectVectorGrouping(tt.vecOp, tt.rangeOp); got != tt.want {
+				t.Errorf("canInjectVectorGrouping() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

From ff9fd69f76d15ee1aa0024d6ac135ca4be870acd Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Mon, 23 Nov 2020 12:04:15 +0100
Subject: [PATCH 2/2] got linted.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
---
 pkg/logql/ast_test.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pkg/logql/ast_test.go b/pkg/logql/ast_test.go
index 6ff709916dbe..e94100bbbb8c 100644
--- a/pkg/logql/ast_test.go
+++ b/pkg/logql/ast_test.go
@@ -3,10 +3,11 @@ package logql
 import (
 	"testing"
 
-	"github.com/grafana/loki/pkg/logql/log"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/loki/pkg/logql/log"
 )
 
 func Test_logSelectorExpr_String(t *testing.T) {