Fixes vector grouping injection. (#2975)

* 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>

* got linted.

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>

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 {

pkg/logql/ast_test.go

@@ -3,11 +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) {
@@ -340,3 +340,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)
+			}
+		})
+	}
+}