MQE: Add support to disable functions+aggregations by name

grafana · Jan 9, 2025 · 8646304 · 8646304
1 parent ce6871c
commit 8646304
Show file tree

Hide file tree

Showing 9 changed files with 219 additions and 144 deletions.
diff --git a/cmd/mimir/config-descriptor.json b/cmd/mimir/config-descriptor.json
@@ -1969,23 +1969,34 @@
             },
             {
               "kind": "field",
-              "name": "enable_vector_vector_binary_comparison_operations",
+              "name": "enable_binary_logical_operations",
               "required": false,
-              "desc": "Enable support for binary comparison operations between two vectors in the Mimir query engine. Only applies if the MQE is in use.",
+              "desc": "Enable support for binary logical operations in the Mimir query engine. Only applies if the MQE is in use.",
               "fieldValue": null,
               "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-vector-vector-binary-comparison-operations",
+              "fieldFlag": "querier.mimir-query-engine.enable-binary-logical-operations",
               "fieldType": "boolean",
               "fieldCategory": "experimental"
             },
             {
               "kind": "field",
-              "name": "enable_vector_scalar_binary_comparison_operations",
+              "name": "enable_one_to_many_and_many_to_one_binary_operations",
               "required": false,
-              "desc": "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.",
+              "desc": "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.",
               "fieldValue": null,
               "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-vector-scalar-binary-comparison-operations",
+              "fieldFlag": "querier.mimir-query-engine.enable-one-to-many-and-many-to-one-binary-operations",
+              "fieldType": "boolean",
+              "fieldCategory": "experimental"
+            },
+            {
+              "kind": "field",
+              "name": "enable_scalars",
+              "required": false,
+              "desc": "Enable support for scalars in the Mimir query engine. Only applies if the MQE is in use.",
+              "fieldValue": null,
+              "fieldDefaultValue": true,
+              "fieldFlag": "querier.mimir-query-engine.enable-scalars",
               "fieldType": "boolean",
               "fieldCategory": "experimental"
             },
@@ -2002,57 +2013,57 @@
             },
             {
               "kind": "field",
-              "name": "enable_binary_logical_operations",
+              "name": "enable_subqueries",
               "required": false,
-              "desc": "Enable support for binary logical operations in the Mimir query engine. Only applies if the MQE is in use.",
+              "desc": "Enable support for subqueries in the Mimir query engine. Only applies if the MQE is in use.",
               "fieldValue": null,
               "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-binary-logical-operations",
+              "fieldFlag": "querier.mimir-query-engine.enable-subqueries",
               "fieldType": "boolean",
               "fieldCategory": "experimental"
             },
             {
               "kind": "field",
-              "name": "enable_scalars",
+              "name": "enable_vector_scalar_binary_comparison_operations",
               "required": false,
-              "desc": "Enable support for scalars in the Mimir query engine. Only applies if the MQE is in use.",
+              "desc": "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.",
               "fieldValue": null,
               "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-scalars",
+              "fieldFlag": "querier.mimir-query-engine.enable-vector-scalar-binary-comparison-operations",
               "fieldType": "boolean",
               "fieldCategory": "experimental"
             },
             {
               "kind": "field",
-              "name": "enable_subqueries",
+              "name": "enable_vector_vector_binary_comparison_operations",
               "required": false,
-              "desc": "Enable support for subqueries in the Mimir query engine. Only applies if the MQE is in use.",
+              "desc": "Enable support for binary comparison operations between two vectors in the Mimir query engine. Only applies if the MQE is in use.",
               "fieldValue": null,
               "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-subqueries",
+              "fieldFlag": "querier.mimir-query-engine.enable-vector-vector-binary-comparison-operations",
               "fieldType": "boolean",
               "fieldCategory": "experimental"
             },
             {
               "kind": "field",
-              "name": "enable_histogram_quantile_function",
+              "name": "disabled_aggregations",
               "required": false,
-              "desc": "Enable support for the histogram_quantile function in the Mimir query engine. Only applies if the MQE is in use.",
+              "desc": "Comma-separated list of aggregations to disable. Queries using a disabled operation will fall back to Prometheus' query engine. Only applies if MQE is in use.",
               "fieldValue": null,
-              "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-histogram-quantile-function",
-              "fieldType": "boolean",
+              "fieldDefaultValue": "",
+              "fieldFlag": "querier.mimir-query-engine.disabled-aggregations",
+              "fieldType": "string",
               "fieldCategory": "experimental"
             },
             {
               "kind": "field",
-              "name": "enable_one_to_many_and_many_to_one_binary_operations",
+              "name": "disabled_functions",
               "required": false,
-              "desc": "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.",
+              "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.",
               "fieldValue": null,
-              "fieldDefaultValue": true,
-              "fieldFlag": "querier.mimir-query-engine.enable-one-to-many-and-many-to-one-binary-operations",
-              "fieldType": "boolean",
+              "fieldDefaultValue": "",
+              "fieldFlag": "querier.mimir-query-engine.disabled-functions",
+              "fieldType": "string",
               "fieldCategory": "experimental"
             }
           ],

diff --git a/cmd/mimir/help-all.txt.tmpl b/cmd/mimir/help-all.txt.tmpl
@@ -2097,12 +2097,14 @@ Usage of ./cmd/mimir/mimir:
     	Maximum number of split (by time) or partial (by shard) queries that will be scheduled in parallel by the query-frontend for a single input query. This limit is introduced to have a fairer query scheduling and avoid a single query over a large time range saturating all available queriers. (default 14)
   -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 operation will fall back to Prometheus' query engine. 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.
   -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
     	[experimental] Enable support for binary logical operations in the Mimir query engine. Only applies if the MQE is in use. (default true)
-  -querier.mimir-query-engine.enable-histogram-quantile-function
-    	[experimental] Enable support for the histogram_quantile function in the Mimir query engine. Only applies if the MQE is in use. (default true)
   -querier.mimir-query-engine.enable-one-to-many-and-many-to-one-binary-operations
     	[experimental] 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. (default true)
   -querier.mimir-query-engine.enable-scalar-scalar-binary-comparison-operations

diff --git a/pkg/querier/engine/config.go b/pkg/querier/engine/config.go
@@ -33,7 +33,7 @@ type Config struct {
 
 	PromQLExperimentalFunctionsEnabled bool `yaml:"promql_experimental_functions_enabled" category:"experimental"`
 
-	MimirQueryEngine streamingpromql.FeatureToggles `yaml:"mimir_query_engine" category:"experimental"`
+	MimirQueryEngine streamingpromql.MQEOpts `yaml:"mimir_query_engine" category:"experimental"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
@@ -73,8 +73,8 @@ func NewPromQLEngineOptions(cfg Config, activityTracker *activitytracker.Activit
 	}
 
 	mqeOpts := streamingpromql.EngineOpts{
-		CommonOpts:     commonOpts,
-		FeatureToggles: cfg.MimirQueryEngine,
+		CommonOpts: commonOpts,
+		MQEOpts:    cfg.MimirQueryEngine,
 	}
 
 	return commonOpts, mqeOpts, cfg.PromQLExperimentalFunctionsEnabled

diff --git a/pkg/streamingpromql/config.go b/pkg/streamingpromql/config.go
@@ -5,32 +5,35 @@ package streamingpromql
 import (
 	"flag"
 
+	"github.com/grafana/dskit/flagext"
 	"github.com/prometheus/prometheus/promql"
 )
 
 type EngineOpts struct {
-	CommonOpts     promql.EngineOpts
-	FeatureToggles FeatureToggles
+	CommonOpts promql.EngineOpts
+	MQEOpts    MQEOpts
 
 	// 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 FeatureToggles struct {
+type MQEOpts struct {
 	EnableAggregationOperations                  bool `yaml:"enable_aggregation_operations" category:"experimental"`
-	EnableVectorVectorBinaryComparisonOperations bool `yaml:"enable_vector_vector_binary_comparison_operations" category:"experimental"`
-	EnableVectorScalarBinaryComparisonOperations bool `yaml:"enable_vector_scalar_binary_comparison_operations" category:"experimental"`
-	EnableScalarScalarBinaryComparisonOperations bool `yaml:"enable_scalar_scalar_binary_comparison_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"`
 	EnableScalars                                bool `yaml:"enable_scalars" category:"experimental"`
+	EnableScalarScalarBinaryComparisonOperations bool `yaml:"enable_scalar_scalar_binary_comparison_operations" category:"experimental"`
 	EnableSubqueries                             bool `yaml:"enable_subqueries" category:"experimental"`
-	EnableHistogramQuantileFunction              bool `yaml:"enable_histogram_quantile_function" category:"experimental"`
-	EnableOneToManyAndManyToOneBinaryOperations  bool `yaml:"enable_one_to_many_and_many_to_one_binary_operations" category:"experimental"`
+	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"`
 }
 
-// EnableAllFeatures enables all features supported by MQE, including experimental or incomplete features.
-var EnableAllFeatures = FeatureToggles{
+// MQEOptsAllFeatures enables all features supported by MQE, including experimental or incomplete features.
+var MQEOptsAllFeatures = MQEOpts{
 	// 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,
@@ -40,17 +43,20 @@ var EnableAllFeatures = FeatureToggles{
 	true,
 	true,
 	true,
-	true,
+	[]string{},
+	[]string{},
 }
 
-func (t *FeatureToggles) RegisterFlags(f *flag.FlagSet) {
+func (t *MQEOpts) 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.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.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.EnableScalarScalarBinaryComparisonOperations, "querier.mimir-query-engine.enable-scalar-scalar-binary-comparison-operations", true, "Enable support for binary comparison operations between two scalars 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.")
 	f.BoolVar(&t.EnableScalars, "querier.mimir-query-engine.enable-scalars", true, "Enable support for scalars in the Mimir query engine. Only applies if the MQE is in use.")
+	f.BoolVar(&t.EnableScalarScalarBinaryComparisonOperations, "querier.mimir-query-engine.enable-scalar-scalar-binary-comparison-operations", true, "Enable support for binary comparison operations between two scalars in the Mimir query engine. Only applies if the MQE is in use.")
 	f.BoolVar(&t.EnableSubqueries, "querier.mimir-query-engine.enable-subqueries", true, "Enable support for subqueries in the Mimir query engine. Only applies if the MQE is in use.")
-	f.BoolVar(&t.EnableHistogramQuantileFunction, "querier.mimir-query-engine.enable-histogram-quantile-function", true, "Enable support for the histogram_quantile function 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.")
+	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.")
 }
diff --git a/pkg/streamingpromql/engine.go b/pkg/streamingpromql/engine.go
@@ -49,7 +49,7 @@ func NewEngine(opts EngineOpts, limitsProvider QueryLimitsProvider, metrics *sta
 		timeout:                  opts.CommonOpts.Timeout,
 		limitsProvider:           limitsProvider,
 		activeQueryTracker:       opts.CommonOpts.ActiveQueryTracker,
-		featureToggles:           opts.FeatureToggles,
+		mqeOpts:                  opts.MQEOpts,
 		noStepSubqueryIntervalFn: opts.CommonOpts.NoStepSubqueryIntervalFn,
 
 		logger: logger,
@@ -69,7 +69,7 @@ type Engine struct {
 	timeout                  time.Duration
 	limitsProvider           QueryLimitsProvider
 	activeQueryTracker       promql.QueryTracker
-	featureToggles           FeatureToggles
+	mqeOpts                  MQEOpts
 	noStepSubqueryIntervalFn func(rangeMillis int64) int64
 
 	logger                                    log.Logger