MQE: Add support to disable functions+aggregations by name

cmd/mimir/config-descriptor.json

@@ -2048,7 +2048,7 @@
               "kind": "field",
               "name": "disabled_aggregations",
               "required": false,
-              "desc": "Comma-separated list of aggregations to disable. Queries using a disabled aggregation will fall back to Prometheus' query engine. Only applies if MQE is in use.",
+              "desc": "Comma-separated list of aggregations to disable support for. Only applies if MQE is in use.",
               "fieldValue": null,
               "fieldDefaultValue": "",
               "fieldFlag": "querier.mimir-query-engine.disabled-aggregations",
@@ -2059,7 +2059,7 @@
               "kind": "field",
               "name": "disabled_functions",
               "required": false,
-              "desc": "Comma-separated list of function names to disable. Queries using a disabled function will fall back to Prometheus' query engine. Only applies if MQE is in use.",
+              "desc": "Comma-separated list of function names to disable support for. Only applies if MQE is in use.",
               "fieldValue": null,
               "fieldDefaultValue": "",
               "fieldFlag": "querier.mimir-query-engine.disabled-functions",

cmd/mimir/help-all.txt.tmpl

@@ -2098,9 +2098,9 @@ Usage of ./cmd/mimir/mimir:
   -querier.max-samples int
     	Maximum number of samples a single query can load into memory. This config option should be set on query-frontend too when query sharding is enabled. (default 50000000)
   -querier.mimir-query-engine.disabled-aggregations comma-separated-list-of-strings
-    	[experimental] Comma-separated list of aggregations to disable. Queries using a disabled aggregation will fall back to Prometheus' query engine. Only applies if MQE is in use.
+    	[experimental] Comma-separated list of aggregations to disable support for. Only applies if MQE is in use.
   -querier.mimir-query-engine.disabled-functions comma-separated-list-of-strings
-    	[experimental] Comma-separated list of function names to disable. Queries using a disabled function will fall back to Prometheus' query engine. Only applies if MQE is in use.
+    	[experimental] Comma-separated list of function names to disable support for. Only applies if MQE is in use.
   -querier.mimir-query-engine.enable-aggregation-operations
     	[experimental] Enable support for aggregation operations in the Mimir query engine. Only applies if the MQE is in use. (default true)
   -querier.mimir-query-engine.enable-binary-logical-operations

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

@@ -1533,15 +1533,13 @@ mimir_query_engine:
   # CLI flag: -querier.mimir-query-engine.enable-vector-vector-binary-comparison-operations
   [enable_vector_vector_binary_comparison_operations: <boolean> | default = true]
 
-  # (experimental) Comma-separated list of aggregations to disable. Queries
-  # using a disabled aggregation will fall back to Prometheus' query engine.
+  # (experimental) Comma-separated list of aggregations to disable support for.
   # Only applies if MQE is in use.
   # CLI flag: -querier.mimir-query-engine.disabled-aggregations
   [disabled_aggregations: <string> | default = ""]
 
-  # (experimental) Comma-separated list of function names to disable. Queries
-  # using a disabled function will fall back to Prometheus' query engine. Only
-  # applies if MQE is in use.
+  # (experimental) Comma-separated list of function names to disable support
+  # for. Only applies if MQE is in use.
   # CLI flag: -querier.mimir-query-engine.disabled-functions
   [disabled_functions: <string> | default = ""]
 ```

pkg/querier/engine/config.go

@@ -33,7 +33,7 @@ type Config struct {
 
 	PromQLExperimentalFunctionsEnabled bool `yaml:"promql_experimental_functions_enabled" category:"experimental"`
 
-	MimirQueryEngine streamingpromql.MQEOpts `yaml:"mimir_query_engine" category:"experimental"`
+	MimirQueryEngine streamingpromql.Features `yaml:"mimir_query_engine" category:"experimental"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
@@ -74,7 +74,7 @@ func NewPromQLEngineOptions(cfg Config, activityTracker *activitytracker.Activit
 
 	mqeOpts := streamingpromql.EngineOpts{
 		CommonOpts: commonOpts,
-		MQEOpts:    cfg.MimirQueryEngine,
+		Features:   cfg.MimirQueryEngine,
 	}
 
 	return commonOpts, mqeOpts, cfg.PromQLExperimentalFunctionsEnabled

pkg/streamingpromql/config.go

@@ -7,18 +7,19 @@ import (
 
 	"github.com/grafana/dskit/flagext"
 	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
 )
 
 type EngineOpts struct {
 	CommonOpts promql.EngineOpts
-	MQEOpts    MQEOpts
+	Features   Features
 
 	// When operating in pedantic mode, we panic if memory consumption is > 0 after Query.Close()
 	// (indicating something was not returned to a pool).
 	Pedantic bool
 }
 
-type MQEOpts struct {
+type Features struct {
 	EnableAggregationOperations                  bool `yaml:"enable_aggregation_operations" category:"experimental"`
 	EnableBinaryLogicalOperations                bool `yaml:"enable_binary_logical_operations" category:"experimental"`
 	EnableOneToManyAndManyToOneBinaryOperations  bool `yaml:"enable_one_to_many_and_many_to_one_binary_operations" category:"experimental"`
@@ -28,12 +29,13 @@ type MQEOpts struct {
 	EnableVectorScalarBinaryComparisonOperations bool `yaml:"enable_vector_scalar_binary_comparison_operations" category:"experimental"`
 	EnableVectorVectorBinaryComparisonOperations bool `yaml:"enable_vector_vector_binary_comparison_operations" category:"experimental"`
 
-	DisabledAggregations flagext.StringSliceCSV `yaml:"disabled_aggregations" category:"experimental"`
-	DisabledFunctions    flagext.StringSliceCSV `yaml:"disabled_functions" category:"experimental"`
+	DisabledAggregations      flagext.StringSliceCSV `yaml:"disabled_aggregations" category:"experimental"`
+	DisabledAggregationsItems []parser.ItemType
+	DisabledFunctions         flagext.StringSliceCSV `yaml:"disabled_functions" category:"experimental"`
 }
 
-// MQEOptsAllFeatures enables all features supported by MQE, including experimental or incomplete features.
-var MQEOptsAllFeatures = MQEOpts{
+// EnableAllFeatures enables all features supported by MQE, including experimental or incomplete features.
+var EnableAllFeatures = Features{
 	// Note that we deliberately use a keyless literal here to force a compilation error if we don't keep this in sync with new fields added to FeatureToggles.
 	true,
 	true,
@@ -44,10 +46,11 @@ var MQEOptsAllFeatures = MQEOpts{
 	true,
 	true,
 	[]string{},
+	[]parser.ItemType{},
 	[]string{},
 }
 
-func (t *MQEOpts) RegisterFlags(f *flag.FlagSet) {
+func (t *Features) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&t.EnableAggregationOperations, "querier.mimir-query-engine.enable-aggregation-operations", true, "Enable support for aggregation operations in the Mimir query engine. Only applies if the MQE is in use.")
 	f.BoolVar(&t.EnableBinaryLogicalOperations, "querier.mimir-query-engine.enable-binary-logical-operations", true, "Enable support for binary logical operations in the Mimir query engine. Only applies if the MQE is in use.")
 	f.BoolVar(&t.EnableOneToManyAndManyToOneBinaryOperations, "querier.mimir-query-engine.enable-one-to-many-and-many-to-one-binary-operations", true, "Enable support for one-to-many and many-to-one binary operations (group_left/group_right) in the Mimir query engine. Only applies if the MQE is in use.")
@@ -57,6 +60,6 @@ func (t *MQEOpts) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&t.EnableVectorScalarBinaryComparisonOperations, "querier.mimir-query-engine.enable-vector-scalar-binary-comparison-operations", true, "Enable support for binary comparison operations between a vector and a scalar in the Mimir query engine. Only applies if the MQE is in use.")
 	f.BoolVar(&t.EnableVectorVectorBinaryComparisonOperations, "querier.mimir-query-engine.enable-vector-vector-binary-comparison-operations", true, "Enable support for binary comparison operations between two vectors in the Mimir query engine. Only applies if the MQE is in use.")
 
-	f.Var(&t.DisabledAggregations, "querier.mimir-query-engine.disabled-aggregations", "Comma-separated list of aggregations to disable. Queries using a disabled aggregation will fall back to Prometheus' query engine. Only applies if MQE is in use.")
-	f.Var(&t.DisabledFunctions, "querier.mimir-query-engine.disabled-functions", "Comma-separated list of function names to disable. Queries using a disabled function will fall back to Prometheus' query engine. Only applies if MQE is in use.")
+	f.Var(&t.DisabledAggregations, "querier.mimir-query-engine.disabled-aggregations", "Comma-separated list of aggregations to disable support for. Only applies if MQE is in use.")
+	f.Var(&t.DisabledFunctions, "querier.mimir-query-engine.disabled-functions", "Comma-separated list of function names to disable support for. Only applies if MQE is in use.")
 }

pkg/streamingpromql/engine.go

@@ -9,15 +9,18 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"slices"
 	"time"
 
 	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
 
 	"github.com/grafana/mimir/pkg/querier/stats"
+	"github.com/grafana/mimir/pkg/streamingpromql/operators/aggregations"
 )
 
 const defaultLookbackDelta = 5 * time.Minute // This should be the same value as github.com/prometheus/prometheus/promql.defaultLookbackDelta.
@@ -44,12 +47,26 @@ func NewEngine(opts EngineOpts, limitsProvider QueryLimitsProvider, metrics *sta
 		return nil, errors.New("enabling delayed name removal not supported by Mimir query engine")
 	}
 
+	// Sort DisabledFunctions to optimise lookups
+	slices.Sort(opts.Features.DisabledFunctions)
+
+	opts.Features.DisabledAggregationsItems = make([]parser.ItemType, 0, len(opts.Features.DisabledAggregations))
+	for _, agg := range opts.Features.DisabledAggregations {
+		item, ok := aggregations.GetAggregationItemType(agg)
+		if !ok {
+			return nil, fmt.Errorf("disabled aggregation '%s' does not exist", agg)
+		}
+		opts.Features.DisabledAggregationsItems = append(opts.Features.DisabledAggregationsItems, item)
+	}
+	// No point sorting DisabledAggregations earlier, as ItemType ints are not in order.
+	slices.Sort(opts.Features.DisabledAggregationsItems)
+
 	return &Engine{
 		lookbackDelta:            lookbackDelta,
 		timeout:                  opts.CommonOpts.Timeout,
 		limitsProvider:           limitsProvider,
 		activeQueryTracker:       opts.CommonOpts.ActiveQueryTracker,
-		mqeOpts:                  opts.MQEOpts,
+		features:                 opts.Features,
 		noStepSubqueryIntervalFn: opts.CommonOpts.NoStepSubqueryIntervalFn,
 
 		logger: logger,
@@ -69,7 +86,7 @@ type Engine struct {
 	timeout                  time.Duration
 	limitsProvider           QueryLimitsProvider
 	activeQueryTracker       promql.QueryTracker
-	mqeOpts                  MQEOpts
+	features                 Features
 	noStepSubqueryIntervalFn func(rangeMillis int64) int64
 
 	logger                                    log.Logger