frontend: Fix downstream queries in experimental subquery spin-off fe…

…ature (#10603)

* frontend: Fix downstream queries in experimental subquery spin-off feature

Massive oversight in #10460. The downstream queries were all just `__query__` instead of the actual query.
I also optimized the mapper so that it groups downstream queries further up the tree, which should reduce the number of queries that need to be run.

* Remove print

pkg/frontend/querymiddleware/astmapper/subquery_spin_off.go

@@ -90,25 +90,11 @@ func (m *subquerySpinOffMapper) MapExpr(expr parser.Expr) (mapped parser.Expr, f
 		// The last argument will typically contain the subquery in an aggregation function
 		// Examples: last_over_time(<subquery>[5m:]) or quantile_over_time(0.5, <subquery>[5m:])
 		if sq, ok := e.Args[lastArgIdx].(*parser.SubqueryExpr); ok {
-			// @ is not supported
-			if sq.StartOrEnd != 0 || sq.Timestamp != nil {
-				return downstreamQuery(expr)
-			}
-
-			// Filter out subqueries with ranges less than 1 hour as they are not worth spinning off.
-			if sq.Range < 1*time.Hour {
-				return downstreamQuery(expr)
-			}
-
-			selectorsCt := countSelectors(sq.Expr)
-
-			// Evaluate constants within the frontend engine
-			if selectorsCt == 0 {
+			canBeSpunOff, isConstant := subqueryCanBeSpunOff(*sq)
+			if isConstant {
 				return expr, false, nil
 			}
-
-			// Filter out subqueries that are just selectors, they are fast enough that they aren't worth spinning off.
-			if selectorsCt == 1 && !isComplexExpr(sq.Expr) {
+			if !canBeSpunOff {
 				return downstreamQuery(expr)
 			}
 
@@ -184,7 +170,8 @@ func isComplexExpr(expr parser.Node) bool {
 func hasSubqueryInChildren(expr parser.Node) bool {
 	switch e := expr.(type) {
 	case *parser.SubqueryExpr:
-		return true
+		canBeSpunOff, _ := subqueryCanBeSpunOff(*e)
+		return canBeSpunOff
 	default:
 		for _, child := range parser.Children(e) {
 			if hasSubqueryInChildren(child) {
@@ -195,6 +182,32 @@ func hasSubqueryInChildren(expr parser.Node) bool {
 	}
 }
 
+func subqueryCanBeSpunOff(sq parser.SubqueryExpr) (spinoff, constant bool) {
+	// @ is not supported
+	if sq.StartOrEnd != 0 || sq.Timestamp != nil {
+		return false, false
+	}
+
+	// Filter out subqueries with ranges less than 1 hour as they are not worth spinning off.
+	if sq.Range < 1*time.Hour {
+		return false, false
+	}
+
+	selectorsCt := countSelectors(sq.Expr)
+
+	// Evaluate constants within the frontend engine
+	if selectorsCt == 0 {
+		return false, true
+	}
+
+	// Filter out subqueries that are just selectors, they are fast enough that they aren't worth spinning off.
+	if selectorsCt == 1 && !isComplexExpr(sq.Expr) {
+		return false, false
+	}
+
+	return true, false
+}
+
 func countSelectors(expr parser.Node) int {
 	switch e := expr.(type) {
 	case *parser.VectorSelector, *parser.MatrixSelector:

pkg/frontend/querymiddleware/astmapper/subquery_spin_off_test.go

@@ -109,9 +109,9 @@ func TestSubquerySpinOffMapper(t *testing.T) {
 		{
 			name:                      "ignore single selector subquery",
 			in:                        `sum(avg_over_time((foo > 1)[3d:1m]) * avg_over_time(foo[3d]))`,
-			out:                       `sum(__downstream_query__{__query__="avg_over_time((foo > 1)[3d:1m])"} * __downstream_query__{__query__="avg_over_time(foo[3d])"})`,
+			out:                       `__downstream_query__{__query__="sum(avg_over_time((foo > 1)[3d:1m]) * avg_over_time(foo[3d]))"}`,
 			expectedSubqueries:        0,
-			expectedDownstreamQueries: 2,
+			expectedDownstreamQueries: 1,
 		},
 		{
 			name: "subquery of aggregation",
@@ -234,6 +234,31 @@ func TestSubquerySpinOffMapper(t *testing.T) {
 			expectedSubqueries:        1,
 			expectedDownstreamQueries: 0,
 		},
+		{
+			name: "map downstream query as top-level as possible",
+			in: `sum by (group_1) (
+      sum_over_time(
+        avg by (group_1) (metric_counter{group_2="1"})[1d:5m] offset 1m
+      )
+    *
+      avg by (group_1) (
+        avg_over_time(metric_counter{group_2="2"}[1d:5m] offset 1m)
+      )
+  *
+    0.083333
+)`,
+			out: `sum by (group_1) (
+      sum_over_time(
+        __subquery_spinoff__{__offset__="1m0s",__query__="avg by (group_1) (metric_counter{group_2=\"1\"})",__range__="24h0m0s",__step__="5m0s"}[1d] offset 1m
+      )
+    *
+	  __downstream_query__{__query__="avg by (group_1) (avg_over_time(metric_counter{group_2=\"2\"}[1d:5m] offset 1m))"}
+  *
+    0.083333
+)`,
+			expectedSubqueries:        1,
+			expectedDownstreamQueries: 1,
+		},
 	} {
 		tt := tt
 

pkg/frontend/querymiddleware/spin_off_subqueries_queryable.go

@@ -67,7 +67,12 @@ func (q *spinOffSubqueriesQuerier) Select(ctx context.Context, _ bool, hints *st
 
 	switch name {
 	case astmapper.DownstreamQueryMetricName:
-		downstreamReq, err := q.req.WithQuery(astmapper.DownstreamQueryLabelName)
+		query, ok := values[astmapper.DownstreamQueryLabelName]
+		if !ok {
+			return storage.ErrSeriesSet(errors.New("missing required labels for downstream query"))
+		}
+
+		downstreamReq, err := q.req.WithQuery(query)
 		if err != nil {
 			return storage.ErrSeriesSet(err)
 		}

pkg/frontend/querymiddleware/spin_off_subqueries_test.go

@@ -54,6 +54,17 @@ func TestSubquerySpinOff_Correctness(t *testing.T) {
 					)`,
 			expectedSkippedReason: "no-subquery",
 		},
+		"subquery max with downstream join": {
+			query: `max_over_time(
+							rate(metric_counter[1m])
+						[2d:1m]
+					)
+					* on (group_1) group_left()
+					max by (group_1)(
+							rate(metric_counter[1m])
+					)`,
+			expectedSpunOffSubqueries: 1,
+		},
 		"subquery max": {
 			query: `max_over_time(
 							rate(metric_counter[1m])
@@ -134,6 +145,21 @@ func TestSubquerySpinOff_Correctness(t *testing.T) {
 					)`,
 			expectedSpunOffSubqueries: 1,
 		},
+		"subquery max with offset shorter than step": {
+			query: `
+sum by (group_1) (
+      sum_over_time(
+        avg by (group_1) (metric_counter{group_2="1"})[1d:5m] offset 1m
+      )
+    *
+      avg by (group_1) (
+        avg_over_time(metric_counter{group_2="2"}[1d:5m] offset 1m)
+      )
+  *
+    0.083333
+)`,
+			expectedSpunOffSubqueries: 1,
+		},
 	}
 
 	queryable := setupSubquerySpinOffTestSeries(t, 2*24*time.Hour)