From 10e9695b21ae8c91598e7d71c25ffb6484041f53 Mon Sep 17 00:00:00 2001
From: Susana Ferreira <susana.ferreira@grafana.com>
Date: Thu, 29 Sep 2022 13:08:05 +0200
Subject: [PATCH] Fix rangemapper with VectorExpr support (#7288)

**What this PR does / why we need it**:
The new expression type VectorExpr introduced recently in
[PR](https://github.com/grafana/loki/pull/7007), is not taken into
account in the query mapping in the AST mapper rangemapper.
The AST mapper shardmapper was updated in
[PR](https://github.com/grafana/loki/pull/7045).

**Special notes for your reviewer**:

**Checklist**
- [ ] Reviewed the `CONTRIBUTING.md` guide
- [ ] Documentation added
- [x] Tests updated
- [x] `CHANGELOG.md` updated
- [ ] Changes that require user attention or interaction to upgrade are
documented in `docs/sources/upgrading/_index.md`
---
 CHANGELOG.md                  |  1 +
 pkg/logql/rangemapper.go      |  4 ++++
 pkg/logql/rangemapper_test.go | 13 +++++++++++++
 pkg/logql/shardmapper_test.go |  5 +++++
 4 files changed, 23 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb1a60a2dba4..33670819a0d6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,7 @@
 * [7263](https://github.com/grafana/loki/pull/7263) **bboreham**: Dependencies: klauspost/compress package to v1.15.11; improves performance.
 
 ##### Fixes
+* [7288](https://github.com/grafana/loki/pull/7288) **ssncferreira**: Fix query mapping in AST mapper `rangemapper` to support the new `VectorExpr` expression.
 * [7040](https://github.com/grafana/loki/pull/7040) **bakunowski**: Remove duplicated `loki_boltdb_shipper` prefix from `tables_upload_operation_total` metric.
 * [6937](https://github.com/grafana/loki/pull/6937) **ssncferreira**: Fix topk and bottomk expressions with parameter <= 0.
 * [6780](https://github.com/grafana/loki/pull/6780) **periklis**:  Attach the panic recovery handler on all HTTP handlers
diff --git a/pkg/logql/rangemapper.go b/pkg/logql/rangemapper.go
index 7bac7e76b2ea..b3746271be82 100644
--- a/pkg/logql/rangemapper.go
+++ b/pkg/logql/rangemapper.go
@@ -156,6 +156,8 @@ func (m RangeMapper) Map(expr syntax.SampleExpr, vectorAggrPushdown *syntax.Vect
 		return e, nil
 	case *syntax.LiteralExpr:
 		return e, nil
+	case *syntax.VectorExpr:
+		return e, nil
 	default:
 		// ConcatSampleExpr and DownstreamSampleExpr are not supported input expression types
 		return nil, errors.Errorf("unexpected expr type (%T) for ASTMapper type (%T) ", expr, m)
@@ -438,6 +440,8 @@ func isSplittableByRange(expr syntax.SampleExpr) bool {
 		return isSplittableByRange(e.SampleExpr) || literalLHS && isSplittableByRange(e.RHS) || literalRHS
 	case *syntax.LabelReplaceExpr:
 		return isSplittableByRange(e.Left)
+	case *syntax.VectorExpr:
+		return false
 	default:
 		return false
 	}
diff --git a/pkg/logql/rangemapper_test.go b/pkg/logql/rangemapper_test.go
index e5ed72f77776..68c442eaa5ce 100644
--- a/pkg/logql/rangemapper_test.go
+++ b/pkg/logql/rangemapper_test.go
@@ -1708,6 +1708,14 @@ func Test_SplitRangeVectorMapping_Noop(t *testing.T) {
 			`sum by (foo) (sum_over_time({app="foo"} | json | unwrap bar [3m])) / sum_over_time({app="foo"} | json | unwrap bar [6m])`,
 			`(sum by (foo) (sum_over_time({app="foo"} | json | unwrap bar [3m])) / sum_over_time({app="foo"} | json | unwrap bar [6m]))`,
 		},
+		{
+			`count_over_time({app="foo"}[3m]) or vector(0)`,
+			`(count_over_time({app="foo"}[3m]) or vector(0.000000))`,
+		},
+		{
+			`sum(last_over_time({app="foo"} | logfmt | unwrap total_count [1d]) by (foo)) or vector(0)`,
+			`(sum(last_over_time({app="foo"} | logfmt | unwrap total_count [1d]) by (foo)) or vector(0.000000))`,
+		},
 
 		// should be noop if literal expression
 		{
@@ -1718,6 +1726,11 @@ func Test_SplitRangeVectorMapping_Noop(t *testing.T) {
 			`5 * 5`,
 			`25`,
 		},
+		// should be noop if VectorExpr
+		{
+			`vector(0)`,
+			`vector(0.000000)`,
+		},
 	} {
 		tc := tc
 		t.Run(tc.expr, func(t *testing.T) {
diff --git a/pkg/logql/shardmapper_test.go b/pkg/logql/shardmapper_test.go
index a33323dc4462..8ae45bea087f 100644
--- a/pkg/logql/shardmapper_test.go
+++ b/pkg/logql/shardmapper_test.go
@@ -263,6 +263,11 @@ func TestMappingStrings(t *testing.T) {
 			in:  `avg_over_time({job=~"myapps.*"} |= "stats" | json busy="utilization" | unwrap busy [5m])`,
 			out: `avg_over_time({job=~"myapps.*"} |= "stats" | json busy="utilization" | unwrap busy [5m])`,
 		},
+		// should be noop if VectorExpr
+		{
+			in:  `vector(0)`,
+			out: `vector(0.000000)`,
+		},
 	} {
 		t.Run(tc.in, func(t *testing.T) {
 			ast, err := syntax.ParseExpr(tc.in)