MQE: fix issues where "remove statically empty expressions" behaves incorrectly (#15014)

#### What this PR does

This PR fixes two issues in the "remove statically empty expressions"
optimisation pass:

* it is ineffective when applied to queries rewritten by Grafana Cloud's
Adaptive Metrics feature
* it can incorrectly determine that an expression is empty if the
expression is a `timestamp()` function call directly on a selector and
that selector contains an offset or @ modifier

#### Which issue(s) this PR fixes or relates to

#14989

#### Checklist

- [x] Tests updated.
- [n/a] Documentation added.
- [x] `CHANGELOG.md` updated - the order of entries should be
`[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`. If changelog entry
is not needed, please add the `changelog-not-needed` label to the PR.
- [n/a]
[`about-versioning.md`](https://github.com/grafana/mimir/blob/main/docs/sources/mimir/configure/about-versioning.md)
updated with experimental features.


<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Medium Risk**
> Touches query-planning optimization logic for `timestamp()`
comparisons; mistakes could change whether parts of queries are skipped,
but the change is narrowly scoped and backed by expanded tests.
> 
> **Overview**
> Fixes the MQE `RemoveStaticallyEmptyExpressions` optimization so it no
longer misclassifies `timestamp()` comparisons as empty when
`timestamp()` wraps a selector with `offset`/`@` (now derives the
earliest possible sample time via the selector’s queried time range),
while still optimizing non-selector `timestamp(<expr>)` comparisons
based only on the evaluation start.
> 
> Makes the optimization effective for Grafana Cloud Adaptive Metrics
rewrites by unwrapping the reserved wrapper functions around
`timestamp()` calls, adds defensive error handling, and expands
unit/integration tests plus a changelog reference update.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
2342b78. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: charleskorn

CHANGELOG.md

@@ -188,7 +188,7 @@
 * [ENHANCEMENT] Query-frontend: Stream JSON encoding directly to the response body to avoid a full-copy allocation of the serialized payload. #14840
 * [ENHANCEMENT] Activity tracker: Added `activity_tracker_unfinished_activities_loaded` metric to report the number of unfinished activities detected on startup. #14860
 * [ENHANCEMENT] Distributor now uses record validation time as Kafka record timestamp to reduce rejections among consumers. #14921
-* [ENHANCEMENT] MQE: Add optimisation pass to optimise away expressions containing comparisons with `timestamp()` that can't produce results due to the query time range. #14989
+* [ENHANCEMENT] MQE: Add optimisation pass to optimise away expressions containing comparisons with `timestamp()` that can't produce results due to the query time range. #14989 #15014
 * [ENHANCEMENT] Distributor: OTLP endpoint now returns partial success (HTTP 200) instead of HTTP 429 when the usage tracker rejects some series due to the active series limit but other series are successfully ingested. The `RejectedDataPoints` field reports the count of distributor-side rejections (usage tracker filtering). #14789
 * [ENHANCEMENT] MQE: Account for memory consumption of labels returned by binary operations in query memory consumption estimate earlier. #15033
 * [ENHANCEMENT] Query-frontend: Log the number of series and samples returned for queries in `query stats` log lines. #15044

pkg/streamingpromql/config.go

@@ -122,13 +122,14 @@ func NewTestEngineOpts() EngineOpts {
 		Logger:   log.NewNopLogger(),
 		Limits:   NewStaticQueryLimitsProvider(),
 
-		EnablePruneToggles:                   true,
-		EnableCommonSubexpressionElimination: true,
-		EnableSubsetSelectorElimination:      true,
-		EnableNarrowBinarySelectors:          true,
-		EnableEliminateDeduplicateAndMerge:   true,
-		EnableReduceMatchers:                 true,
-		EnableProjectionPushdown:             true,
-		EnableMultiAggregation:               true,
+		EnablePruneToggles:                     true,
+		EnableCommonSubexpressionElimination:   true,
+		EnableSubsetSelectorElimination:        true,
+		EnableNarrowBinarySelectors:            true,
+		EnableEliminateDeduplicateAndMerge:     true,
+		EnableReduceMatchers:                   true,
+		EnableProjectionPushdown:               true,
+		EnableMultiAggregation:                 true,
+		EnableRemoveStaticallyEmptyExpressions: true,
 	}
 }

pkg/streamingpromql/optimize/plan/remove_statically_empty_expressions.go

@@ -4,10 +4,12 @@ package plan
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prometheus/prometheus/model/timestamp"
 
 	"github.com/grafana/mimir/pkg/streamingpromql/operators/functions"
 	"github.com/grafana/mimir/pkg/streamingpromql/planning"
@@ -60,7 +62,7 @@ func NewRemoveStaticallyEmptyExpressionsOptimizationPass(reg prometheus.Register
 }
 
 func (s *RemoveStaticallyEmptyExpressionsOptimizationPass) Name() string {
-	return "remove statically empty expressions"
+	return "Remove statically empty expressions"
 }
 
 func (s *RemoveStaticallyEmptyExpressionsOptimizationPass) Apply(ctx context.Context, plan *planning.QueryPlan, maximumSupportedQueryPlanVersion planning.QueryPlanVersion) (*planning.QueryPlan, error) {
@@ -117,7 +119,9 @@ func (s *RemoveStaticallyEmptyExpressionsOptimizationPass) apply(node planning.N
 		}
 	}
 
-	if isAlwaysEmpty(node, params) {
+	if empty, err := isAlwaysEmpty(node, params); err != nil {
+		return nil, false, err
+	} else if empty {
 		noOp := &core.NoOp{NoOpDetails: &core.NoOpDetails{}}
 		return noOp, true, nil
 	}
@@ -127,83 +131,134 @@ func (s *RemoveStaticallyEmptyExpressionsOptimizationPass) apply(node planning.N
 
 // isAlwaysEmpty returns true if node can be statically determined to produce an empty instant
 // vector for the entire query time range described by params.
-func isAlwaysEmpty(node planning.Node, params *planning.QueryParameters) bool {
+func isAlwaysEmpty(node planning.Node, params *planning.QueryParameters) (bool, error) {
 	node = unwrap(node)
 
 	switch node := node.(type) {
 	case *core.NoOp:
-		return true
+		return true, nil
 	case *core.BinaryExpression:
 		return isAlwaysEmptyBinaryExpression(node, params)
 	default:
-		return false
+		return false, nil
 	}
 }
 
-func isAlwaysEmptyBinaryExpression(node *core.BinaryExpression, params *planning.QueryParameters) bool {
-	earliestPossibleTimestampValueInMilliseconds := float64(params.TimeRange.StartT - params.LookbackDelta.Milliseconds())
-
+func isAlwaysEmptyBinaryExpression(node *core.BinaryExpression, params *planning.QueryParameters) (bool, error) {
 	if node.ReturnBool {
-		return false
+		return false, nil
 	}
 
 	switch node.Op {
 	case core.BINARY_LAND:
-		return isAlwaysEmpty(node.LHS, params) || isAlwaysEmpty(node.RHS, params)
+		lhsEmpty, err := isAlwaysEmpty(node.LHS, params)
+		if err != nil {
+			return false, err
+		}
 
-	case core.BINARY_LSS:
-		// timestamp(v) < C: always empty when C <= the earliest value that timestamp() could return
-		// timestamp() returns the value in seconds since the epoch, so we need to convert to milliseconds.
-		if constant, ok := isTimestampComparison(node.LHS, node.RHS); ok {
-			return constant*1000 <= earliestPossibleTimestampValueInMilliseconds
+		if lhsEmpty {
+			return true, nil
 		}
 
+		return isAlwaysEmpty(node.RHS, params)
+
+	case core.BINARY_LSS:
+		// Check for timestamp(v) < C.
+		return isAlwaysEmptyTimestampComparison(node.LHS, node.RHS, false, params)
+
 	case core.BINARY_LTE:
-		// timestamp(v) <= C: always empty when C < the earliest value that timestamp() could return
-		if constant, ok := isTimestampComparison(node.LHS, node.RHS); ok {
-			return constant*1000 < earliestPossibleTimestampValueInMilliseconds
-		}
+		// Check for timestamp(v) <= C.
+		return isAlwaysEmptyTimestampComparison(node.LHS, node.RHS, true, params)
 
 	case core.BINARY_GTR:
-		// C > timestamp(v): equivalent to timestamp(v) < C.
-		if constant, ok := isTimestampComparison(node.RHS, node.LHS); ok {
-			return constant*1000 <= earliestPossibleTimestampValueInMilliseconds
-		}
+		// Check for C > timestamp(v), equivalent to timestamp(v) < C.
+		return isAlwaysEmptyTimestampComparison(node.RHS, node.LHS, false, params)
 
 	case core.BINARY_GTE:
-		// C >= timestamp(v): equivalent to timestamp(v) <= C.
-		if constant, ok := isTimestampComparison(node.RHS, node.LHS); ok {
-			return constant*1000 < earliestPossibleTimestampValueInMilliseconds
-		}
+		// Check for C >= timestamp(v), equivalent to timestamp(v) <= C.
+		return isAlwaysEmptyTimestampComparison(node.RHS, node.LHS, true, params)
 	}
 
-	return false
+	return false, nil
 }
 
-// isTimestampComparison checks whether timestampSide is (or wraps) a timestamp()
-// function call and constantSide is (or wraps) a NumberLiteral. If so, it returns the constant
-// value and true.
-func isTimestampComparison(timestampSide, constantSide planning.Node) (float64, bool) {
-	if !isTimestampCall(timestampSide) {
-		return 0, false
+// isAlwaysEmptyTimestampComparison returns true if timestampSide and constantSide represent
+// a timestamp(...) invocation and number literal respectively, and the value of constantSide
+// is such that the expression timestampSide < constantSide (inclusive=false) or
+// timestampSide <= constantSide (inclusive=true) can never return any results.
+func isAlwaysEmptyTimestampComparison(timestampSide, constantSide planning.Node, inclusive bool, params *planning.QueryParameters) (bool, error) {
+	timestampCall, ok := findTimestampCall(timestampSide)
+	if !ok {
+		return false, nil
 	}
 
-	constantSide = unwrap(constantSide)
-	literal, ok := constantSide.(*core.NumberLiteral)
+	constant, ok := findConstant(constantSide)
 	if !ok {
-		return 0, false
+		return false, nil
+	}
+
+	if len(timestampCall.Args) < 1 {
+		// Should never happen, but check to avoid panicking here.
+		return false, fmt.Errorf("expected at least one argument in call to timestamp(), got %d", len(timestampCall.Args))
 	}
 
-	return literal.Value, true
+	selector, timestampWrapsSelector := timestampCall.Args[0].(*core.VectorSelector)
+
+	// The expression timestamp(X) < C is guaranteed to return no results if the lowest possible
+	// value of timestamp(X) is greater than or equal to C.
+	//
+	// The expression timestamp(X) <= C is guaranteed to return no results if the lowest possible
+	// value of timestamp is greater than C.
+	//
+	// If X is a selector, then timestamp(X) will return the timestamps of the underlying samples, so we need to check
+	// the time range queried to account for the lookback window, offsets and @ modifiers.
+	//
+	// If X is not a selector, then timestamp(X) can only return timestamps of the steps in the query time range.
+
+	var earliestPossibleTimestampValueInMilliseconds float64
+	if timestampWrapsSelector {
+		timeRange, err := selector.QueriedTimeRange(params.TimeRange, params.LookbackDelta)
+		if err != nil {
+			return false, err
+		}
+		earliestPossibleTimestampValueInMilliseconds = float64(timestamp.FromTime(timeRange.MinT))
+	} else {
+		earliestPossibleTimestampValueInMilliseconds = float64(params.TimeRange.StartT)
+	}
+
+	constantInMilliseconds := constant.Value * 1000
+
+	if inclusive {
+		return earliestPossibleTimestampValueInMilliseconds > constantInMilliseconds, nil
+	}
+
+	return earliestPossibleTimestampValueInMilliseconds >= constantInMilliseconds, nil
 }
 
-// isTimestampCall returns true if node is (or wraps) a FunctionCall for the timestamp() function.
+// findTimestampCall returns the function node and true if node is (or wraps) a FunctionCall for the timestamp() function.
 // It unwraps DeduplicateAndMerge, DropName, and StepInvariantExpression layers transparently.
-func isTimestampCall(node planning.Node) bool {
+func findTimestampCall(node planning.Node) (*core.FunctionCall, bool) {
 	node = unwrap(node)
 
-	fc, ok := node.(*core.FunctionCall)
-	return ok && fc.Function == functions.FUNCTION_TIMESTAMP
+	f, ok := node.(*core.FunctionCall)
+	if !ok {
+		return nil, false
+	}
+
+	if f.Function == functions.FUNCTION_TIMESTAMP {
+		return f, true
+	}
+
+	return nil, false
+}
+
+// findConstant returns the number literal node and true if node is (or wraps) a NumberLiteral.
+// It unwraps DeduplicateAndMerge, DropName, and StepInvariantExpression layers transparently.
+func findConstant(node planning.Node) (*core.NumberLiteral, bool) {
+	node = unwrap(node)
+
+	literal, ok := node.(*core.NumberLiteral)
+	return literal, ok
 }
 
 // unwrap removes transparent wrapper nodes returning the innermost non-wrapper node.
@@ -216,6 +271,13 @@ func unwrap(node planning.Node) planning.Node {
 			node = n.Inner
 		case *core.StepInvariantExpression:
 			node = n.Inner
+		case *core.FunctionCall:
+			if (n.Function == functions.FUNCTION_ADAPTIVE_METRICS_RESERVED_1 || n.Function == functions.FUNCTION_ADAPTIVE_METRICS_RESERVED_2) && len(n.Args) > 0 {
+				// The Adaptive Metrics query rewriting can wrap a timestamp() call (eg. expression becomes wrapper(timestamp(...)) < T), so unwrap it.
+				node = n.Args[0]
+			} else {
+				return node
+			}
 		default:
 			return node
 		}