MQE: add subquery support to range vector splitting

CHANGELOG.md

@@ -19,6 +19,7 @@
 * [ENHANCEMENT] Validation: Add an optional `reason` field to `limited_queries` rules, aligning them with `blocked_queries`. When set, the reason is included in the client-facing error and the query-frontend's `"query limited"` log line. #16407
 * [ENHANCEMENT] Compactor: Add the experimental `-compactor.scheduler-client.enable-ring-based-cleanup` option, which when disabled stops a scheduler-mode compactor from running the ring-based background blocks cleaner. #16457
 * [FEATURE] Querier: Add experimental per-tenant limit `-querier.max-blocks-per-store-request` to cap the number of blocks a single store-gateway request may reference. Disabled by default. #16292
+* [FEATURE] MQE: Range vector splitting can now also split subqueries, in addition to range vector selectors. Enable with the experimental `-querier.mimir-query-engine.range-vector-splitting.enable-subquery-splitting` flag, in addition to `-querier.mimir-query-engine.range-vector-splitting.enabled`. Disabled by default. #16444
 * [FEATURE] Validation: Add optional `id`, `note`, `created_by`, `created_at`, and `expires_at` fields to `blocked_queries` and `limited_queries` rules, for tooling to attach ownership/context metadata to a rule. For rules with `expires_at` set, the earliest `expires_at` per tenant and `id` (rules without an `id` are grouped together) is exported as the `cortex_blocked_query_rule_expires_at`/`cortex_limited_query_rule_expires_at` metrics, so an alert can fire on stale rules; this is informational only and never affects enforcement. The query-frontend's `"query blocked"` log line now also includes the matched rule's `id` and whether it is expired, and rate-limited queries are now logged with a new `"query limited"` line carrying the same fields. #16395
 * [BUGFIX] Query-frontend: Wait for the querier ring to be populated during startup, up to 30 seconds, before reporting the query-frontend as ready. Previously a query-frontend could become ready before it had seen any querier in the ring and fail every query it received until the ring was populated. Only applies when remote execution is enabled, and can be disabled with the experimental `-query-frontend.wait-for-querier-ring-on-startup=false`. #16333
 * [BUGFIX] Query-frontend: Fail queries with a clear error, rather than planning them against an invalid maximum supported query plan version, when the querier ring contains only unhealthy queriers. #16333

pkg/streamingpromql/config.go

@@ -86,6 +86,10 @@ type RangeVectorSplittingConfig struct {
 	//  without caching (e.g. possibly if splitting is extended to range queries in the future, or if we add
 	//  parallelisation and just want to use query splitting for that and not cache).
 	IntermediateResultsCache rangevectorsplittingcache.Config `yaml:"intermediate_results_cache" category:"experimental"`
+
+	// EnableSubquerySplitting enables splitting subqueries, in addition to range vector selectors. Requires
+	// Enabled and EngineOpts.EnableCommonSubexpressionElimination.
+	EnableSubquerySplitting bool `yaml:"enable_subquery_splitting" category:"experimental"`
 }
 
 type RangeQuerySplittingAndCachingConfig struct {
@@ -141,6 +145,7 @@ func (o *EngineOpts) RegisterFlags(f *flag.FlagSet) {
 func (c *RangeVectorSplittingConfig) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&c.Enabled, "querier.mimir-query-engine.range-vector-splitting.enabled", false, "Enable splitting function over range vectors queries into smaller blocks for caching.")
 	f.DurationVar(&c.SplitInterval, "querier.mimir-query-engine.range-vector-splitting.split-interval", 2*time.Hour, "Time interval used for splitting function over range vectors queries into cacheable blocks.")
+	f.BoolVar(&c.EnableSubquerySplitting, "querier.mimir-query-engine.range-vector-splitting.enable-subquery-splitting", false, "Enable splitting subqueries, in addition to range vector selectors. Requires -querier.mimir-query-engine.range-vector-splitting.enabled and -querier.mimir-query-engine.enable-common-subexpression-elimination to also be enabled.")
 	c.IntermediateResultsCache.RegisterFlagsWithPrefix(f, "querier.mimir-query-engine.range-vector-splitting.")
 }
 
@@ -216,6 +221,8 @@ func (c *RangeVectorSplittingConfig) Validate() error {
 		if err := c.IntermediateResultsCache.Validate(); err != nil {
 			return errors.Wrap(err, "invalid intermediate results cache config")
 		}
+	} else if c.EnableSubquerySplitting {
+		return fmt.Errorf("range vector splitting subqueries is enabled but range vector splitting is not enabled")
 	}
 	return nil
 }

pkg/streamingpromql/engine.go

@@ -100,7 +100,7 @@ func NewEngineWithCache(opts EngineOpts, metrics *stats.QueryMetrics, planner *Q
 		planning.NODE_TYPE_NUMBER_LITERAL:        planning.NodeMaterializerFunc[*core.NumberLiteral](core.MaterializeNumberLiteral),
 		planning.NODE_TYPE_STRING_LITERAL:        planning.NodeMaterializerFunc[*core.StringLiteral](core.MaterializeStringLiteral),
 		planning.NODE_TYPE_UNARY_EXPRESSION:      planning.NodeMaterializerFunc[*core.UnaryExpression](core.MaterializeUnaryExpression),
-		planning.NODE_TYPE_SUBQUERY:              planning.NodeMaterializerFunc[*core.Subquery](core.MaterializeSubquery),
+		planning.NODE_TYPE_SUBQUERY:              planning.RangeAwareNodeMaterializerFunc[*core.Subquery](core.MaterializeSubquery),
 		planning.NODE_TYPE_DEDUPLICATE_AND_MERGE: planning.NodeMaterializerFunc[*core.DeduplicateAndMerge](core.MaterializeDeduplicateAndMerge),
 		planning.NODE_TYPE_DROP_NAME:             planning.NodeMaterializerFunc[*core.DropName](core.MaterializeDropName),
 		planning.NODE_TYPE_NO_OP:                 planning.NodeMaterializerFunc[*core.NoOp](core.MaterializeNoOp),
@@ -113,7 +113,7 @@ func NewEngineWithCache(opts EngineOpts, metrics *stats.QueryMetrics, planner *Q
 		planning.NODE_TYPE_MULTI_AGGREGATION_GROUP:    planning.NodeMaterializerFunc[*multiaggregation.MultiAggregationGroup](multiaggregation.MaterializeMultiAggregationGroup),
 		planning.NODE_TYPE_MULTI_AGGREGATION_INSTANCE: planning.NodeMaterializerFunc[*multiaggregation.MultiAggregationInstance](multiaggregation.MaterializeMultiAggregationInstance),
 
-		planning.NODE_TYPE_SPLIT_FUNCTION_OVER_RANGE_VECTOR: rangevectorsplitting.NewMaterializer(opts.RangeVectorSplitting.Enabled, opts.RangeVectorSplitting.SplitInterval, opts.Limits, opts.TimeNow, intermediateCache, opts.CommonOpts.Reg, opts.Logger),
+		planning.NODE_TYPE_SPLIT_FUNCTION_OVER_RANGE_VECTOR: rangevectorsplitting.NewMaterializer(opts.RangeVectorSplitting.Enabled, opts.RangeVectorSplitting.SplitInterval, opts.RangeVectorSplitting.EnableSubquerySplitting, opts.Limits, opts.TimeNow, intermediateCache, opts.CommonOpts.Reg, opts.Logger),
 		planning.NODE_TYPE_TIME_RANGE_SPLIT:                 splitandcache.NewTimeRangeSplitMaterializer(opts.RangeQuerySplittingAndCaching.SplitEnabled, opts.CommonOpts.Reg),
 		planning.NODE_TYPE_CACHE: splitandcache.NewCacheMaterializer(
 			opts.RangeQuerySplittingAndCaching.CacheEnabled,

pkg/streamingpromql/optimize/plan/commonsubexpressionelimination/optimization_pass.go

@@ -94,6 +94,13 @@ func (e *OptimizationPass) Apply(ctx context.Context, plan *planning.QueryPlan,
 		return nil, err
 	}
 
+	// Range vector splitting can materialize and execute a nested Subquery/StepInvariantExpression more than
+	// once, once per split block, deduplicate it here.
+	splitSubqueriesDeduplicated, err := e.deduplicateSplitSubqueriesAcrossBlocks(plan.Root)
+	if err != nil {
+		return nil, err
+	}
+
 	e.selectorsInspected.Add(float64(len(paths)))
 	e.duplicateSelectorsEliminated.Add(float64(stats.duplicateSelectorsEliminated))
 	e.subsetSelectorsEliminated.Add(float64(stats.subsetSelectorsEliminated))
@@ -104,6 +111,7 @@ func (e *OptimizationPass) Apply(ctx context.Context, plan *planning.QueryPlan,
 		"selectors_inspected", len(paths),
 		"duplicate_selectors_eliminated", stats.duplicateSelectorsEliminated,
 		"subset_selectors_eliminated", stats.subsetSelectorsEliminated,
+		"split_subqueries_deduplicated", splitSubqueriesDeduplicated,
 	)
 
 	return plan, nil
@@ -929,6 +937,120 @@ func isDuplicateNode(node planning.Node) bool {
 	return isDuplicate
 }
 
+// deduplicateSplitSubqueriesAcrossBlocks finds SplitFunctionCall nodes wrapping a subquery, and deduplicates any
+// Subquery/StepInvariantExpression nested inside that subquery's inner expression.
+// Returns the number of Duplicate nodes introduced.
+func (e *OptimizationPass) deduplicateSplitSubqueriesAcrossBlocks(n planning.Node) (int, error) {
+	introduced := 0
+
+	if splitCall, ok := n.(*rangevectorsplitting.SplitFunctionCall); ok {
+		if splitCall.Inner.ChildCount() != 1 {
+			return 0, fmt.Errorf("expected SplitFunctionCall's inner function call to have exactly one child, got %d", splitCall.Inner.ChildCount())
+		}
+
+		if subquery, isSubquery := unwrapDuplicate(splitCall.Inner.Child(0)).(*core.Subquery); isSubquery {
+			count, err := e.deduplicateAcrossSplitBlocks(subquery.Child(0))
+			if err != nil {
+				return 0, err
+			}
+
+			introduced += count
+		}
+	}
+
+	for i := range n.ChildCount() {
+		count, err := e.deduplicateSplitSubqueriesAcrossBlocks(n.Child(i))
+		if err != nil {
+			return 0, err
+		}
+
+		introduced += count
+	}
+
+	return introduced, nil
+}
+
+func unwrapDuplicate(n planning.Node) planning.Node {
+	switch n := n.(type) {
+	case *Duplicate:
+		return unwrapDuplicate(n.Inner)
+	case *DuplicateFilter:
+		return unwrapDuplicate(n.Inner)
+	default:
+		return n
+	}
+}
+
+// deduplicateAcrossSplitBlocks wraps the child of every core.Subquery/core.StepInvariantExpression in n's subtree in a
+// Duplicate node, at any nesting depth, so different split blocks can safely materialize it more than once (see range_vector_splitting_2h.test:830).
+// Returns the number of Duplicate nodes introduced.
+func (e *OptimizationPass) deduplicateAcrossSplitBlocks(n planning.Node) (int, error) {
+	if isSubqueryOrStepInvariantExpression(n) {
+		if n.ChildCount() != 1 {
+			return 0, fmt.Errorf("expected node of type %s to have exactly one child, got %d", n.NodeType(), n.ChildCount())
+		}
+
+		child := n.Child(0)
+
+		if isDeduplicateOrDeduplicateFilter(child) {
+			// keep recursing since a further nested Subquery/StepInvariantExpression inside it may still need its own Duplicate.
+			return e.deduplicateAcrossSplitBlocks(child)
+		}
+
+		// Result type is always Vector or Scalar in practice (see planning.go's StepInvariantExpr handling).
+		// This is a defensive check in case that invariant is ever broken.
+		if resultType, err := child.ResultType(); err != nil {
+			return 0, err
+		} else if resultType != parser.ValueTypeVector && resultType != parser.ValueTypeScalar {
+			return 0, fmt.Errorf("cannot deduplicate %s node (%s) across split blocks: unexpected result type %s", n.NodeType(), n.Describe(), resultType)
+		}
+
+		introduced, err := e.deduplicateAcrossSplitBlocks(child)
+		if err != nil {
+			return 0, err
+		}
+
+		duplicate := &Duplicate{Inner: child, DuplicateDetails: &DuplicateDetails{}}
+		e.duplicationNodesIntroduced.Inc()
+
+		if err := n.ReplaceChild(0, duplicate); err != nil {
+			return 0, err
+		}
+
+		return introduced + 1, nil
+	}
+
+	introduced := 0
+	for i := range n.ChildCount() {
+		count, err := e.deduplicateAcrossSplitBlocks(n.Child(i))
+		if err != nil {
+			return 0, err
+		}
+
+		introduced += count
+	}
+
+	return introduced, nil
+}
+
+func isSubqueryOrStepInvariantExpression(n planning.Node) bool {
+	switch n.(type) {
+	case *core.Subquery, *core.StepInvariantExpression:
+		return true
+	default:
+		return false
+	}
+}
+
+func isDeduplicateOrDeduplicateFilter(node planning.Node) bool {
+	switch node.(type) {
+	case *Duplicate, *DuplicateFilter:
+		return true
+	default:
+		return false
+	}
+}
+
 type path []pathElement
 
 type pathElement struct {

pkg/streamingpromql/optimize/plan/rangevectorsplitting/node.go

@@ -68,40 +68,61 @@ func (s *SplitFunctionCall) ExpressionPosition() (posrange.PositionRange, error)
 }
 
 func (s *SplitFunctionCall) MinimumRequiredPlanVersion(types.QueryTimeRange) (planning.QueryPlanVersion, error) {
+	if containsSubquery(s.Inner) {
+		return planning.QueryPlanV20, nil
+	}
+
 	return planning.QueryPlanV18, nil
 }
 
+// containsSubquery returns true if n or any of its descendants is a Subquery.
+func containsSubquery(n planning.Node) bool {
+	if _, ok := n.(*core.Subquery); ok {
+		return true
+	}
+
+	for child := range planning.ChildrenIter(n) {
+		if containsSubquery(child) {
+			return true
+		}
+	}
+
+	return false
+}
+
 // limitsProvider provides the tenant limits needed to compute split ranges at materialize time.
 type limitsProvider interface {
 	GetMaxOutOfOrderTimeWindow(ctx context.Context) (time.Duration, error)
 }
 
 type Materializer struct {
-	enabled       bool
-	splitInterval time.Duration
-	limits        limitsProvider
-	timeNow       func() time.Time
-	cache         *cache.CacheFactory
-	logger        log.Logger
+	enabled                 bool
+	splitInterval           time.Duration
+	enableSubquerySplitting bool
+	limits                  limitsProvider
+	timeNow                 func() time.Time
+	cache                   *cache.CacheFactory
+	logger                  log.Logger
 
 	nodesSplit   prometheus.Counter
 	nodesUnsplit *prometheus.CounterVec
 }
 
 var _ planning.NodeMaterializer = &Materializer{}
 
-func NewMaterializer(enabled bool, splitInterval time.Duration, limits limitsProvider, timeNow func() time.Time, cache *cache.CacheFactory, reg prometheus.Registerer, logger log.Logger) *Materializer {
+func NewMaterializer(enabled bool, splitInterval time.Duration, enableSubquerySplitting bool, limits limitsProvider, timeNow func() time.Time, cache *cache.CacheFactory, reg prometheus.Registerer, logger log.Logger) *Materializer {
 	if timeNow == nil {
 		timeNow = time.Now
 	}
 
 	return &Materializer{
-		enabled:       enabled,
-		splitInterval: splitInterval,
-		limits:        limits,
-		timeNow:       timeNow,
-		cache:         cache,
-		logger:        logger,
+		enabled:                 enabled,
+		splitInterval:           splitInterval,
+		enableSubquerySplitting: enableSubquerySplitting,
+		limits:                  limits,
+		timeNow:                 timeNow,
+		cache:                   cache,
+		logger:                  logger,
 		nodesSplit: promauto.With(reg).NewCounter(prometheus.CounterOpts{
 			Name: "cortex_mimir_query_engine_range_vector_splitting_nodes_materialized_split_total",
 			Help: "Total number of range vector splitting nodes materialized as split operators.",
@@ -143,6 +164,12 @@ func (m Materializer) Materialize(ctx context.Context, n planning.Node, material
 		return nil, fmt.Errorf("inner node of split function call does not implement SplitNode: %T", innerNode)
 	}
 
+	if containsSubquery(innerNode) && !m.enableSubquerySplitting {
+		level.Warn(m.logger).Log("msg", "split function node wraps a subquery but subquery splitting is disabled, falling back to unsplit execution; this can happen if subquery splitting is enabled on the query-frontend but not yet on the querier")
+		m.nodesUnsplit.WithLabelValues("subquery_splitting_disabled").Inc()
+		return materializer.FactoryForNode(ctx, s.Inner, timeRange)
+	}
+
 	// The split ranges are computed here, at materialize time, rather than at planning time, because they depend on
 	// the querier's current time and the tenant's out-of-order window. If the ranges turn out not to be worth
 	// splitting (e.g. there's no complete cacheable block, or every block falls within the out-of-order window), fall

pkg/streamingpromql/optimize/plan/rangevectorsplitting/node_test.go

@@ -164,7 +164,7 @@ func TestMaterializer_computeRanges(t *testing.T) {
 	fixedNow := timestamp.Time(100 * hourInMs)
 
 	newMaterializer := func(oooWindow time.Duration) *Materializer {
-		return NewMaterializer(true, 2*time.Hour, staticLimits{oooWindow: oooWindow}, func() time.Time { return fixedNow }, nil, nil, nil)
+		return NewMaterializer(true, 2*time.Hour, true, staticLimits{oooWindow: oooWindow}, func() time.Time { return fixedNow }, nil, nil, nil)
 	}
 
 	// inner builds an inner matrix selector with the given range and offset.