Mimir query engine: add support for unary negation (#9196)

* Add failing tests

* Add feature toggle

* Add support for unary negation of instant vectors

# Conflicts:
#	pkg/streamingpromql/query.go

* Add support for unary negation of scalars

* Add changelog entry

* Address PR feedback: add further test cases

Author: charleskorn

CHANGELOG.md

@@ -32,7 +32,7 @@
 * [CHANGE] Querier: allow wrapping errors with context errors only when the former actually correspond to `context.Canceled` and `context.DeadlineExceeded`. #9175
 * [FEATURE] Alertmanager: Added `-alertmanager.log-parsing-label-matchers` to control logging when parsing label matchers. This flag is intended to be used with `-alertmanager.utf8-strict-mode-enabled` to validate UTF-8 strict mode is working as intended. The default value is `false`. #9173
 * [FEATURE] Alertmanager: Added `-alertmanager.utf8-migration-logging-enabled` to enable logging of tenant configurations that are incompatible with UTF-8 strict mode. The default value is `false`. #9174
-* [FEATURE] Querier: add experimental streaming PromQL engine, enabled with `-querier.query-engine=mimir`. #8422 #8430 #8454 #8455 #8360 #8490 #8508 #8577 #8660 #8671 #8677 #8747 #8850 #8872 #8838 #8911 #8909 #8923 #8924 #8925 #8932 #8933 #8934 #8962 #8986 #8993 #8995 #9008 #9017 #9018 #9019 #9120 #9121 #9136 #9139 #9140 #9145 #9191 #9194
+* [FEATURE] Querier: add experimental streaming PromQL engine, enabled with `-querier.query-engine=mimir`. #8422 #8430 #8454 #8455 #8360 #8490 #8508 #8577 #8660 #8671 #8677 #8747 #8850 #8872 #8838 #8911 #8909 #8923 #8924 #8925 #8932 #8933 #8934 #8962 #8986 #8993 #8995 #9008 #9017 #9018 #9019 #9120 #9121 #9136 #9139 #9140 #9145 #9191 #9194 #9196
 * [FEATURE] Experimental Kafka-based ingest storage. #6888 #6894 #6929 #6940 #6951 #6974 #6982 #7029 #7030 #7091 #7142 #7147 #7148 #7153 #7160 #7193 #7349 #7376 #7388 #7391 #7393 #7394 #7402 #7404 #7423 #7424 #7437 #7486 #7503 #7508 #7540 #7621 #7682 #7685 #7694 #7695 #7696 #7697 #7701 #7733 #7734 #7741 #7752 #7838 #7851 #7871 #7877 #7880 #7882 #7887 #7891 #7925 #7955 #7967 #8031 #8063 #8077 #8088 #8135 #8176 #8184 #8194 #8216 #8217 #8222 #8233 #8503 #8542 #8579 #8657 #8686 #8688 #8703 #8706 #8708 #8738 #8750 #8778 #8808 #8809 #8841 #8842 #8845 #8853 #8886 #8988
   * What it is:
     * When the new ingest storage architecture is enabled, distributors write incoming write requests to a Kafka-compatible backend, and the ingesters asynchronously replay ingested data from Kafka. In this architecture, the write and read path are de-coupled through a Kafka-compatible backend. The write path and Kafka load is a function of the incoming write traffic, the read path load is a function of received queries. Whatever the load on the read path, it doesn't affect the write path.

cmd/mimir/config-descriptor.json

@@ -2043,6 +2043,17 @@
               "fieldFlag": "querier.mimir-query-engine.enable-scalars",
               "fieldType": "boolean",
               "fieldCategory": "experimental"
+            },
+            {
+              "kind": "field",
+              "name": "enable_unary_negation",
+              "required": false,
+              "desc": "Enable support for unary negation in Mimir's query engine. Only applies if the Mimir query engine is in use.",
+              "fieldValue": null,
+              "fieldDefaultValue": true,
+              "fieldFlag": "querier.mimir-query-engine.enable-unary-negation",
+              "fieldType": "boolean",
+              "fieldCategory": "experimental"
             }
           ],
           "fieldValue": null,

cmd/mimir/help-all.txt.tmpl

@@ -1915,6 +1915,8 @@ Usage of ./cmd/mimir/mimir:
     	[experimental] Enable support for ..._over_time functions in Mimir's query engine. Only applies if the Mimir query engine is in use. (default true)
   -querier.mimir-query-engine.enable-scalars
     	[experimental] Enable support for scalars in Mimir's query engine. Only applies if the Mimir query engine is in use. (default true)
+  -querier.mimir-query-engine.enable-unary-negation
+    	[experimental] Enable support for unary negation in Mimir's query engine. Only applies if the Mimir query engine is in use. (default true)
   -querier.minimize-ingester-requests
     	If true, when querying ingesters, only the minimum required ingesters required to reach quorum will be queried initially, with other ingesters queried only if needed due to failures from the initial set of ingesters. Enabling this option reduces resource consumption for the happy path at the cost of increased latency for the unhappy path. (default true)
   -querier.minimize-ingester-requests-hedging-delay duration

docs/sources/mimir/configure/configuration-parameters/index.md

@@ -1537,6 +1537,11 @@ mimir_query_engine:
   # applies if the Mimir query engine is in use.
   # CLI flag: -querier.mimir-query-engine.enable-scalars
   [enable_scalars: <boolean> | default = true]
+
+  # (experimental) Enable support for unary negation in Mimir's query engine.
+  # Only applies if the Mimir query engine is in use.
+  # CLI flag: -querier.mimir-query-engine.enable-unary-negation
+  [enable_unary_negation: <boolean> | default = true]
 ```
 
 ### frontend

pkg/streamingpromql/config.go

@@ -23,6 +23,7 @@ type FeatureToggles struct {
 	EnableOffsetModifier        bool `yaml:"enable_offset_modifier" category:"experimental"`
 	EnableOverTimeFunctions     bool `yaml:"enable_over_time_functions" category:"experimental"`
 	EnableScalars               bool `yaml:"enable_scalars" category:"experimental"`
+	EnableUnaryNegation         bool `yaml:"enable_unary_negation" category:"experimental"`
 }
 
 var overTimeFunctionNames = []string{
@@ -43,6 +44,7 @@ var EnableAllFeatures = FeatureToggles{
 	true,
 	true,
 	true,
+	true,
 }
 
 func (t *FeatureToggles) RegisterFlags(f *flag.FlagSet) {
@@ -51,4 +53,5 @@ func (t *FeatureToggles) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&t.EnableOffsetModifier, "querier.mimir-query-engine.enable-offset-modifier", true, "Enable support for offset modifier in Mimir's query engine. Only applies if the Mimir query engine is in use.")
 	f.BoolVar(&t.EnableOverTimeFunctions, "querier.mimir-query-engine.enable-over-time-functions", true, "Enable support for ..._over_time functions in Mimir's query engine. Only applies if the Mimir query engine is in use.")
 	f.BoolVar(&t.EnableScalars, "querier.mimir-query-engine.enable-scalars", true, "Enable support for scalars in Mimir's query engine. Only applies if the Mimir query engine is in use.")
+	f.BoolVar(&t.EnableUnaryNegation, "querier.mimir-query-engine.enable-unary-negation", true, "Enable support for unary negation in Mimir's query engine. Only applies if the Mimir query engine is in use.")
 }

pkg/streamingpromql/engine_test.go

@@ -51,7 +51,6 @@ func TestUnsupportedPromQLFeatures(t *testing.T) {
 		`count_values("foo", metric{})`:         "'count_values' aggregation with parameter",
 		"rate(metric{}[5m:1m])":                 "PromQL expression type *parser.SubqueryExpr for range vectors",
 		"quantile_over_time(0.4, metric{}[5m])": "'quantile_over_time' function",
-		"-sum(metric{})":                        "PromQL expression type *parser.UnaryExpr for instant vectors",
 	}
 
 	for expression, expectedError := range unsupportedExpressions {
@@ -118,6 +117,17 @@ func TestUnsupportedPromQLFeaturesWithFeatureToggles(t *testing.T) {
 		requireRangeQueryIsUnsupported(t, featureToggles, "2", "scalar values")
 		requireInstantQueryIsUnsupported(t, featureToggles, "2", "scalar values")
 	})
+
+	t.Run("unary negation", func(t *testing.T) {
+		featureToggles := EnableAllFeatures
+		featureToggles.EnableUnaryNegation = false
+
+		requireRangeQueryIsUnsupported(t, featureToggles, "-sum(metric{})", "unary negation of instant vectors")
+		requireInstantQueryIsUnsupported(t, featureToggles, "-sum(metric{})", "unary negation of instant vectors")
+
+		requireRangeQueryIsUnsupported(t, featureToggles, "-(1)", "unary negation of scalars")
+		requireInstantQueryIsUnsupported(t, featureToggles, "-(1)", "unary negation of scalars")
+	})
 }
 
 func requireRangeQueryIsUnsupported(t *testing.T, featureToggles FeatureToggles, expression string, expectedError string) {

pkg/streamingpromql/functions.go

@@ -229,3 +229,14 @@ func instantVectorToScalarOperatorFactory(args []types.Operator, memoryConsumpti
 
 	return operators.NewInstantVectorToScalar(inner, start, end, interval, memoryConsumptionTracker, expressionPosition), nil
 }
+
+func unaryNegationOfInstantVectorOperatorFactory(inner types.InstantVectorOperator, memoryConsumptionTracker *limiting.MemoryConsumptionTracker, expressionPosition posrange.PositionRange) types.InstantVectorOperator {
+	f := functions.FunctionOverInstantVector{
+		SeriesDataFunc:           functions.UnaryNegation,
+		SeriesMetadataFunc:       functions.DropSeriesName,
+		NeedsSeriesDeduplication: true,
+	}
+
+	o := operators.NewFunctionOverInstantVector(inner, memoryConsumptionTracker, f, expressionPosition)
+	return operators.NewDeduplicateAndMerge(o, memoryConsumptionTracker)
+}

pkg/streamingpromql/functions/math.go

@@ -4,6 +4,9 @@ package functions
 
 import (
 	"math"
+
+	"github.com/grafana/mimir/pkg/streamingpromql/limiting"
+	"github.com/grafana/mimir/pkg/streamingpromql/types"
 )
 
 var Abs = FloatTransformationDropHistogramsFunc(math.Abs)
@@ -48,3 +51,14 @@ var Sgn = FloatTransformationDropHistogramsFunc(func(f float64) float64 {
 	// Otherwise, if the value is 0, we should return 0.
 	return f
 })
+
+var UnaryNegation InstantVectorSeriesFunction = func(seriesData types.InstantVectorSeriesData, _ *limiting.MemoryConsumptionTracker) (types.InstantVectorSeriesData, error) {
+	for i := range seriesData.Floats {
+		seriesData.Floats[i].F = -seriesData.Floats[i].F
+	}
+
+	// Prometheus' engine currently leaves histograms as-is for unary negation, so we do the same.
+	// See https://github.com/prometheus/prometheus/pull/14821 for more discussion of this.
+
+	return seriesData, nil
+}

pkg/streamingpromql/operators/unary_negation_scalar.go

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package operators
+
+import (
+	"context"
+
+	"github.com/prometheus/prometheus/promql/parser/posrange"
+
+	"github.com/grafana/mimir/pkg/streamingpromql/types"
+)
+
+type UnaryNegationOfScalar struct {
+	Inner              types.ScalarOperator
+	expressionPosition posrange.PositionRange
+}
+
+var _ types.ScalarOperator = &UnaryNegationOfScalar{}
+
+func NewUnaryNegationOfScalar(inner types.ScalarOperator, expressionPosition posrange.PositionRange) *UnaryNegationOfScalar {
+	return &UnaryNegationOfScalar{
+		Inner:              inner,
+		expressionPosition: expressionPosition,
+	}
+}
+
+func (u *UnaryNegationOfScalar) GetValues(ctx context.Context) (types.ScalarData, error) {
+	values, err := u.Inner.GetValues(ctx)
+	if err != nil {
+		return types.ScalarData{}, err
+	}
+
+	for i := range values.Samples {
+		values.Samples[i].F = -values.Samples[i].F
+	}
+
+	return values, nil
+}
+
+func (u *UnaryNegationOfScalar) ExpressionPosition() posrange.PositionRange {
+	return u.expressionPosition
+}
+
+func (u *UnaryNegationOfScalar) Close() {
+	u.Inner.Close()
+}

pkg/streamingpromql/query.go

@@ -244,6 +244,21 @@ func (q *Query) convertToInstantVectorOperator(expr parser.Expr) (types.InstantV
 		}
 
 		return operators.NewVectorVectorBinaryOperation(lhs, rhs, *e.VectorMatching, e.Op, q.memoryConsumptionTracker, q.annotations, e.PositionRange())
+	case *parser.UnaryExpr:
+		if e.Op != parser.SUB {
+			return nil, compat.NewNotSupportedError(fmt.Sprintf("unary expression with '%s'", e.Op))
+		}
+
+		if !q.engine.featureToggles.EnableUnaryNegation {
+			return nil, compat.NewNotSupportedError("unary negation of instant vectors")
+		}
+
+		inner, err := q.convertToInstantVectorOperator(e.Expr)
+		if err != nil {
+			return nil, err
+		}
+
+		return unaryNegationOfInstantVectorOperatorFactory(inner, q.memoryConsumptionTracker, e.PositionRange()), nil
 
 	case *parser.StepInvariantExpr:
 		// One day, we'll do something smarter here.
@@ -338,6 +353,23 @@ func (q *Query) convertToScalarOperator(expr parser.Expr) (types.ScalarOperator,
 
 	case *parser.Call:
 		return q.convertFunctionCallToScalarOperator(e)
+
+	case *parser.UnaryExpr:
+		if e.Op != parser.SUB {
+			return nil, compat.NewNotSupportedError(fmt.Sprintf("unary expression with '%s'", e.Op))
+		}
+
+		if !q.engine.featureToggles.EnableUnaryNegation {
+			return nil, compat.NewNotSupportedError("unary negation of scalars")
+		}
+
+		inner, err := q.convertToScalarOperator(e.Expr)
+		if err != nil {
+			return nil, err
+		}
+
+		return operators.NewUnaryNegationOfScalar(inner, e.PositionRange()), nil
+
 	case *parser.StepInvariantExpr:
 		// One day, we'll do something smarter here.
 		return q.convertToScalarOperator(e.Expr)